From f60bc1d9ed7b54392b14f4009d424fdf03dc5e7d Mon Sep 17 00:00:00 2001 From: caozx1110 Date: Thu, 9 Jul 2026 18:35:49 +0800 Subject: [PATCH 1/4] refactor: add shared motion tracking common modules --- .../envs/motion_tracking/common/__init__.py | 8 + .../envs/motion_tracking/common/config.py | 170 ++++++ .../common/domain_randomization.py | 186 ++++++ .../{g1 => common}/motion_loader.py | 0 .../numba.py} | 14 +- .../motion_tracking/common/observations.py | 279 +++++++++ .../envs/motion_tracking/common/reset.py | 83 +++ .../envs/motion_tracking/common/rewards.py | 357 +++++++++++ .../motion_tracking/common/terminations.py | 69 +++ .../envs/motion_tracking/common/tracking.py | 576 ++++++++++++++++++ .../envs/motion_tracking/common/transforms.py | 103 ++++ 11 files changed, 1838 insertions(+), 7 deletions(-) create mode 100644 src/unilab/envs/motion_tracking/common/__init__.py create mode 100644 src/unilab/envs/motion_tracking/common/config.py create mode 100644 src/unilab/envs/motion_tracking/common/domain_randomization.py rename src/unilab/envs/motion_tracking/{g1 => common}/motion_loader.py (100%) rename src/unilab/envs/motion_tracking/{g1/motion_tracking_numba.py => common/numba.py} (98%) create mode 100644 src/unilab/envs/motion_tracking/common/observations.py create mode 100644 src/unilab/envs/motion_tracking/common/reset.py create mode 100644 src/unilab/envs/motion_tracking/common/rewards.py create mode 100644 src/unilab/envs/motion_tracking/common/terminations.py create mode 100644 src/unilab/envs/motion_tracking/common/tracking.py create mode 100644 src/unilab/envs/motion_tracking/common/transforms.py diff --git a/src/unilab/envs/motion_tracking/common/__init__.py b/src/unilab/envs/motion_tracking/common/__init__.py new file mode 100644 index 000000000..01d622066 --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/__init__.py @@ -0,0 +1,8 @@ +"""Robot-agnostic motion-tracking engine and owner modules. + +This package holds the motion-tracking task engine (:class:`MotionTrackingEnv` +/ :class:`MotionTrackingDeployEnv`) and its per-concern owner modules (config, +rewards, observations, terminations, transforms, reset, domain randomization, +numba hot path, motion loading). Per-robot profiles live under ``g1/`` and +``x2/`` and only carry robot-specific defaults and thin registry subclasses. +""" diff --git a/src/unilab/envs/motion_tracking/common/config.py b/src/unilab/envs/motion_tracking/common/config.py new file mode 100644 index 000000000..0fe264501 --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/config.py @@ -0,0 +1,170 @@ +"""Configuration dataclasses for the robot-agnostic motion-tracking engine. + +The default field values are kept at the historical G1 profile so that the +registered ``G1MotionTracking`` / ``G1MotionTrackingDeploy`` configs remain +field-identical to the pre-refactor monolith. Per-robot profiles (g1/x2) +override these via subclasses. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base.scene import SceneCfg +from unilab.envs.locomotion.g1.base import G1BaseCfg + +from .rewards import RewardConfig + + +@dataclass +class PoseRandomization: + """Pose randomization ranges for reset.""" + + x: tuple[float, float] = (-0.05, 0.05) + y: tuple[float, float] = (-0.05, 0.05) + z: tuple[float, float] = (-0.01, 0.01) + roll: tuple[float, float] = (-0.1, 0.1) + pitch: tuple[float, float] = (-0.1, 0.1) + yaw: tuple[float, float] = (-0.2, 0.2) + + +@dataclass +class VelocityRandomization: + """Velocity randomization ranges for reset.""" + + x: tuple[float, float] = (-0.5, 0.5) + y: tuple[float, float] = (-0.5, 0.5) + z: tuple[float, float] = (-0.2, 0.2) + roll: tuple[float, float] = (-0.52, 0.52) + pitch: tuple[float, float] = (-0.52, 0.52) + yaw: tuple[float, float] = (-0.78, 0.78) + + +@dataclass +class DomainRand: + """Domain randomization config required by motrix backend hooks.""" + + randomize_base_mass: bool = False + added_mass_range: list[float] = field(default_factory=lambda: [-1.5, 1.5]) + + random_com: bool = False + com_offset_x: list[float] = field(default_factory=lambda: [-0.05, 0.05]) + com_offset_y: list[float] = field(default_factory=lambda: [-0.05, 0.05]) + com_offset_z: list[float] = field(default_factory=lambda: [-0.05, 0.05]) + + randomize_gravity: bool = False + gravity_range: list[list[float]] = field( + default_factory=lambda: [[0.0, 0.0, -9.81], [0.0, 0.0, -9.81]] + ) + + push_robots: bool = False + push_interval: int = 750 + max_force: list[float] = field(default_factory=lambda: [1.0, 1.0, 0.5]) + push_body_name: str | None = None + + randomize_kp: bool = False + kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) + + randomize_kd: bool = False + kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) + + randomize_geom_friction: bool = False + friction_range: list[float] = field(default_factory=lambda: [0.3, 1.2]) + friction_geom_pattern: str = r"^(left|right)_foot[1-7]_collision$" + + randomize_joint_default_pos: bool = False + joint_default_pos_range: list[float] = field(default_factory=lambda: [-0.01, 0.01]) + + +# Backward-compatible alias for the historical (snake_case) class name. +Domain_Rand = DomainRand + + +def _zero_pose_randomization() -> PoseRandomization: + return PoseRandomization( + x=(0.0, 0.0), + y=(0.0, 0.0), + z=(0.0, 0.0), + roll=(0.0, 0.0), + pitch=(0.0, 0.0), + yaw=(0.0, 0.0), + ) + + +def _zero_velocity_randomization() -> VelocityRandomization: + return VelocityRandomization( + x=(0.0, 0.0), + y=(0.0, 0.0), + z=(0.0, 0.0), + roll=(0.0, 0.0), + pitch=(0.0, 0.0), + yaw=(0.0, 0.0), + ) + + +@dataclass +class MotionTrackingCfg(G1BaseCfg): + """Configuration for the motion tracking environment.""" + + scene: SceneCfg = field( + default_factory=lambda: SceneCfg( + model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") + ) + ) + # Kept at the historical single-clip default for backward compatibility. + motion_file: str | list[str] = str( + ASSETS_ROOT_PATH / "motions" / "g1" / "dance1_subject2_part.npz" + ) + # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "gangnam_style.npz") + # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "fight1_subject5_from_csv.npz") #LAFAN + # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "dance_basic_slide_180_R_loop_001__A322_M.npz") #LAFAN + # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "playing_violin_R_003__A327_from_csv.npz") #Seed + anchor_body_name: str = "torso_link" + body_names: tuple[str, ...] = ( + "pelvis", + "left_hip_roll_link", + "left_knee_link", + "left_ankle_roll_link", + "right_hip_roll_link", + "right_knee_link", + "right_ankle_roll_link", + "torso_link", + "left_shoulder_roll_link", + "left_elbow_link", + "left_wrist_yaw_link", + "right_shoulder_roll_link", + "right_elbow_link", + "right_wrist_yaw_link", + ) + sampling_mode: Literal["start", "clip_start", "uniform", "adaptive", "mixed"] = "adaptive" + sampling_start_ratio: float = 0.0 + truncate_on_clip_end: bool = False + max_episode_seconds: float = 10.0 + reward_config: RewardConfig = field(default_factory=RewardConfig) + pose_randomization: PoseRandomization = field(default_factory=PoseRandomization) + velocity_randomization: VelocityRandomization = field(default_factory=VelocityRandomization) + domain_rand: DomainRand = field(default_factory=DomainRand) + joint_position_range: tuple[float, float] = (-0.1, 0.1) + # Termination thresholds + anchor_pos_z_threshold: float = 0.25 + anchor_ori_threshold: float = 0.8 + ee_body_pos_z_threshold: float = 0.25 + ee_body_names: tuple[str, ...] = ( + "left_ankle_roll_link", + "right_ankle_roll_link", + "left_wrist_yaw_link", + "right_wrist_yaw_link", + ) + undesired_contact_z_threshold: float = 0.05 + terminate_on_undesired_contacts: bool = False + numba_acceleration: bool = False + numba_num_threads: int | None = None + + +@dataclass +class MotionTrackingDeployEnvCfg(MotionTrackingCfg): + """Base deploy configuration for motion tracking.""" + + pass diff --git a/src/unilab/envs/motion_tracking/common/domain_randomization.py b/src/unilab/envs/motion_tracking/common/domain_randomization.py new file mode 100644 index 000000000..aa74e50e8 --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/domain_randomization.py @@ -0,0 +1,186 @@ +"""Domain-randomization reset/observation provider for motion tracking.""" + +from __future__ import annotations + +import time +from typing import Any, cast + +import numpy as np + +from unilab.dr import ( + DomainRandomizationCapabilities, + DomainRandomizationProvider, + IntervalRandomizationPlan, + ResetPlan, +) +from unilab.dr.dr_utils import ( + build_common_reset_randomization, + build_interval_push_plan, + validate_common_reset_randomization, + validate_interval_push_support, + zero_actions, +) +from unilab.dr.types import RESET_TERM_GEOM_FRICTION, ResetRandomizationPayload +from unilab.dtype_config import get_global_dtype + +from .reset import build_motion_reference_state + + +class MotionTrackingDomainRandomizationProvider(DomainRandomizationProvider): + def __init__( + self, + *, + base_kp: np.ndarray | None = None, + base_kd: np.ndarray | None = None, + base_geom_friction: np.ndarray | None = None, + foot_geom_ids: np.ndarray | None = None, + ) -> None: + self._base_kp = base_kp + self._base_kd = base_kd + self._base_geom_friction = base_geom_friction + self._foot_geom_ids = foot_geom_ids + self._last_reset_observation_timing_ms: dict[str, float] = {} + + @property + def last_reset_observation_timing_ms(self) -> dict[str, float]: + return dict(self._last_reset_observation_timing_ms) + + def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: + validate_common_reset_randomization( + env, capabilities, base_kp=self._base_kp, base_kd=self._base_kd + ) + validate_interval_push_support(env, capabilities) + if getattr(env.cfg.domain_rand, "randomize_geom_friction", False): + if not capabilities.supports_reset_term(RESET_TERM_GEOM_FRICTION): + raise NotImplementedError( + f"{env._backend.backend_type} backend does not support " + "geom-friction reset randomization" + ) + if ( + self._base_geom_friction is None + or self._foot_geom_ids is None + or self._foot_geom_ids.size == 0 + ): + raise ValueError("randomize_geom_friction=True but provider has no foot geom IDs") + + def build_interval_randomization_plan( + self, env: Any, step_counter: int + ) -> IntervalRandomizationPlan | None: + return build_interval_push_plan(env, step_counter) + + def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: + num_reset = len(env_ids) + motion_frames = env.motion_sampler.sample_frames(env_ids) + motion_data = env.motion_loader.get_motion_at_frame(motion_frames) + qpos, qvel = build_motion_reference_state(env, env_ids, motion_data) + + info_updates = { + "current_actions": zero_actions(num_reset, env._num_action), + "last_actions": zero_actions(num_reset, env._num_action), + } + randomization = build_common_reset_randomization( + env, num_reset, base_kp=self._base_kp, base_kd=self._base_kd + ) + + dr_cfg = env.cfg.domain_rand + if getattr(dr_cfg, "randomize_geom_friction", False): + assert self._base_geom_friction is not None + assert self._foot_geom_ids is not None + payload = randomization or ResetRandomizationPayload() + low, high = dr_cfg.friction_range + scale = np.random.uniform(low, high, size=(num_reset, 1)).astype(np.float64) + geom_friction = np.broadcast_to( + self._base_geom_friction, + (num_reset, *self._base_geom_friction.shape), + ).copy() + geom_friction[:, self._foot_geom_ids, 0] = scale * np.ones( + (1, self._foot_geom_ids.size) + ) + payload.geom_friction = geom_friction + randomization = payload + + if getattr(dr_cfg, "randomize_joint_default_pos", False): + low, high = dr_cfg.joint_default_pos_range + info_updates["default_dof_pos_bias"] = np.random.uniform( + low, high, size=(num_reset, env._num_action) + ).astype(get_global_dtype()) + + return ResetPlan( + env_ids=env_ids, + qpos=qpos, + qvel=qvel, + info_updates=info_updates, + randomization=randomization, + ) + + def build_reset_observation( + self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any] + ) -> dict[str, np.ndarray]: + obs_t0 = time.perf_counter() + + t0 = time.perf_counter() + motion_data = env.motion_loader.get_motion_at_frame( + env.motion_sampler.current_frames[env_ids] + ) + motion_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + linvel = env.get_local_linvel()[env_ids] + linvel_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + gyro = env.get_gyro()[env_ids] + gyro_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + dof_pos = env.get_dof_pos()[env_ids] + dof_pos_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + dof_vel = env.get_dof_vel()[env_ids] + dof_vel_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + robot_body_pos_w, robot_body_quat_w = env._backend.get_body_pose_w_rows( + env_ids, env.body_ids + ) + body_pose_ms = (time.perf_counter() - t0) * 1000.0 + + obs_info = dict(info_updates) + default_dof_pos_bias = info_updates.get("default_dof_pos_bias") + if isinstance(default_dof_pos_bias, np.ndarray): + obs_info["default_dof_pos_bias"] = default_dof_pos_bias + obs_info["env_ids"] = env_ids + + t0 = time.perf_counter() + obs = cast( + dict[str, np.ndarray], + env._compute_obs( + obs_info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ), + ) + compute_obs_ms = (time.perf_counter() - t0) * 1000.0 + + total_ms = (time.perf_counter() - obs_t0) * 1000.0 + getters_ms = linvel_ms + gyro_ms + dof_pos_ms + dof_vel_ms + body_pose_ms + self._last_reset_observation_timing_ms = { + "dr_reset_observation_getters_ms": getters_ms, + "dr_reset_obs_get_motion_ms": motion_ms, + "dr_reset_obs_get_local_linvel_ms": linvel_ms, + "dr_reset_obs_get_gyro_ms": gyro_ms, + "dr_reset_obs_get_dof_pos_ms": dof_pos_ms, + "dr_reset_obs_get_dof_vel_ms": dof_vel_ms, + "dr_reset_obs_get_body_pose_ms": body_pose_ms, + "dr_reset_observation_compute_obs_ms": compute_obs_ms, + "dr_reset_observation_internal_gap_ms": ( + total_ms - motion_ms - getters_ms - compute_obs_ms + ), + } + return obs diff --git a/src/unilab/envs/motion_tracking/g1/motion_loader.py b/src/unilab/envs/motion_tracking/common/motion_loader.py similarity index 100% rename from src/unilab/envs/motion_tracking/g1/motion_loader.py rename to src/unilab/envs/motion_tracking/common/motion_loader.py diff --git a/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py b/src/unilab/envs/motion_tracking/common/numba.py similarity index 98% rename from src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py rename to src/unilab/envs/motion_tracking/common/numba.py index e8b988655..850caa83d 100644 --- a/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py +++ b/src/unilab/envs/motion_tracking/common/numba.py @@ -912,7 +912,7 @@ def _compute_update_state_kernel( ) -class G1MotionTrackingNumbaAccelerator: +class MotionTrackingNumbaAccelerator: """Driver that keeps config-derived arrays and calls the fused kernel.""" def __init__( @@ -976,19 +976,19 @@ def __init__( @classmethod def from_env( cls, env: Any, num_threads: int | None = None - ) -> "G1MotionTrackingNumbaAccelerator": + ) -> "MotionTrackingNumbaAccelerator": if not NUMBA_AVAILABLE: raise RuntimeError( - "G1MotionTracking numba_acceleration=True requires numba. Install it or run " + "MotionTracking numba_acceleration=True requires numba. Install it or run " "through `uv run --with numba ...`; disable numba_acceleration to use the " "numpy path." ) unsupported = unsupported_terms(env._cfg.reward_config.scales) if unsupported: raise ValueError( - "G1MotionTracking Numba accelerator does not support active reward terms " + "MotionTracking Numba accelerator does not support active reward terms " f"{sorted(unsupported)}. Disable numba_acceleration or add these terms to " - "src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py." + "src/unilab/envs/motion_tracking/common/numba.py." ) default_angles = getattr(env, "default_angles", None) if default_angles is None: @@ -1042,9 +1042,9 @@ def _sync_scales(self, scales: Mapping[str, float]) -> None: unsupported = unsupported_terms(scales) if unsupported: raise ValueError( - "G1MotionTracking Numba accelerator does not support active reward terms " + "MotionTracking Numba accelerator does not support active reward terms " f"{sorted(unsupported)}. Disable numba_acceleration or add these terms to " - "src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py." + "src/unilab/envs/motion_tracking/common/numba.py." ) self.scale.fill(0.0) for name, value in scales.items(): diff --git a/src/unilab/envs/motion_tracking/common/observations.py b/src/unilab/envs/motion_tracking/common/observations.py new file mode 100644 index 000000000..65310807c --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/observations.py @@ -0,0 +1,279 @@ +"""Observation construction for motion tracking. + +Holds the robot-agnostic observation builders. The environment classes keep a +thin polymorphic method surface (``_compute_obs`` / ``_build_actor_obs`` / +``_write_body_*``) that delegates here so subclasses can still override obs +layout (SAC critic tail, box object obs, deploy mimic actor) while the core +math lives in one place. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from unilab.dtype_config import get_global_dtype +from unilab.utils.geometry import np_write_relative_anchor_transform_pos_rot6d + + +def actor_obs_dim(n: int) -> int: + return 3 + 6 + 3 + 3 + n * 5 + + +def critic_base_obs_dim(n: int) -> int: + return 3 + 6 + 3 + 3 + n * 5 + + +def mimic_actor_obs_dim(n: int) -> int: + # unitree_rl_lab mimic deploy actor input: + # motion_command(2n), motion_anchor_ori_b(6), gyro(3), joints, actions. + return 6 + 3 + n * 5 + + +def obs_groups_spec(env: Any) -> dict[str, int]: + # Actor: command(2n) + motion_anchor_pos_b(3) + motion_anchor_ori_b(6) + # + linvel(3) + gyro(3) + joint_pos(n) + joint_vel(n) + actions(n) + # Critic mirrors BeyondMimic physical terms without actor observation noise: + # command, motion anchor, robot body pos/ori, linvel, gyro, joints, actions. + n = env._num_action + actor_width = getattr(env, "_actor_obs_width", env._actor_obs_dim(n)) + critic_width = getattr( + env, + "_critic_obs_width", + env._critic_base_obs_dim(n) + len(env._cfg.body_names) * 9, + ) + return {"obs": actor_width, "critic": critic_width} + + +def build_actor_obs( + *, + actor_obs_dim: int, + command: np.ndarray, + motion_anchor_pos_b: np.ndarray, + motion_anchor_ori_b: np.ndarray, + noisy_linvel: np.ndarray, + noisy_gyro: np.ndarray, + noisy_joint_pos_rel: np.ndarray, + noisy_dof_vel: np.ndarray, + last_actions: np.ndarray, +) -> np.ndarray: + num_envs = command.shape[0] + n_action = noisy_joint_pos_rel.shape[1] + actor_obs = np.empty((num_envs, actor_obs_dim), dtype=get_global_dtype()) + offset = 0 + actor_obs[:, offset : offset + command.shape[1]] = command + offset += command.shape[1] + actor_obs[:, offset : offset + 3] = motion_anchor_pos_b + offset += 3 + actor_obs[:, offset : offset + 6] = motion_anchor_ori_b + offset += 6 + actor_obs[:, offset : offset + 3] = noisy_linvel + offset += 3 + actor_obs[:, offset : offset + 3] = noisy_gyro + offset += 3 + actor_obs[:, offset : offset + n_action] = noisy_joint_pos_rel + offset += n_action + actor_obs[:, offset : offset + n_action] = noisy_dof_vel + offset += n_action + actor_obs[:, offset : offset + n_action] = last_actions + return actor_obs + + +def build_mimic_actor_obs( + *, + command: np.ndarray, + motion_anchor_ori_b: np.ndarray, + noisy_gyro: np.ndarray, + noisy_joint_pos_rel: np.ndarray, + noisy_dof_vel: np.ndarray, + last_actions: np.ndarray, +) -> np.ndarray: + """unitree_rl_lab mimic deploy actor layout: 2n + 6 + 3 + n + n + n.""" + return np.concatenate( + [ + command, + motion_anchor_ori_b, + noisy_gyro, + noisy_joint_pos_rel, + noisy_dof_vel, + last_actions, + ], + axis=1, + dtype=get_global_dtype(), + ) + + +def write_body_pos_in_anchor_frame( + anchor_pos: np.ndarray, + anchor_quat: np.ndarray, + body_pos: np.ndarray, + out: np.ndarray, + *, + body_vec_error: np.ndarray, +) -> None: + aw = anchor_quat[:, None, 0] + ax = anchor_quat[:, None, 1] + ay = anchor_quat[:, None, 2] + az = anchor_quat[:, None, 3] + + num_envs, n_body = body_pos.shape[:2] + rel_pos = body_vec_error[:num_envs, :n_body] + + vx = rel_pos[..., 0] + vy = rel_pos[..., 1] + vz = rel_pos[..., 2] + np.subtract(body_pos[..., 0], anchor_pos[:, None, 0], out=vx) + np.subtract(body_pos[..., 1], anchor_pos[:, None, 1], out=vy) + np.subtract(body_pos[..., 2], anchor_pos[:, None, 2], out=vz) + + tx = 2 * (az * vy - ay * vz) + ty = 2 * (ax * vz - az * vx) + tz = 2 * (ay * vx - ax * vy) + + out[..., 0] = vx + aw * tx + az * ty - ay * tz + out[..., 1] = vy + aw * ty + ax * tz - az * tx + out[..., 2] = vz + aw * tz + ay * tx - ax * ty + + +def write_body_ori6_in_anchor_frame( + anchor_quat: np.ndarray, + body_quat: np.ndarray, + out: np.ndarray, +) -> None: + aw = anchor_quat[:, None, 0] + ax = anchor_quat[:, None, 1] + ay = anchor_quat[:, None, 2] + az = anchor_quat[:, None, 3] + bw = body_quat[..., 0] + bx = body_quat[..., 1] + by = body_quat[..., 2] + bz = body_quat[..., 3] + + rw = aw * bw + ax * bx + ay * by + az * bz + rx = aw * bx - ax * bw - ay * bz + az * by + ry = aw * by + ax * bz - ay * bw - az * bx + rz = aw * bz - ax * by + ay * bx - az * bw + + out[..., 0] = 1 - 2 * (ry * ry + rz * rz) + out[..., 1] = 2 * (rx * ry - rw * rz) + out[..., 2] = 2 * (rx * ry + rw * rz) + out[..., 3] = 1 - 2 * (rx * rx + rz * rz) + out[..., 4] = 2 * (rx * rz - rw * ry) + out[..., 5] = 2 * (ry * rz + rw * rx) + + +def compute_obs( + env: Any, + info: dict, + motion_data: Any, + linvel: np.ndarray, + gyro: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, +) -> dict[str, np.ndarray]: + """Compute observations as dict with actor and critic groups.""" + num_envs = linvel.shape[0] + dtype = get_global_dtype() + n_action = dof_pos.shape[1] + n_body = env._n_motion_bodies + + # Get anchor states + anchor_pos_w = motion_data.body_pos_w[:, env.anchor_body_idx] + anchor_quat_w = motion_data.body_quat_w[:, env.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, env.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, env.anchor_body_idx] + + # Motion anchor pose in robot frame + if num_envs == env._num_envs: + motion_anchor_pos_b = env._motion_anchor_pos_b + motion_anchor_ori_b = env._motion_anchor_ori_b + joint_pos_rel = env._joint_pos_rel + zero_actions = env._zero_actions + else: + motion_anchor_pos_b = np.empty((num_envs, 3), dtype=dtype) + motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) + joint_pos_rel = np.empty((num_envs, n_action), dtype=dtype) + zero_actions = np.zeros((num_envs, n_action), dtype=dtype) + np_write_relative_anchor_transform_pos_rot6d( + robot_anchor_pos_w, + robot_anchor_quat_w, + anchor_pos_w, + anchor_quat_w, + motion_anchor_pos_b, + motion_anchor_ori_b, + ) + + # Joint positions and velocities + bias = info.get("default_dof_pos_bias") + effective_default = env.default_angles + bias if bias is not None else env.default_angles + np.subtract(dof_pos, effective_default, out=joint_pos_rel) + last_actions = info.get("current_actions") + if not isinstance(last_actions, np.ndarray): + last_actions = zero_actions + + if num_envs == env._num_envs: + command = env._motion_command + else: + command = np.empty((num_envs, n_action * 2), dtype=dtype) + command[:, :n_action] = motion_data.joint_pos + command[:, n_action : n_action * 2] = motion_data.joint_vel + + # Per-step observation noise on sensor channels (actor only). + # Critic uses the clean originals — asymmetric actor–critic contract. + noise_cfg = env._cfg.noise_config + noise_enabled = noise_cfg.level > 0.0 + if noise_enabled: + linvel_actor = env._obs_noise(linvel, noise_cfg.scale_linvel) + gyro_actor = env._obs_noise(gyro, noise_cfg.scale_gyro) + joint_pos_actor = env._obs_noise(joint_pos_rel, noise_cfg.scale_joint_angle) + dof_vel_actor = env._obs_noise(dof_vel, noise_cfg.scale_joint_vel) + else: + linvel_actor = linvel + gyro_actor = gyro + joint_pos_actor = joint_pos_rel + dof_vel_actor = dof_vel + + # Actor observations (noisy proprioception) + actor_obs = env._build_actor_obs( + command=command, + motion_anchor_pos_b=motion_anchor_pos_b, + motion_anchor_ori_b=motion_anchor_ori_b, + noisy_linvel=linvel_actor, + noisy_gyro=gyro_actor, + noisy_joint_pos_rel=joint_pos_actor, + noisy_dof_vel=dof_vel_actor, + last_actions=last_actions, + ) + + # Critic observations (clean proprioception + privileged body transforms) + critic_obs = np.empty((num_envs, env._critic_obs_width), dtype=dtype) + offset = 0 + critic_obs[:, offset : offset + command.shape[1]] = command + offset += command.shape[1] + critic_obs[:, offset : offset + 3] = motion_anchor_pos_b + offset += 3 + critic_obs[:, offset : offset + 6] = motion_anchor_ori_b + offset += 6 + critic_obs[:, offset : offset + 3] = linvel + offset += 3 + critic_obs[:, offset : offset + 3] = gyro + offset += 3 + critic_obs[:, offset : offset + n_action] = joint_pos_rel + offset += n_action + critic_obs[:, offset : offset + n_action] = dof_vel + offset += n_action + critic_obs[:, offset : offset + n_action] = last_actions + offset += n_action + robot_body_pos_b = critic_obs[:, offset : offset + n_body * 3].reshape(num_envs, n_body, 3) + env._write_body_pos_in_anchor_frame( + robot_anchor_pos_w, robot_anchor_quat_w, robot_body_pos_w, robot_body_pos_b + ) + offset += n_body * 3 + robot_body_ori_b = critic_obs[:, offset : offset + n_body * 6].reshape(num_envs, n_body, 6) + env._write_body_ori6_in_anchor_frame( + robot_anchor_quat_w, robot_body_quat_w, robot_body_ori_b + ) + return {"obs": actor_obs, "critic": critic_obs} diff --git a/src/unilab/envs/motion_tracking/common/reset.py b/src/unilab/envs/motion_tracking/common/reset.py new file mode 100644 index 000000000..bea6b99e2 --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/reset.py @@ -0,0 +1,83 @@ +"""Reset-state construction for motion tracking.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from unilab.dtype_config import get_global_dtype +from unilab.utils.geometry import np_sample_uniform +from unilab.utils.rotation import np_quat_apply, np_quat_from_euler_xyz, np_quat_inv, np_quat_mul + +from .motion_loader import MotionData + + +def build_motion_reference_state( + env: Any, env_ids: np.ndarray, motion_data: MotionData +) -> tuple[np.ndarray, np.ndarray]: + dtype = get_global_dtype() + num_reset = len(env_ids) + + root_pos = motion_data.body_pos_w[:, 0].copy() + root_ori = motion_data.body_quat_w[:, 0].copy() + root_lin_vel = motion_data.body_lin_vel_w[:, 0].copy() + root_ang_vel = motion_data.body_ang_vel_w[:, 0].copy() + joint_pos = motion_data.joint_pos.copy() + joint_vel = motion_data.joint_vel.copy() + + pose_rand = env.cfg.pose_randomization + pose_ranges = [ + (pose_rand.x[0], pose_rand.x[1]), + (pose_rand.y[0], pose_rand.y[1]), + (pose_rand.z[0], pose_rand.z[1]), + (pose_rand.roll[0], pose_rand.roll[1]), + (pose_rand.pitch[0], pose_rand.pitch[1]), + (pose_rand.yaw[0], pose_rand.yaw[1]), + ] + pose_samples = np.array( + [[np.random.uniform(low, high) for low, high in pose_ranges] for _ in range(num_reset)], + dtype=dtype, + ) + root_pos += pose_samples[:, 0:3] + root_ori = np_quat_mul( + np_quat_from_euler_xyz(pose_samples[:, 3], pose_samples[:, 4], pose_samples[:, 5]), + root_ori, + ) + + vel_rand = env.cfg.velocity_randomization + vel_ranges = [ + (vel_rand.x[0], vel_rand.x[1]), + (vel_rand.y[0], vel_rand.y[1]), + (vel_rand.z[0], vel_rand.z[1]), + (vel_rand.roll[0], vel_rand.roll[1]), + (vel_rand.pitch[0], vel_rand.pitch[1]), + (vel_rand.yaw[0], vel_rand.yaw[1]), + ] + vel_samples = np.array( + [[np.random.uniform(low, high) for low, high in vel_ranges] for _ in range(num_reset)], + dtype=dtype, + ) + root_lin_vel += vel_samples[:, :3] + root_ang_vel += vel_samples[:, 3:] + + joint_pos += np_sample_uniform( + env.cfg.joint_position_range[0], + env.cfg.joint_position_range[1], + joint_pos.shape, + dtype=np.float32, + ) + joint_range = env._get_joint_range() + if joint_range is not None: + joint_pos = np.clip(joint_pos, joint_range[:, 0], joint_range[:, 1]) + + qpos = np.tile(env._init_qpos, (num_reset, 1)) + qvel = np.tile(env._init_qvel, (num_reset, 1)) + qpos[:, 0:3] = root_pos + qpos[:, 3:7] = root_ori + qpos[:, 7:] = joint_pos + + qvel[:, 0:3] = root_lin_vel + qvel[:, 3:6] = np_quat_apply(np_quat_inv(root_ori), root_ang_vel) + qvel[:, 6:] = joint_vel + return qpos, qvel diff --git a/src/unilab/envs/motion_tracking/common/rewards.py b/src/unilab/envs/motion_tracking/common/rewards.py new file mode 100644 index 000000000..85f9728f9 --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/rewards.py @@ -0,0 +1,357 @@ +"""Reward configuration and reward functions for motion tracking. + +Reward terms are plain module-level callables ``fn(ctx: RewardContext) -> np.ndarray`` +mirroring :mod:`unilab.envs.locomotion.common.rewards`. Robot-specific terms that +live on env subclasses (box-object / joint-effort terms) are stored in the same +``_reward_fns`` dispatch table as bound methods and are called with the same +``ctx`` argument. + +``RewardContext`` intentionally carries the environment's *preallocated scratch +buffers* (``body_vec_error``, ``joint_error``, ... ``undesired_contact_mask``). +These buffers are env-owned so the hot path and the optional numba kernel run +with zero per-step allocations; the term functions write into them in place. The +op order, constants, clip bounds, and in-place ``out=`` usage are load-bearing — +they must stay bit-identical to preserve numeric parity with the numba path. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping + +import numpy as np + + +@dataclass +class RewardConfig: + """Reward configuration for motion tracking.""" + + scales: dict[str, float] = field( + default_factory=lambda: { + "motion_global_root_pos": 0.5, + "motion_global_root_ori": 0.5, + "motion_body_pos": 1.0, + "motion_body_ori": 1.0, + "motion_body_lin_vel": 1.0, + "motion_body_ang_vel": 1.0, + "motion_ee_body_pos_z": 0.0, + "motion_joint_pos": 0.0, + "motion_joint_vel": 0.0, + "action_rate_l2": -0.1, + "joint_limit": -10.0, + } + ) + # Standard deviations for exponential rewards + std_root_pos: float = 0.3 + std_root_ori: float = 0.4 + std_body_pos: float = 0.3 + std_body_ori: float = 0.4 + std_body_lin_vel: float = 1.0 + std_body_ang_vel: float = 3.14 + std_joint_pos: float = 0.2 + std_joint_vel: float = 1.0 + + +@dataclass +class RewardContext: + """Bundle of everything the motion-tracking reward functions may read. + + Built once per ``_compute_reward`` call. The array fields prefixed as + buffers (``env_error``, ``reward_term``, ``weighted_reward``, ...) are the + environment's preallocated scratch arrays; reward terms write into them in + place to keep the hot path allocation-free and numba-parity exact. + """ + + # ── semantic inputs ────────────────────────────────────────────── + info: dict + motion_data: Any = None + robot_body_pos_w: np.ndarray | None = None + robot_body_quat_w: np.ndarray | None = None + robot_body_lin_vel_w: np.ndarray | None = None + robot_body_ang_vel_w: np.ndarray | None = None + ref_body_pos_w: np.ndarray | None = None # env.body_pos_relative_w + ref_body_quat_w: np.ndarray | None = None # env.body_quat_relative_w + dof_pos: np.ndarray | None = None + dof_vel: np.ndarray | None = None + + # ── config-derived scalars / indices ──────────────────────────── + reward_config: Any = None + anchor_body_idx: int = 0 + ee_body_indices: np.ndarray | None = None + undesired_contact_body_indices: np.ndarray | None = None + joint_lower: np.ndarray | None = None + joint_upper: np.ndarray | None = None + undesired_contact_z_threshold: float = 0.0 + num_envs: int = 0 + + # ── env-owned scratch buffers (zero-alloc hot path) ────────────── + body_vec_error: np.ndarray | None = None + joint_error: np.ndarray | None = None + joint_error_upper: np.ndarray | None = None + env_error: np.ndarray | None = None + env_error2: np.ndarray | None = None + reward_term: np.ndarray | None = None + weighted_reward: np.ndarray | None = None + quat_error_w: np.ndarray | None = None + quat_error_x: np.ndarray | None = None + ee_pos_error_z: np.ndarray | None = None + undesired_contact_mask: np.ndarray | None = None + + +# ── buffered math helpers ──────────────────────────────────────────── + + +def _mean_body_xyz_squared_error( + ctx: RewardContext, reference: np.ndarray, actual: np.ndarray +) -> np.ndarray: + vec_error = ctx.body_vec_error + env_error = ctx.env_error + tmp_error = ctx.reward_term + np.subtract(reference[..., 0], actual[..., 0], out=vec_error[..., 0]) + np.square(vec_error[..., 0], out=vec_error[..., 0]) + np.sum(vec_error[..., 0], axis=1, out=env_error) + np.subtract(reference[..., 1], actual[..., 1], out=vec_error[..., 1]) + np.square(vec_error[..., 1], out=vec_error[..., 1]) + np.sum(vec_error[..., 1], axis=1, out=tmp_error) + env_error += tmp_error + np.subtract(reference[..., 2], actual[..., 2], out=vec_error[..., 2]) + np.square(vec_error[..., 2], out=vec_error[..., 2]) + np.sum(vec_error[..., 2], axis=1, out=tmp_error) + env_error += tmp_error + env_error /= reference.shape[1] + return env_error + + +def _quat_error_magnitude_squared_body( + ctx: RewardContext, q1: np.ndarray, q2: np.ndarray +) -> np.ndarray: + rel_w = ctx.quat_error_w + rel_x = ctx.quat_error_x + # Motion/backend quaternions are unit quaternions, so the relative + # rotation angle only needs abs(dot(q1, q2)). + np.multiply(q1[..., 0], q2[..., 0], out=rel_w) + np.multiply(q1[..., 1], q2[..., 1], out=rel_x) + rel_w += rel_x + np.multiply(q1[..., 2], q2[..., 2], out=rel_x) + rel_w += rel_x + np.multiply(q1[..., 3], q2[..., 3], out=rel_x) + rel_w += rel_x + np.abs(rel_w, out=rel_w) + np.clip(rel_w, 0.0, 1.0, out=rel_w) + np.arccos(rel_w, out=rel_x) + rel_x *= 2.0 + np.square(rel_x, out=rel_x) + return rel_x + + +def _exp_reward_from_error(ctx: RewardContext, error: np.ndarray, std: float) -> np.ndarray: + out = ctx.reward_term + np.divide(error, -(std**2), out=out) + np.exp(out, out=out) + return out + + +# ── reward terms ───────────────────────────────────────────────────── + + +def motion_global_root_pos(ctx: RewardContext) -> np.ndarray: + motion_data = ctx.motion_data + robot_body_pos_w = ctx.robot_body_pos_w + anchor_pos_w = motion_data.body_pos_w[:, ctx.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, ctx.anchor_body_idx] + error = ctx.env_error + np.subtract(anchor_pos_w[:, 0], robot_anchor_pos_w[:, 0], out=error) + np.square(error, out=error) + np.subtract(anchor_pos_w[:, 1], robot_anchor_pos_w[:, 1], out=ctx.reward_term) + np.square(ctx.reward_term, out=ctx.reward_term) + error += ctx.reward_term + np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=ctx.reward_term) + np.square(ctx.reward_term, out=ctx.reward_term) + error += ctx.reward_term + return _exp_reward_from_error(ctx, error, ctx.reward_config.std_root_pos) + + +def motion_global_root_ori(ctx: RewardContext) -> np.ndarray: + motion_data = ctx.motion_data + robot_body_quat_w = ctx.robot_body_quat_w + anchor_quat_w = motion_data.body_quat_w[:, ctx.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, ctx.anchor_body_idx] + np.multiply(anchor_quat_w[:, 0], robot_anchor_quat_w[:, 0], out=ctx.env_error) + np.multiply(anchor_quat_w[:, 1], robot_anchor_quat_w[:, 1], out=ctx.reward_term) + ctx.env_error += ctx.reward_term + np.multiply(anchor_quat_w[:, 2], robot_anchor_quat_w[:, 2], out=ctx.reward_term) + ctx.env_error += ctx.reward_term + np.multiply(anchor_quat_w[:, 3], robot_anchor_quat_w[:, 3], out=ctx.reward_term) + ctx.env_error += ctx.reward_term + np.abs(ctx.env_error, out=ctx.env_error) + np.clip(ctx.env_error, 0.0, 1.0, out=ctx.env_error) + np.arccos(ctx.env_error, out=ctx.env_error) + ctx.env_error *= 2.0 + np.square(ctx.env_error, out=ctx.env_error) + return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_root_ori) + + +def motion_body_pos(ctx: RewardContext) -> np.ndarray: + robot_body_pos_w = ctx.robot_body_pos_w + error = _mean_body_xyz_squared_error(ctx, ctx.ref_body_pos_w, robot_body_pos_w) + return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_pos) + + +def motion_body_ori(ctx: RewardContext) -> np.ndarray: + robot_body_quat_w = ctx.robot_body_quat_w + error = _quat_error_magnitude_squared_body(ctx, ctx.ref_body_quat_w, robot_body_quat_w) + np.sum(error, axis=-1, out=ctx.env_error) + ctx.env_error /= error.shape[1] + return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_body_ori) + + +def motion_body_lin_vel(ctx: RewardContext) -> np.ndarray: + motion_data = ctx.motion_data + robot_body_lin_vel_w = ctx.robot_body_lin_vel_w + error = _mean_body_xyz_squared_error(ctx, motion_data.body_lin_vel_w, robot_body_lin_vel_w) + return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_lin_vel) + + +def motion_body_ang_vel(ctx: RewardContext) -> np.ndarray: + motion_data = ctx.motion_data + robot_body_ang_vel_w = ctx.robot_body_ang_vel_w + error = _mean_body_xyz_squared_error(ctx, motion_data.body_ang_vel_w, robot_body_ang_vel_w) + return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_ang_vel) + + +def motion_ee_body_pos_z(ctx: RewardContext) -> np.ndarray: + robot_body_pos_w = ctx.robot_body_pos_w + np.subtract( + ctx.ref_body_pos_w[:, ctx.ee_body_indices, 2], + robot_body_pos_w[:, ctx.ee_body_indices, 2], + out=ctx.ee_pos_error_z, + ) + np.square(ctx.ee_pos_error_z, out=ctx.ee_pos_error_z) + np.sum(ctx.ee_pos_error_z, axis=-1, out=ctx.env_error) + ctx.env_error /= ctx.ee_pos_error_z.shape[1] + return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_body_pos) + + +def motion_joint_pos(ctx: RewardContext) -> np.ndarray: + motion_data = ctx.motion_data + dof_pos = ctx.dof_pos + np.subtract(motion_data.joint_pos, dof_pos, out=ctx.joint_error) + np.square(ctx.joint_error, out=ctx.joint_error) + np.sum(ctx.joint_error, axis=1, out=ctx.env_error) + ctx.env_error /= dof_pos.shape[1] + return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_joint_pos) + + +def motion_joint_vel(ctx: RewardContext) -> np.ndarray: + motion_data = ctx.motion_data + dof_vel = ctx.dof_vel + np.subtract(motion_data.joint_vel, dof_vel, out=ctx.joint_error) + np.square(ctx.joint_error, out=ctx.joint_error) + np.sum(ctx.joint_error, axis=1, out=ctx.env_error) + ctx.env_error /= dof_vel.shape[1] + return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_joint_vel) + + +def undesired_contacts(ctx: RewardContext) -> np.ndarray: + robot_body_pos_w = ctx.robot_body_pos_w + body_z = robot_body_pos_w[:, ctx.undesired_contact_body_indices, 2] + np.less( + body_z, + ctx.undesired_contact_z_threshold, + out=ctx.undesired_contact_mask, + ) + np.sum(ctx.undesired_contact_mask, axis=-1, out=ctx.env_error) + return ctx.env_error + + +def action_rate_l2(ctx: RewardContext) -> np.ndarray: + info = ctx.info + np.subtract(info["current_actions"], info["last_actions"], out=ctx.joint_error) + np.square(ctx.joint_error, out=ctx.joint_error) + np.sum(ctx.joint_error, axis=1, out=ctx.env_error) + return ctx.env_error + + +def joint_limit(ctx: RewardContext) -> np.ndarray: + dof_pos = ctx.dof_pos + lower = ctx.joint_lower + upper = ctx.joint_upper + if lower is None or upper is None: + ctx.reward_term.fill(0.0) + return ctx.reward_term + + # Compute violation + np.subtract(lower, dof_pos, out=ctx.joint_error) + np.maximum(ctx.joint_error, 0, out=ctx.joint_error) + np.subtract(dof_pos, upper, out=ctx.joint_error_upper) + np.maximum(ctx.joint_error_upper, 0, out=ctx.joint_error_upper) + ctx.joint_error += ctx.joint_error_upper + np.square(ctx.joint_error, out=ctx.joint_error) + np.sum(ctx.joint_error, axis=1, out=ctx.reward_term) + return ctx.reward_term + + +def build_reward_functions() -> dict[str, Callable[[RewardContext], np.ndarray]]: + """Return the robot-agnostic reward-term dispatch table. + + Keys match :data:`unilab.envs.motion_tracking.common.numba.TERM_ORDER`. + """ + return { + "motion_global_root_pos": motion_global_root_pos, + "motion_global_root_ori": motion_global_root_ori, + "motion_body_pos": motion_body_pos, + "motion_body_ori": motion_body_ori, + "motion_body_lin_vel": motion_body_lin_vel, + "motion_body_ang_vel": motion_body_ang_vel, + "motion_ee_body_pos_z": motion_ee_body_pos_z, + "motion_joint_pos": motion_joint_pos, + "motion_joint_vel": motion_joint_vel, + "action_rate_l2": action_rate_l2, + "joint_limit": joint_limit, + "undesired_contacts": undesired_contacts, + } + + +def compute_reward( + ctx: RewardContext, + *, + active_reward_fns: Mapping[str, Callable[[RewardContext], np.ndarray]], + all_reward_fns: Mapping[str, Callable[[RewardContext], np.ndarray]], + scales: Mapping[str, float], + ctrl_dt: float, + enable_log: bool, +) -> np.ndarray: + """Reduce ``scales × fns(ctx)`` into the per-env reward (in place, zero-alloc). + + Uses ``ctx.env_error2`` as the reward accumulator and ``ctx.weighted_reward`` + as the per-term scratch, matching the numba kernel's op order. Logs per-term + means into ``ctx.info["log"]`` every 4th step, then scales by ``ctrl_dt``. + """ + reward = ctx.env_error2 + reward.fill(0.0) + + info = ctx.info + step_count = info.get("steps") + should_log = enable_log and ( + int(step_count[0]) % 4 == 0 if isinstance(step_count, np.ndarray) else True + ) + log = {} if should_log else info.get("log", {}) + + for name, scale in scales.items(): + if scale == 0: + continue + reward_fn = active_reward_fns.get(name) + if reward_fn is None: + if should_log and name in all_reward_fns: + log[f"reward/{name}"] = 0.0 + continue + rew = reward_fn(ctx) + weighted_rew = ctx.weighted_reward + np.multiply(rew, scale, out=weighted_rew) + reward += weighted_rew + if should_log: + log[f"reward/{name}"] = float(np.sum(weighted_rew) / weighted_rew.size) + + info["log"] = log + reward *= ctrl_dt + return reward diff --git a/src/unilab/envs/motion_tracking/common/terminations.py b/src/unilab/envs/motion_tracking/common/terminations.py new file mode 100644 index 000000000..a09bfd90c --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/terminations.py @@ -0,0 +1,69 @@ +"""Termination computation for motion tracking.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from unilab.utils.geometry import np_gravity_z_in_body_from_quat + + +def compute_terminations( + env: Any, + motion_data: Any, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, +) -> np.ndarray: + """Compute termination conditions (writes into ``env._terminated``).""" + terminated = env._terminated + terminated.fill(False) + + # Anchor position error (Z-axis only) + anchor_pos_w = motion_data.body_pos_w[:, env.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, env.anchor_body_idx] + np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=env._env_error) + np.abs(env._env_error, out=env._env_error) + np.greater(env._env_error, env._cfg.anchor_pos_z_threshold, out=env._env_bool) + terminated |= env._env_bool + + # Anchor orientation error (gravity direction). The gravity-z difference + # is bounded by 2 for unit quaternions, so huge thresholds disable this + # termination without doing the per-step math. + if env._cfg.anchor_ori_threshold < 2.0: + anchor_quat_w = motion_data.body_quat_w[:, env.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, env.anchor_body_idx] + motion_gravity_z_b = np_gravity_z_in_body_from_quat(anchor_quat_w) + robot_gravity_z_b = np_gravity_z_in_body_from_quat(robot_anchor_quat_w) + np.subtract(motion_gravity_z_b, robot_gravity_z_b, out=env._env_error) + np.abs(env._env_error, out=env._env_error) + np.greater(env._env_error, env._cfg.anchor_ori_threshold, out=env._env_bool) + terminated |= env._env_bool + + # End-effector position error (Z-axis only) + if env._has_ee_body_indices: + np.subtract( + env.body_pos_relative_w[:, env.ee_body_indices, 2], + robot_body_pos_w[:, env.ee_body_indices, 2], + out=env._ee_pos_error_z, + ) + np.abs(env._ee_pos_error_z, out=env._ee_pos_error_z) + np.greater( + env._ee_pos_error_z, + env._cfg.ee_body_pos_z_threshold, + out=env._ee_terminated, + ) + np.logical_or.reduce(env._ee_terminated, axis=1, out=env._env_bool) + terminated |= env._env_bool + + if env._cfg.terminate_on_undesired_contacts and env._has_undesired_contact_body_indices: + body_z = robot_body_pos_w[:, env.undesired_contact_body_indices, 2] + np.less( + body_z, + env._cfg.undesired_contact_z_threshold, + out=env._undesired_contact_mask, + ) + np.logical_or.reduce(env._undesired_contact_mask, axis=-1, out=env._env_bool) + terminated |= env._env_bool + + return terminated diff --git a/src/unilab/envs/motion_tracking/common/tracking.py b/src/unilab/envs/motion_tracking/common/tracking.py new file mode 100644 index 000000000..c20fb9073 --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/tracking.py @@ -0,0 +1,576 @@ +"""Robot-agnostic motion-tracking engine. + +Holds :class:`MotionTrackingEnv` (the imitation engine, inheriting the shared +``G1BaseEnv`` locomotion base) and :class:`MotionTrackingDeployEnv` (the +unitree_rl_lab mimic actor variant). Per-concern math lives in the owner +modules (``rewards`` / ``observations`` / ``terminations`` / ``transforms`` / +``reset`` / ``domain_randomization`` / ``numba``); the engine keeps only the +thin polymorphic method surface and per-step orchestration. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from unilab.base.backend import create_backend, env_backend_kwargs +from unilab.base.np_env import NpEnvState +from unilab.dtype_config import get_global_dtype +from unilab.envs.locomotion.g1.base import G1BaseEnv + +from . import observations +from .config import MotionTrackingCfg, MotionTrackingDeployEnvCfg +from .domain_randomization import MotionTrackingDomainRandomizationProvider +from .motion_loader import MotionData, MotionLoader, MotionSampler +from .reset import build_motion_reference_state +from .rewards import RewardContext, build_reward_functions, compute_reward +from .terminations import compute_terminations +from .transforms import update_relative_transforms + + +class MotionTrackingEnv(G1BaseEnv): + """Motion Tracking Environment (robot-agnostic imitation engine).""" + + _cfg: MotionTrackingCfg + + def __init__(self, cfg: MotionTrackingCfg, num_envs=1, backend_type="mujoco"): + if not cfg.motion_file: + raise ValueError("motion_file must be specified in config") + + backend = create_backend( + backend_type, + cfg.scene, + num_envs, + cfg.sim_dt, + base_name=cfg.asset.base_name, + push_body_name=cfg.domain_rand.push_body_name, + add_body_sensors=True, + **env_backend_kwargs(cfg), + ) + super().__init__(cfg, backend, num_envs) + + # Resolve body IDs for backend querying and motion-file indexing. + self.body_ids = self._backend.get_body_ids(cfg.body_names) + motion_body_ids = self._backend.get_motion_body_ids(cfg.body_names) + + self.anchor_body_idx = cfg.body_names.index(cfg.anchor_body_name) + + # Get end-effector body indices for termination + self.ee_body_indices = np.array( + [cfg.body_names.index(name) for name in cfg.ee_body_names], dtype=np.int32 + ) + self._has_ee_body_indices = bool(self.ee_body_indices.size) + + # Get non-EE body indices for undesired contact penalty + ee_set = set(cfg.ee_body_names) + self.undesired_contact_body_indices = np.array( + [i for i, name in enumerate(cfg.body_names) if name not in ee_set], + dtype=np.int32, + ) + self._has_undesired_contact_body_indices = bool(self.undesired_contact_body_indices.size) + + # Load motion data + self.motion_loader = MotionLoader(cfg.motion_file, body_indices=motion_body_ids) + self.motion_sampler = MotionSampler( + self.motion_loader, + mode=cfg.sampling_mode, + num_envs=num_envs, + start_ratio=cfg.sampling_start_ratio, + ) + needs_kp_kd = cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd + needs_friction = getattr(cfg.domain_rand, "randomize_geom_friction", False) + base_kp = base_kd = None + if needs_kp_kd: + base_kp, base_kd = backend.get_actuator_gains() + base_geom_friction = None + foot_geom_ids = None + if needs_friction: + import re as _re + + base_geom_friction = backend.get_geom_friction() + geom_names = backend.get_geom_names() + pattern = _re.compile(cfg.domain_rand.friction_geom_pattern) + foot_geom_ids = np.asarray( + [i for i, name in enumerate(geom_names) if name and pattern.match(name)], + dtype=np.int64, + ) + if foot_geom_ids.size == 0: + raise ValueError( + "friction_geom_pattern " + f"'{cfg.domain_rand.friction_geom_pattern}' did not match any geom" + ) + dr_provider = MotionTrackingDomainRandomizationProvider( + base_kp=base_kp, + base_kd=base_kd, + base_geom_friction=base_geom_friction, + foot_geom_ids=foot_geom_ids, + ) + self._init_domain_randomization(dr_provider) + + dtype = get_global_dtype() + n_body = len(cfg.body_names) + self._n_motion_bodies = n_body + self._actor_obs_width = self._actor_obs_dim(self._num_action) + self._critic_base_obs_width = self._critic_base_obs_dim(self._num_action) + self._critic_obs_width = self._critic_base_obs_width + n_body * 9 + self._copy_body_state_w = self._backend.copy_body_state_w + + # Buffers for relative body transforms + self.body_pos_relative_w = np.zeros((num_envs, n_body, 3), dtype=dtype) + self.body_quat_relative_w = np.zeros((num_envs, n_body, 4), dtype=dtype) + self.body_quat_relative_w[:, :, 0] = 1.0 # Initialize to identity quaternion + self._motion_data_buffer = ( + self.motion_loader.make_motion_data_buffer(num_envs) + if hasattr(self.motion_loader, "make_motion_data_buffer") + else None + ) + self._zero_actions = np.zeros((num_envs, self._num_action), dtype=dtype) + self._joint_range = self._backend.get_joint_range() + if self._joint_range is not None: + self._joint_range = np.asarray(self._joint_range, dtype=dtype) + self._joint_lower = self._joint_range[:, 0] + self._joint_upper = self._joint_range[:, 1] + else: + self._joint_lower = None + self._joint_upper = None + self._delta_pos_w = np.empty((num_envs, 3), dtype=dtype) + self._delta_ori_w = np.empty((num_envs, 4), dtype=dtype) + self._motion_anchor_pos_b = np.empty((num_envs, 3), dtype=dtype) + self._motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) + self._motion_command = np.empty((num_envs, self._num_action * 2), dtype=dtype) + self._joint_pos_rel = np.empty((num_envs, self._num_action), dtype=dtype) + self._robot_body_pos_w = np.empty((num_envs, n_body, 3), dtype=dtype) + self._robot_body_quat_w = np.empty((num_envs, n_body, 4), dtype=dtype) + self._robot_body_lin_vel_w = np.empty((num_envs, n_body, 3), dtype=dtype) + self._robot_body_ang_vel_w = np.empty((num_envs, n_body, 3), dtype=dtype) + self._quat_error_w = np.empty((num_envs, n_body), dtype=dtype) + self._quat_error_x = np.empty((num_envs, n_body), dtype=dtype) + self._body_vec_error = np.empty((num_envs, n_body, 3), dtype=dtype) + self._body_vec_tmp = np.empty((num_envs, n_body, 3), dtype=dtype) + self._joint_error = np.empty((num_envs, self._num_action), dtype=dtype) + self._joint_error_upper = np.empty((num_envs, self._num_action), dtype=dtype) + self._env_error = np.empty((num_envs,), dtype=dtype) + self._env_error2 = np.empty((num_envs,), dtype=dtype) + self._reward_term = np.empty((num_envs,), dtype=dtype) + self._weighted_reward = np.empty((num_envs,), dtype=dtype) + self._terminated = np.empty((num_envs,), dtype=bool) + self._env_bool = np.empty((num_envs,), dtype=bool) + self._ee_pos_error_z = np.empty((num_envs, self.ee_body_indices.size), dtype=dtype) + self._ee_terminated = np.empty((num_envs, self.ee_body_indices.size), dtype=bool) + self._undesired_contact_mask = np.empty( + (num_envs, self.undesired_contact_body_indices.size), dtype=bool + ) + + self._enable_reward_log = True + self._init_reward_functions() + self._active_reward_fns = { + name: reward_fn + for name, reward_fn in self._reward_fns.items() + if self._reward_term_is_active(name) + } + self._numba_accelerator = None + if cfg.numba_acceleration: + from unilab.envs.motion_tracking.common.numba import MotionTrackingNumbaAccelerator + + self._numba_accelerator = MotionTrackingNumbaAccelerator.from_env( + self, num_threads=cfg.numba_num_threads + ) + self._clip_end_truncated = np.zeros((num_envs,), dtype=bool) + + def _effective_default_angles(self, env_ids: np.ndarray | None = None) -> np.ndarray: + """Return default_angles with per-episode joint-default-pos bias applied.""" + state = getattr(self, "_state", None) + if state is not None: + bias = state.info.get("default_dof_pos_bias") + if bias is not None: + if env_ids is not None: + return self.default_angles + bias[env_ids] + return self.default_angles + bias + return self.default_angles + + def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: + state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions)) + state.info["current_actions"] = actions + exec_actions = ( + state.info["last_actions"] + if self._cfg.control_config.simulate_action_latency + else actions + ) + bias = state.info.get("default_dof_pos_bias") + base = self.default_angles + bias if bias is not None else self.default_angles + ctrl: np.ndarray = exec_actions * self._cfg.control_config.action_scale + base + return ctrl + + def _resample_reference_state(self, env_ids: np.ndarray) -> None: + motion_frames = self.motion_sampler.sample_frames(env_ids) + motion_data = self.motion_loader.get_motion_at_frame(motion_frames) + qpos, qvel = build_motion_reference_state(self, env_ids, motion_data) + self._backend.set_state(env_ids, qpos, qvel) + + def _refresh_observation_rows( + self, obs: dict[str, np.ndarray], info: dict, env_ids: np.ndarray + ) -> None: + motion_data = self.motion_loader.get_motion_at_frame( + self.motion_sampler.current_frames[env_ids] + ) + row_ids = np.asarray(env_ids, dtype=np.intp) + linvel = self._backend.get_sensor_data_rows(self._cfg.sensor.local_linvel, row_ids) + gyro = self._backend.get_sensor_data_rows(self._cfg.sensor.gyro, row_ids) + dof_pos = self.get_dof_pos()[row_ids] + dof_vel = self.get_dof_vel()[row_ids] + robot_body_pos_w, robot_body_quat_w = self._backend.get_body_pose_w_rows( + row_ids, self.body_ids + ) + + obs_info: dict[str, Any] = {} + current_actions = info.get("current_actions") + if isinstance(current_actions, np.ndarray): + obs_info["current_actions"] = current_actions[env_ids] + obs_info["env_ids"] = env_ids + + refreshed_obs = self._compute_obs( + obs_info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ) + for key, value in refreshed_obs.items(): + if value.shape[0] == len(env_ids): + obs[key][env_ids] = value + else: + obs[key][env_ids] = value[env_ids] + + def _get_body_pose_w(self) -> tuple[np.ndarray, np.ndarray]: + return self._backend.get_body_pose_w(self.body_ids) + + def _get_body_state_w(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + copy_body_state_w = self._copy_body_state_w + if copy_body_state_w is not None: + return copy_body_state_w( + self.body_ids, + self._robot_body_pos_w, + self._robot_body_quat_w, + self._robot_body_lin_vel_w, + self._robot_body_ang_vel_w, + ) + robot_body_pos_w, robot_body_quat_w = self._get_body_pose_w() + robot_body_lin_vel_w, robot_body_ang_vel_w = self._backend.get_body_vel_w(self.body_ids) + return ( + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + ) + + def _get_joint_range(self) -> np.ndarray | None: + return self._joint_range + + def _get_current_motion(self) -> MotionData: + if self._motion_data_buffer is None: + return self.motion_sampler.get_current_motion() + return self.motion_sampler.get_current_motion(self._motion_data_buffer) + + @property + def obs_groups_spec(self) -> dict[str, int]: + return observations.obs_groups_spec(self) + + def _actor_obs_dim(self, n: int) -> int: + return observations.actor_obs_dim(n) + + def _critic_base_obs_dim(self, n: int) -> int: + return observations.critic_base_obs_dim(n) + + def _build_actor_obs( + self, + *, + command: np.ndarray, + motion_anchor_pos_b: np.ndarray, + motion_anchor_ori_b: np.ndarray, + noisy_linvel: np.ndarray, + noisy_gyro: np.ndarray, + noisy_joint_pos_rel: np.ndarray, + noisy_dof_vel: np.ndarray, + last_actions: np.ndarray, + ) -> np.ndarray: + n_action = noisy_joint_pos_rel.shape[1] + return observations.build_actor_obs( + actor_obs_dim=self._actor_obs_dim(n_action), + command=command, + motion_anchor_pos_b=motion_anchor_pos_b, + motion_anchor_ori_b=motion_anchor_ori_b, + noisy_linvel=noisy_linvel, + noisy_gyro=noisy_gyro, + noisy_joint_pos_rel=noisy_joint_pos_rel, + noisy_dof_vel=noisy_dof_vel, + last_actions=last_actions, + ) + + def _init_reward_functions(self): + self._reward_fns = build_reward_functions() + + def _reward_term_is_active(self, name: str) -> bool: + if name == "joint_limit": + return self._joint_lower is not None and self._joint_upper is not None + if name == "undesired_contacts": + return self._has_undesired_contact_body_indices + if name == "motion_ee_body_pos_z": + return self._has_ee_body_indices + return True + + def update_state(self, state: NpEnvState) -> NpEnvState: + self._clip_end_truncated.fill(False) + + # Get current motion data + motion_data = self._get_current_motion() + + # Get robot state + linvel = self.get_local_linvel() + gyro = self.get_gyro() + dof_pos = self.get_dof_pos() + dof_vel = self.get_dof_vel() + + # Get body states + ( + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + ) = self._get_body_state_w() + + numba_accelerator = getattr(self, "_numba_accelerator", None) + if numba_accelerator is not None: + noise_cfg = self._cfg.noise_config + numba_result = numba_accelerator.compute_update_state( + info=state.info, + motion_data=motion_data, + linvel=linvel, + gyro=gyro, + dof_pos=dof_pos, + dof_vel=dof_vel, + robot_body_pos_w=robot_body_pos_w, + robot_body_quat_w=robot_body_quat_w, + robot_body_lin_vel_w=robot_body_lin_vel_w, + robot_body_ang_vel_w=robot_body_ang_vel_w, + ref_body_pos_w=self.body_pos_relative_w, + ref_body_quat_w=self.body_quat_relative_w, + motion_anchor_pos_b=self._motion_anchor_pos_b, + motion_anchor_ori_b=self._motion_anchor_ori_b, + joint_pos_rel=self._joint_pos_rel, + scales=self._cfg.reward_config.scales, + enable_log=self._enable_reward_log, + noise_level=noise_cfg.level, + noise_scale_linvel=noise_cfg.scale_linvel, + noise_scale_gyro=noise_cfg.scale_gyro, + noise_scale_joint_angle=noise_cfg.scale_joint_angle, + noise_scale_joint_vel=noise_cfg.scale_joint_vel, + ) + terminated = numba_result.terminated + reward = numba_result.reward + obs = numba_result.obs + state.info["log"] = numba_result.log + else: + # Compute relative body transforms (for observations and rewards) + self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) + + # Compute terminations + terminated = self._compute_terminations( + motion_data, robot_body_pos_w, robot_body_quat_w + ) + + # Compute reward + reward = self._compute_reward( + state.info, + motion_data, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + ) + + # Compute observations + obs = self._compute_obs( + state.info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ) + + # Update failure statistics for adaptive sampling + self.motion_sampler.update_failure_stats(terminated) + + # Advance motion frames + done_env_ids = self.motion_sampler.step() + if len(done_env_ids) > 0: + if self._cfg.truncate_on_clip_end: + self._clip_end_truncated[done_env_ids] = True + else: + # Match BeyondMimic: clip boundaries are command resampling points, not + # episode boundaries; sync the simulated robot to the new reference. + resample_env_ids = done_env_ids[~terminated[done_env_ids]] + if len(resample_env_ids) > 0: + self._resample_reference_state(resample_env_ids) + self._refresh_observation_rows(obs, state.info, resample_env_ids) + + return state.replace(obs=obs, reward=reward, terminated=terminated) + + def _compute_truncated(self, state: NpEnvState) -> np.ndarray: + truncated = super()._compute_truncated(state) + clip_end_only = getattr(self, "_env_bool", None) + if clip_end_only is None or clip_end_only.shape != (self._num_envs,): + clip_end_only = np.empty((self._num_envs,), dtype=bool) + self._env_bool = clip_end_only + np.logical_not(state.terminated, out=clip_end_only) + np.logical_and(self._clip_end_truncated, clip_end_only, out=clip_end_only) + np.logical_or(truncated, clip_end_only, out=truncated) + return truncated + + def _update_relative_transforms( + self, motion_data, robot_body_pos_w: np.ndarray, robot_body_quat_w: np.ndarray + ): + """Update relative body transforms for tracking.""" + update_relative_transforms(self, motion_data, robot_body_pos_w, robot_body_quat_w) + + def _compute_terminations( + self, + motion_data, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + ) -> np.ndarray: + """Compute termination conditions.""" + return compute_terminations(self, motion_data, robot_body_pos_w, robot_body_quat_w) + + def _write_body_pos_in_anchor_frame( + self, + anchor_pos: np.ndarray, + anchor_quat: np.ndarray, + body_pos: np.ndarray, + out: np.ndarray, + ) -> None: + observations.write_body_pos_in_anchor_frame( + anchor_pos, anchor_quat, body_pos, out, body_vec_error=self._body_vec_error + ) + + def _write_body_ori6_in_anchor_frame( + self, + anchor_quat: np.ndarray, + body_quat: np.ndarray, + out: np.ndarray, + ) -> None: + observations.write_body_ori6_in_anchor_frame(anchor_quat, body_quat, out) + + def _compute_obs( + self, + info: dict, + motion_data, + linvel: np.ndarray, + gyro: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + ) -> dict[str, np.ndarray]: + """Compute observations as dict with actor and critic groups.""" + return observations.compute_obs( + self, + info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ) + + def _compute_reward( + self, + info: dict, + motion_data, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + robot_body_lin_vel_w: np.ndarray, + robot_body_ang_vel_w: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + ) -> np.ndarray: + """Compute reward.""" + ctx = RewardContext( + info=info, + motion_data=motion_data, + robot_body_pos_w=robot_body_pos_w, + robot_body_quat_w=robot_body_quat_w, + robot_body_lin_vel_w=robot_body_lin_vel_w, + robot_body_ang_vel_w=robot_body_ang_vel_w, + ref_body_pos_w=self.body_pos_relative_w, + ref_body_quat_w=self.body_quat_relative_w, + dof_pos=dof_pos, + dof_vel=dof_vel, + reward_config=self._cfg.reward_config, + anchor_body_idx=self.anchor_body_idx, + ee_body_indices=self.ee_body_indices, + undesired_contact_body_indices=self.undesired_contact_body_indices, + joint_lower=self._joint_lower, + joint_upper=self._joint_upper, + undesired_contact_z_threshold=self._cfg.undesired_contact_z_threshold, + num_envs=self._num_envs, + body_vec_error=self._body_vec_error, + joint_error=self._joint_error, + joint_error_upper=self._joint_error_upper, + env_error=self._env_error, + env_error2=self._env_error2, + reward_term=self._reward_term, + weighted_reward=self._weighted_reward, + quat_error_w=self._quat_error_w, + quat_error_x=self._quat_error_x, + ee_pos_error_z=self._ee_pos_error_z, + undesired_contact_mask=self._undesired_contact_mask, + ) + return compute_reward( + ctx, + active_reward_fns=self._active_reward_fns, + all_reward_fns=self._reward_fns, + scales=self._cfg.reward_config.scales, + ctrl_dt=self._cfg.ctrl_dt, + enable_log=self._enable_reward_log, + ) + + +class MotionTrackingDeployEnv(MotionTrackingEnv): + """Deploy-oriented motion tracking env with unitree_rl_lab mimic actor inputs.""" + + _cfg: MotionTrackingDeployEnvCfg + + def _actor_obs_dim(self, n: int) -> int: + return observations.mimic_actor_obs_dim(n) + + def _build_actor_obs( + self, + *, + command: np.ndarray, + motion_anchor_pos_b: np.ndarray, + motion_anchor_ori_b: np.ndarray, + noisy_linvel: np.ndarray, + noisy_gyro: np.ndarray, + noisy_joint_pos_rel: np.ndarray, + noisy_dof_vel: np.ndarray, + last_actions: np.ndarray, + ) -> np.ndarray: + return observations.build_mimic_actor_obs( + command=command, + motion_anchor_ori_b=motion_anchor_ori_b, + noisy_gyro=noisy_gyro, + noisy_joint_pos_rel=noisy_joint_pos_rel, + noisy_dof_vel=noisy_dof_vel, + last_actions=last_actions, + ) diff --git a/src/unilab/envs/motion_tracking/common/transforms.py b/src/unilab/envs/motion_tracking/common/transforms.py new file mode 100644 index 000000000..b1ffb7d69 --- /dev/null +++ b/src/unilab/envs/motion_tracking/common/transforms.py @@ -0,0 +1,103 @@ +"""Relative body-transform computation for motion tracking. + +Fills the environment's ``body_pos_relative_w`` / ``body_quat_relative_w`` +reference buffers each step. The op order and in-place ``out=`` usage are +load-bearing for numba parity and must not change. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + + +def update_relative_transforms( + env: Any, + motion_data: Any, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, +) -> None: + """Update relative body transforms for tracking.""" + # Get anchor states + anchor_pos_w = motion_data.body_pos_w[:, env.anchor_body_idx] + anchor_quat_w = motion_data.body_quat_w[:, env.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, env.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, env.anchor_body_idx] + + # Compute delta transform: keep robot's XY position, use motion's Z height + # and apply yaw-only rotation difference. + delta_pos_w = env._delta_pos_w + delta_pos_w[:] = robot_anchor_pos_w + delta_pos_w[:, 2] = anchor_pos_w[:, 2] + + # Compute yaw-only rotation difference, equivalent to + # np_yaw_quat(np_quat_mul(robot_anchor_quat_w, np_quat_inv(anchor_quat_w))). + delta_ori_w = env._delta_ori_w + rw, rx, ry, rz = ( + robot_anchor_quat_w[:, 0], + robot_anchor_quat_w[:, 1], + robot_anchor_quat_w[:, 2], + robot_anchor_quat_w[:, 3], + ) + aw, ax, ay, az = ( + anchor_quat_w[:, 0], + anchor_quat_w[:, 1], + anchor_quat_w[:, 2], + anchor_quat_w[:, 3], + ) + qw = rw * aw + rx * ax + ry * ay + rz * az + qx = -rw * ax + rx * aw - ry * az + rz * ay + qy = -rw * ay + rx * az + ry * aw - rz * ax + qz = -rw * az - rx * ay + ry * ax + rz * aw + half_yaw = 0.5 * np.arctan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz)) + np.cos(half_yaw, out=delta_ori_w[:, 0]) + delta_ori_w[:, 1:3] = 0.0 + np.sin(half_yaw, out=delta_ori_w[:, 3]) + + dw1 = delta_ori_w[:, 0] + dz1 = delta_ori_w[:, 3] + dw = dw1[:, None] + dz = dz1[:, None] + mw = motion_data.body_quat_w[..., 0] + mx = motion_data.body_quat_w[..., 1] + my = motion_data.body_quat_w[..., 2] + mz = motion_data.body_quat_w[..., 3] + out_quat = env.body_quat_relative_w + out_quat[..., 0] = dw * mw + out_quat[..., 0] -= dz * mz + out_quat[..., 1] = dw * mx + out_quat[..., 1] -= dz * my + out_quat[..., 2] = dw * my + out_quat[..., 2] += dz * mx + out_quat[..., 3] = dw * mz + out_quat[..., 3] += dz * mw + + rel_pos = env._body_vec_error + vx = rel_pos[..., 0] + vy = rel_pos[..., 1] + vz = rel_pos[..., 2] + np.subtract(motion_data.body_pos_w[..., 0], anchor_pos_w[:, None, 0], out=vx) + np.subtract(motion_data.body_pos_w[..., 1], anchor_pos_w[:, None, 1], out=vy) + np.subtract(motion_data.body_pos_w[..., 2], anchor_pos_w[:, None, 2], out=vz) + + yaw_cross = env._env_error + yaw_z2 = env._reward_term + np.multiply(dw1, dz1, out=yaw_cross) + yaw_cross *= 2.0 + np.square(dz1, out=yaw_z2) + yaw_z2 *= 2.0 + yaw_cross_2d = yaw_cross[:, None] + yaw_z2_2d = yaw_z2[:, None] + + out_pos = env.body_pos_relative_w + out_pos[..., 0] = vx + out_pos[..., 0] -= yaw_cross_2d * vy + out_pos[..., 0] -= yaw_z2_2d * vx + out_pos[..., 0] += delta_pos_w[:, None, 0] + out_pos[..., 1] = vy + out_pos[..., 1] += yaw_cross_2d * vx + out_pos[..., 1] -= yaw_z2_2d * vy + out_pos[..., 1] += delta_pos_w[:, None, 1] + out_pos[..., 2] = vz + out_pos[..., 2] += delta_pos_w[:, None, 2] From 0cb244864da1113f344dfae3218ff1fcb32ef100 Mon Sep 17 00:00:00 2001 From: caozx1110 Date: Thu, 9 Jul 2026 18:37:58 +0800 Subject: [PATCH 2/4] refactor: route motion tracking profiles through common engine --- .../envs/motion_tracking/g1/box_tracking.py | 9 +- .../envs/motion_tracking/g1/flip_tracking.py | 23 +- .../motion_tracking/g1/motion_box_loader.py | 2 +- .../envs/motion_tracking/g1/tracking.py | 1484 +---------------- .../envs/motion_tracking/g1/tracking_obs.py | 15 +- .../envs/motion_tracking/x2/flip_tracking.py | 16 +- 6 files changed, 79 insertions(+), 1470 deletions(-) diff --git a/src/unilab/envs/motion_tracking/g1/box_tracking.py b/src/unilab/envs/motion_tracking/g1/box_tracking.py index 5024323bb..48fefad54 100644 --- a/src/unilab/envs/motion_tracking/g1/box_tracking.py +++ b/src/unilab/envs/motion_tracking/g1/box_tracking.py @@ -24,6 +24,7 @@ np_subtract_frame_transforms, ) +from ..common.rewards import RewardContext from .motion_box_loader import BoxMotionData, BoxMotionLoader from .tracking import ( G1MotionTrackingCfg, @@ -344,8 +345,8 @@ def _compute_obs( ) return obs - def _reward_object_position(self, info: dict) -> np.ndarray: - motion_data: BoxMotionData = info["motion_data"] + def _reward_object_position(self, ctx: RewardContext) -> np.ndarray: + motion_data: BoxMotionData = ctx.motion_data if motion_data.object_pos_w is None: return np.zeros((self._num_envs,), dtype=get_global_dtype()) obj_pos_w = self._backend.get_body_pos_w(self._object_body_ids)[:, 0, :] @@ -354,8 +355,8 @@ def _reward_object_position(self, info: dict) -> np.ndarray: np.exp(-error / self._cfg.reward_config.std_object_pos**2), dtype=get_global_dtype() ) - def _reward_object_orientation(self, info: dict) -> np.ndarray: - motion_data: BoxMotionData = info["motion_data"] + def _reward_object_orientation(self, ctx: RewardContext) -> np.ndarray: + motion_data: BoxMotionData = ctx.motion_data if motion_data.object_quat_w is None: return np.zeros((self._num_envs,), dtype=get_global_dtype()) obj_quat_w = self._backend.get_body_quat_w(self._object_body_ids)[:, 0, :] diff --git a/src/unilab/envs/motion_tracking/g1/flip_tracking.py b/src/unilab/envs/motion_tracking/g1/flip_tracking.py index b6515a36d..62a883ed3 100644 --- a/src/unilab/envs/motion_tracking/g1/flip_tracking.py +++ b/src/unilab/envs/motion_tracking/g1/flip_tracking.py @@ -13,6 +13,7 @@ from unilab.base import registry from unilab.base.scene import SceneCfg +from ..common.config import _zero_pose_randomization, _zero_velocity_randomization from .tracking import ( G1MotionTrackingCfg, G1MotionTrackingEnv, @@ -21,28 +22,6 @@ ) -def _zero_pose_randomization() -> PoseRandomization: - return PoseRandomization( - x=(0.0, 0.0), - y=(0.0, 0.0), - z=(0.0, 0.0), - roll=(0.0, 0.0), - pitch=(0.0, 0.0), - yaw=(0.0, 0.0), - ) - - -def _zero_velocity_randomization() -> VelocityRandomization: - return VelocityRandomization( - x=(0.0, 0.0), - y=(0.0, 0.0), - z=(0.0, 0.0), - roll=(0.0, 0.0), - pitch=(0.0, 0.0), - yaw=(0.0, 0.0), - ) - - @dataclass class G1FlipTrackingCfg(G1MotionTrackingCfg): """Config profile for flip tracking clips.""" diff --git a/src/unilab/envs/motion_tracking/g1/motion_box_loader.py b/src/unilab/envs/motion_tracking/g1/motion_box_loader.py index 80a9e9e09..820ff7f14 100644 --- a/src/unilab/envs/motion_tracking/g1/motion_box_loader.py +++ b/src/unilab/envs/motion_tracking/g1/motion_box_loader.py @@ -6,7 +6,7 @@ import numpy as np -from .motion_loader import MotionData, MotionLoader +from ..common.motion_loader import MotionData, MotionLoader @dataclass diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index 4a9b4df98..d69313140 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -1,201 +1,42 @@ -"""G1 Motion Tracking Environment - Motion imitation task.""" +"""G1 Motion Tracking profiles — thin registry subclasses over the shared engine. -from __future__ import annotations +The robot-agnostic engine and owner modules now live in +:mod:`unilab.envs.motion_tracking.common`. This module keeps the G1 registry +entries (``G1MotionTracking`` / ``G1MotionTrackingDeploy``) and re-exports the +historical ``G1*`` / ``Domain_Rand`` / ``_build_motion_reference_state`` symbol +names so existing subclasses and tests keep importing them from ``.tracking``. +""" -import time -from dataclasses import dataclass, field -from typing import Any, Literal, cast +from __future__ import annotations -import numpy as np +from dataclasses import dataclass -from unilab.assets import ASSETS_ROOT_PATH from unilab.base import registry -from unilab.base.backend import create_backend, env_backend_kwargs -from unilab.base.np_env import NpEnvState -from unilab.base.scene import SceneCfg -from unilab.dr import ( - DomainRandomizationCapabilities, - DomainRandomizationProvider, - IntervalRandomizationPlan, - ResetPlan, -) -from unilab.dr.dr_utils import ( - build_common_reset_randomization, - build_interval_push_plan, - validate_common_reset_randomization, - validate_interval_push_support, - zero_actions, -) -from unilab.dr.types import RESET_TERM_GEOM_FRICTION, ResetRandomizationPayload -from unilab.dtype_config import get_global_dtype -from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv -from unilab.utils.geometry import ( - np_gravity_z_in_body_from_quat, - np_sample_uniform, - np_write_relative_anchor_transform_pos_rot6d, -) -from unilab.utils.rotation import ( - np_quat_apply, - np_quat_from_euler_xyz, - np_quat_inv, - np_quat_mul, -) - -from .motion_loader import MotionData, MotionLoader, MotionSampler - - -@dataclass -class RewardConfig: - """Reward configuration for motion tracking.""" - - scales: dict[str, float] = field( - default_factory=lambda: { - "motion_global_root_pos": 0.5, - "motion_global_root_ori": 0.5, - "motion_body_pos": 1.0, - "motion_body_ori": 1.0, - "motion_body_lin_vel": 1.0, - "motion_body_ang_vel": 1.0, - "motion_ee_body_pos_z": 0.0, - "motion_joint_pos": 0.0, - "motion_joint_vel": 0.0, - "action_rate_l2": -0.1, - "joint_limit": -10.0, - } - ) - # Standard deviations for exponential rewards - std_root_pos: float = 0.3 - std_root_ori: float = 0.4 - std_body_pos: float = 0.3 - std_body_ori: float = 0.4 - std_body_lin_vel: float = 1.0 - std_body_ang_vel: float = 3.14 - std_joint_pos: float = 0.2 - std_joint_vel: float = 1.0 - - -@dataclass -class PoseRandomization: - """Pose randomization ranges for reset.""" - - x: tuple[float, float] = (-0.05, 0.05) - y: tuple[float, float] = (-0.05, 0.05) - z: tuple[float, float] = (-0.01, 0.01) - roll: tuple[float, float] = (-0.1, 0.1) - pitch: tuple[float, float] = (-0.1, 0.1) - yaw: tuple[float, float] = (-0.2, 0.2) - - -@dataclass -class VelocityRandomization: - """Velocity randomization ranges for reset.""" - - x: tuple[float, float] = (-0.5, 0.5) - y: tuple[float, float] = (-0.5, 0.5) - z: tuple[float, float] = (-0.2, 0.2) - roll: tuple[float, float] = (-0.52, 0.52) - pitch: tuple[float, float] = (-0.52, 0.52) - yaw: tuple[float, float] = (-0.78, 0.78) - -@dataclass -class Domain_Rand: - """Domain randomization config required by motrix backend hooks.""" - - randomize_base_mass: bool = False - added_mass_range: list[float] = field(default_factory=lambda: [-1.5, 1.5]) - - random_com: bool = False - com_offset_x: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - com_offset_y: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - com_offset_z: list[float] = field(default_factory=lambda: [-0.05, 0.05]) - - randomize_gravity: bool = False - gravity_range: list[list[float]] = field( - default_factory=lambda: [[0.0, 0.0, -9.81], [0.0, 0.0, -9.81]] - ) - - push_robots: bool = False - push_interval: int = 750 - max_force: list[float] = field(default_factory=lambda: [1.0, 1.0, 0.5]) - push_body_name: str | None = None - - randomize_kp: bool = False - kp_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_kd: bool = False - kd_multiplier_range: list[float] = field(default_factory=lambda: [0.9, 1.1]) - - randomize_geom_friction: bool = False - friction_range: list[float] = field(default_factory=lambda: [0.3, 1.2]) - friction_geom_pattern: str = r"^(left|right)_foot[1-7]_collision$" - - randomize_joint_default_pos: bool = False - joint_default_pos_range: list[float] = field(default_factory=lambda: [-0.01, 0.01]) - - -@dataclass -class G1MotionTrackingCfg(G1BaseCfg): - """Configuration for G1 motion tracking environment.""" +from ..common.config import ( + Domain_Rand, + DomainRand, + MotionTrackingCfg, + MotionTrackingDeployEnvCfg, + PoseRandomization, + VelocityRandomization, + _zero_pose_randomization, + _zero_velocity_randomization, +) +from ..common.domain_randomization import MotionTrackingDomainRandomizationProvider +from ..common.reset import build_motion_reference_state +from ..common.rewards import RewardConfig +from ..common.tracking import MotionTrackingDeployEnv, MotionTrackingEnv - scene: SceneCfg = field( - default_factory=lambda: SceneCfg( - model_file=str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") - ) - ) - # Kept at the historical single-clip default for backward compatibility. - motion_file: str | list[str] = str( - ASSETS_ROOT_PATH / "motions" / "g1" / "dance1_subject2_part.npz" - ) - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "gangnam_style.npz") - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "fight1_subject5_from_csv.npz") #LAFAN - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "dance_basic_slide_180_R_loop_001__A322_M.npz") #LAFAN - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "playing_violin_R_003__A327_from_csv.npz") #Seed - anchor_body_name: str = "torso_link" - body_names: tuple[str, ...] = ( - "pelvis", - "left_hip_roll_link", - "left_knee_link", - "left_ankle_roll_link", - "right_hip_roll_link", - "right_knee_link", - "right_ankle_roll_link", - "torso_link", - "left_shoulder_roll_link", - "left_elbow_link", - "left_wrist_yaw_link", - "right_shoulder_roll_link", - "right_elbow_link", - "right_wrist_yaw_link", - ) - sampling_mode: Literal["start", "clip_start", "uniform", "adaptive", "mixed"] = "adaptive" - sampling_start_ratio: float = 0.0 - truncate_on_clip_end: bool = False - max_episode_seconds: float = 10.0 - reward_config: RewardConfig = field(default_factory=RewardConfig) - pose_randomization: PoseRandomization = field(default_factory=PoseRandomization) - velocity_randomization: VelocityRandomization = field(default_factory=VelocityRandomization) - domain_rand: Domain_Rand = field(default_factory=Domain_Rand) - joint_position_range: tuple[float, float] = (-0.1, 0.1) - # Termination thresholds - anchor_pos_z_threshold: float = 0.25 - anchor_ori_threshold: float = 0.8 - ee_body_pos_z_threshold: float = 0.25 - ee_body_names: tuple[str, ...] = ( - "left_ankle_roll_link", - "right_ankle_roll_link", - "left_wrist_yaw_link", - "right_wrist_yaw_link", - ) - undesired_contact_z_threshold: float = 0.05 - terminate_on_undesired_contacts: bool = False - numba_acceleration: bool = False - numba_num_threads: int | None = None +# ── backward-compatible aliases (historical G1* symbol names) ──────── +G1MotionTrackingCfg = MotionTrackingCfg +G1MotionTrackingDomainRandomizationProvider = MotionTrackingDomainRandomizationProvider +_build_motion_reference_state = build_motion_reference_state @registry.envcfg("G1MotionTracking") @dataclass -class G1MotionTrackingEnvCfg(G1MotionTrackingCfg): +class G1MotionTrackingEnvCfg(MotionTrackingCfg): """Registered configuration for G1 motion tracking.""" pass @@ -203,1259 +44,46 @@ class G1MotionTrackingEnvCfg(G1MotionTrackingCfg): @registry.envcfg("G1MotionTrackingDeploy") @dataclass -class G1MotionTrackingDeployEnvCfg(G1MotionTrackingCfg): +class G1MotionTrackingDeployEnvCfg(MotionTrackingDeployEnvCfg): """Registered deploy configuration for G1 motion tracking.""" pass -def _build_motion_reference_state( - env: Any, env_ids: np.ndarray, motion_data: MotionData -) -> tuple[np.ndarray, np.ndarray]: - dtype = get_global_dtype() - num_reset = len(env_ids) - - root_pos = motion_data.body_pos_w[:, 0].copy() - root_ori = motion_data.body_quat_w[:, 0].copy() - root_lin_vel = motion_data.body_lin_vel_w[:, 0].copy() - root_ang_vel = motion_data.body_ang_vel_w[:, 0].copy() - joint_pos = motion_data.joint_pos.copy() - joint_vel = motion_data.joint_vel.copy() - - pose_rand = env.cfg.pose_randomization - pose_ranges = [ - (pose_rand.x[0], pose_rand.x[1]), - (pose_rand.y[0], pose_rand.y[1]), - (pose_rand.z[0], pose_rand.z[1]), - (pose_rand.roll[0], pose_rand.roll[1]), - (pose_rand.pitch[0], pose_rand.pitch[1]), - (pose_rand.yaw[0], pose_rand.yaw[1]), - ] - pose_samples = np.array( - [[np.random.uniform(low, high) for low, high in pose_ranges] for _ in range(num_reset)], - dtype=dtype, - ) - root_pos += pose_samples[:, 0:3] - root_ori = np_quat_mul( - np_quat_from_euler_xyz(pose_samples[:, 3], pose_samples[:, 4], pose_samples[:, 5]), - root_ori, - ) - - vel_rand = env.cfg.velocity_randomization - vel_ranges = [ - (vel_rand.x[0], vel_rand.x[1]), - (vel_rand.y[0], vel_rand.y[1]), - (vel_rand.z[0], vel_rand.z[1]), - (vel_rand.roll[0], vel_rand.roll[1]), - (vel_rand.pitch[0], vel_rand.pitch[1]), - (vel_rand.yaw[0], vel_rand.yaw[1]), - ] - vel_samples = np.array( - [[np.random.uniform(low, high) for low, high in vel_ranges] for _ in range(num_reset)], - dtype=dtype, - ) - root_lin_vel += vel_samples[:, :3] - root_ang_vel += vel_samples[:, 3:] - - joint_pos += np_sample_uniform( - env.cfg.joint_position_range[0], - env.cfg.joint_position_range[1], - joint_pos.shape, - dtype=np.float32, - ) - joint_range = env._get_joint_range() - if joint_range is not None: - joint_pos = np.clip(joint_pos, joint_range[:, 0], joint_range[:, 1]) - - qpos = np.tile(env._init_qpos, (num_reset, 1)) - qvel = np.tile(env._init_qvel, (num_reset, 1)) - qpos[:, 0:3] = root_pos - qpos[:, 3:7] = root_ori - qpos[:, 7:] = joint_pos - - qvel[:, 0:3] = root_lin_vel - qvel[:, 3:6] = np_quat_apply(np_quat_inv(root_ori), root_ang_vel) - qvel[:, 6:] = joint_vel - return qpos, qvel - - -class G1MotionTrackingDomainRandomizationProvider(DomainRandomizationProvider): - def __init__( - self, - *, - base_kp: np.ndarray | None = None, - base_kd: np.ndarray | None = None, - base_geom_friction: np.ndarray | None = None, - foot_geom_ids: np.ndarray | None = None, - ) -> None: - self._base_kp = base_kp - self._base_kd = base_kd - self._base_geom_friction = base_geom_friction - self._foot_geom_ids = foot_geom_ids - self._last_reset_observation_timing_ms: dict[str, float] = {} - - @property - def last_reset_observation_timing_ms(self) -> dict[str, float]: - return dict(self._last_reset_observation_timing_ms) - - def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: - validate_common_reset_randomization( - env, capabilities, base_kp=self._base_kp, base_kd=self._base_kd - ) - validate_interval_push_support(env, capabilities) - if getattr(env.cfg.domain_rand, "randomize_geom_friction", False): - if not capabilities.supports_reset_term(RESET_TERM_GEOM_FRICTION): - raise NotImplementedError( - f"{env._backend.backend_type} backend does not support " - "geom-friction reset randomization" - ) - if ( - self._base_geom_friction is None - or self._foot_geom_ids is None - or self._foot_geom_ids.size == 0 - ): - raise ValueError("randomize_geom_friction=True but provider has no foot geom IDs") - - def build_interval_randomization_plan( - self, env: Any, step_counter: int - ) -> IntervalRandomizationPlan | None: - return build_interval_push_plan(env, step_counter) - - def build_reset_plan(self, env: Any, env_ids: np.ndarray) -> ResetPlan: - num_reset = len(env_ids) - motion_frames = env.motion_sampler.sample_frames(env_ids) - motion_data = env.motion_loader.get_motion_at_frame(motion_frames) - qpos, qvel = _build_motion_reference_state(env, env_ids, motion_data) - - info_updates = { - "current_actions": zero_actions(num_reset, env._num_action), - "last_actions": zero_actions(num_reset, env._num_action), - } - randomization = build_common_reset_randomization( - env, num_reset, base_kp=self._base_kp, base_kd=self._base_kd - ) - - dr_cfg = env.cfg.domain_rand - if getattr(dr_cfg, "randomize_geom_friction", False): - assert self._base_geom_friction is not None - assert self._foot_geom_ids is not None - payload = randomization or ResetRandomizationPayload() - low, high = dr_cfg.friction_range - scale = np.random.uniform(low, high, size=(num_reset, 1)).astype(np.float64) - geom_friction = np.broadcast_to( - self._base_geom_friction, - (num_reset, *self._base_geom_friction.shape), - ).copy() - geom_friction[:, self._foot_geom_ids, 0] = scale * np.ones( - (1, self._foot_geom_ids.size) - ) - payload.geom_friction = geom_friction - randomization = payload - - if getattr(dr_cfg, "randomize_joint_default_pos", False): - low, high = dr_cfg.joint_default_pos_range - info_updates["default_dof_pos_bias"] = np.random.uniform( - low, high, size=(num_reset, env._num_action) - ).astype(get_global_dtype()) - - return ResetPlan( - env_ids=env_ids, - qpos=qpos, - qvel=qvel, - info_updates=info_updates, - randomization=randomization, - ) - - def build_reset_observation( - self, env: Any, env_ids: np.ndarray, info_updates: dict[str, Any] - ) -> dict[str, np.ndarray]: - obs_t0 = time.perf_counter() - - t0 = time.perf_counter() - motion_data = env.motion_loader.get_motion_at_frame( - env.motion_sampler.current_frames[env_ids] - ) - motion_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - linvel = env.get_local_linvel()[env_ids] - linvel_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - gyro = env.get_gyro()[env_ids] - gyro_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - dof_pos = env.get_dof_pos()[env_ids] - dof_pos_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - dof_vel = env.get_dof_vel()[env_ids] - dof_vel_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - robot_body_pos_w, robot_body_quat_w = env._backend.get_body_pose_w_rows( - env_ids, env.body_ids - ) - body_pose_ms = (time.perf_counter() - t0) * 1000.0 - - obs_info = dict(info_updates) - default_dof_pos_bias = info_updates.get("default_dof_pos_bias") - if isinstance(default_dof_pos_bias, np.ndarray): - obs_info["default_dof_pos_bias"] = default_dof_pos_bias - obs_info["env_ids"] = env_ids - - t0 = time.perf_counter() - obs = cast( - dict[str, np.ndarray], - env._compute_obs( - obs_info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ), - ) - compute_obs_ms = (time.perf_counter() - t0) * 1000.0 - - total_ms = (time.perf_counter() - obs_t0) * 1000.0 - getters_ms = linvel_ms + gyro_ms + dof_pos_ms + dof_vel_ms + body_pose_ms - self._last_reset_observation_timing_ms = { - "dr_reset_observation_getters_ms": getters_ms, - "dr_reset_obs_get_motion_ms": motion_ms, - "dr_reset_obs_get_local_linvel_ms": linvel_ms, - "dr_reset_obs_get_gyro_ms": gyro_ms, - "dr_reset_obs_get_dof_pos_ms": dof_pos_ms, - "dr_reset_obs_get_dof_vel_ms": dof_vel_ms, - "dr_reset_obs_get_body_pose_ms": body_pose_ms, - "dr_reset_observation_compute_obs_ms": compute_obs_ms, - "dr_reset_observation_internal_gap_ms": ( - total_ms - motion_ms - getters_ms - compute_obs_ms - ), - } - return obs - - @registry.env("G1MotionTracking", sim_backend="mujoco") @registry.env("G1MotionTracking", sim_backend="motrix") -class G1MotionTrackingEnv(G1BaseEnv): +class G1MotionTrackingEnv(MotionTrackingEnv): """G1 Motion Tracking Environment.""" - _cfg: G1MotionTrackingCfg - - def __init__(self, cfg: G1MotionTrackingCfg, num_envs=1, backend_type="mujoco"): - if not cfg.motion_file: - raise ValueError("motion_file must be specified in config") - - backend = create_backend( - backend_type, - cfg.scene, - num_envs, - cfg.sim_dt, - base_name=cfg.asset.base_name, - push_body_name=cfg.domain_rand.push_body_name, - add_body_sensors=True, - **env_backend_kwargs(cfg), - ) - super().__init__(cfg, backend, num_envs) - - # Resolve body IDs for backend querying and motion-file indexing. - self.body_ids = self._backend.get_body_ids(cfg.body_names) - motion_body_ids = self._backend.get_motion_body_ids(cfg.body_names) - - self.anchor_body_idx = cfg.body_names.index(cfg.anchor_body_name) - - # Get end-effector body indices for termination - self.ee_body_indices = np.array( - [cfg.body_names.index(name) for name in cfg.ee_body_names], dtype=np.int32 - ) - self._has_ee_body_indices = bool(self.ee_body_indices.size) - - # Get non-EE body indices for undesired contact penalty - ee_set = set(cfg.ee_body_names) - self.undesired_contact_body_indices = np.array( - [i for i, name in enumerate(cfg.body_names) if name not in ee_set], - dtype=np.int32, - ) - self._has_undesired_contact_body_indices = bool(self.undesired_contact_body_indices.size) - - # Load motion data - self.motion_loader = MotionLoader(cfg.motion_file, body_indices=motion_body_ids) - self.motion_sampler = MotionSampler( - self.motion_loader, - mode=cfg.sampling_mode, - num_envs=num_envs, - start_ratio=cfg.sampling_start_ratio, - ) - needs_kp_kd = cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd - needs_friction = getattr(cfg.domain_rand, "randomize_geom_friction", False) - base_kp = base_kd = None - if needs_kp_kd: - base_kp, base_kd = backend.get_actuator_gains() - base_geom_friction = None - foot_geom_ids = None - if needs_friction: - import re as _re - - base_geom_friction = backend.get_geom_friction() - geom_names = backend.get_geom_names() - pattern = _re.compile(cfg.domain_rand.friction_geom_pattern) - foot_geom_ids = np.asarray( - [i for i, name in enumerate(geom_names) if name and pattern.match(name)], - dtype=np.int64, - ) - if foot_geom_ids.size == 0: - raise ValueError( - "friction_geom_pattern " - f"'{cfg.domain_rand.friction_geom_pattern}' did not match any geom" - ) - dr_provider = G1MotionTrackingDomainRandomizationProvider( - base_kp=base_kp, - base_kd=base_kd, - base_geom_friction=base_geom_friction, - foot_geom_ids=foot_geom_ids, - ) - self._init_domain_randomization(dr_provider) - - dtype = get_global_dtype() - n_body = len(cfg.body_names) - self._n_motion_bodies = n_body - self._actor_obs_width = self._actor_obs_dim(self._num_action) - self._critic_base_obs_width = self._critic_base_obs_dim(self._num_action) - self._critic_obs_width = self._critic_base_obs_width + n_body * 9 - self._copy_body_state_w = self._backend.copy_body_state_w - - # Buffers for relative body transforms - self.body_pos_relative_w = np.zeros((num_envs, n_body, 3), dtype=dtype) - self.body_quat_relative_w = np.zeros((num_envs, n_body, 4), dtype=dtype) - self.body_quat_relative_w[:, :, 0] = 1.0 # Initialize to identity quaternion - self._motion_data_buffer = ( - self.motion_loader.make_motion_data_buffer(num_envs) - if hasattr(self.motion_loader, "make_motion_data_buffer") - else None - ) - self._zero_actions = np.zeros((num_envs, self._num_action), dtype=dtype) - self._joint_range = self._backend.get_joint_range() - if self._joint_range is not None: - self._joint_range = np.asarray(self._joint_range, dtype=dtype) - self._joint_lower = self._joint_range[:, 0] - self._joint_upper = self._joint_range[:, 1] - else: - self._joint_lower = None - self._joint_upper = None - self._delta_pos_w = np.empty((num_envs, 3), dtype=dtype) - self._delta_ori_w = np.empty((num_envs, 4), dtype=dtype) - self._motion_anchor_pos_b = np.empty((num_envs, 3), dtype=dtype) - self._motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) - self._motion_command = np.empty((num_envs, self._num_action * 2), dtype=dtype) - self._joint_pos_rel = np.empty((num_envs, self._num_action), dtype=dtype) - self._robot_body_pos_w = np.empty((num_envs, n_body, 3), dtype=dtype) - self._robot_body_quat_w = np.empty((num_envs, n_body, 4), dtype=dtype) - self._robot_body_lin_vel_w = np.empty((num_envs, n_body, 3), dtype=dtype) - self._robot_body_ang_vel_w = np.empty((num_envs, n_body, 3), dtype=dtype) - self._quat_error_w = np.empty((num_envs, n_body), dtype=dtype) - self._quat_error_x = np.empty((num_envs, n_body), dtype=dtype) - self._body_vec_error = np.empty((num_envs, n_body, 3), dtype=dtype) - self._body_vec_tmp = np.empty((num_envs, n_body, 3), dtype=dtype) - self._joint_error = np.empty((num_envs, self._num_action), dtype=dtype) - self._joint_error_upper = np.empty((num_envs, self._num_action), dtype=dtype) - self._env_error = np.empty((num_envs,), dtype=dtype) - self._env_error2 = np.empty((num_envs,), dtype=dtype) - self._reward_term = np.empty((num_envs,), dtype=dtype) - self._weighted_reward = np.empty((num_envs,), dtype=dtype) - self._terminated = np.empty((num_envs,), dtype=bool) - self._env_bool = np.empty((num_envs,), dtype=bool) - self._ee_pos_error_z = np.empty((num_envs, self.ee_body_indices.size), dtype=dtype) - self._ee_terminated = np.empty((num_envs, self.ee_body_indices.size), dtype=bool) - self._undesired_contact_mask = np.empty( - (num_envs, self.undesired_contact_body_indices.size), dtype=bool - ) - - self._enable_reward_log = True - self._init_reward_functions() - self._active_reward_fns = { - name: reward_fn - for name, reward_fn in self._reward_fns.items() - if self._reward_term_is_active(name) - } - self._numba_accelerator = None - if cfg.numba_acceleration: - from unilab.envs.motion_tracking.g1.motion_tracking_numba import ( - G1MotionTrackingNumbaAccelerator, - ) - - self._numba_accelerator = G1MotionTrackingNumbaAccelerator.from_env( - self, num_threads=cfg.numba_num_threads - ) - self._clip_end_truncated = np.zeros((num_envs,), dtype=bool) - - def _effective_default_angles(self, env_ids: np.ndarray | None = None) -> np.ndarray: - """Return default_angles with per-episode joint-default-pos bias applied.""" - state = getattr(self, "_state", None) - if state is not None: - bias = state.info.get("default_dof_pos_bias") - if bias is not None: - if env_ids is not None: - return self.default_angles + bias[env_ids] - return self.default_angles + bias - return self.default_angles - - def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: - state.info["last_actions"] = state.info.get("current_actions", np.zeros_like(actions)) - state.info["current_actions"] = actions - exec_actions = ( - state.info["last_actions"] - if self._cfg.control_config.simulate_action_latency - else actions - ) - bias = state.info.get("default_dof_pos_bias") - base = self.default_angles + bias if bias is not None else self.default_angles - ctrl: np.ndarray = exec_actions * self._cfg.control_config.action_scale + base - return ctrl - - def _resample_reference_state(self, env_ids: np.ndarray) -> None: - motion_frames = self.motion_sampler.sample_frames(env_ids) - motion_data = self.motion_loader.get_motion_at_frame(motion_frames) - qpos, qvel = _build_motion_reference_state(self, env_ids, motion_data) - self._backend.set_state(env_ids, qpos, qvel) - - def _refresh_observation_rows( - self, obs: dict[str, np.ndarray], info: dict, env_ids: np.ndarray - ) -> None: - motion_data = self.motion_loader.get_motion_at_frame( - self.motion_sampler.current_frames[env_ids] - ) - row_ids = np.asarray(env_ids, dtype=np.intp) - linvel = self._backend.get_sensor_data_rows(self._cfg.sensor.local_linvel, row_ids) - gyro = self._backend.get_sensor_data_rows(self._cfg.sensor.gyro, row_ids) - dof_pos = self.get_dof_pos()[row_ids] - dof_vel = self.get_dof_vel()[row_ids] - robot_body_pos_w, robot_body_quat_w = self._backend.get_body_pose_w_rows( - row_ids, self.body_ids - ) - - obs_info: dict[str, Any] = {} - current_actions = info.get("current_actions") - if isinstance(current_actions, np.ndarray): - obs_info["current_actions"] = current_actions[env_ids] - obs_info["env_ids"] = env_ids - - refreshed_obs = self._compute_obs( - obs_info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - for key, value in refreshed_obs.items(): - if value.shape[0] == len(env_ids): - obs[key][env_ids] = value - else: - obs[key][env_ids] = value[env_ids] - - def _get_body_pose_w(self) -> tuple[np.ndarray, np.ndarray]: - return self._backend.get_body_pose_w(self.body_ids) - - def _get_body_state_w(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - copy_body_state_w = self._copy_body_state_w - if copy_body_state_w is not None: - return copy_body_state_w( - self.body_ids, - self._robot_body_pos_w, - self._robot_body_quat_w, - self._robot_body_lin_vel_w, - self._robot_body_ang_vel_w, - ) - robot_body_pos_w, robot_body_quat_w = self._get_body_pose_w() - robot_body_lin_vel_w, robot_body_ang_vel_w = self._backend.get_body_vel_w(self.body_ids) - return ( - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - ) - - def _get_joint_range(self) -> np.ndarray | None: - return self._joint_range - - def _get_current_motion(self) -> MotionData: - if self._motion_data_buffer is None: - return self.motion_sampler.get_current_motion() - return self.motion_sampler.get_current_motion(self._motion_data_buffer) - - @property - def obs_groups_spec(self) -> dict[str, int]: - # Actor: command(2n) + motion_anchor_pos_b(3) + motion_anchor_ori_b(6) - # + linvel(3) + gyro(3) + joint_pos(n) + joint_vel(n) + actions(n) - # Critic mirrors BeyondMimic physical terms without actor observation noise: - # command, motion anchor, robot body pos/ori, linvel, gyro, joints, actions. - n = self._num_action - actor_width = getattr(self, "_actor_obs_width", self._actor_obs_dim(n)) - critic_width = getattr( - self, - "_critic_obs_width", - self._critic_base_obs_dim(n) + len(self._cfg.body_names) * 9, - ) - return {"obs": actor_width, "critic": critic_width} - - def _actor_obs_dim(self, n: int) -> int: - return 3 + 6 + 3 + 3 + n * 5 - - def _critic_base_obs_dim(self, n: int) -> int: - return 3 + 6 + 3 + 3 + n * 5 - - def _build_actor_obs( - self, - *, - command: np.ndarray, - motion_anchor_pos_b: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_linvel: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, - ) -> np.ndarray: - num_envs = command.shape[0] - n_action = noisy_joint_pos_rel.shape[1] - actor_obs = np.empty((num_envs, self._actor_obs_dim(n_action)), dtype=get_global_dtype()) - offset = 0 - actor_obs[:, offset : offset + command.shape[1]] = command - offset += command.shape[1] - actor_obs[:, offset : offset + 3] = motion_anchor_pos_b - offset += 3 - actor_obs[:, offset : offset + 6] = motion_anchor_ori_b - offset += 6 - actor_obs[:, offset : offset + 3] = noisy_linvel - offset += 3 - actor_obs[:, offset : offset + 3] = noisy_gyro - offset += 3 - actor_obs[:, offset : offset + n_action] = noisy_joint_pos_rel - offset += n_action - actor_obs[:, offset : offset + n_action] = noisy_dof_vel - offset += n_action - actor_obs[:, offset : offset + n_action] = last_actions - return actor_obs - - def _init_reward_functions(self): - self._reward_fns = { - "motion_global_root_pos": self._reward_motion_global_root_pos, - "motion_global_root_ori": self._reward_motion_global_root_ori, - "motion_body_pos": self._reward_motion_body_pos, - "motion_body_ori": self._reward_motion_body_ori, - "motion_body_lin_vel": self._reward_motion_body_lin_vel, - "motion_body_ang_vel": self._reward_motion_body_ang_vel, - "motion_ee_body_pos_z": self._reward_motion_ee_body_pos_z, - "motion_joint_pos": self._reward_motion_joint_pos, - "motion_joint_vel": self._reward_motion_joint_vel, - "action_rate_l2": self._reward_action_rate_l2, - "joint_limit": self._reward_joint_limit, - "undesired_contacts": self._reward_undesired_contacts, - } - - def _reward_term_is_active(self, name: str) -> bool: - if name == "joint_limit": - return self._joint_lower is not None and self._joint_upper is not None - if name == "undesired_contacts": - return self._has_undesired_contact_body_indices - if name == "motion_ee_body_pos_z": - return self._has_ee_body_indices - return True - - def update_state(self, state: NpEnvState) -> NpEnvState: - self._clip_end_truncated.fill(False) - - # Get current motion data - motion_data = self._get_current_motion() - - # Get robot state - linvel = self.get_local_linvel() - gyro = self.get_gyro() - dof_pos = self.get_dof_pos() - dof_vel = self.get_dof_vel() - - # Get body states - ( - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - ) = self._get_body_state_w() - - numba_accelerator = getattr(self, "_numba_accelerator", None) - if numba_accelerator is not None: - noise_cfg = self._cfg.noise_config - numba_result = numba_accelerator.compute_update_state( - info=state.info, - motion_data=motion_data, - linvel=linvel, - gyro=gyro, - dof_pos=dof_pos, - dof_vel=dof_vel, - robot_body_pos_w=robot_body_pos_w, - robot_body_quat_w=robot_body_quat_w, - robot_body_lin_vel_w=robot_body_lin_vel_w, - robot_body_ang_vel_w=robot_body_ang_vel_w, - ref_body_pos_w=self.body_pos_relative_w, - ref_body_quat_w=self.body_quat_relative_w, - motion_anchor_pos_b=self._motion_anchor_pos_b, - motion_anchor_ori_b=self._motion_anchor_ori_b, - joint_pos_rel=self._joint_pos_rel, - scales=self._cfg.reward_config.scales, - enable_log=self._enable_reward_log, - noise_level=noise_cfg.level, - noise_scale_linvel=noise_cfg.scale_linvel, - noise_scale_gyro=noise_cfg.scale_gyro, - noise_scale_joint_angle=noise_cfg.scale_joint_angle, - noise_scale_joint_vel=noise_cfg.scale_joint_vel, - ) - terminated = numba_result.terminated - reward = numba_result.reward - obs = numba_result.obs - state.info["log"] = numba_result.log - else: - # Compute relative body transforms (for observations and rewards) - self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) - - # Compute terminations - terminated = self._compute_terminations( - motion_data, robot_body_pos_w, robot_body_quat_w - ) - - # Compute reward - reward = self._compute_reward( - state.info, - motion_data, - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - dof_pos, - dof_vel, - ) - - # Compute observations - obs = self._compute_obs( - state.info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - - # Update failure statistics for adaptive sampling - self.motion_sampler.update_failure_stats(terminated) - - # Advance motion frames - done_env_ids = self.motion_sampler.step() - if len(done_env_ids) > 0: - if self._cfg.truncate_on_clip_end: - self._clip_end_truncated[done_env_ids] = True - else: - # Match BeyondMimic: clip boundaries are command resampling points, not - # episode boundaries; sync the simulated robot to the new reference. - resample_env_ids = done_env_ids[~terminated[done_env_ids]] - if len(resample_env_ids) > 0: - self._resample_reference_state(resample_env_ids) - self._refresh_observation_rows(obs, state.info, resample_env_ids) - - return state.replace(obs=obs, reward=reward, terminated=terminated) - - def _compute_truncated(self, state: NpEnvState) -> np.ndarray: - truncated = super()._compute_truncated(state) - clip_end_only = getattr(self, "_env_bool", None) - if clip_end_only is None or clip_end_only.shape != (self._num_envs,): - clip_end_only = np.empty((self._num_envs,), dtype=bool) - self._env_bool = clip_end_only - np.logical_not(state.terminated, out=clip_end_only) - np.logical_and(self._clip_end_truncated, clip_end_only, out=clip_end_only) - np.logical_or(truncated, clip_end_only, out=truncated) - return truncated - - def _update_relative_transforms( - self, motion_data, robot_body_pos_w: np.ndarray, robot_body_quat_w: np.ndarray - ): - """Update relative body transforms for tracking.""" - # Get anchor states - anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] - anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] - - # Compute delta transform: keep robot's XY position, use motion's Z height - # and apply yaw-only rotation difference. - delta_pos_w = self._delta_pos_w - delta_pos_w[:] = robot_anchor_pos_w - delta_pos_w[:, 2] = anchor_pos_w[:, 2] - - # Compute yaw-only rotation difference, equivalent to - # np_yaw_quat(np_quat_mul(robot_anchor_quat_w, np_quat_inv(anchor_quat_w))). - delta_ori_w = self._delta_ori_w - rw, rx, ry, rz = ( - robot_anchor_quat_w[:, 0], - robot_anchor_quat_w[:, 1], - robot_anchor_quat_w[:, 2], - robot_anchor_quat_w[:, 3], - ) - aw, ax, ay, az = ( - anchor_quat_w[:, 0], - anchor_quat_w[:, 1], - anchor_quat_w[:, 2], - anchor_quat_w[:, 3], - ) - qw = rw * aw + rx * ax + ry * ay + rz * az - qx = -rw * ax + rx * aw - ry * az + rz * ay - qy = -rw * ay + rx * az + ry * aw - rz * ax - qz = -rw * az - rx * ay + ry * ax + rz * aw - half_yaw = 0.5 * np.arctan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz)) - np.cos(half_yaw, out=delta_ori_w[:, 0]) - delta_ori_w[:, 1:3] = 0.0 - np.sin(half_yaw, out=delta_ori_w[:, 3]) - - dw1 = delta_ori_w[:, 0] - dz1 = delta_ori_w[:, 3] - dw = dw1[:, None] - dz = dz1[:, None] - mw = motion_data.body_quat_w[..., 0] - mx = motion_data.body_quat_w[..., 1] - my = motion_data.body_quat_w[..., 2] - mz = motion_data.body_quat_w[..., 3] - out_quat = self.body_quat_relative_w - out_quat[..., 0] = dw * mw - out_quat[..., 0] -= dz * mz - out_quat[..., 1] = dw * mx - out_quat[..., 1] -= dz * my - out_quat[..., 2] = dw * my - out_quat[..., 2] += dz * mx - out_quat[..., 3] = dw * mz - out_quat[..., 3] += dz * mw - - rel_pos = self._body_vec_error - vx = rel_pos[..., 0] - vy = rel_pos[..., 1] - vz = rel_pos[..., 2] - np.subtract(motion_data.body_pos_w[..., 0], anchor_pos_w[:, None, 0], out=vx) - np.subtract(motion_data.body_pos_w[..., 1], anchor_pos_w[:, None, 1], out=vy) - np.subtract(motion_data.body_pos_w[..., 2], anchor_pos_w[:, None, 2], out=vz) - - yaw_cross = self._env_error - yaw_z2 = self._reward_term - np.multiply(dw1, dz1, out=yaw_cross) - yaw_cross *= 2.0 - np.square(dz1, out=yaw_z2) - yaw_z2 *= 2.0 - yaw_cross_2d = yaw_cross[:, None] - yaw_z2_2d = yaw_z2[:, None] - - out_pos = self.body_pos_relative_w - out_pos[..., 0] = vx - out_pos[..., 0] -= yaw_cross_2d * vy - out_pos[..., 0] -= yaw_z2_2d * vx - out_pos[..., 0] += delta_pos_w[:, None, 0] - out_pos[..., 1] = vy - out_pos[..., 1] += yaw_cross_2d * vx - out_pos[..., 1] -= yaw_z2_2d * vy - out_pos[..., 1] += delta_pos_w[:, None, 1] - out_pos[..., 2] = vz - out_pos[..., 2] += delta_pos_w[:, None, 2] - - def _compute_terminations( - self, - motion_data, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> np.ndarray: - """Compute termination conditions.""" - terminated = self._terminated - terminated.fill(False) - - # Anchor position error (Z-axis only) - anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] - np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=self._env_error) - np.abs(self._env_error, out=self._env_error) - np.greater(self._env_error, self._cfg.anchor_pos_z_threshold, out=self._env_bool) - terminated |= self._env_bool - - # Anchor orientation error (gravity direction). The gravity-z difference - # is bounded by 2 for unit quaternions, so huge thresholds disable this - # termination without doing the per-step math. - if self._cfg.anchor_ori_threshold < 2.0: - anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] - motion_gravity_z_b = np_gravity_z_in_body_from_quat(anchor_quat_w) - robot_gravity_z_b = np_gravity_z_in_body_from_quat(robot_anchor_quat_w) - np.subtract(motion_gravity_z_b, robot_gravity_z_b, out=self._env_error) - np.abs(self._env_error, out=self._env_error) - np.greater(self._env_error, self._cfg.anchor_ori_threshold, out=self._env_bool) - terminated |= self._env_bool - - # End-effector position error (Z-axis only) - if self._has_ee_body_indices: - np.subtract( - self.body_pos_relative_w[:, self.ee_body_indices, 2], - robot_body_pos_w[:, self.ee_body_indices, 2], - out=self._ee_pos_error_z, - ) - np.abs(self._ee_pos_error_z, out=self._ee_pos_error_z) - np.greater( - self._ee_pos_error_z, - self._cfg.ee_body_pos_z_threshold, - out=self._ee_terminated, - ) - np.logical_or.reduce(self._ee_terminated, axis=1, out=self._env_bool) - terminated |= self._env_bool - - if self._cfg.terminate_on_undesired_contacts and self._has_undesired_contact_body_indices: - body_z = robot_body_pos_w[:, self.undesired_contact_body_indices, 2] - np.less( - body_z, - self._cfg.undesired_contact_z_threshold, - out=self._undesired_contact_mask, - ) - np.logical_or.reduce(self._undesired_contact_mask, axis=-1, out=self._env_bool) - terminated |= self._env_bool - - return terminated - - def _write_body_pos_in_anchor_frame( - self, - anchor_pos: np.ndarray, - anchor_quat: np.ndarray, - body_pos: np.ndarray, - out: np.ndarray, - ) -> None: - aw = anchor_quat[:, None, 0] - ax = anchor_quat[:, None, 1] - ay = anchor_quat[:, None, 2] - az = anchor_quat[:, None, 3] - - num_envs, n_body = body_pos.shape[:2] - rel_pos = self._body_vec_error[:num_envs, :n_body] - - vx = rel_pos[..., 0] - vy = rel_pos[..., 1] - vz = rel_pos[..., 2] - np.subtract(body_pos[..., 0], anchor_pos[:, None, 0], out=vx) - np.subtract(body_pos[..., 1], anchor_pos[:, None, 1], out=vy) - np.subtract(body_pos[..., 2], anchor_pos[:, None, 2], out=vz) - - tx = 2 * (az * vy - ay * vz) - ty = 2 * (ax * vz - az * vx) - tz = 2 * (ay * vx - ax * vy) - - out[..., 0] = vx + aw * tx + az * ty - ay * tz - out[..., 1] = vy + aw * ty + ax * tz - az * tx - out[..., 2] = vz + aw * tz + ay * tx - ax * ty - - def _write_body_ori6_in_anchor_frame( - self, - anchor_quat: np.ndarray, - body_quat: np.ndarray, - out: np.ndarray, - ) -> None: - aw = anchor_quat[:, None, 0] - ax = anchor_quat[:, None, 1] - ay = anchor_quat[:, None, 2] - az = anchor_quat[:, None, 3] - bw = body_quat[..., 0] - bx = body_quat[..., 1] - by = body_quat[..., 2] - bz = body_quat[..., 3] - - rw = aw * bw + ax * bx + ay * by + az * bz - rx = aw * bx - ax * bw - ay * bz + az * by - ry = aw * by + ax * bz - ay * bw - az * bx - rz = aw * bz - ax * by + ay * bx - az * bw - - out[..., 0] = 1 - 2 * (ry * ry + rz * rz) - out[..., 1] = 2 * (rx * ry - rw * rz) - out[..., 2] = 2 * (rx * ry + rw * rz) - out[..., 3] = 1 - 2 * (rx * rx + rz * rz) - out[..., 4] = 2 * (rx * rz - rw * ry) - out[..., 5] = 2 * (ry * rz + rw * rx) - - def _compute_obs( - self, - info: dict, - motion_data, - linvel: np.ndarray, - gyro: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - ) -> dict[str, np.ndarray]: - """Compute observations as dict with actor and critic groups.""" - num_envs = linvel.shape[0] - dtype = get_global_dtype() - n_action = dof_pos.shape[1] - n_body = self._n_motion_bodies - - # Get anchor states - anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] - anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] - - # Motion anchor pose in robot frame - if num_envs == self._num_envs: - motion_anchor_pos_b = self._motion_anchor_pos_b - motion_anchor_ori_b = self._motion_anchor_ori_b - joint_pos_rel = self._joint_pos_rel - zero_actions = self._zero_actions - else: - motion_anchor_pos_b = np.empty((num_envs, 3), dtype=dtype) - motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) - joint_pos_rel = np.empty((num_envs, n_action), dtype=dtype) - zero_actions = np.zeros((num_envs, n_action), dtype=dtype) - np_write_relative_anchor_transform_pos_rot6d( - robot_anchor_pos_w, - robot_anchor_quat_w, - anchor_pos_w, - anchor_quat_w, - motion_anchor_pos_b, - motion_anchor_ori_b, - ) - - # Joint positions and velocities - bias = info.get("default_dof_pos_bias") - effective_default = self.default_angles + bias if bias is not None else self.default_angles - np.subtract(dof_pos, effective_default, out=joint_pos_rel) - last_actions = info.get("current_actions") - if not isinstance(last_actions, np.ndarray): - last_actions = zero_actions - - if num_envs == self._num_envs: - command = self._motion_command - else: - command = np.empty((num_envs, n_action * 2), dtype=dtype) - command[:, :n_action] = motion_data.joint_pos - command[:, n_action : n_action * 2] = motion_data.joint_vel - - # Per-step observation noise on sensor channels (actor only). - # Critic uses the clean originals — asymmetric actor–critic contract. - noise_cfg = self._cfg.noise_config - noise_enabled = noise_cfg.level > 0.0 - if noise_enabled: - linvel_actor = self._obs_noise(linvel, noise_cfg.scale_linvel) - gyro_actor = self._obs_noise(gyro, noise_cfg.scale_gyro) - joint_pos_actor = self._obs_noise(joint_pos_rel, noise_cfg.scale_joint_angle) - dof_vel_actor = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) - else: - linvel_actor = linvel - gyro_actor = gyro - joint_pos_actor = joint_pos_rel - dof_vel_actor = dof_vel - - # Actor observations (noisy proprioception) - actor_obs = self._build_actor_obs( - command=command, - motion_anchor_pos_b=motion_anchor_pos_b, - motion_anchor_ori_b=motion_anchor_ori_b, - noisy_linvel=linvel_actor, - noisy_gyro=gyro_actor, - noisy_joint_pos_rel=joint_pos_actor, - noisy_dof_vel=dof_vel_actor, - last_actions=last_actions, - ) - - # Critic observations (clean proprioception + privileged body transforms) - critic_obs = np.empty((num_envs, self._critic_obs_width), dtype=dtype) - offset = 0 - critic_obs[:, offset : offset + command.shape[1]] = command - offset += command.shape[1] - critic_obs[:, offset : offset + 3] = motion_anchor_pos_b - offset += 3 - critic_obs[:, offset : offset + 6] = motion_anchor_ori_b - offset += 6 - critic_obs[:, offset : offset + 3] = linvel - offset += 3 - critic_obs[:, offset : offset + 3] = gyro - offset += 3 - critic_obs[:, offset : offset + n_action] = joint_pos_rel - offset += n_action - critic_obs[:, offset : offset + n_action] = dof_vel - offset += n_action - critic_obs[:, offset : offset + n_action] = last_actions - offset += n_action - robot_body_pos_b = critic_obs[:, offset : offset + n_body * 3].reshape(num_envs, n_body, 3) - self._write_body_pos_in_anchor_frame( - robot_anchor_pos_w, robot_anchor_quat_w, robot_body_pos_w, robot_body_pos_b - ) - offset += n_body * 3 - robot_body_ori_b = critic_obs[:, offset : offset + n_body * 6].reshape(num_envs, n_body, 6) - self._write_body_ori6_in_anchor_frame( - robot_anchor_quat_w, robot_body_quat_w, robot_body_ori_b - ) - return {"obs": actor_obs, "critic": critic_obs} - - def _compute_reward( - self, - info: dict, - motion_data, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, - robot_body_lin_vel_w: np.ndarray, - robot_body_ang_vel_w: np.ndarray, - dof_pos: np.ndarray, - dof_vel: np.ndarray, - ) -> np.ndarray: - """Compute reward.""" - reward = self._env_error2 - reward.fill(0.0) - cfg = self._cfg.reward_config - - step_count = info.get("steps") - should_log = self._enable_reward_log and ( - int(step_count[0]) % 4 == 0 if isinstance(step_count, np.ndarray) else True - ) - log = {} if should_log else info.get("log", {}) - - # Store motion and robot states in info for reward functions - info["motion_data"] = motion_data - info["robot_body_pos_w"] = robot_body_pos_w - info["robot_body_quat_w"] = robot_body_quat_w - info["robot_body_lin_vel_w"] = robot_body_lin_vel_w - info["robot_body_ang_vel_w"] = robot_body_ang_vel_w - info["reward_ref_body_pos_w"] = self.body_pos_relative_w - info["reward_ref_body_quat_w"] = self.body_quat_relative_w - info["anchor_body_idx"] = self.anchor_body_idx - info["dof_pos"] = dof_pos - info["dof_vel"] = dof_vel - - for name, scale in cfg.scales.items(): - if scale == 0: - continue - reward_fn = self._active_reward_fns.get(name) - if reward_fn is None: - if should_log and name in self._reward_fns: - log[f"reward/{name}"] = 0.0 - continue - rew = reward_fn(info) - weighted_rew = self._weighted_reward - np.multiply(rew, scale, out=weighted_rew) - reward += weighted_rew - if should_log: - log[f"reward/{name}"] = float(np.sum(weighted_rew) / weighted_rew.size) - - info["log"] = log - reward *= self._cfg.ctrl_dt - return reward - - def _mean_body_xyz_squared_error(self, reference: np.ndarray, actual: np.ndarray) -> np.ndarray: - vec_error = self._body_vec_error - env_error = self._env_error - tmp_error = self._reward_term - np.subtract(reference[..., 0], actual[..., 0], out=vec_error[..., 0]) - np.square(vec_error[..., 0], out=vec_error[..., 0]) - np.sum(vec_error[..., 0], axis=1, out=env_error) - np.subtract(reference[..., 1], actual[..., 1], out=vec_error[..., 1]) - np.square(vec_error[..., 1], out=vec_error[..., 1]) - np.sum(vec_error[..., 1], axis=1, out=tmp_error) - env_error += tmp_error - np.subtract(reference[..., 2], actual[..., 2], out=vec_error[..., 2]) - np.square(vec_error[..., 2], out=vec_error[..., 2]) - np.sum(vec_error[..., 2], axis=1, out=tmp_error) - env_error += tmp_error - env_error /= reference.shape[1] - return env_error - - def _quat_error_magnitude_squared_body(self, q1: np.ndarray, q2: np.ndarray) -> np.ndarray: - rel_w = self._quat_error_w - rel_x = self._quat_error_x - # Motion/backend quaternions are unit quaternions, so the relative - # rotation angle only needs abs(dot(q1, q2)). - np.multiply(q1[..., 0], q2[..., 0], out=rel_w) - np.multiply(q1[..., 1], q2[..., 1], out=rel_x) - rel_w += rel_x - np.multiply(q1[..., 2], q2[..., 2], out=rel_x) - rel_w += rel_x - np.multiply(q1[..., 3], q2[..., 3], out=rel_x) - rel_w += rel_x - np.abs(rel_w, out=rel_w) - np.clip(rel_w, 0.0, 1.0, out=rel_w) - np.arccos(rel_w, out=rel_x) - rel_x *= 2.0 - np.square(rel_x, out=rel_x) - return rel_x - - def _exp_reward_from_error(self, error: np.ndarray, std: float) -> np.ndarray: - out = self._reward_term - np.divide(error, -(std**2), out=out) - np.exp(out, out=out) - return out - - # Reward functions - def _reward_motion_global_root_pos(self, info: dict) -> np.ndarray: - motion_data = info["motion_data"] - robot_body_pos_w = info["robot_body_pos_w"] - anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] - error = self._env_error - np.subtract(anchor_pos_w[:, 0], robot_anchor_pos_w[:, 0], out=error) - np.square(error, out=error) - np.subtract(anchor_pos_w[:, 1], robot_anchor_pos_w[:, 1], out=self._reward_term) - np.square(self._reward_term, out=self._reward_term) - error += self._reward_term - np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=self._reward_term) - np.square(self._reward_term, out=self._reward_term) - error += self._reward_term - return self._exp_reward_from_error(error, self._cfg.reward_config.std_root_pos) - - def _reward_motion_global_root_ori(self, info: dict) -> np.ndarray: - motion_data = info["motion_data"] - robot_body_quat_w = info["robot_body_quat_w"] - anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] - np.multiply(anchor_quat_w[:, 0], robot_anchor_quat_w[:, 0], out=self._env_error) - np.multiply(anchor_quat_w[:, 1], robot_anchor_quat_w[:, 1], out=self._reward_term) - self._env_error += self._reward_term - np.multiply(anchor_quat_w[:, 2], robot_anchor_quat_w[:, 2], out=self._reward_term) - self._env_error += self._reward_term - np.multiply(anchor_quat_w[:, 3], robot_anchor_quat_w[:, 3], out=self._reward_term) - self._env_error += self._reward_term - np.abs(self._env_error, out=self._env_error) - np.clip(self._env_error, 0.0, 1.0, out=self._env_error) - np.arccos(self._env_error, out=self._env_error) - self._env_error *= 2.0 - np.square(self._env_error, out=self._env_error) - return self._exp_reward_from_error(self._env_error, self._cfg.reward_config.std_root_ori) - - def _reward_motion_body_pos(self, info: dict) -> np.ndarray: - robot_body_pos_w = info["robot_body_pos_w"] - error = self._mean_body_xyz_squared_error(self.body_pos_relative_w, robot_body_pos_w) - return self._exp_reward_from_error(error, self._cfg.reward_config.std_body_pos) - - def _reward_motion_body_ori(self, info: dict) -> np.ndarray: - robot_body_quat_w = info["robot_body_quat_w"] - error = self._quat_error_magnitude_squared_body( - self.body_quat_relative_w, robot_body_quat_w - ) - np.sum(error, axis=-1, out=self._env_error) - self._env_error /= error.shape[1] - return self._exp_reward_from_error(self._env_error, self._cfg.reward_config.std_body_ori) - - def _reward_motion_body_lin_vel(self, info: dict) -> np.ndarray: - motion_data = info["motion_data"] - robot_body_lin_vel_w = info["robot_body_lin_vel_w"] - error = self._mean_body_xyz_squared_error(motion_data.body_lin_vel_w, robot_body_lin_vel_w) - return self._exp_reward_from_error(error, self._cfg.reward_config.std_body_lin_vel) - - def _reward_motion_body_ang_vel(self, info: dict) -> np.ndarray: - motion_data = info["motion_data"] - robot_body_ang_vel_w = info["robot_body_ang_vel_w"] - error = self._mean_body_xyz_squared_error(motion_data.body_ang_vel_w, robot_body_ang_vel_w) - return self._exp_reward_from_error(error, self._cfg.reward_config.std_body_ang_vel) - - def _reward_motion_ee_body_pos_z(self, info: dict) -> np.ndarray: - robot_body_pos_w = info["robot_body_pos_w"] - np.subtract( - self.body_pos_relative_w[:, self.ee_body_indices, 2], - robot_body_pos_w[:, self.ee_body_indices, 2], - out=self._ee_pos_error_z, - ) - np.square(self._ee_pos_error_z, out=self._ee_pos_error_z) - np.sum(self._ee_pos_error_z, axis=-1, out=self._env_error) - self._env_error /= self._ee_pos_error_z.shape[1] - return self._exp_reward_from_error(self._env_error, self._cfg.reward_config.std_body_pos) - - def _reward_motion_joint_pos(self, info: dict) -> np.ndarray: - motion_data = info["motion_data"] - dof_pos = info["dof_pos"] - np.subtract(motion_data.joint_pos, dof_pos, out=self._joint_error) - np.square(self._joint_error, out=self._joint_error) - np.sum(self._joint_error, axis=1, out=self._env_error) - self._env_error /= dof_pos.shape[1] - return self._exp_reward_from_error(self._env_error, self._cfg.reward_config.std_joint_pos) - - def _reward_motion_joint_vel(self, info: dict) -> np.ndarray: - motion_data = info["motion_data"] - dof_vel = info["dof_vel"] - np.subtract(motion_data.joint_vel, dof_vel, out=self._joint_error) - np.square(self._joint_error, out=self._joint_error) - np.sum(self._joint_error, axis=1, out=self._env_error) - self._env_error /= dof_vel.shape[1] - return self._exp_reward_from_error(self._env_error, self._cfg.reward_config.std_joint_vel) - - def _reward_undesired_contacts(self, info: dict) -> np.ndarray: - robot_body_pos_w = info["robot_body_pos_w"] - body_z = robot_body_pos_w[:, self.undesired_contact_body_indices, 2] - np.less( - body_z, - self._cfg.undesired_contact_z_threshold, - out=self._undesired_contact_mask, - ) - np.sum(self._undesired_contact_mask, axis=-1, out=self._env_error) - return self._env_error - - def _reward_action_rate_l2(self, info: dict) -> np.ndarray: - np.subtract(info["current_actions"], info["last_actions"], out=self._joint_error) - np.square(self._joint_error, out=self._joint_error) - np.sum(self._joint_error, axis=1, out=self._env_error) - return self._env_error - - def _reward_joint_limit(self, info: dict) -> np.ndarray: - dof_pos = info["dof_pos"] - lower = self._joint_lower - upper = self._joint_upper - if lower is None or upper is None: - self._reward_term.fill(0.0) - return self._reward_term - - # Compute violation - np.subtract(lower, dof_pos, out=self._joint_error) - np.maximum(self._joint_error, 0, out=self._joint_error) - np.subtract(dof_pos, upper, out=self._joint_error_upper) - np.maximum(self._joint_error_upper, 0, out=self._joint_error_upper) - self._joint_error += self._joint_error_upper - np.square(self._joint_error, out=self._joint_error) - np.sum(self._joint_error, axis=1, out=self._reward_term) - return self._reward_term + _cfg: MotionTrackingCfg @registry.env("G1MotionTrackingDeploy", sim_backend="mujoco") @registry.env("G1MotionTrackingDeploy", sim_backend="motrix") -class G1MotionTrackingDeployEnv(G1MotionTrackingEnv): +class G1MotionTrackingDeployEnv(MotionTrackingDeployEnv): """Deploy-oriented G1 motion tracking env with unitree_rl_lab mimic actor inputs.""" - _cfg: G1MotionTrackingDeployEnvCfg - - def _actor_obs_dim(self, n: int) -> int: - # unitree_rl_lab mimic deploy actor input: - # motion_command(2n), motion_anchor_ori_b(6), gyro(3), joints, actions. - return 6 + 3 + n * 5 - - def _build_actor_obs( - self, - *, - command: np.ndarray, - motion_anchor_pos_b: np.ndarray, - motion_anchor_ori_b: np.ndarray, - noisy_linvel: np.ndarray, - noisy_gyro: np.ndarray, - noisy_joint_pos_rel: np.ndarray, - noisy_dof_vel: np.ndarray, - last_actions: np.ndarray, - ) -> np.ndarray: - return np.concatenate( - [ - command, - motion_anchor_ori_b, - noisy_gyro, - noisy_joint_pos_rel, - noisy_dof_vel, - last_actions, - ], - axis=1, - dtype=get_global_dtype(), - ) + _cfg: MotionTrackingDeployEnvCfg + + +__all__ = [ + "DomainRand", + "Domain_Rand", + "G1MotionTrackingCfg", + "G1MotionTrackingDeployEnv", + "G1MotionTrackingDeployEnvCfg", + "G1MotionTrackingDomainRandomizationProvider", + "G1MotionTrackingEnv", + "G1MotionTrackingEnvCfg", + "MotionTrackingCfg", + "MotionTrackingDeployEnv", + "MotionTrackingDeployEnvCfg", + "MotionTrackingDomainRandomizationProvider", + "MotionTrackingEnv", + "PoseRandomization", + "RewardConfig", + "VelocityRandomization", + "_build_motion_reference_state", + "_zero_pose_randomization", + "_zero_velocity_randomization", +] diff --git a/src/unilab/envs/motion_tracking/g1/tracking_obs.py b/src/unilab/envs/motion_tracking/g1/tracking_obs.py index 0e04d01f1..b6dfdc157 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/envs/motion_tracking/g1/tracking_obs.py @@ -43,6 +43,7 @@ from unilab.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import NoiseConfig +from ..common.rewards import RewardContext from .tracking import ( Domain_Rand, G1MotionTrackingDomainRandomizationProvider, @@ -308,18 +309,18 @@ def _init_reward_functions(self) -> None: self._reward_fns["joint_acc_l2"] = self._reward_joint_acc_l2 self._reward_fns["joint_torque_l2"] = self._reward_joint_torque_l2 - def _reward_joint_acc_l2(self, info: dict) -> np.ndarray: - dof_vel = info["dof_vel"] - prev_dof_vel = info.get("prev_dof_vel") + def _reward_joint_acc_l2(self, ctx: RewardContext) -> np.ndarray: + dof_vel = ctx.dof_vel + prev_dof_vel = ctx.info.get("prev_dof_vel") if prev_dof_vel is None or prev_dof_vel.shape != dof_vel.shape: return np.zeros((self._num_envs,), dtype=get_global_dtype()) joint_acc = (dof_vel - prev_dof_vel) / self._cfg.ctrl_dt return np.asarray(np.sum(np.square(joint_acc), axis=1), dtype=get_global_dtype()) - def _reward_joint_torque_l2(self, info: dict) -> np.ndarray: - dof_pos = info["dof_pos"] - dof_vel = info["dof_vel"] - last_actions = info.get("last_actions") + def _reward_joint_torque_l2(self, ctx: RewardContext) -> np.ndarray: + dof_pos = ctx.dof_pos + dof_vel = ctx.dof_vel + last_actions = ctx.info.get("last_actions") if last_actions is None: return np.zeros((self._num_envs,), dtype=get_global_dtype()) target_q = ( diff --git a/src/unilab/envs/motion_tracking/x2/flip_tracking.py b/src/unilab/envs/motion_tracking/x2/flip_tracking.py index 742400c19..da7a0ff1c 100644 --- a/src/unilab/envs/motion_tracking/x2/flip_tracking.py +++ b/src/unilab/envs/motion_tracking/x2/flip_tracking.py @@ -9,20 +9,20 @@ from unilab.base import registry from unilab.base.scene import SceneCfg from unilab.envs.locomotion.g1.base import Sensor -from unilab.envs.motion_tracking.g1.flip_tracking import ( +from unilab.envs.motion_tracking.common.config import ( + PoseRandomization, + VelocityRandomization, _zero_pose_randomization, _zero_velocity_randomization, ) -from unilab.envs.motion_tracking.g1.tracking import ( - G1MotionTrackingDeployEnv, - G1MotionTrackingDeployEnvCfg, - PoseRandomization, - VelocityRandomization, +from unilab.envs.motion_tracking.common.tracking import ( + MotionTrackingDeployEnv, + MotionTrackingDeployEnvCfg, ) @dataclass -class X2MotionTrackingCfg(G1MotionTrackingDeployEnvCfg): +class X2MotionTrackingCfg(MotionTrackingDeployEnvCfg): """Base X2 motion-tracking config profile.""" scene: SceneCfg = field( @@ -119,7 +119,7 @@ class X2WallFlipTrackingEnvCfg(X2WallFlipTrackingCfg): @registry.env("X2WallFlipTracking", sim_backend="mujoco") @registry.env("X2WallFlipTracking", sim_backend="motrix") -class X2WallFlipTrackingEnv(G1MotionTrackingDeployEnv): +class X2WallFlipTrackingEnv(MotionTrackingDeployEnv): """X2 wall flip-tracking environment implementation.""" _cfg: X2WallFlipTrackingCfg From f9cfbe34ee0ac47e26c6e9166b7fafd39a4a3706 Mon Sep 17 00:00:00 2001 From: caozx1110 Date: Thu, 9 Jul 2026 21:33:02 +0800 Subject: [PATCH 3/4] test: cover shared motion tracking modules --- .../benchmark_g1_motion_tracking_numba.py | 16 +- tests/envs/test_env_configs.py | 12 +- tests/envs/test_g1_motion_tracking_numba.py | 16 +- tests/envs/test_motion_loader.py | 2 +- tests/envs/test_motion_tracking_rewards.py | 148 ++++++++++++++++++ 5 files changed, 171 insertions(+), 23 deletions(-) create mode 100644 tests/envs/test_motion_tracking_rewards.py diff --git a/benchmark/benchmark_g1_motion_tracking_numba.py b/benchmark/benchmark_g1_motion_tracking_numba.py index 5db0f5beb..7829acf75 100644 --- a/benchmark/benchmark_g1_motion_tracking_numba.py +++ b/benchmark/benchmark_g1_motion_tracking_numba.py @@ -45,8 +45,8 @@ from benchmark.core.device_info import get_device_info_dict, get_device_info_line from unilab.dtype_config import get_global_dtype -from unilab.envs.motion_tracking.g1.motion_tracking_numba import ( - G1MotionTrackingNumbaAccelerator, +from unilab.envs.motion_tracking.common.numba import ( + MotionTrackingNumbaAccelerator, ) from unilab.envs.motion_tracking.g1.tracking import ( G1MotionTrackingCfg, @@ -390,7 +390,7 @@ def compute_numpy(batch: SyntheticBatch) -> tuple[np.ndarray, np.ndarray, dict[s def compute_numba( - batch: SyntheticBatch, accelerator: G1MotionTrackingNumbaAccelerator + batch: SyntheticBatch, accelerator: MotionTrackingNumbaAccelerator ) -> tuple[np.ndarray, np.ndarray, dict[str, float]]: info = {**batch.info, "log": {}} out = accelerator.compute( @@ -451,7 +451,7 @@ def compute_numpy_update_state( def compute_numba_update_state( - batch: SyntheticBatch, accelerator: G1MotionTrackingNumbaAccelerator + batch: SyntheticBatch, accelerator: MotionTrackingNumbaAccelerator ) -> tuple[dict[str, np.ndarray], np.ndarray, np.ndarray, dict[str, float]]: info = {**batch.info, "log": {}} noise_cfg = batch.env._cfg.noise_config @@ -494,7 +494,7 @@ def time_call(fn, *, iters: int, warmup: int) -> tuple[float, float, float]: def check_parity( - batch: SyntheticBatch, accelerator: G1MotionTrackingNumbaAccelerator + batch: SyntheticBatch, accelerator: MotionTrackingNumbaAccelerator ) -> dict[str, float]: reward_np, terminated_np, log_np = compute_numpy(batch) reward_nb, terminated_nb, log_nb = compute_numba(batch, accelerator) @@ -555,7 +555,7 @@ def bench_one( ) ] - compile_driver = G1MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=1) + compile_driver = MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=1) t0 = time.perf_counter() compute_numba(batch, compile_driver) compute_numba_update_state(batch, compile_driver) @@ -565,7 +565,7 @@ def bench_one( for threads in dict.fromkeys([1, *thread_counts]): if threads > max_threads: continue - accelerator = G1MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=threads) + accelerator = MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=threads) numba_mean, numba_min, numba_std = time_call( lambda: compute_numba(batch, accelerator), iters=iters, warmup=warmup ) @@ -597,7 +597,7 @@ def bench_one( (record for record in records if record.path == "numba_accelerator"), key=lambda record: record.mean_ms, ) - best_update_accelerator = G1MotionTrackingNumbaAccelerator.from_env( + best_update_accelerator = MotionTrackingNumbaAccelerator.from_env( batch.env, num_threads=best_numba.threads ) numba_update_state_mean, _, _ = time_call( diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 3d9b0009f..d32609afb 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -478,7 +478,7 @@ def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray] def test_g1_motion_tracking_reset_observation_uses_sparse_body_pose_rows(): - from unilab.envs.motion_tracking.g1.motion_loader import MotionData + from unilab.envs.motion_tracking.common.motion_loader import MotionData from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingDomainRandomizationProvider class FakeBackend: @@ -557,7 +557,7 @@ def compute_obs( def _compute_g1_motion_tracking_obs_stub(env_cls: type): - from unilab.envs.motion_tracking.g1.motion_loader import MotionData + from unilab.envs.motion_tracking.common.motion_loader import MotionData env = cast(Any, object.__new__(env_cls)) env._num_envs = 1 @@ -761,7 +761,7 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_reward_fast_path_matches_reference(): - from unilab.envs.motion_tracking.g1.motion_loader import MotionData + from unilab.envs.motion_tracking.common.motion_loader import MotionData from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig from unilab.utils.rotation import np_quat_error_magnitude @@ -1250,7 +1250,7 @@ def test_g1_motion_tracking_cfg_preserves_legacy_defaults(): def test_g1_motion_tracking_init_delegates_motion_body_ids_to_backend(monkeypatch): from unilab.envs.locomotion.g1.base import G1BaseEnv - from unilab.envs.motion_tracking.g1 import tracking as tracking_module + from unilab.envs.motion_tracking.common import tracking as tracking_module from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingCfg, G1MotionTrackingEnv calls: dict[str, Any] = {} @@ -1331,7 +1331,7 @@ def __init__( assert calls["motion_loader"][0] == "dummy_motion.npz" np.testing.assert_array_equal(calls["motion_loader"][1], np.array([1, 2], dtype=np.int32)) assert calls["motion_sampler"][1:] == ("adaptive", 4, cfg.sampling_start_ratio) - assert calls["dr_provider"] == "G1MotionTrackingDomainRandomizationProvider" + assert calls["dr_provider"] == "MotionTrackingDomainRandomizationProvider" assert calls["reward_init"] is True @@ -1500,7 +1500,7 @@ def _make_g1_motion_tracking_clip_end_stub( step_env_ids: np.ndarray | None = None, ): from unilab.base.np_env import NpEnvState - from unilab.envs.motion_tracking.g1.motion_loader import MotionData + from unilab.envs.motion_tracking.common.motion_loader import MotionData from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv class FakeBackend: diff --git a/tests/envs/test_g1_motion_tracking_numba.py b/tests/envs/test_g1_motion_tracking_numba.py index f4f3d1796..484802a12 100644 --- a/tests/envs/test_g1_motion_tracking_numba.py +++ b/tests/envs/test_g1_motion_tracking_numba.py @@ -6,9 +6,9 @@ import numpy as np import pytest -from unilab.envs.motion_tracking.g1.motion_tracking_numba import ( +from unilab.envs.motion_tracking.common.numba import ( NUMBA_AVAILABLE, - G1MotionTrackingNumbaAccelerator, + MotionTrackingNumbaAccelerator, unsupported_terms, ) from unilab.envs.motion_tracking.g1.tracking import ( @@ -24,11 +24,11 @@ def test_unsupported_terms_only_reports_active_unknown_terms(): def test_numba_accelerator_requires_numba_when_declared(monkeypatch): - import unilab.envs.motion_tracking.g1.motion_tracking_numba as motion_tracking_numba + import unilab.envs.motion_tracking.common.numba as motion_tracking_numba monkeypatch.setattr(motion_tracking_numba, "NUMBA_AVAILABLE", False) with pytest.raises(RuntimeError, match="requires numba"): - G1MotionTrackingNumbaAccelerator.from_env(_make_env(8, include_undesired=False)) + MotionTrackingNumbaAccelerator.from_env(_make_env(8, include_undesired=False)) def test_numba_accelerator_rejects_unsupported_active_reward_terms(): @@ -36,7 +36,7 @@ def test_numba_accelerator_rejects_unsupported_active_reward_terms(): env._cfg.reward_config.scales = {"motion_body_pos": 1.0, "custom": 1.0} if NUMBA_AVAILABLE: with pytest.raises(ValueError, match="does not support active reward terms"): - G1MotionTrackingNumbaAccelerator.from_env(env) + MotionTrackingNumbaAccelerator.from_env(env) @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") @@ -68,7 +68,7 @@ def test_g1_motion_tracking_numba_reward_termination_parity(): ) log_np = dict(info["log"]) - accel = G1MotionTrackingNumbaAccelerator.from_env(env, num_threads=2) + accel = MotionTrackingNumbaAccelerator.from_env(env, num_threads=2) out = accel.compute( info={ "steps": np.zeros((n,), dtype=np.uint32), @@ -148,7 +148,7 @@ def test_g1_motion_tracking_numba_update_state_parity_without_noise(): robot_body_quat_w, ) - accel = G1MotionTrackingNumbaAccelerator.from_env(env, num_threads=2) + accel = MotionTrackingNumbaAccelerator.from_env(env, num_threads=2) out = accel.compute_update_state( info=info, motion_data=motion_data, @@ -182,7 +182,7 @@ def test_g1_motion_tracking_numba_update_state_parity_without_noise(): @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") def test_g1_motion_tracking_numba_term_py_funcs_match_numpy_math(): - from unilab.envs.motion_tracking.g1 import motion_tracking_numba as T + from unilab.envs.motion_tracking.common import numba as T n = 4 env = _make_env(n, include_undesired=True) diff --git a/tests/envs/test_motion_loader.py b/tests/envs/test_motion_loader.py index 31d383576..1b5069faf 100644 --- a/tests/envs/test_motion_loader.py +++ b/tests/envs/test_motion_loader.py @@ -2,7 +2,7 @@ import numpy as np -from unilab.envs.motion_tracking.g1.motion_loader import MotionLoader, MotionSampler +from unilab.envs.motion_tracking.common.motion_loader import MotionLoader, MotionSampler def _write_motion_npz( diff --git a/tests/envs/test_motion_tracking_rewards.py b/tests/envs/test_motion_tracking_rewards.py new file mode 100644 index 000000000..ddc7fabb3 --- /dev/null +++ b/tests/envs/test_motion_tracking_rewards.py @@ -0,0 +1,148 @@ +"""Guard tests for the extracted motion-tracking reward module. + +These lock the ``common.rewards`` term functions and dispatch against +hand-computed values on a small synthetic ``RewardContext`` and guard against +numpy/numba reward-term drift. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np + +from unilab.envs.motion_tracking.common import numba as numba_mod +from unilab.envs.motion_tracking.common import rewards +from unilab.envs.motion_tracking.common.rewards import RewardConfig, RewardContext + + +def _make_ctx(*, scales: dict[str, float] | None = None) -> RewardContext: + """Build a small deterministic RewardContext (2 envs, 3 bodies, 2 joints).""" + rng = np.random.default_rng(7) + num_envs, n_body, n_action = 2, 3, 2 + dtype = np.float64 + + def rq(shape: tuple[int, ...]) -> np.ndarray: + q = rng.normal(size=(*shape, 4)).astype(dtype) + q /= np.linalg.norm(q, axis=-1, keepdims=True) + return q + + reward_config = RewardConfig() + if scales is not None: + reward_config.scales = scales + + motion_data = SimpleNamespace( + joint_pos=rng.normal(size=(num_envs, n_action)).astype(dtype), + joint_vel=rng.normal(size=(num_envs, n_action)).astype(dtype), + body_pos_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), + body_quat_w=rq((num_envs, n_body)), + body_lin_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), + body_ang_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), + ) + + robot_body_pos_w = rng.normal(size=(num_envs, n_body, 3)).astype(dtype) + # Force a deterministic contact pattern for the undesired_contacts term. + robot_body_pos_w[:, :, 2] = np.array([[0.01, 0.20, 0.03], [0.10, 0.02, 0.30]], dtype=dtype) + undesired_idx = np.array([0, 1], dtype=np.int32) + + current_actions = rng.normal(size=(num_envs, n_action)).astype(dtype) + last_actions = rng.normal(size=(num_envs, n_action)).astype(dtype) + + return RewardContext( + info={ + "current_actions": current_actions, + "last_actions": last_actions, + "steps": np.zeros((num_envs,), dtype=np.uint32), + }, + motion_data=motion_data, + robot_body_pos_w=robot_body_pos_w, + robot_body_quat_w=rq((num_envs, n_body)), + robot_body_lin_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), + robot_body_ang_vel_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), + ref_body_pos_w=rng.normal(size=(num_envs, n_body, 3)).astype(dtype), + ref_body_quat_w=rq((num_envs, n_body)), + dof_pos=rng.normal(size=(num_envs, n_action)).astype(dtype), + dof_vel=rng.normal(size=(num_envs, n_action)).astype(dtype), + reward_config=reward_config, + anchor_body_idx=0, + ee_body_indices=np.array([2], dtype=np.int32), + undesired_contact_body_indices=undesired_idx, + joint_lower=np.array([-0.5, -0.5], dtype=dtype), + joint_upper=np.array([0.5, 0.5], dtype=dtype), + undesired_contact_z_threshold=0.05, + num_envs=num_envs, + body_vec_error=np.empty((num_envs, n_body, 3), dtype=dtype), + joint_error=np.empty((num_envs, n_action), dtype=dtype), + joint_error_upper=np.empty((num_envs, n_action), dtype=dtype), + env_error=np.empty((num_envs,), dtype=dtype), + env_error2=np.empty((num_envs,), dtype=dtype), + reward_term=np.empty((num_envs,), dtype=dtype), + weighted_reward=np.empty((num_envs,), dtype=dtype), + quat_error_w=np.empty((num_envs, n_body), dtype=dtype), + quat_error_x=np.empty((num_envs, n_body), dtype=dtype), + ee_pos_error_z=np.empty((num_envs, 1), dtype=dtype), + undesired_contact_mask=np.empty((num_envs, undesired_idx.size), dtype=bool), + ) + + +def test_build_reward_functions_matches_numba_term_order(): + assert set(rewards.build_reward_functions()) == set(numba_mod.TERM_ORDER) + + +def test_motion_joint_pos_term_matches_hand_computed(): + ctx = _make_ctx() + out = rewards.motion_joint_pos(ctx).copy() + error = np.mean(np.square(ctx.motion_data.joint_pos - ctx.dof_pos), axis=1) + expected = np.exp(-error / ctx.reward_config.std_joint_pos**2) + np.testing.assert_allclose(out, expected, rtol=1e-12, atol=1e-12) + + +def test_action_rate_l2_term_matches_hand_computed(): + ctx = _make_ctx() + out = rewards.action_rate_l2(ctx).copy() + expected = np.sum( + np.square(ctx.info["current_actions"] - ctx.info["last_actions"]), axis=1 + ) + np.testing.assert_allclose(out, expected, rtol=1e-12, atol=1e-12) + + +def test_undesired_contacts_term_matches_hand_computed(): + ctx = _make_ctx() + out = rewards.undesired_contacts(ctx).copy() + idx = ctx.undesired_contact_body_indices + expected = np.sum( + ctx.robot_body_pos_w[:, idx, 2] < ctx.undesired_contact_z_threshold, axis=1 + ).astype(np.float64) + np.testing.assert_allclose(out, expected, rtol=1e-12, atol=1e-12) + + +def test_compute_reward_matches_hand_computed_weighted_sum(): + scales = { + "motion_joint_pos": 1.0, + "action_rate_l2": -0.1, + "undesired_contacts": -0.5, + # zero-weighted terms are skipped by compute_reward + "motion_body_pos": 0.0, + } + ctrl_dt = 0.02 + fns = rewards.build_reward_functions() + + # Reference terms computed on an untouched ctx (fresh buffers per call). + joint_pos_ref = rewards.motion_joint_pos(_make_ctx()).copy() + action_rate_ref = rewards.action_rate_l2(_make_ctx()).copy() + undesired_ref = rewards.undesired_contacts(_make_ctx()).copy() + + ctx = _make_ctx(scales=scales) + reward = rewards.compute_reward( + ctx, + active_reward_fns=fns, + all_reward_fns=fns, + scales=scales, + ctrl_dt=ctrl_dt, + enable_log=False, + ).copy() + + expected = ( + 1.0 * joint_pos_ref - 0.1 * action_rate_ref - 0.5 * undesired_ref + ) * ctrl_dt + np.testing.assert_allclose(reward, expected, rtol=1e-12, atol=1e-12) From 68f5637d7350806629a62be7e2c58ccabdc10b15 Mon Sep 17 00:00:00 2001 From: caozx1110 Date: Thu, 9 Jul 2026 22:16:43 +0800 Subject: [PATCH 4/4] fix: satisfy motion tracking refactor type checks --- .../envs/motion_tracking/common/numba.py | 4 +- .../motion_tracking/common/observations.py | 4 +- .../envs/motion_tracking/common/rewards.py | 188 ++++++++++-------- .../envs/motion_tracking/g1/tracking_obs.py | 4 + tests/envs/test_motion_tracking_rewards.py | 8 +- 5 files changed, 117 insertions(+), 91 deletions(-) diff --git a/src/unilab/envs/motion_tracking/common/numba.py b/src/unilab/envs/motion_tracking/common/numba.py index 850caa83d..6a29d6acf 100644 --- a/src/unilab/envs/motion_tracking/common/numba.py +++ b/src/unilab/envs/motion_tracking/common/numba.py @@ -974,9 +974,7 @@ def __init__( self._zero_joint_noise = np.zeros((self.num_envs, self.num_action), dtype=np.float64) @classmethod - def from_env( - cls, env: Any, num_threads: int | None = None - ) -> "MotionTrackingNumbaAccelerator": + def from_env(cls, env: Any, num_threads: int | None = None) -> "MotionTrackingNumbaAccelerator": if not NUMBA_AVAILABLE: raise RuntimeError( "MotionTracking numba_acceleration=True requires numba. Install it or run " diff --git a/src/unilab/envs/motion_tracking/common/observations.py b/src/unilab/envs/motion_tracking/common/observations.py index 65310807c..fd5b0b6ee 100644 --- a/src/unilab/envs/motion_tracking/common/observations.py +++ b/src/unilab/envs/motion_tracking/common/observations.py @@ -273,7 +273,5 @@ def compute_obs( ) offset += n_body * 3 robot_body_ori_b = critic_obs[:, offset : offset + n_body * 6].reshape(num_envs, n_body, 6) - env._write_body_ori6_in_anchor_frame( - robot_anchor_quat_w, robot_body_quat_w, robot_body_ori_b - ) + env._write_body_ori6_in_anchor_frame(robot_anchor_quat_w, robot_body_quat_w, robot_body_ori_b) return {"obs": actor_obs, "critic": critic_obs} diff --git a/src/unilab/envs/motion_tracking/common/rewards.py b/src/unilab/envs/motion_tracking/common/rewards.py index 85f9728f9..83ae009a0 100644 --- a/src/unilab/envs/motion_tracking/common/rewards.py +++ b/src/unilab/envs/motion_tracking/common/rewards.py @@ -101,12 +101,18 @@ class RewardContext: # ── buffered math helpers ──────────────────────────────────────────── +def _required_array(value: np.ndarray | None, name: str) -> np.ndarray: + if value is None: + raise RuntimeError(f"RewardContext.{name} is required for this reward term") + return value + + def _mean_body_xyz_squared_error( ctx: RewardContext, reference: np.ndarray, actual: np.ndarray ) -> np.ndarray: - vec_error = ctx.body_vec_error - env_error = ctx.env_error - tmp_error = ctx.reward_term + vec_error = _required_array(ctx.body_vec_error, "body_vec_error") + env_error = _required_array(ctx.env_error, "env_error") + tmp_error = _required_array(ctx.reward_term, "reward_term") np.subtract(reference[..., 0], actual[..., 0], out=vec_error[..., 0]) np.square(vec_error[..., 0], out=vec_error[..., 0]) np.sum(vec_error[..., 0], axis=1, out=env_error) @@ -125,8 +131,8 @@ def _mean_body_xyz_squared_error( def _quat_error_magnitude_squared_body( ctx: RewardContext, q1: np.ndarray, q2: np.ndarray ) -> np.ndarray: - rel_w = ctx.quat_error_w - rel_x = ctx.quat_error_x + rel_w = _required_array(ctx.quat_error_w, "quat_error_w") + rel_x = _required_array(ctx.quat_error_x, "quat_error_x") # Motion/backend quaternions are unit quaternions, so the relative # rotation angle only needs abs(dot(q1, q2)). np.multiply(q1[..., 0], q2[..., 0], out=rel_w) @@ -145,7 +151,7 @@ def _quat_error_magnitude_squared_body( def _exp_reward_from_error(ctx: RewardContext, error: np.ndarray, std: float) -> np.ndarray: - out = ctx.reward_term + out = _required_array(ctx.reward_term, "reward_term") np.divide(error, -(std**2), out=out) np.exp(out, out=out) return out @@ -156,139 +162,163 @@ def _exp_reward_from_error(ctx: RewardContext, error: np.ndarray, std: float) -> def motion_global_root_pos(ctx: RewardContext) -> np.ndarray: motion_data = ctx.motion_data - robot_body_pos_w = ctx.robot_body_pos_w + robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") anchor_pos_w = motion_data.body_pos_w[:, ctx.anchor_body_idx] robot_anchor_pos_w = robot_body_pos_w[:, ctx.anchor_body_idx] - error = ctx.env_error + error = _required_array(ctx.env_error, "env_error") + reward_term = _required_array(ctx.reward_term, "reward_term") np.subtract(anchor_pos_w[:, 0], robot_anchor_pos_w[:, 0], out=error) np.square(error, out=error) - np.subtract(anchor_pos_w[:, 1], robot_anchor_pos_w[:, 1], out=ctx.reward_term) - np.square(ctx.reward_term, out=ctx.reward_term) - error += ctx.reward_term - np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=ctx.reward_term) - np.square(ctx.reward_term, out=ctx.reward_term) - error += ctx.reward_term + np.subtract(anchor_pos_w[:, 1], robot_anchor_pos_w[:, 1], out=reward_term) + np.square(reward_term, out=reward_term) + error += reward_term + np.subtract(anchor_pos_w[:, 2], robot_anchor_pos_w[:, 2], out=reward_term) + np.square(reward_term, out=reward_term) + error += reward_term return _exp_reward_from_error(ctx, error, ctx.reward_config.std_root_pos) def motion_global_root_ori(ctx: RewardContext) -> np.ndarray: motion_data = ctx.motion_data - robot_body_quat_w = ctx.robot_body_quat_w + robot_body_quat_w = _required_array(ctx.robot_body_quat_w, "robot_body_quat_w") anchor_quat_w = motion_data.body_quat_w[:, ctx.anchor_body_idx] robot_anchor_quat_w = robot_body_quat_w[:, ctx.anchor_body_idx] - np.multiply(anchor_quat_w[:, 0], robot_anchor_quat_w[:, 0], out=ctx.env_error) - np.multiply(anchor_quat_w[:, 1], robot_anchor_quat_w[:, 1], out=ctx.reward_term) - ctx.env_error += ctx.reward_term - np.multiply(anchor_quat_w[:, 2], robot_anchor_quat_w[:, 2], out=ctx.reward_term) - ctx.env_error += ctx.reward_term - np.multiply(anchor_quat_w[:, 3], robot_anchor_quat_w[:, 3], out=ctx.reward_term) - ctx.env_error += ctx.reward_term - np.abs(ctx.env_error, out=ctx.env_error) - np.clip(ctx.env_error, 0.0, 1.0, out=ctx.env_error) - np.arccos(ctx.env_error, out=ctx.env_error) - ctx.env_error *= 2.0 - np.square(ctx.env_error, out=ctx.env_error) - return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_root_ori) + env_error = _required_array(ctx.env_error, "env_error") + reward_term = _required_array(ctx.reward_term, "reward_term") + np.multiply(anchor_quat_w[:, 0], robot_anchor_quat_w[:, 0], out=env_error) + np.multiply(anchor_quat_w[:, 1], robot_anchor_quat_w[:, 1], out=reward_term) + env_error += reward_term + np.multiply(anchor_quat_w[:, 2], robot_anchor_quat_w[:, 2], out=reward_term) + env_error += reward_term + np.multiply(anchor_quat_w[:, 3], robot_anchor_quat_w[:, 3], out=reward_term) + env_error += reward_term + np.abs(env_error, out=env_error) + np.clip(env_error, 0.0, 1.0, out=env_error) + np.arccos(env_error, out=env_error) + env_error *= 2.0 + np.square(env_error, out=env_error) + return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_root_ori) def motion_body_pos(ctx: RewardContext) -> np.ndarray: - robot_body_pos_w = ctx.robot_body_pos_w - error = _mean_body_xyz_squared_error(ctx, ctx.ref_body_pos_w, robot_body_pos_w) + ref_body_pos_w = _required_array(ctx.ref_body_pos_w, "ref_body_pos_w") + robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") + error = _mean_body_xyz_squared_error(ctx, ref_body_pos_w, robot_body_pos_w) return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_pos) def motion_body_ori(ctx: RewardContext) -> np.ndarray: - robot_body_quat_w = ctx.robot_body_quat_w - error = _quat_error_magnitude_squared_body(ctx, ctx.ref_body_quat_w, robot_body_quat_w) - np.sum(error, axis=-1, out=ctx.env_error) - ctx.env_error /= error.shape[1] - return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_body_ori) + ref_body_quat_w = _required_array(ctx.ref_body_quat_w, "ref_body_quat_w") + robot_body_quat_w = _required_array(ctx.robot_body_quat_w, "robot_body_quat_w") + error = _quat_error_magnitude_squared_body(ctx, ref_body_quat_w, robot_body_quat_w) + env_error = _required_array(ctx.env_error, "env_error") + np.sum(error, axis=-1, out=env_error) + env_error /= error.shape[1] + return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_body_ori) def motion_body_lin_vel(ctx: RewardContext) -> np.ndarray: motion_data = ctx.motion_data - robot_body_lin_vel_w = ctx.robot_body_lin_vel_w + robot_body_lin_vel_w = _required_array(ctx.robot_body_lin_vel_w, "robot_body_lin_vel_w") error = _mean_body_xyz_squared_error(ctx, motion_data.body_lin_vel_w, robot_body_lin_vel_w) return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_lin_vel) def motion_body_ang_vel(ctx: RewardContext) -> np.ndarray: motion_data = ctx.motion_data - robot_body_ang_vel_w = ctx.robot_body_ang_vel_w + robot_body_ang_vel_w = _required_array(ctx.robot_body_ang_vel_w, "robot_body_ang_vel_w") error = _mean_body_xyz_squared_error(ctx, motion_data.body_ang_vel_w, robot_body_ang_vel_w) return _exp_reward_from_error(ctx, error, ctx.reward_config.std_body_ang_vel) def motion_ee_body_pos_z(ctx: RewardContext) -> np.ndarray: - robot_body_pos_w = ctx.robot_body_pos_w + ref_body_pos_w = _required_array(ctx.ref_body_pos_w, "ref_body_pos_w") + robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") + ee_body_indices = _required_array(ctx.ee_body_indices, "ee_body_indices") + ee_pos_error_z = _required_array(ctx.ee_pos_error_z, "ee_pos_error_z") + env_error = _required_array(ctx.env_error, "env_error") np.subtract( - ctx.ref_body_pos_w[:, ctx.ee_body_indices, 2], - robot_body_pos_w[:, ctx.ee_body_indices, 2], - out=ctx.ee_pos_error_z, + ref_body_pos_w[:, ee_body_indices, 2], + robot_body_pos_w[:, ee_body_indices, 2], + out=ee_pos_error_z, ) - np.square(ctx.ee_pos_error_z, out=ctx.ee_pos_error_z) - np.sum(ctx.ee_pos_error_z, axis=-1, out=ctx.env_error) - ctx.env_error /= ctx.ee_pos_error_z.shape[1] - return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_body_pos) + np.square(ee_pos_error_z, out=ee_pos_error_z) + np.sum(ee_pos_error_z, axis=-1, out=env_error) + env_error /= ee_pos_error_z.shape[1] + return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_body_pos) def motion_joint_pos(ctx: RewardContext) -> np.ndarray: motion_data = ctx.motion_data - dof_pos = ctx.dof_pos - np.subtract(motion_data.joint_pos, dof_pos, out=ctx.joint_error) - np.square(ctx.joint_error, out=ctx.joint_error) - np.sum(ctx.joint_error, axis=1, out=ctx.env_error) - ctx.env_error /= dof_pos.shape[1] - return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_joint_pos) + dof_pos = _required_array(ctx.dof_pos, "dof_pos") + joint_error = _required_array(ctx.joint_error, "joint_error") + env_error = _required_array(ctx.env_error, "env_error") + np.subtract(motion_data.joint_pos, dof_pos, out=joint_error) + np.square(joint_error, out=joint_error) + np.sum(joint_error, axis=1, out=env_error) + env_error /= dof_pos.shape[1] + return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_joint_pos) def motion_joint_vel(ctx: RewardContext) -> np.ndarray: motion_data = ctx.motion_data - dof_vel = ctx.dof_vel - np.subtract(motion_data.joint_vel, dof_vel, out=ctx.joint_error) - np.square(ctx.joint_error, out=ctx.joint_error) - np.sum(ctx.joint_error, axis=1, out=ctx.env_error) - ctx.env_error /= dof_vel.shape[1] - return _exp_reward_from_error(ctx, ctx.env_error, ctx.reward_config.std_joint_vel) + dof_vel = _required_array(ctx.dof_vel, "dof_vel") + joint_error = _required_array(ctx.joint_error, "joint_error") + env_error = _required_array(ctx.env_error, "env_error") + np.subtract(motion_data.joint_vel, dof_vel, out=joint_error) + np.square(joint_error, out=joint_error) + np.sum(joint_error, axis=1, out=env_error) + env_error /= dof_vel.shape[1] + return _exp_reward_from_error(ctx, env_error, ctx.reward_config.std_joint_vel) def undesired_contacts(ctx: RewardContext) -> np.ndarray: - robot_body_pos_w = ctx.robot_body_pos_w - body_z = robot_body_pos_w[:, ctx.undesired_contact_body_indices, 2] + robot_body_pos_w = _required_array(ctx.robot_body_pos_w, "robot_body_pos_w") + undesired_contact_body_indices = _required_array( + ctx.undesired_contact_body_indices, "undesired_contact_body_indices" + ) + undesired_contact_mask = _required_array(ctx.undesired_contact_mask, "undesired_contact_mask") + env_error = _required_array(ctx.env_error, "env_error") + body_z = robot_body_pos_w[:, undesired_contact_body_indices, 2] np.less( body_z, ctx.undesired_contact_z_threshold, - out=ctx.undesired_contact_mask, + out=undesired_contact_mask, ) - np.sum(ctx.undesired_contact_mask, axis=-1, out=ctx.env_error) - return ctx.env_error + np.sum(undesired_contact_mask, axis=-1, out=env_error) + return env_error def action_rate_l2(ctx: RewardContext) -> np.ndarray: info = ctx.info - np.subtract(info["current_actions"], info["last_actions"], out=ctx.joint_error) - np.square(ctx.joint_error, out=ctx.joint_error) - np.sum(ctx.joint_error, axis=1, out=ctx.env_error) - return ctx.env_error + joint_error = _required_array(ctx.joint_error, "joint_error") + env_error = _required_array(ctx.env_error, "env_error") + np.subtract(info["current_actions"], info["last_actions"], out=joint_error) + np.square(joint_error, out=joint_error) + np.sum(joint_error, axis=1, out=env_error) + return env_error def joint_limit(ctx: RewardContext) -> np.ndarray: - dof_pos = ctx.dof_pos + dof_pos = _required_array(ctx.dof_pos, "dof_pos") + reward_term = _required_array(ctx.reward_term, "reward_term") + joint_error = _required_array(ctx.joint_error, "joint_error") + joint_error_upper = _required_array(ctx.joint_error_upper, "joint_error_upper") lower = ctx.joint_lower upper = ctx.joint_upper if lower is None or upper is None: - ctx.reward_term.fill(0.0) - return ctx.reward_term + reward_term.fill(0.0) + return reward_term # Compute violation - np.subtract(lower, dof_pos, out=ctx.joint_error) - np.maximum(ctx.joint_error, 0, out=ctx.joint_error) - np.subtract(dof_pos, upper, out=ctx.joint_error_upper) - np.maximum(ctx.joint_error_upper, 0, out=ctx.joint_error_upper) - ctx.joint_error += ctx.joint_error_upper - np.square(ctx.joint_error, out=ctx.joint_error) - np.sum(ctx.joint_error, axis=1, out=ctx.reward_term) - return ctx.reward_term + np.subtract(lower, dof_pos, out=joint_error) + np.maximum(joint_error, 0, out=joint_error) + np.subtract(dof_pos, upper, out=joint_error_upper) + np.maximum(joint_error_upper, 0, out=joint_error_upper) + joint_error += joint_error_upper + np.square(joint_error, out=joint_error) + np.sum(joint_error, axis=1, out=reward_term) + return reward_term def build_reward_functions() -> dict[str, Callable[[RewardContext], np.ndarray]]: @@ -327,7 +357,7 @@ def compute_reward( as the per-term scratch, matching the numba kernel's op order. Logs per-term means into ``ctx.info["log"]`` every 4th step, then scales by ``ctrl_dt``. """ - reward = ctx.env_error2 + reward = _required_array(ctx.env_error2, "env_error2") reward.fill(0.0) info = ctx.info @@ -346,7 +376,7 @@ def compute_reward( log[f"reward/{name}"] = 0.0 continue rew = reward_fn(ctx) - weighted_rew = ctx.weighted_reward + weighted_rew = _required_array(ctx.weighted_reward, "weighted_reward") np.multiply(rew, scale, out=weighted_rew) reward += weighted_rew if should_log: diff --git a/src/unilab/envs/motion_tracking/g1/tracking_obs.py b/src/unilab/envs/motion_tracking/g1/tracking_obs.py index b6dfdc157..0a02f47b6 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking_obs.py +++ b/src/unilab/envs/motion_tracking/g1/tracking_obs.py @@ -312,6 +312,8 @@ def _init_reward_functions(self) -> None: def _reward_joint_acc_l2(self, ctx: RewardContext) -> np.ndarray: dof_vel = ctx.dof_vel prev_dof_vel = ctx.info.get("prev_dof_vel") + if dof_vel is None: + raise RuntimeError("RewardContext.dof_vel is required for joint_acc_l2") if prev_dof_vel is None or prev_dof_vel.shape != dof_vel.shape: return np.zeros((self._num_envs,), dtype=get_global_dtype()) joint_acc = (dof_vel - prev_dof_vel) / self._cfg.ctrl_dt @@ -321,6 +323,8 @@ def _reward_joint_torque_l2(self, ctx: RewardContext) -> np.ndarray: dof_pos = ctx.dof_pos dof_vel = ctx.dof_vel last_actions = ctx.info.get("last_actions") + if dof_pos is None or dof_vel is None: + raise RuntimeError("RewardContext.dof_pos and dof_vel are required for joint_torque_l2") if last_actions is None: return np.zeros((self._num_envs,), dtype=get_global_dtype()) target_q = ( diff --git a/tests/envs/test_motion_tracking_rewards.py b/tests/envs/test_motion_tracking_rewards.py index ddc7fabb3..406a00e9b 100644 --- a/tests/envs/test_motion_tracking_rewards.py +++ b/tests/envs/test_motion_tracking_rewards.py @@ -100,9 +100,7 @@ def test_motion_joint_pos_term_matches_hand_computed(): def test_action_rate_l2_term_matches_hand_computed(): ctx = _make_ctx() out = rewards.action_rate_l2(ctx).copy() - expected = np.sum( - np.square(ctx.info["current_actions"] - ctx.info["last_actions"]), axis=1 - ) + expected = np.sum(np.square(ctx.info["current_actions"] - ctx.info["last_actions"]), axis=1) np.testing.assert_allclose(out, expected, rtol=1e-12, atol=1e-12) @@ -142,7 +140,5 @@ def test_compute_reward_matches_hand_computed_weighted_sum(): enable_log=False, ).copy() - expected = ( - 1.0 * joint_pos_ref - 0.1 * action_rate_ref - 0.5 * undesired_ref - ) * ctrl_dt + expected = (1.0 * joint_pos_ref - 0.1 * action_rate_ref - 0.5 * undesired_ref) * ctrl_dt np.testing.assert_allclose(reward, expected, rtol=1e-12, atol=1e-12)