tinyvla_b200/scripts/eval_microduck.py

#!/usr/bin/env python
"""Offline eval + renders of the v4 checkpoint on AlexWortega/microduck-vla.

Zero-shot robot (never in training). For every behavior x split it runs the
model over one episode in true v4 mode — K=3 demo exemplars from ANOTHER episode
of the same behavior in the LM stream, a duck morphology descriptor, robot text —
and renders an mp4: camera frame + per-dim action curves (GT vs predicted
chunks) + running MSE against the zero-prediction baseline.

Data is stored ALREADY NORMALIZED (val with train stats), so predictions and GT
live in the same space; MSE ratio to the zero baseline is the honest metric
(project lesson: absolute errors on a new robot are meaningless).
"""

from __future__ import annotations

import argparse
import io
import json
import tarfile
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch

BEHAVIORS = ["ground_pick", "kick_left", "kick_right", "roller", "roller_crouch",
             "roulade", "sitstand", "stand", "walking"]

DUCK_MORPH = {  # rough MicroDuck descriptor (bipedal toy robot, 14 actuators)
    "arm_dof": 2, "reach_m": 0.08, "gripper_width_m": 0.0, "num_cameras": 2,
    "is_mobile": 1, "control_hz": 50, "joint_lo_mean": -1.5, "joint_hi_mean": 1.5,
    "workspace_x": 0.1, "workspace_y": 0.1, "workspace_z": 0.15, "payload_kg": 0.05,
    "ee_type_parallel": -1, "ee_type_multi": -1, "base_holonomic": -1, "reserved": 0,
}
DUCK_TEXT = "Robot: MicroDuck, a tiny bipedal duck robot with two legs, a neck and a head, 14 joints."


def load_split(root: Path, split: str):
    """episode_index -> list of (frame, cam0_jpg, state, chunk)."""
    eps = defaultdict(list)
    for shard in sorted((root / split).glob("shard-*.tar")):
        with tarfile.open(shard) as t:
            metas = {}
            jpgs = {}
            for m in t.getmembers():
                base, _, comp = m.name.partition(".")
                if comp == "meta.npz":
                    metas[base] = t.extractfile(m).read()
                elif comp == "cam0.jpg":
                    jpgs[base] = t.extractfile(m).read()
            for base, raw in metas.items():
                d = np.load(io.BytesIO(raw))
                ep, fr = int(base.split("_")[1]), int(base.split("_")[2])
                eps[ep].append((fr, jpgs[base], d["state"], d["action_chunk"]))
    for ep in eps:
        eps[ep].sort(key=lambda x: x[0])
    return eps


