Reinforcement learning environments from Snake to bin packing, in JAX

What it does Jumanji is a collection of 22 reinforcement learning environments implemented in JAX. It covers classic games such as Snake, Pac-Man, and Tetris alongside combinatorial challenges including vehicle routing, job-shop scheduling, and bin packing. The project provides a standardized API, wrappers for Gymnasium and DeepMind-Env compatibility, and example training agents to support large-scale RL experimentation.

The interesting bit The suite treats arcade games and industrial NP-hard problems as first-class citizens in the same accelerated framework. By keeping every environment in JAX, Jumanji aims to let researchers iterate faster and run large-scale experiments without leaving the hardware-accelerated context.

Key highlights

• 22 environments spanning logic puzzles, packing problems, and routing tasks, with scalable difficulty.
• Built-in wrappers for dm_env and gym to plug into Stable Baselines3, RLlib, Acme, and Gymnasium.
• Originated at InstaDeep and accepted at ICLR 2024, now developed jointly with the open-source community.
• Includes example agents and documentation aimed at making industrial RL research more accessible.

Verdict Researchers who need batched, hardware-accelerated environment rollouts—particularly for combinatorial optimization—should look here. If you need traditional CPU-bound environment loops and have no use for JAX acceleration, look elsewhere.