gxm.Environment#

class Environment#

Bases: Generic[TEnvironmentState], ABC

Base class for environments in gxm. Environments should inherit from this class and implement the init, step, reset, and num_actions methods.

__init__()#

Methods

__init__()

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

unwrapped

Retrieve the base environment by unwrapping all wrappers.

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.

get_wrapper(wrapper_type)#

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

Parameters:

wrapper_type (type[Environment]) – The type of the wrapper to retrieve.

Return type:

Environment

Returns:

The first wrapper of the specified type.

Raises:

ValueError – If no wrapper of the specified type is found.

has_wrapper(wrapper_type)#

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

Parameters:

wrapper_type (type[Environment]) – The type to check for.

Return type:

bool

Returns:

True if the environment or any of its wrappers is of the specified type, False otherwise.

id: str#

The unique identifier of the environment.

abstract init(key)#

Initialize the environment and return the initial state.

Parameters:

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

Return type:

tuple[TypeVar(TEnvironmentState, bound= EnvironmentState), Timestep]

Returns:

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

observation_space: Space#

The observation space of the environment.

abstract 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 (TypeVar(TEnvironmentState, bound= EnvironmentState)) – The current state of the environment.

Return type:

tuple[TypeVar(TEnvironmentState, bound= EnvironmentState), Timestep]

Returns:

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

abstract 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 (TypeVar(TEnvironmentState, bound= EnvironmentState)) – The current state of the environment.

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

Return type:

tuple[TypeVar(TEnvironmentState, bound= EnvironmentState), Timestep]

Returns:

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

property unwrapped: Environment#

Retrieve the base environment by unwrapping all wrappers.

Returns:

The base environment without any wrappers.