From 761b1d13ea392144d35840d86c0c9e8b3f55cb82 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 26 Aug 2026 19:27:57 +0800 Subject: [PATCH] perf(motion): fuse command metrics in numba (#1317) --- .../tasks/motion_tracking/common/kernels.py | 181 +++++++++++++++++- .../motion_tracking/common/manager_terms.py | 75 ++++---- tests/tasks/test_motion_term_parity.py | 93 +++++++++ 3 files changed, 302 insertions(+), 47 deletions(-) diff --git a/src/unilab/tasks/motion_tracking/common/kernels.py b/src/unilab/tasks/motion_tracking/common/kernels.py index c046bb052..e7a761c9b 100644 --- a/src/unilab/tasks/motion_tracking/common/kernels.py +++ b/src/unilab/tasks/motion_tracking/common/kernels.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from math import atan2, sqrt import numpy as np from numba import config, njit, prange, set_num_threads @@ -28,6 +29,174 @@ def configure_motion_kernel_runtime() -> None: _runtime_configured = True +@njit(inline="always") +def _quat_error_squared( + q1_w: float, + q1_x: float, + q1_y: float, + q1_z: float, + q2_w: float, + q2_x: float, + q2_y: float, + q2_z: float, +) -> float: + """Return the squared shortest-path angular error for two quaternions.""" + rel_w = abs(q2_w * q1_w + q2_x * q1_x + q2_y * q1_y + q2_z * q1_z) + rel_x = -q2_w * q1_x + q2_x * q1_w - q2_y * q1_z + q2_z * q1_y + rel_y = -q2_w * q1_y + q2_x * q1_z + q2_y * q1_w - q2_z * q1_x + rel_z = -q2_w * q1_z - q2_x * q1_y + q2_y * q1_x + q2_z * q1_w + xyz_norm = sqrt(rel_x * rel_x + rel_y * rel_y + rel_z * rel_z) + clipped_w = min(max(rel_w, -1.0), 1.0) + angle = 2.0 * atan2(xyz_norm, clipped_w) + return angle * angle + + +@njit(cache=True, nogil=True, parallel=True) +def update_motion_metrics_kernel( + env_ids: np.ndarray, + anchor_body_idx: int, + motion_body_pos_w: np.ndarray, + robot_body_pos_w: np.ndarray, + motion_body_quat_w: np.ndarray, + robot_body_quat_w: np.ndarray, + motion_body_lin_vel_w: np.ndarray, + robot_body_lin_vel_w: np.ndarray, + motion_body_ang_vel_w: np.ndarray, + robot_body_ang_vel_w: np.ndarray, + body_pos_relative_w: np.ndarray, + body_quat_relative_w: np.ndarray, + motion_joint_pos: np.ndarray, + robot_joint_pos: np.ndarray, + motion_joint_vel: np.ndarray, + robot_joint_vel: np.ndarray, + error_anchor_pos: np.ndarray, + error_anchor_rot: np.ndarray, + error_anchor_lin_vel: np.ndarray, + error_anchor_ang_vel: np.ndarray, + error_body_pos: np.ndarray, + error_body_rot: np.ndarray, + error_body_lin_vel: np.ndarray, + error_body_ang_vel: np.ndarray, + error_joint_pos: np.ndarray, + error_joint_vel: np.ndarray, +) -> None: + """Write all MotionCommand error metrics for the selected environment rows. + + ``env_ids`` is a cold-path-owned all-row index buffer for normal steps and + the reset row list for partial resets. Keeping one kernel for both paths + avoids a second production mathematical implementation while preserving the + manager's row-scoped reset contract. + """ + num_bodies = motion_body_pos_w.shape[1] + num_joints = motion_joint_pos.shape[1] + for row in prange(env_ids.shape[0]): + env_idx = env_ids[row] + anchor_dx = ( + motion_body_pos_w[env_idx, anchor_body_idx, 0] + - robot_body_pos_w[env_idx, anchor_body_idx, 0] + ) + anchor_dy = ( + motion_body_pos_w[env_idx, anchor_body_idx, 1] + - robot_body_pos_w[env_idx, anchor_body_idx, 1] + ) + anchor_dz = ( + motion_body_pos_w[env_idx, anchor_body_idx, 2] + - robot_body_pos_w[env_idx, anchor_body_idx, 2] + ) + error_anchor_pos[env_idx] = sqrt( + anchor_dx * anchor_dx + anchor_dy * anchor_dy + anchor_dz * anchor_dz + ) + error_anchor_rot[env_idx] = sqrt( + _quat_error_squared( + motion_body_quat_w[env_idx, anchor_body_idx, 0], + motion_body_quat_w[env_idx, anchor_body_idx, 1], + motion_body_quat_w[env_idx, anchor_body_idx, 2], + motion_body_quat_w[env_idx, anchor_body_idx, 3], + robot_body_quat_w[env_idx, anchor_body_idx, 0], + robot_body_quat_w[env_idx, anchor_body_idx, 1], + robot_body_quat_w[env_idx, anchor_body_idx, 2], + robot_body_quat_w[env_idx, anchor_body_idx, 3], + ) + ) + + anchor_lin_sq = 0.0 + anchor_ang_sq = 0.0 + for component in range(3): + lin_delta = ( + motion_body_lin_vel_w[env_idx, anchor_body_idx, component] + - robot_body_lin_vel_w[env_idx, anchor_body_idx, component] + ) + ang_delta = ( + motion_body_ang_vel_w[env_idx, anchor_body_idx, component] + - robot_body_ang_vel_w[env_idx, anchor_body_idx, component] + ) + anchor_lin_sq += lin_delta * lin_delta + anchor_ang_sq += ang_delta * ang_delta + error_anchor_lin_vel[env_idx] = sqrt(anchor_lin_sq) + error_anchor_ang_vel[env_idx] = sqrt(anchor_ang_sq) + + body_pos_sum = 0.0 + body_rot_sum = 0.0 + body_lin_sum = 0.0 + body_ang_sum = 0.0 + for body_idx in range(num_bodies): + pos_sq = 0.0 + lin_sq = 0.0 + ang_sq = 0.0 + for component in range(3): + pos_delta = ( + body_pos_relative_w[env_idx, body_idx, component] + - robot_body_pos_w[env_idx, body_idx, component] + ) + lin_delta = ( + motion_body_lin_vel_w[env_idx, body_idx, component] + - robot_body_lin_vel_w[env_idx, body_idx, component] + ) + ang_delta = ( + motion_body_ang_vel_w[env_idx, body_idx, component] + - robot_body_ang_vel_w[env_idx, body_idx, component] + ) + pos_sq += pos_delta * pos_delta + lin_sq += lin_delta * lin_delta + ang_sq += ang_delta * ang_delta + body_pos_sum += sqrt(pos_sq) + body_lin_sum += sqrt(lin_sq) + body_ang_sum += sqrt(ang_sq) + body_rot_sum += sqrt( + _quat_error_squared( + body_quat_relative_w[env_idx, body_idx, 0], + body_quat_relative_w[env_idx, body_idx, 1], + body_quat_relative_w[env_idx, body_idx, 2], + body_quat_relative_w[env_idx, body_idx, 3], + robot_body_quat_w[env_idx, body_idx, 0], + robot_body_quat_w[env_idx, body_idx, 1], + robot_body_quat_w[env_idx, body_idx, 2], + robot_body_quat_w[env_idx, body_idx, 3], + ) + ) + if num_bodies == 0: + error_body_pos[env_idx] = np.nan + error_body_rot[env_idx] = np.nan + error_body_lin_vel[env_idx] = np.nan + error_body_ang_vel[env_idx] = np.nan + else: + inv_bodies = 1.0 / num_bodies + error_body_pos[env_idx] = body_pos_sum * inv_bodies + error_body_rot[env_idx] = body_rot_sum * inv_bodies + error_body_lin_vel[env_idx] = body_lin_sum * inv_bodies + error_body_ang_vel[env_idx] = body_ang_sum * inv_bodies + + joint_pos_sq = 0.0 + joint_vel_sq = 0.0 + for joint_idx in range(num_joints): + pos_delta = motion_joint_pos[env_idx, joint_idx] - robot_joint_pos[env_idx, joint_idx] + vel_delta = motion_joint_vel[env_idx, joint_idx] - robot_joint_vel[env_idx, joint_idx] + joint_pos_sq += pos_delta * pos_delta + joint_vel_sq += vel_delta * vel_delta + error_joint_pos[env_idx] = sqrt(joint_pos_sq) + error_joint_vel[env_idx] = sqrt(joint_vel_sq) + + @njit(cache=True, nogil=True, parallel=True) def termination_anchor_pos_kernel( motion_body_pos_w: np.ndarray, @@ -97,16 +266,7 @@ def reward_motion_body_ori_kernel( y2 = actual[env_idx, body_idx, 2] z2 = actual[env_idx, body_idx, 3] - # Relative rotation actual * conjugate(reference), matching - # np_quat_error_magnitude_squared_batched without materializing arrays. - rel_w = abs(w2 * w1 + x2 * x1 + y2 * y1 + z2 * z1) - rel_x = -w2 * x1 + x2 * w1 - y2 * z1 + z2 * y1 - rel_y = -w2 * y1 + x2 * z1 + y2 * w1 - z2 * x1 - rel_z = -w2 * z1 - x2 * y1 + y2 * x1 + z2 * w1 - xyz_norm = np.sqrt(rel_x * rel_x + rel_y * rel_y + rel_z * rel_z) - clipped_w = min(max(rel_w, -1.0), 1.0) - angle = 2.0 * np.arctan2(xyz_norm, clipped_w) - error += angle * angle + error += _quat_error_squared(w1, x1, y1, z1, w2, x2, y2, z2) out[env_idx] = np.exp(error / denominator) @@ -167,4 +327,5 @@ def reward_motion_body_ang_vel_kernel( "reward_motion_body_ori_kernel", "reward_motion_body_pos_kernel", "termination_anchor_pos_kernel", + "update_motion_metrics_kernel", ] diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index 782072e54..bd4053fc4 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -30,6 +30,7 @@ reward_motion_body_ori_kernel, reward_motion_body_pos_kernel, termination_anchor_pos_kernel, + update_motion_metrics_kernel, ) from .motion_loader import MotionData, MotionLoader, MotionSampler from .observations import write_body_ori6_in_anchor_frame, write_body_pos_in_anchor_frame @@ -196,6 +197,8 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): 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) # Env ids of the most recent scoped (reset-path) compute; None after a # per-step compute. Written by `_update_command`, consumed by # `post_compute` to restrict refresh work to the reset rows. @@ -224,6 +227,10 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): 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_motion_kernel_runtime() + self._update_metrics() def _make_motion_loader( self, @@ -517,43 +524,37 @@ def _refresh_relative_state_rows(self, env_ids: np.ndarray) -> None: self.robot_body_ori_b[env_ids] = robot_body_ori_b def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: - # All error metrics are row-wise functions of the motion/robot buffers; - # on the reset path only the reset rows are recomputed since other rows - # are unchanged since the per-step update. - sel: np.ndarray | slice = slice(None) if env_ids is None else env_ids - self.metrics["error_anchor_pos"][sel] = np.linalg.norm( - self.anchor_pos_w[sel] - self.robot_anchor_pos_w[sel], axis=-1 - ) - self.metrics["error_anchor_rot"][sel] = np.sqrt( - np_quat_error_magnitude_squared_batched( - self.anchor_quat_w[sel], self.robot_anchor_quat_w[sel] - ) - ) - self.metrics["error_anchor_lin_vel"][sel] = np.linalg.norm( - self.anchor_lin_vel_w[sel] - self.robot_anchor_lin_vel_w[sel], axis=-1 - ) - self.metrics["error_anchor_ang_vel"][sel] = np.linalg.norm( - self.anchor_ang_vel_w[sel] - self.robot_anchor_ang_vel_w[sel], axis=-1 - ) - self.metrics["error_body_pos"][sel] = np.linalg.norm( - self.body_pos_relative_w[sel] - self.robot_body_pos_w[sel], axis=-1 - ).mean(axis=-1) - self.metrics["error_body_rot"][sel] = np.sqrt( - np_quat_error_magnitude_squared_batched( - self.body_quat_relative_w[sel], self.robot_body_quat_w[sel] - ) - ).mean(axis=-1) - self.metrics["error_body_lin_vel"][sel] = np.linalg.norm( - self.body_lin_vel_w[sel] - self.robot_body_lin_vel_w[sel], axis=-1 - ).mean(axis=-1) - self.metrics["error_body_ang_vel"][sel] = np.linalg.norm( - self.body_ang_vel_w[sel] - self.robot_body_ang_vel_w[sel], axis=-1 - ).mean(axis=-1) - self.metrics["error_joint_pos"][sel] = np.linalg.norm( - self.joint_pos[sel] - self.robot_joint_pos[sel], axis=-1 - ) - self.metrics["error_joint_vel"][sel] = np.linalg.norm( - self.joint_vel[sel] - self.robot_joint_vel[sel], axis=-1 + # All row-wise metrics are written by one Numba kernel. Passing an + # explicit all-row index buffer for normal steps lets the same kernel + # serve partial-reset rows without retaining a NumPy runtime formula. + rows = self._all_env_ids if env_ids is None else env_ids + update_motion_metrics_kernel( + rows, + self.anchor_body_idx, + self._body_pos_w, + self._robot_body_pos_w, + self._motion_data.body_quat_w, + self._robot_body_quat_w, + self._motion_data.body_lin_vel_w, + self._robot_body_lin_vel_w, + self._motion_data.body_ang_vel_w, + self._robot_body_ang_vel_w, + self.body_pos_relative_w, + self.body_quat_relative_w, + self._motion_data.joint_pos, + self.robot_joint_pos, + self._motion_data.joint_vel, + self.robot_joint_vel, + self.metrics["error_anchor_pos"], + self.metrics["error_anchor_rot"], + self.metrics["error_anchor_lin_vel"], + self.metrics["error_anchor_ang_vel"], + self.metrics["error_body_pos"], + self.metrics["error_body_rot"], + self.metrics["error_body_lin_vel"], + self.metrics["error_body_ang_vel"], + self.metrics["error_joint_pos"], + self.metrics["error_joint_vel"], ) # Sampler statistics are global scalars, so every row tracks them. self.metrics["sampling_entropy"].fill(self.sampler.sampling_entropy) diff --git a/tests/tasks/test_motion_term_parity.py b/tests/tasks/test_motion_term_parity.py index 6c56381f0..18c6ca75b 100644 --- a/tests/tasks/test_motion_term_parity.py +++ b/tests/tasks/test_motion_term_parity.py @@ -320,6 +320,99 @@ def test_motion_hot_kernels_compile_parallel_on_term_construction(body_setup) -> assert get_num_threads() == min(8, config.NUMBA_DEFAULT_NUM_THREADS) +def test_motion_metrics_kernel_matches_numpy_and_scopes_rows() -> None: + rng = np.random.default_rng(1701) + num_envs, num_bodies, num_joints = 257, 12, 29 + anchor_body_idx = 4 + motion_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + relative_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + motion_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + robot_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + relative_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + motion_lin = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_lin = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + motion_ang = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_ang = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + motion_joint_pos = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + robot_joint_pos = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + motion_joint_vel = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + robot_joint_vel = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + inputs = ( + motion_pos, + robot_pos, + motion_quat, + robot_quat, + motion_lin, + robot_lin, + motion_ang, + robot_ang, + relative_pos, + relative_quat, + motion_joint_pos, + robot_joint_pos, + motion_joint_vel, + robot_joint_vel, + ) + snapshots = tuple(value.copy() for value in inputs) + expected = ( + np.linalg.norm(motion_pos[:, anchor_body_idx] - robot_pos[:, anchor_body_idx], axis=-1), + np.sqrt( + np_quat_error_magnitude_squared_batched( + motion_quat[:, anchor_body_idx], robot_quat[:, anchor_body_idx] + ) + ), + np.linalg.norm(motion_lin[:, anchor_body_idx] - robot_lin[:, anchor_body_idx], axis=-1), + np.linalg.norm(motion_ang[:, anchor_body_idx] - robot_ang[:, anchor_body_idx], axis=-1), + np.linalg.norm(relative_pos - robot_pos, axis=-1).mean(axis=-1), + np.sqrt(np_quat_error_magnitude_squared_batched(relative_quat, robot_quat)).mean(axis=-1), + np.linalg.norm(motion_lin - robot_lin, axis=-1).mean(axis=-1), + np.linalg.norm(motion_ang - robot_ang, axis=-1).mean(axis=-1), + np.linalg.norm(motion_joint_pos - robot_joint_pos, axis=-1), + np.linalg.norm(motion_joint_vel - robot_joint_vel, axis=-1), + ) + outputs = tuple(np.full(num_envs, -123.0, dtype=np.float32) for _ in expected) + + def run(rows: np.ndarray) -> None: + kernels.update_motion_metrics_kernel( + rows, + anchor_body_idx, + motion_pos, + robot_pos, + motion_quat, + robot_quat, + motion_lin, + robot_lin, + motion_ang, + robot_ang, + relative_pos, + relative_quat, + motion_joint_pos, + robot_joint_pos, + motion_joint_vel, + robot_joint_vel, + *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=2e-6, atol=1e-5) + 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=2e-6, atol=1e-5) + for actual, snapshot in zip(inputs, snapshots, strict=True): + np.testing.assert_array_equal(actual, snapshot) + assert kernels.update_motion_metrics_kernel.targetoptions["nopython"] is True + assert kernels.update_motion_metrics_kernel.targetoptions["nogil"] is True + assert kernels.update_motion_metrics_kernel.targetoptions["parallel"] is True + assert kernels.update_motion_metrics_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)