SageMaker RLでtic-tac-toe(3目並べ)エージェントの強化学習モデルを作成する

2019.12.24

DA事業本部の貞松です。各所アドベントカレンダーもいよいよ終盤です。

本記事は『機械学習 on AWS Advent Calendar 2019』24日目のエントリーです。
クラスメソッド 機械学習 on AWS Advent Calendar 2019 - Qiita
クラスメソッド 機械学習 on AWS Advent Calendar 2019 | シリーズ | Developers.IO

今回はSageMaker RLを用いてtic-tac-toe(3目並べ)エージェントの強化学習モデルを作成します。

SageMaker RLの概要

SageMaker RLを使用するメリット

通常のSageMakerと同様にマネージドなインフラストラクチャやトレーニングジョブ、エンドポイントのデプロイなどを利用することで、強化学習に関係する処理の構築に集中することができます。
また、AWSから提供されるコンテナには一般的な強化学習ライブラリが含まれており、サンプルのソースコードとJupyter Notebookを利用してそれらの使い方を簡単に試すことができます。

使用可能なフレームワーク

SageMaker RLでは、Apache MXNetとTensorFlow用のコンテナを利用可能です。

使用可能なライブラリ

上述のフレームワーク用コンテナには以下のような強化学習ライブラリが含まれています。

  • Open AI Gym
  • Intel Coach
  • Berkeley Ray RLLib

サポートされているシミュレーション環境

SageMaker RLでは、以下のシミュレーション環境をサポートしています。

  • AWS RoboMakerおよびAmazon Sumerianのファーストパーティシミュレーター
  • RoboschoolやEnergyPlusなどのGym Interfaceを使用して開発されたOpenAI Gym環境およびオープンソースシミュレーション環境
  • Gym Interfaceを使用して自前で開発したシミュレーション環境
  • MATLABやSimulinkなどの商用シミュレーター ※ユーザーが自前でライセンスを管理する必要があります

tic-tac-toeエージェントの強化学習モデルを作成する

毎度おなじみ、amazon-sagemaker-examplesreinforcement_learning のサンプルも用意されています。
その中にtic-tac-toeのサンプルソースコード、ノートブックが用意されているので、こちらに沿って強化学習モデルを作成していきます。

amazon-sagemaker-examples/reinforcement_learning/rl_tic_tac_toe_coach_customEnv/

モデルトレーニングの構成要素

  • 機械学習フレームワーク: MXNet
  • 強化学習ライブラリ: Coach, RLLib
  • 描画ライブラリ: ipywidgets

ノートブックの実行

環境設定

import sagemaker

bucket = sagemaker.Session().default_bucket()
role = sagemaker.get_execution_role()
import os
import numpy as np
import sagemaker
from sagemaker.rl import RLEstimator, RLToolkit, RLFramework
from tic_tac_toe_game import TicTacToeGame

内部で使用されるソースコードの確認

以下、内部で使用するソースコードの内容を確認するコマンドの羅列です。実際のソースコードはリンク先のGitHubページを参照してください。

tic_tac_toe.pyは、学習時に使用されるシミュレーション環境(tic-tac-toeのゲーム環境)が実装されています。

!pygmentize ./src/tic_tac_toe.py

/reinforcement_learning/rl_tic_tac_toe_coach_customEnv/src/tic_tac_toe.py

preset.pyは学習パラメータのプリセットが定義されています。
ここで強化学習エージェントの使用する学習アルゴリズムも設定されており、ここではCliped PPOを使用しています。

!pygmentize ./src/preset.py

/reinforcement_learning/rl_tic_tac_toe_coach_customEnv/src/preset.py

train_coach.pyはトレーニング実行時のエントリーポイントです。

!pygmentize ./src/train-coach.py

/reinforcement_learning/rl_tic_tac_toe_coach_customEnv/src/train-coach.py

学習の設定と実行

sagemaker.rlで実装されているRLEstimatorに対して、学習の為の設定を渡して実行します。
ここではRLのtoolkitとしてCoach、frameworkとしてMXNetを使用しており、学習用のインスタンスはml.m4.xlargeを1台で設定しています。

estimator = RLEstimator(source_dir='src',
                        entry_point="train-coach.py",
                        dependencies=["common/sagemaker_rl"],
                        toolkit=RLToolkit.COACH,
                        toolkit_version='0.11.0',
                        framework=RLFramework.MXNET,
                        role=role,
                        train_instance_count=1,
                        train_instance_type='ml.m4.xlarge',
                        output_path='s3://{}/'.format(bucket),
                        base_job_name='DEMO-rl-tic-tac-toe',
                        hyperparameters={'save_model': 1})

