Skip to content

fix: respect motrix qpos metadata in set_state #227

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TATP-233 merged 1 commit into from
Apr 12, 2026
Merged

fix: respect motrix qpos metadata in set_state #227

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 33 additions & 12 deletions src/unilab/base/backend/motrix_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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 #
Expand Down Expand Up @@ -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)
Expand Down
109 changes: 94 additions & 15 deletions tests/base/test_sim_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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),
Expand All @@ -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
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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
# ---------------------------------------------------------------------------
Expand All @@ -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

Expand Down Expand Up @@ -338,23 +417,23 @@ 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]))
np.testing.assert_allclose(pos_b, 0.0, atol=1e-5)

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, :]
np.testing.assert_allclose(pos_w, bkd.get_base_pos(), atol=1e-5)

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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 = [
Expand Down Expand Up @@ -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])
Expand Down Expand Up @@ -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):
Expand Down
Loading