Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX
Authors: Clément Bonnet
Daniel Luo
Donal John Byrne
Shikha Surana
Sasha Abramowitz
Paul Duckworth
Vincent Coyette
Laurence Illing Midgley
Elshadai Tegegn
Tristan Kalloniatis
Omayma Mahjoub
Matthew Macfarlane
Andries Petrus Smit
Nathan Grinsztajn
Raphael Boige
Cemlyn Neil Waters
Mohamed Ali Ali Mimouni
Ulrich Armel Mbou Sob
Ruan John de Kock
Siddarth Singh
Daniel Furelos-Blanco
Victor Le
Arnu Pretorius
Alexandre Laterre
Venue: ICLR 2024 poster (2024)
Keywords: reinforcement learning
jax
combinatorial
research
Primary area: datasets and benchmarks
TL;DR: We introduce Jumanji, an open-source, diverse suite of RL environments, designed to be fast, flexible, and scalable.

Abstract: Open-source reinforcement learning (RL) environments have played a crucial role in driving progress in the development of AI algorithms.
In modern RL research, there is a need for simulated environments that are performant, scalable, and modular to enable their utilization in a wider range of potential real-world applications.
Therefore, we present Jumanji, a suite of diverse RL environments specifically designed to be fast, flexible, and scalable.
Jumanji provides a suite of environments focusing on combinatorial problems frequently encountered in industry, as well as challenging general decision-making tasks.
By leveraging the efficiency of JAX and hardware accelerators like GPUs and TPUs, Jumanji enables rapid iteration of research ideas and large-scale experimentation, ultimately empowering more capable agents.
Unlike existing RL environment suites, Jumanji is highly customizable, allowing users to tailor the initial state distribution and problem complexity to their needs.
Furthermore, we provide actor-critic baselines for each environment, accompanied by preliminary findings on scaling and generalization scenarios.
Jumanji aims to set a new standard for speed, adaptability, and scalability of RL environments.

Published as a conference paper at ICLR 2024
JUMANJI: A DIVERSE SUITE OF SCALABLE REINFORCEMENT LEARNING ENVIRONMENTS IN JAX
Clément Bonnet
Daniel Luo'
Donal Byrne'
Shikha Surana'
Sasha Abramowitz!
Paul Duckworth'
Vincent Coyette
Laurence I. Midgley²
Elshadai Tegegn*
Tristan Kalloniatis
Omayma Mahjoub!
Matthew Macfarlane
Andries P. Smit!
Nathan Grinsztajn'
Raphael Boige'
Cemlyn N. Waters
Mohamed A. Mimouni!
UIrich A. Mbou Sob!
Ruan de Kock
Siddarth Singh!
Daniel Furelos-Blanco*
Victor Le
Arnu Pretorius
Alexandre Laterre'
InstaDeep
University of Cambridge Uniiversity of Amsterdam *Imperial College London
ABSTRACT
Open-source reinforcement learning (RL) environments have played a crucial role in driving progress the development of Al algorithms. In modern RL research, there is need for simulated environments that are performant, scalable, and modu- lar enable their utilization wider range of potential real-world applications. Therefore, present Jumanji. suite of diverse RL environments specifically designed to be fast, flexible, and scalable. Jumanji provides a suite of environ- ments focusing on combinatorial problems frequently encountered in industry, well as challenging general decision-making tasks. By leveraging the efficiency of JAX and hardware accelerators like GPUs and TPUs, Jumanji enables rapid iteration of research ideas and large-scale experimentation, ultimately empowering more capable agents. Unlike existing RL environment suites, Jumanji highly customizable, allowing users to tailor the initial state distribution and problem complexity their needs. Furthermore, we provide actor-critic baselines for each environment, accompanied by preliminary findings on scaling and generalization scenarios. Jumanji aims to set a new standard for speed, adaptability, and scalability of RL environments.
INTRODUCTION
High-quality datasets and benchmarks are crucial to the development of AI research (Deng et al.. 2009: Krizhevsky et al. 2012 Bellemare et al.) 2013). They allow for coordinated research problems that serve as a measure of progress toward shared goals. However, most currently open- sourced reinforcement learning (RL) environment libraries are not closely tied practical problems. Furthermore, in industrial settings operating at scale, these libraries do not provide sufficient flexibility and scalability to facilitate long-term AI research suitably close to real-world applications.
For RL to be useful in the real world, further research progress is needed. This will require benchmarks that are: (1) fast, i.e. hardware-acceleratee to overcome sequential bottlenecks and allow for faster experiment iteration; (2) flexible by allowing easy customization to capture realistic problem settings
"Correspondence <c1 ement. bonnet . 60gma com>

---

Published as a conference paper at ICLR 2024
2433 Seee 20M4
HP E
1.20
26
BinPack
Cleaner
Connector
CVRP
FlatPack
Game2048
GraphColoring
JobShop
Knapsack
Maze
Minesweeper
MMST
Scren 2010
MultCVRPP
PacMan
Rarehoure
RubiksCube
SiidingTilePuzde
Snake
Sokoban
Sudoku
Tetris
TSP
Figure 1: All 22 environments implemented in Jumanji (in alphabetic order) are divided into three categories. Routing problems: Cleaner, Connector CVRP (Capacitated Vehicle Routing Problem), Maze, MMST (Multiple Minimum Spanning Tree), MultiCVRP (Multiple-vehicle CVRP), PacMan, RobotWarehouse, Snake, Sokoban, and TSP (Travelling Salesman Problem). Packing problems: BinPack, FlatPack, JobShop. Knapsack, and Tetris. Logic games: Game2048, GraphColoring. Minesweeper, RubiksCube, Sliding Tiddnzzilee and Sudoku.
of interest (e.g. intrinsic stochasticity and distribution shift); and (3) scalable, to be able to arbitrarily set the level of difficulty of an environment, ensuring more faithful representation of the challenges that arise in real-world problems.
Although many other popular benchmarks exist such the Arcade Learning Environment (ALE) (Bellemare et al. (213), OpenAl Gym (Brockman et al., 2016), and Gymmax (Lange) 2022), which satisfy subset of the above criteria, none of them have managed to cover all three. We believe it is crucial to fill this gap to help push RL research closer to industrial applications.
In this paper, we introduce Jumanji: diverse suite of fast.flexible and scalable RL environments. Jumanji is organized into three problem categories: routing, packing, and logic. At its core is a set of NP-hard combinatorial optimization problems (COPs) that closely resemble problems encountered in the real world. The environment dynamics of these problems seamlessly scale with complexity allowing for long-term research suitable for real-world industrial applications. Jumanji is written in JAX (Bradbury et al. 2018). to leverage composable transformations with automatic differentiation and the XLA compiler for highly efficient RL systems that run directly on GPU or TPU accelerators. Furthermore, Jumanji promotes flexibility by allowing arbitrary initial state distributions via easily modifiable reset functions and bespoke generators. We empirically demonstrate the capabilities of Jumanji through set of initial experiments. Specifically, we present results training actor-critic agent across all environments, establishing a benchmark useful for future comparisons. We show that Jumanji environments are highly scalable, demonstrating high throughput in both single-device and
2

---

