JaxMARL: Multi-Agent RL Environments in JAX

基准测试在机器学习算法的开发中扮演着重要的角色。例如，强化学习（RL）领域的研究受到可用环境和基准测试的严重影响。然而，RL环境通常在CPU上运行，限制了它们在典型的学术计算中的可扩展性。JAX的最新进展使得更广泛地使用硬件加速以克服这些计算障碍，实现了大规模并行RL训练管道和环境，这对于多智能体强化学习（MARL）研究尤其有用。首先，每个环境步骤必须考虑多个智能体，增加了计算负担；其次，由于非稳态、分散的部分可观察性或其他MARL挑战，样本复杂性增加。在本文中，我们提出了JaxMARL，这是第一个将易用性与GPU启用的效率相结合的开源代码库，支持大量常用的MARL环境以及流行的基准算法。在考虑挂钟时间时，我们的实验表明，与现有方法相比，我们基于JAX的训练管道每次运行的速度高达12500倍。这使得评估变得高效和全面，有可能缓解该领域的评估危机。我们还介绍并基准测试了SMAX，这是流行的星际争霸多智能体挑战的矢量化简化版本，消除了运行星际争霸II游戏引擎的必要性。这不仅实现了GPU加速，而且提供了一个更灵活的MARL环境，释放了自我对战、元学习和其他未来MARL应用的潜力。我们在https://github.com/flairox/jaxmarl上提供了代码。

本文旨在通过JaxMARL解决强化学习环境在CPU上运行的限制，提高多智能体强化学习的可扩展性和效率。

本文提出了JaxMARL框架，结合GPU加速和大量常用多智能体强化学习环境，支持多种基准算法，实现高效、易用的训练流程，并介绍并测试了SMAX，这是一种基于StarCraft Multi-Agent Challenge的简化多智能体强化学习环境，可以在GPU上运行。

JaxMARL框架可以使多智能体强化学习的训练速度提高至现有方法的12500倍，有望解决该领域的评估危机。SMAX提供了一个更灵活的多智能体强化学习环境，可以用于自我博弈、元学习等未来应用。

最近的相关研究包括基于CPU的多智能体强化学习环境和基于GPU的强化学习框架，如TensorFlow和PyTorch。