JaxMARL: multi-agent reinforcement learning that runs entirely in JAX

One library for many multi-agent environments, not one of everything

JaxMARL exists to solve a specific evaluation problem. Its stated aim is one library that lets a researcher run a proposed multi-agent reinforcement learning method across a wide range of tasks and compare it against relevant baselines, without rewriting the environment interface each time. The environment table in the README lists eleven entries: MPE for communication oriented tasks in a multi-agent particle world, Overcooked for fully cooperative human-AI coordination, OvercookedV2 as a partially observable and stochastic extension, Multi-Agent Brax for continuous multi-agent robotic control, Hanabi as a partially observable card game, SMAX, STORM for matrix games rendered as grid worlds, JaxNav for differential drive navigation, Coin Game for social dilemmas, Switch Riddle for debugging, and JaxRobotarium for multi-robot work with open access sim2real.

That breadth is the selling point, but it is also the constraint. Every one of these is a JAX reimplementation rather than a wrapper around an existing simulator, so the environments trade fidelity for vectorisation. The Switch Riddle entry is described as simple and included for debugging, which tells you the authors expected people to instrument these code paths.

The README also points to a blog post and a Colab notebook walkthrough for basic usage. Those two links are the fastest way to see the API shape before installing anything locally.

The baseline set follows CleanRL, one file per algorithm

The baseline side of the project is organised around single file implementations in a baselines directory, a convention the README attributes to CleanRL's philosophy. The listed algorithms are IPPO, MAPPO, IQL, VDN, QMIX, TransfQMIX and SHAQ, with the README table pointing each one at a paper and a source folder. Value based methods IQL, VDN, QMIX and TransfQMIX all live under the QLearning source folder, so reading one file teaches you the pattern for the others.

Configuration is handled with Hydra, with specifics explained in each algorithm's README. Most baseline files contain wandb logging code, and that logging is disabled by default and enabled within the file's config. This is a deliberate choice for reproducibility and clean diffs: a researcher who wants sweep tracking turns it on per file rather than through a global setting that could change behaviour of every run.

The trade-off is that the baselines are reference implementations, not a tuned catalogue of state of the art results. If your comparison needs a specific hyperparameter sweep, expect to write it yourself inside Hydra rather than inherit it.

SMAX removes the StarCraft II engine from the loop

SMAX is the project's own environment rather than a port, described as a vectorised and simplified version of StarCraft Multi-Agent Challenge that removes the need to run the StarCraft II game engine. The README marks it as Novel in the reference column, so it does not map onto an existing paper.

This matters practically. The original SMAC benchmark requires the heavyweight game binary, which makes it awkward to run inside containers, on CI runners, or across a cluster. A pure JAX grid-like micro-management environment steps in the same conceptual space without that dependency, which is why it is often the first environment people try when they want to see whether the JAX toolchain pays off on their hardware.

The honest caveat is that SMAX is not SMAC. It is simplified by construction, so results measured on it do not transfer directly to the full benchmark. The README does not document any correlation study between SMAX and the original environment, so treat it as a development and benchmarking stand-in rather than a substitute for published SMAC numbers.

Installing is a Python 3.11 floor plus extras that pull in the training stack

The packaging metadata is explicit about what you are signing up for. The project name is jaxmarl, the build backend is setuptools, and the version is read dynamically from the jaxmarl package attribute. The interpreter floor is high:

requires-python = ">=3.11"

The base dependency list pins the numerical floor of the whole stack:

dependencies = [
  "jax>=0.4.25",
  "numpy>=1.26",
  "flax",
  "safetensors",
  "chex",
  "matplotlib",
  "pillow",
  "jaxtyping>=0.2.28",
  "gymnax",
]

Training extras are opt in, split across three groups. The algs group carries optax, distrax, flashbax, wandb, hydra-core, omegaconf, tqdm and imageio. The dev group adds pytest, mpe2, ruff, pyright and pre-commit. The mabrax group carries brax at version 0.10.3 or later, which is only needed if you want the Multi-Agent Brax environment.

That split is worth noting: hydra and wandb are not installed by default, so a plain install gives you the environments and their interfaces, not a runnable baseline.

The Makefile wraps a pinned GPU container rather than your local interpreter

