instadeepai/jumanji

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

What it solves

Jumanji provides a scalable suite of reinforcement learning (RL) environments designed to accelerate RL research and development. It addresses the need for high-speed, hardware-accelerated environments that allow for faster iteration and large-scale experimentation while reducing complexity.

How it works

Written entirely in JAX, the library leverages JAX's capabilities for automatic vectorization, parallelization (jax.vmap, jax.pmap), and JIT-compilation (jax.jit). It combines an interface inspired by OpenAI Gym (using a registry and render methods) and DeepMind Environment (using a TimeStep structure).

Who it’s for

It is designed for RL practitioners and researchers who need high-performance environments for testing agents, particularly those focusing on combinatorial problems, logic games, and routing problems.

Highlights

  • Diverse Environment Suite: Includes 22 environments across categories like Logic (e.g., Sudoku, Rubik's Cube), Packing (e.g., BinPack, Tetris), Routing (e.g., TSP, CVRP), and Swarms (Search and Rescue).
  • Hardware Acceleration: Fully JAX-based, enabling massive parallelization and high-speed execution on GPUs/TPUs.
  • Framework Agnostic: Provides wrappers for popular RL libraries such as Acme, Stable Baselines3, RLlib, and Gymnasium.
  • Scalable Difficulty: Environments are designed so their difficulty can be scaled to be arbitrarily hard.

Related

  • Project
    RobertTLange/gymnax

    A collection of reinforcement learning environments reimplemented in JAX to enable high-throughput, vectorized experiments on GPUs and TPUs.

  • Project
    google/brax

    A fast, fully differentiable physics engine written in JAX for large-scale rigid body simulation and reinforcement learning training.

  • Project
    google/tunix

    A JAX-based post-training library for LLMs that provides scalable support for supervised fine-tuning and reinforcement learning, specifically optimized for TPU performance.

  • Project
    instadeepai/Mava

    A distributed multi-agent reinforcement learning research library built in JAX that provides high-performance implementations of state-of-the-art MARL algorithms.

  • Project
    JuliaReinforcementLearning/ReinforcementLearning.jl

    A comprehensive Julia package for reinforcement learning research that provides extensible components for implementing, benchmarking, and reproducing RL algorithms.