refactor: consolidate numpy geometry helpers under unilab.utils

AGENTS.md

@@ -46,7 +46,7 @@ UniLab 是一个 **高性能、模块化、contract 驱动** 的 RL infrastructu
 - backend contract: `src/unilab/base/backend/base.py`
 - training run helpers: `src/unilab/training/run.py`
 - visualization helpers: `src/unilab/visualization/`
-- env shared numeric helpers: `src/unilab/envs/common/rotation.py`, `src/unilab/envs/common/math.py`
+- shared numeric helpers: `src/unilab/utils/rotation.py`, `src/unilab/utils/geometry.py`
 - MLX rotation helpers: `src/unilab/algos/mlx/common/rotation.py`
 - config schema: `src/unilab/structured_configs.py`
 - async runner: `src/unilab/ipc/async_runner.py`

docs/sphinx/source/api_reference/envs/index.md

@@ -15,7 +15,6 @@ can be selected via `uv run train --algo <algo> --task <name> --sim <backend>`.
 locomotion
 manipulation
 motion_tracking
-common
 ```
 
 ```{eval-rst}

docs/sphinx/source/api_reference/utils/index.md

@@ -1,6 +1,7 @@
 # `unilab.utils` — Utilities
 
-Device probing, tensor helpers, support-matrix bookkeeping, NaN guards.
+Device probing, tensor helpers, support-matrix bookkeeping, NaN guards,
+and pure-numpy geometry/rotation helpers shared across envs and scripts.
 
 ```{eval-rst}
 .. autosummary::

scripts/deploy/sim_prototype.py

@@ -42,7 +42,7 @@
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 sys.path.insert(0, str(REPO_ROOT / "src"))
-from unilab.envs.common.rotation import (  # noqa: E402
+from unilab.utils.rotation import (  # noqa: E402
     np_matrix_from_quat,
     np_subtract_frame_transforms,
 )

scripts/manip_loco/diagnose_go2_arm_ik.py

@@ -16,8 +16,8 @@
 
 from unilab.base import registry
 from unilab.base.registry import ensure_registries
-from unilab.envs.common.rotation import np_matrix_from_quat
 from unilab.envs.locomotion.go2_arm.manip_loco import RewardConfig
+from unilab.utils.rotation import np_matrix_from_quat
 
 
 def _parse_vec3(raw: list[float], *, name: str) -> list[float]:

scripts/motion/bones_seed_csv_to_npz.py

@@ -42,7 +42,7 @@
 
 from unilab.assets import ASSETS_ROOT_PATH
 from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors
-from unilab.envs.common.rotation import np_quat_angular_velocity, np_quat_ensure_continuity
+from unilab.utils.rotation import np_quat_angular_velocity, np_quat_ensure_continuity
 
 ROOT_COLUMNS = [
     "Frame",

scripts/motion/csv_to_npz.py

@@ -30,7 +30,7 @@
 
 from unilab.assets import ASSETS_ROOT_PATH
 from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors
-from unilab.envs.common.rotation import np_quat_angular_velocity, np_quat_ensure_continuity
+from unilab.utils.rotation import np_quat_angular_velocity, np_quat_ensure_continuity
 
 
 def quat_slerp(q1: np.ndarray, q2: np.ndarray, t: float) -> np.ndarray:

scripts/motion/x2_csv_to_tracking_npz.py

@@ -20,7 +20,7 @@
 
 from unilab.assets import ASSETS_ROOT_PATH
 from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors
-from unilab.envs.common.rotation import np_quat_angular_velocity, np_quat_ensure_continuity
+from unilab.utils.rotation import np_quat_angular_velocity, np_quat_ensure_continuity
 
 ROOT_QPOS_DIM = 7
 ROOT_QVEL_DIM = 6

src/unilab/envs/locomotion/common/commands.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import np_wrap_to_pi, np_yaw_from_quat
+from unilab.utils.rotation import np_wrap_to_pi, np_yaw_from_quat
 
 
 @dataclass

src/unilab/envs/locomotion/common/dr_provider.py

@@ -26,7 +26,7 @@
     zero_actions,
 )
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import np_quat_mul, np_yaw_to_quat
+from unilab.utils.rotation import np_quat_mul, np_yaw_to_quat
 
 
 class LocomotionDRProvider(DomainRandomizationProvider):

src/unilab/envs/locomotion/go1/joystick.py

@@ -12,7 +12,6 @@
 from unilab.base.np_env import NpEnvState
 from unilab.base.scene import SceneCfg
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import np_quat_mul, np_yaw_to_quat
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import Commands
 from unilab.envs.locomotion.common.domain_rand import DomainRandConfig
@@ -23,6 +22,7 @@
     TerrainSpawnManager,
 )
 from unilab.envs.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv
+from unilab.utils.rotation import np_quat_mul, np_yaw_to_quat
 
 
 @dataclass

src/unilab/envs/locomotion/go1/rough.py

@@ -14,11 +14,6 @@
 from unilab.dr import DomainRandomizationManager, ResetPlan
 from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import (
-    np_quat_apply_inverse,
-    np_quat_from_euler_xyz,
-    np_quat_mul,
-)
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import (
     Commands,
@@ -58,6 +53,11 @@
     random_rough,
     wave_terrain,
 )
+from unilab.utils.rotation import (
+    np_quat_apply_inverse,
+    np_quat_from_euler_xyz,
+    np_quat_mul,
+)
 
 # pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false
 

src/unilab/envs/locomotion/go2/footstand.py

@@ -13,13 +13,13 @@
 from unilab.base.scene import SceneCfg
 from unilab.dr import ResetPlan, ResetRandomizationPayload
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import np_quat_apply, np_quat_apply_inverse
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import Commands
 from unilab.envs.locomotion.common.domain_rand import DomainRandConfig
 from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider
 from unilab.envs.locomotion.common.rewards import RewardContext
 from unilab.envs.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig
+from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse
 
 
 @dataclass

src/unilab/envs/locomotion/go2/rough.py

@@ -12,11 +12,6 @@
 from unilab.dr import DomainRandomizationManager, ResetPlan
 from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import (
-    np_quat_apply_inverse,
-    np_quat_from_euler_xyz,
-    np_quat_mul,
-)
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import (
     apply_heading_yaw_feedback,
@@ -54,6 +49,11 @@
     random_rough,
     wave_terrain,
 )
+from unilab.utils.rotation import (
+    np_quat_apply_inverse,
+    np_quat_from_euler_xyz,
+    np_quat_mul,
+)
 
 # pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false
 

src/unilab/envs/locomotion/go2_arm/base.py

@@ -5,18 +5,14 @@
 import numpy as np
 
 from unilab.base.backend import SimBackend
-from unilab.envs.common.rotation import (
-    np_matrix_from_quat,
-    np_quat_canonicalize,
-    np_quat_inv,
-    np_quat_mul,
-)
 from unilab.envs.locomotion.common.base import (
     ControlConfigBase,
     LocomotionBaseCfg,
     LocomotionBaseEnv,
     Sensor,
 )
+from unilab.utils.geometry import np_quat_orientation_error_local
+from unilab.utils.rotation import np_matrix_from_quat
 
 DEFAULT_LEG_ANGLES = np.asarray(
     [
@@ -241,7 +237,7 @@ def compute_arm_ik_delta(
                         "goal_local_quat and curr_local_quat are required when "
                         "ik.use_orientation=true and ik.orientation_mode='target'"
                     )
-                orn_err = _orientation_error_local(goal_local_quat, curr_local_quat)
+                orn_err = np_quat_orientation_error_local(goal_local_quat, curr_local_quat)
             elif cfg.orientation_mode == "zero_error":
                 orn_err = np.zeros_like(pos_err)
             else:
@@ -264,20 +260,3 @@ def compute_arm_ik_delta(
         if cfg.dq_clip > 0.0:
             dq = np.clip(dq, -cfg.dq_clip, cfg.dq_clip)
         return dq
-
-
-def _normalize_quat(q: np.ndarray) -> np.ndarray:
-    q = np.asarray(q)
-    if q.ndim == 1:
-        q = q[None, :]
-    norm = np.linalg.norm(q, axis=1, keepdims=True)
-    return q / np.clip(norm, 1.0e-8, None)
-
-
-def _orientation_error_local(goal_quat: np.ndarray, curr_quat: np.ndarray) -> np.ndarray:
-    goal = _normalize_quat(goal_quat)
-    curr = _normalize_quat(curr_quat)
-    rel = np_quat_mul(goal, np_quat_inv(curr))
-    rel = np_quat_canonicalize(rel)
-    sign = np.where(rel[:, 0:1] < 0.0, -1.0, 1.0)
-    return rel[:, 1:] * sign

src/unilab/envs/locomotion/go2_arm/manip_loco.py

@@ -12,7 +12,6 @@
 from unilab.base.scene import SceneCfg
 from unilab.dr.types import ResetPlan
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import np_matrix_from_quat, np_quat_from_euler_xyz
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import Commands
 from unilab.envs.locomotion.common.domain_rand import DomainRandConfig
@@ -25,6 +24,13 @@
     Go2ArmSensor,
     build_go2_arm_position_gains,
 )
+from unilab.utils.geometry import (
+    np_cartesian_to_spherical as _cart2sphere,
+)
+from unilab.utils.geometry import (
+    np_spherical_to_cartesian as _sphere2cart,
+)
+from unilab.utils.rotation import np_matrix_from_quat, np_quat_from_euler_xyz
 
 
 def _default_go2_arm_model_file() -> str:
@@ -50,27 +56,6 @@ def _resolve_go2_arm_scene(cfg: "Go2ArmManipLocoCfg") -> SceneCfg:
     return scene
 
 
-def _sphere2cart(sphere: np.ndarray) -> np.ndarray:
-    """Convert (..., 3)[l, phi, theta] to (..., 3)[x, y, z]."""
-    l = sphere[..., 0]
-    phi = sphere[..., 1]
-    theta = sphere[..., 2]
-    x = l * np.cos(phi) * np.cos(theta)
-    y = l * np.sin(theta)
-    z = l * np.sin(phi) * np.cos(theta)
-    return np.stack([x, y, z], axis=-1)
-
-
-def _cart2sphere(cart: np.ndarray) -> np.ndarray:
-    """Convert (..., 3)[x, y, z] to (..., 3)[l, phi, theta]."""
-    cart = np.asarray(cart)
-    l_sq = np.sum(cart**2, axis=-1, keepdims=True)
-    l = np.sqrt(np.maximum(l_sq, 1e-12))
-    phi = np.arctan2(cart[..., 2:3], cart[..., 0:1])
-    theta = np.arcsin(np.clip(cart[..., 1:2] / l, -1.0, 1.0))
-    return np.concatenate([l, phi, theta], axis=-1)
-
-
 @dataclass
 class InitState:
     pos: list[float] = field(default_factory=lambda: [0.0, 0.0, 0.42])

src/unilab/envs/locomotion/go2w/joystick.py

@@ -17,10 +17,6 @@
     zero_actions,
 )
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import (
-    np_quat_mul,
-    np_yaw_to_quat,
-)
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import (
     Commands,
@@ -43,6 +39,10 @@
     compute_go2w_motor_ctrl,
     stack_joint_sensors,
 )
+from unilab.utils.rotation import (
+    np_quat_mul,
+    np_yaw_to_quat,
+)
 
 GO2W_HIP_INDICES = np.asarray([0, 3, 6, 9], dtype=np.int32)
 

src/unilab/envs/locomotion/go2w/rough.py

@@ -12,10 +12,6 @@
 from unilab.dr import DomainRandomizationManager, ResetPlan
 from unilab.dr.dr_utils import zero_actions
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import (
-    np_quat_from_euler_xyz,
-    np_quat_mul,
-)
 from unilab.envs.locomotion.common import rewards
 from unilab.envs.locomotion.common.commands import (
     Commands,
@@ -54,6 +50,10 @@
     random_rough,
     wave_terrain,
 )
+from unilab.utils.rotation import (
+    np_quat_from_euler_xyz,
+    np_quat_mul,
+)
 
 # pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false
 

src/unilab/envs/manipulation/allegro_inhand/rotation.py

@@ -27,7 +27,10 @@
     zero_actions,
 )
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import np_quat_conjugate, np_quat_mul, np_quat_to_axis_angle
+from unilab.utils.geometry import (
+    np_normalize_axis,
+    np_quat_angular_velocity_from_pair,
+)
 
 from .base import AllegroBaseCfg, AllegroBaseEnv
 
@@ -41,15 +44,19 @@ def resolve_grasp_cache_path(cache_path: str) -> epath.Path:
 
 
 def normalize_rotation_axis(rotation_axis: tuple[float, float, float]) -> np.ndarray:
+    # Cast to the training dtype first so the norm and division happen at that
+    # precision, matching the pre-refactor bit-exact behavior for float32 runs.
     axis = np.asarray(rotation_axis, dtype=get_global_dtype())
-    return np.asarray(axis / np.linalg.norm(axis), dtype=get_global_dtype())
+    return np.asarray(np_normalize_axis(axis), dtype=get_global_dtype())
 
 
 def compute_ball_angvel(
     ball_quat: np.ndarray, prev_ball_quat: np.ndarray, ctrl_dt: float
 ) -> np.ndarray:
-    rel_quat = np_quat_mul(ball_quat, np_quat_conjugate(prev_ball_quat))
-    return np.asarray(np_quat_to_axis_angle(rel_quat) / ctrl_dt, dtype=get_global_dtype())
+    return np.asarray(
+        np_quat_angular_velocity_from_pair(ball_quat, prev_ball_quat, ctrl_dt),
+        dtype=get_global_dtype(),
+    )
 
 
 def compute_pd_torques(

src/unilab/envs/manipulation/sharpa_inhand/base.py

@@ -13,7 +13,7 @@
 from unilab.base.np_env import NpEnv, NpEnvState
 from unilab.base.scene import SceneCfg
 from unilab.dtype_config import get_global_dtype
-from unilab.envs.common.rotation import np_quat_apply, np_quat_mul
+from unilab.utils.rotation import np_quat_apply, np_quat_mul
 
 DEFAULT_ACTUATED_JOINT_NAMES: list[str] = [
     "right_thumb_CMC_FE",