JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

Alexander Rutherford (University of Oxford), Benjamin Ellis (University of Oxford), Matteo Gallici (Universitat Politècnica de Catalunya), Jonathan Cook (University of Oxford), Andrei Lupu (University of Oxford), Garðar Ingvarsson (University College London), Timon Willi (University of Oxford), Akbir Khan (University College London), Christian Schroeder de Witt (University of Oxford), Alexandra Souly (University College London), Saptarashmi Bandyopadhyay (University of Maryland), Mikayel Samvelyan (University College London), Minqi Jiang (University College London), Robert Lange (Technical University Berlin), Shimon Whiteson (University of Oxford), Bruno Lacerda (University of Oxford), Nick Hawes (University of Oxford), Tim Rocktäschel (University College London), Chris Lu (University of Oxford), Jakob Foerster (University of Oxford)

Abstract

Benchmarks play an important role in the development of machine learning algorithms, with reinforcement learning (RL) research having been heavily influenced by the available environments. 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-ofuse 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