研究用のRLフレームワークを作ってる話: DebugRLの紹介
強化学習苦手の会Advent Calendar 2020の13日目の記事です。
今回の記事はRLのアルゴリズム開発用のフレームワークを最近がんばって開発してるので紹介したいと思います。 僕は全てのRLフレームワークを触ってるわけではないですが、よく知られてるやつ:
- PFRL
- Autonomous Learning Library
- OpenAI Spinning Up
- rlpyt
- Machina
は一通り触ってきました。
全て深層強化学習用のRLフレームワークであり、大体のDeep RLアルゴリズムの比較が容易に実行できるよう各々何かしらの工夫が凝らされてます。 既存のRLアルゴリズムをロボティクスなどアプリケーションとして利用するだけならかなり便利ですが、新しいアルゴリズム開発をする上では、
など、どれも何かしらの欠点を抱えてる印象でした。 思いついたアルゴリズムを実装し、論文にまとめるまでの流れでは
- 新しいアルゴリズムを簡単に載せられるか？途中で簡単に変更を加えられるか？ (実装、デバッグするとき)
- アルゴリズムの条件を簡単に変更できるか？ (デバッグ、論文書くとき)
- アルゴリズムの特徴が実際に反映されているか定量的に評価できるか？ (論文書くとき)
- 実験結果の管理が容易か？結果は簡単に視覚化できるか？他実験との条件の比較は容易か？ (デバッグ、論文書くとき)
の４つがすごい大事になってきます。 特に３つ目はコードに落としこむのがだいぶめんどくさいですし、AtariやMujocoで評価するのは骨が折れます(何回もSeedを回さないと行けないし、そもそも学習にどちゃくそ時間がかるのでめんどくさい)。 この４つをフレームワークが満たしていれば、アルゴリズム開発がめっちゃ簡単になるやんけ。 そう思って新しいフレームワークDebugRLを書いてます。
DebugRLの紹介 (概要は日本語のREADMEを見てね。)
DebugRLは先程の４つの条件を満たすことを目的にしてます(特に３つ目)。 それぞれについて意識したことをまとめてみました。
新しいアルゴリズムを簡単に載せられるか？途中で簡単に変更を加えられるか？
これはOpenAI Spinningupっぽくすることで実現してます。 既存のフレームワークはAgentクラスを各アルゴリズムで共有させて抽象化を図ることが多いですが、個人的にはこの形式は新しいアルゴリズムの実装に向いておらず、また、可読性も大きく下げてしまうと思っています。 一方でOpenAI Spinningupはやや冗長な書き方をしている部分はありますが、アルゴリズム同士の依存関係が最小であり、かなり読みやすく、また、アルゴリズムの実装や編集がしやすいです。 これを踏まえて、DebugRLでもアルゴリズム同士の関係性を可能な限り減らし、OpenAI Spinningupっぽい書き方を意識することで可読性を上げ、アルゴリズムの実装と編集をしやすくしてます。
アルゴリズムの条件を簡単に変更できるか？
この部分はもうちょっと改良できる余地があると思いますが、現在はハイパラなどを全て辞書形式で管理し、後でTrainsに渡すことで実現してます。(Trainsは後で説明します) こんな感じ。 argparseでハイパラを設定することが多いですが、個人的には管理がしづらく、またハイパラの内容が分かりづらいことが多いので、辞書で統一しました。
アルゴリズムの特徴が実際に反映されているか定量的に評価できるか？
これはDebugRLで一番頑張ったところになります。 この思想自体は[1902.10250] Diagnosing Bottlenecks in Deep Q-learning Algorithmsを参考にしているので真新しいものではないですが、ちゃんとフレームワークとして成立しているものは存在していない気がします。
論文のablation studyなどではアルゴリズムの特徴がきちんと現れているか確認しなければいけませんが、これはRLアルゴリズムだと結構大変です。 例えばTRPOやPPOでは更新前と更新後の方策のKLダイバージェンスを小さくするトリックが入っていますが、それをMujocoやAtariで評価するのは割とめんどくさいです (やる場合は何度も環境をResetして大量のサンプルを集めてKLダイバージェンスを近似する必要がでてきます)
こういったサンプルでの近似評価は、そもそも環境のダイナミクスが行列形式で表せる場合は必要になりません。 RLのデバッグ時はPendulumやMountainCarなどかなり低次元環境でテストすることが多いですが、PendulumもMountainCarもある程度ダイナミクスを離散化すれば行列形式で表すことができます。 ダイナミクスが行列形式で表せる場合、Q値の理論値や定常分布の理論値を計算することは容易ですし、先ほどのKLダイバージェンスの差異もサンプルで近似する必要もなくなります。
これを踏まえて、DebugRLではダイナミクスを行列形式で表したデバッグ環境４つ、GridCraft, MountainCar, Pendulum, CartPoleを用意しています。 どれも真のQ値の計算や定常分布の計算ができるので、先ほどの評価がめっちゃ簡単です。やったね。
RLの研究をしているなら誰しも「あれ、思いついたアルゴリズムで得られるQ値って本当に真のQ値に近いんかな？」って思うことがあります。
通常のGymで評価するのは結構面倒くさいですが、DebugRLなら簡単にできちゃいます。
例えば、離散化されたPendulumでのQ値は下図のようにDQNでもそれっぽいのが学習できてるのが視覚的にわかります。
DQNのような離散行動だけでなく、ロボティクス界隈の欲張りさんのためにも連続行動アルゴリズム(SAC)にも対応させていますし、CNN使ったアルゴリズムを解析したい人のためにPendulumとMountainCarでは画像入力にも対応させています。
実験結果の管理が容易か？結果は簡単に視覚化できるか？他実験との条件の比較は容易か？
これはTrainsと呼ばれる実験管理ツールと連携することで実現しています。
多くのRLライブラリはTensorboardを採用していますが、個人的にはTensorboardは結構重くて使いづらいですし、実験条件の比較もしづらいです。
一方で、Trainsは実験条件の管理が神なので採用しました。
先ほど出た辞書であるOPTIONSをTrainsに渡すことで、実験条件をTrainsに記録させ、他の実験と容易に比較できるようにしています。
その他の基本的な機能であるadd_scalarやadd_histogramなどはTensorboardと一緒ですが、かなり軽いので使いやすいです。
今頑張ってるところ
現在はDebugRLはAtariやMujocoを走らせることができません。 せっかくデバッグしたコードが存在するなら全く同じコードでAtariやMujocoを走らせたいので、現在はDebugRLのAtariやMujocoへの対応を頑張ってます。 おわり。
