Using Gymnasium environments

Jaxnasium's algorithms are written for end-to-end RL training in pure JAX (PureJaxRL). However, this requires environment step and reset functions to be written in JAX and be fully JIT-compatible. While this can greatly improve training times, many environments are not (yet) ported over to JAX.

Jaxnasium provides a GymnasiumWrapper to use gymnasium environments within its training pipeline. It calls back into Python on every step via jax.experimental.io_callback, so the rest of the training loop stays compiled. Each step is then a Python roundtrip, so throughput will naturally be lower compared to native JAX environments, but this wrapper opens the door for many more environments to be used with the same algorithms.

import gymnasium as gym
import jax
from gymnasium.vector import AutoresetMode

from jaxnasium.wrappers import GymnasiumWrapper
from jaxnasium.algorithms import PPO

venv = gym.make_vec(
    "CartPole-v1",
    num_envs=2,
    vectorization_mode="sync",
    vector_kwargs={"autoreset_mode": AutoresetMode.SAME_STEP},
)
env = GymnasiumWrapper(venv)

ppo = PPO(num_envs=2, total_timesteps=100_000)
agent, metrics = ppo.train(jax.random.key(0), env)

eval_env = GymnasiumWrapper(gym.make("CartPole-v1"))
returns = agent.evaluate(jax.random.key(1), eval_env, num_eval_episodes=10)

Training

During training, Jaxnasium algorithms expect a leading num_envs batch axis, typically obtained from VecEnvWrapper, which jax.vmaps the environment step and reset. Gymnasium environments cannot rely on this wrapper, so have to do vectorization on the Gymnasium side, e.g. via gymnasium.make_vec:

venv = gym.make_vec(
    "CartPole-v1",
    num_envs=2,
    vectorization_mode="sync",
    vector_kwargs={"autoreset_mode": AutoresetMode.SAME_STEP},
)
env = GymnasiumWrapper(venv)

ppo = PPO(num_envs=2, total_timesteps=100_000)
agent, metrics = ppo.train(jax.random.key(0), env)

Notes

• autoreset_mode=AutoresetMode.SAME_STEP is preferred. as it matches the Jaxnasium API. With Gymnasium's default mode (NEXT_STEP) or DISABLED, the GymnasiumWrapper resets finished environments itself within the same step, which costs an extra reset call on every step in which an episode ends.
• An explicit vectorization_mode. Some environments (e.g. CartPole-v1) ship their own vectorized implementation, which make_vec picks by default. Not all of these implementations allow individual environments to be reset manually, so the wrapper rejects it; set to "sync" (or "async") instead.
• The algorithm's num_envs must equal the vector environment's. The batch axis now comes from the environment, so training raises an error if they differ, or if the environment is not vectorized at all.

Evaluation

algorithm.evaluate() runs one episode at a time (non-vectorized). A VecEnvWrapper is typically removed automatically, but this is not the case for Gymnasium environments. As such, in contrast to training, pass a plain Gymnasium environment to the GymnasiumWrapper instead.

eval_env = GymnasiumWrapper(gym.make("CartPole-v1"))
returns = agent.evaluate(jax.random.key(1), eval_env, num_eval_episodes=10)

jaxnasium.wrappers.GymnasiumWrapper

Runs CPU-based Gymnasium environments inside a JAX training loop, by calling back into Python on every step. Every step round-trips through Python, so throughput is likely slower compared to native JAX environments; this wrapper aims for compatibility, not speed.

Envpool environments are also possible through this wrapper (more information).

Vectorization (required for training jaxnasium algorithms) Jaxnasium algorithms expect a leading batch axis and as such expect a vectorized environment. For Gymnasium environments, this needs to happen on the Gymnasium side instead of via a jaxnasium VecEnvWrapper. As such, wrap the gymnasium environment with a gymnasium.vector.VectorEnv. Note that, in contrast to training, algorithm.evaluate expects a regular non-vectorized environment.

Example:

import gymnasium
from gymnasium.vector import AsyncVectorEnv, AutoresetMode, SyncVectorEnv

venv = SyncVectorEnv(
    [lambda: gymnasium.make("CartPole-v1")] * 8,
    autoreset_mode=AutoresetMode.SAME_STEP,
)
train_env = GymnasiumWrapper(venv)  # batched, for training
eval_env = GymnasiumWrapper(
    gymnasium.make("CartPole-v1")
)  # unbatched, for evaluation

Arguments:

• env: a Gymnasium environment, or a gymnasium.vector.VectorEnv of them, preferably built with autoreset_mode=AutoresetMode.SAME_STEP.

Warning

Gymnasium info dictionaries are dropped: their keys vary between steps, while a callback has to return a fixed PyTree of fixed shapes. Only _TERMINAL_OBSERVATION is filled in by this wrapper.

step(key: PRNGKeyArray, state: Int[Array, ' num_envs'], action: PyTree[Real[Array, ...]]) -> tuple[TimeStep, Int[Array, ' num_envs']]

