gxm.wrappers.FlattenObservation#

class FlattenObservation(env, unwrap=True)#

Bases: Wrapper

Wrapper that adds a rollout method to the environment.

__init__(env, unwrap=True)#

Methods

__init__(env[, unwrap])

flatten(obs)

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.

env

id

The unique identifier of the environment.

action_space

The action space of the environment.

observation_space

The observation space of the environment.
classmethod flatten(obs)#
Return type:

Array

init(key)#

Initialize the environment and return the initial state.

Parameters:

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

Return type:

tuple[WrapperState, Timestep]

Returns:

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

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 (WrapperState) – The current state of the environment.

Return type:

tuple[WrapperState, 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 (WrapperState) – The current state of the environment.

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

Return type:

tuple[WrapperState, Timestep]

Returns:

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