#!/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()