JaxMARL

安装 | 快速开始 | 环境 | 算法 | 引用

基于JAX的多智能体强化学习库

JaxMARL（基于JAX的多智能体强化学习库）结合了易用性与GPU加速的高效性能，支持多种常用的MARL环境及主流基线算法。我们的目标是提供一个统一库，可在广泛任务中对MARL方法进行系统评估。我们还引入了SMAX环境，这是经典StarCraft多智能体挑战的向量化简化版本，无需运行StarCraft II游戏引擎。

更多细节请参阅我们的博客文章或Colab教程。

环境 🌍

基线算法 🦉

我们遵循CleanRL的单文件实现理念，所有算法实现位于baselines目录。使用Hydra管理配置文件，各算法README包含具体说明。大部分文件包含wandb日志记录代码（默认关闭，可通过配置启用）。

安装 🧗

环境安装 - 安装前请确保已根据硬件加速器配置正确的JAX环境（已测试至JAX 0.4.36版本）。可通过PyPi直接安装JaxMARL环境：

pip install jaxmarl 

算法运行 - 如需运行算法，请按以下步骤安装源码：

1. 克隆仓库：

   git clone https://github.com/FLAIROx/JaxMARL.git && cd JaxMARL
   

2. 安装依赖：

   pip install -e .[algs]
   export PYTHONPATH=./JaxMARL:$PYTHONPATH
   

3. 推荐使用我们的Dockerfile快速启动（具体用法见下文）

开发模式 - 如需运行测试套件，克隆仓库后通过以下命令安装开发依赖： pip install -e .[dev]

快速入门 🚀

我们借鉴了PettingZoo（多智能体强化学习环境库）和Gymnax（纯JAX环境库）的接口设计。可通过Colab教程体验智能体训练，更多示例脚本见教程目录。

基础JaxMARL API 使用示例 🖥️

动作、观测、奖励和终止信号均以智能体名称为键的字典形式传递，支持不同动作/观测空间。终止字典包含额外的"__all__"键表示全局终止状态。我们采用并行结构，每个时间步所有智能体同步执行动作。对于异步游戏（如Hanabi），非行动智能体将接收空动作。

import jax
from jaxmarl import make

key = jax.random.PRNGKey(0)
key, key_reset, key_act, key_step = jax.random.split(key, 4)

# 初始化环境
env = make('MPE_simple_world_comm_v3')

重置环境。

obs, state = env.reset(key_reset)

采样随机动作。

key_act = jax.random.split(key_act, env.num_agents) actions = {agent: env.action_space(agent).sample(key_act[i]) for i, agent in enumerate(env.agents)}

执行步进转换。

obs, state, reward, done, infos = env.step(key_step, state, actions)

### Dockerfile（容器配置文件） 🐋
为帮助实验快速部署，我们提供了[Dockerfile](https://github.com/FLAIROx/JaxMARL/blob/main/Dockerfile)及其配套的[Makefile](https://github.com/FLAIROx/JaxMARL/blob/main/Makefile)。在安装Docker和[Nvidia容器工具包](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/index.html)后，可通过以下命令构建容器：

make build

构建完成后可通过以下命令运行容器：

make run

## 贡献指南 🔨
欢迎贡献！请查看我们的[贡献指南](https://github.com/FLAIROx/JaxMARL/blob/main/CONTRIBUTING.md)了解如何添加环境/算法或提交缺陷报告。如果您正在寻找项目方向，路线图下也列出了一些开发建议 :)

<h2 name="cite" id="cite">引用JaxMARL 📜 </h2>
若您在研究中使用了JaxMARL，请通过以下方式引用：

``` bibtex
@inproceedings{
    flair2024jaxmarl,
    title={JaxMARL: Multi-Agent RL Environments and Algorithms in JAX},
    author={Alexander Rutherford and Benjamin Ellis and Matteo Gallici and Jonathan Cook and Andrei Lupu and Gar{\dh}ar Ingvarsson and Timon Willi and Ravi Hammond and Akbir Khan and Christian Schroeder de Witt and Alexandra Souly and Saptarashmi Bandyopadhyay and Mikayel Samvelyan and Minqi Jiang and Robert Tjarko Lange and Shimon Whiteson and Bruno Lacerda and Nick Hawes and Tim Rockt{\"a}schel and Chris Lu and Jakob Nicolaus Foerster},
    booktitle={The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track},
    year={2024},
}

相关项目 🙌

本项目受到以下库的启发，推荐您查看这些优秀作品：

原生JAX算法：

• Mava：JAX实现的主流MARL算法
• PureJaxRL：PPO的JAX实现，展示端到端JAX强化学习训练

原生JAX环境：

• Gymnax：包含经典控制、bsuite和MinAtar任务的经典RL实现
• Jumanji：涵盖简单游戏到NP难组合问题的多样化环境
• Pgx：JAX实现的国际象棋、围棋、将棋等经典棋盘游戏
• Brax：用JAX编写的全微分物理引擎，支持连续控制任务
• XLand-MiniGrid：受XLand和MiniGrid启发的元强化学习网格世界
• Craftax：JAX实现的Crafter+NetHack游戏

JaxMARL 快速上手指南

环境准备

系统要求

• Python ≥ 3.9
• JAX ≥ 0.4.36（需根据硬件加速器安装对应版本）
• 推荐使用 Linux 系统（支持 GPU 加速）

前置依赖

# 安装 JAX（推荐使用清华镜像加速）
pip install --upgrade pip
pip install jax --index-url https://pypi.tuna.tsinghua.edu.cn/simple

# 安装基础工具（如使用 Docker 可跳过）
sudo apt-get update && sudo apt-get install -y git

安装步骤

仅安装环境（推荐国内用户）

pip install jaxmarl -i https://pypi.tuna.tsinghua.edu.cn/simple

安装完整版（含算法）

# 克隆仓库
git clone https://github.com/FLAIROx/JaxMARL.git && cd JaxMARL

# 安装依赖（国内用户可加 -i 参数指定镜像源）
pip install -e .[algs]
export PYTHONPATH=./JaxMARL:$PYTHONPATH

使用 Docker（推荐快速启动）

# 构建镜像（需提前安装 Docker 和 Nvidia Container Toolkit）
make build

# 运行容器
make run

基本使用

最简环境测试

import jax
from jaxmarl import make

key = jax.random.PRNGKey(0)
key, key_reset, key_act, key_step = jax.random.split(key, 4)

# 初始化环境
env = make('MPE_simple_world_comm_v3')  # 使用 PettingZoo 风格接口

# 重置环境
obs, state = env.reset(key_reset)

# 采样随机动作
key_act = jax.random.split(key_act, env.num_agents)
actions = {agent: env.action_space(agent).sample(key_act[i]) 
           for i, agent in enumerate(env.agents)}

# 执行一步
obs, state, reward, done, infos = env.step(key_step, state, actions)

运行教程示例

# 本地运行 Colab 笔记本（需安装 jupyter）
pip install jupyter
jupyter notebook jaxmarl/tutorials/JaxMARL_Walkthrough.ipynb

📌 提示：

1. 国内用户安装依赖时建议添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 使用清华镜像加速
2. Docker 用户可通过 make run 直接进入预配置环境
3. 环境列表详见官方文档：https://jaxmarl.foersterlab.com/