[論文レビュー] EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine

Jiayi Weng, Min Lin+10 more|arXiv (Cornell University)|Jun 21, 2022

Reinforcement Learning in Robotics被引用数 14

ひとこと要約

EnvPool は、C++ で構築されたスレッドプール実行エンジンを用いて環境シミュレーションを最適化することで、強化学習（RL）の学習を高速化する、並列性に優れた強化学習環境実行エンジンです。MuJoCo では最大 300 万物理ステップ/秒、Atari では最大 100 万フレーム/秒を達成し、ベースラインの Python 実装と比べて最大 19.2 倍の高速化を実現しました。既存の RL ライブラリとシームレスに統合可能で、ラップトップでも Atari Pong や MuJoCo Ant の学習を 5 分未塔で完了できます。

ABSTRACT

There has been significant progress in developing reinforcement learning (RL) training systems. Past works such as IMPALA, Apex, Seed RL, Sample Factory, and others, aim to improve the system's overall throughput. In this paper, we aim to address a common bottleneck in the RL training system, i.e., parallel environment execution, which is often the slowest part of the whole system but receives little attention. With a curated design for paralleling RL environments, we have improved the RL environment simulation speed across different hardware setups, ranging from a laptop and a modest workstation, to a high-end machine such as NVIDIA DGX-A100. On a high-end machine, EnvPool achieves one million frames per second for the environment execution on Atari environments and three million frames per second on MuJoCo environments. When running EnvPool on a laptop, the speed is 2.8x that of the Python subprocess. Moreover, great compatibility with existing RL training libraries has been demonstrated in the open-sourced community, including CleanRL, rl_games, DeepMind Acme, etc. Finally, EnvPool allows researchers to iterate their ideas at a much faster pace and has great potential to become the de facto RL environment execution engine. Example runs show that it only takes five minutes to train agents to play Atari Pong and MuJoCo Ant on a laptop. EnvPool is open-sourced at https://github.com/sail-sg/envpool.

研究の動機と目的

• 強化学習学習システムにおける並列環境実行の非最適化ボトルネックを解消すること。
• ラップトップからハイエンドの DGX-A100 システムまで、多様なハードウェア環境でのシミュレーションスループットを向上させること。
• OpenAI Gym や dm_env API のドロップインリプレースメントを提供することで、最小限のコード変更で研究者がエージェントの学習を高速化できるようにすること。
• 同期実行と非同期実行の両モードをサポートし、ハードウェア利用効率を最大化すること。
• CleanRL や rl_games、DeepMind Acme といった主要な RL 学習ライブラリとの広範な互換性を確保すること。

提案手法

• 複数の RL 環境を並列に実行する C++ ベースのスレッドプール実行エンジンを実装し、Python の GIL や I/O オーバーヘッドを最小限に抑える。
• C++ 側での Python ラッパー最適化により、環境とのインタラクションや観測処理のレイテンシを低減する。
• 標準の OpenAI Gym および dm_env 互換 Python API を通じて、シームレスな統合を可能にする。
• 同期実行と非同期実行の両モードをサポートし、後者はより高いスループットとデータの陳腐化の低減を実現する。
• 開発者がカスタム C++ ベースの RL 環境を最小限の作業で EnvPool に統合できるモジュラーなアーキテクチャを設計する。
• C++ と Python 間での効率的なメモリ管理とゼロコピーデータ転送を活用し、シリアル化のオーバーヘッドを低減する。

実験結果

リサーチクエスチョン

• RQ1 高度に並列化された環境実行エンジンは、RL 学習におけるシミュレーションレイテンシを顕著に低減できるか？
• RQ2 EnvPool のパフォーマンスは、ラップトップからハイエンドの DGX システムまで、さまざまなハードウェア構成でどのようにスケーリングするか？
• RQ3 標準的な RL ベンチマークにおいて、サンプル効率を維持したまま、EnvPool はトレーニングスループットをどの程度向上できるか？
• RQ4 エージェントコードの変更なしに、EnvPool を既存の RL 学習パイプラインに効率的に統合できるか？
• RQ5 環境バッチ処理において、非同期実行は同期実行と比較してどの程度のパフォーマンス向上を達成できるか？

主な発見

• NVIDIA DGX-A100（256 CPU コア）では、Atari で 100 万フレーム/秒、MuJoCo で 300 万物理ステップ/秒を達成し、それぞれベースラインの Python 実装と比べて 14.9 倍および 19.2 倍の高速化を実現した。
• 12 CPU コア搭載のラップトップでも、標準の Python subprocess ベースの環境実行と比べて 2.8 倍の高速化を達成した。
• ラップトップ上での Atari Pong や MuJoCo Ant におけるエンドツーエンドの学習が 5 分未塔で完了し、迅速な研究反復の実用性を示した。
• MuJoCo HalfCheetah-v3 における同じ Acme PPO 学習設定において、EnvPool は DummyVecEnv よりも壁時計時間（wall time）を半分以下に削減しながら、同等の最終エピソードリターンを維持した。
• トレーニング曲線から、num_envs と batch_size の積を一定に保った場合、並列環境数（num_envs）を増やすことでトレーニング時間が顕著に短縮され、サンプル効率に悪影響を与えないことが示された。
• CleanRL や rl_games、Ray、DeepMind Acme といった主要な RL フレームワークとの高い互換性を示し、プラグアンドプレイ統合を可能にした。