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提供了代码。

本文旨在解决强化学习环境在CPU上运行时计算负担大、难以扩展的问题，提出了一种基于JAX的GPU加速的多智能体强化学习训练流水线，支持多种常用的多智能体强化学习环境和基准算法。

本文提出的JaxMARL是第一个结合易用性和GPU加速效率的开源代码库，可以显著提高多智能体强化学习的训练效率，有助于解决该领域的评估危机。

本文的亮点包括：1. 提出了JaxMARL，是第一个结合易用性和GPU加速效率的开源代码库；2. 实验结果表明，相较于现有方法，JaxMARL的训练效率可以提高12500倍；3. 引入并评估了SMAX，一种基于StarCraft Multi-Agent Challenge的简化版本，可以避免运行StarCraft II游戏引擎，提供了更加灵活的多智能体强化学习环境，有助于未来的自我博弈、元学习等应用；4. 开源了代码，提供了多种常用的多智能体强化学习环境和基准算法。

最近在该领域的相关研究包括：1. OpenAI提出的Spinning Up，是一个易用性和可扩展性俱佳的强化学习代码库；2. Google提出的TensorFlow Agents，是一个基于TensorFlow的多智能体强化学习框架；3. Facebook提出的Horizon，是一个基于PyTorch的强化学习平台。