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