diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index daa1f5bd8..49aaa417f 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/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: diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index 5e86a1bba..eac620b60 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/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 diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index a570094f0..be0eaa9a8 100644 --- a/tests/envs/test_env_configs.py +++ b/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