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.

翻译：基准测试在机器学习算法发展中扮演重要角色。例如，强化学习（RL）研究深受可用环境与基准测试的影响。然而，传统RL环境在CPU上运行，受限于典型学术计算资源，难以扩展。JAX的最新进展使得硬件加速技术得到更广泛应用，能够克服这些计算瓶颈，实现大规模并行化RL训练流程与环境。这对多智能体强化学习（MARL）研究尤为关键：首先，每个环境步骤需考虑多个智能体，增加了计算负担；其次，由于非平稳性、分布式部分可观测性或其它MARL挑战，样本复杂度显著提升。本文提出JaxMARL，这是首个兼顾易用性与GPU高效性的开源代码库，支持大量常用MARL环境及主流基线算法。实验表明，从挂钟时间衡量，基于JAX的训练流程单次运行速度比现有方法快达12500倍。这实现了高效且彻底的评估，有望缓解该领域的评估危机。我们还引入并基准测试了SMAX——经典《星际争霸多智能体挑战》（StarCraft Multi-Agent Challenge）的矢量化简化版本，无需运行《星际争霸II》游戏引擎。这不仅实现了GPU加速，还提供了更灵活的MARL环境，为自我对弈、元学习及其他未来MARL应用释放潜力。代码开源地址：https://github.com/flairox/jaxmarl。