From 2d70559eec8b215e3f8c929eb73042ff0e0744dc Mon Sep 17 00:00:00 2001 From: Rohan Bansal Date: Mon, 11 May 2026 19:39:00 -0400 Subject: [PATCH 1/6] add sweep scripts --- configs/inference_profiles.json | 17 +- requirements-modal.txt | 6 +- scripts/exps/calculate_starvation_offline.py | 26 + scripts/exps/modal_alpha_fairness_sweep.py | 618 +++++++++++++++++++ scripts/exps/modal_fairness_sweep.py | 438 +++++++++++++ scripts/exps/modal_starvation_sweep.py | 425 +++++++++++++ scripts/exps/plot_fairness_sweep.py | 300 +++++++++ scripts/exps/plot_starvation_sweep.py | 168 +++++ scripts/modal_serve.py | 11 +- src/armory/scheduling/dynamic_action.py | 44 +- src/armory/scheduling/mirror.py | 12 + src/sims/libero/metrics.py | 173 +++++- 12 files changed, 2218 insertions(+), 20 deletions(-) create mode 100644 scripts/exps/calculate_starvation_offline.py create mode 100644 scripts/exps/modal_alpha_fairness_sweep.py create mode 100644 scripts/exps/modal_fairness_sweep.py create mode 100644 scripts/exps/modal_starvation_sweep.py create mode 100644 scripts/exps/plot_fairness_sweep.py create mode 100644 scripts/exps/plot_starvation_sweep.py 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/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/calculate_starvation_offline.py b/scripts/exps/calculate_starvation_offline.py new file mode 100644 index 0000000..ad11fab --- /dev/null +++ b/scripts/exps/calculate_starvation_offline.py @@ -0,0 +1,26 @@ + +from pathlib import Path +import numpy as np + +episode = Path("/coc/flash7/rbansal66/vvla/0_real_0_failure") +control_hz = 20 # change if this run used a different hz + +costs = np.load(episode / "cost_history.npy") +starved = np.isnan(costs) + +first_non_starved = np.flatnonzero(~starved) +if first_non_starved.size: + first = int(first_non_starved[0]) + post_first_starved = int(starved[first:].sum()) + post_first_observed = int(len(costs) - first) +else: + post_first_starved = 0 + post_first_observed = 0 + +print("observed_steps:", len(costs)) +print("starvation_steps:", int(starved.sum())) +print("starvation_rate:", float(starved.mean())) +print("starvation_seconds:", float(starved.sum() / control_hz)) +print("post_first_starvation_steps:", post_first_starved) +print("post_first_observed_steps:", post_first_observed) +print("post_first_starvation_rate:", post_first_starved / post_first_observed if post_first_observed else 0.0) diff --git a/scripts/exps/modal_alpha_fairness_sweep.py b/scripts/exps/modal_alpha_fairness_sweep.py new file mode 100644 index 0000000..e03b49c --- /dev/null +++ b/scripts/exps/modal_alpha_fairness_sweep.py @@ -0,0 +1,618 @@ +"""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 = 4 +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") +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"}, +} +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)) + + # 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_starvation_vs_fairness(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"] + needed = {"starvation_variance", "mean_starvation", "scheduler", "model", "scenario_id"} + if not needed.issubset(df.columns): + return + df = df.dropna(subset=["starvation_variance", "mean_starvation"]) + + plots_dir.mkdir(parents=True, exist_ok=True) + + for (scenario_id, model), sub in df.groupby(["scenario_id", "model"]): + fig, ax = plt.subplots(figsize=(8, 6)) + + # Baselines: aggregate across seeds → one point per scheduler + for sched in BASELINE_SCHEDULERS: + sub_b = sub[sub["scheduler"] == sched] + if sub_b.empty: + continue + x = float(sub_b["mean_starvation"].mean()) + y = float(sub_b["starvation_variance"].mean()) + xerr = float(sub_b["mean_starvation"].std(ddof=0)) if len(sub_b) > 1 else 0.0 + yerr = float(sub_b["starvation_variance"].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, + ) + + # Dynamic-action: one curve, points sorted by alpha, colored on a gradient + sub_d = sub[sub["scheduler"] == DYNAMIC_SCHEDULER] + if not sub_d.empty: + agg = ( + sub_d.groupby("alpha_requested") + .agg( + mean_starvation=("mean_starvation", "mean"), + starvation_variance=("starvation_variance", "mean"), + starvation_std=("mean_starvation", "std"), + variance_std=("starvation_variance", "std"), + count=("seed", "count"), + ) + .reset_index() + .sort_values("alpha_requested") + ) + xs = agg["mean_starvation"].to_numpy() + ys = agg["starvation_variance"].to_numpy() + alphas = agg["alpha_requested"].to_numpy() + ax.plot(xs, ys, "-", color="0.5", linewidth=1.2, alpha=0.7, zorder=2) + sc = 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, + ) + cbar = fig.colorbar(sc, ax=ax, pad=0.02) + cbar.set_label("alpha", fontsize=10) + 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["variance_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=12) + ax.set_ylabel("Cross-robot starvation variance (lower is fairer)", fontsize=12) + ax.set_title( + f"Starvation vs fairness — scenario={scenario_id}, model={model}", + fontsize=13, fontweight="bold", + ) + ax.set_ylim(bottom=0.0) + ax.set_xlim(left=0.0) + ax.grid(True, alpha=0.3) + ax.legend(loc="upper right", fontsize=9, frameon=False) + plt.tight_layout() + + 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 = "gr00t-n1.7", + # scenarios: str = "1f9s,5f5s", + scenarios: str = "5f20s", + alpha_grid: str = "0.0,0.1,0.2,0.3,0.5,0.7,1.0", + seeds: str = "42", + output_dir: str = "experiments/sweeps/fairness_alpha_gr00t", + 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) + var = result.get("starvation_variance", "") + starv = result.get("mean_starvation", "") + print( + f"{result['status']}: {result['run_id']} " + f"starvation_variance={var if var == '' else f'{float(var):.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/exps/modal_starvation_sweep.py b/scripts/exps/modal_starvation_sweep.py new file mode 100644 index 0000000..0d8e346 --- /dev/null +++ b/scripts/exps/modal_starvation_sweep.py @@ -0,0 +1,425 @@ +"""Run server/client scheduler sweeps on Modal. + +Example: + modal run scripts/modal_sweep.py \ + --schedulers fixed-max-batch,greedy-deadline,round-robin \ + --experiment-configs configs/experiments/mock/short.json \ + --num-robots 1,2,3,4,5,6,7,8,9,10 \ + --server-config configs/server/mock.json \ + --seeds 7,42 \ + --output-dir experiments/sweeps/big_mock +""" + +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-scheduler-sweep" +REMOTE_ROOT = pathlib.Path("/app") +REMOTE_OUTPUT_ROOT = pathlib.Path("/tmp/armory_sweep") +PYTHONPATH = ":".join( + [ + str(REMOTE_ROOT / "src"), + str(REMOTE_ROOT / "src/backends"), + str(REMOTE_ROOT / "packages/armory-client/src"), + ] +) + + +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: + scheduler: str + experiment_config: str # path relative to repo root + num_robots: int + seed: int + + @property + def run_id(self) -> str: + config_name = pathlib.Path(self.experiment_config).stem + return f"scheduler={self.scheduler}__config={config_name}__robots={self.num_robots}__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 _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 _safe_float(value: Any, default: float = 0.0) -> float: + try: + if value is None: + return default + return float(value) + except (TypeError, ValueError): + return default + + +def _build_server_cmd(srv_cfg: dict[str, Any], *, port: int, scheduler: str) -> list[str]: + cmd = [ + sys.executable, + "scripts/serve.py", + "--port", + str(port), + "--env", + srv_cfg.get("env", "LIBERO"), + "--max-batch-size", + str(srv_cfg.get("max_batch_size", 1)), + "--scheduling-algorithm", + scheduler, + f"policy:{srv_cfg.get('policy_type', 'default')}", + ] + for k, v in srv_cfg.get("policy", {}).items(): + cmd += [f"--policy.{k.replace('_', '-')}", str(v)] + return cmd + + +def _expand_experiment_config(exp_cfg: dict[str, Any], num_robots: int) -> dict[str, Any]: + """Return a copy of exp_cfg with robot_0's profile replicated to num_robots robots.""" + robot_template = exp_cfg["robots"]["robot_0"] + expanded = { + **exp_cfg, + "experiment": {**exp_cfg["experiment"], "num_robots": num_robots}, + "robots": {f"robot_{i}": dict(robot_template) for i in range(num_robots)}, + } + return expanded + + +def _build_client_cmd( + *, + port: int, + seed: int, + output_dir: pathlib.Path, + experiment_config_path: pathlib.Path, + max_steps: int, +) -> list[str]: + return [ + sys.executable, + "scripts/run_libero.py", + "--host", + "127.0.0.1", + "--port", + str(port), + "--env", + "mock", + "--overwrite", + "--progress-type", + "logging", + "--max-steps", + str(max_steps), + "--seed", + str(seed), + "--output-dir", + str(output_dir), + "--experiment-config", + str(experiment_config_path), + ] + + +def _summarize_run(output_dir: pathlib.Path, case: SweepCase) -> dict[str, Any]: + summary_path = output_dir / "summary.csv" + results_path = output_dir / "results.csv" + runtime_path = output_dir / "runtime_metadata.json" + server_path = output_dir / "server_metadata.json" + + total_success = 0.0 + overall_starvation_rate = 0.0 + post_first_starvation_rate = 0.0 + if summary_path.exists(): + with summary_path.open() as f: + rows = list(csv.DictReader(f)) + if rows: + total_success = sum(_safe_float(r.get("success")) for r in rows) / len(rows) + starvation_steps = sum(_safe_float(r.get("starvation_steps")) for r in rows) + observed_steps = sum(_safe_float(r.get("observed_steps")) for r in rows) + post_first_starvation_steps = sum( + _safe_float(r.get("post_first_starvation_steps")) for r in rows + ) + post_first_observed_steps = sum( + _safe_float(r.get("post_first_observed_steps")) for r in rows + ) + overall_starvation_rate = starvation_steps / observed_steps if observed_steps else 0.0 + post_first_starvation_rate = ( + post_first_starvation_steps / post_first_observed_steps + if post_first_observed_steps + else 0.0 + ) + + robot_rates: list[float] = [] + if results_path.exists(): + by_robot: dict[str, dict[str, float]] = {} + with results_path.open() as f: + for row in csv.DictReader(f): + robot = str(row.get("robot_idx", "unknown")) + stats = by_robot.setdefault(robot, {"starvation_steps": 0.0, "observed_steps": 0.0}) + stats["starvation_steps"] += _safe_float(row.get("starvation_steps")) + stats["observed_steps"] += _safe_float(row.get("observed_steps")) + robot_rates = [ + stats["starvation_steps"] / stats["observed_steps"] + for stats in by_robot.values() + if stats["observed_steps"] > 0 + ] + + runtime = json.loads(runtime_path.read_text()) if runtime_path.exists() else {} + server = json.loads(server_path.read_text()) if server_path.exists() else {} + sorted_rates = sorted(robot_rates) + tail_count = max(1, int(len(sorted_rates) * 0.1)) if sorted_rates else 0 + + summary = { + "run_id": case.run_id, + "scheduler": case.scheduler, + "experiment_config": case.experiment_config, + "num_robots": case.num_robots, + "seed": case.seed, + "success_rate": total_success, + "starvation_rate": overall_starvation_rate, + "post_first_starvation_rate": post_first_starvation_rate, + "robot_starvation_rate_max": max(robot_rates) if robot_rates else 0.0, + "robot_starvation_rate_std": _safe_float(__import__("statistics").pstdev(robot_rates)) + if len(robot_rates) > 1 + else 0.0, + "robot_starvation_rate_cvar90": sum(sorted_rates[-tail_count:]) / tail_count + if tail_count + else 0.0, + "max_batch_size": server.get("max_batch_size", ""), + "action_horizon": server.get("action_horizon", ""), + "max_steps": runtime.get("max_steps", ""), + "num_trials_per_task": runtime.get("num_trials_per_task", ""), + } + + try: + from sims.libero.metrics import compute_server_timing_health # noqa: PLC0415 + + health = compute_server_timing_health(output_dir) + if health is not None: + summary.update(health) + except Exception: # noqa: BLE001 + pass + + return summary + + +@app.function(image=image, timeout=60 * 60, cpu=4, memory=2048) +def run_case( + case: SweepCase, + *, + server_config: str, + port: int, + max_batch_size_override: int | None = None, + max_steps_override: int | None = None, +) -> dict[str, Any]: + exp_cfg: dict[str, Any] = json.loads((REMOTE_ROOT / case.experiment_config).read_text()) + srv_cfg: dict[str, Any] = json.loads((REMOTE_ROOT / server_config).read_text()) + + if max_batch_size_override is not None: + srv_cfg["max_batch_size"] = max_batch_size_override + if max_steps_override is not None: + exp_cfg["experiment"]["max_steps"] = max_steps_override + + exp_cfg = _expand_experiment_config(exp_cfg, case.num_robots) + max_steps = int(exp_cfg["experiment"]["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)) + (run_dir / "server_config.json").write_text(json.dumps(srv_cfg, indent=2)) + + server_cmd = _build_server_cmd(srv_cfg, port=port, scheduler=case.scheduler) + client_cmd = _build_client_cmd( + port=port, + seed=case.seed, + output_dir=output_dir, + experiment_config_path=saved_exp_config, + max_steps=max_steps, + ) + + 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 * 45) + summary = _summarize_run(output_dir, case) + summary["status"] = "ok" + except Exception as exc: # noqa: BLE001 + summary = { + "run_id": case.run_id, + "scheduler": case.scheduler, + "experiment_config": case.experiment_config, + "num_robots": case.num_robots, + "seed": case.seed, + "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}) + + +@app.local_entrypoint() +def main( + schedulers: str = "fixed-max-batch,greedy-deadline,round-robin", + experiment_configs: str = "configs/experiments/mock/short.json", + num_robots: str = "2,4,6", + server_config: str = "configs/server/mock.json", + seeds: str = "7", + output_dir: str = "experiments/sweeps/mock", + port: int = 8080, + max_batch_size: int | None = None, + max_steps: int | None = None, +) -> None: + """Run the Cartesian product of schedulers, experiment_configs, num_robots, and seeds.""" + 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(scheduler=scheduler, experiment_config=cfg, num_robots=n, seed=seed) + for scheduler in _parse_csv(schedulers) + for cfg in _parse_csv(experiment_configs) + for n in _parse_csv(num_robots, cast=int) + for seed in _parse_csv(seeds, cast=int) + ] + + rows: list[dict[str, Any]] = [] + for result in run_case.map( + cases, + kwargs={ + "server_config": server_config, + "port": port, + "max_batch_size_override": max_batch_size, + "max_steps_override": max_steps, + }, + order_outputs=False, + ): + artifact_bytes = result.pop("artifact_tgz") + 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) + print( + f"{result['status']}: {result['run_id']} " + f"starvation={_safe_float(result.get('starvation_rate')):.3f} " + ) + + 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}") + + suspicious = [r for r in rows if r.get("timing_suspicious")] + if suspicious: + print(f"WARNING: {len(suspicious)} run(s) flagged for suspicious timings:") + for r in suspicious: + print(f" {r['run_id']}: {r.get('timing_flags', '')}") + + sys.path.insert(0, str(pathlib.Path(__file__).parent)) + from plot_sweep import DEFAULT_METRICS, plot_results # noqa: PLC0415 + + timing_metrics = [ + "step_interval_p95_ms", + "inference_p99_ms", + "inbound_p95_ms", + "outbound_p95_ms", + ] + plot_results(latest_csv, out / "plots", metrics=list(DEFAULT_METRICS) + timing_metrics) 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/exps/plot_starvation_sweep.py b/scripts/exps/plot_starvation_sweep.py new file mode 100644 index 0000000..d20e9bf --- /dev/null +++ b/scripts/exps/plot_starvation_sweep.py @@ -0,0 +1,168 @@ +"""Plot scheduler sweep metrics from scripts/modal_sweep.py output. + +Example: + uv run python scripts/plot_sweep.py \ + --results experiments/sweeps/mock/sweep_results.csv \ + --output-dir experiments/sweeps/mock/plots +""" + +from __future__ import annotations + +import argparse +import pathlib + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import pandas as pd + +DEFAULT_METRICS = [ + "starvation_rate", + "post_first_starvation_rate", + "robot_starvation_rate_max", + "success_rate", +] + + +METRIC_LABELS = { + "starvation_rate": "Starvation rate", + "post_first_starvation_rate": "Starvation rate excl. startup", + "robot_starvation_rate_max": "Worst robot starvation rate", + "robot_starvation_rate_std": "Robot starvation std. dev.", + "robot_starvation_rate_cvar90": "Tail robot starvation rate", + "success_rate": "Success rate", + "step_interval_p95_ms": "Step interval p95 (ms)", + "inference_p99_ms": "Inference latency p99 (ms)", + "inbound_p95_ms": "Client→server transport p95 (ms)", + "outbound_p95_ms": "Server→client transport p95 (ms)", +} + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--results", type=pathlib.Path, required=True) + parser.add_argument("--output-dir", type=pathlib.Path, default=None) + parser.add_argument("--x", default="num_robots") + parser.add_argument("--line", default="scheduler") + parser.add_argument("--metrics", default=",".join(DEFAULT_METRICS)) + return parser.parse_args() + + +def _metric_label(metric: str) -> str: + return METRIC_LABELS.get(metric, metric.replace("_", " ").title()) + + +def _wilson_ci(p: float, n: int, z: float = 1.96) -> tuple[float, float]: + """95% Wilson score interval for a proportion p estimated from n observations.""" + if n == 0: + return (p, p) + z2 = z * z + denom = 1 + z2 / n + center = (p + z2 / (2 * n)) / denom + half = z * (p * (1 - p) / n + z2 / (4 * n * n)) ** 0.5 / denom + return (max(0.0, center - half), min(1.0, center + half)) + + +def _plot_metric( + df: pd.DataFrame, + *, + metric: str, + x_col: str, + line_col: str, + output_dir: pathlib.Path, +) -> pathlib.Path: + agg = df.groupby([line_col, x_col])[metric].agg(["mean", "count"]).reset_index() + agg.columns = [line_col, x_col, "value", "n"] + agg = agg.sort_values([line_col, x_col]) + + is_proportion = agg["value"].between(0.0, 1.0).all() + if is_proportion: + ci = agg.apply( + lambda r: pd.Series(_wilson_ci(r["value"], int(r["n"])), index=["lo", "hi"]), axis=1 + ) + agg = pd.concat([agg, ci], axis=1) + + fig, ax = plt.subplots(figsize=(8, 4.8)) + for line_value, group in agg.groupby(line_col): + xs = group[x_col].to_numpy() + ys = group["value"].to_numpy() + (line,) = ax.plot(xs, ys, marker="o", linewidth=2.0, label=str(line_value)) + if is_proportion and (group["n"] > 1).any(): + ax.fill_between( + xs, + group["lo"].to_numpy(), + group["hi"].to_numpy(), + alpha=0.15, + color=line.get_color(), + ) + + ax.set_xlabel(x_col.replace("_", " ").title()) + ax.set_ylabel(_metric_label(metric)) + ax.set_title(_metric_label(metric)) + ax.grid(True, axis="y", alpha=0.25) + ax.legend(title=line_col.replace("_", " ").title()) + fig.tight_layout() + + output_path = output_dir / f"{metric}_by_{x_col}.png" + fig.savefig(output_path, dpi=160) + plt.close(fig) + return output_path + + +def plot_results( + results: pathlib.Path, + output_dir: pathlib.Path | None = None, + *, + x: str = "num_robots", + line: str = "scheduler", + metrics: list[str] | None = None, +) -> None: + output_dir = output_dir or (results.parent / "plots") + output_dir.mkdir(parents=True, exist_ok=True) + if metrics is None: + metrics = list(DEFAULT_METRICS) + + df = pd.read_csv(results) + if "status" in df.columns: + df = df[df["status"] == "ok"].copy() + if df.empty: + raise SystemExit("No successful rows found in results CSV") + + missing = [m for m in metrics if m not in df.columns] + if missing: + raise SystemExit(f"Missing metric column(s): {', '.join(missing)}") + for column in [x, line, *metrics]: + if column not in df.columns: + raise SystemExit(f"Missing required column: {column}") + + for metric in metrics: + df[metric] = pd.to_numeric(df[metric], errors="coerce") + numeric_x = pd.to_numeric(df[x], errors="coerce") + if numeric_x.notna().all(): + df[x] = numeric_x + + written = [ + _plot_metric(df, metric=metric, x_col=x, line_col=line, output_dir=output_dir) + for metric in metrics + ] + print("Wrote plots:") + for path in written: + print(path) + + +def main() -> None: + args = _parse_args() + metrics = [m.strip() for m in args.metrics.split(",") if m.strip()] + plot_results( + args.results, + args.output_dir, + x=args.x, + line=args.line, + metrics=metrics, + ) + + +if __name__ == "__main__": + main() 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/src/armory/scheduling/dynamic_action.py b/src/armory/scheduling/dynamic_action.py index 2c6457d..47bac64 100644 --- a/src/armory/scheduling/dynamic_action.py +++ b/src/armory/scheduling/dynamic_action.py @@ -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 @@ -74,17 +91,22 @@ def _advance_debts(self, now: float) -> None: ) 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/mirror.py b/src/armory/scheduling/mirror.py index 08fcd33..193a64b 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 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, From ccc137e04b10357869d9ae8a0d2f1a10de361358 Mon Sep 17 00:00:00 2001 From: Rohan Bansal Date: Mon, 11 May 2026 22:28:35 -0400 Subject: [PATCH 2/6] add fairness plot changes --- scripts/exps/modal_alpha_fairness_sweep.py | 181 +++++++++++++-------- 1 file changed, 114 insertions(+), 67 deletions(-) diff --git a/scripts/exps/modal_alpha_fairness_sweep.py b/scripts/exps/modal_alpha_fairness_sweep.py index e03b49c..9eeeba9 100644 --- a/scripts/exps/modal_alpha_fairness_sweep.py +++ b/scripts/exps/modal_alpha_fairness_sweep.py @@ -54,6 +54,7 @@ DEFAULT_MAX_STEPS = 200 DEFAULT_TRIALS_PER_ROBOT = 1 + BASELINE_SCHEDULERS = ("fixed-max-batch", "greedy-deadline", "round-robin") DYNAMIC_SCHEDULER = "dynamic-action" @@ -329,6 +330,9 @@ def _summarize_run(output_dir: pathlib.Path, case: SweepCase, horizons: list[int 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 @@ -435,92 +439,135 @@ def _write_rows(path: pathlib.Path, rows: list[dict[str, Any]]) -> None: 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 numpy as np import pandas as pd df = pd.read_csv(results_csv) if df.empty: return df = df[df.get("status", "ok") == "ok"] - needed = {"starvation_variance", "mean_starvation", "scheduler", "model", "scenario_id"} + needed = { + "mean_starvation", "starvation_variance", "min_freshness", + "scheduler", "model", "scenario_id", + } if not needed.issubset(df.columns): return - df = df.dropna(subset=["starvation_variance", "mean_starvation"]) + 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 = plt.subplots(figsize=(8, 6)) - - # Baselines: aggregate across seeds → one point per scheduler - for sched in BASELINE_SCHEDULERS: - sub_b = sub[sub["scheduler"] == sched] - if sub_b.empty: - continue - x = float(sub_b["mean_starvation"].mean()) - y = float(sub_b["starvation_variance"].mean()) - xerr = float(sub_b["mean_starvation"].std(ddof=0)) if len(sub_b) > 1 else 0.0 - yerr = float(sub_b["starvation_variance"].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, - ) + fig, (ax_var, ax_min) = plt.subplots(1, 2, figsize=(13, 5.5)) - # Dynamic-action: one curve, points sorted by alpha, colored on a gradient - sub_d = sub[sub["scheduler"] == DYNAMIC_SCHEDULER] - if not sub_d.empty: - agg = ( - sub_d.groupby("alpha_requested") - .agg( - mean_starvation=("mean_starvation", "mean"), - starvation_variance=("starvation_variance", "mean"), - starvation_std=("mean_starvation", "std"), - variance_std=("starvation_variance", "std"), - count=("seed", "count"), - ) - .reset_index() - .sort_values("alpha_requested") - ) - xs = agg["mean_starvation"].to_numpy() - ys = agg["starvation_variance"].to_numpy() - alphas = agg["alpha_requested"].to_numpy() - ax.plot(xs, ys, "-", color="0.5", linewidth=1.2, alpha=0.7, zorder=2) - sc = 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, - ) - cbar = fig.colorbar(sc, ax=ax, pad=0.02) - cbar.set_label("alpha", fontsize=10) - 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["variance_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=12) - ax.set_ylabel("Cross-robot starvation variance (lower is fairer)", fontsize=12) - ax.set_title( + _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 (α=∞ welfare)", + ) + _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", ) - ax.set_ylim(bottom=0.0) - ax.set_xlim(left=0.0) - ax.grid(True, alpha=0.3) - ax.legend(loc="upper right", fontsize=9, frameon=False) - plt.tight_layout() - 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") @@ -600,11 +647,11 @@ def main( tar.extractall(run_dir) result["artifact_path"] = str(run_dir) rows.append(result) - var = result.get("starvation_variance", "") + mf = result.get("min_freshness", "") starv = result.get("mean_starvation", "") print( f"{result['status']}: {result['run_id']} " - f"starvation_variance={var if var == '' else f'{float(var):.4f}'} " + f"min_freshness={mf if mf == '' else f'{float(mf):.4f}'} " f"mean_starvation={starv if starv == '' else f'{float(starv):.4f}'}" ) From 0bb4e177a90c267c5053cc8b0bf233803d00a4dd Mon Sep 17 00:00:00 2001 From: Rohan Bansal Date: Mon, 11 May 2026 22:52:36 -0400 Subject: [PATCH 3/6] generate plots --- .../edge_cases/1_fast_9_slow.jsonc | 89 ---- .../experiments/edge_cases/1_fast_only.jsonc | 26 - .../edge_cases/4_fast_1_slow.jsonc | 54 -- .../edge_cases/5_fast_5_slow.jsonc | 89 ---- configs/heterogeneous_1fast_11slow.yaml | 20 - configs/heterogeneous_4fast_8slow.yaml | 22 - scripts/exps/calculate_starvation_offline.py | 26 - scripts/exps/modal_alpha_fairness_sweep.py | 15 +- scripts/exps/modal_starvation_sweep.py | 97 +++- scripts/exps/plot_starvation_sweep.py | 33 +- scripts/modal_sweep.py | 472 ------------------ scripts/plot_sweep.py | 191 ------- 12 files changed, 108 insertions(+), 1026 deletions(-) delete mode 100644 configs/experiments/edge_cases/1_fast_9_slow.jsonc delete mode 100644 configs/experiments/edge_cases/1_fast_only.jsonc delete mode 100644 configs/experiments/edge_cases/4_fast_1_slow.jsonc delete mode 100644 configs/experiments/edge_cases/5_fast_5_slow.jsonc delete mode 100644 configs/heterogeneous_1fast_11slow.yaml delete mode 100644 configs/heterogeneous_4fast_8slow.yaml delete mode 100644 scripts/exps/calculate_starvation_offline.py delete mode 100644 scripts/modal_sweep.py delete mode 100644 scripts/plot_sweep.py diff --git a/configs/experiments/edge_cases/1_fast_9_slow.jsonc b/configs/experiments/edge_cases/1_fast_9_slow.jsonc deleted file mode 100644 index 98857e0..0000000 --- a/configs/experiments/edge_cases/1_fast_9_slow.jsonc +++ /dev/null @@ -1,89 +0,0 @@ -{ - "experiment": { - "action_chunk_broker_type": "rtc", - "num_robots": 10, - "trials_per_robot": 2 - }, - "toxiproxy": { - "api_url": "http://127.0.0.1:8474", - "listen_host": "127.0.0.1", - "listen_port_base": 15000, - "server_args": [] - }, - "sampling": { - "default_seed": 7, - "resample_every_requests": 1 - }, - "robots": { - "robot_0": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_1": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_2": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_3": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_4": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_5": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_6": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_7": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_8": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_9": { - "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/experiments/edge_cases/1_fast_only.jsonc b/configs/experiments/edge_cases/1_fast_only.jsonc deleted file mode 100644 index 7c46fdc..0000000 --- a/configs/experiments/edge_cases/1_fast_only.jsonc +++ /dev/null @@ -1,26 +0,0 @@ -{ - "experiment": { - "action_chunk_broker_type": "rtc", - "num_robots": 1, - "trials_per_robot": 2 - }, - "toxiproxy": { - "api_url": "http://127.0.0.1:8474", - "listen_host": "127.0.0.1", - "listen_port_base": 15000, - "server_args": [] - }, - "sampling": { - "default_seed": 7, - "resample_every_requests": 1 - }, - "robots": { - "robot_0": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - } - } -} diff --git a/configs/experiments/edge_cases/4_fast_1_slow.jsonc b/configs/experiments/edge_cases/4_fast_1_slow.jsonc deleted file mode 100644 index c3fa4dd..0000000 --- a/configs/experiments/edge_cases/4_fast_1_slow.jsonc +++ /dev/null @@ -1,54 +0,0 @@ -{ - "experiment": { - "action_chunk_broker_type": "rtc", - "num_robots": 5, - "trials_per_robot": 2 - }, - "toxiproxy": { - "api_url": "http://127.0.0.1:8474", - "listen_host": "127.0.0.1", - "listen_port_base": 15000, - "server_args": [] - }, - "sampling": { - "default_seed": 7, - "resample_every_requests": 1 - }, - "robots": { - "robot_0": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_1": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_2": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_3": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_4": { - "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/experiments/edge_cases/5_fast_5_slow.jsonc b/configs/experiments/edge_cases/5_fast_5_slow.jsonc deleted file mode 100644 index 66e84a8..0000000 --- a/configs/experiments/edge_cases/5_fast_5_slow.jsonc +++ /dev/null @@ -1,89 +0,0 @@ -{ - "experiment": { - "action_chunk_broker_type": "rtc", - "num_robots": 10, - "trials_per_robot": 2 - }, - "toxiproxy": { - "api_url": "http://127.0.0.1:8474", - "listen_host": "127.0.0.1", - "listen_port_base": 15000, - "server_args": [] - }, - "sampling": { - "default_seed": 7, - "resample_every_requests": 1 - }, - "robots": { - "robot_0": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_1": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_2": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_3": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_4": { - "execution_horizon": 4, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_5": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_6": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_7": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_8": { - "execution_horizon": 10, - "uplink_median_ms": 0.0, - "uplink_sigma": 0.0, - "downlink_median_ms": 0.0, - "downlink_sigma": 0.0 - }, - "robot_9": { - "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/heterogeneous_1fast_11slow.yaml deleted file mode 100644 index 8803a67..0000000 --- a/configs/heterogeneous_1fast_11slow.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# Heterogeneous control-rate config for run_real.py. -# -# Maps workstation id (int, matches FleetConfig.robots[].id) to the control_hz -# the piper_client_armory node should declare. Robots not listed here use the -# node's default (CONTROL_HZ in client_node_armory.py). -# -# This config: 1 fast (20 Hz) + 11 slow (10 Hz) = 12 robots. -control_hz: - 14: 20 - 13: 10 - 12: 10 - 11: 10 - 8: 10 - 7: 10 - 6: 10 - 5: 10 - 4: 10 - 3: 10 - 2: 10 - 1: 10 diff --git a/configs/heterogeneous_4fast_8slow.yaml b/configs/heterogeneous_4fast_8slow.yaml deleted file mode 100644 index cde1af2..0000000 --- a/configs/heterogeneous_4fast_8slow.yaml +++ /dev/null @@ -1,22 +0,0 @@ -# Heterogeneous control-rate config for run_real.py. -# -# Maps workstation id (int, matches FleetConfig.robots[].id) to the control_hz -# the piper_client_armory node should declare. Robots not listed here use the -# node's default (CONTROL_HZ in client_node_armory.py). -# -# This config: 4 fast (20 Hz) + 7 slow (10 Hz) = 11 robots. Per request. -# (Note: only 11 entries — drop one robot from the trial selection or extend -# this file if you want all 12.) -control_hz: - 14: 20 - 13: 20 - 12: 20 - 11: 20 - 8: 10 - 7: 10 - 6: 10 - 5: 10 - 4: 10 - 3: 10 - 2: 10 - 1: 10 diff --git a/scripts/exps/calculate_starvation_offline.py b/scripts/exps/calculate_starvation_offline.py deleted file mode 100644 index ad11fab..0000000 --- a/scripts/exps/calculate_starvation_offline.py +++ /dev/null @@ -1,26 +0,0 @@ - -from pathlib import Path -import numpy as np - -episode = Path("/coc/flash7/rbansal66/vvla/0_real_0_failure") -control_hz = 20 # change if this run used a different hz - -costs = np.load(episode / "cost_history.npy") -starved = np.isnan(costs) - -first_non_starved = np.flatnonzero(~starved) -if first_non_starved.size: - first = int(first_non_starved[0]) - post_first_starved = int(starved[first:].sum()) - post_first_observed = int(len(costs) - first) -else: - post_first_starved = 0 - post_first_observed = 0 - -print("observed_steps:", len(costs)) -print("starvation_steps:", int(starved.sum())) -print("starvation_rate:", float(starved.mean())) -print("starvation_seconds:", float(starved.sum() / control_hz)) -print("post_first_starvation_steps:", post_first_starved) -print("post_first_observed_steps:", post_first_observed) -print("post_first_starvation_rate:", post_first_starved / post_first_observed if post_first_observed else 0.0) diff --git a/scripts/exps/modal_alpha_fairness_sweep.py b/scripts/exps/modal_alpha_fairness_sweep.py index 9eeeba9..0c92ee6 100644 --- a/scripts/exps/modal_alpha_fairness_sweep.py +++ b/scripts/exps/modal_alpha_fairness_sweep.py @@ -47,15 +47,15 @@ ) CONTROL_HZ = 20 -# MAX_BATCH_SIZE = 4 -MAX_BATCH_SIZE = 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") +BASELINE_SCHEDULERS = ("fixed-max-batch", "greedy-deadline", "round-robin", "lookahead-actions") DYNAMIC_SCHEDULER = "dynamic-action" MODEL_TO_PROFILE = { @@ -72,6 +72,7 @@ "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" @@ -578,12 +579,12 @@ def _plot_starvation_vs_fairness(results_csv: pathlib.Path, plots_dir: pathlib.P @app.local_entrypoint() def main( # models: str = "pi05,gr00t-n1.7", - models: str = "gr00t-n1.7", + models: str = "pi05", # scenarios: str = "1f9s,5f5s", - scenarios: str = "5f20s", - alpha_grid: str = "0.0,0.1,0.2,0.3,0.5,0.7,1.0", + 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_gr00t", + output_dir: str = "experiments/sweeps/fairness_alpha_sweep_pi05_1f9s", port: int = 8080, max_steps: int = DEFAULT_MAX_STEPS, ) -> None: diff --git a/scripts/exps/modal_starvation_sweep.py b/scripts/exps/modal_starvation_sweep.py index 0d8e346..efba743 100644 --- a/scripts/exps/modal_starvation_sweep.py +++ b/scripts/exps/modal_starvation_sweep.py @@ -15,19 +15,21 @@ import csv import dataclasses import datetime as dt -import io import json import pathlib +import shutil import subprocess import sys -import tarfile from typing import Any import modal APP_NAME = "armory-scheduler-sweep" +ARTIFACTS_VOLUME_NAME = "armory-scheduler-sweep-artifacts" REMOTE_ROOT = pathlib.Path("/app") REMOTE_OUTPUT_ROOT = pathlib.Path("/tmp/armory_sweep") +REMOTE_ARTIFACTS_ROOT = pathlib.Path("/artifacts") +ARTIFACT_SKIP_SUFFIXES = {".mp4", ".parquet", ".npz"} PYTHONPATH = ":".join( [ str(REMOTE_ROOT / "src"), @@ -56,6 +58,8 @@ def _ignore_modal_copy(path: pathlib.Path) -> bool: app = modal.App(APP_NAME) +artifacts_volume = modal.Volume.from_name(ARTIFACTS_VOLUME_NAME, create_if_missing=True) + @dataclasses.dataclass(frozen=True) class SweepCase: @@ -95,16 +99,14 @@ def _run_subprocess( 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 _copy_run_dir(src_root: pathlib.Path, dest_root: pathlib.Path) -> None: + """Copy run_dir into the mounted artifacts volume, skipping bulky binaries.""" + for src in src_root.rglob("*"): + if src.is_dir() or src.suffix in ARTIFACT_SKIP_SUFFIXES: + continue + dst = dest_root / src.relative_to(src_root) + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dst) def _safe_float(value: Any, default: float = 0.0) -> float: @@ -136,14 +138,32 @@ def _build_server_cmd(srv_cfg: dict[str, Any], *, port: int, scheduler: str) -> def _expand_experiment_config(exp_cfg: dict[str, Any], num_robots: int) -> dict[str, Any]: - """Return a copy of exp_cfg with robot_0's profile replicated to num_robots robots.""" + """Return a copy of exp_cfg with robot profiles set for num_robots robots. + + If the config already defines more than one robot explicitly, those profiles are used + as-is and num_robots is ignored (the config is authoritative). + Otherwise robot_0's profile is replicated to fill num_robots robots. + """ + if len(exp_cfg["robots"]) > 1: + actual = len(exp_cfg["robots"]) + return {**exp_cfg, "experiment": {**exp_cfg["experiment"], "num_robots": actual}} robot_template = exp_cfg["robots"]["robot_0"] - expanded = { + return { **exp_cfg, "experiment": {**exp_cfg["experiment"], "num_robots": num_robots}, "robots": {f"robot_{i}": dict(robot_template) for i in range(num_robots)}, } - return expanded + + +def _config_num_robots(cfg_path: str, fallback: int) -> int: + """Read robot count from a local config file when robots are pre-defined, else fallback.""" + try: + n = len(json.loads(pathlib.Path(cfg_path).read_text()).get("robots", {})) + if n > 1: + return n + except Exception: + pass + return fallback def _build_client_cmd( @@ -260,12 +280,19 @@ def _summarize_run(output_dir: pathlib.Path, case: SweepCase) -> dict[str, Any]: return summary -@app.function(image=image, timeout=60 * 60, cpu=4, memory=2048) +@app.function( + image=image, + timeout=60 * 60, + cpu=4, + memory=8192, + volumes={str(REMOTE_ARTIFACTS_ROOT): artifacts_volume}, +) def run_case( case: SweepCase, *, server_config: str, port: int, + stamp: str, max_batch_size_override: int | None = None, max_steps_override: int | None = None, ) -> dict[str, Any]: @@ -335,7 +362,10 @@ def run_case( server_proc.kill() server_proc.wait(timeout=20) - summary["artifact_tgz"] = _tar_directory(run_dir) + dest = REMOTE_ARTIFACTS_ROOT / stamp / case.run_id + _copy_run_dir(run_dir, dest) + artifacts_volume.commit() + summary["artifact_remote_path"] = str(dest) return summary @@ -344,7 +374,7 @@ def _write_rows(path: pathlib.Path, rows: list[dict[str, Any]]) -> None: keys: list[str] = [] for row in rows: for key in row: - if key != "artifact_tgz" and key not in keys: + if key not in keys: keys.append(key) with path.open("w", newline="") as f: writer = csv.DictWriter(f, fieldnames=keys) @@ -370,7 +400,12 @@ def main( stamp = dt.datetime.now(dt.timezone.utc).strftime("%Y%m%d_%H%M%S") # noqa: UP017 artifacts_dir = out / "artifacts" cases = [ - SweepCase(scheduler=scheduler, experiment_config=cfg, num_robots=n, seed=seed) + SweepCase( + scheduler=scheduler, + experiment_config=cfg, + num_robots=_config_num_robots(cfg, n), + seed=seed, + ) for scheduler in _parse_csv(schedulers) for cfg in _parse_csv(experiment_configs) for n in _parse_csv(num_robots, cast=int) @@ -383,23 +418,35 @@ def main( kwargs={ "server_config": server_config, "port": port, + "stamp": stamp, "max_batch_size_override": max_batch_size, "max_steps_override": max_steps, }, order_outputs=False, ): - artifact_bytes = result.pop("artifact_tgz") - 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) print( f"{result['status']}: {result['run_id']} " f"starvation={_safe_float(result.get('starvation_rate')):.3f} " ) + artifacts_dir.mkdir(parents=True, exist_ok=True) + print(f"Downloading artifacts from volume '{ARTIFACTS_VOLUME_NAME}/{stamp}' -> {artifacts_dir}") + subprocess.run( + [ + "modal", + "volume", + "get", + ARTIFACTS_VOLUME_NAME, + stamp, + str(artifacts_dir), + "--force", + ], + check=True, + ) + for row in rows: + row["artifact_path"] = str(artifacts_dir / stamp / row["run_id"]) + sweep_csv = out / f"sweep_results_{stamp}.csv" latest_csv = out / "sweep_results.csv" _write_rows(sweep_csv, rows) diff --git a/scripts/exps/plot_starvation_sweep.py b/scripts/exps/plot_starvation_sweep.py index d20e9bf..c3aca62 100644 --- a/scripts/exps/plot_starvation_sweep.py +++ b/scripts/exps/plot_starvation_sweep.py @@ -65,6 +65,15 @@ def _wilson_ci(p: float, n: int, z: float = 1.96) -> tuple[float, float]: return (max(0.0, center - half), min(1.0, center + half)) +STARVATION_METRICS = { + "starvation_rate", + "post_first_starvation_rate", + "robot_starvation_rate_max", + "robot_starvation_rate_std", + "robot_starvation_rate_cvar90", +} + + def _plot_metric( df: pd.DataFrame, *, @@ -72,12 +81,17 @@ def _plot_metric( x_col: str, line_col: str, output_dir: pathlib.Path, + reduce: str = "mean", ) -> pathlib.Path: - agg = df.groupby([line_col, x_col])[metric].agg(["mean", "count"]).reset_index() - agg.columns = [line_col, x_col, "value", "n"] + if reduce == "min": + agg = df.groupby([line_col, x_col])[metric].agg(["min", "count"]).reset_index() + agg.columns = [line_col, x_col, "value", "n"] + else: + agg = df.groupby([line_col, x_col])[metric].agg(["mean", "count"]).reset_index() + agg.columns = [line_col, x_col, "value", "n"] agg = agg.sort_values([line_col, x_col]) - is_proportion = agg["value"].between(0.0, 1.0).all() + is_proportion = reduce == "mean" and agg["value"].between(0.0, 1.0).all() if is_proportion: ci = agg.apply( lambda r: pd.Series(_wilson_ci(r["value"], int(r["n"])), index=["lo", "hi"]), axis=1 @@ -98,14 +112,18 @@ def _plot_metric( color=line.get_color(), ) + title_suffix = " (best seed)" if reduce == "min" else "" ax.set_xlabel(x_col.replace("_", " ").title()) ax.set_ylabel(_metric_label(metric)) - ax.set_title(_metric_label(metric)) + ax.set_title(_metric_label(metric) + title_suffix) ax.grid(True, axis="y", alpha=0.25) ax.legend(title=line_col.replace("_", " ").title()) fig.tight_layout() - output_path = output_dir / f"{metric}_by_{x_col}.png" + filename = ( + f"{metric}_by_{x_col}_min_seed.png" if reduce == "min" else f"{metric}_by_{x_col}.png" + ) + output_path = output_dir / filename fig.savefig(output_path, dpi=160) plt.close(fig) return output_path @@ -147,6 +165,11 @@ def plot_results( _plot_metric(df, metric=metric, x_col=x, line_col=line, output_dir=output_dir) for metric in metrics ] + written += [ + _plot_metric(df, metric=metric, x_col=x, line_col=line, output_dir=output_dir, reduce="min") + for metric in metrics + if metric in STARVATION_METRICS + ] print("Wrote plots:") for path in written: print(path) diff --git a/scripts/modal_sweep.py b/scripts/modal_sweep.py deleted file mode 100644 index efba743..0000000 --- a/scripts/modal_sweep.py +++ /dev/null @@ -1,472 +0,0 @@ -"""Run server/client scheduler sweeps on Modal. - -Example: - modal run scripts/modal_sweep.py \ - --schedulers fixed-max-batch,greedy-deadline,round-robin \ - --experiment-configs configs/experiments/mock/short.json \ - --num-robots 1,2,3,4,5,6,7,8,9,10 \ - --server-config configs/server/mock.json \ - --seeds 7,42 \ - --output-dir experiments/sweeps/big_mock -""" - -from __future__ import annotations - -import csv -import dataclasses -import datetime as dt -import json -import pathlib -import shutil -import subprocess -import sys -from typing import Any - -import modal - -APP_NAME = "armory-scheduler-sweep" -ARTIFACTS_VOLUME_NAME = "armory-scheduler-sweep-artifacts" -REMOTE_ROOT = pathlib.Path("/app") -REMOTE_OUTPUT_ROOT = pathlib.Path("/tmp/armory_sweep") -REMOTE_ARTIFACTS_ROOT = pathlib.Path("/artifacts") -ARTIFACT_SKIP_SUFFIXES = {".mp4", ".parquet", ".npz"} -PYTHONPATH = ":".join( - [ - str(REMOTE_ROOT / "src"), - str(REMOTE_ROOT / "src/backends"), - str(REMOTE_ROOT / "packages/armory-client/src"), - ] -) - - -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) - -artifacts_volume = modal.Volume.from_name(ARTIFACTS_VOLUME_NAME, create_if_missing=True) - - -@dataclasses.dataclass(frozen=True) -class SweepCase: - scheduler: str - experiment_config: str # path relative to repo root - num_robots: int - seed: int - - @property - def run_id(self) -> str: - config_name = pathlib.Path(self.experiment_config).stem - return f"scheduler={self.scheduler}__config={config_name}__robots={self.num_robots}__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 _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 _copy_run_dir(src_root: pathlib.Path, dest_root: pathlib.Path) -> None: - """Copy run_dir into the mounted artifacts volume, skipping bulky binaries.""" - for src in src_root.rglob("*"): - if src.is_dir() or src.suffix in ARTIFACT_SKIP_SUFFIXES: - continue - dst = dest_root / src.relative_to(src_root) - dst.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(src, dst) - - -def _safe_float(value: Any, default: float = 0.0) -> float: - try: - if value is None: - return default - return float(value) - except (TypeError, ValueError): - return default - - -def _build_server_cmd(srv_cfg: dict[str, Any], *, port: int, scheduler: str) -> list[str]: - cmd = [ - sys.executable, - "scripts/serve.py", - "--port", - str(port), - "--env", - srv_cfg.get("env", "LIBERO"), - "--max-batch-size", - str(srv_cfg.get("max_batch_size", 1)), - "--scheduling-algorithm", - scheduler, - f"policy:{srv_cfg.get('policy_type', 'default')}", - ] - for k, v in srv_cfg.get("policy", {}).items(): - cmd += [f"--policy.{k.replace('_', '-')}", str(v)] - return cmd - - -def _expand_experiment_config(exp_cfg: dict[str, Any], num_robots: int) -> dict[str, Any]: - """Return a copy of exp_cfg with robot profiles set for num_robots robots. - - If the config already defines more than one robot explicitly, those profiles are used - as-is and num_robots is ignored (the config is authoritative). - Otherwise robot_0's profile is replicated to fill num_robots robots. - """ - if len(exp_cfg["robots"]) > 1: - actual = len(exp_cfg["robots"]) - return {**exp_cfg, "experiment": {**exp_cfg["experiment"], "num_robots": actual}} - robot_template = exp_cfg["robots"]["robot_0"] - return { - **exp_cfg, - "experiment": {**exp_cfg["experiment"], "num_robots": num_robots}, - "robots": {f"robot_{i}": dict(robot_template) for i in range(num_robots)}, - } - - -def _config_num_robots(cfg_path: str, fallback: int) -> int: - """Read robot count from a local config file when robots are pre-defined, else fallback.""" - try: - n = len(json.loads(pathlib.Path(cfg_path).read_text()).get("robots", {})) - if n > 1: - return n - except Exception: - pass - return fallback - - -def _build_client_cmd( - *, - port: int, - seed: int, - output_dir: pathlib.Path, - experiment_config_path: pathlib.Path, - max_steps: int, -) -> list[str]: - return [ - sys.executable, - "scripts/run_libero.py", - "--host", - "127.0.0.1", - "--port", - str(port), - "--env", - "mock", - "--overwrite", - "--progress-type", - "logging", - "--max-steps", - str(max_steps), - "--seed", - str(seed), - "--output-dir", - str(output_dir), - "--experiment-config", - str(experiment_config_path), - ] - - -def _summarize_run(output_dir: pathlib.Path, case: SweepCase) -> dict[str, Any]: - summary_path = output_dir / "summary.csv" - results_path = output_dir / "results.csv" - runtime_path = output_dir / "runtime_metadata.json" - server_path = output_dir / "server_metadata.json" - - total_success = 0.0 - overall_starvation_rate = 0.0 - post_first_starvation_rate = 0.0 - if summary_path.exists(): - with summary_path.open() as f: - rows = list(csv.DictReader(f)) - if rows: - total_success = sum(_safe_float(r.get("success")) for r in rows) / len(rows) - starvation_steps = sum(_safe_float(r.get("starvation_steps")) for r in rows) - observed_steps = sum(_safe_float(r.get("observed_steps")) for r in rows) - post_first_starvation_steps = sum( - _safe_float(r.get("post_first_starvation_steps")) for r in rows - ) - post_first_observed_steps = sum( - _safe_float(r.get("post_first_observed_steps")) for r in rows - ) - overall_starvation_rate = starvation_steps / observed_steps if observed_steps else 0.0 - post_first_starvation_rate = ( - post_first_starvation_steps / post_first_observed_steps - if post_first_observed_steps - else 0.0 - ) - - robot_rates: list[float] = [] - if results_path.exists(): - by_robot: dict[str, dict[str, float]] = {} - with results_path.open() as f: - for row in csv.DictReader(f): - robot = str(row.get("robot_idx", "unknown")) - stats = by_robot.setdefault(robot, {"starvation_steps": 0.0, "observed_steps": 0.0}) - stats["starvation_steps"] += _safe_float(row.get("starvation_steps")) - stats["observed_steps"] += _safe_float(row.get("observed_steps")) - robot_rates = [ - stats["starvation_steps"] / stats["observed_steps"] - for stats in by_robot.values() - if stats["observed_steps"] > 0 - ] - - runtime = json.loads(runtime_path.read_text()) if runtime_path.exists() else {} - server = json.loads(server_path.read_text()) if server_path.exists() else {} - sorted_rates = sorted(robot_rates) - tail_count = max(1, int(len(sorted_rates) * 0.1)) if sorted_rates else 0 - - summary = { - "run_id": case.run_id, - "scheduler": case.scheduler, - "experiment_config": case.experiment_config, - "num_robots": case.num_robots, - "seed": case.seed, - "success_rate": total_success, - "starvation_rate": overall_starvation_rate, - "post_first_starvation_rate": post_first_starvation_rate, - "robot_starvation_rate_max": max(robot_rates) if robot_rates else 0.0, - "robot_starvation_rate_std": _safe_float(__import__("statistics").pstdev(robot_rates)) - if len(robot_rates) > 1 - else 0.0, - "robot_starvation_rate_cvar90": sum(sorted_rates[-tail_count:]) / tail_count - if tail_count - else 0.0, - "max_batch_size": server.get("max_batch_size", ""), - "action_horizon": server.get("action_horizon", ""), - "max_steps": runtime.get("max_steps", ""), - "num_trials_per_task": runtime.get("num_trials_per_task", ""), - } - - try: - from sims.libero.metrics import compute_server_timing_health # noqa: PLC0415 - - health = compute_server_timing_health(output_dir) - if health is not None: - summary.update(health) - except Exception: # noqa: BLE001 - pass - - return summary - - -@app.function( - image=image, - timeout=60 * 60, - cpu=4, - memory=8192, - volumes={str(REMOTE_ARTIFACTS_ROOT): artifacts_volume}, -) -def run_case( - case: SweepCase, - *, - server_config: str, - port: int, - stamp: str, - max_batch_size_override: int | None = None, - max_steps_override: int | None = None, -) -> dict[str, Any]: - exp_cfg: dict[str, Any] = json.loads((REMOTE_ROOT / case.experiment_config).read_text()) - srv_cfg: dict[str, Any] = json.loads((REMOTE_ROOT / server_config).read_text()) - - if max_batch_size_override is not None: - srv_cfg["max_batch_size"] = max_batch_size_override - if max_steps_override is not None: - exp_cfg["experiment"]["max_steps"] = max_steps_override - - exp_cfg = _expand_experiment_config(exp_cfg, case.num_robots) - max_steps = int(exp_cfg["experiment"]["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)) - (run_dir / "server_config.json").write_text(json.dumps(srv_cfg, indent=2)) - - server_cmd = _build_server_cmd(srv_cfg, port=port, scheduler=case.scheduler) - client_cmd = _build_client_cmd( - port=port, - seed=case.seed, - output_dir=output_dir, - experiment_config_path=saved_exp_config, - max_steps=max_steps, - ) - - 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 * 45) - summary = _summarize_run(output_dir, case) - summary["status"] = "ok" - except Exception as exc: # noqa: BLE001 - summary = { - "run_id": case.run_id, - "scheduler": case.scheduler, - "experiment_config": case.experiment_config, - "num_robots": case.num_robots, - "seed": case.seed, - "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) - - dest = REMOTE_ARTIFACTS_ROOT / stamp / case.run_id - _copy_run_dir(run_dir, dest) - artifacts_volume.commit() - summary["artifact_remote_path"] = str(dest) - 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 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}) - - -@app.local_entrypoint() -def main( - schedulers: str = "fixed-max-batch,greedy-deadline,round-robin", - experiment_configs: str = "configs/experiments/mock/short.json", - num_robots: str = "2,4,6", - server_config: str = "configs/server/mock.json", - seeds: str = "7", - output_dir: str = "experiments/sweeps/mock", - port: int = 8080, - max_batch_size: int | None = None, - max_steps: int | None = None, -) -> None: - """Run the Cartesian product of schedulers, experiment_configs, num_robots, and seeds.""" - 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( - scheduler=scheduler, - experiment_config=cfg, - num_robots=_config_num_robots(cfg, n), - seed=seed, - ) - for scheduler in _parse_csv(schedulers) - for cfg in _parse_csv(experiment_configs) - for n in _parse_csv(num_robots, cast=int) - for seed in _parse_csv(seeds, cast=int) - ] - - rows: list[dict[str, Any]] = [] - for result in run_case.map( - cases, - kwargs={ - "server_config": server_config, - "port": port, - "stamp": stamp, - "max_batch_size_override": max_batch_size, - "max_steps_override": max_steps, - }, - order_outputs=False, - ): - rows.append(result) - print( - f"{result['status']}: {result['run_id']} " - f"starvation={_safe_float(result.get('starvation_rate')):.3f} " - ) - - artifacts_dir.mkdir(parents=True, exist_ok=True) - print(f"Downloading artifacts from volume '{ARTIFACTS_VOLUME_NAME}/{stamp}' -> {artifacts_dir}") - subprocess.run( - [ - "modal", - "volume", - "get", - ARTIFACTS_VOLUME_NAME, - stamp, - str(artifacts_dir), - "--force", - ], - check=True, - ) - for row in rows: - row["artifact_path"] = str(artifacts_dir / stamp / row["run_id"]) - - 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}") - - suspicious = [r for r in rows if r.get("timing_suspicious")] - if suspicious: - print(f"WARNING: {len(suspicious)} run(s) flagged for suspicious timings:") - for r in suspicious: - print(f" {r['run_id']}: {r.get('timing_flags', '')}") - - sys.path.insert(0, str(pathlib.Path(__file__).parent)) - from plot_sweep import DEFAULT_METRICS, plot_results # noqa: PLC0415 - - timing_metrics = [ - "step_interval_p95_ms", - "inference_p99_ms", - "inbound_p95_ms", - "outbound_p95_ms", - ] - plot_results(latest_csv, out / "plots", metrics=list(DEFAULT_METRICS) + timing_metrics) diff --git a/scripts/plot_sweep.py b/scripts/plot_sweep.py deleted file mode 100644 index c3aca62..0000000 --- a/scripts/plot_sweep.py +++ /dev/null @@ -1,191 +0,0 @@ -"""Plot scheduler sweep metrics from scripts/modal_sweep.py output. - -Example: - uv run python scripts/plot_sweep.py \ - --results experiments/sweeps/mock/sweep_results.csv \ - --output-dir experiments/sweeps/mock/plots -""" - -from __future__ import annotations - -import argparse -import pathlib - -import matplotlib - -matplotlib.use("Agg") - -import matplotlib.pyplot as plt -import pandas as pd - -DEFAULT_METRICS = [ - "starvation_rate", - "post_first_starvation_rate", - "robot_starvation_rate_max", - "success_rate", -] - - -METRIC_LABELS = { - "starvation_rate": "Starvation rate", - "post_first_starvation_rate": "Starvation rate excl. startup", - "robot_starvation_rate_max": "Worst robot starvation rate", - "robot_starvation_rate_std": "Robot starvation std. dev.", - "robot_starvation_rate_cvar90": "Tail robot starvation rate", - "success_rate": "Success rate", - "step_interval_p95_ms": "Step interval p95 (ms)", - "inference_p99_ms": "Inference latency p99 (ms)", - "inbound_p95_ms": "Client→server transport p95 (ms)", - "outbound_p95_ms": "Server→client transport p95 (ms)", -} - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--results", type=pathlib.Path, required=True) - parser.add_argument("--output-dir", type=pathlib.Path, default=None) - parser.add_argument("--x", default="num_robots") - parser.add_argument("--line", default="scheduler") - parser.add_argument("--metrics", default=",".join(DEFAULT_METRICS)) - return parser.parse_args() - - -def _metric_label(metric: str) -> str: - return METRIC_LABELS.get(metric, metric.replace("_", " ").title()) - - -def _wilson_ci(p: float, n: int, z: float = 1.96) -> tuple[float, float]: - """95% Wilson score interval for a proportion p estimated from n observations.""" - if n == 0: - return (p, p) - z2 = z * z - denom = 1 + z2 / n - center = (p + z2 / (2 * n)) / denom - half = z * (p * (1 - p) / n + z2 / (4 * n * n)) ** 0.5 / denom - return (max(0.0, center - half), min(1.0, center + half)) - - -STARVATION_METRICS = { - "starvation_rate", - "post_first_starvation_rate", - "robot_starvation_rate_max", - "robot_starvation_rate_std", - "robot_starvation_rate_cvar90", -} - - -def _plot_metric( - df: pd.DataFrame, - *, - metric: str, - x_col: str, - line_col: str, - output_dir: pathlib.Path, - reduce: str = "mean", -) -> pathlib.Path: - if reduce == "min": - agg = df.groupby([line_col, x_col])[metric].agg(["min", "count"]).reset_index() - agg.columns = [line_col, x_col, "value", "n"] - else: - agg = df.groupby([line_col, x_col])[metric].agg(["mean", "count"]).reset_index() - agg.columns = [line_col, x_col, "value", "n"] - agg = agg.sort_values([line_col, x_col]) - - is_proportion = reduce == "mean" and agg["value"].between(0.0, 1.0).all() - if is_proportion: - ci = agg.apply( - lambda r: pd.Series(_wilson_ci(r["value"], int(r["n"])), index=["lo", "hi"]), axis=1 - ) - agg = pd.concat([agg, ci], axis=1) - - fig, ax = plt.subplots(figsize=(8, 4.8)) - for line_value, group in agg.groupby(line_col): - xs = group[x_col].to_numpy() - ys = group["value"].to_numpy() - (line,) = ax.plot(xs, ys, marker="o", linewidth=2.0, label=str(line_value)) - if is_proportion and (group["n"] > 1).any(): - ax.fill_between( - xs, - group["lo"].to_numpy(), - group["hi"].to_numpy(), - alpha=0.15, - color=line.get_color(), - ) - - title_suffix = " (best seed)" if reduce == "min" else "" - ax.set_xlabel(x_col.replace("_", " ").title()) - ax.set_ylabel(_metric_label(metric)) - ax.set_title(_metric_label(metric) + title_suffix) - ax.grid(True, axis="y", alpha=0.25) - ax.legend(title=line_col.replace("_", " ").title()) - fig.tight_layout() - - filename = ( - f"{metric}_by_{x_col}_min_seed.png" if reduce == "min" else f"{metric}_by_{x_col}.png" - ) - output_path = output_dir / filename - fig.savefig(output_path, dpi=160) - plt.close(fig) - return output_path - - -def plot_results( - results: pathlib.Path, - output_dir: pathlib.Path | None = None, - *, - x: str = "num_robots", - line: str = "scheduler", - metrics: list[str] | None = None, -) -> None: - output_dir = output_dir or (results.parent / "plots") - output_dir.mkdir(parents=True, exist_ok=True) - if metrics is None: - metrics = list(DEFAULT_METRICS) - - df = pd.read_csv(results) - if "status" in df.columns: - df = df[df["status"] == "ok"].copy() - if df.empty: - raise SystemExit("No successful rows found in results CSV") - - missing = [m for m in metrics if m not in df.columns] - if missing: - raise SystemExit(f"Missing metric column(s): {', '.join(missing)}") - for column in [x, line, *metrics]: - if column not in df.columns: - raise SystemExit(f"Missing required column: {column}") - - for metric in metrics: - df[metric] = pd.to_numeric(df[metric], errors="coerce") - numeric_x = pd.to_numeric(df[x], errors="coerce") - if numeric_x.notna().all(): - df[x] = numeric_x - - written = [ - _plot_metric(df, metric=metric, x_col=x, line_col=line, output_dir=output_dir) - for metric in metrics - ] - written += [ - _plot_metric(df, metric=metric, x_col=x, line_col=line, output_dir=output_dir, reduce="min") - for metric in metrics - if metric in STARVATION_METRICS - ] - print("Wrote plots:") - for path in written: - print(path) - - -def main() -> None: - args = _parse_args() - metrics = [m.strip() for m in args.metrics.split(",") if m.strip()] - plot_results( - args.results, - args.output_dir, - x=args.x, - line=args.line, - metrics=metrics, - ) - - -if __name__ == "__main__": - main() From fdc5ff2669dc8c9ac1b22b46faaf8198a45f63ce Mon Sep 17 00:00:00 2001 From: Rohan Bansal Date: Mon, 11 May 2026 23:00:23 -0400 Subject: [PATCH 4/6] fix min_ex horizon renaming --- packages/armory-client/src/armory_client/schemas.py | 2 +- scripts/serve.py | 4 ++-- src/armory/scheduling/base.py | 10 ++++++---- src/armory/scheduling/baselines.py | 4 ++-- src/armory/scheduling/dynamic_action.py | 8 +++----- src/armory/scheduling/lookahead.py | 4 ++-- src/armory/scheduling/lookahead_actions.py | 4 ++-- src/armory/scheduling/mirror.py | 8 ++++---- src/armory/serving/engine.py | 8 +++++--- src/armory/serving/scheduler.py | 6 +++--- src/armory/serving/server.py | 4 ++-- 11 files changed, 32 insertions(+), 30 deletions(-) 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/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 47bac64..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] = {} @@ -86,9 +86,7 @@ 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, deadlines: dict[str, float]) -> float: 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 193a64b..814cf26 100644 --- a/src/armory/scheduling/mirror.py +++ b/src/armory/scheduling/mirror.py @@ -352,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. """ @@ -379,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, ) From 7c3081fc1a943c87091964dda47c31dce4223d34 Mon Sep 17 00:00:00 2001 From: Rohan Bansal Date: Mon, 11 May 2026 23:10:23 -0400 Subject: [PATCH 5/6] refactored savers into common file --- .../src/armory_client/runtime/real_saver.py | 98 +++++----------- .../src/armory_client/runtime/saver_utils.py | 100 ++++++++++++++++ scripts/exps/modal_alpha_fairness_sweep.py | 2 +- src/sims/libero/subscribers/saver.py | 110 ++++++------------ 4 files changed, 162 insertions(+), 148 deletions(-) create mode 100644 packages/armory-client/src/armory_client/runtime/saver_utils.py 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/scripts/exps/modal_alpha_fairness_sweep.py b/scripts/exps/modal_alpha_fairness_sweep.py index 0c92ee6..6fa06be 100644 --- a/scripts/exps/modal_alpha_fairness_sweep.py +++ b/scripts/exps/modal_alpha_fairness_sweep.py @@ -557,7 +557,7 @@ def _plot_starvation_vs_fairness(results_csv: pathlib.Path, plots_dir: pathlib.P ax_min, sub, y_col="min_freshness", y_label="Min freshness = 1 − max(starvation) (higher is fairer)", - title="Mean starvation vs min freshness (α=∞ welfare)", + title="Mean starvation vs min freshness", ) _autoscale_with_pad(ax_min, sub, "min_freshness") 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}") From ded8a76e3d84567311999c257a0dcd9e5f5d6558 Mon Sep 17 00:00:00 2001 From: Rohan Bansal Date: Mon, 11 May 2026 23:11:42 -0400 Subject: [PATCH 6/6] fix folder naming --- configs/experiments/mock/short.json | 17 ---- .../exps/edge_cases_real/1fast_11slow.yaml | 20 +++++ configs/exps/edge_cases_real/4fast_8slow.yaml | 22 +++++ .../exps/edge_cases_sim/1_fast_9_slow.jsonc | 89 +++++++++++++++++++ configs/exps/edge_cases_sim/1_fast_only.jsonc | 26 ++++++ .../exps/edge_cases_sim/4_fast_1_slow.jsonc | 54 +++++++++++ .../exps/edge_cases_sim/5_fast_5_slow.jsonc | 89 +++++++++++++++++++ 7 files changed, 300 insertions(+), 17 deletions(-) delete mode 100644 configs/experiments/mock/short.json create mode 100644 configs/exps/edge_cases_real/1fast_11slow.yaml create mode 100644 configs/exps/edge_cases_real/4fast_8slow.yaml create mode 100644 configs/exps/edge_cases_sim/1_fast_9_slow.jsonc create mode 100644 configs/exps/edge_cases_sim/1_fast_only.jsonc create mode 100644 configs/exps/edge_cases_sim/4_fast_1_slow.jsonc create mode 100644 configs/exps/edge_cases_sim/5_fast_5_slow.jsonc 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/exps/edge_cases_real/1fast_11slow.yaml b/configs/exps/edge_cases_real/1fast_11slow.yaml new file mode 100644 index 0000000..8803a67 --- /dev/null +++ b/configs/exps/edge_cases_real/1fast_11slow.yaml @@ -0,0 +1,20 @@ +# Heterogeneous control-rate config for run_real.py. +# +# Maps workstation id (int, matches FleetConfig.robots[].id) to the control_hz +# the piper_client_armory node should declare. Robots not listed here use the +# node's default (CONTROL_HZ in client_node_armory.py). +# +# This config: 1 fast (20 Hz) + 11 slow (10 Hz) = 12 robots. +control_hz: + 14: 20 + 13: 10 + 12: 10 + 11: 10 + 8: 10 + 7: 10 + 6: 10 + 5: 10 + 4: 10 + 3: 10 + 2: 10 + 1: 10 diff --git a/configs/exps/edge_cases_real/4fast_8slow.yaml b/configs/exps/edge_cases_real/4fast_8slow.yaml new file mode 100644 index 0000000..cde1af2 --- /dev/null +++ b/configs/exps/edge_cases_real/4fast_8slow.yaml @@ -0,0 +1,22 @@ +# Heterogeneous control-rate config for run_real.py. +# +# Maps workstation id (int, matches FleetConfig.robots[].id) to the control_hz +# the piper_client_armory node should declare. Robots not listed here use the +# node's default (CONTROL_HZ in client_node_armory.py). +# +# This config: 4 fast (20 Hz) + 7 slow (10 Hz) = 11 robots. Per request. +# (Note: only 11 entries — drop one robot from the trial selection or extend +# this file if you want all 12.) +control_hz: + 14: 20 + 13: 20 + 12: 20 + 11: 20 + 8: 10 + 7: 10 + 6: 10 + 5: 10 + 4: 10 + 3: 10 + 2: 10 + 1: 10 diff --git a/configs/exps/edge_cases_sim/1_fast_9_slow.jsonc b/configs/exps/edge_cases_sim/1_fast_9_slow.jsonc new file mode 100644 index 0000000..98857e0 --- /dev/null +++ b/configs/exps/edge_cases_sim/1_fast_9_slow.jsonc @@ -0,0 +1,89 @@ +{ + "experiment": { + "action_chunk_broker_type": "rtc", + "num_robots": 10, + "trials_per_robot": 2 + }, + "toxiproxy": { + "api_url": "http://127.0.0.1:8474", + "listen_host": "127.0.0.1", + "listen_port_base": 15000, + "server_args": [] + }, + "sampling": { + "default_seed": 7, + "resample_every_requests": 1 + }, + "robots": { + "robot_0": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_1": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_2": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_3": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_4": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_5": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_6": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_7": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_8": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_9": { + "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/exps/edge_cases_sim/1_fast_only.jsonc b/configs/exps/edge_cases_sim/1_fast_only.jsonc new file mode 100644 index 0000000..7c46fdc --- /dev/null +++ b/configs/exps/edge_cases_sim/1_fast_only.jsonc @@ -0,0 +1,26 @@ +{ + "experiment": { + "action_chunk_broker_type": "rtc", + "num_robots": 1, + "trials_per_robot": 2 + }, + "toxiproxy": { + "api_url": "http://127.0.0.1:8474", + "listen_host": "127.0.0.1", + "listen_port_base": 15000, + "server_args": [] + }, + "sampling": { + "default_seed": 7, + "resample_every_requests": 1 + }, + "robots": { + "robot_0": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + } + } +} diff --git a/configs/exps/edge_cases_sim/4_fast_1_slow.jsonc b/configs/exps/edge_cases_sim/4_fast_1_slow.jsonc new file mode 100644 index 0000000..c3fa4dd --- /dev/null +++ b/configs/exps/edge_cases_sim/4_fast_1_slow.jsonc @@ -0,0 +1,54 @@ +{ + "experiment": { + "action_chunk_broker_type": "rtc", + "num_robots": 5, + "trials_per_robot": 2 + }, + "toxiproxy": { + "api_url": "http://127.0.0.1:8474", + "listen_host": "127.0.0.1", + "listen_port_base": 15000, + "server_args": [] + }, + "sampling": { + "default_seed": 7, + "resample_every_requests": 1 + }, + "robots": { + "robot_0": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_1": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_2": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_3": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_4": { + "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/exps/edge_cases_sim/5_fast_5_slow.jsonc b/configs/exps/edge_cases_sim/5_fast_5_slow.jsonc new file mode 100644 index 0000000..66e84a8 --- /dev/null +++ b/configs/exps/edge_cases_sim/5_fast_5_slow.jsonc @@ -0,0 +1,89 @@ +{ + "experiment": { + "action_chunk_broker_type": "rtc", + "num_robots": 10, + "trials_per_robot": 2 + }, + "toxiproxy": { + "api_url": "http://127.0.0.1:8474", + "listen_host": "127.0.0.1", + "listen_port_base": 15000, + "server_args": [] + }, + "sampling": { + "default_seed": 7, + "resample_every_requests": 1 + }, + "robots": { + "robot_0": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_1": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_2": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_3": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_4": { + "execution_horizon": 4, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_5": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_6": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_7": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_8": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + }, + "robot_9": { + "execution_horizon": 10, + "uplink_median_ms": 0.0, + "uplink_sigma": 0.0, + "downlink_median_ms": 0.0, + "downlink_sigma": 0.0 + } + } +}