Skip to content

Ant

Lerax port of Gymnasium's Ant environment. A quadruped (torso + 4 two-link legs, 8 actuated hinges) is dropped on flat ground and must run in the +x direction without flipping over.

Observation space

Concatenation of qpos, qvel, and clipped external contact forces cfrc_ext[1:] flattened. By default exclude_current_positions_from_observation=True drops qpos[:2] (root x/y) and include_cfrc_ext_in_observation=True includes contact forces. Observation is unbounded (Box(-inf, inf)).

Action space

Box(low, high) taken from the MJX model's actuator_ctrlrange — 8 continuous torques, one per hinge.

Reward

forward_reward + healthy_reward - ctrl_cost - contact_cost where:

  • forward_reward = forward_reward_weight * x_velocity (x velocity computed from torso xpos change over dt)
  • healthy_reward = is_healthy * healthy_reward (default 1.0)
  • ctrl_cost = ctrl_cost_weight * sum(action ** 2) (default weight 0.5)
  • contact_cost = contact_cost_weight * sum(clip(cfrc_ext, -1, 1) ** 2) (default weight 5e-4)

Termination

Terminates when unhealthy if terminate_when_unhealthy=True (default). Healthy means qpos/qvel are finite and torso z is within healthy_z_range=(0.2, 1.0). No built-in truncation.

lerax.env.mujoco.Ant

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

MJX-based ant environment matching Gymnasium's Ant-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, '']

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