arlbench.core.wrappers.flatten_observation

Observation Flattening Wrapper.

Classes

FlattenObservationWrapper(env)

Wraps the given environment to flatten its observations.
class arlbench.core.wrappers.flatten_observation.FlattenObservationWrapper(env)[source]

Bases: Wrapper

Wraps the given environment to flatten its observations.

property observation_space: Box

The flattened observation space of the environment.

Returns:

Flattened obseration space.

Return type:

spaces.Box

reset(rng)[source]
Calls the reset() function of the environment and

flattens the returned observations.

Parameters:

rng (chex.PRNGKey) – Random number generator key.

Returns:

Result of the step function but observations are flattened.

Return type:

tuple[Any, jnp.ndarray]

step(env_state, action, rng)[source]

Calls the step() function of the environment and flattens the returned observations.

Parameters:

  • env_state (Any) – The internal environment state.

  • action (jnp.ndarray) – The actions to take.

  • rng (chex.PRNGKey) – Random number generator key.

Returns:

Result of the step function but observations are flattened.

Return type:

tuple[Any, tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, dict]]

previous

arlbench.core.wrappers

next

arlbench.core.wrappers.wrapper