🤖 AI

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

This paper introduces JaxMARL, an open-source Python library that leverages JAX to provide GPU-accelerated, massively parallel multi-agent reinforcement learning environments and algorithms, achieving significant speedups over existing approaches and offering a flexible, engine-free reimplementation of the StarCraft Multi-Agent Challenge.

Original authors: Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Gardar Ingvarsson, Timon Willi, Ravi Hammond, Akbir Khan, Christian Schroeder de Witt, Alexandra Souly, Saptarashmi Ba

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Gardar Ingvarsson, Timon Willi, Ravi Hammond, Akbir Khan, Christian Schroeder de Witt, Alexandra Souly, Saptarashmi Bandyopadhyay, Mikayel Samvelyan, Minqi Jiang, Robert Tjarko Lange, Shimon Whiteson, Bruno Lacerda, Nick Hawes, Tim Rocktaschel, Chris Lu, Jakob Nicolaus Foerster

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a team of robots how to work together. In the world of Artificial Intelligence, this is called Multi-Agent Reinforcement Learning (MARL). Usually, to teach these robots, researchers have to run thousands of simulations.

Think of the traditional way of doing this as trying to teach a soccer team using a single, slow coach who has to run around the field on foot, giving instructions to one player at a time. It works, but it takes forever. If you want to test 100 different coaching strategies, you might spend months just waiting for the results. This is the "bottleneck" the paper describes: the computers (CPUs) used for these simulations are too slow to handle the massive amount of practice needed to train teams of agents effectively.

Enter JaxMARL.

The authors of this paper built a new tool called JaxMARL. You can think of this as upgrading that single, slow coach into a super-orchestra conductor with a fleet of 10,000 drones.

Here is how they did it and what they found, broken down simply:

1. The Speed Boost (The "Drone Fleet")

Instead of running simulations one by one on a standard computer processor, JaxMARL uses a special programming library called JAX that lets the computer use powerful graphics cards (GPUs)—the same chips used for gaming—to do the math.

The Analogy: Imagine you need to paint 10,000 walls. The old way (CPU) is like hiring one painter who paints one wall, dries, then paints the next. The JaxMARL way is like having a machine that can spray-paint all 10,000 walls at the exact same time.
The Result: The paper claims their system is 14 times faster for a single training run. But when they run many experiments at once (which researchers often do), it is up to 12,500 times faster. This turns a process that used to take weeks into something that takes hours or minutes.

2. The New Playgrounds (Environments)

To train these AI agents, you need "games" or environments. The paper introduces a library of many different games, all rewritten to run on this super-fast system.

SMAX (The StarCraft Replacement): One of the most popular games for training AI teams is based on the video game StarCraft II. However, the original game is heavy and slow because it has to run the entire 3D graphics engine just to train the AI.
- The Fix: The authors created SMAX. Think of this as a "skeleton version" of the game. It keeps all the rules and strategy of StarCraft but strips away the fancy 3D graphics. Because it's just the logic running on a GPU, it is 40,000 times faster than the original game engine.
STORM (The Grid World): They also built a new set of games called STORM. Imagine a board game where players move around a grid collecting coins, but the rules are designed to test how well they cooperate or compete. This helps researchers study how agents learn to work together or trick each other.

3. The Coaches (Algorithms)

Just having a fast playground isn't enough; you need good training methods. The paper also rewrote several famous "coaching strategies" (algorithms like PPO and QMIX) to work with this new fast system. They verified that these new fast coaches produce the same smart results as the old, slow coaches, just much quicker.

4. Why This Matters (The "Evaluation Crisis")

The paper points out a problem in the field: because training is so slow, researchers often only test their new ideas on one or two games. This is like a chef only testing a new recipe on one type of pasta. If the recipe works on spaghetti but fails on penne, the chef doesn't know.

Because JaxMARL is so fast, researchers can now test their ideas on dozens of different games in the time it used to take to test just one. This helps ensure that the AI is actually smart and general, rather than just "memorizing" a single specific game.

Summary

In short, JaxMARL is a free, open-source toolbox that lets researchers train teams of AI agents using the massive power of modern graphics cards. It replaces slow, heavy video game engines with lightweight, lightning-fast versions, allowing scientists to run experiments thousands of times faster than before. This helps them figure out better ways for AI to cooperate, compete, and solve complex problems without getting stuck waiting for computers to catch up.

Technical Summary: JaxMARL

Problem Statement

Reinforcement Learning (RL) research relies heavily on benchmarks to define problems and enable algorithm comparison. However, traditional RL environments typically run on CPUs, creating a significant bottleneck for scalability. This limitation is particularly acute in Multi-Agent Reinforcement Learning (MARL), where the computational burden is compounded by the need to train multiple agents simultaneously, handle non-stationarity, and manage decentralized partial observability. Consequently, experimental runs are often slow, expensive, and limited in scope.

Furthermore, the field faces an "evaluation crisis." Recent analyses indicate that many MARL papers test on only one or two environments, leading to potential overfitting, biased comparisons, and unclear progress markers. While hardware acceleration (e.g., GPUs) offers a solution by enabling massive parallelism, existing implementations often require non-Python codebases (e.g., C++), reducing accessibility for the broader research community where Python is the standard. Although JAX has successfully accelerated single-agent RL (e.g., PureJaxRL), it has not been widely adopted for multi-agent scenarios.

