File size: 2,677 Bytes
87d0bc2 | 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 | """Plot N run-series (comma-separated prefixes/labels) on one scaling figure."""
import json, os, argparse
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from common import wilson_ci
ap = argparse.ArgumentParser()
ap.add_argument("--prefixes", required=True)
ap.add_argument("--labels", required=True)
ap.add_argument("--budgets", nargs="+", type=int, default=[10, 50, 100, 1000])
ap.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2])
ap.add_argument("--out", default="results/ablation.png")
a = ap.parse_args()
os.makedirs(os.path.dirname(a.out), exist_ok=True)
prefixes = a.prefixes.split(","); labels = a.labels.split("|")
colors = ["#2563eb", "#16a34a", "#dc2626", "#9333ea"]
plt.figure(figsize=(7.8, 4.9)); allrows = {}
for (pref, label, color) in zip(prefixes, labels, colors):
rows = []
for b in a.budgets:
per, K, N = [], 0, 0
for s in a.seeds:
p = f"runs/{pref}_b{b:04d}_seed{s}/eval.json"
if not os.path.exists(p): continue
o = json.load(open(p))["overall"]; per.append(o["rate"]); K += o["k"]; N += o["n"]
if not per: continue
per = np.array(per); pr, lo, hi = wilson_ci(K, N)
rows.append(dict(demos=b, per_seed=per.tolist(), pooled_rate=pr, ci_low=lo, ci_high=hi,
pooled_k=K, pooled_n=N, mean=float(per.mean()),
std=float(per.std(ddof=1)) if len(per) > 1 else 0.0))
allrows[pref] = rows
if not rows: continue
xs = [r["demos"] for r in rows]; y = np.array([r["pooled_rate"] for r in rows])
lo = np.array([r["ci_low"] for r in rows]); hi = np.array([r["ci_high"] for r in rows])
plt.errorbar(xs, y, yerr=[y - lo, hi - y], fmt="o-", capsize=4, lw=2, color=color, label=label, zorder=3)
for r in rows:
plt.scatter([r["demos"]] * len(r["per_seed"]), r["per_seed"], color=color, alpha=.3, s=18, zorder=2)
print(f"\n=== {label} ===")
print("Demos pooled 95% CI per-seed mean±std")
for r in rows:
ps = " ".join(f"{v:.1%}" for v in r["per_seed"])
print(f"{r['demos']:5d} {r['pooled_rate']:6.1%} [{r['ci_low']:.1%},{r['ci_high']:.1%}] {ps:22s} {r['mean']:.1%}±{r['std']:.1%}")
plt.xscale("log"); plt.xlabel("number of demonstrations (log scale)"); plt.ylabel("online success rate")
plt.title("Cube-single — proprioception ablation\n3 seeds x 100 rollouts x 5 tasks, 95% Wilson CI")
plt.grid(True, which="both", alpha=.3); plt.legend(); plt.tight_layout()
plt.savefig(a.out, dpi=130); print("\nfigure ->", a.out)
json.dump(allrows, open(a.out.replace(".png", ".json"), "w"), indent=2)
|