From a482a1fc249cc3fb86b12c94eb43625181e2a66e Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Thu, 27 Aug 2026 02:33:29 +0800 Subject: [PATCH] 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], + )