Mava
A JAX research codebase for rapidly training and evaluating distributed multi-agent reinforcement learning systems.
What does this agent do, and when should you use it?
Mava is a research-oriented JAX codebase for multi-agent reinforcement learning, built around single-file implementations and rapid iteration. It implements PPO, Q Learning, SAC, MAT, Sable, and GPO systems across independent-learning, centralized-training/decentralized-execution, and heterogeneous-agent paradigms. Training is launched from Python system files, with Hydra configuration supplied from YAML defaults and command-line overrides. Environment wrappers support both JAX and non-JAX environment suites: JAX environments use the Anakin architecture, while non-JAX environments can use Sebulba. Mava natively writes standardized JSON experiment logs for downstream aggregation and plotting with MARL-eval, and it is explicitly not intended as an importable modular library.
A researcher runs a system file such as python mava/systems/ppo/anakin/ff_ippo.py. The system uses Hydra-managed defaults from mava/configs/ and accepts overrides such as env=lbf or env=rware env/scenario=tiny-4ag. It connects environments through wrappers in mava/wrappers/, using Anakin for end-to-end JIT-compiled training with JAX environments or Sebulba for non-JAX environments. Available runnable implementations include ff_ippo.py, ff_mappo.py, rec_qmix.py, ff_isac.py, mat.py, and ff_sable.py. Runs produce standardized JSON logs intended for downstream processing with MARL-eval.
- A MARL researcher wants to modify `ff_ippo.py` directly and test a new training idea on JAX accelerators.
- A researcher training policies in Level-based Foraging needs to launch a PPO system with `env=lbf`.
- A team studying StarCraft Multi-Agent Challenge, Multi-Robot Warehouse, or Multi-Particle Environments needs existing environment wrappers for repeatable experiments.
- An experimenter comparing discrete-action Q Learning with continuous-action SAC or PPO needs runnable implementations for those algorithm families.
- A research group needs standardized JSON records across MARL runs before aggregating and plotting results with MARL-eval.
What are this agent's strengths and limitations?
- Single-file JAX implementations are designed to be read, changed, and used for rapid MARL research iteration.
- Supports both Podracer architectures: Anakin for JAX environments and Sebulba for non-JAX environments.
- Covers PPO, IQL, QMIX, SAC, MAT, Sable, and GPO, including feed-forward and some recurrent variants.
- Writes standardized JSON logs that connect to MARL-eval aggregation and visualization workflows.
- The project explicitly is not a modular importable library, so adoption generally means working from and modifying repository system files.
- The core runtime depends on JAX, and architecture choice is constrained by whether an environment is written in JAX.
- The algorithm table documents Sebulba support for PPO `ff_ippo.py`; the roadmap still lists adding Sebulba versions of more algorithms.
- Easy scaling across multiple TPUs/GPUs remains an unchecked roadmap item rather than a documented completed capability.
How do you install or deploy this agent?
Requires Python 3.11 or 3.12, Git, network access, and JAX dependencies. Clone and install with uv:git clone https://github.com/instadeepai/Mava.gitcd Mavauv syncsource .venv/bin/activate
For CUDA 12 use uv sync --extra cuda12; for TPU use uv sync --extra tpu. Alternatively, from a virtual environment run pip install -e ".[cuda12]"; omit [cuda12] without a GPU or on Mac. No API credentials are documented.
How do you use this agent?
For a first run: python mava/systems/ppo/anakin/ff_ippo.py. For Level-based Foraging: python mava/systems/ppo/anakin/ff_ippo.py env=lbf. For the Robot Warehouse tiny-4ag scenario: python mava/systems/ppo/anakin/ff_ippo.py env=rware env/scenario=tiny-4ag. Default system configurations are in mava/configs/. Choose Anakin or Sebulba according to whether the target environment is written in JAX.
How does this agent compare with similar options?
Mava says its code philosophy was adapted from PureJaxRL and inspired by CleanRL: all emphasize understandable, research-friendly implementations, while Mava additionally uses small utilities for frequently reused elements such as networks and logging plus Hydra configuration. The README lists JaxMARL as a related project offering accelerated MARL environments with baselines in JAX; Mava focuses on its own training systems and environment wrappers.