"""Evaluate a trained ACT policy on Humanoid-Open-Drawer-v0 in IsaacSim. Mirrors dp_results/rl/online_residual_perstep/eval_residual_in_sim.py but for ACT (no BC DP, no residual). Loads ACT ckpt + dataset_stats.pkl, runs N=20 episodes with shuffled rand inits (same seed scheme as our BC eval). """ import argparse import os import pickle import random import sys from collections import deque from pathlib import Path sys.path.insert(0, "/home/hongyi/PWM/EgoVLA_Release") sys.path.insert(0, "/home/hongyi/PWM/policy_rl/act") from omni.isaac.lab.app import AppLauncher parser = argparse.ArgumentParser() parser.add_argument("--task", default="Humanoid-Open-Drawer-v0") parser.add_argument("--ckpt", required=True, help="Path to ACT policy_*.ckpt or policy_last.ckpt") parser.add_argument("--stats", default=None, help="Path to dataset_stats.pkl. Defaults to same dir as ckpt.") parser.add_argument("--num_episodes", type=int, default=20) parser.add_argument("--n_action_exec", type=int, default=24, help="How many actions from the predicted chunk to execute before re-querying.") parser.add_argument("--img_size", type=int, default=224) parser.add_argument("--rand_mode", choices=["sequential", "shuffled"], default="shuffled") parser.add_argument("--rand_offset", type=int, default=4000, help="Eval rand index range starts here (matches BC + residual eval).") parser.add_argument("--eval_torch_seed", type=int, default=42) parser.add_argument("--table_idx", type=int, default=1) parser.add_argument("--room_idx", type=int, default=1) parser.add_argument("--output_dir", type=str, default="/tmp/eval_act_out") parser.add_argument("--save_video", action="store_true") parser.add_argument("--temporal_agg", action="store_true", help="Aggregate overlapping chunk predictions (ALOHA's eval trick).") AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() args_cli.headless = True args_cli.enable_cameras = True app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app import cv2 import numpy as np import torch import gymnasium as gym import humanoid.tasks # noqa import omni.isaac.lab_tasks # noqa from omni.isaac.lab_tasks.utils import parse_env_cfg from omni.isaac.lab.controllers import DifferentialIKController, DifferentialIKControllerCfg from humanoid.tasks.base_env import BaseEnv from human_plan.ego_bench_eval.utils import ik_step, TASK_INIT_EPISODE, TASK_MAX_HORIZON from policy import ACTPolicy # ACT's wrapper from /home/hongyi/PWM/policy_rl/act/policy.py # ─── seed map matches our BC / residual eval ───────────────────────────────── seed_map = { "Humanoid-Push-Box-v0": 5, "Humanoid-Open-Drawer-v0": 6, "Humanoid-Close-Drawer-v0": 7, "Humanoid-Open-Laptop-v0": 8, "Humanoid-Pour-Balls-v0": 3, } # 38-d upper-body proprio slice indices (must match BC training) _ARM_IDS_LEFT = [4, 8, 12, 16, 20, 22, 24] _ARM_IDS_RIGHT = [5, 9, 13, 17, 21, 23, 25] _HAND_IDS_LEFT = [26, 36, 27, 37, 28, 38, 29, 39, 30, 40, 46, 48] _HAND_IDS_RIGHT = [31, 41, 32, 42, 33, 43, 34, 44, 35, 45, 47, 49] _UPPER_BODY_IDS = _ARM_IDS_LEFT + _ARM_IDS_RIGHT + _HAND_IDS_LEFT + _HAND_IDS_RIGHT # 38 def split_38d(a): return a[:7], a[7:14], a[14:26], a[26:38] def load_act_policy(ckpt_path, stats_path, num_queries=None): """Reconstruct ACTPolicy from the checkpoint. num_queries: if None, auto-detect from the ckpt's `model.pos_table` shape (pos_table is (1, num_queries+2, hidden_dim) due to ACT's 2 special tokens). """ if num_queries is None: _peek = torch.load(ckpt_path, map_location='cpu', weights_only=False) sd = _peek["act_state_dict"] if isinstance(_peek, dict) and "act_state_dict" in _peek else _peek if "model.pos_table" in sd: num_queries = int(sd["model.pos_table"].shape[1] - 2) print(f"[act-eval] auto-detected num_queries={num_queries} from ckpt pos_table", flush=True) else: num_queries = 50 print(f"[act-eval] pos_table not in ckpt; defaulting num_queries=50", flush=True) del _peek, sd # Re-build with the SAME hyperparameters used in imitate_episodes.py policy_config = { 'lr': 1e-5, 'num_queries': num_queries, 'kl_weight': 10, 'hidden_dim': 512, 'dim_feedforward': 3200, 'lr_backbone': 1e-5, 'backbone': 'resnet18', 'enc_layers': 4, 'dec_layers': 7, 'nheads': 8, 'camera_names': ['top'], 'state_dim': 38, } # ACT's build_ACT_model_and_optimizer calls argparse.parse_args() on sys.argv # which requires --policy_class --seed --num_epochs --task_name --ckpt_dir --batch_size --lr. # Stash the real argv, swap with dummy required args during model build, restore after. import sys as _sys real_argv = _sys.argv _sys.argv = [real_argv[0], '--policy_class', 'ACT', '--seed', '0', '--num_epochs', '1', '--task_name', 'sim_open_drawer_100demos', '--ckpt_dir', '/tmp/_act_eval_dummy', '--batch_size', '1', '--lr', '1e-5'] try: policy = ACTPolicy(policy_config) finally: _sys.argv = real_argv # Auto-detect BC vs RL ckpt format. # BC ckpt: raw state_dict (saved by imitate_episodes.py) # RL ckpt: nested {'act_state_dict': ..., 'log_std': ..., 'critic_state_dict': ..., 'iter': ..., 'args': ...} ckpt = torch.load(ckpt_path, map_location='cuda:0', weights_only=False) if isinstance(ckpt, dict) and "act_state_dict" in ckpt: state_dict = ckpt["act_state_dict"] print(f"[act-eval] loaded RL ckpt (iter={ckpt.get('iter', '?')})", flush=True) else: state_dict = ckpt print(f"[act-eval] loaded BC ckpt", flush=True) policy.load_state_dict(state_dict) policy.cuda().eval() with open(stats_path, 'rb') as f: stats = pickle.load(f) return policy, stats, policy_config def main(): device = "cuda:0" ckpt_path = args_cli.ckpt stats_path = args_cli.stats or str(Path(ckpt_path).parent / "dataset_stats.pkl") print(f"[act-eval] ckpt={ckpt_path} stats={stats_path}", flush=True) policy, stats, policy_config = load_act_policy(ckpt_path, stats_path) chunk_size = policy_config['num_queries'] print(f"[act-eval] chunk_size={chunk_size} state_dim={policy_config['state_dim']}", flush=True) pre_q = stats["qpos_mean"]; pre_qs = stats["qpos_std"] pre_a = stats["action_mean"]; pre_as = stats["action_std"] # ─── env init randomization — MATCHES the BC / residual eval scheme ─── task = args_cli.task task_name = task[9:-3] # "Open-Drawer" random.seed(seed_map[task]) torch.manual_seed(args_cli.eval_torch_seed) np.random.seed(args_cli.eval_torch_seed) with open("/home/hongyi/PWM/EgoVLA_Release/init_poses_fixed_set_100traj.pkl", "rb") as f: init_poses = pickle.load(f) episode_list = TASK_INIT_EPISODE[task_name] if args_cli.rand_mode == "shuffled": rand_pool = list(range(10000)); random.Random(seed_map[task]).shuffle(rand_pool) # Off-by-one to MATCH residual eval's reset_with_new_seed which advances # the pointer BEFORE the first episode. So idx[0]=7190 is dropped, episode 0 # starts from idx[1]=1677, matching residual eval logs. eval_idx_seq = [rand_pool[(args_cli.rand_offset + i + 1) % 10000] for i in range(args_cli.num_episodes)] print(f"[act-eval] shuffled idx seq (first 8): {eval_idx_seq[:8]}", flush=True) else: eval_idx_seq = [args_cli.rand_offset + i + 1 for i in range(args_cli.num_episodes)] # ─── env ─── env_cfg = parse_env_cfg(task, num_envs=1) env_cfg.episode_length_s = 60 env_cfg.randomize = True env_cfg.spawn_background = True env_cfg.room_idx = args_cli.room_idx env_cfg.table_idx = args_cli.table_idx env_cfg.randomize_idx = eval_idx_seq[0] env: BaseEnv = gym.make(task, cfg=env_cfg) env.reset() ik_l = DifferentialIKController( DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), num_envs=1, device=env.sim.device) ik_r = DifferentialIKController( DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="pinv"), num_envs=1, device=env.sim.device) lw = torch.zeros(1, ik_l.action_dim, device=env.robot.device) lrc = torch.zeros(1, ik_l.action_dim, device=env.robot.device) rw = torch.zeros(1, ik_r.action_dim, device=env.robot.device) rrc = torch.zeros(1, ik_r.action_dim, device=env.robot.device) action_buf = torch.zeros((1, env.num_actions), device=env.robot.device) out_dir = Path(args_cli.output_dir); out_dir.mkdir(parents=True, exist_ok=True) max_horizon = TASK_MAX_HORIZON.get(task, 400) n_succ = 0 for ep in range(args_cli.num_episodes): seq_name = episode_list[ep % len(episode_list)][0] env.cfg.randomize_idx = eval_idx_seq[ep] env.reset(); ik_l.reset(); ik_r.reset() # warm up to init pose ld = init_poses[task_name][seq_name][0]["left_dof"] rd = init_poses[task_name][seq_name][0]["right_dof"] le = init_poses[task_name][seq_name][0]["left_ee"] re = init_poses[task_name][seq_name][0]["right_ee"] env_results = None for _ in range(100): ik_step(env, ik_l, ik_r, lw, rw, lrc, rrc, le, re, ld, rd, action_buf) env_results = env.step(action_buf) # Eval loop: ACT predicts chunk_size actions, execute n_action_exec, re-query save_path = str(out_dir / f"ep{ep:03d}.mp4") if args_cli.save_video else None frames = [] if save_path else None success = False action_chunk = None # (chunk_size, 38) step_in_chunk = 0 # Optional temporal aggregation buffer if args_cli.temporal_agg: all_time_actions = torch.zeros((max_horizon, max_horizon + chunk_size, 38)).cuda() for step_i in range(max_horizon): rgb = env_results[0]["fixed_rgb"][0].cpu().numpy() qpos = env_results[0]["qpos"][0].cpu().numpy().astype(np.float32) qpos = qpos[_UPPER_BODY_IDS] # → 38-d if save_path: frames.append(rgb[:, :, ::-1].copy()) need_query = (action_chunk is None) or (step_in_chunk >= args_cli.n_action_exec) if need_query: # Resize image to 224x224 + HWC→CHW + normalize like training img_resized = cv2.resize(rgb, (args_cli.img_size, args_cli.img_size), interpolation=cv2.INTER_AREA) img_t = torch.from_numpy(img_resized).permute(2, 0, 1).float() / 255.0 img_t = img_t.unsqueeze(0).unsqueeze(0).cuda() # (B=1, n_cam=1, 3, H, W) qpos_norm = (qpos - pre_q) / pre_qs qpos_t = torch.from_numpy(qpos_norm).float().unsqueeze(0).cuda() # (1, 38) with torch.inference_mode(): a_hat = policy(qpos_t, img_t) # (1, chunk_size, 38) normalized action_chunk = a_hat[0].cpu().numpy() * pre_as + pre_a # un-normalize → (chunk_size, 38) if args_cli.temporal_agg: all_time_actions[step_i, step_i:step_i+chunk_size] = a_hat[0] step_in_chunk = 0 if args_cli.temporal_agg: # Average overlapping action predictions for this step (ALOHA recipe) actions_for_curr = all_time_actions[:, step_i] # (T, 38) in NORMALIZED space # mask out zero rows (those never written) mask = (actions_for_curr.abs().sum(dim=-1) > 0) if mask.any(): valid = actions_for_curr[mask] weights = torch.exp(-0.01 * torch.arange(valid.shape[0], device=valid.device, dtype=valid.dtype)) weights = weights / weights.sum() a_norm = (valid * weights.unsqueeze(-1)).sum(dim=0) # (38,) normalized a = a_norm.cpu().numpy() * pre_as + pre_a else: a = action_chunk[step_in_chunk].copy() else: a = action_chunk[step_in_chunk].copy() a_le, a_re, a_lh, a_rh = split_38d(a) ik_step(env, ik_l, ik_r, lw, rw, lrc, rrc, a_le, a_re, a_lh, a_rh, action_buf) env_results = env.step(action_buf) step_in_chunk += 1 if env_results[0]["success"].sum().item() == 1: success = True if save_path: frames.append(env_results[0]["fixed_rgb"][0].cpu().numpy()[:, :, ::-1].copy()) break if save_path and frames: h, w = frames[0].shape[:2] vw = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), 15, (w, h)) for f in frames: vw.write(f) vw.release() if success: n_succ += 1 tag = "SUCCESS" if success else "FAIL" print(f"[{ep+1}/{args_cli.num_episodes}] seq={seq_name} rand={eval_idx_seq[ep]} " f"{tag} ({step_i+1} steps) cum {n_succ}/{ep+1}", flush=True) print(f"\n[act-eval] {task_name} success rate: {n_succ}/{args_cli.num_episodes} " f"= {n_succ/args_cli.num_episodes*100:.1f}%", flush=True) env.close() simulation_app.close() if __name__ == "__main__": main()