gymnax makes reinforcement learning environments run inside JAX

The problem: CPU environments throttle every experiment

The README opens with a complaint: are you fed up with slow CPU based RL environment processes, and do you want to leverage massive vectorization for high-throughput RL experiments. That is the problem gymnax attacks. In a typical reinforcement learning setup, the environment steps run in Python on the CPU while the policy runs on the accelerator, and the two sides hand data back and forth. The environment side becomes the bottleneck, and it becomes the bottleneck differently for every environment, which is why throughput numbers are hard to compare across papers. gymnax takes a different approach: bring the environment into JAX itself. The library description in the README is Reinforcement Learning Environments in JAX, and the mechanism is that it brings the power of jit and vmap/pmap to the classic gym API. Because the environment becomes a JAX function, you can compile it, batch it and map it over many parameter settings at once. The README states that gymnax allows explicit functional control of environment settings, such as the random seed or hyperparameters, which enables accelerated and parallelised rollouts for different configurations, and gives meta reinforcement learning as the example case. The audience is researchers running large sweeps and meta learning experiments, not someone who wants to play CartPole. The companion project gymnax-blines provides training scripts and checkpoints for PPO and ES, which tells you the intended level of use.

A functional API where reset and step take a key and a params object

The API is the interesting part of the design, because it changes how you think about an environment.

import jax
import gymnax

rng = jax.random.PRNGKey(0)
rng, key_reset, key_act, key_step = jax.random.split(rng, 4)

# Instantiate the environment & its settings.
env, env_params = gymnax.make("Pendulum-v1")

# Reset the environment.
obs, state = env.reset(key_reset, env_params)

# Sample a random action.
action = env.action_space(env_params).sample(key_act)

# Perform the step transition.
n_obs, n_state, reward, done, _ = env.step(key_step, state, action, env_params)

Three details stand out. The environment carries its own randomness through explicit PRNG keys, so there is no hidden global seed and every rollout can be reproduced. The parameters are a separate object returned by gymnax.make, so hyperparameters travel with the environment rather than living inside it. And step returns an observation, a new state, a reward and a done flag, with a fifth value that the example discards. The parameter object is what makes meta learning practical. Because the parameters are plain JAX data, you can vmap over a batch of parameter settings and run the same environment configuration many times over, which is exactly the workflow the README describes for meta reinforcement learning. The action space is also functional. action_space(env_params) returns a space object with a sample method that takes a key, so even picking a random action is part of the jittable graph.

Which environments are actually ported

Coverage is the practical question, and the README answers it with a table. The classic control family is complete in the usual gym naming: Acrobot-v1, Pendulum-v1, CartPole-v1, MountainCar-v0 and MountainCarContinuous-v0, each linked to its source file under gymnax/environments/classic_control/ and each attributed to Brockman et al. 2016. The MinAtar family includes Asterix-MinAtar, Breakout-MinAtar, Freeway-MinAtar and SpaceInvaders-MinAtar, attributed to Young and Tian 2019, with sources under gymnax/environments/minatar/. The bsuite family includes Catch-bsuite and DeepSea-bsuite, attributed to Osband et al. 2019. Beyond those there is a collection of classic and meta RL tasks, which the README describes as classic/meta RL tasks. The table also records where trained checkpoints exist. For the classic control environments, the gymnax-blines repository has both PPO and ES checkpoints, and the table gives the return of each, for example R: 500 for CartPole-v1, R: -130 for Pendulum-v1 and R: -80 for Acrobot-v1. MinAtar has PPO checkpoints with returns ranging from R: 15 for Asterix to R: 131 for SpaceInvaders, and bsuite entries carry PPO and ES checkpoints where available. The last column gives seconds per one million steps on an A100 with 2k parallel environments, and it is the most honest number in the README. Freeway-MinAtar at 0.87 and Asterix-MinAtar at 0.92 are the slowest entries, while CartPole-v1 at 0.05 is the fastest.

Anakin architecture and distributed evolutionary optimisation

