tinyvla2/scripts/eval_incontext_lekiwi.py

#!/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_C_incontext/final"
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")
    _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

    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("=== C-incontext (demo conditioning, 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("\nprior baselines (no demo conditioning): C-diverse 295mm | C-qwen-morph desc=none 294mm")


if __name__ == "__main__":
    main()