| |
| """Multi-seed open-loop test: is the RAPTOR cascade's error bias or sampling noise? |
| Samples S action chunks per anchor (different seeds), scores each sample, the per-anchor mean of |
| samples, and the hold-current-pose baseline, on the anchors of ablation_runs/openloop_compare. |
| python scripts/eval/raptor_multiseed.py <checkpoint> <out_npz> [seeds=8] |
| """ |
| import json, sys, time |
| from pathlib import Path |
| import numpy as np, torch |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| for sub in ("scripts", "scripts/eval"): |
| sys.path.insert(0, str(ROOT / sub)) |
| import trex_ablation_eval as tae |
| import openloop_compare as oc |
| from groot.vla.model.trex_track_force.runtime import delta_base_to_absolute |
| from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA |
|
|
| ckpt, out_npz = Path(sys.argv[1]), Path(sys.argv[2]); S = int(sys.argv[3]) if len(sys.argv) > 3 else 8 |
| A = json.load(open("ablation_runs/openloop_compare/anchors.json")); root = A["dataset_root"]; anchors = [(int(e), float(t)) for e, t in A["anchors"]] |
| device = torch.device("cuda") |
| model = TrexTrackForceVLA.load_lora(str(ckpt)).to(device=device, dtype=torch.bfloat16).eval(); policy = model.action_head |
| transform_eval, _ = tae.load_pipeline(ckpt, training=False) |
| builders = {}; preds, gts, holds = [], [], [] |
| for i, (ep, anchor) in enumerate(anchors): |
| if ep not in builders: builders[ep] = tae.ChunkBuilder(root, ep) |
| b = builders[ep]; gt_eef, _ = oc.ground_truth(b, anchor) |
| raw = b.build(anchor, blocks=1, prompt=oc.episode_task(b), with_deform=False, history_only_video=True) |
| reference = np.asarray(raw["state.eef62"], np.float64).reshape(-1, 62)[0]; raw.pop("action.eef62") |
| collated = transform_eval(tae.unsqueeze_dict_values(dict(raw))); batch = {} |
| for k, v in collated.items(): |
| if torch.is_tensor(v): v = v.to(device=device, dtype=torch.bfloat16) if v.is_floating_point() else v.to(device) |
| batch[k] = v |
| t0 = time.perf_counter(); samples = [] |
| for s in range(S): |
| with torch.inference_mode(): |
| res = policy.sample(batch, seed=500 + i + 1000 * (s + 1), run_force_refinement=True) |
| samples.append(delta_base_to_absolute(reference, b.stats.denormalize_action(res["action_pred"][0, :16, :62].float().cpu().numpy()))) |
| preds.append(np.stack(samples)); gts.append(gt_eef); holds.append(np.repeat(reference[None], 16, 0)) |
| print(f"[{i + 1}/{len(anchors)}] ep{ep} t={anchor:.2f}s {S} samples in {time.perf_counter() - t0:.1f}s", flush=True) |
| P = np.stack(preds); G = np.stack(gts); Hd = np.stack(holds) |
| np.savez(out_npz, pred=P, gt=G, hold=Hd) |
| e_single = oc.errors(P[:, 0].astype(np.float64), G); e_mean = oc.errors(P.mean(1).astype(np.float64), G); e_hold = oc.errors(Hd, G) |
| spread = {m: float(np.mean([oc.errors(P[:, s].astype(np.float64), P.mean(1).astype(np.float64))[m].mean() for s in range(S)])) for m in ("pos", "rot", "hand")} |
| H = np.r_[9:31, 40:62]; Pp = np.r_[0:3, 31:34] |
| cos = lambda a, b: ((a * b).sum(-1) / (np.linalg.norm(a, axis=-1) * np.linalg.norm(b, axis=-1) + 1e-9)) |
| gdh, gdp = G[:, 15, H] - G[:, 0, H], G[:, 15, Pp] - G[:, 0, Pp]; M = P.mean(1) |
| print("\n## Multi-seed test on %s (%d anchors x %d samples)\n" % (ckpt, len(anchors), S)) |
| print("| prediction | EEF pos err [cm] | EEF rot err [deg] | hand RMSE [rad] | hand disp cos | EEF disp cos |"); print("|---|---|---|---|---|---|") |
| for name, e, pr in (("single sample (seed 0)", e_single, P[:, 0]), (f"mean of {S} samples", e_mean, M), ("hold current pose", e_hold, Hd)): |
| dh, dp = pr[:, 15, H] - pr[:, 0, H], pr[:, 15, Pp] - pr[:, 0, Pp] |
| print(f"| {name} | {100 * e['pos'].mean():.2f} | {e['rot'].mean():.2f} | {e['hand'].mean():.4f} | {np.nanmean(cos(dh, gdh)):.2f} | {np.nanmean(cos(dp, gdp)):.2f} |") |
| print(f"\nsample-to-sample spread around the per-anchor mean: pos {100 * spread['pos']:.2f} cm, rot {spread['rot']:.2f} deg, hand {spread['hand']:.4f} rad " |
| f"(compare: GT motion within a chunk {100 * np.linalg.norm(gdp, axis=-1).mean():.2f} cm)") |
|
|
