From 5de05ab9f93d1c6d8ba74d31af60c1f3507dcb1f Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Tue, 25 Aug 2026 23:44:30 +0800 Subject: [PATCH 01/12] perf(motion): move fixed hot terms to numba kernels --- pyproject.toml | 1 + .../tasks/motion_tracking/common/kernels.py | 170 ++++++++++ .../motion_tracking/common/manager_terms.py | 170 ++++++++-- tests/tasks/test_motion_term_parity.py | 314 +++++++++++++++--- uv.lock | 55 +++ 5 files changed, 632 insertions(+), 78 deletions(-) create mode 100644 src/unilab/tasks/motion_tracking/common/kernels.py diff --git a/pyproject.toml b/pyproject.toml index cccfad225..6ef219feb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ license-files = ["LICENSE"] requires-python = ">=3.10,<3.14" dependencies = [ "numpy", + "numba>=0.67", "prettytable>=3.10", "torch==2.9.0 ; sys_platform == 'linux' and platform_machine == 'aarch64'", "torch==2.7.0 ; sys_platform != 'linux' or platform_machine != 'aarch64'", diff --git a/src/unilab/tasks/motion_tracking/common/kernels.py b/src/unilab/tasks/motion_tracking/common/kernels.py new file mode 100644 index 000000000..c046bb052 --- /dev/null +++ b/src/unilab/tasks/motion_tracking/common/kernels.py @@ -0,0 +1,170 @@ +"""Parallel CPU kernels for the fixed motion-tracking hot-term set.""" + +from __future__ import annotations + +import os + +import numpy as np +from numba import config, njit, prange, set_num_threads + +_DEFAULT_MOTION_KERNEL_THREADS = 8 +_runtime_configured = False + +# The OpenMP layer's worker wakeups dominate these short kernels after a long +# backend physics phase. Workqueue has stable sub-millisecond dispatch here. +# Respect an application-provided NUMBA_THREADING_LAYER selection. +if "NUMBA_THREADING_LAYER" not in os.environ: + setattr(config, "THREADING_LAYER", "workqueue") + + +def configure_motion_kernel_runtime() -> None: + """Initialize the task-local Numba worker mask once on the cold path.""" + global _runtime_configured + if _runtime_configured: + return + if "NUMBA_NUM_THREADS" not in os.environ: + available_threads = int(getattr(config, "NUMBA_DEFAULT_NUM_THREADS", os.cpu_count() or 1)) + set_num_threads(min(_DEFAULT_MOTION_KERNEL_THREADS, available_threads)) + _runtime_configured = True + + +@njit(cache=True, nogil=True, parallel=True) +def termination_anchor_pos_kernel( + motion_body_pos_w: np.ndarray, + robot_body_pos_w: np.ndarray, + anchor_body_idx: int, + threshold: float, + out: np.ndarray, +) -> None: + """Write the per-environment anchor-height termination mask.""" + for env_idx in prange(motion_body_pos_w.shape[0]): + error = abs( + motion_body_pos_w[env_idx, anchor_body_idx, 2] + - robot_body_pos_w[env_idx, anchor_body_idx, 2] + ) + out[env_idx] = error > threshold + + +@njit(cache=True, nogil=True, parallel=True) +def reward_motion_body_pos_kernel( + reference: np.ndarray, + actual: np.ndarray, + body_ids: np.ndarray, + std: float, + out: np.ndarray, +) -> None: + """Write the relative body-position exponential reward.""" + num_bodies = body_ids.shape[0] + if num_bodies == 0: + out[:] = np.nan + return + denominator = -(num_bodies * std * std) + for env_idx in prange(reference.shape[0]): + error = reference[env_idx, body_ids[0], 0] * 0 + for body_offset in range(num_bodies): + body_idx = body_ids[body_offset] + dx = reference[env_idx, body_idx, 0] - actual[env_idx, body_idx, 0] + dy = reference[env_idx, body_idx, 1] - actual[env_idx, body_idx, 1] + dz = reference[env_idx, body_idx, 2] - actual[env_idx, body_idx, 2] + error += dx * dx + dy * dy + dz * dz + out[env_idx] = np.exp(error / denominator) + + +@njit(cache=True, nogil=True, parallel=True) +def reward_motion_body_ori_kernel( + reference: np.ndarray, + actual: np.ndarray, + body_ids: np.ndarray, + std: float, + out: np.ndarray, +) -> None: + """Write the relative body-orientation exponential reward.""" + num_bodies = body_ids.shape[0] + if num_bodies == 0: + out[:] = np.nan + return + denominator = -(num_bodies * std * std) + for env_idx in prange(reference.shape[0]): + error = reference[env_idx, body_ids[0], 0] * 0 + for body_offset in range(num_bodies): + body_idx = body_ids[body_offset] + w1 = reference[env_idx, body_idx, 0] + x1 = reference[env_idx, body_idx, 1] + y1 = reference[env_idx, body_idx, 2] + z1 = reference[env_idx, body_idx, 3] + w2 = actual[env_idx, body_idx, 0] + x2 = actual[env_idx, body_idx, 1] + 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 + out[env_idx] = np.exp(error / denominator) + + +@njit(cache=True, nogil=True, parallel=True) +def reward_motion_body_lin_vel_kernel( + reference: np.ndarray, + actual: np.ndarray, + body_ids: np.ndarray, + std: float, + out: np.ndarray, +) -> None: + """Write the global body-linear-velocity exponential reward.""" + num_bodies = body_ids.shape[0] + if num_bodies == 0: + out[:] = np.nan + return + denominator = -(num_bodies * std * std) + for env_idx in prange(reference.shape[0]): + error = reference[env_idx, body_ids[0], 0] * 0 + for body_offset in range(num_bodies): + body_idx = body_ids[body_offset] + dx = reference[env_idx, body_idx, 0] - actual[env_idx, body_idx, 0] + dy = reference[env_idx, body_idx, 1] - actual[env_idx, body_idx, 1] + dz = reference[env_idx, body_idx, 2] - actual[env_idx, body_idx, 2] + error += dx * dx + dy * dy + dz * dz + out[env_idx] = np.exp(error / denominator) + + +@njit(cache=True, nogil=True, parallel=True) +def reward_motion_body_ang_vel_kernel( + reference: np.ndarray, + actual: np.ndarray, + body_ids: np.ndarray, + std: float, + out: np.ndarray, +) -> None: + """Write the global body-angular-velocity exponential reward.""" + num_bodies = body_ids.shape[0] + if num_bodies == 0: + out[:] = np.nan + return + denominator = -(num_bodies * std * std) + for env_idx in prange(reference.shape[0]): + error = reference[env_idx, body_ids[0], 0] * 0 + for body_offset in range(num_bodies): + body_idx = body_ids[body_offset] + dx = reference[env_idx, body_idx, 0] - actual[env_idx, body_idx, 0] + dy = reference[env_idx, body_idx, 1] - actual[env_idx, body_idx, 1] + dz = reference[env_idx, body_idx, 2] - actual[env_idx, body_idx, 2] + error += dx * dx + dy * dy + dz * dz + out[env_idx] = np.exp(error / denominator) + + +__all__ = [ + "configure_motion_kernel_runtime", + "reward_motion_body_ang_vel_kernel", + "reward_motion_body_lin_vel_kernel", + "reward_motion_body_ori_kernel", + "reward_motion_body_pos_kernel", + "termination_anchor_pos_kernel", +] diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index 8e5798fb9..176710e5c 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -23,6 +23,14 @@ np_quat_mul, ) +from .kernels import ( + configure_motion_kernel_runtime, + reward_motion_body_ang_vel_kernel, + reward_motion_body_lin_vel_kernel, + reward_motion_body_ori_kernel, + reward_motion_body_pos_kernel, + termination_anchor_pos_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 @@ -807,7 +815,37 @@ def _validate(self, command_name: str, std: float) -> tuple[MotionCommand, float return _command(self._env, command_name), _positive_std(std, term_name=type(self).__name__) -class motion_relative_body_position_error_exp(_BodyTerm): +class _NumbaBodyTerm(_BodyTerm): + """Shared cold-path setup for the four fixed parallel body reward kernels.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + configure_motion_kernel_runtime() + command = _command(env, self._command_name) + if isinstance(self._body_ids, slice): + body_ids = np.arange(len(command.cfg.body_names), dtype=np.intp) + else: + body_ids = self._body_ids + body_ids.setflags(write=False) + self._kernel_body_ids = body_ids + self._kernel_result = np.empty(self.num_envs, dtype=command.body_pos_relative_w.dtype) + + def _kernel_std(self, scale: float) -> float: + return cast(float, self._kernel_result.dtype.type(scale)) + + +class motion_relative_body_position_error_exp(_NumbaBodyTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + command = _command(env, self._command_name) + reward_motion_body_pos_kernel( + command.body_pos_relative_w, + command.robot_body_pos_w, + self._kernel_body_ids, + self._kernel_std(1.0), + self._kernel_result, + ) + def __call__( self, env: ManagerBasedRlEnv, @@ -817,14 +855,28 @@ def __call__( ) -> np.ndarray: del env, body_names command, scale = self._validate(command_name, std) - error = self._squared_error_3d( - command.body_pos_relative_w[:, self._body_ids], - command.robot_body_pos_w[:, self._body_ids], + reward_motion_body_pos_kernel( + command.body_pos_relative_w, + command.robot_body_pos_w, + self._kernel_body_ids, + self._kernel_std(scale), + self._kernel_result, ) - return self._exp_neg_scaled(error.mean(axis=-1), scale) + return self._kernel_result -class motion_relative_body_orientation_error_exp(_BodyTerm): +class motion_relative_body_orientation_error_exp(_NumbaBodyTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + command = _command(env, self._command_name) + reward_motion_body_ori_kernel( + command.body_quat_relative_w, + command.robot_body_quat_w, + self._kernel_body_ids, + self._kernel_std(1.0), + self._kernel_result, + ) + def __call__( self, env: ManagerBasedRlEnv, @@ -834,14 +886,28 @@ def __call__( ) -> np.ndarray: del env, body_names command, scale = self._validate(command_name, std) - error = np_quat_error_magnitude_squared_batched( - command.body_quat_relative_w[:, self._body_ids], - command.robot_body_quat_w[:, self._body_ids], + reward_motion_body_ori_kernel( + command.body_quat_relative_w, + command.robot_body_quat_w, + self._kernel_body_ids, + self._kernel_std(scale), + self._kernel_result, ) - return self._exp_neg_scaled(error.mean(axis=-1), scale) + return self._kernel_result -class motion_global_body_linear_velocity_error_exp(_BodyTerm): +class motion_global_body_linear_velocity_error_exp(_NumbaBodyTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + command = _command(env, self._command_name) + reward_motion_body_lin_vel_kernel( + command.body_lin_vel_w, + command.robot_body_lin_vel_w, + self._kernel_body_ids, + self._kernel_std(1.0), + self._kernel_result, + ) + def __call__( self, env: ManagerBasedRlEnv, @@ -851,14 +917,28 @@ def __call__( ) -> np.ndarray: del env, body_names command, scale = self._validate(command_name, std) - error = self._squared_error_3d( - command.body_lin_vel_w[:, self._body_ids], - command.robot_body_lin_vel_w[:, self._body_ids], + reward_motion_body_lin_vel_kernel( + command.body_lin_vel_w, + command.robot_body_lin_vel_w, + self._kernel_body_ids, + self._kernel_std(scale), + self._kernel_result, ) - return self._exp_neg_scaled(error.mean(axis=-1), scale) + return self._kernel_result -class motion_global_body_angular_velocity_error_exp(_BodyTerm): +class motion_global_body_angular_velocity_error_exp(_NumbaBodyTerm): + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(cfg, env) + command = _command(env, self._command_name) + reward_motion_body_ang_vel_kernel( + command.body_ang_vel_w, + command.robot_body_ang_vel_w, + self._kernel_body_ids, + self._kernel_std(1.0), + self._kernel_result, + ) + def __call__( self, env: ManagerBasedRlEnv, @@ -868,11 +948,14 @@ def __call__( ) -> np.ndarray: del env, body_names command, scale = self._validate(command_name, std) - error = self._squared_error_3d( - command.body_ang_vel_w[:, self._body_ids], - command.robot_body_ang_vel_w[:, self._body_ids], + reward_motion_body_ang_vel_kernel( + command.body_ang_vel_w, + command.robot_body_ang_vel_w, + self._kernel_body_ids, + self._kernel_std(scale), + self._kernel_result, ) - return self._exp_neg_scaled(error.mean(axis=-1), scale) + return self._kernel_result class motion_relative_body_position_z_error_exp(_BodyTerm): @@ -948,11 +1031,48 @@ def __call__( return np.sum(command.robot_body_pos_w[:, self._body_ids, 2] < threshold, axis=-1) -def bad_anchor_pos_z_only( - env: ManagerBasedRlEnv, command_name: str, threshold: float -) -> np.ndarray: - command = _command(env, command_name) - return np.abs(command.anchor_pos_w[:, 2] - command.robot_anchor_pos_w[:, 2]) > threshold +class bad_anchor_pos_z_only(ManagerTermBase): + """Anchor-height termination backed by a parallel, pre-warmed Numba kernel.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + configure_motion_kernel_runtime() + command_name = cfg.params.get("command_name") + if not isinstance(command_name, str) or not command_name: + raise ValueError(f"{type(self).__name__} requires a non-empty command_name") + self._command_name = command_name + self._result = np.empty(self.num_envs, dtype=np.bool_) + command = _command(env, command_name) + threshold = command.body_pos_w.dtype.type(cfg.params.get("threshold", 0.0)) + termination_anchor_pos_kernel( + command.body_pos_w, + command.robot_body_pos_w, + command.anchor_body_idx, + threshold, + self._result, + ) + + def __call__( + self, + env: ManagerBasedRlEnv, + command_name: str, + threshold: float, + ) -> np.ndarray: + del env + if command_name != self._command_name: + raise ValueError( + f"{type(self).__name__} was bound to '{self._command_name}', got '{command_name}'" + ) + command = _command(self._env, command_name) + threshold_value = command.body_pos_w.dtype.type(threshold) + termination_anchor_pos_kernel( + command.body_pos_w, + command.robot_body_pos_w, + command.anchor_body_idx, + threshold_value, + self._result, + ) + return self._result def bad_anchor_ori( diff --git a/tests/tasks/test_motion_term_parity.py b/tests/tasks/test_motion_term_parity.py index 807803b66..6c56381f0 100644 --- a/tests/tasks/test_motion_term_parity.py +++ b/tests/tasks/test_motion_term_parity.py @@ -1,44 +1,76 @@ -"""Issue #1296: bit-parity tests for the temp-array-eliminated motion tracking -terms. Each optimized term is compared against the naive NumPy expression it -replaced; results must be exactly equal (same op order, only fewer temps). - -``_command`` is monkeypatched to return a stub because the terms type-check -against the real MotionCommand; the stub provides the same buffer attributes. -""" +"""Numerical-contract tests for optimized motion-tracking manager terms.""" from __future__ import annotations +import os from types import SimpleNamespace from typing import Any import numpy as np import pytest +from numba import config, get_num_threads, threading_layer -from unilab.managers import RewardTermCfg +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 def _make_env(command: Any) -> SimpleNamespace: return SimpleNamespace(num_envs=command.num_envs) +def _unit_quat(value: np.ndarray) -> np.ndarray: + value /= np.linalg.norm(value, axis=-1, keepdims=True) + return value + + @pytest.fixture def body_setup(monkeypatch: pytest.MonkeyPatch): rng = np.random.default_rng(42) - num_envs, num_bodies = 8, 12 + num_envs, num_bodies = 257, 12 + anchor_body_idx = 4 + + body_pos_relative_w = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + body_pos_w = body_pos_relative_w.copy() + robot_body_pos_w = body_pos_relative_w + 0.1 * rng.standard_normal( + body_pos_relative_w.shape, dtype=np.float32 + ) + body_quat_relative_w = _unit_quat( + rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32) + ) + robot_body_quat_w = _unit_quat( + body_quat_relative_w + + 0.05 * rng.standard_normal(body_quat_relative_w.shape, dtype=np.float32) + ) + body_lin_vel_w = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_body_lin_vel_w = body_lin_vel_w + 0.2 * rng.standard_normal( + body_lin_vel_w.shape, dtype=np.float32 + ) + body_ang_vel_w = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_body_ang_vel_w = body_ang_vel_w + 0.5 * rng.standard_normal( + body_ang_vel_w.shape, dtype=np.float32 + ) + command = SimpleNamespace( num_envs=num_envs, cfg=SimpleNamespace(body_names=tuple(f"b{i}" for i in range(num_bodies))), - anchor_pos_w=rng.standard_normal((num_envs, 3), dtype=np.float32), - robot_anchor_pos_w=rng.standard_normal((num_envs, 3), dtype=np.float32), + anchor_body_idx=anchor_body_idx, + body_pos_w=body_pos_w, + robot_body_pos_w=robot_body_pos_w, + anchor_pos_w=body_pos_w[:, anchor_body_idx], + robot_anchor_pos_w=robot_body_pos_w[:, anchor_body_idx], joint_pos=rng.standard_normal((num_envs, 29), dtype=np.float32), robot_joint_pos=rng.standard_normal((num_envs, 29), dtype=np.float32), - body_pos_relative_w=rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32), - robot_body_pos_w=rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32), - body_lin_vel_w=rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32), - robot_body_lin_vel_w=rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32), + body_pos_relative_w=body_pos_relative_w, + body_quat_relative_w=body_quat_relative_w, + robot_body_quat_w=robot_body_quat_w, + body_lin_vel_w=body_lin_vel_w, + robot_body_lin_vel_w=robot_body_lin_vel_w, + body_ang_vel_w=body_ang_vel_w, + robot_body_ang_vel_w=robot_body_ang_vel_w, ) - # Snapshot inputs so tests can assert the terms never mutate them. + snapshots = { key: value.copy() for key, value in vars(command).items() if isinstance(value, np.ndarray) } @@ -46,11 +78,38 @@ def body_setup(monkeypatch: pytest.MonkeyPatch): return command, _make_env(command), snapshots +def _reward_cfg(*, body_names: tuple[str, ...] | None = None) -> RewardTermCfg: + params: dict[str, Any] = {"command_name": "motion"} + if body_names is not None: + params["body_names"] = body_names + return RewardTermCfg(func=None, weight=1.0, params=params) + + +def _expected_body_reward( + reference: np.ndarray, + actual: np.ndarray, + body_ids: slice | list[int], + std: float, + *, + orientation: bool, +) -> np.ndarray: + reference = reference[:, body_ids] + actual = actual[:, body_ids] + if orientation: + error = np_quat_error_magnitude_squared_batched(reference, actual) + else: + error = np.square(reference - actual).sum(axis=-1) + return np.exp(-error.mean(axis=-1) / std**2) + + def test_anchor_position_error_exp_bit_parity(body_setup) -> None: command, env, snapshots = body_setup out = mt.motion_global_anchor_position_error_exp(env, "motion", std=0.3) expected = np.exp( - -np.sum(np.square(snapshots["anchor_pos_w"] - snapshots["robot_anchor_pos_w"]), axis=-1) + -np.sum( + np.square(snapshots["anchor_pos_w"] - snapshots["robot_anchor_pos_w"]), + axis=-1, + ) / 0.3**2 ) np.testing.assert_array_equal(out, expected) @@ -66,50 +125,199 @@ def test_joint_position_error_exp_bit_parity(body_setup) -> None: np.testing.assert_array_equal(out, expected) -def test_body_term_error_exp_bit_parity_and_scratch_reuse(body_setup) -> None: +def test_anchor_pos_termination_numba_parity_and_output_reuse(body_setup) -> None: command, env, snapshots = body_setup - cfg = RewardTermCfg(func=None, weight=1.0, params={"command_name": "motion"}) - term = mt.motion_relative_body_position_error_exp(cfg, env) + cfg = TerminationTermCfg( + func=mt.bad_anchor_pos_z_only, + params={"command_name": "motion", "threshold": 0.15}, + ) + term = mt.bad_anchor_pos_z_only(cfg, env) - out = term(env, "motion", std=0.3) - expected = np.exp( - -np.square(snapshots["body_pos_relative_w"] - snapshots["robot_body_pos_w"]) - .sum(axis=-1) - .mean(axis=-1) - / 0.3**2 + out = term(env, "motion", threshold=0.15) + expected = ( + np.abs( + snapshots["body_pos_w"][:, command.anchor_body_idx, 2] + - snapshots["robot_anchor_pos_w"][:, 2] + ) + > 0.15 ) np.testing.assert_array_equal(out, expected) - # Repeat calls reuse scratch and stay correct after the buffers change. - rng = np.random.default_rng(7) - command.body_pos_relative_w[:] = rng.standard_normal( - command.body_pos_relative_w.shape, dtype=np.float32 - ) - out2 = term(env, "motion", std=0.3) - expected2 = np.exp( - -np.square(command.body_pos_relative_w - snapshots["robot_body_pos_w"]) - .sum(axis=-1) - .mean(axis=-1) - / 0.3**2 + out2 = term(env, "motion", threshold=0.3) + assert out2 is out + expected2 = ( + np.abs( + snapshots["body_pos_w"][:, command.anchor_body_idx, 2] + - snapshots["robot_anchor_pos_w"][:, 2] + ) + > 0.3 ) np.testing.assert_array_equal(out2, expected2) + np.testing.assert_array_equal(command.body_pos_w, snapshots["body_pos_w"]) + - # Body-subset selection changes the scratch shape but stays bit-identical. - cfg_sub = RewardTermCfg( - func=None, - weight=1.0, - params={"command_name": "motion", "body_names": ("b0", "b3", "b11")}, - ) - term_sub = mt.motion_global_body_linear_velocity_error_exp(cfg_sub, env) - out_sub = term_sub(env, "motion", std=1.0, body_names=("b0", "b3", "b11")) - ids = [0, 3, 11] - expected_sub = np.exp( - -np.square(snapshots["body_lin_vel_w"][:, ids] - snapshots["robot_body_lin_vel_w"][:, ids]) - .sum(axis=-1) - .mean(axis=-1) - / 1.0**2 - ) - np.testing.assert_array_equal(out_sub, expected_sub) +@pytest.mark.parametrize( + ("term_type", "reference_name", "actual_name", "std", "orientation"), + [ + ( + mt.motion_relative_body_position_error_exp, + "body_pos_relative_w", + "robot_body_pos_w", + 0.3, + False, + ), + ( + mt.motion_relative_body_orientation_error_exp, + "body_quat_relative_w", + "robot_body_quat_w", + 0.4, + True, + ), + ( + mt.motion_global_body_linear_velocity_error_exp, + "body_lin_vel_w", + "robot_body_lin_vel_w", + 1.0, + False, + ), + ( + mt.motion_global_body_angular_velocity_error_exp, + "body_ang_vel_w", + "robot_body_ang_vel_w", + 3.14, + False, + ), + ], +) +def test_numba_body_rewards_match_numpy_and_reuse_output( + body_setup, + term_type, + reference_name: str, + actual_name: str, + std: float, + orientation: bool, +) -> None: + command, env, snapshots = body_setup + term = term_type(_reward_cfg(), env) + + out = term(env, "motion", std=std) + expected = _expected_body_reward( + snapshots[reference_name], + snapshots[actual_name], + slice(None), + std, + orientation=orientation, + ) + np.testing.assert_allclose(out, expected, rtol=2e-6, atol=2e-7) + assert out.dtype == snapshots[reference_name].dtype + first_result = out.copy() + + second_std = std * 1.5 + out2 = term(env, "motion", std=second_std) + assert out2 is out + expected2 = _expected_body_reward( + snapshots[reference_name], + snapshots[actual_name], + slice(None), + second_std, + orientation=orientation, + ) + np.testing.assert_allclose(out2, expected2, rtol=2e-6, atol=2e-7) + assert np.any(first_result != out2) + np.testing.assert_array_equal(getattr(command, reference_name), snapshots[reference_name]) + np.testing.assert_array_equal(getattr(command, actual_name), snapshots[actual_name]) + + +@pytest.mark.parametrize( + ("term_type", "reference_name", "actual_name", "std", "orientation"), + [ + ( + mt.motion_relative_body_position_error_exp, + "body_pos_relative_w", + "robot_body_pos_w", + 0.3, + False, + ), + ( + mt.motion_relative_body_orientation_error_exp, + "body_quat_relative_w", + "robot_body_quat_w", + 0.4, + True, + ), + ( + mt.motion_global_body_linear_velocity_error_exp, + "body_lin_vel_w", + "robot_body_lin_vel_w", + 1.0, + False, + ), + ( + mt.motion_global_body_angular_velocity_error_exp, + "body_ang_vel_w", + "robot_body_ang_vel_w", + 3.14, + False, + ), + ], +) +def test_numba_body_rewards_preserve_body_subset_contract( + body_setup, + term_type, + reference_name: str, + actual_name: str, + std: float, + orientation: bool, +) -> None: + command, env, snapshots = body_setup + body_names = ("b0", "b3", "b11") + term = term_type(_reward_cfg(body_names=body_names), env) + + out = term(env, "motion", std=std, body_names=body_names) + expected = _expected_body_reward( + snapshots[reference_name], + snapshots[actual_name], + [0, 3, 11], + std, + orientation=orientation, + ) + np.testing.assert_allclose(out, expected, rtol=2e-6, atol=2e-7) + assert command.cfg.body_names == tuple(f"b{i}" for i in range(12)) + + +def test_motion_hot_kernels_compile_parallel_on_term_construction(body_setup) -> None: + _, env, _ = body_setup + mt.bad_anchor_pos_z_only( + TerminationTermCfg( + func=mt.bad_anchor_pos_z_only, + params={"command_name": "motion", "threshold": 0.15}, + ), + env, + ) + for term_type in ( + mt.motion_relative_body_position_error_exp, + mt.motion_relative_body_orientation_error_exp, + mt.motion_global_body_linear_velocity_error_exp, + mt.motion_global_body_angular_velocity_error_exp, + ): + term_type(_reward_cfg(), env) + + dispatchers = ( + kernels.termination_anchor_pos_kernel, + kernels.reward_motion_body_pos_kernel, + kernels.reward_motion_body_ori_kernel, + kernels.reward_motion_body_lin_vel_kernel, + kernels.reward_motion_body_ang_vel_kernel, + ) + for dispatcher in dispatchers: + assert dispatcher.targetoptions["nopython"] is True + assert dispatcher.targetoptions["nogil"] is True + assert dispatcher.targetoptions["parallel"] is True + assert dispatcher.signatures + if "NUMBA_THREADING_LAYER" not in os.environ: + assert threading_layer() == "workqueue" + if "NUMBA_NUM_THREADS" not in os.environ: + assert get_num_threads() == min(8, config.NUMBA_DEFAULT_NUM_THREADS) def test_joint_pos_limits_bit_parity() -> None: diff --git a/uv.lock b/uv.lock index e089de414..c57389e98 100644 --- a/uv.lock +++ b/uv.lock @@ -1297,6 +1297,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d3/97/68f80ca3ac4924f250cdfa6e20142a803e5e50fca96ef5148c52ee8c10ea/librt-0.8.1-cp313-cp313-win_arm64.whl", hash = "sha256:924817ab3141aca17893386ee13261f1d100d1ef410d70afe4389f2359fea4f0", size = 52495, upload-time = "2026-02-17T16:12:11.633Z" }, ] +[[package]] +name = "llvmlite" +version = "0.49.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/27/72ae94ea5c8f7349ec1c229d4cd058feb799cbd0833ad6d1b47c919b37b7/llvmlite-0.49.0.tar.gz", hash = "sha256:00f16db782f4a13c78c5804aedc434e46794a77e89999a168f9401106270e50a", size = 194467, upload-time = "2026-08-11T16:26:00.489Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/0d/daceb212c44cad1115b2d05dd55beafe23ff06627344adb4ded0c661bb1a/llvmlite-0.49.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:ee81e96c15a6f870918f1eb60c913551c16aa23defb4f5f1acfa660d6a0aaac2", size = 40479229, upload-time = "2026-08-11T16:22:56.104Z" }, + { url = "https://files.pythonhosted.org/packages/72/2c/eb42378b4f3afc71f9fe172d01f30135dc1d54c7fd95cf76d5445d6f7809/llvmlite-0.49.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:854941c2267fd4fc5b2ce02b8af8ecdffa79fb7784591d3a89370322039ea09f", size = 59890659, upload-time = "2026-08-11T16:23:03.359Z" }, + { url = "https://files.pythonhosted.org/packages/4f/dc/fe880ac1eb93c09b6c9a0539ad18c98778386978a0e20a13a55788044ad2/llvmlite-0.49.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da7b64474ac15ca595efa2644d5c6836638ccf70709fad3aba3fc56a55966928", size = 58344482, upload-time = "2026-08-11T16:23:12.122Z" }, + { url = "https://files.pythonhosted.org/packages/59/f6/5c18be29145cfca1d9e859e55a3c586a8c5a821825017b04c7999cd166c9/llvmlite-0.49.0-cp310-cp310-win_amd64.whl", hash = "sha256:b352c14353330c879e339b8f8d7491d565fe94242697714a24e80bd757202384", size = 41865252, upload-time = "2026-08-11T16:23:20.532Z" }, + { url = "https://files.pythonhosted.org/packages/9f/d0/ab52de2328e97ca96cdf0331a5f774796bddc420a51768f4501193f80cbb/llvmlite-0.49.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:4b0e710880b7cc910392bd6b9f1bbf468fed99b182e4420d51598f36114b3dce", size = 40479230, upload-time = "2026-08-11T16:23:28.744Z" }, + { url = "https://files.pythonhosted.org/packages/1f/80/0989432d12b7c86a6f5f380eb92eca7de779af9b34dedbd311b694d7da8d/llvmlite-0.49.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a8c0fc9d624bdc30a3d2db11eb2fb98f80fb209d20b37604eda516cd9b699cf4", size = 59890659, upload-time = "2026-08-11T16:23:37.346Z" }, + { url = "https://files.pythonhosted.org/packages/58/e9/76859ca36aaa460b6ae0508e01637f0e9bdb9b59faaa4637ade3b94bbcca/llvmlite-0.49.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:20496a5c9fdb8179fb9300e7d19f6782555d98aeeb4a322264aa7fd99f980618", size = 58344482, upload-time = "2026-08-11T16:23:44.199Z" }, + { url = "https://files.pythonhosted.org/packages/7d/49/47cd23e05d52d117b6119871ec299adedc9d8d332a2296964d9b2adc06d9/llvmlite-0.49.0-cp311-cp311-win_amd64.whl", hash = "sha256:6a5b06c1b5fc4ae4c9b169b065f42b719448ef1f873687ef224ef69969b75ec3", size = 41865253, upload-time = "2026-08-11T16:23:50.198Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ae/3f699ebe3590e15e023a6372dd147526fd8ec398aacf9ceb844e854964a8/llvmlite-0.49.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:b541c8fac3450db7574d1f53cf9dff83f285bfed9d69bf81fe71fc2a7d4f97fe", size = 40479231, upload-time = "2026-08-11T16:23:56.773Z" }, + { url = "https://files.pythonhosted.org/packages/be/3c/e97f69c62a2d972066d9a2612ce1f3de313035ac897a5b9f787cad8b55f7/llvmlite-0.49.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6acba646d88abbc87d5c113a3d62c1fbf8b8fee11c6493f516803e30f21ae870", size = 59890658, upload-time = "2026-08-11T16:24:05.451Z" }, + { url = "https://files.pythonhosted.org/packages/69/e6/e942ee08605fc0526ff3854260c384d8315a5830e16c4c2a5aebc14dc9bf/llvmlite-0.49.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4ec8ad805e7515cb8440a690eb3cef4d34acb29eef80b705ec4e1c1ad3c43c68", size = 58344481, upload-time = "2026-08-11T16:24:13.503Z" }, + { url = "https://files.pythonhosted.org/packages/84/49/2a44871cac6b5a2fd4aabd68cfdaf6de9a5c7edb36dee5d47b77bda4b50f/llvmlite-0.49.0-cp312-cp312-win_amd64.whl", hash = "sha256:3a9c9e3af4e214acfefa4f73ebe7bc3fb35854a62b654edb3953f5ae33c08ba3", size = 41865543, upload-time = "2026-08-11T16:24:20.41Z" }, + { url = "https://files.pythonhosted.org/packages/7d/85/0b536a3c59f2636d9dd51dda832b6c1d0ffec37608429dedf128664918f1/llvmlite-0.49.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:039fa4054a06f537fb39248d4472284ca96be311a142ec09e69f95630ab469cc", size = 40479230, upload-time = "2026-08-11T16:24:27.295Z" }, + { url = "https://files.pythonhosted.org/packages/8e/1c/ca8ba47b057b793099784475499771780ec46839f2782f753a7079d23520/llvmlite-0.49.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ddc7aecd4f56397ed6e8f120ec5dcd5a1a8f0e6032ca4af413462792d4dca2e3", size = 59890659, upload-time = "2026-08-11T16:24:35.595Z" }, + { url = "https://files.pythonhosted.org/packages/de/af/9526dfdd33a923f33e29a18b8f9801ee7ee4b7397e88d28192c1024c4a75/llvmlite-0.49.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d3dee64784201b64c13a8df62c48a4f4218858faaa65889866bb29bdc243c038", size = 58344482, upload-time = "2026-08-11T16:24:45.79Z" }, + { url = "https://files.pythonhosted.org/packages/96/7f/9f5afcf6476b228d6b170408f377a0c4f91477fc1fc91f8141088b45bf46/llvmlite-0.49.0-cp313-cp313-win_amd64.whl", hash = "sha256:a1b414dc6b164738ec39dd8987cea73829057b7dd92fc6d91b52838385fc1dd2", size = 41865544, upload-time = "2026-08-11T16:24:53.962Z" }, +] + [[package]] name = "lxml" version = "6.0.4" @@ -1955,6 +1979,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" }, ] +[[package]] +name = "numba" +version = "0.67.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "llvmlite" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7a/90/2544f4e3a61e501d6c9a5418fd4b905323222693d54a02cab0106a0af865/numba-0.67.0.tar.gz", hash = "sha256:cd75aa535b33fa05d9d930b1ae8af9f97a2881e96d72dfb38ec9b78284d9f851", size = 2836515, upload-time = "2026-08-11T23:04:00.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/2e/6e72b3edbb7c7d6b44b2ca9e1b62e91997415d181541ef47fc6957c59bf2/numba-0.67.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:8c0e88acd4341ddf40779db3c0228b9188aca7fcab5f5f3ce9949a1fc71e9a02", size = 2745135, upload-time = "2026-08-11T23:03:08.321Z" }, + { url = "https://files.pythonhosted.org/packages/e5/17/5358f24235ef1a5a80b7e28f3e1baa886c0bcf07dc68557009284e6ba698/numba-0.67.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d6c8e9ba3f9602471e8c6f563ffcce8db8046741f0bafb782a052e41dc6b6861", size = 3821881, upload-time = "2026-08-11T23:03:11.172Z" }, + { url = "https://files.pythonhosted.org/packages/0e/18/2f00694248e32c53812baf3d36a7c656dbdd667c6993087b3da068f74b02/numba-0.67.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:694c81c6560b2b47e5fc1dc39c29175b907adf862d9af0af801453400a022a61", size = 3528397, upload-time = "2026-08-11T23:03:13.107Z" }, + { url = "https://files.pythonhosted.org/packages/7f/39/4175b074929938011bd4b564beb4e0fcffd46252e01f60602b57ffb02b06/numba-0.67.0-cp310-cp310-win_amd64.whl", hash = "sha256:ed333e0af4386294e7f03e550e01411856b6935e717d859225e0a7338c6b6795", size = 2815861, upload-time = "2026-08-11T23:03:15.072Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ed/55ba4e54ee878396de6b18e6533cc4a92fa519e8c82d55cf40f98c0a6831/numba-0.67.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:3fa3d1b27f96f2c0d54513d953d7197886aa1eaa7d2439a0eedc44d993fb181a", size = 2744821, upload-time = "2026-08-11T23:03:17.321Z" }, + { url = "https://files.pythonhosted.org/packages/be/78/3f3c45dbaec3cf02bbb1825731beca50f591227e95143d6bd7a64897641c/numba-0.67.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8c80c847301dc33dc8f84a97a952004023d9a05578ae4512b087176264cc1960", size = 3827182, upload-time = "2026-08-11T23:03:19.684Z" }, + { url = "https://files.pythonhosted.org/packages/a5/24/4e70cb86534283d859c3aea2302da523e41539b98dd6c3c4d0a42af95cda/numba-0.67.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e7a7b0121466f1e9a8a074b0545fe90e16389623abf979b5d7c299dca1294d7e", size = 3532817, upload-time = "2026-08-11T23:03:22.06Z" }, + { url = "https://files.pythonhosted.org/packages/26/4d/23dab7f4233be0fc34f54a169ed85238467cd24d8adf2498e5c12ea19dc7/numba-0.67.0-cp311-cp311-win_amd64.whl", hash = "sha256:cfba1ac34f0363fb1a250a10e97240780d11e05227892f7286b26fbfd0ad58ce", size = 2815700, upload-time = "2026-08-11T23:03:23.812Z" }, + { url = "https://files.pythonhosted.org/packages/0d/58/915cddba90010348ed0444451132fdde9b000bcbaff1582029b5bf115d11/numba-0.67.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:6004d8d5f28d4028687fb2d972d629295b13685943bd2ed5cd8810c3b848e219", size = 2745050, upload-time = "2026-08-11T23:03:25.607Z" }, + { url = "https://files.pythonhosted.org/packages/bb/38/926757caaac18a66f057d7544a63620bf360a07d281c9f7ecadd2aa83963/numba-0.67.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f63d43db06b4756424d6d2484737c902e0ae944a0eec3e8b0b4de2c695b15caa", size = 3884596, upload-time = "2026-08-11T23:03:27.688Z" }, + { url = "https://files.pythonhosted.org/packages/4d/6d/58291dc58da39d98b32db7f044729f6d8d4920cd9622fbab3179b54ff4c4/numba-0.67.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:76d3335aaeffb9dc88309420890e73497a00be08a7530441bc2b58ffe025bfa5", size = 3585290, upload-time = "2026-08-11T23:03:29.684Z" }, + { url = "https://files.pythonhosted.org/packages/6e/63/ab21828b4056afed71f9ecb40f4de26c2c19de731cc001961aca74b79464/numba-0.67.0-cp312-cp312-win_amd64.whl", hash = "sha256:50e2b72406c18cda5dd7431b0082cb85ea94e06c64c33607248fc8bef92cfb81", size = 2815645, upload-time = "2026-08-11T23:03:31.732Z" }, + { url = "https://files.pythonhosted.org/packages/49/dd/bd9fe772f6c84597b76cac229b3f2890f01a2c64fd70e48ceaae10dd65cb/numba-0.67.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:77e1c7173fee57a0d84e006c7e70346689d6cb3e7db503489bae58646b4eff7b", size = 2744872, upload-time = "2026-08-11T23:03:33.649Z" }, + { url = "https://files.pythonhosted.org/packages/a1/1c/c05609739cc41116d36e30cb2b41fb00f126bb52e1b0bac907051ad8a35d/numba-0.67.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9c4953387c77864b596d8296e2cfbdef82b0eea4166ab4864b05d226c51143e0", size = 3892004, upload-time = "2026-08-11T23:03:35.797Z" }, + { url = "https://files.pythonhosted.org/packages/4a/77/a5276ad4178250403e0e2251f3e1f8ac18feac779b0474a8bcb08558490d/numba-0.67.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88f6e0f5cb6c545e158b6ef0496c01b6d6958a7ccc6634a1576a94bbbab29ff2", size = 3591878, upload-time = "2026-08-11T23:03:37.845Z" }, + { url = "https://files.pythonhosted.org/packages/0f/80/d48f0ba7442516ceb5a1585f0c81d3aa531bc96bfcabcd9f8f925768c426/numba-0.67.0-cp313-cp313-win_amd64.whl", hash = "sha256:b68ad5125fe245339cc8dcc036081fc1ea482c5063387b9612a76ccd83dc91cd", size = 2815504, upload-time = "2026-08-11T23:03:39.736Z" }, +] + [[package]] name = "numpy" version = "2.2.6" @@ -3928,6 +3981,7 @@ dependencies = [ { name = "lark" }, { name = "mediapy" }, { name = "ninja", marker = "sys_platform == 'linux'" }, + { name = "numba" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "onnxruntime", version = "1.19.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -3990,6 +4044,7 @@ requires-dist = [ { name = "mujoco-uni-runtime", marker = "extra == 'mujoco'", specifier = "==0.4.0" }, { name = "mujoco-warp", marker = "extra == 'mjwarp'", specifier = "==3.10.0.3" }, { name = "ninja", marker = "sys_platform == 'linux'" }, + { name = "numba", specifier = ">=0.67" }, { name = "numpy" }, { name = "onnxruntime", marker = "python_full_version < '3.11'", specifier = "<1.20" }, { name = "onnxruntime", marker = "python_full_version >= '3.11'", specifier = ">=1.20" }, From 970eacf8e1eda22b2c2f579af840aa535c8c0019 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 26 Aug 2026 00:09:19 +0800 Subject: [PATCH 02/12] perf(body-state): fuse selected body cache copies --- src/unilab/base/backend/body_state.py | 40 ++++++++++ src/unilab/base/backend/mjwarp/backend.py | 23 ++++++ src/unilab/base/backend/motrix/backend.py | 36 +++------ src/unilab/base/backend/motrix/body_state.py | 40 ++++++++++ src/unilab/base/backend/mujoco/backend.py | 16 +++- src/unilab/base/entity.py | 24 +++++- .../motion_tracking/common/manager_terms.py | 24 ++---- tests/base/test_body_state_copy.py | 79 +++++++++++++++++++ tests/base/test_entity_facade.py | 59 ++++++++++++++ tests/base/test_mjwarp_backend.py | 6 ++ 10 files changed, 298 insertions(+), 49 deletions(-) create mode 100644 src/unilab/base/backend/body_state.py create mode 100644 src/unilab/base/backend/motrix/body_state.py create mode 100644 tests/base/test_body_state_copy.py diff --git a/src/unilab/base/backend/body_state.py b/src/unilab/base/backend/body_state.py new file mode 100644 index 000000000..d669e0657 --- /dev/null +++ b/src/unilab/base/backend/body_state.py @@ -0,0 +1,40 @@ +"""Shared host-copy kernel for backend-owned body-state caches.""" + +from __future__ import annotations + +import numpy as np +from numba import njit, prange + + +@njit(cache=True, nogil=True, parallel=True) +def copy_selected_body_state( + source_pos: np.ndarray, + source_quat: np.ndarray, + source_lin_vel: np.ndarray, + source_ang_vel: np.ndarray, + selected_ids: np.ndarray, + out_pos: np.ndarray, + out_quat: np.ndarray, + out_lin_vel: np.ndarray, + out_ang_vel: np.ndarray, +) -> None: + """Copy selected cache columns into four caller-owned state buffers.""" + for env_idx in prange(source_pos.shape[0]): + for output_idx in range(selected_ids.shape[0]): + source_idx = selected_ids[output_idx] + out_pos[env_idx, output_idx, 0] = source_pos[env_idx, source_idx, 0] + out_pos[env_idx, output_idx, 1] = source_pos[env_idx, source_idx, 1] + out_pos[env_idx, output_idx, 2] = source_pos[env_idx, source_idx, 2] + out_quat[env_idx, output_idx, 0] = source_quat[env_idx, source_idx, 0] + out_quat[env_idx, output_idx, 1] = source_quat[env_idx, source_idx, 1] + out_quat[env_idx, output_idx, 2] = source_quat[env_idx, source_idx, 2] + out_quat[env_idx, output_idx, 3] = source_quat[env_idx, source_idx, 3] + out_lin_vel[env_idx, output_idx, 0] = source_lin_vel[env_idx, source_idx, 0] + out_lin_vel[env_idx, output_idx, 1] = source_lin_vel[env_idx, source_idx, 1] + out_lin_vel[env_idx, output_idx, 2] = source_lin_vel[env_idx, source_idx, 2] + out_ang_vel[env_idx, output_idx, 0] = source_ang_vel[env_idx, source_idx, 0] + out_ang_vel[env_idx, output_idx, 1] = source_ang_vel[env_idx, source_idx, 1] + out_ang_vel[env_idx, output_idx, 2] = source_ang_vel[env_idx, source_idx, 2] + + +__all__ = ["copy_selected_body_state"] diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index 2effeb0f2..ac79cacc0 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -34,6 +34,7 @@ ) from unilab.utils.rotation import np_quat_apply_inverse_batched +from ..body_state import copy_selected_body_state from .dependencies import load_mjwarp_dependencies from .materialization import materialize_mjwarp_scene from .playback import run_mjwarp_playback, validate_mjwarp_visual_model @@ -1197,6 +1198,28 @@ def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: mapped = self._mapped_tracked_ids("world-frame body angular velocities", body_ids) return self._tracked_angvel_w_all[:, mapped, :] + def copy_body_state_w( + self, + body_ids: np.ndarray, + out_pos: np.ndarray, + out_quat: np.ndarray, + out_lin_vel: np.ndarray, + out_ang_vel: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + mapped = self._mapped_tracked_ids("world-frame body state", body_ids) + copy_selected_body_state( + self._tracked_pos_w_all, + self._tracked_quat_w_all, + self._tracked_linvel_w_all, + self._tracked_angvel_w_all, + mapped, + out_pos, + out_quat, + out_lin_vel, + out_ang_vel, + ) + return out_pos, out_quat, out_lin_vel, out_ang_vel + def get_body_pos_b(self, body_ids: np.ndarray) -> np.ndarray: del body_ids self._unsupported_body_kinematics("base-frame body positions") diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index b98b06e81..9f9e22d13 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -53,6 +53,7 @@ resolve_system_camera_view, tracking_camera_lookat, ) +from .body_state import copy_selected_motrix_body_state from .playback import run_motrix_playback logger = logging.getLogger(__name__) @@ -328,7 +329,6 @@ def __init__( self._render_offsets_np: np.ndarray | None = None self._render_tracking_camera: MotrixTrackingCamera | None = None self.backend_type = "motrix" - self._link_velocity_cache: np.ndarray | None = None # Pre-cache link objects to avoid repeated get_link() lookups. self._link_cache: dict[int, "mtx.Link"] = {} @@ -1135,31 +1135,15 @@ def copy_body_state_w( out_ang_vel: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ids = self._as_body_ids(body_ids) - poses_w = self._get_link_poses_w(ids) - out_pos[..., 0] = poses_w[..., 0] - out_pos[..., 1] = poses_w[..., 1] - out_pos[..., 2] = poses_w[..., 2] - out_quat[..., 0] = poses_w[..., 6] - out_quat[..., 1] = poses_w[..., 3] - out_quat[..., 2] = poses_w[..., 4] - out_quat[..., 3] = poses_w[..., 5] - - link_velocity_cache = self._ensure_link_velocity_cache() - if self._link_velocity_cache is None or self._link_velocity_cache.shape != ( - self._num_envs, - len(ids), - 6, - ): - self._link_velocity_cache = np.empty( - (self._num_envs, len(ids), 6), dtype=self._np_dtype - ) - np.take(link_velocity_cache, ids, axis=1, out=self._link_velocity_cache) - out_lin_vel[..., 0] = self._link_velocity_cache[..., 0] - out_lin_vel[..., 1] = self._link_velocity_cache[..., 1] - out_lin_vel[..., 2] = self._link_velocity_cache[..., 2] - out_ang_vel[..., 0] = self._link_velocity_cache[..., 3] - out_ang_vel[..., 1] = self._link_velocity_cache[..., 4] - out_ang_vel[..., 2] = self._link_velocity_cache[..., 5] + copy_selected_motrix_body_state( + self._link_poses, + self._ensure_link_velocity_cache(), + ids, + out_pos, + out_quat, + out_lin_vel, + out_ang_vel, + ) return out_pos, out_quat, out_lin_vel, out_ang_vel def get_body_vel_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: diff --git a/src/unilab/base/backend/motrix/body_state.py b/src/unilab/base/backend/motrix/body_state.py new file mode 100644 index 000000000..9d2086521 --- /dev/null +++ b/src/unilab/base/backend/motrix/body_state.py @@ -0,0 +1,40 @@ +"""Temporary Motrix packed-cache body-state copy kernel.""" + +from __future__ import annotations + +import numpy as np +from numba import njit, prange + + +# TODO(#1305, #1308): Delete this kernel after MotrixSim exposes the native +# selected-body fused world-state API tracked by the convergence issue. +@njit(cache=True, nogil=True, parallel=True) +def copy_selected_motrix_body_state( + link_poses: np.ndarray, + link_velocities: np.ndarray, + body_ids: np.ndarray, + out_pos: np.ndarray, + out_quat: np.ndarray, + out_lin_vel: np.ndarray, + out_ang_vel: np.ndarray, +) -> None: + """Copy selected xyzw pose and packed velocity cache columns as wxyz state.""" + for env_idx in prange(link_poses.shape[0]): + for output_idx in range(body_ids.shape[0]): + body_idx = body_ids[output_idx] + out_pos[env_idx, output_idx, 0] = link_poses[env_idx, body_idx, 0] + out_pos[env_idx, output_idx, 1] = link_poses[env_idx, body_idx, 1] + out_pos[env_idx, output_idx, 2] = link_poses[env_idx, body_idx, 2] + out_quat[env_idx, output_idx, 0] = link_poses[env_idx, body_idx, 6] + out_quat[env_idx, output_idx, 1] = link_poses[env_idx, body_idx, 3] + out_quat[env_idx, output_idx, 2] = link_poses[env_idx, body_idx, 4] + out_quat[env_idx, output_idx, 3] = link_poses[env_idx, body_idx, 5] + out_lin_vel[env_idx, output_idx, 0] = link_velocities[env_idx, body_idx, 0] + out_lin_vel[env_idx, output_idx, 1] = link_velocities[env_idx, body_idx, 1] + out_lin_vel[env_idx, output_idx, 2] = link_velocities[env_idx, body_idx, 2] + out_ang_vel[env_idx, output_idx, 0] = link_velocities[env_idx, body_idx, 3] + out_ang_vel[env_idx, output_idx, 1] = link_velocities[env_idx, body_idx, 4] + out_ang_vel[env_idx, output_idx, 2] = link_velocities[env_idx, body_idx, 5] + + +__all__ = ["copy_selected_motrix_body_state"] diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index c6551b532..dcadac119 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -43,6 +43,7 @@ SimBackend, normalize_play_render_mode, ) +from ..body_state import copy_selected_body_state from .playback import run_mujoco_playback @@ -1448,10 +1449,17 @@ def copy_body_state_w( out_ang_vel: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: mapped = self._get_mapped_indices(body_ids) - np.take(self._tracked_pos_w_all, mapped, axis=1, out=out_pos) - np.take(self._tracked_quat_w_all, mapped, axis=1, out=out_quat) - np.take(self._tracked_linvel_w_all, mapped, axis=1, out=out_lin_vel) - np.take(self._tracked_angvel_w_all, mapped, axis=1, out=out_ang_vel) + copy_selected_body_state( + self._tracked_pos_w_all, + self._tracked_quat_w_all, + self._tracked_linvel_w_all, + self._tracked_angvel_w_all, + mapped, + out_pos, + out_quat, + out_lin_vel, + out_ang_vel, + ) return out_pos, out_quat, out_lin_vel, out_ang_vel # ------------------------------------------------------------------ # diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 546167847..4deeaeb85 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -8,9 +8,10 @@ from __future__ import annotations import re -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, NoReturn @@ -26,6 +27,10 @@ NamesCfg = tuple[str, ...] | list[str] | None +BodyStateCopyFn = Callable[ + [np.ndarray, np.ndarray, np.ndarray, np.ndarray], + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], +] @dataclass(frozen=True) @@ -1211,6 +1216,23 @@ def set_joint_position_target( ) self.data.write_ctrl(target, env_ids, actuator_ids=actuator_ids) + def bind_body_state_copy( + self, + body_ids: np.ndarray | Sequence[int] | slice | None = None, + ) -> BodyStateCopyFn: + """Bind entity-local body columns to the backend copy contract on the cold path.""" + if self._body_ids is None: + raise self._capability_error( + "body-state copy", + "body_names were not declared in EntityCfg", + ) + local_ids = self._normalize_local_body_ids(body_ids, capability="body-state copy") + if local_ids.size == 0: + raise ValueError(f"Entity '{self.name}' body-state copy selected no bodies") + backend_ids = np.array(self._body_ids[local_ids], copy=True, dtype=np.int32) + backend_ids.setflags(write=False) + return partial(self._backend.copy_body_state_w, backend_ids) + def write_root_state_to_sim( self, root_state: np.ndarray, diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index 176710e5c..782072e54 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -137,6 +137,7 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): ) self._robot_body_ids = np.asarray(body_ids, dtype=np.intp) self._robot_body_ids.setflags(write=False) + self._copy_robot_body_state = self.robot.bind_body_state_copy(self._robot_body_ids) motion_body_ids = self.robot.motion_body_ids[self._robot_body_ids] self.motion = self._make_motion_loader(cfg.motion_file, motion_body_ids) if self.motion.num_joints != len(self.robot.joint_names): @@ -393,24 +394,11 @@ def _refresh_robot_state( if not force and self._robot_cache_step == step: return if env_ids is None: - body_index = self._robot_body_ids - # Single gather straight into the destination buffers (issue #1296); - # the previous src[:][:, body_index] form materialized two copies. - np.take(self.robot.data.body_link_pos_w, body_index, axis=1, out=self._robot_body_pos_w) - np.take( - self.robot.data.body_link_quat_w, body_index, axis=1, out=self._robot_body_quat_w - ) - np.take( - self.robot.data.body_link_lin_vel_w, - body_index, - axis=1, - out=self._robot_body_lin_vel_w, - ) - np.take( - self.robot.data.body_link_ang_vel_w, - body_index, - axis=1, - out=self._robot_body_ang_vel_w, + self._copy_robot_body_state( + self._robot_body_pos_w, + self._robot_body_quat_w, + self._robot_body_lin_vel_w, + self._robot_body_ang_vel_w, ) else: # Partial-reset path (issue #1295): gather only the reset rows from diff --git a/tests/base/test_body_state_copy.py b/tests/base/test_body_state_copy.py new file mode 100644 index 000000000..338e811fa --- /dev/null +++ b/tests/base/test_body_state_copy.py @@ -0,0 +1,79 @@ +"""Focused parity tests for shared and temporary body-state copy kernels.""" + +from __future__ import annotations + +import numpy as np + +from unilab.base.backend.body_state import copy_selected_body_state +from unilab.base.backend.motrix.body_state import copy_selected_motrix_body_state + + +def _outputs(num_envs: int, num_selected: int) -> tuple[np.ndarray, ...]: + shape = (num_envs, num_selected, 3) + return ( + np.empty(shape, dtype=np.float32), + np.empty((num_envs, num_selected, 4), dtype=np.float32), + np.empty(shape, dtype=np.float32), + np.empty(shape, dtype=np.float32), + ) + + +def test_shared_body_state_copy_kernel_preserves_selection_and_outputs() -> None: + rng = np.random.default_rng(7) + num_envs, num_bodies = 257, 9 + pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + quat = rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32) + lin_vel = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + ang_vel = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + selected = np.asarray([7, 1, 5], dtype=np.intp) + out_pos, out_quat, out_lin_vel, out_ang_vel = _outputs(num_envs, len(selected)) + + copy_selected_body_state( + pos, + quat, + lin_vel, + ang_vel, + selected, + out_pos, + out_quat, + out_lin_vel, + out_ang_vel, + ) + + np.testing.assert_array_equal(out_pos, pos[:, selected]) + np.testing.assert_array_equal(out_quat, quat[:, selected]) + np.testing.assert_array_equal(out_lin_vel, lin_vel[:, selected]) + np.testing.assert_array_equal(out_ang_vel, ang_vel[:, selected]) + assert copy_selected_body_state.targetoptions["nopython"] is True + assert copy_selected_body_state.targetoptions["nogil"] is True + assert copy_selected_body_state.targetoptions["parallel"] is True + assert copy_selected_body_state.signatures + + +def test_temporary_motrix_copy_kernel_converts_xyzw_and_packed_velocity() -> None: + rng = np.random.default_rng(11) + num_envs, num_bodies = 257, 9 + poses = rng.standard_normal((num_envs, num_bodies, 7), dtype=np.float32) + velocities = rng.standard_normal((num_envs, num_bodies, 6), dtype=np.float32) + selected = np.asarray([8, 2, 4], dtype=np.int32) + out_pos, out_quat, out_lin_vel, out_ang_vel = _outputs(num_envs, len(selected)) + + copy_selected_motrix_body_state( + poses, + velocities, + selected, + out_pos, + out_quat, + out_lin_vel, + out_ang_vel, + ) + + np.testing.assert_array_equal(out_pos, poses[:, selected, :3]) + np.testing.assert_array_equal(out_quat[..., 0], poses[:, selected, 6]) + np.testing.assert_array_equal(out_quat[..., 1:], poses[:, selected, 3:6]) + np.testing.assert_array_equal(out_lin_vel, velocities[:, selected, :3]) + np.testing.assert_array_equal(out_ang_vel, velocities[:, selected, 3:]) + assert copy_selected_motrix_body_state.targetoptions["nopython"] is True + assert copy_selected_motrix_body_state.targetoptions["nogil"] is True + assert copy_selected_motrix_body_state.targetoptions["parallel"] is True + assert copy_selected_motrix_body_state.signatures diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index cc863dff6..2cfd10e65 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -157,6 +157,21 @@ def get_body_ang_vel_w(self, ids: np.ndarray) -> np.ndarray: self._check("body angular velocity state") return self.body_ang_vel[:, ids] + def copy_body_state_w( + self, + ids: np.ndarray, + out_pos: np.ndarray, + out_quat: np.ndarray, + out_lin_vel: np.ndarray, + out_ang_vel: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + self._check("body state copy") + np.take(self.body_pos, ids, axis=1, out=out_pos) + np.take(self.body_quat, ids, axis=1, out=out_quat) + np.take(self.body_lin_vel, ids, axis=1, out=out_lin_vel) + np.take(self.body_ang_vel, ids, axis=1, out=out_ang_vel) + return out_pos, out_quat, out_lin_vel, out_ang_vel + def get_body_lin_vel_b(self, ids: np.ndarray) -> np.ndarray: self._check("body-frame linear velocity state") return self.body_lin_vel_b[:, ids] @@ -216,6 +231,50 @@ def test_backend_profiles_materialize_identical_local_entity_contract(backend_ty ) +def test_body_state_copy_binding_freezes_local_to_backend_mapping() -> None: + backend, scene = _scene() + robot = scene["robot"] + copy_body_state = robot.bind_body_state_copy(np.asarray([1, 0], dtype=np.int32)) + before = backend.calls.copy() + out_pos = np.empty((backend.num_envs, 2, 3), dtype=np.float32) + out_quat = np.empty((backend.num_envs, 2, 4), dtype=np.float32) + out_lin_vel = np.empty_like(out_pos) + out_ang_vel = np.empty_like(out_pos) + + result = copy_body_state(out_pos, out_quat, out_lin_vel, out_ang_vel) + + assert result == (out_pos, out_quat, out_lin_vel, out_ang_vel) + np.testing.assert_array_equal(out_pos, backend.body_pos[:, [4, 7]]) + np.testing.assert_array_equal(out_quat, backend.body_quat[:, [4, 7]]) + np.testing.assert_array_equal(out_lin_vel, backend.body_lin_vel[:, [4, 7]]) + np.testing.assert_array_equal(out_ang_vel, backend.body_ang_vel[:, [4, 7]]) + assert backend.calls["body state copy"] == before["body state copy"] + 1 + for capability in ( + "body position state", + "body quaternion state", + "body linear velocity state", + "body angular velocity state", + ): + assert backend.calls[capability] == before[capability] + + +@pytest.mark.parametrize( + ("body_ids", "error_type", "message"), + [ + (np.asarray([], dtype=np.int32), ValueError, "selected no bodies"), + ([2], IndexError, "out of range"), + ([0, 0], ValueError, "contain duplicates"), + ([True], TypeError, "1-D integer array"), + ], +) +def test_body_state_copy_binding_rejects_invalid_local_ids( + body_ids: Any, error_type: type[Exception], message: str +) -> None: + _, scene = _scene() + with pytest.raises(error_type, match=message): + scene["robot"].bind_body_state_copy(body_ids) + + def test_state_read_cache_is_scoped_shared_and_explicitly_invalidated() -> None: backend, scene = _scene() robot = scene["robot"] diff --git a/tests/base/test_mjwarp_backend.py b/tests/base/test_mjwarp_backend.py index db16df03b..12d85d3f7 100644 --- a/tests/base/test_mjwarp_backend.py +++ b/tests/base/test_mjwarp_backend.py @@ -304,6 +304,12 @@ def test_body_state_matches_mujoco_backend() -> None: mujoco_backend.get_body_ang_vel_w(body_ids), atol=atol, ) + expected_state = mjwarp_backend.get_body_state_w(body_ids) + outputs = tuple(np.empty_like(value) for value in expected_state) + result = mjwarp_backend.copy_body_state_w(body_ids, *outputs) + assert result == outputs + for actual, expected in zip(outputs, expected_state, strict=True): + np.testing.assert_allclose(actual, expected, atol=atol) np.testing.assert_allclose( mjwarp_backend.get_body_lin_vel_b(body_ids), mujoco_backend.get_body_lin_vel_b(body_ids), From a2e08c61bccffb524d071b9947626eda295a62f1 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 26 Aug 2026 17:05:45 +0800 Subject: [PATCH 03/12] refactor(motrix): drop temporary body-state kernel --- src/unilab/base/backend/motrix/backend.py | 36 ++++++++++---- src/unilab/base/backend/motrix/body_state.py | 40 ---------------- tests/base/test_body_state_copy.py | 32 +------------ tests/base/test_motrix_backend_options.py | 49 ++++++++++++++++++++ 4 files changed, 76 insertions(+), 81 deletions(-) delete mode 100644 src/unilab/base/backend/motrix/body_state.py diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 9f9e22d13..b98b06e81 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -53,7 +53,6 @@ resolve_system_camera_view, tracking_camera_lookat, ) -from .body_state import copy_selected_motrix_body_state from .playback import run_motrix_playback logger = logging.getLogger(__name__) @@ -329,6 +328,7 @@ def __init__( self._render_offsets_np: np.ndarray | None = None self._render_tracking_camera: MotrixTrackingCamera | None = None self.backend_type = "motrix" + self._link_velocity_cache: np.ndarray | None = None # Pre-cache link objects to avoid repeated get_link() lookups. self._link_cache: dict[int, "mtx.Link"] = {} @@ -1135,15 +1135,31 @@ def copy_body_state_w( out_ang_vel: np.ndarray, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ids = self._as_body_ids(body_ids) - copy_selected_motrix_body_state( - self._link_poses, - self._ensure_link_velocity_cache(), - ids, - out_pos, - out_quat, - out_lin_vel, - out_ang_vel, - ) + poses_w = self._get_link_poses_w(ids) + out_pos[..., 0] = poses_w[..., 0] + out_pos[..., 1] = poses_w[..., 1] + out_pos[..., 2] = poses_w[..., 2] + out_quat[..., 0] = poses_w[..., 6] + out_quat[..., 1] = poses_w[..., 3] + out_quat[..., 2] = poses_w[..., 4] + out_quat[..., 3] = poses_w[..., 5] + + link_velocity_cache = self._ensure_link_velocity_cache() + if self._link_velocity_cache is None or self._link_velocity_cache.shape != ( + self._num_envs, + len(ids), + 6, + ): + self._link_velocity_cache = np.empty( + (self._num_envs, len(ids), 6), dtype=self._np_dtype + ) + np.take(link_velocity_cache, ids, axis=1, out=self._link_velocity_cache) + out_lin_vel[..., 0] = self._link_velocity_cache[..., 0] + out_lin_vel[..., 1] = self._link_velocity_cache[..., 1] + out_lin_vel[..., 2] = self._link_velocity_cache[..., 2] + out_ang_vel[..., 0] = self._link_velocity_cache[..., 3] + out_ang_vel[..., 1] = self._link_velocity_cache[..., 4] + out_ang_vel[..., 2] = self._link_velocity_cache[..., 5] return out_pos, out_quat, out_lin_vel, out_ang_vel def get_body_vel_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: diff --git a/src/unilab/base/backend/motrix/body_state.py b/src/unilab/base/backend/motrix/body_state.py deleted file mode 100644 index 9d2086521..000000000 --- a/src/unilab/base/backend/motrix/body_state.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Temporary Motrix packed-cache body-state copy kernel.""" - -from __future__ import annotations - -import numpy as np -from numba import njit, prange - - -# TODO(#1305, #1308): Delete this kernel after MotrixSim exposes the native -# selected-body fused world-state API tracked by the convergence issue. -@njit(cache=True, nogil=True, parallel=True) -def copy_selected_motrix_body_state( - link_poses: np.ndarray, - link_velocities: np.ndarray, - body_ids: np.ndarray, - out_pos: np.ndarray, - out_quat: np.ndarray, - out_lin_vel: np.ndarray, - out_ang_vel: np.ndarray, -) -> None: - """Copy selected xyzw pose and packed velocity cache columns as wxyz state.""" - for env_idx in prange(link_poses.shape[0]): - for output_idx in range(body_ids.shape[0]): - body_idx = body_ids[output_idx] - out_pos[env_idx, output_idx, 0] = link_poses[env_idx, body_idx, 0] - out_pos[env_idx, output_idx, 1] = link_poses[env_idx, body_idx, 1] - out_pos[env_idx, output_idx, 2] = link_poses[env_idx, body_idx, 2] - out_quat[env_idx, output_idx, 0] = link_poses[env_idx, body_idx, 6] - out_quat[env_idx, output_idx, 1] = link_poses[env_idx, body_idx, 3] - out_quat[env_idx, output_idx, 2] = link_poses[env_idx, body_idx, 4] - out_quat[env_idx, output_idx, 3] = link_poses[env_idx, body_idx, 5] - out_lin_vel[env_idx, output_idx, 0] = link_velocities[env_idx, body_idx, 0] - out_lin_vel[env_idx, output_idx, 1] = link_velocities[env_idx, body_idx, 1] - out_lin_vel[env_idx, output_idx, 2] = link_velocities[env_idx, body_idx, 2] - out_ang_vel[env_idx, output_idx, 0] = link_velocities[env_idx, body_idx, 3] - out_ang_vel[env_idx, output_idx, 1] = link_velocities[env_idx, body_idx, 4] - out_ang_vel[env_idx, output_idx, 2] = link_velocities[env_idx, body_idx, 5] - - -__all__ = ["copy_selected_motrix_body_state"] diff --git a/tests/base/test_body_state_copy.py b/tests/base/test_body_state_copy.py index 338e811fa..dcc43bdd2 100644 --- a/tests/base/test_body_state_copy.py +++ b/tests/base/test_body_state_copy.py @@ -1,11 +1,10 @@ -"""Focused parity tests for shared and temporary body-state copy kernels.""" +"""Focused parity tests for the shared body-state copy kernel.""" from __future__ import annotations import numpy as np from unilab.base.backend.body_state import copy_selected_body_state -from unilab.base.backend.motrix.body_state import copy_selected_motrix_body_state def _outputs(num_envs: int, num_selected: int) -> tuple[np.ndarray, ...]: @@ -48,32 +47,3 @@ def test_shared_body_state_copy_kernel_preserves_selection_and_outputs() -> None assert copy_selected_body_state.targetoptions["nogil"] is True assert copy_selected_body_state.targetoptions["parallel"] is True assert copy_selected_body_state.signatures - - -def test_temporary_motrix_copy_kernel_converts_xyzw_and_packed_velocity() -> None: - rng = np.random.default_rng(11) - num_envs, num_bodies = 257, 9 - poses = rng.standard_normal((num_envs, num_bodies, 7), dtype=np.float32) - velocities = rng.standard_normal((num_envs, num_bodies, 6), dtype=np.float32) - selected = np.asarray([8, 2, 4], dtype=np.int32) - out_pos, out_quat, out_lin_vel, out_ang_vel = _outputs(num_envs, len(selected)) - - copy_selected_motrix_body_state( - poses, - velocities, - selected, - out_pos, - out_quat, - out_lin_vel, - out_ang_vel, - ) - - np.testing.assert_array_equal(out_pos, poses[:, selected, :3]) - np.testing.assert_array_equal(out_quat[..., 0], poses[:, selected, 6]) - np.testing.assert_array_equal(out_quat[..., 1:], poses[:, selected, 3:6]) - np.testing.assert_array_equal(out_lin_vel, velocities[:, selected, :3]) - np.testing.assert_array_equal(out_ang_vel, velocities[:, selected, 3:]) - assert copy_selected_motrix_body_state.targetoptions["nopython"] is True - assert copy_selected_motrix_body_state.targetoptions["nogil"] is True - assert copy_selected_motrix_body_state.targetoptions["parallel"] is True - assert copy_selected_motrix_body_state.signatures diff --git a/tests/base/test_motrix_backend_options.py b/tests/base/test_motrix_backend_options.py index d964b96a5..30317f896 100644 --- a/tests/base/test_motrix_backend_options.py +++ b/tests/base/test_motrix_backend_options.py @@ -312,6 +312,55 @@ def test_motrix_backend_uses_cached_batch_link_velocities() -> None: np.testing.assert_allclose(ang_vel, backend._link_velocities[:, body_ids, 3:]) +def test_motrix_copy_body_state_uses_cached_state_and_reuses_scratch() -> None: + import unilab.base.backend.motrix.backend as mod + + num_envs, num_bodies = 2, 4 + poses = np.arange(num_envs * num_bodies * 7, dtype=np.float32).reshape(num_envs, num_bodies, 7) + velocities = np.arange(num_envs * num_bodies * 6, dtype=np.float32).reshape( + num_envs, num_bodies, 6 + ) + + backend = object.__new__(mod.MotrixBackend) + backend._num_envs = num_envs + backend._np_dtype = np.float32 + backend._link_poses = poses + backend._link_velocities = velocities + backend._link_velocity_cache_valid = True + backend._link_velocity_cache = None + body_ids = np.asarray([3, 1], dtype=np.int32) + shape = (num_envs, len(body_ids), 3) + out_pos = np.empty(shape, dtype=np.float32) + out_quat = np.empty((num_envs, len(body_ids), 4), dtype=np.float32) + out_lin_vel = np.empty(shape, dtype=np.float32) + out_ang_vel = np.empty(shape, dtype=np.float32) + + result = backend.copy_body_state_w( + body_ids, + out_pos, + out_quat, + out_lin_vel, + out_ang_vel, + ) + + assert result[0] is out_pos + assert result[1] is out_quat + assert result[2] is out_lin_vel + assert result[3] is out_ang_vel + np.testing.assert_array_equal(out_pos, poses[:, body_ids, :3]) + np.testing.assert_array_equal(out_quat[:, :, 0], poses[:, body_ids, 6]) + np.testing.assert_array_equal(out_quat[:, :, 1:], poses[:, body_ids, 3:6]) + np.testing.assert_array_equal(out_lin_vel, velocities[:, body_ids, :3]) + np.testing.assert_array_equal(out_ang_vel, velocities[:, body_ids, 3:]) + + first_scratch = backend._link_velocity_cache + velocities += 1000.0 + backend.copy_body_state_w(body_ids, out_pos, out_quat, out_lin_vel, out_ang_vel) + + assert backend._link_velocity_cache is first_scratch + np.testing.assert_array_equal(out_lin_vel, velocities[:, body_ids, :3]) + + def test_motrix_backend_get_body_pose_w_slices_cached_poses_once() -> None: import unilab.base.backend.motrix.backend as mod From 8018643edd33259ac0d6ac725dad418fea367bf1 Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:33:29 +0800 Subject: [PATCH 04/12] perf(motion): fuse command metrics in numba (#1317) (#1320) --- .../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) From 1d18060e2aca746be862069d213b5d6ce6a3e21a Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:52:52 +0800 Subject: [PATCH 05/12] perf(motion): fuse relative transforms in numba (#1318) (#1321) --- .../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) From d09d05aa39123ec1001210500e61b6539cb22caf Mon Sep 17 00:00:00 2001 From: YUFEI JIA <59379871+TATP-233@users.noreply.github.com> Date: Wed, 26 Aug 2026 21:05:03 +0800 Subject: [PATCH 06/12] perf(observations): reduce batch pipeline copies (#1319) (#1322) --- src/unilab/managers/_noise/noise_cfg.py | 15 ++- src/unilab/managers/observation_manager.py | 72 +++++++++++-- .../test_observation_buffers_noise.py | 102 ++++++++++++++++++ 3 files changed, 179 insertions(+), 10 deletions(-) diff --git a/src/unilab/managers/_noise/noise_cfg.py b/src/unilab/managers/_noise/noise_cfg.py index d5cf2d8ec..13383b018 100644 --- a/src/unilab/managers/_noise/noise_cfg.py +++ b/src/unilab/managers/_noise/noise_cfg.py @@ -68,14 +68,21 @@ def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) -> n_min = self._as_array(self.n_min, data.dtype) n_max = self._as_array(self.n_max, data.dtype) - # Generate uniform noise in [0, 1) and scale to [n_min, n_max). + # Generate uniform noise in [0, 1) and transform the generated array + # in place. The pre-existing expression allocated one array for each + # multiply, add, and final data operation; keeping the same float32 + # cast and ufunc order preserves bit-level results while returning the + # transformed noise buffer itself. noise = rng.random(data.shape).astype(data.dtype, copy=False) - noise = noise * (n_max - n_min) + n_min + np.multiply(noise, n_max - n_min, out=noise) + np.add(noise, n_min, out=noise) if self.operation == "add": - return data + noise + np.add(data, noise, out=noise) + return noise elif self.operation == "scale": - return data * noise + np.multiply(data, noise, out=noise) + return noise elif self.operation == "abs": return noise else: diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index 7bf1b4dae..8cfa6af60 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -335,7 +335,7 @@ def _row_env_ids(mask: np.ndarray) -> list[int]: ) # Sanitize (applies to both "warn" and "sanitize" policies). - return np.nan_to_num(tensor, nan=0.0, posinf=0.0, neginf=0.0) + return np.nan_to_num(tensor, copy=False, nan=0.0, posinf=0.0, neginf=0.0) def compute( self, @@ -381,6 +381,17 @@ def compute_group( group_term_names = self._group_obs_term_names[group_name] group_obs: dict[str, np.ndarray] = {} obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name], strict=False) + # In the strict default policy a finite result is by far the common + # case. For concatenated groups, scan the assembled output once and + # only inspect individual slices when an error is actually found; this + # retains the per-term diagnostic while removing repeated full scans + # from the hot path. + defer_error_nan_check = ( + self._group_obs_concatenate[group_name] + and not self._group_obs_temporal[group_name] + and group_cfg.nan_check_per_term + and group_cfg.nan_policy == "error" + ) # Reset path (issue #1259 R2): when no term in this group uses delay or # history buffers, everything downstream of the term call is row # independent, so only the reset rows are processed. Term calls and @@ -409,10 +420,28 @@ def compute_group( # NoiseModel.__call__ likewise returns a new array. obs = self._group_obs_class_instances[group_name][term_name](obs) fresh = True - if not row_scoped and not fresh: - # Terms may return backend/command-owned buffers; copy before the - # in-place clip/scale below. Skipped when noise already produced - # a fresh array (issue #1296). + sanitizes_per_term = group_cfg.nan_check_per_term and group_cfg.nan_policy in ( + "warn", + "sanitize", + ) + exposes_term_output = ( + not self._group_obs_concatenate[group_name] + and term_cfg.delay_max_lag == 0 + and term_cfg.history_length == 0 + ) + if ( + not row_scoped + and not fresh + and ( + term_cfg.clip is not None + or term_cfg.scale is not None + or sanitizes_per_term + or exposes_term_output + ) + ): + # Concatenation and temporal buffers already copy their inputs. + # Only take a defensive copy when this pipeline may mutate the + # term or expose it directly to callers. obs = obs.copy() if row_scoped: # Fresh row copy; safe for the in-place clip/scale below. @@ -425,7 +454,11 @@ def compute_group( np.multiply(obs, scale, out=obs) # Check for NaN/Inf before delay/history buffers (per-term checking). - if group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": + if ( + group_cfg.nan_check_per_term + and group_cfg.nan_policy != "disabled" + and not defer_error_nan_check + ): obs = self._check_and_handle_nans( obs, context=f"{group_name}/{term_name}", @@ -474,6 +507,33 @@ def compute_group( result = np.concatenate( list(group_obs.values()), axis=self._group_obs_concatenate_dim[group_name] ) + if defer_error_nan_check: + finite = np.isfinite(result) + if not finite.all(): + axis = self._group_obs_concatenate_dim[group_name] + axis = axis if axis >= 0 else result.ndim + axis + offset = 0 + for term_name, term_dims in zip( + group_term_names, + self._group_obs_term_dim[group_name], + strict=True, + ): + width = int(term_dims[axis - 1]) + selectors = [slice(None)] * result.ndim + selectors[axis] = slice(offset, offset + width) + selector = tuple(selectors) + if not finite[selector].all(): + # Reuse the established diagnostic path so the + # error still names the first offending term and + # reports reset-row IDs when applicable. + self._check_and_handle_nans( + result[selector], + context=f"{group_name}/{term_name}", + policy=group_cfg.nan_policy, + env_ids=env_ids if row_scoped else None, + ) + break + offset += width # Final check for concatenated result (non-per-term checking). if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": result = self._check_and_handle_nans( diff --git a/tests/managers/test_observation_buffers_noise.py b/tests/managers/test_observation_buffers_noise.py index 0688dbfe7..818fea0a0 100644 --- a/tests/managers/test_observation_buffers_noise.py +++ b/tests/managers/test_observation_buffers_noise.py @@ -95,6 +95,28 @@ def test_noise_configs_use_supplied_generator() -> None: np.testing.assert_array_equal(ConstantNoiseCfg(bias=2.0, operation="abs").apply(data), 2.0) +@pytest.mark.parametrize("operation", ["add", "scale", "abs"]) +def test_uniform_noise_inplace_matches_reference_expression(operation: str) -> None: + data = np.arange(24, dtype=np.float32).reshape(8, 3) + n_min = np.asarray([-0.2, -0.1, -0.05], dtype=np.float32) + n_max = np.asarray([0.3, 0.4, 0.5], dtype=np.float32) + cfg = UniformNoiseCfg(n_min=tuple(n_min), n_max=tuple(n_max), operation=operation) + + reference_rng = np.random.default_rng(1702) + unit = reference_rng.random(data.shape).astype(data.dtype, copy=False) + noise = unit * (n_max - n_min) + n_min + if operation == "add": + expected = data + noise + elif operation == "scale": + expected = data * noise + else: + expected = noise + + actual = cfg.apply(data, rng=np.random.default_rng(1702)) + np.testing.assert_array_equal(actual, expected) + np.testing.assert_array_equal(data, np.arange(24, dtype=np.float32).reshape(8, 3)) + + def test_additive_bias_noise_supports_scalar_terms() -> None: from unilab.managers._noise import NoiseModelWithAdditiveBias @@ -143,6 +165,86 @@ def test_observation_groups_pipeline_order_and_history(fake_env: FakeEnv) -> Non assert manager.get_active_iterable_terms(0)[0][0] == "policy-state" +def test_concatenated_result_owns_each_result_and_protects_term_buffers( + fake_env: FakeEnv, +) -> None: + source = fake_env.obs + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "source": ObservationTermCfg(func=lambda env: env.obs), + "constant": ObservationTermCfg( + func=lambda env: np.ones((env.num_envs, 1), dtype=np.float32) + ), + } + ) + }, + fake_env, + ) + + first = manager.compute(update_history=True)["policy"] + assert isinstance(first, np.ndarray) + first_address = first.ctypes.data + expected_first = np.concatenate((source.copy(), np.ones((fake_env.num_envs, 1))), axis=1) + np.testing.assert_array_equal(first, expected_first) + + fake_env.obs += 100.0 + second = manager.compute(update_history=True)["policy"] + assert isinstance(second, np.ndarray) + assert second.ctypes.data != first_address + np.testing.assert_array_equal(first, expected_first) + np.testing.assert_array_equal(second[:, :2], fake_env.obs) + assert not np.shares_memory(second, fake_env.obs) + + +def test_concatenated_nan_sanitize_does_not_mutate_term_owned_input( + fake_env: FakeEnv, +) -> None: + source = fake_env.obs.copy() + source[0, 0] = np.nan + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={"state": ObservationTermCfg(func=lambda env: source)}, + nan_policy="sanitize", + nan_check_per_term=False, + ) + }, + fake_env, + ) + + result = manager.compute(update_history=True)["policy"] + assert isinstance(result, np.ndarray) + assert np.isfinite(result).all() + assert np.isnan(source[0, 0]) + + +def test_concatenated_nan_error_still_identifies_offending_term(fake_env: FakeEnv) -> None: + def invalid(env: FakeEnv) -> np.ndarray: + result = np.ones((env.num_envs, 1), dtype=np.float32) + result[2, 0] = np.nan + return result + + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "finite": ObservationTermCfg(func=lambda env: env.obs), + "invalid": ObservationTermCfg(func=invalid), + } + ) + }, + fake_env, + ) + + with pytest.raises( + ValueError, + match=r"NaN detected.*'policy/invalid'.*environments: \[2\]", + ): + manager.compute(update_history=True) + + def test_observation_noise_model_delay_and_seed_reproducibility() -> None: cfg = { "policy": ObservationGroupCfg( From 512fc823c591488dfcb287a6ec80f1740b76d8d2 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 27 Aug 2026 01:03:51 +0800 Subject: [PATCH 07/12] perf(env): confine DP collector host compute to the per-rank CPU block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-rank off-policy collectors pin their MuJoCo BatchEnvPool workers to a per-rank CPU block via EnvCfg.cpu_ids, but the collector's host-side compute did not follow: Numba's parallel kernels sized their pool from the host CPU count and drifted across rank boundaries, and the OpenBLAS pool spawned at import kept the host-wide mask. NpEnv.__init__ now applies apply_env_cpu_runtime(cfg.cpu_ids) on the cold path: the process is confined to the block (existing threads pinned individually via /proc/self/task, later threads — including Numba's lazily-launched pool — inherit the mask) and Numba's pool is sized to len(cpu_ids) unless NUMBA_NUM_THREADS is set explicitly. cpu_ids=None keeps the single-rank path bit-identical. Backend-agnostic: any env declaring cpu_ids (e.g. motrix once it grows affinity support) gets the same confinement. --- .../zh_CN/2-user_guide/2-algorithms/3-sac.md | 4 +- src/unilab/base/base.py | 9 +- src/unilab/base/cpu_runtime.py | 98 ++++++ src/unilab/base/np_env.py | 7 + tests/base/test_cpu_runtime.py | 282 ++++++++++++++++++ 5 files changed, 396 insertions(+), 4 deletions(-) create mode 100644 src/unilab/base/cpu_runtime.py create mode 100644 tests/base/test_cpu_runtime.py diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md index bebe85295..374f98a4b 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md @@ -76,7 +76,9 @@ rank 子目录或任何日志文件。 `training.log_dir` 保持原样。 collector 的 CPU 亲和按 rank 自动均分（`cpu_count // world_size` 一段），可用 -`training.dp_collector_cpu_ids` 显式指定。 +`training.dp_collector_cpu_ids` 显式指定。该核区经 `EnvCfg.cpu_ids` 生效：除 +MuJoCo worker 线程逐核绑定外，collector 进程本身（含 Numba 并行 kernel 线程池，池 +大小取核区长度）也被限制在同一核区内，避免跨 rank 抢占。 当前限制： diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py index f6c1d385b..a00e53960 100644 --- a/src/unilab/base/base.py +++ b/src/unilab/base/base.py @@ -40,9 +40,12 @@ class EnvCfg: post_step_forward_sensor: bool = False adaptive_chunk_size: bool = True chunk_size: Optional[int] = None - # Explicit worker CPU affinity for the MuJoCo BatchEnvPool (Linux only). - # ``cpu_ids[i]`` pins pool worker thread ``i`` to one CPU; ``None`` keeps - # the default OS scheduling behavior. + # Explicit CPU block owned by this env's process (Linux affinity only). + # ``cpu_ids[i]`` pins MuJoCo BatchEnvPool worker thread ``i`` to one CPU; + # env construction also confines the owning process to the same block and + # sizes Numba's parallel pool to ``len(cpu_ids)`` so host-side post-step + # compute stays inside the rank's partition. ``None`` keeps the default + # OS scheduling behavior. cpu_ids: Optional[list[int]] = None # ``mjwarp`` owns contact/constraint storage independently from MuJoCo. # Keep its capacity knobs explicit in the task owner configuration so a diff --git a/src/unilab/base/cpu_runtime.py b/src/unilab/base/cpu_runtime.py new file mode 100644 index 000000000..79653f913 --- /dev/null +++ b/src/unilab/base/cpu_runtime.py @@ -0,0 +1,98 @@ +"""Process-local CPU confinement for envs that own an explicit CPU block. + +Multi-rank off-policy data-parallel runs partition host CPUs so each rank's +collector owns one contiguous block (``training.dp_collector_cpu_ids``, routed +into ``EnvCfg.cpu_ids``). The MuJoCo BatchEnvPool already pins its physics +workers to that block, but the collector's host-side compute did not follow: +Numba parallel kernels (``unilab.base.backend.body_state`` and the +motion-tracking kernels) size their pool from the host CPU count and leave +placement to the OS, so they drift across rank boundaries and compete with +sibling ranks' pinned physics workers. + +``apply_env_cpu_runtime`` is the generic env-level counterpart, applied once on +the env-construction cold path — before managers, backend materialization, or +the first Numba parallel call exist: + +- ``os.sched_setaffinity`` pins the calling thread to the block, and every + already-running thread (e.g. BLAS pools spawned at ``import numpy``) is + pinned individually via ``/proc/self/task``; threads spawned later + (including Numba's lazily-launched pool) inherit the mask. +- ``numba.set_num_threads(len(cpu_ids))`` sizes Numba's pool to the block + instead of the host, unless the operator pinned ``NUMBA_NUM_THREADS`` + (mirroring the motion-kernel runtime policy). + +The function is backend-agnostic: ``EnvCfg.cpu_ids`` is the single source of +truth, so any backend whose env declares a block gets the same confinement. +""" + +from __future__ import annotations + +import os +import warnings +from collections.abc import Sequence + +_PROC_TASK_DIR = "/proc/self/task" + + +def _confine_existing_threads(ids: set[int]) -> None: + """Pin already-running threads; later threads inherit the caller's mask. + + Native pools spawned before env construction (OpenBLAS spawns its workers + at ``import numpy``) would otherwise keep the host-wide mask and compete + with sibling ranks' pinned CPUs. Threads may exit between listing and + pinning; those races are ignored. + """ + if not os.path.isdir(_PROC_TASK_DIR): + return + for entry in os.listdir(_PROC_TASK_DIR): + try: + os.sched_setaffinity(int(entry), ids) + except OSError: + continue + + +def apply_env_cpu_runtime(cpu_ids: Sequence[int] | None) -> None: + """Confine this process's host-side compute to the env-owned CPU block. + + Cold path only: call from env construction, before the backend pool, + managers, or any Numba parallel kernel exist. ``None`` (the single-rank + default) is a no-op so the default path stays bit-identical. + + Structural validation (non-empty, unique, non-negative ints) is owned by + ``EnvCfg.validate``; this function fails closed on CPU ids that are not + available to the process. + """ + if cpu_ids is None: + return + ids = {int(cpu_id) for cpu_id in cpu_ids} + + if hasattr(os, "sched_setaffinity") and hasattr(os, "sched_getaffinity"): + available = set(os.sched_getaffinity(0)) + missing = sorted(ids - available) + if missing: + raise ValueError( + f"EnvCfg.cpu_ids entries {missing} are not available to this process " + f"(sched_getaffinity={sorted(available)})" + ) + os.sched_setaffinity(0, ids) + _confine_existing_threads(ids) + else: + warnings.warn( + "EnvCfg.cpu_ids process confinement requires os.sched_setaffinity " + "(Linux); only the Numba thread cap is applied", + stacklevel=2, + ) + + if "NUMBA_NUM_THREADS" in os.environ: + return + from numba import set_num_threads + + try: + set_num_threads(len(ids)) + except ValueError as exc: + # len(ids) <= NUMBA_NUM_THREADS holds whenever the pool default came + # from this host's CPU count; warn instead of failing env construction. + warnings.warn( + f"EnvCfg.cpu_ids Numba thread cap to {len(ids)} rejected: {exc}", + stacklevel=2, + ) diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index b90181605..7d21a7f78 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -14,6 +14,7 @@ from unilab.base.backend import SimBackend from unilab.base.backend.base import BackendPlayRenderPlan from unilab.base.base import ABEnv, EnvCfg, EnvPlayCapabilities +from unilab.base.cpu_runtime import apply_env_cpu_runtime from unilab.base.scene import SceneCfg from unilab.dr import DomainRandomizationManager, DomainRandomizationProvider from unilab.dtype_config import get_global_dtype @@ -111,6 +112,12 @@ class NpEnv(ABEnv): """Backend-agnostic numpy environment base class.""" def __init__(self, cfg: EnvCfg, backend: SimBackend, num_envs: int): + # Cold-path process confinement for envs that own an explicit CPU + # block (multi-rank DP collectors): keeps host-side NumPy/Numba compute + # inside the same CPUs the backend pool workers are pinned to. Runs + # before managers/materialization so Numba's lazily-launched pool + # inherits the confined mask. + apply_env_cpu_runtime(cfg.cpu_ids) self._cfg = cfg self._backend: SimBackend = backend self._num_envs = num_envs diff --git a/tests/base/test_cpu_runtime.py b/tests/base/test_cpu_runtime.py new file mode 100644 index 000000000..06cf65866 --- /dev/null +++ b/tests/base/test_cpu_runtime.py @@ -0,0 +1,282 @@ +"""Tests for env-owned CPU block process confinement (``apply_env_cpu_runtime``). + +Multi-rank off-policy collectors pin their MuJoCo pool workers to a per-rank +CPU block via ``EnvCfg.cpu_ids``; ``NpEnv.__init__`` additionally confines the +owning process to the same block and sizes Numba's parallel pool to it, so +host-side kernels cannot drift onto sibling ranks' CPUs. Unit tests mock the +OS/Numba seams; one subprocess test validates the real placement contract. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from dataclasses import dataclass +from unittest.mock import MagicMock + +import gymnasium as gym +import numba +import numpy as np +import pytest + +import unilab.base.cpu_runtime as cpu_runtime +from unilab.base.base import EnvCfg +from unilab.base.cpu_runtime import apply_env_cpu_runtime +from unilab.base.np_env import NpEnv, NpEnvState + + +def _record_affinity(monkeypatch: pytest.MonkeyPatch, available: set[int]) -> list[tuple]: + calls: list[tuple] = [] + monkeypatch.setattr(os, "sched_getaffinity", lambda _pid: set(available)) + monkeypatch.setattr(os, "sched_setaffinity", lambda pid, ids: calls.append((pid, set(ids)))) + return calls + + +def _record_numba(monkeypatch: pytest.MonkeyPatch) -> list[int]: + calls: list[int] = [] + monkeypatch.setattr(numba, "set_num_threads", lambda n: calls.append(int(n))) + return calls + + +def _record_confine(monkeypatch: pytest.MonkeyPatch) -> list[set[int]]: + calls: list[set[int]] = [] + monkeypatch.setattr( + cpu_runtime, "_confine_existing_threads", lambda ids: calls.append(set(ids)) + ) + return calls + + +def test_none_is_noop(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("NUMBA_NUM_THREADS", raising=False) + affinity_calls = _record_affinity(monkeypatch, {0, 1, 2, 3}) + confine_calls = _record_confine(monkeypatch) + numba_calls = _record_numba(monkeypatch) + + apply_env_cpu_runtime(None) + + assert affinity_calls == [] + assert confine_calls == [] + assert numba_calls == [] + + +def test_applies_affinity_and_numba_cap(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("NUMBA_NUM_THREADS", raising=False) + affinity_calls = _record_affinity(monkeypatch, {0, 1, 2, 3}) + confine_calls = _record_confine(monkeypatch) + numba_calls = _record_numba(monkeypatch) + + apply_env_cpu_runtime([1, 2]) + + assert affinity_calls == [(0, {1, 2})] + assert confine_calls == [{1, 2}] + assert numba_calls == [2] + + +def test_respects_explicit_numba_num_threads(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("NUMBA_NUM_THREADS", "4") + affinity_calls = _record_affinity(monkeypatch, {0, 1, 2, 3}) + confine_calls = _record_confine(monkeypatch) + numba_calls = _record_numba(monkeypatch) + + apply_env_cpu_runtime([1, 2]) + + assert affinity_calls == [(0, {1, 2})] + assert confine_calls == [{1, 2}] + assert numba_calls == [] + + +def test_unavailable_cpu_ids_fail_closed(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("NUMBA_NUM_THREADS", raising=False) + affinity_calls = _record_affinity(monkeypatch, {0, 1}) + confine_calls = _record_confine(monkeypatch) + numba_calls = _record_numba(monkeypatch) + + with pytest.raises(ValueError, match="not available"): + apply_env_cpu_runtime([1, 2]) + + assert affinity_calls == [] + assert confine_calls == [] + assert numba_calls == [] + + +def test_platform_without_affinity_warns_and_caps_numba(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("NUMBA_NUM_THREADS", raising=False) + monkeypatch.delattr(os, "sched_setaffinity") + monkeypatch.delattr(os, "sched_getaffinity") + confine_calls = _record_confine(monkeypatch) + numba_calls = _record_numba(monkeypatch) + + with pytest.warns(UserWarning, match="sched_setaffinity"): + apply_env_cpu_runtime([0, 1]) + + assert confine_calls == [] + assert numba_calls == [2] + + +def test_confine_existing_threads_pins_tasks_and_skips_failures( + monkeypatch: pytest.MonkeyPatch, +): + calls: list[tuple] = [] + + def fake_setaffinity(pid, ids): + if pid == 456: + raise ProcessLookupError + calls.append((pid, set(ids))) + + monkeypatch.setattr(os, "sched_setaffinity", fake_setaffinity) + monkeypatch.setattr(os.path, "isdir", lambda path: path == cpu_runtime._PROC_TASK_DIR) + monkeypatch.setattr(os, "listdir", lambda path: ["123", "456", "789"]) + + cpu_runtime._confine_existing_threads({1, 2}) + + assert calls == [(123, {1, 2}), (789, {1, 2})] + + +def test_confine_existing_threads_without_proc_is_noop(monkeypatch: pytest.MonkeyPatch): + calls: list[tuple] = [] + monkeypatch.setattr(os, "sched_setaffinity", lambda pid, ids: calls.append((pid, set(ids)))) + monkeypatch.setattr(os.path, "isdir", lambda path: False) + + cpu_runtime._confine_existing_threads({0}) + + assert calls == [] + + +# --------------------------------------------------------------------------- +# NpEnv wiring +# --------------------------------------------------------------------------- + + +@dataclass +class _StubCfg(EnvCfg): + max_episode_seconds: float | None = 1.0 + + +class _StubNpEnv(NpEnv): + def __init__(self, cfg: EnvCfg): + backend = MagicMock() + backend.get_scene_model_file.return_value = None + super().__init__(cfg, backend, 1) + + @property + def obs_groups_spec(self) -> dict[str, int]: + return {"obs": 1} + + @property + def action_space(self) -> gym.Space: + return gym.spaces.Box(low=-1.0, high=1.0, shape=(1,), dtype=np.float32) + + def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: + return actions + + def update_state(self, state: NpEnvState) -> NpEnvState: + return state + + +@pytest.mark.parametrize("cpu_ids", (None, [2, 3])) +def test_np_env_init_applies_env_cpu_runtime(monkeypatch: pytest.MonkeyPatch, cpu_ids): + import unilab.base.np_env as np_env_module + + calls: list[list[int] | None] = [] + monkeypatch.setattr( + np_env_module, + "apply_env_cpu_runtime", + lambda value: calls.append(None if value is None else list(value)), + ) + + _StubNpEnv(EnvCfg(cpu_ids=cpu_ids)) + + assert calls == [cpu_ids] + + +# --------------------------------------------------------------------------- +# Real placement contract in a fresh process +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (hasattr(os, "sched_setaffinity") and os.path.isdir("/proc/self/task")), + reason="requires Linux sched affinity and /proc", +) +def test_numba_threads_inherit_confined_block_in_fresh_process(): + script = r""" +import json +import os + +# Production collectors import numpy/numba (through the backend modules) +# before env construction, so mirror that ordering here: the OpenBLAS pool +# spawned at `import numpy` predates the env hook and must be confined +# retroactively, while Numba's pool launches after it and inherits the mask. +import numba # noqa: F401 +import numpy as np + +from unilab.base.cpu_runtime import apply_env_cpu_runtime + +block = sorted(os.sched_getaffinity(0))[:2] +apply_env_cpu_runtime(block) + +from numba import get_num_threads, njit, prange + + +@njit(parallel=True) +def _probe(out): + for i in prange(out.shape[0]): + out[i] = i * 2.0 + + +_probe(np.zeros(256)) + + +def _expand(mask): + cpus = set() + for part in mask.split(","): + if "-" in part: + lo, hi = part.split("-", 1) + cpus.update(range(int(lo), int(hi) + 1)) + elif part: + cpus.add(int(part)) + return sorted(cpus) + + +masks = [] +for tid in os.listdir("/proc/self/task"): + with open(f"/proc/self/task/{tid}/status") as fh: + for line in fh: + if line.startswith("Cpus_allowed_list"): + masks.append(_expand(line.split(":", 1)[1].strip())) + +print( + "RESULT:" + + json.dumps( + { + "block": block, + "affinity": sorted(os.sched_getaffinity(0)), + "numba_threads": get_num_threads(), + "masks": masks, + } + ) +) +""" + env = dict(os.environ) + env.pop("NUMBA_NUM_THREADS", None) + proc = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=300, + env=env, + ) + assert proc.returncode == 0, proc.stderr + result_lines = [line for line in proc.stdout.splitlines() if line.startswith("RESULT:")] + assert len(result_lines) == 1, proc.stdout + payload = json.loads(result_lines[0].removeprefix("RESULT:")) + assert payload["affinity"] == payload["block"] + assert payload["numba_threads"] == len(payload["block"]) + # Every thread in the process must stay inside the block: the main thread, + # the OpenBLAS pool spawned at import (confined retroactively), and Numba's + # pool (inherits the confined mask at its lazy launch). + assert payload["masks"] + for mask in payload["masks"]: + assert set(mask) <= set(payload["block"]) From a482a1fc249cc3fb86b12c94eb43625181e2a66e Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 27 Aug 2026 02:33:29 +0800 Subject: [PATCH 08/12] perf(mjwarp): reduce motion tracking reset latency --- src/unilab/base/backend/mjwarp/backend.py | 55 ++++++++++++++++++++--- tests/base/test_mjwarp_backend.py | 14 ++++++ tests/base/test_mjwarp_cuda_graph.py | 23 +++++++++- tests/base/test_mjwarp_host_cache.py | 35 +++++++++++++++ 4 files changed, 121 insertions(+), 6 deletions(-) diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index ac79cacc0..0e9af8404 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -40,8 +40,33 @@ from .playback import run_mjwarp_playback, validate_mjwarp_visual_model _GRAPH_CAPTURE_MIN_DRIVER = (12, 4) -_RESET_SCRATCH_CAPACITY = 128 -_RESET_SCRATCH_MIN_BATCH_SIZE = 8 * _RESET_SCRATCH_CAPACITY +# Reset scratch storage is deliberately bounded. The original 128-world +# allocation covers the smaller G1 owners, while the 8192-world motion +# tracking owner routinely resets about 250 rows per vector step. Scale the +# cold-path allocation with the world count and cap it at 512 so that this +# workload stays on the sparse route without making an unbounded per-task +# allocation. Keeping the minimum at 128 preserves the #1288 route for the +# 1024/2048-world owners. +_RESET_SCRATCH_CAPACITY = 512 +_RESET_SCRATCH_MIN_CAPACITY = 128 +_RESET_SCRATCH_MIN_BATCH_SIZE = 8 * _RESET_SCRATCH_MIN_CAPACITY +_RESET_SCRATCH_WORLD_FRACTION = 16 + + +def _reset_scratch_capacity_for_batch(num_envs: int) -> int: + """Choose a bounded, power-of-two scratch capacity on the cold path. + + A fixed shape is required by the captured reset graphs. The capacity is + rounded down to a power of two so graph/data shapes stay predictable, and + is never allowed below the proven 128-world route or above the 512-world + memory bound. Small batches retain the eager/full-forward fallback rather + than paying for a second ``mujoco_warp.Data`` instance. + """ + if num_envs < _RESET_SCRATCH_MIN_BATCH_SIZE: + return 0 + target = max(_RESET_SCRATCH_MIN_CAPACITY, num_envs // _RESET_SCRATCH_WORLD_FRACTION) + power_of_two = 1 << (target.bit_length() - 1) + return min(_RESET_SCRATCH_CAPACITY, power_of_two) @contextmanager @@ -232,9 +257,7 @@ def __init__( # A bounded secondary Data avoids running reset-time forward over every # production world when only a small row set terminated. It is built # only for batches large enough to amortize the extra graph and copies. - self._reset_scratch_capacity = ( - _RESET_SCRATCH_CAPACITY if self._num_envs >= _RESET_SCRATCH_MIN_BATCH_SIZE else 0 - ) + self._reset_scratch_capacity = _reset_scratch_capacity_for_batch(self._num_envs) self._reset_scratch_data: Any | None = None self._reset_scratch_mask_device: Any | None = None self._reset_scratch_qpos_staging: np.ndarray | None = None @@ -1190,6 +1213,16 @@ def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: mapped = self._mapped_tracked_ids("world-frame body orientations", body_ids) return self._tracked_quat_w_all[:, mapped, :] + def get_body_pose_w_rows( + self, env_ids: np.ndarray, body_ids: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + """Gather world-frame body pose for selected environments only.""" + rows = np.asarray(env_ids, dtype=np.intp) + mapped = self._mapped_tracked_ids("world-frame body poses", body_ids) + return self._tracked_pos_w_all[rows[:, None], mapped], self._tracked_quat_w_all[ + rows[:, None], mapped + ] + def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: mapped = self._mapped_tracked_ids("world-frame body linear velocities", body_ids) return self._tracked_linvel_w_all[:, mapped, :] @@ -1198,6 +1231,18 @@ def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: mapped = self._mapped_tracked_ids("world-frame body angular velocities", body_ids) return self._tracked_angvel_w_all[:, mapped, :] + def get_body_lin_vel_w_rows(self, env_ids: np.ndarray, body_ids: np.ndarray) -> np.ndarray: + """Gather world-frame body linear velocity for selected rows.""" + rows = np.asarray(env_ids, dtype=np.intp) + mapped = self._mapped_tracked_ids("world-frame body linear velocities", body_ids) + return self._tracked_linvel_w_all[rows[:, None], mapped] + + def get_body_ang_vel_w_rows(self, env_ids: np.ndarray, body_ids: np.ndarray) -> np.ndarray: + """Gather world-frame body angular velocity for selected rows.""" + rows = np.asarray(env_ids, dtype=np.intp) + mapped = self._mapped_tracked_ids("world-frame body angular velocities", body_ids) + return self._tracked_angvel_w_all[rows[:, None], mapped] + def copy_body_state_w( self, body_ids: np.ndarray, diff --git a/tests/base/test_mjwarp_backend.py b/tests/base/test_mjwarp_backend.py index 12d85d3f7..0eb67b0da 100644 --- a/tests/base/test_mjwarp_backend.py +++ b/tests/base/test_mjwarp_backend.py @@ -305,6 +305,20 @@ def test_body_state_matches_mujoco_backend() -> None: atol=atol, ) expected_state = mjwarp_backend.get_body_state_w(body_ids) + row_ids = np.asarray([1, 0], dtype=np.int32) + mjwarp_pose_rows = mjwarp_backend.get_body_pose_w_rows(row_ids, body_ids) + for actual, expected in zip(mjwarp_pose_rows, expected_state[:2], strict=True): + np.testing.assert_allclose(actual, expected[row_ids], atol=atol) + np.testing.assert_allclose( + mjwarp_backend.get_body_lin_vel_w_rows(row_ids, body_ids), + expected_state[2][row_ids], + atol=atol, + ) + np.testing.assert_allclose( + mjwarp_backend.get_body_ang_vel_w_rows(row_ids, body_ids), + expected_state[3][row_ids], + atol=atol, + ) outputs = tuple(np.empty_like(value) for value in expected_state) result = mjwarp_backend.copy_body_state_w(body_ids, *outputs) assert result == outputs diff --git a/tests/base/test_mjwarp_cuda_graph.py b/tests/base/test_mjwarp_cuda_graph.py index e01b6439f..c20a05f59 100644 --- a/tests/base/test_mjwarp_cuda_graph.py +++ b/tests/base/test_mjwarp_cuda_graph.py @@ -6,7 +6,11 @@ import pytest -from unilab.base.backend.mjwarp.backend import MjwarpBackend, _cuda_graph_eligibility +from unilab.base.backend.mjwarp.backend import ( + MjwarpBackend, + _cuda_graph_eligibility, + _reset_scratch_capacity_for_batch, +) class _FakeDevice: @@ -174,6 +178,23 @@ def test_cuda_graph_capture_includes_materialized_reset_scratch() -> None: assert backend._can_use_reset_scratch(5) is False +@pytest.mark.parametrize( + ("num_envs", "expected_capacity"), + [ + (512, 0), + (1024, 128), + (2048, 128), + (4096, 256), + (8192, 512), + (16384, 512), + ], +) +def test_reset_scratch_capacity_scales_with_batch_and_stays_bounded( + num_envs: int, expected_capacity: int +) -> None: + assert _reset_scratch_capacity_for_batch(num_envs) == expected_capacity + + def test_ineligible_cuda_graph_warns_and_uses_eager_operations() -> None: warp = _FakeWarp(driver=(12, 3), mempool=True) mujoco_warp = _FakeMujocoWarp() diff --git a/tests/base/test_mjwarp_host_cache.py b/tests/base/test_mjwarp_host_cache.py index 9c8235ba7..557765ff6 100644 --- a/tests/base/test_mjwarp_host_cache.py +++ b/tests/base/test_mjwarp_host_cache.py @@ -144,3 +144,38 @@ def test_host_reset_routes_small_row_sets_through_scratch() -> None: "scratch-refresh", ] assert not any(event[0] in {"main-forward", "main-refresh"} for event in events) + + +def test_row_body_getters_gather_selected_rows_without_full_batch_reads() -> None: + """MJWarp's partial-reset getters must not materialize the full env batch.""" + backend = object.__new__(MjwarpBackend) + backend._body_id_to_tracked_idx = np.asarray([2, 0, 1], dtype=np.intp) + backend._tracked_pos_w_all = np.arange(4 * 3 * 3, dtype=np.float32).reshape(4, 3, 3) + backend._tracked_quat_w_all = np.arange(4 * 3 * 4, dtype=np.float32).reshape(4, 3, 4) + backend._tracked_linvel_w_all = np.arange(4 * 3 * 3, dtype=np.float32).reshape(4, 3, 3) + 1000 + backend._tracked_angvel_w_all = np.arange(4 * 3 * 3, dtype=np.float32).reshape(4, 3, 3) + 2000 + + def fail_full_getter(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("row getter must not call a full-batch getter") + + backend.get_body_pos_w = fail_full_getter # type: ignore[method-assign] + backend.get_body_quat_w = fail_full_getter # type: ignore[method-assign] + backend.get_body_lin_vel_w = fail_full_getter # type: ignore[method-assign] + backend.get_body_ang_vel_w = fail_full_getter # type: ignore[method-assign] + + rows = np.asarray([3, 1, 3], dtype=np.int32) + body_ids = np.asarray([1, 2], dtype=np.int32) + mapped = np.asarray([0, 1], dtype=np.intp) + expected_index = (rows[:, None], mapped) + + pose_pos, pose_quat = backend.get_body_pose_w_rows(rows, body_ids) + np.testing.assert_array_equal(pose_pos, backend._tracked_pos_w_all[expected_index]) + np.testing.assert_array_equal(pose_quat, backend._tracked_quat_w_all[expected_index]) + np.testing.assert_array_equal( + backend.get_body_lin_vel_w_rows(rows, body_ids), + backend._tracked_linvel_w_all[expected_index], + ) + np.testing.assert_array_equal( + backend.get_body_ang_vel_w_rows(rows, body_ids), + backend._tracked_angvel_w_all[expected_index], + ) From bc586c4af30a34ad6cb4f267dcbe8c5cbbb4e0c0 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 27 Aug 2026 04:16:22 +0800 Subject: [PATCH 09/12] benchmark(env): add MuJoCo pool thread-scaling and env-step phase-CPU probes (#1328) Diagnostic probes for the SAC/MuJoCo single-GPU collector CPU under-utilization report: pool thread-count scaling on the G1 scene, and per-phase wall/CPU attribution of a full task env step. New files only; no behavior change. --- .../env/benchmark_env_step_phase_cpu.py | 168 ++++++++++++++++++ .../benchmark_mujoco_pool_thread_scaling.py | 146 +++++++++++++++ 2 files changed, 314 insertions(+) create mode 100644 scripts/benchmark/env/benchmark_env_step_phase_cpu.py create mode 100644 scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py diff --git a/scripts/benchmark/env/benchmark_env_step_phase_cpu.py b/scripts/benchmark/env/benchmark_env_step_phase_cpu.py new file mode 100644 index 000000000..904d91d1e --- /dev/null +++ b/scripts/benchmark/env/benchmark_env_step_phase_cpu.py @@ -0,0 +1,168 @@ +"""Per-phase wall/CPU attribution for a full task env step (issue #1328). + +Builds a real task env through the same Hydra compose + ``BackendAdapter`` +override path the off-policy collector uses, wraps ``backend.step`` / +``update_state`` / ``_reset_done_envs`` with process-wide CPU-time measurement +(``os.times``), and reports each phase's wall share and the average number of +cores it kept busy. This makes low-parallelism host phases visible next to the +thread-pool physics phase. + +``--cpu-ids 0-31`` additionally injects ``EnvCfg.cpu_ids`` into the env +override (the same key the multi-GPU DP collector path uses), which both pins +the MuJoCo pool workers and confines the process's host-side compute via +``apply_env_cpu_runtime`` — the A/B used in the issue. + +Run: + uv run scripts/benchmark/env/benchmark_env_step_phase_cpu.py + + # pinned A/B: + uv run scripts/benchmark/env/benchmark_env_step_phase_cpu.py --cpu-ids 0-31 + + # tuning: + uv run scripts/benchmark/env/benchmark_env_step_phase_cpu.py \ + --config-group sac --task g1_motion_tracking/mujoco \ + --num-envs 4096 --warmup 20 --iters 150 +""" + +from __future__ import annotations + +import argparse +import os +import time +from collections import defaultdict +from collections.abc import Sequence + +import numpy as np + +REPO_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +) + + +def _cpu_time() -> float: + t = os.times() + return t.user + t.system + + +def _parse_cpu_ids(spec: str) -> list[int]: + ids: list[int] = [] + for part in spec.split(","): + part = part.strip() + if "-" in part: + lo, hi = part.split("-", 1) + ids.extend(range(int(lo), int(hi) + 1)) + elif part: + ids.append(int(part)) + return ids + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--config-group", default="sac", help="conf/ used for compose") + parser.add_argument("--task", default="g1_motion_tracking/mujoco") + parser.add_argument("--num-envs", type=int, default=4096) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=150) + parser.add_argument( + "--cpu-ids", + default=None, + help="Optional env cpu_ids override, e.g. '0-31'; pins the MuJoCo pool " + "and confines host-side compute (sizes the pool to len(cpu_ids))", + ) + args = parser.parse_args(argv) + + import hydra + from omegaconf import OmegaConf + + from unilab.base.config_adapter import BackendAdapter, create_env + from unilab.training import ensure_registries + + ensure_registries() + with hydra.initialize_config_dir( + version_base="1.3", config_dir=os.path.join(REPO_ROOT, "conf", args.config_group) + ): + cfg = hydra.compose( + config_name="config", + overrides=[f"task={args.task}", f"algo.num_envs={args.num_envs}"], + ) + OmegaConf.resolve(cfg) + env_cfg_override = BackendAdapter( + cfg, root_dir=REPO_ROOT, algo_name=str(cfg.algo.algo) + ).build_task_env_cfg_override() + if args.cpu_ids is not None: + env_cfg_override = { + **(env_cfg_override or {}), + "cpu_ids": _parse_cpu_ids(args.cpu_ids), + } + env = create_env(cfg, num_envs=args.num_envs, env_cfg_override=env_cfg_override) + if env.state is None: + env.init_state() + + wall_ms: defaultdict[str, float] = defaultdict(float) + cpu_ms: defaultdict[str, float] = defaultdict(float) + counts: defaultdict[str, int] = defaultdict(int) + + def wrap(name, fn): + def wrapped(*a, **kw): + w0 = time.perf_counter() + c0 = _cpu_time() + out = fn(*a, **kw) + wall_ms[name] += (time.perf_counter() - w0) * 1000.0 + cpu_ms[name] += (_cpu_time() - c0) * 1000.0 + counts[name] += 1 + return out + + return wrapped + + env._backend.step = wrap("backend_step", env._backend.step) + env.update_state = wrap("update_state", env.update_state) + env._reset_done_envs = wrap("reset_done", env._reset_done_envs) + + action_dim = env.action_space.shape[-1] + rng = np.random.default_rng(0) + + def actions(): + return rng.uniform(-0.2, 0.2, size=(args.num_envs, action_dim)).astype(np.float32) + + for _ in range(args.warmup): + env.step(actions()) + wall_ms.clear() + cpu_ms.clear() + counts.clear() + + n_reset = 0 + wall0 = time.perf_counter() + cpu0 = _cpu_time() + for _ in range(args.iters): + state = env.step(actions()) + n_reset += int(np.count_nonzero(state.terminated | state.truncated)) + total_wall = (time.perf_counter() - wall0) * 1000.0 + total_cpu = (_cpu_time() - cpu0) * 1000.0 + + print( + f"pool nthread={env._backend._n_threads} num_envs={args.num_envs} " + f"cpu_ids={'None' if args.cpu_ids is None else args.cpu_ids}" + ) + print(f"iters={args.iters} total_resets={n_reset}") + print(f"{'phase':>16s} {'wall_ms':>9s} {'cpu_ms':>9s} {'cores':>6s} {'wall%':>6s}") + step_wall = total_wall / args.iters + for name in ("backend_step", "update_state", "reset_done"): + w = wall_ms[name] / args.iters + c = cpu_ms[name] / args.iters + print(f"{name:>16s} {w:9.2f} {c:9.2f} {c / w if w else 0:6.2f} {100 * w / step_wall:6.1f}") + other_w = total_wall - sum(wall_ms.values()) + other_c = total_cpu - sum(cpu_ms.values()) + print( + f"{'other(step glue)':>16s} {other_w / args.iters:9.2f} {other_c / args.iters:9.2f} " + f"{(other_c / other_w) if other_w > 0 else 0:6.2f} {100 * other_w / total_wall:6.1f}" + ) + print( + f"{'TOTAL step':>16s} {step_wall:9.2f} {total_cpu / args.iters:9.2f} " + f"{total_cpu / total_wall:6.2f} {100.0:6.1f}" + ) + print(f"steps/s={args.num_envs * args.iters / (total_wall / 1000.0):.0f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py b/scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py new file mode 100644 index 000000000..65ff96343 --- /dev/null +++ b/scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py @@ -0,0 +1,146 @@ +"""MuJoCo BatchEnvPool thread-count scaling probe (issue #1328). + +Steps a raw ``BatchEnvPool`` (no env semantics, no learner) on the G1 flat +scene with several ``nthread`` / ``cpu_ids`` configurations and reports, per +configuration, wall time per ``pool.step`` and the average number of cores the +process kept busy (process CPU time / wall time via ``os.times``). + +Used to separate two effects of the default +``nthread = min(num_envs, 2 * cpu_count)`` pool sizing: + +- thread count vs. pinning (``cpu_ids``): on the reference 16C/32T host the + 32-thread unpinned and pinned rows match, so the 2x-oversubscription loss + comes from the thread count itself; +- the physics scaling ceiling: throughput saturates near the physical core + count (memory-bandwidth bound), so extra threads mostly cost wall time. + +Run: + uv run scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py + + # subset + tuning: + uv run scripts/benchmark/env/benchmark_mujoco_pool_thread_scaling.py \ + --num-envs 4096 --nstep 3 --chunk-size 6 \ + --configs 64:unpinned,32:unpinned,32:pinned,16:pinned +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Sequence + +import numpy as np + +REPO_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +) +DEFAULT_MODEL = os.path.join(REPO_ROOT, "src/unilab/assets/robots/g1/scene_flat.xml") + + +def _cpu_time() -> float: + t = os.times() + return t.user + t.system + + +def build_state(model, nenvs: int) -> np.ndarray: + """Tile the ``stand`` keyframe (or a plain forward) into a full-batch state.""" + import mujoco + + data = mujoco.MjData(model) + key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, "stand") + if key_id >= 0: + mujoco.mj_resetDataKeyframe(model, data, key_id) + mujoco.mj_forward(model, data) + spec = int(mujoco.mjtState.mjSTATE_FULLPHYSICS) + row = np.empty(mujoco.mj_stateSize(model, spec), dtype=np.float64) + mujoco.mj_getState(model, data, row, spec) + return np.tile(row, (nenvs, 1)).copy() + + +def bench_config( + model, + state0: np.ndarray, + *, + nthread: int, + pinned: bool, + nstep: int, + chunk_size: int | None, + warmup: int, + iters: int, +) -> tuple[float, float]: + """Return (wall ms/step, busy cores) for one pool configuration.""" + from mujoco_uni.batch_env import BatchEnvPool + + cpu_ids = list(range(nthread)) if pinned else None + pool = BatchEnvPool(model, nbatch=state0.shape[0], nthread=nthread, cpu_ids=cpu_ids) + nenvs = state0.shape[0] + ctrl = np.zeros((nenvs, nstep, model.nu), dtype=np.float64) + st = state0.copy() + try: + for _ in range(warmup): + st = pool.step(st, nstep=nstep, control=ctrl, chunk_size=chunk_size) + t0 = time.perf_counter() + c0 = _cpu_time() + for _ in range(iters): + st = pool.step(st, nstep=nstep, control=ctrl, chunk_size=chunk_size) + wall_ms = (time.perf_counter() - t0) / iters * 1000.0 + cores = (_cpu_time() - c0) / iters * 1000.0 / wall_ms + return wall_ms, cores + finally: + pool.close() + + +def _parse_configs(spec: str) -> list[tuple[int, bool]]: + out = [] + for item in spec.split(","): + nthread_s, mode = item.strip().split(":") + if mode not in ("pinned", "unpinned"): + raise ValueError(f"unknown config mode {mode!r} in {item!r}") + out.append((int(nthread_s), mode == "pinned")) + return out + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--model", default=DEFAULT_MODEL, help="MuJoCo XML scene path") + parser.add_argument("--num-envs", type=int, default=4096) + parser.add_argument("--nstep", type=int, default=3, help="sim substeps per pool.step") + parser.add_argument("--chunk-size", type=int, default=6) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=30) + parser.add_argument( + "--configs", + default="64:unpinned,32:unpinned,32:pinned,24:pinned,16:pinned,8:pinned", + help="Comma-separated nthread:pinned|unpinned entries", + ) + args = parser.parse_args(argv) + + import mujoco + + model = mujoco.MjModel.from_xml_path(args.model) + state0 = build_state(model, args.num_envs) + print( + f"model={os.path.basename(args.model)} nu={model.nu} nv={model.nv} " + f"nstate={state0.shape[1]} num_envs={args.num_envs} host_cpus={os.cpu_count()} " + f"nstep={args.nstep} chunk_size={args.chunk_size}" + ) + print(f"{'config':>18s} | {'ms/step':>8s} | {'cores':>6s}") + for nthread, pinned in _parse_configs(args.configs): + wall_ms, cores = bench_config( + model, + state0, + nthread=nthread, + pinned=pinned, + nstep=args.nstep, + chunk_size=args.chunk_size, + warmup=args.warmup, + iters=args.iters, + ) + label = f"{nthread}t {'pinned' if pinned else 'unpinned'}" + print(f"{label:>18s} | {wall_ms:8.2f} | {cores:6.1f}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 56032a65a1f3f1738b037d2f5ac08233ddaea13b Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Thu, 27 Aug 2026 14:21:12 +0800 Subject: [PATCH 10/12] fix(base): type-check and test cpu_runtime on non-Linux hosts os.sched_setaffinity/sched_getaffinity are Linux-only, so mypy on darwin rejected the direct attribute access (attr-defined) and the unit tests' monkeypatch.setattr/delattr failed because the attributes do not exist. Resolve the affinity symbols via getattr at call time (identical runtime semantics, still monkeypatchable) and pass raising=False to the test monkeypatch seams so they work whether or not the host exposes them. --- src/unilab/base/cpu_runtime.py | 16 +++++++++++----- tests/base/test_cpu_runtime.py | 18 ++++++++++++------ 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/src/unilab/base/cpu_runtime.py b/src/unilab/base/cpu_runtime.py index 79653f913..4a36aa441 100644 --- a/src/unilab/base/cpu_runtime.py +++ b/src/unilab/base/cpu_runtime.py @@ -42,11 +42,15 @@ def _confine_existing_threads(ids: set[int]) -> None: with sibling ranks' pinned CPUs. Threads may exit between listing and pinning; those races are ignored. """ - if not os.path.isdir(_PROC_TASK_DIR): + # Resolved at call time (not import) so tests can monkeypatch the seam; + # ``getattr`` keeps this checkable on platforms where typeshed hides the + # Linux-only symbol (mypy ``attr-defined`` on darwin). + sched_setaffinity = getattr(os, "sched_setaffinity", None) + if sched_setaffinity is None or not os.path.isdir(_PROC_TASK_DIR): return for entry in os.listdir(_PROC_TASK_DIR): try: - os.sched_setaffinity(int(entry), ids) + sched_setaffinity(int(entry), ids) except OSError: continue @@ -66,15 +70,17 @@ def apply_env_cpu_runtime(cpu_ids: Sequence[int] | None) -> None: return ids = {int(cpu_id) for cpu_id in cpu_ids} - if hasattr(os, "sched_setaffinity") and hasattr(os, "sched_getaffinity"): - available = set(os.sched_getaffinity(0)) + sched_setaffinity = getattr(os, "sched_setaffinity", None) + sched_getaffinity = getattr(os, "sched_getaffinity", None) + if sched_setaffinity is not None and sched_getaffinity is not None: + available = set(sched_getaffinity(0)) missing = sorted(ids - available) if missing: raise ValueError( f"EnvCfg.cpu_ids entries {missing} are not available to this process " f"(sched_getaffinity={sorted(available)})" ) - os.sched_setaffinity(0, ids) + sched_setaffinity(0, ids) _confine_existing_threads(ids) else: warnings.warn( diff --git a/tests/base/test_cpu_runtime.py b/tests/base/test_cpu_runtime.py index 06cf65866..461867562 100644 --- a/tests/base/test_cpu_runtime.py +++ b/tests/base/test_cpu_runtime.py @@ -29,8 +29,12 @@ def _record_affinity(monkeypatch: pytest.MonkeyPatch, available: set[int]) -> list[tuple]: calls: list[tuple] = [] - monkeypatch.setattr(os, "sched_getaffinity", lambda _pid: set(available)) - monkeypatch.setattr(os, "sched_setaffinity", lambda pid, ids: calls.append((pid, set(ids)))) + # raising=False: sched_*affinity is Linux-only, so the attribute may not + # exist on the host running the tests (e.g. macOS dev machines). + monkeypatch.setattr(os, "sched_getaffinity", lambda _pid: set(available), raising=False) + monkeypatch.setattr( + os, "sched_setaffinity", lambda pid, ids: calls.append((pid, set(ids))), raising=False + ) return calls @@ -103,8 +107,8 @@ def test_unavailable_cpu_ids_fail_closed(monkeypatch: pytest.MonkeyPatch): def test_platform_without_affinity_warns_and_caps_numba(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("NUMBA_NUM_THREADS", raising=False) - monkeypatch.delattr(os, "sched_setaffinity") - monkeypatch.delattr(os, "sched_getaffinity") + monkeypatch.delattr(os, "sched_setaffinity", raising=False) + monkeypatch.delattr(os, "sched_getaffinity", raising=False) confine_calls = _record_confine(monkeypatch) numba_calls = _record_numba(monkeypatch) @@ -125,7 +129,7 @@ def fake_setaffinity(pid, ids): raise ProcessLookupError calls.append((pid, set(ids))) - monkeypatch.setattr(os, "sched_setaffinity", fake_setaffinity) + monkeypatch.setattr(os, "sched_setaffinity", fake_setaffinity, raising=False) monkeypatch.setattr(os.path, "isdir", lambda path: path == cpu_runtime._PROC_TASK_DIR) monkeypatch.setattr(os, "listdir", lambda path: ["123", "456", "789"]) @@ -136,7 +140,9 @@ def fake_setaffinity(pid, ids): def test_confine_existing_threads_without_proc_is_noop(monkeypatch: pytest.MonkeyPatch): calls: list[tuple] = [] - monkeypatch.setattr(os, "sched_setaffinity", lambda pid, ids: calls.append((pid, set(ids)))) + monkeypatch.setattr( + os, "sched_setaffinity", lambda pid, ids: calls.append((pid, set(ids))), raising=False + ) monkeypatch.setattr(os.path, "isdir", lambda path: False) cpu_runtime._confine_existing_threads({0}) From 411a6ee3e0e0fe5a852c93f40774a3b4af4b4b0e Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 27 Aug 2026 14:32:13 +0800 Subject: [PATCH 11/12] fix(logging): make collector reward reporting timely Reward displays (tensorboard reward/mean and the terminal logger) lagged badly on off-policy and APPO runs: - collectors sent metrics only every num_envs * 10 env steps, so the reported reward changed just once per ~10 learner iterations; - runners then averaged the last 100 (off-policy) or 50 (APPO) reports, each already a rolling 100-episode mean, delaying the visible curve by ~1000 iterations. Report metrics every collector cycle, keep the runner-side window at the last 10 reports, and bound the per-worker episode reward/length buffers with deque(maxlen=100) instead of lists that grew for the whole run. Co-Authored-By: Claude Fable 5 --- src/unilab/algos/appo/runner.py | 9 +++++---- src/unilab/algos/appo/worker.py | 16 +++++++++------ src/unilab/algos/hora/appo_runner.py | 9 +++++---- src/unilab/algos/hora/appo_worker.py | 16 +++++++++------ .../algos/offpolicy/double_buffer_runner.py | 5 ++++- src/unilab/algos/offpolicy/worker.py | 20 ++++++++++--------- 6 files changed, 45 insertions(+), 30 deletions(-) diff --git a/src/unilab/algos/appo/runner.py b/src/unilab/algos/appo/runner.py index 4473b46ac..370ad4b2b 100644 --- a/src/unilab/algos/appo/runner.py +++ b/src/unilab/algos/appo/runner.py @@ -312,7 +312,10 @@ def learn( ) logger_started = False - reward_history: deque = deque(maxlen=200) + # Recent collector reports; each entry is already the collector's + # rolling 100-episode mean, so a short window keeps the logged + # reward timely without losing smoothing. + reward_history: deque = deque(maxlen=10) latest_reward_components: dict = {} staging_pool = RolloutStagingPool( @@ -396,9 +399,7 @@ def learn( logger.update_staging_pool(staging_pool.active_count, staging_pool.capacity) mean_reward = ( - sum(list(reward_history)[-50:]) / max(len(list(reward_history)[-50:]), 1) - if reward_history - else 0.0 + sum(reward_history) / max(len(reward_history), 1) if reward_history else 0.0 ) last_mean_reward = float(mean_reward) best_mean_reward = max(best_mean_reward, last_mean_reward) diff --git a/src/unilab/algos/appo/worker.py b/src/unilab/algos/appo/worker.py index 6375aed69..35eed27d3 100644 --- a/src/unilab/algos/appo/worker.py +++ b/src/unilab/algos/appo/worker.py @@ -8,7 +8,7 @@ import statistics import sys import time -from collections import defaultdict +from collections import defaultdict, deque from queue import Empty, Full from typing import Any, Dict @@ -235,8 +235,10 @@ def to_float32_np(x): obs_td = TensorDict({"policy": obs_torch}, batch_size=num_envs, device=collector_device) total_steps = 0 - ep_rewards = [] - ep_lengths = [] + # Bounded rolling window of the most recent completed episodes; an + # unbounded list here grows for the entire run. + ep_rewards: deque[float] = deque(maxlen=100) + ep_lengths: deque[int] = deque(maxlen=100) current_ep_rewards = np.zeros(num_envs, dtype=np.float32) current_ep_lengths = np.zeros(num_envs, dtype=np.int32) ep_reward_components = defaultdict(list) @@ -353,15 +355,17 @@ def to_float32_np(x): if k.startswith("reward/"): ep_reward_components[k].append(v) - if metrics_queue is not None and total_steps % (num_envs * 10) == 0: + # Report every env step so learner-side reward and throughput + # displays track the current policy without extra lag. + if metrics_queue is not None: try: msg: dict[str, Any] = { "total_steps": total_steps, } if ep_rewards: - msg["mean_ep_reward"] = statistics.mean(ep_rewards[-100:]) + msg["mean_ep_reward"] = statistics.mean(ep_rewards) msg["mean_ep_length"] = ( - statistics.mean(ep_lengths[-100:]) if ep_lengths else 0.0 + statistics.mean(ep_lengths) if ep_lengths else 0.0 ) if ep_completions > 0: msg["timeout_rate"] = ep_timeouts / ep_completions diff --git a/src/unilab/algos/hora/appo_runner.py b/src/unilab/algos/hora/appo_runner.py index 252591023..3bfb1b965 100644 --- a/src/unilab/algos/hora/appo_runner.py +++ b/src/unilab/algos/hora/appo_runner.py @@ -303,7 +303,10 @@ def learn( f"epochs={learner.num_learning_epochs})" ) - reward_history: deque = deque(maxlen=200) + # Recent collector reports; each entry is already the collector's + # rolling 100-episode mean, so a short window keeps the logged + # reward timely without losing smoothing. + reward_history: deque = deque(maxlen=10) latest_reward_components: dict = {} staging_pool = RolloutStagingPool( capacity=self.staging_pool_size, @@ -379,9 +382,7 @@ def learn( logger.update_staging_pool(staging_pool.active_count, staging_pool.capacity) mean_reward = ( - sum(list(reward_history)[-50:]) / max(len(list(reward_history)[-50:]), 1) - if reward_history - else 0.0 + sum(reward_history) / max(len(reward_history), 1) if reward_history else 0.0 ) last_mean_reward = float(mean_reward) best_mean_reward = max(best_mean_reward, last_mean_reward) diff --git a/src/unilab/algos/hora/appo_worker.py b/src/unilab/algos/hora/appo_worker.py index a5cccb432..0d5ec8002 100644 --- a/src/unilab/algos/hora/appo_worker.py +++ b/src/unilab/algos/hora/appo_worker.py @@ -5,7 +5,7 @@ import statistics import sys import time -from collections import defaultdict +from collections import defaultdict, deque from typing import Any, Dict import numpy as np @@ -227,8 +227,10 @@ def to_float32_np(x): ) total_steps = 0 - ep_rewards = [] - ep_lengths = [] + # Bounded rolling window of the most recent completed episodes; an + # unbounded list here grows for the entire run. + ep_rewards: deque[float] = deque(maxlen=100) + ep_lengths: deque[int] = deque(maxlen=100) current_ep_rewards = np.zeros(num_envs, dtype=np.float32) current_ep_lengths = np.zeros(num_envs, dtype=np.int32) ep_reward_components = defaultdict(list) @@ -365,15 +367,17 @@ def to_float32_np(x): if k.startswith("reward/"): ep_reward_components[k].append(v) - if metrics_queue is not None and total_steps % (num_envs * 10) == 0: + # Report every env step so learner-side reward and throughput + # displays track the current policy without extra lag. + if metrics_queue is not None: try: msg: dict[str, Any] = { "total_steps": total_steps, } if ep_rewards: - msg["mean_ep_reward"] = statistics.mean(ep_rewards[-100:]) + msg["mean_ep_reward"] = statistics.mean(ep_rewards) msg["mean_ep_length"] = ( - statistics.mean(ep_lengths[-100:]) if ep_lengths else 0.0 + statistics.mean(ep_lengths) if ep_lengths else 0.0 ) if ep_completions > 0: msg["timeout_rate"] = ep_timeouts / ep_completions diff --git a/src/unilab/algos/offpolicy/double_buffer_runner.py b/src/unilab/algos/offpolicy/double_buffer_runner.py index f689811ba..ce506b23b 100644 --- a/src/unilab/algos/offpolicy/double_buffer_runner.py +++ b/src/unilab/algos/offpolicy/double_buffer_runner.py @@ -969,7 +969,10 @@ def learn( time.sleep(0.5) - reward_history: deque = deque(maxlen=100) + # Recent collector reports; each entry is already the collector's + # rolling 100-episode mean, so a short window keeps the logged + # reward timely without losing smoothing. + reward_history: deque = deque(maxlen=10) latest_reward_components: dict[str, float] = {} has_logged_reward = False last_buf_log = 0 diff --git a/src/unilab/algos/offpolicy/worker.py b/src/unilab/algos/offpolicy/worker.py index 7bb8e748b..07fa38fca 100644 --- a/src/unilab/algos/offpolicy/worker.py +++ b/src/unilab/algos/offpolicy/worker.py @@ -243,12 +243,15 @@ def _run_collector( replay_buffer.trace_recorder = trace_recorder replay_buffer.trace_thread_time = trace_thread_time replay_buffer.attach_stop_event(stop_event) + from collections import defaultdict, deque + total_steps = 0 - ep_rewards = [] - ep_lengths = [] + # Bounded rolling window of the most recent completed episodes; an + # unbounded list here grows for the entire run. + ep_rewards: deque[float] = deque(maxlen=100) + ep_lengths: deque[int] = deque(maxlen=100) current_ep_rewards = np.zeros(num_envs, dtype=np.float32) current_ep_lengths = np.zeros(num_envs, dtype=np.int32) - from collections import defaultdict ep_reward_components = defaultdict(list) timing_accum_ms: defaultdict[str, float] = defaultdict(float) @@ -454,8 +457,9 @@ def _run_collector( if k.startswith("reward/"): ep_reward_components[k].append(v) - # Send metrics periodically - if metrics_queue is not None and total_steps % (num_envs * 10) == 0: + # Send metrics every collector cycle so learner-side reward and + # throughput displays track the current policy without extra lag. + if metrics_queue is not None: import statistics try: @@ -464,10 +468,8 @@ def _run_collector( "buffer_size": int(replay_buffer.size[0]), } if ep_rewards: - msg["mean_ep_reward"] = statistics.mean(ep_rewards[-100:]) - msg["mean_ep_length"] = ( - statistics.mean(ep_lengths[-100:]) if ep_lengths else 0.0 - ) + msg["mean_ep_reward"] = statistics.mean(ep_rewards) + msg["mean_ep_length"] = statistics.mean(ep_lengths) if ep_lengths else 0.0 # Add mean reward components if ep_reward_components: components_mean = {} From 22fc439ea546d4b38b3add57b50770291aa73811 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 27 Aug 2026 15:36:13 +0800 Subject: [PATCH 12/12] fix(env): keep per-step reward log entries through autoreset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ManagerBasedRlEnv.reset() replaced state.info["log"] with the reset-only extras (Episode_Reward/*), wiping the fresh per-step reward/* entries that _update_state_in_read_phase() had just computed for the current transition. On any step where at least one env resets — with thousands of envs, nearly every step — collectors therefore saw no reward/* keys at all, so the per-term reward components in tensorboard and the terminal logger stayed frozen at one stale value for thousands of iterations (observed as long flat staircases on reward/motion_* etc.). Merge instead of replace on the autoreset path: the pre-reset per-step entries stay, reset extras layer on top. Standalone (non-autoreset) resets are unchanged. Co-Authored-By: Claude Fable 5 --- src/unilab/envs/manager_based_rl_env.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 18369bcfa..07257827f 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -579,6 +579,14 @@ def reset( if self._state is not None: for name, values in reset_obs.items(): self._state.obs[name][ids] = values + if self._autoreset_reset_active: + # Autoreset runs at the tail of step(): keep this step's + # per-step log entries (reward/* etc., computed pre-reset) and + # layer the reset extras (Episode_Reward/* etc.) on top, so + # consumers still see the transition's reward breakdown. + step_log = self._state.info.get("log") + if step_log: + log = {**step_log, **log} self._state.info["log"] = log if not self._autoreset_reset_active: self._state.terminated[ids] = False