diff --git a/configs/experiments/mock/short.json b/configs/experiments/mock/short.json deleted file mode 100644 index 5d2433d..0000000 --- a/configs/experiments/mock/short.json +++ /dev/null @@ -1,17 +0,0 @@ -{ - "experiment": { - "action_chunk_broker_type": "naive_async", - "num_robots": 1, - "trials_per_robot": 1, - "max_steps": 50 - }, - "robots": { - "robot_0": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - } - } -} diff --git a/configs/heterogeneous_1fast_11slow.yaml b/configs/exps/edge_cases_real/1fast_11slow.yaml similarity index 100% rename from configs/heterogeneous_1fast_11slow.yaml rename to configs/exps/edge_cases_real/1fast_11slow.yaml diff --git a/configs/heterogeneous_4fast_8slow.yaml b/configs/exps/edge_cases_real/4fast_8slow.yaml similarity index 100% rename from configs/heterogeneous_4fast_8slow.yaml rename to configs/exps/edge_cases_real/4fast_8slow.yaml diff --git a/configs/experiments/edge_cases/1_fast_9_slow.jsonc b/configs/exps/edge_cases_sim/1_fast_9_slow.jsonc similarity index 100% rename from configs/experiments/edge_cases/1_fast_9_slow.jsonc rename to configs/exps/edge_cases_sim/1_fast_9_slow.jsonc diff --git a/configs/experiments/edge_cases/1_fast_only.jsonc b/configs/exps/edge_cases_sim/1_fast_only.jsonc similarity index 100% rename from configs/experiments/edge_cases/1_fast_only.jsonc rename to configs/exps/edge_cases_sim/1_fast_only.jsonc diff --git a/configs/experiments/edge_cases/4_fast_1_slow.jsonc b/configs/exps/edge_cases_sim/4_fast_1_slow.jsonc similarity index 100% rename from configs/experiments/edge_cases/4_fast_1_slow.jsonc rename to configs/exps/edge_cases_sim/4_fast_1_slow.jsonc diff --git a/configs/experiments/edge_cases/5_fast_5_slow.jsonc b/configs/exps/edge_cases_sim/5_fast_5_slow.jsonc similarity index 100% rename from configs/experiments/edge_cases/5_fast_5_slow.jsonc rename to configs/exps/edge_cases_sim/5_fast_5_slow.jsonc diff --git a/configs/inference_profiles.json b/configs/inference_profiles.json index 7b2bc1c..056406f 100644 --- a/configs/inference_profiles.json +++ b/configs/inference_profiles.json @@ -11,6 +11,21 @@ "2": 0.0765, "3": 0.0835, "4": 0.0918, - "5": 0.1009 + "5": 0.1009, + "6": 0.1150, + "7": 0.1244, + "8": 0.1366, + "9": 0.1535, + "10": 0.1853, + "11": 0.2146, + "12": 0.2267, + "13": 0.2465, + "14": 0.2582, + "15": 0.2754, + "16": 0.3138, + "17": 0.3322, + "18": 0.3482, + "19": 0.3645, + "20": 0.3857 } } \ No newline at end of file diff --git a/packages/armory-client/src/armory_client/runtime/real_saver.py b/packages/armory-client/src/armory_client/runtime/real_saver.py index 83b59fc..8ea568d 100644 --- a/packages/armory-client/src/armory_client/runtime/real_saver.py +++ b/packages/armory-client/src/armory_client/runtime/real_saver.py @@ -20,22 +20,24 @@ import re import time from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass import imageio -import matplotlib - -matplotlib.use("Agg") -import matplotlib.pyplot as plt import numpy as np from typing_extensions import override from armory_client.action_chunkers.action_chunk_broker import ActionChunkBroker from armory_client.runtime import subscriber as _subscriber +from armory_client.runtime.saver_utils import ( + EpisodeSaveData, + Result, + plot_cost_history, + save_action_chunks, + save_actions_left, + save_cost_history_npy, + save_timestamps, +) from armory_client.schemas import ( Action, - ActionChunk, - JSONDataclass, Observation, Timestamp, ) @@ -55,33 +57,6 @@ def _robot_idx_from_id(robot_id: str) -> int: return int(match.group(1)) if match else 0 -@dataclass(frozen=True) -class Result(JSONDataclass): - """Per-episode metadata. Schema matches sims.libero.subscribers.saver.Result.""" - - robot_idx: int - success: bool - steps_taken: int - task_suite_name: str - task_id: int - task_language: str - episode_idx: int - - -@dataclass -class _EpisodeSaveData: - """Snapshot of one episode's data, safe to hand off to a background thread.""" - - timestamps: list[Timestamp] - observations_buffer: dict[int, Observation] - action_chunks: list[ActionChunk] - actions_left_snapshot: list[int] - cost_history: list[float] - success: bool - episode_idx: int - initial_state: np.ndarray | None - - class RealSaver(_subscriber.Subscriber): """Saves real-robot trajectory data; on-disk layout matches the sim Saver.""" @@ -167,7 +142,7 @@ def on_step(self, observation: Observation, action: Action) -> None: @override def on_episode_end(self) -> None: - data = _EpisodeSaveData( + data = EpisodeSaveData( timestamps=self._timestamps, observations_buffer=self._observations_buffer, action_chunks=list(self._action_chunk_broker.action_chunks), @@ -186,21 +161,21 @@ def close(self) -> None: # ── disk writes (same names + formats as sim Saver) ───────── - def _save_all(self, data: _EpisodeSaveData) -> None: + def _save_all(self, data: EpisodeSaveData) -> None: out_folder = self._get_out_folder(data) try: self._save_metadata(out_folder, data) - self._save_timestamps(out_folder, data) - self._save_action_chunks(out_folder, data) + save_timestamps(data.timestamps, out_folder) + save_action_chunks(data.action_chunks, out_folder) if self._save_video_enabled: self._save_video(out_folder, data) self._save_debug_data(out_folder, data) - self._save_actions_left(out_folder, data) + save_actions_left(data.actions_left_snapshot, out_folder) self._save_cost_history(out_folder, data) except Exception: logger.exception("RealSaver: error writing episode %s", out_folder) - def _get_out_folder(self, data: _EpisodeSaveData) -> pathlib.Path: + def _get_out_folder(self, data: EpisodeSaveData) -> pathlib.Path: # Use the snapshot's episode_idx (assigned at on_episode_end time) so # concurrent flushes don't collide on a disk-scan. robot_folder = self._out_dir / str(self._robot_idx) @@ -213,7 +188,7 @@ def _get_out_folder(self, data: _EpisodeSaveData) -> pathlib.Path: out_folder.mkdir(parents=True, exist_ok=True) return out_folder - def _save_metadata(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: + def _save_metadata(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: result = Result( success=data.success, robot_idx=self._robot_idx, @@ -225,13 +200,7 @@ def _save_metadata(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> No ) result.to_json(out_folder / "metadata.json") - def _save_timestamps(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - Timestamp.to_csv(data.timestamps, out_folder / "timestamps.csv") - - def _save_action_chunks(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - ActionChunk.to_parquet(data.action_chunks, out_folder / "action_chunks.parquet") - - def _save_video(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: + def _save_video(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: images = [ obs.image for obs in data.observations_buffer.values() @@ -245,7 +214,7 @@ def _save_video(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: fps=self._control_hz, ) - def _save_debug_data(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: + def _save_debug_data(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: has_noise = any( getattr(chunk, "noise", None) is not None for chunk in data.action_chunks ) @@ -278,29 +247,14 @@ def _save_debug_data(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> np.savez_compressed(debug_data_file, **data_to_save) - def _save_actions_left(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - np.save( - out_folder / "actions_left.npy", - np.array(data.actions_left_snapshot, dtype=np.int32), - ) - - def _save_cost_history(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - costs = np.array(data.cost_history, dtype=np.float64) - np.save(out_folder / "cost_history.npy", costs) - + def _save_cost_history(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: + costs = save_cost_history_npy(data.cost_history, out_folder) if costs.size == 0: return - - steps = np.arange(len(costs)) - fig, ax = plt.subplots(figsize=(10, 4)) - ax.plot(steps, costs, linewidth=0.8, color="steelblue") - ax.set_xlabel("Environment step") - ax.set_ylabel("Cost (s)") - ax.set_title( - f"Cost per step — robot {self._robot_idx} | " - f"{self._task_suite_name} task {self._task_id}" + plot_cost_history( + costs, + out_folder, + robot_idx=self._robot_idx, + task_suite_name=self._task_suite_name, + task_id=self._task_id, ) - ax.grid(True, alpha=0.3) - fig.tight_layout() - fig.savefig(out_folder / "cost_history.png", dpi=150) - plt.close(fig) diff --git a/packages/armory-client/src/armory_client/runtime/saver_utils.py b/packages/armory-client/src/armory_client/runtime/saver_utils.py new file mode 100644 index 0000000..8f67ebf --- /dev/null +++ b/packages/armory-client/src/armory_client/runtime/saver_utils.py @@ -0,0 +1,100 @@ +"""Shared helpers for episode data persistence. + +Used by both the sim ``Saver`` (sims.libero.subscribers.saver) and the +real-robot ``RealSaver`` (armory_client.runtime.real_saver) so they emit the +same on-disk layout for the offline metrics pipeline to consume. +""" + +from __future__ import annotations + +import pathlib +from dataclasses import dataclass + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + +from armory_client.schemas import ( + ActionChunk, + JSONDataclass, + Observation, + Timestamp, +) + + +@dataclass(frozen=True) +class Result(JSONDataclass): + """Per-episode metadata persisted to ``metadata.json``.""" + + robot_idx: int + success: bool + steps_taken: int + task_suite_name: str + task_id: int + task_language: str + episode_idx: int + + +@dataclass +class EpisodeSaveData: + """Snapshot of one episode's data, safe to hand off to a background thread.""" + + timestamps: list[Timestamp] + observations_buffer: dict[int, Observation] + action_chunks: list[ActionChunk] + actions_left_snapshot: list[int] + cost_history: list[float] + success: bool + episode_idx: int + initial_state: np.ndarray | None + + +def save_timestamps(timestamps: list[Timestamp], out_folder: pathlib.Path) -> None: + Timestamp.to_csv(timestamps, out_folder / "timestamps.csv") + + +def save_action_chunks( + action_chunks: list[ActionChunk], out_folder: pathlib.Path +) -> None: + ActionChunk.to_parquet(action_chunks, out_folder / "action_chunks.parquet") + + +def save_actions_left( + actions_left_snapshot: list[int], out_folder: pathlib.Path +) -> None: + np.save( + out_folder / "actions_left.npy", + np.array(actions_left_snapshot, dtype=np.int32), + ) + + +def save_cost_history_npy( + cost_history: list[float], out_folder: pathlib.Path +) -> np.ndarray: + costs = np.array(cost_history, dtype=np.float64) + np.save(out_folder / "cost_history.npy", costs) + return costs + + +def plot_cost_history( + costs: np.ndarray, + out_folder: pathlib.Path, + robot_idx: int, + task_suite_name: str, + task_id: int, +) -> None: + steps = np.arange(len(costs)) + fig, ax = plt.subplots(figsize=(10, 4)) + ax.plot(steps, costs, linewidth=0.8, color="steelblue") + ax.set_xlabel("Environment step") + ax.set_ylabel("Cost (s)") + ax.set_title( + f"Cost per step — robot {robot_idx} | " + f"{task_suite_name} task {task_id}" + ) + ax.grid(True, alpha=0.3) + fig.tight_layout() + fig.savefig(out_folder / "cost_history.png", dpi=150) + plt.close(fig) diff --git a/packages/armory-client/src/armory_client/schemas.py b/packages/armory-client/src/armory_client/schemas.py index c20a41d..261eef1 100644 --- a/packages/armory-client/src/armory_client/schemas.py +++ b/packages/armory-client/src/armory_client/schemas.py @@ -256,7 +256,7 @@ def from_http_metadata(cls, payload: dict[str, Any]) -> "ServerMetadata": scheduler_kwargs: dict | None = None # Minimum action-index gap between serves for a robot (GPU worker throttle). - min_ex: int = 10 + min_execution_horizon: int = 0 # Set by Modal when running behind a tunnel; clients should use this for WebSocket tunnel_url: str | None = None diff --git a/requirements-modal.txt b/requirements-modal.txt index 797c467..3f1fb8e 100644 --- a/requirements-modal.txt +++ b/requirements-modal.txt @@ -73,6 +73,8 @@ async-lru==2.3.0 # via # armory # jupyterlab +asyncssh==2.21.0 + # via armory attrs==26.1.0 # via # aiohttp @@ -151,7 +153,7 @@ cloudpickle==2.1.0 # armory # gym # gymnasium -cmake==4.3.1 +cmake==3.30.0 # via # armory # lerobot @@ -171,6 +173,7 @@ contourpy==1.3.3 cryptography==46.0.7 # via # armory + # asyncssh # google-auth cycler==0.12.1 # via @@ -1563,6 +1566,7 @@ typing-extensions==4.15.0 # aiosignal # anyio # armory + # asyncssh # beautifulsoup4 # chex # dash diff --git a/scripts/exps/modal_alpha_fairness_sweep.py b/scripts/exps/modal_alpha_fairness_sweep.py new file mode 100644 index 0000000..6fa06be --- /dev/null +++ b/scripts/exps/modal_alpha_fairness_sweep.py @@ -0,0 +1,666 @@ +"""Sweep alpha for dynamic-action vs three baselines, on a list of (n_fast, n_slow) scenarios. + +For each (scenario, model): + - fixed-max-batch, greedy-deadline, round-robin: single point each (seeds averaged) + - dynamic-action: a curve over alpha = 0..1 + +Plot output: + - one figure per (scenario, model) + - x = mean starvation rate, y = Jain's index on freshness + - baselines render as labeled scatter points; dynamic-action renders as a connected curve + +Example: + modal run scripts/experiments/modal_alpha_fairness_sweep.py \ + --models pi05,gr00t-n1.7 \ + --scenarios 1f9s,5f5s \ + --alpha-grid 0.0,0.1,0.2,0.3,0.5,0.7,1.0 \ + --seeds 7,42 \ + --output-dir experiments/sweeps/fairness_alpha +""" + +from __future__ import annotations + +import csv +import dataclasses +import datetime as dt +import io +import json +import pathlib +import re +import subprocess +import sys +import tarfile +import threading +from typing import Any + +import modal + +APP_NAME = "armory-fairness-alpha-sweep" +REMOTE_ROOT = pathlib.Path("/app") +REMOTE_OUTPUT_ROOT = pathlib.Path("/tmp/armory_fairness_alpha_sweep") +PYTHONPATH = ":".join( + [ + str(REMOTE_ROOT / "src"), + str(REMOTE_ROOT / "src/backends"), + str(REMOTE_ROOT / "packages/armory-client/src"), + ] +) + +CONTROL_HZ = 20 +MAX_BATCH_SIZE = 5 +# MAX_BATCH_SIZE = 20 +FAST_HORIZON = 4 +SLOW_HORIZON = 10 +DEFAULT_MAX_STEPS = 200 +DEFAULT_TRIALS_PER_ROBOT = 1 + + +BASELINE_SCHEDULERS = ("fixed-max-batch", "greedy-deadline", "round-robin", "lookahead-actions") +DYNAMIC_SCHEDULER = "dynamic-action" + +MODEL_TO_PROFILE = { + "pi05": "l40s_pi05", + "gr00t-n1.7": "l40s_gr00t", +} + +MODEL_TO_ENUM_NAME = { + "pi05": "PI05", + "gr00t-n1.7": "GROOT_N17", +} + +BASELINE_STYLE = { + "fixed-max-batch": {"marker": "s", "color": "#1f77b4"}, + "greedy-deadline": {"marker": "^", "color": "#2ca02c"}, + "round-robin": {"marker": "D", "color": "#d62728"}, + "lookahead-actions": {"marker": "o", "color": "#9467bd"}, +} +DYNAMIC_CMAP = "viridis" + + +def _ignore_modal_copy(path: pathlib.Path) -> bool: + parts = set(path.parts) + return bool(parts & {".git", ".venv", ".ruff_cache", ".pytest_cache", "__pycache__"}) + + +image = ( + modal.Image.debian_slim(python_version="3.11") + .apt_install("git") + .pip_install_from_requirements("requirements-modal-mock.txt") + .workdir(str(REMOTE_ROOT)) + .env({"PYTHONPATH": PYTHONPATH, "MPLBACKEND": "Agg"}) + .add_local_dir("packages", str(REMOTE_ROOT / "packages"), copy=True, ignore=_ignore_modal_copy) + .add_local_dir("src", str(REMOTE_ROOT / "src"), copy=True, ignore=_ignore_modal_copy) + .add_local_dir("configs", str(REMOTE_ROOT / "configs"), copy=True, ignore=_ignore_modal_copy) + .add_local_dir("scripts", str(REMOTE_ROOT / "scripts"), copy=True, ignore=_ignore_modal_copy) +) + +app = modal.App(APP_NAME) + + +@dataclasses.dataclass(frozen=True) +class Scenario: + """A heterogeneity scenario: n_fast fast robots + n_slow slow robots.""" + + n_fast: int + n_slow: int + + @property + def n_total(self) -> int: + return self.n_fast + self.n_slow + + @property + def scenario_id(self) -> str: + return f"{self.n_fast}f{self.n_slow}s" + + def horizons(self) -> list[int]: + return [FAST_HORIZON] * self.n_fast + [SLOW_HORIZON] * self.n_slow + + +@dataclasses.dataclass(frozen=True) +class SweepCase: + model: str + scheduler: str + scenario_id: str + n_fast: int + n_slow: int + seed: int + # alpha is only meaningful for dynamic-action; None for the three baselines + alpha: float | None + + @property + def run_id(self) -> str: + alpha_part = "" if self.alpha is None else f"__alpha={self.alpha:.3f}" + return ( + f"model={self.model}__scheduler={self.scheduler}" + f"__scenario={self.scenario_id}{alpha_part}__seed={self.seed}" + ) + + +def _parse_csv(value: str, *, cast=str) -> list[Any]: + return [cast(item.strip()) for item in value.split(",") if item.strip()] + + +def _parse_scenarios(value: str) -> list[Scenario]: + """Parse strings like '1f9s,5f5s' into Scenario objects.""" + pattern = re.compile(r"^\s*(\d+)f(\d+)s\s*$") + out: list[Scenario] = [] + for token in value.split(","): + if not token.strip(): + continue + m = pattern.match(token) + if not m: + raise ValueError( + f"Bad scenario token {token!r}; expected 'fs' (e.g. '1f9s')" + ) + out.append(Scenario(n_fast=int(m.group(1)), n_slow=int(m.group(2)))) + return out + + +def _build_experiment_config(horizons: list[int], max_steps: int) -> dict[str, Any]: + return { + "experiment": { + "action_chunk_broker_type": "naive_async", + "num_robots": len(horizons), + "trials_per_robot": DEFAULT_TRIALS_PER_ROBOT, + "max_steps": max_steps, + }, + "robots": { + f"robot_{i}": { + "execution_horizon": int(h), + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0, + } + for i, h in enumerate(horizons) + }, + } + + +def _build_server_cmd( + *, model: str, scheduler: str, port: int, alpha: float | None +) -> list[str]: + profile = MODEL_TO_PROFILE[model] + pre_policy = [ + sys.executable, + "scripts/serve.py", + "--port", + str(port), + "--env", + "LIBERO", + "--model", + MODEL_TO_ENUM_NAME[model], + "--max-batch-size", + str(MAX_BATCH_SIZE), + "--scheduling-algorithm", + scheduler, + ] + if scheduler == DYNAMIC_SCHEDULER: + if alpha is None: + raise ValueError("dynamic-action requires alpha") + pre_policy += ["--alpha", str(alpha)] + post_policy = [ + "policy:mock", + "--policy.action-horizon", + "10", + "--policy.action-dim", + "7", + "--policy.profile", + profile, + ] + return pre_policy + post_policy + + +def _build_client_cmd( + *, port: int, seed: int, output_dir: pathlib.Path, experiment_config_path: pathlib.Path +) -> list[str]: + return [ + sys.executable, + "scripts/run_libero.py", + "--host", + "127.0.0.1", + "--port", + str(port), + "--env", + "mock", + "--overwrite", + "--progress-type", + "logging", + "--seed", + str(seed), + "--output-dir", + str(output_dir), + "--experiment-config", + str(experiment_config_path), + ] + + +def _stream_to_log_and_stdout(stream, log_file, prefix: str) -> None: + """Tee a subprocess text stream line-by-line to both ``log_file`` and ``sys.stdout``. + + Routing subprocess output via this helper makes it visible in the Modal + container log (which only captures the function's own stdout/stderr) while + still preserving the per-run log files in the artifact tarball. + """ + for raw in stream: + log_file.write(raw) + log_file.flush() + line = raw if raw.endswith("\n") else raw + "\n" + sys.stdout.write(f"[{prefix}] {line}") + sys.stdout.flush() + + +def _run_subprocess( + args: list[str], + *, + log_path: pathlib.Path, + prefix: str, + timeout_s: int | None = None, +) -> None: + log_path.parent.mkdir(parents=True, exist_ok=True) + log_file = log_path.open("w") + proc = subprocess.Popen( + args, + cwd=REMOTE_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env={ + **{k: v for k, v in __import__("os").environ.items()}, + **dict(PYTHONPATH=PYTHONPATH, MPLBACKEND="Agg"), + }, + ) + reader = threading.Thread( + target=_stream_to_log_and_stdout, + args=(proc.stdout, log_file, prefix), + daemon=True, + ) + reader.start() + try: + rc = proc.wait(timeout=timeout_s) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + reader.join(timeout=5) + log_file.close() + raise + reader.join(timeout=5) + log_file.close() + if rc != 0: + raise subprocess.CalledProcessError(rc, args) + + +def _tar_directory(path: pathlib.Path) -> bytes: + def compact_filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: + if pathlib.Path(info.name).suffix in {".mp4", ".parquet", ".npz"}: + return None + return info + + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w:gz") as tar: + tar.add(path, arcname=path.name, filter=compact_filter) + return buffer.getvalue() + + +def _summarize_run(output_dir: pathlib.Path, case: SweepCase, horizons: list[int]) -> dict[str, Any]: + from sims.libero.metrics import ( # noqa: PLC0415 + compute_fairness_metrics, + compute_starvation_variance_series, + ) + + fairness = compute_fairness_metrics(output_dir) + summary: dict[str, Any] = { + "run_id": case.run_id, + "model": case.model, + "scheduler": case.scheduler, + "scenario_id": case.scenario_id, + "n_fast": case.n_fast, + "n_slow": case.n_slow, + "n_total": case.n_fast + case.n_slow, + "alpha_requested": case.alpha, + "seed": case.seed, + "horizons": json.dumps(horizons), + } + if fairness is not None: + summary["alpha_observed"] = fairness.get("alpha") + summary["jain_freshness"] = fairness["jain_freshness"] + summary["jain_starvation"] = fairness["jain_starvation"] + rates = fairness["starvation_rate"] + if rates: + summary["mean_starvation"] = float(sum(rates) / len(rates)) + summary["max_starvation"] = float(max(rates)) + summary["min_starvation"] = float(min(rates)) + # α=∞ welfare: max-min freshness, equivalently 1 - max(starvation). + # Asymmetric: only the worst-off robot's freshness shows up. + summary["min_freshness"] = 1.0 - float(max(rates)) + + # Pull cross-robot starvation variance from the same source as the + # starvation_variance_over_time plot so the sweep summary matches its + # final value exactly (actions_left<=0 on a wall-clock canvas, not + # cost_history NaNs aggregated per episode). + series = compute_starvation_variance_series(output_dir) + if series is not None: + summary["starvation_variance"] = series["final_starvation_variance"] + return summary + + +# @app.function(image=image, timeout=60 * 60, cpu=25, memory=16384) +@app.function(image=image, timeout=60 * 60, cpu=25, memory=16384) +def run_case(case: SweepCase, *, port: int, max_steps: int) -> dict[str, Any]: + horizons = [FAST_HORIZON] * case.n_fast + [SLOW_HORIZON] * case.n_slow + exp_cfg = _build_experiment_config(horizons, max_steps) + + run_dir = REMOTE_OUTPUT_ROOT / case.run_id + output_dir = run_dir / "output" + log_dir = run_dir / "logs" + run_dir.mkdir(parents=True, exist_ok=True) + log_dir.mkdir(parents=True, exist_ok=True) + + saved_exp_config = run_dir / "experiment_config.json" + (run_dir / "case.json").write_text(json.dumps(dataclasses.asdict(case), indent=2)) + saved_exp_config.write_text(json.dumps(exp_cfg, indent=2)) + + server_cmd = _build_server_cmd( + model=case.model, scheduler=case.scheduler, port=port, alpha=case.alpha + ) + client_cmd = _build_client_cmd( + port=port, seed=case.seed, output_dir=output_dir, experiment_config_path=saved_exp_config + ) + + server_log = log_dir / "server.log" + server_log_file = server_log.open("w") + server_proc = subprocess.Popen( + server_cmd, + cwd=REMOTE_ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env={ + **{k: v for k, v in __import__("os").environ.items()}, + **dict(PYTHONPATH=PYTHONPATH, MPLBACKEND="Agg"), + }, + ) + server_reader = threading.Thread( + target=_stream_to_log_and_stdout, + args=(server_proc.stdout, server_log_file, "server"), + daemon=True, + ) + server_reader.start() + try: + _run_subprocess( + client_cmd, + log_path=log_dir / "client.log", + prefix="client", + timeout_s=60 * 30, + ) + summary = _summarize_run(output_dir, case, horizons) + summary["status"] = "ok" + except Exception as exc: # noqa: BLE001 + summary = { + "run_id": case.run_id, + "model": case.model, + "scheduler": case.scheduler, + "scenario_id": case.scenario_id, + "n_fast": case.n_fast, + "n_slow": case.n_slow, + "n_total": case.n_fast + case.n_slow, + "alpha_requested": case.alpha, + "seed": case.seed, + "horizons": json.dumps(horizons), + "status": "failed", + "error": repr(exc), + } + finally: + server_proc.terminate() + try: + server_proc.wait(timeout=20) + except subprocess.TimeoutExpired: + server_proc.kill() + server_proc.wait(timeout=20) + server_reader.join(timeout=5) + server_log_file.close() + + summary["artifact_tgz"] = _tar_directory(run_dir) + return summary + + +def _write_rows(path: pathlib.Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + keys: list[str] = [] + for row in rows: + for key in row: + if key != "artifact_tgz" and key not in keys: + keys.append(key) + with path.open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=keys) + writer.writeheader() + for row in rows: + writer.writerow({key: row.get(key, "") for key in keys}) + + +def _plot_one_yaxis(ax, sub, *, y_col: str, y_label: str, title: str): + """Scatter (mean_starvation, y_col) onto ``ax`` with baselines + dynamic-action curve. + + Returns the dynamic-action scatter handle (for a shared colorbar) or None. + """ + import numpy as np # noqa: PLC0415 + + handle = None + for sched in BASELINE_SCHEDULERS: + sub_b = sub[sub["scheduler"] == sched].dropna(subset=["mean_starvation", y_col]) + if sub_b.empty: + continue + x = float(sub_b["mean_starvation"].mean()) + y = float(sub_b[y_col].mean()) + xerr = float(sub_b["mean_starvation"].std(ddof=0)) if len(sub_b) > 1 else 0.0 + yerr = float(sub_b[y_col].std(ddof=0)) if len(sub_b) > 1 else 0.0 + style = BASELINE_STYLE[sched] + ax.errorbar( + x, y, xerr=xerr, yerr=yerr, + marker=style["marker"], markersize=11, color=style["color"], + linestyle="none", capsize=3, label=sched, zorder=4, + ) + + sub_d = sub[sub["scheduler"] == DYNAMIC_SCHEDULER].dropna(subset=["mean_starvation", y_col]) + if not sub_d.empty: + agg = ( + sub_d.groupby("alpha_requested") + .agg( + mean_starvation=("mean_starvation", "mean"), + y_mean=(y_col, "mean"), + starvation_std=("mean_starvation", "std"), + y_std=(y_col, "std"), + count=("seed", "count"), + ) + .reset_index() + .sort_values("alpha_requested") + ) + xs = agg["mean_starvation"].to_numpy() + ys = agg["y_mean"].to_numpy() + alphas = agg["alpha_requested"].to_numpy() + ax.plot(xs, ys, "-", color="0.5", linewidth=1.2, alpha=0.7, zorder=2) + handle = ax.scatter( + xs, ys, c=alphas, cmap=DYNAMIC_CMAP, s=70, + edgecolors="black", linewidths=0.6, zorder=3, + label=f"{DYNAMIC_SCHEDULER} (alpha sweep)", + vmin=0.0, vmax=1.0, + ) + n_seeds = int(agg["count"].max()) if not agg.empty else 1 + if n_seeds > 1: + xerr = (agg["starvation_std"].fillna(0.0) / np.sqrt(n_seeds)).to_numpy() + yerr = (agg["y_std"].fillna(0.0) / np.sqrt(n_seeds)).to_numpy() + ax.errorbar( + xs, ys, xerr=xerr, yerr=yerr, + fmt="none", ecolor="0.6", capsize=2, alpha=0.6, zorder=2, + ) + + ax.set_xlabel("Mean starvation rate (lower is better)", fontsize=11) + ax.set_ylabel(y_label, fontsize=11) + ax.set_title(title, fontsize=12, fontweight="bold") + ax.grid(True, alpha=0.3) + ax.legend(loc="best", fontsize=8, frameon=False) + return handle + + +def _autoscale_with_pad(ax, sub, y_col: str, pad_frac: float = 0.12) -> None: + """Tighten x/y limits to the data range with a fractional padding.""" + s = sub.dropna(subset=["mean_starvation", y_col]) + if s.empty: + return + xs = s["mean_starvation"].to_numpy(dtype=float) + ys = s[y_col].to_numpy(dtype=float) + x_lo, x_hi = float(xs.min()), float(xs.max()) + y_lo, y_hi = float(ys.min()), float(ys.max()) + x_pad = max((x_hi - x_lo) * pad_frac, 1e-6) + y_pad = max((y_hi - y_lo) * pad_frac, 1e-6) + ax.set_xlim(x_lo - x_pad, x_hi + x_pad) + ax.set_ylim(y_lo - y_pad, y_hi + y_pad) + + +def _plot_starvation_vs_fairness(results_csv: pathlib.Path, plots_dir: pathlib.Path) -> None: + """One PNG per (scenario, model) with two panels: variance and min-freshness.""" + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + import pandas as pd + + df = pd.read_csv(results_csv) + if df.empty: + return + df = df[df.get("status", "ok") == "ok"] + needed = { + "mean_starvation", "starvation_variance", "min_freshness", + "scheduler", "model", "scenario_id", + } + if not needed.issubset(df.columns): + return + df = df.dropna(subset=["mean_starvation"]) + + plots_dir.mkdir(parents=True, exist_ok=True) + + for (scenario_id, model), sub in df.groupby(["scenario_id", "model"]): + fig, (ax_var, ax_min) = plt.subplots(1, 2, figsize=(13, 5.5)) + + _plot_one_yaxis( + ax_var, sub, + y_col="starvation_variance", + y_label="Cross-robot starvation variance (lower is fairer)", + title="Mean starvation vs variance", + ) + ax_var.set_ylim(bottom=0.0) + ax_var.set_xlim(left=0.0) + + handle = _plot_one_yaxis( + ax_min, sub, + y_col="min_freshness", + y_label="Min freshness = 1 − max(starvation) (higher is fairer)", + title="Mean starvation vs min freshness", + ) + _autoscale_with_pad(ax_min, sub, "min_freshness") + + if handle is not None: + cbar = fig.colorbar(handle, ax=[ax_var, ax_min], pad=0.02, fraction=0.03) + cbar.set_label("scheduler alpha", fontsize=10) + + fig.suptitle( + f"Starvation vs fairness — scenario={scenario_id}, model={model}", + fontsize=13, fontweight="bold", + ) + safe_model = model.replace(".", "_").replace("/", "_") + out = plots_dir / f"starvation_vs_fairness__{scenario_id}__{safe_model}.png" + fig.savefig(out, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"Wrote {out}") + + +@app.local_entrypoint() +def main( + # models: str = "pi05,gr00t-n1.7", + models: str = "pi05", + # scenarios: str = "1f9s,5f5s", + scenarios: str = "1f9s", + alpha_grid: str = "0.0,0.25,0.5,0.75,1.0", + seeds: str = "42", + output_dir: str = "experiments/sweeps/fairness_alpha_sweep_pi05_1f9s", + port: int = 8080, + max_steps: int = DEFAULT_MAX_STEPS, +) -> None: + out = pathlib.Path(output_dir) + stamp = dt.datetime.now(dt.timezone.utc).strftime("%Y%m%d_%H%M%S") # noqa: UP017 + artifacts_dir = out / "artifacts" + + model_list = _parse_csv(models) + scenario_list = _parse_scenarios(scenarios) + alpha_list = _parse_csv(alpha_grid, cast=float) + seed_list = _parse_csv(seeds, cast=int) + + cases: list[SweepCase] = [] + for scenario in scenario_list: + for model in model_list: + for seed in seed_list: + # baselines: one per scheduler + for sched in BASELINE_SCHEDULERS: + cases.append( + SweepCase( + model=model, + scheduler=sched, + scenario_id=scenario.scenario_id, + n_fast=scenario.n_fast, + n_slow=scenario.n_slow, + seed=seed, + alpha=None, + ) + ) + # dynamic-action: one per alpha + for alpha in alpha_list: + cases.append( + SweepCase( + model=model, + scheduler=DYNAMIC_SCHEDULER, + scenario_id=scenario.scenario_id, + n_fast=scenario.n_fast, + n_slow=scenario.n_slow, + seed=seed, + alpha=alpha, + ) + ) + + print( + f"Submitting {len(cases)} cases " + f"(scenarios={[s.scenario_id for s in scenario_list]} " + f"models={model_list} alphas={alpha_list} seeds={seed_list})" + ) + + rows: list[dict[str, Any]] = [] + for result in run_case.map( + cases, + kwargs={"port": port, "max_steps": max_steps}, + order_outputs=False, + ): + artifact_bytes = result.pop("artifact_tgz", None) + if artifact_bytes is not None: + run_dir = artifacts_dir / result["run_id"] + run_dir.mkdir(parents=True, exist_ok=True) + with tarfile.open(fileobj=io.BytesIO(artifact_bytes), mode="r:gz") as tar: + tar.extractall(run_dir) + result["artifact_path"] = str(run_dir) + rows.append(result) + mf = result.get("min_freshness", "") + starv = result.get("mean_starvation", "") + print( + f"{result['status']}: {result['run_id']} " + f"min_freshness={mf if mf == '' else f'{float(mf):.4f}'} " + f"mean_starvation={starv if starv == '' else f'{float(starv):.4f}'}" + ) + + sweep_csv = out / f"sweep_results_{stamp}.csv" + latest_csv = out / "sweep_results.csv" + _write_rows(sweep_csv, rows) + _write_rows(latest_csv, rows) + print(f"Wrote {latest_csv}") + print(f"Wrote {sweep_csv}") + + _plot_starvation_vs_fairness(latest_csv, out / "plots") diff --git a/scripts/exps/modal_fairness_sweep.py b/scripts/exps/modal_fairness_sweep.py new file mode 100644 index 0000000..049ce46 --- /dev/null +++ b/scripts/exps/modal_fairness_sweep.py @@ -0,0 +1,438 @@ +"""Run a fairness-vs-heterogeneity sweep on Modal. + +Sweeps n_fast (number of "fast" robots in a 15-robot fleet) for each combination of +(model, scheduler, seed) and reports Jain's index on per-robot starvation rate. + +Heterogeneity definition: + - 15 robots total, partitioned into n_fast fast + (15 - n_fast) slow robots + - fast robots: execution_horizon = FAST_HORIZON (4) + - slow robots: execution_horizon = SLOW_HORIZON (10) + - n_fast is the only knob; per-robot horizons are fixed regardless of model + +Plot output: + - one figure per model (pi05, gr00t) + - x = n_fast, y = Jain's freshness index, one line per scheduler + +Example: + modal run scripts/modal_fairness_sweep.py \ + --models pi05,gr00t-n1.7 \ + --schedulers fixed-max-batch,greedy-deadline,round-robin,dynamic-action \ + --n-fast-grid 0,1,3,5,7,10,12,14,15 \ + --seeds 7,42 \ + --output-dir experiments/sweeps/fairness_het +""" + +from __future__ import annotations + +import csv +import dataclasses +import datetime as dt +import io +import json +import pathlib +import subprocess +import sys +import tarfile +from typing import Any + +import modal + +APP_NAME = "armory-fairness-sweep" +REMOTE_ROOT = pathlib.Path("/app") +REMOTE_OUTPUT_ROOT = pathlib.Path("/tmp/armory_fairness_sweep") +PYTHONPATH = ":".join( + [ + str(REMOTE_ROOT / "src"), + str(REMOTE_ROOT / "src/backends"), + str(REMOTE_ROOT / "packages/armory-client/src"), + ] +) + +NUM_ROBOTS = 15 +CONTROL_HZ = 20 +MAX_BATCH_SIZE = 4 +FAST_HORIZON = 4 +SLOW_HORIZON = 10 +ALPHA_FOR_DYNAMIC = 1.0 +DEFAULT_MAX_STEPS = 200 +DEFAULT_TRIALS_PER_ROBOT = 1 + +MODEL_TO_PROFILE = { + "pi05": "l40s_pi05", + "gr00t-n1.7": "l40s_gr00t", +} + +# tyro accepts the ModelFamily enum *name*, not its value +MODEL_TO_ENUM_NAME = { + "pi05": "PI05", + "gr00t-n1.7": "GROOT_N17", +} + +SCHEDULER_DISPLAY = { + "fixed-max-batch": "fixed-max-batch", + "greedy-deadline": "greedy-deadline", + "round-robin": "round-robin", + "dynamic-action": f"dynamic-action (α={ALPHA_FOR_DYNAMIC})", +} + + +def _ignore_modal_copy(path: pathlib.Path) -> bool: + parts = set(path.parts) + return bool(parts & {".git", ".venv", ".ruff_cache", ".pytest_cache", "__pycache__"}) + + +image = ( + modal.Image.debian_slim(python_version="3.11") + .apt_install("git") + .pip_install_from_requirements("requirements-modal-mock.txt") + .workdir(str(REMOTE_ROOT)) + .env({"PYTHONPATH": PYTHONPATH, "MPLBACKEND": "Agg"}) + .add_local_dir("packages", str(REMOTE_ROOT / "packages"), copy=True, ignore=_ignore_modal_copy) + .add_local_dir("src", str(REMOTE_ROOT / "src"), copy=True, ignore=_ignore_modal_copy) + .add_local_dir("configs", str(REMOTE_ROOT / "configs"), copy=True, ignore=_ignore_modal_copy) + .add_local_dir("scripts", str(REMOTE_ROOT / "scripts"), copy=True, ignore=_ignore_modal_copy) +) + +app = modal.App(APP_NAME) + + +@dataclasses.dataclass(frozen=True) +class SweepCase: + model: str + scheduler: str + n_fast: int + seed: int + + @property + def run_id(self) -> str: + return ( + f"model={self.model}__scheduler={self.scheduler}" + f"__nfast={self.n_fast:02d}__seed={self.seed}" + ) + + +def _parse_csv(value: str, *, cast=str) -> list[Any]: + return [cast(item.strip()) for item in value.split(",") if item.strip()] + + +def _compute_horizons(n_fast: int, n_total: int = NUM_ROBOTS) -> list[int]: + """Per-robot execution horizons: n_fast at FAST_HORIZON, the rest at SLOW_HORIZON.""" + n_slow = n_total - n_fast + return [FAST_HORIZON] * n_fast + [SLOW_HORIZON] * n_slow + + +def _build_experiment_config(horizons: list[int], max_steps: int) -> dict[str, Any]: + return { + "experiment": { + "action_chunk_broker_type": "naive_async", + "num_robots": len(horizons), + "trials_per_robot": DEFAULT_TRIALS_PER_ROBOT, + "max_steps": max_steps, + }, + "robots": { + f"robot_{i}": { + "execution_horizon": int(h), + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0, + } + for i, h in enumerate(horizons) + }, + } + + +def _build_server_cmd(*, model: str, scheduler: str, port: int) -> list[str]: + profile = MODEL_TO_PROFILE[model] + # All top-level Args flags must come BEFORE the policy:mock subcommand, + # otherwise tyro treats them as belonging to the policy subcommand. + pre_policy = [ + sys.executable, + "scripts/serve.py", + "--port", + str(port), + "--env", + "LIBERO", + "--model", + MODEL_TO_ENUM_NAME[model], + "--max-batch-size", + str(MAX_BATCH_SIZE), + "--scheduling-algorithm", + scheduler, + ] + if scheduler == "dynamic-action": + pre_policy += ["--alpha", str(ALPHA_FOR_DYNAMIC)] + post_policy = [ + "policy:mock", + "--policy.action-horizon", + "10", + "--policy.action-dim", + "7", + "--policy.profile", + profile, + ] + return pre_policy + post_policy + + +def _build_client_cmd( + *, port: int, seed: int, output_dir: pathlib.Path, experiment_config_path: pathlib.Path +) -> list[str]: + return [ + sys.executable, + "scripts/run_libero.py", + "--host", + "127.0.0.1", + "--port", + str(port), + "--env", + "mock", + "--overwrite", + "--progress-type", + "logging", + "--seed", + str(seed), + "--output-dir", + str(output_dir), + "--experiment-config", + str(experiment_config_path), + ] + + +def _run_subprocess(args: list[str], *, log_path: pathlib.Path, timeout_s: int | None = None) -> None: + log_path.parent.mkdir(parents=True, exist_ok=True) + with log_path.open("w") as log_file: + result = subprocess.run( + args, + cwd=REMOTE_ROOT, + stdout=log_file, + stderr=subprocess.STDOUT, + text=True, + timeout=timeout_s, + env={ + **{k: v for k, v in __import__("os").environ.items()}, + **dict(PYTHONPATH=PYTHONPATH, MPLBACKEND="Agg"), + }, + ) + if result.returncode != 0: + raise subprocess.CalledProcessError(result.returncode, args) + + +def _tar_directory(path: pathlib.Path) -> bytes: + def compact_filter(info: tarfile.TarInfo) -> tarfile.TarInfo | None: + if pathlib.Path(info.name).suffix in {".mp4", ".parquet", ".npz"}: + return None + return info + + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w:gz") as tar: + tar.add(path, arcname=path.name, filter=compact_filter) + return buffer.getvalue() + + +def _summarize_run(output_dir: pathlib.Path, case: SweepCase, horizons: list[int]) -> dict[str, Any]: + from sims.libero.metrics import compute_fairness_metrics # noqa: PLC0415 + + fairness = compute_fairness_metrics(output_dir) + summary: dict[str, Any] = { + "run_id": case.run_id, + "model": case.model, + "scheduler": case.scheduler, + "n_fast": case.n_fast, + "seed": case.seed, + "horizons": json.dumps(horizons), + } + if fairness is not None: + summary["alpha"] = fairness.get("alpha") + summary["jain_freshness"] = fairness["jain_freshness"] + summary["jain_starvation"] = fairness["jain_starvation"] + rates = fairness["starvation_rate"] + if rates: + summary["mean_starvation"] = float(sum(rates) / len(rates)) + summary["max_starvation"] = float(max(rates)) + summary["min_starvation"] = float(min(rates)) + return summary + + +@app.function(image=image, timeout=60 * 60, cpu=4, memory=16384) +def run_case(case: SweepCase, *, port: int, max_steps: int) -> dict[str, Any]: + horizons = _compute_horizons(case.n_fast) + exp_cfg = _build_experiment_config(horizons, max_steps) + + run_dir = REMOTE_OUTPUT_ROOT / case.run_id + output_dir = run_dir / "output" + log_dir = run_dir / "logs" + run_dir.mkdir(parents=True, exist_ok=True) + log_dir.mkdir(parents=True, exist_ok=True) + + saved_exp_config = run_dir / "experiment_config.json" + (run_dir / "case.json").write_text(json.dumps(dataclasses.asdict(case), indent=2)) + saved_exp_config.write_text(json.dumps(exp_cfg, indent=2)) + + server_cmd = _build_server_cmd(model=case.model, scheduler=case.scheduler, port=port) + client_cmd = _build_client_cmd( + port=port, seed=case.seed, output_dir=output_dir, experiment_config_path=saved_exp_config + ) + + server_log = log_dir / "server.log" + with server_log.open("w") as log_file: + server_proc = subprocess.Popen( + server_cmd, + cwd=REMOTE_ROOT, + stdout=log_file, + stderr=subprocess.STDOUT, + text=True, + env={ + **{k: v for k, v in __import__("os").environ.items()}, + **dict(PYTHONPATH=PYTHONPATH, MPLBACKEND="Agg"), + }, + ) + try: + _run_subprocess(client_cmd, log_path=log_dir / "client.log", timeout_s=60 * 30) + summary = _summarize_run(output_dir, case, horizons) + summary["status"] = "ok" + except Exception as exc: # noqa: BLE001 + summary = { + "run_id": case.run_id, + "model": case.model, + "scheduler": case.scheduler, + "n_fast": case.n_fast, + "seed": case.seed, + "horizons": json.dumps(horizons), + "status": "failed", + "error": repr(exc), + } + finally: + server_proc.terminate() + try: + server_proc.wait(timeout=20) + except subprocess.TimeoutExpired: + server_proc.kill() + server_proc.wait(timeout=20) + + summary["artifact_tgz"] = _tar_directory(run_dir) + return summary + + +def _write_rows(path: pathlib.Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + keys: list[str] = [] + for row in rows: + for key in row: + if key != "artifact_tgz" and key not in keys: + keys.append(key) + with path.open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=keys) + writer.writeheader() + for row in rows: + writer.writerow({key: row.get(key, "") for key in keys}) + + +def _plot_fairness_curves(results_csv: pathlib.Path, plots_dir: pathlib.Path) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + import numpy as np + import pandas as pd + + df = pd.read_csv(results_csv) + if df.empty: + return + df = df[df.get("status", "ok") == "ok"] + if "jain_freshness" not in df.columns: + return + df = df.dropna(subset=["jain_freshness"]) + + plots_dir.mkdir(parents=True, exist_ok=True) + schedulers = sorted(df["scheduler"].unique()) + color_cycle = plt.cm.tab10(np.linspace(0, 1, max(len(schedulers), 2))) + + for model in sorted(df["model"].unique()): + sub_model = df[df["model"] == model] + fig, ax = plt.subplots(figsize=(9, 5)) + for color, sched in zip(color_cycle, schedulers): + sub = sub_model[sub_model["scheduler"] == sched] + if sub.empty: + continue + agg = ( + sub.groupby("n_fast")["jain_freshness"] + .agg(["mean", "std", "count"]) + .reset_index() + .sort_values("n_fast") + ) + yerr = agg["std"].fillna(0.0) / agg["count"].clip(lower=1).pow(0.5) + ax.errorbar( + agg["n_fast"], + agg["mean"], + yerr=yerr, + marker="o", + linewidth=1.6, + capsize=3, + color=color, + label=SCHEDULER_DISPLAY.get(sched, sched), + ) + ax.set_xlabel("Heterogeneity (n_fast / 15)", fontsize=12) + ax.set_ylabel("Jain's index on freshness rate", fontsize=12) + ax.set_title(f"Fairness vs heterogeneity — {model}", fontsize=13, fontweight="bold") + ax.set_ylim(0, 1.02) + ax.set_xlim(-0.5, NUM_ROBOTS + 0.5) + ax.grid(True, alpha=0.3) + ax.legend(loc="lower left", fontsize=9, frameon=False) + plt.tight_layout() + out = plots_dir / f"jains_vs_het__{model.replace('.', '_').replace('/', '_')}.png" + fig.savefig(out, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"Wrote {out}") + + +@app.local_entrypoint() +def main( + models: str = "pi05,gr00t-n1.7", + schedulers: str = "fixed-max-batch,greedy-deadline,round-robin,dynamic-action", + n_fast_grid: str = "0,1,3,5,7,10,12,14,15", + seeds: str = "7,42,123", + output_dir: str = "experiments/sweeps/fairness_het_2", + port: int = 8080, + max_steps: int = DEFAULT_MAX_STEPS, +) -> None: + out = pathlib.Path(output_dir) + stamp = dt.datetime.now(dt.timezone.utc).strftime("%Y%m%d_%H%M%S") # noqa: UP017 + artifacts_dir = out / "artifacts" + + cases = [ + SweepCase(model=m, scheduler=s, n_fast=n, seed=seed) + for m in _parse_csv(models) + for s in _parse_csv(schedulers) + for n in _parse_csv(n_fast_grid, cast=int) + for seed in _parse_csv(seeds, cast=int) + ] + print(f"Submitting {len(cases)} cases") + + rows: list[dict[str, Any]] = [] + for result in run_case.map( + cases, + kwargs={"port": port, "max_steps": max_steps}, + order_outputs=False, + ): + artifact_bytes = result.pop("artifact_tgz", None) + if artifact_bytes is not None: + run_dir = artifacts_dir / result["run_id"] + run_dir.mkdir(parents=True, exist_ok=True) + with tarfile.open(fileobj=io.BytesIO(artifact_bytes), mode="r:gz") as tar: + tar.extractall(run_dir) + result["artifact_path"] = str(run_dir) + rows.append(result) + jain = result.get("jain_freshness", "") + print( + f"{result['status']}: {result['run_id']} " + f"jain_freshness={jain if jain == '' else f'{float(jain):.4f}'}" + ) + + sweep_csv = out / f"sweep_results_{stamp}.csv" + latest_csv = out / "sweep_results.csv" + _write_rows(sweep_csv, rows) + _write_rows(latest_csv, rows) + print(f"Wrote {latest_csv}") + print(f"Wrote {sweep_csv}") + + _plot_fairness_curves(latest_csv, out / "plots") diff --git a/scripts/modal_sweep.py b/scripts/exps/modal_starvation_sweep.py similarity index 100% rename from scripts/modal_sweep.py rename to scripts/exps/modal_starvation_sweep.py diff --git a/scripts/exps/plot_fairness_sweep.py b/scripts/exps/plot_fairness_sweep.py new file mode 100644 index 0000000..a3474ab --- /dev/null +++ b/scripts/exps/plot_fairness_sweep.py @@ -0,0 +1,300 @@ +"""Post-process a fairness sweep produced by scripts/modal_fairness_sweep.py. + +Walks the sweep's artifact directory, recomputes per-cell fairness metrics +(Jain's freshness, Jain's starvation, cross-robot starvation variance) from +each run's per-robot starvation rates, writes a richer CSV, and emits two +plots per model: Jain's vs n_fast and starvation variance vs n_fast. + +Example: + uv run python scripts/plot_fairness_sweep.py \ + --sweep-dir experiments/sweeps/fairness_het_2 +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys +from typing import Any + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "src")) +sys.path.insert(0, str(REPO_ROOT / "src/backends")) +sys.path.insert(0, str(REPO_ROOT / "packages/armory-client/src")) + +from sims.libero.metrics import compute_fairness_metrics # noqa: E402 + + +def _find_output_dir(artifact_dir: pathlib.Path) -> pathlib.Path | None: + """Locate the run's output directory inside an artifact tree. + + Modal sweep tarballs nest as //output/, but some local + layouts collapse to /output/ — handle both. + """ + for candidate in ( + artifact_dir / "output", + artifact_dir / artifact_dir.name / "output", + ): + if candidate.exists(): + return candidate + return None + + +def _per_cell_metrics( + output_dir: pathlib.Path, n_fast: int +) -> dict[str, Any] | None: + fairness = compute_fairness_metrics(output_dir) + if fairness is None: + return None + rates = fairness["starvation_rate"] + robot_idx = fairness.get("robot_idx") or list(range(len(rates))) + if not rates: + return None + rates_arr = np.asarray(rates, dtype=float) + # Convention from _compute_horizons: robots [0..n_fast-1] are fast, rest are slow. + fast_mask = np.asarray([int(i) < n_fast for i in robot_idx]) + fast_rates = rates_arr[fast_mask] + slow_rates = rates_arr[~fast_mask] + return { + "alpha": fairness.get("alpha"), + "jain_freshness": float(fairness["jain_freshness"]), + "jain_starvation": float(fairness["jain_starvation"]), + "starvation_var": float(np.var(rates_arr)), + "starvation_std": float(np.std(rates_arr)), + "starvation_max_minus_min": float(rates_arr.max() - rates_arr.min()), + "starvation_mean": float(rates_arr.mean()), + "starvation_max": float(rates_arr.max()), + "starvation_min": float(rates_arr.min()), + "fast_cohort_mean_starvation": float(fast_rates.mean()) if fast_rates.size else float("nan"), + "slow_cohort_mean_starvation": float(slow_rates.mean()) if slow_rates.size else float("nan"), + "n_robots_observed": int(rates_arr.size), + } + + +def _load_case(artifact_dir: pathlib.Path) -> dict[str, Any] | None: + case_path = artifact_dir / "case.json" + nested = artifact_dir / artifact_dir.name / "case.json" + if case_path.exists(): + return json.loads(case_path.read_text()) + if nested.exists(): + return json.loads(nested.read_text()) + return None + + +def collect_metrics(sweep_dir: pathlib.Path) -> pd.DataFrame: + artifacts_dir = sweep_dir / "artifacts" + rows: list[dict[str, Any]] = [] + for artifact in sorted(artifacts_dir.iterdir()): + if not artifact.is_dir(): + continue + output_dir = _find_output_dir(artifact) + case = _load_case(artifact) + if output_dir is None or case is None: + print(f"[skip] {artifact.name}: missing output/ or case.json") + continue + metrics = _per_cell_metrics(output_dir, int(case.get("n_fast", 0))) + if metrics is None: + print(f"[skip] {artifact.name}: no fairness metrics (run may have failed)") + continue + rows.append({**case, **metrics, "run_id": artifact.name}) + df = pd.DataFrame(rows) + if df.empty: + return df + df = df.sort_values(["model", "scheduler", "n_fast", "seed"]).reset_index(drop=True) + return df + + +def _plot_metric_vs_nfast( + df: pd.DataFrame, + metric: str, + *, + ylabel: str, + title_prefix: str, + out_path: pathlib.Path, + ylim: tuple[float, float] | None = None, + legend_loc: str = "best", +) -> None: + schedulers = sorted(df["scheduler"].unique()) + color_cycle = plt.cm.tab10(np.linspace(0, 1, max(len(schedulers), 2))) + + fig, ax = plt.subplots(figsize=(9, 5)) + for color, sched in zip(color_cycle, schedulers): + sub = df[df["scheduler"] == sched] + if sub.empty: + continue + agg = ( + sub.groupby("n_fast")[metric] + .agg(["mean", "std", "count"]) + .reset_index() + .sort_values("n_fast") + ) + yerr = agg["std"].fillna(0.0) / agg["count"].clip(lower=1).pow(0.5) + ax.errorbar( + agg["n_fast"], + agg["mean"], + yerr=yerr, + marker="o", + linewidth=1.6, + capsize=3, + color=color, + label=sched, + ) + ax.set_xlabel("Heterogeneity (n_fast out of 15)", fontsize=12) + ax.set_ylabel(ylabel, fontsize=12) + ax.set_title(title_prefix, fontsize=13, fontweight="bold") + if ylim is not None: + ax.set_ylim(*ylim) + n_max = int(df["n_fast"].max()) if not df.empty else 15 + ax.set_xlim(-0.5, n_max + 0.5) + ax.grid(True, alpha=0.3) + ax.legend(loc=legend_loc, fontsize=9, frameon=False) + plt.tight_layout() + out_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"Wrote {out_path}") + + +def _plot_cohort_starvation(df: pd.DataFrame, model: str, out_path: pathlib.Path) -> None: + """Two-panel plot: fast-cohort starvation (top) and slow-cohort starvation (bottom).""" + sub_model = df[df["model"] == model] + if sub_model.empty: + return + schedulers = sorted(sub_model["scheduler"].unique()) + color_cycle = plt.cm.tab10(np.linspace(0, 1, max(len(schedulers), 2))) + + fig, (ax_fast, ax_slow) = plt.subplots(2, 1, figsize=(10, 8), sharex=True) + for color, sched in zip(color_cycle, schedulers): + sub = sub_model[sub_model["scheduler"] == sched] + for ax, col in ((ax_fast, "fast_cohort_mean_starvation"), (ax_slow, "slow_cohort_mean_starvation")): + cell = sub.dropna(subset=[col]) + if cell.empty: + continue + agg = ( + cell.groupby("n_fast")[col] + .agg(["mean", "std", "count"]) + .reset_index() + .sort_values("n_fast") + ) + yerr = agg["std"].fillna(0.0) / agg["count"].clip(lower=1).pow(0.5) + ax.errorbar( + agg["n_fast"], + agg["mean"], + yerr=yerr, + marker="o", + linewidth=1.6, + capsize=3, + color=color, + label=sched, + ) + + n_max = int(sub_model["n_fast"].max()) + for ax, title in ((ax_fast, "Fast cohort (horizon=4)"), (ax_slow, "Slow cohort (horizon=10)")): + ax.set_ylabel("Mean starvation rate", fontsize=11) + ax.set_xlim(-0.5, n_max + 0.5) + ax.set_ylim(0, 1) + ax.grid(True, alpha=0.3) + ax.set_title(title, fontsize=11) + ax_slow.set_xlabel("Heterogeneity (n_fast out of 15)", fontsize=12) + ax_fast.legend(loc="upper left", fontsize=9, frameon=False) + fig.suptitle( + f"Trade-off: fast vs slow cohort starvation — {model}", + fontsize=13, fontweight="bold", + ) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + out_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out_path, dpi=150, bbox_inches="tight") + plt.close(fig) + print(f"Wrote {out_path}") + + +def make_plots(df: pd.DataFrame, plots_dir: pathlib.Path) -> None: + if df.empty: + print("No data to plot") + return + plots_dir.mkdir(parents=True, exist_ok=True) + for model in sorted(df["model"].unique()): + model_slug = model.replace(".", "_").replace("/", "_") + sub = df[df["model"] == model] + _plot_metric_vs_nfast( + sub, + "jain_freshness", + ylabel="Jain's index on freshness rate", + title_prefix=f"Fairness vs heterogeneity — {model}", + out_path=plots_dir / f"jains_vs_het__{model_slug}.png", + ylim=(0, 1.02), + legend_loc="lower left", + ) + _plot_metric_vs_nfast( + sub, + "starvation_var", + ylabel="Cross-robot variance of starvation rate", + title_prefix=f"Starvation variance vs heterogeneity — {model}", + out_path=plots_dir / f"starvation_var_vs_het__{model_slug}.png", + legend_loc="upper left", + ) + _plot_metric_vs_nfast( + sub, + "starvation_mean", + ylabel="Mean starvation rate (all 15 robots)", + title_prefix=f"Aggregate starvation vs heterogeneity — {model}", + out_path=plots_dir / f"mean_starvation_vs_het__{model_slug}.png", + ylim=(0, 1), + legend_loc="upper left", + ) + _plot_cohort_starvation( + df, + model, + plots_dir / f"cohort_starvation_vs_het__{model_slug}.png", + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--sweep-dir", + type=pathlib.Path, + required=True, + help="Sweep directory (must contain artifacts/ subdir).", + ) + parser.add_argument( + "--output-csv", + type=pathlib.Path, + default=None, + help="Output CSV path (default: /sweep_metrics_recomputed.csv).", + ) + parser.add_argument( + "--plots-dir", + type=pathlib.Path, + default=None, + help="Output plots dir (default: /plots).", + ) + args = parser.parse_args() + + sweep_dir = args.sweep_dir + if not (sweep_dir / "artifacts").exists(): + raise SystemExit(f"No artifacts/ subdir in {sweep_dir}") + + output_csv = args.output_csv or (sweep_dir / "sweep_metrics_recomputed.csv") + plots_dir = args.plots_dir or (sweep_dir / "plots") + + df = collect_metrics(sweep_dir) + if df.empty: + raise SystemExit("No usable runs found.") + df.to_csv(output_csv, index=False) + print(f"Wrote {output_csv} ({len(df)} rows)") + + make_plots(df, plots_dir) + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_sweep.py b/scripts/exps/plot_starvation_sweep.py similarity index 100% rename from scripts/plot_sweep.py rename to scripts/exps/plot_starvation_sweep.py diff --git a/scripts/modal_serve.py b/scripts/modal_serve.py index f881e7d..c75b2f4 100644 --- a/scripts/modal_serve.py +++ b/scripts/modal_serve.py @@ -14,10 +14,12 @@ GPU = "h100" REGION = "us-east" -ENV_MODE = "REAL_SORT_LEGOS" +ENV_MODE = "LIBERO" MAX_BATCH_SIZE = 4 PORT = 8080 MODEL = "PI05" +SCHEDULING_ALGORITHM = "greedy-deadline" +ALPHA = 1.0 REPO_ROOT = pathlib.Path(__file__).parent.parent @@ -99,9 +101,10 @@ def generate_requirements() -> None: } ) .add_local_python_source( - "armory", "armory_client", "openpi", "openpi_client", "libero", "gr00t" + "armory", "armory_client", "openpi", "openpi_client", "libero", "gr00t", "openpi_adapter", "gr00t_adapter" ) .add_local_dir(str(REPO_ROOT / "scripts"), remote_path="/root/scripts") + .add_local_dir(str(REPO_ROOT / "configs"), remote_path="/root/configs") ) @@ -133,6 +136,10 @@ def startup(self) -> None: str(MAX_BATCH_SIZE), "--port", str(PORT), + "--scheduling-algorithm", + SCHEDULING_ALGORITHM, + "--alpha", + str(ALPHA), ] def _stream_logs(proc: subprocess.Popen) -> None: diff --git a/scripts/serve.py b/scripts/serve.py index cda39e3..25e4721 100644 --- a/scripts/serve.py +++ b/scripts/serve.py @@ -68,7 +68,7 @@ class Args: alpha: float = 1.0 - min_ex: int = 0 + min_execution_horizon: int = 0 lookahead_horizon_ms: int = 500 lookahead_timestep_ms: int = 50 @@ -126,7 +126,7 @@ def main(args: Args) -> None: args, action_horizon_steps=resolved.metadata.action_horizon ) resolved.metadata.scheduler_kwargs = scheduler_kwargs - resolved.metadata.min_ex = args.min_ex + resolved.metadata.min_execution_horizon = args.min_execution_horizon server = PolicyServer( metadata=resolved.metadata, diff --git a/src/armory/scheduling/base.py b/src/armory/scheduling/base.py index f227cb1..b0bbdb8 100644 --- a/src/armory/scheduling/base.py +++ b/src/armory/scheduling/base.py @@ -24,18 +24,18 @@ def __init__( self, batch_queue: mp.Queue, max_batch_size: int = 1, - min_ex: int = 0, + min_execution_horizon: int = 0, ): self._batch_queue = batch_queue self._max_batch_size = max_batch_size # Mirror the engine's _should_serve gate. The engine drops a request - # whose action_index_start is not at least min_ex past what was last + # whose action_index_start is not at least min_execution_horizon past what was last # served; if the scheduler doesn't apply the same gate, it keeps # emitting batches the engine will reject. Each rejected batch returns # an empty ResponseBatch which still pops in_flight, freeing the # GreedyDeadline gate to emit again — a tight loop that buries the # GPU's batch_queue. - self._min_ex = min_ex + self._min_ex = min_execution_horizon self.latency_tracker = EMALatencyTracker() self.mirror = Mirror(self.latency_tracker) @@ -75,7 +75,9 @@ def schedule(self) -> list[SchedulerDecision]: started_at = time.time() next_avail = self.mirror.next_time_server_available() in_flight = self.mirror.in_flight_batches_count - candidates = self.mirror.schedulable_requests(self._latest_requests, min_ex=self._min_ex) + candidates = self.mirror.schedulable_requests( + self._latest_requests, min_execution_horizon=self._min_ex + ) candidate_ids = [r.robot_id for r in candidates] deadlines = self.mirror.deadlines() if self.mirror.robots else {} diff --git a/src/armory/scheduling/baselines.py b/src/armory/scheduling/baselines.py index 9c0ec12..e82ae1b 100644 --- a/src/armory/scheduling/baselines.py +++ b/src/armory/scheduling/baselines.py @@ -130,9 +130,9 @@ def __init__( self, batch_queue: mp.Queue, max_batch_size: int = 1, - min_ex: int = 0, + min_execution_horizon: int = 0, ): - super().__init__(batch_queue, max_batch_size, min_ex=min_ex) + super().__init__(batch_queue, max_batch_size, min_execution_horizon=min_execution_horizon) self._rr_index: int = 0 self._rr_robot_order: list[str] = [] diff --git a/src/armory/scheduling/dynamic_action.py b/src/armory/scheduling/dynamic_action.py index 2c6457d..82f4ccb 100644 --- a/src/armory/scheduling/dynamic_action.py +++ b/src/armory/scheduling/dynamic_action.py @@ -11,11 +11,11 @@ def __init__( self, batch_queue: mp.Queue, max_batch_size: int = 1, - min_ex: int = 0, + min_execution_horizon: int = 0, *, alpha: float = 0.0, ): - super().__init__(batch_queue, max_batch_size, min_ex=min_ex) + super().__init__(batch_queue, max_batch_size, min_execution_horizon=min_execution_horizon) self._alpha = max(0.0, alpha) self._service_debt: dict[str, float] = {} self._demand_rate: dict[str, float] = {} @@ -36,22 +36,29 @@ def reset_robot(self, robot_id: str) -> None: self._service_debt.pop(robot_id, None) self._demand_rate.pop(robot_id, None) - def get_next_batches(self) -> list[list[SlotRequest]]: - # return if there are any batches in the queue or no schedulable requests - if not self._batch_queue.empty() or not (candidates := self.schedulable_requests): - return [] + def get_next_batches( + self, candidates: list[SlotRequest] + ) -> tuple[list[list[SlotRequest]], dict[str, Any]]: + if self.mirror.in_flight_batches_count > 0: + return [], {"reason": "server_busy"} + if not candidates: + return [], {"reason": "no_candidates"} # advance debts to reflect the demand rate now = time.time() self._advance_debts(now) + deadlines = self.mirror.deadlines() # sort the candidates by deadline and robot id - ordered = sorted(candidates, key=lambda r: (self._infer_deadline(r), r.robot_id)) + ordered = sorted( + candidates, + key=lambda r: (self._infer_deadline(r, deadlines), r.robot_id), + ) # get the maximum batch size max_size = min(self._max_batch_size, len(ordered)) # get the best batch by scoring the batches best_batch = max( (tuple(ordered[:k]) for k in range(1, max_size + 1)), - key=lambda b: self._score(b, now), + key=lambda b: self._score(b, now, deadlines), ) # update the service debt for the best batch @@ -60,7 +67,17 @@ def get_next_batches(self) -> list[list[SlotRequest]]: self._service_debt[r.robot_id] = max( 0.0, self._service_debt.get(r.robot_id, 0.0) - 1.0 ) - return [list(best_batch)] + chosen = list(best_batch) + notes = { + "rule": "alpha_fair_dynamic_action", + "alpha": self._alpha, + "max_batch_size": self._max_batch_size, + "chosen_batch_size": len(chosen), + "infer_deadlines": {r.robot_id: self._infer_deadline(r, deadlines) for r in ordered}, + "service_debt": dict(self._service_debt), + "demand_rate": dict(self._demand_rate), + } + return [chosen], notes def _advance_debts(self, now: float) -> None: # return if the time elapsed is less than or equal to 0 @@ -69,22 +86,25 @@ def _advance_debts(self, now: float) -> None: return # advance the debts for each robot based on the demand rate for robot_id, rate in self._demand_rate.items(): - self._service_debt[robot_id] = ( - self._service_debt.get(robot_id, 0.0) + elapsed * rate - ) + self._service_debt[robot_id] = self._service_debt.get(robot_id, 0.0) + elapsed * rate self._last_advance = now - def _infer_deadline(self, request: SlotRequest) -> float: + def _infer_deadline(self, request: SlotRequest, deadlines: dict[str, float]) -> float: # return the deadline for the request (from greedy deadline scheduler) - return self._deadlines.get( + return deadlines.get( request.robot_id, request.deadline ) - self.latency_tracker.action_latency(request.robot_id) - def _score(self, batch: tuple[SlotRequest, ...], now: float) -> tuple: + def _score( + self, + batch: tuple[SlotRequest, ...], + now: float, + deadlines: dict[str, float], + ) -> tuple: # infer latency for the batch infer_latency = max(self.latency_tracker.infer_latency(len(batch)), 1e-6) # get the earliest deadline for the batch - earliest = min(self._infer_deadline(r) for r in batch) + earliest = min(self._infer_deadline(r, deadlines) for r in batch) # check if the batch fits within the earliest deadline fits = int(infer_latency <= earliest - now) # get the base score for the batch diff --git a/src/armory/scheduling/lookahead.py b/src/armory/scheduling/lookahead.py index a8927cf..1e6e51d 100644 --- a/src/armory/scheduling/lookahead.py +++ b/src/armory/scheduling/lookahead.py @@ -15,14 +15,14 @@ def __init__( self, batch_queue: mp.Queue, max_batch_size: int = 1, - min_ex: int = 0, + min_execution_horizon: int = 0, *, horizon_ms: int = 1000, timestep_ms: int = 50, action_horizon_steps: int = 10, control_hz: int = 20, ) -> None: - super().__init__(batch_queue, max_batch_size, min_ex=min_ex) + super().__init__(batch_queue, max_batch_size, min_execution_horizon=min_execution_horizon) assert timestep_ms > 0, "timestep_ms must be positive" assert horizon_ms > 0, "horizon_ms must be positive" assert action_horizon_steps > 0, "action_horizon_steps must be positive" diff --git a/src/armory/scheduling/lookahead_actions.py b/src/armory/scheduling/lookahead_actions.py index cdf187f..91c7bee 100644 --- a/src/armory/scheduling/lookahead_actions.py +++ b/src/armory/scheduling/lookahead_actions.py @@ -169,7 +169,7 @@ def __init__( self, batch_queue: mp.Queue, max_batch_size: int = 1, - min_ex: int = 0, + min_execution_horizon: int = 0, *, horizon: float = 1.0, max_depth: int = 5, @@ -177,7 +177,7 @@ def __init__( step_budget_nodes: int = 32, scheduling_buffer: float = 0.01, ) -> None: - super().__init__(batch_queue, max_batch_size, min_ex=min_ex) + super().__init__(batch_queue, max_batch_size, min_execution_horizon=min_execution_horizon) self.horizon = horizon self.max_depth = max_depth self.max_in_flight = max_in_flight diff --git a/src/armory/scheduling/mirror.py b/src/armory/scheduling/mirror.py index 08fcd33..814cf26 100644 --- a/src/armory/scheduling/mirror.py +++ b/src/armory/scheduling/mirror.py @@ -158,6 +158,18 @@ def deadline(self) -> float: step = self.steps[-1] while step.next_action_step <= self.max_overall_action_step: + # If no chunk covers next_action_step, advance_step can never + # increment it (action_is_available stays False forever) and the + # loop spins indefinitely. Treat the gap as the stall point and + # return the current step time. + # NOTE Rohan: hack from Claude. fix properly + if not any( + chunk.action_index_start + <= step.next_action_step + <= chunk.action_index_start + chunk.execution_horizon - 1 + for chunk in self.chunks + ): + return step.time step = self.advance_step(step) return step.time @@ -340,9 +352,9 @@ def next_time_server_available(self) -> float: def schedulable_requests( self, requests: dict[RobotID, SlotRequest], - min_ex: int = 0, + min_execution_horizon: int = 0, ) -> list[SlotRequest]: - """Filter requests whose next-chunk start is at least ``min_ex`` past + """Filter requests whose next-chunk start is at least ``min_execution_horizon`` past the last queued chunk. Mirrors the engine's _should_serve gate so the scheduler doesn't emit batches the engine will drop. """ @@ -367,11 +379,11 @@ def schedulable_requests( if robot.get_latest_control_step_before(obs_cutoff) is None: continue # - + _, action_index_start = self._next_chunk_context(robot_id, dispatch_time) if ( len(robot.chunks) == 0 - or action_index_start > robot.chunks[-1].action_index_start + min_ex + or action_index_start > robot.chunks[-1].action_index_start + min_execution_horizon ): schedulable_requests.append(request) # else: diff --git a/src/armory/serving/engine.py b/src/armory/serving/engine.py index b8bf66b..b6061de 100644 --- a/src/armory/serving/engine.py +++ b/src/armory/serving/engine.py @@ -57,7 +57,7 @@ def __init__( gpu_out_ep: str, ready_event: Event, log_queue: mp.Queue | None = None, - min_ex: int = 10, + min_execution_horizon: int = 10, ) -> None: self.policy_factory = policy_factory self.max_batch_size = max_batch_size @@ -67,7 +67,7 @@ def __init__( self.gpu_out_ep = gpu_out_ep self.ready_event = ready_event self.log_queue = log_queue - self._min_ex = min_ex + self._min_ex = min_execution_horizon def run(self) -> None: signal.signal(signal.SIGINT, signal.SIG_IGN) @@ -165,7 +165,9 @@ def run(self) -> None: for sd, action_dict, chunk_id in zip(slot_datas, actions, chunk_ids, strict=True) ] - self._update_state(slot_requests, slot_datas, actions) # NOTE from Rohan: this was originally slot_reqs + self._update_state( + slot_requests, slot_datas, actions + ) # NOTE from Rohan: this was originally slot_reqs # Send responses directly to WS — not via scheduler result_sock.send_pyobj( diff --git a/src/armory/serving/scheduler.py b/src/armory/serving/scheduler.py index 76133da..0f4924c 100644 --- a/src/armory/serving/scheduler.py +++ b/src/armory/serving/scheduler.py @@ -61,7 +61,7 @@ def __init__( scheduler_kwargs: dict | None, ready_event: Event, log_queue: mp.Queue | None = None, - min_ex: int = 0, + min_execution_horizon: int = 0, ) -> None: self.sched_in_ep = sched_in_ep self.result_ep = result_ep @@ -72,7 +72,7 @@ def __init__( self.scheduler_kwargs = scheduler_kwargs self.ready_event = ready_event self.log_queue = log_queue - self.min_ex = min_ex + self.min_execution_horizon = min_execution_horizon def run(self) -> None: signal.signal(signal.SIGINT, signal.SIG_IGN) @@ -104,7 +104,7 @@ def run(self) -> None: scheduler = cls( self.batch_queue, max_batch_size=self.max_batch_size, - min_ex=self.min_ex, + min_execution_horizon=self.min_execution_horizon, **extra_kwargs, ) diff --git a/src/armory/serving/server.py b/src/armory/serving/server.py index 97508b0..a07a31f 100644 --- a/src/armory/serving/server.py +++ b/src/armory/serving/server.py @@ -254,7 +254,7 @@ def _start_backend( socket_addresses["gpu_out_ep"], gpu_ready, log_queue, - min_ex=metadata.min_ex, + min_execution_horizon=metadata.min_execution_horizon, ).run, daemon=True, ) @@ -270,7 +270,7 @@ def _start_backend( scheduler_kwargs, sched_ready, log_queue, - min_ex=metadata.min_ex, + min_execution_horizon=metadata.min_execution_horizon, ).run, daemon=True, ) diff --git a/src/sims/libero/metrics.py b/src/sims/libero/metrics.py index 106990d..c45f55d 100644 --- a/src/sims/libero/metrics.py +++ b/src/sims/libero/metrics.py @@ -1137,14 +1137,20 @@ def generate_starvation_tail_metrics_plot(output_path: pathlib.Path) -> None: logger.info(f"Saved {plots_dir / 'starvation_tail_metrics.png'}") -def generate_starvation_variance_plot( +def compute_starvation_variance_series( output_path: pathlib.Path, control_hz: float | None = None -) -> None: - """Plot cumulative starvation rate per robot and its cross-robot variance.""" +) -> dict | None: + """Cross-robot variance of cumulative starvation rate over wall-clock time. + + Starvation here matches the ``actions_left <= 0`` definition used by the + starvation_variance_over_time plot (queue depth at the control step), + aligned onto a shared wall-clock canvas via ``_build_actions_left_matrix``. + + Returns ``None`` when no per-robot actions_left data is available. + """ robots, matrix, _, control_hz, _ = _build_actions_left_matrix(output_path, control_hz) if matrix.size == 0: - logger.warning("No actions_left.npy data found for starvation variance plot") - return + return None valid_mask = ~np.isnan(matrix) starved_mask = valid_mask & (matrix <= 0) @@ -1158,6 +1164,31 @@ def generate_starvation_variance_plot( ) starvation_variance = np.nanvar(cumulative_rates, axis=0) time_seconds = np.arange(matrix.shape[1], dtype=float) / max(control_hz, 1.0) + final_variance = float(starvation_variance[-1]) if starvation_variance.size else float("nan") + return { + "robots": robots, + "cumulative_rates": cumulative_rates, + "cumulative_starved": cumulative_starved, + "cumulative_observed": cumulative_observed, + "starvation_variance": starvation_variance, + "time_seconds": time_seconds, + "control_hz": float(control_hz), + "final_starvation_variance": final_variance, + } + + +def generate_starvation_variance_plot( + output_path: pathlib.Path, control_hz: float | None = None +) -> None: + """Plot cumulative starvation rate per robot and its cross-robot variance.""" + series = compute_starvation_variance_series(output_path, control_hz) + if series is None: + logger.warning("No actions_left.npy data found for starvation variance plot") + return + robots = series["robots"] + cumulative_rates = series["cumulative_rates"] + starvation_variance = series["starvation_variance"] + time_seconds = series["time_seconds"] fig, (ax_rates, ax_var) = plt.subplots( 2, @@ -1223,6 +1254,137 @@ def generate_starvation_variance_plot( logger.info(f"Saved {plots_dir / 'starvation_variance_over_time.png'}") +def generate_per_robot_starvation_rate_gif( + output_path: pathlib.Path, + control_hz: float | None = None, + fps: int = 15, + max_frames: int = 150, +) -> None: + """Animate per-robot cumulative starvation rate as stacked line + bar chart. + + Top panel: per-robot lines revealed over wall-clock time (same data as the + top panel of starvation_variance_over_time.png). + Bottom panel: per-robot bar chart whose heights track each robot's + cumulative starvation rate at the current frame, with an overall + weighted-average reference line that updates per frame. + """ + series = compute_starvation_variance_series(output_path, control_hz) + if series is None: + logger.warning("No actions_left.npy data found for per-robot starvation GIF") + return + import matplotlib.animation as animation # noqa: PLC0415 + + robots = series["robots"] + cumulative_rates = series["cumulative_rates"] + cumulative_starved = series["cumulative_starved"] + cumulative_observed = series["cumulative_observed"] + time_seconds = series["time_seconds"] + n_cols = cumulative_rates.shape[1] + if n_cols < 2: + logger.warning("Not enough timesteps for per-robot starvation GIF") + return + + # Downsample frame indices: include t=0 and the final frame, evenly spaced. + n_frames = min(max_frames, n_cols) + frame_indices = np.unique(np.linspace(0, n_cols - 1, n_frames).astype(int)) + + plot_order = np.argsort([int(robot) for robot in robots]) + colors = plt.cm.tab20(np.linspace(0, 1, max(len(robots), 2))) + ordered_robots = [robots[i] for i in plot_order] + bar_labels = [str(r) for r in ordered_robots] + + fig, (ax_line, ax_bar) = plt.subplots( + 2, 1, figsize=(max(8, 1.2 * len(robots)), 8), gridspec_kw={"height_ratios": [2, 1.5]} + ) + + lines = [] + for color_idx, row_idx in enumerate(plot_order): + robot = robots[row_idx] + (line,) = ax_line.plot( + [], + [], + linewidth=1.5, + color=colors[color_idx % len(colors)], + label=f"robot_{robot}", + ) + lines.append((row_idx, line)) + ax_line.set_xlim(0, time_seconds[-1]) + ax_line.set_ylim(0, 1) + ax_line.set_xlabel("Wall-clock time (s)", fontsize=12) + ax_line.set_ylabel("Cumulative starvation rate", fontsize=12) + ax_line.grid(True, alpha=0.3) + ax_line.legend( + loc="upper right", + ncol=min(max(1, len(robots)), 5), + fontsize=8, + frameon=False, + ) + title = ax_line.set_title("", fontsize=13, fontweight="bold") + time_marker = ax_line.axvline(0.0, color="black", linewidth=1.0, alpha=0.5) + + bar_colors = [colors[i % len(colors)] for i in range(len(plot_order))] + bars = ax_bar.bar( + bar_labels, + np.zeros(len(plot_order)), + color=bar_colors, + edgecolor="black", + alpha=0.85, + ) + bar_texts = [ + ax_bar.text( + bar.get_x() + bar.get_width() / 2.0, + 0.0, + "0.0%", + ha="center", + va="bottom", + fontsize=8, + ) + for bar in bars + ] + overall_line = ax_bar.axhline( + 0.0, color="red", linestyle="--", linewidth=2, label="Overall: 0.0%" + ) + overall_legend = ax_bar.legend(loc="upper right", fontsize=9) + ax_bar.set_xlabel("Robot index", fontsize=12) + ax_bar.set_ylabel("Cumulative starvation rate", fontsize=12) + ax_bar.set_ylim(0, 1.0) + ax_bar.grid(axis="y", alpha=0.3) + + def update(frame_col: int): + upto = frame_col + 1 + for row_idx, line in lines: + line.set_data(time_seconds[:upto], cumulative_rates[row_idx, :upto]) + t_now = time_seconds[frame_col] + time_marker.set_xdata([t_now, t_now]) + title.set_text(f"Per-robot cumulative starvation rate — t = {t_now:.1f}s") + + col_rates = cumulative_rates[:, frame_col] + ordered_rates = col_rates[plot_order] + for bar, text, rate in zip(bars, bar_texts, ordered_rates): + height = float(rate) if np.isfinite(rate) else 0.0 + bar.set_height(height) + text.set_y(height + 0.01) + text.set_text("--" if not np.isfinite(rate) else f"{height * 100:.1f}%") + + total_starved = float(np.nansum(cumulative_starved[:, frame_col])) + total_observed = float(np.nansum(cumulative_observed[:, frame_col])) + overall = total_starved / total_observed if total_observed > 0 else 0.0 + overall_line.set_ydata([overall, overall]) + overall_legend.get_texts()[0].set_text(f"Overall: {overall * 100:.1f}%") + return [ln for _, ln in lines] + list(bars) + bar_texts + [overall_line, title, time_marker] + + anim = animation.FuncAnimation( + fig, update, frames=frame_indices.tolist(), interval=1000 / max(fps, 1), blit=False + ) + fig.tight_layout() + plots_dir = output_path / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + out = plots_dir / "per_robot_starvation_rate.gif" + anim.save(out, writer=animation.PillowWriter(fps=fps)) + plt.close(fig) + logger.info(f"Saved {out}") + + def generate_jains_starvation_over_time_plot( output_path: pathlib.Path, control_hz: float | None = None ) -> None: @@ -1891,6 +2053,7 @@ def generate_all_plots(output_path: pathlib.Path) -> None: generate_starvation_plot, generate_starvation_tail_metrics_plot, generate_starvation_variance_plot, + generate_per_robot_starvation_rate_gif, # slow (~5-10s per run); run manually if needed generate_jains_starvation_over_time_plot, generate_staleness_plot, generate_batch_size_plot, diff --git a/src/sims/libero/subscribers/saver.py b/src/sims/libero/subscribers/saver.py index 889f2e3..c5cc98b 100644 --- a/src/sims/libero/subscribers/saver.py +++ b/src/sims/libero/subscribers/saver.py @@ -3,26 +3,27 @@ import logging import pathlib import time - -import imageio -import matplotlib - -matplotlib.use("Agg") -import dataclasses from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass from typing import TYPE_CHECKING -import matplotlib.pyplot as plt +import dataclasses +import imageio import numpy as np from typing_extensions import override from armory_client.action_chunkers.action_chunk_broker import ActionChunkBroker from armory_client.runtime import subscriber as _subscriber +from armory_client.runtime.saver_utils import ( + EpisodeSaveData, + Result, + plot_cost_history, + save_action_chunks, + save_actions_left, + save_cost_history_npy, + save_timestamps, +) from armory_client.schemas import ( Action, - ActionChunk, - JSONDataclass, Observation, Timestamp, ) @@ -35,31 +36,6 @@ logger = logging.getLogger(__name__) -@dataclass(frozen=True) -class Result(JSONDataclass): - robot_idx: int - success: bool - steps_taken: int - task_suite_name: str - task_id: int - task_language: str - episode_idx: int - - -@dataclass -class _EpisodeSaveData: - """Snapshot of all data needed to persist one episode, safe to hand off to a thread.""" - - timestamps: list[Timestamp] - observations_buffer: dict[int, Observation] - action_chunks: list[ActionChunk] - actions_left_snapshot: list[int] - cost_history: list[float] - current_success: bool - episode_idx: int - initial_state: np.ndarray | None - - class Saver(_subscriber.Subscriber): """Saves episode data by offloading I/O to a background thread pool.""" @@ -127,14 +103,14 @@ def on_step(self, observation: Observation, action: Action) -> None: @override def on_episode_end(self) -> None: - data = _EpisodeSaveData( + data = EpisodeSaveData( timestamps=self._timestamps, observations_buffer=self._observations_buffer, # Shallow-copy the broker list in case it gets reset between episodes. action_chunks=list(self._action_chunk_broker.action_chunks), actions_left_snapshot=self._actions_left_snapshot, cost_history=self._cost_history, - current_success=self._environment.current_success, + success=self._environment.current_success, episode_idx=self._environment.episode_idx, initial_state=self._environment.current_initial_state, ) @@ -144,35 +120,39 @@ def on_episode_end(self) -> None: def close(self) -> None: self._executor.shutdown(wait=True) - def _save_all(self, data: _EpisodeSaveData) -> None: + def _save_all(self, data: EpisodeSaveData) -> None: out_folder, dir_episode_idx = self._get_out_folder(data) data = dataclasses.replace(data, episode_idx=dir_episode_idx) self._save_metadata(out_folder, data) - self._save_timestamps(out_folder, data) - self._save_action_chunks(out_folder, data) + logger.info(f"Saving timestamps to {out_folder / 'timestamps.csv'}") + save_timestamps(data.timestamps, out_folder) + logger.info(f"Saving action chunks to {out_folder}") + save_action_chunks(data.action_chunks, out_folder) if self._save_video_enabled: self._save_video(out_folder, data) self._save_debug_data(out_folder, data) - self._save_actions_left(out_folder, data) + path = out_folder / "actions_left.npy" + save_actions_left(data.actions_left_snapshot, out_folder) + logger.info(f"Saved actions_left to {path}") self._save_cost_history(out_folder, data) - def _get_out_folder(self, data: _EpisodeSaveData) -> tuple[pathlib.Path, int]: + def _get_out_folder(self, data: EpisodeSaveData) -> tuple[pathlib.Path, int]: robot_folder = self._out_dir / str(self._robot_idx) pathlib.Path(robot_folder).mkdir(parents=True, exist_ok=True) existing = list(robot_folder.iterdir()) next_idx = max([int(p.name.split("_")[0]) for p in existing if p.is_dir()], default=-1) + 1 - success_str = "success" if data.current_success else "failure" + success_str = "success" if data.success else "failure" out_folder = ( robot_folder / f"{next_idx}_{self._task_suite_name}_{self._task_id}_{success_str}" ) pathlib.Path(out_folder).mkdir(parents=True, exist_ok=True) return pathlib.Path(out_folder), next_idx - def _save_metadata(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: + def _save_metadata(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: logger.info(f"Saving metadata to {out_folder / 'metadata.json'}") result = Result( - success=data.current_success, + success=data.success, robot_idx=self._robot_idx, steps_taken=len(data.timestamps), task_suite_name=self._task_suite_name, @@ -182,15 +162,7 @@ def _save_metadata(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> No ) result.to_json(out_folder / "metadata.json") - def _save_timestamps(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - logger.info(f"Saving timestamps to {out_folder / 'timestamps.csv'}") - Timestamp.to_csv(data.timestamps, out_folder / "timestamps.csv") - - def _save_action_chunks(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - logger.info(f"Saving action chunks to {out_folder}") - ActionChunk.to_parquet(data.action_chunks, out_folder / "action_chunks.parquet") - - def _save_video(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: + def _save_video(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: logger.info(f"Saving video to {out_folder / 'out.mp4'}") images = [obs.image for obs in data.observations_buffer.values()] imageio.mimwrite( @@ -199,7 +171,7 @@ def _save_video(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: fps=self._control_hz, # NOTE: saving in control hz fps for now ) - def _save_debug_data(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: + def _save_debug_data(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: """Save debug data as a single .npz file with observations, noise, and actions.""" # Check if we have noise data has_noise = any(chunk.noise is not None for chunk in data.action_chunks) @@ -249,29 +221,17 @@ def _save_debug_data(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> np.savez_compressed(debug_data_file, **data_to_save) logger.info(f"Saved {len(data.action_chunks)} chunks to {debug_data_file}") - def _save_actions_left(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - path = out_folder / "actions_left.npy" - np.save(path, np.array(data.actions_left_snapshot, dtype=np.int32)) - logger.info(f"Saved actions_left to {path}") - - def _save_cost_history(self, out_folder: pathlib.Path, data: _EpisodeSaveData) -> None: - costs = np.array(data.cost_history, dtype=np.float64) + def _save_cost_history(self, out_folder: pathlib.Path, data: EpisodeSaveData) -> None: npy_path = out_folder / "cost_history.npy" - np.save(npy_path, costs) + costs = save_cost_history_npy(data.cost_history, out_folder) logger.info(f"Saved cost_history to {npy_path}") plot_path = out_folder / "cost_history.png" - steps = np.arange(len(costs)) - fig, ax = plt.subplots(figsize=(10, 4)) - ax.plot(steps, costs, linewidth=0.8, color="steelblue") - ax.set_xlabel("Environment step") - ax.set_ylabel("Cost (s)") - ax.set_title( - f"Cost per step — robot {self._robot_idx} | " - f"{self._task_suite_name} task {self._task_id}" + plot_cost_history( + costs, + out_folder, + robot_idx=self._robot_idx, + task_suite_name=self._task_suite_name, + task_id=self._task_id, ) - ax.grid(True, alpha=0.3) - fig.tight_layout() - fig.savefig(plot_path, dpi=150) - plt.close(fig) logger.info(f"Saved cost_history plot to {plot_path}")