Two research directions are explicitly supported, and they are the reason this library is not just a vectorised gym. The first is the Anakin sub-architecture proposed in the Podracer paper by Hessel et al. 2021. The README explains that executing both the environment and the policy on the accelerator facilitates this architecture, which places the environment step on the same device as the policy so that the inner loop of an RL algorithm stays in compiled code. The second is highly distributed evolutionary optimisation, with evosax by the same author named as the tool. Evolution strategies need many independent evaluations of a candidate policy, which is the shape of problem that vectorised environments suit well. The repository ships example notebooks that show these workflows: examples/01_anakin.ipynb, examples/02_evolution.ipynb, examples/03_meta_a2c.ipynb and examples/getting_started.ipynb. The getting started notebook is also published as a Colab, which is the quickest way to see the API without installing anything. The design intent is visible in the module names. A package directory per family, gymnax/environments/classic_control/, gymnax/environments/minatar/ and gymnax/environments/bsuite/, each environment in its own file, and the whole thing wired together through a single gymnax.make entry point.

Installing gymnax and running the first episode

The package is on PyPI as gymnax, and the dependency list in setup.py is short.

pip install gymnax

The runtime requirements are jax, jaxlib, chex, flax, pyyaml and gym>=0.26, with bsuite and minatar as the test requirements. The gym dependency is there for compatibility with the action space objects rather than because the environments themselves are gym based. The setup.py reads its version from gymnax/_version.py with a regular expression, and raises a RuntimeError if the version string cannot be found. It also builds a download URL for the tagged source archive from that same version string. What you should see when you run the example is a reset that returns an observation and a state, an action sampled from the action space, and a step that returns the next observation, the next state, a reward and a done flag. If the environment or the policy is not behaving, the first thing to check is the JAX version rather than the environment implementation, because the ported code depends on JAX APIs that change between releases. The README also points at the gymnax-blines repository for the training scripts and checkpoints, and at the Colab notebook for a guided introduction.

The limits: ported environments only, and a licence classifier that disagrees

Two limitations deserve to be stated plainly. The first is coverage. gymnax only contains environments that somebody has ported to JAX. An environment you need that is not in the table does not exist in this library, and porting one means rewriting its step function in JAX with explicit randomness, which is the work the library saves you for the environments it already covers. There is also a subtlety in the model class, since JAX arrays are immutable and the environment state has to be threaded through the step function rather than mutated in place. The second is a discrepancy in the packaging metadata. The licence field on the repository says Apache-2.0, and the README badge reads Apache2.0, while the classifiers in setup.py list the OSI approved MIT License. The LICENSE file at the root of the tree and the badge are the more authoritative sources, and the classifier looks like a leftover, but if your compliance tooling reads the package metadata on PyPI rather than the repository, it will report the wrong licence for a package that is installed as a dependency. Neither limitation makes the library a poor choice for its target workload. They do mean it is a research tool with a specific scope rather than a general environment registry.

Apache 2 licensed, a version 1.0 line after a long gap

The version history tells a story about the pace of the work. v0.0.8 came out on 2024-04-21 with a fix for plotting requirements. Then v1.0.0 on 2026-08-16, and v1.0.1 on 2026-10-18. The repository is not archived and the last push was on 2026-10-18, the same day as the 1.0.1 release, so the project is being worked on. The default branch is env-dev rather than main, which suggests active environment work happens on a separate line. Moving to a 1.0 line is the significant signal. In the Python ecosystem that usually means the author considers the public API stable, so pin to 1.0.1 rather than tracking the branch, and read CHANGELOG.md before upgrading across the 0.0.8 to 1.0.0 boundary, because that jump spans a long period and probably contains API changes. The repository layout is conventional for this kind of project: gymnax/ for the package, tests/, examples/ with the notebooks, docs/, a setup.py rather than a pyproject.toml, a CHANGELOG.md and a CONTRIBUTING.md, with a Codecov configuration for coverage reporting and black as the code formatter.

Editorial conclusion

Adopt gymnax when your reinforcement learning research depends on many parallel rollouts, on meta learning over environment parameters, or on running the policy and the environment on the same accelerator. Do not adopt it if you need an environment that has not been ported, since coverage is limited to the classic control, MinAtar, bsuite and meta RL tasks listed in the repository. Before starting, check the setup.py classifier against the Apache-2.0 licence in the repository, and confirm your JAX version matches the environment implementations you intend to use.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. README
4. Releases
5. RobertTLange/gymnax on GitHub