# JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

arXiv: https://arxiv.org/abs/2311.10090 · HTML: https://mlsanity.com/paper/2311.10090/

## Key facts

- Cited by 63 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 80 references that are in the map.
- First posted on arXiv 16 Nov 2023.

## Builds on (most cited first, 30 of 53)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Games](https://mlsanity.com/paper/2103.01955/) (2021)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [QMIX: Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1803.11485/) (2018)
- [The StarCraft Multi-Agent Challenge](https://mlsanity.com/paper/1902.04043/) (2019)
- [Learning to Communicate with Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1605.06676/) (2016)
- [Value-Decomposition Networks For Cooperative Multi-Agent Learning](https://mlsanity.com/paper/1706.05296/) (2017)
- [Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2003.08839/) (2020)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [On the Utility of Learning about Humans for Human-AI Coordination](https://mlsanity.com/paper/1910.05789/) (2019)
- [Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?](https://mlsanity.com/paper/2011.09533/) (2020)
- [The Complexity of Decentralized Control of Markov Decision Processes](https://mlsanity.com/paper/1301.3836/) (2013)
- [Learning with Opponent-Learning Awareness](https://mlsanity.com/paper/1709.04326/) (2017)
- [PettingZoo: Gym for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2009.14471/) (2020)
- [Benchmarking Multi-Agent Deep Reinforcement Learning Algorithms in Cooperative Tasks](https://mlsanity.com/paper/2006.07869/) (2020)
- ["Other-Play" for Zero-Shot Coordination](https://mlsanity.com/paper/2003.02979/) (2020)
- [The Hanabi Challenge: A New Frontier for AI Research](https://mlsanity.com/paper/1902.00506/) (2019)
- [Trust Region Policy Optimisation in Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2109.11251/) (2021)
- [Multiagent Cooperation and Competition with Deep Reinforcement Learning](https://mlsanity.com/paper/1511.08779/) (2015)
- [FACMAC: Factored Multi-Agent Centralised Policy Gradients](https://mlsanity.com/paper/2003.06709/) (2020)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Behaviour Suite for Reinforcement Learning](https://mlsanity.com/paper/1908.03568/) (2019)
- [SMACv2: An Improved Benchmark for Cooperative Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2212.07489/) (2022)
- [Stabilising Experience Replay for Deep Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/1702.08887/) (2017)
- [Shared Experience Actor-Critic for Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2006.07169/) (2020)

## Built on by (most cited first, 30 of 63)

- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS](https://mlsanity.com/paper/2408.01584/) (2024)
- [1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities](https://mlsanity.com/paper/2503.14858/) (2025)
- [BenchMARL: Benchmarking Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2312.01472/) (2023)
- [Multi-agent Reinforcement Learning: A Comprehensive Survey](https://mlsanity.com/paper/2312.10256/) (2023)
- [Learning to Reason at the Frontier of Learnability](https://mlsanity.com/paper/2502.12272/) (2025)
- [Scaling Opponent Shaping to High Dimensional Games](https://mlsanity.com/paper/2312.12568/) (2023)
- [The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Design](https://mlsanity.com/paper/2406.17949/) (2024)
- [HyperMARL: Adaptive Hypernetworks for Multi-Agent RL](https://mlsanity.com/paper/2412.04233/) (2024)
- [MEAL: A Benchmark for Continual Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2506.14990/) (2025)
- [minimax: Efficient Baselines for Autocurricula in JAX](https://mlsanity.com/paper/2311.12716/) (2023)
- [JaxLife: An Open-Ended Agentic Simulator](https://mlsanity.com/paper/2409.00853/) (2024)
- [ROTATE: Regret-driven Open-ended Training for Ad Hoc Teamwork](https://mlsanity.com/paper/2505.23686/) (2025)
- [VS-Bench: Evaluating VLMs for Strategic Abilities in Multi-Agent Environments](https://mlsanity.com/paper/2506.02387/) (2025)
- [An Extended Benchmarking of Multi-Agent Reinforcement Learning Algorithms in Complex Fully Cooperative Tasks](https://mlsanity.com/paper/2502.04773/) (2025)
- [Partner Modelling Emerges in Recurrent Agents (But Only When It Matters)](https://mlsanity.com/paper/2505.17323/) (2025)
- [Capability-Aware Shared Hypernetworks for Flexible Heterogeneous Multi-Robot Coordination](https://mlsanity.com/paper/2501.06058/) (2025)
- [The Yokai Learning Environment: Tracking Beliefs Over Space and Time](https://mlsanity.com/paper/2508.12480/) (2025)
- [Self-Supervised Goal-Reaching Results in Multi-Agent Cooperation and Exploration](https://mlsanity.com/paper/2509.10656/) (2025)
- [gfnx: Fast and Scalable Library for Generative Flow Networks in JAX](https://mlsanity.com/paper/2511.16592/) (2025)
- [Subequivariant Reinforcement Learning in 3D Multi-Entity Physical Environments](https://mlsanity.com/paper/2407.12505/) (2024)
- [RobocupGym: A challenging continuous control benchmark in Robocup](https://mlsanity.com/paper/2407.14516/) (2024)
- [Multi-Agent Craftax: Benchmarking Open-Ended Multi-Agent Reinforcement Learning at the Hyperscale](https://mlsanity.com/paper/2511.04904/) (2025)
- [JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning](https://mlsanity.com/paper/2512.06102/) (2025)
- [Shaping Zero-Shot Coordination via State Blocking](https://mlsanity.com/paper/2605.11688/) (2026)
- [Benchmarking the Limits of In-Context Reinforcement Learning for Ad-Hoc Teamwork](https://mlsanity.com/paper/2605.24423/) (2026)
- [Benchmarking Open-Ended Multi-Agent Coordination in Language Agents](https://mlsanity.com/paper/2606.08340/) (2026)
- [CRAX: Fast Safe Reinforcement Learning Benchmarking](https://mlsanity.com/paper/2606.20376/) (2026)
- [A Batch Sequential Halving Algorithm without Performance Degradation](https://mlsanity.com/paper/2406.00424/) (2024)
- [EvoRL: A GPU-accelerated Framework for Evolutionary Reinforcement Learning](https://mlsanity.com/paper/2501.15129/) (2025)

## Abstract (arXiv)

Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, we present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge, which removes the need to run the StarCraft II game engine. This not only enables GPU acceleration, but also provides a more flexible MARL environment, unlocking the potential for self-play, meta-learning, and other future applications in MARL. The code is available at https://github.com/flairox/jaxmarl.