File size: 7,119 Bytes
5a2e445 1bb5587 5a2e445 65feb04 5a2e445 1bb5587 5a2e445 1bb5587 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 | #!/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.data.eval_utils import StateAdapter
from tinyvla.modeling_tinyvla import TinyVLAPolicy
CKPT = "outputs/tv2_demoonly/step_10000"
DS = "lekiwi_cleanup"
ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}"
K = 4
@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)
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 = _sa(item["observation.state"])
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("=== DEMO-ONLY (demos are the ONLY identity cue, no FT) on held-out LeKiwi ===")
print(f"REAL LeKiwi demos (K={K}) endpoint {run(get_real_support, None):.1f}mm")
print(f"GARBAGE/noise demos endpoint {run(get_wrong_support, None):.1f}mm")
print(f"OTHER-ROBOT (jaco) demos 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()
|