Google Colabで生成した動画をインライン表示する

「つくりながら学ぶ！深層強化学習」という本の中で、CartPoleなどで動画を生成して表示しているものがあります。

Google Colab上で実行している場合、本の例ではそのまま動画で表示することができないのですが、Colab(Jupyter)のセル上で動画表示可能にすることができるようにしました。

といっても、こちらの「ChainerRL を Colaboratory で動かす - Qiita」を参考にさせていただいただけなのですが。

下記のように必要なツールやライブラリを入れて、HTMLへの出力(Javascriptでのアニメーション化)を利用します。

必要なツールやライブラリのインストール

!apt -qq install xvfb ffmpeg> /dev/null
!pip -q install gym
!pip -q install JSAnimation
!pip -q install pyglet==1.2.4
!pip -q install pyopengl
!pip -q install pyvirtualdisplay
!apt -qq install python-opengl

出力用の仮想ディスプレイを生成

from pyvirtualdisplay import Display
display = Display(visible=0, size=(1024, 768))
display.start()
import os
os.environ["DISPLAY"] = ":" + str(display.display) + "." + str(display.screen)

必要なPythonのライブラリをimport

import matplotlib
import gym
import JSAnimation
import pyglet
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

CartPoleの動きをランダムで行わせる

frames = []
env = gym.make('CartPole-v0')
observation = env.reset()

for step in range(0, 200):
  frames.append(env.render(mode='rgb_array'))
  action = np.random.choice(2)
  observation, reward, done, info = env.step(action)

動画をフレームごとに生成して表示

from JSAnimation.IPython_display import display_animation
from matplotlib import animation
from IPython.display import HTML

plt.figure(figsize=(frames[0].shape[1]/72.0, frames[0].shape[0]/72.0), dpi=72)
patch = plt.imshow(frames[0])
plt.axis('off')
  
def animate(i):
  patch.set_data(frames[i])
    
anim = matplotlib.animation.FuncAnimation(plt.gcf(), animate, frames=len(frames), interval=50)
HTML(anim.to_jshtml())