1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162 | from collections.abc import Callable
import os
import os.path
from pathlib import Path
from gymnasium import error, logger
import numpy as np
from lit_sb3.vec_env.base_vec_env import VecEnv, VecEnvObs, VecEnvStepReturn, VecEnvWrapper
from lit_sb3.vec_env.dummy_vec_env import DummyVecEnv
from lit_sb3.vec_env.subproc_vec_env import SubprocVecEnv
class VecVideoRecorder(VecEnvWrapper):
"""Wraps a VecEnv or VecEnvWrapper object to record rendered image as mp4 video.
It requires ffmpeg or avconv to be installed on the machine.
Note: for now it only allows to record one video and all videos
must have at least two frames.
The video recorder code was adapted from Gymnasium v1.0.
Args:
venv:
video_folder: Where to save videos
record_video_trigger: Function that defines when to start recording.
The function takes the current number of step,
and returns whether we should start recording or not.
video_length: Length of recorded videos
name_prefix: Prefix to the video name
"""
video_name: str
video_path: str
def __init__(
self,
venv: VecEnv,
video_folder: str | Path,
record_video_trigger: Callable[[int], bool],
video_length: int = 200,
name_prefix: str = "rl-video",
) -> None:
VecEnvWrapper.__init__(self, venv)
self.env = venv
# Temp variable to retrieve metadata
temp_env = venv
# Unwrap to retrieve metadata dict
# that will be used by gym recorder
while isinstance(temp_env, VecEnvWrapper):
temp_env = temp_env.venv
if isinstance(temp_env, (DummyVecEnv, SubprocVecEnv)):
metadata = temp_env.get_attr("metadata")[0]
else: # pragma: no cover # assume gym interface
metadata = temp_env.metadata
self.env.metadata = metadata
assert self.env.render_mode == "rgb_array", f"The render_mode must be 'rgb_array', not {self.env.render_mode}"
self.frames_per_sec = self.env.metadata.get("render_fps", 30)
self.record_video_trigger = record_video_trigger
self.video_folder = str(Path(video_folder).resolve())
# Create output folder if needed
Path(self.video_folder).mkdir(parents=True, exist_ok=True)
self.name_prefix = name_prefix
self.step_id = 0
self.video_length = video_length
self.recording = False
self.recorded_frames: list[np.ndarray] = []
try:
import moviepy # noqa: F401 # type: ignore[import-untyped]
except ImportError as e: # pragma: no cover
raise error.DependencyNotInstalled("MoviePy is not installed, run `pip install 'gymnasium[other]'`") from e
def reset(self) -> VecEnvObs:
obs = self.venv.reset()
if self._video_enabled():
self._start_video_recorder()
return obs
def _start_video_recorder(self) -> None:
# Update video name and path
self.video_name = f"{self.name_prefix}-step-{self.step_id}-to-step-{self.step_id + self.video_length}.mp4"
self.video_path = str(Path(self.video_folder) / self.video_name)
self._start_recording()
self._capture_frame()
def _video_enabled(self) -> bool:
return self.record_video_trigger(self.step_id)
def step_wait(self) -> VecEnvStepReturn:
obs, rewards, dones, infos = self.venv.step_wait()
self.step_id += 1
if self.recording:
self._capture_frame()
if len(self.recorded_frames) > self.video_length:
print(f"Saving video to {self.video_path}")
self._stop_recording()
elif self._video_enabled():
self._start_video_recorder()
return obs, rewards, dones, infos
def _capture_frame(self) -> None:
assert self.recording, "Cannot capture a frame, recording wasn't started."
frame = self.env.render()
if isinstance(frame, np.ndarray):
self.recorded_frames.append(frame)
else:
self._stop_recording()
logger.warn(
f"Recording stopped: expected type of frame returned by render to be a numpy array, got instead {type(frame)}."
)
def close(self) -> None:
"""Closes the wrapper then the video recorder."""
VecEnvWrapper.close(self)
if self.recording: # pragma: no cover
self._stop_recording()
def _start_recording(self) -> None:
"""Start a new recording. If it is already recording, stops the current recording before starting the new.
one.
"""
if self.recording: # pragma: no cover
self._stop_recording()
self.recording = True
def _stop_recording(self) -> None:
"""Stop current recording and saves the video."""
assert self.recording, "_stop_recording was called, but no recording was started"
if len(self.recorded_frames) == 0: # pragma: no cover
logger.warn("Ignored saving a video as there were zero frames to save.")
else:
from moviepy.video.io.ImageSequenceClip import ImageSequenceClip # type: ignore[import-untyped]
clip = ImageSequenceClip(self.recorded_frames, fps=self.frames_per_sec)
clip.write_videofile(self.video_path)
del clip
del self.recorded_frames
self.recorded_frames = []
self.recording = False
def __del__(self) -> None:
"""Warn the user in case last video wasn't saved."""
if len(self.recorded_frames) > 0: # pragma: no cover
logger.warn("Unable to save last video! Did you call close()?")
|