Blame for PWM_code/policy_rl/eval_in_sim.py

"""
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()