Stream camera video frames instead of buffering episodes

isaaclab_arena/tests/test_camera_observation_video_recorder.py

@@ -5,10 +5,11 @@
 
 """Unit tests for CameraObsVideoRecorder.
 
-No Isaac Sim or GPU required. The moviepy encoder is mocked so tests run
-fast and on CPU-only machines.
+No Isaac Sim or GPU required. The moviepy encoder is replaced by a stand-in that encodes
+nothing and only counts the frames it is handed, so tests run fast and on CPU-only machines.
 """
 
+import contextlib
 import gymnasium as gym
 import os
 import shutil
@@ -72,6 +73,46 @@ def _configure_step(env: _StubEnv, done_envs: list[int] | None = None, n_envs: i
     env._step_return = (obs, None, terminated, truncated, None)
 
 
+class _FakeVideoWriter:
+    """Stand-in for moviepy's FFMPEG_VideoWriter that counts frames and touches its file.
+
+    Creating the file mirrors ffmpeg, so tests can assert that a partial episode's file is
+    removed rather than left behind.
+    """
+
+    def __init__(self, filename, size, fps, **kwargs):
+        self.filename = filename
+        self.size = size
+        self.fps = fps
+        self.frames_written = 0
+        self.closed = False
+        with open(filename, "wb"):
+            pass
+
+    def write_frame(self, frame):
+        self.frames_written += 1
+
+    def close(self):
+        self.closed = True
+
+
+@contextlib.contextmanager
+def _patched_writers():
+    """Replace the encoder with ``_FakeVideoWriter`` and yield the instances created."""
+    instances: list[_FakeVideoWriter] = []
+
+    def factory(*args, **kwargs):
+        writer = _FakeVideoWriter(*args, **kwargs)
+        instances.append(writer)
+        return writer
+
+    with patch(
+        "isaaclab_arena.video.camera_observation_video_recorder.FFMPEG_VideoWriter",
+        side_effect=factory,
+    ):
+        yield instances
+
+
 # ---------------------------------------------------------------------------
 # Tests
 # ---------------------------------------------------------------------------
@@ -80,25 +121,40 @@ def _configure_step(env: _StubEnv, done_envs: list[int] | None = None, n_envs: i
 def test_video_files_written_on_termination(tmp_path):
     """A file per camera is written when an env terminates."""
     env = _make_env()
-    with patch("isaaclab_arena.video.camera_observation_video_recorder.ImageSequenceClip") as mock_clip_cls:
+    with _patched_writers() as writers:
         recorder = CameraObsVideoRecorder(env, video_folder=str(tmp_path))
 
         _configure_step(env)
-        recorder.step(None)  # accumulate one frame
+        recorder.step(None)  # stream one frame
 
         _configure_step(env, done_envs=[0])
-        recorder.step(None)  # env 0 terminates → flush
+        recorder.step(None)  # env 0 terminates → finalise
 
-        written_paths = [c.args[0] for c in mock_clip_cls.return_value.write_videofile.call_args_list]
-        assert len(written_paths) == len(CAMERAS)
+        finalised = [writer.filename for writer in writers if writer.closed]
+        assert len(finalised) == len(CAMERAS)
         for cam in CAMERAS:
-            assert os.path.join(str(tmp_path), f"robot-cam-env0-{cam}-episode-0.mp4") in written_paths
+            assert os.path.join(str(tmp_path), f"robot-cam-env0-{cam}-episode-0.mp4") in finalised
+
+
+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)
 
 
 def test_episode_counter_increments_per_env(tmp_path):
     """Each env tracks its own episode count independently via the env's centralized index."""
     env = _make_env()
-    with patch("isaaclab_arena.video.camera_observation_video_recorder.ImageSequenceClip"):
+    with _patched_writers():
         recorder = CameraObsVideoRecorder(env, video_folder=str(tmp_path))
 
         _configure_step(env)
@@ -123,10 +179,8 @@ def test_episode_counter_increments_per_env(tmp_path):
 def test_multiple_episodes_produce_sequential_filenames(tmp_path):
     """Consecutive episodes for an env are named episode-0, episode-1, ..."""
     env = _make_env()
-    written_paths = []
 
-    with patch("isaaclab_arena.video.camera_observation_video_recorder.ImageSequenceClip") as mock_clip_cls:
-        mock_clip_cls.return_value.write_videofile.side_effect = lambda path, **_: written_paths.append(path)
+    with _patched_writers() as writers:
         recorder = CameraObsVideoRecorder(env, video_folder=str(tmp_path))
 
         for _ in range(3):
@@ -135,60 +189,70 @@ def test_multiple_episodes_produce_sequential_filenames(tmp_path):
             _configure_step(env, done_envs=[0], n_envs=1)
             recorder.step(None)
 
+    written_paths = [writer.filename for writer in writers]
     for episode in range(3):
         for cam in CAMERAS:
             assert os.path.join(str(tmp_path), f"robot-cam-env0-{cam}-episode-{episode}.mp4") in written_paths
 
 
 def test_partial_episode_dropped_on_close(tmp_path):
-    """Frames accumulated without termination are silently discarded on close()."""
+    """Frames streamed without a termination leave no file behind on close()."""
     env = _make_env()
-    with patch("isaaclab_arena.video.camera_observation_video_recorder.ImageSequenceClip") as mock_clip_cls:
+    with _patched_writers() as writers:
         recorder = CameraObsVideoRecorder(env, video_folder=str(tmp_path))
 
         _configure_step(env)
-        recorder.step(None)  # accumulate frames, no termination
+        recorder.step(None)  # stream frames, no termination
 
         recorder.close()
 
-        mock_clip_cls.return_value.write_videofile.assert_not_called()
+        # Every encoder was shut down and its incomplete file removed.
+        assert writers and all(writer.closed for writer in writers)
+        assert list(tmp_path.iterdir()) == []
 
 
 def test_no_video_written_for_empty_episode(tmp_path):
-    """An env terminating with no buffered frames writes no video; its episode index still advances."""
+    """An env terminating with no recorded frames writes no video; its episode index still advances."""
     env = _make_env()
-    with patch("isaaclab_arena.video.camera_observation_video_recorder.ImageSequenceClip") as mock_clip_cls:
+    with _patched_writers() as writers:
         recorder = CameraObsVideoRecorder(env, video_folder=str(tmp_path))
 
         # Terminate on the very first step — no prior frames were recorded.
         _configure_step(env, done_envs=[0])
         recorder.step(None)
 
-        # No video for the empty episode, but the env's centralized index still advanced past it
-        # (so a later episode's video number stays in lockstep with the per-episode results record).
-        mock_clip_cls.return_value.write_videofile.assert_not_called()
+        # No encoder was ever opened for the empty episode, but the env's centralized index still
+        # advanced past it (so a later episode's video number stays in lockstep with the
+        # per-episode results record).
+        assert not any(writer.filename.endswith("env0-front-episode-0.mp4") for writer in writers)
         assert env.get_episode_index(0) == 1
 
 
-def test_post_reset_frame_not_appended(tmp_path):
-    """The obs on a terminal step (post-reset) is not buffered for the next episode."""
+def test_post_reset_frame_not_recorded(tmp_path):
+    """The obs on a terminal step (post-reset) is not recorded into either episode."""
     env = _make_env()
-    with patch("isaaclab_arena.video.camera_observation_video_recorder.ImageSequenceClip"):
+    with _patched_writers() as writers:
         recorder = CameraObsVideoRecorder(env, video_folder=str(tmp_path))
 
         _configure_step(env)
-        recorder.step(None)  # 1 frame buffered for both envs
+        recorder.step(None)  # 1 frame recorded for both envs
 
         _configure_step(env, done_envs=[0])
         recorder.step(None)  # env 0 terminates; post-reset frame discarded
 
-        # env 0's buffer is empty (flushed and post-reset frame discarded)
-        for cam in CAMERAS:
-            assert recorder.buffers[cam][0] == []
-
-        # env 1 accumulated 2 frames (neither step was terminal for it)
+        writer_by_filename = {writer.filename: writer for writer in writers}
         for cam in CAMERAS:
-            assert len(recorder.buffers[cam][1]) == 2
+            # env 0's episode 0 closed holding only the single pre-termination frame, and no
+            # encoder was opened for its next episode.
+            env0_episode0 = writer_by_filename[os.path.join(str(tmp_path), f"robot-cam-env0-{cam}-episode-0.mp4")]
+            assert env0_episode0.frames_written == 1
+            assert env0_episode0.closed
+            assert os.path.join(str(tmp_path), f"robot-cam-env0-{cam}-episode-1.mp4") not in writer_by_filename
+
+            # env 1 recorded 2 frames (neither step was terminal for it) and is still open.
+            env1_episode0 = writer_by_filename[os.path.join(str(tmp_path), f"robot-cam-env1-{cam}-episode-0.mp4")]
+            assert env1_episode0.frames_written == 2
+            assert not env1_episode0.closed
 
 
 @pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg not available")

isaaclab_arena/video/camera_observation_video_recorder.py

@@ -5,9 +5,10 @@
 
 """Gym wrapper that records one mp4 per (env, camera, episode) in ``obs['camera_obs']``.
 
-Frames are flushed to disk each time an environment resets (terminated or
-truncated), so each output file corresponds to exactly one complete episode.
-Partial episodes cut off by ``num_steps`` are discarded on ``close()``.
+Each frame is streamed straight to a per-(env, camera) ffmpeg encoder as it arrives, and
+the file is finalised when that environment resets (terminated or truncated), so each
+output file corresponds to exactly one complete episode. Partial episodes cut off by
+``num_steps`` are deleted on ``close()``.
 
 Output filename: ``<name_prefix>-env<N>-<camera_name>-episode-<E>.mp4``
 
@@ -16,11 +17,8 @@
 mounted camera mp4s (what the policy actually sees) are written
 together when ``--record_viewport_video --record_camera_video`` is set.
 
-Memory note: each env buffers raw uint8 frames for its current episode before
-encoding.  Buffers are cleared after each episode is written to disk, so peak
-RAM is N×L×H×W×C bytes where L is max episode length, not the full rollout.
-For 10 envs, 500-step episodes, 512×512×3 frames that is ~3.8 GB of raw frames
-— the encoded mp4s are far smaller (H.264 compresses ~100:1).
+Memory note: This class uses incremental ffmpeg encoding to avoid storing the raw frames in
+memory, which uses substantial amounts of RAM.
 """
 
 from __future__ import annotations
@@ -32,7 +30,7 @@
 import torch
 from dataclasses import dataclass
 
-from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
+from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter
 
 CAMERA_OBS_GROUP_KEY = "camera_obs"
 
@@ -91,11 +89,22 @@ def _sanitize_cam_key(camera_name: str) -> str:
     return camera_name.replace("/", "_").replace(os.sep, "_")
 
 
+@dataclass
+class EpisodeVideoWriter:
+    """The open ffmpeg encoder for one (env, camera, episode) and the file it is writing."""
+
+    writer: FFMPEG_VideoWriter
+    """Encoder consuming raw frames; closing it finalises the mp4."""
+
+    path: str
+    """Destination the encoder is writing to, kept so a partial episode can be deleted."""
+
+
 class CameraObsVideoRecorder(gym.Wrapper):
     """Record one mp4 per (env, camera, episode) in ``obs['camera_obs']``.
 
     Cameras are batched as ``[N_envs, H, W, C]``.  Each env is recorded
-    independently; its buffer is flushed when that env resets (terminated
+    independently; its encoder is finalised when that env resets (terminated
     or truncated), producing one file per completed episode:
     ``<name_prefix>-env<N>-<camera_name>-episode-<E>.mp4``.
     """
@@ -113,8 +122,9 @@ def __init__(
         self.name_prefix = name_prefix
         self.fps = fps if fps is not None else int(env.metadata.get("render_fps", 30))
 
-        # camera_name -> list of per-env frame lists: buffers[camera_name][env_idx] = [frame, ...]
-        self.buffers: dict[str, list[list[np.ndarray]]] = {}
+        # camera_name -> one entry per env, holding that env's open encoder for its current
+        # episode, or None while no episode is in progress.
+        self.writers: dict[str, list[EpisodeVideoWriter | None]] = {}
 
     def step(self, action):
         result = self.env.step(action)
@@ -133,37 +143,60 @@ def step(self, action):
             done_set = set(done_envs)
 
             for camera_name, frames in cam_obs.items():
-                if camera_name not in self.buffers:
-                    self.buffers[camera_name] = [[] for _ in range(n_envs)]
+                if camera_name not in self.writers:
+                    self.writers[camera_name] = [None] * n_envs
                 for env_idx in range(n_envs):
                     if env_idx not in done_set:
-                        self.buffers[camera_name][env_idx].append(_to_uint8(frames[env_idx]))
+                        self._write_frame(camera_name, env_idx, _to_uint8(frames[env_idx]))
 
             if done_envs:
-                self._flush_envs(done_envs)
+                self._finish_envs(done_envs)
 
         return result
 
-    def _flush_envs(self, env_ids: list[int]) -> None:
+    def _write_frame(self, camera_name: str, env_idx: int, frame: np.ndarray) -> None:
+        """Append one frame to this (env, camera)'s episode video, opening the encoder if needed."""
+        assert frame.ndim == 3 and frame.shape[2] == 3, (
+            f"Camera '{camera_name}' produced a frame of shape {frame.shape}; expected (H, W, 3) RGB."
+            " The encoder is configured for 3-channel input."
+        )
+        episode_writer = self.writers[camera_name][env_idx]
+        if episode_writer is None:
+            # The env's counter still names the episode now in progress; it is advanced on reset,
+            # inside env.step.
+            episode_num = self.unwrapped.get_episode_index(env_idx)
+            path = os.path.join(
+                self.video_folder,
+                format_episode_video_filename(self.name_prefix, env_idx, camera_name, episode_num),
+            )
+            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),
+                path=path,
+            )
+            self.writers[camera_name][env_idx] = episode_writer
+        episode_writer.writer.write_frame(frame)
+
+    def _finish_envs(self, env_ids: list[int]) -> None:
+        """Finalise the mp4 of every stream open for each env that just reset."""
         for env_idx in env_ids:
