ただし、「IMPALA」といった、現行の分散強化学習アーキテクチャにはいくつかのボトルネックがあり、Googleが発表した「SEED RL」ではこれを解消しているという。
この新しい分散強化学習アーキテクチャでは、モデルでの推論を集中させてバッチ処理レイヤーを導入することで、計算コスト/通信量を削減し、数千台のマシンに拡張可能とのこと。
・モデルの推論を一括してLearner側のGPU/TPUで実行
環境内で経験をサンプリングするActorは、モデルで推論を実行して次のアクションを予測する。Actor自身で推論モデルを更新しつつ、経験が蓄積されるとこれをLearnerに送信。ここでモデルが最適化されてActorのモデルにも更新をかける。Actor側での推論はCPUを使用して実行されるのが通常だ。また、Learnerで更新されたモデルのパラメーターは頻繁にActorに送信される。
SEED RLでは、モデルでの推論を一括してLearner側のGPU/TPUで実行し、パラメーターと状態をローカルに蓄えておくことで、効率よく学習を行う。
・CPUでのパフォーマンスがIMPALAの3~4倍に
非同期のSEED RLでは、ActorとLearnerのポリシー(報酬を最大化するための方策)のズレが大きいため、「V-trace」アルゴリズムによってこれを補正。反復学習を用いたQ-learning(将来的な報酬が最大になるようアクションを選択していく)手法「R2D2」なども統合されているようだ。一般的に利用される学習環境でベンチマークテストしたところ、たとえばDeepMind Labでは、64コアのCloud TPUで毎秒240万フレームを達成。これはIMPALAの80倍の改善とのこと。また、CPUでのパフォーマンスもIMPALAの3~4倍改善され、大幅な学習効率の向上が示された。
Learnerが推論を一括して行うとのアーキテクチャは、GPU/TPUのスケーリングが容易で、Actorの数もスケーリング可能。自動運転や空調管理など現実世界を扱った最適化にも対応しやすいだろう。
参照元:Massively Scaling Reinforcement Learning with SEED RL/ Google AI Blog