robometer/scripts/benchmark_progress_mark_local.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import gc
import json
import os
import re
from collections import defaultdict
from pathlib import Path
from typing import Optional

import av
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import torch

from robometer.data.dataset_types import ProgressSample, Trajectory
from robometer.evals.eval_server import compute_batch_outputs
from robometer.utils.save import load_model_from_hf
from robometer.utils.setup_utils import setup_batch_collator

matplotlib.use("Agg")
av.logging.set_level(av.logging.ERROR)


REPO_ROOT = Path(__file__).resolve().parents[1]
PROJECT_ROOT = Path(__file__).resolve().parents[2]
WORKSPACE_ROOT = Path(__file__).resolve().parents[3]
DEFAULT_MODEL_PATH = PROJECT_ROOT / "models" / "Robometer-4B"
DEFAULT_VIDEOS_ROOT = WORKSPACE_ROOT / "Videos"
DEFAULT_OUTPUT_ROOT = REPO_ROOT / "Benchmark_Eval"
DEFAULT_CHUNK = "chunk-000"
DEFAULT_CAMERA_DIR = "observation.images.wrist_image_left"
DEFAULT_MAX_FRAMES = 128
DEFAULT_MIN_FRAMES = 32
FRAME_BACKOFF_CANDIDATES = [512, 384, 320, 256, 224, 192, 160, 144, 128, 112, 96, 80, 64, 48, 32, 24, 16, 8]
DEFAULT_INFERENCE_MODE = "frame_steps"
DEFAULT_PREFIX_SAMPLE_FRAMES = 4
DEFAULT_PREFIX_BATCH_SIZE = 1
DEFAULT_DESCRIPTION_EPISODES_PATH = WORKSPACE_ROOT / "Robo-Dopamine" / "Evaluation" / "Benchmark_Eval" / "Description" / "episodes.jsonl"


def normalize_episode_name(name: str) -> str:
    name = name.strip()
    if name.endswith(".mp4"):
        return name
    if name.startswith("episode_"):
        return f"{name}.mp4"
    if name.isdigit():
        return f"episode_{int(name):06d}.mp4"
    return f"episode_{name}.mp4"


def chunk_to_dir_name(chunk: str) -> str:
    m = re.fullmatch(r"chunk-(\d+)(?:_(.+))?", chunk)
    if not m:
        return f"Chunk{chunk.split('-')[-1]}"
    suffix = m.group(2)
    return f"Chunk{m.group(1)}_{suffix}" if suffix else f"Chunk{m.group(1)}"


def parse_episode_input(token: str) -> str:
    return normalize_episode_name(token)


