#!/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()