Skip to content

Rescale Action

lerax.wrapper.RescaleAction

Bases: AbstractPureTransformActionWrapper[Float[Array, ' ...'], MaskType, StateType, Float[Array, ' ...'], ObsType, MaskType]

Affine rescaling of a box action to a different range.

Note

Only compatible with Box action spaces.

Attributes:

Name Type Description
env AbstractEnvLike[StateType, Float[Array, ' ...'], ObsType, MaskType]

The environment to wrap.
action_space Box

The action space of the wrapper.

Parameters:

Name Type Description Default
env AbstractEnvLike[StateType, Float[Array, ' ...'], ObsType, MaskType]

The environment to wrap.

 required

Raises:

Type Description
ValueError

If the environment's action space is not a Box.

name property 

name: str

Return the name of the environment

observation_space property 

observation_space: AbstractSpace[ObsType, Any]

unwrapped property 

unwrapped: AbstractEnv

Return the wrapped environment

env instance-attribute 

env: AbstractEnvLike[
    StateType, Float[Array, " ..."], ObsType, MaskType
] = env

func instance-attribute 

func: Callable[
    [Float[Array, " ..."]], Float[Array, " ..."]
] = rescale

mask_func instance-attribute 

mask_func: Callable[[MaskType], MaskType] = lambda x: x

action_space instance-attribute 

action_space: Box = action_space

initial 

initial(*, key: Key) -> TransformActionState[StateType]

action_mask 

action_mask(
    state: TransformActionState[StateType], *, key: Key
) -> WrapperMaskType | None

transition 

transition(
    state: TransformActionState[StateType],
    action: WrapperActType,
    *,
    key: Key,
) -> TransformActionState[StateType]

observation 

observation(
    state: TransformActionState[StateType], *, key: Key
) -> ObsType

reward 

reward(
    state: TransformActionState[StateType],
    action: WrapperActType,
    next_state: TransformActionState[StateType],
    *,
    key: Key,
) -> Float[Array, ""]

terminal 

terminal(
    state: TransformActionState[StateType], *, key: Key
) -> Bool[Array, ""]

truncate 

truncate(
    state: TransformActionState[StateType],
) -> Bool[Array, ""]

state_info 

state_info(state: TransformActionState[StateType]) -> dict

transition_info 

transition_info(
    state: TransformActionState[StateType],
    action: WrapperActType,
    next_state: TransformActionState[StateType],
) -> dict

default_renderer 

default_renderer() -> AbstractRenderer

Return the default renderer for the wrapped environment

render 

render(state: WrapperStateType, renderer: AbstractRenderer)

Render a frame from a state

render_states 

render_states(
    states: Sequence[StateType],
    renderer: AbstractRenderer | Literal["auto"] = "auto",
    dt: float = 0.0,
)

Render a sequence of frames from multiple states.

Parameters:

Name Type Description Default
states Sequence[StateType]

A sequence of environment states to render.

 required
renderer AbstractRenderer | Literal['auto']

The renderer to use for rendering. If "auto", uses the default renderer.

 'auto'
dt float

The time delay between rendering each frame, in seconds.

 0.0

render_stacked 

render_stacked(
    states: StateType,
    renderer: AbstractRenderer | Literal["auto"] = "auto",
    dt: float = 0.0,
)

Render multiple frames from stacked states.

Stacked states are typically batched states stored in a pytree structure.

Parameters:

Name Type Description Default
states StateType

A pytree of stacked environment states to render.

 required
renderer AbstractRenderer | Literal['auto']

The renderer to use for rendering. If "auto", uses the default renderer.

 'auto'
dt float

The time delay between rendering each frame, in seconds.

 0.0

reset 

reset(*, key: Key) -> tuple[StateType, ObsType, dict]

Wrap the functional logic into a Gym API reset method.

Parameters:

Name Type Description Default
key Key

A JAX PRNG key for any stochasticity in the reset.

 required

Returns:

Type Description
tuple[StateType, ObsType, dict]

A tuple of the initial state, initial observation, and additional info.

step 

step(
    state: StateType, action: ActType, *, key: Key
) -> tuple[
    StateType,
    ObsType,
    Float[Array, ""],
    Bool[Array, ""],
    Bool[Array, ""],
    dict,
]

Wrap the functional logic into a Gym API step method.

Parameters:

Name Type Description Default
state StateType

The current environment state.

 required
action ActType

The action to take.

 required
key Key

A JAX PRNG key for any stochasticity in the step.

 required

Returns:

Type Description
tuple[StateType, ObsType, Float[Array, ''], Bool[Array, ''], Bool[Array, ''], dict]

A tuple of the next state, observation, reward, terminal flag, truncate flag, and additional info.

__init__ 

__init__(
    env: AbstractEnvLike[
        StateType, Float[Array, " ..."], ObsType, MaskType
    ],
    min: Float[Array, " ..."] = jnp.array(-1.0),
    max: Float[Array, " ..."] = jnp.array(1.0),
)