Research map / Multi-agent systems and games / Cooperative multi-agent RL
JaxMARL: Multi-Agent RL Environments and Algorithms in JAX
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces JaxMARL, an open-source JAX library of GPU-accelerated multi-agent RL environments (including new SMAX and STORM suites) and baseline algorithms (IPPO, MAPPO, IQL, VDN, QMIX), reporting large wall-clock speedups over CPU-based PyMARL and related implementations.
Problem
MARL research is slow and under-evaluated because most environments run on CPU and papers typically test on very few domains, limiting scalability on academic compute.
Approach
Reimplement popular MARL environments and baseline algorithms end-to-end in JAX so environments and training run on the same accelerator, enabling jit/vmap parallelism across many seeds, plus a SMAC-like JAX environment (SMAX) that avoids the StarCraft II engine.
Key results
- JaxMARL PPO training pipeline is about 14x faster than PyMARL for a single MPE run and up to 12,500x when many training runs are vectorized.
- SMAX training is reported as about 31x faster than SMAC for a single IPPO run on the same wall-clock comparison and about 40,000x faster when multiple runs are vectorized on a single NVIDIA 2080.
- Per Table 3, JaxMARL environments achieve up to ~1e7 steps per second with 10k parallel envs (e.g., MPE Simple Spread 4.0e7), far above CPU originals (e.g., 8.3e4 for MPE Simple Spread).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use PPO instead of Q-Learning | JaxMARL · SMAX, Overcooked, and MPE | normalized score | Helped |
| Use PPO instead of Q-Learning | JaxMARL on SMAX · SMAX | training time: 6 times faster | Helped |
| Use PPO instead of Q-Learning | JaxMARL on Overcooked · Overcooked | training time: 10 times faster | Helped |
| Run the environment and training pipeline in JAX | JaxMARL IPPO · MPE | training speedup relative to PyMARL: 14x | Helped |
| Vectorize multiple training runs | JaxMARL IPPO · MPE | training speedup relative to PyMARL: 12,500x | Helped |
| Use SMAX with JaxMARL IPPO instead of SMAC with PyMARL | IPPO training · SMAX / SMAC | training speedup: 31x | Helped |
| Vectorize 1024 training runs | JaxMARL QMIX · MPE | training speedup per agent relative to PyMARL: 21,500x | Helped |
| Use IPPO instead of Q-Learning | JaxMARL · Overcooked | normalized score | Helped |
| Use MAPPO instead of IPPO | JaxMARL · SMAX | win rate | Helped |
| Use IPPO instead of MAPPO | JaxMARL · Hanabi | score | Mixed |
| Increase the number of vectorized environments | JaxMARL environment stepping · MPE Simple Spread | steps per second: 5.5 × 10^3 -> 4.0 × 10^7 | Helped |
| Increase the number of vectorized environments | JaxMARL environment stepping · Overcooked | steps per second: 3.6 × 10^3 -> 1.7 × 10^7 | Helped |
Surprising findings
- PPO-based methods (MAPPO, IPPO) outperform Q-learning methods (IQL, VDN, QMIX) across the benchmarked environments, both in performance and training runtime.
- Speedups are less pronounced for off-policy, value-based methods due to replay-ratio imbalance as parallel environments increase.
Limitations
- Speedups are less pronounced for off-policy value-based methods, and Q-learning replay buffers consume large GPU memory.
- End-to-end JAX implementations struggle with environments having a variable number of agents or massive observation sizes.
- Environments largely re-implement or draw from existing suites and do not yet push the boundaries of MARL capabilities.
Experimental setup
Datasets: SMAC scenarios SMACv2 scenarios MPE Overcooked Hanabi MABrax (Multi-Agent MuJoCo derivative) Coin Game Switch Riddle
Benchmarks: SMAX MPE Overcooked Hanabi MABrax STORM Coin Game
Metrics: steps per second wall clock training time speedup relative to PyMARL episode return win rate interquartile mean mean normalized score
Baselines: PyMARL E-PyMARL MARLLib SMAC PettingZoo Overcooked-AI Hanabi Learning Environment
Builds on: PureJaxRL PyMARL CleanRL Gymnax SMACv2
Closest papers by meaning
- Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX (ICLR 2024 · similarity 0.89)
- FightLadder: A Benchmark for Competitive Multi-Agent Reinforcement Learning (ICML 2024 · similarity 0.89)
- HyperMARL: Adaptive Hypernetworks for Multi-Agent RL (NeurIPS 2025 · similarity 0.89)
- Dr. MAS: Stable Reinforcement Learning for Multi-Agent LLM Systems (NeurIPS 2026 · similarity 0.88)
- XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX (NeurIPS 2024 · similarity 0.88)
- SRL: Scaling Distributed Reinforcement Learning to Over Ten Thousand Cores (ICLR 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.