Skip to content

Inverted Double Pendulum

Lerax port of Gymnasium's Inverted Double Pendulum environment. A MuJoCo cart with two un-actuated hinge links stacked on top. The agent applies horizontal cart force and must keep the tip near the upright position.

Observation space

9-dim float vector: [cart_x, sin(theta1), sin(theta2), cos(theta1), cos(theta2), clip(qvel, -10, 10)[0..2], clip(qfrc_constraint, -10, 10)[0]]. Unbounded Box.

Action space

Box(low, high) from the model's actuator_ctrlrange — a single continuous force applied to the cart.

Reward

alive_bonus - dist_penalty - vel_penalty with

  • alive_bonus = (tip_y > 1) * healthy_reward (default healthy_reward = 10.0)
  • dist_penalty = 0.01 * tip_x ** 2 + (tip_y - 2) ** 2
  • vel_penalty = 1e-3 * qvel[1] ** 2 + 5e-3 * qvel[2] ** 2

tip_x, tip_y are the x and z of site_xpos[0].

Termination

Terminates when the tip site z drops to <= 1. No built-in truncation.

lerax.env.mujoco.InvertedDoublePendulum

Bases: AbstractMujocoEnv[Float[Array, '...'], Float[Array, '...']]

MJX-based inverted double pendulum environment matching Gymnasium's InvertedDoublePendulum-v5.

unwrapped property

unwrapped: Self

Return the unwrapped environment

action_mask

action_mask(
    state: MujocoEnvState, *, key: Key[Array, ""]
) -> None

transition

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

truncate

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

state_info

state_info(state: MujocoEnvState) -> dict

default_renderer

default_renderer() -> MujocoRenderer

render

render(state: MujocoEnvState, 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[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.

model_from_path staticmethod

model_from_path(xml_file: str | Path) -> mujoco.MjModel