perf: optimize issue 673 reset_done paths

src/unilab/base/backend/motrix/backend.py

@@ -60,6 +60,16 @@ def _first_scalar(value: Any) -> float:
     return float(arr.reshape(-1)[0])
 
 
+def _contiguous_slice(indices: np.ndarray) -> slice | None:
+    if indices.size == 0:
+        return None
+    start = int(indices[0])
+    stop = start + int(indices.size)
+    if np.array_equal(indices, np.arange(start, stop, dtype=indices.dtype)):
+        return slice(start, stop)
+    return None
+
+
 @dataclass
 class _MotrixSceneContext:
     model: "mtx.SceneModel"
@@ -185,6 +195,7 @@ def __init__(
         self._body_floatingbase = self._body.floatingbase
         self._joint_dof_pos_indices = np.asarray(self._model.joint_dof_pos_indices, dtype=np.intp)
         self._joint_dof_vel_indices = np.asarray(self._model.joint_dof_vel_indices, dtype=np.intp)
+        self._joint_dof_pos_slice = _contiguous_slice(self._joint_dof_pos_indices)
         position_actuators: list["mtx.PositionActuator"] = []
         for actuator in self._model.actuators:
             if actuator.typ == "position":
@@ -212,6 +223,11 @@ def __init__(
                 joint_pos_idx.append(int(joint.dof_pos_index))
             if len(joint_pos_idx) == int(self._model.num_actuators):
                 self._actuator_joint_pos_indices = np.asarray(joint_pos_idx, dtype=np.intp)
+        self._actuator_joint_pos_slice = (
+            _contiguous_slice(self._actuator_joint_pos_indices)
+            if self._actuator_joint_pos_indices is not None
+            else None
+        )
         self._default_actuator_kp = np.zeros((self.num_actuators,), dtype=np.float32)
         self._default_actuator_kd = np.zeros((self.num_actuators,), dtype=np.float32)
         for actuator in self._position_actuators:
@@ -599,16 +615,22 @@ def set_state(
             self.num_actuators
         ):
             # Fully-actuated model: hold every joint at its reset position (unchanged).
-            ctrl = qpos_motrix[:, self._joint_dof_pos_indices]
+            if self._joint_dof_pos_slice is not None:
+                ctrl = qpos_motrix[:, self._joint_dof_pos_slice]
+            else:
+                ctrl = qpos_motrix[:, self._joint_dof_pos_indices]
         elif self._actuator_joint_pos_indices is not None:
             # Under-actuated / parallel model: hold only the actuated joints.
-            ctrl = qpos_motrix[:, self._actuator_joint_pos_indices]
+            if self._actuator_joint_pos_slice is not None:
+                ctrl = qpos_motrix[:, self._actuator_joint_pos_slice]
+            else:
+                ctrl = qpos_motrix[:, self._actuator_joint_pos_indices]
         else:
             ctrl = np.zeros((len(env_indices), self.num_actuators), dtype=self._np_dtype)
         data_slice.actuator_ctrls = np.ascontiguousarray(ctrl)
 
         self._model.forward_kinematic(data_slice)
-        self._refresh_link_pose_cache(env_indices)
+        self._refresh_link_pose_cache(env_indices, data_slice=data_slice)
         self._invalidate_link_velocity_cache()
 
     def get_dr_capabilities(self) -> DomainRandomizationCapabilities:
@@ -985,13 +1007,17 @@ def _motrix_qpos_to_mujoco(self, qpos: np.ndarray) -> np.ndarray:
             qpos_mujoco[..., quat_indices] = qpos[..., quat_indices[[3, 0, 1, 2]]]
         return qpos_mujoco
 
-    def _refresh_link_pose_cache(self, env_indices: np.ndarray | None = None) -> None:
+    def _refresh_link_pose_cache(
+        self, env_indices: np.ndarray | None = None, data_slice: Any | None = None
+    ) -> None:
         if env_indices is None:
             self._link_poses = self._model.get_link_poses(self._data)
         else:
-            mask = np.zeros(self._num_envs, dtype=bool)
-            mask[env_indices] = True
-            self._link_poses[env_indices] = self._model.get_link_poses(self._data[mask])
+            if data_slice is None:
+                mask = np.zeros(self._num_envs, dtype=bool)
+                mask[env_indices] = True
+                data_slice = self._data[mask]
+            self._link_poses[env_indices] = self._model.get_link_poses(data_slice)
 
     def _refresh_link_velocity_cache(self, env_indices: np.ndarray | None = None) -> None:
         if env_indices is None:

src/unilab/envs/motion_tracking/g1/tracking.py

@@ -446,7 +446,9 @@ def build_reset_observation(
         dof_vel_ms = (time.perf_counter() - t0) * 1000.0
 
         t0 = time.perf_counter()
-        all_pos_w, all_quat_w = env._get_body_pose_w()
+        robot_body_pos_w, robot_body_quat_w = env._backend.get_body_pose_w_rows(
+            env_ids, env.body_ids
+        )
         body_pose_ms = (time.perf_counter() - t0) * 1000.0
 
         obs_info = dict(info_updates)
@@ -465,8 +467,8 @@ def build_reset_observation(
                 gyro,
                 dof_pos,
                 dof_vel,
-                all_pos_w[env_ids],
-                all_quat_w[env_ids],
+                robot_body_pos_w,
+                robot_body_quat_w,
             ),
         )
         compute_obs_ms = (time.perf_counter() - t0) * 1000.0

tests/envs/test_env_configs.py

@@ -477,6 +477,85 @@ def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]
     np.testing.assert_array_equal(env._backend.calls[0], np.array([1, 3], dtype=np.int32))
 
 
+def test_g1_motion_tracking_reset_observation_uses_sparse_body_pose_rows():
+    from unilab.envs.motion_tracking.g1.motion_loader import MotionData
+    from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingDomainRandomizationProvider
+
+    class FakeBackend:
+        def __init__(self) -> None:
+            self.row_calls: list[tuple[np.ndarray, np.ndarray]] = []
+
+        def get_body_pose_w_rows(
+            self, env_ids: np.ndarray, body_ids: np.ndarray
+        ) -> tuple[np.ndarray, np.ndarray]:
+            self.row_calls.append((env_ids.copy(), body_ids.copy()))
+            rows = len(env_ids)
+            bodies = len(body_ids)
+            return (
+                np.full((rows, bodies, 3), 2.0, dtype=np.float32),
+                np.full((rows, bodies, 4), 3.0, dtype=np.float32),
+            )
+
+        def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+            raise AssertionError("reset observation should use sparse body pose rows")
+
+    class FakeMotionLoader:
+        def get_motion_at_frame(self, frames: np.ndarray) -> MotionData:
+            rows = len(frames)
+            return MotionData(
+                joint_pos=np.zeros((rows, 2), dtype=np.float32),
+                joint_vel=np.zeros((rows, 2), dtype=np.float32),
+                body_pos_w=np.zeros((rows, 2, 3), dtype=np.float32),
+                body_quat_w=np.zeros((rows, 2, 4), dtype=np.float32),
+                body_lin_vel_w=np.zeros((rows, 2, 3), dtype=np.float32),
+                body_ang_vel_w=np.zeros((rows, 2, 3), dtype=np.float32),
+            )
+
+    class FakeMotionSampler:
+        current_frames = np.array([10, 11, 12, 13], dtype=np.int32)
+
+    env = SimpleNamespace(
+        _backend=FakeBackend(),
+        body_ids=np.array([1, 3], dtype=np.int32),
+        motion_loader=FakeMotionLoader(),
+        motion_sampler=FakeMotionSampler(),
+        get_local_linvel=lambda: np.zeros((4, 3), dtype=np.float32),
+        get_gyro=lambda: np.zeros((4, 3), dtype=np.float32),
+        get_dof_pos=lambda: np.zeros((4, 2), dtype=np.float32),
+        get_dof_vel=lambda: np.zeros((4, 2), dtype=np.float32),
+    )
+
+    captured: dict[str, np.ndarray] = {}
+
+    def compute_obs(
+        obs_info,
+        motion_data,
+        linvel,
+        gyro,
+        dof_pos,
+        dof_vel,
+        robot_body_pos_w,
+        robot_body_quat_w,
+    ):
+        del obs_info, motion_data, linvel, gyro, dof_pos, dof_vel
+        captured["robot_body_pos_w"] = robot_body_pos_w
+        captured["robot_body_quat_w"] = robot_body_quat_w
+        return {"obs": np.zeros((2, 1), dtype=np.float32)}
+
+    env._compute_obs = compute_obs
+    provider = G1MotionTrackingDomainRandomizationProvider()
+    env_ids = np.array([1, 3], dtype=np.int32)
+
+    obs = provider.build_reset_observation(env, env_ids, {})
+
+    assert obs["obs"].shape == (2, 1)
+    assert len(env._backend.row_calls) == 1
+    np.testing.assert_array_equal(env._backend.row_calls[0][0], env_ids)
+    np.testing.assert_array_equal(env._backend.row_calls[0][1], env.body_ids)
+    assert captured["robot_body_pos_w"].shape == (2, 2, 3)
+    assert captured["robot_body_quat_w"].shape == (2, 2, 4)
+
+
 def _compute_g1_motion_tracking_obs_stub(env_cls: type):
     from unilab.envs.motion_tracking.g1.motion_loader import MotionData