"""Backbone-neutral trajectory contract for Route A. The module deliberately stops before model-specific normalization. Every backbone therefore receives the same selected RGB frames, labels and masks. """ from __future__ import annotations import hashlib import json from dataclasses import dataclass, replace from pathlib import Path from typing import Any, Iterable, Literal, Sequence import numpy as np import torch from PIL import Image TemporalStrategy = Literal["forward", "reverse", "rewind", "different_task"] def deterministic_frame_indices( frame_count: int, target_frames: int = 8 ) -> tuple[np.ndarray, np.ndarray]: """Return official-style rint(linspace) indices and a real-frame mask.""" if frame_count <= 0 or target_frames <= 0: raise ValueError("frame_count and target_frames must be positive") if frame_count >= target_frames: indices = np.rint( np.linspace(0, frame_count - 1, target_frames) ).astype(np.int64) indices[0] = 0 indices[-1] = frame_count - 1 return indices, np.ones(target_frames, dtype=np.bool_) real = np.arange(frame_count, dtype=np.int64) padding = np.full(target_frames - frame_count, frame_count - 1, dtype=np.int64) valid = np.concatenate( [ np.ones(frame_count, dtype=np.bool_), np.zeros(target_frames - frame_count, dtype=np.bool_), ] ) return np.concatenate([real, padding]), valid def temporal_permutation(length: int, strategy: TemporalStrategy) -> np.ndarray: """Return a permutation shared by frames, labels, indices and masks.""" if length <= 0: raise ValueError("length must be positive") if strategy in ("forward", "different_task"): return np.arange(length, dtype=np.int64) if strategy == "reverse": return np.arange(length - 1, -1, -1, dtype=np.int64) if strategy == "rewind": split = length // 2 return np.concatenate( [ np.arange(split, dtype=np.int64), np.arange(length - 1, split - 1, -1, dtype=np.int64), ] ) raise ValueError(f"unsupported temporal strategy: {strategy}") def letterbox_rgb(frame: np.ndarray, canvas_size: int = 256) -> np.ndarray: """Deterministically resize an RGB uint8 frame onto a square black canvas.""" if ( frame.ndim != 3 or frame.shape[-1] != 3 or frame.dtype != np.uint8 or canvas_size <= 0 ): raise ValueError("frame must be uint8 [H,W,3] and canvas_size must be positive") height, width = frame.shape[:2] if height <= 0 or width <= 0: raise ValueError("frame dimensions must be positive") scale = min(canvas_size / width, canvas_size / height) resized_width = max(1, min(canvas_size, round(width * scale))) resized_height = max(1, min(canvas_size, round(height * scale))) resized = Image.fromarray(frame, mode="RGB").resize( (resized_width, resized_height), Image.Resampling.BILINEAR ) canvas = Image.new("RGB", (canvas_size, canvas_size), (0, 0, 0)) offset = ( (canvas_size - resized_width) // 2, (canvas_size - resized_height) // 2, ) canvas.paste(resized, offset) return np.asarray(canvas, dtype=np.uint8) def _c51_targets(values: np.ndarray, bins: int) -> np.ndarray: if bins < 2: raise ValueError("progress bins must be at least 2") positions = np.clip(values, 0.0, 1.0) * (bins - 1) lower = np.floor(positions).astype(np.int64) upper = np.ceil(positions).astype(np.int64) output = np.zeros((len(values), bins), dtype=np.float32) rows = np.arange(len(values)) output[rows, lower] += upper - positions output[rows, upper] += positions - lower exact = lower == upper output[rows[exact], lower[exact]] = 1.0 return output @dataclass(frozen=True) class TrajectoryContract: trajectory_id: str task: str data_source: str quality_label: str frames_uint8: np.ndarray source_frame_indices: np.ndarray frame_valid_mask: np.ndarray progress: np.ndarray progress_distribution: np.ndarray progress_loss_mask: np.ndarray success_labels: np.ndarray success_loss_mask: np.ndarray strategy: TemporalStrategy metadata: dict[str, Any] def validate(self) -> None: frame_count = self.frames_uint8.shape[0] if self.frames_uint8.dtype != np.uint8 or self.frames_uint8.ndim != 4: raise ValueError("frames_uint8 must be uint8 [T,H,W,3]") if self.frames_uint8.shape[-1] != 3: raise ValueError("frames_uint8 must have three RGB channels") aligned = ( self.source_frame_indices, self.frame_valid_mask, self.progress, self.progress_distribution, self.progress_loss_mask, self.success_labels, self.success_loss_mask, ) if any(value.shape[0] != frame_count for value in aligned): raise ValueError("all frame-level fields must have the same leading size") if np.any(self.progress < 0) or np.any(self.progress > 1): raise ValueError("progress values must be in [0,1]") if np.any(self.progress_loss_mask & ~self.frame_valid_mask): raise ValueError("padding may not enter progress loss") if np.any(self.success_loss_mask & ~self.frame_valid_mask): raise ValueError("padding may not enter success loss") @dataclass(frozen=True) class PreferenceContract: trajectory_a: TrajectoryContract trajectory_b: TrajectoryContract preference_label: float strategy: str def swapped(self) -> PreferenceContract: return replace( self, trajectory_a=self.trajectory_b, trajectory_b=self.trajectory_a, preference_label=1.0 - self.preference_label, ) def prepare_trajectory( *, trajectory_id: str, frames: np.ndarray, task: str, data_source: str, quality_label: str, partial_success: float | None = None, strategy: TemporalStrategy = "forward", target_frames: int = 8, canvas_size: int = 256, progress_bins: int = 10, success_cutoff: float = 1.0, metadata: dict[str, Any] | None = None, ) -> TrajectoryContract: """Materialize one auditable trajectory from an official sampler trajectory.""" if frames.ndim != 4 or frames.shape[-1] != 3 or frames.dtype != np.uint8: raise ValueError("frames must be uint8 [T,H,W,3]") if not trajectory_id or not task or not data_source: raise ValueError("trajectory_id, task and data_source are required") if partial_success is not None and not 0.0 <= partial_success <= 1.0: raise ValueError("partial_success must be in [0,1]") if not 0.0 <= success_cutoff <= 1.0: raise ValueError("success_cutoff must be in [0,1]") selected_indices, frame_valid = deterministic_frame_indices( len(frames), target_frames ) permutation = temporal_permutation(target_frames, strategy) selected_indices = selected_indices[permutation] frame_valid = frame_valid[permutation] selected_frames = np.stack( [letterbox_rgb(frames[index], canvas_size) for index in selected_indices] ) denominator = max(len(frames) - 1, 1) progress = selected_indices.astype(np.float32) / denominator if partial_success is not None: progress *= partial_success if strategy == "different_task": progress.fill(0.0) quality = quality_label.lower() dense_progress_allowed = quality not in {"failure", "failed", "suboptimal"} progress_loss_mask = frame_valid & dense_progress_allowed if strategy == "different_task": progress_loss_mask = frame_valid.copy() success_labels = ( (progress >= success_cutoff) & (quality == "successful") & frame_valid ).astype(np.float32) success_loss_mask = frame_valid.copy() audit_metadata = dict(metadata or {}) audit_metadata.update( { "source_frame_count": int(len(frames)), "selected_frame_indices": selected_indices.tolist(), "frame_valid_mask": frame_valid.tolist(), "temporal_strategy": strategy, "canvas_size": canvas_size, "letterbox_fill_rgb": [0, 0, 0], } ) result = TrajectoryContract( trajectory_id=trajectory_id, task=task, data_source=data_source, quality_label=quality_label, frames_uint8=selected_frames, source_frame_indices=selected_indices, frame_valid_mask=frame_valid, progress=progress, progress_distribution=_c51_targets(progress, progress_bins), progress_loss_mask=progress_loss_mask, success_labels=success_labels, success_loss_mask=success_loss_mask, strategy=strategy, metadata=audit_metadata, ) result.validate() return result def adapt_official_trajectory( trajectory: Any, *, strategy: TemporalStrategy, progress_supervised: bool, canvas_size: int = 256, progress_bins: int = 10, ) -> TrajectoryContract: """Adapt an official Robometer sampler trajectory without changing its labels. The sampler must expose its selected source indices and padding mask in ``trajectory.metadata``. This is intentionally strict: silently re-sampling here would invalidate the matched-backbone comparison. """ metadata = dict(trajectory.metadata or {}) if "selected_frame_indices" not in metadata or "frame_valid_mask" not in metadata: raise ValueError("official sampler metadata lacks frozen frame indices/mask") frames = np.asarray(trajectory.frames) if frames.ndim != 4 or frames.shape[-1] != 3 or frames.dtype != np.uint8: raise ValueError("official sampler frames must be uint8 [T,H,W,3]") selected_indices = np.asarray(metadata["selected_frame_indices"], dtype=np.int64) frame_valid = np.asarray(metadata["frame_valid_mask"], dtype=np.bool_) if len(selected_indices) != len(frames) or len(frame_valid) != len(frames): raise ValueError("official sampler frame metadata is not aligned") target_progress = trajectory.target_progress if target_progress is None or len(target_progress) != len(frames): raise ValueError("official target_progress is missing or misaligned") first_target = target_progress[0] if isinstance(first_target, torch.Tensor) or np.asarray(first_target).ndim > 0: progress_distribution = np.stack( [np.asarray(value, dtype=np.float32) for value in target_progress] ) if progress_distribution.shape != (len(frames), progress_bins): raise ValueError("official discrete progress target has unexpected shape") support = np.linspace(0.0, 1.0, progress_bins, dtype=np.float32) progress = progress_distribution @ support else: progress = np.asarray(target_progress, dtype=np.float32) progress_distribution = _c51_targets(progress, progress_bins) predict_mask = trajectory.predict_last_frame_mask if predict_mask is None: predict_valid = np.ones(len(frames), dtype=np.bool_) else: predict_valid = np.asarray(predict_mask, dtype=np.float32) > 0 progress_loss_mask = frame_valid & predict_valid & progress_supervised success_labels = np.asarray(trajectory.success_label, dtype=np.float32) if success_labels.shape != (len(frames),): raise ValueError("official success labels are missing or misaligned") metadata.update( { "adapter": "official_robometer_sampler_v1", "temporal_strategy": strategy, "canvas_size": canvas_size, "letterbox_fill_rgb": [0, 0, 0], } ) result = TrajectoryContract( trajectory_id=trajectory.id, task=trajectory.task, data_source=trajectory.data_source, quality_label=trajectory.quality_label, frames_uint8=np.stack( [letterbox_rgb(frame, canvas_size) for frame in frames] ), source_frame_indices=selected_indices, frame_valid_mask=frame_valid, progress=progress, progress_distribution=progress_distribution, progress_loss_mask=progress_loss_mask, success_labels=success_labels, success_loss_mask=frame_valid.copy(), strategy=strategy, metadata=metadata, ) result.validate() return result def make_preference_contract( chosen: TrajectoryContract, rejected: TrajectoryContract, *, strategy: str, chosen_first: bool = True, ) -> PreferenceContract: if chosen_first: return PreferenceContract(chosen, rejected, 1.0, strategy) return PreferenceContract(rejected, chosen, 0.0, strategy) class CosmosRBMCollator: """Stack the frozen contract without invoking a backbone processor.""" schema_version = "cosmos_rbm_batch_v1" @staticmethod def _stack_trajectories(items: Sequence[TrajectoryContract]) -> dict[str, Any]: for item in items: item.validate() return { "frames_uint8": torch.from_numpy( np.stack([item.frames_uint8 for item in items]) ), "source_frame_indices": torch.from_numpy( np.stack([item.source_frame_indices for item in items]) ), "frame_valid_mask": torch.from_numpy( np.stack([item.frame_valid_mask for item in items]) ), "progress": torch.from_numpy(np.stack([item.progress for item in items])), "progress_distribution": torch.from_numpy( np.stack([item.progress_distribution for item in items]) ), "progress_loss_mask": torch.from_numpy( np.stack([item.progress_loss_mask for item in items]) ), "success_labels": torch.from_numpy( np.stack([item.success_labels for item in items]) ), "success_loss_mask": torch.from_numpy( np.stack([item.success_loss_mask for item in items]) ), "trajectory_ids": [item.trajectory_id for item in items], "tasks": [item.task for item in items], "data_sources": [item.data_source for item in items], "quality_labels": [item.quality_label for item in items], "strategies": [item.strategy for item in items], "metadata": [item.metadata for item in items], } def __call__( self, samples: Sequence[TrajectoryContract | PreferenceContract] ) -> dict[str, Any]: trajectories = [item for item in samples if isinstance(item, TrajectoryContract)] preferences = [item for item in samples if isinstance(item, PreferenceContract)] output: dict[str, Any] = { "schema_version": self.schema_version, "num_trajectories": len(trajectories), "num_preferences": len(preferences), } if trajectories: output["trajectory_inputs"] = self._stack_trajectories(trajectories) if preferences: output["preference_inputs"] = { "trajectory_a": self._stack_trajectories( [item.trajectory_a for item in preferences] ), "trajectory_b": self._stack_trajectories( [item.trajectory_b for item in preferences] ), "preference_labels": torch.tensor( [item.preference_label for item in preferences], dtype=torch.float32 ), "strategies": [item.strategy for item in preferences], } return output def validate_span_bounds( bounds: Sequence[tuple[int, int]], frame_count: int, sequence_length: int ) -> None: """Validate the one-frame/one-nonoverlapping-token-span contract.""" if len(bounds) != frame_count: raise ValueError("visual span count must match frame count") previous_stop = -1 for start, stop in bounds: if start < 0 or stop <= start or stop > sequence_length: raise ValueError("visual token span is empty or out of bounds") if start < previous_stop: raise ValueError("visual token spans overlap") previous_stop = stop def worker_invariant_sample_ids( ids: Iterable[str], *, seed: int, worker_count: int ) -> list[str]: """Simulate round-robin worker sharding while preserving global sample order.""" if worker_count <= 0: raise ValueError("worker_count must be positive") materialized = list(ids) unique = list(dict.fromkeys(materialized)) if len(unique) != len(materialized): raise ValueError("sample ids must be unique") ranked = sorted( unique, key=lambda value: hashlib.sha256(f"{seed}:{value}".encode()).digest(), ) shards = [list(enumerate(ranked))[worker::worker_count] for worker in range(worker_count)] return [value for _, value in sorted(item for shard in shards for item in shard)] def audit_group_overlap(train_ids: Iterable[str], eval_ids: Iterable[str]) -> set[str]: return set(train_ids) & set(eval_ids) def dump_batch_audit(batch: dict[str, Any], output: Path) -> None: """Write metadata and tensor shapes only; never dump trajectory pixels.""" def describe(value: Any) -> Any: if isinstance(value, torch.Tensor): return {"shape": list(value.shape), "dtype": str(value.dtype)} if isinstance(value, dict): return {key: describe(item) for key, item in value.items()} return value output.parent.mkdir(parents=True, exist_ok=True) output.write_text(json.dumps(describe(batch), indent=2) + "\n", encoding="utf-8")