Upload tinyvla2/scripts/eval_physical_ai.py with huggingface_hub

tinyvla2/scripts/eval_physical_ai.py

@@ -0,0 +1,493 @@
+#!/usr/bin/env python
+"""Evaluate the physical_ai_ft (GR00T) checkpoint — the run that was never measured.
+
+Everything here exists because this checkpoint can produce a plausible-looking
+number for the wrong reason. Three specific traps, each with a gate or control:
+
+1. The published weights carry torch.compile's `expert._orig_mod.` prefix, and
+   LeRobot loads with strict=False while FlowMatchingExpert zero-inits action_out.
+   Loaded naively the model returns its own initial Gaussian noise. Gate: strict
+   load of the rekeyed copy; Control B re-runs the broken path as a labelled row.
+
+2. Action semantics are per-DIM, not per-family: r1_pro's `action[3:6]` IS
+   `state[6:9]`, gr1 mixes 12 absolute dims with deltas, and 17 of gr1's 44 dims
+   are constant. So the trivial baseline is searched per dim, in raw units, not
+   assumed. Getting this wrong on r1_pro (hold-still 0.895 vs predict-mean 3.79)
+   would manufacture a 4x win out of nothing.
+
+3. The realized training mixture is ~97% r1_pro, because the stream samples a
+   dataset per EPISODE and r1_pro episodes are ~37x longer. Per-family numbers are
+   printed next to that table or they will be misread.
+
+Usage:
+    python scripts/eval_physical_ai.py [--preflight] [--datasets-per-family 8]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+CKPT = "/home/alexw/tinyvla/outputs/physical_ai_ft_fixed"
+RAW_CKPT = "/home/alexw/tinyvla_data/b200/outputs/physical_ai_ft/final"
+SPEC = "/home/alexw/tinyvla/configs/physical_ai_stream.yaml"
+FAMILY = {10: "gr1", 11: "bi_panda_grip", 12: "bi_panda_hand",
+          13: "single_panda", 14: "r1_pro", 15: "g1"}
+DEAD_STD = 1e-4
+EVAL_SEED = 20260830  # deliberately not the training seed (42)
+
+
+# ----------------------------------------------------------------- calibration
+
+def classify_dims(state_raw, action_raw):
+    """Per action dim: ABS (tracks a state dim), DELTA, or DEAD (constant).
+
+    Searched in RAW units and over lags 0..2, because the z-score hides exact
+    identities and commanded targets lead the measured state.
+    """
+    T, A = action_raw.shape
+    S = state_raw.shape[1]
+    std_a = action_raw.std(0)
+    kind = np.full(A, "delta", dtype=object)
+    amap = np.full(A, -1, dtype=int)
+    lag = np.zeros(A, dtype=int)
+
+    dead = std_a < DEAD_STD
+    kind[dead] = "dead"
+
+    for j in range(A):
+        if dead[j]:
+            continue
+        best = (np.inf, -1, 0)
+        for lg in (0, 1, 2):
+            a = action_raw[: T - lg, j]
+            sl = state_raw[lg:, :]
+            rms = np.sqrt(((sl - a[:, None]) ** 2).mean(0))
+            i = int(np.argmin(rms))
+            if rms[i] < best[0]:
+                best = (float(rms[i]), i, lg)
+        rms, i, lg = best
+        if i < 0:
+            continue
+        a = action_raw[: T - lg, j]
+        s = state_raw[lg:, i]
+        if s.std() < 1e-9 or a.std() < 1e-9:
+            continue
+        corr = float(np.corrcoef(a, s)[0, 1])
+        if corr > 0.95 and rms < 0.5 * std_a[j]:
+            kind[j], amap[j], lag[j] = "abs", i, lg
+    return kind, amap, lag, std_a
+
+
+def displacement(A, kind, amap, ref):
+    """(chunk, adim) actions -> displacement, per the dim's semantics.
+
+    ABS   : residual against the hold-still reference  (trivial baseline = hold still)
+    DELTA : integrated                                  (trivial baseline = zero motion)
+    In both cases the trivial baseline maps to d == 0, which is what makes the
+    ratio comparable across families.
+    """
+    d = np.array(A, dtype=np.float64, copy=True)
+    is_abs = kind == "abs"
+    is_del = kind == "delta"
+    d[:, is_abs] = A[:, is_abs] - ref[is_abs][None, :]
+    d[:, is_del] = np.cumsum(A[:, is_del], axis=0)
+    return d
+
+
+def chunk_err(pred, gt, kind, amap, ref, sel):
+    """Mean-over-steps L2 of the displacement error, restricted to dims `sel`."""
+    dp = displacement(pred, kind, amap, ref)[:, sel]
+    dg = displacement(gt, kind, amap, ref)[:, sel]
+    return float(np.linalg.norm(dp - dg, axis=1).mean()), float(np.linalg.norm(dg, axis=1).mean())
+
+
+# --------------------------------------------------------------------- helpers
+
+def make_batch(items, tok, cfg, device="cuda", state_override=None, image_override=None,
+               task_override=None, emb_override=None, padding="max_length"):
+    tasks = task_override if task_override is not None else [it["task"] for it in items]
+    t = tok(tasks, padding=padding, truncation=True, max_length=48, return_tensors="pt")
+    cam0 = torch.stack([it["observation.images.cam0"] for it in items])
+    cam1 = torch.stack([it["observation.images.cam1"] for it in items])
+    if image_override is not None:
+        cam0, cam1 = image_override
+    st = torch.stack([it["observation.state"] for it in items])
+    if state_override is not None:
+        st = state_override
+    emb = torch.stack([it["embodiment_id"] for it in items])
+    if emb_override is not None:
+        emb = emb_override
+    return {
+        "observation.images.cam0": cam0.to(device),
+        "observation.images.cam1": cam1.to(device),
+        "observation.state": st.to(device),
+        "observation.language.tokens": t["input_ids"].to(device),
+        "observation.language.attention_mask": t["attention_mask"].bool().to(device),
+        "embodiment_id": emb.to(device),
+    }
+
+
+def predict_mean(pol, batch, n_seeds: int):
+    """Average the flow ODE over n_seeds draws of the initial noise.
+
+    predict_action_chunk starts from torch.randn, and control C measured seed
+    sensitivity S ~ 1.2 on this checkpoint: two draws disagree with each other
+    about as much as either disagrees with ground truth. The expert itself is fine
+    (flow loss L/L0 = 0.12) — the 10-step Euler integration simply does not
+    contract the initial noise. Since the metric is L2, the conditional MEAN is the
+    optimal point estimate, so averaging is the right estimator, not a fudge.
+    """
+    acc = None
+    for k in range(n_seeds):
+        torch.manual_seed(1000 + k)
+        with torch.autocast("cuda", torch.bfloat16):
+            p = pol.predict_action_chunk(batch).float()
+        acc = p if acc is None else acc + p
+    return (acc / n_seeds).cpu().numpy()
+
+
+def bootstrap_ci(per_ep, n=2000, seed=0):
+    """CI over EPISODES — frames inside an episode are not independent."""
+    if len(per_ep) < 2:
+        return (float("nan"), float("nan"))
+    rng = np.random.default_rng(seed)
+    a = np.asarray(per_ep, dtype=float)
+    means = [a[rng.integers(0, len(a), len(a))].mean() for _ in range(n)]
+    return (float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5)))
+
+
+def realized_mixture(specs, metas):
+    """Sample share actually seen, given the stream draws a DATASET PER EPISODE."""
+    rows = []
+    for si, s in enumerate(specs):
+        m = metas.get(si)
+        if m is None:
+            continue
+        rows.append((s["embodiment_id"], s["weight"], m["samples_per_ep"], m["n_eps"]))
+    tot_w = sum(r[1] for r in rows) or 1.0
+    fam = defaultdict(lambda: [0.0, 0.0, 0])
+    for emb, w, spe, n_eps in rows:
+        share = (w / tot_w) * spe
+        fam[emb][0] += share
+        fam[emb][1] += w / tot_w
+        fam[emb][2] += n_eps
+    norm = sum(v[0] for v in fam.values()) or 1.0
+    return {k: (v[0] / norm, v[1], v[2]) for k, v in fam.items()}
+
+
+# ------------------------------------------------------------------------ main
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--preflight", action="store_true", help="gates and tables only")
+    ap.add_argument("--datasets-per-family", type=int, default=8)
+    ap.add_argument("--episodes", type=int, default=3)
+    ap.add_argument("--starts", type=int, default=12)
+    ap.add_argument("--seeds", type=int, default=4,
+                    help="average the flow ODE over N noise draws (control C measured S~1.2)")
+    ap.add_argument("--out", default="/home/alexw/tinyvla/outputs/physical_ai_eval.json")
+    args = ap.parse_args()
+
+    from safetensors.torch import load_file
+    from transformers import AutoTokenizer
+
+    from tinyvla.data.hub_stream import EvalEpisodeStream
+    from tinyvla.modeling_tinyvla import TinyVLAPolicy
+
+    specs_all = yaml.safe_load(open(SPEC))["datasets"]
+
+    # ---------------------------------------------------------------- GATE 1
+    pol = TinyVLAPolicy.from_pretrained(CKPT)
+    sd = load_file(f"{CKPT}/model.safetensors")
+    pol.load_state_dict(sd, strict=True)
+    aow = pol.expert.action_out.weight.abs().mean().item()
+    assert aow > 1e-3, "expert.action_out at zero init — expert never trained"
+    cfg = pol.config
+    pol = pol.cuda().eval()
+    print(f"GATE  rekey {len(sd)}/{len(pol.state_dict())} strict=True OK | "
+          f"action_out |w|={aow:.4f} | action_dim={pol.action_dim} "
+          f"max_action={cfg.max_action_dim} max_state={cfg.max_state_dim}")
+    tok = AutoTokenizer.from_pretrained(cfg.lm_model_name)
+
+    # ------------------------------------------------- dataset selection
+    rng = np.random.default_rng(EVAL_SEED)
+    by_fam = defaultdict(list)
+    for s in specs_all:
+        by_fam[s["embodiment_id"]].append(s)
+    chosen = []
+    for emb in sorted(by_fam):
+        pool = by_fam[emb]
+        k = min(args.datasets_per_family if emb != 14 else 12, len(pool))
+        idx = rng.choice(len(pool), size=k, replace=False)
+        chosen += [pool[i] for i in sorted(idx)]
+    print(f"selected {len(chosen)} datasets across {len(by_fam)} families "
+          f"(seed {EVAL_SEED}, not the training seed)")
+
+    stream = EvalEpisodeStream(chosen, chunk=cfg.chunk_size, image_size=cfg.image_size,
+                               max_state_dim=cfg.max_state_dim,
+                               max_action_dim=cfg.max_action_dim, shuffle_buffer=1)
+
+    # ------------------------------------------------- GATE 2/3 + calibration
+    metas, calib = {}, {}
+    for si, s in enumerate(chosen):
+        try:
+            m = stream._load_meta(si)
+            eps = m["episodes"]
+            metas[si] = {"n_eps": len(eps),
+                         "samples_per_ep": float(np.mean([e[1] for e in eps[:200]])) / m["stride"],
+                         "stride": m["stride"], "adim": m["action_dim"],
+                         "stats_src": m.get("stats_src", "?")}
+            sraw, araw, _ = stream.episode_raw(si)
+            kind, amap, lag, std_a = classify_dims(sraw, araw)
+            calib[si] = dict(kind=kind, amap=amap, lag=lag, std_a=std_a)
+            print(f"  [{si:>3}] {s['prefix'][:44] or s['repo_id'][-30:]:46} "
+                  f"emb={s['embodiment_id']} adim={m['action_dim']:>3} "
+                  f"abs={int((kind=='abs').sum()):>3} del={int((kind=='delta').sum()):>3} "
+                  f"dead={int((kind=='dead').sum()):>3}", flush=True)
+        except Exception as e:
+            print(f"  [{si:>3}] {s.get('prefix','')[:44]:46} CALIB FAIL {type(e).__name__}: {str(e)[:60]}",
+                  flush=True)
+
+    mix = realized_mixture(chosen, metas)
+    print("\nREALIZED MIXTURE (stream draws a dataset PER EPISODE, so share = weight x episode length)")
+    for emb in sorted(mix):
+        share, wshare, n_eps = mix[emb]
+        print(f"  {FAMILY[emb]:14} realized {share*100:5.1f}%   intended(weight) {wshare*100:5.1f}%   "
+              f"episodes {n_eps}")
+    draws_8gpu = 8 * 256 * 60000 / max(1e-9, np.mean([m["samples_per_ep"] for m in metas.values()]))
+    print("HOLD-OUT  P(episode never drawn), 8-GPU / 1-GPU scenario (no log disambiguates):")
+    for emb in sorted(mix):
+        share, _, n_eps = mix[emb]
+        lam8 = draws_8gpu * share / max(1, n_eps)
+        print(f"  {FAMILY[emb]:14} P(unseen) {math.exp(-lam8)*100:5.1f}% / {math.exp(-lam8/8)*100:5.1f}%")
+
+    if args.preflight:
+        return
+
+    # ------------------------------------------------------------ main pass
+    print("\n=== main pass ===", flush=True)
+    results = defaultdict(lambda: defaultdict(list))
+    per_ep_json = []
+    subsample = []  # kept for the control block
+
+    for si, s in enumerate(chosen):
+        if si not in calib:
+            continue
+        c = calib[si]
+        emb = s["embodiment_id"]
+        m = stream._load_meta(si)
+        eps = m["episodes"]
+        n_ep = 1 if emb == 14 else args.episodes
+        n_st = 24 if emb == 14 else args.starts
+        pick = rng.choice(len(eps), size=min(n_ep, len(eps)), replace=False)
+        for e_i in pick:
+            ep_idx, length, task = eps[e_i]
+            n_rows = length // m["stride"]
+            hi = n_rows - cfg.chunk_size - 1
+            if hi <= 0:
+                continue
+            if emb == 14:  # prefix-decode: starts confined to the episode head (labelled bias)
+                hi = min(hi, 600)
+            starts = np.unique(np.linspace(0, hi, n_st).astype(int))
+            try:
+                items = stream.episode_samples(si, ep_idx, task, starts)
+            except Exception as ex:
+                print(f"  ep fail {s.get('prefix','')[:36]} {type(ex).__name__}: {str(ex)[:60]}", flush=True)
+                continue
+            if not items:
+                continue
+
+            preds = []
+            for k in range(0, len(items), 8):
+                b = make_batch(items[k:k + 8], tok, cfg)
+                preds.append(predict_mean(pol, b, args.seeds))
+            pred = np.concatenate(preds, 0)
+
+            ad = m["action_dim"]
+            sm, ss = m["stats"]["observation.state"]
+            am, as_ = m["stats"]["action"]
+            kind, amap = c["kind"], c["amap"]
+            for subset in ("abs", "delta"):
+                sel = (kind == subset)
+                if not sel.any():
+                    continue
+                e_l, f0_l, f1_l, f2_l = [], [], [], []
+                for t_i, it in enumerate(items):
+                    A = it["action"][:, :ad].numpy().astype(np.float64)
+                    P = pred[t_i][:, :ad].astype(np.float64)
+                    s_norm = it["observation.state"].numpy()
+                    s_raw = s_norm[:len(sm)] * np.maximum(ss, 1e-6) + sm
+                    ref = np.zeros(ad)
+                    for j in range(ad):
+                        if kind[j] == "abs" and amap[j] >= 0 and amap[j] < len(s_raw):
+                            ref[j] = (s_raw[amap[j]] - am[j]) / max(as_[j], 1e-6)
+                    err, floor_zero = chunk_err(P, A, kind, amap, ref, sel)
+                    b0, _ = chunk_err(np.zeros_like(A), A, kind, amap, ref, sel)
+                    b1, _ = chunk_err(np.tile(ref, (len(A), 1)), A, kind, amap, ref, sel)
+                    b2, _ = chunk_err(np.tile(A[0], (len(A), 1)), A, kind, amap, ref, sel)
+                    e_l.append(err); f0_l.append(b0); f1_l.append(b1); f2_l.append(b2)
+                key = (emb, subset)
+                results[key]["err"].append(np.mean(e_l))
+                results[key]["b0"].append(np.mean(f0_l))
+                results[key]["b1"].append(np.mean(f1_l))
+                results[key]["b2"].append(np.mean(f2_l))
+                results[key]["ndim"].append(int(sel.sum()))
+                per_ep_json.append(dict(emb=int(emb), subset=subset, dataset=s.get("prefix", s["repo_id"]),
+                                        episode=int(ep_idx), err=float(np.mean(e_l)),
+                                        b0=float(np.mean(f0_l)), b1=float(np.mean(f1_l)),
+                                        b2=float(np.mean(f2_l)), ndim=int(sel.sum())))
+            if len(subsample) < 24:
+                subsample.append((si, items[:4], c, m))
+            print(f"  {FAMILY[emb]:14} {s.get('prefix','')[:34]:36} ep{ep_idx:<6} "
+                  f"n={len(items)}", flush=True)
+
+    # ------------------------------------------------------------- report
+    print("\n=== RESULTS (displacement error / trivial floor; <1 beats the floor) ===")
+    print(f"{'family':16}{'subset':7}{'nep':>4}{'dims':>5}{'err':>8}{'B0':>8}{'B1':>8}{'B2':>8}"
+          f"{'ratio':>8}   95% CI (bootstrap over episodes)")
+    summary = {}
+    for (emb, subset), v in sorted(results.items()):
+        err = np.mean(v["err"]); b0 = np.mean(v["b0"]); b1 = np.mean(v["b1"]); b2 = np.mean(v["b2"])
+        floor = min(b0, b1)
+        ratios = [e / max(min(x, y), 1e-9) for e, x, y in zip(v["err"], v["b0"], v["b1"])]
+        lo, hi = bootstrap_ci(ratios)
+        tag = "  [IN-TRAINING]" if emb == 14 else ""
+        print(f"{FAMILY[emb]:16}{subset:7}{len(v['err']):>4}{int(np.mean(v['ndim'])):>5}"
+              f"{err:>8.3f}{b0:>8.3f}{b1:>8.3f}{b2:>8.3f}{err/floor:>8.2f}   "
+              f"[{lo:.2f}, {hi:.2f}]{tag}")
+        summary[f"{FAMILY[emb]}/{subset}"] = dict(err=err, b0=b0, b1=b1, b2=b2,
+                                                  ratio=err / floor, ci=[lo, hi],
+                                                  n_ep=len(v["err"]))
+        if abs(b2 - floor) / max(floor, 1e-9) < 0.1:
+            print(f"{'':16}^^ oracle repeat-A0 ({b2:.3f}) is within 10% of the floor "
+                  f"({floor:.3f}) — this subset carries almost no predictable signal")
+
+    # ------------------------------------------------------------ controls
+    print("\n=== CONTROLS ===", flush=True)
+    ctrl = {}
+
+    # C: ODE seed sensitivity
+    si, items, c, m = subsample[0]
+    b = make_batch(items, tok, cfg)
+    torch.manual_seed(1)
+    with torch.autocast("cuda", torch.bfloat16):
+        p1 = pol.predict_action_chunk(b).float().cpu().numpy()
+    torch.manual_seed(2)
+    with torch.autocast("cuda", torch.bfloat16):
+        p2 = pol.predict_action_chunk(b).float().cpu().numpy()
+    gt = torch.stack([it["action"] for it in items]).numpy()[:, :, :p1.shape[-1]]
+    S = float(np.abs(p1 - p2).mean() / max(np.abs(p1 - gt).mean(), 1e-9))
+    a1 = predict_mean(pol, b, args.seeds)
+    torch.manual_seed(77)
+    a2 = predict_mean(pol, b, args.seeds)
+    S_avg = float(np.abs(a1 - a2).mean() / max(np.abs(a1 - gt).mean(), 1e-9))
+    ctrl["ode_seed_sensitivity"] = S
+    ctrl["ode_seed_sensitivity_averaged"] = S_avg
+    print(f"  C  ODE seed sensitivity: single draw S = {S:.3f}, "
+          f"{args.seeds}-seed average S = {S_avg:.3f}   (headline uses the average)")
+
+    # E: flow-matching training loss vs the trivial predictor
+    e_rows = {}
+    for si, items, c, m in subsample[:12]:
+        emb = chosen[si]["embodiment_id"]
+        b = make_batch(items, tok, cfg)
+        b["action"] = torch.stack([it["action"] for it in items]).cuda()
+        b["action_dim_mask"] = torch.stack([it["action_dim_mask"] for it in items]).cuda()
+        b["action_is_pad"] = torch.stack([it["action_is_pad"] for it in items]).cuda()
+        torch.manual_seed(0)
+        with torch.autocast("cuda", torch.bfloat16):
+            loss, _ = pol.forward(b)
+        A = b["action"][:, :, : m["action_dim"]]
+        l0 = float((torch.randn_like(A) - A).pow(2).mean())
+        e_rows.setdefault(emb, []).append(float(loss) / l0)
+    ctrl["flow_loss_ratio"] = {FAMILY[k]: float(np.mean(v)) for k, v in e_rows.items()}
+    print("  E  flow loss L/L0 :", {k: round(v, 3) for k, v in ctrl["flow_loss_ratio"].items()},
+          "  (<1 = expert learned something, independent of the ODE and the metric)")
+
+    # D3/D4: in-distribution donor swap — is vision used at all, or is this a proprio regressor?
+    def ratio_for(items, c, m, state_override=None, image_override=None, emb_override=None):
+        b = make_batch(items, tok, cfg, state_override=state_override,
+                       image_override=image_override, emb_override=emb_override)
+        pr = predict_mean(pol, b, args.seeds)
+        ad = m["action_dim"]
+        sm, ss = m["stats"]["observation.state"]
+        am, as_ = m["stats"]["action"]
+        kind, amap = c["kind"], c["amap"]
+        sel = (kind != "dead")
+        num, den = [], []
+        for t_i, it in enumerate(items):
+            A = it["action"][:, :ad].numpy().astype(np.float64)
+            P = pr[t_i][:, :ad].astype(np.float64)
+            s_raw = it["observation.state"].numpy()[:len(sm)] * np.maximum(ss, 1e-6) + sm
+            ref = np.zeros(ad)
+            for j in range(ad):
+                if kind[j] == "abs" and 0 <= amap[j] < len(s_raw):
+                    ref[j] = (s_raw[amap[j]] - am[j]) / max(as_[j], 1e-6)
+            e, f = chunk_err(P, A, kind, amap, ref, sel)
+            b0, _ = chunk_err(np.zeros_like(A), A, kind, amap, ref, sel)
+            b1, _ = chunk_err(np.tile(ref, (len(A), 1)), A, kind, amap, ref, sel)
+            num.append(e); den.append(min(b0, b1))
+        return float(np.mean(num) / max(np.mean(den), 1e-9))
+
+    d_rows = defaultdict(dict)
+    for pos, (si, items, c, m) in enumerate(subsample[:8]):
+        emb = chosen[si]["embodiment_id"]
+        donor = subsample[(pos + 1) % len(subsample)][1]
+        n = min(len(items), len(donor))
+        it_, dn_ = items[:n], donor[:n]
+        base = ratio_for(it_, c, m)
+        img = (torch.stack([d["observation.images.cam0"] for d in dn_]),
+               torch.stack([d["observation.images.cam1"] for d in dn_]))
+        d3 = ratio_for(it_, c, m, image_override=img)
+        d4 = ratio_for(it_, c, m, state_override=torch.stack([d["observation.state"] for d in dn_]))
+        d6 = ratio_for(it_, c, m, emb_override=torch.full((n,), 11 if emb != 11 else 10).long())
+        d_rows[emb] = dict(base=base, d3_img=d3, d4_state=d4, d6_emb=d6)
+    ctrl["ablation"] = {FAMILY[k]: {kk: round(vv, 3) for kk, vv in v.items()} for k, v in d_rows.items()}
+    print("  D  donor-swap ablation (ratio; higher = input mattered):")
+    for k, v in ctrl["ablation"].items():
+        print(f"       {k:14} base {v['base']:.3f} | images {v['d3_img']:.3f} "
+              f"| state {v['d4_state']:.3f} | wrong emb id {v['d6_emb']:.3f}")
+
+    # B: the broken (un-rekeyed) path, as shipped
+    broken = TinyVLAPolicy.from_pretrained(RAW_CKPT).cuda().eval()
+    si, items, c, m = subsample[0]
+    b = make_batch(items, tok, cfg)
+    with torch.autocast("cuda", torch.bfloat16):
+        pb = broken.predict_action_chunk(b).float().cpu().numpy()
+    ad = m["action_dim"]
+    A = torch.stack([it["action"] for it in items]).numpy()[:, :, :ad]
+    ctrl["broken_err"] = float(np.linalg.norm(pb[:, :, :ad] - A, axis=2).mean())
+    with torch.autocast("cuda", torch.bfloat16):
+        pg = pol.predict_action_chunk(b).float().cpu().numpy()
+    ctrl["fixed_err"] = float(np.linalg.norm(pg[:, :, :ad] - A, axis=2).mean())
+    print(f"  B  un-rekeyed (as published) err {ctrl['broken_err']:.3f} vs rekeyed "
+          f"{ctrl['fixed_err']:.3f}   (equal => the rename did not take effect)")
+    del broken
+    torch.cuda.empty_cache()
+
+    # F: tokenizer padding fidelity (training used max_length=48)
+    si, items, c, m = subsample[0]
+    r_max = ratio_for(items, c, m)
+    b_long = make_batch(items, tok, cfg, padding=True)
+    with torch.autocast("cuda", torch.bfloat16):
+        _ = pol.predict_action_chunk(b_long)
+    ctrl["padding_note"] = "headline uses padding=max_length(48), as train_fast.Collate did"
+    print(f"  F  padding: headline uses max_length=48 (training convention); ratio {r_max:.3f}")
+
+    Path(args.out).write_text(json.dumps(
+        {"summary": summary, "controls": ctrl, "per_episode": per_ep_json,
+         "mixture": {FAMILY[k]: v[0] for k, v in mix.items()}}, indent=2))
+    print(f"\nwrote {args.out}")
+
+
+if __name__ == "__main__":
+    main()