RobertTLange/gymnax

RL Environments in JAX 🌍

What it solves

gymnax addresses the bottleneck of slow, CPU-based reinforcement learning (RL) environments. By reimplementing classic RL environments in JAX, it allows researchers to run environments and policies on the same accelerator (GPU/TPU), eliminating the communication overhead between the CPU and GPU and enabling massive vectorization.

How it works

It brings the classic OpenAI Gym API to the JAX ecosystem. Because it is written in JAX, environments are fully functional and compatible with JAX primitives like jit (just-in-time compilation), vmap (vectorized mapping), and pmap (parallel mapping). This allows for high-throughput rollouts across thousands of environment instances simultaneously and the use of lax.scan to compile entire episode loops into a single accelerator operation.

Who it’s for

It is designed for RL researchers and developers who need high-throughput experiments, those working on meta-RL (where different environment configurations must be parallelized), and users of evolutionary optimization strategies.

Highlights

  • High-Throughput Vectorization: Leverages vmap and pmap for massive parallel rollouts.
  • Accelerator-Native: Executes both the environment and the policy on the GPU/TPU to facilitate architectures like the Anakin sub-architecture.
  • Broad Environment Support: Includes implementations of classic control, bsuite, MinAtar, and various meta-RL tasks.
  • Functional Control: Allows explicit control over random seeds and hyperparameters for accelerated parallelization.
  • Visualization Tools: Includes a Visualizer for generating GIF animations of agent behavior.

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

    A high-performance JAX library implementing over 30 Evolution Strategies for efficient, vectorized neuroevolution on hardware accelerators.

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