Stream camera video frames instead of buffering episodes - #1000
Stream camera video frames instead of buffering episodes#1000alexmillane wants to merge 1 commit into
Conversation
alexmillane
left a comment
alexmillane
left a comment
There was a problem hiding this comment.
Self review 1.
| def test_frames_are_streamed_not_buffered(tmp_path): | ||
| """Every frame reaches the encoder as it arrives, and no frame list is retained.""" | ||
| env = _make_env() | ||
| with _patched_writers() as writers: | ||
| recorder = CameraObsVideoRecorder(env, video_folder=str(tmp_path)) | ||
|
|
||
| for _ in range(3): | ||
| _configure_step(env) | ||
| recorder.step(None) | ||
|
|
||
| # One open encoder per (env, camera), each already handed all three frames. | ||
| assert len(writers) == len(CAMERAS) * 2 | ||
| assert all(writer.frames_written == 3 for writer in writers) | ||
| assert not hasattr(recorder, "buffers") |
There was a problem hiding this comment.
Superfluous. Remove.
| No Isaac Sim or GPU required. The moviepy encoder is replaced by a stand-in that records | ||
| the frames it is handed, so tests run fast and on CPU-only machines. |
There was a problem hiding this comment.
The stand up doesn't actually record, it just tracks how many frames have been passed to be recorded.
| assert recorder.buffers[cam][0] == [] | ||
|
|
||
| # env 1 accumulated 2 frames (neither step was terminal for it) | ||
| by_path = {writer.filename: writer for writer in writers} |
There was a problem hiding this comment.
writer_by_filename.
| Memory note: frames are handed to ffmpeg one at a time and never accumulated, so host RAM | ||
| is independent of episode length. What stays resident is one frame per open stream plus | ||
| each encoder's internal state; ``num_envs × num_cameras`` encoders run concurrently, each | ||
| pinned to a single thread to bound that state. Buffering whole episodes instead costs | ||
| ``num_envs × L × H × W × C`` bytes — for 10 envs of 1000-step episodes with three | ||
| 1280×720 cameras that is ~83 GB, against mp4s of a few MB each. |
There was a problem hiding this comment.
Change this to
Memory note: This class uses incremental ffmpeg encoding the avoid storing the raw frames in memory, which uses substantial amounts of RAM.
|
|
||
|
|
||
| @dataclass | ||
| class _EpisodeVideoWriter: |
There was a problem hiding this comment.
Remove leading underscore.
| # The env's counter still names the episode now in progress; it is advanced on reset, | ||
| # inside env.step. Sharing the env's index keeps the filename's episode number in | ||
| # lockstep with the per-episode results record's ``episode_in_env``. |
| # One thread per encoder: frames arrive far slower than a single x264 thread encodes, | ||
| # and num_envs x num_cameras encoders run at once, so this bounds their combined state. |
There was a problem hiding this comment.
Change to: "We use one thread because frames arrive slower than a single thread is able to encode.
CameraObsVideoRecorder held every frame of an episode in memory before encoding, costing num_envs x episode_length x H x W x C bytes. Frames are now written to a per-(env, camera) ffmpeg encoder as they arrive. Measured on canned_food_in_bin with 3x 1280x720 cameras: peak host RAM at 6 envs drops 60.5 GB -> 13.1 GB, with wall clock unchanged. Signed-off-by: alex <amillane@nvidia.com>
78173ad to
443177f
Compare
Greptile SummaryThe PR replaces per-episode raw camera-frame buffering with incremental encoding through one ffmpeg writer per environment and camera, substantially reducing peak host memory.
Confidence Score: 5/5The PR appears safe to merge, with no concrete blocking or independently actionable non-blocking issue identified. Episode counters are sampled when each current episode’s writer opens, completed streams are finalized on reset, partial streams are removed on shutdown, and the updated tests cover the principal lifecycle transitions. Important Files Changed
Sequence DiagramsequenceDiagram
participant Env
participant Recorder as CameraObsVideoRecorder
participant Writer as FFMPEG_VideoWriter
Env->>Recorder: step() observations
alt Environment is active
Recorder->>Writer: lazily open episode writer
Recorder->>Writer: write_frame(frame)
else Environment reset
Recorder->>Writer: close and finalize completed episode
end
alt Wrapper closes during partial episode
Recorder->>Writer: close
Recorder->>Recorder: delete partial mp4
end
Reviews (1): Last reviewed commit: "Stream camera frames to the encoder inst..." | Re-trigger Greptile |
![arena-review-bot[bot]](https://avatars.githubusercontent.com/u/195794629?s=60&v=4){kind=small}
| height, width, _ = frame.shape | ||
| # We use one thread because frames arrive slower than a single thread is able to encode. | ||
| episode_writer = EpisodeVideoWriter( | ||
| writer=FFMPEG_VideoWriter(path, size=(width, height), fps=self.fps, threads=1), |
There was a problem hiding this comment.
🔵 Concurrent encoders now scale with num_envs
Nice RAM win. One thing worth sanity-checking: we now keep a live ffmpeg process open per (env, camera) for the whole episode, so peak concurrent encoders is num_envs × num_cameras — whereas before we ran at most one encode at a time, at reset. At the 6-env measurement that's 18 processes; a larger sweep (e.g. 64 envs × 3 cams) is ~192 persistent subprocesses/pipes. Since the threads=1 note says frames arrive slowly, CPU is presumably fine — have you confirmed the process/FD count holds up at the higher env counts you'll actually run?
🤖 Isaac Lab-Arena Review BotSummaryThis PR switches Findings🔵 Improvement — Test CoverageGood. Tests are CPU-only (no Isaac Sim), so the inner/outer sim pattern does not apply; the moviepy encoder is replaced by a counting stand-in and there is a VerdictShip it |
2 participants
Summary
Encode camera videos incrementally so host RAM no longer scales with episode length.
Detailed description
CameraObsVideoRecorderstored every frame of an episode before encoding, costingnum_envs × episode_length × H × W × C. For robolab tasks (3× 1280x720 droid cameras, 1000-step episodes, 50Hz) that is 8.3 GB per env per episode.FFMPEG_VideoWriteras they arrive.canned_food_in_bin, peak host RAM drops 8.32 GB → 0.72 GB. Wall clock unchanged.Before and after experiment.