File size: 8,721 Bytes
5a2e445 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | #!/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()
|