Multi-agent RL in JAX, minus the StarCraft II binary

What it does

It wraps a broad set of multi-agent environments and baseline algorithms into one JAX-native framework. It covers communication, cooperation, and robotic control tasks—from Overcooked and Hanabi to multi-agent Brax and JaxNav—alongside SMAX, a vectorised rewrite of the StarCraft Multi-Agent Challenge that removes the StarCraft II engine dependency. Observations, actions, and rewards flow as per-agent dictionaries through an API inspired by PettingZoo and Gymnax.

The interesting bit

The baseline algorithms follow CleanRL’s single-file philosophy, meaning each method fits in one readable script with Hydra-managed configs. Because both environments and training loops stay in JAX, the whole pipeline can live on the GPU without the overhead of an external game engine.

Key highlights

• Eleven environments including MPE, OvercookedV2, Coin Game, JaxNav, and the StarCraft-free SMAX.
• Baseline implementations of IPPO, MAPPO, IQL, VDN, QMIX, TransfQMIX, SHAQ, and PQN-VDN.
• End-to-end JAX execution with dictionary-based agent APIs and parallel step semantics.
• JAX versions up to 0.4.36 are tested; Docker environment provided for reproducibility.
• wandb logging hooks are built in but disabled by default.

Caveats

• Asynchronous games like Hanabi use dummy actions for idle agents rather than native turn-based stepping.
• The library is tested only up to JAX 0.4.36, so compatibility with newer releases is unclear.
• Baseline algorithm files include wandb logging code that must be enabled via config if desired.

Verdict

MARL researchers who want a unified benchmark suite and readable baseline code in one JAX stack should start here. Single-agent practitioners or teams wedded to non-JAX simulators will not find much to use.