Wrappers

Wrappers modify the behaviour of an environment and can be stacked freely.

import jaxnasium as jym

env = jym.make("Breakout-MinAtar")
env = jym.FlattenObservationWrapper(env)
env = jym.LogWrapper(env)

Base class

jaxnasium.Wrapper

Base class for all wrappers.

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

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

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

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

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

action_space: Space | PyTree[Space] property

observation_space: Space | PyTree[Space] property

multi_agent: bool property

__getattr__(name)

state_type property

Environment wrappers

jaxnasium.VecEnvWrapper

Wrapper to vectorize environments. Simply calls jax.vmap on the reset and step methods of the environment. The number of environmnents is determined by the leading axis of the inputs to the reset and step methods, as if you would call jax.vmap directly.

We use a wrapper instead of jax.vmap in each algorithm directly to control where the vectorization happens. This allows other wrappers to act on the vectorized environment, e.g. NormalizeVecObsWrapper and NormalizeVecRewardWrapper.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__getattr__(name)

__check_init__()

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

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

jaxnasium.LogWrapper

Log the episode returns and lengths. Modeled after the LogWrapper in PureJaxRL.

This wrapper inserts episode returns and lengths into the info dictionary of the TimeStep object. The returned_episode_returns and returned_episode_lengths are the returns and lengths of the last completed episode.

After collecting a trajectory of n steps and collecting all the info dicts, the episode returns may be collected via:

return_values = jax.tree.map(
    lambda x: x[data["returned_episode"]], data["returned_episode_returns"]
)

Arguments:

• _env: Environment to wrap.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

reset(key: PRNGKeyArray) -> tuple[Observation, LogEnvState]

step(key: PRNGKeyArray, state: LogEnvState, action: PyTree[int | float | Array]) -> tuple[TimeStep, LogEnvState]

jaxnasium.NormalizeVecObsWrapper

Normalize the observations of the environment via running mean and variance. This wrapper acts on vectorized environments and in turn should be wrapped within a VecEnvWrapper.

Arguments:

• _env: Environment to wrap.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__getattr__(name)

__check_init__()

update_state_and_get_obs(obs, state: NormalizeVecObsState)

reset(key: PRNGKeyArray) -> tuple[Observation, NormalizeVecObsState]

step(key: PRNGKeyArray, state: NormalizeVecObsState, action: PyTree[int | float | Array]) -> tuple[TimeStep, NormalizeVecObsState]

jaxnasium.NormalizeVecRewardWrapper

Normalize the rewards of the environment via running mean and variance. This wrapper acts on vectorized environments and in turn should be wrapped within a VecEnvWrapper.

Arguments:

• _env: Environment to wrap.
• gamma: Discount factor for the rewards.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__getattr__(name)

gamma: float = 0.99 class-attribute instance-attribute

__check_init__()

reset(key: PRNGKeyArray) -> tuple[Observation, NormalizeVecRewState]

step(key: PRNGKeyArray, state: NormalizeVecRewState, action: PyTree[int | float | Array]) -> tuple[TimeStep, NormalizeVecRewState]

jaxnasium.TransformObservationWrapper

Transform observations with a function.

func (and observation_space_fn) is applied to each observation leaf. terminal observations (in info) are also transformed, and action masks are untouched.

Arguments:

• env: Environment to wrap.
• func: Function applied to every observation leaf, or a pytree of functions with the same structure as the observation (one function per leaf).
• observation_space_fn: Function giving the new space of every observation leaf, or a pytree of them (same structure). If None, the observation space is kept, so pass it whenever func would change the observation space shape or dtype.

state_type property

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

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

action_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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

func: Callable | PyTree[Callable] instance-attribute

observation_space_fn: Callable | PyTree[Callable] | None = None class-attribute instance-attribute

reset(key: PRNGKeyArray) -> tuple[Observation, TEnvState]

step(key: PRNGKeyArray, state: TEnvState, action: PyTree[int | float | Array]) -> tuple[TimeStep, TEnvState]

observation(observation: Observation) -> Observation

observation_space: Space | PyTree[Space] property

jaxnasium.FlattenObservationWrapper

Flatten the observations of the environment. Flattens each observation leaf to a vector, keeping the PyTree structure of the observation.

Arguments:

• env: Environment to wrap.

state_type property

step(key: PRNGKeyArray, state: TEnvState, action: PyTree[int | float | Array]) -> tuple[TimeStep, TEnvState]

reset(key: PRNGKeyArray) -> tuple[Observation, TEnvState]

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

