MDP terms

Autodoc for shared manager terms under wbc_mjlab.env.mdp. Design rules and term catalogs: Shared MDP (Shared MDP).

Jump to: Actions · Motion command · RSI / adaptive sampling · Observations · Rewards · Terminations · Events (assistive wrench)

Actions

class wbc_mjlab.env.mdp.actions.ReferenceJointPositionActionCfg

Bases: BaseActionCfg

Joint position targets as motion reference plus scaled residual actions.

Implements Zest-style control: q_cmd = q_ref + scale * a_t (minus encoder bias).

command_name: str = 'motion'

build(env: ManagerBasedRlEnv) → ReferenceJointPositionAction

class wbc_mjlab.env.mdp.actions.ReferenceJointPositionAction

Bases: BaseAction

Control joints via position targets relative to the motion reference.

cfg: ReferenceJointPositionActionCfg

apply_actions() → None

Motion command

class wbc_mjlab.env.mdp.commands.MotionLoader

Bases: object

Load a motion library from a bundled NPZ or a dataset directory (npz/*.npz).

class wbc_mjlab.env.mdp.commands.MotionCommand

Bases: CommandTerm

Multi-clip motion playback with RSI resampling and assistive-wrench state.

Loads an NPZ motion library, streams reference kinematics each step, and resamples start frames at episode boundaries according to RsiCfg. Observation and reward terms read reference state from this command via command_name="motion".

cfg: MotionCommandCfg

property bin_count: int

property ref_base_height: torch.Tensor

Anchor height relative to env origin (z_I r̂_IB).

property ref_base_lin_vel_b: torch.Tensor

Reference anchor linear velocity in anchor frame (B v̂_IB).

property ref_base_ang_vel_b: torch.Tensor

Reference anchor angular velocity in anchor frame (B ω̂_IB).

property ref_gravity_b: torch.Tensor

Reference gravity in anchor frame (B ĝ_I).

property ref_base_lin_acc_b: torch.Tensor

property ref_base_ang_acc_b: torch.Tensor

property tracked_joint_pos: torch.Tensor

Reference joint positions for actuated / tracked DoFs (absolute).

property tracked_joint_vel: torch.Tensor

Reference joint velocities for actuated / tracked DoFs (absolute).

property command: torch.Tensor

Legacy stacked WBC reference vector (same layout as default ref obs terms).

Prefer configuring individual reference observation terms in wbc_env_cfg. Kept for ONNX metadata, wbc_command_dim, and legacy deploy bundles.

property wbc_command_dim: int

property joint_pos: torch.Tensor

property joint_vel: torch.Tensor

property body_pos_w: torch.Tensor

property body_quat_w: torch.Tensor

property body_lin_vel_w: torch.Tensor

property body_ang_vel_w: torch.Tensor

property anchor_pos_w: torch.Tensor

property anchor_quat_w: torch.Tensor

property anchor_lin_vel_w: torch.Tensor

property anchor_ang_vel_w: torch.Tensor

property anchor_lin_acc_w: torch.Tensor

property anchor_ang_acc_w: torch.Tensor

property robot_joint_pos: torch.Tensor

property robot_joint_vel: torch.Tensor

property robot_body_pos_w: torch.Tensor

property robot_body_quat_w: torch.Tensor

property robot_body_lin_vel_w: torch.Tensor

property robot_body_ang_vel_w: torch.Tensor

property robot_anchor_pos_w: torch.Tensor

property robot_anchor_quat_w: torch.Tensor

property robot_anchor_lin_vel_w: torch.Tensor

property robot_anchor_ang_vel_w: torch.Tensor

update_relative_body_poses() → None

compute(dt: float) → None

set_viewer_task_id(task_id: str | None) → None

update_viewer_gui(env_idx: int) → None

create_gui(name: str, server: viser.ViserServer, get_env_idx: Callable[[], int], on_change: Callable[[], None] | None = None, request_action: Callable[[str, Any], None] | None = None) → None

on_viewer_pause(paused: bool) → None

apply_gui_reset(env_ids: torch.Tensor) → bool

reset_to_frame(env_ids: torch.Tensor, frame: int) → None

class wbc_mjlab.env.mdp.commands.MotionCommandCfg

Bases: MotionCommandCfg

Config for MotionCommand (set per robot in <robot>_base_cfg).

motion_file: str

Path to converted NPZ library (from wbc-mjlab-data-to-npz).

anchor_body_name: str

Body used for anchor-frame errors and assistive wrench.

body_names: tuple[str, ...]

Keybodies tracked in rewards, RSI, and critic observations.

entity_name: str

Scene entity name for the robot (usually "robot").

actuated_joint_names: tuple[str, ...] = ()

If set, joint tracking metrics/RSI use only these DoFs (subset of the robot).

pose_range: dict[str, tuple[float, float]]

Domain randomization ranges for reference root pose at resample.

velocity_range: dict[str, tuple[float, float]]

Domain randomization ranges for reference root velocity at resample.

joint_position_range: tuple[float, float] = (-0.52, 0.52)

Domain randomization range for joint position offsets at resample.

rsi: RsiCfg

Reference-state initialization / adaptive bin sampling.

assistive_wrench_enabled: bool = True

Whether assistive wrench curriculum is active.

assistive_beta_max: float = 0.6

Max assistive gain β.

assistive_eta: float = 0.8

Assistive wrench curriculum exponent.

class VizCfg

Bases: object

mode: Literal['ghost', 'frames'] = 'ghost'

ghost_color: tuple[float, float, float, float] = (0.45, 0.6, 0.9, 0.5)

viz: VizCfg

build(env: ManagerBasedRlEnv) → MotionCommand

RSI / adaptive sampling

class wbc_mjlab.env.mdp.sampling.RsiCfg

Bases: object

Reference-state initialization (RSI) and adaptive bin sampling.

sampling_mode: Literal['adaptive', 'uniform', 'start'] = 'adaptive'

strategy: Literal['binary_failure', 'similarity_ema'] = 'similarity_ema'

similarity_terms: tuple[AdaptiveSimilarityTermCfg, ...]

bin_width_s: float = 4.0

uniform_ratio: float = 0.15

alpha: float = 0.005

temperature_base: float = 1.0

similarity_norm_by_remaining_clip: bool = False

min_bin_span_ratio: float = 0.0

persist_failure_levels: bool = False

failure_levels_filename: str = 'rsi_bin_stats.npz'

similarity_from_rewards: bool = False

tracking_reward_prefix: str = 'motion_'

class wbc_mjlab.env.mdp.sampling.AdaptiveSimilarityTermCfg

Bases: object

One exp-kernel term in the per-step RSI similarity score s_k.

term: Literal['joint_pos', 'anchor_pos', 'anchor_ori', 'body_pos', 'body_ori', 'body_lin_vel', 'body_ang_vel']

weight: float = 1.0

std: float | None = None

class wbc_mjlab.env.mdp.sampling.TrackingSimilarityState

Bases: object

Robot vs reference errors for one similarity step.

tracked_joint_pos_error: torch.Tensor

anchor_pos_error: torch.Tensor

anchor_ori_error: torch.Tensor

body_pos_error: torch.Tensor

body_ori_error: torch.Tensor

body_lin_vel_error: torch.Tensor

body_ang_vel_error: torch.Tensor

wbc_mjlab.env.mdp.sampling.joint_pos_similarity_preset() → tuple[AdaptiveSimilarityTermCfg, ...]

Hand-tuned RSI similarity: joint position only.

wbc_mjlab.env.mdp.sampling.keybody_similarity_preset() → tuple[AdaptiveSimilarityTermCfg, ...]

Hand-tuned RSI similarity: anchor + keybody tracking terms.

wbc_mjlab.env.mdp.sampling.sample_adaptive_bins(bin_failure_levels: torch.Tensor, valid_bin_indices: torch.Tensor, *, segment_length: torch.Tensor, segment_start_idx: torch.Tensor, bin_width_frames: int, temperature_base: float, uniform_ratio: float, num_samples: int, device: torch.device | str) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Draw (trajectory, bin, frame) from failure-weighted bin distribution.

wbc_mjlab.env.mdp.sampling.save_rsi_bin_stats(path: str | Path, command: MotionCommand) → Path

Write adaptive RSI bin failure levels to path (.npz).

wbc_mjlab.env.mdp.sampling.load_rsi_bin_stats(path: str | Path, command: MotionCommand, *, strict: bool = False) → bool

Restore bin failure levels from path. Returns False if the file is missing.

Observations

wbc_mjlab.env.mdp.observations.ref_base_height(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference anchor height relative to env origin (z_I r̂_IB).

wbc_mjlab.env.mdp.observations.ref_anchor_pos_w(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference anchor xyz position relative to env origin (SE reference command).

wbc_mjlab.env.mdp.observations.ref_base_lin_vel_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference anchor linear velocity in anchor frame (B v̂_IB).

wbc_mjlab.env.mdp.observations.ref_base_ang_vel_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference anchor angular velocity in anchor frame (B ω̂_IB).

wbc_mjlab.env.mdp.observations.ref_gravity_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference gravity in anchor frame (B ĝ_I).

wbc_mjlab.env.mdp.observations.ref_joint_pos(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference joint positions for tracked DoFs (absolute).

wbc_mjlab.env.mdp.observations.ref_joint_vel(env: ManagerBasedRlEnv, command_name: str, asset_cfg: SceneEntityCfg = mjlab.managers.scene_entity_config.SceneEntityCfg) → torch.Tensor

Reference joint velocities (critic privileged).

wbc_mjlab.env.mdp.observations.ref_anchor_ori_6d(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference anchor orientation as 6D rotation matrix columns (world frame).

wbc_mjlab.env.mdp.observations.motion_anchor_pos_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Full xyz anchor tracking error in robot anchor frame (non-SE layouts).

wbc_mjlab.env.mdp.observations.motion_anchor_ori_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

wbc_mjlab.env.mdp.observations.motion_anchor_pos_error_w(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

World-frame anchor position tracking error (ref − robot).

wbc_mjlab.env.mdp.observations.motion_anchor_ori_error(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Anchor orientation tracking error as axis-angle (3); ‖·‖ = quat_error_magnitude.

wbc_mjlab.env.mdp.observations.robot_body_pos_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

wbc_mjlab.env.mdp.observations.robot_body_ori_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

wbc_mjlab.env.mdp.observations.ref_body_pos_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference keybody positions in the robot anchor frame.

wbc_mjlab.env.mdp.observations.ref_body_ori_b(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Reference keybody orientations in the robot anchor frame (6D rotation columns).

wbc_mjlab.env.mdp.observations.motion_segment_phase(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Normalized clip progress in [0, 1] (privileged critic feature).

wbc_mjlab.env.mdp.observations.motion_tracking_step_rewards(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Unweighted per-step motion_* reward values (privileged critic feature).

Rewards

Tracking kernels and regularizers. Shared optional params are documented on the module docstring of wbc_mjlab.env.mdp.rewards.

wbc_mjlab.env.mdp.rewards.tracking_std_from_sigma(sigma: float, *, dim: int = 1) → float

Map per-DoF σ to aggregate std with std = 2σ√dim (Table S4 convention).

wbc_mjlab.env.mdp.rewards.motion_global_anchor_position_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float, *, kappa: float = 1.0) → torch.Tensor

Exponential reward for anchor position tracking in world frame.

wbc_mjlab.env.mdp.rewards.motion_global_anchor_orientation_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float, *, kappa: float = 1.0) → torch.Tensor

Exponential reward for anchor orientation tracking in world frame.

wbc_mjlab.env.mdp.rewards.motion_anchor_linear_velocity_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float, *, kappa: float = 1.0) → torch.Tensor

Exponential reward for anchor linear velocity tracking (world frame).

wbc_mjlab.env.mdp.rewards.motion_anchor_angular_velocity_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float, *, kappa: float = 1.0) → torch.Tensor

Exponential reward for anchor angular velocity tracking (world frame).

wbc_mjlab.env.mdp.rewards.motion_joint_position_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float | None = None, sigma_per_joint: float | None = None, per_joint: bool = False, *, kappa: float = 1.0, asset_cfg: SceneEntityCfg = mjlab.managers.scene_entity_config.SceneEntityCfg) → torch.Tensor

Joint position tracking with optional per-joint exponential kernels.

Default: exp(-κ Σ e_i² / std²) with std = 2σ√n when sigma_per_joint is set. per_joint=True: mean over joints of exp(-κ e_i² / (2σ)²).

wbc_mjlab.env.mdp.rewards.motion_joint_velocity_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float | None = 1.0, sigma_per_joint: float | None = None, per_joint: bool = False, *, kappa: float = 1.0, asset_cfg: SceneEntityCfg = mjlab.managers.scene_entity_config.SceneEntityCfg) → torch.Tensor

Joint velocity tracking; same kernel options as motion_joint_position_error_exp.

wbc_mjlab.env.mdp.rewards.motion_relative_body_position_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float | None = None, body_names: tuple[str, ...] | None = None, sigma_per_keybody: float | None = None, per_keybody: bool = False, body_error_aggregate: str = 'mean', *, kappa: float = 1.0) → torch.Tensor

Exponential reward for keybody position tracking (anchor-relative).

wbc_mjlab.env.mdp.rewards.motion_relative_body_orientation_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float | None = None, body_names: tuple[str, ...] | None = None, sigma_per_keybody: float | None = None, per_keybody: bool = False, body_error_aggregate: str = 'mean', *, kappa: float = 1.0) → torch.Tensor

Exponential reward for keybody orientation tracking (anchor-relative).

wbc_mjlab.env.mdp.rewards.motion_global_body_linear_velocity_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float | None = None, body_names: tuple[str, ...] | None = None, sigma_per_keybody: float | None = None, per_keybody: bool = False, body_error_aggregate: str = 'mean', *, kappa: float = 1.0) → torch.Tensor

Exponential reward for keybody linear velocity tracking (world frame).

wbc_mjlab.env.mdp.rewards.motion_global_body_angular_velocity_error_exp(env: ManagerBasedRlEnv, command_name: str, std: float | None = None, body_names: tuple[str, ...] | None = None, sigma_per_keybody: float | None = None, per_keybody: bool = False, body_error_aggregate: str = 'mean', *, kappa: float = 1.0) → torch.Tensor

Exponential reward for keybody angular velocity tracking (world frame).

wbc_mjlab.env.mdp.rewards.action_rate_l1(env: ManagerBasedRlEnv) → torch.Tensor

Penalize action changes with L1 (sum of absolute deltas).

wbc_mjlab.env.mdp.rewards.actuator_torque_soft_limit(env: ManagerBasedRlEnv, soft_ratio: float = 0.95, asset_cfg: SceneEntityCfg = mjlab.managers.scene_entity_config.SceneEntityCfg) → torch.Tensor

Penalize actuator forces approaching MuJoCo torque limits (normalized soft margin).

wbc_mjlab.env.mdp.rewards.feet_slip(env: ManagerBasedRlEnv, sensor_name: str, asset_cfg: SceneEntityCfg = mjlab.managers.scene_entity_config.SceneEntityCfg) → torch.Tensor

Penalize foot planar velocity while in contact (slip).

wbc_mjlab.env.mdp.rewards.angular_momentum_penalty(env: ManagerBasedRlEnv, sensor_name: str, *, axes: str = 'xy') → torch.Tensor

Penalize whole-body angular momentum (roll/pitch by default).

Terminations

wbc_mjlab.env.mdp.terminations.bad_anchor_pos(env: ManagerBasedRlEnv, command_name: str, threshold: float) → torch.Tensor

Terminate when anchor body position drifts beyond threshold (m).

wbc_mjlab.env.mdp.terminations.bad_anchor_pos_z_only(env: ManagerBasedRlEnv, command_name: str, threshold: float) → torch.Tensor

Terminate when anchor height (z) drifts beyond threshold (m).

wbc_mjlab.env.mdp.terminations.bad_anchor_ori(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg, command_name: str, threshold: float) → torch.Tensor

Terminate when projected-gravity z mismatch exceeds threshold.

wbc_mjlab.env.mdp.terminations.bad_motion_body_pos(env: ManagerBasedRlEnv, command_name: str, threshold: float, body_names: tuple[str, ...] | None = None) → torch.Tensor

Terminate when any keybody position error exceeds threshold (m).

Parameters:

• env – Manager-based RL env.

• command_name – Motion command term name.

• threshold – Max allowed L2 position error per body.

• body_names – Subset of command bodies (default: all).

wbc_mjlab.env.mdp.terminations.bad_motion_body_pos_z_only(env: ManagerBasedRlEnv, command_name: str, threshold: float, body_names: tuple[str, ...] | None = None) → torch.Tensor

Terminate when any keybody height error exceeds threshold (m).

Typically used for end-effector bodies (EE height termination).

wbc_mjlab.env.mdp.terminations.excessive_contact_force(env: ManagerBasedRlEnv, sensor_name: str, force_threshold: float = 10.0) → torch.Tensor

Terminate when any contact slot on sensor_name exceeds force_threshold (N).

wbc_mjlab.env.mdp.terminations.excessive_keybody_ground_contact_force(env: ManagerBasedRlEnv, sensor_name: str, body_names: tuple[str, ...], body_slot_order: tuple[str, ...], force_threshold: float) → torch.Tensor

Terminate when listed keybodies exceed force_threshold on the ground sensor.

Events (assistive wrench)

class wbc_mjlab.env.mdp.assistive_wrench.AssistiveWrenchEvent

Bases: object

Apply a bin-coupled virtual spatial wrench at the anchor body each step.

wbc_mjlab.env.mdp.assistive_wrench.assistive_wrench_force(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Privileged observation: assistive force applied at the anchor (world frame).

wbc_mjlab.env.mdp.assistive_wrench.assistive_wrench_torque(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Privileged observation: assistive torque applied at the anchor (world frame).

wbc_mjlab.env.mdp.assistive_wrench.assistive_wrench_gain(env: ManagerBasedRlEnv, command_name: str) → torch.Tensor

Privileged observation: per-episode assistive gain β.