Methodology

The authors present JaxMARL, an open-source, Python-based library that leverages the JAX framework to provide GPU-accelerated, massively parallel MARL environments and algorithms.

Core Technical Approach

JAX Integration: The library utilizes JAX's vmap operation to vectorize functions across input dimensions, allowing thousands of environment instances to run in parallel on a single GPU. It also employs Just-In-Time (JIT) compilation to optimize performance.
Co-location: By implementing both the environment dynamics and the agent training pipelines in JAX, the system eliminates the latency associated with copying data between CPU and GPU memory.
Interface Design: The library is designed with a simple interface inspired by PettingZoo and Gymnax, adhering to the "CleanRL" philosophy of providing clear, single-file implementations.

Environment Suite

JaxMARL includes a diverse set of environments covering continuous and discrete dynamics, variable agent counts, and cooperative, competitive, and mixed-incentive settings:

New Environments:
- SMAX: A JAX-based approximate reimplementation of the StarCraft Multi-Agent Challenge (SMAC/SMACv2). Unlike the original, which relies on the heavy StarCraft II game engine, SMAX uses lightweight dynamics, removes the need for the game engine, and supports competitive self-play and high customizability.
- STORM (Spatial-Temporal Representations of Matrix Games): An extension of matrix games (like the Prisoner's Dilemma) into grid-world scenarios. It introduces temporally extended actions, partial observability, and stochasticity to study social dilemmas and opponent shaping.
Existing Environments: JAX reimplementations of Hanabi, Overcooked, Multi-Agent Brax (MABrax), Multi-Agent Particle Environment (MPE), Coin Game, and Switch Riddle.

Algorithm Implementation

The library implements popular MARL baselines compatible with JAX's JIT and vmap capabilities:

Policy Gradient: Independent PPO (IPPO) and Multi-Agent PPO (MAPPO), supporting both feed-forward and RNN policies with parameter sharing.
Q-Learning: Independent Q-Learning (IQL), Value Decomposition Networks (VDN), and QMIX, implemented to ensure consistency with the PyMARL codebase.

Key Contributions

First Open-Source JAX MARL Library: JaxMARL is the first library to combine GPU-enabled efficiency with support for a wide range of popular MARL environments and baseline algorithms in a Python-native format.
New Benchmark Suites: The introduction of SMAX (a faster, more flexible alternative to SMAC) and STORM (a grid-world extension of matrix games) expands the available tooling for MARL research.
Algorithm Implementations: Providing JAX-native versions of IPPO, MAPPO, QMIX, VDN, and IQL.
Comprehensive Benchmarking: The authors provide thorough benchmarks comparing the speed and correctness of their implementations against existing CPU-based repositories (e.g., PyMARL, SMAC).
Evaluation Recommendations: The paper proposes standard minimal sets of evaluation environments for different research settings (CTDE, Zero-Shot Coordination, General-Sum, Continuous Actions) and provides scripts for large-scale evaluation using Inter-Quartile Mean (IQM) statistics to mitigate bias.

Results

The authors demonstrate significant performance gains and correctness equivalence:

Speed:
- In wall-clock time, the JAX-based training pipeline is approximately 14 times faster than existing approaches (e.g., PyMARL) for single training runs.
- When multiple training runs are vectorized, the speedup reaches up to 12,500x (e.g., training 1,024 seeds in ~198 seconds vs. over an hour for a single run in PyMARL).
- SMAX training is reported to be 40,000x faster than equivalent SMAC implementations when vectorized on a single NVIDIA 2080.
- Environment step rates (steps per second) for random actions are orders of magnitude higher than CPU counterparts (e.g., MPE Simple Spread reaches $4.0 \times 10^7$ steps/sec with 10k environments).
Performance:
- In multi-environment comparisons (SMAX, Overcooked, MPE), PPO-based methods (IPPO/MAPPO) showed clear advantages over Q-Learning methods.
- Training curves and final scores for algorithms like IPPO and MAPPO on Hanabi and MPE matched or exceeded results from prior work (e.g., PyTorch/C++ implementations), verifying the correctness of the JAX implementations.
Correctness: The authors verified that their environment transitions and rewards match original implementations (e.g., Overcooked-AI, PettingZoo, Hanabi Learning Environment) within tight tolerances ( $1 \times 10^{-4}$ ).

Significance

The paper argues that JaxMARL addresses critical bottlenecks in MARL research:

Alleviating the Evaluation Crisis: By drastically reducing the computational cost of training and evaluation, the library enables researchers to test algorithms across a broad set of environments rather than a few, reducing the risk of overfitting and enabling more robust conclusions.
Democratizing Hardware Acceleration: By providing a Python-based, easy-to-use interface, JaxMARL makes GPU-accelerated MARL accessible to researchers who cannot or do not wish to write complex C++ code.
Enabling New Research Directions: The speed and flexibility of the library (particularly with SMAX) unlock potential for research areas such as self-play, meta-learning, and large-scale population-based training that were previously computationally prohibitive.

The authors conclude that while challenges remain (such as handling variable agent numbers efficiently and the specific limitations of off-policy methods in JAX), JaxMARL represents a significant step forward in making MARL research faster, more thorough, and more accessible.