Jumanji turns 22 reinforcement learning problems into JAX arrays

Twenty-two environments that live in JAX rather than beside it

Jumanji is a set of reinforcement learning environments written in JAX, and the writing is the point rather than an implementation detail. In a conventional environment, the reset and step functions run as Python and NumPy code, and the agent runs on a GPU while the environment runs on a CPU. That split is the usual bottleneck in RL research, since each environment step is a round trip. Here the environment itself is written in JAX, so it can be jitted and vmapped like any other array program, and the README frames the goal as high speed environments that allow faster iteration and large scale experimentation while reducing complexity. The suite covers 22 environments, ranging from simple games to NP-hard combinatorial problems, with names such as BinPack, CVRP, GraphColoring, JobShop, Knapsack, RobotWarehouse, RubiksCube, Sokoban, Sudoku, Tetris, TSP, Level-Based Foraging and Search and Rescue among them. The research behind it was published at ICLR 2024.

Difficulty is scaled by parameter, not by writing a new task

One of the four stated goals is to provide environments whose difficulty can be scaled to be arbitrarily hard. That is a design decision with real consequences for experiment design. Rather than training on a fixed instance size, you change a parameter and the same environment becomes a harder problem, which is how you study scaling behaviour, sample efficiency or curriculum effects without maintaining several separate environment implementations. The two companion goals are a simple, well-tested API for JAX-based environments and making RL research more accessible, and the fourth is to close the gap between research and industrial applications by covering problems of the kind that show up in industry. The trade-off is that a scaled combinatorial problem is not the same thing as a realistic simulation. These environments are deliberately abstract, and a result on a scaled Knapsack or JobShop instance transfers to a real logistics problem only through the modelling choices you make yourself.

dm_env and gym wrappers are the bridge out of JAX

The suite includes a set of wrappers whose stated purpose is to connect to existing RL frameworks and libraries, and the README names Acme, Stable Baselines3, RLlib, Gymnasium and DeepMind's dm_env as the targets. Two of them are singled out: a `dm_env` wrapper and a `gym` wrapper. This is the pragmatic part of the design. If you already have an agent written against Gymnasium or dm_env, the wrapper lets you point it at a Jumanji environment and get the same interface, which means you can benchmark against your existing pipeline without rewriting the agent. The cost of that convenience is that a non-JAX agent still has to call out to the environment, and the environment is fast precisely when it can be batched inside JAX. Anyone using the wrapper path should check what the wrapper actually does with batching, because that is where the speed advantage either survives or disappears.

Installing Jumanji is a Python 3.10 plus JAX dependency set

The package is built with hatchling and published to PyPI, and the project metadata pins the runtime floor. The interesting part for anyone integrating it into an existing environment is the dependency list in pyproject.toml.

requires-python = ">=3.10"

dependencies = [
  "chex>=0.1.3",
  "dm-env>=1.5",
  "esquilax>=2.0.0",
  "gymnasium>=1.0",
  "huggingface-hub",
  "jax>=0.4.27", # jnp.clip(x, min=, max=) signature (Array API); a_min/a_max removed in 0.8.0
  "matplotlib>=3.8.0",
  "numpy>=1.19.5",
  "Pillow>=9.0.0",
  "requests>=2.32",
  "typing-extensions>=4.0.0",
]

JAX is a hard requirement rather than an optional extra, and the version floor is high. The comment attached to the jax pin is a small warning sign for teams that track their dependency versions: it records that jnp.clip changed its signature for Array API compatibility and that the a_min and a_max arguments were removed in 0.8.0, so the pin exists to stay clear of that break. Since JAX releases move quickly, expect to align the rest of your stack with it rather than around it. The repository also carries a uv.lock, so a uv based project can pin the resolved set rather than re-resolving it.

Training extras are grouped under a train extra

The learning-oriented dependencies are kept out of the core install and grouped under an optional extra named train, which is a sensible split if you only want to run environments.

[project.optional-dependencies]
train = [
  "distrax>=0.1.5",
  "dm-haiku",
  "hydra-core==1.3.4",
  "neptune-client==0.16.15",
  "optax>=0.1.4",
  "rlax>=0.1.4",
  "tensorboardX==2.5.1",
  "tqdm>=4.64.1",
]

The composition says a lot about the intended workflow. distrax and optax are the distribution and optimizer layer, dm-haiku and rlax are the Haiku network library and the RL primitives, hydra-core==1.3.4 is pinned exactly because Hydra is the config system the examples are written against, and neptune-client is there for experiment tracking. Two of the pins are exact versions, which is a hint that the training examples are tied to specific Hydra and Neptune behaviour. The examples directory follows the same split, with a training notebook, a load_checkpoints notebook and a script to visualize a random agent, so the fastest way to see a task in action is that last script rather than a notebook.

Where Jumanji is the wrong choice

There are clear situations where this suite is the wrong tool, and they are not edge cases. First, if your agent does not want to live in JAX. The wrappers make it possible, but they exist as a concession, and a research group that wants a native JAX agent, a JAX training loop and a JAX policy network gets the benefit while everyone else takes a detour. Second, if you need a domain simulation. The catalogue runs from games to combinatorial problems, and none of the 22 environments model a robot, a chemical process or a user interface, so anyone whose problem is a physical or continuous-control domain has to write the environment themselves, which means writing it in JAX and accepting the same setup cost. Third, scale of adoption: the suite is broad but not deep, and the README does not describe a registry with versioned environment identifiers, even though the environments table in the documentation has a registered version column, so stability of a specific task across releases is something a reader has to check per environment rather than assume.

Apache 2.0 licensing with a recent release line

The package declares the Apache Software License in its classifiers and points at the LICENSE file in the repository root, which is the standard permissive licence with an explicit patent grant. The project metadata also marks the distribution as Production/Stable under the development status classifier, and the intended audience is listed as Science/Research and Developers. On the release side, the tagged versions are v1.1.2 from 2026-07-31, v1.1.1 from 2025-06-18 and v1.1.0 from 2024-11-22, and the last push was on 2026-10-16, so the code is moving while the release tags advance less often. A paper tied to a specific environment version may therefore need a check against the current main branch, since the difference between a tag and the working tree is where behavioural changes are most likely to hide. Documentation lives on a separate site built with mkdocs, with the mkdocs.yml and docs/ directories in the repository.

Editorial conclusion

Jumanji suits reinforcement learning researchers who want environments that run on accelerator hardware and problems whose difficulty can be turned up without writing a new task, and it suits teams already comfortable with JAX. It does not suit a team that wants a mature third party environment catalogue, a non-JAX toolchain, or a domain model of its own business, and the wrappers exist precisely because non-JAX pipelines are a compromise rather than the intended path. Before adopting it, check that the specific environment you need is in the registered suite, run the random agent visualisation script in examples to see the problem and its observation shapes, and read the environment API docs to confirm the difficulty parameters behave the way your experiment assumes. The last push was on 2026-10-16, and the current release is v1.1.2 from 2026-07-31.

Frequently asked questions

Official sources

1. instadeepai/jumanji on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases