eval.py

"""
eval.py — rollout evaluation harness for Diffusion Policy (low-dim M1 + image M2).

Loads a trained checkpoint, runs N rollouts in a live robosuite environment, and
reports the success rate. Obs handling and the receding-horizon loop live in
obs_pipeline.py (one shared implementation for train + eval); the checkpoint's
hparams carry obs_mode + camera/crop config so eval reconstructs the pipeline
without any CLI drift.

Usage:
    python eval.py --checkpoint runs/Square_.../best.ckpt --task Square --n-rollouts 50
    python eval.py --checkpoint runs/LiftImg_.../best.ckpt --task Lift --save-gifs 3

Image mode needs the mesa/osmesa offscreen render prefix:
    MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa \
    LD_LIBRARY_PATH=<conda-env>/lib:$LD_LIBRARY_PATH python eval.py ...
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from datasets import Normalizer
from obs_pipeline import LowDimObsPipeline, ImageObsPipeline, run_rollout


TASK_ENV_MAP = {"Lift": "Lift", "Can": "PickPlaceCan", "Square": "NutAssemblySquare"}
PROPRIO_KEYS = ("robot0_eef_pos", "robot0_eef_quat", "robot0_gripper_qpos")


# --------------------------------------------------------------------------- #
# Checkpoint loading
# --------------------------------------------------------------------------- #
def _apply_ema(net, ckpt):
    """Apply EMA shadow weights if safe (power schedule or sufficiently trained)."""
    ema_d = ckpt.get("ema")
    if ema_d is None:
        print("No EMA in checkpoint — using raw weights.")
        return
    if "power" in ema_d:
        for k, p in net.named_parameters():
            if k in ema_d["shadow"]:
                p.data.copy_(ema_d["shadow"][k].to(p.dtype))
        print(f"EMA weights applied (power-schedule, step={ema_d['_step']}).")
        return
    N = ckpt.get("global_step", 0)
    decay = ema_d["decay"]
    if decay ** N > 0.5:
        print(f"WARNING: EMA shadow is {decay**N*100:.0f}% initial random weights — "
              f"using raw model weights instead.")
        return
    for k, p in net.named_parameters():
        if k in ema_d["shadow"]:
            p.data.copy_(ema_d["shadow"][k].to(p.dtype))
    print("EMA weights applied.")


def load_policy(ckpt_path: str, device: torch.device, use_ema: bool = True):
    """Return (net, normalizer, obs_keys, hparams). Dispatches on hparams.obs_mode."""
    from train import DiffusionPolicyNet, ImageDiffusionPolicy

    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    hp = ckpt["hparams"]
    obs_mode = hp.get("obs_mode", "lowdim")
    normalizer = Normalizer.from_state_dict(ckpt["normalizer"])

    if obs_mode == "image":
        net = ImageDiffusionPolicy(
            cameras=hp["cameras"],
            proprio_dim=ckpt["obs_dim"],   # proprio dim stored here for image mode
            act_dim=ckpt["act_dim"],
            obs_horizon=hp["obs_horizon"],
            pred_horizon=hp["pred_horizon"],
            down_dims=tuple(hp["down_dims"]),
            n_timesteps=hp["n_timesteps"],
            num_kp=hp.get("num_kp", 32),
            normalizer=normalizer,
        ).to(device)
    else:
        net = DiffusionPolicyNet(
            obs_dim=ckpt["obs_dim"], obs_horizon=hp["obs_horizon"],
            obs_cond_dim=hp["obs_cond_dim"], act_dim=ckpt["act_dim"],
            pred_horizon=hp["pred_horizon"], down_dims=tuple(hp["down_dims"]),
            n_timesteps=hp["n_timesteps"],
        ).to(device)

    net.load_state_dict(ckpt["model_state"])
    if use_ema:
        _apply_ema(net, ckpt)
    else:
        print("Using raw model weights (no EMA).")
    net.normalizer = normalizer  # image policy needs it for un-normalization
    net.eval()
    return net, normalizer, ckpt["obs_keys"], hp


# --------------------------------------------------------------------------- #
# Environment
# --------------------------------------------------------------------------- #
def make_env(task: str, seed: int = 0, obs_mode: str = "lowdim",
             cameras=("agentview", "robot0_eye_in_hand"), size: int = 84,
             render_offscreen: bool = False, horizon: int = 500):
    import robosuite as suite
    env_name = TASK_ENV_MAP.get(task, task)
    image = obs_mode == "image"
    offscreen = image or render_offscreen
    return suite.make(
        env_name=env_name, robots="Panda",
        has_renderer=False, has_offscreen_renderer=offscreen,
        use_camera_obs=image, use_object_obs=True,
        camera_names=list(cameras) if offscreen else None,
        camera_heights=size, camera_widths=size,
        reward_shaping=False, control_freq=20,
        ignore_done=False, horizon=horizon, seed=seed,
    )


# --------------------------------------------------------------------------- #
# Pixel-alignment paranoia check (Phase 6/7)
# --------------------------------------------------------------------------- #
def verify_pixel_alignment(dataset_path: str, env, cameras, tol: float = 2 / 255) -> None:
    """Replay demo_0's states[0] through THIS env and compare to the stored frame 0.

    Catches residual train/eval render drift (driver/renderer). Aborts loudly on fail.
    """
    import h5py
    from scripts.regenerate_image_obs import extract_frame_obs

    with h5py.File(dataset_path, "r") as f:
        demos = sorted(f["data"].keys(), key=lambda n: int(n.split("_")[-1]))
        d0 = f["data"][demos[0]]
        state0 = np.asarray(d0["states"][0])
        stored = {c: np.asarray(d0["obs"][f"{c}_image"][0]) for c in cameras}

    env.reset()
    env.sim.set_state_from_flattened(state0)
    env.sim.forward()
    live = extract_frame_obs(env, list(cameras))
    for c in cameras:
        diff = np.abs(live[f"{c}_image"].astype(np.float32) - stored[c].astype(np.float32)).mean() / 255.0
        if diff > tol:
            raise SystemExit(
                f"PIXEL ALIGNMENT FAILED for {c}: mean|Δ|={diff:.4f} > {tol:.4f}. "
                f"Eval renders differ from the training dataset — aborting."
            )
        print(f"  pixel alignment {c}: mean|Δ|={diff:.5f}  OK")


# --------------------------------------------------------------------------- #
# Eval loop
# --------------------------------------------------------------------------- #
def _build_pipeline(obs_mode, obs_keys, hp, normalizer, device):
    if obs_mode == "image":
        return ImageObsPipeline(hp["cameras"], normalizer, hp["crop"],
                                hp["obs_horizon"], PROPRIO_KEYS, device)
    return LowDimObsPipeline(obs_keys, normalizer, hp["obs_horizon"])


def _build_predict_fn(net, normalizer, n_ddim_steps, device, obs_mode):
    if obs_mode == "image":
        return lambda p: net.predict_action(*p.get_cond_inputs(), n_ddim_steps=n_ddim_steps)
    return lambda p: normalizer.unnormalize_action(
        net.predict(p.get_obs_tensor(device), n_ddim_steps=n_ddim_steps).squeeze(0).cpu().numpy()
    )


def evaluate(args: argparse.Namespace) -> dict:
    device = torch.device(args.device)

    print(f"Loading checkpoint: {args.checkpoint}")
    net, normalizer, obs_keys, hp = load_policy(args.checkpoint, device, use_ema=not args.no_ema)
    obs_mode = hp.get("obs_mode", "lowdim")
    action_horizon = args.action_horizon or hp["action_horizon"]
    n_ddim_steps = args.n_ddim_steps or hp["n_ddim_steps"]
    cameras = hp.get("cameras", ["agentview", "robot0_eye_in_hand"])
    img_size = hp.get("img_size", 84)
    print(f"Task: {args.task}  |  obs_mode={obs_mode}  |  {args.n_rollouts} rollouts")

    results_dir = Path(args.results_dir) / Path(args.checkpoint).parent.name
    results_dir.mkdir(parents=True, exist_ok=True)

    pipeline = _build_pipeline(obs_mode, obs_keys, hp, normalizer, device)
    predict_fn = _build_predict_fn(net, normalizer, n_ddim_steps, device, obs_mode)
    save_gifs = args.save_gifs > 0

    successes, ep_lengths = [], []
    for i in range(args.n_rollouts):
        env = make_env(args.task, seed=args.seed + i, obs_mode=obs_mode,
                       cameras=cameras, size=img_size,
                       render_offscreen=save_gifs, horizon=args.horizon)
        if i == 0 and obs_mode == "image" and args.dataset:
            verify_pixel_alignment(args.dataset, env, cameras)

        render_cam = "agentview" if (save_gifs and i < args.save_gifs) else None
        success, length, frames = run_rollout(
            env, pipeline, predict_fn, max_steps=args.horizon,
            action_horizon=action_horizon, render_cam=render_cam, render_size=512,
        )
        env.close()

        successes.append(int(success))
        ep_lengths.append(length)
        tag = "SUCCESS" if success else "fail"
        print(f"  rollout {i+1:3d}/{args.n_rollouts}  {tag}  ({length} steps)")

        if render_cam and frames:
            _save_gif(frames, results_dir / f"{args.task.lower()}_{obs_mode}_{tag}_{i+1}.gif")

    success_rate = float(np.mean(successes))
    avg_len = float(np.mean(ep_lengths))
    summary = {
        "task": args.task, "obs_mode": obs_mode, "checkpoint": args.checkpoint,
        "n_rollouts": args.n_rollouts, "success_rate": success_rate,
        "avg_ep_len": avg_len, "successes": successes,
    }
    print(f"\nSuccess rate: {success_rate*100:.1f}%  ({sum(successes)}/{args.n_rollouts})")
    print(f"Avg episode length: {avg_len:.0f} steps")

    with open(results_dir / "eval_results.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(f"Results saved: {results_dir/'eval_results.json'}")
    return summary


def _save_gif(frames: list, path: Path, fps: int = 20) -> None:
    import imageio
    imageio.mimsave(str(path), [np.asarray(f) for f in frames], fps=fps, loop=0)
    print(f"    saved {path.name}  ({len(frames)} frames)")


# --------------------------------------------------------------------------- #
# Args
# --------------------------------------------------------------------------- #
def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument("--checkpoint", required=True, help="Path to .ckpt file")
    p.add_argument("--task", default="Lift", help="Lift / Can / Square")
    p.add_argument("--n-rollouts", type=int, default=50)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--horizon", type=int, default=500, help="Max env steps per rollout")
    p.add_argument("--save-gifs", type=int, default=0, help="Save GIFs for the first N rollouts")
    p.add_argument("--dataset", default=None,
                   help="Image dataset path for the one-time pixel-alignment check")
    p.add_argument("--results-dir", default="results")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--no-ema", action="store_true", help="Use raw model weights, not EMA")
    p.add_argument("--n-ddim-steps", type=int, default=0, help="0 = checkpoint default")
    p.add_argument("--action-horizon", type=int, default=0, help="0 = checkpoint default")
    return p.parse_args()


if __name__ == "__main__":
    evaluate(parse_args())