Mava - 基于JAX的高效多智能体强化学习框架

基于JAX的分布式多智能体强化学习

欢迎使用Mava！🦁

安装 | 快速入门

Mava提供了简化的代码，用于在多智能体强化学习（MARL）中快速迭代想法，并提供了JAX实现的MARL算法，可以通过JAX的pmap轻松实现跨设备并行化。Mava是源自InstaDeep研究团队的项目。

如果您想加入我们的努力，请随时联系、提出问题或阅读我们的贡献指南（或者只需点个星标🌟以了解最新进展）！

概述 🦜

Mava目前为MARL研究提供以下构建块：

🥑 MARL算法实现：实现了遵循集中式训练与分散式执行（CTDE）和分散式训练与分散式执行（DTDE）MARL范式的多智能体PPO系统。
🍬 环境包装器：用于将Jumanji环境映射到与Mava兼容的环境的示例包装器。目前，我们支持[机器人仓库][jumanji_rware]和[基于等级的觅食][jumanji_lbf]，并计划在不久的将来支持更多环境。我们最近还添加了对[JaxMARL][jaxmarl]中SMAX环境的支持。
🎓 教育材料：[快速入门笔记本][quickstart]，展示如何使用Mava并突出基于JAX的MARL的附加价值。
🧪 统计稳健评估：Mava原生支持记录到符合[Gorsane et al. (2022)][toward_standard_eval]建议标准的json文件。这使得可以使用[MARL-eval][marl_eval]库中的工具轻松进行下游实验绘图和聚合。

性能和速度 🚀

SMAX

为了比较Mava的稳定性与其他基于JAX的基准算法，我们在广泛的[SMAX][smax]任务上训练了Mava的循环IPPO和MAPPO系统。在所有情况下，我们没有重新运行基准，而是直接采用JaxMARL技术报告中的最终胜率结果。有关SMAX实验的完整结果，请参见以下页面。

Mava循环IPPO和MAPPO在3s5z、6h_vs_8z和3s5z_vs_3s6z SMAX任务上的表现。

机器人仓库

以下所有实验均在配备8GB内存的NVIDIA Quadro RTX 4000 GPU上进行。

为了展示端到端JAX-based MARL系统和JAX-based环境的实用性，我们比较了Mava与[EPyMARL][epymarl]在简单的[机器人仓库][rware]（RWARE）任务上的训练总时间。我们的目标是说明使用端到端Jax-based系统可能带来的速度提升，而不一定致力于实现最佳性能。对于EPyMARL，我们使用Papoudakis et al. (2020)推荐的超参数，而对于Mava，我们进行了基本的网格搜索。在这两种情况下，系统都使用16个向量化环境训练至2000万总环境步数。

Mava前馈MAPPO在tiny-2ag、tiny-4ag和small-4ag RWARE任务上的表现。

📌 关于收敛性能差异的重要说明

为了从基于JAX的系统提供的挂钟速度提升中受益，环境也需要用JAX编写。正因如此，Mava并不使用与EPyMARL完全相同版本的RWARE环境，而是使用[Jumanji][jumanji_rware]中的JAX实现版本RWARE，名为RobotWarehouse。底层环境逻辑的一个显著区别是，RobotWarehouse不会尝试解决智能体碰撞问题，而是在智能体发生碰撞时终止回合。在我们的实验中，这似乎使环境变得更具挑战性。因此，我们展示了Mava在Jumanji上有无碰撞终止条件的表现，图例中用"w/o collision"表示。更详细的讨论请参见以下页面。

基于等级的觅食

Mava还支持[Jumanji][jumanji_lbf]的LBF。我们评估了Mava的循环MAPPO系统在LBF上的表现，并与[EPyMARL][epymarl]进行了比较（我们为EPyMARL使用了原始的LBF），在2个和4个智能体的设置下训练了2000万个时间步。两个系统都使用16个向量化环境进行训练。对于EPyMARL系统，我们使用NVIDIA A100 GPU，对于Mava系统，我们使用具有4GB内存的GeForce RTX 3050笔记本GPU。为了展示Mava如何适应不同的硬件，我们还在TPU v3-8上训练了Mava系统。我们计划很快发布Mava所有算法在各种LBF场景下的全面性能基准。

Mava循环MAPPO在2s-8x8-2p-2f-coop和15x15-4p-3fz基于等级的觅食任务上的表现。

🧨 使用向量化环境的每秒步数实验

此外，我们通过展示随着并行环境数量增加的每秒步数来说明Mava的速度。这些每秒步数扩展图是使用标准笔记本GPU计算的，具体是使用6GB内存的RTX-3060 GPU。

Mava每秒步数随向量化环境增加的扩展和2000万环境步骤的总训练运行时间。

代码理念 🧘

