Paper: Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX (ICLR-2024)

Reviewer comment (Experiments):
I do not see any major weakness. One issue is that Figure 3 does not seem like a fair comparison with GPU.

Author response:
We agree with the reviewer that the performance obtained with A100s should be closer to that of TPUs, but the formers were not available at the time of running the experiments. We would like to clarify that the experiment illustrated in Figure 3 is meant to highlight the increase in throughput as one scales the available hardware, and not to showcase differences between TPUs and GPUs.

Reviewer comment (Experiments):
In particular TPU-v4s should be compared with A100s instead of RTX 2080 Super.

Author response:
We agree with the reviewer that the performance obtained with A100s should be closer to that of TPUs, but the formers were not available at the time of running the experiments. We would like to clarify that the experiment illustrated in Figure 3 is meant to highlight the increase in throughput as one scales the available hardware, and not to showcase differences between TPUs and GPUs.

Reviewer comment (Experiments):
I am curious why the authors chose A2C as the training algorithm instead of P.

Author response:
Indeed, we chose the A2C algorithm for our benchmark instead of, say, PPO or other more advanced algorithms. We purposely wanted to show how simple algorithms, when combined with neural network architectures that account for the symmetries of the considered problem, can achieve reasonable performances. This benchmark can then be used by researchers to build better-performing algorithms and use our A2C performance as a reference.

Reviewer comment (Novelty):
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.

Author response:
We have adjusted the text in the "Combinatorial Optimization Problems" paragraph of the related work section to better emphasize how COPs can be tackled with RL. We included the works of [Kool et al., 2018; Hottung et al., 2022] which proposed a transformer architecture to efficiently treat COPs as RL episodes.

Reviewer comment (Presentation):
Having some sort of UML or diagram would be of great help to understanding the API.

Author response:
We thank the reviewer for the feedback and agree that a UML diagram would improve the clarity of the code base and provide useful insight for new developers working with Jumanji. We will provide this as part of the appendix and in the repository documentation.

Reviewer comment (Experiments):
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).

Author response:
Although a random policy is expected to perform poorly, we have decided to include this in our plots as we have repeatedly found that the random policy score is a useful lower bound for debugging new algorithms. An agent whose neural network parameterization and/or algorithm is incorrect could perform close to the random policy in expectation. Therefore, it helps to assess algorithm implementations and we think it will be useful to the community as well.

Reviewer comment (Presentation):
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).

Author response:
We agree with the reviewer that the reward scales are generally arbitrary, and do not bring in itself more information than the relative improvements we observe during the training. However, in some environments (e.g. TSP or CVRP where the reward is the negative tour length, JSSP where it’s the negative makespan), they directly correspond to the underlying problem objective and align with previous work, so having it could make the comparison easier in general.

Reviewer comment (Evaluation):
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).

Author response:
We thank the reviewer for their feedback and agree that this would be insightful to the reader. We have amended section C.2 of the appendix with such a table, and we hope it helps better understand the experiment.

Reviewer comment (Experiments):
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.

Author response:
Since we can’t compare the throughput of a Jumanji environment with a non-JAX version of it, we can at least demonstrate the difference in throughput between running the environment on a CPU (Figure 7b) and on a GPU (Figure 3a). We observe that running the environment on an accelerator brings orders of magnitude speed-ups compared to having the simulation on a CPU.

Reviewer comment (Reproducibility):
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?

Author response:
Gamma is usually specified as part of the problem (i.e. the MDP). However, most environment implementations (e.g. in Gym) do not include gamma. Therefore, we have adhered to the standard practice and have not included it. Moreover, all Jumanji environments have a limited horizon, which allows for undiscounted values (gamma = 1). Nonetheless, we mention gamma in Section 3.1 when properly defining the RL objective.

Reviewer comment (Writing):
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)?

Author response:
We thank the reviewer for their feedback and have accordingly updated the **State** paragraph of section 3.2 for better clarity. The reviewer is correct in that the state is a pytree. Additionally, the key is needed in order to handle stochasticity in wrappers like the `JumanjiToDMEnvWrapper`. If the environment is not stochastic, i.e. if it does not use the key, then it does not have to split it. On the contrary, it will split the key when any stochastic computation is needed within the step function.

Reviewer comment (Reproducibility):
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.

Author response:
We agree and therefore aim to use the version control in the environments sparingly, only increasing the version after changes that affect the environment's behavior/performance. This policy is in line with that followed in Gymnasium (formerly, OpenAI Gym), where there are at most a handful of versions for each environment.

Reviewer comment (Experiments):
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).

Author response:
We emphasize that Fig 7.b displays the scaling of the environment performance on a 2-core CPU and not a TPU.

Reviewer comment (Presentation):
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.

Author response:
The higher the number of parallel environments, the greater the parallelization on hardware accelerators. As a reference, Brax [Freeman et al., 2021] trained with a batch size of 2048 and experimented with a number of parallel environments ranging from $2^7$ to $2^{16}$.

Reviewer comment (Novelty):
A lack of a new research results and novel approaches. But it’s totally expected from such kind of more engineering oriented projects.

Author response:
We agree with the reviewer that this work does not present new research results or approaches. Instead, our contribution focuses on utilizing advances in both hardware architecture and the high-performance computing library JAX to provide a test bed of hardware-accelerated environments. This contribution is impactful to the research community as it enables fast iteration of research and, due to the efficiency of our JAX environments, removes the need for large compute clusters to carry out meaningful RL research, as is the case with the majority of RL benchmarks.

Reviewer comment (Novelty):
What are the most important research challenges do you expect Jumanji will help to address?

Author response:
Within the COP community, recent works [Kwon et al., 2020; Hottung et al., 2022; Grinsztajn et al., 2023; Chalumeau et al., 2023] have tackled problems such as generalization to bigger problem sizes and different problem distributions. Follow-up works could benefit from using Jumanji to train and evaluate their methods on a wide distribution of problems. Additionally, we would like to invite researchers and engineers working on real-world applications to use the flexibility of Jumanji to tackle the very problems that hinder RL applications.

**References**:

- Kwon et al. POMO: Policy Optimization with Multiple Optima for Reinforcement Learning. NeurIPS (2020).
- Hottung et al. Efficient Active Search for Combinatorial Optimization Problems. ICLR (2022).
- Grinsztajn et al. Winner Takes It All: Training Performant RL Populations for Combinatorial Optimization NeurIPS (2023).
- Chalumeau et al. Combinatorial Optimization with Policy Adaptation using Latent Space Search. NeurIPS (2023).