estimator.fit()

estimator.fitで学習を実行するとSageMakerのトレーニングジョブが作成、実行されます。
学習実行時には以下のように、学習のステータスや学習途中の状況が出力されます。

2019-12-24 09:27:26 Starting - Starting the training job...
2019-12-24 09:27:28 Starting - Launching requested ML instances......
2019-12-24 09:28:32 Starting - Preparing the instances for training......
2019-12-24 09:29:56 Downloading - Downloading input data
2019-12-24 09:29:56 Training - Training image download completed. Training in progress..

### 中略

## Starting to improve simple_rl_graph task index 0
Training> Name=main_level/agent, Worker=0, Episode=1, Total reward=-0.9, Steps=14, Training iteration=0
Training> Name=main_level/agent, Worker=0, Episode=2, Total reward=0.3, Steps=25, Training iteration=0
Training> Name=main_level/agent, Worker=0, Episode=3, Total reward=0, Steps=30, Training iteration=0
Training> Name=main_level/agent, Worker=0, Episode=4, Total reward=-1.1, Steps=34, Training iteration=0
Training> Name=main_level/agent, Worker=0, Episode=5, Total reward=-1.1, Steps=38, Training iteration=0

### 中略

## agent: Starting evaluation phase
Testing> Name=main_level/agent, Worker=0, Episode=9007, Total reward=1, Steps=50000, Training iteration=24
Testing> Name=main_level/agent, Worker=0, Episode=9007, Total reward=1, Steps=50000, Training iteration=24
Testing> Name=main_level/agent, Worker=0, Episode=9007, Total reward=-2.0, Steps=50000, Training iteration=24
Testing> Name=main_level/agent, Worker=0, Episode=9007, Total reward=-2.0, Steps=50000, Training iteration=24
Testing> Name=main_level/agent, Worker=0, Episode=9007, Total reward=-2.0, Steps=50000, Training iteration=24
## agent: Finished evaluation phase. Success rate = 0.0, Avg Total Reward = -0.8
ONNX correction applied to discrete PPO agent.
2019-12-24 09:41:45,868 sagemaker-containers INFO     Reporting training SUCCESS

エンドポイントにデプロイ

学習済みモデルをエンドポイントにデプロイします。
RLExtimatorのdeployにentry_pointとしてdeploy-coach.pyを渡して実行します。

!pygmentize ./src/deploy-coach.py

/reinforcement_learning/rl_tic_tac_toe_coach_customEnv/src/deploy-coach.py

predictor = estimator.deploy(initial_instance_count=1, 
                             instance_type='ml.m4.xlarge', 
                             entry_point='deploy-coach.py')

学習済みモデルの動作確認

ipywidgetsを使用してJupyter Notebook上にインタラクティブなビジュアル(GUI的なもの)を描画できます。

https://ipywidgets.readthedocs.io/en/latest/

これを利用して学習済みの強化学習エージェントとtic-tac-toeで対戦できるGUIが実装されています。

/reinforcement_learning/rl_tic_tac_toe_coach_customEnv/tic_tac_toe_game.py

t = TicTacToeGame(predictor)
t.start()

実際に実行してみると3マス×3マスの格子が描画され、エージェントが先行で×を置き、プレイヤーが後攻で○を置いていきます。

先行の学習済みのエージェントは必ず初手で角に×を置いてきます。 そこで後攻のプレイヤーが初手で真ん中に○を置きつつ、その後はエージェントが2個×を並べてくるのを防ぎ続けると何とか引き分けに持ち込めます。

後攻のプレイヤーが初手で真ん中に○を置かなかった場合は、その後どこに○を置いても確実に負けます。強い。

エンドポイントの削除

一通りの実行、確認が完了したらエンドポイントを削除しておきましょう。

predictor.delete_endpoint()

まとめ

SageMaker RLを使用してtic-tac-toe(3目並べ)エージェントの強化学習モデルを作成してみました。
環境設定からトレーニングジョブ実行、エンドポイントへのデプロイまで、トータル30分以内で完全に定石を抑えた強化学習エージェントが仕上がりました。
サンプルや使用されているライブラリのリソースを参考にして、他のゲームや問題に対する適用や、他のアルゴリズムに使用やパラメータチューニング等を行うことができそうです。
re:Invent 2019にて、来シーズンに向けたアップデートが多数発表されたDeepRacerもSageMaker RL + RoboMakerの組み合わせが使用されているので、そのあたりも絡めてSageMaker RLをうまく活用して強化学習の理解を深めていきたいです。