| |
| """Run the GR00T (physical_ai_ft) checkpoint on held-out LeKiwi. |
| |
| This is a deliberately unfair test, and the point is to measure HOW unfair. |
| |
| The checkpoint is `conditioning: id` + `action_space: native`. Its action head was |
| trained on six GR00T families whose native layouts are 44/14/24/12/23/26-43 dims, |
| each addressed by an embodiment id in 10..15. LeKiwi (SO-100 arm on a 3-wheel kiwi |
| base, 9 dims: 6 joint targets in degrees + x_mm/y_mm/theta) was never assigned an |
| id and its dim k has no reason to mean what any GR00T robot's dim k means. |
| |
| So the model is given every advantage available: |
| - each trained id 10..15 is tried, and the BEST is reported (oracle id selection) |
| - id 0 is reported too: rows 0..9 are the stale C-scaled ids this run never |
| touched, so it is the "no matching identity" reference |
| - the same 4-seed ODE averaging and the same per-dim ABS/DELTA baseline search |
| used for the in-distribution evaluation |
| |
| Preprocessing mirrors HubEpisodeStream: decimate to 10 Hz, per-dataset z-score, |
| zero-pad state to 256 and actions to 64. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| from eval_physical_ai import bootstrap_ci, chunk_err, classify_dims, predict_mean |
|
|
| CKPT = "/home/alexw/tinyvla/outputs/physical_ai_ft_fixed" |
| NAME, ROOT = "lekiwi_cleanup", str(Path("~/tinyvla_data/lekiwi/lekiwi_cleanup").expanduser()) |
| TARGET_HZ = 10.0 |
| IDS = [10, 11, 12, 13, 14, 15, 0] |
|
|
|
|
| @torch.no_grad() |
| def main(): |
| from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| from tinyvla.modeling_tinyvla import TinyVLAPolicy |
|
|
| pol = TinyVLAPolicy.from_pretrained(CKPT) |
| pol.load_state_dict(load_file(f"{CKPT}/model.safetensors"), strict=True) |
| cfg = pol.config |
| pol = pol.cuda().eval() |
| tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) |
| print(f"GATE strict load OK | action_dim={pol.action_dim} " |
| f"max_state={cfg.max_state_dim} max_action={cfg.max_action_dim}") |
|
|
| meta = LeRobotDatasetMetadata(NAME, root=ROOT) |
| stride = max(1, round(meta.fps / TARGET_HZ)) |
| ds = LeRobotDataset(NAME, root=ROOT, video_backend="torchcodec", |
| delta_timestamps={"action": [t / TARGET_HZ for t in range(cfg.chunk_size)]}) |
| sm = np.asarray(meta.stats["observation.state"]["mean"], dtype=np.float64) |
| ss = np.maximum(np.asarray(meta.stats["observation.state"]["std"], dtype=np.float64), 1e-6) |
| am = np.asarray(meta.stats["action"]["mean"], dtype=np.float64) |
| as_ = np.maximum(np.asarray(meta.stats["action"]["std"], dtype=np.float64), 1e-6) |
| ad = len(am) |
| print(f"LeKiwi: {meta.total_episodes} eps, {meta.fps} fps -> stride {stride} ({TARGET_HZ} Hz), " |
| f"action_dim {ad}") |
|
|
| |
| hf = ds.reader.hf_dataset.with_format("numpy") |
| s_all = np.asarray(hf["observation.state"], dtype=np.float64)[::stride] |
| a_all = np.asarray(hf["action"], dtype=np.float64)[::stride] |
| kind, amap, lag, std_a = classify_dims(s_all[:4000], a_all[:4000]) |
| names = meta.features["action"].get("names") or list(range(ad)) |
| names = names.get("motors", names) if isinstance(names, dict) else names |
| print("dim classification:", |
| ", ".join(f"{names[j]}={kind[j]}" + (f"<-s{amap[j]}" if kind[j] == "abs" else "") |
| for j in range(ad))) |
|
|
| imk = sorted(k for k in meta.features if k.startswith("observation.images")) |
| prim = next((k for k in imk if "wrist" not in k), imk[0]) |
| wrist = next((k for k in imk if "wrist" in k), None) |
|
|
| def rs(x): |
| return torch.nn.functional.interpolate(x[None].float(), size=(cfg.image_size, cfg.image_size), |
| mode="bilinear", align_corners=False)[0] |
|
|
| |
| test_eps = range(max(0, meta.total_episodes - 12), meta.total_episodes) |
| samples = [] |
| for ep in test_eps: |
| s = int(meta.episodes["dataset_from_index"][ep]) |
| e = int(meta.episodes["dataset_to_index"][ep]) |
| for idx in range(s, max(s + 1, e - cfg.chunk_size * stride), 90): |
| item = ds[idx] |
| A = item["action"].numpy().astype(np.float64)[:, :ad] |
| if A.shape[0] < cfg.chunk_size: |
| continue |
| st_raw = item["observation.state"].numpy().astype(np.float64)[:ad] |
| st_n = (st_raw - sm) / ss |
| samples.append({ |
| "ep": ep, |
| "cam0": rs(item[prim]), |
| "cam1": rs(item[wrist]) if wrist else torch.zeros(3, cfg.image_size, cfg.image_size), |
| "state": torch.nn.functional.pad(torch.tensor(st_n, dtype=torch.float32), |
| (0, cfg.max_state_dim - len(st_n))), |
| "A": (A - am) / as_, |
| "ref": np.array([(st_raw[amap[j]] - am[j]) / as_[j] if kind[j] == "abs" and amap[j] >= 0 |
| else 0.0 for j in range(ad)]), |
| "task": item.get("task") or "", |
| }) |
| print(f"held-out samples: {len(samples)} from {len(list(test_eps))} episodes\n") |
|
|
| def run(emb_id): |
| preds = [] |
| for k in range(0, len(samples), 8): |
| chunk = samples[k:k + 8] |
| t = tok([c["task"] for c in chunk], padding="max_length", truncation=True, |
| max_length=48, return_tensors="pt") |
| b = {"observation.images.cam0": torch.stack([c["cam0"] for c in chunk]).cuda(), |
| "observation.images.cam1": torch.stack([c["cam1"] for c in chunk]).cuda(), |
| "observation.state": torch.stack([c["state"] for c in chunk]).cuda(), |
| "observation.language.tokens": t["input_ids"].cuda(), |
| "observation.language.attention_mask": t["attention_mask"].bool().cuda(), |
| "embodiment_id": torch.full((len(chunk),), emb_id, dtype=torch.long).cuda()} |
| preds.append(predict_mean(pol, b, 4)) |
| return np.concatenate(preds, 0) |
|
|
| print(f"{'emb id':>8} {'subset':7}{'dims':>5}{'err':>9}{'B0':>9}{'B1':>9}{'B2':>9}{'ratio':>8}" |
| f" 95% CI") |
| best = {} |
| for emb in IDS: |
| pred = run(emb) |
| for subset in ("abs", "delta"): |
| sel = (kind == subset) |
| if not sel.any(): |
| continue |
| per_ep = {} |
| for i, c in enumerate(samples): |
| P = pred[i][:, :ad].astype(np.float64) |
| e, _ = chunk_err(P, c["A"], kind, amap, c["ref"], sel) |
| b0, _ = chunk_err(np.zeros_like(c["A"]), c["A"], kind, amap, c["ref"], sel) |
| b1, _ = chunk_err(np.tile(c["ref"], (len(c["A"]), 1)), c["A"], kind, amap, c["ref"], sel) |
| b2, _ = chunk_err(np.tile(c["A"][0], (len(c["A"]), 1)), c["A"], kind, amap, c["ref"], sel) |
| per_ep.setdefault(c["ep"], []).append((e, b0, b1, b2)) |
| rows = np.array([np.mean(v, axis=0) for v in per_ep.values()]) |
| err, b0, b1, b2 = rows.mean(0) |
| floor = min(b0, b1) |
| ratios = rows[:, 0] / np.maximum(rows[:, 1:3].min(1), 1e-9) |
| lo, hi = bootstrap_ci(ratios) |
| tag = " (untrained id row)" if emb == 0 else "" |
| print(f"{emb:>8} {subset:7}{int(sel.sum()):>5}{err:>9.3f}{b0:>9.3f}{b1:>9.3f}{b2:>9.3f}" |
| f"{err/floor:>8.2f} [{lo:.2f}, {hi:.2f}]{tag}") |
| k = subset |
| if k not in best or err / floor < best[k][1]: |
| best[k] = (emb, err / floor, b2 / floor) |
|
|
| print("\nORACLE id selection (best of 10..15, i.e. maximally generous to the checkpoint):") |
| for subset, (emb, r, b2r) in best.items(): |
| note = " <- oracle repeat-A0 is BETTER, so the model loses to a dynamics-free predictor" \ |
| if b2r < r else "" |
| print(f" {subset:6} best id {emb}: ratio {r:.2f} (oracle repeat-A0 {b2r:.2f}){note}") |
| print("\nratio < 1.0 beats the trivial baseline (hold still for abs dims, zero motion for delta)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
