Google、数千台のマシンに拡張可能な分散強化学習アーキテクチャ「SEED RL」発表！

強化学習では、ゲームなどの環境を探索しながら最適な行動を学習していく。複数の「Actor」による経験を「Learner」に送信し、モデルを更新していく強化学習の分散アーキテクチャでは、学習効率が高められる。

ただし、「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の数もスケーリング可能。自動運転や空調管理など現実世界を扱った最適化にも対応しやすいだろう。