diff --git a/src/unilab/base/backend/motrix_backend.py b/src/unilab/base/backend/motrix_backend.py index 40540825e..7da498099 100644 --- a/src/unilab/base/backend/motrix_backend.py +++ b/src/unilab/base/backend/motrix_backend.py @@ -76,6 +76,14 @@ def __init__( self._data = mtx.SceneData(self._model, batch=[num_envs]) # pyright: ignore[reportPossiblyUnbound] self._body = self._model.get_body(base_name) self._body_link = self._model.get_link(base_name) + 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._floating_base_quat_indices: tuple[np.ndarray, ...] = tuple( + np.asarray(floating_base.dof_pos_indices[3:7], dtype=np.intp) + for floating_base in getattr(self._model, "floating_bases", []) + if len(floating_base.dof_pos_indices) >= 7 + ) self._default_base_mass_override = np.asarray( self._body_link.get_mass_override(self._data) ).copy() @@ -121,9 +129,7 @@ def num_actuators(self) -> int: @property def num_dof_vel(self) -> int: - # model.num_dof_vel includes 6 floating-base DOFs; exclude them - # to match get_dof_vel() which returns joint velocities only. - return int(self._model.num_dof_vel) - 6 + return int(len(self._joint_dof_vel_indices)) def get_actuator_ctrl_range(self) -> np.ndarray: arr: np.ndarray = np.array(self._model.actuator_ctrl_limits, dtype=self._np_dtype) @@ -185,9 +191,7 @@ def set_state( qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, ) -> None: - # Convert quaternion from mujoco (wxyz) to motrix (xyzw) - qpos_motrix = qpos.copy() - qpos_motrix[:, 3:7] = qpos[:, [4, 5, 6, 3]] + qpos_motrix = self._mujoco_qpos_to_motrix(qpos) # Create mask for batch operation mask = np.zeros(self._num_envs, dtype=bool) @@ -200,8 +204,8 @@ def set_state( data_slice.set_dof_vel(qvel) data_slice.set_dof_pos(qpos_motrix, self._model) - # Set control to joint positions (actuator target for PD control) - ctrl = qpos_motrix[:, 7:] + # Actuator targets follow the model's joint-position ordering only. + ctrl = qpos_motrix[:, self._joint_dof_pos_indices] data_slice.actuator_ctrls = np.ascontiguousarray(ctrl) self._model.forward_kinematic(data_slice) @@ -226,16 +230,26 @@ def get_play_capabilities(self) -> BackendPlayCapabilities: # ------------------------------------------------------------------ # def get_base_pos(self) -> np.ndarray: - return np.asarray(self._body.floatingbase.get_translation(self._data)) + if self._body_floatingbase is not None: + return np.asarray(self._body_floatingbase.get_translation(self._data)) + return np.asarray(self._body_link.get_pose(self._data))[:, :3] def get_base_quat(self) -> np.ndarray: - return self._xyzw_to_wxyz(np.asarray(self._body.floatingbase.get_rotation(self._data))) + if self._body_floatingbase is not None: + quat = np.asarray(self._body_floatingbase.get_rotation(self._data)) + else: + quat = np.asarray(self._body_link.get_rotation(self._data)) + return self._xyzw_to_wxyz(quat) def get_base_lin_vel(self) -> np.ndarray: - return np.asarray(self._body.floatingbase.get_global_linear_velocity(self._data)) + if self._body_floatingbase is not None: + return np.asarray(self._body_floatingbase.get_global_linear_velocity(self._data)) + return np.asarray(self._body_link.get_linear_velocity(self._data)) def get_base_ang_vel(self) -> np.ndarray: - return np.asarray(self._body.floatingbase.get_global_angular_velocity(self._data)) + if self._body_floatingbase is not None: + return np.asarray(self._body_floatingbase.get_global_angular_velocity(self._data)) + return np.asarray(self._body_link.get_angular_velocity(self._data)) # ------------------------------------------------------------------ # # DOF state # @@ -328,6 +342,13 @@ def _xyzw_to_wxyz(self, q: np.ndarray) -> np.ndarray: """motrix xyzw → wxyz""" return q[..., [3, 0, 1, 2]] + def _mujoco_qpos_to_motrix(self, qpos: np.ndarray) -> np.ndarray: + """Convert every MuJoCo freejoint quaternion slice from wxyz to xyzw.""" + qpos_motrix = np.array(qpos, copy=True) + for quat_indices in self._floating_base_quat_indices: + qpos_motrix[:, quat_indices] = qpos[:, quat_indices[[1, 2, 3, 0]]] + return qpos_motrix + def _refresh_link_pose_cache(self, env_indices: np.ndarray | None = None) -> None: if env_indices is None: self._link_poses = self._model.get_link_poses(self._data) diff --git a/tests/base/test_sim_backend.py b/tests/base/test_sim_backend.py index 4c954c47c..5a06500ed 100644 --- a/tests/base/test_sim_backend.py +++ b/tests/base/test_sim_backend.py @@ -7,6 +7,8 @@ uv run pytest -m slow tests/base/test_sim_backend.py -v """ +from typing import Any, cast + import numpy as np import pytest @@ -26,6 +28,7 @@ def _xml(robot: str, scene: str = "scene_flat.xml") -> str: ] _G1 = dict(model_file=_xml("g1"), base_name="pelvis") +_ALLEGRO = dict(model_file=_xml("allegro_hand", "scene.xml"), base_name="palm") NUM_ENVS = 2 SIM_DT = 0.005 @@ -40,6 +43,12 @@ def _shape(arr: np.ndarray, *expected: int) -> None: assert arr.shape == expected, f"expected shape {expected}, got {arr.shape}" +def _mujoco_module() -> Any: + import mujoco + + return cast(Any, mujoco) + + def _unit_quat(q: np.ndarray, tag: str = "") -> None: np.testing.assert_allclose( np.linalg.norm(q, axis=-1), @@ -58,13 +67,24 @@ def _identity_qpos_mujoco(nq: int, xyz=(0.0, 0.0, 0.8)) -> np.ndarray: def _mujoco_expected_dof_dims(model) -> tuple[int, int]: - import mujoco + mujoco = _mujoco_module() if model.njnt > 0 and int(model.jnt_type[0]) == int(mujoco.mjtJoint.mjJNT_FREE): return model.nq - 7, model.nv - 6 return model.nq, model.nv +def _allegro_state() -> tuple[np.ndarray, np.ndarray]: + qpos = np.zeros((1, 23), dtype=np.float64) + qvel = np.zeros((1, 22), dtype=np.float64) + qpos[0, :16] = np.linspace(-0.2, 0.2, 16) + qpos[0, 16:19] = np.array([-0.01, 0.02, 0.16]) + qpos[0, 19:23] = np.array([0.92387953, 0.0, 0.38268343, 0.0]) + qvel[0, :16] = np.linspace(-0.15, 0.15, 16) + qvel[0, 16:22] = np.array([0.1, -0.2, 0.05, 0.3, -0.1, 0.2]) + return qpos, qvel + + # --------------------------------------------------------------------------- # MuJoCo — basic, 3 robots, no body sensors # --------------------------------------------------------------------------- @@ -226,7 +246,7 @@ def test_dof_pos_finite_after_step(self, bkd): @pytest.mark.slow def test_mujoco_backend_discards_visual_assets(): - import mujoco + mujoco = _mujoco_module() from unilab.base.backend.mujoco_backend import MuJoCoBackend @@ -242,7 +262,7 @@ def test_mujoco_backend_discards_visual_assets(): @pytest.mark.slow def test_mujoco_backend_fixed_base_dof_views_do_not_skip_first_joint(): - import mujoco + mujoco = _mujoco_module() from unilab.base.backend.mujoco_backend import MuJoCoBackend @@ -258,6 +278,65 @@ def test_mujoco_backend_fixed_base_dof_views_do_not_skip_first_joint(): _unit_quat(bkd.get_base_quat(), "MuJoCo fixed-base quat") +@pytest.mark.slow +def test_motrix_backend_fixed_base_base_views_are_available(): + from unilab.base.backend.motrix_backend import MotrixBackend + + pytest.importorskip("motrixsim") + bkd = MotrixBackend(_ALLEGRO["model_file"], NUM_ENVS, SIM_DT, base_name=_ALLEGRO["base_name"]) + _shape(bkd.get_base_pos(), NUM_ENVS, 3) + _shape(bkd.get_base_quat(), NUM_ENVS, 4) + np.testing.assert_allclose(bkd.get_base_lin_vel(), 0.0, atol=1e-8) + np.testing.assert_allclose(bkd.get_base_ang_vel(), 0.0, atol=1e-8) + _unit_quat(bkd.get_base_quat(), "Motrix fixed-base quat") + + +@pytest.mark.slow +def test_motrix_backend_fixed_base_set_state_matches_mujoco_for_hand_and_ball(): + from unilab.base.backend.motrix_backend import MotrixBackend + from unilab.base.backend.mujoco_backend import MuJoCoBackend + + pytest.importorskip("motrixsim") + mj = MuJoCoBackend( + _ALLEGRO["model_file"], + NUM_ENVS, + SIM_DT, + base_name=_ALLEGRO["base_name"], + add_body_sensors=True, + ) + mx = MotrixBackend( + _ALLEGRO["model_file"], + NUM_ENVS, + SIM_DT, + base_name=_ALLEGRO["base_name"], + add_body_sensors=True, + ) + qpos, qvel = _allegro_state() + env_idx = np.array([0]) + + mj.set_state(env_idx, qpos, qvel) + mx.set_state(env_idx, qpos, qvel) + + np.testing.assert_allclose(mx.get_dof_pos()[0], qpos[0, :16], atol=1e-6) + np.testing.assert_allclose(mx.get_dof_vel()[0], qvel[0, :16], atol=1e-6) + np.testing.assert_allclose(np.asarray(mx.data.actuator_ctrls[0]), qpos[0, :16], atol=1e-6) + np.testing.assert_allclose(mx.get_base_pos(), mj.get_base_pos(), atol=2e-3) + np.testing.assert_allclose(np.abs(mx.get_base_quat()), np.abs(mj.get_base_quat()), atol=2e-3) + + mj_ball_id = _mj_body_id(mj.model, "ball") + mx_ball_id = _mx_link_id(mx.model, "ball") + np.testing.assert_allclose( + mx.get_body_pos_w(np.array([mx_ball_id])), + mj.get_body_pos_w(np.array([mj_ball_id])), + atol=2e-3, + ) + np.testing.assert_allclose( + mx.get_body_quat_w(np.array([mx_ball_id])), + mj.get_body_quat_w(np.array([mj_ball_id])), + atol=2e-3, + ) + + # --------------------------------------------------------------------------- # MuJoCo — body sensors, G1 only # --------------------------------------------------------------------------- @@ -282,9 +361,9 @@ def body_ids(self, bkd): return np.array(bkd._tracked_body_ids[:2]) def _bname(self, bkd, bid: int) -> str: - import mujoco + mujoco = _mujoco_module() - return mujoco.mj_id2name(bkd.model, mujoco.mjtObj.mjOBJ_BODY, bid) + return cast(str, mujoco.mj_id2name(bkd.model, mujoco.mjtObj.mjOBJ_BODY, bid)) # world frame @@ -338,7 +417,7 @@ def test_get_sensor_data_unknown_raises(self, bkd): def test_base_body_pos_b_is_zero(self, bkd): """Base body position relative to itself must be [0,0,0].""" - import mujoco + mujoco = _mujoco_module() pelvis_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, "pelvis") pos_b = bkd.get_body_pos_b(np.array([pelvis_id])) @@ -346,7 +425,7 @@ def test_base_body_pos_b_is_zero(self, bkd): def test_base_body_pos_w_matches_get_base_pos(self, bkd): """get_body_pos_w for the base body must equal get_base_pos().""" - import mujoco + mujoco = _mujoco_module() pelvis_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, "pelvis") pos_w = bkd.get_body_pos_w(np.array([pelvis_id]))[:, 0, :] @@ -354,7 +433,7 @@ def test_base_body_pos_w_matches_get_base_pos(self, bkd): def test_base_body_quat_b_is_identity(self, bkd): """Base body quaternion relative to itself must be identity [1,0,0,0].""" - import mujoco + mujoco = _mujoco_module() pelvis_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, "pelvis") quat_b = bkd.get_body_quat_b(np.array([pelvis_id]))[:, 0, :] # (N, 4) @@ -578,13 +657,13 @@ def test_base_body_pos_b_is_zero(self, bkd): def _mj_body_id(mj_model, name: str) -> int: - import mujoco + mujoco = _mujoco_module() - return mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_BODY, name) + return int(mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_BODY, name)) def _mx_link_id(mx_model, name: str) -> int: - return mx_model.get_link_index(name) + return int(mx_model.get_link_index(name)) @pytest.mark.slow @@ -693,7 +772,7 @@ def synced(self): @pytest.fixture def body_pairs(self, synced): """选前 3 个 body（pelvis + 两个髋关节），分别返回 mujoco IDs 和 motrixsim link IDs。""" - import mujoco + mujoco = _mujoco_module() mj, mx = synced names = [ @@ -830,7 +909,7 @@ def test_get_init_qvel(self, bkd): np.testing.assert_array_equal(qvel, 0.0) def test_get_body_ids(self, bkd): - import mujoco + mujoco = _mujoco_module() base_name = mujoco.mj_id2name(bkd.model, mujoco.mjtObj.mjOBJ_BODY, 1) ids = bkd.get_body_ids([base_name]) @@ -922,8 +1001,8 @@ def backends(self): from unilab.base.backend.motrix_backend import MotrixBackend from unilab.base.backend.mujoco_backend import MuJoCoBackend - mj = MuJoCoBackend(**_G1, num_envs=NUM_ENVS, sim_dt=SIM_DT) - mx = MotrixBackend(**_G1, num_envs=NUM_ENVS, sim_dt=SIM_DT) + mj = MuJoCoBackend(_G1["model_file"], NUM_ENVS, SIM_DT, base_name=_G1["base_name"]) + mx = MotrixBackend(_G1["model_file"], NUM_ENVS, SIM_DT, base_name=_G1["base_name"]) return mj, mx def test_num_actuators_match(self, backends):