gxm.Dynamics#

class Dynamics(*args, **kwargs)#

Bases: Protocol[TDynamicsState, TStep]

Base class for world models in gxm.

Dynamics define state transitions: given an action, they transition to a new state and produce a step output. They have no notion of episodes, rewards, or termination — those are added by Environment.

All Environment instances are also Dynamics instances, so any function typed dynamics: Dynamics can accept an environment directly.

__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 dynamics and return the initial state.

reset(key, state)

Reset the dynamics to an initial state.

step(key, state, action)

Advance the dynamics by one step 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.
action_space: Space#

The action space of the dynamics.

get_wrapper(wrapper_type)#

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

Parameters:

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

Return type:

Dynamics

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 dynamics or any of its wrappers is of a specific type.

Parameters:

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

Return type:

bool

Returns:

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

id: str#

The unique identifier of the dynamics.

abstract init(key)#

Initialize the dynamics and return the initial state.

Parameters:

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

Return type:

tuple[TypeVar(TDynamicsState, bound= DynamicsState), TypeVar(TStep, bound= Step, covariant=True)]

Returns:

A tuple of the initial state and the initial step output.

observation_space: Space#

The observation space of the dynamics.

abstract reset(key, state)#

Reset the dynamics to an initial state.

Parameters:

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

  • state (TypeVar(TDynamicsState, bound= DynamicsState)) – The current state.

Return type:

tuple[TypeVar(TDynamicsState, bound= DynamicsState), TypeVar(TStep, bound= Step, covariant=True)]

Returns:

A tuple of the reset state and the initial step output.

abstract step(key, state, action)#

Advance the dynamics by one step given an action.

Parameters:

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

  • state (TypeVar(TDynamicsState, bound= DynamicsState)) – The current state.

  • action (Any) – The action to apply.

Return type:

tuple[TypeVar(TDynamicsState, bound= DynamicsState), TypeVar(TStep, bound= Step, covariant=True)]

Returns:

A tuple of the new state and the resulting step output.

property unwrapped: Dynamics#

Retrieve the base dynamics by unwrapping all wrappers.

Returns:

The base dynamics without any wrappers.