Published as a conference paper at ICLR 2024
multi-device setting. Finally, we illustrate the flexibility of environments by customizing initial state distributions to study generalization in a real-world problem setting.
The main contributions of this paper are as follows:
1. We introduce Jumanji: an open-source and diverse suite of industry-inspired RL environ- ments, that are fast, flexible, and scalable.
2. We provide baseline actor-critic agents for all environments.
3. We present initial experiments demonstrating that, unlike existing RL benchmarks, Jumanji environments offer a high degree of scalability and flexibility.
2 RELATED WORK
Benchmark environments have been pivotal in the development and evaluation of RL algorithms. OpenAl Gym (Brockman et al. 2016). with its diverse task suite and user-friendly API, has become a benchmarking standard in the field. Other libraries, such as DMLab (Beattie al. (016) for complex 3D environments, and Mujoco (Todorov et al.] (012)) for high-fidelity physics simulations, have enabled researchers push the boundaries of agent capabilities. However, despite the significant contributions, these libraries have limitations in computational efficiency and scalability.
Hlardware-accelrratee Environments A common approach to increasing environment throughput through parallelization of the environment itself. Prior work such as EnvPool (Weng et al. 2022) utilizes multiple CPU cores and C++ based threading of multiple instances of an environment in order to expedite the bottleneck of sequential simulation steps. GPU-accelerated environments like Nvidia's
CuL.E (Dalton et al.) 2019; Dalton & frosio 2020) and Isaac Gym (Makoviychuk et al. (021)) take different approach, leveraging the parallel processing capabilities of GPUs. CuLE provides a CUDA port of ALE (Bellemare et al. (203),, rendering frames directly GPU, whilst Isaac Gym provides an accelerated alternative to Mujoco. Although these environments offer significant efficiency gains, they are strictly limited to GPUs and cannot be readily extended to other hardware accelerators (e.g. TPUs). JÄX (Bradbury et al. 2018) is a numerical computing library that leverages automatic differentiation, vectorization, parallelization, and an XLA compiler for device-agnostic optimization. JAX is utilized in RL environments such as Brax (Freeman et al.] [02T). a differentiable physics engine, Pgx (Koyamada et al., 2202).. collection board game simulators, Gymnax (Lange 2022). a library of popular RL environments re-implemented in JAX, JaxMARL (Rutherford et al. (023), a
collection of commonly used MARL environments, and Craftax -Cattx--CCasiccccc (Matthews al. 2024) a
re- re-implementation in JAX of the open-ended environment Crafter. These environments represent significant advance in efficiency but are limited in scope and flexibility.
Combinatorial Optimization Problems (COPs) COPs present a significant area of research in RL, with many real-world and industrial applications. Examples include the Traveling Salesman Problem, Bin Packing, the Capacitated Vehicle Routing Problem, and the Knapsack Problem (Bengio
et al. 2021). While there have been substantial advances in related software, such as Google OR Tools Google 2023), there is noticeable gap in support for RL-based COP solutions. Libraries such OR RL Benchmarks (Balaji et al. 2019) and OR-Gym (Hubbs al. 2020)) provide COP environments that adhere to the standard Gym API, however, they are restricted to run on CPU, making it difficult parallelize and scale.
Benchmark Diversity Multiple RL benchmarks have been proposed to facilitate agent general- ization, such as Procgen (Cobbe et al., 2201)) OpenSpiel (Lanctot et al.) 2020), BSuite (Osband) al. 2020). and Xlannd (Team et al. (021). Whilst providing challenging scenarios for training. Procgen and OpenSpiel do not inherently support scaling of the environment dynamics, and unlike Jumanji, none are designed to utilize hardware accelerators. Lastly, whilst no single benchmark suite can handle all situations, the ability to extend and create environments is crucial. Unity ML Agents Juliani et al. 2020) is extendable and parallelizable via the Unity game engine, but not optimized for accelerators. Gymnasium aims to standardize Gym (Brockman al. (0166) but does not directly provide a base for new environments. Most libraries rely on hard-coded components that require users develop new extensions. Jumanji, however, follows composition- -based design,
3

---

Published as a conference paper at ICLR 2024
allowing for easy modifications of initial state distributions, reset behaviors, level generators, and rendering.
JUMANJI
We introduce Jumanji (v1 0), suite of 22 JAX-based environments, visualized in Figure I These diverse problems rely on variety of geometries, including grids, graphs, and sets. The environments organized into three problem categories: routing, packing, and logic. Many of these are NP-hard COPs inspired by real-world industry settings. Jumanji leverages JAX to significantly accelerate and parallelize simulation steps while remaining flexible and allowing for scalable problem complexity.
This section provides an overview of the library, first introducing the RL setting and then Jumanji's application programming interface (API).
3.1 RL PRELIMINARIES
Each Jumanji environment is structured Markov decision process (MDP) (Puterman 1994). M = (S, A, P. 7). where is the state space, A is the action space, μ is the initial state distribution, defines the environment transition dynamics, R is the reward function, and y the discount factor. We can generate trajectories from an MDP by rolling out the environment dynamics. That is, at time step t. an action a, transitions the environment from the current state s; to the next state 81+1 1 as defined by the environment dynamics P, resulting in reward The objective of an agent is often to maximize the discounted expected return, given by E Bo, [E 7 (rts at)], where y is the discount factor and æ the agent's policy.
3.2 API
Jumanji's interface is lightweight, flexible, and capable of representing diverse of RL problems. draws inspiration from OpenAl Gym (Brockman et al. 2016). dm_env (Muldal et al. 2019), and Brax (Freeman al. 2021). It is flexible in three ways: (1) allowing customization of the initial state distribution via generators, (2) custom visualization via environment viewers, and (3) custom reward functions. Below, we introduce the key components of the Jumanji API, including the environment interface, state, observation, generators, specs, and registry.
Environment Interface The Environment interface defines the blueprint for Jumanji environments. Each environment must contain the following methods: reset, step, observation spec, and action_ secton_sspes spec. The API allows for optional render and animate methods to visualize a state sequence of states. For a code snippet demonstrating how to create new environment by extending the environment interface, see Appendix [A.4 Here, we provide code to instantiate an environment from the Jumanji registry, reset, step, and (optionally) render it:
import jax import jumanji
2
Instantiate Jumanji environment from the registry
env jumanji.a....... make Snake-vl
6
7 Reset the environment
key random. PRNGKey Saandm..RRNKKe/.. (0)
state, timestep jax. jit (env.reset. (key)
10
11
Sample an action and take an environment step action env. acv.aacion_spe(() gcneratte_vauuc state, timestep = jax. jit lenv.sten (state, action)
is
(Optional) Render the envi ronment state env. render (state)
16
4

---

Published as a conference paper at ICLR 2024
State The State is a pytree (e.g. dataclass namedtuple) that contains all the required information to transition the environment's 's dynamics for a given action. This is a design choice, and as such, Jumanji environments are stateless i.e. the reset and step methods are functionally pure. This allows Jumanji to leverage JAX's transformations (j grad, vmap, and pmap) to make the environments highly scalable. Every state includes pseudorandom number generator (PRNG) key, which is used during potential auto-reset and in the case of stochastic transitions.
TimeStep and Observation The Ti eimeStep contains all the information available to an agent about the state of the environment at a time step, namely; the Step__ypp (first, mid, or last), reward, discuun, observation. and extras. As such, it is based on the dm_env env TimeStep but with an additional extras field, where environment metrics can be logged that are neither available to the agent nor part of the state. For Jumanji environments, the Observation is a JAX pytree, making it amenable multiple data types.
Generators For given environment, generator is used define the initial state distribution. Jumanji provides flexibility by allowing the use of custom generators, enabling users to define an initial state distribution specific to their problem. In most environments, the reset method calls a generator that returns the initial state. The generator can be specified upon environment instantiation. The user can choose from a set of pre-existing generators or implement their own generator. If not specified, default generator is used.
Specs Inspired by dm env, Jumanji specifications define the tree structure, shape, data type, and optionally the bounds of the observations, actions, rewards, and discounts. Jumanji's Spec class more general than its dm_env counterpart, allowing for nested structures. This is achieved by implementing each spec as a nested JAX pytree composed of a set of primitive specs (Array, BoundedArray, DiscreteArray tree while each non-leaf node is itself another nested Spec object. Environments in Jumanji have their action space described spec, which means although current environments have discrete actions, Jumanji supports both discrete and continuous action spaces.
Mul t iD screteArrayyyy which form the leaves of the
Registry Jumanji keeps strict versioning of its environments for the purpose of reproducibility. This is achieved through a registry of standard environments with their respective configurations. For each environment, version suffix is appended, e.g. "Snake-v1". When significant changes are made to environments, the version number incremented.
4 JUMANJI BENCHMARK
In this section, we first describe a highly efficient method for training RL agents in Jumanji envi- ronments. Secondly, we provide details of standard actor-critic baseline agents. Finally, we present experiments demonstrating the speed and parallelization of Jumanji environments.
4.1 EFFICIENT TRAINING
As described in Section 3.2 Jumanji environments designed to be stateless, allowing Jumanji to take full advantage of JAX's transformations. JAX-based stateless environments provide multiple benefits. Firstly, can JIT-compile the full training loop of an agent. This often includes rolling out the environment to generate trajectories, and then separately updating the parameters of the agent based on that experience. We provide a throughput ablation in appendix D.2 to demonstrate the speed-up that arises when removing data transfer between host and device. Secondly, JAX's grad allows for efficient computation of gradients using automatic differentiation. Thirdly, JAX's vectorization (vmap) can be used to generate rollouts and compute parameter updates in parallel on a single device. Finally, JAX's process parallelism (pmap) can be leveraged to parallelize the computations across multiple devices, where gradients are aggregated across devices using amean. Hessel et al. (202)) propose the Anakin architecture for exactly this setting, with an emphasis on maximizing the utilization of TPU pods, although their approach is general and also applies to multi-GPU RL training. All the experiments in this paper are implemented using this efficient Anakin framework.
5

---

Published as a conference paper at ICLR 2024
BinPack
Cicaner
Connector
CVRP
FlatFack
24me
60
-24
2300
300
033
035
375
100
1X
5.2
GraphColoring
-50
JobSinop
Knapsack
Minesweeper
MMST
1WA MM
-20
-000
2.5
NURICVPP
Pachaa
RobetWarehoure
RublisSuuee
10.3
Siddnngtllehrzze
Smake
100
50
0.5
-2.
100
Sokabon
Sudokk
1009
Tatris
TSP
2008
Brerowerr cape
33
18
Enviccmment Steps
Ewveenwet Staps
Envtorment
Random
Actor-Critic
Optimal Performance
Figure 2: Learning curves from training an actor-critic agent (A2C) in blue compared to a masked random policy in red on all 22 Jumanji environments. When possible to determine, the optimal expected return is shown with dashed line. Experiments were run with three different seeds, with the median represented as a blue curve and the min/max as the shaded region.
4.2 ACTOR-CRITIC BASELINE
Jumanji provides an implementation of an A2C (Mnih et al. 2016) agent, built using the DeepMind JAX ecosystem (Babuschkin et al.) (020). Since Jumanji environments use different geometries (e.g. images, sets, etc.), the agent relies on environment- speciffe neural networks, e.g. image inputs may be fed to convolution neural network while vermutation-equivrria problems may use transformer architecture (Vaswani al. 2017). To promote research using Jumanji, we open-source the algorithm, the training pipeline, checkpoints, and the aforementioned actor-critic networks which are compatible with any algorithms relying on a policy or state-value function. Appendix B provides further details on these environment-speat specifc networks.
To benchmark Jumanji environments, we provide learning curves of our A2C implementation on all 22 environments. We compare algorithm to the optimal performance (where possible to determine) and a random policy that uniformly samples actions from the set of valid actions. Note, that the optimal performances for TSP, CVRP, and Knapsack are taken from (Kwon al. 2020). The experiments were performed on a TPUv3-8 using the Anakin framework. We refer the reader to Appendix [.1 for more details on the training. In Figure 2 we show the learning curves the A2C agent on each of the registered environments (i.e. the default configurations), along with the random baseline and optimal performance. The experiments were run three times for each the 22 registered environments, with the median represented as blue curve and the min/max as the shaded region. Although our standard A2C agent improves upon the random baseline, optimality gaps remain in most environments (i.e., the differences between the A2C and optimal performance) highlighting challenges in solving combinatorial problems.
4.3 ENVIRONMENT PARALLELIZATIE EXPERIMENTS
We present an initial experiment demonstrating the speed of Jumanji environments as we parallelize the step function. Figure 3al shows how the throughput of the environment step function increases

---

Published as a conference paper at ICLR 2024
10
Krapsack
Epoch Duration (sec)
Minsswepper
104
TSP
10
DVRP
Bikding Tike Puzzie
10a
Steps per Second
Connector
10a
10"
Sokoban
Snaka
10a
McCCVPPP
TPU-128
TPU-64
TPLL33
TPU-16
1PU8
GPU
CPU
10
Paeman
RobotWharehouss
1.00
MMST
RubiksCube
Performance
0.95
JobShop
10
CraphColoring Game2040
Teiris
0.05
128 TPU-32
TPU:16 TPU.8
GPU
10²
BinPack
0.80
2"
2
2
0.75
10
12
Number of Parallel Environments
Time (hr)
(a) Effective environment steps per second
(b) A2C training on Connector
Figure 3: Analysis of environment parallelization. (a): Scaling of the effective environment steps per second (throughput) for each registered environment as the number parallel environments increases, on 8-core TPU-y4. The legend is ordered by the throughput of the rightmost data point. The results on GPU and CPU are presented in Appendix C22 (b): Training of an A2C agent on the Connector environment on CPU, GPU (RTX 2080 super), and TPU-v4 with number of cores varying from 8 to 128. Each training is run for 255M steps. Full training takes weeks on CPU, which is why it is not visible on the bottom plot. Performance denotes the proportion of wires connected (an optimal policy would reach 1.0). See Appendix C.2 for further details.
with the number of environments run in parallel, on TPUv4-8. To study parallelization on different hardware, we run similar experiment GPU (V100) and CPU in Appendix C.2 We compute the number of steps per second by averaging in parallel. The cost of the reset function is enviromment-depenie
consecutive actions each taken on 500 environments
and in some cases expensive, therefore, we focus on the scaling properties of the step function and do not reset the environments. Refer Appendix E for discussion parallelizing environments.
To quantify the benefits of device parallelization, we train an A2C agent on the combinatorial Connector environment varying the hardware. specifically, CPU, GPU (RTX 2080 super), and TPU with 8 to 128 cores. Figure 3bbshows the approximately linear scaling of convergence speed when increasing TPU cores, demonstrating efficient parallelization across devices. For example, we reach 92% of the optimal performance on Connector-v2 in roughly 1.5h with a TPU-128, compared 1.5h with a TPU-8.
FLEXIBILITY AND SCALABILITY IN JUMANJI
Jumanji is designed with flexibility and scalability at its In this section, present initial experiments demonstrating these two key properties. In Section 5.]] we demonstrate the flexibility of Jumanji environments, by implementing specific initial state distributions via custom generators. We first discuss why flexibility is required for building robust RL agents for real-world settings, and showcase it with an initial experiment using multiple generators. In Section 5.22 we discuss how problem complexity is scalable in Jumanji and present experiments demonstrating its impacts on agent performance.
5.1 FLEXIBILITY IN JUMANJI
Discussion of Jumanji Generators Training agents wide range of data distributions has been demonstrated to enhance their robustness towards real-world scenarios (Bossek et al.) 2019
Cobbe et al. 2019 Taiga al. 2023). The flexibility to define the initial state distribution provided
7

---

Published as a conference paper at ICLR 2024
by generators offers two significant advantages. Firstly, it enables users train agents desired data distributions, by creating one or multiple custom generators to sample from. This is especially useful for combinatorial problems where there is canonical instance distribution. Secondly, the environment dynamics are agnostic to the generator, this allows us to maintain consistent dynamics while having the flexibility alter the initial state distribution. This facilitates experimentation and analysis on different initial state configurations, enhancing our ability to understand a given agent's 's behavior across various scenarios.
Cluster
ineer Compression
Explasion
Cluster
Linear Compression
Explosion
(a) TSP generators
(b) Connector generators
Figure 4: Samples from three custom generators: cluster (left), linear compression (middle), and explosion (right), for TSP and Connector environments. In TSP, gray dots represent cities. In Connector, the node pairs be connected are depicted in the same color, and the large and small blocks indicate the starting and ending nodes, respectively.
Example Generators Here, we provide illustrative examples of possible generators for two envi- ronments, Specifically, we showcase three different generators for the Traveling Salesman Problem (TSP) and Connector environments in Figure In both environments, node coordinates, with the objective being to form a brief cycle (TSP) or connect all same-type node pairs without overlap (Connector). Problem instances cluster, linear compression, and explosion. The cluster generator allocates points within a specified radius and center point, the linear compression generator randomly aligns the points along 2D line within the space, and the explosion generator pushes the points away from given reference point in the space.
instance consists of 2D
created using the different generators:
Experiments using Multiple Generators We train two A2C agents on different initial state distributions for a TSP environment comprised of 50 cities and evaluate the generalization capabilities of the result- ing agents. Specifically, one agent is trained using random uniform instances, while the second agent samples from a combination of the uniform genera- tor and the three previously introduced TSP genera- tors. To evaluate our agents, we create two datascts from the VLSI TSP Benchmark Dataset (Rohe) that contain real-world problem instances. During training,
Single Generator
7.5
Four Generators
Tour Length
7.0
6.5
102 problem instances evaluate the agent's performance whilst at test time, we use a larger dataset of 020 instances.
0.5
1.0
2.0
2.5
Environment Steps
le8
Figure 5 shows the learning curves of the two A2C agents (smaller tour length is better). At test time, on the larger set of problem instances, the agent trained on single uniform generator achieved a mean tour length of 6.090 (+000, whereas the agent trained on the four generators attained better average tour length of 5.815 (+-0.025). These results demonstrate that the model trained on broader data distribution, facilitated by the inclusion of multiple generators, outperforms the model trained with a unseen, real-world test set and thus, shows better generalization capabilities. For further experimental details, please refer to Appendix
Figure 5: Learning curves of two agents train- ing TSP, sampling from single uniform generator versus sampling from four gener- ators (uniform, cluster, explosion, and com- pression). Lower tour length is better.
single generator
https:// hwwp//wwwww/ math. www.:hawwwelloccttaeeeeeoo h:crss///www.nncccaaaaeeoooooooe /wwiiiinndxx..ee html
8

---

Published as a conference paper at ICLR 2024
1.00
1.0
Rabio of Connections
0.95
Utlization
Ratio of Solved Cubes
0.8
0.90-
Ge10, Nez
Ce3. Ne5
0.85
6m10 N-10
N-10
C-3. N-7
Gm12,
N=12
N-20
0.80
0.0
0.5
1.0
1.5
2.0
2.5
0.75
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.5
1.0
1.5
2.0
2.5
Environment Steps
1c8
Environment Steps
108
Environment Steps
1c8
(a) Connector
(b) BinPack
(c) RubiksCube
Figure 6: The learning curves of the A2C agent on varying difficulty levels three different Jumanji environments. In Connector, the size of the grid (G) and the number of node pairs to be connected (N) are varied. In BinPack, the number of items (N) is varied. RubiksCube, the size the cube (C) and the number of scrambles made from solved Rubiks Cube (N) are varied.
5.2 SCALABILITY JUMANJI
Jumanji environments are scalable: each environment is equipped with one or more adjustable variables (such as the number of cities for TSP or the number of node pairs and grid size for Connector) that allow for arbitrary complexity. This flexibility is key attribute of Jumanji environments, as enables users to control the problem complexity and examine its impact on the agent's 's performance.
Experiments varying Problem Complexity To showcase scalability, we investigate the scaling properties of an A2C agent on three Jumanji environments: 1) Connector, where we vary the grid size and the number of node pair that need connecting; 2) BinPack, we vary the number of items to pack; 3) RubiksCube, we vary the size of the cube and the minimal number of actions required for solution. Appendix C.4 provides additional details these experiments.
We present learning curves of the three environments with varying difficulty settings in Figure We observe strong degradation in performance across all environments we increase the problem complexity. For example, in RubiksCube, the hardest problem setting we experiment with leads the A2C agent a complete failure learn.
This experiment provides a proof of concept into scaling the complexity of Jumanji environments. It highlights how Jumanji can be used study scaling properties of agents.
6 CONCLUSION
For RL research to be useful real-world applications, challenging new benchmark environments required. To this end, introduce Jumanji, an open-source and diverse suite of industry-inspine RL. environments that are fast, flexible, and scalable. Written in JAX, Jumanji environments can be parallelized and seamlessly scale with hardware (see Section 4.3p. Flexibility is provided by allowing users to define custom initial-state distributions via generators (see Section 3.2 annd..... At the heart of Jumanji
of NP-hard COPs with scalable environment dynamics that Tacilitate industry-scale research. While Jumanji provides industry-inspired environments, capturing the full complexity of industry situations within single benchmark remains challenging task. Nor can a single benchmark cover the full range of possible industry situations.
Jumanji is open-source, lightweight, and easy to extend. We welcome contributions from the community. Current environments all have discrete actions, yet Jumanji supports both discrete and continuous actions. Similarly, Jumanji supports multi-agent environments but only contains single- agent implementations. Future work will expand the library to include multi-agent implementations, environments with continuous actions, and more environments representative of real-world problems, such in the life sciences, agriculture, logistics, and beyond. By providing diverse suite of tasks, Jumanji aims to inspire future research toward RL agents that can learn to solve a wide range of important problems.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Research supported with Cloud TPUs from Google's TPU Research Cloud (TRC). We would like to thank the many open-source contributors who have contributed to the library, including Ugo Okoroafor, Randy Brown, Ole Jorgensen, Danila Kurganov, and Marta Wolinska. We also thank Thomas Barrett, Matthew Morris, Cyprien Courtot, Edan Toledo, Paul Caron, and Ian Davies for their help in the initial design of the library. Laurence Illing Midgley acknowledges support from José Miguel Hernández-Lobato's Turing AI Fellowship under grant EP/VO23756/1.
REFERENCES
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Jurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios. John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiee, Luyu Wang. Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL
com/ deepmind
Bharathan Balaji, Jordan Bell-Masterson, Enes Bilgin, Andreas Damianou, Pablo Moreno Garcia, Arpit Jain, Runfci Luo, Alvaro Maggiar, Balakrishnan Narayanaswamy, and Chun Ye. ORL: Reinforcement Learning Benchmarks for Online Stochastic Optimization Problems, 2019.
Charles Beattie, Joel Z. Leibo. Denis Teplyashin, Tom Ward, Mareus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, Julian Schrittwieser, Keith Anderson, Sarah York, Max Cant, Adam Cain, Adrian Bolton, Stephen Gaffney, Helen King, Demis Hassabis, Shane Legg, and Stig Petersen. Deep Lab, 2016.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The Arcade Learning Environment: An Evaluation Platform for General Agents. Journal of Artificial Intelligence Research, 47:253-279, jun 2013.
Yoshua Bengio, Andrea Lodi, and Antoine Prouvost. Machine learning for combinatorial optimization: methodological tour d' d'horizon. European Journal of Operational Research, 290(2):405 421, 2021.
Jakob Bossek, Pascal Kerschke, Aneta Neumann, Markus Wagner, Frank Neumann, and Heike Trautmann. Evolving diverse tsp instances by means of novel and creative mutation operators. In Proceedings of the 15th ACM/SIGEVO Conference Foundations of Genetic Algorithms,, pp.
58-71, 2019.
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. JAX: composable transformations of Python+NumPP programs, 2018. URL
http:/ /ttp::/ggtttuu.uu giggtb......... com wwt.c/////oooo///ooooo
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and
Wojciech Zaremba. OpenAl Gym, 2016.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging Procedural Generation to Benchmark Reinforcement Learning. arXiv preprint arXiv:1912.01588: 2019.
Steven Dalton and iuri frosio. Accelerating Reinforcement Learning through GPU Atari Em ulation. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 19773-19782. Curran As- sociates, Inc., 2020. URL httpt:tt///..m. htttpss://prcceediinnt//// nearggssseerrie neurips pss cc pawpp/ppppppeeem httpcc//ppp////2///11e/.. B4d78a 21
Steven Dalton, Iuri Frosio, and Michael Garland. GPU-Accelerated Atari Emulation for Reinforce- ment Learning, 2019.
10

---

Published as a conference paper at ICLR 2024
Jia Deng, Wei Dong. Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large- scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248 258, 2009. doi: 10. 1009/CVV..2099.52068.2
C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax A Differentiable Physics Engine for Large Scale Rigid Body Simulation, 2021. URL
Google. OR-Tools Google Optimization Tools.
2023. Google's s software suite for combinatorial optimization.
Arthur Guez, Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sébastien Racanière, Théophane Weber, David Raposo, Adam Santoro, Laurent Orseau, Tom Eccles, Greg Wayne, David Silver, and Timothy Lillicrap. An investigation of model-free planning, 2019.
Matteo Hessel, Manuel Kroiss, Aidan Clark, lurii Kemaev, John Quan, Thomas Keck, Fabio Viola, and Hado van Hasselt. Podracer architectures for scalable reinforcement learning. arXiv preprini arXiv:2104.06277 2021.
Christian D. Hubbs, Hector D. Perez, Owais Sarwar, Nikolaos V. Sahinidis, Ignacio E. Grossmann, and John M. Wassick. OR-Gym: A Reinforcement Learning Library for Operations Research Problems, 2020.
Arthur Juliani, Vincent-Pierre Berges, Ervin Teng, Andrew Cohen, Jonathan Harper, Chris Elion, Chris Goy, Yuan Gao, Hunter Henry, Marwan Mattar, and Danny Lange. Unity: A general platform for intelligent agents. arXiv preprint arXiv:1809.0262. 0822700999027.... 2020.
Wouter Kool, Herke van Hoof, and Max Welling. Attention, learn to solve routing problems!, 2018.
Sotetsu Koyamada, Shinri Okano, Soichiro Nishimori, Yu Murata, Keigo Habara, Haruka Kita, and Shin Ishii. Pgx: Hardware-acceleratee parallel game simulation for reinforcement learning. arXiv preprint 1.33::23003.75033 2023.
Alex Krizhevsky. Ilya Sutskever, and Geoffrey E Hinton. ImageNet Classification with Deep Convolutional Neural Networks. In F. Pereira, C.J. Burges, L. Bottou, and K.Q. Wein- berger (eds.). Advances in Neural Information Processing Systems, volume 25. Curran Asso- ciates, Inc., 2012. URL https: http::::s///pccceeeeeeee neurips. cc/paper files/paper/ 201 2 ff le/ c3 99899999988888 od 0oweeeeeo666666666666666666
Yeong-Dae Kwon, Jinho Choo, Byoungjip Kim, Iljoo Yoon, Youngjune Gwon, and Seungjai Min. Pomo: Policy optimization with multiple optima for reinforcement learning. Advances Neural Information Processing Systems, 33:21 88-21198. 2020.
Marc Lanctot, Edward Lockhart, Jean- Baptiste Lespiau, Vinicius Zambaldi, Satyaki Upadhyay, Julien Pérolat, Sriram Srinivasan, Finbarr Timbers, Karl Tuyls, Shayegan Omidshafiei, Daniel Hennes, Dustin Morrill, Paul Muller, Timo Ewalds, Ryan Faulkner, János Kramár, Bart De Vylder, Brennan Saeta, James Bradbury, David Ding, Sebastian Borgeaud, Matthew Lai, Julian Schrit- twieser, Thomas Anthony, Edward Hughes, Ivo Danihelka, and Jonah Ryan Ryan-Davis. OpenSpiel: A Framework for Reinforcement Learning in Games, 2020.
Robert Tjarko Lange. gymnax: A JAX -baed Reinforcement Learning Environment Library, 2022.
URL httt::////ithtu com/
Ryan Lowe, Yi I Wu, Aviv Tamar, Jean Harb, OpenAl Pieter Abbeel, and Igor Mordatch. Multi-agent actor-critic for mixed cooperative-compettive environments. Advances in neural information processing systems, 30, 2017.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo. Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, and Gavriel State. Isaac Gym: High Performance GPU-Based Physies Simulation For Robot Learning, 2021.
Michael Matthews, Michael Beukman, Benjamin Ellis, Mikayel Samvelyan, Matthew Jackson, Samuel Coward, and Jakob Foerster. Craftax: A lightning- fast benchmark for open-ended rein- forcement learning, 2024.
11

---

Published as a conference paper at ICLR 2024
Volodymyr Mnih, Adrià Puigdomènech Badia, Mchdi Mirza, Alex Graves, Timothy P. Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous Methods for Deep Reinforcement Learning, 2016.
Alistair Muldal, Yotam Doron, John Aslanides, Tim Harley, Tom Ward, and Siqi Liu. dm_env: A Python interface for reinforcement learning environments, 2019. URL http: /gitt:://githb.... com/
deepmind _nv
Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, Benjamin Van Roy, Richard Sutton, David Silver, and Hado Van Hasselt. Behaviour suite for reinforcement learning, 2020.
Georgios Papoudakis, Filippos Christianos, Lukas Schäfer, and Stefano V. Albrecht. Benchmarking multi-agent deep reinforcement learning algorithms in cooperative tasks. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks (NeurIPS), 2021. URL
www.://orgggb/b/eeoo/oooo 07869
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming John Wiley & Sons, 1994.
Andre Rohe. Vlsi data sets. https///www.chh...aeerroooooooo/ html
Alexander Rutherford, Benjamin Ellis, Matteo Gallici, Jonathan Cook, Andrei Lupu, Gardar Ing- varsson, Timon Willi, Akbir Khan, Christian Schroeder de Witt, Alexandra Souly, Saptarashmi Bandyopadhyay. Mikayel Samvelyan. Minqi Jiang, Robert Tjarko Lange, Shimon Whiteson, Bruno Lacerda, Nick Hawes, Tim Rocktaschel, Chris Lu, and Jakob Nicolaus Foerster. Jaxmarl: Multi-agent rl environments in jax, 2023.
Max-Philipp B. Schrader.
gym-sokoban.
sokobao 2018.
Adrien Ali Taiga, Rishabh Agarwal, Jesse Farebrother, Aaron Courville, and Marc Bellemare. Investigating multi-task pretraining and generalization in reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL /ttttps://ppnreeee
net/ Det/Eoruu?idd-SEe SER
Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, Nat McAleese, Nathalie Bradley-Schmieg, Nathaniel Wong, Nicolas Porcel, Roberta Raileanu, Steph Hughes-Fitt, Valentin Dalibard, and Wojeiech Marian Czarnecki. Open-Ended Learning Leads to Generally Capable Agents, 2021.
Emanuel Todorov, Tom Frez, and Yuval Tassa. MuloCo: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012. doi: 10.1 10.1109/ROSS201 2012..28611699
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL File331ee225577eee1lb000dddddLd3 Papp:es3aa.caaaaaaaaaap..ppppee
www.:neuuriicoooo neurips. files//aper/20177
Théophane Weher, Sébastien Racanière, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pascanu, Peter Battaglia. Demis Hassabis, David Silver, and Daan Wierstra. Imagination- augmented agents for deep reinforcement learning. 2018
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. Envpool: A highly parallel reinforcement learning environment execution engine. 2022.
12

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
The Jumanji Library
14
Logic Environments
A.2
Packing Environments
15
A.3
Routing Environments
17
A.4
Extending the library
19
A.5
Multi-agent RL with Jumanji
21
B
Networks
Experiments
23
Actor-Critic Baseline
23
C.2 Environment Parallelization Experiments
23
Flexibility Experiments
24
Scalability Experiments
25
D Roll out the Environment
27
D. Animate Episode
27
D.2 Placement of the Environment on Device
28
Discussion on Parallelization
29
Parallelizing the dynamics
29
E.2
Parallelizing auto-reset during training
29
License
30
13

---

Published as a conference paper at ICLR 2024
A THE JUMANJI LIBRARY
The Jumanji library contains 22 diverse RL environments designed to fast, flexible, and scalable. These environments are organized into three categories: logic, packing, and routing, and Table 1 shows the environments in each category. The following subsections provide a more detailed description of each environment and are sorted by the different environment categories.
Table I: Jumanji Environments.
Logic
Packing
Routing
Game2048
BinPack
Cleaner
GraphColoring
FlatPack
Connector
Minesweeper
JobShop
CVRP
RubiksCube
Knapsack
Maze
SlidingTilen Puzzulzle
Tetris
MMST
Sudoku
MultiCVRP
PacMan
RobotWarehouse
Snake
Sokoban
TSP
A.1 LOGIC ENVIRONMENTS
GAME2048
Component
Description
2305 Sowce 2834
Objective
Reach high-valued tile, aiming to surpass 2048.
2
Observation
Board, action mask, and step count.
16
Action
Up (0), right (1), down (2), or left (3).
Reset
4x4 grid with single cell being either or
Reward
Sum of merged cells upon taking an action.
2
Versions
Game 8v
GRAPHICOLORING
Component
Description
Objective
Color graph vertices without adjacent vertices sharing the
same color.
Observation
Graph, colors of the vertices, action mask, and current node,
Action
Integer to represent unique color.
Reset
Graph with 20 nodes and a 0.8 edge probability.
Reward
Negative the number of unique colors used for vertices. GraphColor corcooloonggv
Versions
EINESWEEPER
Component
Description
Objective
Clear the board without detonating any mines.
Observation
Board, action mask, number of mines, and step count.
Action
Coordinates of the square to explore.
Reset
Uniformly samples 10 mines in a 10x 10 grid.
Reward
reward for a square without mine, and 0 otherwise.
Versions
Mi swweeper--
14

---

Published as a conference paper at ICLR 2024
RUBIKSCUBE
Component
Description
Objective
Match all stickers each face to single color.
Observation
View of the current cube state and the step count.
Action
Tuple representing: face, depth, and direction of the turn.
Reset
Applies number of scrambles to a 3x3 solved cube.
Reward
Reward of 1 for solving the cube and otherwise 0. Solved ratio equal to 1.0. Rubi ksCube-vO, Rubi paabkkcccbeeprrtl---crraitt
Optimal Return
Versions
SLIDINGTILEPUZZLE
Component
Description
Objective
Move tiles adjacent empty slot until the puzzle is sorted.
Observation
Puzzle, position of the empty tile, and step count.
Action
Up (0), right (I), down (2), or left (3).
Reset
Applies number of random swaps
5x5 solved puzzle.
Reward
Reward of 1 for newly correct tiles and -1 for newly wrong ones.
Optimal Return
prop_ crop_correctly_plane Slidingf11lePuzzlee
ratio equal to 1.0.
Versions
SUDOKU
Component
Description
Objective
Fill N x N grid with digits I to N in each column, row, and subgrid.
Observation
Board, and the action mask.
Action
Tuple representing the square coordinates and the digit. Uniformly samples puzzle database.
Reset
Reward
Reward if the board correctly solved, and 0 otherwise.
Optimal Return
Solved ratio equal to 1.0.
Versions
Sudoku-vO, Sudoku very--asy-yyyyyee
A.2 PACKING ENVIRONMENTS
BINPACK
Component
Description
Objective
Pack boxes into a container to minimize empty space.
Observation
Available space, of unpacked items, and action mask.
Action
Tuple representing the EMS (space) ID and the item ID.
Reset
Randomly splits container into different items.
Reward
Volume utilization of the container (between 0.0 and 1.0).
Optimal Return
Volume utilization equal 1.0.
Versions
BinPack-v2
15

---

Published as a conference paper at ICLR 2024
FLATPACK
Component
Description
Objective
2D version of BinPack, place all the available blocks grid.
Observation
Current grid, available blocks.
Action
Block to place, rotation to make, coordinates on the grid.
Reset
Random instances guaranteed to be solvable.
Reward
Dense, fraction of the grid covered by the block.
Versions
Fl at Plat.ace Pack-v0
JoBSHOP
Component
Description
Objective
Minimize the time needed to process all the jobs.
Observation
Machines, operation details for each job, and action mask.
Action
Array containing job ID for each machine.
Reset
Instances with a number of jobs, machines, operations, and max. duration.
Reward
Reward of -1 at each time step if not terminating. JobShop-vD
Versions
KNAPSACK (KOOL ET AL.. 2018
Component
Description
Objective
Maximize value by packing items within weight constraint.
Observation
Weights, value, and packed status of the items.
Action
Integer to represent the next item to pack.
Reset
Uniformly samples item weights & values from unit square.
Reward
Sum of the values of the items packed in the bag.
Versions
Kn Knapsack-v.
TETRIS
Component
Description
Objective
Score maximum points by clearing lines in Tetris.
Observation
Grid, Tetromino, and action mask.
Action
Tuple denoting the x-position and rotation of the block.
Reset
Randomly samples Tetrominos from predefined list.
Reward
Proportional to the number of lines cleared.
Versions
Tetris-
16

---

Published as a conference paper at ICLR 2024
A.3 ROUTING ENVIRONMENTS
CLEANER
Component
Description
Objective
Clean many tiles as possible in given time budget.
Observation
Grid, agent location, action mask, and step count.
Action
Array denoting an action (left, right, up, down) per agent.
Reset
Randomly generates the structure of the grid.
Reward
The number of tiles cleaned at each timestep.
Versions
Cle cleaner-vC
CONNECTOR
Component
Description
Objective
Connect each start and end position and minimize steps.
Observation
Grid, action mask, and step count.
Action
Array with an action (left, right, up, down, no-op) per agent,
Reset
Uniform randomly places start and end positions on the grid.
Reward
for connecting agents, - -0.03 for non-connected agents.
Optimal Return
Ratio of connections equal 1.0.
Versions
Connector-v2
CVRP (KOOL AL. 2018)
Component
Description
Objective
Find shortest route visiting each city once and covering the demands.
Observation
Coordinates, demands, current position, and vehicle capacity, Integer representing the next city depot visit.
Action
Reset
Uniformly samples coordinates and demands.
Reward
Negative tour length of the route.
Versions
CVRP-vZ
MAZE
Component
Description
EHE
Objective
Reach the single target cell.
Observation
Maze, agent and target position, action mask, and step count.
Action
Up (0), right (1), down (2), or left (3).
Reset
Randomly generates the structure of the maze.
Reward
for reaching target, otherwise.
Versions
Maze-v0
17

---

Published as a conference paper at ICLR 2024
MMST
Component
Description
Objective
Connect all same-type nodes without using the same utility
nodes.
Observation
Node types, adjacency matrix, action mask, and current agent position.
Action
Integer array represent the node per agent.
Reset
Randomly splits the graph into multiple sub-graphs.
Reward
10 for valid connection, for no connection, for invalid
action.
Versions
MMST-v0
MULTICVRP
Component
Description
Objective
Route multiple agents to satisfy the demands of all cities. Coordinates, demands, time windows, penalties agent locations, local times, and capacities.
Observation
Action
Integer array represent the city for each agent. Uniformly samples coordinates and demands.
Reset
Reward
The negative tour length of all agents combined.
Versions
Mult Multicui 1ECRR-VV
PACMAN
Component
Description
Objective
Collect all the pellets and avoid the ghosts.
Observation
Grid, agent location, ghost location, pellet location,
power-pellet locations. and time left for scatter state.
Action
Up (0), right (1), down (2), left (3), or no-op (4). Deterministic generator start state of the game.
Reset
Reward
+10 for each pellet, +20 for a power pellet, +200 for ghost.
Versions
P PacMan-vO
ROBOTWAREHOUSE (PAPOUDAKIS ET AL.
2021
Component
Description
Objective
Deliver many requested shelves in given time budget.
Observation
View of other agents & shelves, action mask, and step count. No-op (0), forward (1), left (2), right (3), toggle load (4) per agent.
Action
Reset
Randomly places agents on the grid and uniformly selects shelves.
Reward
Number of shelves delivered during the timestep.
Versions
WobotWaarhoone
18

---

Published as a conference paper at ICLR 2024
SNAKE
Component
Description
Objective
Collect as many fruits without colliding with its own body.
Observation
Grid, action mask, and step count.
Action
Up (0), right (1), down (2), or left (3).
Reset
Randomly places snake's head position and fruit on the grid.
Reward
Reward is +1 upon collection of fruit and 0 otherwise.
Optimal Return
Equal 143.
Versions
Snake-vl
SOKOBAN (WEBER ET AL..
2018 GUEZ ET AL.. 2019 SCHRADER 2018
Component
Description
Objective
Box-pushing environment where the goal is place all boxes on their targets. Grid, and step count.
Observation
Action
Up (0), right (1), down (2), or left (3).
Reset
Sample from a dataset (Guez et al., 2019)
Reward
-0.1 for each step, +1/-1 for each correct/incorreet box, +10 success. Sokoban-vD
Versions
TSP (KOOL ET AL.
2018)
Component
Description
Objective
Find shortest route, visit each city once, and return to starting
city.
Observation
Coordinates, current position, trajectory, and action mask. Integer representing the next city visit. Uniformly samples coordinates from a unit square.
Action
Reset
Reward
Negative tour length.
Versions
TSP vvl
A.4 EXTENDING THE LIBRARY
Below is template of how to extend Jumanji's Environment interface to create a new environ- ment:
from typing import Tuple, NamedTuple
import chex
from chex import dataclass
5
import jax
from jumanji import specs
from jumanji. import Knvi ronment.
from jumanji. eihanji.types import Timestep, restart, termination, transition
10
8dataclass
13
class State:
14
key: chex. PRNGKey
19

---

Published as a conference paper at ICLR 2024
class Observation (NamedTuplc)
20
class MyEnv (Knvi ronment [State])
def init (self):
25
def reset
self, key: chex. FRNGKey
Tuple [State, TimeStep [observation]]: observation Observation
30
state State (key-key, ...)
timestep = restart (observation) return state, timestep
34
def step(
35
self, state: State, action:
[oberve((bsrvvttinn]] next_state (key ...) done
39
zeward next
oex:_observttion Observation
ncxt_ next_timestep jax 1ax. lond
43
done, termination,
transition, reward, next_observation,
48
return next_state, next_tinestep
52
def observation_spec (self) -> specs. Spec (pec[obeevvation) obs_ spec
return obs_sped
def (peioonspee(eelf) specs.Specs
action_spec = return actior_spec
After implementing a new environment, one may optionally add it to the registry using jumanji. register. This allows users then instantiate the newly implemented environment with jumanji make.
from jumanji. registration import register
2 register Ldd="MyEnv-vir entry_point="... :MyEn""""" kwargs-...)
env jumanji makee ("MyEnv-v0")
Wrappers Jumanji has several wrappers that one to convert Jumanji environment the API of one's 's choice. For instance, one can use Juman www.niiiToMMMtteeeee apper to make a di__env env environment, Jumanj ToGymkrrpper convert it the gym API.
20

---

Published as a conference paper at ICLR 2024
A.5 MULTI-AGENT RL WITH JUMANJI
Jumanji includes several environments whose action space is multi-dimensional, such as Robot Warehouse, an existing multi-agent environment (Papoudakis et al. 2021), well as Connector, MultiCVRP, MMST and Cleaner. These can be seen as homogeneous multi-agent environments where each agent is responsible for scalar action. For instance, Connector has N heads that need to connect to their target. One can see it as a single-agent environment with an action of shape (N, (one value per head), or multi-agent environment where each head its own agent and outputs a scalar action. This is called homogeneous multi-agent environment because each agent has the same observation and action shapes. However, Jumanji does lack a true multi-agent training algorithm, the intention is only to provide reasonable benchmarks. Instead, when training agent on multi-agent environment, Jumanji trains in the style of centralized training with centralized exccution (Lowe et al. (017)) and treats the environment single-agent one.
B NETWORKS
We provide an implementation of an advantage actor-critic (A2C) agent running on each environment in Jumanji. To do so, the algorithm is made agnostic to the environment and only assumes an anvironment-specific actor-critic network that, given an observation, outputs a policy over actions and value of the current state. This means these networks can be used for any actor-critic algorithm that uses state-value functions (c.g. A2C, PPO, TRPO, ctc).
Each environment comes with its own set of symmetries such as invariance to index permutations and is represented using a specific geometry, e.g. grids/images, sets, etc. Some symmetry groups are very large. e.g. TSP is invariant permutation of the node indices; such permutations form group of size n!! where n is the number of cities. It is inconceivable to hope to statistically learn a good policy for each of these permutations. Instead, we include geometrical biases the network architectures to get these symmetries for free. For instance, we make the TSP policy network equivariant to city permutations and the critic network invariant to those permutations.
We open-source actor-critic networks for each environment in along with their configs in jumanji/tres www.iiinnnoonnnoooeee /env. We list below some of the symmetries existing in each environment and describe what network is implemented as consequence.
BinPack The observation is composed of two sets: the spaces (EMS) and the items to pack in the container. Therefore, we use independent self-attention layer for each set and then use cross- attention between each set based on whether an item fits in the corresponding space EMS). Then, embeddings of both sets are joined using outer product to ensure permutation equivariance (or invariance for the critic) within both items and spaces.
Cleaner The grid is first copied times where is the number of agents. Each agent version of the grid where it is colored differently from the other agents. Then, a CNN is vmap-ed over the different grids and outputs feature vectors that are all passed through MLP output logits for each agent. The network is equivariant to permutations of agent locations.
The network is similar to Cleaner's in processing the agent grids independently via CNN first. Then, as opposed to Cleaner, the feature vectors are passed to transformer so that each agent attend one another in permutation-eqqiiarrine
Connector
way.
The network is adapted from Kool et al. (0018) to have a transformer encoder part that encodes all non-visited nodes and decoder that includes the current position to determine the next action. The important symmetry to respect here is equivariance to permutations of nodes.
CVRP
The observation is permutation invariant with respect to the order of blocks. Therefore, we use sequence model (transformer) to process all the blocks and obtain a permutation-equivarint policy. The grid is processed using a small U-net.
FlatPack
The observation being an image, we use a CNN with valid padding to prevent modeling the board edge the same way empty cells.
Game2048
21

---

Published as a conference paper at ICLR 2024
GraphColoring The observation contains nodes and colors. They represented as sequences and the graph-coloring problem is invariant to permutations of both nodes and colors. Consequently, the provided network uses transformer blocks alternating between self-attention on each sequence and cross- cross-attention between these sequences. Alternatively, a GNN could be implemented instead given the graph structure of the problem.
JobShop The observation contains two sequences with respect to which the problem is invariant to permutations of indices: jobs and machines. Each job itself a sequence of operations that have to run a specific machine. The problem is also invariant to renaming these indices by permuting machines. To leverage these symmetries, we build a network that does cross-attention between the operations and the machines they have to run for each job (parallel across jobs). Then, these operations sequences are reduced (averaged) to provide a single job embedding for each job. Cross-attention between jobs and machines leads to the action distribution (resp. value estimation) in a permutation-cquivarian (resp. invariant) way.
The problem invariant to permutations of the items. The implemented network is a transformer that is also adapted from Kool et al. (2018) and uses self-attention on the remaining
Knapsack
items.
Maze The observation is grid/image, implement CNN to process the grid before passing through an MIP to obtain the action logits or value estimate.
Minesweeper Same as the Maze, Minesweeper has a grid observation processed by CNN.
The observation contains information about two sequences: the different agents and the nodes on the graph. The problem is invariant to permutations of agent IDs and node indices. Hence, we implement& transformer network that alternates between self-attention layers each sequence and cross-attentin between the agents and their nodes.
MMST
MultiCVRP
The vehicles and the customers are first encoded. Then a series of self-attentio and cross-attention is used on both sequences.
PacMan We use a CNN to process the grid image. We then concatenate the grid embedding with diverse observation features like the agent's position and the ghosts'. A final MLP head projects these embeddings to a value (critic) or logits (actor).
Rarehhore
The observation contains feature vector for each agent. They are processed as sequence by a transformer to be equivariant with respect to permutations of agents.
RubiksCube The cube is just flattened and then passed through an MLP. The network would probably benefit from more symmetry seemnetry-preseringg architecture for this environment. Yet, it is not obvious how to do so.
SlidingTilePuzzle logits.
We use CNN to process the grid and then an MLP to project to value and
Snake The observation an image with 5 feature maps. Therefore, we process with a CNN
before using MLP to output logits value estimate.
Sokoban A CNN processes the grid, and final MLP head projects the grid embedding and the step count value or logits.
Sudoku The environment has many symmetries, including permutations of digits (e.g. 3 and 6 are swapped), permutations of columns within 3-column group, etc. We design network that is equivariant to the first symmetry. We flatten the grid and use this feature vector for each digit. We then do self-attention in the digit sequence to respect the permutation equivariance and then transpose back the grid dimension.
22

---

Published as a conference paper at ICLR 2024
Tetris The observation is composed of the grid and the tetromino. We process the former with CNN to which we concatenate the flattened tetromino processed by an MLP. Last, an MLP head outputs action logits value estimate.
TSP The network is almost equivalent to CVRP but does not include depot node. The symmetry that is respected is the permutation to node indices.
C
EXPERIMENTS
C.1
ACTOR-CRITIC BASELINE
Name
Version
Wall-clock time
BinPack
"BinPack-v2"
30h
Cleaner
"Cleaner-vO"
11h
Connector
"Connector-v2"
35h
CVRP
"CVRP-v1"
3h
FlatPack
"FlatPacX-vO"
48h
Game2048
"Game2048-v1"
3h
GraphColoring
"Graphcoloring-vo"
7.5h
JobShop
"JobShop-v0"
20 min
Knapsack
Knapsack-v1"
min
Maze
"Maze-v0"
min
Minesweeper
"Minesweeper-vO"
3h
MMST
"MMST-V0"
6h
MultiCVRP
"XulticVRE-v0"
1.5h
PacMan
"FacMan-v0"
4.5h
RohotWharchouse
kobotwarehoosse-o00
6.5h
RubiksCube
3h
SidiingTieePuzzte TilePuzzee
"S "SLidinngTilePuzzee=000
40 min
Snake
"Snake-v1"
1h
Sokoban
"Sokobar.vv0"
3,5h
Sudoku
"Sudook--vory--aayy-o0
3,5h
Tetris
"Ictris-v0"
3,5h
TSP
"ISF-v1"
2.5h
Table 2:
Correspondance between the name reported in figure 3a legend and the environment version.
We train actor-critic agent each environment different seeds in figure 2 the registered versions displayed in table 2 the open-sourced networks described in section B and the configs available on GitHub well.
For this,
We use the train. script blob/ma i juwann:/traaiiinggeeenneee
py that alternates between evaluating the agent and training it Each training epoch consists of a number of num_learner _learne_sttess of collecting stetep on total cotal_batch_____ size environments in parallel. If multiple devices are available, the batch of environments is split between the accelerators, on which trajectories are collected directly using a local copy of the model. After collecting trajectories, we compute an A2C loss and update the parameters.
The A2C loss is weighted mixing of three terms: policy gradient term -A(r) log no (T), a critic term 2(r)² and an entropy bonus -H(T)(TT)
We release checkpoints for all the agents on Hugging Face Hub.
C.2 ENVIRONMENT PARALLELIZATION EXPERIMENTS
This section provides more details on the parallelization experiments described in section 4.3 The first experiment aims to demonstrate how the environment dynamics be parallelized to merrase the steps throughput, and the second experiment shows how a full training pipeline can benefit from parallelization by reducing the time optimality.
23

---

Published as a conference paper at ICLR 2024
Environment Parallelization This experiment shows how the raw environments' speed increases with the number of environments that are run in parallel. The number environments, gradually increasing from 128 to 8192, is equally divided among 8 TPUv4 cores. To evaluate the speed of the environment dynamics, only the duration of the actual step function has been considered. Starting from a generated initial state, the same action is applied 50 times in all the parallel environments. This procedure is run 500 times to form an epoch. The total number of steps ran in an epoch (50 multiplied by 500 multiplied by the number of parallel environments) is divided by the epoch duration obtain the average number of steps per second. The throughput which is reported in figure Ba is the average over the second epoch run, the first epoch being longer as biased by the JIT compilation time which is only run once and thus is not representative of the final environment dynamies performance. The same experiment is executed on GPU (V100) and CPU as well, and the results are shown Figure 7 Table 2 contains the version of each environment in which speed was measured in this experiment.
10
Mise
16
Knapsark
Cloenee
Seoddann
Mersivupp
Slops per Second
Stiing Th: Puzze
Steps Second
OViP
Surobaan
t0
10
Stoing Te Paace
Grape Colring
JosShap
Gomee2588
10
Jos5ivo
10²
Colringg
Pacnan
10"
10"
Suinay
SrPuck
StePukk
PurPak
PlaPack
10
2
212
21
2
2
24:
2
21
2
Number Parallel Environments
Number of Parallel Environments
(a) GPU
(b) CPU
Figure 7: Scaling of the effective environment steps per second (throughput) for each registered environment the number of parallel environments increases (a) GPU (Tesla V100) and (b) CPU (2 cores).
Hardware scaling of A2C training This experiment shows the impact of increasing the training resources of the Jumanji A2C agent on the Connector-v2 environment. The training was done on different hardware: CPU with 8 cores, GPU (RTX 2080 super) and TPU-v4 with a varying number of cores, i.e. 16, 32, 64, and 128 cores. The training run for 1000 epochs 100 learning steps in which 256 trajectories of length 10 are sampled. The sampling of the trajectories is split across the available devices, but the number of environment steps sampled per epoch is the same for all the training settings. The A2C agent is run without normalizing advantages, with discount factor of 1, a bootstrapping factor of 0.95, and learning rate of 2 x 10 10-4, We demonstrate almost linear scaling hardware table by plotting training convergence speed as function of the number of TPU cores.
TPU cores
8
16
32
64
128
Convergence time (h)
12.8
4.2
3.2
2.0
1.4
Table 3: Time reach 92% performance as a function of the TPU cores. This table completes the experiment presented in figure 3b
C.3 FLEXIBILITY EXPERIMENTS
This section provides additional details on the experiments conducted using multiple generators for the TSP problem (5.1)
24

---

Published as a conference paper at ICLR 2024
Network The networks of the two A2C agents in this experiment are identical, and additionally, it is the same network used in the baseline experiment for the TSP environment (section 4.2). Full details of the network used in this experiment can be found in Appendix B
Training Procedure The A2C agent trained the single uniform generator was trained in an identical manner as the A2C agent for the TSP baselines experiment, this includes the same training hyper-parametens (e.g, the sequence length, batch size, and so forth). The only notable difference between the training process of the two agents is that the baseline TSP agent is trained on 50 cities whereas the agent for this experiment is trained 20 cities. The A2C agent trained on the combination of the uniform generator and the three custom generators (cluster, linear compression, and explosion). The custom generators were implemented by inheriting the abstract generator class from their environment and then modifying the call method to return instances with the desired initial state distribution. This A2C agent is trained in similar manner with the following difference. The batch of data used to update the single generator agent contains the agent's 's trajectories on only uniform instances, whereas, the four generators agent's s batch of data consists of its trajectories on uniform, cluster, linear compression, and explosion instances. In the former, the agent interacts solely in the environment with the uniform generator, whereas in the latter, the agent sequentially interacts with four environments each with specific (uniform or custom) generator. Lastly, both of the agents were trained for 300 million environment steps.
Evaluation Dataset The dataset was created using the VLSI TSP Benchmark Dataset. There are 102 TSP instances in the VLSI dataset with instance size (i.e., number of cities) ranging from 131 to 744710. Since we conduct this experiment with a TSP environment with 50 cities, we randomly sample 50 cities from each of the 102 instances to obtain unseen, real-world validation dataset of 102 instances each with 50 cities. This dataset is used to evaluate the A2C agents during training. Additionally, we create larger dataset with 1020 instances by randomly sampling 50 cities from each of the 102 instances 10 times. This larger dataset is used to create a lower-variance estimate of the performance of both agents test time.
presents the results of the two A2C agents on the unseen, real-world instances, both during training and at test time, and it can be seen that the agent trained on a broader initial state dis- tribution (i.e., four generators agent) outperformed the agent trained solely on uniform instances. To fur- ther analyze the two agents, they were also evaluated on random set of uniform instances during train- ing, and Figure g depicts the learning curves of the agents. The aim of this evaluation was determine the impact of training on a broader data distribution in-training distribution instances. Even though it does appear that the learning curve of the four generators agent lower/better than the single generator agent's 's curve, this difference in per- formance is not significant. Therefore, these results suggest that training on wider data distribution does not lead to any significant improvement or degrada- tion of performance when evaluating on instances seen during training.
Results Section 5.1
7.0
Single Generator
6.8
Four Generators
Tour Length
6.6
6.4
6.2
6.0
0.s
1.0
1'5
2.0
2.5
3.0
when evaluating
Environment Steps
Figure 8: Learning curves of two agents train- ing TSP and evaluated on instances sam- pled from a uniform generator. One agent (blue) samples from single uniform gener- ator versus the other agent (orange) samples from four generators (uniform, cluster, explo- sion, and compression). Lower tour length is better.
C.4 SCALABILITY EXPERIMENTS
This section provides additional details on the experiments conducted with varying degrees of problem complexity for the Connector, BinPack, and RubiksCube environments 5.2).
The agents in the scalability experiments are actor-critic networks which are fully defined in Appendix(B [B] Therefore, the agents for each environment have identical networks with the sole difference being the complexity level of the environment they trained on.
Networks
25

---

Published as a conference paper at ICLR 2024
The training process is identical to the baselines experiments, use the same training prper-parameters and use the described in Appendix C.I The only difference is that for the purpose of the scalability experiments, we instantiate the Connector, BinPack and RubiksCube environments with different configurations (whereas for the baselines experiment, the default environment parameters configuration is used). The code snippets below show the instantiation of the different environments for the three problems.
Training
py script in the same manner as
Connector setup
from jumanji. environments import Connector from jumanji.
2
htw.://mmtttstotttttooooooeeeee generator import Randommal
Connector httpaccccecccccccccccccaaaaaalla (ggp:e/aattttogddddd_sseeeeeee
2))_agents=))
env_2 2
Connector otg.rrreeeeaaaaaaaaiaaaaaaaaaana
(gridssiiii
num_agents- 10))
5
env 3 CCnnector htgerreeeeettaaaoaaaaaaaaaaeeeer nun_agents- 12))
env_ Connector
nun_ägents- 14))
BinPack setup
from jumanji enviennmnnss import BinPack
wwin_iicc...._aeeeeeeaaaae from jumanji. environments. wwwirooments..ackingg bin pack. generator import RandomGenerator
BinPack
5
generator RandomGeneeattrrr
max_num_itentt ittms=itess. max num ems-15, split num same items=2,
obs "oo__num_emm11,
env 2 BinPack 10
generaaoor-aandodGeeeeee
bes_nuumiitems=22 max max_num__eeteeeeaaaaaaaaaae qum_ens=40, split_nma same iitem==,,
obs ____umm__mm_00
15
env_ 3 = BinPack(
generator- generve.orr-annmmeeeerrrrtror
max_num_itens -ammmmtttmsm300 max_num_ems-60, split_num same tems=2,
19
20
obs _ ems=50,
RubiksCube setup
from jumanji. environments import RubiksCube
from jumanji. wtw.:ttss.ooo.............. cube. generator import ScramblingCenerator
Rubi kaCube
generator-Scrame ingeeneratorrrr (cube_size-3, nun_scrambl commnbls__aee Beterrnnbis__o__rseettt--
env
RubiksCube( genereattrs (cube size=s, scrambles reset=3),
10 env 1
RubiksCube
num_scraablee__oorretttt
env
RubiksCuba
(cbbaa mun_scrambles.on_res
env : RubiksCube
26

---

Published as a conference paper at ICLR 2024
generator- ing.mmrrttrnneeeerarrrroo.o (cube_size-3, nun_scrantles_on_ree
19
env_
Rubi kaCube
20
generator
Snnerrator-Scramb.o. Scramb www.mmiiiggeeeegegoooooo (cube_ size=4, num sum_sccambles on on btt.eiee_______ee___
Table 4: Results of the scalability experiments.
Final Performance
Environment
Configuration
Complexity Level
= 10, N = 2
Easy
0.983
G 10, 10
Medium
0.885
Connector
N 12
Medium
0.855
14,
Difficult
0.785
N 10
Easy
0.993
BinPack
N
Medium
0.934
N 30
Difficult
0.870
C= N
Easy
0.987
C=4, N
Medium
0.829
C= N
Medium
0.838
RubiksCube
3= N7
Difficult
0.213
N=
Difficult
0.000
4=4,N=5
Difficult
0.000
Results We further describe the results obtained from the scalability experiments in Table This table shows the different configurations for each environment along with its qualitative definition of the complexity level and shows the final performance obtained by the A2C agent for each environment and configuration. It can be seen that with the increasing complexity of the environment, the agent performance worsens.
D ROLL OUT THE ENVIRONMENT
D.1 ANIMATE AN EPISODE
Below is code example of how to take random actions in the BinPack environment and an- imate episode. This code can run in notebook e.g. Google Colab. Please see the load ses__chhcccinnntt ipynb notebook on httts::// to load pre-trained agents jumanji/blos http///mminnaammaeeeeeeeme or roll out a random policy.
com/ instadeepai
Spip install --quiet (train)
&natplot otlib notebook import jax
5
import jumanji
6
from jumanji. teriiiiiannnnnnnn import networks
env jumanji. make ("BinFack-v2")
reset_fa jax. jit (env.resse
10
step_fn jax. jit (env.step) randon_policy networks makeerandoo htt.rradddoooolccyiiiiiie (env. unwrapped)
jjax.jit def select_randdm_acte (observation, key):
15
"*"Call "random pollccy which expects a batch of observations.
Matched__oserrattie = tax.tree_ute tree_mapp((
27

---

Published as a conference paper at ICLR 2024
lambda x: [None], observation,
19
return random_poil i ey (batched_obseervatine key) .qqueze (axis-0)
20
key jax. random. www.ee.ee.....eyy........ (0) state, timestep reset in (key)
states [state]
24
Goop until the episode is done. while not (amestep.last())
Select action.
action_key, key jax random.... coriido.spppitt.. (key)
29
action
(t obwervatiiocrc. action_key) step in the environment.
state, timestep step fn (state, action)
states states. append
34
env. an mate (states, interval-100)
D.2 PLACEMENT OF THE ENVIRONMENT ON DEVICE
Effective Steps per Sec.
10t
10a
10
GPU (L 414
GPU
GPU
CAU
Figure 9: Throughput ablation of the full Anakin GPU setup. The experiment uses the Connector environment and is run V100 GPU with batch size of 2048. The pipeline alternates between phases of acting (policy inference and environment step) and learning (back (back- propagation). CPU: The whole Jumanji pipeline is run the CPU. GPU (L): The inference and environment step remain the CPU but the data sent to the GPU for the learning step (L). GPU (L + I): Only the environment step done on CPU. The inference (I) and the learning step done the GPU. GPU (L I + E): The Anakin implementation. The whole training pipeline is performed on GPU, now including the environment step (E).
Implementing environments in Jax achieves its fullest potential when using the Anakin (Hessel et al. 2021) architecture for fully optimizing device accelerators. The pipeline consists of synchronous execution of actor inference (action selection), environment step, and learner step propaarropgtttion all on the device (e.g. GPU, TPU), High throughput is achieved by removing any host-device communication (e.g. CPU to GPU) during the process. We run an ablation study in figure to study the speed-up that comes with implementing the environment in Jax and running on the device.
We use the Connector environment with batch size of 2048. We observe that most of the speed gain arises from running the environment on GPU and not transferring data between the host and the device.
28

---

Published as a conference paper at ICLR 2024
E DISCUSSION ON PARALLELIZATION
E.1 PARALLELIZING THE DYNAMICS
framework within Jumanji is agnostic to using discrete or continuous actions, all of the implemented environments use discrete action space. Having discrete actions often leads to the impossibility of fully parallelizing the dynamics, i.e. the environment step function. For instance, in Sudoku, there are 18 possible actions faces and 3 different rotations). Each action leads to doing a different operation on the cube (3D array). When the step function is vmap-ed, the discrete choice of selecting the rotation to perform as a function of the action is not parallelizable within the SIMD (single instruction, multiple data) paradigm. Therefore, vmap-ing the dynamics leads to transforming the conditional branching to a select XLA statement. This means each the 18 actions is performed for the whole batch and then the correct rotation is selected based on the action.
Although the
When executed, the code below shows how JAX's jax. . lax. cond transformed into a sel select when the function is vmap-ed.
import jax import jax. numpy as inp
3
def [x, bool_): :
return jax. lod. (bool_, lambda a+100, lambda x)
7 args = ampparrry(()) float), array (False)
print wtt.xllla_oopuaatiiioo (f) (*args) as as_hlo_text ())
print (a-nt("--. ("---")
print
jax. awll_ccopuuattteeeee (jax. vmap (2))
12
jap. array ( [C), float), jap array (False)),
J.as tla__hlo_text
Because the vmap-ed dynamics have to generate all possible actions for the whole batch, it may explain why environments end up being slower than expected hardware accelerator. Yet, the use of hardware-acceleratee environments really shines when training a neural network as the policy since avoid transferring data between the CPU and the accelerator.
E.2 PARALLELIZING AUTO-RESET DURING TRAINING
During training, we roll out few steps batch of parallel environments with automatic reset behavior. This means any of the environments that reaches a terminal state is automatically reset to an initial state (with discount of O).
In Jumanji, we implement an environment wrapper called Aut Wrapper to do this auto- reset automatically. This wrapper first calls the environment step function, then checks it reaches a terminal state and if so. it resets the environment. Similar to explained above, this conditional statement is not parallelizable when used with vmap. As a consequence, if the wrapper is vmap-ed (for instance by wrapping it into Jumanji's s Vmapärapperr) both branches (resetting and not resetting) will executed on all environments across the batch, at every timestep. This may very slow if the environment reset is slow function, which is the case for Rubikscube where the reset function is literally 100 times as computationally heavy as the step function.
An alternative to calling the reset function at every step is to CmapAAttoResetin Wrapper that equivalent to vmap-ing the auto-reset behavior but ac- tually only vmaps the step function and then loop the environments to reset the ones that reach terminal state. This way. none of the states in the batch has terminated, the wrapper will not call reset once, compared to the previous wrapper which would still call reset the whole batch.
Jumanji's
29

---

Published as a conference paper at ICLR 2024
F LICENSE
Jumanji released under an Apache License 2.0.
30


Reviews:
[{"review_id": "DhIWRKrdHo", "reviewer": "Reviewer_7YMa", "summary": "Jumanji is a suite of scalable reinforcement learning environments designed for RL research with industrial applications. It provides a collection of environments that are fast, flexible, and scalable, focusing on combinatorial problems and decision-making tasks. Jumanji leverages JAX and hardware accelerators to facilitate rapid research iteration and large-scale experiments. It stands out from existing RL environments by offering customizable initial state distributions and problem complexities and includes actor-critic baselines for benchmarking. The paper demonstrates Jumanji's high scalability and flexibility through experiments, positioning it as a tool to advance RL research.", "questions": "What are the most important research challenges do you expect Jumanji will help to address?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Good paper and an important engineering contribution to an area of research in NP-hard combinatorial optimization problems (COPs). Solid design and software engineering work to make Jumanji modular, scalable, and fast and to fully unlock the power of hardware acceleration. The set of environments and tasks is complimentary in some sense to continuous control Jax-based training environments created by Google Brax team and will help to advance research in the area combinatorial problems and decision-making tasks.", "weaknesses": "A lack of a new research results and novel approaches. But it’s totally expected from such kind of more engineering oriented projects.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NXGgKm93sx", "reviewer": "Reviewer_S6or", "summary": "This paper presents a Jax based RL environment suite called Jumanji. The 18 environments focus on combinatorial optimization problems, designed to be fast, flexible, and scalable. They also provide an A2C benchmark and examples to motivate these problems.", "questions": "- How important is hardware flexibility? Are TPUs widely used outside google?\n- Gamma is put in the MDP formalism of Jumanji. Although this can be seen both in and outside of the tuple, is there any explicit representation of it in the software? I.e. in the Jumanji environments, clearly all the other elements of the tuple are required to be defined for a functioning environment, but is the gamma represented?\n- It would be beneficial to give more of an explanation of the state, just another sentence or so, explaining (perhaps with an example) what it is and contains. I assume it is a pytree (since the observation is), but is the key element required? Does step have to split the key necessarily if it doesn’t use it (small details like this could go in the appendix)?\n- Environment version control is mentioned, but how often are changes made that increment this version? Version control is nice, but if there are hundreds of versions, it isn’t a panacea.\n- Appendix C2 demonstrates weak (sometimes negative) scaling on CPU. Why is this the case? I would expect some speedup up to the 8 cores (assuming you are mapping across all cores, jax by default will just work with 1 (https://github.com/google/jax/issues/5022). \n- Why does figure 3a start at 2^7 environments? The on many of the environments doesn’t seem as impressive as it could if this started at 2^0 perhaps\n- Why is it called Jumanji?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- RL for many years has struggled with good environment code maintenance and support, and it is good to see this problem continue to be addressed \n- The code base seems to be well designed and documented, the doc strings are generally informative and type hints are present.", "weaknesses": "- NP Hard optimization style problems have seen some interest in RL, but are not as common in literature, it would be beneficial to have more citations justifying their uses or explain more how common RL problems can be rethought into the COP formalism\n- Having some sort of UML or diagram would be of great help to understanding the API.\n- I don’t think random policy adds anything in Figure 2. It is expected that random does poorly and I’m not sure it adds much (given the trends of the curves, the impression of learning comes across)\n- I’m not sure how much the y-axis labels matter in Figure 2 given how much clutter they add. A lot of these environments are not super common (and even in common Atari environments human normalised performance is increasingly common as a metric since the actual scores don’t mean much to most people). As long as they are all linear axes, and the optimal performance is there, all that matters is that the lines are going up (since this isn’t an algorithm paper, this figure is just showing things can learn in your environments).\n- A plot like Fig 3(b) with number of TPUs vs. time to reach a certain performance could make a good figure (for the appendix at the very least)\n- If CPU is not visible on the plot, I would just leave it off the labels and keep the text remark\n- Although there are a lot of different environments implemented, it would be beneficial to have a point of comparison. As the authors note, there has been a fair amount of work in high performance environments already. Even if you can’t make a 1 to 1 comparison (because the environments are not the same), finding something of comparable complexity and having a figure in the appendix would help to ground the speedups.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Jd33b4rzRz", "reviewer": "Reviewer_VTas", "summary": "The authors propose Jumanji, a diverse set of accelerated environments written in JAX focused on NP-hard combinatorial optimization problems (COPs). Jumanji is fully open-source, fast, flexible, and scalable, covering 18 environments such as TSP (Travelling Salesman Problem). The authors also present A2C learning curves in these 18 environments to demonstrate end-to-end learning. Interestingly, Jumanji can tune the difficulties of the environments, showing that these environments can get exponentially more difficult to solve.", "questions": "I am curious why the authors chose A2C as the training algorithm instead of P", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* Open-source accelerated environments in COPs: most of the accelerated environments are in robotics (e.g., NVIDIA's isaacgym or Google's brax), but I like the authors specific focus on NP-hard optimization problems.\n* Optimal performance in some games: I like the authors added the reference optimal performance in some of the 18 environments.", "weaknesses": "I do not see any major weakness. One issue is that Figure 3 does not seem like a fair comparison with GPU. In particular TPU-v4s should be compared with A100s instead of RTX 2080 Super.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yeazBWAtMo", "reviewer": "Reviewer_d6hE", "summary": "Suite of Combinatorial Optimization benchmarks in JAX. Some experiments with an AC algorithm", "questions": "What is the contribution of this paper?\n\nWouldn’t it make more sense to remain Gym-compliant in providing a Gym-JAX-CO implementation? This remains implicit, and is not explained.\n\nWould a wrapper be possible for a Gym API? Could you use stable baselines unchanged?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "New benchmarks are always good, especially CO, where fewer benchmarks are available.", "weaknesses": "No comparisons to other benchmarks or implementations. It is a sympathetic and perhaps substantial effort, but lacks elements that would achieve wide adaptations. The software engineering is there, the science is unclear.\n\nThere are Gym-JAX environments, and there are CO-Gym implementations (OR-Gym).  \n\nIt appears that Jumanji does not follow the Gym interface. Stable Baselines algorithm are therefore not a drop in plugin.\n\nExplain clearly the difference with a Gym interface. Why this choice? \n\nCarrying around explicit state deviates from an RL principle, that the environment has the state, and not the agent. \n\nExperimental validation with limited algorithms. No comparison to other benchmarks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)