tinyvla2/scripts/eval_libero.py

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