gymnax

library0.0.9pypypi✓ verified 83d ago

gymnax provides JAX-ready implementations of classic reinforcement learning environments (e.g., CartPole, Acrobot, MinAtar) with fully functional and batched vectorization. Current version 0.0.9, with periodic releases.

ai-ml

pip install gymnax

github ↗homepage ↗

INSTALL

IMPORT

SIG · GYMNAX

gymnax

ai-mlpythonv0.0.9

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

gymnax

✓ import gymnax

✗ from gymnax import make

Old API used gymnax.environments; now import gymnax and use gymnax.make

EnvParams

✓ from gymnax.environments import EnvParams

Create a gymnax environment and reset it.

import jax import gymnax rng = jax.random.PRNGKey(0) rng, key_reset = jax.random.split(rng) env, env_params = gymnax.make('CartPole-v1') obs, state = env.reset(key_reset, env_params) print(obs)

Debug

Known issues

breakingEnvironment names changed between 0.0.4 and 0.0.5. Use 'CartPole-v1' (not 'CartPole-v0').

fix

Update environment name strings to Gym-style versions (e.g., 'CartPole-v1', 'Acrobot-v1').

affects: <0.0.5

deprecatedDirect imports from `gymnax.environments` (e.g., `CartPole`) are deprecated in favor of `gymnax.make`. New environments should be accessed through the registry.

fix

Use `gymnax.make('env_name')` instead of importing environment classes directly.

affects: >=0.0.8

gotchagymnax expects JAX-compatible inputs; using plain NumPy arrays may cause silent errors or performance penalties. Always convert NumPy arrays to JAX arrays via `jax.numpy.array`.

fix

Wrap NumPy arrays with `jnp.array()` before passing to gymnax functions.

affects: all

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'gymnax'

gymnax not installed or installed in a different Python environment.

fix

Install via `pip install gymnax` and ensure you're using the correct Python/conda environment (e.g., `pip list | grep gymnax`).

ValueError: Invalid environment name: 'CartPole-v0'

Environment version string outdated; gymnax uses Gym-style versioning (e.g., 'CartPole-v1').

fix

Use environment names from the gymnax registry: e.g., 'CartPole-v1', 'Acrobot-v1', 'MinAtar/Asterix-v1'.

AttributeError: module 'gymnax' has no attribute 'make'

Old version of gymnax (<0.0.5) where `gymnax.make` did not exist.

fix

Upgrade gymnax: `pip install --upgrade gymnax`. Alternatively, use old API: `from gymnax.environments import CartPole`.

Upgrade

Version history

0.0.9latest on PyPI · released May 23, 2025

Audit

Dependencies

jaxrequiredCore dependency for autograd and JIT compilation

Agent activity

19 hits · last 30 days

node

Amazon

OpenAI (training)

Resources

githubhttps://github.com/RobertTLange/gymnax ↗

homepagehttps://github.com/RobertTLange/gymnax ↗