Stable Baselinesを使ってスーパーマリオブラザーズ1-1をクリアするまで

はじめに

強化学習でゲームを解くことに興味があってちょっとずつ勉強している強化学習苦手勢です。Advent Calender 用に何かゲームの攻略について書いてみることにしました。OpenAI が公開している Stable Baselines の紹介とそれを使ってスーパーマリオブラザーズ 1-1 をクリアするところまでやりましたという内容です。OpenAI Gym / Baselines 深層学習・強化学習 人工知能プログラミング 実践入門を参考にました。

Stable Baselines

Stable Baselines は OpenAI が公開している強化学習アルゴリズムの実装セットです。多くの有名な手法を利用しやすい形で提供しています。強化学習アルゴリズムは実装時に仕込んでしまったバグの検出が難しく、学習が上手くいかない時に問題を切り分けにくくしています。Stable Baselines のような標準化された実装を使用することで実装上の問題に気を取られることがなくなります。

使用方法もかなり簡単です。最小のコードでは以下のようになります。

stable_baseline_sample.py

import gym

from stable_baselines.common.policies import MlpPolicy
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines import PPO2

env = gym.make('CartPole-v1') # CartPole環境の作成
env = DummyVecEnv(env)

model = PPO2(MlpPolicy, env, verbose=1) # Proximal Policy Optimization で学習を行うためのオブジェクトを作成。
model.learn(total_timesteps=10000) # 10000ステップ分の経験で学習を行う

Stable Baselines のインストールは Python3.7 で作った仮想環境に以下のコマンドで導入しました。
Microsoft MPI が必要です。

pip install stable-baselines[mpi]
pip install tensorflow==1.14.0
pip install pyqt5
pip install imageio

スーパーマリオブラザーズ環境の準備

9-3. vcoptでスーパーマリオ1-1をクリアするを参考にしてスーパーマリオブラザーズ環境を用意しました。

Stable Baselines は OpenAI Gym(OpenAIが提供する強化学習用の環境を集めたライブラリ) と互換な形式の環境を要求します。今回は https://github.com/Kautenja/gym-super-mario-bros のものを使用しました。

また、上記「9-3. vcoptでスーパーマリオ1-1をクリアする」にならって nes-py も導入しました。これはボタン入力のラッパーを提供してくれています。

ファミコンのコントローラーには A、B、←、↑、→、↓、START、SELECT のボタンがあり、可能な行動の種類はそれぞれのボタンについてオンまたはオフのいずれかを取るため $2^8=256$ 通りとなります。
しかし、スーパーマリオブラザーズのゲームプレイ上意味のある行動は

• 何もしない
• 右へ歩く (→)
• 右へダッシュする (→ + B)
• 右へジャンプする (→ + A)
• 右へダッシュジャンプする (→ + A + B)
• その場でジャンプする (A)
• 左へ歩く (←)
• 左へダッシュする (← + B)
• 左へジャンプする (← + A)
• 左へダッシュジャンプする (← + A + B)
• しゃがむ、土管に入る (↓)
• ツタを登る (↑)

の 12 種類に限られる (ステージによってはもっと少なくなる) ため、256 通りのボタンのオンオフの組み合わせを直接行動とするよりは 12 種類の中からどれか 1 つを行動として選ぶという方が合理的です。nes-py の JoypadSpace はその橋渡しをしてくれます。

gym-super-mario-bros も nes-py も pip でインストールできます。

pip install gym-super-mario-bros
pip install nes-py

ランダム行動をさせてみました。

マリオが跳んだり走ったりできるようになりましたね。

学習/推論

学習用のコードに入る前に以下のような utils.py を用意しておきます。

utils.py

import os
import datetime

import numpy as np
import pytz
import gym
from stable_baselines.results_plotter import ts2xy
from stable_baselines.bench.monitor import load_results

log_dir = "./logs/"
os.makedirs(log_dir, exist_ok=True)

best_mean_reward = -np.inf
nupdates = 1


def callback(_locals, _globals):
    global nupdates
    global best_mean_reward

    if (nupdates + 1) % 10 == 0:
        x, y = ts2xy(load_results(log_dir), "timesteps")
        if len(y) > 0:
            mean_reward = np.mean(y[-10:])
            update_model = mean_reward > best_mean_reward
            if update_model:
                best_mean_reward = mean_reward
                _locals["self"].model.save("mario_model")
            now = datetime.datetime.now(pytz.timezone("Asia/Tokyo"))
            print(f"time: {now}, nupdates: {nupdates}, mean: {mean_reward:.2f}, best_mean: {best_mean_reward}, model_update: {update_model}")
    nupdates += 1
    return True


class CustomRewardAndDoneEnv(gym.Wrapper):
    def __init__(self, env):
        super(CustomRewardAndDoneEnv, self).__init__(env)

    def step(self, action):
        state, reward, done, info = self.env.step(action)

        reward = reward / 10

        if info["life"] < 2:
            done = True
        
        return state, reward, done, info

