scripts/infer_single_arm_singleview_single_segment.py

"""
Single-segment full-length autoregressive rollout for single-arm singleview DreamDojo.

For each episode:
  - Start from the first GT frame.
  - Rollout autoregressively (chunk_size actions per step) for the ENTIRE episode
    length — no GT reset between segments. Pure autoregressive drift test.
  - GT video and actions come from the dataset pipeline (properly normalized/resized).

Output matches the benchmark format: full_gt.mp4, full_pred.mp4, full_merged.mp4, metrics.json.
"""

import argparse
import json
import sys
from pathlib import Path

import mediapy
import numpy as np
import piq
import torch
import torchvision

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "models" / "DreamDojo"))


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoints-dir", type=str, required=True)
    parser.add_argument("--experiment", type=str, default="dreamdojo_2b_480_640_single_arm_sv")
    parser.add_argument("--dataset-path", type=str, required=True)
    parser.add_argument("--save-dir", type=str, required=True)
    parser.add_argument("--chunk-size", type=int, default=12)
    parser.add_argument("--num-episodes", type=int, default=1)
    parser.add_argument("--save-fps", type=int, default=15)
    parser.add_argument("--output-dir", type=str, default=None)
    return parser.parse_args()


def build_model(args):
    from cosmos_predict2.action_conditioned_config import ActionConditionedSetupArguments
    from cosmos_predict2.config import MODEL_CHECKPOINTS
    from cosmos_predict2._src.predict2.inference.video2world import Video2WorldInference
    from cosmos_predict2.cache_runtime import build_cache_runtime_config

    setup_args = ActionConditionedSetupArguments(
        model="2B/robot/action-cond",
        config_file="cosmos_predict2/_src/predict2/action/configs/action_conditioned/config.py",
        checkpoints_dir=args.checkpoints_dir,
        experiment=args.experiment,
        num_frames=13,
        dataset_path=args.dataset_path,
        save_dir=args.save_dir,
        output_dir=args.output_dir or args.save_dir,
        num_samples=1,
        data_split="full",
        single_base_index=False,
    )

    checkpoints_dir = Path(args.checkpoints_dir)
    last_checkpoint_file = checkpoints_dir / "latest_checkpoint.txt"
    if not last_checkpoint_file.exists():
        parent_file = checkpoints_dir.parent / "latest_checkpoint.txt"
        if parent_file.exists():
            checkpoints_dir = checkpoints_dir.parent
            last_checkpoint_file = parent_file

    if not last_checkpoint_file.exists():
        raise FileNotFoundError(f"Could not find latest_checkpoint.txt in {args.checkpoints_dir} or its parent.")

    with open(last_checkpoint_file) as f:
        last_checkpoint = f.read().strip()
    checkpoint_iter_dir = checkpoints_dir / last_checkpoint

    from examples.action_conditioned import resolve_checkpoint_path
    checkpoint_path = resolve_checkpoint_path(checkpoint_iter_dir)

    checkpoint = MODEL_CHECKPOINTS[setup_args.model_key]
    experiment = setup_args.experiment or checkpoint.experiment
    cache_config = build_cache_runtime_config(setup_args)

    video2world_cli = Video2WorldInference(
        experiment_name=experiment,
        ckpt_path=checkpoint_path,
        s3_credential_path="",
        context_parallel_size=setup_args.context_parallel_size,
        config_file=setup_args.config_file,
        experiment_opts=[],
        cache_config=cache_config,
    )

    return video2world_cli, checkpoint_iter_dir.name


def build_dataset(args):
    """Build dataset with single_base_index=False to access all base_indices per episode."""
    from groot_dreams.dataloader import MultiVideoActionDataset

    dataset = MultiVideoActionDataset(
        num_frames=13,
        dataset_path=args.dataset_path,
        data_split="full",
        single_base_index=False,
        restrict_len=None,
        deterministic_uniform_sampling=False,
    )
    return dataset


