OpenAI Gymを使ったAIスーパーマリオの強化学習

からあげさん（ Twitter: @karaage0703 ）が呼びかけていた、強化学習を使ってAIにスーパーマリオをクリアさせようという企画をみかけて、とてもおもしろそうなので挑戦しました。ただ単に実行するだけ、というわけにはいかず自分なりの工夫が必要だったので、試行錯誤の記録を残します。

AIマリオは、OpenAI gymを基盤として公開されているもので、GitHubにレポジトリがあります。
github.com

f:id:WindVoice:20220102175206p:plain

からあげさんのJupyter Notebookでのお試し

最初は、からあげさんがこちらのURLで記事にされていたのを見かけたのがきっかけでした。書かれている通り、AIがスーパーマリオをプレイするというのは私にはとても魅力的に見えます。子供の頃、8-4まで必死になってプレイしました。私にとっては珍しく最後までプレイしたアクションゲームでした。

GitHubにJupyter Notebookを用意していただいているので、まずはこれをコピーしてはじめました。確かに実行はできるのですが、どこに工夫の余地があるのかとっかかりがわからず（PPOとは？）、変更した場合にうまくいっているのかどうかの判断基準もわからないため、しばらく考え込んでしまいました。

別の方が用意している、DDQN版でチャレンジ

からあげさんのNotebookと大本が同じ、別のかたの取り組みがあることがわかったので、そちらを読んでみました。使っている学習アルゴリズムはDDQN（Double Deep Q-Network）とありました。過去にUinityでQ学習を使う簡単なものを実装してみたことがあり、こちらのほうが何をやっているのか分かりそうでした。環境、行動、報酬、学習という枠があり改造するにも少し考えるとっかかりがあります。ということでこちらを基点に触っていくことにしました。

まずは実行

まずは手習いということで実行してみるとたしかに動くのですが、学習を最小限（マリオ105人分＝105エピソード）しかすることができず、学習成果があるのか、よくわかりません。かといって単純にループを増やすと、マリオ320人分あたりでCUDA（GPUを使った機械学習用のライブラリ）がメモリ不足でエラーとなり、学習を続けることができませんでした。

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 15.90 GiB total capacity; 14.51 GiB already allocated; 39.75 MiB free; 14.85 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

安定したループを回すために①

強化学習では、All You Need Is Killみたいなもので、とにかく長時間トライ＆エラーを繰り返し続けるマリオを作る必要があります。メモリ不足の原因を探るために、まずはGPUの状態を調べます。Colab環境にはnvidia-smiというコマンドがあり、これで状況を表示できます。CLIなので、実行時は ! をつけておく必要があります。

!nvidia-smi

f:id:WindVoice:20220102175321p:plain

この場合、GPUはTesla P100であることがわかりますし、GPUメモリは約16GBあることがわかります。マリオ300人分ほどがトライ＆エラーしたあと同じコマンドを実行すると、メモリの使用量が14GBくらいになっていました。どうやらここがメモリ不足に陥っています。メモリを消費しそうな原因を探してゆくと、全てのマリオが共有している記憶があり（Mario.memory）、これはキューであり一定以上には増えないようになっているのですが、それでもGPUのメモリが少ないためあふれているようです。たぶん、最初にこのコードを書いた方の環境はもっとGPUのメモリが豊富だったのでしょう。いくらか試してmaxlen=100000とあるところを最終的には57800まで減らしました。これで散っていったマリオの記憶を学習して次のマリオが熟練することができます。まさにAll You Need Is Kill状態です。

class Mario(Mario):
    def __init__(self, state_dim, action_dim, save_dir):
        super().__init__(state_dim, action_dim, save_dir)
        self.memory = deque(maxlen=57800)
        self.batch_size = 32

安定したループを回すために②

これで一応学習ループは回るようになったのですが、Google Colabフリー版で実行していると、不意に環境がリセットされることがあります。せっかく1000人のマリオが頑張ったのに、その記憶がなくなってしまい、また最初からやり直し、ということが発生しました。そもそも、強化学習は相当まとまった量の学習を続ける必要があり（元の版を作った方の解説によると、1-1をクリアするには4万マリオほど必要とのこと）、24時間以上かかりそうです。

途中の状態を保存して学習再開ができないと、不意の中断で心が折れてしまいます。このため、３つのことをしました。

Google Driveの接続
学習状態の保存
学習結果の復元

保存については最初からMario.save()がありましたのでこれを使いました。
しかしどういうわけかMario.load()がなく、ここだけ自作しました。

LOAD_PATH = '/content/drive/My Drive/gym_mario/checkpoints/2022-01-01T00-00-00/mario_net_xxx.chkpt'
# checkpointからデータを読み込み
loaded_data = torch.load(LOAD_PATH)
# marioオブジェクトに値をロード
mario.net.load_state_dict(loaded_data['model'])
mario.exploration_rate = loaded_data['exploration_rate']

マリオのニューラルネットワークはMarioNetクラス（なんとオシャレな名前！）に実装されており、Mario.netから参照されます。ニューラルネットワークの復元はMario.net.load_stete_dictで行います。exploration_rate（マリオの行動を決めるとき、ランダムに行動を決めるか報酬が高そうな行動を取るか、の大きく２択になるのですが、ランダムな行動を取る割合を示す値）だけはニューラルネットワークとは別に保存されており、これも復元します。他のパラメータを保存しておきたい場合も、Mario.save()と合わせて追加するだけでOKなはずです。

その他細かな変更

マリオがクリアしたときにはぜひ知りたいので（というのも、高速で学習を回すために肝心のプレイ画面がでないため）、クリアがわかるように学習ループの中にメッセージを追加しました。

        # ゲームが終了したかどうかを確認
        if done or info["flag_get"]:
            if info["flag_get"]:
                # ゴールできたら知りたいですよね
                print(f"<<< Mario get the flag. GOOOOOAL! >>>")
            break

ひたすら学習

さて、安定的に実行できるようになったのでひたすら回すべし……　となるのですが、無料版は数時間ごとに止められてしまうこともあり、色々調べてColab Proを買うことにしました。月1180円くらいで（ほぼ）安定的にGPU インスタンスが使えます。良いGPUは高価なので、自分で買うよりこちらのほうがお安いと思っています。機械学習の環境を自分で作るのも大変ですし、継続的に機械学習をやってみるつもりなら、おすすめです。

途中経過の記録

約10000マリオのころ。落ち着きのないプレイ。
youtu.be

約15000マリオのころ。
youtu.be

約25000マリオのころ。だんだん動きが洗練されてきます。
youtu.be

約40000マリオ。当初の目標回数です。
なお、この動画は何度も実行して良いものを使っています。実際のクリア率は0.5%ほどです。
youtu.be

今後について

もう少し継続的に学習を回して、なんとかクリア率1%超えを狙いたいところです。

今回、マリオの行動の選択肢は最小限（右に移動するか、Aボタンを押してジャンプ）しかなく、1-1以外のステージに進んでも、クリア不可能になることが想定されます。そもそもBボタンダッシュができません。また、行動の選択肢を増やすと学習は初めからやりなおしになります。これは選択肢を増やすとニューラルネットワークを大きくしなければならないためです。またそれだけ多くの学習が必要になるはずです。その後どうするかは、いくつか試してから決めようと思っています。