Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

tests/test_video_recorder.py

@@ -0,0 +1,79 @@
+"""Tests for video recording with mediapy."""
+
+from pathlib import Path
+from unittest.mock import Mock
+
+import mediapy as media
+import numpy as np
+import torch
+
+
+def _make_mock_env(num_envs: int = 1):
+  """Create a mock environment that produces random RGB frames."""
+  env = Mock()
+  env.render_mode = "rgb_array"
+  env.metadata = {"render_fps": 30}
+  env.render.return_value = np.random.default_rng().integers(
+    0, 255, (num_envs, 64, 64, 3), dtype=np.uint8
+  )
+  env.step.return_value = (
+    torch.zeros(num_envs),  # obs
+    torch.zeros(num_envs),  # reward
+    torch.zeros(num_envs, dtype=torch.bool),  # terminated
+    torch.zeros(num_envs, dtype=torch.bool),  # truncated
+    {},  # info
+  )
+  env.close.return_value = None
+  env.unwrapped = env
+  return env
+
+
+def test_step_trigger_writes_video(tmp_path: Path):
+  """VideoRecorder writes a readable mp4 when the step trigger fires."""
+  from mjlab.utils.wrappers.video_recorder import VideoRecorder
+
+  env = _make_mock_env()
+  recorder = VideoRecorder(
+    env,
+    video_folder=tmp_path,
+    step_trigger=lambda step: step == 0,
+    video_length=5,
+    disable_logger=True,
+  )
+
+  action = torch.zeros(1)
+  for _ in range(6):
+    recorder.step(action)
+
+  recorder.close()
+
+  videos = list(tmp_path.glob("*.mp4"))
+  assert len(videos) == 1
+
+  # Verify the file is a valid video readable by mediapy.
+  frames = media.read_video(str(videos[0]))
+  assert len(frames) == 5
+  assert frames[0].shape == (64, 64, 3)
+
+
+def test_accepts_string_path(tmp_path: Path):
+  """VideoRecorder accepts a string path for video_folder."""
+  from mjlab.utils.wrappers.video_recorder import VideoRecorder
+
+  env = _make_mock_env()
+  folder = str(tmp_path / "vids")
+  recorder = VideoRecorder(
+    env,
+    video_folder=folder,
+    step_trigger=lambda step: step == 0,
+    video_length=3,
+    disable_logger=True,
+  )
+
+  action = torch.zeros(1)
+  for _ in range(4):
+    recorder.step(action)
+
+  recorder.close()
+
+  assert list(Path(folder).glob("*.mp4"))