mujoco

Modular play utilities for LAV2 policy deployment.

This module provides reusable components for playing trained policies in MuJoCo environments with keyboard and gamepad control, supporting various action and command modes.

Design follows lav2.controller.run pattern.

类：

名称	描述
MuJoCoPlayConfig	Configuration for MuJoCo policy play session.
MuJoCoPlayRunner	Modular runner for MuJoCo policy play sessions with keyboard/gamepad control.

MuJoCoPlayConfig dataclass

MuJoCoPlayConfig(checkpoint_path: str, params: VehicleParams = VehicleParams(), decimation: int = 2, action_mode: str = 'cmd_motor_speeds', command_mode: str = 'tgt_pos', initial_target: ndarray = (lambda: np.array([1.0, 1.0, 1.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]))(), mjcf_path: str = str(files('lav2.assets').joinpath('mjcf/scene.xml')), enable_gamepad: bool = False, trajectory: BaseTrajectory | None = None, pos_inc: tuple[float, float, float] = (1.0, 1.0, 0.5), vel_limits: tuple[float, float, float] = (5.0, 5.0, 5.0), att_limits: tuple[float, float, float] = (1.0, 1.0, np.pi), br_limits: tuple[float, float, float] = (2.0, 2.0, 2.0))

Configuration for MuJoCo policy play session.

属性：

名称	类型	描述
checkpoint_path	str	Path to the policy checkpoint (JIT or ONNX).
params	VehicleParams	Vehicle parameters (auto-created if None).
decimation	int	Control decimation factor.
action_mode	str	Action mode - 'cmd_pos', 'cmd_vel', 'cmd_ctatt', 'cmd_ctbr', 'cmd_ctbm', 'cmd_motor_speeds'.
command_mode	str	Command mode - 'tgt_pos', 'tgt_vel', 'tgt_ctatt', 'tgt_ctbr'.
initial_target	ndarray	Initial target [x,y,z,vx,vy,vz,roll,pitch,yaw,wx,wy,wz].
mjcf_path	str	Path to MuJoCo XML scene file.
enable_gamepad	bool	Enable gamepad input.
pos_inc	tuple[float, float, float]	Positional increment for keyboard control (dx, dy, dz).
vel_limits	tuple[float, float, float]	Velocity limits [vx, vy, vz].
att_limits	tuple[float, float, float]	Attitude limits [roll, pitch, yaw].
br_limits	tuple[float, float, float]	Body rate limits [wx, wy, wz].

方法：

名称	描述
__post_init__	Validate configuration values after initialization.

__post_init__

__post_init__()

Validate configuration values after initialization.

引发：

类型	描述
ValueError	If action_mode or command_mode is not a valid option.

MuJoCoPlayRunner

MuJoCoPlayRunner(cfg: MuJoCoPlayConfig, policy_runner: JITPolicyRunner | ONNXPolicyRunner, obs_extractor: Callable[[MjData, ndarray, ndarray], ndarray] | None = None)

Modular runner for MuJoCo policy play sessions with keyboard/gamepad control.

Provides reusable observation extraction, control loop, and input handling for LAV2 flight control in MuJoCo. Supports multiple action modes following the lav2.controller.run design pattern.

Initialize the play runner.

参数：

名称	类型	描述	默认
cfg	MuJoCoPlayConfig	Play cfguration.	必需
policy_runner	JITPolicyRunner | ONNXPolicyRunner	Policy runner (JIT or ONNX).	必需
obs_extractor	Callable[[MjData, ndarray, ndarray], ndarray] | None	Custom observation extraction function. If None, uses default flight observation extractor. Signature: (d: MjData, target: ndarray, last_action: ndarray) -> obs	None

方法：

名称	描述
key_callback	Handle keyboard input to update target position (tgt_pos mode only).
gamepad_callback	Handle gamepad input following lav2.controller.run pattern.
process_actions	Process policy action and compute RPM commands (control loop).
apply_actions	Apply rotor dynamics and update MuJoCo actuators (simulation loop).
control_callback	Control loop callback for MuJoCo viewer (respects decimation).
load_callback	Load MuJoCo model and data.
run	Launch the MuJoCo viewer and start the play session.

key_callback

key_callback(keycode: int)

Handle keyboard input to update target position (tgt_pos mode only).

Keys 1-8 adjust target position and heading. Note: Currently only supports tgt_pos command mode.

gamepad_callback

gamepad_callback(key: Key) -> None

Handle gamepad input following lav2.controller.run pattern.

参数：

名称	类型	描述	默认
key	Key	Gamepad key event (axis or button).	必需

process_actions

process_actions(m: MjModel, d: MjData)

Process policy action and compute RPM commands (control loop).

This method is called at the decimated control rate.

参数：

名称	类型	描述	默认
m	MjModel	MuJoCo model.	必需
d	MjData	MuJoCo data.	必需

apply_actions

apply_actions(m: MjModel, d: MjData)

Apply rotor dynamics and update MuJoCo actuators (simulation loop).

This runs at every simulation step, not subject to decimation.

参数：

名称	类型	描述	默认
m	MjModel	MuJoCo model.	必需
d	MjData	MuJoCo data.	必需

control_callback

control_callback(m: MjModel, d: MjData)

Control loop callback for MuJoCo viewer (respects decimation).

Executes policy inference and command processing, but applies rotor dynamics separately to respect simulation decimation.

参数：

名称	类型	描述	默认
m	MjModel	MuJoCo model.	必需
d	MjData	MuJoCo data.	必需

load_callback

load_callback() -> tuple[mujoco.MjModel, mujoco.MjData]

Load MuJoCo model and data.

返回：

类型	描述
tuple[MjModel, MjData]	tuple[mujoco.MjModel, mujoco.MjData]: Tuple of (model, data).

run

run()

Launch the MuJoCo viewer and start the play session.