v3/analyze_release.py

#!/usr/bin/env python3
"""Build the release summary and figure from raw Speedup Patch audit outputs."""

from __future__ import annotations

import csv
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


ROOT = Path(__file__).resolve().parent
RESULTS = ROOT / "results"
campaign = json.loads((RESULTS / "libero_adm_summary.json").read_text())
traces = json.loads((RESULTS / "libero_adm_traces.json").read_text())
paper = json.loads((RESULTS / "paper_artifact_audit.json").read_text())
rows = list(csv.DictReader((RESULTS / "libero_adm_runs.csv").open()))
for row in rows:
    for k in (
        "trainable_parameters",
        "state_mse",
        "state_rmse",
        "eef_rmse",
        "displacement_spearman",
        "samples",
        "seed",
    ):
        row[k] = float(row[k])

adm = campaign["mean_results"]["official_adm_ar_model"]
mlp = campaign["mean_results"]["parameter_matched_mlp"]

long_horizon = {}
for h in range(7, 11):
    long_horizon[str(h)] = {}
    for model in ("official_adm_ar_model", "parameter_matched_mlp"):
        vals = [
            traces[f"{model}_seed_{seed}"]["by_horizon"][h - 1]["eef_rmse"]
            for seed in (42, 43, 44)
        ]
        long_horizon[str(h)][model] = float(np.mean(vals))

derived = {
    "official_adm_real_libero": {
        "state_mse_relative_reduction_percent": 100 * (mlp["state_mse"] - adm["state_mse"]) / mlp["state_mse"],
        "eef_rmse_relative_reduction_percent": 100 * (mlp["eef_rmse"] - adm["eef_rmse"]) / mlp["eef_rmse"],
        "spearman_absolute_gain": adm["displacement_spearman"] - mlp["displacement_spearman"],
        "adm_wins_state_mse_seeds": sum(
            r["state_mse"]
            < next(
                q["state_mse"]
                for q in rows
                if q["seed"] == r["seed"] and q["model"] == "parameter_matched_mlp"
            )
            for r in rows
            if r["model"] == "official_adm_ar_model"
        ),
        "adm_wins_eef_rmse_long_horizons": sum(
            v["official_adm_ar_model"] < v["parameter_matched_mlp"]
            for v in long_horizon.values()
        ),
        "long_horizon_means": long_horizon,
    },
    "paper_artifact_findings": {
        "bigym_success_recomputes": abs(
            paper["claim_2_bigym"]["recomputed_sup_success"] - .6725
        )
        < 1e-12,
        "libero_pi05_registered_values_recompute": (
            abs(paper["claim_3_libero"]["pi05_sup_success"] - .973) < 1e-12
            and round(paper["claim_3_libero"]["pi05_sup_speedup"], 2) == 1.35
        ),
        "libero_vla_registered_values_recompute": (
            abs(paper["claim_3_libero"]["vla_sup_success"] - .9365) < 1e-12
            and round(paper["claim_3_libero"]["vla_sup_speedup"], 2) == 1.34
        ),
        "real_sup_registered_values_recompute": (
            round(paper["claim_4_real"]["sup_success"], 3) == .611
            and round(paper["claim_4_real"]["sup_speedup_from_displayed_steps"], 2)
            == 2.17
        ),
        "figure7_adm_stronger_all_suites": (
            paper["claim_6_figure7"]["adm_stronger_negative_suites"] == 4
        ),
    },
}
(RESULTS / "derived_release_summary.json").write_text(
    json.dumps(derived, indent=2, sort_keys=True) + "\n"
)

plt.rcParams.update({"font.size": 11, "axes.titlesize": 13})
fig, axes = plt.subplots(2, 2, figsize=(15, 9))

# Panel A: paired per-seed state MSE.
ax = axes[0, 0]
for seed in (42, 43, 44):
    a = next(r for r in rows if r["seed"] == seed and r["model"] == "official_adm_ar_model")
    m = next(r for r in rows if r["seed"] == seed and r["model"] == "parameter_matched_mlp")
    ax.plot(["MLP", "official ADM"], [m["state_mse"] * 1e4, a["state_mse"] * 1e4], marker="o", label=f"seed {seed}")