def get_episode_plan(dataset, chunk_size):
    """
    For each episode, collect data_ids whose actions cover the full episode.
    Each data_id provides chunk_size actions. We step through the episode
    in increments of chunk_size timesteps.
    """
    episodes = {}

    global_offset = 0
    for ds_idx, ds in enumerate(dataset.datasets):
        lerobot_ds = ds.lerobot_dataset
        for local_idx, (traj_id, base_index) in enumerate(lerobot_ds.all_steps):
            key = (ds_idx, int(traj_id))
            if key not in episodes:
                episodes[key] = []
            episodes[key].append((global_offset + local_idx, int(base_index)))
        global_offset += len(ds)

    delta_indices = dataset.datasets[0].lerobot_dataset.modality_configs["video"].delta_indices
    timestep_interval = delta_indices[1] - delta_indices[0]
    stride_raw = chunk_size * timestep_interval

    plan = []
    for key, steps in episodes.items():
        ds_idx, traj_id = key
        steps_sorted = sorted(steps, key=lambda x: x[1])
        if not steps_sorted:
            continue

        segment_indices = []
        next_base = 0
        for global_id, base_idx in steps_sorted:
            if base_idx >= next_base:
                segment_indices.append(global_id)
                next_base = base_idx + stride_raw

        if segment_indices:
            traj_length = int(dataset.datasets[ds_idx].lerobot_dataset.trajectory_lengths[
                np.where(dataset.datasets[ds_idx].lerobot_dataset.trajectory_ids == traj_id)[0][0]
            ])
            plan.append({
                "ds_idx": ds_idx,
                "traj_id": int(traj_id),
                "traj_length": traj_length,
                "segment_data_ids": segment_indices,
            })

    return plan


