| |
| """Auto-eval for C-qwen-morph (morphology via shared-Qwen text encode, LM unfrozen) |
| on held-out LeKiwi. Compares correct vs wrong vs no robot-description text, and |
| against prior baselines (C-diverse numeric-MLP 295mm; C-mega text-in-prompt 522mm). |
| """ |
|
|
| 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 |
|
|
| CKPT = "outputs/tv2_C_qwen_morph/final" |
| DS = "lekiwi_cleanup" |
| ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" |
|
|
|
|
| @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") |
| _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 tok_desc(text): |
| if text is None: |
| ids = torch.zeros(1, cfg.morph_text_max_len, dtype=torch.long) |
| mask = torch.zeros(1, cfg.morph_text_max_len, dtype=torch.bool) |
| mask[0, 0] = True |
| return ids, mask |
| t = tok([text], padding="max_length", truncation=True, |
| max_length=cfg.morph_text_max_len, return_tensors="pt") |
| return t["input_ids"], t["attention_mask"].bool() |
|
|
| def run(desc_text): |
| ids, mask = tok_desc(desc_text) |
| 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 "" |
| 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 = _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(), |
| "morph_text_ids": ids.cuda(), "morph_text_mask": mask.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-qwen-morph (LM unfrozen, morph via shared Qwen) on held-out LeKiwi ===") |
| zero_floor = None |
| conditions = [ |
| ("desc=lekiwi (correct)", prompts["lekiwi"]), |
| ("desc=so101 (wrong, non-mobile arm)", prompts["so101"]), |
| ("desc=none", None), |
| ] |
| for label, text in conditions: |
| err, zf = run(text) |
| if zero_floor is None: |
| zero_floor = zf |
| print(f"{label:38} endpoint {err:.1f}mm") |
| print(f"\nzero-floor: {zero_floor:.1f}mm") |
| print("prior baselines: C-diverse numeric-MLP 295.0mm | C-mega text-in-prompt(frozen LM) 298.5mm best") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
