diff --git a/src/unilab/base/backend/__init__.py b/src/unilab/base/backend/__init__.py index 5f51a8bb8..95d126dbd 100644 --- a/src/unilab/base/backend/__init__.py +++ b/src/unilab/base/backend/__init__.py @@ -157,6 +157,8 @@ def create_backend( return cast(SimBackend, MuJoCoBackend(scene, num_envs, sim_dt, **kwargs)) if backend_type == "mjwarp": MjwarpBackend = _load_mjwarp_backend() + if body_state_required: + kwargs["add_body_sensors"] = True if position_actuator_gains is not None: raise ValueError( "mjwarp does not accept position_actuator_gains in the host compatibility " diff --git a/src/unilab/base/backend/mjwarp/backend.py b/src/unilab/base/backend/mjwarp/backend.py index f749333ed..38b390c45 100644 --- a/src/unilab/base/backend/mjwarp/backend.py +++ b/src/unilab/base/backend/mjwarp/backend.py @@ -12,7 +12,7 @@ import time from collections.abc import Callable, Sequence from os import PathLike -from typing import Any +from typing import Any, NoReturn import numpy as np @@ -29,6 +29,7 @@ IntervalRandomizationPlan, ResetRandomizationPayload, ) +from unilab.utils.rotation import np_quat_apply_inverse_batched from .dependencies import load_mjwarp_dependencies from .materialization import materialize_mjwarp_scene @@ -55,6 +56,7 @@ def __init__( push_body_name: str | None = None, nconmax: int | None = None, njmax: int | None = None, + add_body_sensors: bool = False, **unexpected_kwargs: Any, ) -> None: if unexpected_kwargs: @@ -64,6 +66,10 @@ def __init__( raise ValueError(f"num_envs must be a positive integer, got {num_envs!r}") if float(sim_dt) <= 0.0: raise ValueError(f"sim_dt must be positive, got {sim_dt!r}") + if not isinstance(add_body_sensors, bool): + raise TypeError( + f"MjwarpBackend add_body_sensors must be bool, got {type(add_body_sensors).__name__}" + ) nconmax = self._require_capacity(nconmax, name="nconmax", default=512) njmax = self._require_capacity(njmax, name="njmax", default=512) if push_body_name is not None: @@ -80,7 +86,7 @@ def __init__( "host or select the mujoco backend." ) - scene_context = materialize_mjwarp_scene(scene) + scene_context = materialize_mjwarp_scene(scene, add_body_sensors=add_body_sensors) self._scene_cleanup_handle = scene_context.cleanup_handle self.scene_model_file = scene_context.diagnostic_model_file self.scene_visual_model_file = str(scene.visual_model_file or scene.model_file) @@ -92,11 +98,19 @@ def __init__( self._base_name = base_name self._nconmax = nconmax self._njmax = njmax + self._add_body_sensors = add_body_sensors + self._tracked_body_names = scene_context.tracked_body_names self._mujoco = deps.mujoco self._mujoco_warp = deps.mujoco_warp self._warp = deps.warp - self._cpu_model = deps.mujoco.MjModel.from_xml_path(scene_context.source_model_file) + try: + self._cpu_model = deps.mujoco.MjModel.from_xml_path(scene_context.source_model_file) + finally: + # The materialized source (fragment merge and/or injected tracking + # sensors) is only needed to compile the model; release the + # temporary files immediately like the MuJoCo backend does. + self.cleanup_scene_assets() self._cpu_model.opt.timestep = self._sim_dt self._device_model = deps.mujoco_warp.put_model(self._cpu_model) self._device_data = deps.mujoco_warp.make_data( @@ -163,6 +177,11 @@ def __init__( self._ctrl_staging = np.zeros((self._num_envs, self._nu), dtype=np.float32) self._reset_mask_host = np.zeros((self._num_envs,), dtype=np.bool_) self._reset_mask_device = deps.warp.zeros(self._num_envs, dtype=bool) + # Tracked-body views are zero-copy slices of _sensor_cache; they must be + # bound before the first forward barrier refreshes the cache below. + self._body_id_to_tracked_idx: np.ndarray | None = None + if self._add_body_sensors: + self._bind_tracked_body_state() # Begin from explicit model defaults, run a forward barrier, and cache # the resulting sensors/kinematics. This avoids an uninitialized host # cache before NpEnv's first selected-row reset. @@ -237,6 +256,76 @@ def _bind_joint_range(self) -> np.ndarray | None: joint_range = np.asarray(self._cpu_model.jnt_range, dtype=np.float32)[mask] return None if joint_range.size == 0 else joint_range.copy() + def _bind_tracked_body_state(self) -> None: + """Bind zero-copy tracked-body views into the per-step sensor cache. + + Sensor columns follow the ``tracked_body_names`` insertion order from + the cold-path injection; body ids are rebuilt from the compiled model + because MjSpec compilation can reorder bodies (same reasoning as the + MuJoCo backend). + """ + names = self._tracked_body_names + if not names: + raise ValueError( + "mjwarp add_body_sensors requires at least one named body in the model" + ) + body_type = self._mujoco.mjtObj.mjOBJ_BODY + tracked_ids = [self._mujoco.mj_name2id(self._cpu_model, body_type, name) for name in names] + missing = [name for name, body_id in zip(names, tracked_ids, strict=True) if body_id < 0] + if missing: + raise ValueError( + "Injected mjwarp body tracking sensors reference bodies missing from " + f"the compiled model: {missing}" + ) + mapping = np.full(self._nbody, -1, dtype=np.intp) + for index, body_id in enumerate(tracked_ids): + mapping[body_id] = index + self._body_id_to_tracked_idx = mapping + self._tracked_pos_w_all = self._tracked_sensor_view("track_pos_w", 3) + self._tracked_quat_w_all = self._tracked_sensor_view("track_quat_w", 4) + self._tracked_linvel_w_all = self._tracked_sensor_view("track_linvel_w", 3) + self._tracked_angvel_w_all = self._tracked_sensor_view("track_angvel_w", 3) + + def _tracked_sensor_view(self, prefix: str, dim: int) -> np.ndarray: + count = len(self._tracked_body_names) + addresses = [] + for name in self._tracked_body_names: + sensor_name = f"{prefix}_{name}" + try: + address, sensor_dim = self._sensor_slots[sensor_name] + except KeyError as exc: + raise ValueError( + f"Injected mjwarp tracking sensor {sensor_name!r} is missing from the " + "compiled model" + ) from exc + if sensor_dim != dim: + raise ValueError( + f"Injected mjwarp tracking sensor {sensor_name!r} has dim {sensor_dim}; " + f"expected {dim}" + ) + addresses.append(address) + first = addresses[0] + if addresses != [first + index * dim for index in range(count)]: + raise ValueError( + f"Injected mjwarp tracking sensors {prefix}_* are not one contiguous " + "sensor block in tracked-body order" + ) + return self._sensor_cache[:, first : first + count * dim].reshape( + self._num_envs, count, dim + ) + + def _mapped_tracked_ids(self, operation: str, body_ids: np.ndarray) -> np.ndarray: + mapping = self._body_id_to_tracked_idx + if mapping is None: + self._unsupported_body_kinematics(operation) + mapped = mapping[np.asarray(body_ids, dtype=np.intp)] + if np.any(mapped < 0): + raise ValueError( + f"mjwarp {operation} received body ids without injected tracking sensors: " + f"{np.asarray(body_ids)[mapped < 0].tolist()}" + ) + return mapped + # ------------------------------------------------------------------ # # Explicit host-cache barriers # # ------------------------------------------------------------------ # @@ -771,27 +860,28 @@ def get_dof_pos(self) -> np.ndarray: def get_dof_vel(self) -> np.ndarray: return self._qvel_cache[:, self._root_qvel_dim :] - def _unsupported_body_kinematics(self, operation: str) -> None: + def _unsupported_body_kinematics(self, operation: str) -> NoReturn: raise NotImplementedError( f"mjwarp host_numpy profile does not expose {operation}; the G1 host adapter " - "supports only base, dof, and configured sensor cache reads." + "supports base, dof, and configured sensor cache reads, plus tracked body " + "kinematics when constructed with body_state_required/add_body_sensors." ) def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: - del body_ids - self._unsupported_body_kinematics("world-frame body positions") + mapped = self._mapped_tracked_ids("world-frame body positions", body_ids) + return self._tracked_pos_w_all[:, mapped, :] def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: - del body_ids - self._unsupported_body_kinematics("world-frame body orientations") + mapped = self._mapped_tracked_ids("world-frame body orientations", body_ids) + return self._tracked_quat_w_all[:, mapped, :] def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: - del body_ids - self._unsupported_body_kinematics("world-frame body linear velocities") + mapped = self._mapped_tracked_ids("world-frame body linear velocities", body_ids) + return self._tracked_linvel_w_all[:, mapped, :] def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: - del body_ids - self._unsupported_body_kinematics("world-frame body angular velocities") + mapped = self._mapped_tracked_ids("world-frame body angular velocities", body_ids) + return self._tracked_angvel_w_all[:, mapped, :] def get_body_pos_b(self, body_ids: np.ndarray) -> np.ndarray: del body_ids @@ -802,12 +892,20 @@ def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: self._unsupported_body_kinematics("base-frame body orientations") def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - del body_ids - self._unsupported_body_kinematics("base-frame body linear velocities") + # Analytical per the SimBackend contract (#1254): world-frame velocity + # rotated into each body's own frame, matching MuJoCoBackend. + mapped = self._mapped_tracked_ids("base-frame body linear velocities", body_ids) + return np_quat_apply_inverse_batched( + self._tracked_quat_w_all[:, mapped, :], + self._tracked_linvel_w_all[:, mapped, :], + ) def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - del body_ids - self._unsupported_body_kinematics("base-frame body angular velocities") + mapped = self._mapped_tracked_ids("base-frame body angular velocities", body_ids) + return np_quat_apply_inverse_batched( + self._tracked_quat_w_all[:, mapped, :], + self._tracked_angvel_w_all[:, mapped, :], + ) def get_sensor_data(self, name: str) -> np.ndarray: try: diff --git a/src/unilab/base/backend/mjwarp/materialization.py b/src/unilab/base/backend/mjwarp/materialization.py index 34b5ed5af..9410bb07c 100644 --- a/src/unilab/base/backend/mjwarp/materialization.py +++ b/src/unilab/base/backend/mjwarp/materialization.py @@ -10,20 +10,21 @@ class _TemporarySceneCleanup: - """Own the one temporary XML created while merging scene fragments.""" + """Own the temporary XMLs created while materializing one scene.""" - def __init__(self, path: str) -> None: - self._path = path + def __init__(self, *paths: str) -> None: + self._paths = paths self._cleaned = False def cleanup(self) -> None: if self._cleaned: return self._cleaned = True - try: - os.remove(self._path) - except FileNotFoundError: - pass + for path in self._paths: + try: + os.remove(path) + except FileNotFoundError: + pass @dataclass(frozen=True) @@ -33,14 +34,24 @@ class MjwarpSceneContext: source_model_file: str diagnostic_model_file: str cleanup_handle: Any | None = None + tracked_body_names: tuple[str, ...] = () -def materialize_mjwarp_scene(scene: SceneCfg) -> MjwarpSceneContext: +def materialize_mjwarp_scene( + scene: SceneCfg, + *, + add_body_sensors: bool = False, +) -> MjwarpSceneContext: """Resolve a flat/fragments scene before CUDA model upload. Height-field terrain construction is intentionally rejected in the first correctness profile. The rejection happens before model upload so an unsupported owner cannot silently fall back to a different terrain path. + + When ``add_body_sensors`` is set, world-frame body tracking sensors are + injected into the resolved model on the cold path (same helper as the + MuJoCo backend) so the host profile can serve body kinematics from its + per-step sensor cache without extra device transfers. """ if scene is None or not scene.model_file: raise ValueError("MjwarpBackend requires SceneCfg.model_file") @@ -49,22 +60,33 @@ def materialize_mjwarp_scene(scene: SceneCfg) -> MjwarpSceneContext: "mjwarp host_numpy profile does not support generated terrain or height-field " "scanners; select a flat owner YAML or a backend with terrain support." ) + temp_paths: list[str] = [] if not scene.fragment_files: - return MjwarpSceneContext( - source_model_file=str(scene.model_file), - diagnostic_model_file=str(scene.model_file), + source_model_file = str(scene.model_file) + else: + # This is intentionally in a cold-path-only module. The shared XML + # composition helper is not a sibling runtime backend dependency. + from unilab.base.backend.mujoco.xml import materialize_scene_fragments + + source_model_file = materialize_scene_fragments( + str(scene.model_file), + fragment_files=scene.fragment_files, ) + temp_paths.append(source_model_file) - # This is intentionally in a cold-path-only module. The shared XML - # composition helper is not a sibling runtime backend dependency. - from unilab.base.backend.mujoco.xml import materialize_scene_fragments + tracked_body_names: tuple[str, ...] = () + if add_body_sensors: + from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors + + source_model_file, _tracked_body_ids, valid_bnames = inject_mujoco_tracking_sensors( + source_model_file + ) + temp_paths.append(source_model_file) + tracked_body_names = tuple(valid_bnames) - materialized = materialize_scene_fragments( - str(scene.model_file), - fragment_files=scene.fragment_files, - ) return MjwarpSceneContext( - source_model_file=materialized, + source_model_file=source_model_file, diagnostic_model_file=str(scene.model_file), - cleanup_handle=_TemporarySceneCleanup(materialized), + cleanup_handle=_TemporarySceneCleanup(*temp_paths) if temp_paths else None, + tracked_body_names=tracked_body_names, ) diff --git a/tests/base/test_mjwarp_backend.py b/tests/base/test_mjwarp_backend.py index fb87e62a2..4107a8ca2 100644 --- a/tests/base/test_mjwarp_backend.py +++ b/tests/base/test_mjwarp_backend.py @@ -138,3 +138,95 @@ def test_g1_walk_flat_owner_one_step( assert np.isfinite(state.obs["obs"]).all() assert np.isfinite(state.obs["critic"]).all() assert np.isfinite(state.reward).all() + + +def test_body_state_matches_mujoco_backend() -> None: + """Tracked body kinematics on mjwarp match the MuJoCo backend on identical state.""" + import mujoco + + from unilab.base.backend.mujoco.backend import MuJoCoBackend + + _require_cuda_mjwarp() + num_envs = 2 + sim_dt = 0.02 / 3.0 + scene = _scene() + mujoco_backend = MuJoCoBackend( + scene, + num_envs, + sim_dt, + base_name="pelvis", + add_body_sensors=True, + ) + mujoco_backend.materialize() + mjwarp_backend = create_backend( + "mjwarp", + scene, + num_envs, + sim_dt, + base_name="pelvis", + body_state_required=True, + ) + + rng = np.random.default_rng(0) + qpos = np.tile(mjwarp_backend.get_keyframe_qpos("stand"), (num_envs, 1)) + qvel = rng.uniform(-0.5, 0.5, size=(num_envs, int(mujoco_backend.model.nv))).astype(np.float32) + rows = np.arange(num_envs, dtype=np.int32) + mujoco_backend.set_state(rows, qpos, qvel) + mjwarp_backend.set_state(rows, qpos, qvel) + + body_names = ["pelvis", "left_hip_pitch_link", "left_knee_link"] + body_ids = np.asarray( + [ + mujoco.mj_name2id(mujoco_backend.model, mujoco.mjtObj.mjOBJ_BODY, name) + for name in body_names + ], + dtype=np.int32, + ) + assert (body_ids > 0).all() + + atol = 2e-3 + np.testing.assert_allclose( + mjwarp_backend.get_body_pos_w(body_ids), + mujoco_backend.get_body_pos_w(body_ids), + atol=atol, + ) + np.testing.assert_allclose( + mjwarp_backend.get_body_quat_w(body_ids), + mujoco_backend.get_body_quat_w(body_ids), + atol=atol, + ) + np.testing.assert_allclose( + mjwarp_backend.get_body_lin_vel_w(body_ids), + mujoco_backend.get_body_lin_vel_w(body_ids), + atol=atol, + ) + np.testing.assert_allclose( + mjwarp_backend.get_body_ang_vel_w(body_ids), + mujoco_backend.get_body_ang_vel_w(body_ids), + atol=atol, + ) + np.testing.assert_allclose( + mjwarp_backend.get_body_lin_vel_b(body_ids), + mujoco_backend.get_body_lin_vel_b(body_ids), + atol=atol, + ) + np.testing.assert_allclose( + mjwarp_backend.get_body_ang_vel_b(body_ids), + mujoco_backend.get_body_ang_vel_b(body_ids), + atol=atol, + ) + + +def test_body_state_untracked_body_ids_fail_closed() -> None: + """Body ids outside the injected tracking set raise instead of wrapping around.""" + _require_cuda_mjwarp() + backend = create_backend( + "mjwarp", + _scene(), + 1, + 0.02 / 3.0, + base_name="pelvis", + body_state_required=True, + ) + with pytest.raises(ValueError, match="without injected tracking sensors"): + backend.get_body_pos_w(np.asarray([0], dtype=np.int32)) diff --git a/tests/base/test_motrix_backend_options.py b/tests/base/test_motrix_backend_options.py index aeefa158c..d964b96a5 100644 --- a/tests/base/test_motrix_backend_options.py +++ b/tests/base/test_motrix_backend_options.py @@ -828,7 +828,37 @@ def __init__(self, scene: SceneCfg, num_envs: int, sim_dt: float, **kwargs: Any) assert captured["kwargs"]["add_body_sensors"] is True -@pytest.mark.parametrize("backend_type", ["drake", "mjwarp"]) +@pytest.mark.parametrize("body_state_required", [False, True]) +def test_create_backend_maps_body_state_request_inside_mjwarp_adapter( + monkeypatch, body_state_required: bool +) -> None: + import unilab.base.backend as backend_factory + from unilab.base.scene import SceneCfg + + captured: dict[str, Any] = {} + + class FakeMjwarpBackend: + def __init__(self, scene: SceneCfg, num_envs: int, sim_dt: float, **kwargs: Any) -> None: + captured["kwargs"] = kwargs + + monkeypatch.setattr(backend_factory, "_load_mjwarp_backend", lambda: FakeMjwarpBackend) + + backend_factory.create_backend( + "mjwarp", + SceneCfg(model_file="model.xml"), + num_envs=1, + sim_dt=0.01, + body_state_required=body_state_required, + ) + + assert "body_state_required" not in captured["kwargs"] + if body_state_required: + assert captured["kwargs"]["add_body_sensors"] is True + else: + assert "add_body_sensors" not in captured["kwargs"] + + +@pytest.mark.parametrize("backend_type", ["drake"]) def test_create_backend_keeps_body_state_request_out_of_native_state_adapters( monkeypatch, backend_type: str ) -> None: