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

class wbc_mjlab.env.mdp.actions.DefaultOffsetJointPositionActionCfg

Bases: JointPositionActionCfg

Default-relative joint targets; seed a from q_ref on RSI resample.

Live command stays deployable: q_cmd = q_default + scale · a. After motion resample, a is set to (q_ref - q_default) / scale so the first PD target matches the clip. Deploy uses the same formula with measured joints instead of q_ref.

command_name: str = 'motion'

build(env: ManagerBasedRlEnv) → DefaultOffsetJointPositionAction

class wbc_mjlab.env.mdp.actions.DefaultOffsetJointPositionAction

Bases: JointPositionAction

q_cmd = q_default + scale · a with reset warm-start from q_ref.

cfg: DefaultOffsetJointPositionActionCfg

seed_from_motion_reference(env_ids: torch.Tensor) → None

Set raw actions so q_default + scale · a = q_ref (resample only).

Motion command

class wbc_mjlab.env.mdp.commands.MotionLoader

Bases: object

Converted motion library: a bundled NPZ or a dataset directory of npz/*.npz clips.

The NPZ path of the command, kept for converted datasets and bundles; a dataset with source clips is played online (wbc_mjlab.motion.library.OnlineMotion).

frame(steps: torch.Tensor) → MotionFrame

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

reference() → MotionFrame

Reference at the current time_steps: computed once per change of the frame indices (any write to time_steps bumps its version), then shared by every term.

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 | torch.Tensor, env_ids: torch.Tensor | None = None) → 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

a dataset folder of source clips (CSV / PKL, at its root or in raw/) played online, or converted NPZ (a .npz bundle or a folder of npz/*.npz).

Type:

Motion library

motion_source: Literal['auto', 'online', 'npz'] = 'auto'

online when the folder has source clips, else its converted NPZ.

Type:

auto

mirror: bool = False

add a left/right mirrored copy of every clip (virtual segments, no storage).

Type:

Online

debias_z: bool = False

shift each clip vertically so its lowest foot sits at the robot’s sole height.

Type:

Online

motion_storage: Literal['device', 'host'] = 'device'

source frames on the env device, or in pinned host memory (huge datasets).

Type:

Online

source_format: str | None = None

source format name (motion_formats); inferred from the extension if unset.

Type:

Online

source_fps: float | None = None

the format’s / file’s).

Type:

Online

Type:

override the source frame rate (default

robot_id: str = ''

Robot of the motion (symmetry for mirror, feet for debias_z); set per robot.

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.

ground_reference: bool = False

Shift each episode’s reference so its start-frame anchor XY sits on the env origin. Needed on generated terrain, where env origins are sub-terrain centers and clips recorded far from (0, 0) would otherwise start off-tile.

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, body_names: tuple[str, ...] | None = None) → torch.Tensor

Reference keybody positions in the robot anchor frame.

Parameters:

body_names – Subset of motion command bodies (default: all), in the given order.

wbc_mjlab.env.mdp.observations.ref_body_ori_b(env: ManagerBasedRlEnv, command_name: str, body_names: tuple[str, ...] | None = None) → torch.Tensor

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

Parameters:

body_names – Subset of motion command bodies (default: all), in the given order.

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 β.