【強化学習】OpenAI Gym を Google Colab上で描画する方法 (2020.6版)

# 追記: 2022/1/2
この記事で紹介している方法のうちの1つの`gym.wrappers.Monitor`が`gym=0.20.0`で非推奨になりましたので、代替手法を調べて新しい記事を書きました。
（その他の手法は変更なし。また、`gnwrapper.Monitor`は代替手法に対応済みのため、そのまま利用できます。）

https://zenn.dev/ymd_h/articles/36eff0de1ee5ed


# 0. はじめに

[Google Colab](https://colab.research.google.com/)上で[OpenAI Gym](https://github.com/openai/gym) を描画する方法を調べたのでメモ。

参考にしたサイト群

* [ColaboratoryでOpenAI gym](http://bcl.sci.yamaguchi-u.ac.jp/~jun/home/180828-openai-colaboratory/)
* [ChainerRL を Colaboratory で動かす](https://qiita.com/ikeyasu/items/ec3c88ce13a2d5e41f26)
* [OpenAI GymをJupyter notebookで動かすときの注意点一覧](https://qiita.com/mittyantest/items/2ea07b0f8d5976f29da9)
* [How to run OpenAI Gym .render() over a server](https://stackoverflow.com/questions/40195740/how-to-run-openai-gym-render-over-a-server)
* [Rendering OpenAI Gym Envs on Binder and Google Colab](https://davidrpugh.github.io/stochastic-expatriate-descent/openai/binder/google-colab/2020/04/16/remote-rendering-gym-envs.html)

# 1. 課題

 `gym.Env` の `render()` メソッドで環境を表示しようとする際に`NoSuchDisplayException`エラーが出る。



``` python
import gym
env = gym.make('CartPole-v1')
env.reset()
env.render()
```
```
NoSuchDisplayException                    Traceback (most recent call last)
<ipython-input-3-74ea9519f385> in <module>()
      2 env = gym.make('CartPole-v1')
      3 env.reset()
----> 4 env.render()
```


# 2. 対策

調べた限り、Colab上でGymの描画機能を利用する方法は3通りあることがわかった。
どの方法も長短あり、1つに絞ることができなかったので、3種類とも記載する。

## 2.1 共通の準備

3種類いずれの方法でも、X11の仮想ディスプレイである[Xvfb](https://www.x.org/releases/X11R7.6/doc/man/man1/Xvfb.1.xhtml)を利用するので、インストールする。

``` python
!apt update
!apt install xvfb
```

([Dockerイメージ](https://hub.docker.com/_/python)などで独自にJupyter Notebookを起動させる際には、OpenGL関連も必要なため、 `apt install python-opengl` とする。)


さらに、Xvfbを Google Colab (Jupyter Notebook)上から利用するために、[PyVirtualDisplay](https://github.com/ponty/PyVirtualDisplay)を利用する。

``` python
!pip install pyvirtualdisplay

from pyvirtualdisplay import Display

d = Display()
d.start()
```

`"DISPLAY"` 環境変数に `{ディスプレイ番号}.{スクリーン番号}`を設定するという記述があるサイトもあったが、[不要だとPyVirtualDisplayの作者に教えてもらった](https://github.com/ponty/PyVirtualDisplay/issues/54)。

曰く、スクリーン番号は、複数ディスプレイがある状況で利用する値で、PyVirtualDisplayでは1つしか画面を生成しないので`0`固定であり、かつスクリーン番号を書かないと自動的に`0`と解釈されるためだと。([StackOverflow](https://askubuntu.com/a/432257)参照)

もっと言えば、 `pyvirtualdisplay.Display.start()` の中で環境変数を設定しているので、外部から変更することは必要ないとのことであった。
(少なくとも、2020年6月18日時点の最新版である1.3.2で確認済み)


## 2.2 方法1

1つ目はシンプルに[matplotlib](https://matplotlib.org/)で画面データを描画しては消しを繰り返す方法である。

あまり速くない、かつ1回しか表示されないことがデメリットだが、描画データを保持せず上書きし続けるので、描画データが長くなっても対応できる方法である。

```python
import gym
from IPython import display
from pyvirtualdisplay import Display
import matplotlib.pyplot as plt

d = Display()
d.start()

env = gym.make('CartPole-v1')

o = env.reset()

img = plt.imshow(env.render('rgb_array'))
for _ in range(100):
    o, r, d, i = env.step(env.action_space.sample()) # 本当はDNNからアクションを入れる

    display.clear_output(wait=True)
    img.set_data(env.render('rgb_array'))
    plt.axis('off')
    display.display(plt.gcf())

    if d:
        env.reset()
```


## 2.3 方法2

2つ目は、`matplotlib.animation.FuncAnimation`を使ってアニメーションを表示する方法である。

描画画面を繰り返し表示することができ、フレームごとの表示速度を自由に設定できる一方、描画データを保持しておく必要があるためメモリーを多く必要とし、表示する画面サイズや表示枚数を調整しないとメモリーエラーを起こしうる。
（長ーい学習の途中で、エラーを出されると・・・。）

```python
import gym
from IPython import display
from pyvirtualdisplay import Display
import matplotlib.pyplot as plt
from matplotlib import animation


d = Display()
d.start()

env = gym.make('CartPole-v1')

o = env.reset()

img = []
for _ in range(100):
    o, r, d, i = env.step(env.action_space.sample()) # 本当はDNNからアクションを入れる

    display.clear_output(wait=True)
    img.append(env.render('rgb_array'))

    if d:
        env.reset()

dpi = 72
interval = 50 # ms

plt.figure(figsize=(img[0].shape[1]/dpi,img[0].shape[0]/dpi),dpi=dpi)
patch = plt.imshow(img[0])
plt.axis=('off')
animate = lambda i: patch.set_data(img[i])
ani = animation.FuncAnimation(plt.gcf(),animate,frames=len(img),interval=interval)
display.display(display.HTML(ani.to_jshtml()))
```

## 2.4 方法3

最後は、`gym.wrappers.Monitor` を利用して描画データを動画として保存する方法である。
`render()`メソッドは不要で、`step(action)`メソッドを呼び出す際に自動で保存される。

```python
import base64
import io
import gym
from gym.wrappers import Monitor
from IPython import display
from pyvirtualdisplay import Display

d = Display()
d.start()

env = Monitor(gym.make('CartPole-v1'),'./')

o = env.reset()

for _ in range(100):
    o, r, d, i = env.step(env.action_space.sample()) # 本当はDNNからアクションを入れる

    if d:
        env.reset()

for f in env.videos:
    video = io.open(f[0], 'r+b').read()
    encoded = base64.b64encode(video)

    display.display(display.HTML(data="""
        <video alt="test" controls>
        <source src="data:video/mp4;base64,{0}" type="video/mp4" />
        </video>
        """.format(encoded.decode('ascii'))))
```



# 3. ライブラリ: Gym-Notebook-Wrapper

上記の方法を毎回書くのは面倒なため、ライブラリ化した。 

* [Gym-Notebook-Wrapper | GitLab.com](https://gitlab.com/ymd_h/gym-notebook-wrapper)
* [Gym-Notebook-Wrapper | GitHub.com (ミラーレポジトリ)](https://github.com/yamada-github-account/gym-notebook-wrapper)
* [Gym-Notebook-Wrapper | PyPI](https://pypi.org/project/gym-notebook-wrapper/)

## 3.1 インストール

PyPIに公開しているので、 `pip install gym-notebook-wrapper` でインストールできる。

``` shell
!apt update && apt install xvfb
!pip install gym-notebook-wrapper
```

もちろん、Google Colab以外でも利用できるが、 Xvfb を利用するためLinuxが前提。


## 3.2 使い方

gym-notebook-wrapper だと、長いしハイフン(`-`)が入っているので、 インポートできるモジュール名は `gnwrapper` にしてある。

* 方法1 → `gnwrapper.Animation`
* 方法2 → `gnwrapper.LoopAnimation`
* 方法3 → `gnwrapper.Monitor`


### 3.2.1 `gnwrapper.Animation` (= 2.2 方法1)

``` python
import gnwrapper
import gym

env = gnwrapper.Animation(gym.make('CartPole-v1')) # Xvfbが起動される

o = env.reset()

for _ in range(100):
    o, r, d, i = env.step(env.action_space.sample()) # 本当はDNNからアクションを入れる
    env.render() # ここで、前の描画を消し、新しいステップの描画を行う。
    if d:
        env.reset()
```

### 3.2.2 `gnwrapper.LoopAnimation` (= 2.3 方法2)

``` python
import gnwrapper
import gym

env = gnwrapper.LoopAnimation(gym.make('CartPole-v1')) # Xvfbが起動される

o = env.reset()

for _ in range(100):
    o, r, d, i = env.step(env.action_space.sample()) # 本当はDNNからアクションを入れる
    env.render() # ここで、描画データを保存する
    if d:
        env.reset()

env.display() # ここで、保存した描画データをアニメーションとして表示する
```

### 3.2.3 `gnwrapper.Monitor` (= 2.4 方法3)

``` python
import gnwrapper
import gym

env = gnwrapper.Monitor(gym.make('CartPole-v1'),directory="./") # Xvfbが起動される

o = env.reset()

for _ in range(100):
    o, r, d, i = env.step(env.action_space.sample()) # 本当はDNNからアクションを入れる
    if d:
        env.reset()

env.display() # ここで、ビデオとして保存した描画データを表示する
```

# 4. 最後に

ネット上に色々記載されている情報を整理して、OpenAI GymをGoogle Colab上で描画する方法を3種類まとめた。
何度か実際に走らせて確認したコードのはずだけど、コピペミスとかしてたらすみません。

Gym-Notebook-Wrapperはまだまだ荒削りでバグもあるかもしれないので、何かあれば気軽に [issue](https://gitlab.com/ymd_h/gym-notebook-wrapper/-/issues) を立ててもらえると嬉しい。