File size: 1,937 Bytes
987ed1b | 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 | import gym
import numpy as np
from diffusion_policy.real_world.video_recorder import VideoRecorder
class VideoRecordingWrapper(gym.Wrapper):
def __init__(self,
env,
video_recoder: VideoRecorder,
mode='rgb_array',
file_path=None,
steps_per_render=1,
**kwargs
):
"""
When file_path is None, don't record.
"""
super().__init__(env)
self.mode = mode
self.render_kwargs = kwargs
self.steps_per_render = steps_per_render
self.file_path = file_path
self.video_recoder = video_recoder
self.step_count = 0
def reset(self, **kwargs):
obs = super().reset(**kwargs)
self.frames = list()
self.step_count = 0
self.video_recoder.stop()
if self.file_path is not None \
and ((self.step_count % self.steps_per_render) == 0):
if not self.video_recoder.is_ready():
self.video_recoder.start(self.file_path)
frame = self.env.render(
mode=self.mode, **self.render_kwargs)
assert frame.dtype == np.uint8
self.video_recoder.write_frame(frame)
return obs
def step(self, action):
result = super().step(action)
self.step_count += 1
if self.file_path is not None \
and ((self.step_count % self.steps_per_render) == 0):
if not self.video_recoder.is_ready():
self.video_recoder.start(self.file_path)
frame = self.env.render(
mode=self.mode, **self.render_kwargs)
assert frame.dtype == np.uint8
self.video_recoder.write_frame(frame)
return result
def render(self, mode='rgb_array', **kwargs):
if self.video_recoder.is_ready():
self.video_recoder.stop()
return self.file_path
|