diff --git a/AGENTS.md b/AGENTS.md index c96a0e92c..501dd26ca 100644 --- a/AGENTS.md +++ b/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` diff --git a/docs/sphinx/source/api_reference/envs/common.md b/docs/sphinx/source/api_reference/envs/common.md deleted file mode 100644 index 35ff7a466..000000000 --- a/docs/sphinx/source/api_reference/envs/common.md +++ /dev/null @@ -1,13 +0,0 @@ -# `unilab.envs.common` - -Rotation / math helpers used by every env. - -```{eval-rst} -.. automodule:: unilab.envs.common.math - :members: -``` - -```{eval-rst} -.. automodule:: unilab.envs.common.rotation - :members: -``` diff --git a/docs/sphinx/source/api_reference/envs/index.md b/docs/sphinx/source/api_reference/envs/index.md index 2b1babc66..6665cd124 100644 --- a/docs/sphinx/source/api_reference/envs/index.md +++ b/docs/sphinx/source/api_reference/envs/index.md @@ -15,7 +15,6 @@ can be selected via `uv run train --algo --task --sim `. locomotion manipulation motion_tracking -common ``` ```{eval-rst} diff --git a/docs/sphinx/source/api_reference/utils/index.md b/docs/sphinx/source/api_reference/utils/index.md index 3b5c461da..99e6e4270 100644 --- a/docs/sphinx/source/api_reference/utils/index.md +++ b/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:: diff --git a/scripts/deploy/sim_prototype.py b/scripts/deploy/sim_prototype.py index e8db273f6..7043b9bd2 100644 --- a/scripts/deploy/sim_prototype.py +++ b/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, ) diff --git a/scripts/manip_loco/diagnose_go2_arm_ik.py b/scripts/manip_loco/diagnose_go2_arm_ik.py index a1a03d821..601eb70b7 100644 --- a/scripts/manip_loco/diagnose_go2_arm_ik.py +++ b/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]: diff --git a/scripts/motion/bones_seed_csv_to_npz.py b/scripts/motion/bones_seed_csv_to_npz.py index 3fc10ecc0..3031cce4d 100644 --- a/scripts/motion/bones_seed_csv_to_npz.py +++ b/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", diff --git a/scripts/motion/csv_to_npz.py b/scripts/motion/csv_to_npz.py index e3d2f5ddc..2794a7114 100644 --- a/scripts/motion/csv_to_npz.py +++ b/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: diff --git a/scripts/motion/x2_csv_to_tracking_npz.py b/scripts/motion/x2_csv_to_tracking_npz.py index ba2515349..3892eaba2 100644 --- a/scripts/motion/x2_csv_to_tracking_npz.py +++ b/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 diff --git a/src/unilab/envs/common/__init__.py b/src/unilab/envs/common/__init__.py deleted file mode 100644 index a05e8faba..000000000 --- a/src/unilab/envs/common/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Shared environment-owned numeric helpers.""" diff --git a/src/unilab/envs/common/math.py b/src/unilab/envs/common/math.py deleted file mode 100644 index e96d7f105..000000000 --- a/src/unilab/envs/common/math.py +++ /dev/null @@ -1,13 +0,0 @@ -from __future__ import annotations - -import numpy as np - - -def np_sample_uniform( - lower: float | np.ndarray, - upper: float | np.ndarray, - size: tuple[int, ...], - dtype=np.float32, -) -> np.ndarray: - """Sample uniformly from [lower, upper] with output dtype.""" - return np.random.uniform(lower, upper, size).astype(dtype) diff --git a/src/unilab/envs/locomotion/common/commands.py b/src/unilab/envs/locomotion/common/commands.py index 9a2069d2f..e14f32252 100644 --- a/src/unilab/envs/locomotion/common/commands.py +++ b/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 diff --git a/src/unilab/envs/locomotion/common/dr_provider.py b/src/unilab/envs/locomotion/common/dr_provider.py index 2c9db268c..24f75ab54 100644 --- a/src/unilab/envs/locomotion/common/dr_provider.py +++ b/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): diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index 9a1734a0b..6a5985fd9 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/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 diff --git a/src/unilab/envs/locomotion/go1/rough.py b/src/unilab/envs/locomotion/go1/rough.py index 1bbb084bc..9f67c2305 100644 --- a/src/unilab/envs/locomotion/go1/rough.py +++ b/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 diff --git a/src/unilab/envs/locomotion/go2/footstand.py b/src/unilab/envs/locomotion/go2/footstand.py index 956233403..24d876279 100644 --- a/src/unilab/envs/locomotion/go2/footstand.py +++ b/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 diff --git a/src/unilab/envs/locomotion/go2/rough.py b/src/unilab/envs/locomotion/go2/rough.py index f37952ce8..0d7d5975e 100644 --- a/src/unilab/envs/locomotion/go2/rough.py +++ b/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 diff --git a/src/unilab/envs/locomotion/go2_arm/base.py b/src/unilab/envs/locomotion/go2_arm/base.py index 087923395..c54e13718 100644 --- a/src/unilab/envs/locomotion/go2_arm/base.py +++ b/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 diff --git a/src/unilab/envs/locomotion/go2_arm/manip_loco.py b/src/unilab/envs/locomotion/go2_arm/manip_loco.py index f93a48440..08194b4c4 100644 --- a/src/unilab/envs/locomotion/go2_arm/manip_loco.py +++ b/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]) diff --git a/src/unilab/envs/locomotion/go2w/joystick.py b/src/unilab/envs/locomotion/go2w/joystick.py index 1b7c2082f..43059e83e 100644 --- a/src/unilab/envs/locomotion/go2w/joystick.py +++ b/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) diff --git a/src/unilab/envs/locomotion/go2w/rough.py b/src/unilab/envs/locomotion/go2w/rough.py index f27506801..ab96b61ba 100644 --- a/src/unilab/envs/locomotion/go2w/rough.py +++ b/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 diff --git a/src/unilab/envs/manipulation/allegro_inhand/rotation.py b/src/unilab/envs/manipulation/allegro_inhand/rotation.py index d1597434e..74173a807 100644 --- a/src/unilab/envs/manipulation/allegro_inhand/rotation.py +++ b/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( diff --git a/src/unilab/envs/manipulation/sharpa_inhand/base.py b/src/unilab/envs/manipulation/sharpa_inhand/base.py index 6b330ec4e..bdace32ee 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/base.py +++ b/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", diff --git a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py b/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py index bdbc5c81f..2081bf95e 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py +++ b/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py @@ -10,7 +10,6 @@ from unilab.base.np_env import NpEnvState from unilab.dr import ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization -from unilab.envs.common.rotation import np_quat_error_magnitude from unilab.envs.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, SharpaDomainRandConfig, @@ -22,6 +21,7 @@ SharpaInhandRotationDRProvider, SharpaInhandRotationEnv, ) +from unilab.utils.rotation import np_quat_error_magnitude def _default_sharpa_grasp_domain_rand() -> SharpaDomainRandConfig: diff --git a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py b/src/unilab/envs/manipulation/sharpa_inhand/rotation.py index 542298ef9..04285d40b 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py +++ b/src/unilab/envs/manipulation/sharpa_inhand/rotation.py @@ -30,12 +30,6 @@ ResetRandomizationPayload, ) from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( - np_quat_apply, - np_quat_conjugate, - np_quat_mul, - np_quat_to_axis_angle, -) from unilab.envs.manipulation.sharpa_inhand.base import ( SharpaInhandBaseCfg, SharpaInhandBaseEnv, @@ -43,6 +37,13 @@ resolve_grasp_cache_file, sample_scale_grasp_caches, ) +from unilab.utils.geometry import np_sample_uniform_quaternion +from unilab.utils.rotation import ( + np_quat_apply, + np_quat_conjugate, + np_quat_mul, + np_quat_to_axis_angle, +) @dataclass @@ -76,18 +77,9 @@ def sample_random_quaternion(num_envs: int) -> np.ndarray: num_envs: Number of quaternions to sample. Returns: - Quaternion array with shape ``(num_envs, 4)``. + Quaternion array with shape ``(num_envs, 4)`` in float64. """ - u1 = np.random.rand(num_envs) - u2 = np.random.rand(num_envs) * 2.0 * np.pi - u3 = np.random.rand(num_envs) * 2.0 * np.pi - - q1 = np.sqrt(1.0 - u1) * np.sin(u2) - q2 = np.sqrt(1.0 - u1) * np.cos(u2) - q3 = np.sqrt(u1) * np.sin(u3) - q4 = np.sqrt(u1) * np.cos(u3) - - return np.stack([q4, q1, q2, q3], axis=1).astype(np.float64) + return np_sample_uniform_quaternion(num_envs).astype(np.float64) class SharpaInhandRotationDRProvider(DomainRandomizationProvider): diff --git a/src/unilab/envs/manipulation/stewart/balance.py b/src/unilab/envs/manipulation/stewart/balance.py index 63e91f57e..75552146f 100644 --- a/src/unilab/envs/manipulation/stewart/balance.py +++ b/src/unilab/envs/manipulation/stewart/balance.py @@ -23,7 +23,8 @@ from unilab.dr.provider import DomainRandomizationProvider from unilab.dr.types import DomainRandomizationCapabilities, ResetPlan from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( +from unilab.utils.geometry import np_roll_pitch_from_quat +from unilab.utils.rotation import ( np_quat_apply_batched, np_quat_apply_inverse, np_quat_conjugate_batched, @@ -120,13 +121,8 @@ def _ball_home_z(cfg: StewartBalanceCfg) -> float: def _roll_pitch_from_quat(quat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Extract roll/pitch (rad) from a wxyz quaternion via its rotation matrix.""" - w, x, y, z = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3] - r20 = 2.0 * (x * z - w * y) - r21 = 2.0 * (y * z + w * x) - r22 = 1.0 - 2.0 * (x * x + y * y) - roll = np.arctan2(r21, r22) - pitch = np.arctan2(-r20, np.sqrt(np.clip(r21 * r21 + r22 * r22, 0.0, None))) + """Extract roll/pitch (rad) from a wxyz quaternion, cast to float32.""" + roll, pitch = np_roll_pitch_from_quat(quat) return roll.astype(np.float32), pitch.astype(np.float32) diff --git a/src/unilab/envs/motion_tracking/g1/box_tracking.py b/src/unilab/envs/motion_tracking/g1/box_tracking.py index ebf78600e..5024323bb 100644 --- a/src/unilab/envs/motion_tracking/g1/box_tracking.py +++ b/src/unilab/envs/motion_tracking/g1/box_tracking.py @@ -13,8 +13,8 @@ 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.math import np_sample_uniform -from unilab.envs.common.rotation import ( +from unilab.utils.geometry import np_sample_uniform +from unilab.utils.rotation import ( np_matrix_from_quat, np_quat_apply, np_quat_error_magnitude, diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index eac620b60..151519922 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -28,14 +28,18 @@ ) from unilab.dr.types import RESET_TERM_GEOM_FRICTION, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype -from unilab.envs.common.math import np_sample_uniform -from unilab.envs.common.rotation import ( +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 unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv from .motion_loader import MotionData, MotionLoader, MotionSampler @@ -273,64 +277,6 @@ def _build_motion_reference_state( return qpos, qvel -def _gravity_z_in_body_from_quat_w(quat_w: np.ndarray) -> np.ndarray: - """Z component of world gravity ``[0, 0, -1]`` expressed in body frame.""" - return 2.0 * (quat_w[:, 1] * quat_w[:, 1] + quat_w[:, 2] * quat_w[:, 2]) - 1.0 - - -def _write_motion_anchor_transform( - robot_anchor_pos_w: np.ndarray, - robot_anchor_quat_w: np.ndarray, - anchor_pos_w: np.ndarray, - anchor_quat_w: np.ndarray, - out_pos: np.ndarray, - out_ori6: np.ndarray, -) -> None: - aw = robot_anchor_quat_w[:, 0] - ax = robot_anchor_quat_w[:, 1] - ay = robot_anchor_quat_w[:, 2] - az = robot_anchor_quat_w[:, 3] - - vx = anchor_pos_w[:, 0] - robot_anchor_pos_w[:, 0] - vy = anchor_pos_w[:, 1] - robot_anchor_pos_w[:, 1] - vz = anchor_pos_w[:, 2] - robot_anchor_pos_w[:, 2] - - qx = -ax - qy = -ay - qz = -az - tx = 2 * (qy * vz - qz * vy) - ty = 2 * (qz * vx - qx * vz) - tz = 2 * (qx * vy - qy * vx) - out_pos[:, 0] = vx + aw * tx + qy * tz - qz * ty - out_pos[:, 1] = vy + aw * ty + qz * tx - qx * tz - out_pos[:, 2] = vz + aw * tz + qx * ty - qy * tx - - bw = anchor_quat_w[:, 0] - bx = anchor_quat_w[:, 1] - by = anchor_quat_w[:, 2] - bz = anchor_quat_w[:, 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 - - xx = rx * rx - yy = ry * ry - zz = rz * rz - xy = rx * ry - xz = rx * rz - yz = ry * rz - wx = rw * rx - wy = rw * ry - wz = rw * rz - out_ori6[:, 0] = 1 - 2 * (yy + zz) - out_ori6[:, 1] = 2 * (xy - wz) - out_ori6[:, 2] = 2 * (xy + wz) - out_ori6[:, 3] = 1 - 2 * (xx + zz) - out_ori6[:, 4] = 2 * (xz - wy) - out_ori6[:, 5] = 2 * (yz + wx) - - class G1MotionTrackingDomainRandomizationProvider(DomainRandomizationProvider): def __init__( self, @@ -1001,8 +947,8 @@ def _compute_terminations( 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 = _gravity_z_in_body_from_quat_w(anchor_quat_w) - robot_gravity_z_b = _gravity_z_in_body_from_quat_w(robot_anchor_quat_w) + 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) @@ -1127,7 +1073,7 @@ def _compute_obs( 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) - _write_motion_anchor_transform( + np_write_relative_anchor_transform_pos_rot6d( robot_anchor_pos_w, robot_anchor_quat_w, anchor_pos_w, diff --git a/src/unilab/utils/geometry.py b/src/unilab/utils/geometry.py new file mode 100644 index 000000000..019360092 --- /dev/null +++ b/src/unilab/utils/geometry.py @@ -0,0 +1,214 @@ +"""Standalone numpy helpers for coordinate frames and geometry. + +Reusable pure-numpy helpers for rotations, coordinate-frame transforms, +and geometric conversions. Kept dtype-agnostic and side-effect-free so +that env / task / backend code can compose them without carrying task +policy inside a shared module. Complements the vectorized quaternion +primitives in :mod:`unilab.utils.rotation`. +""" + +from __future__ import annotations + +import numpy as np + +from unilab.utils.rotation import ( + np_quat_canonicalize, + np_quat_conjugate, + np_quat_inv, + np_quat_mul, + np_quat_to_axis_angle, +) + + +def np_sample_uniform( + lower: float | np.ndarray, + upper: float | np.ndarray, + size: tuple[int, ...], + dtype=np.float32, +) -> np.ndarray: + """Sample uniformly from ``[lower, upper]`` and cast to ``dtype``.""" + return np.random.uniform(lower, upper, size).astype(dtype) + + +def np_quat_normalize(q: np.ndarray) -> np.ndarray: + """L2-normalize quaternion(s), clamping tiny norms to 1e-8 to avoid divide-by-zero.""" + q = np.asarray(q) + if q.ndim == 1: + norm = float(np.linalg.norm(q)) + return q / max(norm, 1.0e-8) + norm = np.linalg.norm(q, axis=-1, keepdims=True) + return q / np.clip(norm, 1.0e-8, None) + + +def np_normalize_axis(axis: np.ndarray | tuple[float, ...] | list[float]) -> np.ndarray: + """Return a unit-length copy of a rotation axis vector. Raises on zero norm.""" + axis = np.asarray(axis) + norm = float(np.linalg.norm(axis)) + if norm <= 0.0: + raise ValueError(f"axis must be non-zero, got {axis!r}") + return axis / norm + + +def np_roll_pitch_from_quat(quat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Roll and pitch (rad) from a w-first quaternion, computed via rotation-matrix rows. + + Supports either ``(4,)`` or ``(..., 4)`` inputs. Returned arrays match + the caller's leading shape and dtype (no upcast). + """ + w = quat[..., 0] + x = quat[..., 1] + y = quat[..., 2] + z = quat[..., 3] + r20 = 2.0 * (x * z - w * y) + r21 = 2.0 * (y * z + w * x) + r22 = 1.0 - 2.0 * (x * x + y * y) + roll = np.arctan2(r21, r22) + pitch = np.arctan2(-r20, np.sqrt(np.clip(r21 * r21 + r22 * r22, 0.0, None))) + return roll, pitch + + +def np_gravity_z_in_body_from_quat(quat_w: np.ndarray) -> np.ndarray: + """Z component of world gravity ``[0, 0, -1]`` expressed in body frame. + + Equivalent to ``np_quat_apply_inverse(quat_w, [0, 0, -1])[..., 2]`` but + computed directly from quaternion components to skip the intermediate. + """ + return 2.0 * (quat_w[..., 1] * quat_w[..., 1] + quat_w[..., 2] * quat_w[..., 2]) - 1.0 + + +def np_quat_orientation_error_local(goal_quat: np.ndarray, curr_quat: np.ndarray) -> np.ndarray: + """Local-frame orientation error as the signed xyz of the relative quaternion. + + Both inputs are re-normalized to unit length. Returns the imaginary + part of ``goal * inv(curr)`` after canonicalization, giving a + signed-axis-scaled error suitable for PD-style feedback terms. + """ + goal = np_quat_normalize(goal_quat) + curr = np_quat_normalize(curr_quat) + if goal.ndim == 1: + goal = goal[None, :] + if curr.ndim == 1: + curr = curr[None, :] + 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 + + +def np_quat_angular_velocity_from_pair( + quat: np.ndarray, prev_quat: np.ndarray, dt: float +) -> np.ndarray: + """Angular velocity from two consecutive quaternions via axis-angle diff / dt.""" + rel = np_quat_mul(quat, np_quat_conjugate(prev_quat)) + return np_quat_to_axis_angle(rel) / dt + + +def np_sample_uniform_quaternion(num_samples: int) -> np.ndarray: + """Sample uniformly random unit quaternions (w-first) via Shoemake (1992). + + Returns an ``(num_samples, 4)`` array in float64 (leaves any downstream + dtype conversion to the caller). + """ + u1 = np.random.rand(num_samples) + u2 = np.random.rand(num_samples) * 2.0 * np.pi + u3 = np.random.rand(num_samples) * 2.0 * np.pi + + r1 = np.sqrt(1.0 - u1) + r2 = np.sqrt(u1) + q1 = r1 * np.sin(u2) + q2 = r1 * np.cos(u2) + q3 = r2 * np.sin(u3) + q4 = r2 * np.cos(u3) + + return np.stack([q4, q1, q2, q3], axis=1) + + +def np_spherical_to_cartesian(sphere: np.ndarray) -> np.ndarray: + """Convert ``(..., 3)[l, phi, theta]`` to ``(..., 3)[x, y, z]``. + + Uses the go2-arm spherical convention: ``phi`` sweeps in the x-z plane + from the positive x axis, and ``theta`` measures elevation toward + positive y. + """ + length = sphere[..., 0] + phi = sphere[..., 1] + theta = sphere[..., 2] + x = length * np.cos(phi) * np.cos(theta) + y = length * np.sin(theta) + z = length * np.sin(phi) * np.cos(theta) + return np.stack([x, y, z], axis=-1) + + +def np_cartesian_to_spherical(cart: np.ndarray) -> np.ndarray: + """Convert ``(..., 3)[x, y, z]`` to ``(..., 3)[l, phi, theta]`` (inverse of + :func:`np_spherical_to_cartesian`).""" + cart = np.asarray(cart) + l_sq = np.sum(cart**2, axis=-1, keepdims=True) + length = 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] / length, -1.0, 1.0)) + return np.concatenate([length, phi, theta], axis=-1) + + +def np_write_relative_anchor_transform_pos_rot6d( + source_anchor_pos_w: np.ndarray, + source_anchor_quat_w: np.ndarray, + target_anchor_pos_w: np.ndarray, + target_anchor_quat_w: np.ndarray, + out_pos: np.ndarray, + out_rot6d: np.ndarray, +) -> None: + """Fused frame-transform + 6D rotation flatten, writing in place. + + Computes the position of ``target_anchor`` in ``source_anchor``'s frame + (written to ``out_pos`` of shape ``(N, 3)``) and the relative rotation + ``conj(source_anchor) * target_anchor`` expressed as the flattened + first two columns of its rotation matrix (written to ``out_rot6d`` of + shape ``(N, 6)``). No intermediate quaternion arrays are allocated. + + This is the numpy analogue of the numba kernel + :func:`unilab.utils.numba_geometry.write_relative_anchor_transform_at`. + """ + aw = source_anchor_quat_w[:, 0] + ax = source_anchor_quat_w[:, 1] + ay = source_anchor_quat_w[:, 2] + az = source_anchor_quat_w[:, 3] + + vx = target_anchor_pos_w[:, 0] - source_anchor_pos_w[:, 0] + vy = target_anchor_pos_w[:, 1] - source_anchor_pos_w[:, 1] + vz = target_anchor_pos_w[:, 2] - source_anchor_pos_w[:, 2] + + qx = -ax + qy = -ay + qz = -az + tx = 2 * (qy * vz - qz * vy) + ty = 2 * (qz * vx - qx * vz) + tz = 2 * (qx * vy - qy * vx) + out_pos[:, 0] = vx + aw * tx + qy * tz - qz * ty + out_pos[:, 1] = vy + aw * ty + qz * tx - qx * tz + out_pos[:, 2] = vz + aw * tz + qx * ty - qy * tx + + bw = target_anchor_quat_w[:, 0] + bx = target_anchor_quat_w[:, 1] + by = target_anchor_quat_w[:, 2] + bz = target_anchor_quat_w[:, 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 + + xx = rx * rx + yy = ry * ry + zz = rz * rz + xy = rx * ry + xz = rx * rz + yz = ry * rz + wx = rw * rx + wy = rw * ry + wz = rw * rz + out_rot6d[:, 0] = 1 - 2 * (yy + zz) + out_rot6d[:, 1] = 2 * (xy - wz) + out_rot6d[:, 2] = 2 * (xy + wz) + out_rot6d[:, 3] = 1 - 2 * (xx + zz) + out_rot6d[:, 4] = 2 * (xz - wy) + out_rot6d[:, 5] = 2 * (yz + wx) diff --git a/src/unilab/envs/common/rotation.py b/src/unilab/utils/rotation.py similarity index 100% rename from src/unilab/envs/common/rotation.py rename to src/unilab/utils/rotation.py diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index be0eaa9a8..3d9b0009f 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -654,13 +654,13 @@ def test_g1_motion_tracking_critic_uses_clean_beyondmimic_aligned_terms(): def test_g1_motion_tracking_anchor_frame_writers_match_reference(): - from unilab.envs.common.rotation import ( + from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.utils.rotation import ( np_matrix_from_quat, np_quat_apply, np_quat_inv, np_quat_mul, ) - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv rng = np.random.default_rng(123) num_envs = 4 @@ -704,8 +704,8 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_relative_transform_fast_path_matches_reference(): - from unilab.envs.common.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.utils.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat rng = np.random.default_rng(321) num_envs = 4 @@ -761,9 +761,9 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_reward_fast_path_matches_reference(): - from unilab.envs.common.rotation import np_quat_error_magnitude from unilab.envs.motion_tracking.g1.motion_loader import MotionData from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig + from unilab.utils.rotation import np_quat_error_magnitude rng = np.random.default_rng(456) num_envs = 3 diff --git a/tests/test_sharpa.py b/tests/test_sharpa.py index 4f5059305..5b4618f02 100644 --- a/tests/test_sharpa.py +++ b/tests/test_sharpa.py @@ -5,7 +5,6 @@ import numpy as np import pytest -from unilab.envs.common.rotation import np_quat_apply from unilab.envs.manipulation.sharpa_inhand.base import SharpaDomainRandConfig from unilab.envs.manipulation.sharpa_inhand.grasp_gen import ( SharpaInhandRotationGraspCfg, @@ -15,6 +14,7 @@ SharpaInhandRotationEnv, sample_random_quaternion, ) +from unilab.utils.rotation import np_quat_apply def test_sharpa_gravity_direction_randomization_matches_rotated_gravity( diff --git a/tests/utils/test_math_utils.py b/tests/utils/test_math_utils.py index 2193f1aed..da9c4effa 100644 --- a/tests/utils/test_math_utils.py +++ b/tests/utils/test_math_utils.py @@ -1,10 +1,10 @@ -"""Tests for quaternion helpers in unilab.envs.common.rotation.""" +"""Tests for quaternion helpers in unilab.utils.rotation.""" from __future__ import annotations import numpy as np -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_matrix_first_two_cols_from_quat, np_matrix_from_quat, np_quat_angular_velocity, diff --git a/tests/utils/test_utils_package_policy.py b/tests/utils/test_utils_package_policy.py index b424cc414..d3f58e317 100644 --- a/tests/utils/test_utils_package_policy.py +++ b/tests/utils/test_utils_package_policy.py @@ -4,7 +4,15 @@ import unilab.utils ALLOWED_UTILS_API = {"get_default_device", "to_numpy", "to_torch"} -ALLOWED_UTILS_MODULES = {"__init__", "device", "nan_guard", "support_matrix", "tensor"} +ALLOWED_UTILS_MODULES = { + "__init__", + "device", + "geometry", + "nan_guard", + "rotation", + "support_matrix", + "tensor", +} REMOVED_UTILS_SHIMS = { "algo_utils", "device_utils",