CustomRewardAndDoneEnv は報酬と終了条件を修正するためのクラスです。gym-super-mario-bros では直前のマリオの位置より右側に移動していれば +1 の報酬が得られる形になっていますが、報酬が大きすぎない方がよいとOpenAI Gym / Baselines 深層学習・強化学習 人工知能プログラミング 実践入門に書いてあったのためこれを 1/10 にスケールしています。また、終了条件をゲームオーバー時ではなくライフが初期値の 2 から減った時としています。

callback は学習途中で行いたい動作を定義する関数です。ここではログの表示と累積報酬の平均値が過去の最高値を更新した時にモデルを保存する動作を入れています。

次に学習のメインとなるコードです。

mario.py

from nes_py.wrappers import JoypadSpace
import gym_super_mario_bros
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT, COMPLEX_MOVEMENT
from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines import PPO2
from stable_baselines.common import set_global_seeds
from stable_baselines.bench import Monitor
from baselines.common.retro_wrappers import *

from utils import log_dir, callback, CustomRewardAndDoneEnv

env = gym_super_mario_bros.make('SuperMarioBros-1-1-v0')
env = JoypadSpace(env, SIMPLE_MOVEMENT)
env = CustomRewardAndDoneEnv(env)
env = StochasticFrameSkip(env, n=4, stickprob=0.25)
env = Downsample(env, 2)
env = Rgb2gray(env)
env = FrameStack(env, 4)
env = ScaledFloatFrame(env)
env = Monitor(env, log_dir, allow_early_resets=True)

env.seed(0)
set_global_seeds(1)
env = DummyVecEnv([lambda: env])

model = PPO2("CnnPolicy", env, verbose=0)
model.learn(total_timesteps=1000000, callback=callback)

env = gym_super_mario_bros.make('SuperMarioBros-1-1-v0')
env = JoypadSpace(env, SIMPLE_MOVEMENT)
env = CustomRewardAndDoneEnv(env)
env = StochasticFrameSkip(env, n=4, stickprob=0.25)
env = Downsample(env, 2)
env = Rgb2gray(env)
env = FrameStack(env, 4)
env = ScaledFloatFrame(env)
env = Monitor(env, log_dir, allow_early_resets=True)

環境をラップしてラップしてラップする部分です。環境が返してくる状態 (ここではゲームの画面) をエージェントが利用しやすくするように前処理を行う、JoypadSpaceのようにエージェントから受け取った行動を環境が扱えるように変換する、utils.py で実装した CustomRewardAndDoneEnv のように報酬と終了条件に手を加えるなど、環境とエージェントの間のやりとりを円滑化します。自分で書くと面倒なものもあるので、Stable Baseline で提供されているのはありがたいです。それぞれどのような処理をしているか簡単に触れておきます。

StochasticFrameSkip

n フレームごとに行動を決定し、n フレームの間その行動を持続するようにするラッパー。

Downsample

画像サイズを小さくするラッパー。

Rgb2gray

RGB 画像をグレースケールに変換するラッパー。

FrameStack

直近 k フレームをまとめて 1 つの状態としてエージェントに与えるラッパー。動きを捉えるのに有効。

ScaledFloatFrame

画素値を 255 で割って 0～1 に正規化するラッパー。

Monitor

monitor.csv にログを書き出すラッパー。

行動、状態、報酬のやり取りに伴って何かやりたいことがある場合はそれ用のラッパーを書くというのが流儀のようです。

実際に学習に必要な部分は次の 2 行です。

model = PPO2("CnnPolicy", env, verbose=0)
model.learn(total_timesteps=128000, callback=callback)

学習の結果がこちら。

どうも途中で敵に当たってしまうようです。seed をいくつか変えてやってみてもやはり途中でやられてしまっています。

PPO2 には設定可能なハイパーパラメータが列挙されています。その中で今回のスーパーマリオブラザーズに関係ありそうな部分をいじってみました。

learning_rate

学習率です。デフォルトは 0.00025。下げることで学習を安定化させられることがあるようなので、1/10 の 0.000025 に落としました

gamma

割引率です。デフォルトは 0.99。0.8～0.999 程度に設定するのが一般的なようで、1 に近いほど遠い将来に得られる報酬を大きく評価します。スーパーマリオブラザーズではあまり長期的な視点は必要ないと思われるので 0.9 に設定しました。

n_steps

学習を行う前に何ステップ分の経験を収集するか。デフォルトは 128。十分な経験を収集したうえで学習を行うため 1024 に設定しました。

nminibatches

学習時のミニバッチサイズです。デフォルトは4。デフォルトよりは多少大きくしておいた方がよかろうと考えとりあえず 32 に設定しました。

noptepochs

収集した経験を用いて何ステップの学習を行うかです。デフォルトは 4。これも少し大きめに 16 に設定しました。

total_timesteps は 1000000 に設定して学習を開始しましたが、200000 程度で打ち切りました。 

というわけでクリアできました。ここまででハイパーパラメータをいじりつつ半日程度です。

なお、PPO2 は確率的に行動を決定するため、乱数の seed 値によってクリアできたりできなかったりします。学習済みのモデルを用いて 100 通りの seed で実験してみると 5% ～ 20% のクリア率となっていました (ぶれがあるのはちゃんと乱数を固定できていないためだと思います……)。半分程度のマリオはゴール少し前の穴を越えられておらず難所であることがわかりました。