JaxMARL: Multi-Agent RL Environments in JAX

要約

ベンチマークは、機械学習アルゴリズムの開発において重要な役割を果たします。
たとえば、強化学習 (RL) の研究は、利用可能な環境やベンチマークに大きな影響を受けてきました。
ただし、RL 環境は伝統的に CPU 上で実行されるため、一般的なアカデミック コンピューティングではスケーラビリティが制限されます。
JAX の最近の進歩により、これらの計算上のハードルを克服するためにハードウェア アクセラレーションを広く使用できるようになり、大規模な並列 RL トレーニング パイプラインと環境が可能になりました。
これは、マルチエージェント強化学習 (MARL) 研究に特に役立ちます。
まず第一に、環境の各ステップで複数のエージェントを考慮する必要があり、計算負荷が増加します。第二に、非定常性、分散部分可観測性、またはその他の MARL の課題によりサンプルの複雑さが増加します。
このペーパーでは、使いやすさと GPU 対応の効率性を組み合わせた初のオープンソース コード ベースである JaxMARL について紹介します。また、一般的に使用される多数の MARL 環境と一般的なベースライン アルゴリズムをサポートしています。
実時間を考慮すると、JAX ベースのトレーニング パイプラインは実行あたり、既存のアプローチよりも最大 12500 倍高速であることが実験で示されています。
これにより、効率的かつ徹底した評価が可能となり、現場の評価危機を軽減できる可能性があります。
また、人気のある StarCraft Multi-Agent Challenge のベクトル化された簡易バージョンである SMAX の紹介とベンチマークも行います。これにより、StarCraft II ゲーム エンジンを実行する必要がなくなります。
これにより、GPU アクセラレーションが可能になるだけでなく、より柔軟な MARL 環境が提供され、セルフプレイ、メタ学習、および MARL のその他の将来のアプリケーションの可能性が解き放たれます。
コードは https://github.com/flairox/jaxmarl で提供されています。

要約(オリジナル)

Benchmarks play an important role in the development of machine learning algorithms. For example, research in reinforcement learning (RL) has been heavily influenced by available environments and benchmarks. However, RL environments are traditionally run on the CPU, limiting their scalability with typical academic compute. Recent advancements in JAX have enabled the wider use of hardware acceleration to overcome these computational hurdles, enabling massively parallel RL training pipelines and environments. This is particularly useful for multi-agent reinforcement learning (MARL) research. First of all, multiple agents must be considered at each environment step, adding computational burden, and secondly, the sample complexity is increased due to non-stationarity, decentralised partial observability, or other MARL challenges. In this paper, we present JaxMARL, the first open-source code base that combines ease-of-use with GPU enabled efficiency, and supports a large number of commonly used MARL environments as well as popular baseline algorithms. When considering wall clock time, our experiments show that per-run our JAX-based training pipeline is up to 12500x faster than existing approaches. This enables efficient and thorough evaluations, with the potential to alleviate the evaluation crisis of the field. We also introduce and benchmark SMAX, a vectorised, simplified version of the popular StarCraft Multi-Agent Challenge, which removes the need to run the StarCraft II game engine. This not only enables GPU acceleration, but also provides a more flexible MARL environment, unlocking the potential for self-play, meta-learning, and other future applications in MARL. We provide code at https://github.com/flairox/jaxmarl.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google