def main():
    args = parse_args()

    from cosmos_oss.init import init_environment, cleanup_environment
    init_environment()
    torch.enable_grad(False)

    print("Building model...")
    video2world_cli, iter_name = build_model(args)

    print("Building dataset...")
    dataset = build_dataset(args)

    print("Planning episodes...")
    plan = get_episode_plan(dataset, args.chunk_size)
    total_episodes = len(plan)
    num_episodes = min(args.num_episodes, total_episodes)
    print(f"Total episodes: {total_episodes}, testing: {num_episodes}")

    save_root = Path(args.save_dir) / iter_name
    save_root.mkdir(parents=True, exist_ok=True)

    all_psnr, all_ssim, all_lpips = [], [], []

    for ep_idx in range(num_episodes):
        ep_info = plan[ep_idx]
        traj_id = ep_info["traj_id"]
        ep_save_dir = save_root / f"episode_{traj_id:06d}"

        if (ep_save_dir / "full_pred.mp4").exists():
            print(f"[{ep_idx}] episode_{traj_id:06d} already exists, skipping.")
            continue

        num_chunks = len(ep_info["segment_data_ids"])
        print(f"[{ep_idx}] traj_id={traj_id}, traj_length={ep_info['traj_length']}, "
              f"chunks={num_chunks}")

        if num_chunks == 0:
            print("  No chunks, skipping.")
            continue

        ep_save_dir.mkdir(parents=True, exist_ok=True)

        # Get first frame from first segment (GT, properly resized by dataset pipeline)
        first_sample = dataset[ep_info["segment_data_ids"][0]]
        img_array = first_sample["video"].transpose(0, 1)[:1]  # (1, C, H, W) uint8

        # Collect GT video frames and actions from all segments
        gt_frames = []
        chunk_videos = []
        first_round = True

        for chunk_idx, data_id in enumerate(ep_info["segment_data_ids"]):
            sample = dataset[data_id]

            # GT frames for this chunk
            video_tensor = sample["video"]  # (C, T, H, W)
            gt_video_chunk = video_tensor.permute(1, 2, 3, 0).numpy()  # (T, H, W, C)
            gt_frames.append(gt_video_chunk)

            # Actions for this chunk (already normalized by dataset pipeline)
            actions = sample["action"][:args.chunk_size]
            if isinstance(actions, torch.Tensor):
                actions = actions.numpy()

            if actions.shape[0] != args.chunk_size:
                print(f"    chunk {chunk_idx}: only {actions.shape[0]} actions (need {args.chunk_size}), stopping.")
                break

            lam_video = sample.get("lam_video", None)
            current_lam_video = None
            if lam_video is not None and len(lam_video) >= args.chunk_size * 2:
                current_lam_video = lam_video[:args.chunk_size * 2]

            # Prepare input frame
            if not first_round:
                img_tensor = torchvision.transforms.functional.to_tensor(img_array).unsqueeze(0) * 255.0
            else:
                img_tensor = img_array
            first_round = False

            num_video_frames = actions.shape[0] + 1
            vid_input = torch.cat(
                [img_tensor, torch.zeros_like(img_tensor).repeat(num_video_frames - 1, 1, 1, 1)], dim=0
            )
            vid_input = vid_input.to(torch.uint8)
            vid_input = vid_input.unsqueeze(0).permute(0, 2, 1, 3, 4)

            video = video2world_cli.generate_vid2world(
                prompt="",
                input_path=vid_input,
                action=torch.from_numpy(actions).float()
                if isinstance(actions, np.ndarray)
                else actions,
                guidance=0,
                num_video_frames=num_video_frames,
                num_latent_conditional_frames=1,
                resolution="480,640",
                seed=chunk_idx,
                negative_prompt="The video captures a scene with low visual quality, blurring, jittering, or distortion.",
                lam_video=current_lam_video,
            )

            video_normalized = (video - (-1)) / (1 - (-1))
            video_clamped = (
                (torch.clamp(video_normalized[0], 0, 1) * 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy()
            )

            # Use last predicted frame as next input (NO GT reset — pure autoregressive)
            img_array = video_clamped[-1]
            chunk_videos.append(video_clamped)

            print(f"    chunk {chunk_idx+1}/{num_chunks} done")

        if not chunk_videos:
            continue

        # Concatenate: first chunk full, subsequent chunks drop conditioning frame
        chunk_list = [chunk_videos[0]] + [
            chunk_videos[i][:args.chunk_size] for i in range(1, len(chunk_videos))
        ]
        concat_pred = np.concatenate(chunk_list, axis=0)

        # GT: similarly concatenate
        gt_list = [gt_frames[0]] + [
            gt_frames[i][:args.chunk_size] for i in range(1, len(gt_frames))
        ]
        concat_gt = np.concatenate(gt_list, axis=0)

        min_len = min(len(concat_pred), len(concat_gt))
        concat_pred = concat_pred[:min_len]
        concat_gt = concat_gt[:min_len]

        mediapy.write_video(str(ep_save_dir / "full_pred.mp4"), concat_pred, fps=args.save_fps)
        mediapy.write_video(str(ep_save_dir / "full_gt.mp4"), concat_gt, fps=args.save_fps)
        concat_merged = np.concatenate([concat_gt, concat_pred], axis=2)
        mediapy.write_video(str(ep_save_dir / "full_merged.mp4"), concat_merged, fps=args.save_fps)

        x_batch = torch.clamp(torch.from_numpy(concat_pred.copy()) / 255.0, 0, 1).permute(0, 3, 1, 2)
        y_batch = torch.clamp(torch.from_numpy(concat_gt.copy()) / 255.0, 0, 1).permute(0, 3, 1, 2)
        psnr_val = piq.psnr(x_batch, y_batch).mean().item()
        ssim_val = piq.ssim(x_batch, y_batch).mean().item()
        lpips_val = piq.LPIPS()(x_batch, y_batch).mean().item()

        with open(ep_save_dir / "metrics.json", "w") as f:
            json.dump({
                "psnr": psnr_val, "ssim": ssim_val, "lpips": lpips_val,
                "num_chunks": len(chunk_videos),
                "total_frames_pred": len(concat_pred),
                "total_frames_gt": ep_info["traj_length"],
                "trajectory_id": traj_id,
                "mode": "single_segment_full_rollout",
            }, f, indent=2)

        all_psnr.append(psnr_val)
        all_ssim.append(ssim_val)
        all_lpips.append(lpips_val)
        print(f"  -> {len(concat_pred)} frames, PSNR={psnr_val:.2f}, SSIM={ssim_val:.4f}, LPIPS={lpips_val:.4f}")

    if all_psnr:
        summary = {
            "mean_psnr": sum(all_psnr) / len(all_psnr),
            "mean_ssim": sum(all_ssim) / len(all_ssim),
            "mean_lpips": sum(all_lpips) / len(all_lpips),
            "num_episodes_processed": len(all_psnr),
            "mode": "single_segment_full_rollout",
        }
        with open(save_root / "all_summary.json", "w") as f:
            json.dump(summary, f, indent=2)
        print(f"\n=== Summary ({len(all_psnr)} episodes) ===")
        print(f"PSNR: {summary['mean_psnr']:.3f}")
        print(f"SSIM: {summary['mean_ssim']:.4f}")
        print(f"LPIPS: {summary['mean_lpips']:.4f}")

    cleanup_environment()


if __name__ == "__main__":
    main()