def load_json(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def load_episode_catalog(episode_tasks_path: Path, annotations_path: Path) -> list[str]:
    if episode_tasks_path.exists():
        data = load_json(episode_tasks_path)
        catalog = []
        for ep in data.get("episodes", []):
            episode = ep.get("episode")
            if isinstance(episode, str) and episode.strip():
                catalog.append(episode.strip())
        if catalog:
            return catalog
    data = load_json(annotations_path)
    return [ep["episode"] for ep in data.get("episodes", []) if ep.get("episode")]


def load_episode_record(annotations_path: Path, episode_mp4: str) -> dict:
    data = load_json(annotations_path)
    for ep in data.get("episodes", []):
        if ep.get("episode") == episode_mp4:
            return ep
    raise ValueError(f"Episode not found in annotations: {episode_mp4}")


def load_instruction(episode_tasks_path: Path, episode_mp4: str) -> str:
    if episode_tasks_path.exists():
        data = load_json(episode_tasks_path)
        for ep in data.get("episodes", []):
            if ep.get("episode") != episode_mp4:
                continue
            for task in ep.get("tasks") or []:
                if isinstance(task, str) and task.strip():
                    return task.strip()

    m = re.fullmatch(r"episode_(\d+)\.mp4", episode_mp4)
    if m and DEFAULT_DESCRIPTION_EPISODES_PATH.exists():
        target_idx = int(m.group(1))
        with DEFAULT_DESCRIPTION_EPISODES_PATH.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                if obj.get("episode_index") != target_idx:
                    continue
                for task in obj.get("tasks") or []:
                    if isinstance(task, str) and task.strip():
                        return task.strip()
                break
    return "robot manipulation task"


def extract_progress_pairs(episode_record: dict) -> list[tuple[str, int, int]]:
    marks = [a for a in episode_record.get("annotations", []) if a.get("type") == "progress_mark"]
    camera_ref = episode_record.get("camera_reference")
    if camera_ref:
        ref_marks = [m for m in marks if m.get("camera") == camera_ref]
        if ref_marks:
            marks = ref_marks

    grouped: dict[str, list[int]] = defaultdict(list)
    for mark in marks:
        frame = mark.get("frame")
        if isinstance(frame, int):
            grouped[str(mark.get("group_label", "UNK"))].append(frame)

    pairs: list[tuple[str, int, int]] = []
    for label in sorted(grouped):
        frames = sorted(set(grouped[label]))
        if len(frames) >= 2:
            pairs.append((label, frames[0], frames[1]))
    return pairs


def load_progress_mark_catalog(annotations_path: Path) -> list[str]:
    data = load_json(annotations_path)
    catalog = []
    for ep in data.get("episodes", []):
        episode = ep.get("episode")
        if isinstance(episode, str) and episode.strip() and extract_progress_pairs(ep):
            catalog.append(episode.strip())
    return catalog


def parse_episode_selection(token: str, catalog: list[str], progress_catalog: list[str], require_progress_marks: bool = True) -> list[str]:
    aliases = {
        "fisrt10": "first10",
        "frist10": "first10",
        "first 10": "first10",
    }
    token = token.strip()
    token_lower = aliases.get(token.lower(), token.lower())

    if token_lower == "first10":
        return (progress_catalog if require_progress_marks else catalog)[:10]
    if token_lower == "all":
        return progress_catalog if require_progress_marks else catalog

    plus_match = re.fullmatch(r"(.+)\+(\d+)", token)
    if plus_match:
        start_episode = parse_episode_input(plus_match.group(1).strip())
        count = int(plus_match.group(2))
        if count <= 0:
            raise ValueError("Count after '+' must be >= 1.")
        if start_episode not in catalog:
            raise ValueError(f"Invalid episode selection: {plus_match.group(1).strip()}.")
        if require_progress_marks and start_episode not in progress_catalog:
            raise ValueError(f"{start_episode[:-4]} does not have usable progress_mark pairs.")
        base_catalog = progress_catalog if require_progress_marks else catalog
        start_idx = base_catalog.index(start_episode)
        return base_catalog[start_idx:start_idx + count]

    selected = []
    for part in token.split(","):
        part = part.strip()
        if not part:
            continue
        episode_mp4 = parse_episode_input(part)
        if episode_mp4 not in catalog:
            raise ValueError(f"Invalid episode selection: {part}.")
        if require_progress_marks and episode_mp4 not in progress_catalog:
            raise ValueError(f"{episode_mp4[:-4]} does not have usable progress_mark pairs.")
        selected.append(episode_mp4)

    return selected


def select_progress_window(
    start_episode_raw: str,
    count: int,
    catalog: list[str],
    progress_catalog: list[str],
    require_progress_marks: bool = True,
) -> list[str]:
    start_episode = parse_episode_input(start_episode_raw)
    if count <= 0:
        raise ValueError("Count must be >= 1.")
    if start_episode not in catalog:
        raise ValueError(f"Invalid episode selection: {start_episode_raw}.")
    if require_progress_marks and start_episode not in progress_catalog:
        raise ValueError(f"{start_episode[:-4]} does not have usable progress_mark pairs.")
    base_catalog = progress_catalog if require_progress_marks else catalog
    start_idx = base_catalog.index(start_episode)
    return base_catalog[start_idx:start_idx + count]


def result_path_for(output_root: Path, episode_mp4: str) -> Path:
    return output_root / episode_mp4[:-4] / "results.json"


def chunk_output_dir_name(chunk: str, tag: str | None) -> str:
    base = chunk_to_dir_name(chunk)
    return f"{base}_{tag}" if tag else base


def parse_chunk_selection(chunk_arg: str) -> list[str]:
    chunks = [part.strip() for part in chunk_arg.split(",") if part.strip()]
    if not chunks:
        raise ValueError("At least one chunk must be provided.")
    return chunks


def cleanup_cuda_memory() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def is_cuda_oom_error(exc: BaseException) -> bool:
    if isinstance(exc, torch.OutOfMemoryError):
        return True
    return "out of memory" in str(exc).lower()


def build_frame_retry_schedule(max_frames: int, min_frames: int, adaptive: bool) -> list[int]:
    max_frames = max(1, int(max_frames))
    min_frames = max(1, min(int(min_frames), max_frames))
    if not adaptive or max_frames == min_frames:
        return [max_frames]

    schedule = [max_frames]
    for candidate in FRAME_BACKOFF_CANDIDATES:
        if min_frames <= candidate < max_frames:
            schedule.append(candidate)
    if schedule[-1] != min_frames:
        schedule.append(min_frames)
    return schedule


def resolve_episode_list(
    episode_arg: Optional[str],
    count: int,
    catalog: list[str],
    progress_catalog: list[str],
    require_progress_marks: bool,
) -> list[str]:
    if episode_arg:
        token = episode_arg.strip()
        token_lower = token.lower()
        if count > 1:
            return select_progress_window(token, count, catalog, progress_catalog, require_progress_marks=require_progress_marks)
        if token_lower in {"first10", "all"} or "," in token or "+" in token:
            return parse_episode_selection(token, catalog, progress_catalog, require_progress_marks=require_progress_marks)
        episode_mp4 = parse_episode_input(token)
        if episode_mp4 not in catalog:
            raise ValueError(f"Invalid episode selection: {token}.")
        if require_progress_marks and episode_mp4 not in progress_catalog:
            raise ValueError(f"{episode_mp4[:-4]} does not have usable progress_mark pairs.")
        return [episode_mp4]

    base_catalog = progress_catalog if require_progress_marks else catalog
    print(f"[EPISODES] {len(base_catalog)} episodes available")
    user_selection = input("Episodes (000287 / 000287+5 / 000287,000330 / first10 / all): ").strip()
    aliases = {"fisrt10": "first10", "frist10": "first10", "first 10": "first10"}
    token_lower = aliases.get(user_selection.lower(), user_selection.lower())
    if token_lower not in {"first10", "all"} and "," not in user_selection and "+" not in user_selection:
        count_raw = input("How many from this start? (default 5): ").strip()
        window_count = int(count_raw) if count_raw else 5
        return select_progress_window(user_selection, window_count, catalog, progress_catalog, require_progress_marks=require_progress_marks)
    return parse_episode_selection(user_selection, catalog, progress_catalog, require_progress_marks=require_progress_marks)


def load_video_frames_with_indices(
    video_path: Path,
    *,
    fps: float,
    max_frames: int,
    required_frames: list[int],
) -> tuple[np.ndarray, list[int], int, float]:
    all_frames, native_fps = load_all_video_frames(video_path)
    sampled_frames, sampled_indices = sample_video_frames_with_indices(
        all_frames,
        native_fps=native_fps,
        fps=fps,
        max_frames=max_frames,
        required_frames=required_frames,
    )
    return sampled_frames, sampled_indices, len(all_frames), native_fps


def load_all_video_frames(video_path: Path) -> tuple[list[np.ndarray], float]:
    all_frames: list[np.ndarray] = []
    with av.open(str(video_path)) as container:
        stream = container.streams.video[0]
        rate = stream.average_rate or stream.guessed_rate
        native_fps = float(rate) if rate is not None else 1.0
        for frame in container.decode(stream):
            all_frames.append(frame.to_ndarray(format="rgb24"))

    if not all_frames:
        raise RuntimeError(f"Could not extract frames from video: {video_path}")
    return all_frames, native_fps


def sample_video_frames_with_indices(
    all_frames: list[np.ndarray],
    *,
    native_fps: float,
    fps: float,
    max_frames: int,
    required_frames: list[int],
) -> tuple[np.ndarray, list[int]]:
    total_frames = len(all_frames)
    if fps <= 0:
        fps = native_fps

    if native_fps > 0:
        desired_frames = int(round(total_frames * (fps / native_fps)))
    else:
        desired_frames = total_frames
    desired_frames = max(1, min(desired_frames, total_frames, max_frames))

    if desired_frames == total_frames:
        sampled_indices = list(range(total_frames))
    else:
        base_indices = np.linspace(0, total_frames - 1, desired_frames, dtype=int).tolist()
        snapped_indices = list(base_indices)
        used_slots: set[int] = set()

        # Snap the nearest sampling slots onto the required progress_mark frames.
        for frame_idx in sorted(set(required_frames)):
            clamped = max(0, min(total_frames - 1, int(frame_idx)))
            slot = min(
                (idx for idx in range(desired_frames) if idx not in used_slots),
                key=lambda idx: (abs(base_indices[idx] - clamped), idx),
            )
            snapped_indices[slot] = clamped
            used_slots.add(slot)

        sampled_indices = sorted(set(snapped_indices))
        if len(sampled_indices) < desired_frames:
            for idx in base_indices:
                if idx not in sampled_indices:
                    sampled_indices.append(idx)
                    if len(sampled_indices) == desired_frames:
                        break
        if len(sampled_indices) < desired_frames:
            for idx in range(total_frames):
                if idx not in sampled_indices:
                    sampled_indices.append(idx)
                    if len(sampled_indices) == desired_frames:
                        break
        sampled_indices = sorted(sampled_indices[:desired_frames])

    sampled_frames = np.stack([all_frames[idx] for idx in sampled_indices], axis=0)
    return sampled_frames, sampled_indices


def extract_frame(video_path: Path, frame_idx: int) -> Optional[np.ndarray]:
    with av.open(str(video_path)) as container:
        stream = container.streams.video[0]
        for idx, frame in enumerate(container.decode(stream)):
            if idx == frame_idx:
                return frame.to_ndarray(format="rgb24")
    return None


class RobometerLocalRunner:
    def __init__(self, model_path: str, device: Optional[torch.device] = None):
        if device is None:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.device = device
        self.model_path = model_path
        self.exp_config = None
        self.tokenizer = None
        self.processor = None
        self.reward_model = None
        self.batch_collator = None
        self.is_discrete = False
        self.num_bins = 10
        self._load_model_components()

    def _load_model_components(self) -> None:
        print(f"Loading model from {self.model_path} ...")

        self.exp_config, self.tokenizer, self.processor, self.reward_model = load_model_from_hf(
            model_path=self.model_path,
            device=self.device,
        )
        self.reward_model.eval()
        self.batch_collator = setup_batch_collator(
            self.processor,
            self.tokenizer,
            self.exp_config,
            is_eval=True,
        )

        loss_config = getattr(self.exp_config, "loss", None)
        self.is_discrete = (
            getattr(loss_config, "progress_loss_type", "l2").lower() == "discrete"
            if loss_config
            else False
        )
        self.num_bins = (
            getattr(loss_config, "progress_discrete_bins", None)
            or getattr(self.exp_config.model, "progress_discrete_bins", 10)
        )

    def reload_model(self) -> None:
        if self.reward_model is not None:
            del self.reward_model
        cleanup_cuda_memory()
        self._load_model_components()

    @torch.inference_mode()
    def _run_progress_samples(self, progress_samples: list[ProgressSample]) -> tuple[list[list[float]], list[list[float]]]:
        batch = self.batch_collator(progress_samples)

        progress_inputs = batch["progress_inputs"]
        for key, value in progress_inputs.items():
            if hasattr(value, "to"):
                progress_inputs[key] = value.to(self.device)

        results = compute_batch_outputs(
            self.reward_model,
            self.tokenizer,
            progress_inputs,
            sample_type="progress",
            is_discrete_mode=self.is_discrete,
            num_bins=self.num_bins,
        )

        progress_pred = results.get("progress_pred", []) or []
        outputs_success = results.get("outputs_success", {})
        success_probs = outputs_success.get("success_probs", []) if outputs_success else []
        return progress_pred, success_probs

    @torch.inference_mode()
    def _compute_rewards_whole_trajectory(self, video_frames: np.ndarray, task: str) -> tuple[np.ndarray, np.ndarray]:
        traj = Trajectory(
            frames=video_frames,
            frames_shape=tuple(video_frames.shape),
            task=task,
            id="0",
            metadata={"subsequence_length": int(video_frames.shape[0])},
            video_embeddings=None,
        )
        progress_sample = ProgressSample(trajectory=traj, sample_type="progress")
        progress_pred, success_probs = self._run_progress_samples([progress_sample])

        progress_array = (
            np.array(progress_pred[0], dtype=np.float32)
            if progress_pred and len(progress_pred) > 0
            else np.array([], dtype=np.float32)
        )

        success_array = (
            np.array(success_probs[0], dtype=np.float32)
            if success_probs and len(success_probs) > 0
            else np.array([], dtype=np.float32)
        )

        return progress_array, success_array

    @torch.inference_mode()
    def _compute_rewards_frame_steps(
        self,
        video_frames: np.ndarray,
        task: str,
        *,
        prefix_sample_frames: int,
        prefix_batch_size: int,
    ) -> tuple[np.ndarray, np.ndarray]:
        total_frames = int(video_frames.shape[0])
        if total_frames <= 0:
            return np.array([], dtype=np.float32), np.array([], dtype=np.float32)

        prefix_sample_frames = max(1, int(prefix_sample_frames))
        batch_size = max(1, int(prefix_batch_size))

        prefix_samples: list[ProgressSample] = []
        metadata = {"subsequence_length": total_frames}
        for prefix_len in range(1, total_frames + 1):
            indices = np.linspace(0, prefix_len - 1, prefix_sample_frames, dtype=int)
            sub_frames = video_frames[indices]
            traj = Trajectory(
                frames=sub_frames,
                frames_shape=tuple(sub_frames.shape),
                task=task,
                id="0",
                metadata=metadata,
                video_embeddings=None,
            )
            prefix_samples.append(ProgressSample(trajectory=traj, sample_type="progress"))

        progress_values: list[float] = []
        success_values: list[float] = []
        cursor = 0

        while cursor < len(prefix_samples):
            current_samples = prefix_samples[cursor:cursor + batch_size]
            try:
                progress_pred, success_probs = self._run_progress_samples(current_samples)
            except RuntimeError as exc:
                if not is_cuda_oom_error(exc):
                    raise
                if batch_size == 1:
                    raise
                next_batch_size = max(1, batch_size // 2)
                print(
                    f"[OOM] frame_steps batch_size={batch_size} failed; retrying with batch_size={next_batch_size}"
                )
                cleanup_cuda_memory()
                batch_size = next_batch_size
                continue

            for sub_pred in progress_pred:
                if not sub_pred:
                    raise RuntimeError("Robometer returned empty progress predictions for a frame-step prefix")
                progress_values.append(float(sub_pred[-1]))

            for idx in range(len(current_samples)):
                sub_success = success_probs[idx] if idx < len(success_probs) else []
                success_values.append(float(sub_success[-1]) if sub_success else np.nan)

            cursor += len(current_samples)

        success_array = np.array(success_values, dtype=np.float32)
        if success_array.size > 0 and np.isnan(success_array).all():
            success_array = np.array([], dtype=np.float32)

        return np.array(progress_values, dtype=np.float32), success_array

    @torch.inference_mode()
    def compute_rewards_per_frame(
        self,
        video_frames: np.ndarray,
        task: str,
        *,
        inference_mode: str = DEFAULT_INFERENCE_MODE,
        prefix_sample_frames: int = DEFAULT_PREFIX_SAMPLE_FRAMES,
        prefix_batch_size: int = DEFAULT_PREFIX_BATCH_SIZE,
    ) -> tuple[np.ndarray, np.ndarray]:
        if inference_mode == "frame_steps":
            return self._compute_rewards_frame_steps(
                video_frames=video_frames,
                task=task,
                prefix_sample_frames=prefix_sample_frames,
                prefix_batch_size=prefix_batch_size,
            )
        if inference_mode == "whole":
            return self._compute_rewards_whole_trajectory(video_frames=video_frames, task=task)
        raise ValueError(f"Unsupported inference_mode: {inference_mode}")


def compute_rewards_per_frame_local(
    model_path: str,
    video_frames: np.ndarray,
    task: str,
    device: Optional[torch.device] = None,
) -> tuple[np.ndarray, np.ndarray]:
    runner = RobometerLocalRunner(model_path=model_path, device=device)
    return runner.compute_rewards_per_frame(video_frames=video_frames, task=task)


def map_frame_to_sample(frame_idx: int, sampled_indices: list[int]) -> tuple[int, int]:
    nearest_pos = min(range(len(sampled_indices)), key=lambda idx: abs(sampled_indices[idx] - frame_idx))
    return nearest_pos, sampled_indices[nearest_pos]


def build_mark_points(
    frame_pairs: list[tuple[str, int, int]],
    sampled_indices: list[int],
    progress_pred: np.ndarray,
) -> list[dict]:
    points = []
    for label, f0, f1 in frame_pairs:
        for suffix, frame_idx in (("1", f0), ("2", f1)):
            sample_pos, sampled_frame = map_frame_to_sample(frame_idx, sampled_indices)
            progress = float(progress_pred[sample_pos])
            points.append(
                {
                    "group": label,
                    "tag": f"{label}{suffix}",
                    "original_frame": frame_idx,
                    "sample_index": sample_pos,
                    "sampled_frame": sampled_frame,
                    "frame_error": sampled_frame - frame_idx,
                    "progress": progress,
                }
            )
    return points


def write_overview_plot(
    output_dir: Path,
    sampled_indices: list[int],
    progress_pred: np.ndarray,
    mark_points: list[dict],
) -> None:
    fig, ax = plt.subplots(figsize=(10.5, 4.8))
    ax.plot(sampled_indices, progress_pred, linewidth=2.0, color="#444444")

    palette = plt.get_cmap("tab10")
    grouped: dict[str, list[dict]] = defaultdict(list)
    for point in mark_points:
        grouped[point["group"]].append(point)

    for idx, label in enumerate(sorted(grouped)):
        pts = sorted(grouped[label], key=lambda item: item["original_frame"])
        color = palette(idx % 10)
        ax.plot(
            [pt["sampled_frame"] for pt in pts],
            [pt["progress"] for pt in pts],
            marker="o",
            linewidth=1.8,
            markersize=7,
            color=color,
        )
        for pt in pts:
            ax.text(
                pt["sampled_frame"],
                pt["progress"] + 0.03,
                f"{pt['tag']} {pt['progress']:.3f}",
                color=color,
                fontsize=9,
                ha="center",
            )

    ax.set_xlabel("Original Frame")
    ax.set_ylabel("Progress")
    ax.set_ylim(-0.05, 1.05)
    ax.set_title("Robometer Progress Curve")
    ax.grid(True, linestyle="--", alpha=0.35)
    fig.tight_layout()
    fig.savefig(output_dir / "overview.png", dpi=180)
    plt.close(fig)


def write_group_cards(
    output_dir: Path,
    video_path: Path,
    sampled_indices: list[int],
    progress_pred: np.ndarray,
    mark_points: list[dict],
) -> None:
    groups_dir = output_dir / "groups"
    groups_dir.mkdir(parents=True, exist_ok=True)
    palette = plt.get_cmap("tab10")

    grouped: dict[str, list[dict]] = defaultdict(list)
    for point in mark_points:
        grouped[point["group"]].append(point)

    for idx, label in enumerate(sorted(grouped)):
        pts = sorted(grouped[label], key=lambda item: item["original_frame"])
        color = palette(idx % 10)
        first = pts[0]
        second = pts[-1]
        img0 = extract_frame(video_path, first["original_frame"])
        img1 = extract_frame(video_path, second["original_frame"])

        frame_gap = abs(second["original_frame"] - first["original_frame"])
        x_min = max(0, min(first["original_frame"], second["original_frame"]) - max(25, frame_gap))
        x_max = max(first["original_frame"], second["original_frame"]) + max(25, frame_gap)

        local_x = []
        local_y = []
        for x, y in zip(sampled_indices, progress_pred, strict=False):
            if x_min <= x <= x_max:
                local_x.append(x)
                local_y.append(float(y))
        if not local_x:
            local_x = sampled_indices
            local_y = progress_pred.tolist()

        fig = plt.figure(figsize=(9.5, 6.5))
        gs = fig.add_gridspec(2, 2, height_ratios=[1.45, 1.0])
        ax_img0 = fig.add_subplot(gs[0, 0])
        ax_img1 = fig.add_subplot(gs[0, 1])
        ax_curve = fig.add_subplot(gs[1, :])

        for ax, img, title in (
            (ax_img0, img0, f"{first['tag']} frame {first['original_frame']}"),
            (ax_img1, img1, f"{second['tag']} frame {second['original_frame']}"),
        ):
            ax.axis("off")
            ax.set_title(title, fontsize=10)
            if img is None:
                ax.text(0.5, 0.5, "frame read failed", ha="center", va="center", fontsize=9)
            else:
                ax.imshow(img)

        ax_curve.plot(local_x, local_y, marker="o", linewidth=1.6, markersize=4, color="#666666")
        ax_curve.plot(
            [pt["sampled_frame"] for pt in pts],
            [pt["progress"] for pt in pts],
            marker="o",
            linewidth=2.0,
            markersize=7,
            color=color,
        )
        for pt in pts:
            ax_curve.text(
                pt["sampled_frame"],
                pt["progress"] + 0.03,
                f"{pt['tag']} {pt['progress']:.3f}",
                color=color,
                fontsize=9,
                ha="center",
            )

        ax_curve.set_xlim(x_min, x_max)
        ax_curve.set_ylim(-0.05, 1.05)
        ax_curve.set_xlabel("Original Frame")
        ax_curve.set_ylabel("Progress")
        ax_curve.set_title(f"{label} Group")
        ax_curve.grid(True, linestyle="--", alpha=0.35)
        ax_curve.text(
            0.015,
            0.97,
            "\n".join(
                f"{pt['tag']}: frame={pt['original_frame']}, sampled={pt['sampled_frame']}, progress={pt['progress']:.3f}"
                for pt in pts
            ),
            transform=ax_curve.transAxes,
            va="top",
            ha="left",
            fontsize=9,
            bbox={"facecolor": "white", "edgecolor": "#cccccc", "alpha": 0.92, "boxstyle": "round,pad=0.35"},
        )
        fig.tight_layout()
        fig.savefig(groups_dir / f"{label}.png", dpi=180)
        plt.close(fig)


def save_results(output_dir: Path, payload: dict) -> None:
    with (output_dir / "results.json").open("w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)
        f.write("\n")


def run_episode(
    *,
    episode_mp4: str,
    chunk: str,
    fps: float,
    max_frames: int,
    min_frames: int,
    adaptive_max_frames: bool,
    inference_mode: str,
    prefix_sample_frames: int,
    prefix_batch_size: int,
    runner: RobometerLocalRunner,
    overwrite: bool,
    videos_root_base: Path,
    allow_no_marks: bool,
    tag: str | None = None,
) -> None:
    direct_chunk_dir = videos_root_base / chunk
    if direct_chunk_dir.exists():
        videos_root = direct_chunk_dir
    else:
        chunk_filtered = f"{chunk}_filtered"
        videos_root = videos_root_base / chunk_filtered
    annotations_path = videos_root / "annotations.json"
    episode_tasks_path = videos_root / "episode_tasks.json"
    video_path = videos_root / DEFAULT_CAMERA_DIR / episode_mp4
    output_dir = DEFAULT_OUTPUT_ROOT / chunk_output_dir_name(chunk, tag) / episode_mp4[:-4]

    if not overwrite and (output_dir / "results.json").exists():
        print(f"[SKIP] {episode_mp4[:-4]} already done")
        return

    episode_record = load_episode_record(annotations_path, episode_mp4)
    instruction = load_instruction(episode_tasks_path, episode_mp4)
    frame_pairs = extract_progress_pairs(episode_record)
    if not frame_pairs and not allow_no_marks:
        raise ValueError(f"{episode_mp4} has no usable progress_mark pairs")
    if not video_path.exists():
        raise FileNotFoundError(f"Video not found: {video_path}")

    required_frames = [frame for _, f0, f1 in frame_pairs for frame in (f0, f1)]

    print(f"[RUN] {episode_mp4[:-4]}")
    print(f"Loading frames from {video_path} ...")
    all_frames, native_fps = load_all_video_frames(video_path)
    total_frames = len(all_frames)
    retry_schedule = (
        build_frame_retry_schedule(max_frames, min_frames, adaptive_max_frames)
        if inference_mode == "whole"
        else [max_frames]
    )
    print(
        f"[MODE] {inference_mode}"
        + (
            f" (prefix_sample_frames={prefix_sample_frames}, prefix_batch_size={prefix_batch_size})"
            if inference_mode == "frame_steps"
            else ""
        )
    )

    progress_pred: np.ndarray | None = None
    success_probs: np.ndarray | None = None
    sampled_indices: list[int] | None = None
    used_max_frames = retry_schedule[0]

    for attempt_idx, frame_budget in enumerate(retry_schedule, start=1):
        frames, sampled_indices = sample_video_frames_with_indices(
            all_frames,
            native_fps=native_fps,
            fps=fps,
            max_frames=frame_budget,
            required_frames=required_frames,
        )
        print(
            f"Loaded {total_frames} total frames; sampled {len(frames)} frames at fps={fps:g} "
            f"(max_frames={frame_budget}, try {attempt_idx}/{len(retry_schedule)})"
        )
        try:
            progress_pred, success_probs = runner.compute_rewards_per_frame(
                video_frames=frames,
                task=instruction,
                inference_mode=inference_mode,
                prefix_sample_frames=prefix_sample_frames,
                prefix_batch_size=prefix_batch_size,
            )
            used_max_frames = frame_budget
            break
        except RuntimeError as exc:
            if inference_mode != "whole" or not is_cuda_oom_error(exc):
                raise
            del frames
            cleanup_cuda_memory()
            if attempt_idx == len(retry_schedule):
                raise
            next_budget = retry_schedule[attempt_idx]
            print(
                f"[OOM] {episode_mp4[:-4]} hit CUDA OOM at max_frames={frame_budget}; "
                f"retrying with max_frames={next_budget}"
            )
            runner.reload_model()

    if progress_pred is None or success_probs is None or sampled_indices is None:
        raise RuntimeError(f"Failed to compute rewards for {episode_mp4[:-4]}")
    if progress_pred.size == 0:
        raise RuntimeError("Robometer returned empty progress predictions")
    if progress_pred.size != len(sampled_indices):
        raise RuntimeError(
            f"Progress length mismatch: got {progress_pred.size} predictions for {len(sampled_indices)} sampled frames"
        )

    mark_points = build_mark_points(frame_pairs, sampled_indices, progress_pred)
    output_dir.mkdir(parents=True, exist_ok=True)

    np.save(str(output_dir / "progress.npy"), progress_pred)
    np.save(str(output_dir / "success_probs.npy"), success_probs)

    payload = {
        "episode": episode_mp4,
        "chunk": chunk,
        "instruction": instruction,
        "video_path": str(video_path),
        "num_total_frames": total_frames,
        "native_fps": native_fps,
        "sample_fps": float(fps),
        "num_sampled_frames": len(sampled_indices),
        "max_frames_used": used_max_frames,
        "inference_mode": inference_mode,
        "prefix_sample_frames": int(prefix_sample_frames),
        "sampled_original_frame_indices": sampled_indices,
        "progress_marks": [{"group": label, "frame_pair": [f0, f1]} for label, f0, f1 in frame_pairs],
        "mark_points": mark_points,
        "progress_min": float(np.min(progress_pred)),
        "progress_max": float(np.max(progress_pred)),
        "progress_mean": float(np.mean(progress_pred)),
    }
    save_results(output_dir, payload)
    write_overview_plot(output_dir, sampled_indices, progress_pred, mark_points)
    write_group_cards(output_dir, video_path, sampled_indices, progress_pred, mark_points)

    print(f"[OK] Saved results to {output_dir}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Run Robometer local inference aligned to progress_mark frames.")
    parser.add_argument("--episode", default=None, help="000287 / episode_000287 / episode_000287.mp4")
    parser.add_argument(
        "--count",
        type=int,
        default=1,
        help="when --episode is set, run this many progress_mark episodes from that start",
    )
    parser.add_argument("--chunk", default=DEFAULT_CHUNK)
    parser.add_argument("--videos-root", type=Path, default=DEFAULT_VIDEOS_ROOT)
    parser.add_argument("--model-path", default=str(DEFAULT_MODEL_PATH))
    parser.add_argument("--fps", type=float, default=3.0)
    parser.add_argument("--max-frames", type=int, default=DEFAULT_MAX_FRAMES)
    parser.add_argument("--min-frames", type=int, default=DEFAULT_MIN_FRAMES)
    parser.add_argument(
        "--inference-mode",
        choices=["frame_steps", "whole"],
        default=DEFAULT_INFERENCE_MODE,
        help="frame_steps: stable dense curve via prefix inference; whole: one forward pass on the full sampled trajectory",
    )
    parser.add_argument("--prefix-sample-frames", type=int, default=DEFAULT_PREFIX_SAMPLE_FRAMES)
    parser.add_argument("--tag", type=str, default=None,
                        help="Optional suffix appended to chunk output dir (e.g. 'original' → Chunk000_original)")
    parser.add_argument("--prefix-batch-size", type=int, default=DEFAULT_PREFIX_BATCH_SIZE)
    parser.add_argument(
        "--adaptive-max-frames",
        dest="adaptive_max_frames",
        action="store_true",
        help="on CUDA OOM, retry an episode with smaller frame budgets",
    )
    parser.add_argument(
        "--no-adaptive-max-frames",
        dest="adaptive_max_frames",
        action="store_false",
        help="disable dynamic frame-budget retry",
    )
    parser.add_argument("--overwrite", action="store_true", help="rerun episodes even if results.json already exists")
    parser.add_argument("--allow-no-marks", action="store_true", help="allow episodes without progress_mark pairs")
    parser.set_defaults(adaptive_max_frames=True)
    args = parser.parse_args()

    if "CUDA_VISIBLE_DEVICES" not in os.environ:
        user_gpu = input("GPU id (default 5): ").strip()
        os.environ["CUDA_VISIBLE_DEVICES"] = user_gpu if user_gpu else "5"

    chunks = parse_chunk_selection(args.chunk)
    runner = RobometerLocalRunner(model_path=args.model_path)
    ran_any = False

    for chunk in chunks:
        direct_chunk_dir = args.videos_root / chunk
        if direct_chunk_dir.exists():
            videos_root = direct_chunk_dir
        else:
            chunk_filtered = f"{chunk}_filtered"
            videos_root = args.videos_root / chunk_filtered
        annotations_path = videos_root / "annotations.json"
        episode_tasks_path = videos_root / "episode_tasks.json"
        catalog = load_episode_catalog(episode_tasks_path, annotations_path)
        progress_catalog = load_progress_mark_catalog(annotations_path)
        require_progress_marks = not bool(args.allow_no_marks)

        episode_list = resolve_episode_list(
            episode_arg=args.episode,
            count=args.count,
            catalog=catalog,
            progress_catalog=progress_catalog,
            require_progress_marks=require_progress_marks,
        )

        output_root = DEFAULT_OUTPUT_ROOT / chunk_output_dir_name(chunk, args.tag)
        if not args.overwrite:
            filtered = []
            for episode_mp4 in episode_list:
                if result_path_for(output_root, episode_mp4).exists():
                    print(f"[SKIP] {episode_mp4[:-4]} already done")
                    continue
                filtered.append(episode_mp4)
            episode_list = filtered

        if not episode_list:
            print(f"[OK] {chunk}: nothing to run.")
            continue

        ran_any = True
        print(f"[CHUNK] {chunk}")
        print(f"[QUEUE] {', '.join(ep[:-4] for ep in episode_list)}")
        for episode_mp4 in episode_list:
            run_episode(
                episode_mp4=episode_mp4,
                chunk=chunk,
                fps=float(args.fps),
                max_frames=int(args.max_frames),
                min_frames=int(args.min_frames),
                adaptive_max_frames=bool(args.adaptive_max_frames),
                inference_mode=str(args.inference_mode),
                prefix_sample_frames=int(args.prefix_sample_frames),
                prefix_batch_size=int(args.prefix_batch_size),
                runner=runner,
                overwrite=args.overwrite,
                videos_root_base=args.videos_root,
                allow_no_marks=bool(args.allow_no_marks),
                tag=args.tag,
            )

    if not ran_any:
        print("[OK] Nothing to run.")


if __name__ == "__main__":
    main()