gymnax
A JAX-based library providing accelerated reinforcement learning environments with full compatibility to the classic gym API.
What is gymnax?
Gymnax is a JAX-based library that provides accelerated reinforcement learning environments compatible with the OpenAI Gym API. It solves the problem of slow CPU-based environment simulation by leveraging JAX's just-in-time compilation and vectorization capabilities, enabling massively parallel rollouts for faster RL experimentation.
Reinforcement learning researchers and practitioners who need high-throughput environment simulation for tasks like meta-learning, evolutionary optimization, or large-scale policy evaluation, particularly those already working within the JAX ecosystem.
Developers choose Gymnax for its unique combination of full Gym API compatibility with JAX-native acceleration, allowing them to easily port existing workflows while gaining orders-of-magnitude speed improvements through batched and compiled environment execution.
Overview
RL Environments in JAX 🌍
Use Cases
Best For
- Running high-throughput RL experiments with massive environment parallelization
- Meta-reinforcement learning research requiring explicit control over environment parameters
- Evolutionary strategy optimization where population evaluation benefits from vectorization
- Benchmarking RL algorithms with accelerated classic control and MinAtar environments
- Educational purposes for learning JAX-based RL with a familiar gym interface
- Implementing the Anakin architecture for fully accelerated agent-environment loops
Not Ideal For
- Projects relying on environments not implemented in JAX, such as complex 3D simulations or proprietary robotics suites
- Teams with existing RL workflows deeply integrated with TensorFlow or PyTorch without JAX interoperability needs
- Applications requiring real-time environment interaction without upfront compilation overhead, like interactive demos or rapid prototyping with diverse, unsupported environments
Pros & Cons
Pros
Leverages JAX's jit, vmap, and pmap for compiled, batched rollouts, enabling massive parallelization—benchmarks show 1M steps in under 0.1 seconds on an A100 for classic control tasks.
Maintains a drop-in compatible interface with reset and step functions, easing adoption for users experienced with OpenAI Gym without sacrificing JAX benefits.
Provides fine-grained control over random seeds and environment parameters via env_params, facilitating reproducible research and meta-RL experiments as highlighted in the examples.
Includes a Visualizer class that generates GIF animations from state sequences, covering classic_control and MinAtar environments for easy result sharing.
Cons
Focuses on reimplementations of classic control, bsuite, and MinAtar—lacks support for modern, complex environments like full Atari suites or physics simulators found in Brax, limiting scope for broader RL research.
Requires familiarity with JAX's functional programming model and accelerator setup; the README assumes users are comfortable with jit and vmap, which can be a barrier for those new to the ecosystem.
Some components, like the RolloutWrapper for batch evaluation, are marked as 'work-in-progress' in the README, indicating instability or incomplete functionality for production use.
Frequently Asked Questions
Related Projects
ALXGoogle Research
OryxProbabilistic reasoning and statistical analysis in TensorFlow
BRAXMassively parallel rigidbody physics simulation on accelerator hardware.
MctxMonte Carlo tree search in JAX
Found a gem we're missing?
Open-Awesome is built by the community, for the community. Submit a project, suggest an awesome list, or help improve the catalog on GitHub.