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.

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