Research map / Reinforcement learning / Visual and representation RL
Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces Jumanji, a JAX-based suite of 22 fast, flexible, scalable RL environments (mostly NP-hard combinatorial problems) with A2C baselines, scaling and generalization experiments.
Problem
Existing RL environment suites are not closely tied to practical industrial problems and lack hardware acceleration, flexibility in initial state distributions, and scalable difficulty.
Approach
A diverse suite of 22 JAX-implemented, stateless RL environments organized into routing, packing and logic categories, with composable generators for initial-state distributions and tunable problem complexity, trained end-to-end on accelerators with the Anakin framework.
Key results
- A2C improves over masked random policy on all 22 environments but optimality gaps remain in most.
- Training on four generators yields mean tour length 5.815 vs 6.090 for a single uniform generator on 1,020 VLSI TSP instances.
- On Connector-v2, 92% of optimal performance is reached in about 1.5h on TPU-128 versus 11.5h on TPU-8.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train with four generators instead of only the uniform generator (ablation) | A2C on TSP with 50 cities · VLSI TSP test set (1,020 instances) | mean tour length: 6.090 ;(\pm 0.029) -> 5.815 ;(\pm 0.025) | Helped |
| Increase device parallelism from TPU-8 to TPU-128 | A2C on Connector-v2 · Connector-v2 | time to reach 92% of optimal performance: 11.5h -> 1.5h | Helped |
| Increase problem complexity (ablation) | A2C · Connector, BinPack, and RubiksCube | agent performance | Hurt |
| Use the hardest tested RubiksCube problem setting (ablation) | A2C · RubiksCube | learning performance | Hurt |
Surprising findings
- In RubiksCube, the hardest setting leads the A2C agent to completely fail to learn.
Limitations
- Only single-agent implementations; multi-agent training algorithm not provided
- All current environments have discrete actions
- Baseline agent leaves substantial optimality gaps on most environments
Experimental setup
Datasets: VLSI TSP Benchmark Dataset
Benchmarks: Jumanji 22 environments VLSI TSP instances (1020 test)
Metrics: episode return effective environment steps per second mean tour length proportion of wires connected
Baselines: masked random policy optimal return (TSP, CVRP, Knapsack from prior work) Gymnax EnvPool CuLE Isaac Gym
Builds on: A2C Anakin JAX Gymnax dm_env
Closest papers by meaning
- JaxMARL: Multi-Agent RL Environments and Algorithms in JAX (NeurIPS 2024 · similarity 0.89)
- Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning (ICML 2024 · similarity 0.87)
- Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks (ICLR 2025 · similarity 0.87)
- XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX (NeurIPS 2024 · similarity 0.86)
- SRL: Scaling Distributed Reinforcement Learning to Over Ten Thousand Cores (ICLR 2024 · similarity 0.86)
- NAVIX: Scaling MiniGrid Environments with JAX (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.