Mava中的当前代码改编自[PureJaxRL][purejaxrl]，它提供了高质量的单文件实现，具有研究友好的特性。反过来，PureJaxRL的灵感来自[CleanRL][cleanrl]的代码理念。顺着这种易用和易懂的强化学习代码库的脉络，Mava并非设计为模块化库，也不意味着被导入使用。我们的仓库专注于实现的简单性和清晰性，同时利用JAX提供的优势，如pmap和vmap，使其成为研究人员和实践者构建的优秀资源。

安装 🎬

目前Mava并不意味着作为库安装，而是作为一个研究工具使用。

你可以通过克隆仓库并使用pip安装来使用Mava，如下所示：

git clone https://github.com/instadeepai/mava.git
cd mava
pip install -e .

我们已在Python 3.11和3.12上测试过Mava，但早期版本也可能适用。注意，由于JAX的安装取决于您的硬件加速器，我们建议用户明确安装正确的JAX版本（请参阅官方安装指南）。有关更详细的安装指南，包括Docker构建和虚拟环境，请参阅我们的详细安装指南。

快速开始 ⚡

要开始训练您的第一个Mava系统，只需运行其中一个系统文件。例如：

python mava/systems/ff_ippo.py

Mava使用Hydra进行配置管理。要查看我们的默认系统配置，请查看mava/configs/目录。Hydra的一个优点是可以在配置yaml文件中设置配置，或者在运行时从终端动态覆盖配置。例如，要在LBF环境上运行系统，上述代码可以简单地修改如下：

python mava/systems/ff_ippo.py env=lbf

也可以通过在终端进行以下配置更新来运行不同的场景：

python mava/systems/ff_ippo.py env=rware env/scenario=tiny-4ag

此外，我们还提供了一个[快速入门notebook][quickstart]，可用于快速创建和训练您的第一个多智能体系统。

高级用法 👽

Mava可用于各种高级系统。例如，我们演示如何将一个PPO系统的经验数据记录到FlashbaxVault中。然后，这个vault可以轻松集成到离线MARL系统中，如OG-MARL中的系统。更多信息请参见高级README。

贡献 🤝

请阅读我们的贡献文档，了解如何提交拉取请求、我们的贡献者许可协议和社区指南的详细信息。

路线图 🛤️

我们计划按以下增量迭代扩展Mava：

🌴 支持更多环境。
🔁 更强大的循环系统。
🌳 支持非基于JAX的环境。
🦾 支持离线策略算法。
🎛 连续动作空间环境和算法。

请继续关注我们开发的下一阶段！

TensorFlow 2 Mava:

Mava最初是用Tensorflow 2编写的。对基于TF2的框架和系统的支持现已完全废弃。如果您仍想使用它，请安装Mava的v0.1.3版本（即pip install id-mava==0.1.3）。

另请参阅 🔎

InstaDeep的JAX中的MARL生态系统。 特别是，我们建议用户查看以下姊妹仓库：

🔌 OG-MARL：JAX中离线MARL的数据集和基准。
🌴 Jumanji：JAX中多样化的可扩展强化学习环境套件。
😎 Matrax：JAX中的矩阵游戏集合。
⚡ Flashbax：JAX中加速的重放缓冲区。
📈 [MARL-eval][marl_eval]：MARL的标准化实验数据聚合和可视化。

相关。 其他与JAX中加速MARL相关的库。

🦊 JaxMARL：JAX中加速的MARL环境和基准。
🌀 [DeepMind Anakin][anakin_paper]：用于大规模训练RL智能体的Anakin podracing架构。
♟️ Pgx：经典棋类游戏（如国际象棋、围棋和将棋）的JAX实现。
🔼 Minimax：RL自动课程基准的JAX实现。

引用Mava 📚

如果您在工作中使用Mava，请引用随附的[技术报告][Paper]：

@article{dekock2023mava,
    title={Mava: a research library for distributed multi-agent reinforcement learning in JAX},
    author={Ruan de Kock and Omayma Mahjoub and Sasha Abramowitz and Wiem Khlifi and Callum Rhys Tilbury
    and Claude Formanek and Andries P. Smit and Arnu Pretorius},
    year={2023},
    journal={arXiv preprint arXiv:2107.01460},
    url={https://arxiv.org/pdf/2107.01460.pdf},
}

致谢 🙏

我们要感谢所有为Mava之前的TF版本做出贡献的作者：Kale-ab Tessera、St John Grimbly、Kevin Eloff、Siphelele Danisa、Lawrence Francis、Jonathan Shock、Herman Kamper、Willie Brink、Herman Engelbrecht、Alexandre Laterre、Karim Beguir。他们的贡献可以在我们的TF技术报告中找到。

Mava的开发得到了Google TPU Research Cloud (TRC)提供的Cloud TPU支持 🌤。