| import gym | |
| import numpy as np | |
| class VideoWrapper(gym.Wrapper): | |
| def __init__(self, | |
| env, | |
| mode='rgb_array', | |
| enabled=True, | |
| steps_per_render=1, | |
| **kwargs | |
| ): | |
| super().__init__(env) | |
| self.mode = mode | |
| self.enabled = enabled | |
| self.render_kwargs = kwargs | |
| self.steps_per_render = steps_per_render | |
| self.frames = list() | |
| self.step_count = 0 | |
| def reset(self, **kwargs): | |
| obs = super().reset(**kwargs) | |
| self.frames = list() | |
| self.step_count = 1 | |
| if self.enabled: | |
| frame = self.env.render( | |
| mode=self.mode, **self.render_kwargs) | |
| assert frame.dtype == np.uint8 | |
| self.frames.append(frame) | |
| return obs | |
| def step(self, action): | |
| result = super().step(action) | |
| self.step_count += 1 | |
| if self.enabled and ((self.step_count % self.steps_per_render) == 0): | |
| frame = self.env.render( | |
| mode=self.mode, **self.render_kwargs) | |
| assert frame.dtype == np.uint8 | |
| self.frames.append(frame) | |
| return result | |
| def render(self, mode='rgb_array', **kwargs): | |
| return self.frames | |