追記: 2022/1/2

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

0. はじめに

Google Colab上でOpenAI Gym を描画する方法を調べたのでメモ。

参考にしたサイト群

1. 課題

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

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を利用するので、インストールする。

!apt update
!apt install xvfb

(Dockerイメージなどで独自にJupyter Notebookを起動させる際には、OpenGL関連も必要なため、 apt install python-opengl とする。)

さらに、Xvfbを Google Colab (Jupyter Notebook)上から利用するために、PyVirtualDisplayを利用する。

!pip install pyvirtualdisplay

from pyvirtualdisplay import Display

d = Display()
d.start()

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

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

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

2.2 方法1

1つ目はシンプルにmatplotlibで画面データを描画しては消しを繰り返す方法である。

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

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を使ってアニメーションを表示する方法である。

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

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)メソッドを呼び出す際に自動で保存される。

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

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

3.1 インストール

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

!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)

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)

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)

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 を立ててもらえると嬉しい。

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