scripts/infer_single_arm_multiview_ctrlworld.py

"""
Full-episode GT-anchored rollout using Ctrl-World on single_arm/multiview dataset.

For each episode:
  1. Compute interact_num based on episode length to match original duration.
  2. Use GT actions (replay mode) — no policy model needed.
  3. Generate video frames autoregressively using Ctrl-World.
  4. Save predicted video + metrics for benchmark evaluation.

Requires: prepare_ctrlworld_single_arm_multiview.py to be run first.
"""

import sys
import os
import importlib
import json
import datetime
from pathlib import Path
from argparse import ArgumentParser

import numpy as np
import torch
import einops
import mediapy
import piq

ctrl_world_dir = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "models", "Ctrl-World"
)
sys.path.insert(0, ctrl_world_dir)

from models.pipeline_ctrl_world import CtrlWorldDiffusionPipeline
from models.ctrl_world import CrtlWorld
from decord import VideoReader, cpu
from accelerate import Accelerator


DATASET_EXAMPLE_BASE = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "models", "Ctrl-World", "dataset_example"
)


class CtrlWorldAgent:
    def __init__(self, args):
        args.val_model_path = args.ckpt_path
        self.args = args
        self.accelerator = Accelerator()
        self.device = self.accelerator.device
        self.dtype = args.dtype

        self.model = CrtlWorld(args)
        self.model.load_state_dict(torch.load(args.val_model_path))
        self.model.to(self.device).to(self.dtype)
        self.model.eval()
        print("Ctrl-World model loaded")

        with open(args.data_stat_path, "r") as f:
            data_stat = json.load(f)
            self.state_p01 = np.array(data_stat["state_01"])[None, :]
            self.state_p99 = np.array(data_stat["state_99"])[None, :]

    def normalize_bound(self, data, data_min, data_max, clip_min=-1, clip_max=1, eps=1e-8):
        ndata = 2 * (data - data_min) / (data_max - data_min + eps) - 1
        return np.clip(ndata, clip_min, clip_max)

    def get_traj_info(self, episode_id, start_idx=0, steps=8):
        val_dataset_dir = self.args.val_dataset_dir
        annotation_path = f"{val_dataset_dir}/annotation/val/{episode_id}.json"
        with open(annotation_path) as f:
            anno = json.load(f)
            length = anno["video_length"]

        frames_ids = np.arange(start_idx, start_idx + steps)
        max_ids = np.ones_like(frames_ids) * (length - 1)
        frames_ids = np.min([frames_ids, max_ids], axis=0).astype(int)

        instruction = anno["texts"][0]
        car_action = np.array(anno["states"])
        car_action = car_action[frames_ids]
        joint_pos = np.array(anno["joints"])
        joint_pos = joint_pos[frames_ids]

        video_dict = []
        video_latent = []
        for vid_info in anno["videos"]:
            video_path = f"{val_dataset_dir}/{vid_info['video_path']}"
            vr = VideoReader(video_path, ctx=cpu(0), num_threads=2)
            actual_video_len = len(vr)
            if length > actual_video_len:
                length = actual_video_len
                frames_ids = np.clip(frames_ids, 0, length - 1)
            try:
                true_video = vr.get_batch(range(length)).asnumpy()
            except:
                true_video = vr.get_batch(range(length)).numpy()
            true_video = true_video[frames_ids]
            video_dict.append(true_video)

            device = self.device
            true_video_t = torch.from_numpy(true_video).to(self.dtype).to(device)
            x = true_video_t.permute(0, 3, 1, 2) / 255.0 * 2 - 1
            vae = self.model.pipeline.vae
            with torch.no_grad():
                latents = []
                for i in range(0, len(x), 32):
                    batch = x[i:i + 32]
                    latent = vae.encode(batch).latent_dist.sample().mul_(vae.config.scaling_factor)
                    latents.append(latent)
                x = torch.cat(latents, dim=0)
            video_latent.append(x)

        return car_action, joint_pos, video_dict, video_latent, instruction

    def forward_wm(self, action_cond, video_latent_true, video_latent_cond, his_cond=None, text=None):
        args = self.args
        image_cond = video_latent_cond

        action_cond = self.normalize_bound(action_cond, self.state_p01, self.state_p99)
        action_cond = torch.tensor(action_cond).unsqueeze(0).to(self.device).to(self.dtype)

        with torch.no_grad():
            if text is not None:
                text_token = self.model.action_encoder(
                    action_cond, text, self.model.tokenizer, self.model.text_encoder
                )
            else:
                text_token = self.model.action_encoder(action_cond)
            pipeline = self.model.pipeline

            _, latents = CtrlWorldDiffusionPipeline.__call__(
                pipeline,
                image=image_cond,
                text=text_token,
                width=args.width,
                height=int(args.height * 3),
                num_frames=args.num_frames,
                history=his_cond,
                num_inference_steps=args.num_inference_steps,
                decode_chunk_size=args.decode_chunk_size,
                max_guidance_scale=args.guidance_scale,
                fps=args.fps,
                motion_bucket_id=args.motion_bucket_id,
                mask=None,
                output_type="latent",
                return_dict=False,
                frame_level_cond=True,
            )
        latents = einops.rearrange(latents, "b f c (m h) (n w) -> (b m n) f c h w", m=3, n=1)

        # Decode GT
        true_video = torch.stack(video_latent_true, dim=0)
        decoded_video = []
        bsz, frame_num = true_video.shape[:2]
        true_video_flat = true_video.flatten(0, 1)
        for i in range(0, true_video_flat.shape[0], args.decode_chunk_size):
            chunk = true_video_flat[i:i + args.decode_chunk_size] / pipeline.vae.config.scaling_factor
            decoded_video.append(pipeline.vae.decode(chunk, num_frames=chunk.shape[0]).sample)
        true_video_dec = torch.cat(decoded_video, dim=0)
        true_video_dec = true_video_dec.reshape(bsz, frame_num, *true_video_dec.shape[1:])
        true_video_dec = ((true_video_dec / 2.0 + 0.5).clamp(0, 1) * 255)
        true_video_dec = true_video_dec.detach().to(torch.float32).cpu().numpy().transpose(0, 1, 3, 4, 2).astype(np.uint8)

        # Decode predicted
        decoded_video = []
        bsz, frame_num = latents.shape[:2]
        x = latents.flatten(0, 1)
        for i in range(0, x.shape[0], args.decode_chunk_size):
            chunk = x[i:i + args.decode_chunk_size] / pipeline.vae.config.scaling_factor
            decoded_video.append(pipeline.vae.decode(chunk, num_frames=chunk.shape[0]).sample)
        videos = torch.cat(decoded_video, dim=0)
        videos = videos.reshape(bsz, frame_num, *videos.shape[1:])
        videos = ((videos / 2.0 + 0.5).clamp(0, 1) * 255)
        videos = videos.detach().to(torch.float32).cpu().numpy().transpose(0, 1, 3, 4, 2).astype(np.uint8)

        return true_video_dec, videos, latents


