tinyvla2/scripts/eval_canonical.py

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