Adding the possibility to have a prefix on the name of the video.
@@ -78,11 +78,13 @@ class VideoWrapper(base.EnvironmentWrapper):
environment: dm_env.Environment,
*,
path: str = '~/acme',
+ filename: str = '',
process_path: Callable[[str, str], str] = paths.process_path,
record_every: int = 100,
frame_rate: int = 30):
super(VideoWrapper, self).__init__(environment)
self._path = process_path(path, 'videos')
+ self._filename = filename
self._record_every = record_every
self._frame_rate = frame_rate
self._frames = []
@@ -95,7 +97,8 @@ class VideoWrapper(base.EnvironmentWrapper):
def _write_frames(self):
"""Writes frames to video."""
if self._counter % self._record_every == 0:
- path = os.path.join(self._path, '{:04d}.html'.format(self._counter))
+ path = os.path.join(self._path,
+ f'{self._filename}_{self._counter:04d}.html')
video = _make_animation(self._frames, self._frame_rate).to_html5_video()
with open(path, 'w') as f:
