わずか 1 秒で 100 万フレームのデータ! AI専門家のヤン・シュイチェン氏のチームが強化学習に関する新たな研究を開発し、そのコードをオープンソース化した。

[[434376]]

当時、Dota 2 をプレイする OpenAI Five のトレーニングには 50,000 個を超える CPU コアが必要でした。ちょうど昨日、Yan Shuicheng 氏のチームが最新の強化学習トレーニング環境エンジンをリリースし、256 コア CPU の実行速度が直接 100 万フレーム/秒に達しました。ノートパソコンの i7-8750H でも 1 秒あたり 50,000 フレームに到達できます。

近年、深層強化学習 (Deep RL) は急速な進歩を遂げており、2015 年に Nature で発表された DQN から、後に Lee Sedol や Ke Jie などの世界チャンピオンを破った AlphaGo シリーズ、StarCraft を代表する AlphaStar や Dota 2 を代表する OpenAI Five などの複雑なゲームまで、多くの広範囲にわたる成果が生まれています。

アルゴリズムの進歩に加えて、最も重要なのは、近年の深層強化学習エージェントのトレーニング速度とスループットの大幅な向上です。当時、DQN では単純な Atari ゲームを実行するのに 1 週​​間以上かかりましたが、現在では RL システムは非常に高いスループットを実現し、複雑なゲームやシナリオでトレーニングできます。

現在、世界最大の強化学習トレーニング システムはすべて分散トレーニング方式を採用しています。たとえば、OpenAI Five はトレーニングに 50,000 個以上の CPU コアと数千個の GPU を使用しています。

オープンAIファイブ

同様に、AlphaStar の各トレーニング エージェントは、同時に実行されている 10,000 を超える StarCraft ゲーム エンジンに接続されています。ゲームインタラクションデータからの下流反復学習のトレーニングを高速化するには、これらのトレーニングシステムで RL 環境ゲームエンジンを実行するために大量の CPU リソースを使用する必要があります。

アルファスター

このような膨大なリソース消費により、大規模な強化学習トレーニング システムの敷居は非常に高くなります。学術界ではリソースが限られているため、大規模なゲームや複雑な RL 環境での研究を行うことは困難です。学術界だけでなく、産業界の研究機関でも、CPU レベルの使用率 10,000 コアを達成することはかなりの負担となります。

Sea AI Lab の研究者は、ゲーム エンジンを含む RL 環境が RL トレーニング システム全体の中で最も遅い部分であり、データ供給側での位置がシステム全体のスループットの上限を決定することに気づきました。

しかし、このコンポーネントは研究者から十分な注目を集めていません。現在、RL環境の並列実行に最も一般的に使用されている方法は、Pythonマルチプロセスを使用して単純なトレーニング環境の並列処理を実行するgym.vector_envです。インターフェースは変更されませんが、Pythonの制限により、高いスループットを実現するために、より多くのCPUリソースを使用して環境をシミュレートすることしかできません。

Sea AI Lab は、RL 環境のシミュレーション パフォーマンスと CPU 利用効率を向上させるために、高度に並列化された RL 環境エンジン ソリューション EnvPool を提供しています。

この RL 環境エンジンは、最下層の C++ スレッド プールを使用して複数の RL 環境インスタンスを非同期的に実行し、並列化の効率を大幅に向上させます。 EnvPool は、単一の NVIDIA DGX A100 上の 256 個の CPU コアを使用して、古典的な RL 環境シミュレーター Atari ゲームで 1 秒あたり 100 万フレームという驚異的な実行速度を実現します。

EnvPool システムの概要

このような高いスループットは、研究者が最もよく使用する実行エンジンである gym.vector_env の約 14 倍です。これは、同じハードウェア リソースを使用して、リソース利用効率を桁違いに向上できること、または逆に、EnvPool を使用すると、使用するリソースを桁違いに少なくして、非常に高いスループットを実現できることを意味します。

研究者にとって、EnvPool は便利で使いやすい Python インターフェースを提供します。下図に示すように、最も単純な同期実行モードのインターフェースは Gym API とまったく同じです。単一環境の gym/dm_env と比較すると、EnvPool は単一環境と対話するための元の API をバッチ対話に拡張す​​るだけです。

