mlsanityData to 7 Oct 2026

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

arXiv 2306.09884 ↗First posted 16 Jun 2023cs.LG cs.AIcited by 52

Not on the map. Why?

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.

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

Builds on

18 of 33 references.

References · 18

  1. 2017Attention Is All You Needcited by 59,932
  2. 2017Multi-Agent Actor-Critic for Mixed Cooperative-Competitive Environmentscited by 1,519
  3. 2021Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learningcited by 882
  4. 2018Attention, Learn to Solve Routing Problems!cited by 646
  5. 2018Machine Learning for Combinatorial Optimization: a Methodological Tour d'Horizoncited by 612
  6. 2019Leveraging Procedural Generation to Benchmark Reinforcement Learningcited by 478
  7. 2021Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulationcited by 375
  8. 2020POMO: Policy Optimization with Multiple Optima for Reinforcement Learningcited by 341
  9. 2018Unity: A General Platform for Intelligent Agentscited by 232
  10. 2020Benchmarking Multi-Agent Deep Reinforcement Learning Algorithms in Cooperative Taskscited by 207
  11. 2019OpenSpiel: A Framework for Reinforcement Learning in Gamescited by 180
  12. 2019Behaviour Suite for Reinforcement Learningcited by 129
6 more
  1. 2019An investigation of model-free planningcited by 64
  2. 2022EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Enginecited by 58
  3. 2020OR-Gym: A Reinforcement Learning Library for Operations Research Problemscited by 50
  4. 2021Podracer architectures for scalable Reinforcement Learningcited by 37
  5. 2019ORL: Reinforcement Learning Benchmarks for Online Stochastic Optimization Problemscited by 21
  6. 2019Accelerating Reinforcement Learning through GPU Atari Emulationcited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2023: 5; 2024: 15; 2025: 16; 2026: 16.

52 papers.

Sort:

By area:

  1. 2024Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learningcited by 80
  2. 2024Simplifying Deep Temporal Difference Learningcited by 71
  3. 2023RL4CO: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmarkcited by 70
  4. 2023XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAXcited by 54
  5. 2024GPUDrive: Data-driven, multi-agent driving simulation at 1 million FPScited by 43
  6. 2024POGEMA: A Benchmark Platform for Cooperative Multi-Agent Pathfindingcited by 20
  7. 2023QDax: A Library for Quality-Diversity and Population-based Algorithms with Hardware Accelerationcited by 16
  8. 2024The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Designcited by 11
  9. 2024Sable: a Performant, Efficient and Scalable Sequence Model for MARLcited by 10
  10. 2025FrontierCO: Real-World and Large-Scale Evaluation of Machine Learning Solvers for Combinatorial Optimizationcited by 10
  11. 2023minimax: Efficient Baselines for Autocurricula in JAXcited by 9
  12. 2024Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancycited by 7
38 more
  1. 2024PARCO: Parallel AutoRegressive Models for Multi-Agent Combinatorial Optimizationcited by 7
  2. 2025SocialJax: An Evaluation Suite for Multi-agent Reinforcement Learning in Sequential Social Dilemmascited by 7
  3. 2025ROTATE: Regret-driven Open-ended Training for Ad Hoc Teamworkcited by 6
  4. 2025Trust-Region Twisted Policy Improvementcited by 5
  5. 2025Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAXcited by 5
  6. 2024SPO: Sequential Monte Carlo Policy Optimisationcited by 4
  7. 2025PuzzleJAX: A Benchmark for Reasoning and Learningcited by 4
  8. 2025BuilderBench: The Building Blocks of Intelligent Agentscited by 4
  9. 2024Discovering Minimal Reinforcement Learning Environmentscited by 3
  10. 2024Beyond the Boundaries of Proximal Policy Optimizationcited by 3
  11. 2025Twice Sequential Monte Carlo for Tree Searchcited by 3
  12. 2026A Reference Architecture of Reinforcement Learning Frameworkscited by 3
  13. 2025A2Perf: Real-World Autonomous Agents Benchmarkcited by 2
  14. 2025Unsupervised Partner Design Enables Robust Ad-hoc Teamworkcited by 2
  15. 2025CAMAR: Continuous Actions Multi-Agent Routingcited by 2
  16. 2025Multi-Agent Craftax: Benchmarking Open-Ended Multi-Agent Reinforcement Learning at the Hyperscalecited by 2
  17. 2024Deep Dive into Model-free Reinforcement Learning for Biological and Robotic Systems: Theory and Practicecited by 1
  18. 2024A Batch Sequential Halving Algorithm without Performance Degradationcited by 1
  19. 2025Investigating Memory in Model-Free RL with POPGym Arcadecited by 1
  20. 2026gym-invmgmt: An Open Benchmarking Framework for Inventory Management Methodscited by 1
  21. 2026CRAX: Fast Safe Reinforcement Learning Benchmarkingcited by 1
  22. 2026Position: RL Researchers Need to Distinguish Between Solving Simulators and Using Simulators as a Proxycited by 1
  23. 2023AI planning in the imagination: High-level planning on learned abstract search spacescited by 0
  24. 2024Multi-Agent Environments for Vehicle Routing Problemscited by 0
  25. 2024A Research Agenda for Usability and Generalisation in Reinforcement Learningcited by 0
  26. 2025ABCDWaveNet: Advancing Robust Road Ponding Detection in Fog through Dynamic Frequency-Spatial Synergycited by 0
  27. 2025Oryx: a Scalable Sequence Model for Many-Agent Coordination in Offline MARLcited by 0
  28. 2025Ludax: A GPU-Accelerated Domain Specific Language for Board Gamescited by 0
  29. 2026KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learningcited by 0
  30. 2026JaxARC: A High-Performance JAX-based Environment for Abstraction and Reasoning Researchcited by 0
  31. 2026Particle Monte Carlo Tree Searchcited by 0
  32. 2026Rethinking Ratio-Based Trust Regions for Policy Optimization in Multi-Agent Reinforcement Learningcited by 0
  33. 2026Self-Supervised On-Policy Reinforcement Learning via Contrastive Proximal Policy Optimisationcited by 0
  34. 2026Contract-Based Compositional Shielding for Safe Multi-Agent Reinforcement Learningcited by 0
  35. 2026RENEW: Towards Learning World Models and Repairing Model Exploitation from Preferencescited by 0
  36. 2026DungeonBench: A Benchmark for Rules-Rich Tactical Reasoning in Dungeons & Dragons Combatcited by 0
  37. 2026Out-of-Distribution Generalisation with Sequence Models in Offline Multi-Agent Reinforcement Learningcited by 0
  38. 2026PlayTrain: An Efficient Reinforcement Learning Framework for LLM-Generated Adaptable JavaScript Gamescited by 0

50 most cited of 52 All 52 (JSON)