from __future__ import annotations
from typing import TYPE_CHECKING, cast
import torch
from mjlab.entity import Entity
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.sensor import ContactSensor
from mjlab.utils.lab_api.math import (
matrix_from_quat,
quat_apply_inverse,
quat_box_minus,
subtract_frame_transforms,
)
from wbc_mjlab.env.mdp.torque_envelope import torque_speed_limits
from .commands import MotionCommand
if TYPE_CHECKING:
from mjlab.envs import ManagerBasedRlEnv
_DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
def _motion_command(env: ManagerBasedRlEnv, command_name: str) -> MotionCommand:
return cast(MotionCommand, env.command_manager.get_term(command_name))
# --- Actor reference features (configurable obs terms; were stacked in MotionCommand.command) ---
# SE actor layouts: ``apply_se_actor`` in ``presets/se_actor.py``.
[docs]
def ref_base_height(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference anchor height relative to env origin (z_I r̂_IB)."""
return _motion_command(env, command_name).ref_base_height
[docs]
def ref_anchor_pos_w(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference anchor xyz position relative to env origin (SE reference command)."""
command = _motion_command(env, command_name)
return command.anchor_pos_w - env.scene.env_origins
[docs]
def ref_base_lin_vel_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference anchor linear velocity in anchor frame (B v̂_IB)."""
return _motion_command(env, command_name).ref_base_lin_vel_b
[docs]
def ref_base_ang_vel_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference anchor angular velocity in anchor frame (B ω̂_IB)."""
return _motion_command(env, command_name).ref_base_ang_vel_b
[docs]
def ref_gravity_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference gravity in anchor frame (B ĝ_I)."""
return _motion_command(env, command_name).ref_gravity_b
[docs]
def ref_joint_pos(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference joint positions for tracked DoFs (absolute)."""
return _motion_command(env, command_name).tracked_joint_pos
def _body_index(asset: Entity, body_name: str) -> int:
return asset.body_names.index(body_name)
def torso_ang_vel(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
body_name: str = "torso_link",
) -> torch.Tensor:
"""Torso angular velocity in the torso frame (Zest Table S3: T ω_IT)."""
asset: Entity = env.scene[asset_cfg.name]
idx = _body_index(asset, body_name)
quat_w = asset.data.body_link_quat_w[:, idx]
ang_vel_w = asset.data.body_link_ang_vel_w[:, idx]
return quat_apply_inverse(quat_w, ang_vel_w)
def torso_projected_gravity(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
body_name: str = "torso_link",
) -> torch.Tensor:
"""Gravity direction in the torso frame (Zest Table S3: T g_I)."""
asset: Entity = env.scene[asset_cfg.name]
idx = _body_index(asset, body_name)
quat_w = asset.data.body_link_quat_w[:, idx]
return quat_apply_inverse(quat_w, asset.data.gravity_vec_w)
[docs]
def ref_joint_vel(
env: ManagerBasedRlEnv,
command_name: str,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Reference joint velocities (critic privileged)."""
command = _motion_command(env, command_name)
if asset_cfg.joint_ids is not None:
return command.joint_vel[:, asset_cfg.joint_ids]
return command.tracked_joint_vel
[docs]
def ref_anchor_ori_6d(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference anchor orientation as 6D rotation matrix columns (world frame)."""
command = _motion_command(env, command_name)
mat = matrix_from_quat(command.anchor_quat_w)
return mat[..., :2].reshape(mat.shape[0], -1)
def ref_base_lin_acc_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference anchor linear acceleration in anchor frame (critic privileged)."""
return _motion_command(env, command_name).ref_base_lin_acc_b
def ref_base_ang_acc_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference anchor angular acceleration in anchor frame (critic privileged)."""
return _motion_command(env, command_name).ref_base_ang_acc_b
# --- Critic privileged keybody / anchor-relative features ---
def _body_lin_vel_in_anchor_frame(
anchor_quat_w: torch.Tensor,
body_lin_vel_w: torch.Tensor,
) -> torch.Tensor:
num_envs, num_bodies, _ = body_lin_vel_w.shape
quat = anchor_quat_w[:, None, :].expand(num_envs, num_bodies, 4).reshape(-1, 4)
vel = body_lin_vel_w.reshape(-1, 3)
vel_b = quat_apply_inverse(quat, vel)
return vel_b.view(num_envs, num_bodies * 3)
def _body_ang_vel_in_anchor_frame(
anchor_quat_w: torch.Tensor,
body_ang_vel_w: torch.Tensor,
) -> torch.Tensor:
num_envs, num_bodies, _ = body_ang_vel_w.shape
quat = anchor_quat_w[:, None, :].expand(num_envs, num_bodies, 4).reshape(-1, 4)
vel = body_ang_vel_w.reshape(-1, 3)
vel_b = quat_apply_inverse(quat, vel)
return vel_b.view(num_envs, num_bodies * 3)
[docs]
def motion_anchor_pos_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Full xyz anchor tracking error in robot anchor frame (non-SE layouts)."""
command = _motion_command(env, command_name)
pos, _ = subtract_frame_transforms(
command.robot_anchor_pos_w,
command.robot_anchor_quat_w,
command.anchor_pos_w,
command.anchor_quat_w,
)
return pos.view(env.num_envs, -1)
def motion_anchor_pos_z_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Z-only anchor tracking error in robot anchor frame (SE measurement)."""
return motion_anchor_pos_b(env, command_name)[:, 2:3]
[docs]
def motion_anchor_ori_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
command = _motion_command(env, command_name)
_, ori = subtract_frame_transforms(
command.robot_anchor_pos_w,
command.robot_anchor_quat_w,
command.anchor_pos_w,
command.anchor_quat_w,
)
mat = matrix_from_quat(ori)
return mat[..., :2].reshape(mat.shape[0], -1)
[docs]
def motion_anchor_pos_error_w(
env: ManagerBasedRlEnv, command_name: str
) -> torch.Tensor:
"""World-frame anchor position tracking error (ref − robot)."""
command = _motion_command(env, command_name)
return command.anchor_pos_w - command.robot_anchor_pos_w
[docs]
def motion_anchor_ori_error(
env: ManagerBasedRlEnv, command_name: str
) -> torch.Tensor:
"""Anchor orientation tracking error as axis-angle (3); ‖·‖ = quat_error_magnitude."""
command = _motion_command(env, command_name)
return quat_box_minus(command.anchor_quat_w, command.robot_anchor_quat_w)
[docs]
def robot_body_pos_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
command = _motion_command(env, command_name)
num_bodies = len(command.cfg.body_names)
pos_b, _ = subtract_frame_transforms(
command.robot_anchor_pos_w[:, None, :].repeat(1, num_bodies, 1),
command.robot_anchor_quat_w[:, None, :].repeat(1, num_bodies, 1),
command.robot_body_pos_w,
command.robot_body_quat_w,
)
return pos_b.view(env.num_envs, -1)
[docs]
def robot_body_ori_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
command = _motion_command(env, command_name)
num_bodies = len(command.cfg.body_names)
_, ori_b = subtract_frame_transforms(
command.robot_anchor_pos_w[:, None, :].repeat(1, num_bodies, 1),
command.robot_anchor_quat_w[:, None, :].repeat(1, num_bodies, 1),
command.robot_body_pos_w,
command.robot_body_quat_w,
)
mat = matrix_from_quat(ori_b)
return mat[..., :2].reshape(mat.shape[0], -1)
[docs]
def ref_body_pos_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference keybody positions in the robot anchor frame."""
command = _motion_command(env, command_name)
num_bodies = len(command.cfg.body_names)
pos_b, _ = subtract_frame_transforms(
command.robot_anchor_pos_w[:, None, :].repeat(1, num_bodies, 1),
command.robot_anchor_quat_w[:, None, :].repeat(1, num_bodies, 1),
command.body_pos_w,
command.body_quat_w,
)
return pos_b.view(env.num_envs, -1)
[docs]
def ref_body_ori_b(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference keybody orientations in the robot anchor frame (6D rotation columns)."""
command = _motion_command(env, command_name)
num_bodies = len(command.cfg.body_names)
_, ori_b = subtract_frame_transforms(
command.robot_anchor_pos_w[:, None, :].repeat(1, num_bodies, 1),
command.robot_anchor_quat_w[:, None, :].repeat(1, num_bodies, 1),
command.body_pos_w,
command.body_quat_w,
)
mat = matrix_from_quat(ori_b)
return mat[..., :2].reshape(mat.shape[0], -1)
def motion_body_lin_vel(
env: ManagerBasedRlEnv, command_name: str
) -> torch.Tensor:
"""Actual keybody linear velocities in the robot anchor frame."""
command = _motion_command(env, command_name)
return _body_lin_vel_in_anchor_frame(
command.robot_anchor_quat_w, command.robot_body_lin_vel_w
)
def motion_body_ang_vel(
env: ManagerBasedRlEnv, command_name: str
) -> torch.Tensor:
"""Actual keybody angular velocities in the robot anchor frame."""
command = _motion_command(env, command_name)
return _body_ang_vel_in_anchor_frame(
command.robot_anchor_quat_w, command.robot_body_ang_vel_w
)
def ref_body_lin_vel(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference keybody linear velocities in the robot anchor frame."""
command = _motion_command(env, command_name)
return _body_lin_vel_in_anchor_frame(
command.robot_anchor_quat_w, command.body_lin_vel_w
)
def ref_body_ang_vel(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Reference keybody angular velocities in the robot anchor frame."""
command = _motion_command(env, command_name)
return _body_ang_vel_in_anchor_frame(
command.robot_anchor_quat_w, command.body_ang_vel_w
)
def keybody_contact_forces(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
"""Log-scaled contact forces for tracked keybodies (critic privileged)."""
sensor: ContactSensor = env.scene[sensor_name]
sensor_data = sensor.data
assert sensor_data.force is not None
forces_flat = sensor_data.force.flatten(start_dim=1)
return torch.sign(forces_flat) * torch.log1p(torch.abs(forces_flat))
# --- Actuation (G1 Unitree envelope) ---
def joint_mechanical_power(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Per-joint mechanical power τ·ω (actuator generalized force)."""
asset: Entity = env.scene[asset_cfg.name]
ids = asset_cfg.joint_ids
tau = asset.data.qfrc_actuator[:, ids]
qd = asset.data.joint_vel[:, ids]
return tau * qd
def joint_torque_envelope_ratio(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""|τ| / max admissible torque magnitude at current velocity (≈1 at envelope)."""
from wbc_mjlab.robots.g1.envelope import g1_joint_envelope_tensors
asset: Entity = env.scene[asset_cfg.name]
ids = asset_cfg.joint_ids
tau = asset.data.qfrc_actuator[:, ids]
qd = asset.data.joint_vel[:, ids]
envl = g1_joint_envelope_tensors(env, asset_cfg)
tau_low, tau_high = torque_speed_limits(qd, envl)
cap = torch.maximum(tau_high, -tau_low).clamp(min=1.0e-6)
return (tau.abs() / cap).flatten(start_dim=1)
def joint_torque_envelope_margin(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Normalized distance to torque envelope: 0 on boundary, negative if inside."""
from wbc_mjlab.robots.g1.envelope import g1_joint_envelope_tensors
asset: Entity = env.scene[asset_cfg.name]
ids = asset_cfg.joint_ids
tau = asset.data.qfrc_actuator[:, ids]
qd = asset.data.joint_vel[:, ids]
envl = g1_joint_envelope_tensors(env, asset_cfg)
tau_low, tau_high = torque_speed_limits(qd, envl)
cap = torch.maximum(tau_high, -tau_low).clamp(min=1.0e-6)
margin_high = (tau_high - tau) / cap
margin_low = (tau - tau_low) / cap
return torch.minimum(margin_high, margin_low).flatten(start_dim=1)
[docs]
def motion_segment_phase(env: ManagerBasedRlEnv, command_name: str) -> torch.Tensor:
"""Normalized clip progress in ``[0, 1]`` (privileged critic feature)."""
command = _motion_command(env, command_name)
seg_start = command.motion.segment_start_idx[command.trajectory_ids]
seg_len = command.motion.segment_length[command.trajectory_ids].float().clamp(min=1.0)
local_step = (command.time_steps - seg_start).float()
phase = local_step / seg_len
return phase.clamp(0.0, 1.0).unsqueeze(-1)
[docs]
def motion_tracking_step_rewards(
env: ManagerBasedRlEnv, command_name: str
) -> torch.Tensor:
"""Unweighted per-step ``motion_*`` reward values (privileged critic feature)."""
if not hasattr(env, "reward_manager"):
n_terms = sum(1 for name in env.cfg.rewards if name.startswith("motion_"))
return torch.zeros(env.num_envs, n_terms, device=env.device)
command = _motion_command(env, command_name)
if not hasattr(command, "_critic_motion_reward_indices"):
reward_manager = env.reward_manager
command._critic_motion_reward_indices = [
idx
for idx, name in enumerate(reward_manager._term_names)
if name.startswith("motion_")
]
indices = command._critic_motion_reward_indices
return env.reward_manager._step_reward[:, indices]
