Add interactive drive recording options

integrations/omnidreams/omnidreams/interactive_drive/README.md

@@ -467,6 +467,28 @@ Controls (apply in all three modes):
 - `X` exit scene (return to the scene selector; HUD mode only)
 - `Esc` quit
 
+Manifests can enable rollout recording with either backend:
+
+```yaml
+recording_enabled: true
+recording_dir: recordings
+recording_hotkey: F9
+recording_auto_start: false
+```
+
+Press the configured hotkey once to start and again to stop. Set
+`recording_auto_start: true` to start recording as soon as each rollout begins.
+Relative `recording_dir` values are written under the flashdreams repository
+root; absolute paths are used as-is. Each saved recording writes
+`first_frame.png`, `prompt.txt`, `metadata.json`, `hdmap.mp4`, and
+`inferred.mp4` under a timestamped subdirectory. `first_frame.png` is copied
+from the first retained inferred frame so it matches the inferred video. To
+bound host memory during long auto-start rollouts, each stream keeps the most
+recent 600 frames; `metadata.json` records any dropped frame counts.
+Raster-only recordings contain the HD-map/raster stream, save `first_frame.png`
+from that stream, and omit `inferred.mp4` because there is no world-model
+output.
+
 The browser control hint is static today, so it does not confirm every keydown
 visually. If the world-model backend is still producing a chunk, input can be
 accepted before the visual response arrives.

integrations/omnidreams/omnidreams/interactive_drive/app.py

@@ -16,6 +16,7 @@
     KeyboardState,
 )
 from omnidreams.interactive_drive.presenter import SlangPyPresenter
