#!/usr/bin/env python """Closed-loop LIBERO evaluation for canonical-schema TinyVLA checkpoints. Adapts env observations to the canonical schema the policy was trained on (cam0/cam1, padded normalized state, embodiment_id) and unnormalizes the predicted actions with the LIBERO dataset stats. Usage: python scripts/eval_libero.py --checkpoint outputs/libero_ft/final \ --suite libero_spatial --episodes 20 [--refresh-s 1.0] """ from __future__ import annotations import argparse from pathlib import Path import torch def make_normalizer(repo_id="HuggingFaceVLA/libero"): from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata meta = LeRobotDatasetMetadata(repo_id) stats = meta.stats def norm(key, x): s = stats[key] mean = torch.as_tensor(s["mean"], dtype=torch.float32, device=x.device) std = torch.as_tensor(s["std"], dtype=torch.float32, device=x.device).clamp(min=1e-6) return (x - mean) / std def unnorm_action(x): s = stats["action"] mean = torch.as_tensor(s["mean"], dtype=torch.float32, device=x.device) std = torch.as_tensor(s["std"], dtype=torch.float32, device=x.device) return x * std + mean return norm, unnorm_action, meta @torch.no_grad() def main(): parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--suite", default="libero_spatial", choices=["libero_spatial", "libero_object", "libero_goal", "libero_10", "libero_90"]) parser.add_argument("--episodes", type=int, default=20) parser.add_argument("--embodiment-id", type=int, default=2) parser.add_argument("--refresh-s", type=float, default=None, help="if set, refresh the semantic latent only every N seconds (dual-rate mode)") parser.add_argument("--max-steps", type=int, default=520) parser.add_argument("--n-action-steps", type=int, default=None, help="execute only first N actions of each chunk before re-planning") parser.add_argument("--action-repeat", type=int, default=1, help="env steps per predicted action (dataset 10fps vs env 20Hz -> 2)") parser.add_argument("--save-video-dir", type=Path, default=None, help="save per-episode mp4s of the agentview camera here") args = parser.parse_args() import numpy as np from lerobot.envs.factory import make_env, make_env_config from transformers import AutoTokenizer from tinyvla.modeling_tinyvla import TinyVLAPolicy policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() cfg = policy.config if args.n_action_steps: cfg.n_action_steps = args.n_action_steps tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) norm, unnorm_action, meta = make_normalizer() fps = meta.fps env_cfg = make_env_config("libero", task=args.suite) envs_dict = make_env(env_cfg, n_envs=1) task_envs = envs_dict[args.suite] # {task_id: vec_env} from scipy.spatial.transform import Rotation def to_canonical(obs, task_text, latent=None): imgs = {} for slot, key in (("cam0", "image"), ("cam1", "image2")): x = torch.as_tensor(np.asarray(obs["pixels"][key])) if x.dim() == 4: # (1, H, W, C) x = x[0] # robosuite renders 180-degree rotated relative to the recorded dataset x = x.flip(0).flip(1) x = x.permute(2, 0, 1).float() / 255.0 if x.shape[-2:] != (cfg.image_size, cfg.image_size): x = torch.nn.functional.interpolate( x[None], size=(cfg.image_size, cfg.image_size), mode="bilinear", align_corners=False )[0] imgs[slot] = x # dataset convention (OpenVLA-style): [eef pos(3), eef axis-angle(3), gripper qpos(2)] rs = obs["robot_state"] pos = np.asarray(rs["eef"]["pos"]).flatten() quat = np.asarray(rs["eef"]["quat"]).flatten() # robosuite: (x, y, z, w) rotvec = Rotation.from_quat(quat).as_rotvec() # canonicalize antipodal representation to match dataset convention # (dataset uses rotvec with positive x-component, ~+pi for downward gripper) if rotvec[0] < 0: theta = np.linalg.norm(rotvec) if theta > 1e-6: rotvec = rotvec * (theta - 2 * np.pi) / theta grip = np.asarray(rs["gripper"]["qpos"]).flatten() state = torch.tensor(np.concatenate([pos, rotvec, grip]), dtype=torch.float32) state = norm("observation.state", state) state = torch.nn.functional.pad(state, (0, cfg.max_state_dim - state.shape[-1])) t = tok([task_text], padding=True, truncation=True, max_length=cfg.tokenizer_max_length, return_tensors="pt") batch = { "observation.images.cam0": imgs["cam0"][None].cuda(), "observation.images.cam1": imgs["cam1"][None].cuda(), "observation.state": state[None].cuda(), "observation.language.tokens": t["input_ids"].cuda(), "observation.language.attention_mask": t["attention_mask"].bool().cuda(), "embodiment_id": torch.tensor([args.embodiment_id], device="cuda"), } if latent is not None: batch["semantic_latent"] = latent return batch action_dim = meta.features["action"]["shape"][0] total, succ = 0, 0 per_task = {} task_ids = sorted(task_envs.keys()) ep_plan = [(tid, i) for i in range((args.episodes + len(task_ids) - 1) // len(task_ids)) for tid in task_ids] ep_plan = ep_plan[: args.episodes] if args.save_video_dir: args.save_video_dir.mkdir(parents=True, exist_ok=True) for ep, (tid, rep) in enumerate(ep_plan): env = task_envs[tid] obs, info = env.reset(seed=1000 + rep) frames = [] if args.save_video_dir else None try: task_text = env.get_attr("task_description")[0] except Exception: task_text = getattr(getattr(env, "envs", [None])[0], "task_description", "") policy.reset() latent = None last_refresh = -1e9 done = False step_i = 0 ep_succ = False while not done and step_i < args.max_steps: t_now = step_i / fps batch = to_canonical(obs, task_text) if args.refresh_s is not None: if t_now - last_refresh >= args.refresh_s: latent = policy._semantic_latent(batch) last_refresh = t_now batch["semantic_latent"] = latent act_norm = policy.select_action(batch) # (1, max_action_dim) normalized padded act = unnorm_action(act_norm[0, :action_dim].cpu()).clamp(-1, 1) for _ in range(args.action_repeat): try: obs, reward, terminated, truncated, info = env.step(act.numpy()[None]) except ValueError: # stepped into env's internal horizon done = True break done = bool(terminated[0] or truncated[0]) if info.get("is_success") is not None: ep_succ = ep_succ or bool(np.asarray(info["is_success"]).flatten()[0]) if frames is not None: frames.append(np.asarray(obs["pixels"]["image"])[0][::-1, ::-1]) step_i += 1 if done or step_i >= args.max_steps: done = done or step_i >= args.max_steps break total += 1 succ += int(ep_succ) per_task.setdefault(task_text[:50], []).append(int(ep_succ)) print(f"ep {ep}: {'SUCCESS' if ep_succ else 'fail'} ({step_i} steps) | {task_text[:60]}") if frames: import imageio.v2 as imageio tag = "succ" if ep_succ else "fail" path = args.save_video_dir / f"ep{ep:02d}_task{tid}_{tag}.mp4" imageio.mimwrite(path, frames, fps=20, quality=7) print(f"\n=== {args.suite} | {args.checkpoint} | refresh={args.refresh_s} ===") print(f"success rate: {succ}/{total} = {succ/total:.1%}") for t, v in per_task.items(): print(f" {sum(v)}/{len(v)} {t}") if __name__ == "__main__": main()