func: Callable | PyTree[Callable] instance-attribute

observation_space_fn: Callable | PyTree[Callable] | None = None class-attribute instance-attribute

observation(observation: Observation) -> Observation

__init__(env: Environment)

jaxnasium.FrameStackWrapper

Stack the last num_frames observations along a new leading axis, like gymnasium.wrappers.FrameStackObservation.

Each observation leaf of shape shape becomes (num_frames, *shape), ordered from oldest to newest, so obs[-1] is always the current frame. At the start of an episode the stack is filled by repeating the first frame. Action masks are not stacked, as they only apply to the current step.

The wrapped environment resets itself, so episode boundaries are detected from terminated/truncated and the stack is refilled there. The terminal observation in info is stacked as well, so bootstrapping on truncation sees the full stack.

Apply this before vectorization (VecEnvWrapper); the algorithms add that themselves.

Arguments:

• _env: Environment to wrap.
• num_frames: Number of frames to stack. Default is 4.

state_type property

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

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

action_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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__getattr__(name)

num_frames: int = eqx.field(static=True, default=4) class-attribute instance-attribute

__check_init__()

reset(key: PRNGKeyArray) -> tuple[Observation, FrameStackState]

step(key: PRNGKeyArray, state: FrameStackState, action: PyTree[int | float | Array]) -> tuple[TimeStep, FrameStackState]

observation_space: Space | PyTree[Space] property

jaxnasium.TransformRewardWrapper

Transform the rewards of the environment using a given function.

Arguments:

• _env: Environment to wrap.
• transform_fn: Function to transform the rewards.

state_type property

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

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

transform_fn: Callable instance-attribute

step(key: PRNGKeyArray, state: TEnvState, action: PyTree[int | float | Array]) -> tuple[TimeStep, TEnvState]

jaxnasium.ScaleRewardWrapper

Scale the rewards of the environment by a given factor.

Arguments:

• _env: Environment to wrap.
• scale: Factor to scale the rewards by.

state_type property

step(key: PRNGKeyArray, state: TEnvState, action: PyTree[int | float | Array]) -> tuple[TimeStep, TEnvState]

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

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

scale: float = scale instance-attribute

__init__(env: Environment, scale: float = 1.0)

transform_fn = lambda r: r * scale instance-attribute

jaxnasium.DiscreteActionWrapper

Wrapper to convert continuous actions to discrete actions.

Arguments:

• _env: Environment to wrap.
• num_actions: Number of discrete actions to convert to.

state_type property

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

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

num_actions: int instance-attribute

step(key: PRNGKeyArray, state: TEnvState, action: PyTree[int | Int[Array, ' num_actions']]) -> tuple[TimeStep, TEnvState]

action_space: PyTree[Discrete | MultiDiscrete] property

original_action_space: Space | PyTree[Space] property

Return the original action space of the environment. This is useful for algorithms that need to know the original action space.

jaxnasium.FlattenActionSpaceWrapper

Wrapper to convert (PyTrees of) (multi-)discrete action spaces to a single discrete action space. This grows the action space (significantly for large action spaces), but allows to use algorithms that only support discrete action spaces.

First flattens each MultiDiscrete action space to a single discrete action space, then combines possibly remaining discrete action spaces to a single discrete action space.

Arguments:

• _env: Environment to wrap.

state_type property

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

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

step(key: PRNGKeyArray, state: TEnvState, action: int) -> tuple[TimeStep, TEnvState]

action_space: Discrete property

original_action_space: Space property

Return the original action space of the environment.

jaxnasium.StackActionSpaceWrapper

Wrapper to stack homogeneous discrete action spaces into a single MultiDiscrete.

As an example:

(Discrete(20), Discrete(20), Discrete(20))  # -> MultiDiscrete([20, 20, 20])
(Discrete(10), Discrete(10), Discrete(10), Discrete(5)) -> [MultiDiscrete([10, 10, 10]), Discrete(5)]

This may be useful because Jaxnasium algorithms will vmap homogenous MultiDiscrete spaces, while homogeneous Discrete spaces in a pytree are simply jax.tree.mapped over. Combining them into a single MultiDiscrete may then allow for more performance during runtime and compilation. Noteably, the environment could also define their action space as a single MultiDiscrete, but this wrapper allows for a more flexible definition of the action space, while preserving performance benefits.

Continuous spaces are not stacked, and are passed through as a no-op. NOTE: continuous could be stacked as well, it is currently just not implemented.

