Mava

<p align="center"> <a href="docs/images/mava_logos/mava_full_logo.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/275c62c6-7a82-4d27-8920-948179affba1.png" alt="Mava标志" width="50%"/> </a> </p> <h2 align="center"> <p>基于JAX的分布式多智能体强化学习</p> </h2> <div align="center"> <a href="https://www.python.org/doc/versions/"> <img src="https://img.shields.io/badge/python-3.12-blue" alt="Python版本"> </a> <a href="https://www.python.org/doc/versions/"> <img src="https://img.shields.io/badge/python-3.11-blue" alt="Python版本"> </a> <a href="https://github.com/instadeepai/Mava/blob/main/LICENSE"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/c01d466e-0261-416c-a360-5f42b3ab251b.svg" alt="许可证" /> </a> <a href="https://github.com/psf/black"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/9aec10fe-2fdb-44ff-b213-9d8225965c00.svg" alt="代码风格" /> </a> <a href="http://mypy-lang.org/"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/4497b5c1-7214-4bdf-9593-05fc8b6201e6.svg" alt="MyPy" /> </a> <a href="https://arxiv.org/pdf/2107.01460.pdf"> <img src="https://img.shields.io/badge/预印本-ArXiv-red" alt="ArXiv"> </a> <a href="https://colab.research.google.com/github/instadeepai/Mava/blob/develop/examples/Quickstart.ipynb" target="_parent"><img src="https://yellow-cdn.veclightyear.com/835a84d5/9f97b602-0cf4-4d22-9b24-841e7897bb2a.svg" alt="在Colab中打开"/></a> </div>

欢迎使用Mava！🦁

安装 | 快速入门

</div>

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实验的完整结果，请参见以下页面。

<p align="center"> <a href="docs/images/smax_results/legend.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/cc27ae2c-7516-4d58-88a6-40a45e52b3ad.png" alt="图例" width="40%"/> </a> </p> <p align="center"> <a href="docs/images/smax_results/3s5z.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/d6ee3a92-d4a5-4ad4-bd34-37da3a6c3fbb.png" alt="Mava ff mappo tiny 2ag" width="30%" style="display:inline-block; margin-right: 10px;"/> </a> <a href="docs/images/smax_results/6h_vs_8z.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/17d7bfc0-83b1-4392-b5e6-7e56d8ab73bc.png" alt="Mava ff mappo tiny 4ag" width="30%" style="display:inline-block; margin-right: 10px;"/> </a> <a href="docs/images/smax_results/3s5z_vs_3s6z.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/920078f6-cf47-45f9-b8dd-791602de31ff.png" alt="Mava ff mappo small 4ag" width="30%" style="display:inline-block; margin-right: 10px;"/> </a> <br> <div style="text-align:center; margin-top: 10px;"> Mava循环IPPO和MAPPO在<code>3s5z</code>、<code>6h_vs_8z</code>和<code>3s5z_vs_3s6z</code> SMAX任务上的表现。</div> </p>

机器人仓库

以下所有实验均在配备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万总环境步数。

<p align="center"> <a href="docs/images/rware_results/ff_mappo/main_readme/legend.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/81edffbb-3fb3-460a-896e-6605c351b00e.png" alt="图例" width="80%"/> </a> </p> <p align="center"> <a href="docs/images/rware_results/ff_mappo/main_readme/tiny-2ag-1.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/718c0a7d-713f-4fd3-be3e-17166ac60777.png" alt="Mava ff mappo tiny 2ag" width="30%" style="display:inline-block; margin-right: 10px;"/> </a> <a href="docs/images/rware_results/ff_mappo/main_readme/tiny-4ag-1.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/3d7efe16-f40a-45b1-a0f1-dfb708dac2d4.png" alt="Mava ff mappo tiny 4ag" width="31%" style="display:inline-block; margin-right: 10px;"/> </a> <a href="docs/images/rware_results/ff_mappo/main_readme/small-4ag-1.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/56142b3a-a630-44cc-804d-a7ae59ed8670.png" alt="Mava ff mappo small 4ag" width="31%" style="display:inline-block; margin-right: 10px;"/> </a> <br> <div style="text-align:center; margin-top: 10px;"> Mava前馈MAPPO在<code>tiny-2ag</code>、<code>tiny-4ag</code>和<code>small-4ag</code> RWARE任务上的表现。</div> </p>

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

为了从基于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场景下的全面性能基准。

<p align="center"> <a href="docs/images/lbf_results/legend_rec_mappo.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/42ad42fb-83cd-4b87-add8-90bdd4abec55.png" alt="图例" width="60%"/> </a> </p> <p align="center"> <a href="docs/images/lbf_results/2s-8x8-2p-2f-coop_rec_mappo.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/bc79f215-5297-40c9-aa7a-fd3d7a4881fd.png" alt="Mava ff mappo tiny 2ag" width="30%" style="display:inline-block; margin-right: 10px;"/> </a> <a href="docs/images/lbf_results/15x15-4p-3f_rec_mappo.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/0b5e5ec4-dbca-4dc4-94a9-619637da7868.png" alt="Mava ff mappo small 4ag" width="30%" style="display:inline-block; margin-right: 10px;"/> </a> <br> <div style="text-align:center; margin-top: 10px;"> Mava循环MAPPO在<code>2s-8x8-2p-2f-coop</code>和<code>15x15-4p-3fz</code>基于等级的觅食任务上的表现。</div> </p>

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

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

<p align="center"> <a href="docs/images/speed_results/mava_sps_results.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/dbeed90f-6dcf-47d4-8417-04ab7aa33f51.png" alt="Mava sps" width="55%"/> </a> <a href="docs/images/speed_results/ff_mappo_speed_comparison.png"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/7237eff5-39ec-409a-b02b-ba920718847b.png" alt="Mava ff mappo速度比较" width="39.33%" style="display:inline-block; margin-right: 10px;"/> </a> <br> <div style="text-align:center; margin-top: 10px;"> Mava每秒步数随向量化环境增加的扩展和2000万环境步骤的总训练运行时间。</div> </p>

代码理念 🧘

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支持 🌤。