reset(key: PRNGKeyArray) -> tuple[Observation, Int[Array, ' num_envs']]

EnvPool

Since EnvPool can create environments with the Gymnasium standard, the GymnasiumWrapper can trivially be used to speed up training times on EnvPool supported environments:

import envpool

venv = envpool.make_gymnasium("CartPole-v1", num_envs=2, seed=0)
env = GymnasiumWrapper(venv)

EnvPool resets finished environments on the next step, so the wrapper resets them itself, as described above. Seed the environments when creating them: EnvPool does not take per-reset seeds into account, so runs are not seeded by the key passed to reset.

EnvPool's XLA interface

The EnvPoolXLAWrapper steps EnvPool's C++ environments directly from the compiled program, via EnvPool's XLA interface. Python is then only called on steps where an episode ended somewhere in the batch, to reset those environments. This is typically faster than the GymnasiumWrapper, especially when episodes rarely end.

import envpool
from jaxnasium.wrappers import EnvPoolXLAWrapper

venv = envpool.make_gymnasium("LunarLander-v2", num_envs=16, seed=0)
env = EnvPoolXLAWrapper(venv)

ppo = PPO(num_envs=16, total_timesteps=500_000)
agent, metrics = ppo.train(jax.random.key(0), env)

jaxnasium.wrappers.EnvPoolXLAWrapper

Runs an EnvPool environment through EnvPool's XLA interface. Compared to GymnasiumWrapper(envpool_env), regular steps run EnvPool's C++ environments directly from the compiled program, without calling back into Python. Python is only called on steps where an episode ended somewhere in the batch, to reset those environments.

import envpool

venv = envpool.make_gymnasium("LunarLander-v2", num_envs=8, seed=0)
env = EnvPoolXLAWrapper(venv)

Arguments:

• env: a bare EnvPool environment from envpool.make_gymnasium(..., num_envs=n). Gymnasium wrappers around it would be bypassed by the XLA step, and are therefore rejected.
• autoreset_mode: "same_step" (default) or "next_step", see above.

Autoreset behaviour

EnvPool itself resets a finished environment on the next step, which ignores its action and returns the first observation with a reward of 0. Jaxnasium by default expect same step reset behaviour, so we default to that here; however, since we have to do that manually, we step into Python and essentially sacrafice performance in exchange for keeping the trainsitions clean of observations on which we shouldn't train.

Warning

• Float64 observations or actions (e.g. EnvPool's MuJoCo) require JAX's 64-bit mode.
• As is the case for GymnasiumWrapper: training only (evaluate with a non-vectorized environment), and seeds are set when creating the EnvPool environment.

state_type property

step_env(key: PRNGKeyArray, state: Any, action: PyTree[Real[Array, ...]]) -> tuple[TimeStep, Any]

reset_env(key: PRNGKeyArray) -> tuple[Observation, Any]

action_space: Space | PyTree[Space] property

observation_space: Space | PyTree[Space] property

auto_reset(key: PRNGKeyArray, timestep_step: TimeStep, state_step: TEnvState) -> tuple[TimeStep, TEnvState]

Auto-resets the environment when the episode is terminated or truncated.

Given a step timestep and state, this function will auto-reset the environment and return the new timestep and state when the episode is terminated or truncated. Inserts the original observation in info to bootstrap correctly on truncated episodes.

Arguments:

• key: JAX PRNG key.
• timestep_step: The timestep returned by the step_env method.
• state_step: The state returned by the step_env method.

Returns: A tuple of the new timestep and state with the state and observation reset to a new initial state and observation when the episode is terminated or truncated. The original observation is inserted in info to bootstrap correctly on truncated episodes.

sample_action(key: PRNGKeyArray) -> PyTree[Real[Array, ...]]

sample_observation(key: PRNGKeyArray) -> Observation

multi_agent: bool property

agent_structure: PyTreeDef property

Returns the structure of the agent space. In single-agent environments this is simply a PyTreeDef(). However, for multi-agent environments this is a PyTreeDef((, x num_agents)).

_env: Any = env instance-attribute

__check_init__()

__getattr__(name)

_callback_env: _CallbackEnv = _CallbackEnv(env, num_envs, mode) instance-attribute

_internal_num_envs: int | None property

Size of the batch axis. This is None for a regular (unbatched) environment.

is_vectorized: bool property

Whether the wrapped environment is a gymnasium.vector.VectorEnv, in that case, its observations, rewards and flags carry a leading num_envs axis.

_require_env_keys(key: PRNGKeyArray, name: str) -> None

Limitations

Warning

• Info dictionaries are dropped, since a callback must return a fixed structure of fixed shapes. Only _TERMINAL_OBSERVATION is provided.
• Only flat spaces (Box, Discrete and MultiDiscrete) are supported.
• Training runs cannot be vmapped. Running several seeds or hyperparameters at once with jax.vmap would make every copy step the same Python environments, so it raises an error.
• Only a small subset of environments that adhere to the Gymnasium API have been tested.