github.com/instadeepai/jumanji ↗

About instadeepai/jumanji

Jumanji is a diverse suite of 22 reinforcement learning environments implemented in JAX, developed by InstaDeep and the open-source community. The library is designed to enable hardware-accelerated RL research by providing high-speed, scalable environments that support JAX features like automatic vectorization with vmap and pmap, as well as JIT compilation. The environments span multiple categories including logic puzzles like 2048 and Sudoku, packing problems like bin packing and job shop scheduling, routing challenges such as the traveling salesman problem and vehicle routing, and multi-agent scenarios including robot warehouse and search and rescue operations.

The project provides a clean API that combines elements from both OpenAI Gym and DeepMind Environment interfaces, making it familiar to RL practitioners. Each environment returns a TimeStep structure containing step type, reward, discount, observation, and extras fields for logging metrics. Jumanji includes wrappers to integrate with popular frameworks like Stable Baselines3, RLlib, Acme, and Gymnasium, allowing users to easily adapt the environments to their preferred training infrastructure.

Beyond the environments themselves, Jumanji offers example implementations including a vanilla actor-critic agent and environment-specific neural network architectures to demonstrate how to train agents on these tasks. The library maintains strict versioning for reproducibility and was accepted as a research paper at ICLR 2024. The project emphasizes making RL research more accessible and helping close the gap between academic research and industrial applications through difficulty-scalable environments suitable for large-scale experimentation.