mlsanityData to 14 Oct 2026

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

arXiv 2311.10090 ↗First posted 16 Nov 2023cs.LG cs.AI cs.MAcited by 63

Not on the map. Why?

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.

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

53 of 80 references.

References · 53

  1. 2017Proximal Policy Optimization Algorithmscited by 15,215
  2. 2016Asynchronous Methods for Deep Reinforcement Learningcited by 2,258
  3. 2015Deep Reinforcement Learning with Double Q-learningcited by 1,745
  4. 2017Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environmentscited by 1,558
  5. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  6. 2021Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learningcited by 998
  7. 2021The Surprising Effectiveness of PPO in Cooperative, Multi-Agent Gamescited by 942
  8. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 807
  9. 2018QMIX: Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learningcited by 708
  10. 2019The StarCraft Multi-Agent Challengecited by 597
  11. 2016Learning to Communicate with Deep Multi-Agent Reinforcement Learningcited by 575
  12. 2017Value-Decomposition Networks For Cooperative Multi-Agent Learningcited by 503
38 more
  1. 2020Monotonic Value Function Factorisation for Deep Multi-Agent Reinforcement Learningcited by 485
  2. 2021Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulationcited by 407
  3. 2019On the Utility of Learning about Humans for Human-AI Coordinationcited by 349
  4. 2020Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?cited by 294
  5. 2013The Complexity of Decentralized Control of Markov Decision Processescited by 231
  6. 2017Learning with Opponent-Learning Awarenesscited by 222
  7. 2020PettingZoo: Gym for Multi-Agent Reinforcement Learningcited by 220
  8. 2020Benchmarking Multi-Agent Deep Reinforcement Learning Algorithms in Cooperative Taskscited by 212
  9. 2020"Other-Play" for Zero-Shot Coordinationcited by 198
  10. 2019The Hanabi Challenge: A New Frontier for AI Researchcited by 197
  11. 2021Trust Region Policy Optimisation in Multi-Agent Reinforcement Learningcited by 186
  12. 2015Multiagent Cooperation and Competition with Deep Reinforcement Learningcited by 184
  13. 2020FACMAC: Factored Multi-Agent Centralised Policy Gradientscited by 178
  14. 2021A Survey of Zero-shot Generalisation in Deep Reinforcement Learningcited by 170
  15. 2019Behaviour Suite for Reinforcement Learningcited by 132
  16. 2022SMACv2: An Improved Benchmark for Cooperative Multi-Agent Reinforcement Learningcited by 122
  17. 2017Stabilising Experience Replay for Deep Multi-Agent Reinforcement Learningcited by 117
  18. 2020Shared Experience Actor-Critic for Multi-Agent Reinforcement Learningcited by 116
  19. 2023Structured State Space Models for In-Context Reinforcement Learningcited by 111
  20. 2022Discovered Policy Optimisationcited by 104
  21. 2021Scalable Evaluation of Multi-Agent Reinforcement Learning with Melting Potcited by 99
  22. 2022VMAS: A Vectorized Multi-Agent Simulator for Collective Robot Learningcited by 68
  23. 2022EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Enginecited by 62
  24. 2023POPGym: Benchmarking Partially Observable Reinforcement Learningcited by 62
  25. 2023Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learningcited by 51
  26. 2022Towards a Standardised Performance Evaluation Protocol for Cooperative MARLcited by 49
  27. 2019Simplified Action Decoder for Deep Multi-Agent Reinforcement Learningcited by 48
  28. 2022Model-Free Opponent Shapingcited by 44
  29. 2022Melting Pot 2.0cited by 40
  30. 2022COLA: Consistent Learning with Opponent-Learning Awarenesscited by 38
  31. 2021MALib: A Parallel Framework for Population-based Multi-agent Reinforcement Learningcited by 23
  32. 2022MARLlib: A Scalable and Efficient Multi-agent Reinforcement Learning Librarycited by 21
  33. 2022Proximal Learning With Opponent-Learning Awarenesscited by 20
  34. 2022Adversarial Cheap Talkcited by 16
  35. 2023JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for tradingcited by 14
  36. 2020Generating and Adapting to Diverse Ad-Hoc Cooperation Agents in Hanabicited by 13
  37. 2021Centralized Model and Exploration Policy for Multi-Agent RLcited by 12
  38. 2022Trust Region Bounds for Decentralized PPO Under Non-stationaritycited by 11

50 most cited of 53 All 53 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2023: 4; 2024: 11; 2025: 21; 2026: 27.

63 papers.

Sort:

