feat(simulation): opt-in ground-truth object pose stream for eval scoring

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

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()