#!/usr/bin/env python """Auto-eval for in-context demo conditioning on held-out LeKiwi. The key test: give the model K=3 REAL (obs, action) example pairs from LeKiwi's OWN data at test time — NO gradient update, NO fine-tuning — and see if that beats zero-shot with no examples, wrong examples, or the numeric/text descriptor alone. """ 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.modeling_tinyvla import TinyVLAPolicy CKPT = "outputs/tv2_richslow/step_10000" DS = "lekiwi_cleanup" ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" K = 3 @torch.no_grad() def main(): pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() cfg = pol.config tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) 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") 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) def img_at(idx, ep, ep_start): it = ds[idx] im = torch.nn.functional.interpolate(it[prim][None], size=(256, 256), mode="bilinear")[0] raw = quantile_normalize(store.chunk_for(ep, idx - ep_start), q01, q99) act = torch.from_numpy(raw[:, :7]).float() act = torch.nn.functional.pad(act, (0, cfg.max_action_dim - 7)) return im, act # REAL LeKiwi demo pairs from early episodes (support pool), disjoint from test support_pool_eps = range(0, min(20, m.total_episodes - 15)) test = range(max(0, m.total_episodes - 15), m.total_episodes) def get_real_support(k): imgs, acts = [], [] for _ in range(k): ep = int(np.random.choice(list(support_pool_eps))) s = int(m.episodes["dataset_from_index"][ep]) e = int(m.episodes["dataset_to_index"][ep]) idx = int(np.random.randint(s, max(s + 1, e - 1))) im, act = img_at(idx, ep, s) imgs.append(im) acts.append(act) return torch.stack(imgs), torch.stack(acts) def get_wrong_support(k): # random noise images + random actions — a garbage support set control return torch.rand(k, 3, 256, 256), torch.randn(k, 50, cfg.max_action_dim) * 0.3 # control: demos from a DIFFERENT ROBOT (jaco) — if the model truly extracts # robot-specific dynamics from exemplars, these should be worse than LeKiwi's own def get_otherrobot_support(k): from lerobot.datasets.lerobot_dataset import LeRobotDataset as _LD, LeRobotDatasetMetadata as _LM from tinyvla.data.canonical import CanonicalChunkStore as _CS jr = "/home/alexw/tinyvla_data/heldout/heldout_jaco_play" jm = _LM("heldout_jaco_play", root=jr) jds = _LD("heldout_jaco_play", root=jr, delta_timestamps={"action": [t / jm.fps for t in range(50)]}, video_backend="torchcodec") jstore = _CS("heldout_jaco_play", src_fps=jm.fps, chunk=50) jst = jstore.compute_stats(200) jq01, jq99 = np.asarray(jst["q01"]), np.asarray(jst["q99"]) jimk = sorted(kk for kk in jds.meta.features if kk.startswith("observation.images"))[0] imgs, acts = [], [] for _ in range(k): ep = int(np.random.randint(0, 50)) s = int(jm.episodes["dataset_from_index"][ep]); e = int(jm.episodes["dataset_to_index"][ep]) idx = int(np.random.randint(s, max(s + 1, e - 1))) it = jds[idx] im = torch.nn.functional.interpolate(it[jimk][None], size=(256, 256), mode="bilinear")[0] raw = quantile_normalize(jstore.chunk_for(ep, idx - s), jq01, jq99) act = torch.nn.functional.pad(torch.from_numpy(raw[:, :7]).float(), (0, cfg.max_action_dim - 7)) imgs.append(im); acts.append(act) return torch.stack(imgs), torch.stack(acts) def run(support_fn, desc_text): ids_t = tok([desc_text] if desc_text else [""], padding="max_length", truncation=True, max_length=cfg.morph_text_max_len, return_tensors="pt") errs = [] 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 "" t = tok([task], 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 = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) sup_img, sup_act = support_fn(K) 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(), "morph_text_ids": ids_t["input_ids"].cuda(), "morph_text_mask": ids_t["attention_mask"].bool().cuda(), "support_images": sup_img[None].cuda(), "support_actions": sup_act[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) return np.mean(errs) print("=== RICH-SLOW (demos in LM attention stream, no FT) on held-out LeKiwi ===") print(f"REAL LeKiwi demos (K={K}) + desc=lekiwi endpoint {run(get_real_support, prompts['lekiwi']):.1f}mm") print(f"REAL LeKiwi demos (K={K}) + desc=none endpoint {run(get_real_support, None):.1f}mm") print(f"WRONG/garbage demos + desc=lekiwi endpoint {run(get_wrong_support, prompts['lekiwi']):.1f}mm") print(f"WRONG/garbage demos + desc=none endpoint {run(get_wrong_support, None):.1f}mm") print(f"OTHER-ROBOT (jaco) demos + desc=none endpoint {run(get_otherrobot_support, None):.1f}mm") print("\nprior baselines: C-diverse 295mm | C-qwen-morph desc=none 294mm | v1 demo(same-task) real=290.6 garbage=288.5 (dead)") if __name__ == "__main__": main()