File size: 2,626 Bytes
af8c481 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """Collect <runs_dir>/*/eval.json into one markdown table (plus baselines).
usage: python scripts/collect.py [runs_dir=runs] [baseline_json=results/ablation_baseline/eval_all.json]"""
from __future__ import annotations
import json, sys
from pathlib import Path
root = Path(sys.argv[1] if len(sys.argv) > 1 else "runs")
rows = []
for run in sorted(root.glob("*/eval.json")):
r = json.load(open(run))
ol = r["openloop_action_mse"]; sp = r["openloop_action_mse_split"]
cfg = r["config"]; tr = json.load(open(run.parent / "train_args.json")) if (run.parent / "train_args.json").exists() else {}
rows.append({
"run": run.parent.name,
"params_M": r["params_total"] / 1e6,
"d/L": f"{cfg['dim']}/{cfg['num_layers']}",
"future_state": cfg["predict_future_state"],
"train_frac": tr.get("train_fraction"),
"heldout": r.get("heldout_fraction"),
"cascade": ol.get("cascade", [float("nan")])[0],
"coarse_only(τ=.4)": ol.get("coarse_only", [float("nan")])[0],
"coarse_to_zero": ol.get("coarse_to_zero", [float("nan")])[0],
"cascade_arm": sp.get("cascade_arm", float("nan")),
"cascade_hand": sp.get("cascade_hand", float("nan")),
"fstate_mse": r.get("openloop_future_state_mse", {}).get("cascade", float("nan")),
"act_loss@τ.2": r["forward_tau_losses"]["0.2"]["action_loss"][0],
})
base = json.load(open(sys.argv[2] if len(sys.argv) > 2 else Path(__file__).resolve().parents[1] / "results/ablation_baseline/eval_all.json"))
for name, r in base.items():
ol = r["openloop_action_mse"]; sp = r["openloop_action_mse_split"]
rows.append({
"run": f"baseline:{name}", "params_M": 5686.4 if name != "noforce" else 5638.5, "d/L": "3072/30 (Wan2.2 LoRA r4)",
"future_state": "video", "train_frac": None, "heldout": None,
"cascade": ol.get("cascade", [float("nan")])[0], "coarse_only(τ=.4)": ol.get("coarse_only", [float("nan")])[0],
"coarse_to_zero": float("nan"), "cascade_arm": sp.get("cascade_arm", float("nan")), "cascade_hand": sp.get("cascade_hand", float("nan")),
"fstate_mse": float("nan"), "act_loss@τ.2": r["forward_tau_losses"]["0.2"]["action_loss"][0],
})
keys = list(rows[0].keys())
print("| " + " | ".join(keys) + " |")
print("|" + "---|" * len(keys))
for row in rows:
cells = []
for k in keys:
v = row[k]
cells.append(f"{v:.4f}" if isinstance(v, float) and k != "params_M" else (f"{v:.1f}" if isinstance(v, float) else str(v)))
print("| " + " | ".join(cells) + " |")
json.dump(rows, open(root / "summary.json", "w"), indent=1)
|