Arguments:

• _env: Environment to wrap.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__getattr__(name)

__check_init__()

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

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

action_space: Space | PyTree[Space] property

observation_space: Space | PyTree[Space] property

original_action_space: Space | PyTree[Space] property

Return the original action space of the environment.

jaxnasium.MetaParamsWrapper

Adds params as to the step and reset signature similar to Gymnax

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

reset(key, params: dict)

step(key, state, action, params: dict)

Third-party library wrappers

These are applied automatically by jaxnasium.make when an environment comes from an external suite. You only need them directly if you construct the third-party environment yourself.

jaxnasium.wrappers.BraxWrapper

Wrapper for Brax environments to transform them into the Jaxnasium environment interface.

Note: Brax environments would typically be wrapped with a VmapWrapper, EpisodeWrapper and AutoResetWrapper VmapWrapper is not included here, as it is replaced by Jaxnasium's VecEnvWrapper. The effects of EpisodeWrapper (truncation) and AutoResetWrapper are merged into this wrapper.

Arguments:

• _env: Brax environment.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.GymnaxWrapper

Wrapper for Gymnax environments to transform them into the Jaxnasium environment interface.

Arguments:

• _env: Gymnax environment.
• handle_truncation: If True, the wrapper will reimplement the autoreset behavior to include truncated information and the terminal_observation in the info dictionary. If False, the wrapper will mirror the Gymnax behavior by ignoring truncations. Default=True.
• _params: Gymnax environment parameters. If None, the environment's default_params are used.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.JaxMARLWrapper

Wrapper for JaxMARL environments to transform them into the Jaxnasium environment interface.

Arguments:

• _env: JaxMARL environment.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.JumanjiWrapper

Wrapper for Jumanji environments to transform them into the Jaxnasium environment interface. Multi-agent Jumanji envs (with num_agents) become multi-agent, with agents agent_{i}.

Arguments:

• _env: Jumanji environment.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.NavixWrapper

Wrapper for Navix environments to transform them into the Jaxnasium environment interface.

Arguments:

• _env: Navix environment.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.PgxWrapper

Wrapper for Pgx environments to transform them into the Jaxnasium environment interface.

Arguments:

• _env: Pgx environment.
• self_play: Use a single model for all players in multi-agent environments.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.xMinigridWrapper

Wrapper for xMinigrid environments to transform them into the Jaxnasium environment interface.

Arguments:

• _env: xMinigrid environment.
• _params: xMinigrid environment parameters.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.OctaxWrapper

Wrapper for Octax CHIP-8 environments.

Octax stores the CHIP-8 display as (width, height), so observations are transposed to the usual channels-first image layout (frame_skip, height, width). This mirrors their GymnaxWrapper.

Arguments:

• _env: Octax environment.

state_type property

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

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

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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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

__check_init__()

__getattr__(name)

jaxnasium.wrappers.PlaygroundWrapper

state_type property

step(key: PRNGKeyArray, state: BraxWrapperState, action: float) -> tuple[TimeStep, BraxWrapperState]

reset(key: PRNGKeyArray) -> tuple[Observation, BraxWrapperState]

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

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

observation_space: Any 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, ...]]

Convenience method to sample a random action from the environment's action space. While one could use self.action_space.sample(key), this method additionally works on composite action spaces.

sample_observation(key: PRNGKeyArray) -> Observation

Convenience method to sample a random observation from the environment's observation space. While one could use self.observation_space.sample(key), this method additionally works on composite observation spaces.

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 instance-attribute

__check_init__()

__getattr__(name)

max_episode_steps: int = 1000 class-attribute instance-attribute

_observation_shape: tuple[int, ...] = eqx.field(static=True, init=False) class-attribute instance-attribute

__post_init__()

Utility functions

jaxnasium.is_wrapped(wrapped_env: Environment, wrapper_class: type | str) -> bool

Check if the environment is wrapped with a specific wrapper class.

jaxnasium.unwrap_to(wrapped_env: Environment, wrapper_class: type) -> Environment

The environment inside the outermost wrapper_class, dropping everything around it.

jaxnasium.remove_wrapper(wrapped_env: Environment, wrapper_class: type) -> Environment

Remove every wrapper_class from the environment, keeping the rest of the stack.

jaxnasium.insert_wrapper(wrapped_env: Environment, new_wrapper: Callable[[Environment], Environment], *, inner_wrapper: type) -> Environment

Insert a wrapper into a stack of wrappers at the position of a given inner wrapper.