""" Evaluate a Diffusion Policy checkpoint on the Humanoid-Open-Drawer-v0 IsaacLab env. Runs N episodes, reports success rate. Optionally saves rollout videos. """ 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/external/diffusion_policy") from omni.isaac.lab.app import AppLauncher parser = argparse.ArgumentParser() parser.add_argument("--task", type=str, default="Humanoid-Open-Drawer-v0") parser.add_argument("--num_episodes", type=int, default=20) parser.add_argument("--ckpt", type=str, required=True) parser.add_argument("--output_dir", type=str, default="/home/hongyi/PWM/policy_rl/eval_out") parser.add_argument("--save_video", action="store_true") parser.add_argument("--n_action_exec", type=int, default=8, help="Execute this many actions from each DP chunk before re-querying.") parser.add_argument("--img_size", type=int, default=96) parser.add_argument("--room_idx", type=int, default=1) parser.add_argument("--table_idx", type=int, default=2) parser.add_argument("--rand_offset", type=int, default=4000, help="Sequential randomize_idx start (rand_mode=sequential), " "or shift into the shuffled pool (rand_mode=shuffled).") parser.add_argument("--rand_mode", choices=["sequential", "shuffled"], default="sequential", help="sequential: walk randomize_idx by +1 (matches " "train/collection); shuffled: per-task seeded shuffle " "of [0,9999], spans full 100x100 placement grid.") parser.add_argument("--eval_torch_seed", type=int, default=42, help="Seed for torch RNG before the eval rollout loop. " "Makes DDPM action sampling deterministic across runs " "of the same ckpt+rand_offset. Set to a different " "value to sample a different DDPM rollout trace.") parser.add_argument("--proprio_mode", choices=["full50", "upper38"], default="full50", help="full50: sim qpos as-is (50-dim, matches old BC). " "upper38: slice sim qpos via _UPPER_BODY_IDS (matches " "tableHF-style 38-dim replays).") 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 tqdm import gymnasium as gym import humanoid.tasks # noqa import omni.isaac.lab_tasks # noqa import omni.usd 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 diffusion_policy.workspace.train_diffusion_unet_hybrid_workspace import TrainDiffusionUnetHybridWorkspace # Upper-body qpos indices (same as replay_hdf5_to_diffsynth.py _UPPER_BODY_IDS). _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 seed_map = { "Humanoid-Open-Drawer-v0": 6, "Humanoid-Close-Drawer-v0": 7, "Humanoid-Open-Laptop-v0": 8, "Humanoid-Stack-Can-v0": 9, "Humanoid-Push-Box-v0": 10, "Humanoid-Flip-Mug-v0": 11, "Humanoid-Sort-Cans-v0": 12, "Humanoid-Pour-Balls-v0": 13, "Humanoid-Insert-Cans-v0": 14, "Humanoid-Unload-Cans-v0": 15, "Humanoid-Stack-Can-Into-Drawer-v0": 16, } def load_dp_policy(ckpt_path: str, device: str = "cuda:0"): """Load DP from a checkpoint (uses saved cfg + state).""" import dill payload = torch.load(open(ckpt_path, "rb"), pickle_module=dill, map_location="cpu") cfg = payload["cfg"] ws = TrainDiffusionUnetHybridWorkspace(cfg) ws.load_payload(payload, exclude_keys=None, include_keys=None) # Prefer EMA model policy = ws.ema_model if cfg.training.use_ema else ws.model policy.to(device) policy.eval() return policy def prep_obs(rgb_hist: deque, qpos_hist: deque, img_size: int, device: str): """Build DP obs_dict: image (1, To, 3, H, W) + proprio (1, To, 50).""" imgs = [] for rgb in rgb_hist: img = cv2.resize(rgb, (img_size, img_size), interpolation=cv2.INTER_AREA) img = np.transpose(img, (2, 0, 1)).astype(np.float32) / 255.0 imgs.append(img) imgs = np.stack(imgs, axis=0) # (To, 3, H, W) proprio = np.stack(list(qpos_hist), axis=0).astype(np.float32) # (To, 50) obs_dict = { "image": torch.from_numpy(imgs).unsqueeze(0).to(device), "proprio": torch.from_numpy(proprio).unsqueeze(0).to(device), } return obs_dict def split_action_38d(action: np.ndarray): """action: (38,) -> (left_ee(7,), right_ee(7,), left_qpos(12,), right_qpos(12,)).""" return action[:7], action[7:14], action[14:26], action[26:38] def run_episode(env, env_cfg_reset, left_ik, right_ik, left_ik_w, right_ik_w, left_ik_r, right_ik_r, action_buf, init_poses, task_name, seq_name, policy, n_obs_steps, n_action_exec, img_size, max_horizon, save_video_path=None, proprio_mode="full50"): env_cfg_reset() env.reset() left_ik.reset(); right_ik.reset() # Init pose warmup (same as collection) 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"] for _ in range(100): ik_step(env, left_ik, right_ik, left_ik_w, right_ik_w, left_ik_r, right_ik_r, le, re, ld, rd, action_buf) env_results = env.step(action_buf) rgb_hist = deque(maxlen=n_obs_steps) qpos_hist = deque(maxlen=n_obs_steps) rgb = env_results[0]["fixed_rgb"][0].cpu().numpy() qpos = env_results[0]["qpos"][0].cpu().numpy().astype(np.float32) if proprio_mode == "upper38": qpos = qpos[_UPPER_BODY_IDS] for _ in range(n_obs_steps): rgb_hist.append(rgb) qpos_hist.append(qpos) frames = [] if save_video_path else None pending_actions = deque() device = next(policy.parameters()).device result = False for step_i in range(max_horizon): if save_video_path: frames.append(rgb[:, :, ::-1].copy()) if not pending_actions: obs_dict = prep_obs(rgb_hist, qpos_hist, img_size, device) with torch.inference_mode(): out = policy.predict_action(obs_dict) actions = out["action"][0].cpu().numpy() # (n_action_steps, 38) for a in actions[:n_action_exec]: pending_actions.append(a) a = pending_actions.popleft() a_le, a_re, a_lh, a_rh = split_action_38d(a) ik_step(env, left_ik, right_ik, left_ik_w, right_ik_w, left_ik_r, right_ik_r, a_le, a_re, a_lh, a_rh, action_buf) env_results = env.step(action_buf) rgb = env_results[0]["fixed_rgb"][0].cpu().numpy() qpos = env_results[0]["qpos"][0].cpu().numpy().astype(np.float32) if proprio_mode == "upper38": qpos = qpos[_UPPER_BODY_IDS] rgb_hist.append(rgb) qpos_hist.append(qpos) if env_results[0]["success"].sum().item() == 1: result = True if save_video_path: frames.append(rgb[:, :, ::-1].copy()) break if save_video_path and frames: h, w = frames[0].shape[:2] out = cv2.VideoWriter(save_video_path, cv2.VideoWriter_fourcc(*"mp4v"), 15, (w, h)) for f in frames: out.write(f) out.release() return result, (step_i + 1) def main(): Path(args_cli.output_dir).mkdir(parents=True, exist_ok=True) # Seed torch (and numpy) so DDPM action sampling is deterministic across # runs of the same ckpt + rand_offset. The env's initial-state seed and # rand_idx are already deterministic (per-task seeded shuffle); this just # nails down the DP policy's stochasticity. print(f"[seed] eval_torch_seed = {args_cli.eval_torch_seed}", flush=True) torch.manual_seed(args_cli.eval_torch_seed) torch.cuda.manual_seed_all(args_cli.eval_torch_seed) np.random.seed(args_cli.eval_torch_seed) print(f"Loading DP checkpoint: {args_cli.ckpt}", flush=True) policy = load_dp_policy(args_cli.ckpt, device="cuda:0") cfg = None # re-load cfg to read horizon/n_obs_steps import dill payload = torch.load(open(args_cli.ckpt, "rb"), pickle_module=dill, map_location="cpu") cfg = payload["cfg"] n_obs_steps = cfg.n_obs_steps print(f"DP: n_obs_steps={n_obs_steps}, horizon={cfg.horizon}, n_action_steps={cfg.n_action_steps}", flush=True) task = args_cli.task task_name = task[9:-3] random.seed(seed_map[task]) rand_idxes = list(range(10000)) random.shuffle(rand_idxes) 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] max_horizon = TASK_MAX_HORIZON[task] 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: BaseEnv = gym.make(task, cfg=env_cfg) # Randomize-idx schedule: sequential (default, matches train/collection) # or shuffled (spans full 100x100 placement grid). if args_cli.rand_mode == "shuffled": rand_pool = list(range(10000)) random.Random(seed_map[task]).shuffle(rand_pool) eval_idx_seq = [rand_pool[(args_cli.rand_offset + i) % 10000] for i in range(args_cli.num_episodes)] print(f"[rand] shuffled randomize_idx schedule (first 8): " f"{eval_idx_seq[:8]}", flush=True) rand_pos = [0] curr_random_idx = [eval_idx_seq[0]] def reset_with_new_seed(): rand_pos[0] += 1 if rand_pos[0] < len(eval_idx_seq): curr_random_idx[0] = eval_idx_seq[rand_pos[0]] env.cfg.randomize_idx = curr_random_idx[0] else: print(f"[rand] sequential randomize_idx starting at {args_cli.rand_offset}", flush=True) curr_random_idx = [args_cli.rand_offset] def reset_with_new_seed(): curr_random_idx[0] += 1 env.cfg.randomize_idx = curr_random_idx[0] env.cfg.randomize_idx = curr_random_idx[0] env.reset() left_ik = DifferentialIKController( DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), num_envs=1, device=env.sim.device, ) right_ik = DifferentialIKController( DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="pinv"), num_envs=1, device=env.sim.device, ) left_ik_w = torch.zeros(1, left_ik.action_dim, device=env.robot.device) left_ik_r = torch.zeros(1, left_ik.action_dim, device=env.robot.device) right_ik_w = torch.zeros(1, right_ik.action_dim, device=env.robot.device) right_ik_r = torch.zeros(1, right_ik.action_dim, device=env.robot.device) action_buf = torch.zeros((1, env.num_actions), device=env.robot.device) n_success = 0 episode_log = [] for ep_i in range(args_cli.num_episodes): seq_name = episode_list[ep_i % len(episode_list)][0] # First run to a tmp path, then rename by outcome. tmp_path = None if args_cli.save_video: tmp_path = os.path.join(args_cli.output_dir, f"_ep{ep_i:03d}_tmp.mp4") ok, steps = run_episode( env, reset_with_new_seed, left_ik, right_ik, left_ik_w, right_ik_w, left_ik_r, right_ik_r, action_buf, init_poses, task_name, seq_name, policy, n_obs_steps, args_cli.n_action_exec, args_cli.img_size, max_horizon, tmp_path, proprio_mode=args_cli.proprio_mode, ) if args_cli.save_video and tmp_path is not None and os.path.exists(tmp_path): tag = "success" if ok else "fail" final_name = f"{tag}_ep{ep_i:03d}.mp4" os.rename(tmp_path, os.path.join(args_cli.output_dir, final_name)) if ok: n_success += 1 tag = "SUCCESS" if ok else "FAIL" print(f"[{ep_i+1}/{args_cli.num_episodes}] seq={seq_name} rand={curr_random_idx[0]} {tag} ({steps} steps) cum {n_success}/{ep_i+1}", flush=True) episode_log.append(dict(ep_i=ep_i, seq_name=seq_name, success=ok, steps=steps)) print(f"\n============================================================", flush=True) print(f"DP Open-Drawer eval: {n_success}/{args_cli.num_episodes} = {n_success/args_cli.num_episodes*100:.1f}%", flush=True) print(f"============================================================", flush=True) import json with open(os.path.join(args_cli.output_dir, "eval_results.json"), "w") as f: json.dump(dict( ckpt=args_cli.ckpt, task=task, num_episodes=args_cli.num_episodes, n_success=n_success, success_rate=n_success / args_cli.num_episodes, episodes=episode_log, ), f, indent=2) env.close() if __name__ == "__main__": main() simulation_app.close()