JaxMARL: Multi-Agent RL Environments in JAX
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.
6556f0ff2de0aee1a11fab3f
Alexander Rutherford
6556f0ff2de0aee1a11fab40
Benjamin Ellis
6556f0ff2de0aee1a11fab41
Matteo Gallici
claimed_verified
2025-06-02T08:03:18.145Z
637f520c7119bd030dfee863
/avatars/e3ff0d3c7baf74063854087218cb240a.svg
Matteo
user
mttga
6556f0ff2de0aee1a11fab42
Jonathan Cook
6556f0ff2de0aee1a11fab43
Andrei Lupu
admin_assigned
2023-11-17T16:10:34.543Z
64c3ffbb17bd80604501fc28
/avatars/cfe60c2c6b5656cb541357877aa757dc.svg
Andrei Lupu
user
iamexploit
6556f0ff2de0aee1a11fab44
Gardar Ingvarsson
claimed_verified
2023-11-18T12:06:08.645Z
6309f082450a9475bd123617
/avatars/5622cc7a4c2da68c320672792029785d.svg
Garðar Ingvarsson
user
gardarjuto
6556f0ff2de0aee1a11fab45
Timon Willi
6556f0ff2de0aee1a11fab46
Akbir Khan
admin_assigned
2023-11-17T16:10:48.810Z
62a0788fd066f825a0164c68
/avatars/a565c173c82dea26a5a61370ce42b81f.svg
Akbir Khan
user
akbir
6556f0ff2de0aee1a11fab47
Christian Schroeder de Witt
admin_assigned
2023-11-17T16:10:56.297Z
62cd901ff27cc94a7594f7f9
/avatars/239d4a2a850e3bb819dac8a1ca3bbd77.svg
Christian Schroeder de Witt
user
csdw
6556f0ff2de0aee1a11fab48
Alexandra Souly
admin_assigned
2023-11-17T16:11:02.831Z
6426d309bc4f1d51f547a264
/avatars/ac43cb7d23b1cdcb85b9ba068279021d.svg
Alexandra Souly
user
szantamano
6556f0ff2de0aee1a11fab49
Saptarashmi Bandyopadhyay
admin_assigned
2023-11-17T16:11:16.428Z
63ba1de3205688cd2f95a00c
/avatars/0fd966f6cd783129cfda19191ce824b1.svg
Saptarashmi Bandyopadhyay
user
saptab
6556f0ff2de0aee1a11fab4a
Mikayel Samvelyan
admin_assigned
2023-11-17T16:11:23.207Z
64f650648afa37b1248c8792
https://cdn-avatars.huggingface.co/v1/production/uploads/64f650648afa37b1248c8792/jqYcQWIdLrnBXASnVBUp-.jpeg
Mikayel Samvelyan
user
samvelyan
6556f0ff2de0aee1a11fab4b
Minqi Jiang
admin_assigned
2023-11-17T16:12:34.235Z
630368320547362a22a21798
/avatars/035b468584e43b4fbbf4c4aa1d2bbac7.svg
Minqi Jiang
user
minqi
6556f0ff2de0aee1a11fab4c
Robert Tjarko Lange
6556f0ff2de0aee1a11fab4d
Shimon Whiteson
6556f0ff2de0aee1a11fab4e
Bruno Lacerda
6556f0ff2de0aee1a11fab4f
Nick Hawes
6556f0ff2de0aee1a11fab50
Tim Rocktaschel
6556f0ff2de0aee1a11fab51
Chris Lu
claimed_verified
2024-05-13T07:50:42.164Z
64cbfc1d3576a06fa252736d
/avatars/7345beafa590d5c9d5ac80cb33a9f0ab.svg
Chris Lu
user
chrlu
6556f0ff2de0aee1a11fab52
Jakob Nicolaus Foerster
JaxMARL: Multi-Agent RL Environments in JAX
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.
JaxMARL is an open-source codebase for reinforcement learning that uses hardware acceleration to enable fast, parallel training of multi-agent reinforcement learning environments, and introduces a vectorized version of the StarCraft Multi-Agent Challenge.
reinforcement learning JAX multi-agent reinforcement learning GPU acceleration StarCraft Multi-Agent Challenge MARL environments baseline algorithms self-play meta-learning