OpenAI Gym用ラッパーの自作方法

@abstractmethod
def observation(self, observation):
    raise NotImplementedError

from gym import ObservationWrapper


class TransformObservation(ObservationWrapper):
    r"""Transform the observation via an arbitrary function.
    Example::
        >>> import gym
        >>> env = gym.make('CartPole-v1')
        >>> env = TransformObservation(env, lambda obs: obs + 0.1*np.random.randn(*obs.shape))
        >>> env.reset()
        array([-0.08319338,  0.04635121, -0.07394746,  0.20877492])
    Args:
        env (Env): environment
        f (callable): a function that transforms the observation
    """

    def __init__(self, env, f):
        super().__init__(env)
        assert callable(f)
        self.f = f

    def observation(self, observation):
        return self.f(observation)

@abstractmethod
def reward(self, reward):
    raise NotImplementedError

from gym import RewardWrapper


class TransformReward(RewardWrapper):
    r"""Transform the reward via an arbitrary function.
    Example::
        >>> import gym
        >>> env = gym.make('CartPole-v1')
        >>> env = TransformReward(env, lambda r: 0.01*r)
        >>> env.reset()
        >>> observation, reward, done, info = env.step(env.action_space.sample())
        >>> reward
        0.01
    Args:
        env (Env): environment
        f (callable): a function that transforms the reward
    """

    def __init__(self, env, f):
        super().__init__(env)
        assert callable(f)
        self.f = f

    def reward(self, reward):
        return self.f(reward)

@abstractmethod
def action(self, action):
    raise NotImplementedError

@abstractmethod
def reverse_action(self, action):
    raise NotImplementedError

import numpy as np
from gym import ActionWrapper
from gym.spaces import Box


class ClipAction(ActionWrapper):
    r"""Clip the continuous action within the valid bound."""

    def __init__(self, env):
        assert isinstance(env.action_space, Box)
        super().__init__(env)

    def action(self, action):
        return np.clip(action, self.action_space.low, self.action_space.high)

!apt update
!apt install xvfb
!pip install pyvirtualdisplay

from pyvirtualdisplay import Display

d = Display()
d.start()

import gym
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib.animation import PillowWriter
from IPython.display import HTML

class GymVideo(gym.Wrapper):
    def __init__(self, env):
        super().__init__(env)
        self.frames = []

    def step(self, ac, monitor_frame=False):
        if monitor_frame:
            self.frames.append(self.env.render(mode='rgb_array'))
        return self.env.step(ac)

    def save_video(self, save_name=None, display_inline=False):
        """
        params:
            save_name      : default None. If you set save_name, gif animation will be saved.
            display_inline : You should use when you set save_name.
                             if you set display_inline True, gif animation will be displayed inline.
        """
        patch = plt.imshow(self.frames[0], cmap='gray')
        plt.axis('off')
        
        def animate(i):
            patch.set_data(self.frames[i])
        
        anim = animation.FuncAnimation(plt.gcf(), animate, frames=len(self.frames), interval=50)

        if save_name != None:
            if display_inline != True:
                anim.save(save_name, writer=PillowWriter(fps=60))
            else:
                anim.save(save_name, writer=PillowWriter(fps=60))
                display(HTML(anim.to_jshtml(default_mode='once')))
        else:
            display(HTML(anim.to_jshtml(default_mode='once')))
        plt.close()