refactor: split motion tracking common modules

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(

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.
+"""

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