From acdb6f5a0e6b70e420a327d7e529d00c12c1bd41 Mon Sep 17 00:00:00 2001 From: Jerrybery Date: Mon, 7 Sep 2026 13:55:22 +0800 Subject: [PATCH] feat(simulation): opt-in ground-truth object pose stream for eval scoring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MujocoSimModule gains publish_ground_truth (default off) and ground_truth_hz (20 Hz). When enabled, world poses of every free-joint scene body (robot root excluded — odom already covers it) publish on a new gt_object_poses Out[PoseStamped] stream, frame_id = body name. Ground truth is privileged scoring data for the evals framework (issue #3594): never consumed by the agent. Part 1 of 2; part 2 adds the eval-side gt_store() and predicate library. --- dimos/simulation/engines/mujoco_sim_module.py | 73 ++++++++++++ .../engines/test_mujoco_sim_module.py | 110 +++++++++++++++++- 2 files changed, 182 insertions(+), 1 deletion(-) diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 87ff5a0a21..e90dde801b 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -18,6 +18,8 @@ - camera streams (Out ports), replacing ``MujocoCamera`` - joint state via shared memory, consumed by ``ShmMujocoAdapter`` inside ``ControlCoordinator`` +- optionally, ground-truth world poses of free-joint scene bodies + (``publish_ground_truth``), for eval scoring only This avoids the prior pattern of sharing engines via a global in-process registry, which was fragile when ``WorkerManager`` places the adapter and @@ -84,6 +86,29 @@ def _find_sensor_slice(model: mujoco.MjModel, *names: str, dim: int = 3) -> slic _RX180 = R.from_euler("x", 180, degrees=True) +_MJJNT_FREE = int(mujoco.mjtJoint.mjJNT_FREE) # type: ignore[attr-defined] + + +def _resolve_gt_bodies(model: mujoco.MjModel, root_qpos_adr: int | None) -> list[tuple[str, int]]: + """(body name, qpos adr) for every free-joint body except the robot root. + + The robot root pose already flows on ``odom``; ground truth covers the + remaining scene bodies so eval scorers can check physical outcomes. + """ + bodies: list[tuple[str, int]] = [] + for joint_id in range(model.njnt): + if int(model.jnt_type[joint_id]) != _MJJNT_FREE: + continue + qpos_adr = int(model.jnt_qposadr[joint_id]) + if root_qpos_adr is not None and qpos_adr == root_qpos_adr: + continue + body_id = int(model.jnt_bodyid[joint_id]) + raw_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id) or f"body_{body_id}" + # Attached submodels get a leading "/" namespace separator on newer + # MuJoCo; strip it so frame_id stays a stable plain body name. + bodies.append((raw_name.lstrip("/"), qpos_adr)) + return bodies + def _pose_matrix( position: NDArray[np.float64], rotation: NDArray[np.float64] @@ -251,6 +276,11 @@ class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): spawn_z: float | None = None spawn_yaw: float | None = None reset_joint_positions: list[float] | None = None + # Opt-in privileged ground truth for eval scoring: world poses of all + # free-joint scene bodies (robot root excluded, odom covers it). Never + # consumed by the agent. Off = zero overhead, zero behavior change. + publish_ground_truth: bool = False + ground_truth_hz: float = 20.0 headless: bool = False dof: int = 7 @@ -332,6 +362,10 @@ class MujocoSimModule( # root. Published every step; consumers like the viser viewer use # this to translate the robot in world space. odom: Out[PoseStamped] + # Ground-truth world poses of free-joint scene bodies, throttled to + # config.ground_truth_hz. Only publishes when config.publish_ground_truth + # is set; frame_id carries the MuJoCo body name. Eval scoring only. + gt_object_poses: Out[PoseStamped] tf: Out[TFMessage] def __init__(self, **kwargs: Any) -> None: @@ -360,6 +394,11 @@ def __init__(self, **kwargs: Any) -> None: self._root_base_qpos_adr: int | None = None self._root_spawn_clearance_z: float | None = None + # (body name, qpos adr) pairs for ground-truth publishing, resolved + # once at start. Stays empty unless config.publish_ground_truth. + self._gt_bodies: list[tuple[str, int]] = [] + self._gt_last_publish_monotonic = 0.0 + @property def _camera_link(self) -> str: return f"{self.config.camera_name}_link" @@ -551,6 +590,14 @@ def add_camera( self._imu_base_qpos_slice = None self._root_spawn_clearance_z = self._compute_root_spawn_clearance_z() + if self.config.publish_ground_truth: + self._gt_bodies = _resolve_gt_bodies(self._engine.model, self._root_base_qpos_adr) + logger.info( + "MujocoSimModule: ground-truth publishing enabled", + bodies=[name for name, _ in self._gt_bodies], + hz=self.config.ground_truth_hz, + ) + # Wire SHM bridge hooks. self._sim_hooks = _WholeBodySimHooks( self._shm, @@ -806,6 +853,32 @@ def _publish_shm_and_lcm(self, engine: MujocoEngine) -> None: ) ) + # Ground truth - throttled; ``_gt_bodies`` is empty unless + # config.publish_ground_truth, so the disabled path costs one check. + if self._gt_bodies: + now = time.monotonic() + if now - self._gt_last_publish_monotonic >= 1.0 / self.config.ground_truth_hz: + self._gt_last_publish_monotonic = now + ts = time.time() + for name, qpos_adr in self._gt_bodies: + body_pos = data.qpos[qpos_adr : qpos_adr + 3] + body_quat = data.qpos[qpos_adr + 3 : qpos_adr + 7] # (w, x, y, z) + self.gt_object_poses.publish( + PoseStamped( + ts=ts, + frame_id=name, + position=Vector3( + float(body_pos[0]), float(body_pos[1]), float(body_pos[2]) + ), + orientation=Quaternion( + float(body_quat[1]), + float(body_quat[2]), + float(body_quat[3]), + float(body_quat[0]), + ), # PoseStamped uses x,y,z,w + ) + ) + # IMU - only if MJCF declared the sensors. if ( self._imu_quat_slice is None diff --git a/dimos/simulation/engines/test_mujoco_sim_module.py b/dimos/simulation/engines/test_mujoco_sim_module.py index 3fa20048b5..470d3a5cc5 100644 --- a/dimos/simulation/engines/test_mujoco_sim_module.py +++ b/dimos/simulation/engines/test_mujoco_sim_module.py @@ -26,7 +26,11 @@ from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.simulation.engines.mujoco_engine import CameraFrame, MujocoEngine -from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule, MujocoSimModuleConfig +from dimos.simulation.engines.mujoco_sim_module import ( + MujocoSimModule, + MujocoSimModuleConfig, + _resolve_gt_bodies, +) class _FakeData: @@ -564,3 +568,107 @@ def test_publish_loop_pacing_is_independent_of_frame_timestamp_magnitude(base_ts assert elapsed >= (len(frame_ts) - 1) / fps finally: module.stop() + + +def _make_post_step_module(config: MujocoSimModuleConfig) -> MujocoSimModule: + """Module with just enough state for _publish_shm_and_lcm().""" + module = MujocoSimModule() + module.config = config + module._root_base_qpos_adr = None + module._imu_quat_slice = None + module._imu_base_qpos_slice = None + module._imu_gyro_slice = None + module._imu_accel_slice = None + + class _FakeShm: + def signal_ready(self, *, num_joints: int, arm_joints: int) -> None: + pass + + def signal_stop(self) -> None: + pass + + def cleanup(self) -> None: + pass + + module._shm = _FakeShm() + return module + + +def test_gt_poses_not_published_when_disabled() -> None: + module = _make_post_step_module(MujocoSimModuleConfig(dof=2)) + try: + assert module.config.publish_ground_truth is False + published: list[Any] = [] + module.gt_object_poses.subscribe(published.append) + + module._publish_shm_and_lcm(_FakeEngine) + + assert published == [] + assert module._gt_bodies == [] + finally: + module.stop() + + +def test_gt_poses_published_and_throttled_when_enabled() -> None: + module = _make_post_step_module( + MujocoSimModuleConfig(dof=2, publish_ground_truth=True, ground_truth_hz=20.0) + ) + try: + module._gt_bodies = [("box_0", 0)] + published: list[Any] = [] + module.gt_object_poses.subscribe(published.append) + + module._publish_shm_and_lcm(_FakeEngine) + + assert len(published) == 1 + pose = published[0] + assert pose.frame_id == "box_0" + assert pose.position.to_numpy() == pytest.approx([0.0, 0.0, 0.75]) + assert ( + pose.orientation.x, + pose.orientation.y, + pose.orientation.z, + pose.orientation.w, + ) == pytest.approx((0.0, 0.0, 0.0, 1.0)) + + # Second step inside the 50ms throttle window is dropped. + module._publish_shm_and_lcm(_FakeEngine) + assert len(published) == 1 + + # Once the interval elapses, the next step publishes again. + module._gt_last_publish_monotonic -= 1.0 + module._publish_shm_and_lcm(_FakeEngine) + assert len(published) == 2 + finally: + module.stop() + + +@pytest.mark.mujoco +def test_resolve_gt_bodies_excludes_robot_root(tmp_path: Path) -> None: + scene_xml = tmp_path / "scene.xml" + robot_xml = tmp_path / "robot.xml" + _write_scene_xml(scene_xml) + _write_robot_xml(robot_xml) + + chair_001 = _scene_entity("chair_001") + chair_001["initial_pose"]["x"] = -1.0 # type: ignore[index] + + module = MujocoSimModule( + scene_xml=scene_xml, + robot_mjcf=robot_xml, + scene_entities=[_scene_entity("chair_000"), chair_001], + ) + try: + model = module._compose_model() + + bodies = _resolve_gt_bodies(model, root_qpos_adr=0) + + assert [name for name, _ in bodies] == ["entity:chair_000", "entity:chair_001"] + for _, qpos_adr in bodies: + assert qpos_adr > 0 + assert _resolve_gt_bodies(model, root_qpos_adr=None) == [ + ("base", 0), + *bodies, + ] + finally: + module.stop()