Rutherford, A;
Ellis, B;
Gallici, M;
Cook, J;
Lupu, A;
Ingvarsson, G;
Willi, T;
... Foerster, J; + view all
(2024)
JaxMARL: Multi-Agent RL Environments and Algorithms in JAX.
In: Globersons, Amir and Mackey, Lester and Belgrave, Danielle and Fan, Angela and Paquet, Ulrich and Tomczak, Jakub M and Zhang, Cheng, (eds.)
Advances in Neural Information Processing Systems.
NeurIPS
Preview |
Text
NeurIPS-2024-jaxmarl-multi-agent-rl-environments-and-algorithms-in-jax-Paper-Datasets_and_Benchmarks_Track.pdf - Published Version Download (4MB) | Preview |
Abstract
Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, we present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation 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. The code is available at https://github.com/flairox/jaxmarl.
| Type: | Proceedings paper |
|---|---|
| Title: | JaxMARL: Multi-Agent RL Environments and Algorithms in JAX |
| Event: | NeurIPS 2024 |
| Open access status: | An open access version is available from UCL Discovery |
| DOI: | 10.52202/079017-1612 |
| Publisher version: | https://papers.nips.cc/paper_files/paper/2024 |
| Language: | English |
| Additional information: | This version is the author accepted manuscript. For information on re-use, please refer to the publisher’s terms and conditions. |
| Keywords: | Reinforcement learning, multi-agent, multi-agent reinforcement learning, jax |
| UCL classification: | UCL UCL > Provost and Vice Provost Offices > UCL BEAMS UCL > Provost and Vice Provost Offices > UCL BEAMS > Faculty of Engineering Science > Dept of Computer Science |
| URI: | https://discovery.ucl.ac.uk/id/eprint/10216730 |
Archive Staff Only
 |
View Item |
