# Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX

arXiv: https://arxiv.org/abs/2306.09884 · HTML: https://mlsanity.com/paper/2306.09884/

## Key facts

- Cited by 58 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 33 references that are in the map.
- First posted on arXiv 16 Jun 2023.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environments](https://mlsanity.com/paper/1706.02275/) (2017)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [Attention, Learn to Solve Routing Problems!](https://mlsanity.com/paper/1803.08475/) (2018)
- [Machine Learning for Combinatorial Optimization: a Methodological Tour d'Horizon](https://mlsanity.com/paper/1811.06128/) (2018)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [POMO: Policy Optimization with Multiple Optima for Reinforcement Learning](https://mlsanity.com/paper/2010.16011/) (2020)
- [Unity: A General Platform for Intelligent Agents](https://mlsanity.com/paper/1809.02627/) (2018)
- [Benchmarking Multi-Agent Deep Reinforcement Learning Algorithms in Cooperative Tasks](https://mlsanity.com/paper/2006.07869/) (2020)
- [OpenSpiel: A Framework for Reinforcement Learning in Games](https://mlsanity.com/paper/1908.09453/) (2019)
- [Behaviour Suite for Reinforcement Learning](https://mlsanity.com/paper/1908.03568/) (2019)
- [An investigation of model-free planning](https://mlsanity.com/paper/1901.03559/) (2019)
- [EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine](https://mlsanity.com/paper/2206.10558/) (2022)
- [OR-Gym: A Reinforcement Learning Library for Operations Research Problems](https://mlsanity.com/paper/2008.06319/) (2020)
- [Podracer architectures for scalable Reinforcement Learning](https://mlsanity.com/paper/2104.06272/) (2021)
- [ORL: Reinforcement Learning Benchmarks for Online Stochastic Optimization Problems](https://mlsanity.com/paper/1911.10641/) (2019)
- [Accelerating Reinforcement Learning through GPU Atari Emulation](https://mlsanity.com/paper/1907.08467/) (2019)

## Built on by (most cited first, 30 of 58)

- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Simplifying Deep Temporal Difference Learning](https://mlsanity.com/paper/2407.04811/) (2024)
- [RL4CO: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmark](https://mlsanity.com/paper/2306.17100/) (2023)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPS](https://mlsanity.com/paper/2408.01584/) (2024)
- [POGEMA: A Benchmark Platform for Cooperative Multi-Agent Pathfinding](https://mlsanity.com/paper/2407.14931/) (2024)
- [QDax: A Library for Quality-Diversity and Population-based Algorithms with Hardware Acceleration](https://mlsanity.com/paper/2308.03665/) (2023)
- [The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Design](https://mlsanity.com/paper/2406.17949/) (2024)
- [FrontierCO: Real-World and Large-Scale Evaluation of Machine Learning Solvers for Combinatorial Optimization](https://mlsanity.com/paper/2505.16952/) (2025)
- [Benchmarking Massively Parallelized Multi-Task Reinforcement Learning for Robotics Tasks](https://mlsanity.com/paper/2507.23172/) (2025)
- [minimax: Efficient Baselines for Autocurricula in JAX](https://mlsanity.com/paper/2311.12716/) (2023)
- [Sable: a Performant, Efficient and Scalable Sequence Model for MARL](https://mlsanity.com/paper/2410.01706/) (2024)
- [PC-Gym: Benchmark Environments For Process Control Problems](https://mlsanity.com/paper/2410.22093/) (2024)
- [SocialJax: An Evaluation Suite for Multi-agent Reinforcement Learning in Sequential Social Dilemmas](https://mlsanity.com/paper/2503.14576/) (2025)
- [Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy](https://mlsanity.com/paper/2407.07333/) (2024)
- [PARCO: Parallel AutoRegressive Models for Multi-Agent Combinatorial Optimization](https://mlsanity.com/paper/2409.03811/) (2024)
- [Trust-Region Twisted Policy Improvement](https://mlsanity.com/paper/2504.06048/) (2025)
- [ROTATE: Regret-driven Open-ended Training for Ad Hoc Teamwork](https://mlsanity.com/paper/2505.23686/) (2025)
- [BuilderBench: The Building Blocks of Intelligent Agents](https://mlsanity.com/paper/2510.06288/) (2025)
- [SPO: Sequential Monte Carlo Policy Optimisation](https://mlsanity.com/paper/2402.07963/) (2024)
- [Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX](https://mlsanity.com/paper/2510.01764/) (2025)
- [PuzzleJAX: A Benchmark for Reasoning and Learning](https://mlsanity.com/paper/2508.16821/) (2025)
- [Discovering Minimal Reinforcement Learning Environments](https://mlsanity.com/paper/2406.12589/) (2024)
- [Beyond the Boundaries of Proximal Policy Optimization](https://mlsanity.com/paper/2411.00666/) (2024)
- [Investigating Memory in Model-Free RL with POPGym Arcade](https://mlsanity.com/paper/2503.01450/) (2025)
- [Unsupervised Partner Design Enables Robust Ad-hoc Teamwork](https://mlsanity.com/paper/2508.06336/) (2025)
- [Twice Sequential Monte Carlo for Tree Search](https://mlsanity.com/paper/2511.14220/) (2025)
- [A Reference Architecture of Reinforcement Learning Frameworks](https://mlsanity.com/paper/2603.06413/) (2026)
- [A2Perf: Real-World Autonomous Agents Benchmark](https://mlsanity.com/paper/2503.03056/) (2025)
- [CAMAR: Continuous Actions Multi-Agent Routing](https://mlsanity.com/paper/2508.12845/) (2025)

## Abstract (arXiv)

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.