RobertTLange/gymnax

RL Environments in JAX 🌍

What it solves

Gymnax addresses the bottleneck of slow, CPU-based reinforcement learning (RL) environment processes. It allows researchers to run environment transitions and policies on accelerators (GPUs/TPUs) to achieve high-throughput experiments and massive vectorization.

How it works

It reimplements classic RL environments (such as those from OpenAI Gym) using JAX. By doing so, it brings JAX primitives like jit (just-in-time compilation), vmap (vectorized mapping), and pmap (parallel mapping) to the environment API. This allows for the execution of both the environment and the policy on the same accelerator, enabling fast parallelized rollouts and functional control over environment settings like random seeds and hyperparameters.

Who it’s for

It is designed for RL researchers and developers who need high-performance environment simulation, specifically those working on meta-RL, evolutionary optimization, or distributed architectures like the Anakin sub-architecture.

Highlights

  • Accelerated Environments: Includes implementations of classic control, bsuite, MinAtar, and various meta-RL tasks.
  • JAX Integration: Full support for jit, vmap, and pmap for massive vectorization across random keys or environment parameters.
  • Differentiable Transitions: Opt-in support for JAX gradients in continuous-control environments (e.g., Pendulum, MountainCarContinuous).
  • Fast Rollouts: Supports lax.scan for compiling entire episode loops into a single accelerator operation.
  • Auto-reset: Implements an automatic reset behavior after terminal or truncated steps.

Related

  • Project
    instadeepai/jumanji

    A diverse suite of scalable, hardware-accelerated reinforcement learning environments written in JAX for high-speed RL research and experimentation.

  • Project
    bold-lab-ai/JaxMARL

    A JAX-native library for Multi-Agent Reinforcement Learning that provides GPU-accelerated environments and baseline algorithms for efficient benchmarking.

  • Project
    RobertTLange/evosax

    evosax is a JAX‑based library that implements 30+ Evolution Strategies (e.g., CMA‑ES, OpenAI‑ES, ARS, SimpleGA) with a vectorised `ask‑eval‑tell` API, enabling fast, hardware‑accelerated black‑box optimisation for neuro‑evolution, RL, and hyper‑parameter search.

  • Project
    google/brax

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

  • 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.