Multi-agent RL experiments that compile end-to-end

What it does

Mava is a research-oriented codebase for multi-agent reinforcement learning built entirely on JAX. It bundles single-file implementations of algorithms like PPO, QMIX, and SAC with wrappers for several popular MARL environment suites. The project is explicitly designed as a clone-and-modify research tool rather than a stable library, favoring rapid iteration over semantic-versioned convenience.

The interesting bit

The codebase splits training across two distribution architectures borrowed from the Podracer lineage: Anakin, which JIT-compiles the entire training loop when environments are written in JAX, and Sebulba, which pairs accelerators with CPU-bound environments that resist compilation. Researchers keep the same algorithm code but swap execution strategy depending on whether their environment is JAX-native or not.

Key highlights

• Single-file algorithm implementations (IPPO, MAPPO, MAT, Sable, and others) meant to be edited directly.
• Native support for standardized JSON evaluation logging compatible with the MARL-eval library.
• Wrappers for JAX-based environments (Jumanji, JaxMARL, Matrax) and non-JAX fallbacks (SMAC-lite, robotic warehouse).
• Explicitly not a pip-installable library; the authors expect users to clone the repo and work inside the source tree.

Caveats

• The README warns that Mava is “not meant to be installed as a library,” so expect breaking changes and a workflow that assumes you are editing source files.
• Only a subset of algorithms (notably some PPO variants) support the Sebulba architecture for non-JAX environments; most implementations currently require JAX-native environments.

Verdict

Researchers who want to iterate quickly on MARL ideas in JAX will find Mava’s compilation-friendly loops and hackable file structure genuinely productive. Production engineers looking for a stable, versioned dependency should look elsewhere.