File size: 4,548 Bytes
5a2e445 63eab78 5a2e445 63eab78 5a2e445 63eab78 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 | #!/usr/bin/env python
"""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()
|