arlbench.core.wrappers package¶

Submodules¶

arlbench.core.wrappers.flatten_observation module¶

Observation Flattening Wrapper.

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

arlbench.core.wrappers.wrapper module¶

Wrapper class.

class arlbench.core.wrappers.wrapper.Wrapper(env)[source]¶

Bases: object

Base class for ARLBench wrappers.

Module contents¶

class arlbench.core.wrappers.FlattenObservationWrapper(env)[source]¶