バッチ取得状態により、アルゴリズムは推論のためにデータを直接 GPU に送信することが容易になります。これにより、アルゴリズムの実装の難しさが大幅に軽減されます。環境から出力されるデータは、すでに GPU および TPU 処理に適したバッチ形式になっているため、GPU/TPU の並列効率をより有効に活用できます。

非同期モードでは、EnvPool は元の step 関数を send 関数と recv 関数に分割します。send 関数を呼び出してアクションを環境に送信した後、スレッドは現在の対話型環境が新しい状態を返すのを待つ必要はなく、スレッド プールで実行するために EnvPool に引き渡されます。

実行を完了した環境は、新しい状態をキューに入れて、recv 関数を通じてバッチで取得します。同期モードと比較して、非同期モードでは、環境の合計数と各インタラクションのバッチ サイズを指定する必要があります。

非同期方式では、アルゴリズムと環境は、お互いを待つことなく、時間的に効果的に重なり合うことができます。

Sea AI Lab の研究者は、TPU v3-8 仮想マシンと NVIDIA DGX-A100 モデルという 2 つの異なるマシン モデルで、EnvPool とその他の RL 環境実行エンジンのスループットを評価および比較しました。 TPU 仮想マシンには 96 個の CPU コアと 2 つの NUMA ノードがあり、NVIDIA DGX-A100 マシンには 256 個の CPU コアと 8 つの NUMA ノードがあります。

EnvPool と比較される RL 環境シミュレーション システムには、シンプルな Python for ループのマルチ環境実行、業界で最も一般的に使用されている gym.vector_env、そして昨年リリースされたばかりで研究者に知られている最速の環境シミュレーション システムである Sample Factory などがあります。

DGX-100

TPU-VM

結果から、EnvPool はさまざまなモデルやさまざまなワーカー数で大きな利点があることがわかります。特に、同じ 256 個の CPU コアを使用する NVIDIA DGX A100 では、EnvPool のパフォーマンスは市場で最も一般的に使用されている gym.vector_env の 13.3 倍で、驚異的な 1 秒あたり 100 万フレームで実行されます。このスループットでは、10^9 フレーム (10 億) のデータを生成するのにわずか 17 分しかかかりません。

EnvPool は、比較的リソースが少ない設定でも非常に優れたパフォーマンスを発揮します。12 コアで gym.vector_env を使用して Atari を実行すると、18,000 FPS しか達成できませんが、EnvPool は CPU をより効率的に使用して 50,000 FPS を達成できます。これは、広く使用されているこのベンチマークの約 3 倍です。

プロジェクト紹介

R&DチームSea AI Lab（SAIL）はシンガポールのSEAグループに所属しています。2020年末に設立され、Yan Shuicheng氏が率いています。最先端の画期的な基礎研究に重点を置いています。

現在、EnvPool は GitHub でオープンソース化されています。

https://github.com/sail-sg/envpool

これは内部イテレーションの 2 番目のバージョンです。最初のバージョンの内部実装と比較して、オープンソースの 2 番目のバージョンは開発者 API の簡素化に重点を置いており、コミュニティ開発者がさまざまなゲームや RL トレーニング環境に簡単にアクセスできるようになります。

C++ エンジンに接続する場合、開発者は単一の環境の実行ロジックを定義するだけで済みます。EnvPool は分散実行を担当し、バッチ インタラクション API を提供するため、新しい RL 環境やゲーム エンジンに簡単に接続して、すぐに並列アクセラレーションを実現できます。

EnvPool は現在、集中的に開発中です。今後のコア アップデートには、より多くの RL 環境 (最近発表された継続的なアクション スペース用のフリーフォーオール Mujoco を含む) へのアクセスが含まれ、DeepMind の Acme へのアクセスや、EnvPool の同じ最初の作成者による人気の RL トレーニング ライブラリ Tianshou へのアクセスなど、ユーザーが EnvPool を使用して既存のオープン ソース RL トレーニング ライブラリを簡単に高速化できる例が提供されます。

さらに、EnvPool の成果は NVIDIA GTC 2021 カンファレンスでも発表されるよう招待されました。