def behavior_of(ep: int) -> str:
    return BEHAVIORS[min(ep // 54, len(BEHAVIORS) - 1)]


def decode_img(jpg: bytes, size: int = 256) -> torch.Tensor:
    from PIL import Image

    a = np.asarray(Image.open(io.BytesIO(jpg)).convert("RGB"), dtype=np.uint8)
    x = torch.from_numpy(a.copy()).permute(2, 0, 1).float() / 255.0
    return x


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", type=Path, required=True)
    ap.add_argument("--data", type=Path, default=Path("/workspace/microduck"))
    ap.add_argument("--out", type=Path, default=Path("/workspace/renders"))
    ap.add_argument("--dims", type=int, nargs="+", default=[0, 1, 2, 7, 8, 9])
    args = ap.parse_args()

    from transformers import AutoTokenizer
    from tinyvla.modeling_tinyvla import TinyVLAPolicy
    from tinyvla.modules.embodiment import MORPH_FIELDS

    torch.backends.cuda.enable_cudnn_sdp(False)
    pol = TinyVLAPolicy.from_pretrained(args.ckpt).cuda().eval()
    cfg = pol.config
    tok = AutoTokenizer.from_pretrained(cfg.lm_model_name)

    sc = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3,
          "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416,
          "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2}
    morph = torch.tensor([DUCK_MORPH.get(f, 0) * sc.get(f, 1) for f in MORPH_FIELDS],
                         dtype=torch.float32)[None].cuda()
    mt = tok([DUCK_TEXT], padding="max_length", truncation=True,
             max_length=cfg.morph_text_max_len, return_tensors="pt")

    args.out.mkdir(parents=True, exist_ok=True)
    A = 14
    metrics = {}

    for split in ("train", "validation"):
        eps = load_split(args.data, split)
        by_beh = defaultdict(list)
        for ep in sorted(eps):
            by_beh[behavior_of(ep)].append(ep)

        for beh in BEHAVIORS:
            pool = by_beh.get(beh, [])
            if len(pool) < 2:
                continue
            ep, sup_ep = pool[0], pool[1]  # eval первый, демо из второго
            frames = eps[ep]
            sup_frames = eps[sup_ep]

            # K=3 support из другого эпизода того же поведения
            sidx = np.linspace(0, len(sup_frames) - 1, 3).astype(int)
            sup_img = torch.stack([decode_img(sup_frames[i][1]) for i in sidx])
            sup_act = torch.stack([
                torch.nn.functional.pad(torch.from_numpy(sup_frames[i][3]),
                                        (0, cfg.max_action_dim - A)) for i in sidx])

            t_task = tok([f"MicroDuck: perform {beh.replace('_', ' ')}"], padding="max_length",
                         truncation=True, max_length=cfg.tokenizer_max_length, return_tensors="pt")

            preds, gts = [], []
            for fr, jpg, state, chunk in frames:
                img = decode_img(jpg)
                st = torch.nn.functional.pad(torch.from_numpy(state), (0, cfg.max_state_dim - 61))
                b = {"observation.images.cam0": img[None].cuda(),
                     "observation.images.cam1": torch.zeros_like(img)[None].cuda(),
                     "observation.state": st[None].cuda(),
                     "observation.language.tokens": t_task["input_ids"].cuda(),
                     "observation.language.attention_mask": t_task["attention_mask"].bool().cuda(),
                     "morph_text_ids": mt["input_ids"].cuda(),
                     "morph_text_mask": mt["attention_mask"].bool().cuda(),
                     "morphology": morph,
                     "support_images": sup_img[None].cuda(),
                     "support_actions": sup_act[None].cuda(),
                     "embodiment_id": torch.tensor([0]).cuda()}
                with torch.no_grad(), torch.autocast("cuda", torch.bfloat16):
                    pr = pol.predict_action_chunk(b)[0].float().cpu().numpy()[:, :A]
                preds.append(pr)
                gts.append(chunk)

            preds = np.stack(preds)          # (T, 50, 14)
            gts = np.stack(gts)
            mse = float(np.mean((preds - gts) ** 2))
            mse_zero = float(np.mean(gts ** 2))  # данные нормализованы: 0 = среднее
            metrics[f"{split}/{beh}"] = {"mse": mse, "mse_zero": mse_zero,
                                         "ratio": mse / max(mse_zero, 1e-9),
                                         "episode": ep, "steps": len(frames)}
            print(f"{split:10} {beh:14} ep{ep:4d} mse={mse:.4f} zero={mse_zero:.4f} "
                  f"ratio={mse/max(mse_zero,1e-9):.3f}", flush=True)

            render(args.out / f"{split}_{beh}.mp4", frames, preds, gts, args.dims,
                   f"{beh} [{split}] ep{ep} | MSE {mse:.3f} vs zero {mse_zero:.3f}")

    (args.out / "metrics.json").write_text(json.dumps(metrics, indent=2))
    n_ok = sum(1 for m in metrics.values() if m["ratio"] < 1.0)
    print(f"\nитого: {n_ok}/{len(metrics)} комбинаций лучше zero-baseline")


def render(path, frames, preds, gts, dims, title):
    import matplotlib
    matplotlib.use("Agg")
    import imageio.v2 as imageio
    import matplotlib.pyplot as plt
    from PIL import Image

    T = len(frames)
    out = []
    for t in range(T):
        fig, axes = plt.subplots(1, 2, figsize=(10, 4.2), dpi=80,
                                 gridspec_kw={"width_ratios": [1, 1.4]})
        img = Image.open(io.BytesIO(frames[t][1])).convert("RGB")
        axes[0].imshow(img); axes[0].axis("off")
        axes[0].set_title(f"t={t}/{T}", fontsize=9)
        ax = axes[1]
        # GT: сплошные линии первых шагов каждого чанка (реально исполненная траектория)
        gt_traj = gts[:, 0, :]          # (T, 14): первый шаг каждого чанка
        xs = np.arange(T)
        for i, d in enumerate(dims):
            ax.plot(xs, gt_traj[:, d] + i * 2.5, lw=1.0, color="k", alpha=0.7)
            # предсказанный чанк из текущего t: 50 шагов @50Гц = 5 obs-шагов вперёд
            px = t + np.linspace(0, 5, preds.shape[1])
            ax.plot(px, preds[t, :, d] + i * 2.5, lw=1.4, color="tab:red", alpha=0.9)
        ax.axvline(t, color="tab:blue", lw=0.8)
        ax.set_yticks([i * 2.5 for i in range(len(dims))])
        ax.set_yticklabels([f"dim{d}" for d in dims], fontsize=7)
        ax.set_xlim(0, T + 5); ax.set_xlabel("obs step (10 Hz)", fontsize=8)
        ax.set_title("чёрное = GT, красное = предсказанный чанк (1 c)", fontsize=8)
        fig.suptitle(title, fontsize=10)
        fig.tight_layout()
        fig.canvas.draw()
        w, h = fig.canvas.get_width_height()
        out.append(np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
                   .reshape(h, w, 4)[..., :3].copy())
        plt.close(fig)
    imageio.mimwrite(path, out, fps=10, quality=7)


if __name__ == "__main__":
    main()