diff --git a/egomimic/eval/bimanual_cartesian_eval.py b/egomimic/eval/bimanual_cartesian_eval.py index a16a298b2..ca6628eeb 100644 --- a/egomimic/eval/bimanual_cartesian_eval.py +++ b/egomimic/eval/bimanual_cartesian_eval.py @@ -311,7 +311,9 @@ def _deinterpolate(chunk: np.ndarray, rows: int) -> np.ndarray: index = np.linspace(0, chunk.shape[-2] - 1, rows).round().astype(int) return chunk[..., index, :] - def _arc_pred_time_indexed(self, prediction: torch.Tensor, embodiment_id: int): + def _arc_pred_time_indexed( + self, prediction: torch.Tensor, embodiment_id: int, *, source_batch=None + ): """Prediction as a time-indexed (B, T, 14) numpy chunk. Baseline: unnormalized pose chunk at tokenizer resolution (no @@ -343,6 +345,8 @@ def _arc_pred_for_arcmatch( prediction: torch.Tensor, ground_truth: np.ndarray, embodiment_id: int, + *, + source_batch=None, ): """Prediction geometry for shared-D arcmatch. @@ -350,7 +354,9 @@ def _arc_pred_for_arcmatch( Arc subclass returns token waypoints unless any arm has ``L_gt < D``. """ del ground_truth - return self._arc_pred_time_indexed(prediction, embodiment_id) + return self._arc_pred_time_indexed( + prediction, embodiment_id, source_batch=source_batch + ) def _extra_metrics( self, @@ -375,7 +381,7 @@ def _extra_metrics( if ground_truth is None or len(ground_truth) == 0: return {} arcmatch_pred = self._arc_pred_for_arcmatch( - prediction, ground_truth, embodiment_id + prediction, ground_truth, embodiment_id, source_batch=source_batch ) if isinstance(arcmatch_pred, list): predictions = arcmatch_pred @@ -394,7 +400,9 @@ def _extra_metrics( min_distance_unit=self.arcmatch_distance, ) ) - decoded = self._arc_pred_time_indexed(prediction, embodiment_id) + decoded = self._arc_pred_time_indexed( + prediction, embodiment_id, source_batch=source_batch + ) ti_pred = [sample for sample in decoded] values.update(dtw_metrics(ti_pred, truth, max_samples=self.dtw_max_samples)) rows = min(decoded.shape[-2], ground_truth.shape[-2]) @@ -417,13 +425,15 @@ def _extra_metrics( for name, value in values.items() } - def _viz_source(self, actions: torch.Tensor, embodiment_id: int) -> torch.Tensor: + def _viz_source( + self, actions: torch.Tensor, embodiment_id: int, *, source_batch=None + ) -> torch.Tensor: """Rows to hand the revert transforms. Identity for a time-indexed run. A subclass whose action space is NOT a stack of poses must override this. The revert transforms rotate AND translate every row they are - given, so any row that is not a pose comes out as a position -- see - ArcBimanualCartesianEval, where row M is a velocity. + given. A representation with additional clock channels must decode into + physical poses first, with its source-specific measured anchor. """ del embodiment_id return actions @@ -491,8 +501,12 @@ def _maybe_log_overlay( pred_native = self._native(pred_normalized, embodiment_id).detach().cpu() gt_native = self._native(gt_normalized, embodiment_id).detach().cpu() # Everything downstream treats these rows as poses, so convert first. - pred_native = self._viz_source(pred_native, embodiment_id) - gt_native = self._viz_source(gt_native, embodiment_id) + pred_native = self._viz_source( + pred_native, embodiment_id, source_batch=source_batch + ) + gt_native = self._viz_source( + gt_native, embodiment_id, source_batch=source_batch + ) # Unnormalize the obs pose and drop an optional T_obs=1 axis. The revert # transforms expect per-sample obs of shape (D,), not (T_obs, D). diff --git a/egomimic/eval/open_loop_stream.py b/egomimic/eval/open_loop_stream.py new file mode 100644 index 000000000..bfc3a7155 --- /dev/null +++ b/egomimic/eval/open_loop_stream.py @@ -0,0 +1,1760 @@ +"""Episode-level open-loop segment simulation for bimanual cartesian policies. + +This evaluator measures a policy over an entire recorded episode rather than +averaging independent action chunks. Validation samples are observations at +known episode/frame indices. We cache the prediction made from each +observation, then replay the episode at validation end: + +* execute the first ``execute_fraction`` of a baseline control chunk, or cap + ARC by an exact waypoint prefix (default) or a mode-aware interpolated + distance prefix; +* compare those control-frequency commands with the ground-truth commands; +* advance to the observation at the resulting frame; +* repeat until the episode ends. + +The next observation is the recorded observation at the next boundary. This +is an oracle-observation open-loop segment rollout: it measures compounding +segment error while keeping the evaluation deterministic and comparable +between baseline and ARC representations. +""" + +from __future__ import annotations + +import json +import math +from collections import defaultdict +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist + +from egomimic.eval.bimanual_cartesian_eval import ( + BimanualCartesianEval, + overlay_annotation_fields, +) +from egomimic.eval.distance_budget_dtw import ( + ARC_DISTANCE_SEMANTICS, + METRIC_FRAME_KEY, + METRIC_VERSION, + per_arm_cumulative_distance, + score_distance_dtw_episode, + summarize_distance_dtw, +) +from egomimic.eval.eval import EvaluationDataRequirements +from egomimic.eval.video import EvalVideo +from egomimic.pl_utils.pl_data_utils import DEFAULT_VALID_GROUP +from egomimic.rldb.zarr.arc_length_tokenizer import ( + _rotation_to_ypr, + _ypr_to_rotation, + cumulative_rotation_length, + slerp_pair_ypr, +) +from egomimic.rldb.zarr.cartesian_arc import ( + TokenizeBimanualArcLengthCartesian, + bimanual_arc_token_shape, + validate_clock_mode, + validate_streams, +) + +XYZ_COLS = (0, 1, 2, 7, 8, 9) +YPR_COLS = (3, 4, 5, 10, 11, 12) +GRIP_COLS = (6, 13) +PAIRED_COLS = XYZ_COLS + GRIP_COLS +ARC_EXECUTION_CAP_MODES = ("waypoints", "distance") +ARC_VIDEO_TRAJECTORY_CAP_MODES = ( + "execution_horizon", + "distance", +) +METRIC_DIMS = { + "mse": 14, + "xyz_mse": len(XYZ_COLS), + "ypr_mse": len(YPR_COLS), + "grip_mse": len(GRIP_COLS), + "paired_mse": len(PAIRED_COLS), +} + + +def gt_chunk_in_anchor_frame( + rows: np.ndarray, frames: np.ndarray, anchor: np.ndarray +) -> np.ndarray: + """Recorded poses of later frames, expressed in the anchor frame's wrist frame. + + ``rows`` holds each later frame's own first ground-truth row (its commanded + pose in that frame's wrist frame), ``frames`` their EEF-to-world anchors + (m, 2, 4, 4) and ``anchor`` the replanning frame's (2, 4, 4). The result is + an (m, 14) chunk directly comparable with a prediction decoded at the + anchor, for any m: an ARC chunk executes a fixed distance, so its length in + frames is not bounded by the 100-row preserved ground-truth copy. + """ + from scipy.spatial.transform import Rotation + + rows = np.asarray(rows, dtype=np.float64) + frames = np.asarray(frames, dtype=np.float64) + anchor = np.asarray(anchor, dtype=np.float64) + out = rows.copy() + for arm, offset in enumerate((0, 7)): + rotation, origin = frames[:, arm, :3, :3], frames[:, arm, :3, 3] + anchor_rotation, anchor_origin = anchor[arm, :3, :3], anchor[arm, :3, 3] + world = np.einsum("mij,mj->mi", rotation, rows[:, offset : offset + 3]) + out[:, offset : offset + 3] = (world + origin - anchor_origin) @ anchor_rotation + relative = np.einsum( + "ji,mjk,mkl->mil", + anchor_rotation, + rotation, + _ypr_to_rotation(rows[:, offset + 3 : offset + 6]).as_matrix(), + ) + out[:, offset + 3 : offset + 6] = _rotation_to_ypr( + Rotation.from_matrix(relative) + ) + return out + + +def validate_arc_execution_cap_mode(mode: str) -> str: + """Normalize the two supported ARC execution-cap semantics.""" + value = str(mode).strip().lower() + if value not in ARC_EXECUTION_CAP_MODES: + raise ValueError( + "arc_execution_cap_mode must be one of " + f"{ARC_EXECUTION_CAP_MODES}, got {mode!r}" + ) + return value + + +def validate_arc_video_trajectory_cap_mode(mode: str) -> str: + """Normalize ARC overlay capping independently of metric semantics.""" + + value = str(mode).strip().lower() + if value not in ARC_VIDEO_TRAJECTORY_CAP_MODES: + raise ValueError( + "arc_video_trajectory_cap_mode must be one of " + f"{ARC_VIDEO_TRAJECTORY_CAP_MODES}, got {mode!r}" + ) + return value + + +def executed_control_steps(control_horizon: int, execute_fraction: float) -> int: + """Return the positive control-step prefix executed from each chunk.""" + + horizon = int(control_horizon) + fraction = float(execute_fraction) + if horizon < 1: + raise ValueError("control_horizon must be positive") + if not 0.0 < fraction <= 1.0: + raise ValueError("execute_fraction must be in (0, 1]") + return max(1, min(horizon, int(math.ceil(horizon * fraction)))) + + +def executed_arc_waypoints(num_waypoints: int, execute_fraction: float) -> int: + """Return an exact integral M-based ARC execution prefix.""" + M = int(num_waypoints) + fraction = float(execute_fraction) + if M < 2: + raise ValueError("num_waypoints must be at least two") + if not 0.0 < fraction <= 1.0: + raise ValueError("execute_fraction must be in (0, 1]") + scaled = M * fraction + count = int(round(scaled)) + if not math.isclose(scaled, count, rel_tol=0.0, abs_tol=1e-9): + raise ValueError( + "M-based ARC execution requires M * execute_fraction to be an " + f"integer, got {M} * {fraction} = {scaled}" + ) + if count < 2: + raise ValueError( + f"M-based ARC execution must retain at least two waypoints, got {count}" + ) + return count + + +def _distance_boundary(cumulative: np.ndarray, target: float) -> float: + """Fractional source-row index of a cap, or the final available row.""" + reached = np.flatnonzero(cumulative >= target) + if not len(reached): + return float(len(cumulative) - 1) + upper = int(reached[0]) + if upper == 0: + return 0.0 + delta = cumulative[upper] - cumulative[upper - 1] + alpha = (target - cumulative[upper - 1]) / delta if delta > 1e-12 else 1.0 + return float(upper - 1 + np.clip(alpha, 0.0, 1.0)) + + +def _boundary_rows(boundary: float) -> int: + return max(2, int(math.ceil(boundary - 1e-12)) + 1) + + +def _waypoint_at(waypoints: np.ndarray, boundary: float) -> np.ndarray: + lower = min(int(math.floor(boundary)), len(waypoints) - 2) + alpha = boundary - lower + result = waypoints[lower] + alpha * (waypoints[lower + 1] - waypoints[lower]) + for columns in (slice(3, 6), slice(10, 13)): + result[columns] = slerp_pair_ypr( + waypoints[lower, columns], waypoints[lower + 1, columns], np.array([alpha]) + )[0] + return result + + +def truncate_cartesian_trajectory_by_stream_distance( + trajectory: np.ndarray, + max_translation_distance: float, + max_rotation_distance: float | None = None, + max_gripper_distance: float | None = None, +) -> np.ndarray: + """Cap each independent stream at its own distance budget.""" + value = np.asarray(trajectory, dtype=np.float64) + if value.ndim != 2 or value.shape[1] != 14: + raise ValueError( + f"cartesian trajectory must have shape (T, 14), got {value.shape}" + ) + for cap in (max_translation_distance, max_rotation_distance, max_gripper_distance): + if cap is not None and (not math.isfinite(float(cap)) or cap <= 0): + raise ValueError("trajectory distance caps must be positive and finite") + if len(value) < 2: + return value.copy() + cumulative = per_arm_cumulative_distance(value[:, XYZ_COLS]) + translation_cumulatives = [cumulative[:, arm] for arm in range(2)] + boundaries = [ + _distance_boundary(values, max_translation_distance) + for values in translation_cumulatives + ] + rotation_cumulatives = None + rotation_boundaries = None + if max_rotation_distance is not None: + rotation_cumulatives = [ + cumulative_rotation_length(value[:, offset + 3 : offset + 6]) + for offset in (0, 7) + ] + rotation_boundaries = [ + _distance_boundary(values, max_rotation_distance) + for values in rotation_cumulatives + ] + gripper_cumulatives = None + gripper_boundaries = None + if max_gripper_distance is not None: + gripper_cumulatives = [ + np.concatenate(([0.0], np.cumsum(np.abs(np.diff(value[:, offset + 6]))))) + for offset in (0, 7) + ] + gripper_boundaries = [ + _distance_boundary(values, max_gripper_distance) + for values in gripper_cumulatives + ] + streams = [] + for arm, offset in enumerate((0, 7)): + columns = ( + list(range(offset, offset + 7)) + if max_rotation_distance is None + else [offset, offset + 1, offset + 2, offset + 6] + ) + streams.append((columns, boundaries[arm])) + if rotation_boundaries is not None: + streams.append( + (list(range(offset + 3, offset + 6)), rotation_boundaries[arm]) + ) + if gripper_boundaries is not None: + streams.append(([offset + 6], gripper_boundaries[arm])) + count = min(_boundary_rows(boundary) for _, boundary in streams) + result = value[:count].copy() + for columns, boundary in streams: + terminal = _waypoint_at(value, boundary) + end = _boundary_rows(boundary) - 1 + result[end:, columns] = terminal[columns] + return result + + +class OpenLoopSimEval(BimanualCartesianEval): + """Open-loop episode simulation using the canonical arrival-clock ARC codec.""" + + def data_requirements(self): + return EvaluationDataRequirements( + ordered=True, + complete_episodes=True, + max_episodes=self.limit_val_episodes, + sample_id_key="episode_hash", + frame_index_key="frame_index", + source_fps=1.0 / self.control_dt, + required_keys=(self.initial_action_key,) + if self.action_mode == "arc" + else (), + ) + + def trainer_overrides(self): + return {"limit_val_batches": 1.0, "num_sanity_val_steps": 0} + + def __init__( + self, + *, + action_key: str = "actions_cartesian", + ground_truth_action_key: str = "actions_cartesian_untokenized", + execute_fraction: float = 0.30, + control_horizon: int = 100, + control_dt: float = 1.0 / 30.0, + action_mode: str = "auto", + arc_execution_cap_mode: str = "waypoints", + arc_video_trajectory_cap_mode: str = "execution_horizon", + min_distance_unit: float = 0.40, + rotation_distance_unit: float = 1.7453292519943295, + resampled_vector_length: int = 100, + clock_mode: str = "velocity", + streams: str = "full_separation", + log_step: int | None = None, + gripper_distance_unit: float | None = None, + gripper_stream_mode: str = "events", + results_path: str | None = None, + records_dump_path: str | None = None, + video_overlay_mode: str = "executed_chunk", + trajectory_snapshot_path: str | None = None, + video_only: bool = False, + distance_dtw_enabled: bool = False, + dtw_max_cells: int = 50_000_000, + dtw_max_prediction_steps: int = 30_000, + dtw_max_segment_steps: int | None = None, + require_episode_start: bool = True, + limit_val_episodes: int | None = None, + requires_ordered_validation: bool = True, + limit_val_batches: int | float | None = None, + deterministic_seed: int = 420042, + obs_pose_key: str = "observations.state.ee_pose", + image_key: str = "observations.images.front_img_1", + viz_func: Mapping | None = None, + revert_transforms: Mapping | None = None, + video_output_dir: str | None = None, + episode_dump_dir: str | None = None, + video_chunk_frames: int = 1000, + max_episode_frames: int = 6000, + viz_every_n_epochs: int = 1, + viz_max_batches: int | None = None, + initial_action_key: str = "arc.current_action", + **kwargs, + ): + retired = { + "arc_chunking_mode", + "velocity_mode", + "velocity_layout", + "velocity_norm", + "token_layout", + "e1_speed_floors", + "projection_offset_m", + "projection_min_speed", + "duration_timing", + "speed_encoding", + "tick_prefix", + } + if retired.intersection(kwargs): + raise ValueError( + f"Retired ARC options: {sorted(retired.intersection(kwargs))}; use clock_mode and streams" + ) + mode = str(action_mode) + if mode not in ("auto", "baseline", "arc"): + raise ValueError("action_mode must be auto, baseline, or arc") + if float(control_dt) <= 0: + raise ValueError("control_dt must be positive") + if not math.isfinite(float(min_distance_unit)) or float(min_distance_unit) <= 0: + raise ValueError("min_distance_unit must be positive and finite") + if rotation_distance_unit is not None and ( + not math.isfinite(float(rotation_distance_unit)) + or float(rotation_distance_unit) <= 0 + ): + raise ValueError("rotation_distance_unit must be positive and finite") + if gripper_distance_unit is not None and ( + not math.isfinite(float(gripper_distance_unit)) + or float(gripper_distance_unit) <= 0 + ): + raise ValueError("gripper_distance_unit must be positive and finite") + validate_clock_mode(clock_mode) + validate_streams(streams) + if mode != "baseline" and arc_execution_cap_mode != "waypoints": + raise ValueError( + "Arrival-clock ARC supports waypoint-prefix execution only" + ) + self.execute_fraction = float(execute_fraction) + self.control_horizon = int(control_horizon) + self.control_dt = float(control_dt) + self.action_mode = mode + self.initial_action_key = str(initial_action_key) + self.arc_execution_cap_mode = validate_arc_execution_cap_mode( + arc_execution_cap_mode + ) + self.arc_video_trajectory_cap_mode = validate_arc_video_trajectory_cap_mode( + arc_video_trajectory_cap_mode + ) + self.ground_truth_action_key = str(ground_truth_action_key) + self.min_distance_unit = float(min_distance_unit) + self.rotation_distance_unit = ( + None if rotation_distance_unit is None else float(rotation_distance_unit) + ) + self.gripper_distance_unit = ( + None if gripper_distance_unit is None else float(gripper_distance_unit) + ) + if gripper_stream_mode not in ("distance", "events"): + raise ValueError("gripper_stream_mode must be 'distance' or 'events'") + self.gripper_stream_mode = gripper_stream_mode + if gripper_stream_mode == "events": + if self.arc_execution_cap_mode != "waypoints": + raise ValueError( + "event gripper streams require waypoint-prefix execution" + ) + self.gripper_distance_unit = None + self.resampled_vector_length = int(resampled_vector_length) + self.clock_mode, self.streams = clock_mode, streams + self.execute_arc_waypoints = ( + executed_arc_waypoints(self.resampled_vector_length, self.execute_fraction) + if mode == "arc" + else None + ) + self.log_step = None if log_step is None else int(log_step) + if self.log_step is not None and self.log_step < 0: + raise ValueError("log_step must be nonnegative") + self.results_path = Path(results_path) if results_path else None + # Opt-in: one .npz per scored episode with the executed predicted and ground-truth control steps + # (frame index, prediction, ground truth) for plotting. None (default) writes nothing. + self.episode_dump_dir = Path(episode_dump_dir) if episode_dump_dir else None + # Optional diagnostic: every validation frame's raw prediction, GT and + # metric frame, so decoding can be re-run offline (e.g. jitter studies). + self.records_dump_path = Path(records_dump_path) if records_dump_path else None + if video_overlay_mode not in ("per_frame", "executed_chunk"): + raise ValueError("video_overlay_mode must be per_frame or executed_chunk") + self.video_overlay_mode = video_overlay_mode + self.trajectory_snapshot_path = ( + Path(trajectory_snapshot_path) if trajectory_snapshot_path else None + ) + if ( + self.trajectory_snapshot_path is not None + and self.trajectory_snapshot_path.suffix != ".npz" + ): + raise ValueError("trajectory_snapshot_path must end in .npz") + self._trajectory_snapshot_written = False + self._video_seen_frames = set() + self._video_last_frame = {} + self.video_only = bool(video_only) + self.distance_dtw_enabled = bool(distance_dtw_enabled) + self.dtw_max_cells = int(dtw_max_cells) + self.dtw_max_prediction_steps = int(dtw_max_prediction_steps) + self.dtw_max_segment_steps = ( + None if dtw_max_segment_steps is None else int(dtw_max_segment_steps) + ) + if self.dtw_max_segment_steps is not None and self.dtw_max_segment_steps < 1: + raise ValueError("dtw_max_segment_steps must be positive when set") + if self.dtw_max_cells < 1 or self.dtw_max_prediction_steps < 1: + raise ValueError("DTW resource limits must be positive") + self.require_episode_start = bool(require_episode_start) + self.limit_val_episodes = ( + None if limit_val_episodes is None else int(limit_val_episodes) + ) + if self.limit_val_episodes is not None and self.limit_val_episodes < 1: + raise ValueError("limit_val_episodes must be positive") + self.requires_ordered_validation = bool(requires_ordered_validation) + if limit_val_batches is not None: + raise ValueError( + "open_loop_sim is episode-based; use limit_val_episodes instead " + "of limit_val_batches" + ) + self._records: list[dict[str, Any]] = [] + self.last_results = None + self._arc_tokenizer = None + self._metric_device = torch.device("cpu") + + # Reuse the graph evaluator's normalizer binding, deterministic model, + # embodiment resolution, metric-group namespacing, and episode-aware + # video buffering. Chunk-level ARC metrics remain deliberately disabled + # here because this evaluator scores the executed control prefix. + # Existing ABC experiment blocks are merged into a selected evaluator + # config by Hydra. Consume their legacy ARC-only knobs so selecting + # this evaluator does not fail on an unrelated ``action_horizon`` (or + # accidentally enable chunk/video metrics). + for ignored_key in ( + "arc_metrics", + "include_reconstruction_loss", + "arcmatch_distance", + "arcmatch_points", + "arc_chunk_rows", + "action_horizon", + "rot_lever_m", + "dtw_max_samples", + ): + kwargs.pop(ignored_key, None) + super().__init__( + action_key=action_key, + obs_pose_key=obs_pose_key, + image_key=image_key, + viz_func=viz_func, + revert_transforms=revert_transforms, + arc_metrics=False, + deterministic_seed=deterministic_seed, + limit_val_batches=None, + video_output_dir=video_output_dir, + video_chunk_frames=video_chunk_frames, + max_episode_frames=max_episode_frames, + viz_every_n_epochs=viz_every_n_epochs, + viz_max_batches=viz_max_batches, + **kwargs, + ) + self._video_enabled = bool(self.viz_func) + if self.video_only and not self._video_enabled: + raise ValueError("video_only requires a configured visualization function") + self.execute_steps = executed_control_steps( + self.control_horizon, self.execute_fraction + ) + + def on_validation_start(self): + if self._video_enabled: + EvalVideo.on_validation_start(self) + self._chunk_states = {} + self._records = [] + self.last_results = None + self._trajectory_snapshot_written = False + self._video_seen_frames = set() + self._video_last_frame = {} + if self.model is not None: + try: + self._metric_device = next(self.model.parameters()).device + except StopIteration: + pass + + def _decoded_video_predictions( + self, + prediction: torch.Tensor, + embodiment_id: int, + max_steps: int, + initial_actions=None, + ) -> tuple[torch.Tensor, list[int]]: + """Decode only each frame's independently executed prediction prefix.""" + + native = self._native(prediction, embodiment_id).detach().cpu().numpy() + decoded_with_steps = [ + self._decode_prediction_with_steps( + sample, + initial_action=None if initial_actions is None else initial_actions[i], + max_steps=max_steps, + ) + for i, sample in enumerate(native) + ] + lengths = [steps for _, steps in decoded_with_steps] + width = max(lengths) + decoded = [] + for value, steps in decoded_with_steps: + if steps < width: + value = np.concatenate( + (value, np.repeat(value[-1:], width - steps, axis=0)), axis=0 + ) + decoded.append(value) + return ( + torch.from_numpy(np.stack(decoded).astype(np.float32, copy=False)), + lengths, + ) + + def _cap_arc_video_trajectories( + self, + prediction: torch.Tensor, + ground_truth: torch.Tensor, + prefix_lengths: list[int], + ) -> tuple[torch.Tensor, torch.Tensor, list[int], list[int]]: + """Cap ARC overlays at the selected translation and rotation prefixes.""" + + cap = self.execute_fraction * self.min_distance_unit + rotation_distance_unit = getattr(self, "rotation_distance_unit", None) + pred_np = prediction.detach().cpu().numpy() + gt_np = ground_truth.detach().cpu().numpy() + pred_values = [] + gt_values = [] + pred_lengths = [] + gt_lengths = [] + for index, steps in enumerate(prefix_lengths): + rotation_cap = ( + None + if rotation_distance_unit is None or self.streams == "translation" + else self.execute_fraction * rotation_distance_unit + ) + gripper_cap = ( + None + if self.gripper_distance_unit is None + else self.execute_fraction * self.gripper_distance_unit + ) + pred = truncate_cartesian_trajectory_by_stream_distance( + pred_np[index, :steps], + cap, + rotation_cap, + gripper_cap, + ) + gt = truncate_cartesian_trajectory_by_stream_distance( + gt_np[index, :steps], + cap, + rotation_cap, + gripper_cap, + ) + pred_values.append(pred) + gt_values.append(gt) + pred_lengths.append(len(pred)) + gt_lengths.append(len(gt)) + + width = max(max(pred_lengths), max(gt_lengths)) + + def _pad(values: list[np.ndarray]) -> torch.Tensor: + padded = [] + for value in values: + if len(value) < width: + value = np.concatenate( + (value, np.repeat(value[-1:], width - len(value), axis=0)), + axis=0, + ) + padded.append(value) + return torch.from_numpy(np.stack(padded).astype(np.float32, copy=False)) + + return _pad(pred_values), _pad(gt_values), pred_lengths, gt_lengths + + def _video_fps(self, source_fps=30): + """Full-frame video uses the control clock, not the DDP shard rate.""" + return max(1, round(1.0 / float(getattr(self, "control_dt", 1.0 / source_fps)))) + + def _collective_video_enabled(self, batch_idx=0): + cadence = self.viz_every_n_epochs + return ( + cadence > 0 + and (getattr(self.trainer, "current_epoch", 0) + 1) % cadence == 0 + and (self.viz_max_batches is None or batch_idx < self.viz_max_batches) + ) + + def _collect_open_loop_video( + self, *, source_id, source_batch, prediction, embodiment_id, embodiment_name + ): + """Gather interleaved DDP frames; only this experiment's rank zero renders. + + Lightning's ordered DistributedSampler assigns index i*world+rank. + Interleaving rank batches reconstructs dataset order without assuming + episode IDs are lexicographically ordered. Sampler padding is deduplicated. + """ + + def cpu(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu() + if isinstance(value, dict): + return {k: cpu(v) for k, v in value.items()} + if isinstance(value, (tuple, list)): + return type(value)(cpu(v) for v in value) + return value + + # Wrist-camera observations are not used by this front-camera overlay. + batch = { + k: v + for k, v in source_batch.items() + if not str(k).startswith("observations.images.") or k == self.image_key + } + payload = dict( + source_id=str(source_id), + source_batch=cpu(batch), + prediction=cpu(prediction), + embodiment_id=embodiment_id, + embodiment_name=embodiment_name, + ) + distributed = dist.is_available() and dist.is_initialized() + world = dist.get_world_size() if distributed else 1 + rank = dist.get_rank() if distributed else 0 + parts = [None] * world + if distributed: + dist.all_gather_object(parts, payload) + else: + parts[0] = payload + if rank != 0: + return + + if not hasattr(self, "_video_seen_frames"): + self._video_seen_frames = set() + self._video_last_frame = {} + for part in parts: + if (part["source_id"], part["embodiment_name"]) != ( + str(source_id), + embodiment_name, + ): + raise RuntimeError( + "Validation ranks disagree on video source/embodiment" + ) + for index in range(max(len(part["prediction"]) for part in parts)): + for part in parts: + n = len(part["prediction"]) + if index >= n: + continue + + def one(value): + if ( + isinstance(value, (torch.Tensor, np.ndarray)) + and value.ndim + and len(value) == n + ): + return value[index : index + 1] + if isinstance(value, (tuple, list)) and len(value) == n: + return value[index : index + 1] + return value + + item = {k: one(v) for k, v in part["source_batch"].items()} + episode = str( + self._batch_values(item["episode_hash"], 1, "episode_hash")[0] + ) + frame = int( + self._batch_values(item["frame_index"], 1, "frame_index")[0] + ) + group = self._validation_group or DEFAULT_VALID_GROUP + episode_key = (group, str(source_id), episode) + key = (*episode_key, frame) + if key in self._video_seen_frames: + continue + expected = self._video_last_frame.get(episode_key, -1) + 1 + if frame != expected: + raise RuntimeError( + f"Full-frame video is missing/out of order: {episode_key}, " + f"expected frame {expected}, got {frame}" + ) + self._maybe_log_open_loop_video( + source_id=source_id, + source_batch=item, + prediction=part["prediction"][index : index + 1], + embodiment_id=embodiment_id, + embodiment_name=embodiment_name, + ) + self._video_seen_frames.add(key) + self._video_last_frame[episode_key] = frame + + def _maybe_log_open_loop_video( + self, + *, + source_id: str, + source_batch: Mapping, + prediction: torch.Tensor, + embodiment_id: int, + embodiment_name: str, + ) -> None: + """Render one frame per validation sample and buffer by episode hash. + + Like the metric path, each video frame shows only the independently + executed prediction prefix and its matching control-frequency ground + truth. ARC uses the configured waypoint- or distance-based cap before + its decoded timing determines the frame width. + """ + + if not getattr(self, "_video_enabled", False): + return + viz_partial = self.viz_func.get(embodiment_name) + if viz_partial is None or self.obs_pose_key not in source_batch: + return + if getattr(self, "video_overlay_mode", "executed_chunk") == "executed_chunk": + self._log_executed_chunk_frame( + source_id=source_id, + source_batch=source_batch, + prediction=prediction, + embodiment_id=embodiment_id, + embodiment_name=embodiment_name, + viz_partial=viz_partial, + ) + return + target_key = ( + self.ground_truth_action_key + if self.ground_truth_action_key in source_batch + else self.action_key + ) + gt_native = ( + self._native_key(source_batch[target_key], target_key, embodiment_id) + .detach() + .cpu() + ) + if gt_native.ndim != 3 or gt_native.shape[0] != prediction.shape[0]: + raise ValueError( + "open_loop_sim video ground truth must be batched as (B, T, 14), " + f"got {tuple(gt_native.shape)}" + ) + pred_native, prefix_lengths = self._decoded_video_predictions( + prediction, + embodiment_id, + int(gt_native.shape[1]), + self._initial_actions(source_batch, len(prediction)), + ) + width = int(pred_native.shape[1]) + gt_prefixes = [] + for index, steps in enumerate(prefix_lengths): + value = gt_native[index, :steps] + if steps < width: + value = torch.cat((value, value[-1:].repeat(width - steps, 1)), dim=0) + gt_prefixes.append(value) + gt_native = torch.stack(gt_prefixes).to(dtype=pred_native.dtype) + + prediction_lengths = list(prefix_lengths) + ground_truth_lengths = list(prefix_lengths) + is_arc_video = self.action_mode == "arc" or ( + self.action_mode == "auto" + and self._is_arc_prediction(prediction[0].detach().cpu().numpy()) + ) + if ( + is_arc_video + and getattr( + self, + "arc_video_trajectory_cap_mode", + "execution_horizon", + ) + == "distance" + ): + ( + pred_native, + gt_native, + prediction_lengths, + ground_truth_lengths, + ) = self._cap_arc_video_trajectories(pred_native, gt_native, prefix_lengths) + obs_pose_native = ( + self._native_pose(source_batch[self.obs_pose_key], embodiment_id) + .detach() + .cpu() + ) + if obs_pose_native.ndim == 3: + if "observation_frame_indices" in source_batch: + obs_pose_native = obs_pose_native[:, -1] + elif obs_pose_native.shape[1] == 1: + obs_pose_native = obs_pose_native.squeeze(1) + pred_camframe = self._revert_to_camframe( + actions=pred_native, + obs_pose=obs_pose_native, + embodiment_name=embodiment_name, + ) + gt_camframe = self._revert_to_camframe( + actions=gt_native, + obs_pose=obs_pose_native, + embodiment_name=embodiment_name, + ) + if pred_camframe is None or gt_camframe is None: + return + + self._write_trajectory_snapshot( + source_id=source_id, + source_batch=source_batch, + prediction=pred_camframe, + ground_truth=gt_camframe, + prediction_length=prediction_lengths[0], + ground_truth_length=ground_truth_lengths[0], + ) + + images = source_batch[self.image_key] + if images.ndim == 5: + images = images[:, -1 if "observation_frame_indices" in source_batch else 0] + images = images.detach().cpu() + if images.ndim == 4 and images.shape[1] in (1, 3): + images = images.permute(0, 2, 3, 1) + flat_predictions = { + f"{embodiment_name}_{self.action_key}": pred_camframe, + } + flat_batch = { + self.image_key: images, + self.action_key: gt_camframe, + "embodiment": source_batch["embodiment"].detach().cpu(), + } + if "intrinsics" in source_batch: + flat_batch["intrinsics"] = source_batch["intrinsics"].detach().cpu() + flat_batch.update( + overlay_annotation_fields( + viz_partial, {**source_batch, "source": source_id} + ) + ) + try: + frames = viz_partial(predictions=flat_predictions, batch=flat_batch) + except Exception as exc: # noqa: BLE001 -- overlays are best effort + if self.complete_video_episodes: + raise RuntimeError( + f"Complete-episode overlay failed for {embodiment_name}" + ) from exc + print( + f"[OpenLoopSimEval] skipped {embodiment_name} overlay: {exc}", + flush=True, + ) + return + frames = np.asarray(frames) + if frames.dtype != np.uint8: + frames = np.clip(frames, 0, 255).astype(np.uint8) + if frames.ndim == 3: + frames = frames[None] + frame_tensor = torch.from_numpy(frames) + group = self._validation_group or DEFAULT_VALID_GROUP + buf_key = (group, embodiment_name) + self._record_video_frames(buf_key, frame_tensor, source_batch) + + def _frame_image(self, source_batch: Mapping) -> torch.Tensor: + images = source_batch[self.image_key] + if images.ndim == 5: + images = images[:, -1 if "observation_frame_indices" in source_batch else 0] + images = images.detach().cpu() + if images.ndim == 4 and images.shape[1] in (1, 3): + images = images.permute(0, 2, 3, 1) + return images + + def _log_executed_chunk_frame( + self, + *, + source_id, + source_batch: Mapping, + prediction: torch.Tensor, + embodiment_id: int, + embodiment_name: str, + viz_partial, + ) -> None: + """Hold one executed chunk's overlay for the frames it spans, then replan. + + At a replanning frame the prediction is decoded exactly as the open-loop + metric executes it: the first execute_fraction of the waypoints, every + stream on its own clock, stopping when the first moving stream runs + out. That chunk of n frames (no fixed frame count) stays on screen for + the next n frames while the video advances. The GT overlay is the + recorded trajectory over those same n frames, in the replanning frame's + wrist frame, so both lines cover identical time. Frames are buffered + until the chunk ends because its GT needs the later frames. + """ + episode = str( + self._batch_values(source_batch["episode_hash"], 1, "episode_hash")[0] + ) + frame = int( + self._batch_values(source_batch["frame_index"], 1, "frame_index")[0] + ) + if METRIC_FRAME_KEY not in source_batch: + raise ValueError( + f"executed_chunk videos need {METRIC_FRAME_KEY} (euler wrist-frame data)" + ) + key = (self._validation_group or DEFAULT_VALID_GROUP, embodiment_name) + states = self.__dict__.setdefault("_chunk_states", {}) + state = states.get(key) + if state is not None and (state["episode"] != episode or frame >= state["end"]): + self._flush_executed_chunk(key) + state = None + metric = np.asarray( + torch.as_tensor(source_batch[METRIC_FRAME_KEY]).detach().cpu(), + dtype=np.float64, + )[0] + target_key = ( + self.ground_truth_action_key + if self.ground_truth_action_key in source_batch + else self.action_key + ) + gt_row = ( + self._native_key(source_batch[target_key], target_key, embodiment_id) + .detach() + .cpu() + .numpy()[0, 0] + ) + if state is None: + native = self._native(prediction, embodiment_id).detach().cpu().numpy()[0] + initial_actions = self._initial_actions(source_batch, 1) + decoded, steps = self._decode_prediction_with_steps( + native, + initial_action=None if initial_actions is None else initial_actions[0], + max_steps=getattr(self, "dtw_max_prediction_steps", 100_000), + ) + obs_pose = ( + self._native_pose(source_batch[self.obs_pose_key], embodiment_id) + .detach() + .cpu() + ) + if obs_pose.ndim == 3: + if "observation_frame_indices" in source_batch: + obs_pose = obs_pose[:, -1] + elif obs_pose.shape[1] == 1: + obs_pose = obs_pose.squeeze(1) + state = dict( + episode=episode, + end=frame + steps, + viz_partial=viz_partial, + prediction=decoded, + obs_pose=obs_pose, + anchor=metric, + frames=[], + ) + states[key] = state + state["frames"].append( + dict( + image=self._frame_image(source_batch), + gt_row=gt_row, + metric=metric, + embodiment=source_batch["embodiment"].detach().cpu(), + intrinsics=source_batch.get("intrinsics"), + annotations=overlay_annotation_fields( + viz_partial, {**source_batch, "source": source_id} + ), + ) + ) + + def _flush_executed_chunk(self, key) -> None: + state = self.__dict__.get("_chunk_states", {}).pop(key, None) + if not state or not state["frames"]: + return + group, embodiment_name = key + frames = state["frames"] + gt = gt_chunk_in_anchor_frame( + [item["gt_row"] for item in frames], + [item["metric"] for item in frames], + state["anchor"], + ) + + def batched(values): + return torch.from_numpy(np.asarray(values, dtype=np.float32)[None]) + + pred_camframe = self._revert_to_camframe( + # Episode end may arrive before this planned chunk finishes. + # Match the metric: render only frames that were actually executed. + actions=batched(state["prediction"][: len(frames)]), + obs_pose=state["obs_pose"], + embodiment_name=embodiment_name, + ) + gt_camframe = self._revert_to_camframe( + actions=batched(gt), + obs_pose=state["obs_pose"], + embodiment_name=embodiment_name, + ) + if pred_camframe is None or gt_camframe is None: + return + out_dir = self._group_video_dir(group, embodiment_name) + for item in frames: + flat_batch = { + self.image_key: item["image"], + self.action_key: gt_camframe, + "embodiment": item["embodiment"], + **item["annotations"], + } + if item["intrinsics"] is not None: + flat_batch["intrinsics"] = ( + torch.as_tensor(item["intrinsics"]).detach().cpu() + ) + try: + rendered = state["viz_partial"]( + predictions={f"{embodiment_name}_{self.action_key}": pred_camframe}, + batch=flat_batch, + ) + except Exception as exc: # noqa: BLE001 -- overlays are best effort + print( + f"[OpenLoopSimEval] skipped {embodiment_name} overlay: {exc}", + flush=True, + ) + return + rendered = np.asarray(rendered) + if rendered.dtype != np.uint8: + rendered = np.clip(rendered, 0, 255).astype(np.uint8) + if rendered.ndim == 3: + rendered = rendered[None] + self._buffer_per_episode( + key, + out_dir, + list(torch.from_numpy(rendered)), + [state["episode"]] * len(rendered), + ) + + def _flush_all_executed_chunks(self) -> None: + for key in list(self.__dict__.get("_chunk_states", {})): + self._flush_executed_chunk(key) + + def _write_trajectory_snapshot( + self, + *, + source_id: str, + source_batch: Mapping, + prediction, + ground_truth, + prediction_length: int, + ground_truth_length: int, + ) -> None: + """Save the first rendered executed prefix for speed/shape diagnostics.""" + + path = getattr(self, "trajectory_snapshot_path", None) + if path is None or getattr(self, "_trajectory_snapshot_written", False): + return + + def _numpy(value) -> np.ndarray: + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + return np.asarray(value) + + pred = _numpy(prediction) + gt = _numpy(ground_truth) + if pred.ndim != 3 or gt.ndim != 3 or pred.shape[0] < 1 or gt.shape[0] < 1: + raise ValueError( + "trajectory snapshot expects batched (B, T, 14) camera-frame arrays" + ) + pred_steps = min(int(prediction_length), int(pred.shape[1])) + gt_steps = min(int(ground_truth_length), int(gt.shape[1])) + if pred_steps < 1 or gt_steps < 1: + raise ValueError("trajectory snapshot execution prefix is empty") + + batch_size = int(pred.shape[0]) + episode = self._batch_values( + source_batch["episode_hash"], batch_size, "episode_hash" + )[0] + frame = -1 + if "frame_index" in source_batch: + frame = int( + self._batch_values( + source_batch["frame_index"], batch_size, "frame_index" + )[0] + ) + + path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + path, + prediction=pred[0, :pred_steps].astype(np.float32, copy=False), + ground_truth=gt[0, :gt_steps].astype(np.float32, copy=False), + control_dt=np.asarray(self.control_dt, dtype=np.float64), + source_id=np.asarray(str(source_id)), + episode_hash=np.asarray(str(episode)), + frame_index=np.asarray(frame, dtype=np.int64), + action_mode=np.asarray(str(self.action_mode)), + arc_execution_cap_mode=np.asarray(str(self.arc_execution_cap_mode)), + ) + self._trajectory_snapshot_written = True + + def _log_wandb_videos(self) -> None: + """Upload only the first episode MP4 for each val loop/panel.""" + + if not self._written_paths: + return + experiment = self._wandb_logger() + if experiment is None: + return + import wandb # type: ignore + + first: dict[tuple[str, str], str] = {} + for group, embodiment_name, path in self._written_paths: + first.setdefault((group, embodiment_name), path) + payload = {} + for (group, embodiment_name), path in first.items(): + prefix = ( + "Val_video" if group == DEFAULT_VALID_GROUP else f"Val_video_{group}" + ) + payload[f"{prefix}/{embodiment_name}"] = wandb.Video( + path, fps=self._video_fps(), format="mp4" + ) + experiment.log(payload, step=self._resolved_log_step()) + + def _resolved_log_step(self) -> int: + log_step = getattr(self, "log_step", None) + if log_step is not None: + return int(log_step) + return int(getattr(self.trainer, "global_step", 0)) + + @staticmethod + def _batch_values(value, batch_size: int, label: str) -> list: + if isinstance(value, torch.Tensor): + if value.ndim == 0: + return [value.item()] * batch_size + if int(value.shape[0]) != batch_size: + raise ValueError( + f"{label} has batch dimension {tuple(value.shape)}, expected " + f"{batch_size}" + ) + return [item.item() if item.ndim == 0 else item for item in value] + if isinstance(value, np.ndarray): + if value.ndim == 0: + return [value.item()] * batch_size + if int(value.shape[0]) != batch_size: + raise ValueError( + f"{label} has batch dimension {value.shape}, expected {batch_size}" + ) + return [item.item() if np.ndim(item) == 0 else item for item in value] + if isinstance(value, (list, tuple)): + if len(value) != batch_size: + raise ValueError( + f"{label} has {len(value)} entries, expected {batch_size}" + ) + return list(value) + return [value] * batch_size + + def _native_key(self, value: torch.Tensor, key: str, embodiment_id: int): + if self.normalizer is None: + raise RuntimeError("open_loop_sim evaluator data context was not bound") + return self._clock_seconds( + self.normalizer.unnormalize({key: value}, embodiment_id).get(key, value) + ) + + def _native(self, normalized, embodiment_id): + return self._clock_seconds(super()._native(normalized, embodiment_id)) + + def _clock_seconds(self, value): + # Keep native encoded clocks intact; only the shared codec unlogs them. + return value + + def _is_arc_prediction(self, prediction): + return prediction.ndim == 2 and prediction.shape == bimanual_arc_token_shape( + self.resampled_vector_length, self.clock_mode, self.streams + ) + + def _codec(self): + if self._arc_tokenizer is None: + self._arc_tokenizer = TokenizeBimanualArcLengthCartesian( + min_distance_unit=self.min_distance_unit, + rotation_distance_unit=self.rotation_distance_unit, + gripper_distance_unit=self.gripper_distance_unit, + gripper_stream_mode=self.gripper_stream_mode, + resampled_vector_length=self.resampled_vector_length, + dt=self.control_dt, + clock_mode=self.clock_mode, + streams=self.streams, + ) + return self._arc_tokenizer + + def _initial_actions(self, source_batch, batch_size): + """Read native same-frame metadata, never normalized proprio by inference.""" + value = source_batch.get(self.initial_action_key) + if value is None: + if self.action_mode != "arc": + return None + raise ValueError(f"ARC evaluation requires {self.initial_action_key}") + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + value = np.asarray(value, dtype=np.float64) + if value.shape != (batch_size, 14) or not np.isfinite(value).all(): + raise ValueError("ARC initial actions must be finite (B,14) native poses") + return value + + def _decode_prediction_with_steps( + self, + prediction: np.ndarray, + *, + initial_action=None, + max_steps: int | None = None, + ) -> tuple[np.ndarray, int]: + is_arc = self._is_arc_prediction(prediction) + if self.action_mode == "arc" and not is_arc: + raise ValueError( + "open_loop_sim action_mode='arc' received a non-ARC prediction " + f"with shape {prediction.shape}" + ) + if self.action_mode == "baseline" and is_arc: + raise ValueError( + "open_loop_sim action_mode='baseline' received an ARC prediction" + ) + if not is_arc: + if prediction.ndim != 2 or prediction.shape[1] != 14: + raise ValueError( + "baseline open_loop_sim predictions must have shape (T, 14), " + f"got {prediction.shape}" + ) + if prediction.shape[0] < self.execute_steps: + raise ValueError( + f"baseline prediction has only {prediction.shape[0]} control " + f"steps, needs {self.execute_steps}" + ) + steps = self.execute_steps + if max_steps is not None: + steps = min(steps, int(max_steps)) + return prediction[:steps].copy(), steps + + count = executed_arc_waypoints(len(prediction), self.execute_fraction) + partial = prediction[:count] + codec = self._codec() + steps = codec.control_steps( + partial, initial_action=initial_action, max_steps=max_steps + ) + return codec.detokenize(partial, steps, initial_action=initial_action), steps + + def _decode_prediction( + self, prediction: np.ndarray, *, initial_action=None + ) -> np.ndarray: + return self._decode_prediction_with_steps( + prediction, initial_action=initial_action + )[0] + + def _append_source_records(self, source_id: str, source_batch, prediction): + if not isinstance(prediction, torch.Tensor) or prediction.ndim != 3: + raise ValueError( + f"open_loop_sim expects batched predictions, got {type(prediction)} " + f"with shape {getattr(prediction, 'shape', None)}" + ) + batch_size = int(prediction.shape[0]) + self._metric_device = prediction.device + if "episode_hash" not in source_batch or "frame_index" not in source_batch: + raise KeyError( + "open_loop_sim requires episode_hash and frame_index in every " + f"validation source ({source_id!r})" + ) + embodiment_id, label = self._embodiment(source_batch) + episode_hashes = self._batch_values( + source_batch["episode_hash"], batch_size, "episode_hash" + ) + frame_indices = self._batch_values( + source_batch["frame_index"], batch_size, "frame_index" + ) + pred_native = self._native(prediction, embodiment_id).detach().cpu().numpy() + initial_actions = self._initial_actions(source_batch, batch_size) + target_key = ( + self.ground_truth_action_key + if self.ground_truth_action_key in source_batch + else self.action_key + ) + target_value = self._native_key( + source_batch[target_key], target_key, embodiment_id + ) + target_native = target_value.detach().cpu().numpy() + metric_anchors = None + if getattr(self, "distance_dtw_enabled", False): + if METRIC_FRAME_KEY not in source_batch: + raise ValueError( + f"Distance DTW needs {METRIC_FRAME_KEY} from updated EEF transforms" + ) + metric_anchors = source_batch[METRIC_FRAME_KEY] + if isinstance(metric_anchors, torch.Tensor): + metric_anchors = metric_anchors.detach().cpu().numpy() + metric_anchors = np.asarray(metric_anchors) + if metric_anchors.shape != (batch_size, 2, 4, 4): + raise ValueError("DTW anchors must have shape (B,2,4,4)") + if target_native.ndim != 3 or target_native.shape[0] != batch_size: + raise ValueError( + f"open_loop_sim ground truth {target_key!r} must be batched, got " + f"{target_native.shape}" + ) + for index in range(batch_size): + episode = str(episode_hashes[index]) + frame = int(frame_indices[index]) + if not episode or frame < 0: + raise ValueError( + f"invalid open_loop_sim episode/frame: {episode!r}/{frame}" + ) + prediction_value = pred_native[index] + self._records.append( + { + "group": self._validation_group or DEFAULT_VALID_GROUP, + "source": str(source_id), + "label": label, + "episode": episode, + "frame": frame, + # ARC keeps its complete token and native ground-truth + # window because its timing payload determines a variable + # control-frame stride at episode replay time. + "prediction": np.asarray(prediction_value, dtype=np.float32), + **( + {"arc.current_action": initial_actions[index].copy()} + if initial_actions is not None + else {} + ), + "ground_truth": np.asarray(target_native[index], dtype=np.float32), + **( + {METRIC_FRAME_KEY: metric_anchors[index].copy()} + if metric_anchors is not None + else {} + ), + } + ) + + @torch.inference_mode() + def on_validation_step(self, batch, batch_idx, dataloader_idx=0): + del dataloader_idx + result = self._forward_deterministic(batch) + for source_id, source_batch in batch.items(): + prediction = result[source_id]["pred_action"] + if not getattr(self, "video_only", False): + self._append_source_records(source_id, source_batch, prediction) + if getattr( + self, "_video_enabled", False + ) and self._collective_video_enabled(batch_idx): + embodiment_id, embodiment_name = self._embodiment(source_batch) + self._collect_open_loop_video( + source_id=source_id, + source_batch=source_batch, + prediction=prediction, + embodiment_id=embodiment_id, + embodiment_name=embodiment_name, + ) + return {} + + @staticmethod + def _merge_records(states: list[list[dict[str, Any]]]) -> list[dict[str, Any]]: + unique: dict[tuple[str, str, str, int], dict[str, Any]] = {} + for state in states: + for record in state: + key = ( + record.get("group", DEFAULT_VALID_GROUP), + record["source"], + record["episode"], + int(record["frame"]), + ) + unique.setdefault(key, record) + return list(unique.values()) + + def _all_records(self) -> list[dict[str, Any]]: + local = self._records + if not dist.is_available() or not dist.is_initialized(): + return self._merge_records([local]) + states = [None] * dist.get_world_size() + dist.all_gather_object(states, local) + return self._merge_records(states) + + def _score_episode(self, records: list[dict[str, Any]]) -> dict[str, Any]: + records = sorted(records, key=lambda record: int(record["frame"])) + frames = [int(record["frame"]) for record in records] + if len(set(frames)) != len(frames): + raise RuntimeError("open_loop_sim received duplicate episode frames") + if self.require_episode_start and frames[0] != 0: + raise RuntimeError( + f"open_loop_sim requires complete episodes from frame 0; " + f"episode {records[0]['episode']!r} starts at frame {frames[0]}" + ) + expected = set(range(frames[0], frames[-1] + 1)) + missing = sorted(expected.difference(frames)) + if missing: + raise RuntimeError( + f"open_loop_sim requires every validation frame in episode " + f"{records[0]['episode']!r}; missing {len(missing)} frames, " + f"first missing={missing[0]}" + ) + by_frame = {frame: record for frame, record in zip(frames, records)} + end_frame = frames[-1] + 1 + + sq = { + "mse": 0.0, + "xyz_mse": 0.0, + "ypr_mse": 0.0, + "grip_mse": 0.0, + "paired_mse": 0.0, + } + # Second regime, same segments: the whole token's tick prefix (see the class docstring). + executed = segments = 0 + segment_control_steps = [] + dump = [] if getattr(self, "episode_dump_dir", None) is not None else None + cursor = frames[0] + while cursor < end_frame: + record = by_frame.get(cursor) + if record is None: + raise RuntimeError( + f"open_loop_sim has no observation at frame {cursor}" + ) + remaining = end_frame - cursor + ground_truth = record["ground_truth"] + if ground_truth.ndim != 2 or ground_truth.shape[1] != 14: + raise ValueError( + "open_loop_sim ground truth must be a control-frequency " + f"(T, 14) trajectory, got {ground_truth.shape}" + ) + from_episode = METRIC_FRAME_KEY in record and self._is_arc_prediction( + np.asarray(record["prediction"]) + ) + prediction, n = self._decode_prediction_with_steps( + record["prediction"], + initial_action=record.get("arc.current_action"), + # An ARC chunk spans a fixed distance, not a fixed frame count: + # with metric frames its ground truth comes from the episode + # itself, so only the episode end bounds the executed chunk. + max_steps=remaining + if from_episode + else min(len(ground_truth), remaining), + ) + if from_episode: + later = [by_frame[cursor + i] for i in range(n)] + ground_truth = gt_chunk_in_anchor_frame( + [item["ground_truth"][0] for item in later], + [item[METRIC_FRAME_KEY] for item in later], + record[METRIC_FRAME_KEY], + ) + error = prediction[:n] - ground_truth[:n] + sq["mse"] += float(np.square(error).sum()) + sq["xyz_mse"] += float(np.square(error[:, XYZ_COLS]).sum()) + sq["ypr_mse"] += float(np.square(error[:, YPR_COLS]).sum()) + sq["grip_mse"] += float(np.square(error[:, GRIP_COLS]).sum()) + sq["paired_mse"] += float(np.square(error[:, PAIRED_COLS]).sum()) + if dump is not None: + dump.append((cursor, prediction[:n], ground_truth[:n])) + executed += n + segments += 1 + segment_control_steps.append(n) + cursor += n + if dump: + self.episode_dump_dir.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + self.episode_dump_dir / f"{records[0]['episode']}.npz", + segment_start=np.asarray([c for c, _, _ in dump], dtype=np.int64), + frame=np.concatenate( + [c + np.arange(len(p)) for c, p, _ in dump] + ).astype(np.int64), + prediction=np.concatenate([p for _, p, _ in dump]).astype(np.float32), + ground_truth=np.concatenate([g for _, _, g in dump]).astype(np.float32), + control_dt=np.asarray(self.control_dt, dtype=np.float64), + ) + + episode_length = end_frame - frames[0] + denominators = {key: executed * dims for key, dims in METRIC_DIMS.items()} + return { + "group": records[0].get("group", DEFAULT_VALID_GROUP), + "episode": records[0]["episode"], + "source": records[0]["source"], + "label": records[0]["label"], + "executed_steps": executed, + "segments": segments, + "segment_control_steps": segment_control_steps, + "coverage": executed / max(episode_length, 1), + "metrics": {key: sq[key] / max(denominators[key], 1) for key in sq}, + **( + {"distance_dtw": score_distance_dtw_episode(self, records)} + if getattr(self, "distance_dtw_enabled", False) + else {} + ), + } + + @staticmethod + def _summarize_episodes(episodes: list[dict[str, Any]]) -> dict[str, Any]: + if not episodes: + raise RuntimeError("open_loop_sim received no validation episodes") + + total_steps = sum(item["executed_steps"] for item in episodes) + total_segments = sum(item["segments"] for item in episodes) + micro = {} + dimensions = { + key: METRIC_DIMS[key] + for key in ("mse", "xyz_mse", "ypr_mse", "grip_mse", "paired_mse") + } + for key, dims in dimensions.items(): + micro[key] = sum( + item["metrics"][key] * item["executed_steps"] * dims + for item in episodes + ) / max(total_steps * dims, 1) + labels = defaultdict(list) + for item in episodes: + labels[item["label"]].append(item) + return { + "episodes": len(episodes), + "executed_control_steps": total_steps, + **( + {"distance_dtw": summarize_distance_dtw(episodes)} + if "distance_dtw" in episodes[0] + else {} + ), + "segments": total_segments, + "coverage": float(np.mean([item["coverage"] for item in episodes])), + "micro": micro, + "macro": { + key: float(np.mean([item["metrics"][key] for item in episodes])) + for key in micro + }, + "per_label": { + label: { + "episodes": len(items), + "executed_control_steps": sum( + item["executed_steps"] for item in items + ), + "coverage": float(np.mean([item["coverage"] for item in items])), + "micro": { + key: float( + np.average( + [item["metrics"][key] for item in items], + weights=[item["executed_steps"] for item in items], + ) + ) + for key in micro + }, + } + for label, items in labels.items() + }, + "episode_results": episodes, + } + + def _compute_results(self, records: list[dict[str, Any]]) -> dict[str, Any]: + grouped = defaultdict(list) + for record in records: + grouped[ + ( + record.get("group", DEFAULT_VALID_GROUP), + record["source"], + record["episode"], + ) + ].append(record) + if self.limit_val_episodes is not None: + grouped_by_source = defaultdict(list) + for key, group in grouped.items(): + grouped_by_source[key[:2]].append((key[2], group)) + grouped = defaultdict(list) + for (validation_group, source), episode_groups in grouped_by_source.items(): + for episode, group in sorted(episode_groups)[: self.limit_val_episodes]: + grouped[(validation_group, source, episode)] = group + episodes = [self._score_episode(group) for group in grouped.values()] + if not episodes: + raise RuntimeError("open_loop_sim received no validation episodes") + by_group = defaultdict(list) + for item in episodes: + by_group[item["group"]].append(item) + results = self._summarize_episodes(episodes) + results.update( + { + "execute_fraction": self.execute_fraction, + "distance_dtw_enabled": getattr(self, "distance_dtw_enabled", False), + "distance_dtw_metric_version": ( + METRIC_VERSION + if getattr(self, "distance_dtw_enabled", False) + else None + ), + "execute_control_steps": ( + self.execute_steps if self.action_mode != "arc" else None + ), + "arc_execution_cap_mode": ( + self.arc_execution_cap_mode if self.action_mode == "arc" else None + ), + "arc_episode_progress_semantics": ARC_DISTANCE_SEMANTICS, + "arc_rotation_distance_unit": getattr( + self, "rotation_distance_unit", None + ), + "arc_rotation_clock": ( + "per_arm_independent" + if getattr(self, "rotation_distance_unit", None) is not None + else None + ), + "execute_arc_waypoints": ( + self.execute_arc_waypoints + if self.action_mode == "arc" + and self.arc_execution_cap_mode == "waypoints" + else None + ), + "execute_arc_distance_m": ( + self.execute_fraction * self.min_distance_unit + if self.action_mode == "arc" + and self.arc_execution_cap_mode == "distance" + else None + ), + "arc_distance_semantics": ( + ARC_DISTANCE_SEMANTICS + if self.action_mode == "arc" + and self.arc_execution_cap_mode == "distance" + else None + ), + "replan_stride_mode": ( + f"arc_{self.arc_execution_cap_mode}_timing" + if self.action_mode == "arc" + else "fixed_control_frames" + ), + "control_horizon": self.control_horizon, + "control_dt": self.control_dt, + "clock_mode": self.clock_mode, + "streams": self.streams, + "arc_metric_regime": "waypoint_prefix" + if self.action_mode == "arc" + else None, + "limit_val_episodes": self.limit_val_episodes, + "groups": sorted(by_group), + "per_group": { + group: self._summarize_episodes(items) + for group, items in sorted(by_group.items()) + }, + # Same summary restricted to one embodiment (episode label), + # so cotraining reports robot and human separately. + "per_group_embodiment": { + group: { + label: { + k: v + for k, v in self._summarize_episodes( + [i for i in items if i["label"] == label] + ).items() + if k != "episode_results" + } + for label in sorted({i["label"] for i in items}) + } + for group, items in sorted(by_group.items()) + }, + } + ) + return results + + def _metric_tensors(self, results: dict[str, Any]) -> dict[str, torch.Tensor]: + metrics = {} + summaries = [] + for group, summary in results["per_group"].items(): + section = "Valid" if group == DEFAULT_VALID_GROUP else f"Valid_{group}" + per_embodiment = results.get("per_group_embodiment", {}).get(group, {}) + # W&B groups panels by the first path segment, so each embodiment + # gets its own section (Valid_yam_bimanual/...), scored on that + # embodiment's episodes only. With several embodiments the pooled + # summary moves to Valid_all_embodiments so no chart mixes them + # under an unlabeled name. + if len(per_embodiment) > 1: + summaries.append((f"{section}_all_embodiments/open_loop_sim", summary)) + summaries.extend( + (f"{section}_{label}/open_loop_sim", item) + for label, item in per_embodiment.items() + ) + else: + summaries.append((f"{section}/open_loop_sim", summary)) + for prefix, summary in summaries: + if "distance_dtw" in summary: + dtw = summary["distance_dtw"] + for name, key in { + "XYZ_MSE": "xyz_mse", + "Episode_XYZ_MSE": "episode_xyz_mse", + "GT_Frames": "gt_frames", + "Predicted_Samples": "predicted_samples", + "Segments": "segments", + "GT_Coverage": "gt_coverage", + "Prediction_Coverage": "prediction_coverage", + "Duration_Ratio": "duration_ratio", + "Clipped_Segments": "clipped_segments", + }.items(): + metrics[f"{prefix}/Distance_DTW/{name}"] = dtw[key] + metrics.update( + { + f"{prefix}/MSE": summary["micro"]["mse"], + f"{prefix}/Episode_MSE": summary["macro"]["mse"], + f"{prefix}/XYZ_MSE": summary["micro"]["xyz_mse"], + f"{prefix}/YPR_MSE": summary["micro"]["ypr_mse"], + f"{prefix}/Grip_MSE": summary["micro"]["grip_mse"], + f"{prefix}/Paired_MSE": summary["micro"]["paired_mse"], + } + ) + metrics.update( + { + f"{prefix}/Coverage": summary["coverage"], + f"{prefix}/Episodes": summary["episodes"], + f"{prefix}/Executed_Control_Steps": summary[ + "executed_control_steps" + ], + f"{prefix}/Segments": summary["segments"], + } + ) + return { + key: torch.tensor(float(value), dtype=torch.float32).to(self._metric_device) + for key, value in metrics.items() + } + + def on_validation_end(self): + if getattr(self, "video_only", False): + self.last_results = None + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return None + if self.trainer is not None and not getattr( + self.trainer, "is_global_zero", True + ): + return None + self._flush_all_executed_chunks() + EvalVideo.on_validation_end(self) + return None + + records = self._all_records() + results = self._compute_results(records) + self.last_results = results + # Every rank participates in episode/frame gathering before rank-zero + # logging. Returning early here would strand rank zero in a collective. + # Held executed-chunk overlays are flushed first so the last chunk of + # every episode reaches the gathered buffers. + if getattr(self, "_video_enabled", False): + self._flush_all_executed_chunks() + EvalVideo.on_validation_end(self) + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return results + if getattr(self, "records_dump_path", None) is not None: + self.records_dump_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(records, self.records_dump_path) + if self.trainer is not None and not getattr( + self.trainer, "is_global_zero", True + ): + return results + # This hook is called from LightningModule.on_validation_end(). Calling + # LightningModule.log_dict() here recursively enters Lightning's + # validation-hook guard and raises ``MisconfigurationException``. + # ``_all_records`` has already merged every rank, so direct logger + # logging on global zero is both sufficient and deterministic. Keep + # values scalar so this works for W&B, TensorBoard, and LoggerCollection + # without asking Lightning to infer validation-step semantics. + if self.trainer is not None: + logger = getattr(self.trainer, "logger", None) + if logger is not None: + metrics = { + key: float(value.detach().cpu().item()) + for key, value in self._metric_tensors(results).items() + } + logger.log_metrics(metrics, step=self._resolved_log_step()) + if self.results_path is not None: + self.results_path.parent.mkdir(parents=True, exist_ok=True) + self.results_path.write_text(json.dumps(results, indent=2) + "\n") + return results