Google Colab上でgym-retroのソニックを学習

以前迷路の学習を方策勾配法でやってみて、それをこちらにまとめた

これと同じ方法をgym-retroに適用してソニックの学習を試してみた。ちなみに先に結果を書いておくと、スコアをちゃんと取れるようになるほどちゃんと学習させることはできていない。ただ、右に行くことに報酬を与えることにより、その行動ばかりとるようにはなるので学習自体はできている模様。

ソニックを実行するためにはromのデータを入手する必要があるが、その方法についてはこちらがわかりやすい

OpenAI Retro Contestの「Gym Retro Integration」でソニック・ザ・ヘッジホッグをプレイする - おおかみ山

今回動かしてみたコードの全体はこちら

machine-learning-samples/sonic_policy_simple.py at master · y-kamiya/machine-learning-samples · GitHub

env.step()が返すdoneやrewardの値はこちらのようなjsonをscenario.jsonとしてrom.mdと同じディレクトリに置くことで設定できる

{
  "done": {
    "variables": {
      "lives": {
        "op": "equal",
        "reference": 2
      }
    }
  },
  "reward": {
    "variables": {
      "x": {
        "reward": 0.001
      },
      "score": {
        "reward": 1.0
      }
    }
  }
}

ライフは１減った時点でdoneとし、x方向に移動した場合とscoreを得た時点でrewardが入るような設定になっている。

また、rewardについては更にgymやgym-retroのbaselinesに含まれるこちらのwrapperを利用して小さな値になるようにしてある。

class RewardScaler(gym.RewardWrapper):
    """
    Bring rewards to a reasonable scale for PPO.
    This is incredibly important and effects performance
    drastically.
    """
    def reward(self, reward):
        return reward * REWARD_RATE

他にもいくつかwrapperを使っているがそれらはすべて以下からloadしたりコピーして持ってきて使っている

retro-baselines/sonic_util.py at master · openai/retro-baselines · GitHub

baselines/atari_wrappers.py at master · openai/baselines · GitHub

ちなみに全体的に参考にしたのはこちらのブログと論文

OpenAI Retro Contest Day 3 – Tristan Sokol – Medium

https://s3-us-west-2.amazonaws.com/openai-assets/research-covers/retro-contest/gotta_learn_fast_report.pdf

ただし、モデルのハイパーパラメータは上記に合わせたわけでなく適当なまま。

とりあえず動くようになった段階で、1 episode = 2500 stepssで400 episodes分学習させてみたところ丸一日くらいかかった。（Core i7, GPUなし）

これでは試行錯誤しづらいため、gpuを使える環境としてGoogle Colabolatoryを利用してみた。使うのは初めてだったが簡単に導入できた。google colabについてはググればたくさん記事が出てくるので割愛。

以下のようなスクリプトを作って毎回のインスタンス起動時に楽できるようにしておいた。

machine-learning-samples/setup_retro_google_colab.sh at master · y-kamiya/machine-learning-samples · GitHub

google colabのインスタンスにつなげたら以下のようにすればpytorch+gym-retroの環境構築し、学習を実行できる

# 最初にrom.mdをupload
from google.colab import files
files.upload()

!git clone https://github.com/y-kamiya/machine-learning-samples.git
!chmod +x machine-learning-samples/python3/reinforcement/sonic/setup_retro_google_colab.sh
!./machine-learning-samples/python3/reinforcement/sonic/setup_retro_google_colab.sh

%cd machine-learning-samples/python3/reinforcement/sonic
!python sonic_policy_simple.py  --episodes 400 --steps 2500 > log

files.download("log")
files.download("model_state_sonic.dat")

これで実行したところ丸一日程度かかってたものが、２時間程度で終わるようになった。