perf(mjwarp): reduce g1 motion tracking reset latency

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,

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

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()

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],
+    )