#!/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()