Jumanji is a diverse suite of 22 scalable reinforcement learning (RL) environments written in JAX, designed for hardware-accelerated research. It aims to make RL research more accessible, facilitate industrial applications, and enable arbitrarily difficult problems by providing a simple, well-tested API and high-speed environments for faster iteration and large-scale experimentation.

How It Works

Jumanji leverages JAX for its core functionality, enabling automatic vectorization (jax.vmap), parallelization (jax.pmap), and Just-In-Time (JIT) compilation (jax.jit). This approach allows for significant speedups and efficient scaling across hardware accelerators like GPUs and TPUs. The environments are designed to be compatible with popular RL frameworks through dm_env and Gymnasium wrappers, offering a familiar interface inspired by both OpenAI Gym and DeepMind Environment.

Quick Start & Requirements

• Install via pip: pip install -U jumanji or from GitHub: pip install git+https://github.com/instadeepai/jumanji.git
• Requires Python 3.10-3.12. JAX installation depends on hardware; follow the official JAX guide.
• Rendering requires Matplotlib and a GUI backend (e.g., python3-tk on Linux).
• Quickstart and advanced usage guides are available.

Highlighted Details

• Features 22 diverse environments, including logic puzzles (2048, Rubik's Cube), packing problems (BinPack, Knapsack), and routing challenges (CVRP, TSP).
• Environments are versioned for reproducibility, similar to OpenAI Gym.
• Includes example agents (random, A2C) and environment-specific network implementations for training inspiration.
• Accepted at ICLR 2024.

Maintenance & Community

• Developed jointly with the open-source community.
• Contribution guidelines and an issue tracker are available.
• Sister repositories include Qdax, Evojax, Brax, Gymnax, and Pgx.

Licensing & Compatibility

• The repository does not explicitly state a license in the README.

Limitations & Caveats

The example agents provided are for inspiration and not maintained to a production standard. The README does not specify licensing details, which could impact commercial use or closed-source integration.