tinyvla2/scripts/eval_robocasa_all.py

#!/usr/bin/env python
"""How does the best checkpoint do on SIMILAR robots (arms like the training ones)
vs the exotic held-out ones? Reports endpoint error AND the zero-prediction floor
for each, since the floor differs a lot per robot (it is what "predicting mean
motion" achieves) — only the ratio to floor is comparable across robots.
"""

from __future__ import annotations

import numpy as np
import torch
import yaml
from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
from transformers import AutoTokenizer

from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize
from tinyvla.data.eval_utils import StateAdapter
from tinyvla.modeling_tinyvla import TinyVLAPolicy
from tinyvla.modules.embodiment import MORPH_FIELDS

import sys
CKPT = sys.argv[1] if len(sys.argv) > 1 else "outputs/tv2_robocasa/final"
_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}

# (label, dataset name, root, morph_key, in_training?)
TARGETS = [
    ("jaco (held-out ARM)", "heldout_jaco_play", "~/tinyvla_data/heldout/heldout_jaco_play", "jaco", False),
    ("ur5 (in training)", "div_ur5", "~/tinyvla_data/diverse/div_ur5", "ur5", True),
    ("xarm (in training)", "div_xarm", "~/tinyvla_data/diverse/div_xarm", "xarm", True),
    ("dlr_edan (in training)", "div_dlr_edan", "~/tinyvla_data/diverse/div_dlr_edan", "dlr_edan", True),
    ("stretch (in training, mobile)", "heldout_cmu_stretch", "~/tinyvla_data/heldout/heldout_cmu_stretch", "hello_stretch", True),
    ("LeKiwi (held-out MOBILE)", "lekiwi_cleanup", "~/tinyvla_data/lekiwi/lekiwi_cleanup", "lekiwi", False),
    ("RoboCasa (NEW domain)", "robocasa365", "~/tinyvla_data/robocasa365", "panda_omron", True),
]


@torch.no_grad()
def main():
    from pathlib import Path

    pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval()
    cfg = pol.config
    tok = AutoTokenizer.from_pretrained(cfg.lm_model_name)
    desc = yaml.safe_load(open("configs/morphology/descriptors.yaml"))

    print(f"=== {CKPT} across robots (endpoint err vs zero-floor) ===")
    rows = []
    for label, name, root, mkey, in_train in TARGETS:
        root = str(Path(root).expanduser())
        try:
            m = LeRobotDatasetMetadata(name, root=root)
            ds = LeRobotDataset(name, root=root,
                                delta_timestamps={"action": [t / m.fps for t in range(50)]},
                                video_backend="torchcodec")
            _sa = StateAdapter(ds.meta, cfg.max_state_dim)
            store = CanonicalChunkStore(name, src_fps=m.fps, chunk=50)
            st = store.compute_stats()
            q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"])
            span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6))
            mid = 0.5 * (q01 + q99)
            imk = sorted(k for k in ds.meta.features if k.startswith("observation.images"))
            prim = next((k for k in imk if any(s in k for s in ("front", "base", "top", "image"))), imk[0])
            morph = torch.tensor([desc[mkey].get(f, 0) * _SC.get(f, 1) for f in MORPH_FIELDS],
                                 dtype=torch.float32)
            n_eps = ds.num_episodes
            test = range(max(0, n_eps - 12), n_eps)
            errs, floors = [], []
            for ep in test:
                s = int(m.episodes["dataset_from_index"][ep])
                e = int(m.episodes["dataset_to_index"][ep])
                for idx in range(s, e - 1, 40):
                    item = ds[idx]
                    t = tok([item.get("task") or ""], padding=True, truncation=True,
                            max_length=48, return_tensors="pt")
                    img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0]
                    stt = _sa(item["observation.state"])
                    b = {"observation.images.cam0": img[None].cuda(),
                         "observation.images.cam1": torch.zeros_like(img)[None].cuda(),
                         "observation.state": stt[None].cuda(),
                         "observation.language.tokens": t["input_ids"].cuda(),
                         "observation.language.attention_mask": t["attention_mask"].bool().cuda(),
                         "morphology": morph[None].cuda(),
                         "embodiment_id": torch.tensor([0]).cuda()}
                    with torch.autocast("cuda", torch.bfloat16):
                        pr = pol.predict_action_chunk(b)[0].cpu().float().numpy()
                    gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid
                    pu = pr[:, :7] * span / 2 + mid
                    gp = np.cumsum(gu[:, :3], 0)
                    pp = np.cumsum(pu[:, :3], 0)
                    errs.append(np.linalg.norm(pp[-1] - gp[-1]) * 1000)
                    floors.append(np.linalg.norm(gp[-1]) * 1000)
            err, floor = float(np.mean(errs)), float(np.mean(floors))
            rows.append((label, in_train, err, floor, err / floor))
            print(f"{label:32} {'[train]' if in_train else '[HELD-OUT]':11} "
                  f"endpoint {err:6.1f}mm  floor {floor:6.1f}mm  ratio {err/floor:.2f}")
        except Exception as ex:
            print(f"{label:32} FAILED {type(ex).__name__}: {str(ex)[:70]}")

    print("\nratio < 1.0 = better than predicting mean motion")
    tr = [r for r in rows if r[1]]
    ho = [r for r in rows if not r[1]]
    if tr:
        print(f"mean ratio, robots IN training : {np.mean([r[4] for r in tr]):.2f}")
    if ho:
        print(f"mean ratio, HELD-OUT robots    : {np.mean([r[4] for r in ho]):.2f}")


if __name__ == "__main__":
    main()