JaxMARL: Multi-Agent RL Environments in JAX
Authors: name: Alexander Rutherford; name: Benjamin Ellis; name: Matteo Gallici; name: Jonathan Cook; name: Andrei Lupu; name: Gardar Ingvarsson; name: Timon Willi; name: Akbir Khan; name: Christian Schroeder de Witt; name: Alexandra Souly; name: Saptarashmi Bandyopadhyay; name: Mikayel Samvelyan; name: Minqi Jiang; name: Robert Tjarko Lange; name: Shimon Whiteson; name: Bruno Lacerda; name: Nick Hawes; name: Tim Rocktaschel; name: Chris Lu; name: Jakob Nicolaus Foerster
Published: 2023-11-16 18:58:43

Abstract: Benchmarks play an important role in the development of machine learning
algorithms. For example, research in reinforcement learning (RL) has been
heavily influenced by available environments and benchmarks. However, RL
environments are traditionally run on the CPU, limiting their scalability with
typical academic compute. Recent advancements in JAX have enabled the wider use
of hardware acceleration to overcome these computational hurdles, enabling
massively parallel RL training pipelines and environments. This is particularly
useful for multi-agent reinforcement learning (MARL) research. First of all,
multiple agents must be considered at each environment step, adding
computational burden, and secondly, the sample complexity is increased due to
non-stationarity, decentralised partial observability, or other MARL
challenges. In this paper, we present JaxMARL, the first open-source code base
that combines ease-of-use with GPU enabled efficiency, and supports a large
number of commonly used MARL environments as well as popular baseline
algorithms. When considering wall clock time, our experiments show that per-run
our JAX-based training pipeline is up to 12500x faster than existing
approaches. This enables efficient and thorough evaluations, with the potential
to alleviate the evaluation crisis of the field. We also introduce and
benchmark SMAX, a vectorised, simplified version of the popular StarCraft
Multi-Agent Challenge, which removes the need to run the StarCraft II game
engine. This not only enables GPU acceleration, but also provides a more
flexible MARL environment, unlocking the potential for self-play,
meta-learning, and other future applications in MARL. We provide code at
https://github.com/flairox/jaxmarl.