diff --git a/isaaclab_arena/tests/test_camera_observation_video_recorder.py b/isaaclab_arena/tests/test_camera_observation_video_recorder.py index 815f56ffde..21efd37b95 100644 --- a/isaaclab_arena/tests/test_camera_observation_video_recorder.py +++ b/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") diff --git a/isaaclab_arena/video/camera_observation_video_recorder.py b/isaaclab_arena/video/camera_observation_video_recorder.py index 842775ad13..f52e05506a 100644 --- a/isaaclab_arena/video/camera_observation_video_recorder.py +++ b/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: ``-env--episode-.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: ``-env--episode-.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()