eval_simpler.py

#!/usr/bin/env python3
"""SimplerEnv evaluation entry point for STAIR.

Usage (server):
    python stair/scripts/eval_simpler.py \
        --checkpoint runs/stair_stage2_v1/checkpoints/step-0010000.pt \
        --policy_setup widowx_bridge \
        --task widowx_put_eggplant_in_basket \
        --num_episodes 50
"""
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

_STAIR_DIR = Path(__file__).resolve().parent.parent
if str(_STAIR_DIR) not in sys.path:
    sys.path.insert(0, str(_STAIR_DIR))

import draccus
import numpy as np


@dataclass
class EvalConfig:
    checkpoint: Path = Path("runs/stair_stage2_v1/checkpoints/step-0010000.pt")
    policy_setup: str = "widowx_bridge"
    task: str = "widowx_put_eggplant_in_basket"
    num_episodes: int = 50
    max_episode_steps: int = 80
    action_scale: float = 1.0
    num_inference_steps: int = 10
    unnorm_key: Optional[str] = None
    results_dir: Path = Path("runs/eval_results")
    action_model_type: str = "DiT-B"


def _build_policy(cfg: EvalConfig):
    from vla import load_vla
    from stair.deploy.stair_policy_sim import STAIRSimPolicy

    model = load_vla(
        cfg.checkpoint,
        load_for_training=False,
        action_model_type=cfg.action_model_type,
        stage=2,
    )
    return STAIRSimPolicy(
        model=model,
        policy_setup=cfg.policy_setup,
        action_scale=cfg.action_scale,
        num_inference_steps=cfg.num_inference_steps,
        unnorm_key=cfg.unnorm_key,
    )


def _run_episode(env, policy, max_steps: int) -> bool:
    from simpler_env.utils.env.observation_utils import get_image_from_maniskill2_obs_dict

    task = env.get_language_instruction()
    obs, _ = env.reset()
    image = get_image_from_maniskill2_obs_dict(env, obs)
    policy.reset(task)
    done = False
    for _ in range(max_steps):
        _, action = policy.step(image, task)
        if bool(action["terminate_episode"][0] > 0):
            break
        obs, _, done, truncated, _ = env.step(
            np.concatenate([action["world_vector"], action["rot_axangle"], action["gripper"]])
        )
        if done or truncated:
            break
        image = get_image_from_maniskill2_obs_dict(env, obs)
    return bool(done)


@draccus.wrap()
def evaluate(cfg: EvalConfig) -> None:
    import json
    import simpler_env

    policy = _build_policy(cfg)
    env = simpler_env.make(cfg.task)
    successes = [_run_episode(env, policy, cfg.max_episode_steps) for _ in range(cfg.num_episodes)]
    env.close()

    sr = float(np.mean(successes))
    print(f"{cfg.task}: success_rate={sr:.3f}  (n={len(successes)})")

    cfg.results_dir.mkdir(parents=True, exist_ok=True)
    out = cfg.results_dir / f"{cfg.task}.json"
    out.write_text(json.dumps({"task": cfg.task, "success_rate": sr, "n": len(successes)}, indent=2))
    print(f"Saved {out}")


if __name__ == "__main__":
    evaluate()