def compute_metrics(pred_frames, gt_frames):
    """Compute PSNR, SSIM, LPIPS between pred and gt frame arrays."""
    x = torch.clamp(torch.from_numpy(pred_frames.copy()) / 255.0, 0, 1).permute(0, 3, 1, 2)
    y = torch.clamp(torch.from_numpy(gt_frames.copy()) / 255.0, 0, 1).permute(0, 3, 1, 2)
    psnr_val = piq.psnr(x, y).mean().item()
    ssim_val = piq.ssim(x, y).mean().item()
    lpips_val = piq.LPIPS()(x, y).mean().item()
    return {"psnr": psnr_val, "ssim": ssim_val, "lpips": lpips_val}


def run_episode(agent, episode_id, save_dir, input_save_dir=None):
    """Run full-episode replay generation for one episode."""
    args = agent.args
    pred_step = args.num_frames  # 5
    num_history = args.num_history  # 6

    # Read annotation to get episode length
    anno_path = f"{args.val_dataset_dir}/annotation/val/{episode_id}.json"
    with open(anno_path) as f:
        anno = json.load(f)
    episode_length = anno["video_length"]

    # Calculate interact_num to cover full episode
    # Each iteration uses frames[i*4 : i*4+5], so max end = (interact_num-1)*4+5 <= episode_length
    interact_num = (episode_length - 1) // (pred_step - 1)
    if interact_num < 1:
        print(f"  Episode {episode_id} too short ({episode_length} frames), skipping")
        return None

    total_steps_needed = (interact_num - 1) * (pred_step - 1) + pred_step

    # Get trajectory info
    eef_gt, joint_pos_gt, video_dict, video_latents, instruction = agent.get_traj_info(
        episode_id, start_idx=0, steps=min(total_steps_needed, episode_length)
    )

    # Initialize history buffers
    his_cond = []
    his_eef = []
    first_latent = torch.cat([v[0] for v in video_latents], dim=1).unsqueeze(0)
    for _ in range(num_history * 4):
        his_cond.append(first_latent)
        his_eef.append(eef_gt[0:1])

    video_to_save_pred = []
    video_to_save_gt = []
    history_idx = [0, 0, -8, -6, -4, -2]

    for i in range(interact_num):
        start_id = int(i * (pred_step - 1))
        end_id = start_id + pred_step

        if end_id > len(eef_gt):
            break

        video_latent_true = [v[start_id:end_id] for v in video_latents]
        cartesian_pose = eef_gt[start_id:end_id]

        # Prepare history
        his_pose = np.concatenate([his_eef[idx] for idx in history_idx], axis=0)
        action_cond = np.concatenate([his_pose, cartesian_pose], axis=0)
        his_cond_input = torch.cat([his_cond[idx] for idx in history_idx], dim=0).unsqueeze(0)
        current_latent = his_cond[-1]

        # Forward world model
        true_videos, pred_videos, predicted_latents = agent.forward_wm(
            action_cond, video_latent_true, current_latent,
            his_cond=his_cond_input,
            text=instruction if args.text_cond else None,
        )

        # Update history
        his_eef.append(cartesian_pose[pred_step - 1:pred_step])
        his_cond.append(
            torch.cat([v[pred_step - 1] for v in predicted_latents], dim=1).unsqueeze(0)
        )

        # Collect frames — all 3 views
        # true_videos/pred_videos shape: (3_views, frames, H, W, 3)
        if i == interact_num - 1:
            video_to_save_pred.append(pred_videos)
            video_to_save_gt.append(true_videos)
        else:
            video_to_save_pred.append(pred_videos[:, :pred_step - 1])
            video_to_save_gt.append(true_videos[:, :pred_step - 1])

        if (i + 1) % 10 == 0:
            print(f"    Step {i+1}/{interact_num}")

    if not video_to_save_pred:
        return None

    # Concatenate — shape: (3_views, total_frames, H, W, 3)
    concat_pred = np.concatenate(video_to_save_pred, axis=1)
    concat_gt = np.concatenate(video_to_save_gt, axis=1)
    num_views = concat_pred.shape[0]
    min_len = min(concat_pred.shape[1], concat_gt.shape[1])
    concat_pred = concat_pred[:, :min_len]
    concat_gt = concat_gt[:, :min_len]

    # Build all-views strips: (frames, H, 3*W, 3)
    pred_strip = np.concatenate([concat_pred[v] for v in range(num_views)], axis=2)
    gt_strip = np.concatenate([concat_gt[v] for v in range(num_views)], axis=2)
    # Full pred: GT top + Pred bottom, all views side by side
    full_pred = np.concatenate([gt_strip, pred_strip], axis=1)

    # Save output (predictions)
    ep_save_dir = Path(save_dir) / f"episode_{episode_id}"
    ep_save_dir.mkdir(parents=True, exist_ok=True)

    mediapy.write_video(str(ep_save_dir / "full_pred.mp4"), full_pred, fps=5)
    mediapy.write_video(str(ep_save_dir / "pred_all_views.mp4"), pred_strip, fps=5)
    mediapy.write_video(str(ep_save_dir / "gt_all_views.mp4"), gt_strip, fps=5)

    # Save input (GT video, actions, metadata)
    if input_save_dir is not None:
        ep_input_dir = Path(input_save_dir) / f"episode_{episode_id}"
        ep_input_dir.mkdir(parents=True, exist_ok=True)

        mediapy.write_video(str(ep_input_dir / "full_gt.mp4"), gt_strip, fps=5)

        # Save all 3 views GT videos from annotation source
        anno_path = f"{args.val_dataset_dir}/annotation/val/{episode_id}.json"
        with open(anno_path) as f:
            anno_data = json.load(f)
        for vid_idx, vid_info in enumerate(anno_data["videos"]):
            src_video = Path(args.val_dataset_dir) / vid_info["video_path"]
            dst_video = ep_input_dir / f"view_{vid_idx}.mp4"
            if src_video.exists() and not dst_video.exists():
                import shutil
                shutil.copy2(str(src_video), str(dst_video))

        # Save actions and states
        input_meta = {
            "episode_id": episode_id,
            "instruction": instruction,
            "num_frames": min_len,
            "episode_length_original": episode_length,
            "fps": 5,
            "interact_num": interact_num,
            "pred_step": pred_step,
            "mode": "replay",
            "states": anno_data["states"][:min_len],
            "joints": anno_data["joints"][:min_len],
        }
        with open(ep_input_dir / "metadata.json", "w") as f:
            json.dump(input_meta, f, indent=2)

    # Metrics (per-view)
    view_names = ["exterior_1_left", "exterior_2_left", "wrist_left"]
    w_per_view = concat_pred.shape[3]
    per_view_metrics = {}
    for v_i in range(num_views):
        per_view_metrics[view_names[v_i]] = compute_metrics(concat_pred[v_i], concat_gt[v_i])
    avg_psnr = np.mean([m["psnr"] for m in per_view_metrics.values()])
    avg_ssim = np.mean([m["ssim"] for m in per_view_metrics.values()])
    avg_lpips = np.mean([m["lpips"] for m in per_view_metrics.values()])
    metrics = {"psnr": float(avg_psnr), "ssim": float(avg_ssim), "lpips": float(avg_lpips)}
    metrics["per_view"] = per_view_metrics
    metrics.update({
        "episode_id": episode_id,
        "instruction": instruction,
        "num_frames_pred": min_len,
        "num_frames_gt": episode_length,
        "interact_num": interact_num,
        "mode": "replay",
    })
    with open(ep_save_dir / "metrics.json", "w") as f:
        json.dump(metrics, f, indent=2)

    return metrics


