gxm.Environment#

class Environment(*args, **kwargs)#

Bases: Dynamics[TEnvironmentState, Timestep], Protocol[TEnvironmentState]

Base class for RL environments in gxm.

Extends Dynamics with episode structure: each step returns a Timestep that includes reward, termination, and truncation signals. Environments should inherit from this class and implement init, step, and reset.

__init__(*args, **kwargs)#

Methods

__init__(*args, **kwargs)

get_wrapper(wrapper_type)

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

has_wrapper(wrapper_type)

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

init(key)

Initialize the environment and return the initial state.

reset(key, state)

Reset the environment to its initial state.

step(key, state, action)

Perform a step in the environment given an action.

Attributes

unwrapped

Retrieve the base dynamics by unwrapping all wrappers.

id

The unique identifier of the dynamics.

action_space

The action space of the dynamics.

observation_space

The observation space of the dynamics.
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= DynamicsState), Timestep]

Returns:

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

abstract reset(key, state)#

Reset the environment to its initial state.

Parameters:

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

  • state (TypeVar(TEnvironmentState, bound= DynamicsState)) – The current state of the environment.

Return type:

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

Returns:

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

abstract step(key, state, action)#

Perform a step in the environment given an action.

Parameters:

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

  • state (TypeVar(TEnvironmentState, bound= DynamicsState)) – The current state of the environment.

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

Return type:

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

Returns:

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