「AllenAct」で深層強化学習を試す
AllenNLPでおなじみのAllenAI研究所が、2020年9月に発表したPyTorchの深層強化学習ライブラリ「AllenAct」で、チュートリアルの迷路課題の実装を、AllenActの構造を追いながら解説します。
AllenActとは
AllenActのキャッチコピーは、「An open source framework for research in Embodied AI」です。”Embodied AI”とは、日本語で”身体性人工知能”と訳されます。
「AIって身体があり、環境とのインタラクションがないと本当の知性に到達しないよね」って概念を強調していますが、要は深層強化学習のライブラリと思って大丈夫（と思います）。
本記事ではこのAllenActの最初に紹介されている、「迷路のチュートリアル」を、AllenActの構造を追いながら解説します。
本記事でぜひ、追体験してみてください。
深層強化学習の実装と実行
1. 環境準備
1.1 マシン環境
マシン環境にはAzureのDSVM（Data Science Virtual Machine）（Ubuntu 18.04）を使用しました。こちらはAzureポータルから立ち上げるだけで、すぐにPyTorchやGPU環境を利用することができます。
VMタイプには「NC6（1時間約120円）」を使用しました。
※AzureでのDeepLearning環境作成方法
https://azure-recipe.kc-cloud.jp/2018/11/adsvm-01/
1.2 AllenACTのインストール
作成したVMにSSHで入って、仮想環境を確認し、py37_pytorchを起動します。
conda info -e
conda activate py37_pytorch
なお、PyTorchのバージョンは1.4.0でした。
（PyTorchはバージョン1.0から1.4までと1.5以上で、動作する関数が割と変わります。1.4までで動作していたものが、1.5以降では動作しないことも多いので注意が必要です）
続いて、今回の取り組みを実施するフォルダとして「study」を作成します。
mkdir study
cd study
AllenActをインストールします。インストールページ
https://allenact.org/installation/installation-allenact/
の案内のgit cloneは 、git clone git@github.com:allenai/allenact.git
と記載されていますが、authがうまくいかない？ので、
通常のHTTPSからgitをcloneします。
以下のようにコマンドを実行して、gitをcloneし、requriremetn.txtの内容をpip インストールします。
git clone https://github.com/allenai/allenact.git
cd allenact
pip install requirements.txt
jupyter notebook
これでAllenActのインストール完了です。続いて、Jupyter Notebookを起動します。SSHでポート8888をつなげておき、Jupyter Notebookにアクセスします。
なお以降ずっと、このフォルダ「allenact」直下で作業をします。
2. チュートリアル実験のConfigファイルを実装
2.1 課題概要
今回解きたい課題は、「Navigation in Minigrid」と呼ばれる迷路です。
これは、https://github.com/maximecb/gym-minigridの課題環境を利用しています。
今回はその中でも、 dMiniGrid-Empty-Random-5x5-v0 を使用します。
これは3×3（外枠の壁を入れて5×5）の領域内でゴールの「緑色タイル」まで、赤い「エージェント」を操作する課題です。
この迷路課題の特徴は、エージェントには方向があり、エージェントの真横方向より前方しか観測できない（部分観測型）という点です。
上記の図で色が少しうすい色の領域が、「その時点で観測できてる範囲」になります。
2.2 実装手順
まず、著者のGitHubから「NavigationMiniGrid_jp.ipynb」を持ってきて、フォルダ「allenact」の直下に置いてください。
2.2 パッケージのimport
最初のセルで、必要なパッケージをimportします。
3. 深層強化学習の設定を記載したクラスを作成開始
3.1 実験環境の定義
それでは、main.pyから呼び出す、深層強化学習の設定を記載したクラスを作成します。
クラス名は、 MiniGridTutorialExperimentConfig()とします。
なお、ここからはAllenActのチュートリアル「Tutorial: Navigation in MiniGrid」
https://allenact.org/tutorials/minigrid-tutorial/
をベースにしていますが、元のチュートリアルが分かりにくいので、本記事では説明の構成をかなり変えています。
それでは、JupyterNotebookのセルに以下を書き始めてください。
[1] 上記では、クラスMiniGridTutorialExperimentConfig()を設定し、@classmethodデコレータ（クラスのメソッドをクラスメソッドに変換して、インスタンスを作成しなくても使えるメソッドにしてくれる）で、関数tag()を定義します。これは任意のタグになります。
[2] 続いて、どんな実験環境（Environment）を使用するのかを@staticmethodデコレータ（クラスのメソッドをクラスメソッドに変換して、インスタンスを作成しなくても使えるメソッドにしてくれる）を用いた、関数 make_env() を定義します。
@classmethodと@staticmethodの違いは、@classmethodは必ず引数が必要になる点だけで（今回は cls）、その他の性質は同じです。
ReturnしているEmptyRandomEnv5x5()は最初のimportで実施した、 from gym_minigrid.envs import EmptyRandomEnv5x5
を指します。これは、別ライブラリで作られていいる迷路の実験環境クラスです。
この実験環境クラス（Environment）は、AllenActの構成図（以下）
の「Environment」を意味しています。
このクラス MiniGridTutorialExperimentConfig()内に引き続き、必要な関数を定義していきます。
3.2 TaskSamplerの定義
[3] 次に、上の図での、Task実行時のデータのサンプルを実施する「TaskSampler」を定義します。
まず、 def make_sampler_fun()でサンプル関数を定義します。今回は、 from plugins.minigrid_plugin.minigrid_tasks import MiniGridTaskSampler 、で定義したTaskSamplerを使用します。
今回は予め用意されているMiniGridTaskSampler を利用しますが、オリジナルの環境ではTaskSamplerを自分で実装します。
なおMiniGridTaskSamplerの具体的な実装はこちらとなります。
[4] 続いて、TaskSamplerで、訓練時、検証時、テスト時にどのようにデータをサンプルするのかを定義します。そのために、
train_task_sampler_args() 、valid_task_sampler_args() 、test_task_sampler_args() を定義します。
これらの実装はほとんどおまじない的です。
こんな感じに実装するのか程度の理解で良いかと思います。
[3]、[4] のTaskSamplerたちの実装は次の通りです。
[5] TaskSamplerの、train_task_sampler_args() 、valid_task_sampler_args() 、test_task_sampler_args() で呼び出される、具体的な関数 self._get_sampler_args() を定義します。
この関数は、訓練時（train）は、後ほど定義する、トータルstep数の間、ずっとstepし続けます（ゴールや失敗時には自動で環境をリセット）。 そのため、何回タスクが実施されるかは不明です。トータルsteps数を満たすまでタスクを繰り返します。
検証時とテスト時は、途中保存のタイミングで、max_tasks回、タスクを実行し、その平均の結果を求めます （検証は20回、テスト時は40回）。
3.3 Sensorの定義
続いて、AllenActの概念図（再掲以下）の「Sensor」を定義します。
Sensorは深層強化学習のディープラーニングモデルにデータを与える役割を果たします。
[6] 状態観測のSensorを設定します。今回Sensorに使用するクラスEGOcentricMiniGridSensor() はフォルダ「plugins」に事前に用意されています。
今回は事前に用意されているので良いですが、これを実装するのも面倒そうです。
なお具体的な実装内容はこちらの通りです。
EgocentricMiniGridSensor() の引数view_channels=3は、迷路の各タイルの状態を設定します。
view_channels=3の場合は、
・タイルのタイプ（壁や床やドアやゴールなど11種類）
・タイルの色（赤や緑など6種類）
・タイルの状態（open, closed, lockedの3種類）
の3種類をセンサーで観測します。どんなタイプや色があるのか、 さらなる詳細はプログラムのNotebookの最後に記載しています。
EgocentricMiniGridSensor() の引数agent_view_size=5は、エージェントが観測できる範囲（横・縦）の長さを示します。
agent_view_size=5だと、以下のような観測範囲です。
＊ ＊ ＊ ＊ ＊
＊ ＊ ＊ ＊ ＊
＊ ＊ A ＊ ＊
＊ ＊ ＊ ＊ ＊
＊ ＊ ＊ ＊ ＊
ただしこの範囲のうち、以下の図のように、自分（Actor）が向いている真横と前方のみが見えます(部分観測)。
以下の図はAが上を向いているときの観測範囲となります。ただし、壁などが途中にあると、その先は観測できません。
＊ ＊ ＊ ＊ ＊
＊ ＊ ＊ ＊ ＊
＊ ＊ A ＊ ＊
ｘ ｘ ｘ ｘ ｘ
ｘ ｘ ｘ ｘ ｘ
なお、エージェントができる行動（Action）は、”left”, “right”, “forward”で、それぞれ、「その場で左に向く」、「その場で右に向く」、「1歩前に進む」の3タイプです。
これを実装すると次の通りです。
3.4 ActorCriticModelの定義
続いて、深層強化学習のディープラーニング・ネットワークとして使用する、ActorCritiModelを定義します。
[7] 今回は、予め用意された MiniGridSimpleConvRNN() クラスを使用します。こちらの実装は以下の通りです。
元々、AllenActに RNNActorCritic()が用意されていて、これを迷路用にチューニングしています。
RNNActorCritic()はリカレントニューラルネットワークRNNを使用したメモリを搭載したタイプのActor-Criticのディープラーニングモデルです。
RNN型のメモリを使用することで、今回の迷路課題のように部分観測であり、過去の自分の状態と観測した状態の情報を保持できるようにしています。
3.5 訓練・検証・テストの設定
最後に訓練や検証の設定を実施します。
[8] 訓練や検証、テストでのサブプロセスの設定をします。ここで定義したサブプロセス分、同時並列に深層強化学習が実行されます。
今回は、訓練時は128個、検証とテスト時は16個とします。またGPUの使用もここで設定します。今回はGPUは使用せず、CPUのみで実施しています。
[9] 最後に関数 training_pipeline() を定義し、深層強化学習の訓練・検証・テストのパイプラインを組み立てます。
今回は深層強化学習の手法としては、PPOを使用します。この関数内で同時にPPOの種々設定を与えます。
ネットワークパラメータの最適化にはAdamを使用し、学習率を線形に小さくすることにします。
関数 TraniningPipeline()の引数にsave_interval=10000 を与え、訓練が1万ステップ進むと、そのタイミングで、その時点でのパラメータで、検証とテストが実施されます。
また学習率のスケジューリングを線形減衰に設定し、訓練ステップの間線形に低下させるようにします。
訓練ステップ数は引数 ppo_steps=int(150000) とし、15万ステップとしています。
以上の内容を実装すると次の通りです。
4. 深層強化学習の実施
4.1 Pythonファイルの作成
以上の[1]～[9]により、深層強化学習の実施設定クラスMiniGridTutorialExperimentConfig()が完成しました。
このままJupyterNotebookで実行したかったのですが、このConfigクラスを実行する際には、「allenact」フォルダのmain.pyにあるように、 OnPolicyRunner().start_train() を実行する必要があります。
この関数に引数がかなり多いので、JupyterNotebookで実行するには不便です。そこで、JupyterNotebook上では、インラインで（Shell コマンドの実行）で、main.pyを実行をすることにします。
はじめに実行するためのファイル準備をします。
[10]まず、実行フォルダとして、フォルダ「allenact」の下にフォルダ「minigrid_jp」を作成します。
そしてその中に、上記のMiniGridTutorialExperimentConfig()クラスと、その前のimport文たち、これらを記したpythonファイル「minigrid_tutorial_jp.py」を作成して、配置します。
4.2 実行
[11] 最後にJupyterNotebook上でmain.pyを実行します。チュートリアルではコア数を8に設定しています。自身の現在の環境でのコア数を確認しておきます。
今回のマシンでは8個のCPUコアがなく、6個でした。そこで、4コアを使用することにします。
訓練の実行は次の通りです。以下をJupyterNotebookのセルに記載して実行します。各引数の意味は以下のコメントの通りです。
4.3 訓練結果をTensorBoardで確認
上記を実施すると、3分程度で学習が終了します。
[12] 訓練結果をTensorBoardで確認します。
JupyterNotebookのセルで以下を実行します。
!tensorboard --logdir minigrid_jp
--logdir には、今回の実行フォルダである minigrid_jp を設定します。SSHでポート6006を通しておき、http://localhost:6006/ にアクセスします。
するとTensorBoardで、訓練のステップ数によるメトリクスの変化と、各タイミングでの検証の結果が表示されます。
見づらいので、左端のSmoothingは0にしておきます。
上記の結果を見ると、
●ステップ数が15万回
●訓練では2万ステップあたりから成功（ゴールにエージェントが100ステップ以内にたどり着く）ようになっていること
●訓練でtotal_lossが0に近づいていること
●検証（valid）でゴールまでのエピソード数がだんだん低下し、2万ステップあたりではゴールに8エピソードかかっていたのが、最後は4エピソード以下あたりになっていること
●検証（valid）で2万ステップあたりからゴールできるようになっていること
が分かります。
※なお、検証のステップ数が1万回に設定したのに、1万回ちょうどからプロットされていないのは、並列に4コアで128サブプロセスを実行しているため、1024環境ごとに実行されることになり、1万回ちょうどではなく、1万回を超えた時点で検証が実施されるためです。
4.4 テスト
[13] テストの実行と結果確認をします。
テストを実行する際には、引数 -t に、訓練時のチェックポイント時刻（フォルダ名）を与えます。
このチェックポイントの時刻（フォルダ名）は、訓練を実行した際に、
09/13 20:01:41 INFO: valid worker 0 loading checkpoint from minigrid_jp/checkpoints/MiniGridTutorial_for_Japanese/2020-09-13_19-58-21/exp_MiniGridTutorial_for_Japanese__stage_00__steps_000000133120.pt [engine.py: 299]などと、記載されています。/MiniGridTutorial_for_Japanese/のあとのパス名です。
JupyterNotebook上で以下のように記載して、テストを実行します。
テストの結果を、再度TensorBoardを立ち上げて確認します。
!tensorboard --logdir minigrid_jp
http://localhost:6006/ にアクセスします。
上記の結果を見ると、横軸がステップ数で15万、縦軸が各メトリクスです。
各ステップ数のときのActor-Criticモデルで迷路課題をテストした結果を求めているので、最初の方は、ep_length（ゴールにかかるエピソード数）が大きいですが、最終的には4以下あたりに収束します。
今回は最終的に3.8あたりにたどり着きました。理論的には3.75エピソードが最高パフォーマンスです（その場で左や右に向きを変えるのも1エピソードにカウントされます）。
以上により、ほぼ完ぺきに迷路課題を解く、深層強化学習が実施できました。
まとめ
最後にAllenActでの深層強化学習の手順をまとめます。
[1] 関数tag()を定義
[2] 実験環境（Environment）を定義（既に用意されているものや自作環境を使用）
[3] TaskSamplerを定義（実験環境に合わせて、自分で実装の必要あり）
※今回は予め用意されているMiniGridTaskSampler()を利用
[4] TaskSamplerで、訓練時、検証時、テスト時のサンプル関数を定義
[5] サンプル関数の詳細を定義
[6] 状態観測のSensorを定義（自分で実装の必要あり）
※今回は予め用意されているEgocentricMiniGridSensor()を利用
[7] ActorCriticModelの定義（ある程度のひな形があるが、カスタマイズの必要あり）
※今回は予め用意されているMiniGridSimpleConvRNN() クラスを使用
[8] 訓練・検証・テストの並列実行のサブプロセス数を設定
[9] 訓練・検証・テストのパイプラインを設定
[10] 実行フォルダとPythonファイルを作成
[11] 訓練を実行
[12] 訓練結果をTensorBoardで確認
[13] テストを実行、確認
となります。
自分で作る必要のあるクラスが多く、他の深層強化学習パッケージよりも扱いづらい印象を私は受けました。
しかしその分カスタマイズが効くとも言えます。
また、PPO以外のアルゴリズムとしては、 DD-PPO、 A2Cが実装されているようです。まだあまりバリエーションは多くありません。
最後に、再掲となりますが、本記事の実装コードは以下に置いています。
以上、ご一読いただき、ありがとうございました。