def _import_cache_module(backend, module_name):
    project_root = os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))
    if project_root not in sys.path:
        sys.path.insert(0, project_root)
    return importlib.import_module(f"methods.cache_strategy.{backend}.{module_name}")


def _import_pruning_module(backend, module_name):
    project_root = os.path.normpath(os.path.join(os.path.dirname(__file__), ".."))
    if project_root not in sys.path:
        sys.path.insert(0, project_root)
    return importlib.import_module(f"methods.prunning.{backend}.{module_name}")


def main():
    parser = ArgumentParser()
    parser.add_argument("--svd_model_path", type=str, required=True)
    parser.add_argument("--clip_model_path", type=str, required=True)
    parser.add_argument("--ckpt_path", type=str, required=True)
    parser.add_argument("--subset", choices=["makovian", "non_makovian"], required=True)
    parser.add_argument("--dataset_example_dir", type=str, default=DATASET_EXAMPLE_BASE)
    parser.add_argument("--dataset_meta_info_path", type=str,
                        default="./models/Ctrl-World/dataset_meta_info")
    parser.add_argument("--save_dir", type=str, default=None)
    parser.add_argument("--input_save_dir", type=str, default=None)
    parser.add_argument("--num_episodes", type=int, default=None)
    parser.add_argument("--num_inference_steps", type=int, default=None,
                        help="Override number of denoising steps (default: use model config, typically 50).")

    # Cache backend args (mutually exclusive)
    _import_cache_module("WorldCache", "config").add_worldcache_args(parser)
    _import_cache_module("DiCache", "config").add_dicache_args(parser)
    _import_cache_module("FasterCache", "config").add_fastercache_args(parser)

    # Pruning/sparse attention backend args
    _import_pruning_module("SiTo", "config").add_sito_args(parser)
    _import_pruning_module("importance_token_merge", "config").add_itm_args(parser)

    args = parser.parse_args()

    from methods.cache_strategy.ctrl_world_utils import validate_ctrl_world_backend_args
    validate_ctrl_world_backend_args(args)

    # Determine paths
    subset_name = f"single_arm_multiview_{args.subset}"
    val_dataset_dir = os.path.join(args.dataset_example_dir, subset_name)

    if args.save_dir is None:
        args.save_dir = (
            f"/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/"
            f"video_gen_physics/sampling_dataset/dense/single_arm/output/multiview/"
            f"ctrlworld/{args.subset}"
        )

    if args.input_save_dir is None:
        args.input_save_dir = (
            f"/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/"
            f"video_gen_physics/sampling_dataset/dense/single_arm/input/multiview/"
            f"ctrlworld/{args.subset}"
        )

    # Build model args from config
    sys.path.insert(0, ctrl_world_dir)
    from config import wm_args
    model_args = wm_args(task_type="replay")
    model_args.svd_model_path = args.svd_model_path
    model_args.clip_model_path = args.clip_model_path
    model_args.ckpt_path = args.ckpt_path
    model_args.val_model_path = args.ckpt_path
    model_args.val_dataset_dir = val_dataset_dir
    model_args.dataset_meta_info_path = args.dataset_meta_info_path
    model_args.data_stat_path = os.path.join(args.dataset_meta_info_path, "droid_subset", "stat.json")
    model_args.__post_init__()
    # Override val_dataset_dir after __post_init__ since it may reset
    model_args.val_dataset_dir = val_dataset_dir

    if args.num_inference_steps is not None:
        model_args.num_inference_steps = args.num_inference_steps
        print(f"[Override] num_inference_steps = {args.num_inference_steps}")

    # Create agent
    agent = CtrlWorldAgent(model_args)

    # Enable cache backend on UNet (mutually exclusive, validated above)
    if getattr(args, "use_worldcache", False):
        adapter = _import_cache_module("WorldCache", "adapter")
        adapter.enable_worldcache(
            agent.model.unet,
            num_steps=model_args.num_inference_steps,
            rel_l1_thresh=args.worldcache_rel_l1_thresh,
            ret_ratio=args.worldcache_ret_ratio,
            probe_depth=args.worldcache_probe_depth,
            motion_sensitivity=args.worldcache_motion_sensitivity,
            hf_enabled=args.worldcache_hf_enabled,
            hf_thresh=args.worldcache_hf_thresh,
            saliency_enabled=args.worldcache_saliency_enabled,
            saliency_weight=args.worldcache_saliency_weight,
            osi_enabled=args.worldcache_osi_enabled,
            dynamic_decay=args.worldcache_dynamic_decay,
        )
        print(f"[Cache] WorldCache enabled: thresh={args.worldcache_rel_l1_thresh}, "
              f"ret_ratio={args.worldcache_ret_ratio}, probe_depth={args.worldcache_probe_depth}")
    if getattr(args, "use_dicache", False):
        adapter = _import_cache_module("DiCache", "adapter")
        adapter.enable_dicache(
            agent.model.unet,
            num_steps=model_args.num_inference_steps,
            rel_l1_thresh=args.dicache_rel_l1_thresh,
            ret_ratio=args.dicache_ret_ratio,
            probe_depth=args.dicache_probe_depth,
        )
        print(f"[Cache] DiCache enabled: thresh={args.dicache_rel_l1_thresh}, "
              f"ret_ratio={args.dicache_ret_ratio}, probe_depth={args.dicache_probe_depth}")
    if getattr(args, "use_fastercache", False):
        adapter = _import_cache_module("FasterCache", "adapter")
        adapter.enable_fastercache(
            agent.model.unet,
            start_step=args.fastercache_start_step,
            model_interval=args.fastercache_model_interval,
            block_interval=args.fastercache_block_interval,
            first_layers_fp=2,
        )
        print(f"[Cache] FasterCache enabled: start_step={args.fastercache_start_step}, "
              f"model_interval={args.fastercache_model_interval}, block_interval={args.fastercache_block_interval}")
    if getattr(args, "use_sito", False):
        if getattr(args, "sito_spatiotemporal_hold", False):
            st_hold = _import_pruning_module("SiTo", "spatiotemporal_hold")
            st_hold.enable_spatiotemporal_hold(
                agent.model.unet,
                keep_ratio=args.sito_st_keep_ratio,
                max_downsample_ratio=args.sito_max_downsample_ratio,
                recompute_every=(args.sito_plan_recompute_every or 999999),
            )
        else:
            adapter = _import_pruning_module("SiTo", "adapter")
            adapter.enable_sito(
                agent.model.unet,
                start_layer_idx=args.sito_start_layer_idx or 0,
                prune_ratio=args.sito_prune_ratio,
                patch_h=args.sito_patch_h,
                patch_w=args.sito_patch_w,
                noise_alpha=args.sito_noise_alpha,
                sim_beta=args.sito_sim_beta,
                max_downsample_ratio=args.sito_max_downsample_ratio,
                plan_recompute_every=getattr(args, "sito_plan_recompute_every", 0),
            )
    if getattr(args, "use_itm", False):
        itm_state = {}
        if getattr(args, "itm_spatiotemporal_hold", False):
            st_hold = _import_pruning_module("SiTo", "spatiotemporal_hold")
            st_hold.enable_spatiotemporal_hold(
                agent.model.unet,
                keep_ratio=args.itm_st_keep_ratio,
                max_downsample_ratio=args.itm_max_downsample_ratio,
                recompute_every=(args.itm_plan_recompute_every or 999999),
                similarity_recover=True,
            )
        elif getattr(args, "itm_block_hold", False):
            block_hold = _import_pruning_module("importance_token_merge", "block_hold")
            itm_state["prune_from_step"] = args.itm_prune_from_step
            itm_state["merge_from_step"] = args.itm_merge_from_step
            block_hold.enable_itm_block_hold(
                agent.model.unet,
                itm_state=itm_state,
                start_layer_idx=args.itm_start_layer_idx or 0,
                compress_ratio=args.itm_compress_ratio,
                max_downsample_ratio=args.itm_max_downsample_ratio,
                self_importance=True,
                plan_recompute_every=getattr(args, "itm_plan_recompute_every", 0),
            )
            # Block-hold with self-importance derives token scores from the
            # hidden states, so it does NOT need classifier-free guidance and
            # can stay at guidance_scale=1 (single batch) — this is what keeps
            # it faster than dense instead of paying the 2x CFG cost.
        else:
            adapter = _import_pruning_module("importance_token_merge", "adapter")
            adapter.enable_itm(
                agent.model.unet,
                itm_state=itm_state,
                start_layer_idx=args.itm_start_layer_idx or 0,
                compress_ratio=args.itm_compress_ratio,
                prune_from_step=args.itm_prune_from_step,
                merge_from_step=args.itm_merge_from_step,
                merge_attn=args.itm_merge_attn,
                merge_crossattn=args.itm_merge_crossattn,
                merge_mlp=args.itm_merge_mlp,
                max_downsample_ratio=args.itm_max_downsample_ratio,
            )
            if model_args.guidance_scale <= 1.0:
                model_args.guidance_scale = 2.0
                print(f"[ITM] guidance_scale overridden to {model_args.guidance_scale} (ITM requires > 1.0)")

    # Find all episodes
    anno_dir = Path(val_dataset_dir) / "annotation" / "val"
    episode_ids = sorted([f.stem for f in anno_dir.glob("*.json")])
    if args.num_episodes:
        episode_ids = episode_ids[:args.num_episodes]

    print(f"\nSubset: {args.subset}")
    print(f"Dataset dir: {val_dataset_dir}")
    print(f"Output dir: {args.save_dir}")
    print(f"Input dir: {args.input_save_dir}")
    print(f"Episodes to process: {len(episode_ids)}")
    print(f"pred_step={model_args.num_frames}, num_history={model_args.num_history}")

    all_metrics = []
    for ep_idx, episode_id in enumerate(episode_ids):
        ep_save_path = Path(args.save_dir) / f"episode_{episode_id}" / "full_pred.mp4"
        if ep_save_path.exists():
            print(f"[{ep_idx+1}/{len(episode_ids)}] Episode {episode_id} already done, skipping")
            continue

        print(f"[{ep_idx+1}/{len(episode_ids)}] Episode {episode_id}")
        metrics = run_episode(agent, episode_id, args.save_dir, input_save_dir=args.input_save_dir)
        if metrics:
            all_metrics.append(metrics)
            print(f"  -> {metrics['num_frames_pred']} frames, "
                  f"PSNR={metrics['psnr']:.2f}, SSIM={metrics['ssim']:.4f}, LPIPS={metrics['lpips']:.4f}")

    # Save summary
    if all_metrics:
        summary = {
            "mean_psnr": sum(m["psnr"] for m in all_metrics) / len(all_metrics),
            "mean_ssim": sum(m["ssim"] for m in all_metrics) / len(all_metrics),
            "mean_lpips": sum(m["lpips"] for m in all_metrics) / len(all_metrics),
            "num_episodes": len(all_metrics),
            "subset": args.subset,
        }
        summary_path = Path(args.save_dir) / "all_summary.json"
        summary_path.parent.mkdir(parents=True, exist_ok=True)
        with open(summary_path, "w") as f:
            json.dump(summary, f, indent=2)
        print(f"\n=== Summary ({len(all_metrics)} episodes) ===")
        print(f"PSNR: {summary['mean_psnr']:.3f}")
        print(f"SSIM: {summary['mean_ssim']:.4f}")
        print(f"LPIPS: {summary['mean_lpips']:.4f}")


if __name__ == "__main__":
    main()