gxm.wrappers.StackObservations#

class StackObservations(env, n_stack, padding='reset')#

Bases: Wrapper[StackObservationsState]

Wrapper that stacks the observation along a new axis.

__init__(env, n_stack, padding='reset')#

Methods

__init__(env, n_stack[, padding])

get_wrapper(wrapper_type)

Retrieve the first wrapper of a specific type from the environment.

has_wrapper(wrapper_type)

Check if the environment or any of its wrappers is of a specific type.

init(key)

Initialize the environment and return the initial state.

reset(key, env_state)

Reset the environment to its initial state.

step(key, env_state, action)

Perform a step in the environment given an action.

Attributes

unwrap

unwrapped

Retrieve the base environment by unwrapping all wrappers.

num_stack

padding

env

id

The unique identifier of the environment.

action_space

The action space of the environment.

observation_space

The observation space of the environment.
action_space: Space#

The action space of the environment.

env: Environment#
id: str#

The unique identifier of the environment.

init(key)#

Initialize the environment and return the initial state.

Parameters:

key (Array) – A JAX random key for any stochastic initialization.

Return type:

tuple[StackObservationsState, Timestep]

Returns:

A tuple containing the initial environment state and the initial timestep.

num_stack: int#
observation_space: Space#

The observation space of the environment.

padding: str#
reset(key, env_state)#

Reset the environment to its initial state.

Parameters:

  • key (Array) – A JAX random key for any stochasticity in the environment.

  • env_state (StackObservationsState) – The current state of the environment.

Return type:

tuple[StackObservationsState, Timestep]

Returns:

A tuple containing the reset environment state and the initial timestep.

step(key, env_state, action)#

Perform a step in the environment given an action.

Parameters:

  • key (Array) – A JAX random key for any stochasticity in the environment.

  • env_state (StackObservationsState) – The current state of the environment.

  • action (Any) – The action to take in the environment.

Return type:

tuple[StackObservationsState, Timestep]

Returns:

A tuple containing the new environment state and the resulting timestep.