sail-sg/envpool

何を解決するか

EnvPoolは、遅い環境シミュレーションによって引き起こされる強化学習（RL）のパフォーマンスボトルネックに対処します。従来のPythonベースのベクトル環境は、サブプロセスによる高コストなオーバーヘッドを抱えており、経験収集のスループットが制限されています。EnvPoolは高性能なC++ベースの実行エンジンを提供し、標準的なPython実装よりもはるかに高いフレーム毎秒（FPS）で環境を大規模並列化できます。

動作方法

EnvPoolはC++で実装されており、pybind11を使用してPythonインターフェースを提供しています。環境プールを管理し、スレッドプールを利用して並列実行を行います。主に2つの相互作用モードをサポートしています：

• 同期API：Gymnasiumやdm_envのAPIに近い設計で、stepおよびreset関数はデフォルトで複数環境のバッチを扱います。
• 非同期API：step関数をsendとrecvの操作に分割し、ユーザーがアクションを入力し、状態を取得する際にバッチ全体の完了を待たずに処理を進められるようにします。これにより、さらに効率が向上します。

XLAをJAX jit関数経由で統合し、Stable-Baselines3、Tianshou、CleanRLなどの主要なRLライブラリと互換性があります。

対象ユーザー

経験データを迅速に大量に収集してエージェントを訓練する必要がある強化学習の研究者や開発者向けです。特にAtari、MuJoCo、あるいはさまざまなロボットシミュレータなど、計算コストの高い環境を使用している方々に適しています。

特徴

• 極めて高いスループット：高スペックハードウェア（DGX-A100）で、Atariで約100万FPS、MuJoCoで約300万ステップ/秒を達成可能。
• 広範な環境サポート：Atari、MuJoCo、Box2D、DeepMind Control Suite、ViZDoomなど多数をサポート。
• クロスプラットフォーム：Linux、macOS、Windowsに対応。
• 柔軟なAPI：Gymnasiumおよびdm_envと互換性あり。同期・非同期の両実行モードをサポート。
• 組み込みレンダリング：バッチ処理されたrgb_array出力とhuman表示モードをサポート。
• C++拡張性：カスタムC++環境を統合するための開発者APIを提供。