ax.set_ylabel("test state MSE × 10⁴ (lower is better)")
ax.set_title("A. Real LIBERO any-step prediction")
ax.legend(frameon=False)
ax.grid(axis="y", alpha=.25)

# Panel B: horizon curves.
ax = axes[0, 1]
for model, label, color in (
    ("official_adm_ar_model", "official ADM", "#1675a9"),
    ("parameter_matched_mlp", "parameter-matched MLP", "#c6651a"),
):
    means, stds = [], []
    for h in range(1, 11):
        vals = [
            traces[f"{model}_seed_{seed}"]["by_horizon"][h - 1]["eef_rmse"]
            for seed in (42, 43, 44)
        ]
        means.append(np.mean(vals))
        stds.append(np.std(vals))
    x = np.arange(1, 11)
    ax.plot(x, means, marker="o", label=label, color=color)
    ax.fill_between(x, np.asarray(means) - stds, np.asarray(means) + stds, alpha=.16, color=color)
ax.axvspan(6.5, 10.5, color="#eef2e5", zorder=-1, label="long-horizon audit")
ax.set_xlabel("action-sequence horizon")
ax.set_ylabel("EEF RMSE (lower is better)")
ax.set_title("B. ADM advantage emerges at horizons 7–10")
ax.legend(frameon=False)
ax.grid(alpha=.25)

# Panel C: exact Figure 7 annotations.
ax = axes[1, 0]
f7 = paper["claim_6_figure7"]["values_read_from_figure_annotations"]
suites = ["object", "spatial", "long", "goal"]
x = np.arange(4)
ax.bar(x - .18, [f7[s]["adm"] for s in suites], width=.36, label="ADM", color="#1675a9")
ax.bar(x + .18, [f7[s]["mlp"] for s in suites], width=.36, label="MLP", color="#c6651a")
ax.set_xticks(x, [s.title() for s in suites])
ax.set_ylabel("Spearman correlation (more negative is stronger)")
ax.set_title("C. Exact annotated values from paper Figure 7")
ax.legend(frameon=False)
ax.grid(axis="y", alpha=.25)

# Panel D: registered-value recomputation.
ax = axes[1, 1]
ax.axis("off")
lines = [
    ("BiGym ACT + SuP", "success 0.6725 → 0.67", "speed 2.01 source-only"),
    ("LIBERO π0.5 + SuP", "success 0.973", f"speed {paper['claim_3_libero']['pi05_sup_speedup']:.4f} → 1.35"),
    ("LIBERO VLA + SuP", f"success {paper['claim_3_libero']['vla_sup_success']:.4f} → 0.937", f"speed {paper['claim_3_libero']['vla_sup_speedup']:.4f} → 1.34"),
    ("Real robot SuP", f"success {paper['claim_4_real']['sup_success']:.4f} → 0.611", f"speed {paper['claim_4_real']['sup_speedup_from_displayed_steps']:.4f} → 2.17"),
    ("Efficiency", "5.12M vs 4B", "781.25× fewer trainable params"),
]
ax.text(.02, .96, "D. Exact source-artifact audit", fontsize=13, fontweight="bold", va="top")
y = .82
for title, left, right in lines:
    ax.text(.02, y, title, fontweight="bold", fontsize=10)
    ax.text(.40, y, left, fontsize=9.5)
    ax.text(.76, y, right, fontsize=9.5)
    y -= .15
ax.text(
    .02,
    .03,
    "Corrections: 5.12M is total SuP trainable parameters in Table 3, not scheduler-only.\n"
    "Displayed rounded BiGym lengths do not uniquely recover the reported 2.01×.",
    fontsize=10,
    color="#7b2f20",
)

fig.suptitle(
    "Speedup Patch v3 audit: official ADM on real LIBERO + exact source arithmetic",
    fontsize=17,
    fontweight="bold",
)
fig.tight_layout(rect=(0, 0, 1, .96))
fig.savefig(RESULTS / "speedup_patch_v3.png", dpi=180)
fig.savefig(RESULTS / "speedup_patch_v3.svg")
print(json.dumps(derived, indent=2, sort_keys=True))