The repository ships a Dockerfile based on the NVIDIA JAX image and a Makefile that turns the container into the primary workflow. The Dockerfile pins jax and jaxlib to the versions already present in the base image rather than resolving them fresh:

    RUN pip freeze | grep -iE '^(jax|jaxlib|jax-cuda13)' > /tmp/jax-pins.txt && \
    pip install --no-cache-dir --constraint /tmp/jax-pins.txt -e .[algs,dev] && \
    git config --global --add safe.directory /home/${MYUSER} && \
    git config --global --add safe.directory /home/${MYUSER}/jaxmarl/environments/robotarium && \
    git config --global --add safe.directory /home/${MYUSER}/jaxmarl/environments/robotarium/jaxrobotarium/robotarium_python_simulator && \
    git submodule update --init --recursive --depth 1 && \
    pip install --no-cache-dir --constraint /tmp/jax-pins.txt -e jaxmarl/environments/robotarium && \
    pip install --no-cache-dir --constraint /tmp/jax-pins.txt -e jaxmarl/environments/robotarium/jaxrobotarium/robotarium_python_simulator && \
    rm /tmp/jax-pins.txt

The pin is the point. Installing jaxmarl without the constraint file can drag in a jax build that disagrees with the CUDA runtime baked into the image, which is the classic source of unexplained device errors on GPU machines. The Makefile detects nvidia-smi and adds --gpus all only when it is present:

build:
	DOCKER_BUILDKIT=1 docker build --build-arg USE_CUDA=$(USE_CUDA) --build-arg MYUSER=$(MYUSER) --build-arg UID=$(ID) --tag $(IMAGE) --progress=plain ${PWD}/.

run:
	$(DOCKER_RUN) /bin/bash

test:
	$(DOCKER_RUN) /bin/bash -c "pytest ./tests/"

local-test:
	pytest ./tests/ -v

Two further targets exist for CI shaped work. workflow-test runs the suite with JAX_PLATFORMS=cpu so the CUDA plugin does not segfault on a runner without a GPU driver, and run-baseline-set launches baselines/run_minimal_baseline_set.py in a detached container with a mounted workspace. The container also sets XLA_PYTHON_CLIENT_PREALLOCATE=false and TF_FORCE_GPU_ALLOW_GROWTH=true, and reserves 20G of shared memory, which matters for JAX workloads on shared runners.

Where the JAX-only bet bites back

The clearest limitation is structural. Because the environments are rewritten in JAX rather than wrapped, any simulator you already trust, and any model you already trained, has to be re-expressed inside this stack. The pyproject dependency list has no torch, no stable-baselines3 and no gymnasium, so mixing JaxMARL with an existing PyTorch codebase means two ecosystems in one process and one device placement scheme per library.

A second limit is the algorithms available. The listed baselines cover independent learners, value decomposition methods and two variants of QMIX. There is no communications, no MAPPO variant with recurrent policies and no trust region method in the visible table. If your comparison needs a specific policy class, you write it yourself, and you own its correctness.

A third limit is version coupling. The package requires jax at 0.4.25 or later and numpy at 1.26 or later, and the container deliberately pins jax to the base image versions. Mixing a JaxMARL checkout into an environment with a different jax build is where breakage usually starts. Type checking is deliberately narrow as well: the pyright configuration covers only three files, jaxmarl/environments/multi_agent_env.py, jaxmarl/environments/spaces.py and jaxmarl/wrappers/baselines.py, with typeCheckingMode set to basic. Elsewhere in the codebase you get no static guarantees.

A young project on an Apache-2.0 licence

Maintenance facts come from the release history. The most recent release, v0.2.0, was published on 2026-09-10, and the last push to the repository was on 2026-10-22, some weeks after the release. Before that, v0.1.0 arrived on 2025-06-01 and v0.0.7 on 2024-12-19. The repository is not archived, and the codebase moved forward within the last six months, so the project is being worked on.

The version numbers themselves suggest an early library. Three tagged releases, one of them below 0.1.0 and the newest at 0.2.0, indicate that the public interface can still move. The repository carries a CHANGELOG.md, which is where you would look before pinning a commit.

Licensing is Apache-2.0, and the pyproject metadata points at the LICENSE file rather than embedding an SPDX string. Apache-2.0 is permissive and carries an explicit patent grant, which is the reason it is common across research code that universities later commercialise. Check the LICENSE file yourself for the obligations that attach to redistribution of modified versions, rather than relying on a summary.

Editorial conclusion

JaxMARL suits research groups that want every environment and baseline behind one JAX API and can live inside a single Python 3.11 environment, and it suits teams that need SMAX without installing the StarCraft II engine. It is the wrong choice for work that depends on environments outside the eleven listed, on native JAX or GPU libraries that the pinned dependency set will not accommodate, or on codebases that cannot absorb a dependency tree this tightly coupled to a single ecosystem. Before adopting it, check that the environment you need exists under jaxmarl/environments, confirm the baseline you want to compare against is present under baselines, and read the CHANGELOG.md to see what changed in the v0.2.0 release dated 2026-09-10.

Frequently asked questions

Official sources

1. bold-lab-ai/JaxMARL on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases