gxm.wrappers.Discretize#

class Discretize(wrapped, actions, unwrap=True)#

Bases: Wrapper[Any, TStep]

Wrapper that discretizes a continuous action space. Maps a discrete set of actions to the continuous action space of the environment. The actions are specified as a list of continuous actions \(A\). The action space of the wrapped environment is then \(\{0, 1, \ldots, |A|-1\}\).

>>> import gxm
>>> from gxm.wrappers import Discretize
>>> env = make("Gymnasium/Pendulum-v1")
>>> actions = jnp.array([-2.0, 0.0, 2.0])
>>> env = Discretize(env, actions)

The actions passed to the Discretize wrapper need to be of shape \((|A|, D)\), where \(|A|\) is the number of discrete actions and \(D\) is the dimensionality of the continuous action space of the wrapped environment.

__init__(wrapped, actions, unwrap=True)#
Parameters:

  • wrapped (Dynamics[Any, TypeVar(TStep, bound= Step, covariant=True)]) – The dynamics to wrap.

  • actions (Any) – The discrete set of actions to map to.

  • unwrap (bool) – Whether to unwrap the environment or treat it as part of the base environment.

Methods

__init__(wrapped, actions[, unwrap])

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

unwrap

unwrapped

Retrieve the base dynamics by unwrapping all wrappers.

wrapped

actions

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.

actions: Any#
id: str#

The unique identifier of the dynamics.

init(key)#

Initialize the dynamics and return the initial state.

Parameters:

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

Return type:

tuple[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.

reset(key, state)#

Reset the dynamics to an initial state.

Parameters:

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

  • state (DynamicsState) – The current state.

Return type:

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

Returns:

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

step(key, state, action)#

Advance the dynamics by one step given an action.

Parameters:

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

  • state (DynamicsState) – The current state.

  • action (Any) – The action to apply.

Return type:

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

Returns:

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

wrapped: Dynamics[Any, TypeVar(TStep, bound= Step, covariant=True)]#