DUICUO

Google は大規模強化学習のフレームワークをオープンソース化しました。

深層強化学習(DRL)は、深層学習分野において最も急速に成長している研究分野の一つです。DRLは、AlphaGo、Dota 2 Five、AlphaStarなど、近年のAIにおける主要なマイルストーンのいくつかを担っており、人間の知能に最も近い分野であると考えられています。しかし、多くの進歩にもかかわらず、DRL手法の実世界への導入は、依然として大規模な人工知能(AI)ラボに限定されています。これは、DRLアーキテクチャが不釣り合いに大量のトレーニングを必要とするため、計算コストが高く、ほとんどの組織にとって現実的ではないことが一因です。最近、Google Researchは、スケーラブルに拡張可能なDRLモデルの新しいアーキテクチャであるSEED RLを提案する論文を発表しました。

[[352276]]

DRLモデルを現実世界に実装する際の課題は、そのアーキテクチャに直接関係しています。DRLは本質的に、ランタイム環境、モデル推論、モデルトレーニング、リプレイバッファなど、様々なタスクを網羅しています。最新のDRLアーキテクチャの多くは、これらのタスクに計算リソースを効率的に割り当てることができず、実装コストが法外に高くなっています。AIハードウェアアクセラレータなどのコンポーネントは、これらの制限の一部に対処するのに役立ってきましたが、その効果は限定的でした。近年、新たなアーキテクチャが登場し、市場で最も成功しているDRL実装の多くに採用されています。

IMPALAにインスパイアされた

現在のDRLアーキテクチャにおいて、IMPALAはこの分野における新たな標準を確立します。DeepMindが2018年の研究論文で最初に提案したIMPALAは、数値計算向けに特別に設計されたアクセラレータを活用したモデルを導入し、長年教師あり学習が享受してきた速度と効率性を最大限に活用しています。IMPALAの中核を成すのは、同時実行性と並列化を最大化するために一般的に用いられるアクターベースモデルです。

IMPALAベースのDRLエージェントのアーキテクチャは、アクターと学習者という2つの主要コンポーネントで構成されています。このモデルでは、アクターは通常CPU上で実行され、環境内での行動とモデルからの推論による次の行動の予測を反復的に行います。アクターは推論モデルのパラメータを頻繁に更新し、十分な数の観測データを収集した後、観測データと行動軌跡を学習者に送信することで、学習者を最適化します。このアーキテクチャでは、学習者は数百台のマシンに分散された推論からの入力を用いて、GPU上でモデルを学習します。計算の観点から見ると、IMPALAアーキテクチャはGPUを用いて学習者の学習を加速させることができ、アクターは多数のマシンにまたがってスケーリングできます。

> 出典: https://ai.googleblog.com/2020/03/massively-scaling-reinforcement.html

IMPALAはDRLアーキテクチャにおける新たな標準を確立しました。しかし、このモデルにはいくつかの固有の制限があります。

• ニューラルネットワーク推論にCPUを使用する:参加者マシンは通常CPUベースです。モデルの計算負荷が増加すると、推論に費やされる時間が環境ステップに必要な計算時間を超え始めます。解決策としては参加者数を増やすことが挙げられますが、これはコストを増加させ、融合に影響を与えます。

• 非効率的なリソース利用:参加者は、環境ステップと推論ステップという2つのタスクを交互に実行します。これらの2つのタスクの計算要件はしばしば異なるため、リソース利用が非効率になったり、参加者の行動が遅くなったりします。

• 帯域幅要件:モデルパラメータ、ループ状態、および観測値は、参加者と学習者間で受け渡されます。さらに、メモリベースのモデルは大きな状態を送信するため、帯域幅の要件が増加します。

Google は、IMPALA アクター モデルにヒントを得て、DRL モデルのスケーリングにおける前身のモデルのいくつかの制限に対処する新しいアーキテクチャを開発しました。

シードRL

全体的に、GoogleのSEED RLアーキテクチャはIMPALAと非常に似ていますが、DeepMindモデルの主要な限界のいくつかに対処するためのいくつかのバリエーションが導入されています。SEED RLでは、ニューラルネットワークの推論は学習者によって専用ハードウェア(GPUまたはTPU)上で集中的に実行されます。これにより、モデルのパラメータと状態がローカルに保持されることで推論が高速化され、データ転送のボトルネックが回避されます。観測データは環境ステップごとに送信されます。

学習者は解決策を与えられ、学習者は推論を行い、参加者にアクションを返します。この巧妙な解決策は、IMPALAのようなモデルの推論限界に対処しますが、遅延の問題が生じる可能性があります。

レイテンシの影響を最小限に抑えるため、SEED RLはメッセージングとストリーミングにgPRCを採用しています。具体的には、SEED RLはストリーミングRPCを活用し、参加者から学習者への接続をオープンに保ち、メタデータは一度だけ送信します。さらに、このフレームワークには、複数の参加者からの推論呼び出しを効率的にバッチ処理するバッチモジュールが含まれています。

> 出典: https://ai.googleblog.com/2020/03/massively-scaling-reinforcement.html

IMPALA アーキテクチャを詳しく調べると、次の 3 つの基本的なスレッド タイプがあることがわかります。

1. 推論

2. データのプリフェッチ

3. トレーニング

推論スレッドは、観測値、報酬、エピソード終了フラグのバッチを受け取ります。これらはループ状態を読み込み、データを推論TPUコアに送信します。推論TPUコアは、サンプリングされたアクションと新しい繰り返し状態を受け取り、最新の繰り返し状態を保存しながら、アクションを参加者に送り返します。軌跡が完全に展開されると、FIFOキューまたはリプレイバッファに追加され、データプリフェッチスレッドによってサンプリングされます。最後に、軌跡はデバイスバッファにプッシュされ、参加している各TPUコアで使用されます。トレーニングスレッド(Pythonのメインスレッド)は、プリフェッチされた軌跡を受け取り、トレーニング済みのTPUコアを使用して勾配を計算し、すべてのTPUコア(推論とトレーニング)にわたってモデルに勾配を同期的に適用します。推論コアとトレーニングコアの比率は、スループットと使用率を最大化するために調整できます。

> 出典: https://arxiv.org/abs/1910.06591

SEED RLアーキテクチャは、学習器を数万コア、参加者数を数万マシンまで拡張できるため、学習器を最大限に活用し、毎秒数百万フレームの学習を可能にします。SEED RLはTensorFlow 2 APIをベースとしており、そのパフォーマンスはTPUによって高速化されます。

SEED RLを評価するために、GoogleはAcadem学習環境、DeepMind Lab環境、そして最近リリースされたGoogle Research Football環境といった一般的なDRLベンチマーク環境を使用しました。いずれの環境でも結果は目覚ましいものでした。例えば、DeepMind Lab環境では、SEED RLは64個のCloud TPUコアを使用して毎秒240万フレームを達成しました。これは、従来の最先端分散エージェントであるIMPALAの80倍の速度です。速度とCPU使用率の向上も確認されました。

> 出典: https://ai.googleblog.com/2020/03/massively-scaling-reinforcement.html

SEED RLは、DRLモデルを大幅に改良したもので、大規模なスケールアップが可能です。Google Researchは、オリジナルのSEED RLアーキテクチャをGitHubでオープンソース化しました。近い将来、これが多くの実用的なDRL実装の基盤モデルとなることは間違いありません。