#!/usr/bin/env python """Physical-space canonical-action eval (variants A/B/C, TinyVLA-2 headline metric). Compares A/B/C in ONE fair space: canonical base-frame EE deltas, unnormalized to physical units — position error (mm), rotation error (deg), gripper error ([0,1]). - B/C predict canonical directly → unnormalize with dataset canonical stats. - A predicts NATIVE actions → for SO101 map joint predictions through FK to EE deltas; for EE-native sources A's native deltas are already comparable (convention aside). - Held-out embodiment: A/B have no ID row → oracle over trained IDs (best), making any C win conservative. C uses the written descriptor (zero-shot by construction). Usage: python scripts/eval_canonical.py --checkpoint outputs/tv2_C_morph_canon/final \ --dataset heldout_jaco_play --root ~/tinyvla_data/heldout/heldout_jaco_play \ --morph-key jaco --episodes 20 """ from __future__ import annotations import argparse from pathlib import Path import numpy as np import torch @torch.no_grad() def main(): ap = argparse.ArgumentParser() ap.add_argument("--checkpoint", type=Path, required=True) ap.add_argument("--dataset", required=True) ap.add_argument("--root", required=True) ap.add_argument("--morph-key", default=None, help="descriptor key for variant C zero-shot") ap.add_argument("--oracle-ids", type=int, default=8, help="A/B: try IDs 0..N-1, report best") ap.add_argument("--episodes", type=int, default=20) ap.add_argument("--ep-start", type=int, default=0, help="first episode index (few-shot: eval on held-out test split after FT episodes)") ap.add_argument("--stride", type=int, default=30) args = ap.parse_args() import yaml from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata from scipy.spatial.transform import Rotation from transformers import AutoTokenizer from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize from tinyvla.modeling_tinyvla import TinyVLAPolicy from tinyvla.modules.embodiment import MORPH_FIELDS policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() cfg = policy.config tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) chunk = cfg.chunk_size meta = LeRobotDatasetMetadata(args.dataset, root=args.root) ds = LeRobotDataset(args.dataset, root=args.root, delta_timestamps={"action": [t / meta.fps for t in range(chunk)]}, video_backend="torchcodec") store = CanonicalChunkStore(args.dataset, src_fps=ds.fps, chunk=chunk) stats = store.compute_stats() q01, q99 = np.asarray(stats["q01"]), np.asarray(stats["q99"]) # morphology descriptor (variant C) morph = None if cfg.conditioning == "morph" and args.morph_key: raw = yaml.safe_load(open("configs/morphology/descriptors.yaml"))[args.morph_key] _sc = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3, "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416, "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2} morph = torch.tensor([raw.get(f, 0) * _sc.get(f, 1) for f in MORPH_FIELDS], dtype=torch.float32) image_keys = sorted(k for k in ds.meta.features if k.startswith("observation.images")) def to_batch(item, emb_id, morph_vec): import re from tinyvla.data.mixture import PRIMARY_RE, WRIST_RE prim = next((k for k in image_keys if PRIMARY_RE.search(k)), image_keys[0]) wr = next((k for k in image_keys if k != prim and WRIST_RE.search(k)), None) def img(x): if x.shape[-2:] != (cfg.image_size, cfg.image_size): x = torch.nn.functional.interpolate(x[None], size=(cfg.image_size, cfg.image_size), mode="bilinear", align_corners=False)[0] return x cam0 = img(item[prim]) cam1 = img(item[wr]) if wr else torch.zeros_like(cam0) state = item["observation.state"].float() state = torch.nn.functional.pad(state, (0, cfg.max_state_dim - state.shape[-1])) t = tok([item.get("task") or ""], padding=True, truncation=True, max_length=cfg.tokenizer_max_length, return_tensors="pt") b = {"observation.images.cam0": cam0[None].cuda(), "observation.images.cam1": cam1[None].cuda(), "observation.state": state[None].cuda(), "observation.language.tokens": t["input_ids"].cuda(), "observation.language.attention_mask": t["attention_mask"].bool().cuda(), "embodiment_id": torch.tensor([emb_id], device="cuda")} if morph_vec is not None: b["morphology"] = morph_vec[None].cuda() return b # native action stats (for variant A unnormalization) — physical EE deltas native_stats = ds.meta.stats.get("action", {}) nat_mean = np.asarray(native_stats.get("mean", np.zeros(7))) nat_std = np.asarray(native_stats.get("std", np.ones(7))) is_native = cfg.action_space == "native" def _unnorm_canon(x): span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) mid = 0.5 * (q01 + q99) return x[:, :7] * span / 2 + mid def _integrate(deltas): """(T,6+) physical per-step EE deltas -> (T,3) cumulative positions, list of cumulative rotations. Amplifies per-step differences into a trajectory with real dynamic range (per-step motion is tiny).""" pos = np.cumsum(deltas[:, :3], axis=0) R = Rotation.identity() rots = [] for k in range(len(deltas)): R = Rotation.from_rotvec(deltas[k, 3:6]) * R rots.append(R) return pos, rots def canonical_phys_err(pred_norm, gt_canon): """Integrated-trajectory error (the per-step delta floor ~3.6mm has no dynamic range — see control). Returns: ep_mm : endpoint position error over the 5s chunk (mm) path_mm: mean cumulative-position error along the chunk (mm) ep_deg : endpoint cumulative-rotation error (deg) """ gu = _unnorm_canon(gt_canon) if is_native: d = pred_norm.shape[-1] pu = pred_norm[:, : min(d, 7)] * nat_std[: min(d, 7)] + nat_mean[: min(d, 7)] if pu.shape[-1] < 7: pu = np.concatenate([pu, gu[:, pu.shape[-1]:7]], axis=1) else: pu = _unnorm_canon(pred_norm) gp, gr = _integrate(gu) pp, pr = _integrate(pu) ep_mm = np.linalg.norm(pp[-1] - gp[-1]) * 1000 path_mm = np.linalg.norm(pp - gp, axis=1).mean() * 1000 ep_deg = np.degrees((pr[-1] * gr[-1].inv()).magnitude()) return ep_mm, path_mm, ep_deg pos_mm = np.linalg.norm(pu[:, :3] - gu[:, :3], axis=1).mean() * 1000 # rotation error: geodesic between rotvec deltas rp = Rotation.from_rotvec(pu[:, 3:6]) rg = Rotation.from_rotvec(gu[:, 3:6]) rot_deg = np.degrees((rp * rg.inv()).magnitude()).mean() grip = np.abs(pu[:, 6] - gu[:, 6]).mean() return pos_mm, rot_deg, grip eps = list(range(args.ep_start, min(args.ep_start + args.episodes, ds.num_episodes))) # candidate embodiment conditionings if cfg.conditioning == "morph": candidates = [("morph", morph)] else: candidates = [(f"id{i}", i) for i in range(args.oracle_ids)] best = None for label, cand in candidates: errs = [] for ep in eps: start = int(ds.meta.episodes["dataset_from_index"][ep]) end = int(ds.meta.episodes["dataset_to_index"][ep]) for idx in range(start, end - 1, args.stride): gt = quantile_normalize(store.chunk_for(ep, idx - start), q01, q99) if cfg.conditioning == "morph": b = to_batch(ds[idx], 0, cand) else: b = to_batch(ds[idx], cand, None) pred = policy.predict_action_chunk(b)[0].cpu().numpy() errs.append(canonical_phys_err(pred, gt)) errs = np.array(errs) m = errs.mean(0) if best is None or m[0] < best[1][0]: best = (label, m) print(f" {label}: endpoint {m[0]:.1f}mm path {m[1]:.1f}mm rot {m[2]:.1f}deg") print(f"\n=== {args.dataset} | {args.checkpoint.name} | cond={cfg.conditioning} ===") print(f"BEST ({best[0]}): endpoint {best[1][0]:.1f}mm path {best[1][1]:.1f}mm rot {best[1][2]:.1f}deg") if __name__ == "__main__": main()