TF2RL: TensorFlow2.x 向け強化学習ライブラリ


友人の @ohtake_i が開発している強化学習向けのライブラリTF2RLの紹介記事。
私もちょこちょこお手伝い (issue対応・PR作成) させてもらってます。
(拙作のリプレイバッファライブラリcpprbもTF2RLのために作ったと言っても過言ではない。)

1. はじめに

名前のとおり、TensorFlow 2系で書かれた強化学習向けライブラリです。
TensorFlow 1系は、Sessionだとかplaceholderだとか、とっつきにくい部分がありましたが、TensorFlow 2系になってかなり書きやすくなっているので、「TensorFlowは書きにくいからPyTorch一択」と思っている人こそ見てもらいたい。

2. 実装済みアルゴリズム

最新の状況は、GitHubのREADME に記載されています。
2020年8月26日現在、以下のアルゴリズムが実装済みです。(随時まだまだ増やしていく予定です。)

また、一部のアルゴリズムは、ApeXGAEにも対応しています。

3. インストール方法

TensorFlowのバージョンは、2.0から最新の2.3まで対応しているはずです。
TensorFlowのバージョンの関係から、デフォルトではTensorFlowを依存ライブラリとしてインストールしないので、pipなりcondaなりでインストールしておいてください。
(もちろんGPU版でもOKです。2.1以降はPyPIのバイナリはCPU/GPU版の区別がなくなりましたので、今後気にする機会も減っていくとは思いますが)

pipの場合
pip install tensorflow
condaの場合
conda install -c anaconda tensorflow

TF2RLは、PyPIに公開しているので、pipでインストールすることが可能です。

pip install tf2rl

4. 利用方法

以下が、READMEに記載のDDPGを利用したコード例です。
アルゴリズムのエージェントを構築し、環境(gym.Env)とともにTrainerに渡してあげるとアルゴリズムに沿って学習します。

DDPGでPendulumの例
import gym
from tf2rl.algos.ddpg import DDPG
from tf2rl.experiments.trainer import Trainer


parser = Trainer.get_argument()
parser = DDPG.get_argument(parser)
args = parser.parse_args()

env = gym.make("Pendulum-v0")
test_env = gym.make("Pendulum-v0")
policy = DDPG(
    state_shape=env.observation_space.shape,
    action_dim=env.action_space.high.size,
    gpu=-1,  # Run on CPU. If you want to run on GPU, specify GPU number
    memory_capacity=10000,
    max_action=env.action_space.high[0],
    batch_size=32,
    n_warmup=500)
trainer = Trainer(policy, env, args, test_env=test_env)
trainer()

学習結果は、TensorBoardで確認することができます。

tensorboard --logdir results

いくつかのパラメータは、argparse経由でプログラム実行時のコマンドラインのオプションとして渡すことが可能です。

5. 課題と今後

元々スクリプト&コマンド実行を前提としていたので、 Trainerargparseと強く結合していて、Google Colab等のノートブック環境でうまく実行できません。 (Trainer 以外のエージェントはだいたい問題なく動くので、モデルだけ利用&学習のためのループはスクラッチで記述することは可能です。)
そのうち、Trainerにメスを入れてなんとかしたいと思っています。

おわりに

何故かアクティブにフィードバックをくれるのは、中華系と思われる方々が多いようです。
日本人ユーザーももっと増えてくれって思っているので、使ってみてフィードバック(issue, PR)をくれるととても嬉しいです。