80 lines
1.9 KiB
Python
80 lines
1.9 KiB
Python
|
"""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"))
|