Add interactive drive recording options#301
Conversation
Greptile SummaryAdds rollout recording support to the OmniDreams interactive-drive demo. A new
Confidence Score: 5/5Safe to merge; all recording I/O is non-fatal and teardown is handled in every exit path. The feature is well-contained: recording failures in both start() and stop() are caught and logged rather than propagated, the frame buffer is bounded, and the recorder is cleanly closed on reset, OOB respawn, loop-end, and exceptions via app.py's finally block. The only finding is a minor session-count numbering inconsistency when a start() attempt fails, which has no runtime impact on correctness. recording.py — session counter increment ordering; otherwise no files require special attention. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[app.py: _build_recorder] --> B[run_main_loop called per rollout]
B --> C{recorder.auto_start?}
C -->|yes| D[recorder.start]
C -->|no| E[wait for hotkey]
D --> F[per-frame: record_frame]
E --> F
F --> G{event?}
G -->|toggle hotkey| H[recorder.toggle]
H --> F
G -->|reset requested| I[close reason=reset]
I --> B
G -->|OOB respawn| J[close reason=respawn]
J --> B
G -->|presenter closes| K[close reason=loop-end]
K --> L[app.py finally: close reason=scene-end no-op]
Reviews (5): Last reviewed commit: "Save raster recording first frame" | Re-trigger Greptile |
| def stop(self, *, reason: str) -> Path | None: | ||
| session = self._active_session | ||
| if session is None: | ||
| return None | ||
| self._active_session = None | ||
| metadata = { | ||
| "scene_id": self._scene.scene_id, | ||
| "scene_path": str(self._scene.scene_path), | ||
| "fps": self._fps, | ||
| "start_time_utc": session.start_time_utc, | ||
| "stop_time_utc": datetime.now(timezone.utc).isoformat(), | ||
| "reason": reason, | ||
| "hdmap_frames": len(session.hdmap_frames), | ||
| "inferred_frames": len(session.inferred_frames), | ||
| } | ||
| (session.output_dir / "metadata.json").write_text( | ||
| json.dumps(metadata, indent=2, sort_keys=True) + "\n", | ||
| encoding="utf-8", | ||
| ) | ||
| _write_video(session.hdmap_frames, session.output_dir / "hdmap.mp4", self._fps) | ||
| _write_video( | ||
| session.inferred_frames, | ||
| session.output_dir / "inferred.mp4", | ||
| self._fps, | ||
| ) | ||
| self._closed_paths.append(session.output_dir) | ||
| print( | ||
| "[recording] saved " | ||
| f"dir={session.output_dir} " | ||
| f"hdmap_frames={len(session.hdmap_frames)} " | ||
| f"inferred_frames={len(session.inferred_frames)}", | ||
| flush=True, | ||
| ) | ||
| return session.output_dir |
There was a problem hiding this comment.
Unhandled exception from video write crashes the session
_active_session is cleared on line 131 before the video files are written. If _write_video raises (disk full, codec error, mediapy bug), the exception propagates through close() → _close_recording() → run_main_loop → the app.py finally block, terminating the entire interactive-drive session. The partial bundle directory (with metadata.json and prompt.txt) is also never added to closed_paths, so callers have no record of the attempt.
Wrapping the file-writing section in a try/except that logs and skips addition to closed_paths — or at least catching and logging the exception so the app continues — would make recording failures non-fatal.
| candidate = base | ||
| suffix = 1 | ||
| while candidate.exists(): | ||
| suffix += 1 | ||
| candidate = root / f"{base.name}-{suffix:02d}" |
There was a problem hiding this comment.
The collision-avoidance loop initialises
suffix = 1 but increments it to 2 before first use, so the first disambiguated directory produced is ...-02 while ...-01 is never generated. Starting suffix at 0 closes the gap.
| candidate = base | |
| suffix = 1 | |
| while candidate.exists(): | |
| suffix += 1 | |
| candidate = root / f"{base.name}-{suffix:02d}" | |
| candidate = base | |
| suffix = 0 | |
| while candidate.exists(): | |
| suffix += 1 | |
| candidate = root / f"{base.name}-{suffix:02d}" |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
liruilong940607
left a comment
There was a problem hiding this comment.
I like this functionality!
We should consider expose these functionalities in the doc instead of burying them here in this config. Now that interactive demo has many controls now -- for example switching between RGB and HDMap, the controls that supported on the job sticker etc -- we might consider open another doc page [could be another level under omnidreams] to detail these user-facing instructions.
| candidate = base | ||
| suffix = 1 | ||
| while candidate.exists(): | ||
| suffix += 1 | ||
| candidate = root / f"{base.name}-{suffix:02d}" |
84a3838 to
75bf982
Compare
| def start(self) -> None: | ||
| if self.is_recording: | ||
| return | ||
| self._session_count += 1 | ||
| start_time = datetime.now(timezone.utc) | ||
| output_dir = self._next_output_dir(start_time) | ||
| output_dir.mkdir(parents=True, exist_ok=False) | ||
| (output_dir / "prompt.txt").write_text(self._scene.prompt, encoding="utf-8") | ||
| self._active_session = _RecordingSession( | ||
| output_dir=output_dir, | ||
| start_time_utc=start_time.isoformat(), | ||
| hdmap_frames=[], | ||
| inferred_frames=[], | ||
| ) | ||
| print(f"[recording] started dir={output_dir}", flush=True) |
There was a problem hiding this comment.
start() I/O failures are not caught and crash the session
output_dir.mkdir() and prompt.txt write are not wrapped in a try/except. A permissions error, disk-full condition, or race on the output directory will raise and propagate through run_main_loop into the app.py try block, terminating the entire interactive-drive session — the same class of problem that was addressed for stop(). When triggered by a user hotkey mid-session or by auto_start at rollout begin, this would end the session with a Python traceback and no recording written, even though the world model was running fine. Wrapping the I/O in a try/except with a logger warning (matching the stop() pattern) would make start failures non-fatal.
| def record_frame(self, frame: PresentedFrame) -> None: | ||
| session = self._active_session | ||
| if session is None: | ||
| return | ||
| session.hdmap_frames.append(_as_rgb_host_uint8(frame.rgb_host_uint8)) | ||
| if frame.model_rgb_host_uint8 is not None: | ||
| inferred_frame = _as_rgb_host_uint8(frame.model_rgb_host_uint8) | ||
| if not session.first_frame_written: | ||
| Image.fromarray(inferred_frame).save( | ||
| session.output_dir / "first_frame.png" | ||
| ) | ||
| session.first_frame_written = True | ||
| session.inferred_frames.append(inferred_frame) |
There was a problem hiding this comment.
Unbounded in-memory frame accumulation
hdmap_frames and inferred_frames grow without bound for the entire lifetime of a recording session. With recording_auto_start: true in the perf manifest, recording begins at rollout start and runs until reset, OOB respawn, or loop end. At a typical interactive-drive resolution (e.g. 640×360) and 10 fps, a 10-minute continuous rollout accumulates roughly 4 GB across both streams before stop() finally calls np.stack and writes the videos. On a system that was already GPU-memory-heavy this will OOM the host. There is no frame cap, sliding window, or background-write path. Consider either capping the buffer (dropping the oldest frames) or streaming frames directly to the video file as they arrive.
Thanks @liruilong940607 . Missed this feedback the other day. |
Add Interactive Drive Recording Options
Summary
Adds rollout recording support to the OmniDreams interactive-drive demo.
This PR lets world-model manifests enable recording, optionally auto-start recording when a rollout begins, and save the captured bundle under a repository-root-relative output directory by default.
Changes
RecordingConfigandInteractiveDriveRecorderfor writing recording bundles.recording_enabledrecording_dirrecording_hotkeyrecording_auto_startrecording_dirvalues from the flashdreams repository root instead of the manifest/config directory.recording_dirvalues as-is.AppConfig,InteractiveDriveApp,KeyboardState, andrun_main_loop.recording_auto_start: true.first_frame.pngprompt.txtmetadata.jsonhdmap.mp4inferred.mp4first_frame.pngfrom the first recorded inferred frame so it matchesinferred.mp4, instead of using the input RGB frame.recording_hotkey: F9andrecording_auto_start: true.Notes
This PR supports configured single-key recording hotkeys such as
F9orr. Modifier combinations such asAlt+Rstill need separate modifier-state plumbing in the browser and SlangPy event paths.Tests
Local verification:
Result: