ifaamas [AAMAS ’24]Proc. of the 23rd International Conference on Autonomous
Agents and Multiagent Systems (AAMAS 2024)May 6 – 10, 2024 Auckland, New
ZealandN. Alechina, V. Dignum, M. Dastani, J.S. Sichman (eds.) 2024 2024 768

# JaxMARL: Multi-Agent RL Environments in JAX

Alexander Rutherford1111Core Contributor222Equal Contribution. Corresponding
Author<EMAIL_ADDRESS>Full authorship contribution
statements appear in the end of the document (Section 8)., Benjamin
Ellis1111Core Contributor222Equal Contribution. Corresponding Author:
<EMAIL_ADDRESS>Full authorship contribution statements appear
in the end of the document (Section 8)., Matteo Gallici2111Core
Contributor222Equal Contribution. Corresponding Author:
<EMAIL_ADDRESS>Full authorship contribution statements appear
in the end of the document (Section 8)., Jonathan Cook1111Core Contributor,
Andrei Lupu1111Core Contributor, Garðar Ingvarsson3111Core Contributor,
Timon Willi1111Core Contributor, Akbir Khan3, Christian Schroeder de Witt1,
Alexandra Souly3, Saptarashmi Bandyopadhyay4,
Mikayel Samvelyan3, Minqi Jiang3, Robert Tjarko Lange5, Shimon Whiteson1,
Bruno Lacerda1, Nick Hawes1,
Tim Rocktäschel3, Chris Lu1111Core Contributor222Equal Contribution.
Corresponding Author<EMAIL_ADDRESS>Full authorship
contribution statements appear in the end of the document (Section 8)., Jakob
Nicolaus Foerster1
1University of Oxford 2Universitat Politècnica de Catalunya 3UCL 4University
of Maryland 5Technical University Berlin

###### 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.

###### Key words and phrases:

Multi-Agent Reinforcement Learning, JAX, Benchmarks

##  1\. Introduction

Benchmarks play a pivotal role in the development of new single and multi-
agent reinforcement learning (MARL) algorithms by defining problems, enabling
comparisons, and focusing efforts. For example, in recent years, Go and Chess
drove the development of MuZero Schrittwieser et al. (2020) while
decentralised StarCraft Micromanagement Foerster et al. (2017) and later the
StarCraft Multi-Agent Challenge (SMAC, Samvelyan et al., 2019) resulted in the
development of algorithms such as QMIX Rashid et al. (2020), a popular MARL
technique.

Data transfer between the CPU (where the environment is simulated) and the GPU
(where the agents are evaluated) is a crucial bottleneck for simulation speed.
Simulation speed in turn is vital for progress in reinforcement learning (RL)
because RL algorithms often require a large number of environment
interactions. This problem is even worse in MARL, where non-stationarity and
decentralised partial observability greatly worsen the sample complexity
Bernstein et al. (2002). Hardware acceleration and parallelisation are crucial
to alleviating this, but current acceleration and parallelisation methods are
typically not implemented in Python, reducing their accessibility for most
machine learning researchers Shacklett et al. (2023); Weng et al. (2022). For
example, the extremely efficient Hanabi library Hu and Foerster (2020) from
Meta-AI research is implemented in C++ and has seen relatively little adoption
by the community. However, recent advances in JAX Bradbury et al. (2018) have
opened up new possibilities for using Python code directly with hardware
accelerators, enabling the wider use of massively parallel RL training
pipelines and environments.

Figure 1. JaxMARL’s philosophy. JaxMARL combines a wide range of environments
with ease of use and evaluation speed.

(a) MPE

(b) Overcooked

(c) Multi-Agent Brax

(d) STORM

(e) Hanabi

(f) Switch Riddle

(g) Coin Game

(h) SMAX

Figure 2. JaxMARL environments. We provide vectorised implementations of a
wide range of environments from different MARL settings.

The JAX Bradbury et al. (2018) library provides composable function
transformations, allowing for automatic vectorisation, device parallelisation,
automatic differentiation and just-in-time (JIT) compilation with XLA Sabne
(2020), for device-agnostic optimisation. Using JAX, both the environment
rollouts and model training can happen on a hardware accelerator (such as a
GPU or TPU), removing the cost of data transfer between devices and allowing
for significant parallelisation. Recently, PureJaxRL Lu et al. (2022a, 2023b)
has demonstrated the power of this end-to-end JAX-based approach; running both
the environment and the model training on a GPU yields a 4000x speedup over a
“traditional” pipeline with a GPU-trained policy but a CPU-based environment.

These accelerations could substantially advance RL and MARL research by
quickening the testing and iteration of ideas. Furthermore, they lower
computational hurdles for in-depth MARL research, enabling researchers to
utilise billions of frames and extract more performance from single GPUs.

Alongside the current computational issues faced by MARL researchers, recent
work also highlights issues with the evaluation standards and use of
benchmarks in the MARL community. In particular, MARL papers typically only
test on a few domains. Of the 75 recent MARL papers analysed by Gorsane et al.
(2022), 50% used only one evaluation environment and a further 30% used only
two. While SMAC and MPE Lowe et al. (2017), the two most used environments,
have various tasks or maps, the lack of a standard set raises the risk of
biased comparisons and incorrect conclusions. This leads to environment
overfitting and unclear progress markers.

Instead, novel MARL methods should be tested on a wide range of domains to
accurately evaluate their limits and enable better comparisons. The likely
issue preventing this is the lack of a unified codebase and the computational
burden of further evaluation.

This paper presents JaxMARL, a Python library that for the first time brings
together JAX implementations of eight common MARL environments under one API.
We additionally provide JAX implementations for four state-of-the-art
algorithms, allowing for end-to-end JAX-based training pipelines in a similar
fashion to PureJaxRL. As outlined in Figure 1, we present a library with end-
to-end hardware-accelerated training, simple Python implementations, and a
broad range of MARL environments. By alleviating computational constraints,
JaxMARL allows rapid evaluation of novel methods across a broad set of
domains, and hence has the potential to be a powerful tool to address MARL’s
evaluation crisis. Specifically, we find that JaxMARL achieves over 12500x
speedup compared to “conventional” aproaches.

We also create SMAX, a JAX-based simplification of the centralised training
with decentralised execution (CTDE) benchmarks SMAC Samvelyan et al. (2019)
and SMACv2 Ellis et al. (2022). SMAX features simplified dynamics, greater
flexibility and a more sophisticated but fully-decentralised heuristic AI,
while retaining the high-dimensional observation space, complex unit type
interactions and procedural scenario generation that lend SMAC and SMACv2 much
of their difficulty.

As shown in Figure 2, in addition to SMAX, our library includes the most
popular environments from several MARL settings. For centralised training with
decentralised execution (CTDE), we include the Multi-Agent Particle
Environments (MPE) Lowe et al. (2017), and Multi-Agent Brax (MABrax).
Meanwhile, for zero-shot coordination (ZSC) and ad-hoc teamplay, we include
Hanabi and Overcooked. Lastly, from the general-sum literature, we include the
CoinGame and Spatial-Temporal Representations of Matrix Games (STORM), a
representation of matrix games as grid-world scenarios with temporally
extended actions. JaxMARL provides the first JAX implementation of these
environments and unifies them in a single codebase.

We additionally provide JAX implementations of Independent PPO (IPPO) Schulman
et al. (2017); de Witt et al. (2020), QMIX, VDN Sunehag et al. (2017) and
Independent $Q$-Learning (IQL) Mnih et al. (2015), four of the most common
MARL algorithms, allowing new techniques to be easily benchmarked against
existing practices. We will extend this list before the camera-ready copy,
e.g. with the popular MAPPO Yu et al. (2022) algorithm.

##  2\. Background

###  2.1. Hardware Accelerated Environments

JAX enables the use of Python code with any hardware accelerator, allowing
researchers to write hardware-accelerated code easily. Within the RL
community, writing environment code in JAX has gained recent popularity. This
brings two chief advantages: firstly, environments written in JAX can be very
easily parallelised by using JAX’s vmap operation, which vectorises a function
across an input dimension, and secondly writing the environment in JAX allows
the agent and environment to be co-located on the GPU, which eliminates the
time taken to copy between CPU and GPU memory. Combined, these two factors
bring significant increases in training speed, with PureJaxRL Lu et al.
(2022a) achieving a $4000$x speedup over traditional training in single-agent
settings.

###  2.2. SMAC

StarCraft is a popular environment for testing RL algorithms. It typically
features features a centralised controller issuing commands to balance
_micromanagement_ , the low-level control of individual units, and
_macromanagement_ , the high level plans for economy and resource management.

SMAC Samvelyan et al. (2019), instead, focuses on decentralised unit
micromanagement across a range of scenarios divided into three broad
categories: _symmetric_ , where each side has the same units, _asymmetric_ ,
where the enemy team has more units, and _micro-trick_ , which are scenarios
designed specifically to feature a particular StarCraft micromanagement
strategy. SMACv2 Ellis et al. (2022) demonstrates that open-loop policies can
be effective on SMAC and adds additional randomly generated scenarios to
rectify SMAC’s lack of stochasticity. However, both of these environments rely
on running the full game of StarCraft II, which severely increases their CPU
and memory requirements. SMAClite Michalski et al. (2023) attempts to
alleviate this computational burden by recreating the SMAC environment
primarily in NumPy, with some core components written in C++. While this is
much more lightweight than SMAC, it cannot be run on a GPU and therefore
cannot be parallelised effectively with typical academic hardware, which
commonly has very few CPU cores compared to industry clusters.

##  3\. JaxMARL

We present JaxMARL, a library containing simple and accessible JAX
implementations of popular MARL environments and algorithms. JAX enables
significant acceleration and parallelisation over existing implementations. To
the best of our knowledge, JaxMARLis the first open source library that
provides JAX-based implementations of a wide range of MARL environments and
baselines.

###  3.1. API

The interface of JaxMARL is inspired by PettingZoo Terry et al. (2021) and
Gymnax. We designed it to be a simple and easy-to-use interface for a wide-
range of MARL problems. An example of instantiating an environment from
JaxMARL’s registry and executing one transition is presented in Figure 3.

⬇

1import jax

2from jaxmarl import make

3

4key = jax.random.PRNGKey(0)

5key, key_reset, key_act, key_step = jax.random.split(key, 4)

6

7# Initialise and reset the environment.

8env = make(’MPE_simple_world_comm_v3’)

9obs, state = env.reset(key_reset)

10

11# Sample random actions.

12key_act = jax.random.split(key_act, env.num_agents)

13actions = {agent: env.action_space(agent).sample(key_act[i]) \

14 for i, agent in enumerate(env.agents)}

15

16# Perform the step transition.

17obs, state, reward, done, infos = env.step(key_step, state, actions)

Figure 3. An example of JaxMARL’s API, which is flexible and easy-to-use.

As JAX’s JIT compilation requires pure functions, our `step` method has two
additional inputs compared to PettingZoo’s. The `state` object stores the
environment’s internal state and is updated with each call to `step`, before
being passed to subsequent calls. Meanwhile, `key_step` is a pseudo-random
key, consumed by JAX functions that require stochasticity. This key is
separated from the internal state for clarity.

Similar to PettingZoo, the remaining inputs and outputs are dictionaries keyed
by agent names, allowing for differing action and observation spaces. However,
as JAX’s JIT compilation requires arrays to have static shapes, the total
number of agents in an environment cannot vary during an episode. Thus, we do
not use PettingZoo’s agent iterator. Instead, the maximum number of agents is
set upon environment instantiation and any agents that terminate before the
end of an episode pass dummy actions thereafter. As asynchronous termination
is possible, we signal the end of an episode using a special `"__all__"` key
within `done`. The same dummy action approach is taken for environments where
agents act asynchronously (e.g. turn-based games).

To ensure clarity and reproducibility, we keep strict registration of
environments with suffixed version numbers, for example “MPE Simple Spread
V3”. Whenever JaxMARL environments correspond to existing CPU-based
implementations, the version numbers match.

###  3.2. Environments

JaxMARL contains a diverse range of environments, all implemented in JAX. We
also introduce SMAX, a SMAC-like environment implemented entirely in JAX. In
this section we introduce these environments and provide details on their
implementations.

#### SMAX

The StarCraft Multi-Agent Challenge (SMAC) is a popular benchmark but has a
number of shortcomings. First, as noted and addressed in prior work Ellis et
al. (2022), it is not sufficiently stochastic to require complex closed-loop
policies. Additionally, SMAC relies on StarCraft II as a simulator. While this
allows SMAC to use the wide range of units, objects and terrain available in
StarCraft II, running an entire instance of StarCraft II is slow Michalski et
al. (2023) and memory intensive. StarCraft II runs on the CPU and therefore
SMAC’s parallelisation is severely limited with typical academic compute.

Table 1. SMAX scenarios. The first section corresponds to SMAC scenarios, while the second corresponds to SMACv2. Scenario | Ally Units | Enemy Units | Start Positions
---|---|---|---
2s3z | 2 stalkers and 3 zealots | 2 stalkers and 3 zealots | Fixed
3s5z | 3 stalkers and 5 zealots | 3 stalkers and 5 zealots | Fixed
5m_vs_6m | 5 marines | 6 marines | Fixed
10m_vs_11m | 10 marines | 11 marines | Fixed
27m_vs_30m | 27 marines | 30 marines | Fixed
3s5z_vs_3s6z | 3 stalkers and 5 zealots | 3 stalkers and 6 zealots | Fixed
3s_vs_5z | 3 stalkers | 5 zealots | Fixed
6h_vs_8z | 6 hydralisks | 8 zealots | Fixed
smacv2_5_units | 5 uniformly randomly chosen | 5 uniformly randomly chosen | SMACv2-style
smacv2_10_units | 10 uniformly randomly chosen | 10 uniformly randomly chosen | SMACv2-style
smacv2_20_units | 20 uniformly randomly chosen | 20 uniformly randomly chosen | SMACv2-style

Using the StarCraft II game engine constrains environment design. For example,
StarCraft II groups units into three races and does not allow units of
different races on the same team, limiting the variety of scenarios that can
be generated. Secondly, SMAC does not support a competitive self-play setting
without significant engineering work. The purpose of SMAX is to address these
limitations. It provides access to a SMAC-like, hardware-accelerated,
customisable environment that supports self-play and custom unit types.

Units in SMAX are modelled as circles in a two-dimensional continuous space.
SMAX makes a number of additional simplifications to the dynamics of StarCraft
II, details of which are given in Appendix A.1.

SMAX also features a different, and more sophisticated, heuristic AI. The
heuristic in SMAC simply moves to a fixed location Michalski et al. (2023),
attacking any enemies it encounters along the way, and the heuristic in SMACv2
globally pursues the nearest agent. Thus the SMAC AI often does not
aggressively pursue enemies that run away, and cannot generalise to the SMACv2
start positions, whereas the SMACv2 heuristic AI conditions on global
information and is exploitable because of its tendency to flip-flop between
two similarly close enemies. SMAC’s heuristic AI must be coded in the map
editor, which does not provide a simple coding interface.

In contrast, SMAX features a decentralised heuristic AI that can effectively
find enemies without requiring the global information of the SMACv2 heuristic.
This guarantees that in principle a 50% win rate is always achievable by
copying the decentralised heuristic policy exactly. This means any win-rate
below 50% represents a concrete failure to learn.

SMAX scenarios incorporate both a number of the original scenarios from SMAC
and scenarios similar to those found in SMACv2. The latter sample units
uniformly across all SMAX unit types (stalker, zealot, hydralisk, zergling,
marine, marauder) and ensure fairness by having identical team composition for
the enemy and ally teams. We provide more details on SMAX in Appendix A.1.

#### Overcooked

Inspired by the popular videogame of the same name, Overcooked is commonly
used for assessing fully cooperative and fully observable Human-AI task
performance. The aim is to quickly prepare and deliver soup, which involves
putting three onions in a pot, cooking the soup, and serving it into bowls.
Two agents, or cooks, must coordinate to effectively divide the tasks to
maximise their common reward signal. Our implementation mimics the original
from Overcooked-AI (Carroll et al., 2019), including all five original layouts
and a simple method for creating additional ones. For a discussion on the
limitations of the Overcooked-AI environment, see Lauffer et al. (2023).

#### Hanabi

Hanabi is a fully cooperative partially observable multiplayer card game,
where players can observe other players’ cards but not their own. To win, the
team must play a series of cards in a specific order while sharing only a
limited amount of information between players. As reasoning about the beliefs
and intentions of other agents is central to performance, it is a common
benchmark for ZSC and ad-hoc teamplay research. Our implementation is inspired
by the Hanabi Learning Environment Bard et al. (2020) and includes custom
configurations for varying game settings, such as the number of colours/ranks,
number of players, and number of hint tokens. Compared to the Hanabi Learning
Environment, which is written in C++ and split over dozens of files, our
implementation is a single easy-to-read Python file, which simplifies
interfacing with the library and running experiments.

#### Multi-Agent Particle Environments (MPE)

The multi-agent particle environments feature a 2D world with simple physics
where particle agents can move, communicate, and interact with fixed
landmarks. Each specific environment varies the format of the world and the
agents’ abilities, creating a diverse set of tasks that include both
competitive and cooperative settings. We implement all the MPE scenarios
featured in the PettingZoo library and the transitions of our implementation
map exactly to theirs. We additionally include a fully cooperative predator-
prey variant of simple tag, presented in Peng et al. (2021). The code is
structured to allow for straightforward extensions, enabling further tasks to
be added.

#### Multi-Agent Brax (MABrax)

MABrax is a derivative of Multi-Agent MuJoCo Peng et al. (2021), an extension
of the MuJoCo Gym environment Todorov et al. (2012) that is commonly used for
benchmarking continuous multi-agent robotic control. Our implementation
utilises BraxFreeman et al. (2021) as the underlying physics engine and
includes five of Multi-Agent MuJoCo’s multi-agent factorisation tasks, where
each agent controls a subset of the joints and only observes the local state.
The included tasks, illustrated in Figure 2, are: `ant_4x2`,
`halfcheetah_6x1`, `hopper_3x1`, `humanoid_9|8`, and `walker2d_2x3`. The task
descriptions mirror those from Gymnasium-Robotics de Lazcano et al. (2023).

Table 2. Benchmark results for JAX-based MARL environments (steps-per-second) when taking random actions. All environments are significantly faster than existing CPU implementations. Environment | Original, 1 Env | Jax, 1 Env | Jax, 100 Envs | Jax, 10k Envs | Maximum Speedup
---|---|---|---|---|---
MPE Simple Spread | $8.34\text{\times}{10}^{4}$ | $5.48\text{\times}{10}^{3}$ | $5.24\text{\times}{10}^{5}$ | $3.99\text{\times}{10}^{7}$ | $4.78\text{\times}{10}^{2}$
MPE Simple Reference | $1.46\text{\times}{10}^{5}$ | $5.24\text{\times}{10}^{3}$ | $4.85\text{\times}{10}^{5}$ | $3.35\text{\times}{10}^{7}$ | $2.29\text{\times}{10}^{2}$
Switch Riddle | $2.69\text{\times}{10}^{4}$ | $6.24\text{\times}{10}^{3}$ | $7.92\text{\times}{10}^{5}$ | $6.68\text{\times}{10}^{7}$ | $2.48\text{\times}{10}^{3}$
Hanabi | $2.10\text{\times}{10}^{3}$ | $1.36\text{\times}{10}^{3}$ | $1.05\text{\times}{10}^{5}$ | $5.02\text{\times}{10}^{6}$ | $2.39\text{\times}{10}^{3}$
Overcooked | $1.91\text{\times}{10}^{3}$ | $3.59\text{\times}{10}^{3}$ | $3.04\text{\times}{10}^{5}$ | $1.69\text{\times}{10}^{7}$ | $8.85\text{\times}{10}^{3}$
MABrax Ant 4x2 | $1.77\text{\times}{10}^{3}$ | $2.70\text{\times}{10}^{2}$ | $1.81\text{\times}{10}^{4}$ | $7.62\text{\times}{10}^{5}$ | $4.31\text{\times}{10}^{2}$
Starcraft 2s3z | $8.31\text{\times}{10}^{1}$ | $5.37\text{\times}{10}^{2}$ | $4.53\text{\times}{10}^{4}$ | $2.71\text{\times}{10}^{6}$ | $3.26\text{\times}{10}^{4}$
Starcraft 27m vs 30m | $2.73\text{\times}{10}^{1}$ | $1.45\text{\times}{10}^{2}$ | $1.12\text{\times}{10}^{4}$ | $1.90\text{\times}{10}^{5}$ | $6.96\text{\times}{10}^{3}$
STORM | – | $2.48\text{\times}{10}^{3}$ | $1.75\text{\times}{10}^{5}$ | $1.46\text{\times}{10}^{7}$ | –
Coin Game | $1.97\text{\times}{10}^{4}$ | $4.67\text{\times}{10}^{3}$ | $4.06\text{\times}{10}^{5}$ | $4.03\text{\times}{10}^{7}$ | $2.05\text{\times}{10}^{3}$

#### Coin Game

Coin Game is a two-player grid-world environment which emulates social
dilemmas such as the iterated prisoner’s dilemma Snyder (1971). Used as a
benchmark for the general-sum setting, it expands on simpler social dilemmas
by adding a high-dimensional state. Two players, ‘red’ and ‘blue’ move in a
grid world and are each awarded 1 point for collecting any coin. However,
‘red’ loses 2 points if ‘blue’ collects a red coin and vice versa. Thus, if
both agents ignore colour when collecting coins their expected reward is 0.
Further details are provided in Appendix A.2.

#### Spatial-Temporal Representations of Matrix Games (STORM)

Inspired by the “in the Matrix” games in Melting Pot 2.0 Agapiou et al.
(2022), the STORM (Khan et al., 2022) environment expands on matrix games by
representing them as grid-world scenarios. Agents collect resources which
define their strategy during interactions and are rewarded based on a pre-
specified payoff matrix. This allows for the embedding of fully cooperative,
competitive or general-sum games, such as the prisoner’s dilemma Snyder
(1971). Thus, STORM can be used for studying paradigms such as opponent
shaping, where agents act with the intent to change other agents’ learning
dynamics, which has been empirically shown to lead to more prosocial outcomes
(Foerster et al., 2018; Khan et al., 2022; Lu et al., 2022b; Zhao et al.,
2022). Compared to the Coin Game or matrix games, the grid-world setting
presents a variety of new challenges such as partial observability, multi-step
agent interactions, temporally-extended actions, and longer time horizons.
Unlike the “in the Matrix” games from Melting Pot, STORM features
stochasticity, increasing the difficulty Ellis et al. (2022). A further
environment specification is provided in Appendix A.3.

#### Switch Riddle

Originally used to illustrate the Differentiable Inter-Agent Learning
algorithm Foerster et al. (2016), Switch Riddle is a simple cooperative
communication environment that we include as a debugging tool. $n$ prisoners
held by a warden can secure their release by collectively ensuring that each
has passed through a room with a light bulb and a switch. Each day, a prisoner
is chosen at random to enter this room. They have three choices: do nothing,
signal to the next prisoner by toggling the light, or inform the warden they
think all prisoners have been in the room. The game ends when a prisoner
informs the warden or the maximum time steps are reached. The rewards are +1
if the prisoner informs the warden, and all prisoners have been in the room,
-1 if the prisoner informs the warden before all prisoners have taken their
turn, and 0 otherwise, including when the maximum time steps are reached. We
benchmark using the implementation from Zhang et al. (2022).

###  3.3. Algorithms

In this section, we present our re-implementation of four well known MARL
baseline algorithms using JAX. The primary objective of these baselines is to
provide a structured framework for developing MARL algorithms leveraging the
advantages of the JaxMARL environments. All of the training pipelines are
fully compatible with JAX’s JIT and VMAP functions, resulting in a significant
acceleration of both the training and metric evaluation processes. This
enables parallelisation of training across various seeds and hyperparameters
on a single machine in parallel. We follow the CleanRL philosophy of providing
clear, single-file implementations Huang et al. (2022).

#### IPPO

Our Independent PPO (IPPO) Schulman et al. (2017); de Witt et al. (2020)
implementation is based on PureJaxRL Lu et al. (2022a), with parameter sharing
across homogeneous agents. We provide both feed-forward and RNN versions.

####  $Q$-learning Methods

Our $Q$-Learning baselines, including Independent $Q$-Learning (IQL) Tampuu et
al. (2017), Value Decomposition Networks (VDN) Sunehag et al. (2018), and QMIX
Rashid et al. (2018), have been implemented in accordance with the PyMARL
codebase Rashid et al. (2018) to ensure consistency with published results and
enable direct comparisons with PyTorch. Our baselines natively support
aggregating trajectories from batched environments, simplifying
parallelisation. This approach is more convenient than managing environments
on distinct threads and subsequently aggregating results, as done in PyMARL.
We provide a brief overview of the implemented baselines in the Appendix.

##  4\. Results

In our results, we aim to demonstrate the speed and correctness of our
environments and algorithms.In several cases, minor changes to the
environments mean that our environments do not exactly match the originals on
a step-by-step level. We therefore demonstrate the correctness in different
ways for each environment and discuss each separately. By combining this
evidence, we demonstrate that our library provides overall correct and far
quicker baselines on a wide range of sufficiently correct and easily-
modifiable environments.

(a) Hanabi

(b) MABrax Ant

(c) Overcooked

(d) Starcraft 2s3z

Figure 4. Speedup of four JaxMARL environments compared to singled-threaded
CPU-based implementations.

(a) MPE Simple Spread Returns

(b) MPE Simple Spread Returns

(c) MPE Training Speed

(d) SMAX Training Speed

Figure 5. IPPO Speed and Performance in JaxMARL compared to MARLLIB and PyMARL
in SMAX and MPE. Return results were averaged across 3 seeds. Performance
results show 1 seed collected on the hardware described in Section 4.1.

###  4.1. Environment Speed

We measure the performance of our environments in steps per second when using
random actions and compare to the original environments in Table 2 and Figure
4. All results were collected on a single NVIDIA A100 GPU and AMD EPYC 7763
64-core processor. Environments were rolled out for 1000 sequential steps.
Many environments have comparable performance to JaxMARL when comparing single
environments, but the ease of parallelisation with Jax allows for more
efficient scaling compared to CPU-based environments. For example, MPE Simple
Spread’s JAX implementation is 2̃0x slower than the original when comparing a
single environment, but even when only running $100$ environments in parallel,
the JAX environment is already over $6$x faster. When considering $10000$
environments, the JAX versions are much faster, achieving speedups of up to
$8500$x over the single-threaded environment (in the case of Overcooked).
Running this many environments in parallel using CPU environments would
require a large CPU cluster and sophisticated communication mechanisms. This
engineering is typically beyond the resources of academic labs, and therefore
JaxMARL can unlock new research directions for such institutions.

###  4.2. Algorithm Speed

We investigate the speed of our IPPO implementation in Figure 5. By
vectorising over agents, it is possible to train a vast number of agents in a
fraction of the time it takes to train a single agent without hardware
acceleration. For MPE, it is possible to train 1024 teams in $198.4$ seconds,
which is less than $0.2$ seconds per teams of agents. A single run of
MARLLIB’s IPPO implementation on the same hardware takes around $2435.7$
seconds on average. This represents a speedup of over $12500$x.

Our JAX-based $Q$-learning algorithms also offer significant speed advantages.
In Figure 6(a), training a single IQL, VDN, or QMIX policy in MPE takes $\sim
130$ seconds while using PyMarl takes over an hour. Training 1024 QMIX
learners in a batch requires $1670$ seconds, which translates to $1.6$ seconds
per learner, indicating a $2700$x speedup. This speedup is not as large as for
IPPO because $Q$-learning baselines are typically trained with fewer parallel
environments. In our experiments, we used 8 parallel environments for
$Q$-learning compared to the 25 or 64 used for PPO. This difference is due to
$Q$-learners benefiting more from a buffer with trajectories collected by
different policies, resulting in a more frequent policy update, rather than
collecting many trajectories with the same policy in parallel.

For SMAX, we compare our vectorised IPPO baseline to the MAPPO implementation
provided in Sun et al. (2023). MAPPO utilises an RNN and IPPO uses a feed
forward network. This was run on a machine with a 64-core CPU and NVIDIA
2080Ti GPU. Additionally, as discussed in Section 3.2, SMAC and SMAX are
different environments. These caveats aside, the differences in performance
are so striking that we believe this clearly demonstrates the advantages of
our approach. We trained 512 SMAX teams on 2s3z in under 33 minutes, whereas a
single training run of PyTorch IPPO implementation takes 44 hours on average.
This is roughly a $40000$x speedup.

(a) Simple Spread Training Time

(b) Simple Spread Returns

(c) Speaker-Listener Returns

(d) QMIX Training Speed

Figure 6. Performance and speed of JaxMARL $Q$-Learning baselines compared to
PyMARL on MPE. Our implementations match PyMARL’s returns, while being over
$2000$x faster to train

###  4.3. Algorithm Correctness

We verify the correctness of our algorithm implementations by comparing to
baselines from other libraries on the MPE Simple Spread and Simple Speaker
Listener environments. For IPPO we report the mean return across $3$ seeds in
Figure 5(b). Results were collected on the same hardware as listed in Section
4.1. Our IPPO implementation obtains the same performance as MARLLIB and runs
$250$x quicker, taking only ten seconds to train.

For the $Q$-learning algorithms, we verify the correctness by comparing with
PyMARL implementations of the same algorithms on the MPE Simple Spread and
Simple Speaker Listener environments. IQL, VDN and QMIX all obtain the same or
better results than their PyMARL counterparts. The returns are from greedy
policies and averaged across 8 runs. The hyperparameters used are from the
PyMARL library.

###  4.4. Environment Correctness

#### MPE

Our MPE environment corresponds exactly to the PettingZoo implementation. We
validate this for each environment using a uniform-random policy on $1000$
rollouts, ensuring all observations and rewards are within a tolerance of
$1\times 10^{-4}$ at each transition. This tolerance accounts for non-
determinism due to running floating point computation on the GPU. The
correspondence is also shown through the performance of IPPO in Figure 5(b)
and the $Q$-learning algorithms in Figures 6(b) and 6(c) respectively, as the
performance of these algorithms is inline with existing baselines Yu et al.
(2022). We additionally report training performance for IQL on the remaining
MPE environments in Appendix C.2.

#### Overcooked

The transition dynamics of our Overcooked implementation match those of the
Overcooked-AI implementation. We demonstrate this by training an IPPO policy
on our implementation and evaluating the policy on both our Overcooked
implementation and the original at regular intervals. Results are illustrated
in Figure 7(a) and performance is similar, demonstrating their equivalence.

#### SMAX

SMAX and SMAC are different environments. However, we demonstrate some
similarity between them by comparing our IPPO and MAPPO implementations
against MAPPO results on SMAC, using the implementation from Sun et al.
(2023). We show this in Figure 8. SMAX and SMAC have different opponent
policies and dynamics, which makes this comparison more qualitative than
precise. We describe the differences between the two in more depth in in the
supplementary material. However, despite these differences, the environments
seem similarly difficult, with some environments being more difficult in SMAC,
and some more difficult in SMAX. This is shown in Figure 8 and in the
supplementary material.

#### MABrax

As Brax differs subtly from MuJoCo, MABrax does not correspond to MAMuJoCo but
the learning dynamics are qualitatively similar. To demonstrate this, we
report mean training return across 10 seeds for IPPO on `ant_4x2` in Figure
7(b), and our results are in line with the performance of TRPO reported in
Kuba et al. (2021). We report the performance of IPPO on HalfCheetah and
Walker in Appendix C.1, the results are also in line with TRPO.

(a) Overcooked

(b) MABrax Ant

(c) 2 Player Hanabi

Figure 7. JaxMARL IPPO baseline results. These results correspond to similar baselines and therefore demonstrate the correctness of our implementations. Figure 8. SMAX IPPO and MAPPO baselines compared to MAPPO in SMAC. Table 3. Recommended Minimal Environment Evaluations for different research settings Setting | Recommended Environments
---|---
CTDE | SMAX (all scenarios), Hanabi (2-5 players), Overcooked
Zero-shot Coordination | Hanabi (2 players), Overcooked (5 basic scenarios)
General-Sum | STORM (iterated prisoner’s dilemma), STORM (matching pennies)
Cooperative Continuous Actions | MABrax

#### Hanabi

Our implementation does not correspond exactly to the Hanabi Learning
Environment as we use a subtly different observation space, with the reasoning
given in Appendix A.4. To demonstrate qualitative similarity, we train IPPO on
Hanabi in self-play with 2 players, with the mean test return across 3 seeds
reported in Figure 7(c).

#### STORM, Coin Game & Switch Riddle

STORM differs from Melting Pot 2.0 significantly, making direct comparisons
challenging, with differences discussed in Appendix A.3. Furthermore, STORM
and Coin Game are general-sum games, so the environment returns of IPPO in
self-play would not be a good indicator of performance. Switch Riddle is a
simple diagnostic environment – we do not use it for thorough evaluations.

##  5\. Evaluation Recommendations

Previous work Gorsane et al. (2022) has found significant differences in the
evaluation protocols between MARL research works. We identify four main
research areas that would benefit from our library: cooperative centralised
training with decentralised execution (CTDE) Foerster et al. (2016), zero-shot
coordination Hu et al. (2020), general-sum games, and cooperative continuous
action methods.

To aid comparisons between methods, we recommend standard _minimal_ sets of
evaluation environments for each of these settings in Table 3. It’s important
to note that these are _minimal_ and we encourage as broad an evaluation as
possible. For example, in the zero-shot coordination setting, all methods
should be able to evaluate on Hanabi and Overcooked. However, it may also be
possible to evaluate such methods on the SMACv2 settings of SMAX. Similarly,
SMAX could be used to evaluate two-player zero-sum methods by training in
self-play. For some settings, such as continuous action environments and
general-sum games, there is only one difficult environment. We encourage
further development of JAX-based environments in these settings to improve the
quality of evaluation.

##  6\. Related Work

Several open-source libraries exist for both MARL algorithms and environments.
The popular library PyMARL Samvelyan et al. (2019) provides PyTorch
implementations of QMIX, VDN and IQL and integrates easily with SMAC. E-PyMARL
Papoudakis et al. (2021) extends this by adding the actor-critic algorithms
MADDPG Lowe et al. (2017), MAA2C Mnih et al. (2016), IA2C Mnih et al. (2016),
and MAPPO, and supports the SMAC, Gym Brockman et al. (2016), Robot Warehouse
Christianos et al. (2020), Level-Based Foraging Christianos et al. (2020), and
MPE environments. Recently released MARLLib Hu et al. (2022) is instead based
on the open-source RL library RLLib Liang et al. (2018) and combines a wide
range of competitive, cooperative and mixed environments with a broad set of
baseline algorithms. Meanwhile, MALib Zhou et al. (2023) focuses on
population-based MARL across a wide range of environments. However, none of
these frameworks feature hardware-accelerated environments and thus lack the
associated performance benefits.

There has also been a recent proliferation of hardware-accelerated and JAX-
based RL environments. Isaac gym Makoviychuk et al. (2021) provides a GPU-
accelerated simulator for a range of robotics platforms and CuLE Dalton and
frosio (2020) is a CUDA reimplementation of the Atari Learning Environment
Bellemare et al. (2013). Both of these environments are GPU-specific and
cannot be extended to other hardware accelerators. Madrona Shacklett et al.
(2023) is an extensible game-engine written in C++ that allows for GPU
acceleration and parallelisation across environments. However, it requires
environment code to be written in C++, limiting its accessibility. VMAS
Bettini et al. (2022) provides a vectorized 2D physics engine written in
PyTorch and a set of challenging multi-robot scenarios, including those from
the MPE environment. For RL environments implemented in JAX, Jumanji Bonnet et
al. (2023) features mostly single-agent environments with a strong focus on
combinatorial problems. The authors also provide an actor-critic baseline in
addition to random actions. PGX Koyamada et al. (2023) includes several board-
game environments written in JAX. Gymnax Lange (2022) provides JAX
implementations of the BSuite Osband et al. (2019), classic continuous
control, MinAtar Young and Tian (2019) and other assorted environments.
Gymnax’s sister-library, gymnax-baselines, provides PPO and ES baselines.
Further extensions to Gymnax Lu et al. (2023a) also include POPGym
environments Morad et al. (2023). Brax Freeman et al. (2021) reimplements the
MuJoCo simulator in JAX and also provides a PPO implementation as a baseline.
Jax-LOB Frey et al. (2023) implements a vectorized limit order book as an RL
environment that runs on the accelerator. Perhaps the most similar to our work
is Mava Pretorius et al. (2021), which provides a MAPPO baseline, as well as
integration with the Robot Warehouse environment. However, none of these
libraries combine a range of JAX-based MARL environments with both value-based
and actor-critic baselines.

Broadly, no other work provides implementations of a wide range of hardware-
accelerated MARL environments, while also implementing value-based and actor-
critic baselines. Secondly, no other JAX simplification of SMAC exists. All
other versions are either tied to the StarCraft II simulator or not hardware
accelerated.

##  7\. Conclusion

Hardware acceleration offers important opportunities for MARL research by
lowering computational barriers, increasing the speed at which ideas can be
iterated, and allowing for more thorough evaluation. We present JaxMARL, an
open-source library of popular MARL environments and baseline algorithms
implemented in JAX. We combine ease of use with hardware accelerator enabled
efficiency to give significant speed-ups compared to traditional CPU-based
implementations. Furthermore, by bringing together a wide range of MARL
environments under one codebase, we have the potential to help alleviate
issues with MARL’s evaluation standards. We hope that JaxMARL will help
advance MARL by improving the ability of academic labs to conduct research
with thorough, fast, and effective evaluations.

##  8\. Author Contributions

This project is a large-scale effort spanning many labs and contributors.

AR11footnotemark: 122footnotemark: 2 led the design of the JaxMARL API and
interface the implementation of IPPO and MPE environments. BE11footnotemark:
122footnotemark: 2 led the design and implementation of the SMAX environments
and IPPO evaluations. AR and BE also led the writing of this manuscript.
MG11footnotemark: 122footnotemark: 2 led the implementation of the off-policy
MARL algorithms, their evaluations., and the implementation of the Switch
Riddle environment.

JC11footnotemark: 1 led the implementation of the Hanabi environment and
heavily assisted with benchmarking and verifying its performance.
AL11footnotemark: 1 led the implementation of the Overcooked environments.
GI11footnotemark: 1 led the implementation of the Multi-Agent Brax
environments. TW11footnotemark: 1 led the implementation of the STORM
environments. AK and AS worked on the STORM environments. CSW led the
implementation of the Predator-Prey environment.

CSW, SB, MS, MJ, and RL provided invaluable discussions for project planning
and implementations across the project. SB helped initiate the project plan.
MS worked on the Multi-Agent Brax environments. MJ worked on the Overcooked
and Hanabi environments. RL assisted with the design of the API and testing
infrastructure.

SW, BL, NH, and TR provided invaluable feedback on the project, manuscript,
and results.

CL11footnotemark: 122footnotemark: 2 initiated the project and led the
organizational and planning efforts, speed-based benchmarking, and Coin Game
implementation.

JF is the primary advisor for the project.

## References

  * (1)
  * Agapiou et al. (2022) John P Agapiou, Alexander Sasha Vezhnevets, Edgar A Duéñez-Guzmán, Jayd Matyas, Yiran Mao, Peter Sunehag, Raphael Köster, Udari Madhushani, Kavya Kopparapu, Ramona Comanescu, et al. 2022\.  Melting Pot 2.0.  _arXiv preprint arXiv:2211.13746_ (2022).
  * Bard et al. (2020) Nolan Bard, Jakob N Foerster, Sarath Chandar, Neil Burch, Marc Lanctot, H Francis Song, Emilio Parisotto, Vincent Dumoulin, Subhodeep Moitra, Edward Hughes, et al. 2020\.  The hanabi challenge: A new frontier for ai research.  _Artificial Intelligence_ 280 (2020), 103216.
  * Bellemare et al. (2013) M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. 2013.  The Arcade Learning Environment: An Evaluation Platform for General Agents.  _Journal of Artificial Intelligence Research_ 47 (jun 2013), 253–279.
  * Bernstein et al. (2002) Daniel S Bernstein, Robert Givan, Neil Immerman, and Shlomo Zilberstein. 2002.  The complexity of decentralized control of Markov decision processes.  _Mathematics of operations research_ 27, 4 (2002), 819–840.
  * Bettini et al. (2022) Matteo Bettini, Ryan Kortvelesy, Jan Blumenkamp, and Amanda Prorok. 2022.  VMAS: A Vectorized Multi-Agent Simulator for Collective Robot Learning.  _The 16th International Symposium on Distributed Autonomous Robotic Systems_ (2022).
  * Bonnet et al. (2023) Clément Bonnet, Daniel Luo, Donal Byrne, Shikha Surana, Vincent Coyette, Paul Duckworth, Laurence I. Midgley, Tristan Kalloniatis, Sasha Abramowitz, Cemlyn N. Waters, Andries P. Smit, Nathan Grinsztajn, Ulrich A. Mbou Sob, Omayma Mahjoub, Elshadai Tegegn, Mohamed A. Mimouni, Raphael Boige, Ruan de Kock, Daniel Furelos-Blanco, Victor Le, Arnu Pretorius, and Alexandre Laterre. 2023.  Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX.  arXiv:2306.09884 [cs.LG] https://arxiv.org/abs/2306.09884
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018.  _JAX: composable transformations of Python+NumPy programs_.  http://github.com/google/jax
  * Brockman et al. (2016) Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. 2016.  OpenAI Gym.  arXiv:arXiv:1606.01540
  * Carroll et al. (2019) Micah Carroll, Rohin Shah, Mark K Ho, Tom Griffiths, Sanjit Seshia, Pieter Abbeel, and Anca Dragan. 2019.  On the utility of learning about humans for human-ai coordination.  _Advances in neural information processing systems_ 32 (2019).
  * Christianos et al. (2020) Filippos Christianos, Lukas Schäfer, and Stefano V Albrecht. 2020.  Shared Experience Actor-Critic for Multi-Agent Reinforcement Learning. In _Advances in Neural Information Processing Systems (NeurIPS)_.
  * Dalton and frosio (2020) Steven Dalton and iuri frosio. 2020.  Accelerating Reinforcement Learning through GPU Atari Emulation. In _Advances in Neural Information Processing Systems_ , H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (Eds.), Vol. 33. Curran Associates, Inc., 19773–19782.  https://proceedings.neurips.cc/paper/2020/file/e4d78a6b4d93e1d79241f7b282fa3413-Paper.pdf
  * de Lazcano et al. (2023) Rodrigo de Lazcano, Kallinteris Andreas, Jun Jet Tai, Seungjae Ryan Lee, and Jordan Terry. 2023.  _Gymnasium Robotics_.  http://github.com/Farama-Foundation/Gymnasium-Robotics
  * de Witt et al. (2020) Christian Schroeder de Witt, Tarun Gupta, Denys Makoviichuk, Viktor Makoviychuk, Philip H. S. Torr, Mingfei Sun, and Shimon Whiteson. 2020.  Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?  https://doi.org/10.48550/arXiv.2011.09533 arXiv:2011.09533 [cs].
  * Ellis et al. (2022) Benjamin Ellis, Skander Moalla, Mikayel Samvelyan, Mingfei Sun, Anuj Mahajan, Jakob N Foerster, and Shimon Whiteson. 2022.  SMACv2: An improved benchmark for cooperative multi-agent reinforcement learning.  _arXiv preprint arXiv:2212.07489_ (2022).
  * Foerster et al. (2016) Jakob Foerster, Ioannis Alexandros Assael, Nando de Freitas, and Shimon Whiteson. 2016.  Learning to Communicate with Deep Multi-Agent Reinforcement Learning. In _Advances in Neural Information Processing Systems_ , D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (Eds.), Vol. 29. Curran Associates, Inc.  https://proceedings.neurips.cc/paper_files/paper/2016/file/c7635bfd99248a2cdef8249ef7bfbef4-Paper.pdf
  * Foerster et al. (2018) Jakob Foerster, Richard Y Chen, Maruan Al-Shedivat, Shimon Whiteson, Pieter Abbeel, and Igor Mordatch. 2018.  Learning with Opponent-Learning Awareness. In _Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems_. 122–130.
  * Foerster et al. (2017) Jakob Foerster, Nantas Nardelli, Gregory Farquhar, Triantafyllos Afouras, Philip HS Torr, Pushmeet Kohli, and Shimon Whiteson. 2017.  Stabilising experience replay for deep multi-agent reinforcement learning. In _International conference on machine learning_. PMLR, 1146–1155.
  * Freeman et al. (2021) C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. 2021.  _Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation_.  http://github.com/google/brax
  * Frey et al. (2023) Sascha Frey, Kang Li, Peer Nagy, Silvia Sapora, Chris Lu, Stefan Zohren, Jakob Foerster, and Anisoara Calinescu. 2023.  JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for trading.  _arXiv preprint arXiv:2308.13289_ (2023).
  * Gorsane et al. (2022) Rihab Gorsane, Omayma Mahjoub, Ruan de Kock, Roland Dubb, Siddarth Singh, and Arnu Pretorius. 2022.  Towards a Standardised Performance Evaluation Protocol for Cooperative MARL.  _arXiv preprint arXiv:2209.10485_ (2022).
  * Hu and Foerster (2020) Hengyuan Hu and Jakob N Foerster. 2020.  Simplified Action Decoder for Deep Multi-Agent Reinforcement Learning. In _International Conference on Learning Representations_.  https://openreview.net/forum?id=B1xm3RVtwB
  * Hu et al. (2020) Hengyuan Hu, Adam Lerer, Alex Peysakhovich, and Jakob Foerster. 2020.  “other-play” for zero-shot coordination. In _International Conference on Machine Learning_. PMLR, 4399–4410.
  * Hu et al. (2022) Siyi Hu, Yifan Zhong, Minquan Gao, Weixun Wang, Hao Dong, Zhihui Li, Xiaodan Liang, Xiaojun Chang, and Yaodong Yang. 2022.  MARLlib: Extending RLlib for Multi-agent Reinforcement Learning.  _arXiv preprint arXiv:2210.13708_ (2022).
  * Huang et al. (2022) Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and João G.M. Araújo. 2022.  CleanRL: High-quality Single-file Implementations of Deep Reinforcement Learning Algorithms.  _Journal of Machine Learning Research_ 23, 274 (2022), 1–18.  http://jmlr.org/papers/v23/21-1342.html
  * Ierusalimschy (2006) Roberto Ierusalimschy. 2006.  _Programming in lua_.  Roberto Ierusalimschy.
  * Khan et al. (2022) Akbir Khan, Newton Kwan, Timon Willi, Chris Lu, Andrea Tacchetti, and Jakob Nicolaus Foerster. 2022.  Context and History Aware Other-Shaping.  (2022).
  * Koyamada et al. (2023) Sotetsu Koyamada, Shinri Okano, Soichiro Nishimori, Yu Murata, Keigo Habara, Haruka Kita, and Shin Ishii. 2023.  Pgx: Hardware-accelerated Parallel Game Simulators for Reinforcement Learning.  _arXiv preprint arXiv:2303.17503_ (2023).
  * Kuba et al. (2021) Jakub Grudzien Kuba, Ruiqing Chen, Muning Wen, Ying Wen, Fanglei Sun, Jun Wang, and Yaodong Yang. 2021.  Trust region policy optimisation in multi-agent reinforcement learning.  _arXiv preprint arXiv:2109.11251_ (2021).
  * Lange (2022) Robert Tjarko Lange. 2022.  _gymnax: A JAX-based Reinforcement Learning Environment Library_.  http://github.com/RobertTLange/gymnax
  * Lauffer et al. (2023) Niklas Lauffer, Ameesh Shah, Micah Carroll, Michael D Dennis, and Stuart Russell. 2023.  Who Needs to Know? Minimal Knowledge for Optimal Coordination. In _International Conference on Machine Learning_. PMLR, 18599–18613.
  * Liang et al. (2018) Eric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gonzalez, Michael Jordan, and Ion Stoica. 2018.  RLlib: Abstractions for distributed reinforcement learning. In _International conference on machine learning_. PMLR, 3053–3062.
  * Lowe et al. (2017) Ryan Lowe, Yi Wu, Aviv Tamar, Jean Harb, Pieter Abbeel, and Igor Mordatch. 2017.  Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments.  _Neural Information Processing Systems (NIPS)_ (2017).
  * Lu et al. (2022a) Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. 2022a.  Discovered policy optimisation.  _Advances in Neural Information Processing Systems_ 35 (2022), 16455–16468.
  * Lu et al. (2023a) Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. 2023a.  Structured State Space Models for In-Context Reinforcement Learning.  _arXiv e-prints_ (2023), arXiv–2303.
  * Lu et al. (2022b) Christopher Lu, Timon Willi, Christian A Schroeder De Witt, and Jakob Foerster. 2022b.  Model-Free Opponent Shaping. In _International Conference on Machine Learning_. PMLR, 14398–14411.
  * Lu et al. (2023b) Chris Lu, Timon Willi, Alistair Letcher, and Jakob Nicolaus Foerster. 2023b.  Adversarial cheap talk. In _International Conference on Machine Learning_. PMLR, 22917–22941.
  * Makoviychuk et al. (2021) Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, and Gavriel State. 2021.  Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning.
  * Michalski et al. (2023) Adam Michalski, Filippos Christianos, and Stefano V Albrecht. 2023.  SMAClite: A Lightweight Environment for Multi-Agent Reinforcement Learning.  _arXiv preprint arXiv:2305.05566_ (2023).
  * Mnih et al. (2016) Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. 2016.  Asynchronous methods for deep reinforcement learning. In _International conference on machine learning_. PMLR, 1928–1937.
  * Mnih et al. (2015) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. 2015\.  Human-level control through deep reinforcement learning.  _nature_ 518, 7540 (2015), 529–533.
  * Morad et al. (2023) Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. 2023.  POPGym: Benchmarking Partially Observable Reinforcement Learning. In _The Eleventh International Conference on Learning Representations_.  https://openreview.net/forum?id=chDrutUTs0K
  * Osband et al. (2019) Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, et al. 2019\.  Behaviour suite for reinforcement learning.  _arXiv preprint arXiv:1908.03568_ (2019).
  * Papoudakis et al. (2021) Georgios Papoudakis, Filippos Christianos, Lukas Schäfer, and Stefano V. Albrecht. 2021.  Benchmarking Multi-Agent Deep Reinforcement Learning Algorithms in Cooperative Tasks. In _Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks (NeurIPS)_.  http://arxiv.org/abs/2006.07869
  * Peng et al. (2021) Bei Peng, Tabish Rashid, Christian Schroeder de Witt, Pierre-Alexandre Kamienny, Philip Torr, Wendelin Böhmer, and Shimon Whiteson. 2021.  Facmac: Factored multi-agent centralised policy gradients.  _Advances in Neural Information Processing Systems_ 34 (2021), 12208–12221.
  * Pretorius et al. (2021) Arnu Pretorius, Kale ab Tessera, Andries P. Smit, Kevin Eloff, Claude Formanek, St John Grimbly, Siphelele Danisa, Lawrence Francis, Jonathan Shock, Herman Kamper, Willie Brink, Herman Engelbrecht, Alexandre Laterre, and Karim Beguir. 2021.  Mava: A Research Framework for Distributed Multi-Agent Reinforcement Learning.  _arXiv preprint arXiv:2107.01460_ (2021).  https://arxiv.org/pdf/2107.01460.pdf
  * Rashid et al. (2020) Tabish Rashid, Mikayel Samvelyan, Christian Schroeder De Witt, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson. 2020.  Monotonic value function factorisation for deep multi-agent reinforcement learning.  _The Journal of Machine Learning Research_ 21, 1 (2020), 7234–7284.
  * Rashid et al. (2018) Tabish Rashid, Mikayel Samvelyan, Christian Schroeder, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson. 2018.  Qmix: Monotonic value function factorisation for deep multi-agent reinforcement learning. In _International conference on machine learning_. PMLR, 4295–4304.
  * Sabne (2020) Amit Sabne. 2020.  XLA : Compiling Machine Learning for Peak Performance.
  * Samvelyan et al. (2019) Mikayel Samvelyan, Tabish Rashid, Christian Schroeder De Witt, Gregory Farquhar, Nantas Nardelli, Tim GJ Rudner, Chia-Man Hung, Philip HS Torr, Jakob Foerster, and Shimon Whiteson. 2019.  The starcraft multi-agent challenge.  _arXiv preprint arXiv:1902.04043_ (2019).
  * Schrittwieser et al. (2020) Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. 2020\.  Mastering atari, go, chess and shogi by planning with a learned model.  _Nature_ 588, 7839 (2020), 604–609.
  * Schulman et al. (2017) John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ (2017).
  * Shacklett et al. (2023) Brennan Shacklett, Luc Guy Rosenzweig, Zhiqiang Xie, Bidipta Sarkar, Andrew Szot, Erik Wijmans, Vladlen Koltun, Dhruv Batra, and Kayvon Fatahalian. 2023.  An Extensible, Data-Oriented Architecture for High-Performance, Many-World Simulation.  _ACM Trans. Graph._ 42, 4 (2023).
  * Snyder (1971) Glenn H Snyder. 1971.  ” Prisoner’s Dilemma” and” Chicken” Models in International Politics.  _International Studies Quarterly_ 15, 1 (1971), 66–103.
  * Sun et al. (2023) Mingfei Sun, Sam Devlin, Jacob Beck, Katja Hofmann, and Shimon Whiteson. 2023.  Trust region bounds for decentralized ppo under non-stationarity. In _Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems_. 5–13.
  * Sunehag et al. (2017) Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojciech Marian Czarnecki, Vinicius Zambaldi, Max Jaderberg, Marc Lanctot, Nicolas Sonnerat, Joel Z Leibo, Karl Tuyls, et al. 2017\.  Value-decomposition networks for cooperative multi-agent learning.  _arXiv preprint arXiv:1706.05296_ (2017).
  * Sunehag et al. (2018) Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojciech Marian Czarnecki, Vinicius Zambaldi, Max Jaderberg, Marc Lanctot, Nicolas Sonnerat, Joel Z Leibo, Karl Tuyls, et al. 2018\.  Value-Decomposition Networks For Cooperative Multi-Agent Learning Based On Team Reward. In _Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems, (AAMAS 2018)_ , Vol. 3. 2085–2087.
  * Tampuu et al. (2017) Ardi Tampuu, Tambet Matiisen, Dorian Kodelja, Ilya Kuzovkin, Kristjan Korjus, Juhan Aru, Jaan Aru, and Raul Vicente. 2017.  Multiagent cooperation and competition with deep reinforcement learning.  _PloS one_ 12, 4 (2017), e0172395.
  * Terry et al. (2021) J Terry, Benjamin Black, Nathaniel Grammel, Mario Jayakumar, Ananth Hari, Ryan Sullivan, Luis S Santos, Clemens Dieffendahl, Caroline Horsch, Rodrigo Perez-Vicente, et al. 2021\.  Pettingzoo: Gym for multi-agent reinforcement learning.  _Advances in Neural Information Processing Systems_ 34 (2021), 15032–15043.
  * Todorov et al. (2012) Emanuel Todorov, Tom Erez, and Yuval Tassa. 2012.  MuJoCo: A physics engine for model-based control. In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_. IEEE, 5026–5033.  https://doi.org/10.1109/IROS.2012.6386109
  * Van Hasselt et al. (2016) Hado Van Hasselt, Arthur Guez, and David Silver. 2016.  Deep reinforcement learning with double q-learning. In _Proceedings of the AAAI conference on artificial intelligence_ , Vol. 30.
  * Weng et al. (2022) Jiayi Weng, Min Lin, Shengyi Huang, Bo Liu, Denys Makoviichuk, Viktor Makoviychuk, Zichen Liu, Yufan Song, Ting Luo, Yukun Jiang, Zhongwen Xu, and Shuicheng Yan. 2022.  EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine. In _Advances in Neural Information Processing Systems_ , S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (Eds.), Vol. 35. Curran Associates, Inc., 22409–22421.  https://proceedings.neurips.cc/paper_files/paper/2022/file/8caaf08e49ddbad6694fae067442ee21-Paper-Datasets_and_Benchmarks.pdf
  * Young and Tian (2019) Kenny Young and Tian Tian. 2019.  MinAtar: An Atari-Inspired Testbed for Thorough and Reproducible Reinforcement Learning Experiments.  _arXiv preprint arXiv:1903.03176_ (2019).
  * Yu et al. (2022) Chao Yu, Akash Velu, Eugene Vinitsky, Jiaxuan Gao, Yu Wang, Alexandre Bayen, and Yi Wu. 2022.  The surprising effectiveness of ppo in cooperative multi-agent games.  _Advances in Neural Information Processing Systems_ 35 (2022), 24611–24624.
  * Zhang et al. (2022) Qizhen Zhang, Chris Lu, Animesh Garg, and Jakob Foerster. 2022.  Centralized Model and Exploration Policy for Multi-Agent RL. In _Proceedings of the 21st International Conference on Autonomous Agents and Multiagent Systems_. 1500–1508.
  * Zhao et al. (2022) Stephen Zhao, Chris Lu, Roger B Grosse, and Jakob Foerster. 2022.  Proximal Learning With Opponent-Learning Awareness.  _Advances in Neural Information Processing Systems_ 35 (2022), 26324–26336.
  * Zhou et al. (2023) Ming Zhou, Ziyu Wan, Hanjing Wang, Muning Wen, Runzhe Wu, Ying Wen, Yaodong Yang, Yong Yu, Jun Wang, and Weinan Zhang. 2023.  MALib: A Parallel Framework for Population-based Multi-agent Reinforcement Learning.  _Journal of Machine Learning Research_ 24, 150 (2023), 1–12.  http://jmlr.org/papers/v24/22-0169.html

##  Appendix A Further Details on Environments

###  A.1. SMAX

Observations in SMAX are structured similarly to SMAC. Each agent observes the
health, previous action, position, weapon cooldown and unit type of all allies
and enemies in its sight range. Like SMACv2Ellis et al. (2022), we use the
sight and attack ranges as prescribed by StarCraft II rather than the fixed
values used in SMAC.

SMAX and SMAC have different returns. SMAC’s reward function, like SMAX’s, is
split into two parts: one part for depleting enemy health, and another for
winning the episode. However, in SMAC, the part which rewards depleting enemy
health scales with the number of agents. This is most clearly demonstrated in
27m_vs_30m, where a random policy gets a return of around $10$ out of a
maximum of $20$ because almost all the reward is for depleting enemy health or
killing agents, rather than winning the episode. In SMAX, however, 50% of the
total return is always for depleting enemy health, and 50% for winning. Unlike
StarCraft II, where all actions happen in a randomised order in the game loop,
some actions in SMAX are simultaneous, meaning draws are possible. In this
case both teams get $0$ reward.

Like SMAC, each environment step in SMAX consists of eight individual time
ticks. SMAX uses a discrete action space, consisting of movement in the four
cardinal directions, a stop action, and a shoot action per enemy.

SMAX makes three notable simplifications of the StarCraft II dynamics to
reduce complexity. First, zerg units do not regenerate health. This health
regeneration is slow at $0.38$ health per second, and so likely has little
impact on the game. Protoss units also do not have shields. Shields only
recharge after 10 seconds out of combat, and therefore are unlikely to
recharge during a single micromanagement task. Protoss units have additional
health to compensate for their lost shields. Finally, the available unit types
are reduced compared to SMAC. SMAX has no medivac, colossus or baneling units.
Each of these unit types has special mechanics that were left out for the sake
of simplicity. For the SMACv2 scenarios, the start positions are generated as
in SMACv2, with the small difference that the ‘surrounded’ start positions now
treat allies and enemies identically, rather than always spawning allies in
the middle of the map. This symmetry guarantees that a 50% win rate is always
achievable.

Collisions are handled by moving agents to their desired location first and
then pushing them out from one another.

###  A.2. Coin Game

Two agents, ‘red’ and ‘blue’, move in a wrap-around grid and collect red and
blue coloured coins. When an agent collects any coin, the agent receives a
reward of $1$. However, when ‘red’ collects a blue coin, ‘blue’ receives a
reward of $-2$ and vice versa. Once a coin is collected, a new coin of the
same colour appears at a random location within the grid. If a coin is
collected by both agents simultaneously, the coin is duplicated and both
agents collect it. Episodes are of a set length.

###  A.3. Spatial-Temporal Representations of Matrix Games (STORM)

This environment features directional agents within an 8x8 grid-world with a
restricted field of view. Agents cannot move backwards or share the same
location. Collisions are resolved by either giving priority to the stationary
agent or randomly if both are moving. Agents collect two unique resources:
cooperate and defect coins. Once an agent picks up any coin, the agent’s
colour shifts, indicating its readiness to interact. The agents can then
release an interact beam directly ahead; when this beam intersects with
another ready agent, both are rewarded based on the specific matrix game
payoff matrix. The agents’ coin collections determine their strategies. For
instance, if an agent has 1 cooperate coin and 3 defect coins, there’s a 25%
likelihood of the agent choosing to cooperate. After an interaction, the two
agents involved are frozen for five steps, revealing their coin collections to
surrounding agents. After five steps, they respawn in a new location, with
their coin count set back to zero. Once an episode concludes, the coin
placements are shuffled. This grid-based approach to matrix games can be
adapted for n-player versions. While STORM is inspired by MeltingPot 2.0,
there are noteworthy differences:

  * •

Meltingpot uses pixel-based observations while we allow for direct grid
access.

  * •

Meltingpot’s grid size is typically 23x15, while ours is 8x8.

  * •

Meltingpot features walls within its layout, ours does not.

  * •

Our environment introduces stochasticity by shuffling the coin placements,
which remain static in Meltingpot.

  * •

Our agents begin with an empty coin inventory, making it easier for them to
adopt pure cooperate or defect tactics, unlike in Meltingpot where they start
with one of each coin.

  * •

MeltingPot is implemented in Lua (Ierusalimschy, 2006) where as ours is a
vectorized implementation in Jax.

We deem the coin shuffling especially crucial because even large environments
representing POMDPs, such as SMAC, can be solved without the need for memory
if they lack sufficient randomness (Ellis et al., 2022).

###  A.4. Hanabi

There are a few details that differ between our Hanabi implementation and the
original Hanabi Learning Environment (HLE). The most notable of these is how
we choose to represent card knowledge information in the agents’ observation.
In the HLE, card knowledge is observed as a colour/rank if there has been an
explicit hint about a given card. As a separate feature, implicit card
knowledge is represented as possible colours/ranks if there has not been an
explicit hint that indicates a given card is not that colour/rank. We, on the
other hand, combine implicit and explicit card knowledge, by only maintaining
a representation of implicit card knowledge, which reduces to explicit card
knowledge in the event an explicit hint is given about a card. This is because
all possible colours/ranks are represented as 1s, whilst all ruled out
colours/ranks are represented as 0s. By giving an explicit hint, all but one
colour/rank are ruled out, leaving a one-hot encoding of the explicit card
knowledge. We implement card knowledge this way, because knowledge updates are
implemented via tensor calculus using JAX Numpy arrays of fixed shape and data
type.

##  Appendix B Value-Based MARL Methods and Implementation details

Key features of our framework include parameter sharing, a recurrent neural
network (RNN) for agents, an epsilon-greedy exploration strategy with linear
decay, a uniform experience replay buffer, and the incorporation of Double
Deep $Q$-Learning (DDQN) Van Hasselt et al. (2016) techniques to enhance
training stability.

Unlike PyMARL, we use the Adam optimizer as the default optimization
algorithm. Below is an introduction to common value-based MARL methods.

IQL (Independent $Q$-Learners) is a straightforward adaptation of Deep
$Q$-Learning to multi-agent scenarios. It features multiple $Q$-Learner agents
that operate independently, optimizing their individual returns. This approach
follows a decentralized learning and decentralized execution pipeline.

VDN (Value Decomposition Networks) extends $Q$-Learning to multi-agent
scenarios with a centralized-learning-decentralized-execution framework.
Individual agents approximate their own action’s $Q$-Value, which is then
summed during training to compute a jointed $Q_{tot}$ for the global state-
action pair. Back-propagation of the global DDQN loss in respect to a global
team reward optimizes the factorization of the jointed $Q$-Value.

QMIX improves upon VDN by relaxing the full factorization requirement. It
ensures that a global $argmax$ operation on the total $Q$-Value ($Q_{tot}$) is
equivalent to individual $argmax$ operations on each agent’s $Q$-Value. This
is achieved using a feed-forward neural network as the mixing network, which
combines agent network outputs to produce $Q_{tot}$ values. The global DDQN
loss is computed using a single shared reward function and is back-propagated
through the mixer network to the agents’ parameters. Hypernetworks generate
the mixing network’s weights and biases, ensuring non-negativity using an
absolute activation function. These hypernetworks are two-layered multi-layer
perceptrons with ReLU non-linearity.

##  Appendix C Training Results

###  C.1. MABrax

The performance of IPPO on HalfCheeta and Walker is reported in Figure 9, with
hyperparameters reported in Table 4.

(a) HalfCheetah

(b) Walker

Figure 9. Performance of IPPO on MABrax Tasks

###  C.2. MPE

Performance of $Q$-Learning baselines in all the MPE scenarios are reported in
Figure 10. The upper row represents cooperative scenarios, with results for
all our $Q$-learning baselines reported. The bottom row refers to competitive
scenarios, and results for IQL are divided by agent types. Hyperparameters are
given in Table 10

###  C.3. SMAX

The performance of IPPO in SMAX versus MAPPO in SMAC is shown in Figure 11
while the performance of our $Q$-learning baselines is reported in Figure 12.
We do not report them together because their hyperparameters were tuned over a
different number of timesteps. Hyperparameters for IPPO and the $Q$-learning
methods are given in Tables 6 and 10 respectively.

Figure 10. $Q$-Learning Baselines in all MPE scenarios. Where no algorithm
names are given, the results represent IQL. Figure 11. IPPO and MAPPO in SMAX
versus MAPPO in SMAC for all SMAC maps.

Figure 12. Performance of $Q$-Learning Baselines for all SMAX scenarios

##  Appendix D Hyperparameters

Value | Ant | HalfCheetah | Walker
---|---|---|---
VF_COEF | 4.5 | 0.14 | 1.9
ENT_COEF | $2\times 10^{-6}$ | $4.5\times 10^{-3}$ | $1\times 10^{-3}$
LR | $1\times 10^{-3}$ | $6\times 10^{-4}$ | $7\times 10^{-3}$
NUM_ENVS | 64 | – | –
NUM_STEPS | 300 | – | –
TOTAL_TIMESTEPS | $1\times 10^{8}$ | – | –
NUM_MINIBATCHES | 4 | – | –
GAMMA | 0.99 | – | –
GAE_LAMBDA | 1.0 | – | –
CLIP_EPS | 0.2 | – | –
MAX_GRAD_NORM | 0.5 | – | –
ACTIVATION | tanh | – | –
ANNEAL_LR | True | – | –
Table 4. MABrax Hyperparameters, where – indicates repeated parameters

Hyperparameter | Value
---|---
LR | $0.0005$
NUM_ENVS | $25$
NUM_STEPS | $128$
TOTAL_TIMESTEPS | $1\times 10^{6}$
UPDATE_EPOCHS | $5$
NUM_MINIBATCHES | $2$
GAMMA | $0.99$
GAE_LAMBDA | $1.0$
CLIP_EPS | $0.3$
ENT_COEF | $0.01$
VF_COEF | $1.0$
MAX_GRAD_NORM | $0.5$
ACTIVATION | tanh
ANNEAL_LR | True
Table 5. Hyperparameters for MPE IPPO

Hyperparameter | Value
---|---
LR | $0.004$
NUM_ENVS | $64$
NUM_STEPS | $128$
TOTAL_TIMESTEPS | $1\times 10^{7}$
UPDATE_EPOCHS | $2$
NUM_MINIBATCHES | $2$
GAMMA | $0.99$
GAE_LAMBDA | $0.95$
CLIP_EPS | $0.2$
SCALE_CLIP_EPS | False
ENT_COEF | $0.0$
VF_COEF | $0.5$
MAX_GRAD_NORM | $0.5$
ACTIVATION | relu
Table 6. Hyperparameters for SMAX IPPO

Hyperparameter | Value
---|---
LR | $5\times 10^{-4}$
NUM_ENVS | $1024$
NUM_STEPS | $128$
TOTAL_TIMESTEPS | $1\times 10^{10}$
UPDATE_EPOCHS | $4$
NUM_MINIBATCHES | $4$
GAMMA | $0.99$
GAE_LAMBDA | $0.95$
CLIP_EPS | $0.2$
ENT_COEF | $0.01$
VF_COEF | $0.5$
MAX_GRAD_NORM | $0.5$
ACTIVATION | relu
ANNEAL_LR | True
NUM_FC_LAYERS | 2
LAYER_WIDTH | 512
Table 7. Hyperparameters for Hanabi IPPO

Hyperparameter | Value
---|---
LR | $2.5\times 10^{-4}$
NUM_ENVS | $16$
NUM_STEPS | $128$
TOTAL_TIMESTEPS | $5\times 10^{6}$
UPDATE_EPOCHS | $4$
NUM_MINIBATCHES | $4$
GAMMA | $0.99$
GAE_LAMBDA | $0.95$
CLIP_EPS | $0.2$
ENT_COEF | $0.01$
VF_COEF | $0.5$
MAX_GRAD_NORM | $0.5$
ACTIVATION | tanh
ANNEAL_LR | True
NUM_EVALS | $16$
Table 8. Hyperparameters for Overcooked IPPO

Hyperparameter | Value
---|---
NUM_ENVS | $8$
NUM_STEPS | $25$
BUFFER_SIZE | $5000$
BUFFER_BATCH_SIZE | $32$
TOTAL_TIMESTEPS | $2\times 10^{6}$
AGENT_HIDDEN_DIM | $64$
AGENT_INIT_SCALE | $2.0$
EPSILON_START | $1.0$
EPSILON_FINISH | $0.05$
EPSILON_ANNEAL_TIME | $100000$
MIXER_EMBEDDING_DIM* | $32$
MIXER_HYPERNET_HIDDEN_DIM* | $64$
MIXER_INIT_SCALE* | $0.00001$
MAX_GRAD_NORM | $25$
TARGET_UPDATE_INTERVAL | $200$
LR | $0.005$
EPS_ADAM | $0.001$
WEIGHT_DECAY_ADAM | $0.00001$
GAMMA | $0.9$
NUM_TEST_EPISODES | $32$
TEST_INTERVAL | $50000$

Table 9. Hyperparameters for MPE $Q$-Learning Algorithms

(* Parameters specific to QMix.)

Hyperparameter | Value
---|---
NUM_ENVS | $8$
NUM_STEPS | $100$
BUFFER_SIZE | $3000$
BUFFER_BATCH_SIZE | $32$
TOTAL_TIMESTEPS | $2\times 10^{7}$
AGENT_HIDDEN_DIM | $256$
AGENT_INIT_SCALE | $1.0$
EPSILON_START | $1.0$
EPSILON_FINISH | $0.05$
EPSILON_ANNEAL_TIME | $100000$
MIXER_EMBEDDING_DIM* | $64$
MIXER_HYPERNET_HIDDEN_DIM* | $256$
MIXER_INIT_SCALE* | $0.001$
MAX_GRAD_NORM | $10$
TARGET_UPDATE_INTERVAL | $200$
LR | $0.001$
EPS_ADAM | $0.00001$
WEIGHT_DECAY_ADAM | $1\times 10^{-6}$
GAMMA | $0.99$
NUM_TEST_EPISODES | $32$
TEST_INTERVAL | $1\times 10^{5}$

Table 10. Hyperparameters for SMAX $Q$-Learning Algorithms

(* Parameters specific to QMix.)