-            # The Arena env has already advanced its per-env episode counter for this reset (within
-            # env.step, before it returned), so the just-finished episode's index is one behind the
-            # current count. Sharing the env's index keeps the filename's episode number in lockstep
-            # with the per-episode results record's ``episode_in_env``.
-            episode_num = self.unwrapped.get_episode_index(env_idx) - 1
-            for camera_name, env_frame_lists in self.buffers.items():
-                frames = env_frame_lists[env_idx]
-                if not frames:
+            for env_writers in self.writers.values():
+                episode_writer = env_writers[env_idx]
+                if episode_writer is None:
                     continue
-                path = os.path.join(
-                    self.video_folder,
-                    format_episode_video_filename(self.name_prefix, env_idx, camera_name, episode_num),
-                )
-                clip = ImageSequenceClip(list(frames), fps=self.fps)
-                clip.write_videofile(path, logger=None, audio=False)
-                del clip
-                env_frame_lists[env_idx] = []
+                episode_writer.writer.close()
+                env_writers[env_idx] = None
 
     def close(self) -> None:
-        # Partial episodes (cut off by num_steps rather than a real reset) are discarded.
+        # Partial episodes (cut off by num_steps rather than a real reset) are discarded: the
+        # encoder is shut down and the incomplete file it was writing is removed.
+        for env_writers in self.writers.values():
+            for env_idx, episode_writer in enumerate(env_writers):
+                if episode_writer is None:
+                    continue
+                episode_writer.writer.close()
+                if os.path.exists(episode_writer.path):
+                    os.remove(episode_writer.path)
+                env_writers[env_idx] = None
         self.env.close()