From a88aefe16ad73b82a3369a781ed4d3ef991f69e7 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 26 Aug 2026 19:48:53 +0800 Subject: [PATCH] perf(motion): fuse relative transforms in numba (#1318) --- .../tasks/motion_tracking/common/kernels.py | 291 +++++++++++++++++- .../motion_tracking/common/manager_terms.py | 114 +------ .../motion_tracking/common/observations.py | 66 ---- .../motion_tracking/common/transforms.py | 103 ------- .../tasks/motion_tracking/g1/manager_terms.py | 47 +-- src/unilab/utils/geometry.py | 63 ---- tests/tasks/test_motion_term_parity.py | 152 ++++++++- 7 files changed, 465 insertions(+), 371 deletions(-) delete mode 100644 src/unilab/tasks/motion_tracking/common/observations.py delete mode 100644 src/unilab/tasks/motion_tracking/common/transforms.py diff --git a/src/unilab/tasks/motion_tracking/common/kernels.py b/src/unilab/tasks/motion_tracking/common/kernels.py index e7a761c9b..d722b76c8 100644 --- a/src/unilab/tasks/motion_tracking/common/kernels.py +++ b/src/unilab/tasks/motion_tracking/common/kernels.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from math import atan2, sqrt +from math import atan2, cos, sin, sqrt import numpy as np from numba import config, njit, prange, set_num_threads @@ -51,6 +51,293 @@ def _quat_error_squared( return angle * angle +@njit(inline="always") +def _quat_mul_components( + w1: float, + x1: float, + y1: float, + z1: float, + w2: float, + x2: float, + y2: float, + z2: float, +) -> tuple[float, float, float, float]: + """Hamilton product for scalar quaternion components.""" + return ( + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ) + + +@njit(inline="always") +def _quat_apply_inverse_components( + w: float, + x: float, + y: float, + z: float, + vx: float, + vy: float, + vz: float, +) -> tuple[float, float, float]: + """Rotate one vector by the inverse of a unit quaternion.""" + qx = -x + qy = -y + qz = -z + tx = 2.0 * (qy * vz - qz * vy) + ty = 2.0 * (qz * vx - qx * vz) + tz = 2.0 * (qx * vy - qy * vx) + return ( + vx + w * tx + qy * tz - qz * ty, + vy + w * ty + qz * tx - qx * tz, + vz + w * tz + qx * ty - qy * tx, + ) + + +@njit(inline="always") +def _quat_to_rot6d_components( + w: float, + x: float, + y: float, + z: float, +) -> tuple[float, float, float, float, float, float]: + """Flatten the first two rotation-matrix columns.""" + xx = x * x + yy = y * y + zz = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + return ( + 1.0 - 2.0 * (yy + zz), + 2.0 * (xy - wz), + 2.0 * (xy + wz), + 1.0 - 2.0 * (xx + zz), + 2.0 * (xz - wy), + 2.0 * (yz + wx), + ) + + +@njit(cache=True, nogil=True, parallel=True) +def update_motion_relative_state_kernel( + env_ids: np.ndarray, + anchor_body_idx: int, + motion_body_pos_local_w: np.ndarray, + motion_body_pos_w: np.ndarray, + motion_body_quat_w: np.ndarray, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + body_pos_relative_w: np.ndarray, + body_quat_relative_w: np.ndarray, + motion_anchor_pos_b: np.ndarray, + motion_anchor_ori_b: np.ndarray, + robot_body_pos_b: np.ndarray, + robot_body_ori_b: np.ndarray, +) -> None: + """Write all MotionCommand relative transforms for selected rows.""" + num_bodies = motion_body_pos_local_w.shape[1] + for row in prange(env_ids.shape[0]): + env_idx = env_ids[row] + motion_anchor_x = motion_body_pos_local_w[env_idx, anchor_body_idx, 0] + motion_anchor_y = motion_body_pos_local_w[env_idx, anchor_body_idx, 1] + motion_anchor_z = motion_body_pos_local_w[env_idx, anchor_body_idx, 2] + motion_anchor_w = motion_body_quat_w[env_idx, anchor_body_idx, 0] + motion_anchor_qx = motion_body_quat_w[env_idx, anchor_body_idx, 1] + motion_anchor_qy = motion_body_quat_w[env_idx, anchor_body_idx, 2] + motion_anchor_qz = motion_body_quat_w[env_idx, anchor_body_idx, 3] + robot_anchor_x = robot_body_pos_w[env_idx, anchor_body_idx, 0] + robot_anchor_y = robot_body_pos_w[env_idx, anchor_body_idx, 1] + robot_anchor_z = robot_body_pos_w[env_idx, anchor_body_idx, 2] + robot_anchor_w = robot_body_quat_w[env_idx, anchor_body_idx, 0] + robot_anchor_qx = robot_body_quat_w[env_idx, anchor_body_idx, 1] + robot_anchor_qy = robot_body_quat_w[env_idx, anchor_body_idx, 2] + robot_anchor_qz = robot_body_quat_w[env_idx, anchor_body_idx, 3] + + yaw_w, yaw_x, yaw_y, yaw_z = _quat_mul_components( + robot_anchor_w, + robot_anchor_qx, + robot_anchor_qy, + robot_anchor_qz, + motion_anchor_w, + -motion_anchor_qx, + -motion_anchor_qy, + -motion_anchor_qz, + ) + half_yaw = 0.5 * atan2( + 2.0 * (yaw_w * yaw_z + yaw_x * yaw_y), + 1.0 - 2.0 * (yaw_y * yaw_y + yaw_z * yaw_z), + ) + delta_w = cos(half_yaw) + delta_z = sin(half_yaw) + yaw_cross = 2.0 * delta_w * delta_z + yaw_z2 = 2.0 * delta_z * delta_z + + anchor_vx = motion_body_pos_w[env_idx, anchor_body_idx, 0] - robot_anchor_x + anchor_vy = motion_body_pos_w[env_idx, anchor_body_idx, 1] - robot_anchor_y + anchor_vz = motion_body_pos_w[env_idx, anchor_body_idx, 2] - robot_anchor_z + anchor_pos_x, anchor_pos_y, anchor_pos_z = _quat_apply_inverse_components( + robot_anchor_w, + robot_anchor_qx, + robot_anchor_qy, + robot_anchor_qz, + anchor_vx, + anchor_vy, + anchor_vz, + ) + motion_anchor_pos_b[env_idx, 0] = anchor_pos_x + motion_anchor_pos_b[env_idx, 1] = anchor_pos_y + motion_anchor_pos_b[env_idx, 2] = anchor_pos_z + rel_w, rel_x, rel_y, rel_z = _quat_mul_components( + robot_anchor_w, + -robot_anchor_qx, + -robot_anchor_qy, + -robot_anchor_qz, + motion_anchor_w, + motion_anchor_qx, + motion_anchor_qy, + motion_anchor_qz, + ) + anchor_rot6d = _quat_to_rot6d_components(rel_w, rel_x, rel_y, rel_z) + for component in range(6): + motion_anchor_ori_b[env_idx, component] = anchor_rot6d[component] + + for body_idx in range(num_bodies): + body_motion_w = motion_body_quat_w[env_idx, body_idx, 0] + body_motion_x = motion_body_quat_w[env_idx, body_idx, 1] + body_motion_y = motion_body_quat_w[env_idx, body_idx, 2] + body_motion_z = motion_body_quat_w[env_idx, body_idx, 3] + out_w, out_x, out_y, out_z = _quat_mul_components( + delta_w, + 0.0, + 0.0, + delta_z, + body_motion_w, + body_motion_x, + body_motion_y, + body_motion_z, + ) + body_quat_relative_w[env_idx, body_idx, 0] = out_w + body_quat_relative_w[env_idx, body_idx, 1] = out_x + body_quat_relative_w[env_idx, body_idx, 2] = out_y + body_quat_relative_w[env_idx, body_idx, 3] = out_z + + vx = motion_body_pos_local_w[env_idx, body_idx, 0] - motion_anchor_x + vy = motion_body_pos_local_w[env_idx, body_idx, 1] - motion_anchor_y + vz = motion_body_pos_local_w[env_idx, body_idx, 2] - motion_anchor_z + relative_x = vx + relative_x -= yaw_cross * vy + relative_x -= yaw_z2 * vx + relative_x += robot_anchor_x + relative_y = vy + relative_y += yaw_cross * vx + relative_y -= yaw_z2 * vy + relative_y += robot_anchor_y + body_pos_relative_w[env_idx, body_idx, 0] = relative_x + body_pos_relative_w[env_idx, body_idx, 1] = relative_y + body_pos_relative_w[env_idx, body_idx, 2] = vz + motion_anchor_z + + robot_vx = robot_body_pos_w[env_idx, body_idx, 0] - robot_anchor_x + robot_vy = robot_body_pos_w[env_idx, body_idx, 1] - robot_anchor_y + robot_vz = robot_body_pos_w[env_idx, body_idx, 2] - robot_anchor_z + robot_pos_x, robot_pos_y, robot_pos_z = _quat_apply_inverse_components( + robot_anchor_w, + robot_anchor_qx, + robot_anchor_qy, + robot_anchor_qz, + robot_vx, + robot_vy, + robot_vz, + ) + robot_body_pos_b[env_idx, body_idx, 0] = robot_pos_x + robot_body_pos_b[env_idx, body_idx, 1] = robot_pos_y + robot_body_pos_b[env_idx, body_idx, 2] = robot_pos_z + body_robot_w = robot_body_quat_w[env_idx, body_idx, 0] + body_robot_x = robot_body_quat_w[env_idx, body_idx, 1] + body_robot_y = robot_body_quat_w[env_idx, body_idx, 2] + body_robot_z = robot_body_quat_w[env_idx, body_idx, 3] + robot_rel_w, robot_rel_x, robot_rel_y, robot_rel_z = _quat_mul_components( + robot_anchor_w, + -robot_anchor_qx, + -robot_anchor_qy, + -robot_anchor_qz, + body_robot_w, + body_robot_x, + body_robot_y, + body_robot_z, + ) + robot_rot6d = _quat_to_rot6d_components( + robot_rel_w, + robot_rel_x, + robot_rel_y, + robot_rel_z, + ) + for component in range(6): + robot_body_ori_b[env_idx, body_idx, component] = robot_rot6d[component] + + +@njit(cache=True, nogil=True, parallel=True) +def update_object_relative_state_kernel( + env_ids: np.ndarray, + robot_anchor_pos_w: np.ndarray, + robot_anchor_quat_w: np.ndarray, + object_pos_w: np.ndarray, + object_quat_w: np.ndarray, + object_lin_vel_w: np.ndarray, + object_state_b: np.ndarray, +) -> None: + """Write BoxMotionCommand object pose and velocity for selected rows.""" + for row in prange(env_ids.shape[0]): + env_idx = env_ids[row] + anchor_x = robot_anchor_pos_w[env_idx, 0] + anchor_y = robot_anchor_pos_w[env_idx, 1] + anchor_z = robot_anchor_pos_w[env_idx, 2] + anchor_w = robot_anchor_quat_w[env_idx, 0] + anchor_qx = robot_anchor_quat_w[env_idx, 1] + anchor_qy = robot_anchor_quat_w[env_idx, 2] + anchor_qz = robot_anchor_quat_w[env_idx, 3] + pos_x, pos_y, pos_z = _quat_apply_inverse_components( + anchor_w, + anchor_qx, + anchor_qy, + anchor_qz, + object_pos_w[env_idx, 0] - anchor_x, + object_pos_w[env_idx, 1] - anchor_y, + object_pos_w[env_idx, 2] - anchor_z, + ) + object_state_b[env_idx, 0] = pos_x + object_state_b[env_idx, 1] = pos_y + object_state_b[env_idx, 2] = pos_z + rel_w, rel_x, rel_y, rel_z = _quat_mul_components( + anchor_w, + -anchor_qx, + -anchor_qy, + -anchor_qz, + object_quat_w[env_idx, 0], + object_quat_w[env_idx, 1], + object_quat_w[env_idx, 2], + object_quat_w[env_idx, 3], + ) + rot6d = _quat_to_rot6d_components(rel_w, rel_x, rel_y, rel_z) + for component in range(6): + object_state_b[env_idx, 3 + component] = rot6d[component] + vel_x, vel_y, vel_z = _quat_apply_inverse_components( + anchor_w, + anchor_qx, + anchor_qy, + anchor_qz, + object_lin_vel_w[env_idx, 0], + object_lin_vel_w[env_idx, 1], + object_lin_vel_w[env_idx, 2], + ) + object_state_b[env_idx, 9] = vel_x + object_state_b[env_idx, 10] = vel_y + object_state_b[env_idx, 11] = vel_z + + @njit(cache=True, nogil=True, parallel=True) def update_motion_metrics_kernel( env_ids: np.ndarray, @@ -328,4 +615,6 @@ def reward_motion_body_ang_vel_kernel( "reward_motion_body_pos_kernel", "termination_anchor_pos_kernel", "update_motion_metrics_kernel", + "update_motion_relative_state_kernel", + "update_object_relative_state_kernel", ] diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index bd4053fc4..26e85651d 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -5,7 +5,6 @@ import dataclasses import math from dataclasses import dataclass, field -from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np @@ -13,9 +12,6 @@ from unilab.envs.mdp.actions import JointPositionAction, JointPositionActionCfg from unilab.managers import CommandTerm, CommandTermCfg, ManagerTermBase, ManagerTermBaseCfg from unilab.managers.scene_entity_config import SceneEntityCfg -from unilab.utils.geometry import ( - np_write_relative_anchor_transform_pos_rot6d, -) from unilab.utils.rotation import ( np_quat_apply_inverse, np_quat_error_magnitude_squared_batched, @@ -31,10 +27,9 @@ reward_motion_body_pos_kernel, termination_anchor_pos_kernel, update_motion_metrics_kernel, + update_motion_relative_state_kernel, ) from .motion_loader import MotionData, MotionLoader, MotionSampler -from .observations import write_body_ori6_in_anchor_frame, write_body_pos_in_anchor_frame -from .transforms import update_relative_transforms if TYPE_CHECKING: from unilab.base.entity import Entity @@ -191,11 +186,6 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): self.robot_body_pos_b = np.empty_like(self._body_pos_w) self.robot_body_ori_b = np.empty((self.num_envs, num_bodies, 6), dtype=dtype) self.joint_default_bias = np.zeros((self.num_envs, num_joints), dtype=dtype) - self._delta_pos_w = np.empty((self.num_envs, 3), dtype=dtype) - self._delta_ori_w = np.empty((self.num_envs, 4), dtype=dtype) - self._body_vec_error = np.empty_like(self._body_pos_w) - self._env_error = np.empty(self.num_envs, dtype=dtype) - self._reward_term = np.empty(self.num_envs, dtype=dtype) self._robot_cache_step = -1 self._all_env_ids = np.arange(self.num_envs, dtype=np.int32) self._all_env_ids.setflags(write=False) @@ -226,10 +216,10 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): self.metrics[name] = np.zeros(self.num_envs, dtype=dtype) self._refresh_motion() self._refresh_robot_state(force=True) - self._refresh_relative_state() - # Compile and execute the metrics kernel on the cold path so the first - # measured command-manager step contains no Numba dispatch/JIT latency. + # Configure and compile both fused kernels on the cold path so the first + # measured manager step contains no Numba worker/JIT initialization. configure_motion_kernel_runtime() + self._refresh_relative_state() self._update_metrics() def _make_motion_loader( @@ -428,101 +418,23 @@ def _refresh_robot_state( self._robot_cache_step = step def _refresh_relative_state(self, env_ids: np.ndarray | None = None) -> None: - if env_ids is not None: - self._refresh_relative_state_rows(env_ids) - return - update_relative_transforms( - self, - self._motion_data, + rows = self._all_env_ids if env_ids is None else env_ids + update_motion_relative_state_kernel( + rows, + self.anchor_body_idx, + self._motion_data.body_pos_w, + self._body_pos_w, + self._motion_data.body_quat_w, self._robot_body_pos_w, self._robot_body_quat_w, - ) - np_write_relative_anchor_transform_pos_rot6d( - self.robot_anchor_pos_w, - self.robot_anchor_quat_w, - self.anchor_pos_w, - self.anchor_quat_w, + self.body_pos_relative_w, + self.body_quat_relative_w, self.motion_anchor_pos_b, self.motion_anchor_ori_b, - ) - write_body_pos_in_anchor_frame( - self.robot_anchor_pos_w, - self.robot_anchor_quat_w, - self._robot_body_pos_w, self.robot_body_pos_b, - body_vec_error=self._body_vec_error, - ) - write_body_ori6_in_anchor_frame( - self.robot_anchor_quat_w, - self._robot_body_quat_w, self.robot_body_ori_b, ) - def _refresh_relative_state_rows(self, env_ids: np.ndarray) -> None: - """Row-scoped variant of `_refresh_relative_state` for partial resets. - - Computes the same transforms on the gathered reset rows and scatters the - results back; untouched rows keep their per-step values. - """ - num_rows = len(env_ids) - num_bodies = self._robot_body_pos_w.shape[1] - dtype = self._body_pos_w.dtype - robot_pos_rows = self._robot_body_pos_w[env_ids] - robot_quat_rows = self._robot_body_quat_w[env_ids] - motion_rows = SimpleNamespace( - body_pos_w=self._motion_data.body_pos_w[env_ids], - body_quat_w=self._motion_data.body_quat_w[env_ids], - ) - # `update_relative_transforms` reads/writes these attributes on the env; - # a namespace with row-sized scratch keeps the shared buffers untouched. - scratch = SimpleNamespace( - anchor_body_idx=self.anchor_body_idx, - _delta_pos_w=np.empty((num_rows, 3), dtype=dtype), - _delta_ori_w=np.empty((num_rows, 4), dtype=dtype), - body_quat_relative_w=np.empty((num_rows, num_bodies, 4), dtype=dtype), - body_pos_relative_w=np.empty((num_rows, num_bodies, 3), dtype=dtype), - _body_vec_error=np.empty((num_rows, num_bodies, 3), dtype=dtype), - _env_error=np.empty(num_rows, dtype=dtype), - _reward_term=np.empty(num_rows, dtype=dtype), - ) - update_relative_transforms(scratch, motion_rows, robot_pos_rows, robot_quat_rows) - self.body_pos_relative_w[env_ids] = scratch.body_pos_relative_w - self.body_quat_relative_w[env_ids] = scratch.body_quat_relative_w - - anchor_idx = self.anchor_body_idx - robot_anchor_pos_rows = robot_pos_rows[:, anchor_idx] - robot_anchor_quat_rows = robot_quat_rows[:, anchor_idx] - anchor_pos_rows = self._body_pos_w[env_ids][:, anchor_idx] - anchor_quat_rows = motion_rows.body_quat_w[:, anchor_idx] - motion_anchor_pos_b = np.empty((num_rows, 3), dtype=dtype) - motion_anchor_ori_b = np.empty((num_rows, 6), dtype=dtype) - np_write_relative_anchor_transform_pos_rot6d( - robot_anchor_pos_rows, - robot_anchor_quat_rows, - anchor_pos_rows, - anchor_quat_rows, - motion_anchor_pos_b, - motion_anchor_ori_b, - ) - self.motion_anchor_pos_b[env_ids] = motion_anchor_pos_b - self.motion_anchor_ori_b[env_ids] = motion_anchor_ori_b - robot_body_pos_b = np.empty((num_rows, num_bodies, 3), dtype=dtype) - write_body_pos_in_anchor_frame( - robot_anchor_pos_rows, - robot_anchor_quat_rows, - robot_pos_rows, - robot_body_pos_b, - body_vec_error=scratch._body_vec_error, - ) - self.robot_body_pos_b[env_ids] = robot_body_pos_b - robot_body_ori_b = np.empty((num_rows, num_bodies, 6), dtype=dtype) - write_body_ori6_in_anchor_frame( - robot_anchor_quat_rows, - robot_quat_rows, - robot_body_ori_b, - ) - self.robot_body_ori_b[env_ids] = robot_body_ori_b - def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: # All row-wise metrics are written by one Numba kernel. Passing an # explicit all-row index buffer for normal steps lets the same kernel diff --git a/src/unilab/tasks/motion_tracking/common/observations.py b/src/unilab/tasks/motion_tracking/common/observations.py deleted file mode 100644 index 6332c4ac5..000000000 --- a/src/unilab/tasks/motion_tracking/common/observations.py +++ /dev/null @@ -1,66 +0,0 @@ -"""In-place body-frame writers shared by motion manager terms.""" - -from __future__ import annotations - -import numpy as np - - -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, num_bodies = body_pos.shape[:2] - rel_pos = body_vec_error[:num_envs, :num_bodies] - 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) - - -__all__ = ["write_body_ori6_in_anchor_frame", "write_body_pos_in_anchor_frame"] diff --git a/src/unilab/tasks/motion_tracking/common/transforms.py b/src/unilab/tasks/motion_tracking/common/transforms.py deleted file mode 100644 index 31b4f31a6..000000000 --- a/src/unilab/tasks/motion_tracking/common/transforms.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Shared 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 stable observation and reward semantics. -""" - -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] diff --git a/src/unilab/tasks/motion_tracking/g1/manager_terms.py b/src/unilab/tasks/motion_tracking/g1/manager_terms.py index 32e34cbb8..f9764680b 100644 --- a/src/unilab/tasks/motion_tracking/g1/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/g1/manager_terms.py @@ -9,16 +9,13 @@ from unilab.managers import ManagerTermBase, ManagerTermBaseCfg from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.tasks.motion_tracking.common.kernels import update_object_relative_state_kernel from unilab.tasks.motion_tracking.common.manager_terms import ( MotionCommand, MotionCommandCfg, MotionJointPositionAction, ) -from unilab.utils.geometry import np_write_relative_anchor_transform_pos_rot6d -from unilab.utils.rotation import ( - np_quat_apply_inverse, - np_quat_error_magnitude_squared_batched, -) +from unilab.utils.rotation import np_quat_error_magnitude_squared_batched from .motion_box_loader import BoxMotionData, BoxMotionLoader @@ -115,37 +112,15 @@ def _resample_command(self, env_ids: np.ndarray) -> None: self.object.write_root_state_to_sim(object_state, env_ids=env_ids) def _refresh_object_state(self, env_ids: np.ndarray | None = None) -> None: - if env_ids is None: - np_write_relative_anchor_transform_pos_rot6d( - self.robot_anchor_pos_w, - self.robot_anchor_quat_w, - self.object.data.root_link_pos_w, - self.object.data.root_link_quat_w, - self._object_obs_b[:, :3], - self._object_obs_b[:, 3:9], - ) - self._object_obs_b[:, 9:12] = np_quat_apply_inverse( - self.robot_anchor_quat_w, - self.object.data.root_link_lin_vel_w, - ) - return - num_rows = len(env_ids) - dtype = self._object_obs_b.dtype - pos_b = np.empty((num_rows, 3), dtype=dtype) - rot_b = np.empty((num_rows, 6), dtype=dtype) - np_write_relative_anchor_transform_pos_rot6d( - self.robot_anchor_pos_w[env_ids], - self.robot_anchor_quat_w[env_ids], - self.object.data.root_link_pos_w[env_ids], - self.object.data.root_link_quat_w[env_ids], - pos_b, - rot_b, - ) - self._object_obs_b[env_ids, :3] = pos_b - self._object_obs_b[env_ids, 3:9] = rot_b - self._object_obs_b[env_ids, 9:12] = np_quat_apply_inverse( - self.robot_anchor_quat_w[env_ids], - self.object.data.root_link_lin_vel_w[env_ids], + rows = self._all_env_ids if env_ids is None else env_ids + update_object_relative_state_kernel( + rows, + self.robot_anchor_pos_w, + self.robot_anchor_quat_w, + self.object.data.root_link_pos_w, + self.object.data.root_link_quat_w, + self.object.data.root_link_lin_vel_w, + self._object_obs_b, ) def post_compute(self) -> None: diff --git a/src/unilab/utils/geometry.py b/src/unilab/utils/geometry.py index 5dc933d3f..5358f4f6d 100644 --- a/src/unilab/utils/geometry.py +++ b/src/unilab/utils/geometry.py @@ -148,66 +148,3 @@ def np_cartesian_to_spherical(cart: np.ndarray) -> np.ndarray: 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 helper computes the relative anchor transform used by motion tracking. - """ - 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/tests/tasks/test_motion_term_parity.py b/tests/tasks/test_motion_term_parity.py index 18c6ca75b..b0d1e2b18 100644 --- a/tests/tasks/test_motion_term_parity.py +++ b/tests/tasks/test_motion_term_parity.py @@ -13,7 +13,15 @@ from unilab.managers import RewardTermCfg, TerminationTermCfg from unilab.tasks.motion_tracking.common import kernels from unilab.tasks.motion_tracking.common import manager_terms as mt -from unilab.utils.rotation import np_quat_error_magnitude_squared_batched +from unilab.utils.rotation import ( + np_matrix_first_two_cols_from_quat, + np_quat_apply_batched, + np_quat_apply_inverse_batched, + np_quat_error_magnitude_squared_batched, + np_quat_inv, + np_quat_mul_batched, + np_yaw_quat, +) def _make_env(command: Any) -> SimpleNamespace: @@ -413,6 +421,148 @@ def run(rows: np.ndarray) -> None: assert kernels.update_motion_metrics_kernel.signatures +def test_motion_relative_state_kernel_matches_numpy_and_scopes_rows() -> None: + rng = np.random.default_rng(1818) + num_envs, num_bodies = 257, 12 + anchor_body_idx = 4 + motion_pos_local = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + env_origins = rng.standard_normal((num_envs, 1, 3), dtype=np.float32) + motion_pos_world = motion_pos_local + env_origins + motion_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + robot_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + inputs = (motion_pos_local, motion_pos_world, motion_quat, robot_pos, robot_quat) + snapshots = tuple(value.copy() for value in inputs) + + motion_anchor_pos_local = motion_pos_local[:, anchor_body_idx] + motion_anchor_quat = motion_quat[:, anchor_body_idx] + robot_anchor_pos = robot_pos[:, anchor_body_idx] + robot_anchor_quat = robot_quat[:, anchor_body_idx] + delta_pos = robot_anchor_pos.copy() + delta_pos[:, 2] = motion_anchor_pos_local[:, 2] + delta_quat = np_yaw_quat( + np_quat_mul_batched(robot_anchor_quat, np_quat_inv(motion_anchor_quat)) + ) + expected_body_pos_relative = np_quat_apply_batched( + delta_quat[:, None], + motion_pos_local - motion_anchor_pos_local[:, None], + ) + expected_body_pos_relative += delta_pos[:, None] + expected_body_quat_relative = np_quat_mul_batched(delta_quat[:, None], motion_quat) + expected_motion_anchor_pos = np_quat_apply_inverse_batched( + robot_anchor_quat, + motion_pos_world[:, anchor_body_idx] - robot_anchor_pos, + ) + expected_motion_anchor_ori = np_matrix_first_two_cols_from_quat( + np_quat_mul_batched(np_quat_inv(robot_anchor_quat), motion_anchor_quat) + ) + expected_robot_body_pos = np_quat_apply_inverse_batched( + robot_anchor_quat[:, None], + robot_pos - robot_anchor_pos[:, None], + ) + expected_robot_body_ori = np_matrix_first_two_cols_from_quat( + np_quat_mul_batched(np_quat_inv(robot_anchor_quat)[:, None], robot_quat) + ) + expected = ( + expected_body_pos_relative, + expected_body_quat_relative, + expected_motion_anchor_pos, + expected_motion_anchor_ori, + expected_robot_body_pos, + expected_robot_body_ori, + ) + outputs = tuple(np.full(value.shape, -123.0, dtype=np.float32) for value in expected) + output_addresses = tuple(value.ctypes.data for value in outputs) + + def run(rows: np.ndarray) -> None: + kernels.update_motion_relative_state_kernel( + rows, + anchor_body_idx, + motion_pos_local, + motion_pos_world, + motion_quat, + robot_pos, + robot_quat, + *outputs, + ) + + selected = np.asarray([0, 3, 128, 256], dtype=np.int32) + run(selected) + untouched = np.ones(num_envs, dtype=bool) + untouched[selected] = False + for actual, reference in zip(outputs, expected, strict=True): + np.testing.assert_allclose(actual[selected], reference[selected], rtol=3e-6, atol=2e-6) + np.testing.assert_array_equal(actual[untouched], -123.0) + + run(np.arange(num_envs, dtype=np.int32)) + for actual, reference in zip(outputs, expected, strict=True): + np.testing.assert_allclose(actual, reference, rtol=3e-6, atol=2e-6) + for actual, snapshot in zip(inputs, snapshots, strict=True): + np.testing.assert_array_equal(actual, snapshot) + assert tuple(value.ctypes.data for value in outputs) == output_addresses + assert kernels.update_motion_relative_state_kernel.targetoptions["nopython"] is True + assert kernels.update_motion_relative_state_kernel.targetoptions["nogil"] is True + assert kernels.update_motion_relative_state_kernel.targetoptions["parallel"] is True + assert kernels.update_motion_relative_state_kernel.signatures + + +def test_object_relative_state_kernel_matches_numpy_and_scopes_rows() -> None: + rng = np.random.default_rng(1819) + num_envs = 257 + anchor_pos = rng.standard_normal((num_envs, 3), dtype=np.float32) + anchor_quat = _unit_quat(rng.standard_normal((num_envs, 4), dtype=np.float32)) + object_pos = rng.standard_normal((num_envs, 3), dtype=np.float32) + object_quat = _unit_quat(rng.standard_normal((num_envs, 4), dtype=np.float32)) + object_lin_vel = rng.standard_normal((num_envs, 3), dtype=np.float32) + inputs = (anchor_pos, anchor_quat, object_pos, object_quat, object_lin_vel) + snapshots = tuple(value.copy() for value in inputs) + expected = np.concatenate( + ( + np_quat_apply_inverse_batched(anchor_quat, object_pos - anchor_pos), + np_matrix_first_two_cols_from_quat( + np_quat_mul_batched(np_quat_inv(anchor_quat), object_quat) + ), + np_quat_apply_inverse_batched(anchor_quat, object_lin_vel), + ), + axis=-1, + ) + output = np.full(expected.shape, -123.0, dtype=np.float32) + output_address = output.ctypes.data + + selected = np.asarray([0, 3, 128, 256], dtype=np.int32) + kernels.update_object_relative_state_kernel( + selected, + anchor_pos, + anchor_quat, + object_pos, + object_quat, + object_lin_vel, + output, + ) + untouched = np.ones(num_envs, dtype=bool) + untouched[selected] = False + np.testing.assert_allclose(output[selected], expected[selected], rtol=3e-6, atol=2e-6) + np.testing.assert_array_equal(output[untouched], -123.0) + + kernels.update_object_relative_state_kernel( + np.arange(num_envs, dtype=np.int32), + anchor_pos, + anchor_quat, + object_pos, + object_quat, + object_lin_vel, + output, + ) + np.testing.assert_allclose(output, expected, rtol=3e-6, atol=2e-6) + for actual, snapshot in zip(inputs, snapshots, strict=True): + np.testing.assert_array_equal(actual, snapshot) + assert output.ctypes.data == output_address + assert kernels.update_object_relative_state_kernel.targetoptions["nopython"] is True + assert kernels.update_object_relative_state_kernel.targetoptions["nogil"] is True + assert kernels.update_object_relative_state_kernel.targetoptions["parallel"] is True + assert kernels.update_object_relative_state_kernel.signatures + + def test_joint_pos_limits_bit_parity() -> None: rng = np.random.default_rng(123) joint_pos = rng.standard_normal((8, 5), dtype=np.float32)