Skip to content

Gymnax

lerax.compatibility.gymnax.GymnaxToLeraxEnv

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

Wrapper of a Gymnax environment to make it compatible with Lerax.

Note

For the sake of simplicity, truncation is not handled and always set to False. To keep the API consistent, info returned by step is always an empty dict.

Attributes:

Name Type Description
action_space AbstractSpace

Action space of the environment.
observation_space AbstractSpace

Observation space of the environment.
env gym.Environment

Gymnax environment being wrapped.
params gym.EnvParams

Parameters for the Gymnax environment.

Parameters:

Name Type Description Default
env gym.Environment

Gymnax environment to wrap.

 required
params gym.EnvParams

Parameters for the Gymnax environment.

 required

env instance-attribute 

env: gym.Environment = env

params instance-attribute 

params: gym.EnvParams = params

action_space instance-attribute 

action_space: AbstractSpace = gymnax_space_to_lerax_space(
    env.action_space(params)
)

observation_space instance-attribute 

observation_space: AbstractSpace = (
    gymnax_space_to_lerax_space(
        env.observation_space(params)
    )
)

name property 

name: str

unwrapped property 

unwrapped: Self

Return the unwrapped environment

__init__ 

__init__(env: gym.Environment, params: gym.EnvParams)

initial 

initial(*, key: Key) -> GymnaxEnvState

action_mask 

action_mask(state: GymnaxEnvState, *, key: Key) -> None

transition 

transition(
    state: GymnaxEnvState, action: Array, *, key: Key
) -> GymnaxEnvState

observation 

observation(state: GymnaxEnvState, *, key: Key) -> Array

reward 

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

terminal 

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

truncate 

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

state_info 

state_info(state: GymnaxEnvState) -> dict

transition_info 

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

render 

render(state: GymnaxEnvState, renderer: AbstractRenderer)

default_renderer 

default_renderer() -> AbstractRenderer

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.

lerax.compatibility.gymnax.LeraxToGymnaxEnv

Bases: gym.Environment[LeraxEnvState[StateType], LeraxEnvParams]

Wrapper of an Lerax environment to make it compatible with Gymnax.

Note

Since Gymnax does not have a truncation concept, truncation and termination are combined into a single "done" signal.

Attributes:

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

Lerax environment being wrapped.
state StateType

Current state of the environment.

Parameters:

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

Lerax environment to wrap.

 required

state instance-attribute 

state: StateType

env instance-attribute 

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

name property 

name: str

default_params property 

default_params: LeraxEnvParams

__init__ 

__init__(env: AbstractEnv[StateType, Array, Array, Any])

step_env 

step_env(
    key: Key,
    state: LeraxEnvState[StateType],
    action: ArrayLike,
    params: LeraxEnvParams,
) -> tuple[
    Array,
    LeraxEnvState[StateType],
    Float[Array, ""],
    Bool[Array, ""],
    dict,
]

reset_env 

reset_env(
    key: Key, params: LeraxEnvParams
) -> tuple[Array, LeraxEnvState[StateType]]

get_obs 

get_obs(
    state: LeraxEnvState[StateType],
    key: Key,
    params: LeraxEnvParams,
) -> Array

is_terminal 

is_terminal(
    state: LeraxEnvState[StateType], params: LeraxEnvParams
) -> Bool[Array, ""]

observation_space 

observation_space(
    params: LeraxEnvParams,
) -> gym_spaces.Space

action_space 

action_space(params: LeraxEnvParams) -> gym_spaces.Space

state_space 

state_space(params: LeraxEnvParams) -> gym_spaces.Space

lerax.compatibility.gymnax.gymnax_space_to_lerax_space 

gymnax_space_to_lerax_space(
    space: gymnax.environments.spaces.Space,
) -> AbstractSpace

Returns a Lerax space corresponding to the given Gymnax space.

Parameters:

Name Type Description Default
space gymnax.environments.spaces.Space

Gymnax space to convert.

 required

Returns:

Type Description
lerax.space.AbstractSpace

The corresponding Lerax space.

lerax.compatibility.gymnax.lerax_to_gymnax_space 

lerax_to_gymnax_space(
    space: lerax.space.AbstractSpace,
) -> gym_spaces.Space

Returns a Gymnax space corresponding to the given Lerax space.

Parameters:

Name Type Description Default
space lerax.space.AbstractSpace

Lerax space to convert.

 required

Returns:

Type Description
gymnax.environments.spaces.Space

The corresponding Gymnax space.