ColaboratoryでOpenAI Gymを使う

はじめに †

ColaboratoryはGoogleが提供している無料で使えるJupyter Notebook環境です。 TensorFlow用のプロセッサーTPUが使えるので、深層学習もできます。

ここでは、Colaboratory上でOpenAIが提供する強化学習ライブラリーのGymを用いて深層強化学習を実行する方法を説明します。 OpenAIが用意している強化学習アルゴリズムのライブラリーのBaselinesは使いにくいので、ここではBaselinesを使いやすくしたライブラリーであるStable Baselinesを使います。

次の環境で確認しました。

• TensorFlow-GPU 1.15
• OpenAI Gym 0.15.3
• Stable Baselines 2.2.1

環境構築 †

!apt-get -y -q install cuda-libraries-dev-10-0 > /dev/null
!apt-get -y -q install xvfb freeglut3-dev ffmpeg > /dev/null
!npm install -g localtunnel

!pip install pyglet
!pip install pyopengl
!pip install pyvirtualdisplay
!pip install tensorflow-gpu==1.15
!pip install gym
!pip install stable-baselines

関数 †

100エピソードの平均ステップ数を評価する関数（CartPole用） †

def evaluate(env_name, model, num_steps=1000):
  episode_rewards = [0.0]
  local_env = gym.make(env_name)
  obs = local_env.reset()
  for i in range(num_steps):
      # _states are only useful when using LSTM policies
      action, _states = model.predict(obs)
      obs, reward, done, info = local_env.step(action)

      # Stats
      episode_rewards[-1] += reward
      if done:
          obs = local_env.reset()
          episode_rewards.append(0.0)
  # Compute mean reward for the last 100 episodes
  mean_100ep_reward = round(np.mean(episode_rewards[-100:]), 1)
  print("Mean reward:", mean_100ep_reward, "Num episodes:", len(episode_rewards))

  return mean_100ep_reward

1エピソード実行する様子を画面に出力する関数 †

def show(env_name, model, num_steps=1000):
  local_env = gym.make(env_name)
  obs = local_env.reset()
  for step in range(num_steps):
    plt.imshow(local_env.render(mode='rgb_array'))
    display.clear_output(wait=True)
    display.display(plt.gcf())

    # _states are only useful when using LSTM policies
    action, _states = model.predict(obs)
    obs, reward, done, info = local_env.step(action)

    if done:
      local_env.close()
      break
  print('Num steps:', step)

  return step

1エピソード実行する様子を動画に記録する関数 †

def record(env_name, model, num_steps=1000):
  from gym import wrappers

  local_env = gym.make(env_name)
  wrap_env = wrappers.Monitor(local_env, './', force=True)
  obs = wrap_env.reset()
  for step in range(num_steps):
    # _states are only useful when using LSTM policies
    action, _states = model.predict(obs)
    obs, reward, done, info = wrap_env.step(action)
    if done:
      wrap_env.close()
      break
  print('Num steps:', step)

  return step

動画ファイルをダウンロードする関数 †

def download():
  from google.colab import files
  import glob

  for file in glob.glob('openaigym.video.*.mp4'):
    files.download(file)

ここからメイン †

%tensorflow_version 1.x
import os
import gym
import numpy as np
from stable_baselines.common.policies import MlpPolicy, CnnPolicy
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines import PPO2

# TensorBoard用ログ
log_dir = './logs/'
os.makedirs(log_dir, exist_ok=True)

# ディスプレイ
from pyvirtualdisplay import Display
display = Display(visible=0, size=(1024, 768))
display.start()
import os
os.environ["DISPLAY"] = ":" + str(display.display) + "." + str(display.screen)

from IPython import display
import matplotlib.pyplot as plt
%matplotlib inline


# 環境
#env_name = 'MsPacman-v0'
env_name = 'CartPole-v0'
ori_env = gym.make(env_name)
# vectorized environments allow to easily multiprocess training
# we demonstrate its usefulness in the next examples
env = DummyVecEnv([lambda: ori_env])  # The algorithms require a vectorized environment to run


# PPO2エージェント
model = PPO2(MlpPolicy, env,  verbose=1, tensorboard_log=log_dir)  # 入力がセンサー値のとき
#model = PPO2(CnnPolicy, env,  verbose=1, tensorboard_log=log_dir)  # 入力が画像のとき

# 保存されたエージェントを読み込む（2回目以後）
#model = PPO2.load('model')
#model.set_env(env)

# 学習前のエージェントを評価する（CartPole用）
mean_reward_before_train = evaluate(env_name, model, num_steps=10000)

# 学習前のランダムに動くエージェントの動きを確認する（デバッグ用）
show(env_name, model, num_steps=1000)

# 学習前のランダムに動くエージェントを動画に記録する
num_steps = record(env_name, model, num_steps=10000)

# 動画を記録するのに時間がかかるので少し待つ
!sleep 10

# 動画ファイルをダウンロードする
download()

ダウンロードに失敗する場合は、まだファイルができていない可能性があります。 このセルをもう一度実行すると、ダウンロードできることがあります。

# エージェントに学習させる
model.learn(total_timesteps=100000)

TensorFlowが実行される度に、次のようなログが出力されます。

--------------------------------------
| approxkl           | 4.4409087e-05 |
| clipfrac           | 0.0           |
| ep_rewmean         | nan           |
| eplenmean          | nan           |
| explained_variance | 0.245         |
| fps                | 806           |
| nupdates           | 782           |
| policy_entropy     | 0.52703226    |
| policy_loss        | 0.00044246577 |
| serial_timesteps   | 100096        |
| time_elapsed       | 128           |
| total_timesteps    | 100096        |
| value_loss         | 3.9141965     |
--------------------------------------

必要ない場合は、エージェントを生成するときのverboseを0にします。

# 学習後のエージェントを評価する（CartPole用）
mean_reward = evaluate(env_name, model, num_steps=10000)

# 学習後のエージェントの動きを確認する（デバッグ用）
show(env_name, model, num_steps=1000)

# 学習後のランダムに動くエージェントを動画に記録する
num_steps = record(env_name, model, num_steps=10000)

# 動画を記録するのに時間がかかるので少し待つ
!sleep 10

# 動画ファイルをダウンロードする
download()

# 学習したエージェントをファイルに保存する
model.save('model')

# エージェントを削除する（デバッグ用）
#del model

TensorBoard †

# TensorBoardとLocalTunnelを起動する
get_ipython().system_raw('tensorboard --logdir {} --host 0.0.0.0 --port 6006 &'.format(log_dir))
get_ipython().system_raw('lt --port 6006 > url.txt 2>&1 &')

# TensorBoardが起動するの時間がかかるので少し待つ
!sleep 10

# 次で出力されるURLにアクセス（ランタイムを再起動すると無効になる）
!cat url.txt

学習のログをTensorBoardで確認できます。

参考にした資料 †

OpenAI Gym on Colaboratory †

• Jun Nishi: ColaboratoryでOpenAI gym
• @enmaru: Google Colaboratoryの90分セッション切れ対策【自動接続】 - Qiita

Stable-Baelines †

• Stable Baselines, a Fork of OpenAI Baselines - Getting Started
• Stable Baselines入門 / エージェントの保存と読み込み - note

TensorBoard on Colaboratoey †

• @kazuyakitahara: Google Colabratory 上でTensorBoardを使う方法 - Qiita
• npaka: Stable Baselines入門 / 学習の監視 - note