class VideoRecorder():
"""
Use ImageEncoder gym.wrappers.monitoring.video_recorder because it make really nice videos using .mp4 and ffmpeg
"""
def start(self, output_file, observation_shape, fps=30):
"""
:param output_file:
:param observation_shape:
:param fps:
:return:
"""
height = observation_shape[2]
width = observation_shape[1]
pixFmt = observation_shape[0]
frame_shape = (height, width, pixFmt)
self._image_encoder = ImageEncoder(output_file, frame_shape, fps, fps)
def add_frame(self, observation):
"""
:param observation:
:return:
"""
self._image_encoder.capture_frame(observation.swapaxes(0, 2))
def close(self):
self._image_encoder.close()
def __del__(self):
# Release everything if job is finished
self.close()
def record_video(buffer, file_path, fps):
"""
Records the given image buffer to a video file.
:param array_like buffer: the list containing the sequential image frames to be recorded.
:param str file_path: the path to the video file.
:param int fps: the video frames-per-second.
:return:
"""
# creates video encoder and adds each frame in the buffer
video_recorder = ImageEncoder(file_path, buffer[0].shape, fps)
for frame in buffer:
video_recorder.capture_frame(frame)
video_recorder.close()
def animate(self, act_fn, nsteps, **kwargs):
"""act_fn could be a list of functions for each agent in the environemnt that we can control"""
if not isinstance(act_fn, list):
act_fn = [act_fn for _ in range(len(self.agents))]
assert len(act_fn) == len(self.agents)
encoder = None
vid_loc = kwargs.pop('vid', None)
obs = self.reset()
frame = self.render(**kwargs)
if vid_loc:
fps = kwargs.pop('fps', 30)
encoder = ImageEncoder(vid_loc, frame.shape, fps)
try:
encoder.capture_frame(frame)
except error.InvalidFrame as e:
print('Invalid video frame, {}'.format(e))
rew = np.zeros((len(self.agents)))
traj_info_list = []
for step in range(nsteps):
a = list(map(lambda afn, o: afn(o), act_fn, obs))
obs, r, done, info = self.step(a)
rew += r
if info:
traj_info_list.append(info)
frame = self.render(**kwargs)
if vid_loc:
try:
encoder.capture_frame(frame)
except error.InvalidFrame as e:
print('Invalid video frame, {}'.format(e))
if done:
break
traj_info = stack_dict_list(traj_info_list)
return rew, traj_info
class EpisodeRecorder(gym.Wrapper):
"""Wrapper for recording episodes to MP4 or jpeg-frames"""
def __init__(self, env, path, fps=60, header=None):
"""Initialize the wrapper
Args:
env: The env to wrap and record
path: The location where to place recordings
fps: The FPS to encode the MP4 at
header: Optional prefix for the output file name
"""
super(EpisodeRecorder, self).__init__(env)
self.fps = fps
self.path = path
self.header = header
self.encoder = None
os.makedirs(path, exist_ok=True)
def _get_output_name(self):
"""Gets a unique name for the output video, starting from 00000.mp4
(Handles collisions with parallel ENVs)
"""
base_name = self.header + "_" if self.header is not None else ""
i = 0
while True:
try:
path = os.path.join(self.path,
f"{base_name}{str(i).zfill(5)}.mp4")
pathlib.Path(path).touch(exist_ok=False)
return path
except FileExistsError:
i += 1
continue
def _init_new(self, shape):
self._close_last()
self.cur_path = self._get_output_name()
self.encoder = ImageEncoder(self.cur_path, shape, self.fps)
def _dump_frame(self, frame):
self.encoder.capture_frame(frame)
def _close_last(self):
if self.encoder is not None:
self.encoder.close()
self.encoder = None
def close(self):
self._close_last()
self.env.close()
def reset(self):
obs = self.env.reset()
# Create a new output
self._init_new(obs.shape)
self._dump_frame(obs)
return obs
def step(self, action):
obs, reward, done, info = self.env.step(action)
self._dump_frame(obs)
return obs, reward, done, info