+from omnidreams.interactive_drive.recording import InteractiveDriveRecorder
 from omnidreams.interactive_drive.runtime.loop import (
     LoopConfig,
     PresenterBackend,
@@ -82,7 +83,10 @@ def __init__(
         """
         self._config = config
         self._backend = backend
-        self._keyboard = KeyboardState()
+        self._keyboard = KeyboardState(
+            recording_enabled=config.recording.enabled,
+            recording_hotkey=config.recording.hotkey,
+        )
         if config.backend == "omnidreams":
             self._keyboard.set_view_mode("model_rgb")
         if presenter is None:
@@ -438,60 +442,77 @@ def run_scene(self) -> None:
         # world model..."); subsequent rollouts come from a manual reset or
         # OOB respawn, so switch the indicator to "Resetting..." for those.
         loading_status = self._loading_status_message
-        while not self._presenter.should_close:
-            simulation = EgoVehicleKinematics(
-                initial_state=state_from_initial_pose(
-                    initial_rig_to_world=self._scene.initial_rig_to_world,
-                    initial_yaw_rad=self._scene.initial_yaw_rad,
-                    # Start each rollout at a fixed 10 m/s so the ego is
-                    # already rolling on initial load (and after a manual
-                    # reset / OOB respawn), instead of launching at the
-                    # clip's full recorded speed.
-                    initial_speed_mps=10.0,
-                ),
-                vehicle_config=self._config.vehicle,
-                ground_snapper=self._ground_snapper,
-                initial_timestamp_us=self._scene.initial_timestamp_us,
-                map_bounds=self._map_bounds,
-                oob_margin_m=self._config.oob_margin_m,
-                oob_warning_zone_m=self._config.oob_warning_zone_m,
-            )
-            # Publish the freshly-built initial state up front so read-side
-            # speed readouts (the HUD speed digit, the browser ``/state``
-            # endpoint) reflect a reset / respawn immediately. Without this
-            # the last telemetry from the previous rollout would linger on
-            # screen through the "Resetting..." window until the new rollout
-            # requested its first chunk -- the "reset doesn't reset the
-            # displayed speed" symptom.
-            self._keyboard.update_telemetry(simulation.current_state)
-            input_backend = KeyboardInputBackend(self._keyboard)
-            reset_requested = run_main_loop(
-                presenter=self._presenter,
-                runtime_controls=self._keyboard,
-                initial_presented_frame=loading_frame,
-                input_backend=input_backend,
-                simulation=simulation,
-                pipeline=self._pipeline,
-                config=LoopConfig(
-                    initial_chunk_size=self._config.chunk.initial_chunk_frames,
-                    chunk_size=self._config.chunk.chunk_frames,
-                    frame_interval_s=self._config.chunk.frame_interval_s,
-                    oob_warn_proximity=self._config.oob_warn_proximity,
-                    oob_respawn_proximity=self._config.oob_respawn_proximity,
-                    oob_respawn_debounce_chunks=(
-                        self._config.oob_respawn_debounce_chunks
+        recorder = self._build_recorder()
+        try:
+            while not self._presenter.should_close:
+                simulation = EgoVehicleKinematics(
+                    initial_state=state_from_initial_pose(
+                        initial_rig_to_world=self._scene.initial_rig_to_world,
+                        initial_yaw_rad=self._scene.initial_yaw_rad,
+                        # Start each rollout at a fixed 10 m/s so the ego is
+                        # already rolling on initial load (and after a manual
+                        # reset / OOB respawn), instead of launching at the
+                        # clip's full recorded speed.
+                        initial_speed_mps=10.0,
                     ),
-                ),
-                loading_status=loading_status,
-            )
-            if not reset_requested:
-                break
-            self._pipeline.reset()
-            loading_status = self._resetting_status_message
-            # Paint the reset indicator at once, before the next rollout's
-            # setup, so a reset shows on screen the instant it's requested
-            # rather than after the rebuild completes.
-            self._present_loading_once(loading_status)
+                    vehicle_config=self._config.vehicle,
+                    ground_snapper=self._ground_snapper,
+                    initial_timestamp_us=self._scene.initial_timestamp_us,
+                    map_bounds=self._map_bounds,
+                    oob_margin_m=self._config.oob_margin_m,
+                    oob_warning_zone_m=self._config.oob_warning_zone_m,
+                )
+                # Publish the freshly-built initial state up front so read-side
+                # speed readouts (the HUD speed digit, the browser ``/state``
+                # endpoint) reflect a reset / respawn immediately. Without this
+                # the last telemetry from the previous rollout would linger on
+                # screen through the "Resetting..." window until the new rollout
+                # requested its first chunk -- the "reset doesn't reset the
+                # displayed speed" symptom.
+                self._keyboard.update_telemetry(simulation.current_state)
+                input_backend = KeyboardInputBackend(self._keyboard)
+                reset_requested = run_main_loop(
+                    presenter=self._presenter,
+                    runtime_controls=self._keyboard,
+                    initial_presented_frame=loading_frame,
+                    input_backend=input_backend,
+                    simulation=simulation,
+                    pipeline=self._pipeline,
+                    config=LoopConfig(
+                        initial_chunk_size=self._config.chunk.initial_chunk_frames,
+                        chunk_size=self._config.chunk.chunk_frames,
+                        frame_interval_s=self._config.chunk.frame_interval_s,
+                        oob_warn_proximity=self._config.oob_warn_proximity,
+                        oob_respawn_proximity=self._config.oob_respawn_proximity,
+                        oob_respawn_debounce_chunks=(
+                            self._config.oob_respawn_debounce_chunks
+                        ),
+                    ),
+                    loading_status=loading_status,
+                    recorder=recorder,
+                )
+                if not reset_requested:
+                    break
+                self._pipeline.reset()
+                loading_status = self._resetting_status_message
+                # Paint the reset indicator at once, before the next rollout's
+                # setup, so a reset shows on screen the instant it's requested
+                # rather than after the rebuild completes.
+                self._present_loading_once(loading_status)
+        finally:
+            if recorder is not None:
+                recorder.close(reason="scene-end")
+
+    def _build_recorder(self) -> InteractiveDriveRecorder | None:
+        if not self._config.recording.enabled:
+            return None
+        if self._scene is None:
+            return None
+        return InteractiveDriveRecorder(
+            self._config.recording,
+            scene=self._scene,
+            fps=self._config.chunk.fps,
+        )
 
     def _present_loading_once(self, loading_status: Callable[[], str]) -> None:
         """Render a single loading-overlay frame immediately (used on reset)."""

integrations/omnidreams/omnidreams/interactive_drive/cli.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import argparse
+import warnings
 from dataclasses import replace
 from pathlib import Path
 
@@ -21,8 +22,13 @@
     WorldModelProfileConfig,
 )
 from omnidreams.interactive_drive.log import configure_logging
+from omnidreams.interactive_drive.recording import RecordingConfig
 from omnidreams.interactive_drive.synthetic_scene import build_synthetic_scene_to_temp
-from omnidreams.interactive_drive.world_model.manifest import load_world_model_manifest
+from omnidreams.interactive_drive.world_model.manifest import (
+    RecordingManifest,
+    load_recording_manifest,
+    load_world_model_manifest,
+)
 from omnidreams.scenes import local_scene_archive_path
 
 # Package root (from this file's location) so packaged-asset defaults below
@@ -31,6 +37,26 @@
 # ``$FLASHDREAMS_CACHE_DIR/omnidreams-scenes/`` (shared with the webrtc server).
 _PACKAGE_ROOT = Path(__file__).resolve().parent
 _CONFIGS_ROOT = _PACKAGE_ROOT / "configs"
+_DEFAULT_RECORDING_DIR_NAME = "recordings"
+
+
+def _find_flashdreams_root(start: Path) -> Path:
+    for candidate in (start, *start.parents):
+        if (candidate / "pyproject.toml").is_file() and (
+            candidate / "integrations" / "omnidreams"
+        ).is_dir():
+            return candidate
+    fallback = Path.cwd().resolve()
+    warnings.warn(
+        "Could not locate the flashdreams repository root from "
+        f"{start}; resolving relative recording_dir values from {fallback}.",
+        RuntimeWarning,
+        stacklevel=2,
+    )
+    return fallback
+
+
+_FLASHDREAMS_ROOT = _find_flashdreams_root(_PACKAGE_ROOT)
 
 # Default scene UUID staged by ``omnidreams-prepare`` (clear-weather base
 # archive in nvidia/omni-dreams-scenes).
@@ -69,6 +95,21 @@ def resolve_manifest_path(path: str | Path) -> Path:
     return cwd_path
 
 
+def _resolve_recording_output_dir(
+    raw_dir: Path | None, *, enabled: bool
+) -> Path | None:
+    if not enabled:
+        return None
+    path = (
+        Path(_DEFAULT_RECORDING_DIR_NAME)
+        if raw_dir is None
+        else Path(raw_dir).expanduser()
+    )
+    if path.is_absolute():
+        return path
+    return (_FLASHDREAMS_ROOT / path).resolve()
+
+
 def build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         description="Single-process flashdreams driving demo"
@@ -148,8 +189,10 @@ def build_parser() -> argparse.ArgumentParser:
         type=Path,
         default=None,
         help=(
-            "Omnidreams pipeline manifest (YAML). Accepts a path or a bundled "
-            "config filename such as example_world_model_perf.yaml."
+            "Interactive-drive manifest (YAML). For the omnidreams backend this "
+            "also configures the world-model pipeline; for raster, recording_* "
+            "fields are honored when present. Accepts a path or a bundled config "
+            "filename such as example_world_model_perf.yaml."
         ),
     )
     parser.add_argument(
@@ -369,6 +412,18 @@ def _oob_kwargs(args: argparse.Namespace) -> dict[str, float | int]:
     return overrides
 
 
+def _recording_config_from_manifest(manifest: RecordingManifest) -> RecordingConfig:
+    return RecordingConfig(
+        enabled=manifest.enabled,
+        output_dir=_resolve_recording_output_dir(
+            manifest.dir,
+            enabled=manifest.enabled,
+        ),
+        hotkey=manifest.hotkey,
+        auto_start=manifest.auto_start,
+    )
+
+
 def main() -> None:
     """Stand-alone entry point for ``python -m omnidreams.interactive_drive.cli``.
 
@@ -448,6 +503,13 @@ def prepare_config_and_backend(
 
     backend: RenderBackend
     if config.backend == "raster":
+        if config.manifest_path is not None:
+            config = replace(
+                config,
+                recording=_recording_config_from_manifest(
+                    load_recording_manifest(config.manifest_path)
+                ),
+            )
         backend = RasterRenderBackend(
             chunk=config.chunk, raster=config.raster, bev=config.bev
         )
@@ -468,6 +530,17 @@ def prepare_config_and_backend(
                     height=manifest.resolution_wh[1],
                 ),
             )
+        config = replace(
+            config,
+            recording=_recording_config_from_manifest(
+                RecordingManifest(
+                    enabled=manifest.recording_enabled,
+                    dir=manifest.recording_dir,
+                    hotkey=manifest.recording_hotkey,
+                    auto_start=manifest.recording_auto_start,
+                )
+            ),
+        )
         backend = WorldModelRenderBackend(
             manifest=manifest,
             chunk=config.chunk,

integrations/omnidreams/omnidreams/interactive_drive/config.py

@@ -7,6 +7,8 @@
 from pathlib import Path
 from typing import Literal
 
+from omnidreams.interactive_drive.recording import RecordingConfig
+
 BackendName = Literal["raster", "omnidreams"]
 ViewMode = Literal["rgb", "model_rgb"]
 ComputeDeviceName = Literal["automatic", "cuda", "vulkan"]
@@ -110,6 +112,7 @@ class AppConfig:
     vehicle: VehicleConfig = VehicleConfig()
     world_model_profile: WorldModelProfileConfig = WorldModelProfileConfig()
     world_model_offload_text_encoder: bool = False
+    recording: RecordingConfig = RecordingConfig()
     bev: BevConfig = BevConfig()
     # OOB thresholds plumbed to LoopConfig (overridable via CLI --oob-*).
     # Match alpasim's driver-side proximity: warn > 0.6, respawn >= 2.0

integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_perf.yaml

@@ -58,3 +58,9 @@ native_dit_attention_backend: cudnn # auto | cudnn | sparge | sage3 | sage3_fp8
 # or uncomment and set an absolute or manifest-relative path below.
 native_vae_encoder: disabled             # disabled | fp8
 # native_vae_fp8_state_path: /path/to/lightvae_fp8_state.pt
+
+
+recording_enabled: true
+recording_dir: recordings
+recording_hotkey: F9
+recording_auto_start: true

integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py

@@ -5,6 +5,7 @@
 import time
 
 from omnidreams.interactive_drive.input.backend import InputBackend, SampledInput
+from omnidreams.interactive_drive.recording import normalize_recording_hotkey
 from omnidreams.interactive_drive.types import (
     ControlSnapshot,
     DriverCommand,
@@ -24,12 +25,17 @@ class KeyboardState:
     referencing the per-scene simulation object.
     """
 
-    def __init__(self) -> None:
+    def __init__(
+        self, *, recording_enabled: bool = False, recording_hotkey: str = "f9"
+    ) -> None:
         self._lock = threading.Lock()
         self._pressed: set[str] = set()
         self._view_mode = "rgb"
         self._drive_command: DriverCommand | None = None
         self._reset_pending = False
+        self._recording_enabled = bool(recording_enabled)
+        self._recording_hotkey = normalize_recording_hotkey(recording_hotkey)
+        self._recording_toggle_pending = False
         # Rising-edge "exit the current scene and return to the scene
         # selector" request, set by a wheel/controller's bound exit button
         # (the HUD's ``x`` key calls the presenter directly). The presenter
@@ -57,6 +63,27 @@ def request_reset(self) -> None:
         with self._lock:
             self._reset_pending = True
 
+    def request_recording_toggle(self) -> None:
+        with self._lock:
+            if self._recording_enabled:
+                self._recording_toggle_pending = True
+
+    def consume_recording_toggle_request(self) -> bool:
+        with self._lock:
+            pending = self._recording_toggle_pending
+            self._recording_toggle_pending = False
+            return pending
+
+    @property
+    def recording_enabled(self) -> bool:
+        with self._lock:
+            return self._recording_enabled
+
+    @property
+    def recording_hotkey(self) -> str:
+        with self._lock:
+            return self._recording_hotkey
+
     def request_exit_scene(self) -> None:
         """Request a return to the scene selector from a bound device button."""
         with self._lock: