32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
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"))
|