JaxMARL: JAXにおけるマルチエージェント強化学習環境

要旨

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

English

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.

JaxMARL: JAXにおけるマルチエージェント強化学習環境

JaxMARL: Multi-Agent RL Environments in JAX

要旨

Support