Skip to content

Gymnasium

lerax.compatibility.gym.GymToLeraxEnv

Bases: AbstractEnv[GymEnvState, Array, Array, None]

Wrap a Gymnasium environment for Lerax.

Note

io_callback makes reset and step slower than native JAX and prevents vmapped rollout. Gymnasium info is discarded because its shape is unknown. Call methods in order because state objects omit required internal state.

Parameters:

Name Type Description Default
env gym.Env

Gymnasium environment to wrap.

 required

Attributes:

Name Type Description
name str

Name of the environment.
action_space AbstractSpace

Action space of the environment.
observation_space AbstractSpace

Observation space of the environment.
env gym.Env

The original Gymnasium environment.

name class-attribute 

name: str = 'GymnasiumEnv'

action_space instance-attribute 

action_space: AbstractSpace = gym_space_to_lerax_space(
    env.action_space
)

observation_space instance-attribute 

observation_space: AbstractSpace = gym_space_to_lerax_space(
    env.observation_space
)

__init__ 

__init__(env: gym.Env)

initial 

initial(
    *args: Any, key: Key[Array, ""], **kwargs: Any
) -> GymEnvState

Call the Gymnasium reset method.

Note

The key generates a reproducible seed unless one is provided.

Parameters:

Name Type Description Default
*args Any

Positional arguments to pass to env.reset.

 ()
key Key[Array, '']

JAX PRNG key used to generate a seed when absent.

 required
**kwargs Any

Keyword arguments to pass to env.reset. If seed is provided, it overrides the generated seed.

 {}

Returns:

Type Description
GymEnvState

The initial environment state.

transition 

transition(
    state: GymEnvState,
    action: Array,
    *,
    key: Key[Array, ""],
) -> GymEnvState

Call the Gymnasium step method through io_callback.

The state is ignored, so call order matters.

Parameters:

Name Type Description Default
state GymEnvState

Current environment state.

 required
action Array

Action to take.

 required
key Key[Array, '']

Unused.

 required

Returns:

Type Description
GymEnvState

Next environment state.

observation 

observation(
    state: GymEnvState, *, key: Key[Array, ""]
) -> Array

Return the stored Gymnasium observation.

Parameters:

Name Type Description Default
state GymEnvState

Current environment state.

 required

Returns:

Type Description
Array

Stored observation.

reward 

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

Return the reward stored in the next state.

Parameters:

Name Type Description Default
state GymEnvState

Current environment state.

 required
action Array

Action taken.

 required
next_state GymEnvState

Next environment state.

 required

Returns:

Type Description
Float[Array, '']

Transition reward.

terminal 

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

Return the stored Gymnasium terminated flag.

Parameters:

Name Type Description Default
state GymEnvState

Current environment state.

 required

Returns:

Type Description
Bool[Array, '']

Whether the state is terminal.

truncate 

truncate(state: GymEnvState) -> Bool[Array, '']

Return the stored Gymnasium truncated flag.

Parameters:

Name Type Description Default
state GymEnvState

Current environment state.

 required

Returns:

Type Description
Bool[Array, '']

Whether the state is truncated.

state_info 

state_info(state: GymEnvState) -> dict

Return empty info to keep JIT shapes stable.

Parameters:

Name Type Description Default
state GymEnvState

Current environment state.

 required

Returns:

Type Description
dict

Empty info.

transition_info 

transition_info(
    state: GymEnvState,
    action: Array,
    next_state: GymEnvState,
) -> dict

Return empty info to keep JIT shapes stable.

Parameters:

Name Type Description Default
state GymEnvState

Current environment state.

 required
action Array

Action taken.

 required
next_state GymEnvState

Next environment state.

 required

Returns:

Type Description
dict

Empty info.

default_renderer 

default_renderer() -> AbstractRenderer

Reject unsupported Gymnasium rendering.

Raises:

Type Description
NotImplementedError

Always.

render 

render(state: GymEnvState, renderer: AbstractRenderer)

Reject unsupported Gymnasium rendering.

Raises:

Type Description
NotImplementedError

Always.

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[Array, ""]
) -> tuple[StateType, ObsType, dict]

Wrap the functional logic into a Gym API reset method.

Parameters:

Name Type Description Default
key Key[Array, '']

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[Array, ""],
) -> 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[Array, '']

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.

close 

close()

lerax.compatibility.gym.LeraxToGymEnv

Bases: gym.Env

Wrap a Lerax environment for Gymnasium.

Run Lerax in Python with internal environment state and a PRNG key.

Attributes:

Name Type Description
metadata dict

Metadata for the Gym environment.
action_space gym.Space

Action space of the environment.
observation_space gym.Space

Observation space of the environment.
render_mode str | None

Render mode for the environment.
env AbstractEnv[StateType, Array, Array, Any]

The Lerax environment to wrap.
state StateType

Current state of the Lerax environment.
key Key[Array, '']

PRNG key for the environment.

Parameters:

Name Type Description Default
env AbstractEnv[StateType, Array, Array, Any]

Lerax environment to wrap.

 required
render_mode Literal['human'] | None

Render mode for the environment.

 None

metadata class-attribute instance-attribute 

metadata: dict = {'render_modes': ['human']}

state instance-attribute 

state: StateType

key instance-attribute 

key: Key[Array, ''] = jr.key(0)

env instance-attribute 

env: AbstractEnv[StateType, Array, Array, Any] = env

action_space instance-attribute 

action_space: gym.Space = lerax_to_gym_space(
    env.action_space
)

observation_space instance-attribute 

observation_space: gym.Space = lerax_to_gym_space(
    env.observation_space
)

render_mode class-attribute instance-attribute 

render_mode: str | None = render_mode

__init__ 

__init__(
    env: AbstractEnv[StateType, Array, Array, Any],
    render_mode: Literal["human"] | None = None,
)

reset 

reset(
    *, seed: int | None = None, options: dict | None = None
)

step 

step(action)

render 

render()

Reject unsupported rendering.

Raises:

Type Description
NotImplementedError

Always.

close 

close()

Implement the Gymnasium Env interface as a no-op.

lerax.compatibility.gym.gym_space_to_lerax_space 

gym_space_to_lerax_space(
    space: gymnasium.Space,
) -> AbstractSpace

Convert a Gymnasium space to a Lerax space.

Parameters:

Name Type Description Default
space gymnasium.Space

Gymnasium space to convert.

 required

Returns:

Type Description
lerax.space.AbstractSpace

Corresponding Lerax space.

lerax.compatibility.gym.lerax_to_gym_space 

lerax_to_gym_space(
    space: lerax.space.AbstractSpace,
) -> gym.Space

Convert a Lerax space to a Gymnasium space.

Parameters:

Name Type Description Default
space lerax.space.AbstractSpace

Lerax space to convert.

 required

Returns:

Type Description
gymnasium.Space

Corresponding Gymnasium space.