| |
| """Auto-eval for C-mega on held-out LeKiwi: text-prompt x numeric-descriptor ablation. |
| |
| Compares (numeric descriptor correct/wrong) x (text prompt correct/none) to see |
| whether the natural-language robot description adds anything beyond the C-scheme |
| numeric descriptor, zero-shot, no fine-tuning. Also reports the zero-floor and the |
| prior C-diverse baseline (295mm, numeric-only) for context. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| 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.data.eval_utils import StateAdapter |
| from tinyvla.modeling_tinyvla import TinyVLAPolicy |
| from tinyvla.modules.embodiment import MORPH_FIELDS |
|
|
| CKPT = "outputs/tv2_C_mega/final" |
| DS = "lekiwi_cleanup" |
| ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" |
| _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} |
|
|
|
|
| def mvec(d): |
| return torch.tensor([d.get(f, 0) * _SC.get(f, 1) for f in MORPH_FIELDS], dtype=torch.float32) |
|
|
|
|
| @torch.no_grad() |
| def main(): |
| 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")) |
| prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) |
|
|
| m = LeRobotDatasetMetadata(DS, root=ROOT) |
| ds = LeRobotDataset(DS, 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(DS, 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 "front" in k or "base" in k), imk[0]) |
|
|
| def integ(d): |
| return np.cumsum(d[:, :3], 0) |
|
|
| n_eps = ds.num_episodes |
| test = range(max(0, n_eps - 15), n_eps) |
|
|
| def run(morph, prompt_prefix): |
| errs, zf = [], [] |
| 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, 30): |
| item = ds[idx] |
| task = item.get("task") or "" |
| text = f"{prompt_prefix} {task}" if prompt_prefix else task |
| t = tok([text], 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 |
| errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) |
| zf.append(np.linalg.norm(integ(gu)[-1]) * 1000) |
| return np.mean(errs), np.mean(zf) |
|
|
| print(f"=== C-mega on held-out LeKiwi (n_eps_test={len(list(test))}) ===") |
| zero_floor = None |
| conditions = [ |
| ("numeric=lekiwi + text=lekiwi", mvec(desc["lekiwi"]), prompts["lekiwi"]), |
| ("numeric=lekiwi + text=none ", mvec(desc["lekiwi"]), None), |
| ("numeric=none + text=lekiwi", torch.zeros(16), prompts["lekiwi"]), |
| ("numeric=none + text=none ", torch.zeros(16), None), |
| ("numeric=so101(wrong,non-mobile) + text=lekiwi", mvec(desc["so101"]), prompts["lekiwi"]), |
| ] |
| for label, morph, prompt in conditions: |
| err, zf = run(morph, prompt) |
| if zero_floor is None: |
| zero_floor = zf |
| print(f"{label:48} endpoint {err:.1f}mm") |
| print(f"\nzero-floor: {zero_floor:.1f}mm") |
| print("prior C-diverse baseline (numeric-only, no mega training): 295.0mm") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