By area:

  1. 2024Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learningcited by 88
  2. 2024GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPScited by 45
  3. 20251000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilitiescited by 45
  4. 2023BenchMARL: Benchmarking Multi-Agent Reinforcement Learningcited by 33
  5. 2023Multi-agent Reinforcement Learning: A Comprehensive Surveycited by 33
  6. 2025Learning to Reason at the Frontier of Learnabilitycited by 23
  7. 2023Scaling Opponent Shaping to High Dimensional Gamescited by 12
  8. 2024The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Designcited by 11
  9. 2024HyperMARL: Adaptive Hypernetworks for Multi-Agent RLcited by 11
  10. 2025MEAL: A Benchmark for Continual Multi-Agent Reinforcement Learningcited by 11
  11. 2023minimax: Efficient Baselines for Autocurricula in JAXcited by 10
  12. 2024JaxLife: An Open-Ended Agentic Simulatorcited by 6
38 more
  1. 2025ROTATE: Regret-driven Open-ended Training for Ad Hoc Teamworkcited by 6
  2. 2025VS-Bench: Evaluating VLMs for Strategic Abilities in Multi-Agent Environmentscited by 6
  3. 2025An Extended Benchmarking of Multi-Agent Reinforcement Learning Algorithms in Complex Fully Cooperative Taskscited by 5
  4. 2025Partner Modelling Emerges in Recurrent Agents (But Only When It Matters)cited by 4
  5. 2025Capability-Aware Shared Hypernetworks for Flexible Heterogeneous Multi-Robot Coordinationcited by 3
  6. 2025The Yokai Learning Environment: Tracking Beliefs Over Space and Timecited by 3
  7. 2025Self-Supervised Goal-Reaching Results in Multi-Agent Cooperation and Explorationcited by 3
  8. 2025gfnx: Fast and Scalable Library for Generative Flow Networks in JAXcited by 3
  9. 2024Subequivariant Reinforcement Learning in 3D Multi-Entity Physical Environmentscited by 2
  10. 2024RobocupGym: A challenging continuous control benchmark in Robocupcited by 2
  11. 2025Multi-Agent Craftax: Benchmarking Open-Ended Multi-Agent Reinforcement Learning at the Hyperscalecited by 2
  12. 2025JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learningcited by 2
  13. 2026Shaping Zero-Shot Coordination via State Blockingcited by 2
  14. 2026Benchmarking the Limits of In-Context Reinforcement Learning for Ad-Hoc Teamworkcited by 2
  15. 2026Benchmarking Open-Ended Multi-Agent Coordination in Language Agentscited by 2
  16. 2026CRAX: Fast Safe Reinforcement Learning Benchmarkingcited by 2
  17. 2024A Batch Sequential Halving Algorithm without Performance Degradationcited by 1
  18. 2025EvoRL: A GPU-accelerated Framework for Evolutionary Reinforcement Learningcited by 1
  19. 2025Fixing Incomplete Value Function Decomposition for Multi-Agent Reinforcement Learningcited by 1
  20. 2025Video Game Level Design as a Multi-Agent Reinforcement Learning Problemcited by 1
  21. 2025Automata-Conditioned Cooperative Multi-Agent Reinforcement Learningcited by 1
  22. 2025When Assisting One Disempowers Anothercited by 1
  23. 2026Learning to Coordinate via Quantum Entanglement in Multi-Agent Reinforcement Learningcited by 1
  24. 2026Decentralized Diffusion Policy Learning for Enhanced Exploration in Cooperative Multi-agent Reinforcement Learningcited by 1
  25. 2026Crazyflow: An Accurate, GPU-Accelerated, Differentiable Drone Simulator in JAXcited by 1
  26. 2026JaxAHT: A JAX-Based Library for Ad Hoc Teamworkcited by 1
  27. 2024MARL-LNS: Cooperative Multi-agent Reinforcement Learning via Large Neighborhoods Searchcited by 0
  28. 2024FlickerFusion: Intra-trajectory Domain Generalizing Multi-Agent RLcited by 0
  29. 2024A Research Agenda for Usability and Generalisation in Reinforcement Learningcited by 0
  30. 2025Oryx: a Scalable Sequence Model for Many-Agent Coordination in Offline MARLcited by 0
  31. 2025Assistax: A Multi-Agent Hardware-Accelerated Reinforcement Learning Benchmark for Assistive Roboticscited by 0
  32. 2025Inter-Agent Relative Representations for Multi-Agent Option Discoverycited by 0
  33. 2026TABX: A High-Throughput Sandbox Battle Simulator for Multi-Agent Reinforcement Learningcited by 0
  34. 2026Recurrent Structural Policy Gradient for Partially Observable Mean Field Gamescited by 0
  35. 2026MOSAIC: A Universal Agent-Level Interface for Cross-Paradigm Agent Mixing and Human-AI Collaborationcited by 0
  36. 2026Strategically Robust Multi-Agent Reinforcement Learning with Linear Function Approximationcited by 0
  37. 2026Automatic Generation of High-Performance RL Environmentscited by 0
  38. 2026NASimJax: A GPU-Accelerated Policy Learning Framework for Penetration Testingcited by 0

50 most cited of 63 All 63 (JSON)