From f74dcd93c1a04e5ea11e446416cf40e00c1a270d Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 24 Mar 2026 11:41:42 -0700 Subject: [PATCH 01/41] initial commit - mujoco camera rendering --- dimos/hardware/manipulators/sim/adapter.py | 10 +- dimos/manipulation/blueprints.py | 62 +++ dimos/manipulation/planning/spec/config.py | 7 +- dimos/msgs/sensor_msgs/PointCloud2.py | 3 + dimos/perception/object_scene_registration.py | 3 +- dimos/robot/all_blueprints.py | 1 + dimos/simulation/engines/mujoco_engine.py | 165 ++++++- dimos/simulation/mujoco/camera.py | 463 ++++++++++++++++++ docs/plans/mujoco_camera_plan.md | 250 ++++++++++ 9 files changed, 954 insertions(+), 10 deletions(-) create mode 100644 dimos/simulation/mujoco/camera.py create mode 100644 docs/plans/mujoco_camera_plan.md diff --git a/dimos/hardware/manipulators/sim/adapter.py b/dimos/hardware/manipulators/sim/adapter.py index 3979ce98c5..39cf6ffc5e 100644 --- a/dimos/hardware/manipulators/sim/adapter.py +++ b/dimos/hardware/manipulators/sim/adapter.py @@ -23,7 +23,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from dimos.simulation.engines.mujoco_engine import MujocoEngine +from dimos.simulation.engines.mujoco_engine import MujocoEngine, get_or_create_engine from dimos.simulation.manipulators.sim_manip_interface import SimManipInterface if TYPE_CHECKING: @@ -41,11 +41,13 @@ def __init__( dof: int = 7, address: str | None = None, headless: bool = True, + engine: MujocoEngine | None = None, **_: Any, ) -> None: - if address is None: - raise ValueError("address (MJCF XML path) is required for sim_mujoco adapter") - engine = MujocoEngine(config_path=Path(address), headless=headless) + if engine is None: + if address is None: + raise ValueError("address (MJCF XML path) is required for sim_mujoco adapter") + engine = get_or_create_engine(config_path=Path(address), headless=headless) # Detect gripper from engine joints gripper_idx = None diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 8110166042..cf0d76fab9 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -456,6 +456,67 @@ def _make_piper_config( ) +# --------------------------------------------------------------------------- +# Sim perception: MuJoCo camera replaces RealSense for sim-based pick-and-place +# --------------------------------------------------------------------------- +# Both the sim adapter and camera resolve the same MujocoEngine via the registry +# (keyed by MJCF path), so they share physics state. +# The engine is created lazily when the adapter connects / camera starts. + +from dimos.control.blueprints._hardware import XARM7_SIM_PATH, sim_xarm7 +from dimos.simulation.mujoco.camera import MujocoCamera +from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode + +xarm_perception_sim = ( + autoconnect( + PickAndPlaceModule.blueprint( + robots=[ + _make_xarm7_config( + "arm", + joint_prefix="arm_", + coordinator_task="traj_arm", + add_gripper=True, + gripper_hardware_id="arm", + tf_extra_links=["link7"], + ), + ], + planning_timeout=10.0, + enable_viz=True, + ), + MujocoCamera.blueprint( + address=str(XARM7_SIM_PATH), + camera_name="wrist_camera", + base_frame_id="link7", + ), + ObjectSceneRegistrationModule.blueprint(target_frame="world"), + ControlCoordinator.blueprint( + tick_rate=100.0, + publish_joint_state=True, + joint_state_frame_id="coordinator", + hardware=[sim_xarm7("arm", headless=False, gripper=True)], + tasks=[ + TaskConfig( + name="traj_arm", + type="trajectory", + joint_names=[f"arm_joint{i + 1}" for i in range(7)], + priority=10, + ), + ], + ), + RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode()), + FoxgloveBridge.blueprint(), + ) + .transports( + { + ("joint_state", JointState): LCMTransport( + "/coordinator/joint_state", JointState + ), + } + ) + .global_config(viewer="foxglove") +) + + __all__ = [ "PIPER_GRIPPER_COLLISION_EXCLUSIONS", "XARM_GRIPPER_COLLISION_EXCLUSIONS", @@ -465,4 +526,5 @@ def _make_piper_config( "xarm7_planner_coordinator_agent", "xarm_perception", "xarm_perception_agent", + "xarm_perception_sim", ] diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index 80cf248f08..f8dcb8b548 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -16,7 +16,6 @@ from __future__ import annotations -from collections.abc import Iterable, Sequence from pathlib import Path from pydantic import Field @@ -65,7 +64,7 @@ class RobotModelConfig(ModuleConfig): velocity_limits: list[float] | None = None auto_convert_meshes: bool = False xacro_args: dict[str, str] = Field(default_factory=dict) - collision_exclusion_pairs: Iterable[tuple[str, str]] = () + collision_exclusion_pairs: list[tuple[str, str]] = Field(default_factory=list) # Motion constraints for trajectory generation max_velocity: float = 1.0 max_acceleration: float = 2.0 @@ -74,9 +73,9 @@ class RobotModelConfig(ModuleConfig): coordinator_task_name: str | None = None gripper_hardware_id: str | None = None # TF publishing for extra links (e.g., camera mount) - tf_extra_links: Sequence[str] = () + tf_extra_links: list[str] = Field(default_factory=list) # Home/observe joint configuration for go_home skill - home_joints: Iterable[float] | None = None + home_joints: list[float] | None = None # Pre-grasp offset distance in meters (along approach direction) pre_grasp_offset: float = 0.10 diff --git a/dimos/msgs/sensor_msgs/PointCloud2.py b/dimos/msgs/sensor_msgs/PointCloud2.py index 67af1c5ac3..f9c3874119 100644 --- a/dimos/msgs/sensor_msgs/PointCloud2.py +++ b/dimos/msgs/sensor_msgs/PointCloud2.py @@ -100,6 +100,9 @@ def __getstate__(self) -> dict[str, object]: # Remove non-picklable objects del state["_pcd_tensor"] state["_pcd_legacy_cache"] = None + # Remove cached_property entries that hold Open3D C++ objects + for key in ("axis_aligned_bounding_box", "oriented_bounding_box", "bounding_box_dimensions"): + state.pop(key, None) return state def __setstate__(self, state: dict[str, object]) -> None: diff --git a/dimos/perception/object_scene_registration.py b/dimos/perception/object_scene_registration.py index 3be2db4b47..5660fcbccb 100644 --- a/dimos/perception/object_scene_registration.py +++ b/dimos/perception/object_scene_registration.py @@ -70,8 +70,9 @@ def __init__( self, target_frame: str = "map", prompt_mode: YoloePromptMode = YoloePromptMode.LRPC, + **kwargs, # type: ignore[no-untyped-def] ) -> None: - super().__init__() + super().__init__(**kwargs) self._target_frame = target_frame self._prompt_mode = prompt_mode self._object_db = ObjectDB() diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 03a1f47f2a..74b7b87da0 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -91,6 +91,7 @@ "unitree-go2-webrtc-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_webrtc_keyboard_teleop:unitree_go2_webrtc_keyboard_teleop", "xarm-perception": "dimos.manipulation.blueprints:xarm_perception", "xarm-perception-agent": "dimos.manipulation.blueprints:xarm_perception_agent", + "xarm-perception-sim": "dimos.manipulation.blueprints:xarm_perception_sim", "xarm6-planner-only": "dimos.manipulation.blueprints:xarm6_planner_only", "xarm7-planner-coordinator": "dimos.manipulation.blueprints:xarm7_planner_coordinator", "xarm7-planner-coordinator-agent": "dimos.manipulation.blueprints:xarm7_planner_coordinator_agent", diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index df8359746a..9bddf0f838 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -16,12 +16,15 @@ from __future__ import annotations +import dataclasses import threading import time from typing import TYPE_CHECKING import mujoco import mujoco.viewer as viewer # type: ignore[import-untyped,import-not-found] +import numpy as np +from numpy.typing import NDArray from dimos.simulation.engines.base import SimulationEngine from dimos.simulation.utils.xml_parser import JointMapping, build_joint_mappings @@ -35,6 +38,74 @@ logger = setup_logger() +@dataclasses.dataclass +class CameraConfig: + """Configuration for a camera to be rendered inside the sim loop.""" + + name: str + """Camera name as defined in the MJCF XML.""" + width: int = 640 + height: int = 480 + fps: float = 15.0 + + +@dataclasses.dataclass +class CameraFrame: + """Thread-safe container for a rendered camera frame.""" + + rgb: NDArray[np.uint8] + """RGB image, shape (H, W, 3).""" + depth: NDArray[np.float32] + """Depth image in meters, shape (H, W).""" + cam_pos: NDArray[np.float64] + """Camera world position, shape (3,).""" + cam_mat: NDArray[np.float64] + """Camera rotation matrix (flattened 3x3), shape (9,).""" + fovy: float + """Vertical field of view in degrees.""" + timestamp: float + """Time of rendering.""" + + +# --------------------------------------------------------------------------- +# Engine registry — allows adapter and camera to share the same engine instance +# by MJCF path. +# --------------------------------------------------------------------------- +_engine_registry: dict[str, MujocoEngine] = {} +_engine_registry_lock = threading.Lock() + + +def get_or_create_engine( + config_path: "Path", + headless: bool = True, + cameras: list[CameraConfig] | None = None, +) -> MujocoEngine: + """Return the shared MujocoEngine for *config_path*, creating one if needed. + + If an engine already exists for the resolved path and *cameras* are supplied, + any **new** camera configs are appended (idempotent by name). + """ + from pathlib import Path + + key = str(Path(config_path).resolve()) + with _engine_registry_lock: + if key in _engine_registry: + engine = _engine_registry[key] + # Merge new camera configs (by name) + if cameras: + existing_names = {c.name for c in engine._camera_configs} + for cam in cameras: + if cam.name not in existing_names: + engine._camera_configs.append(cam) + return engine + + engine = MujocoEngine( + config_path=Path(config_path), headless=headless, cameras=cameras + ) + _engine_registry[key] = engine + return engine + + class MujocoEngine(SimulationEngine): """ MuJoCo simulation engine. @@ -44,7 +115,12 @@ class MujocoEngine(SimulationEngine): - applies control commands """ - def __init__(self, config_path: Path, headless: bool) -> None: + def __init__( + self, + config_path: Path, + headless: bool, + cameras: list[CameraConfig] | None = None, + ) -> None: super().__init__(config_path=config_path, headless=headless) xml_path = self._resolve_xml_path(config_path) @@ -76,6 +152,11 @@ def __init__(self, config_path: Path, headless: bool) -> None: self._joint_position_targets[i] = current_pos self._joint_positions[i] = current_pos + # Camera rendering state (renderers created in sim thread) + self._camera_configs = cameras or [] + self._camera_frames: dict[str, CameraFrame] = {} + self._camera_lock = threading.Lock() + def _resolve_xml_path(self, config_path: Path) -> Path: if config_path is None: raise ValueError("config_path is required for MuJoCo simulation loading") @@ -177,6 +258,64 @@ def _sim_loop(self) -> None: logger.info(f"{self.__class__.__name__}: sim loop started") dt = 1.0 / self._control_frequency + # Camera renderers — created in sim thread (MuJoCo thread-safety). + # Checked each tick so cameras added after connect() are picked up. + cam_renderers: dict[ + str, tuple[CameraConfig, int, mujoco.Renderer, mujoco.Renderer, float, float] + ] = {} # name -> (cfg, cam_id, rgb_r, depth_r, interval, last_render_time) + + def _init_new_cameras() -> None: + """Create renderers for any camera configs not yet initialized.""" + for cfg in self._camera_configs: + if cfg.name in cam_renderers: + continue + cam_id = mujoco.mj_name2id( + self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name + ) + if cam_id < 0: + logger.warning(f"Camera '{cfg.name}' not found in MJCF, skipping") + continue + rgb_renderer = mujoco.Renderer( + self._model, height=cfg.height, width=cfg.width + ) + depth_renderer = mujoco.Renderer( + self._model, height=cfg.height, width=cfg.width + ) + depth_renderer.enable_depth_rendering() + interval = 1.0 / cfg.fps if cfg.fps > 0 else float("inf") + cam_renderers[cfg.name] = ( + cfg, cam_id, rgb_renderer, depth_renderer, interval, 0.0, + ) + logger.info( + f"Camera '{cfg.name}' renderer created " + f"({cfg.width}x{cfg.height} @ {cfg.fps}fps)" + ) + + def _render_cameras(now: float) -> None: + _init_new_cameras() + for name in cam_renderers: + cfg, cam_id, rgb_r, depth_r, interval, last_t = cam_renderers[name] + if now - last_t < interval: + continue + cam_renderers[name] = (cfg, cam_id, rgb_r, depth_r, interval, now) + + rgb_r.update_scene(self._data, camera=cam_id) + rgb = rgb_r.render().copy() + + depth_r.update_scene(self._data, camera=cam_id) + depth = depth_r.render().copy() + + frame = CameraFrame( + rgb=rgb, + depth=depth.astype(np.float32), + cam_pos=self._data.cam_xpos[cam_id].copy(), + cam_mat=self._data.cam_xmat[cam_id].copy(), + fovy=float(self._model.cam_fovy[cam_id]), + timestamp=now, + ) + with self._camera_lock: + self._camera_frames[cfg.name] = frame + def _step_once(sync_viewer: bool) -> None: loop_start = time.time() self._apply_control() @@ -184,6 +323,7 @@ def _step_once(sync_viewer: bool) -> None: if sync_viewer: m_viewer.sync() self._update_joint_state() + _render_cameras(loop_start) elapsed = time.time() - loop_start sleep_time = dt - elapsed @@ -200,6 +340,11 @@ def _step_once(sync_viewer: bool) -> None: while m_viewer.is_running() and not self._stop_event.is_set(): _step_once(sync_viewer=True) + # Clean up renderers + for _cfg, _cam_id, rgb_r, depth_r, _interval, _t in cam_renderers.values(): + rgb_r.close() + depth_r.close() + logger.info(f"{self.__class__.__name__}: sim loop stopped") @property @@ -327,7 +472,25 @@ def get_joint_range(self, joint_index: int) -> tuple[float, float] | None: ) return None + def read_camera(self, camera_name: str) -> CameraFrame | None: + """Read the latest rendered frame for a camera (thread-safe). + + Returns None if the camera hasn't rendered yet or doesn't exist. + """ + with self._camera_lock: + return self._camera_frames.get(camera_name) + + def get_camera_fovy(self, camera_name: str) -> float | None: + """Get vertical field of view for a named camera, in degrees.""" + cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, camera_name) + if cam_id < 0: + return None + return float(self._model.cam_fovy[cam_id]) + __all__ = [ + "CameraConfig", + "CameraFrame", "MujocoEngine", + "get_or_create_engine", ] diff --git a/dimos/simulation/mujoco/camera.py b/dimos/simulation/mujoco/camera.py new file mode 100644 index 0000000000..5526c33202 --- /dev/null +++ b/dimos/simulation/mujoco/camera.py @@ -0,0 +1,463 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MuJoCo simulation camera module. + +Drop-in replacement for RealSenseCamera in manipulation blueprints. +Renders RGB + depth from a named camera in a shared MujocoEngine. +""" + +from __future__ import annotations + +import math +import threading +import time + +import numpy as np +import reactivex as rx +from pydantic import Field + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import Out +from dimos.hardware.sensors.camera.spec import ( + OPTICAL_ROTATION, + DepthCameraConfig, + DepthCameraHardware, +) +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.simulation.engines.mujoco_engine import ( + CameraConfig, + CameraFrame, + MujocoEngine, + get_or_create_engine, +) +from dimos.spec import perception +from dimos.utils.logging_config import setup_logger +from dimos.utils.reactive import backpressure + +logger = setup_logger() + + +def _default_identity_transform() -> Transform: + return Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + ) + + +class MujocoCameraConfig(ModuleConfig, DepthCameraConfig): + """Configuration for the MuJoCo simulation camera.""" + + address: str = "" + """MJCF XML path — same as the sim adapter's address. + Used to look up the shared MujocoEngine from the registry.""" + headless: bool = False + """Whether to run MuJoCo without a viewer window.""" + camera_name: str = "wrist_camera" + """Camera name as defined in the MJCF XML.""" + width: int = 640 + height: int = 480 + fps: int = 15 + base_frame_id: str = "link7" + base_transform: Transform | None = Field(default_factory=_default_identity_transform) + # MuJoCo renders color+depth from same virtual camera — alignment is a no-op + align_depth_to_color: bool = True + enable_depth: bool = True + enable_pointcloud: bool = False + pointcloud_fps: float = 5.0 + camera_info_fps: float = 1.0 + + +class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.DepthCamera): + """Simulated depth camera that renders from a MujocoEngine. + + Implements the same output ports and TF chain as RealSenseCamera so it can + be used as a drop-in replacement in manipulation blueprints. + + The engine is resolved automatically from the registry via ``address`` + (the same MJCF path used by the sim_mujoco adapter). Alternatively, + call ``set_engine()`` before ``start()``. + """ + + color_image: Out[Image] + depth_image: Out[Image] + pointcloud: Out[PointCloud2] + camera_info: Out[CameraInfo] + depth_camera_info: Out[CameraInfo] + + default_config = MujocoCameraConfig + + def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] + super().__init__(*args, **kwargs) + self._engine: MujocoEngine | None = None + self._running = False + self._thread: threading.Thread | None = None + self._color_camera_info: CameraInfo | None = None + self._depth_camera_info: CameraInfo | None = None + + # -- Engine injection (called before start) -- + + def set_engine(self, engine: MujocoEngine) -> None: + """Inject the shared MujocoEngine reference.""" + self._engine = engine + + # -- DepthCameraHardware interface -- + + @property + def _camera_link(self) -> str: + return f"{self.config.camera_name}_link" + + @property + def _color_frame(self) -> str: + return f"{self.config.camera_name}_color_frame" + + @property + def _color_optical_frame(self) -> str: + return f"{self.config.camera_name}_color_optical_frame" + + @property + def _depth_frame(self) -> str: + return f"{self.config.camera_name}_depth_frame" + + @property + def _depth_optical_frame(self) -> str: + return f"{self.config.camera_name}_depth_optical_frame" + + @rpc + def get_color_camera_info(self) -> CameraInfo | None: + return self._color_camera_info + + @rpc + def get_depth_camera_info(self) -> CameraInfo | None: + return self._depth_camera_info + + @rpc + def get_depth_scale(self) -> float: + # MuJoCo depth is already in meters + return 1.0 + + # -- Intrinsics -- + + def _build_camera_info(self) -> None: + """Compute camera intrinsics from the MuJoCo model (pinhole, no distortion).""" + if self._engine is None: + return + fovy_deg = self._engine.get_camera_fovy(self.config.camera_name) + if fovy_deg is None: + logger.error(f"Camera '{self.config.camera_name}' not found in MJCF") + return + + h = self.config.height + w = self.config.width + fovy_rad = math.radians(fovy_deg) + fy = h / (2.0 * math.tan(fovy_rad / 2.0)) + fx = fy # square pixels + cx = w / 2.0 + cy = h / 2.0 + + K = [fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0] + P = [fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0] + D = [0.0, 0.0, 0.0, 0.0, 0.0] + + info = CameraInfo( + height=h, + width=w, + distortion_model="plumb_bob", + D=D, + K=K, + P=P, + frame_id=self._color_optical_frame, + ) + # Color and depth share the same virtual camera + self._color_camera_info = info + self._depth_camera_info = CameraInfo( + height=h, + width=w, + distortion_model="plumb_bob", + D=D, + K=K, + P=P, + frame_id=self._color_optical_frame, + ) + + # -- Lifecycle -- + + @rpc + def start(self) -> None: + if self._engine is None and self.config.address: + from pathlib import Path + + self._engine = get_or_create_engine( + config_path=Path(self.config.address), + headless=self.config.headless, + cameras=[ + CameraConfig( + name=self.config.camera_name, + width=self.config.width, + height=self.config.height, + fps=self.config.fps, + ) + ], + ) + if self._engine is None: + raise RuntimeError( + "MujocoCamera: set address (MJCF path) in config or call set_engine()" + ) + + logger.info( + f"MujocoCamera: engine resolved, connected={self._engine.connected}, " + f"cameras={[c.name for c in self._engine._camera_configs]}" + ) + + self._build_camera_info() + + self._running = True + self._thread = threading.Thread(target=self._publish_loop, daemon=True) + self._thread.start() + + # Periodic camera_info publishing + interval_sec = 1.0 / self.config.camera_info_fps + self._disposables.add( + rx.interval(interval_sec).subscribe( + on_next=lambda _: self._publish_camera_info(), + on_error=lambda e: logger.error(f"CameraInfo publish error: {e}"), + ) + ) + + # Optional pointcloud generation + if self.config.enable_pointcloud and self.config.enable_depth: + pc_interval = 1.0 / self.config.pointcloud_fps + self._disposables.add( + backpressure(rx.interval(pc_interval)).subscribe( + on_next=lambda _: self._generate_pointcloud(), + on_error=lambda e: logger.error(f"Pointcloud error: {e}"), + ) + ) + + @rpc + def stop(self) -> None: + self._running = False + if self._thread and self._thread.is_alive(): + self._thread.join(timeout=2.0) + self._thread = None + self._color_camera_info = None + self._depth_camera_info = None + super().stop() + + # -- Publishing -- + + def _publish_loop(self) -> None: + """Poll engine for rendered frames and publish at configured FPS.""" + interval = 1.0 / self.config.fps + last_timestamp = 0.0 + published_count = 0 + + logger.info(f"MujocoCamera publish loop started (camera={self.config.camera_name})") + + # Wait for engine to connect (adapter may not have started yet) + while self._running and self._engine is not None and not self._engine.connected: + time.sleep(0.1) + + if not self._running: + return + + logger.info("MujocoCamera: engine connected, polling for frames") + + while self._running and self._engine is not None: + try: + frame = self._engine.read_camera(self.config.camera_name) + if frame is None or frame.timestamp <= last_timestamp: + time.sleep(interval * 0.5) + continue + + last_timestamp = frame.timestamp + ts = time.time() + + # Publish RGB + color_img = Image( + data=frame.rgb, + format=ImageFormat.RGB, + frame_id=self._color_optical_frame, + ts=ts, + ) + self.color_image.publish(color_img) + + # Publish depth (float32 meters) + if self.config.enable_depth: + depth_img = Image( + data=frame.depth, + format=ImageFormat.DEPTH, + frame_id=self._color_optical_frame, + ts=ts, + ) + self.depth_image.publish(depth_img) + + # Publish TF (world -> camera from MuJoCo pose) + self._publish_tf(ts, frame) + + published_count += 1 + if published_count == 1: + logger.info( + f"MujocoCamera: first frame published " + f"(rgb={frame.rgb.shape}, depth={frame.depth.shape})" + ) + + # Rate limit + elapsed = time.time() - ts + sleep_time = interval - elapsed + if sleep_time > 0: + time.sleep(sleep_time) + except Exception as e: + logger.error(f"MujocoCamera publish error: {e}") + time.sleep(1.0) + + def _publish_camera_info(self) -> None: + ts = time.time() + if self._color_camera_info: + self._color_camera_info.ts = ts + self.camera_info.publish(self._color_camera_info) + if self._depth_camera_info: + self._depth_camera_info.ts = ts + self.depth_camera_info.publish(self._depth_camera_info) + + def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: + transforms = [] + + # world -> camera_link (from MuJoCo camera world pose, updated each frame) + if frame is not None: + from scipy.spatial.transform import Rotation as R + + # MuJoCo camera frame: X=right, Y=up, Z=back (looks along -Z) + # ROS camera_link: X=forward, Y=left, Z=up + # OPTICAL_ROTATION converts camera_link -> optical (X=right, Y=down, Z=forward) + # We need: cam_xmat @ correction @ OPTICAL_ROTATION = cam_xmat @ Rx(180) + # So correction = Rx(180) @ OPTICAL_ROTATION.inv() + _MJ_TO_CAMLINK = R.from_quat([0.5, -0.5, -0.5, -0.5]) # xyzw + mj_rot = R.from_matrix(frame.cam_mat.reshape(3, 3)) + q_xyzw = (mj_rot * _MJ_TO_CAMLINK).as_quat() + transforms.append( + Transform( + translation=Vector3( + float(frame.cam_pos[0]), + float(frame.cam_pos[1]), + float(frame.cam_pos[2]), + ), + rotation=Quaternion( + float(q_xyzw[0]), + float(q_xyzw[1]), + float(q_xyzw[2]), + float(q_xyzw[3]), + ), + frame_id="world", + child_frame_id=self._camera_link, + ts=ts, + ) + ) + elif self.config.base_transform is not None: + # Fallback: static base_frame_id -> camera_link + transforms.append( + Transform( + translation=self.config.base_transform.translation, + rotation=self.config.base_transform.rotation, + frame_id=self.config.base_frame_id, + child_frame_id=self._camera_link, + ts=ts, + ) + ) + + # camera_link -> depth_frame (identity — same virtual camera) + transforms.append( + Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id=self._camera_link, + child_frame_id=self._depth_frame, + ts=ts, + ) + ) + + # depth_frame -> depth_optical_frame + transforms.append( + Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=OPTICAL_ROTATION, + frame_id=self._depth_frame, + child_frame_id=self._depth_optical_frame, + ts=ts, + ) + ) + + # camera_link -> color_frame (identity — co-located in sim) + transforms.append( + Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id=self._camera_link, + child_frame_id=self._color_frame, + ts=ts, + ) + ) + + # color_frame -> color_optical_frame + transforms.append( + Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=OPTICAL_ROTATION, + frame_id=self._color_frame, + child_frame_id=self._color_optical_frame, + ts=ts, + ) + ) + + self.tf.publish(*transforms) + + def _generate_pointcloud(self) -> None: + """Generate pointcloud from latest depth frame (optional, for visualization).""" + if self._engine is None or self._color_camera_info is None: + return + frame = self._engine.read_camera(self.config.camera_name) + if frame is None: + return + try: + color_img = Image( + data=frame.rgb, + format=ImageFormat.RGB, + frame_id=self._color_optical_frame, + ts=frame.timestamp, + ) + depth_img = Image( + data=frame.depth, + format=ImageFormat.DEPTH, + frame_id=self._color_optical_frame, + ts=frame.timestamp, + ) + pcd = PointCloud2.from_rgbd( + color_image=color_img, + depth_image=depth_img, + camera_info=self._color_camera_info, + depth_scale=1.0, + ) + pcd = pcd.voxel_downsample(0.005) + self.pointcloud.publish(pcd) + except Exception as e: + logger.error(f"Pointcloud generation error: {e}") + + +__all__ = ["MujocoCamera", "MujocoCameraConfig"] diff --git a/docs/plans/mujoco_camera_plan.md b/docs/plans/mujoco_camera_plan.md new file mode 100644 index 0000000000..f8a6b9b395 --- /dev/null +++ b/docs/plans/mujoco_camera_plan.md @@ -0,0 +1,250 @@ +# MuJoCo Camera Module Plan + +## Goal + +Create a `MujocoCamera` module that can drop in for `RealSenseCamera` in manipulation +blueprints, enabling sim-based pick-and-place without real hardware. + +## Current Wiring (real hardware) + +``` +PickAndPlaceModule + ↑ joint_state (LCM) +RealSenseCamera(base_frame_id="link7", base_transform=...) + ↓ color_image, depth_image, camera_info +ObjectSceneRegistrationModule(target_frame="world") +``` + +`ObjectSceneRegistrationModule` consumes exactly three inputs: +- `color_image: In[Image]` — RGB +- `depth_image: In[Image]` — DEPTH16 or float32 +- `camera_info: In[CameraInfo]` — intrinsics (K matrix) + +## Threading Constraint + +MuJoCo's C API is **not thread-safe**. `mujoco.Renderer` must be created and used in +the **same thread** as `mj_step()`. The existing nav camera (mujoco_process.py) solves +this by rendering inside the sim loop. We must follow the same pattern. + +`MujocoEngine._sim_loop()` currently only does: apply_control → mj_step → sync_viewer → update_state. +Rendering must be added **inside this loop**. + +--- + +## Basic Plan + +### 1. Extend `MujocoEngine` with camera rendering + +**File:** `dimos/simulation/engines/mujoco_engine.py` + +Add rendering capability to the sim loop: + +- New init params: `cameras: list[CameraConfig]` (name, width, height, fps) +- Create `mujoco.Renderer` instances in `_sim_loop()` (same thread as `mj_step`) +- Rate-limit rendering per camera (e.g., 15 FPS vs 500 Hz physics) +- After each render, write RGB + depth arrays into a **locked buffer** +- Expose `read_camera(camera_name) -> (rgb, depth, cam_pos, cam_mat, fovy)` that + returns the latest buffered frame (thread-safe copy) + +``` +_sim_loop: + create renderers (same thread) + while running: + apply_control() + mj_step() + for each camera (if due): + renderer.update_scene(data, camera=cam_id) + rgb = rgb_renderer.render() + depth = depth_renderer.render() + cam_pos = data.cam_xpos[cam_id] + cam_mat = data.cam_xmat[cam_id] + write to locked buffer + update_joint_state() +``` + +### 2. Create `MujocoCamera` module + +**File:** `dimos/simulation/mujoco/camera.py` + +A `Module` implementing `DepthCameraHardware` + `perception.DepthCamera`: + +```python +class MujocoCameraConfig(ModuleConfig, DepthCameraConfig): + camera_name: str = "wrist_camera" # MuJoCo camera name in MJCF + width: int = 640 + height: int = 480 + fps: int = 15 + base_frame_id: str = "link7" + base_transform: Transform | None = None + enable_depth: bool = True + enable_pointcloud: bool = False + pointcloud_fps: float = 5.0 + camera_info_fps: float = 1.0 + depth_scale: float = 1.0 # MuJoCo depth is in meters + +class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.DepthCamera): + color_image: Out[Image] + depth_image: Out[Image] + pointcloud: Out[PointCloud2] + camera_info: Out[CameraInfo] + depth_camera_info: Out[CameraInfo] +``` + +**Lifecycle:** +- `set_engine(engine: MujocoEngine)` — called before `start()`, stores engine ref +- `start()` — compute intrinsics from engine's `model.cam_fovy[cam_id]`, start + polling thread that calls `engine.read_camera()` at configured FPS +- Polling thread: read buffered frame → wrap as `Image` → publish on output ports +- Publish TF: `base_frame_id → camera_link → optical frames` (same chain as RealSense) + +### 3. Refactor `SimMujocoAdapter` to accept external engine + +**File:** `dimos/hardware/manipulators/sim/adapter.py` + +Currently creates `MujocoEngine` internally. Add option to accept an existing one: + +```python +class SimMujocoAdapter(SimManipInterface): + def __init__(self, dof=7, address=None, headless=True, engine=None, **_): + if engine is None: + engine = MujocoEngine(config_path=Path(address), headless=headless) + # ... rest unchanged +``` + +### 4. Blueprint wiring + +**File:** `dimos/manipulation/blueprints.py` + +Create a new `xarm_perception_sim` blueprint: + +```python +# Shared engine created once +_engine = MujocoEngine(config_path=Path("path/to/xarm_scene.xml"), headless=False) + +xarm_perception_sim = autoconnect( + PickAndPlaceModule.blueprint( + robots=[_make_xarm7_config(...)], + ... + ), + MujocoCamera.blueprint( + camera_name="wrist_camera", + base_frame_id="link7", + base_transform=_XARM_PERCEPTION_CAMERA_TRANSFORM, + _engine=_engine, # injected reference + ), + ObjectSceneRegistrationModule.blueprint(target_frame="world"), + ControlCoordinator.blueprint( + hardware=[HardwareComponent(adapter_type="sim_mujoco", ...)], + ... + ), +) +``` + +### 5. MJCF scene file + +Create `data/sim_assets/xarm7/scene.xml` with: +- xArm7 robot model (arm + gripper) +- `wrist_camera` camera attached to link7 (matching the real RealSense mount) +- Table + objects for pick-and-place +- Lighting + +--- + +## Improvements & Simplifications + +### A. Depth format: use meters directly (skip DEPTH16 conversion) + +RealSense outputs DEPTH16 (uint16, mm). ObjectSceneRegistrationModule converts it: +```python +depth_meters = depth_data.astype(np.float32) / 1000.0 +``` + +MuJoCo depth is already in meters (float32). We can publish as `ImageFormat.DEPTH32F` +directly. ObjectSceneRegistrationModule already handles both formats (line 281-289): +```python +if depth_image.format == ImageFormat.DEPTH16: + depth_data = depth_data.astype(np.float32) / 1000.0 +``` + +**Simplification:** No artificial uint16 conversion needed. Just publish float32 depth. + +### B. Skip extrinsics — color and depth are co-located + +RealSense has physically separate color/depth sensors with extrinsics between them. +MuJoCo renders both from the same virtual camera. So: + +- Color and depth share the same intrinsics +- No color-to-depth alignment needed (`align_depth_to_color` is a no-op) +- TF chain simplifies: `camera_link = color_frame = depth_frame` (single camera origin) +- `depth_camera_info` = `color_camera_info` (identical) + +### C. Intrinsics from model — no calibration + +RealSense reads intrinsics from hardware via `pyrealsense2`. MuJoCo camera params +are known exactly from the model: + +```python +fovy = model.cam_fovy[cam_id] # vertical FOV in degrees +fy = height / (2 * tan(radians(fovy) / 2)) +fx = fy # square pixels +cx, cy = width / 2, height / 2 +D = [0, 0, 0, 0, 0] # no distortion +``` + +Compute once at `start()`, publish periodically. No runtime calibration. + +### D. Consider rendering in `mj_forward()` instead of `mj_step()` + +For camera-only updates (no physics needed), `mj_forward()` computes kinematics +without advancing the simulation. This could be useful if we want to render from +a snapshot without waiting for the next physics tick. However, for the shared-engine +approach this isn't needed — rendering is already in the sim loop. + +### E. Pointcloud generation — reuse existing utility + +`dimos/simulation/mujoco/depth_camera.py` already has `depth_image_to_point_cloud()`. +But for ObjectSceneRegistrationModule, we don't need to generate pointclouds in the +camera module at all — the perception module generates its own per-object pointclouds +from depth + camera_info. Only enable `pointcloud` output if needed for visualization. + +### F. Engine injection pattern + +Rather than passing engine through Pydantic config (not serializable), use a +post-construction setter: + +```python +camera = dimos.deploy(MujocoCamera, camera_name="wrist_camera", ...) +camera.set_engine(engine) # inject before start() +``` + +Or use a lightweight engine registry keyed by MJCF path: + +```python +# In engine +_engine_registry: dict[str, MujocoEngine] = {} + +# Camera config just needs the MJCF path +class MujocoCameraConfig(ModuleConfig): + address: str # same MJCF path as adapter — looks up shared engine +``` + +This keeps blueprints clean and avoids non-serializable config fields. + +--- + +## Summary + +| Step | File | What | +|------|------|------| +| 1 | `simulation/engines/mujoco_engine.py` | Add camera rendering to sim loop | +| 2 | `simulation/mujoco/camera.py` | New module (same interface as RealSense) | +| 3 | `hardware/manipulators/sim/adapter.py` | Accept external engine | +| 4 | `manipulation/blueprints.py` | New `xarm_perception_sim` blueprint | +| 5 | `data/sim_assets/xarm7/scene.xml` | MuJoCo scene with camera + objects | + +Simplifications vs RealSense: +- No extrinsics / depth alignment +- No distortion +- Float32 depth directly (no DEPTH16 conversion) +- Intrinsics computed from model, not hardware +- Pointcloud generation optional (perception handles it) From ea255e44c2cf824d4bf514df8b9c956b49817809 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 24 Mar 2026 11:41:59 -0700 Subject: [PATCH 02/41] pick and place test --- dimos/manipulation/blueprints.py | 94 +++++++------ dimos/manipulation/pick_and_place_module.py | 15 ++- dimos/msgs/sensor_msgs/PointCloud2.py | 6 +- dimos/simulation/engines/mujoco_engine.py | 25 ++-- .../manipulators/sim_manip_interface.py | 2 +- dimos/simulation/mujoco/camera.py | 127 +++++++----------- 6 files changed, 121 insertions(+), 148 deletions(-) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index cf0d76fab9..7d508c7fac 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -183,6 +183,8 @@ def _make_xarm7_config( add_gripper: bool = False, gripper_hardware_id: str | None = None, tf_extra_links: list[str] | None = None, + pre_grasp_offset: float = 0.10, + home_joints: list[float] | None = None, ) -> RobotModelConfig: """Create XArm7 robot config. @@ -226,8 +228,8 @@ def _make_xarm7_config( coordinator_task_name=coordinator_task, gripper_hardware_id=gripper_hardware_id, tf_extra_links=tf_extra_links or [], - # Home configuration: arm extended forward, elbow up (safe observe pose) - home_joints=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + pre_grasp_offset=pre_grasp_offset, + home_joints=home_joints or [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], ) @@ -467,53 +469,49 @@ def _make_piper_config( from dimos.simulation.mujoco.camera import MujocoCamera from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode -xarm_perception_sim = ( - autoconnect( - PickAndPlaceModule.blueprint( - robots=[ - _make_xarm7_config( - "arm", - joint_prefix="arm_", - coordinator_task="traj_arm", - add_gripper=True, - gripper_hardware_id="arm", - tf_extra_links=["link7"], - ), - ], - planning_timeout=10.0, - enable_viz=True, - ), - MujocoCamera.blueprint( - address=str(XARM7_SIM_PATH), - camera_name="wrist_camera", - base_frame_id="link7", - ), - ObjectSceneRegistrationModule.blueprint(target_frame="world"), - ControlCoordinator.blueprint( - tick_rate=100.0, - publish_joint_state=True, - joint_state_frame_id="coordinator", - hardware=[sim_xarm7("arm", headless=False, gripper=True)], - tasks=[ - TaskConfig( - name="traj_arm", - type="trajectory", - joint_names=[f"arm_joint{i + 1}" for i in range(7)], - priority=10, - ), - ], - ), - RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode()), - FoxgloveBridge.blueprint(), - ) - .transports( - { - ("joint_state", JointState): LCMTransport( - "/coordinator/joint_state", JointState +xarm_perception_sim = autoconnect( + PickAndPlaceModule.blueprint( + robots=[ + _make_xarm7_config( + "arm", + joint_prefix="arm_", + coordinator_task="traj_arm", + add_gripper=True, + gripper_hardware_id="arm", + tf_extra_links=["link7"], + pre_grasp_offset=0.05, + # Observation pose: arm raised, camera above table looking down + home_joints=[0.0, -0.3, 0.0, 1.0, 0.0, 0.5, 0.0], ), - } - ) - .global_config(viewer="foxglove") + ], + planning_timeout=10.0, + enable_viz=True, + ), + MujocoCamera.blueprint( + address=str(XARM7_SIM_PATH), + camera_name="wrist_camera", + base_frame_id="link7", + ), + ObjectSceneRegistrationModule.blueprint(target_frame="world"), + ControlCoordinator.blueprint( + tick_rate=100.0, + publish_joint_state=True, + joint_state_frame_id="coordinator", + hardware=[sim_xarm7("arm", headless=False, gripper=True)], + tasks=[ + TaskConfig( + name="traj_arm", + type="trajectory", + joint_names=[f"arm_joint{i + 1}" for i in range(7)], + priority=10, + ), + ], + ), + RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode()), +).transports( + { + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + } ) diff --git a/dimos/manipulation/pick_and_place_module.py b/dimos/manipulation/pick_and_place_module.py index 81e7bcf2d3..07f1805fbf 100644 --- a/dimos/manipulation/pick_and_place_module.py +++ b/dimos/manipulation/pick_and_place_module.py @@ -303,8 +303,19 @@ def _generate_grasps_for_pick( return None c = det.center - grasp_pose = Pose(Vector3(c.x, c.y, c.z), Quaternion.from_euler(Vector3(0.0, math.pi, 0.0))) - logger.info(f"Heuristic grasp for '{object_name}' at ({c.x:.3f}, {c.y:.3f}, {c.z:.3f})") + # Depth cameras see the surface, not the volumetric center. + # Offset Z downward by half the detected height so the gripper + # reaches the object center rather than hovering above it. + z_offset = det.size.z * 0.5 if det.size.z > 0 else 0.0 + grasp_z = c.z - z_offset + grasp_pose = Pose( + Vector3(c.x, c.y, grasp_z), + Quaternion.from_euler(Vector3(0.0, math.pi, 0.0)), + ) + logger.info( + f"Heuristic grasp for '{object_name}' at ({c.x:.3f}, {c.y:.3f}, {grasp_z:.3f})" + f" (surface={c.z:.3f}, z_offset={z_offset:.3f})" + ) return [grasp_pose] @skill diff --git a/dimos/msgs/sensor_msgs/PointCloud2.py b/dimos/msgs/sensor_msgs/PointCloud2.py index f9c3874119..2903420b15 100644 --- a/dimos/msgs/sensor_msgs/PointCloud2.py +++ b/dimos/msgs/sensor_msgs/PointCloud2.py @@ -101,7 +101,11 @@ def __getstate__(self) -> dict[str, object]: del state["_pcd_tensor"] state["_pcd_legacy_cache"] = None # Remove cached_property entries that hold Open3D C++ objects - for key in ("axis_aligned_bounding_box", "oriented_bounding_box", "bounding_box_dimensions"): + for key in ( + "axis_aligned_bounding_box", + "oriented_bounding_box", + "bounding_box_dimensions", + ): state.pop(key, None) return state diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 9bddf0f838..3cb46fa9f0 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -76,7 +76,7 @@ class CameraFrame: def get_or_create_engine( - config_path: "Path", + config_path: Path, headless: bool = True, cameras: list[CameraConfig] | None = None, ) -> MujocoEngine: @@ -99,9 +99,7 @@ def get_or_create_engine( engine._camera_configs.append(cam) return engine - engine = MujocoEngine( - config_path=Path(config_path), headless=headless, cameras=cameras - ) + engine = MujocoEngine(config_path=Path(config_path), headless=headless, cameras=cameras) _engine_registry[key] = engine return engine @@ -269,22 +267,21 @@ def _init_new_cameras() -> None: for cfg in self._camera_configs: if cfg.name in cam_renderers: continue - cam_id = mujoco.mj_name2id( - self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name - ) + cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name) if cam_id < 0: logger.warning(f"Camera '{cfg.name}' not found in MJCF, skipping") continue - rgb_renderer = mujoco.Renderer( - self._model, height=cfg.height, width=cfg.width - ) - depth_renderer = mujoco.Renderer( - self._model, height=cfg.height, width=cfg.width - ) + rgb_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) + depth_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) depth_renderer.enable_depth_rendering() interval = 1.0 / cfg.fps if cfg.fps > 0 else float("inf") cam_renderers[cfg.name] = ( - cfg, cam_id, rgb_renderer, depth_renderer, interval, 0.0, + cfg, + cam_id, + rgb_renderer, + depth_renderer, + interval, + 0.0, ) logger.info( f"Camera '{cfg.name}' renderer created " diff --git a/dimos/simulation/manipulators/sim_manip_interface.py b/dimos/simulation/manipulators/sim_manip_interface.py index 07e56c5afd..abcfcc96f7 100644 --- a/dimos/simulation/manipulators/sim_manip_interface.py +++ b/dimos/simulation/manipulators/sim_manip_interface.py @@ -206,7 +206,7 @@ def write_gripper_position(self, position: float) -> bool: position = max(jlo, min(jhi, position)) if jhi != jlo: t = (position - jlo) / (jhi - jlo) - ctrl_value = chi - t * (chi - clo) + ctrl_value = clo + t * (chi - clo) else: ctrl_value = clo self._engine.set_position_target(self._gripper_idx, ctrl_value) diff --git a/dimos/simulation/mujoco/camera.py b/dimos/simulation/mujoco/camera.py index 5526c33202..16abf671ab 100644 --- a/dimos/simulation/mujoco/camera.py +++ b/dimos/simulation/mujoco/camera.py @@ -24,15 +24,13 @@ import threading import time -import numpy as np -import reactivex as rx from pydantic import Field +import reactivex as rx from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out from dimos.hardware.sensors.camera.spec import ( - OPTICAL_ROTATION, DepthCameraConfig, DepthCameraHardware, ) @@ -338,96 +336,61 @@ def _publish_camera_info(self) -> None: self.depth_camera_info.publish(self._depth_camera_info) def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: - transforms = [] - - # world -> camera_link (from MuJoCo camera world pose, updated each frame) - if frame is not None: - from scipy.spatial.transform import Rotation as R - - # MuJoCo camera frame: X=right, Y=up, Z=back (looks along -Z) - # ROS camera_link: X=forward, Y=left, Z=up - # OPTICAL_ROTATION converts camera_link -> optical (X=right, Y=down, Z=forward) - # We need: cam_xmat @ correction @ OPTICAL_ROTATION = cam_xmat @ Rx(180) - # So correction = Rx(180) @ OPTICAL_ROTATION.inv() - _MJ_TO_CAMLINK = R.from_quat([0.5, -0.5, -0.5, -0.5]) # xyzw - mj_rot = R.from_matrix(frame.cam_mat.reshape(3, 3)) - q_xyzw = (mj_rot * _MJ_TO_CAMLINK).as_quat() - transforms.append( - Transform( - translation=Vector3( - float(frame.cam_pos[0]), - float(frame.cam_pos[1]), - float(frame.cam_pos[2]), - ), - rotation=Quaternion( - float(q_xyzw[0]), - float(q_xyzw[1]), - float(q_xyzw[2]), - float(q_xyzw[3]), - ), - frame_id="world", - child_frame_id=self._camera_link, - ts=ts, - ) - ) - elif self.config.base_transform is not None: - # Fallback: static base_frame_id -> camera_link - transforms.append( - Transform( - translation=self.config.base_transform.translation, - rotation=self.config.base_transform.rotation, - frame_id=self.config.base_frame_id, - child_frame_id=self._camera_link, - ts=ts, - ) - ) + """Publish TF for the camera frame chain. + + For MuJoCo sim cameras we publish world -> optical_frame directly, + skipping the intermediate camera_link/color_frame chain. This avoids + convention mismatches between MuJoCo (X=right, Y=up, -Z=look) and + ROS camera_link (X=forward, Y=left, Z=up). + + MuJoCo optical = cam_xmat @ Rx(180): + X_opt = cam_X (right in image) + Y_opt = -cam_Y (down in image) + Z_opt = -cam_Z (into scene = look direction) + """ + if frame is None: + return - # camera_link -> depth_frame (identity — same virtual camera) - transforms.append( - Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=self._camera_link, - child_frame_id=self._depth_frame, - ts=ts, - ) + from scipy.spatial.transform import Rotation as R + + # MuJoCo cam frame -> optical frame: flip Y and Z (Rx 180°) + _RX180 = R.from_euler("x", 180, degrees=True) + mj_rot = R.from_matrix(frame.cam_mat.reshape(3, 3)) + optical_rot = mj_rot * _RX180 + q = optical_rot.as_quat() # xyzw + + pos = Vector3( + float(frame.cam_pos[0]), + float(frame.cam_pos[1]), + float(frame.cam_pos[2]), ) + rot = Quaternion(float(q[0]), float(q[1]), float(q[2]), float(q[3])) - # depth_frame -> depth_optical_frame - transforms.append( + # Publish world -> all optical/link frames (all co-located in sim) + self.tf.publish( Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=OPTICAL_ROTATION, - frame_id=self._depth_frame, - child_frame_id=self._depth_optical_frame, + translation=pos, + rotation=rot, + frame_id="world", + child_frame_id=self._color_optical_frame, ts=ts, - ) - ) - - # camera_link -> color_frame (identity — co-located in sim) - transforms.append( + ), Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=self._camera_link, - child_frame_id=self._color_frame, + translation=pos, + rotation=rot, + frame_id="world", + child_frame_id=self._depth_optical_frame, ts=ts, - ) - ) - - # color_frame -> color_optical_frame - transforms.append( + ), Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=OPTICAL_ROTATION, - frame_id=self._color_frame, - child_frame_id=self._color_optical_frame, + translation=pos, + rotation=rot, + frame_id="world", + child_frame_id=self._camera_link, ts=ts, - ) + ), ) - self.tf.publish(*transforms) - def _generate_pointcloud(self) -> None: """Generate pointcloud from latest depth frame (optional, for visualization).""" if self._engine is None or self._color_camera_info is None: From 47b66547699e6acbfd1051cf52e9d7e75d34984a Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 11:43:07 -0700 Subject: [PATCH 03/41] fix: dublicates - remove o3d cache --- dimos/msgs/sensor_msgs/PointCloud2.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/dimos/msgs/sensor_msgs/PointCloud2.py b/dimos/msgs/sensor_msgs/PointCloud2.py index 23e09d2630..2903420b15 100644 --- a/dimos/msgs/sensor_msgs/PointCloud2.py +++ b/dimos/msgs/sensor_msgs/PointCloud2.py @@ -100,9 +100,6 @@ def __getstate__(self) -> dict[str, object]: # Remove non-picklable objects del state["_pcd_tensor"] state["_pcd_legacy_cache"] = None - # Remove cached_property entries that hold unpicklable Open3D types - state.pop("oriented_bounding_box", None) - state.pop("axis_aligned_bounding_box", None) # Remove cached_property entries that hold Open3D C++ objects for key in ( "axis_aligned_bounding_box", From 96e9ee96c490cf2790602eca253940601382db21 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 12:44:50 -0700 Subject: [PATCH 04/41] misc: fix path and comments --- dimos/simulation/engines/mujoco_engine.py | 8 - .../{mujoco => manipulators}/camera.py | 17 -- docs/plans/mujoco_camera_plan.md | 250 ------------------ 3 files changed, 275 deletions(-) rename dimos/simulation/{mujoco => manipulators}/camera.py (94%) delete mode 100644 docs/plans/mujoco_camera_plan.md diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 3cb46fa9f0..1a2bd42e17 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -43,7 +43,6 @@ class CameraConfig: """Configuration for a camera to be rendered inside the sim loop.""" name: str - """Camera name as defined in the MJCF XML.""" width: int = 640 height: int = 480 fps: float = 15.0 @@ -54,18 +53,11 @@ class CameraFrame: """Thread-safe container for a rendered camera frame.""" rgb: NDArray[np.uint8] - """RGB image, shape (H, W, 3).""" depth: NDArray[np.float32] - """Depth image in meters, shape (H, W).""" cam_pos: NDArray[np.float64] - """Camera world position, shape (3,).""" cam_mat: NDArray[np.float64] - """Camera rotation matrix (flattened 3x3), shape (9,).""" fovy: float - """Vertical field of view in degrees.""" timestamp: float - """Time of rendering.""" - # --------------------------------------------------------------------------- # Engine registry — allows adapter and camera to share the same engine instance diff --git a/dimos/simulation/mujoco/camera.py b/dimos/simulation/manipulators/camera.py similarity index 94% rename from dimos/simulation/mujoco/camera.py rename to dimos/simulation/manipulators/camera.py index 16abf671ab..09b178257d 100644 --- a/dimos/simulation/mujoco/camera.py +++ b/dimos/simulation/manipulators/camera.py @@ -64,12 +64,8 @@ class MujocoCameraConfig(ModuleConfig, DepthCameraConfig): """Configuration for the MuJoCo simulation camera.""" address: str = "" - """MJCF XML path — same as the sim adapter's address. - Used to look up the shared MujocoEngine from the registry.""" headless: bool = False - """Whether to run MuJoCo without a viewer window.""" camera_name: str = "wrist_camera" - """Camera name as defined in the MJCF XML.""" width: int = 640 height: int = 480 fps: int = 15 @@ -336,23 +332,10 @@ def _publish_camera_info(self) -> None: self.depth_camera_info.publish(self._depth_camera_info) def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: - """Publish TF for the camera frame chain. - - For MuJoCo sim cameras we publish world -> optical_frame directly, - skipping the intermediate camera_link/color_frame chain. This avoids - convention mismatches between MuJoCo (X=right, Y=up, -Z=look) and - ROS camera_link (X=forward, Y=left, Z=up). - - MuJoCo optical = cam_xmat @ Rx(180): - X_opt = cam_X (right in image) - Y_opt = -cam_Y (down in image) - Z_opt = -cam_Z (into scene = look direction) - """ if frame is None: return from scipy.spatial.transform import Rotation as R - # MuJoCo cam frame -> optical frame: flip Y and Z (Rx 180°) _RX180 = R.from_euler("x", 180, degrees=True) mj_rot = R.from_matrix(frame.cam_mat.reshape(3, 3)) diff --git a/docs/plans/mujoco_camera_plan.md b/docs/plans/mujoco_camera_plan.md deleted file mode 100644 index f8a6b9b395..0000000000 --- a/docs/plans/mujoco_camera_plan.md +++ /dev/null @@ -1,250 +0,0 @@ -# MuJoCo Camera Module Plan - -## Goal - -Create a `MujocoCamera` module that can drop in for `RealSenseCamera` in manipulation -blueprints, enabling sim-based pick-and-place without real hardware. - -## Current Wiring (real hardware) - -``` -PickAndPlaceModule - ↑ joint_state (LCM) -RealSenseCamera(base_frame_id="link7", base_transform=...) - ↓ color_image, depth_image, camera_info -ObjectSceneRegistrationModule(target_frame="world") -``` - -`ObjectSceneRegistrationModule` consumes exactly three inputs: -- `color_image: In[Image]` — RGB -- `depth_image: In[Image]` — DEPTH16 or float32 -- `camera_info: In[CameraInfo]` — intrinsics (K matrix) - -## Threading Constraint - -MuJoCo's C API is **not thread-safe**. `mujoco.Renderer` must be created and used in -the **same thread** as `mj_step()`. The existing nav camera (mujoco_process.py) solves -this by rendering inside the sim loop. We must follow the same pattern. - -`MujocoEngine._sim_loop()` currently only does: apply_control → mj_step → sync_viewer → update_state. -Rendering must be added **inside this loop**. - ---- - -## Basic Plan - -### 1. Extend `MujocoEngine` with camera rendering - -**File:** `dimos/simulation/engines/mujoco_engine.py` - -Add rendering capability to the sim loop: - -- New init params: `cameras: list[CameraConfig]` (name, width, height, fps) -- Create `mujoco.Renderer` instances in `_sim_loop()` (same thread as `mj_step`) -- Rate-limit rendering per camera (e.g., 15 FPS vs 500 Hz physics) -- After each render, write RGB + depth arrays into a **locked buffer** -- Expose `read_camera(camera_name) -> (rgb, depth, cam_pos, cam_mat, fovy)` that - returns the latest buffered frame (thread-safe copy) - -``` -_sim_loop: - create renderers (same thread) - while running: - apply_control() - mj_step() - for each camera (if due): - renderer.update_scene(data, camera=cam_id) - rgb = rgb_renderer.render() - depth = depth_renderer.render() - cam_pos = data.cam_xpos[cam_id] - cam_mat = data.cam_xmat[cam_id] - write to locked buffer - update_joint_state() -``` - -### 2. Create `MujocoCamera` module - -**File:** `dimos/simulation/mujoco/camera.py` - -A `Module` implementing `DepthCameraHardware` + `perception.DepthCamera`: - -```python -class MujocoCameraConfig(ModuleConfig, DepthCameraConfig): - camera_name: str = "wrist_camera" # MuJoCo camera name in MJCF - width: int = 640 - height: int = 480 - fps: int = 15 - base_frame_id: str = "link7" - base_transform: Transform | None = None - enable_depth: bool = True - enable_pointcloud: bool = False - pointcloud_fps: float = 5.0 - camera_info_fps: float = 1.0 - depth_scale: float = 1.0 # MuJoCo depth is in meters - -class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.DepthCamera): - color_image: Out[Image] - depth_image: Out[Image] - pointcloud: Out[PointCloud2] - camera_info: Out[CameraInfo] - depth_camera_info: Out[CameraInfo] -``` - -**Lifecycle:** -- `set_engine(engine: MujocoEngine)` — called before `start()`, stores engine ref -- `start()` — compute intrinsics from engine's `model.cam_fovy[cam_id]`, start - polling thread that calls `engine.read_camera()` at configured FPS -- Polling thread: read buffered frame → wrap as `Image` → publish on output ports -- Publish TF: `base_frame_id → camera_link → optical frames` (same chain as RealSense) - -### 3. Refactor `SimMujocoAdapter` to accept external engine - -**File:** `dimos/hardware/manipulators/sim/adapter.py` - -Currently creates `MujocoEngine` internally. Add option to accept an existing one: - -```python -class SimMujocoAdapter(SimManipInterface): - def __init__(self, dof=7, address=None, headless=True, engine=None, **_): - if engine is None: - engine = MujocoEngine(config_path=Path(address), headless=headless) - # ... rest unchanged -``` - -### 4. Blueprint wiring - -**File:** `dimos/manipulation/blueprints.py` - -Create a new `xarm_perception_sim` blueprint: - -```python -# Shared engine created once -_engine = MujocoEngine(config_path=Path("path/to/xarm_scene.xml"), headless=False) - -xarm_perception_sim = autoconnect( - PickAndPlaceModule.blueprint( - robots=[_make_xarm7_config(...)], - ... - ), - MujocoCamera.blueprint( - camera_name="wrist_camera", - base_frame_id="link7", - base_transform=_XARM_PERCEPTION_CAMERA_TRANSFORM, - _engine=_engine, # injected reference - ), - ObjectSceneRegistrationModule.blueprint(target_frame="world"), - ControlCoordinator.blueprint( - hardware=[HardwareComponent(adapter_type="sim_mujoco", ...)], - ... - ), -) -``` - -### 5. MJCF scene file - -Create `data/sim_assets/xarm7/scene.xml` with: -- xArm7 robot model (arm + gripper) -- `wrist_camera` camera attached to link7 (matching the real RealSense mount) -- Table + objects for pick-and-place -- Lighting - ---- - -## Improvements & Simplifications - -### A. Depth format: use meters directly (skip DEPTH16 conversion) - -RealSense outputs DEPTH16 (uint16, mm). ObjectSceneRegistrationModule converts it: -```python -depth_meters = depth_data.astype(np.float32) / 1000.0 -``` - -MuJoCo depth is already in meters (float32). We can publish as `ImageFormat.DEPTH32F` -directly. ObjectSceneRegistrationModule already handles both formats (line 281-289): -```python -if depth_image.format == ImageFormat.DEPTH16: - depth_data = depth_data.astype(np.float32) / 1000.0 -``` - -**Simplification:** No artificial uint16 conversion needed. Just publish float32 depth. - -### B. Skip extrinsics — color and depth are co-located - -RealSense has physically separate color/depth sensors with extrinsics between them. -MuJoCo renders both from the same virtual camera. So: - -- Color and depth share the same intrinsics -- No color-to-depth alignment needed (`align_depth_to_color` is a no-op) -- TF chain simplifies: `camera_link = color_frame = depth_frame` (single camera origin) -- `depth_camera_info` = `color_camera_info` (identical) - -### C. Intrinsics from model — no calibration - -RealSense reads intrinsics from hardware via `pyrealsense2`. MuJoCo camera params -are known exactly from the model: - -```python -fovy = model.cam_fovy[cam_id] # vertical FOV in degrees -fy = height / (2 * tan(radians(fovy) / 2)) -fx = fy # square pixels -cx, cy = width / 2, height / 2 -D = [0, 0, 0, 0, 0] # no distortion -``` - -Compute once at `start()`, publish periodically. No runtime calibration. - -### D. Consider rendering in `mj_forward()` instead of `mj_step()` - -For camera-only updates (no physics needed), `mj_forward()` computes kinematics -without advancing the simulation. This could be useful if we want to render from -a snapshot without waiting for the next physics tick. However, for the shared-engine -approach this isn't needed — rendering is already in the sim loop. - -### E. Pointcloud generation — reuse existing utility - -`dimos/simulation/mujoco/depth_camera.py` already has `depth_image_to_point_cloud()`. -But for ObjectSceneRegistrationModule, we don't need to generate pointclouds in the -camera module at all — the perception module generates its own per-object pointclouds -from depth + camera_info. Only enable `pointcloud` output if needed for visualization. - -### F. Engine injection pattern - -Rather than passing engine through Pydantic config (not serializable), use a -post-construction setter: - -```python -camera = dimos.deploy(MujocoCamera, camera_name="wrist_camera", ...) -camera.set_engine(engine) # inject before start() -``` - -Or use a lightweight engine registry keyed by MJCF path: - -```python -# In engine -_engine_registry: dict[str, MujocoEngine] = {} - -# Camera config just needs the MJCF path -class MujocoCameraConfig(ModuleConfig): - address: str # same MJCF path as adapter — looks up shared engine -``` - -This keeps blueprints clean and avoids non-serializable config fields. - ---- - -## Summary - -| Step | File | What | -|------|------|------| -| 1 | `simulation/engines/mujoco_engine.py` | Add camera rendering to sim loop | -| 2 | `simulation/mujoco/camera.py` | New module (same interface as RealSense) | -| 3 | `hardware/manipulators/sim/adapter.py` | Accept external engine | -| 4 | `manipulation/blueprints.py` | New `xarm_perception_sim` blueprint | -| 5 | `data/sim_assets/xarm7/scene.xml` | MuJoCo scene with camera + objects | - -Simplifications vs RealSense: -- No extrinsics / depth alignment -- No distortion -- Float32 depth directly (no DEPTH16 conversion) -- Intrinsics computed from model, not hardware -- Pointcloud generation optional (perception handles it) From b1413b61e47d0719ab6f1014b20b580cad3477a8 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 12:45:27 -0700 Subject: [PATCH 05/41] fix: pre-commit issues --- dimos/manipulation/planning/spec/config.py | 1 - dimos/simulation/engines/mujoco_engine.py | 1 + dimos/simulation/manipulators/camera.py | 1 + 3 files changed, 2 insertions(+), 1 deletion(-) diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index a6a6015bbb..f8dcb8b548 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -16,7 +16,6 @@ from __future__ import annotations -from collections.abc import Sequence from pathlib import Path from pydantic import Field diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 1a2bd42e17..dba0ff43f4 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -59,6 +59,7 @@ class CameraFrame: fovy: float timestamp: float + # --------------------------------------------------------------------------- # Engine registry — allows adapter and camera to share the same engine instance # by MJCF path. diff --git a/dimos/simulation/manipulators/camera.py b/dimos/simulation/manipulators/camera.py index 09b178257d..50eb345d92 100644 --- a/dimos/simulation/manipulators/camera.py +++ b/dimos/simulation/manipulators/camera.py @@ -336,6 +336,7 @@ def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: return from scipy.spatial.transform import Rotation as R + # MuJoCo cam frame -> optical frame: flip Y and Z (Rx 180°) _RX180 = R.from_euler("x", 180, degrees=True) mj_rot = R.from_matrix(frame.cam_mat.reshape(3, 3)) From b30ad6b4f9a746070b3bb070dd33daed147f1b97 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 13:08:26 -0700 Subject: [PATCH 06/41] fix: camrenderer dataclass + unregister method --- dimos/simulation/engines/mujoco_engine.py | 143 +++++++++++++--------- 1 file changed, 83 insertions(+), 60 deletions(-) diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index dba0ff43f4..3fbf183952 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -60,6 +60,18 @@ class CameraFrame: timestamp: float +@dataclasses.dataclass +class _CameraRendererState: + """Mutable state for a single camera renderer (sim-thread only).""" + + cfg: CameraConfig + cam_id: int + rgb_renderer: mujoco.Renderer + depth_renderer: mujoco.Renderer + interval: float + last_render_time: float = 0.0 + + # --------------------------------------------------------------------------- # Engine registry — allows adapter and camera to share the same engine instance # by MJCF path. @@ -97,6 +109,14 @@ def get_or_create_engine( return engine +def unregister_engine(engine: MujocoEngine) -> None: + """Remove an engine from the registry (called on disconnect).""" + with _engine_registry_lock: + keys_to_remove = [k for k, v in _engine_registry.items() if v is engine] + for k in keys_to_remove: + del _engine_registry[k] + + class MujocoEngine(SimulationEngine): """ MuJoCo simulation engine. @@ -240,71 +260,77 @@ def disconnect(self) -> bool: if self._sim_thread and self._sim_thread.is_alive(): self._sim_thread.join(timeout=2.0) self._sim_thread = None + unregister_engine(self) return True except Exception as e: logger.error(f"{self.__class__.__name__}: disconnect() failed: {e}") return False + def _init_new_cameras(self, cam_renderers: dict[str, _CameraRendererState]) -> None: + """Create renderers for any camera configs not yet initialized. + + Must be called from the sim thread (MuJoCo thread-safety). + """ + for cfg in self._camera_configs: + if cfg.name in cam_renderers: + continue + cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name) + if cam_id < 0: + logger.warning(f"Camera '{cfg.name}' not found in MJCF, skipping") + continue + rgb_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) + depth_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) + depth_renderer.enable_depth_rendering() + interval = 1.0 / cfg.fps if cfg.fps > 0 else float("inf") + cam_renderers[cfg.name] = _CameraRendererState( + cfg=cfg, + cam_id=cam_id, + rgb_renderer=rgb_renderer, + depth_renderer=depth_renderer, + interval=interval, + ) + logger.info( + f"Camera '{cfg.name}' renderer created ({cfg.width}x{cfg.height} @ {cfg.fps}fps)" + ) + + def _render_cameras(self, now: float, cam_renderers: dict[str, _CameraRendererState]) -> None: + """Render all due cameras and store frames. Must be called from sim thread.""" + self._init_new_cameras(cam_renderers) + for state in cam_renderers.values(): + if now - state.last_render_time < state.interval: + continue + state.last_render_time = now + + state.rgb_renderer.update_scene(self._data, camera=state.cam_id) + rgb = state.rgb_renderer.render().copy() + + state.depth_renderer.update_scene(self._data, camera=state.cam_id) + depth = state.depth_renderer.render().copy() + + frame = CameraFrame( + rgb=rgb, + depth=depth.astype(np.float32), + cam_pos=self._data.cam_xpos[state.cam_id].copy(), + cam_mat=self._data.cam_xmat[state.cam_id].copy(), + fovy=float(self._model.cam_fovy[state.cam_id]), + timestamp=now, + ) + with self._camera_lock: + self._camera_frames[state.cfg.name] = frame + + @staticmethod + def _close_cam_renderers(cam_renderers: dict[str, _CameraRendererState]) -> None: + for state in cam_renderers.values(): + state.rgb_renderer.close() + state.depth_renderer.close() + def _sim_loop(self) -> None: logger.info(f"{self.__class__.__name__}: sim loop started") dt = 1.0 / self._control_frequency # Camera renderers — created in sim thread (MuJoCo thread-safety). # Checked each tick so cameras added after connect() are picked up. - cam_renderers: dict[ - str, tuple[CameraConfig, int, mujoco.Renderer, mujoco.Renderer, float, float] - ] = {} # name -> (cfg, cam_id, rgb_r, depth_r, interval, last_render_time) - - def _init_new_cameras() -> None: - """Create renderers for any camera configs not yet initialized.""" - for cfg in self._camera_configs: - if cfg.name in cam_renderers: - continue - cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name) - if cam_id < 0: - logger.warning(f"Camera '{cfg.name}' not found in MJCF, skipping") - continue - rgb_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) - depth_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) - depth_renderer.enable_depth_rendering() - interval = 1.0 / cfg.fps if cfg.fps > 0 else float("inf") - cam_renderers[cfg.name] = ( - cfg, - cam_id, - rgb_renderer, - depth_renderer, - interval, - 0.0, - ) - logger.info( - f"Camera '{cfg.name}' renderer created " - f"({cfg.width}x{cfg.height} @ {cfg.fps}fps)" - ) - - def _render_cameras(now: float) -> None: - _init_new_cameras() - for name in cam_renderers: - cfg, cam_id, rgb_r, depth_r, interval, last_t = cam_renderers[name] - if now - last_t < interval: - continue - cam_renderers[name] = (cfg, cam_id, rgb_r, depth_r, interval, now) - - rgb_r.update_scene(self._data, camera=cam_id) - rgb = rgb_r.render().copy() - - depth_r.update_scene(self._data, camera=cam_id) - depth = depth_r.render().copy() - - frame = CameraFrame( - rgb=rgb, - depth=depth.astype(np.float32), - cam_pos=self._data.cam_xpos[cam_id].copy(), - cam_mat=self._data.cam_xmat[cam_id].copy(), - fovy=float(self._model.cam_fovy[cam_id]), - timestamp=now, - ) - with self._camera_lock: - self._camera_frames[cfg.name] = frame + cam_renderers: dict[str, _CameraRendererState] = {} def _step_once(sync_viewer: bool) -> None: loop_start = time.time() @@ -313,7 +339,7 @@ def _step_once(sync_viewer: bool) -> None: if sync_viewer: m_viewer.sync() self._update_joint_state() - _render_cameras(loop_start) + self._render_cameras(loop_start, cam_renderers) elapsed = time.time() - loop_start sleep_time = dt - elapsed @@ -330,11 +356,7 @@ def _step_once(sync_viewer: bool) -> None: while m_viewer.is_running() and not self._stop_event.is_set(): _step_once(sync_viewer=True) - # Clean up renderers - for _cfg, _cam_id, rgb_r, depth_r, _interval, _t in cam_renderers.values(): - rgb_r.close() - depth_r.close() - + self._close_cam_renderers(cam_renderers) logger.info(f"{self.__class__.__name__}: sim loop stopped") @property @@ -483,4 +505,5 @@ def get_camera_fovy(self, camera_name: str) -> float | None: "CameraFrame", "MujocoEngine", "get_or_create_engine", + "unregister_engine", ] From 1265f53115849144a7fde57f9c12f2943d8c1ef1 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 13:08:40 -0700 Subject: [PATCH 07/41] fix: mujococamera path --- dimos/manipulation/blueprints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 076a47e26d..1572fd0b22 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -514,7 +514,7 @@ def _make_piper_config( # The engine is created lazily when the adapter connects / camera starts. from dimos.control.blueprints._hardware import XARM7_SIM_PATH, sim_xarm7 -from dimos.simulation.mujoco.camera import MujocoCamera +from dimos.simulation.manipulators.camera import MujocoCamera from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode xarm_perception_sim = autoconnect( From ccceb456666e3ebc944242944608d119b377645d Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 13:16:37 -0700 Subject: [PATCH 08/41] fix: pre-commit errors --- dimos/simulation/manipulators/camera.py | 68 ++++++++++--------------- 1 file changed, 26 insertions(+), 42 deletions(-) diff --git a/dimos/simulation/manipulators/camera.py b/dimos/simulation/manipulators/camera.py index 50eb345d92..3809bfe2d5 100644 --- a/dimos/simulation/manipulators/camera.py +++ b/dimos/simulation/manipulators/camera.py @@ -26,6 +26,7 @@ from pydantic import Field import reactivex as rx +from scipy.spatial.transform import Rotation as R from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig @@ -52,6 +53,8 @@ logger = setup_logger() +_RX180 = R.from_euler("x", 180, degrees=True) + def _default_identity_transform() -> Transform: return Transform( @@ -71,7 +74,7 @@ class MujocoCameraConfig(ModuleConfig, DepthCameraConfig): fps: int = 15 base_frame_id: str = "link7" base_transform: Transform | None = Field(default_factory=_default_identity_transform) - # MuJoCo renders color+depth from same virtual camera — alignment is a no-op + # MuJoCo renders color+depth from same virtual camera, alignment is a no-op align_depth_to_color: bool = True enable_depth: bool = True enable_pointcloud: bool = False @@ -103,17 +106,12 @@ def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] self._engine: MujocoEngine | None = None self._running = False self._thread: threading.Thread | None = None - self._color_camera_info: CameraInfo | None = None - self._depth_camera_info: CameraInfo | None = None - - # -- Engine injection (called before start) -- + self._camera_info_base: CameraInfo | None = None def set_engine(self, engine: MujocoEngine) -> None: """Inject the shared MujocoEngine reference.""" self._engine = engine - # -- DepthCameraHardware interface -- - @property def _camera_link(self) -> str: return f"{self.config.camera_name}_link" @@ -136,19 +134,16 @@ def _depth_optical_frame(self) -> str: @rpc def get_color_camera_info(self) -> CameraInfo | None: - return self._color_camera_info + return self._camera_info_base @rpc def get_depth_camera_info(self) -> CameraInfo | None: - return self._depth_camera_info + return self._camera_info_base @rpc def get_depth_scale(self) -> float: - # MuJoCo depth is already in meters return 1.0 - # -- Intrinsics -- - def _build_camera_info(self) -> None: """Compute camera intrinsics from the MuJoCo model (pinhole, no distortion).""" if self._engine is None: @@ -170,7 +165,7 @@ def _build_camera_info(self) -> None: P = [fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0] D = [0.0, 0.0, 0.0, 0.0, 0.0] - info = CameraInfo( + self._camera_info_base = CameraInfo( height=h, width=w, distortion_model="plumb_bob", @@ -179,19 +174,6 @@ def _build_camera_info(self) -> None: P=P, frame_id=self._color_optical_frame, ) - # Color and depth share the same virtual camera - self._color_camera_info = info - self._depth_camera_info = CameraInfo( - height=h, - width=w, - distortion_model="plumb_bob", - D=D, - K=K, - P=P, - frame_id=self._color_optical_frame, - ) - - # -- Lifecycle -- @rpc def start(self) -> None: @@ -251,12 +233,9 @@ def stop(self) -> None: if self._thread and self._thread.is_alive(): self._thread.join(timeout=2.0) self._thread = None - self._color_camera_info = None - self._depth_camera_info = None + self._camera_info_base = None super().stop() - # -- Publishing -- - def _publish_loop(self) -> None: """Poll engine for rendered frames and publish at configured FPS.""" interval = 1.0 / self.config.fps @@ -323,22 +302,27 @@ def _publish_loop(self) -> None: time.sleep(1.0) def _publish_camera_info(self) -> None: + base = self._camera_info_base + if base is None: + return ts = time.time() - if self._color_camera_info: - self._color_camera_info.ts = ts - self.camera_info.publish(self._color_camera_info) - if self._depth_camera_info: - self._depth_camera_info.ts = ts - self.depth_camera_info.publish(self._depth_camera_info) + color_info = CameraInfo( + height=base.height, + width=base.width, + distortion_model=base.distortion_model, + D=base.D, + K=base.K, + P=base.P, + frame_id=base.frame_id, + ts=ts, + ) + self.camera_info.publish(color_info) + self.depth_camera_info.publish(color_info) def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: if frame is None: return - from scipy.spatial.transform import Rotation as R - - # MuJoCo cam frame -> optical frame: flip Y and Z (Rx 180°) - _RX180 = R.from_euler("x", 180, degrees=True) mj_rot = R.from_matrix(frame.cam_mat.reshape(3, 3)) optical_rot = mj_rot * _RX180 q = optical_rot.as_quat() # xyzw @@ -377,7 +361,7 @@ def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: def _generate_pointcloud(self) -> None: """Generate pointcloud from latest depth frame (optional, for visualization).""" - if self._engine is None or self._color_camera_info is None: + if self._engine is None or self._camera_info_base is None: return frame = self._engine.read_camera(self.config.camera_name) if frame is None: @@ -398,7 +382,7 @@ def _generate_pointcloud(self) -> None: pcd = PointCloud2.from_rgbd( color_image=color_img, depth_image=depth_img, - camera_info=self._color_camera_info, + camera_info=self._camera_info_base, depth_scale=1.0, ) pcd = pcd.voxel_downsample(0.005) From 0fb50cc1f0961704ca679ddeca6cc109bd29c208 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 17:37:28 -0700 Subject: [PATCH 09/41] set threading.event --- dimos/simulation/manipulators/camera.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/dimos/simulation/manipulators/camera.py b/dimos/simulation/manipulators/camera.py index 3809bfe2d5..92513bdcef 100644 --- a/dimos/simulation/manipulators/camera.py +++ b/dimos/simulation/manipulators/camera.py @@ -104,7 +104,7 @@ class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.D def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] super().__init__(*args, **kwargs) self._engine: MujocoEngine | None = None - self._running = False + self._stop_event = threading.Event() self._thread: threading.Thread | None = None self._camera_info_base: CameraInfo | None = None @@ -204,7 +204,7 @@ def start(self) -> None: self._build_camera_info() - self._running = True + self._stop_event.clear() self._thread = threading.Thread(target=self._publish_loop, daemon=True) self._thread.start() @@ -229,7 +229,7 @@ def start(self) -> None: @rpc def stop(self) -> None: - self._running = False + self._stop_event.set() if self._thread and self._thread.is_alive(): self._thread.join(timeout=2.0) self._thread = None @@ -245,19 +245,23 @@ def _publish_loop(self) -> None: logger.info(f"MujocoCamera publish loop started (camera={self.config.camera_name})") # Wait for engine to connect (adapter may not have started yet) - while self._running and self._engine is not None and not self._engine.connected: - time.sleep(0.1) - - if not self._running: + deadline = time.monotonic() + 30.0 + while not self._stop_event.is_set() and self._engine is not None and not self._engine.connected: + if time.monotonic() > deadline: + logger.error("MujocoCamera: timed out waiting for engine to connect") + return + self._stop_event.wait(timeout=0.1) + + if self._stop_event.is_set(): return logger.info("MujocoCamera: engine connected, polling for frames") - while self._running and self._engine is not None: + while not self._stop_event.is_set() and self._engine is not None: try: frame = self._engine.read_camera(self.config.camera_name) if frame is None or frame.timestamp <= last_timestamp: - time.sleep(interval * 0.5) + self._stop_event.wait(timeout=interval * 0.5) continue last_timestamp = frame.timestamp From 3d400c2d7deb785c7feb2f1811766806043f6923 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 17:38:47 -0700 Subject: [PATCH 10/41] feat: clear registry method --- dimos/simulation/engines/mujoco_engine.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 3fbf183952..fe056fce5b 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -117,6 +117,12 @@ def unregister_engine(engine: MujocoEngine) -> None: del _engine_registry[k] +def _clear_registry() -> None: + """Clear the engine registry (for test teardown only).""" + with _engine_registry_lock: + _engine_registry.clear() + + class MujocoEngine(SimulationEngine): """ MuJoCo simulation engine. From bd0660dfe2345d27892ff037c4d6a66e7ef0194f Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 17:59:58 -0700 Subject: [PATCH 11/41] fix: dataclass --- dimos/simulation/engines/mujoco_engine.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index fe056fce5b..d601968ef5 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -16,7 +16,7 @@ from __future__ import annotations -import dataclasses +from dataclasses import dataclass import threading import time from typing import TYPE_CHECKING @@ -38,7 +38,7 @@ logger = setup_logger() -@dataclasses.dataclass +@dataclass class CameraConfig: """Configuration for a camera to be rendered inside the sim loop.""" @@ -48,7 +48,7 @@ class CameraConfig: fps: float = 15.0 -@dataclasses.dataclass +@dataclass class CameraFrame: """Thread-safe container for a rendered camera frame.""" @@ -60,7 +60,7 @@ class CameraFrame: timestamp: float -@dataclasses.dataclass +@dataclass class _CameraRendererState: """Mutable state for a single camera renderer (sim-thread only).""" From 33c4689e8c3ce49c1a6dca4e22bba9e38ed89943 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 18:00:11 -0700 Subject: [PATCH 12/41] feat: add tests --- dimos/simulation/manipulators/camera.py | 6 +- .../manipulators/test_mujoco_camera.py | 325 ++++++++++++++++++ 2 files changed, 330 insertions(+), 1 deletion(-) create mode 100644 dimos/simulation/manipulators/test_mujoco_camera.py diff --git a/dimos/simulation/manipulators/camera.py b/dimos/simulation/manipulators/camera.py index 92513bdcef..ff73b23890 100644 --- a/dimos/simulation/manipulators/camera.py +++ b/dimos/simulation/manipulators/camera.py @@ -246,7 +246,11 @@ def _publish_loop(self) -> None: # Wait for engine to connect (adapter may not have started yet) deadline = time.monotonic() + 30.0 - while not self._stop_event.is_set() and self._engine is not None and not self._engine.connected: + while ( + not self._stop_event.is_set() + and self._engine is not None + and not self._engine.connected + ): if time.monotonic() > deadline: logger.error("MujocoCamera: timed out waiting for engine to connect") return diff --git a/dimos/simulation/manipulators/test_mujoco_camera.py b/dimos/simulation/manipulators/test_mujoco_camera.py new file mode 100644 index 0000000000..3b8ff8596b --- /dev/null +++ b/dimos/simulation/manipulators/test_mujoco_camera.py @@ -0,0 +1,325 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for MujocoCamera and engine registry.""" + +from __future__ import annotations + +import math +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from dimos.simulation.engines.mujoco_engine import ( + CameraConfig, + CameraFrame, + MujocoEngine, + _clear_registry, + _engine_registry, + get_or_create_engine, + unregister_engine, +) +from dimos.simulation.manipulators.camera import MujocoCamera +from dimos.simulation.manipulators.test_sim_adapter import _patch_mujoco_engine + +ARM_DOF = 7 + + +@pytest.fixture(autouse=True) +def clean_registry(): + _clear_registry() + yield + _clear_registry() + + +def _make_camera_frame( + cam_pos: list[float] | None = None, + cam_mat: list[float] | None = None, +) -> CameraFrame: + """Create a CameraFrame with sensible defaults.""" + return CameraFrame( + rgb=np.zeros((480, 640, 3), dtype=np.uint8), + depth=np.ones((480, 640), dtype=np.float32), + cam_pos=np.array(cam_pos or [1.0, 2.0, 3.0]), + cam_mat=np.array(cam_mat or np.eye(3).flatten()), + fovy=45.0, + timestamp=1.0, + ) + + +def _make_mock_engine(fovy: float = 45.0) -> MagicMock: + mock_engine = MagicMock(spec=MujocoEngine) + mock_engine.get_camera_fovy.return_value = fovy + mock_engine.connected = True + mock_engine._camera_configs = [] + mock_engine.read_camera.return_value = _make_camera_frame() + return mock_engine + + +@pytest.fixture +def mock_engine() -> MagicMock: + return _make_mock_engine() + + +@pytest.fixture +def camera_with_mock_engine(mock_engine: MagicMock): + cam = MujocoCamera(camera_name="wrist_camera") + cam.set_engine(mock_engine) + yield cam + cam.stop() + + +# --------------------------------------------------------------------------- +# 1. Engine Registry +# --------------------------------------------------------------------------- + + +@pytest.mark.mujoco +class TestEngineRegistry: + def test_creates_new(self): + patches = _patch_mujoco_engine(ARM_DOF) + for p in patches: + p.start() + try: + engine = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) + assert engine is not None + assert len(_engine_registry) == 1 + finally: + for p in patches: + p.stop() + + def test_returns_same_instance(self): + patches = _patch_mujoco_engine(ARM_DOF) + for p in patches: + p.start() + try: + e1 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) + e2 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) + assert e1 is e2 + assert len(_engine_registry) == 1 + finally: + for p in patches: + p.stop() + + def test_merges_new_cameras(self): + patches = _patch_mujoco_engine(ARM_DOF) + for p in patches: + p.start() + try: + e1 = get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_a")], + ) + get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_b")], + ) + names = {c.name for c in e1._camera_configs} + assert names == {"cam_a", "cam_b"} + finally: + for p in patches: + p.stop() + + def test_deduplicates_cameras(self): + patches = _patch_mujoco_engine(ARM_DOF) + for p in patches: + p.start() + try: + get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_a")], + ) + get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_a")], + ) + engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) + cam_names = [c.name for c in engine._camera_configs] + assert cam_names.count("cam_a") == 1 + finally: + for p in patches: + p.stop() + + def test_unregister_removes(self): + patches = _patch_mujoco_engine(ARM_DOF) + for p in patches: + p.start() + try: + engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) + assert len(_engine_registry) == 1 + unregister_engine(engine) + assert len(_engine_registry) == 0 + finally: + for p in patches: + p.stop() + + def test_unregister_noop_unknown(self): + mock = MagicMock(spec=MujocoEngine) + unregister_engine(mock) # should not raise + + +# --------------------------------------------------------------------------- +# 2. Camera Intrinsics +# --------------------------------------------------------------------------- + + +@pytest.mark.mujoco +class TestCameraIntrinsics: + def test_fovy_45(self, camera_with_mock_engine: MujocoCamera): + cam = camera_with_mock_engine + cam._build_camera_info() + info = cam._camera_info_base + assert info is not None + + expected_fy = 480.0 / (2.0 * math.tan(math.radians(45.0) / 2.0)) + # K = [fx, 0, cx, 0, fy, cy, 0, 0, 1] + assert info.K[0] == pytest.approx(expected_fy, abs=0.01) # fx + assert info.K[4] == pytest.approx(expected_fy, abs=0.01) # fy + assert info.K[2] == pytest.approx(320.0) # cx + assert info.K[5] == pytest.approx(240.0) # cy + + def test_fovy_90(self, mock_engine: MagicMock): + mock_engine.get_camera_fovy.return_value = 90.0 + cam = MujocoCamera(camera_name="wrist_camera") + cam.set_engine(mock_engine) + cam._build_camera_info() + info = cam._camera_info_base + assert info is not None + # tan(45°) = 1.0, so fy = 480 / 2 = 240 + assert info.K[0] == pytest.approx(240.0, abs=0.01) + assert info.K[4] == pytest.approx(240.0, abs=0.01) + cam.stop() + + def test_no_engine_noop(self): + cam = MujocoCamera(camera_name="wrist_camera") + cam._build_camera_info() + assert cam._camera_info_base is None + cam.stop() + + def test_unknown_camera(self, mock_engine: MagicMock): + mock_engine.get_camera_fovy.return_value = None + cam = MujocoCamera(camera_name="nonexistent") + cam.set_engine(mock_engine) + cam._build_camera_info() + assert cam._camera_info_base is None + cam.stop() + + def test_distortion_and_frame_id(self, camera_with_mock_engine: MujocoCamera): + cam = camera_with_mock_engine + cam._build_camera_info() + info = cam._camera_info_base + assert info is not None + assert info.D == [0.0, 0.0, 0.0, 0.0, 0.0] + assert info.distortion_model == "plumb_bob" + assert info.frame_id == "wrist_camera_color_optical_frame" + + +# --------------------------------------------------------------------------- +# 3. Camera Lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.mujoco +class TestMujocoCameraLifecycle: + def test_start_no_engine_raises(self): + cam = MujocoCamera(camera_name="cam", address="") + try: + with pytest.raises(RuntimeError, match="set address"): + cam.start() + finally: + cam.stop() + + def test_start_creates_thread(self, camera_with_mock_engine: MujocoCamera): + cam = camera_with_mock_engine + cam._engine.connected = False # thread will wait, won't spin + # Patch rx.interval to avoid spawning scheduler threads that leak + with patch("dimos.simulation.manipulators.camera.rx.interval", return_value=MagicMock()): + cam.start() + assert cam._thread is not None + assert cam._thread.is_alive() + cam.stop() + assert cam._thread is None + + def test_stop_sets_event(self, camera_with_mock_engine: MujocoCamera): + cam = camera_with_mock_engine + cam._engine.connected = False + with patch("dimos.simulation.manipulators.camera.rx.interval", return_value=MagicMock()): + cam.start() + cam.stop() + assert cam._stop_event.is_set() + assert cam._thread is None + + def test_set_engine_injection(self, mock_engine: MagicMock): + cam = MujocoCamera(camera_name="cam") + assert cam._engine is None + cam.set_engine(mock_engine) + assert cam._engine is mock_engine + cam.stop() + + def test_frame_name_properties(self): + cam = MujocoCamera(camera_name="wrist_camera") + assert cam._camera_link == "wrist_camera_link" + assert cam._color_frame == "wrist_camera_color_frame" + assert cam._color_optical_frame == "wrist_camera_color_optical_frame" + assert cam._depth_frame == "wrist_camera_depth_frame" + assert cam._depth_optical_frame == "wrist_camera_depth_optical_frame" + cam.stop() + + def test_depth_scale(self): + cam = MujocoCamera(camera_name="cam") + assert cam.get_depth_scale() == 1.0 + cam.stop() + + +# --------------------------------------------------------------------------- +# 4. TF Publishing +# --------------------------------------------------------------------------- + + +@pytest.mark.mujoco +class TestTFPublishing: + def test_publish_tf_correct_frames(self, camera_with_mock_engine: MujocoCamera): + cam = camera_with_mock_engine + mock_tf = MagicMock() + frame = _make_camera_frame(cam_pos=[1.0, 2.0, 3.0]) + + with patch.object(type(cam), "tf", new_callable=lambda: property(lambda self: mock_tf)): + cam._publish_tf(ts=0.0, frame=frame) + + mock_tf.publish.assert_called_once() + transforms = mock_tf.publish.call_args.args + assert len(transforms) == 3 + + child_ids = {t.child_frame_id for t in transforms} + assert child_ids == { + "wrist_camera_color_optical_frame", + "wrist_camera_depth_optical_frame", + "wrist_camera_link", + } + for t in transforms: + assert t.frame_id == "world" + assert t.translation.x == pytest.approx(1.0) + assert t.translation.y == pytest.approx(2.0) + assert t.translation.z == pytest.approx(3.0) + + def test_publish_tf_none_noop(self, camera_with_mock_engine: MujocoCamera): + cam = camera_with_mock_engine + mock_tf = MagicMock() + + with patch.object(type(cam), "tf", new_callable=lambda: property(lambda self: mock_tf)): + cam._publish_tf(ts=0.0, frame=None) + + mock_tf.publish.assert_not_called() From ec1b592b1ace7ad455d08765be41c9e35c84da63 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 18:57:04 -0700 Subject: [PATCH 13/41] remove trivial tests --- .../manipulators/test_mujoco_camera.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/dimos/simulation/manipulators/test_mujoco_camera.py b/dimos/simulation/manipulators/test_mujoco_camera.py index 3b8ff8596b..c8d74a9333 100644 --- a/dimos/simulation/manipulators/test_mujoco_camera.py +++ b/dimos/simulation/manipulators/test_mujoco_camera.py @@ -166,9 +166,6 @@ def test_unregister_removes(self): for p in patches: p.stop() - def test_unregister_noop_unknown(self): - mock = MagicMock(spec=MujocoEngine) - unregister_engine(mock) # should not raise # --------------------------------------------------------------------------- @@ -203,12 +200,6 @@ def test_fovy_90(self, mock_engine: MagicMock): assert info.K[4] == pytest.approx(240.0, abs=0.01) cam.stop() - def test_no_engine_noop(self): - cam = MujocoCamera(camera_name="wrist_camera") - cam._build_camera_info() - assert cam._camera_info_base is None - cam.stop() - def test_unknown_camera(self, mock_engine: MagicMock): mock_engine.get_camera_fovy.return_value = None cam = MujocoCamera(camera_name="nonexistent") @@ -262,13 +253,6 @@ def test_stop_sets_event(self, camera_with_mock_engine: MujocoCamera): assert cam._stop_event.is_set() assert cam._thread is None - def test_set_engine_injection(self, mock_engine: MagicMock): - cam = MujocoCamera(camera_name="cam") - assert cam._engine is None - cam.set_engine(mock_engine) - assert cam._engine is mock_engine - cam.stop() - def test_frame_name_properties(self): cam = MujocoCamera(camera_name="wrist_camera") assert cam._camera_link == "wrist_camera_link" @@ -278,10 +262,6 @@ def test_frame_name_properties(self): assert cam._depth_optical_frame == "wrist_camera_depth_optical_frame" cam.stop() - def test_depth_scale(self): - cam = MujocoCamera(camera_name="cam") - assert cam.get_depth_scale() == 1.0 - cam.stop() # --------------------------------------------------------------------------- From e36d2c5cdd8291e1b8bfa6a9a9c60adfb2b8a937 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 18:58:06 -0700 Subject: [PATCH 14/41] misc: pytest fixes --- dimos/robot/all_blueprints.py | 1 + dimos/simulation/manipulators/test_sim_adapter.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 473a9b9d52..27bd48d038 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -138,6 +138,7 @@ "mock-b1-connection-module": "dimos.robot.unitree.b1.connection", "module-a": "dimos.robot.unitree.demo_error_on_name_conflicts", "module-b": "dimos.robot.unitree.demo_error_on_name_conflicts", + "mujoco-camera": "dimos.simulation.manipulators.camera", "navigation-module": "dimos.robot.unitree.rosnav", "navigation-skill-container": "dimos.agents.skills.navigation", "object-db-module": "dimos.perception.detection.moduleDB", diff --git a/dimos/simulation/manipulators/test_sim_adapter.py b/dimos/simulation/manipulators/test_sim_adapter.py index 8f253229f0..6e00e11bec 100644 --- a/dimos/simulation/manipulators/test_sim_adapter.py +++ b/dimos/simulation/manipulators/test_sim_adapter.py @@ -23,6 +23,7 @@ import pytest from dimos.hardware.manipulators.sim.adapter import SimMujocoAdapter, register +from dimos.simulation.engines.mujoco_engine import _clear_registry from dimos.simulation.utils.xml_parser import JointMapping ARM_DOF = 7 @@ -99,6 +100,12 @@ def _patch_mujoco_engine(n_joints: int): class TestSimMujocoAdapter: """Tests for SimMujocoAdapter with and without gripper.""" + @pytest.fixture(autouse=True) + def _clean_engine_registry(self): + _clear_registry() + yield + _clear_registry() + @pytest.fixture def adapter_with_gripper(self): """SimMujocoAdapter with ARM_DOF arm joints + 1 gripper joint.""" From db5f7d0c6a8feb2f0020bde8237382a0a36949f2 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 19:07:51 -0700 Subject: [PATCH 15/41] pre-commit fixes --- dimos/simulation/manipulators/test_mujoco_camera.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/dimos/simulation/manipulators/test_mujoco_camera.py b/dimos/simulation/manipulators/test_mujoco_camera.py index c8d74a9333..b88ba64fa3 100644 --- a/dimos/simulation/manipulators/test_mujoco_camera.py +++ b/dimos/simulation/manipulators/test_mujoco_camera.py @@ -167,7 +167,6 @@ def test_unregister_removes(self): p.stop() - # --------------------------------------------------------------------------- # 2. Camera Intrinsics # --------------------------------------------------------------------------- @@ -263,7 +262,6 @@ def test_frame_name_properties(self): cam.stop() - # --------------------------------------------------------------------------- # 4. TF Publishing # --------------------------------------------------------------------------- From d1e8c69518a08529607009411523a22636eb051e Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 23:35:06 -0700 Subject: [PATCH 16/41] fix: greptile issues --- dimos/manipulation/blueprints.py | 5 +--- dimos/simulation/engines/mujoco_engine.py | 29 +++++++++++++---------- dimos/simulation/manipulators/camera.py | 5 +--- 3 files changed, 19 insertions(+), 20 deletions(-) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 1572fd0b22..ee2a2fde8d 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -230,7 +230,7 @@ def _make_xarm7_config( gripper_hardware_id=gripper_hardware_id, tf_extra_links=tf_extra_links or [], pre_grasp_offset=pre_grasp_offset, - home_joints=home_joints or [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + home_joints=home_joints or [0.0] * 7, ) @@ -506,9 +506,7 @@ def _make_piper_config( ) -# --------------------------------------------------------------------------- # Sim perception: MuJoCo camera replaces RealSense for sim-based pick-and-place -# --------------------------------------------------------------------------- # Both the sim adapter and camera resolve the same MujocoEngine via the registry # (keyed by MJCF path), so they share physics state. # The engine is created lazily when the adapter connects / camera starts. @@ -528,7 +526,6 @@ def _make_piper_config( gripper_hardware_id="arm", tf_extra_links=["link7"], pre_grasp_offset=0.05, - # Observation pose: arm raised, camera above table looking down home_joints=[0.0, -0.3, 0.0, 1.0, 0.0, 0.5, 0.0], ), ], diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index d601968ef5..d38beb2963 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -17,6 +17,7 @@ from __future__ import annotations from dataclasses import dataclass +from pathlib import Path import threading import time from typing import TYPE_CHECKING @@ -31,8 +32,6 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from pathlib import Path - from dimos.msgs.sensor_msgs.JointState import JointState logger = setup_logger() @@ -90,21 +89,25 @@ def get_or_create_engine( If an engine already exists for the resolved path and *cameras* are supplied, any **new** camera configs are appended (idempotent by name). """ - from pathlib import Path - - key = str(Path(config_path).resolve()) + key = str(config_path.resolve()) with _engine_registry_lock: if key in _engine_registry: engine = _engine_registry[key] - # Merge new camera configs (by name) + if engine._headless != headless: + logger.warning( + f"get_or_create_engine: ignoring headless={headless} — " + f"existing engine for '{key}' was created with headless={engine._headless}" + ) + # Merge new camera configs (by name), guarded against sim-thread iteration if cameras: - existing_names = {c.name for c in engine._camera_configs} - for cam in cameras: - if cam.name not in existing_names: - engine._camera_configs.append(cam) + with engine._camera_lock: + existing_names = {c.name for c in engine._camera_configs} + for cam in cameras: + if cam.name not in existing_names: + engine._camera_configs.append(cam) return engine - engine = MujocoEngine(config_path=Path(config_path), headless=headless, cameras=cameras) + engine = MujocoEngine(config_path=config_path, headless=headless, cameras=cameras) _engine_registry[key] = engine return engine @@ -277,7 +280,9 @@ def _init_new_cameras(self, cam_renderers: dict[str, _CameraRendererState]) -> N Must be called from the sim thread (MuJoCo thread-safety). """ - for cfg in self._camera_configs: + with self._camera_lock: + configs_snapshot = list(self._camera_configs) + for cfg in configs_snapshot: if cfg.name in cam_renderers: continue cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name) diff --git a/dimos/simulation/manipulators/camera.py b/dimos/simulation/manipulators/camera.py index ff73b23890..ecede34fdd 100644 --- a/dimos/simulation/manipulators/camera.py +++ b/dimos/simulation/manipulators/camera.py @@ -197,10 +197,7 @@ def start(self) -> None: "MujocoCamera: set address (MJCF path) in config or call set_engine()" ) - logger.info( - f"MujocoCamera: engine resolved, connected={self._engine.connected}, " - f"cameras={[c.name for c in self._engine._camera_configs]}" - ) + logger.info(f"MujocoCamera: engine resolved, connected={self._engine.connected}") self._build_camera_info() From 4ce2d647868993ddedb157ca37e9d20a49c2fadc Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 23:37:47 -0700 Subject: [PATCH 17/41] fix: pytest issues --- dimos/simulation/engines/mujoco_engine.py | 5 ----- .../manipulators/test_mujoco_camera.py | 16 ---------------- 2 files changed, 21 deletions(-) diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index d38beb2963..f7037710ee 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -70,11 +70,6 @@ class _CameraRendererState: interval: float last_render_time: float = 0.0 - -# --------------------------------------------------------------------------- -# Engine registry — allows adapter and camera to share the same engine instance -# by MJCF path. -# --------------------------------------------------------------------------- _engine_registry: dict[str, MujocoEngine] = {} _engine_registry_lock = threading.Lock() diff --git a/dimos/simulation/manipulators/test_mujoco_camera.py b/dimos/simulation/manipulators/test_mujoco_camera.py index b88ba64fa3..7c4da9001d 100644 --- a/dimos/simulation/manipulators/test_mujoco_camera.py +++ b/dimos/simulation/manipulators/test_mujoco_camera.py @@ -82,10 +82,6 @@ def camera_with_mock_engine(mock_engine: MagicMock): cam.stop() -# --------------------------------------------------------------------------- -# 1. Engine Registry -# --------------------------------------------------------------------------- - @pytest.mark.mujoco class TestEngineRegistry: @@ -167,10 +163,6 @@ def test_unregister_removes(self): p.stop() -# --------------------------------------------------------------------------- -# 2. Camera Intrinsics -# --------------------------------------------------------------------------- - @pytest.mark.mujoco class TestCameraIntrinsics: @@ -217,10 +209,6 @@ def test_distortion_and_frame_id(self, camera_with_mock_engine: MujocoCamera): assert info.frame_id == "wrist_camera_color_optical_frame" -# --------------------------------------------------------------------------- -# 3. Camera Lifecycle -# --------------------------------------------------------------------------- - @pytest.mark.mujoco class TestMujocoCameraLifecycle: @@ -262,10 +250,6 @@ def test_frame_name_properties(self): cam.stop() -# --------------------------------------------------------------------------- -# 4. TF Publishing -# --------------------------------------------------------------------------- - @pytest.mark.mujoco class TestTFPublishing: From 220cee3446ca46d507227f14ea687d3b360bcd4c Mon Sep 17 00:00:00 2001 From: ruthwikdasyam <63036454+ruthwikdasyam@users.noreply.github.com> Date: Sat, 28 Mar 2026 06:38:44 +0000 Subject: [PATCH 18/41] CI code cleanup --- dimos/simulation/engines/mujoco_engine.py | 1 + dimos/simulation/manipulators/test_mujoco_camera.py | 4 ---- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index f7037710ee..1cb61ebf45 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -70,6 +70,7 @@ class _CameraRendererState: interval: float last_render_time: float = 0.0 + _engine_registry: dict[str, MujocoEngine] = {} _engine_registry_lock = threading.Lock() diff --git a/dimos/simulation/manipulators/test_mujoco_camera.py b/dimos/simulation/manipulators/test_mujoco_camera.py index 7c4da9001d..2aee68bc9c 100644 --- a/dimos/simulation/manipulators/test_mujoco_camera.py +++ b/dimos/simulation/manipulators/test_mujoco_camera.py @@ -82,7 +82,6 @@ def camera_with_mock_engine(mock_engine: MagicMock): cam.stop() - @pytest.mark.mujoco class TestEngineRegistry: def test_creates_new(self): @@ -163,7 +162,6 @@ def test_unregister_removes(self): p.stop() - @pytest.mark.mujoco class TestCameraIntrinsics: def test_fovy_45(self, camera_with_mock_engine: MujocoCamera): @@ -209,7 +207,6 @@ def test_distortion_and_frame_id(self, camera_with_mock_engine: MujocoCamera): assert info.frame_id == "wrist_camera_color_optical_frame" - @pytest.mark.mujoco class TestMujocoCameraLifecycle: def test_start_no_engine_raises(self): @@ -250,7 +247,6 @@ def test_frame_name_properties(self): cam.stop() - @pytest.mark.mujoco class TestTFPublishing: def test_publish_tf_correct_frames(self, camera_with_mock_engine: MujocoCamera): From 5cd81143bae73fe244811f42728546ddfa2ea3b1 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Fri, 27 Mar 2026 23:59:24 -0700 Subject: [PATCH 19/41] xarm7 updated mujoco env --- data/.lfs/xarm7.tar.gz | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/data/.lfs/xarm7.tar.gz b/data/.lfs/xarm7.tar.gz index 8e2cfa368a..897f052bb8 100644 --- a/data/.lfs/xarm7.tar.gz +++ b/data/.lfs/xarm7.tar.gz @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:47dd79f13845ae6a35368345b7443a9190c7584d548caddd9c3eae224442c6fc -size 3280557 +oid sha256:c97e2283c0a726afd48e91172f84605765b8af8ace7ac107b810a8d11869bc99 +size 1606344 From 359cdebfb6582fed87191e2978c110a341c4c01b Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Sat, 28 Mar 2026 10:50:51 -0700 Subject: [PATCH 20/41] fix: engine registered when started from mujococamera --- dimos/simulation/manipulators/camera.py | 12 ++++++++++++ dimos/simulation/manipulators/test_mujoco_camera.py | 2 ++ 2 files changed, 14 insertions(+) diff --git a/dimos/simulation/manipulators/camera.py b/dimos/simulation/manipulators/camera.py index ecede34fdd..3449e960bb 100644 --- a/dimos/simulation/manipulators/camera.py +++ b/dimos/simulation/manipulators/camera.py @@ -197,6 +197,18 @@ def start(self) -> None: "MujocoCamera: set address (MJCF path) in config or call set_engine()" ) + # Ensure camera config is registered (handles the set_engine() path) + cam_cfg = CameraConfig( + name=self.config.camera_name, + width=self.config.width, + height=self.config.height, + fps=self.config.fps, + ) + with self._engine._camera_lock: + existing_names = {c.name for c in self._engine._camera_configs} + if cam_cfg.name not in existing_names: + self._engine._camera_configs.append(cam_cfg) + logger.info(f"MujocoCamera: engine resolved, connected={self._engine.connected}") self._build_camera_info() diff --git a/dimos/simulation/manipulators/test_mujoco_camera.py b/dimos/simulation/manipulators/test_mujoco_camera.py index 2aee68bc9c..c7d65aed16 100644 --- a/dimos/simulation/manipulators/test_mujoco_camera.py +++ b/dimos/simulation/manipulators/test_mujoco_camera.py @@ -18,6 +18,7 @@ import math from pathlib import Path +import threading from unittest.mock import MagicMock, patch import numpy as np @@ -65,6 +66,7 @@ def _make_mock_engine(fovy: float = 45.0) -> MagicMock: mock_engine.get_camera_fovy.return_value = fovy mock_engine.connected = True mock_engine._camera_configs = [] + mock_engine._camera_lock = threading.Lock() mock_engine.read_camera.return_value = _make_camera_frame() return mock_engine From cb9ac93715e910d36085ee47b904e43f601c13af Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Sat, 28 Mar 2026 11:40:39 -0700 Subject: [PATCH 21/41] feat: mcp control --- dimos/manipulation/blueprints.py | 8 ++++++++ dimos/robot/all_blueprints.py | 1 + 2 files changed, 9 insertions(+) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index ee2a2fde8d..4d6225c9db 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -560,6 +560,13 @@ def _make_piper_config( ) +xarm_perception_sim_agent = autoconnect( + xarm_perception_sim, + McpServer.blueprint(), + McpClient.blueprint(system_prompt=_MANIPULATION_AGENT_SYSTEM_PROMPT), +) + + __all__ = [ "PIPER_GRIPPER_COLLISION_EXCLUSIONS", "XARM_GRIPPER_COLLISION_EXCLUSIONS", @@ -570,4 +577,5 @@ def _make_piper_config( "xarm_perception", "xarm_perception_agent", "xarm_perception_sim", + "xarm_perception_sim_agent", ] diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 27bd48d038..c97d7d0d55 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -92,6 +92,7 @@ "xarm-perception": "dimos.manipulation.blueprints:xarm_perception", "xarm-perception-agent": "dimos.manipulation.blueprints:xarm_perception_agent", "xarm-perception-sim": "dimos.manipulation.blueprints:xarm_perception_sim", + "xarm-perception-sim-agent": "dimos.manipulation.blueprints:xarm_perception_sim_agent", "xarm6-planner-only": "dimos.manipulation.blueprints:xarm6_planner_only", "xarm7-planner-coordinator": "dimos.manipulation.blueprints:xarm7_planner_coordinator", "xarm7-planner-coordinator-agent": "dimos.manipulation.blueprints:xarm7_planner_coordinator_agent", From a0d0c4a2d09a80e3e96445acb733643616ffe0b5 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Sat, 28 Mar 2026 14:46:07 -0700 Subject: [PATCH 22/41] fix: planner DOF mismatch, tune home joints --- dimos/manipulation/blueprints.py | 2 +- dimos/manipulation/manipulation_module.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 4d6225c9db..5f0b505574 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -526,7 +526,7 @@ def _make_piper_config( gripper_hardware_id="arm", tf_extra_links=["link7"], pre_grasp_offset=0.05, - home_joints=[0.0, -0.3, 0.0, 1.0, 0.0, 0.5, 0.0], + home_joints=[0.0, 0.0, 0.0, 0.0, 0.0, -0.7, 0.0], ), ], planning_timeout=10.0, diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 25d1983a86..a5dc6cc3e2 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -471,6 +471,14 @@ def _plan_path_only( if start is None: return self._fail("No joint state") + # Trim goal to planner DOF (e.g. strip gripper joint from coordinator state) + planner_dof = len(start.position) + if len(goal.position) > planner_dof: + goal = JointState( + name=list(goal.name[:planner_dof]) if goal.name else [], + position=list(goal.position[:planner_dof]), + ) + result = self._planner.plan_joint_path( world=self._world_monitor.world, robot_id=robot_id, From db6fd4258cc0e00735f8cc71e1f7970a8feefc95 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 30 Mar 2026 10:50:53 -0700 Subject: [PATCH 23/41] fix: updated mujoco-camera to match latest module def --- dimos/robot/all_blueprints.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 14d81ce98c..cb7782fd8f 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -139,6 +139,7 @@ "mock-b1-connection-module": "dimos.robot.unitree.b1.connection.MockB1ConnectionModule", "module-a": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleA", "module-b": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleB", + "mujoco-camera": "dimos.simulation.manipulators.camera.MujocoCamera", "navigation-module": "dimos.robot.unitree.rosnav.NavigationModule", "navigation-skill-container": "dimos.agents.skills.navigation.NavigationSkillContainer", "object-db-module": "dimos.perception.detection.moduleDB.ObjectDBModule", From ea839564ed63b4bb705903028d5c12ebfa2ccbc7 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 30 Mar 2026 16:13:59 -0700 Subject: [PATCH 24/41] misc: simcamera module path change --- dimos/manipulation/blueprints.py | 2 +- dimos/robot/all_blueprints.py | 2 +- .../{manipulators/camera.py => sensors/mujoco_camera.py} | 0 .../{manipulators => sensors}/test_mujoco_camera.py | 6 +++--- 4 files changed, 5 insertions(+), 5 deletions(-) rename dimos/simulation/{manipulators/camera.py => sensors/mujoco_camera.py} (100%) rename dimos/simulation/{manipulators => sensors}/test_mujoco_camera.py (97%) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 5f0b505574..fffa57dbd1 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -512,7 +512,7 @@ def _make_piper_config( # The engine is created lazily when the adapter connects / camera starts. from dimos.control.blueprints._hardware import XARM7_SIM_PATH, sim_xarm7 -from dimos.simulation.manipulators.camera import MujocoCamera +from dimos.simulation.sensors.mujoco_camera import MujocoCamera from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode xarm_perception_sim = autoconnect( diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index cb7782fd8f..bd31e9a518 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -139,7 +139,7 @@ "mock-b1-connection-module": "dimos.robot.unitree.b1.connection.MockB1ConnectionModule", "module-a": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleA", "module-b": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleB", - "mujoco-camera": "dimos.simulation.manipulators.camera.MujocoCamera", + "mujoco-camera": "dimos.simulation.sensors.mujoco_camera.MujocoCamera", "navigation-module": "dimos.robot.unitree.rosnav.NavigationModule", "navigation-skill-container": "dimos.agents.skills.navigation.NavigationSkillContainer", "object-db-module": "dimos.perception.detection.moduleDB.ObjectDBModule", diff --git a/dimos/simulation/manipulators/camera.py b/dimos/simulation/sensors/mujoco_camera.py similarity index 100% rename from dimos/simulation/manipulators/camera.py rename to dimos/simulation/sensors/mujoco_camera.py diff --git a/dimos/simulation/manipulators/test_mujoco_camera.py b/dimos/simulation/sensors/test_mujoco_camera.py similarity index 97% rename from dimos/simulation/manipulators/test_mujoco_camera.py rename to dimos/simulation/sensors/test_mujoco_camera.py index c7d65aed16..fe09fc1890 100644 --- a/dimos/simulation/manipulators/test_mujoco_camera.py +++ b/dimos/simulation/sensors/test_mujoco_camera.py @@ -33,8 +33,8 @@ get_or_create_engine, unregister_engine, ) -from dimos.simulation.manipulators.camera import MujocoCamera from dimos.simulation.manipulators.test_sim_adapter import _patch_mujoco_engine +from dimos.simulation.sensors.mujoco_camera import MujocoCamera ARM_DOF = 7 @@ -223,7 +223,7 @@ def test_start_creates_thread(self, camera_with_mock_engine: MujocoCamera): cam = camera_with_mock_engine cam._engine.connected = False # thread will wait, won't spin # Patch rx.interval to avoid spawning scheduler threads that leak - with patch("dimos.simulation.manipulators.camera.rx.interval", return_value=MagicMock()): + with patch("dimos.simulation.sensors.mujoco_camera.rx.interval", return_value=MagicMock()): cam.start() assert cam._thread is not None assert cam._thread.is_alive() @@ -233,7 +233,7 @@ def test_start_creates_thread(self, camera_with_mock_engine: MujocoCamera): def test_stop_sets_event(self, camera_with_mock_engine: MujocoCamera): cam = camera_with_mock_engine cam._engine.connected = False - with patch("dimos.simulation.manipulators.camera.rx.interval", return_value=MagicMock()): + with patch("dimos.simulation.sensors.mujoco_camera.rx.interval", return_value=MagicMock()): cam.start() cam.stop() assert cam._stop_event.is_set() From 774fe96362b9796634d2923a1485b31f7563417d Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 30 Mar 2026 16:34:30 -0700 Subject: [PATCH 25/41] fix: mypy taskconfig import --- dimos/manipulation/blueprints.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 579ed44dd0..c327ac5c25 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -500,6 +500,7 @@ def _make_piper_config( # The engine is created lazily when the adapter connects / camera starts. from dimos.control.blueprints._hardware import XARM7_SIM_PATH, sim_xarm7 +from dimos.control.coordinator import TaskConfig as TaskConfig from dimos.simulation.sensors.mujoco_camera import MujocoCamera from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode From 66fb4542c9a5876f7016aeff48c15f773477d0bd Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 31 Mar 2026 12:55:00 -0700 Subject: [PATCH 26/41] fix: remove set engine --- dimos/simulation/sensors/mujoco_camera.py | 7 +------ dimos/simulation/sensors/test_mujoco_camera.py | 6 +++--- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/dimos/simulation/sensors/mujoco_camera.py b/dimos/simulation/sensors/mujoco_camera.py index 3449e960bb..76b226ef2e 100644 --- a/dimos/simulation/sensors/mujoco_camera.py +++ b/dimos/simulation/sensors/mujoco_camera.py @@ -89,8 +89,7 @@ class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.D be used as a drop-in replacement in manipulation blueprints. The engine is resolved automatically from the registry via ``address`` - (the same MJCF path used by the sim_mujoco adapter). Alternatively, - call ``set_engine()`` before ``start()``. + (the same MJCF path used by the sim_mujoco adapter). """ color_image: Out[Image] @@ -108,10 +107,6 @@ def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] self._thread: threading.Thread | None = None self._camera_info_base: CameraInfo | None = None - def set_engine(self, engine: MujocoEngine) -> None: - """Inject the shared MujocoEngine reference.""" - self._engine = engine - @property def _camera_link(self) -> str: return f"{self.config.camera_name}_link" diff --git a/dimos/simulation/sensors/test_mujoco_camera.py b/dimos/simulation/sensors/test_mujoco_camera.py index fe09fc1890..cc36cbe8d7 100644 --- a/dimos/simulation/sensors/test_mujoco_camera.py +++ b/dimos/simulation/sensors/test_mujoco_camera.py @@ -79,7 +79,7 @@ def mock_engine() -> MagicMock: @pytest.fixture def camera_with_mock_engine(mock_engine: MagicMock): cam = MujocoCamera(camera_name="wrist_camera") - cam.set_engine(mock_engine) + cam._engine = mock_engine yield cam cam.stop() @@ -182,7 +182,7 @@ def test_fovy_45(self, camera_with_mock_engine: MujocoCamera): def test_fovy_90(self, mock_engine: MagicMock): mock_engine.get_camera_fovy.return_value = 90.0 cam = MujocoCamera(camera_name="wrist_camera") - cam.set_engine(mock_engine) + cam._engine = mock_engine cam._build_camera_info() info = cam._camera_info_base assert info is not None @@ -194,7 +194,7 @@ def test_fovy_90(self, mock_engine: MagicMock): def test_unknown_camera(self, mock_engine: MagicMock): mock_engine.get_camera_fovy.return_value = None cam = MujocoCamera(camera_name="nonexistent") - cam.set_engine(mock_engine) + cam._engine = mock_engine cam._build_camera_info() assert cam._camera_info_base is None cam.stop() From 9cce0ab2c93cbe3ad7bfe12596f0aa55cb435f75 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 31 Mar 2026 12:58:10 -0700 Subject: [PATCH 27/41] fix: removing all cache entities --- dimos/msgs/sensor_msgs/PointCloud2.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/dimos/msgs/sensor_msgs/PointCloud2.py b/dimos/msgs/sensor_msgs/PointCloud2.py index 2903420b15..f9641e9c39 100644 --- a/dimos/msgs/sensor_msgs/PointCloud2.py +++ b/dimos/msgs/sensor_msgs/PointCloud2.py @@ -100,13 +100,10 @@ def __getstate__(self) -> dict[str, object]: # Remove non-picklable objects del state["_pcd_tensor"] state["_pcd_legacy_cache"] = None - # Remove cached_property entries that hold Open3D C++ objects - for key in ( - "axis_aligned_bounding_box", - "oriented_bounding_box", - "bounding_box_dimensions", - ): - state.pop(key, None) + # Remove all cached_property entries + for key in list(state): + if isinstance(getattr(type(self), key, None), functools.cached_property): + del state[key] return state def __setstate__(self, state: dict[str, object]) -> None: From 84c263e8ef175ad97da6637b443d3dbb81a6a730 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 31 Mar 2026 13:48:54 -0700 Subject: [PATCH 28/41] test: conftest addition --- dimos/simulation/conftest.py | 104 ++++++++++++++++++ .../manipulators/test_sim_adapter.py | 81 +------------- .../simulation/sensors/test_mujoco_camera.py | 14 +-- 3 files changed, 106 insertions(+), 93 deletions(-) create mode 100644 dimos/simulation/conftest.py diff --git a/dimos/simulation/conftest.py b/dimos/simulation/conftest.py new file mode 100644 index 0000000000..f5f5bb33bf --- /dev/null +++ b/dimos/simulation/conftest.py @@ -0,0 +1,104 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared test fixtures for simulation tests.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from dimos.simulation.engines.mujoco_engine import _clear_registry +from dimos.simulation.utils.xml_parser import JointMapping + +ARM_DOF = 7 + + +def _make_joint_mapping(name: str, idx: int) -> JointMapping: + """Create a JointMapping for a simple revolute joint.""" + return JointMapping( + name=name, + joint_id=idx, + actuator_id=idx, + qpos_adr=idx, + dof_adr=idx, + tendon_qpos_adrs=(), + tendon_dof_adrs=(), + ) + + +def _make_gripper_mapping(name: str, idx: int) -> JointMapping: + """Create a JointMapping for a tendon-driven gripper.""" + return JointMapping( + name=name, + joint_id=None, + actuator_id=idx, + qpos_adr=None, + dof_adr=None, + tendon_qpos_adrs=(idx, idx + 1), + tendon_dof_adrs=(idx, idx + 1), + ) + + +def _patch_mujoco_engine(n_joints: int = ARM_DOF): + """Patch only the MuJoCo C-library and filesystem boundaries. + + Mocks ``_resolve_xml_path``, ``MjModel.from_xml_path``, ``MjData``, and + ``build_joint_mappings`` — the rest of ``MujocoEngine.__init__`` runs as-is. + """ + mappings = [_make_joint_mapping(f"joint{i}", i) for i in range(ARM_DOF)] + if n_joints > ARM_DOF: + mappings.append(_make_gripper_mapping(f"joint{ARM_DOF}", ARM_DOF)) + + fake_model = MagicMock() + fake_model.opt.timestep = 0.002 + fake_model.nu = n_joints + fake_model.nq = n_joints + fake_model.njnt = n_joints + fake_model.actuator_ctrlrange = np.array( + [[-6.28, 6.28]] * ARM_DOF + ([[0.0, 255.0]] if n_joints > ARM_DOF else []) + ) + fake_model.jnt_range = np.array( + [[-6.28, 6.28]] * ARM_DOF + ([[0.0, 0.85]] if n_joints > ARM_DOF else []) + ) + fake_model.jnt_qposadr = np.arange(n_joints) + + fake_data = MagicMock() + fake_data.qpos = np.zeros(n_joints + 4) # extra for tendon qpos addresses + fake_data.actuator_length = np.zeros(n_joints) + + patches = [ + patch( + "dimos.simulation.engines.mujoco_engine.MujocoEngine._resolve_xml_path", + return_value=Path("/fake/scene.xml"), + ), + patch( + "dimos.simulation.engines.mujoco_engine.mujoco.MjModel.from_xml_path", + return_value=fake_model, + ), + patch("dimos.simulation.engines.mujoco_engine.mujoco.MjData", return_value=fake_data), + patch("dimos.simulation.engines.mujoco_engine.build_joint_mappings", return_value=mappings), + ] + return patches + + +@pytest.fixture(autouse=True) +def clean_engine_registry(): + """Clear the engine registry before and after each simulation test.""" + _clear_registry() + yield + _clear_registry() diff --git a/dimos/simulation/manipulators/test_sim_adapter.py b/dimos/simulation/manipulators/test_sim_adapter.py index 6e00e11bec..c9a8a29474 100644 --- a/dimos/simulation/manipulators/test_sim_adapter.py +++ b/dimos/simulation/manipulators/test_sim_adapter.py @@ -16,96 +16,17 @@ from __future__ import annotations -from pathlib import Path from unittest.mock import MagicMock, patch -import numpy as np import pytest from dimos.hardware.manipulators.sim.adapter import SimMujocoAdapter, register -from dimos.simulation.engines.mujoco_engine import _clear_registry -from dimos.simulation.utils.xml_parser import JointMapping - -ARM_DOF = 7 - - -def _make_joint_mapping(name: str, idx: int) -> JointMapping: - """Create a JointMapping for a simple revolute joint.""" - return JointMapping( - name=name, - joint_id=idx, - actuator_id=idx, - qpos_adr=idx, - dof_adr=idx, - tendon_qpos_adrs=(), - tendon_dof_adrs=(), - ) - - -def _make_gripper_mapping(name: str, idx: int) -> JointMapping: - """Create a JointMapping for a tendon-driven gripper.""" - return JointMapping( - name=name, - joint_id=None, - actuator_id=idx, - qpos_adr=None, - dof_adr=None, - tendon_qpos_adrs=(idx, idx + 1), - tendon_dof_adrs=(idx, idx + 1), - ) - - -def _patch_mujoco_engine(n_joints: int): - """Patch only the MuJoCo C-library and filesystem boundaries. - - Mocks ``_resolve_xml_path``, ``MjModel.from_xml_path``, ``MjData``, and - ``build_joint_mappings`` — the rest of ``MujocoEngine.__init__`` runs as-is. - """ - mappings = [_make_joint_mapping(f"joint{i}", i) for i in range(ARM_DOF)] - if n_joints > ARM_DOF: - mappings.append(_make_gripper_mapping(f"joint{ARM_DOF}", ARM_DOF)) - - fake_model = MagicMock() - fake_model.opt.timestep = 0.002 - fake_model.nu = n_joints - fake_model.nq = n_joints - fake_model.njnt = n_joints - fake_model.actuator_ctrlrange = np.array( - [[-6.28, 6.28]] * ARM_DOF + ([[0.0, 255.0]] if n_joints > ARM_DOF else []) - ) - fake_model.jnt_range = np.array( - [[-6.28, 6.28]] * ARM_DOF + ([[0.0, 0.85]] if n_joints > ARM_DOF else []) - ) - fake_model.jnt_qposadr = np.arange(n_joints) - - fake_data = MagicMock() - fake_data.qpos = np.zeros(n_joints + 4) # extra for tendon qpos addresses - fake_data.actuator_length = np.zeros(n_joints) - - patches = [ - patch( - "dimos.simulation.engines.mujoco_engine.MujocoEngine._resolve_xml_path", - return_value=Path("/fake/scene.xml"), - ), - patch( - "dimos.simulation.engines.mujoco_engine.mujoco.MjModel.from_xml_path", - return_value=fake_model, - ), - patch("dimos.simulation.engines.mujoco_engine.mujoco.MjData", return_value=fake_data), - patch("dimos.simulation.engines.mujoco_engine.build_joint_mappings", return_value=mappings), - ] - return patches +from dimos.simulation.conftest import ARM_DOF, _patch_mujoco_engine class TestSimMujocoAdapter: """Tests for SimMujocoAdapter with and without gripper.""" - @pytest.fixture(autouse=True) - def _clean_engine_registry(self): - _clear_registry() - yield - _clear_registry() - @pytest.fixture def adapter_with_gripper(self): """SimMujocoAdapter with ARM_DOF arm joints + 1 gripper joint.""" diff --git a/dimos/simulation/sensors/test_mujoco_camera.py b/dimos/simulation/sensors/test_mujoco_camera.py index cc36cbe8d7..c6877195f6 100644 --- a/dimos/simulation/sensors/test_mujoco_camera.py +++ b/dimos/simulation/sensors/test_mujoco_camera.py @@ -24,28 +24,20 @@ import numpy as np import pytest +from dimos.simulation.conftest import _patch_mujoco_engine from dimos.simulation.engines.mujoco_engine import ( CameraConfig, CameraFrame, MujocoEngine, - _clear_registry, _engine_registry, get_or_create_engine, unregister_engine, ) -from dimos.simulation.manipulators.test_sim_adapter import _patch_mujoco_engine from dimos.simulation.sensors.mujoco_camera import MujocoCamera ARM_DOF = 7 -@pytest.fixture(autouse=True) -def clean_registry(): - _clear_registry() - yield - _clear_registry() - - def _make_camera_frame( cam_pos: list[float] | None = None, cam_mat: list[float] | None = None, @@ -84,7 +76,6 @@ def camera_with_mock_engine(mock_engine: MagicMock): cam.stop() -@pytest.mark.mujoco class TestEngineRegistry: def test_creates_new(self): patches = _patch_mujoco_engine(ARM_DOF) @@ -164,7 +155,6 @@ def test_unregister_removes(self): p.stop() -@pytest.mark.mujoco class TestCameraIntrinsics: def test_fovy_45(self, camera_with_mock_engine: MujocoCamera): cam = camera_with_mock_engine @@ -209,7 +199,6 @@ def test_distortion_and_frame_id(self, camera_with_mock_engine: MujocoCamera): assert info.frame_id == "wrist_camera_color_optical_frame" -@pytest.mark.mujoco class TestMujocoCameraLifecycle: def test_start_no_engine_raises(self): cam = MujocoCamera(camera_name="cam", address="") @@ -249,7 +238,6 @@ def test_frame_name_properties(self): cam.stop() -@pytest.mark.mujoco class TestTFPublishing: def test_publish_tf_correct_frames(self, camera_with_mock_engine: MujocoCamera): cam = camera_with_mock_engine From 349eddc259e3719071375f169e9f2ad179b33e4d Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 31 Mar 2026 14:14:56 -0700 Subject: [PATCH 29/41] misc: imports + structlog logging --- dimos/simulation/sensors/mujoco_camera.py | 27 ++++++++++++----------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/dimos/simulation/sensors/mujoco_camera.py b/dimos/simulation/sensors/mujoco_camera.py index 76b226ef2e..237a6987d4 100644 --- a/dimos/simulation/sensors/mujoco_camera.py +++ b/dimos/simulation/sensors/mujoco_camera.py @@ -23,6 +23,8 @@ import math import threading import time +from pathlib import Path +from typing import Any from pydantic import Field import reactivex as rx @@ -100,8 +102,8 @@ class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.D default_config = MujocoCameraConfig - def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] - super().__init__(*args, **kwargs) + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) self._engine: MujocoEngine | None = None self._stop_event = threading.Event() self._thread: threading.Thread | None = None @@ -145,7 +147,7 @@ def _build_camera_info(self) -> None: return fovy_deg = self._engine.get_camera_fovy(self.config.camera_name) if fovy_deg is None: - logger.error(f"Camera '{self.config.camera_name}' not found in MJCF") + logger.error("Camera not found in MJCF", camera_name=self.config.camera_name) return h = self.config.height @@ -173,8 +175,6 @@ def _build_camera_info(self) -> None: @rpc def start(self) -> None: if self._engine is None and self.config.address: - from pathlib import Path - self._engine = get_or_create_engine( config_path=Path(self.config.address), headless=self.config.headless, @@ -204,7 +204,7 @@ def start(self) -> None: if cam_cfg.name not in existing_names: self._engine._camera_configs.append(cam_cfg) - logger.info(f"MujocoCamera: engine resolved, connected={self._engine.connected}") + logger.info("MujocoCamera engine resolved", connected=self._engine.connected) self._build_camera_info() @@ -217,7 +217,7 @@ def start(self) -> None: self._disposables.add( rx.interval(interval_sec).subscribe( on_next=lambda _: self._publish_camera_info(), - on_error=lambda e: logger.error(f"CameraInfo publish error: {e}"), + on_error=lambda e: logger.error("CameraInfo publish error", error=e), ) ) @@ -227,7 +227,7 @@ def start(self) -> None: self._disposables.add( backpressure(rx.interval(pc_interval)).subscribe( on_next=lambda _: self._generate_pointcloud(), - on_error=lambda e: logger.error(f"Pointcloud error: {e}"), + on_error=lambda e: logger.error("Pointcloud error", error=e), ) ) @@ -246,7 +246,7 @@ def _publish_loop(self) -> None: last_timestamp = 0.0 published_count = 0 - logger.info(f"MujocoCamera publish loop started (camera={self.config.camera_name})") + logger.info("MujocoCamera publish loop started", camera=self.config.camera_name) # Wait for engine to connect (adapter may not have started yet) deadline = time.monotonic() + 30.0 @@ -300,8 +300,9 @@ def _publish_loop(self) -> None: published_count += 1 if published_count == 1: logger.info( - f"MujocoCamera: first frame published " - f"(rgb={frame.rgb.shape}, depth={frame.depth.shape})" + "MujocoCamera first frame published", + rgb_shape=frame.rgb.shape, + depth_shape=frame.depth.shape, ) # Rate limit @@ -310,7 +311,7 @@ def _publish_loop(self) -> None: if sleep_time > 0: time.sleep(sleep_time) except Exception as e: - logger.error(f"MujocoCamera publish error: {e}") + logger.error("MujocoCamera publish error", error=e) time.sleep(1.0) def _publish_camera_info(self) -> None: @@ -400,7 +401,7 @@ def _generate_pointcloud(self) -> None: pcd = pcd.voxel_downsample(0.005) self.pointcloud.publish(pcd) except Exception as e: - logger.error(f"Pointcloud generation error: {e}") + logger.error("Pointcloud generation error", error=e) __all__ = ["MujocoCamera", "MujocoCameraConfig"] From ac9144f1cbafd88ede49b7629fe0b06ab59903e6 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 31 Mar 2026 14:15:25 -0700 Subject: [PATCH 30/41] test: added fixture, removed dublicate codes --- dimos/simulation/conftest.py | 27 ++++ .../manipulators/test_sim_adapter.py | 41 ++---- .../simulation/sensors/test_mujoco_camera.py | 118 ++++++------------ 3 files changed, 77 insertions(+), 109 deletions(-) diff --git a/dimos/simulation/conftest.py b/dimos/simulation/conftest.py index f5f5bb33bf..6a12c3d1cc 100644 --- a/dimos/simulation/conftest.py +++ b/dimos/simulation/conftest.py @@ -102,3 +102,30 @@ def clean_engine_registry(): _clear_registry() yield _clear_registry() + + +@pytest.fixture +def patched_mujoco_engine(): + """Start MuJoCo engine patches for the test, stop on teardown. + + Usage: + def test_something(patched_mujoco_engine): + engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) + """ + patches = _patch_mujoco_engine(ARM_DOF) + for p in patches: + p.start() + yield + for p in patches: + p.stop() + + +@pytest.fixture +def patched_mujoco_engine_with_gripper(): + """Same as patched_mujoco_engine but with ARM_DOF+1 joints (gripper).""" + patches = _patch_mujoco_engine(ARM_DOF + 1) + for p in patches: + p.start() + yield + for p in patches: + p.stop() diff --git a/dimos/simulation/manipulators/test_sim_adapter.py b/dimos/simulation/manipulators/test_sim_adapter.py index c9a8a29474..1413004fee 100644 --- a/dimos/simulation/manipulators/test_sim_adapter.py +++ b/dimos/simulation/manipulators/test_sim_adapter.py @@ -21,48 +21,25 @@ import pytest from dimos.hardware.manipulators.sim.adapter import SimMujocoAdapter, register -from dimos.simulation.conftest import ARM_DOF, _patch_mujoco_engine +from dimos.simulation.conftest import ARM_DOF class TestSimMujocoAdapter: """Tests for SimMujocoAdapter with and without gripper.""" @pytest.fixture - def adapter_with_gripper(self): + def adapter_with_gripper(self, patched_mujoco_engine_with_gripper): """SimMujocoAdapter with ARM_DOF arm joints + 1 gripper joint.""" - patches = _patch_mujoco_engine(ARM_DOF + 1) - for p in patches: - p.start() - try: - adapter = SimMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml", headless=True) - finally: - for p in patches: - p.stop() - return adapter + return SimMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml", headless=True) @pytest.fixture - def adapter_no_gripper(self): + def adapter_no_gripper(self, patched_mujoco_engine): """SimMujocoAdapter with ARM_DOF arm joints, no gripper.""" - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - try: - adapter = SimMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml", headless=True) - finally: - for p in patches: - p.stop() - return adapter - - def test_address_required(self): - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - try: - with pytest.raises(ValueError, match="address"): - SimMujocoAdapter(dof=ARM_DOF, address=None) - finally: - for p in patches: - p.stop() + return SimMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml", headless=True) + + def test_address_required(self, patched_mujoco_engine): + with pytest.raises(ValueError, match="address"): + SimMujocoAdapter(dof=ARM_DOF, address=None) def test_gripper_detected(self, adapter_with_gripper): assert adapter_with_gripper._gripper_idx == ARM_DOF diff --git a/dimos/simulation/sensors/test_mujoco_camera.py b/dimos/simulation/sensors/test_mujoco_camera.py index c6877195f6..a67f673966 100644 --- a/dimos/simulation/sensors/test_mujoco_camera.py +++ b/dimos/simulation/sensors/test_mujoco_camera.py @@ -24,7 +24,6 @@ import numpy as np import pytest -from dimos.simulation.conftest import _patch_mujoco_engine from dimos.simulation.engines.mujoco_engine import ( CameraConfig, CameraFrame, @@ -77,82 +76,47 @@ def camera_with_mock_engine(mock_engine: MagicMock): class TestEngineRegistry: - def test_creates_new(self): - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - try: - engine = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) - assert engine is not None - assert len(_engine_registry) == 1 - finally: - for p in patches: - p.stop() - - def test_returns_same_instance(self): - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - try: - e1 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) - e2 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) - assert e1 is e2 - assert len(_engine_registry) == 1 - finally: - for p in patches: - p.stop() - - def test_merges_new_cameras(self): - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - try: - e1 = get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_a")], - ) - get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_b")], - ) - names = {c.name for c in e1._camera_configs} - assert names == {"cam_a", "cam_b"} - finally: - for p in patches: - p.stop() - - def test_deduplicates_cameras(self): - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - try: - get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_a")], - ) - get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_a")], - ) - engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) - cam_names = [c.name for c in engine._camera_configs] - assert cam_names.count("cam_a") == 1 - finally: - for p in patches: - p.stop() - - def test_unregister_removes(self): - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - try: - engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) - assert len(_engine_registry) == 1 - unregister_engine(engine) - assert len(_engine_registry) == 0 - finally: - for p in patches: - p.stop() + def test_creates_new(self, patched_mujoco_engine): + engine = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) + assert engine is not None + assert len(_engine_registry) == 1 + + def test_returns_same_instance(self, patched_mujoco_engine): + e1 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) + e2 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) + assert e1 is e2 + assert len(_engine_registry) == 1 + + def test_merges_new_cameras(self, patched_mujoco_engine): + e1 = get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_a")], + ) + get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_b")], + ) + names = {c.name for c in e1._camera_configs} + assert names == {"cam_a", "cam_b"} + + def test_deduplicates_cameras(self, patched_mujoco_engine): + get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_a")], + ) + get_or_create_engine( + config_path=Path("/fake/scene.xml"), + cameras=[CameraConfig(name="cam_a")], + ) + engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) + cam_names = [c.name for c in engine._camera_configs] + assert cam_names.count("cam_a") == 1 + + def test_unregister_removes(self, patched_mujoco_engine): + engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) + assert len(_engine_registry) == 1 + unregister_engine(engine) + assert len(_engine_registry) == 0 class TestCameraIntrinsics: From 7f57e54ca7b2a6ca30c06fc4fecaf89a0c3ea2ad Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 31 Mar 2026 14:47:40 -0700 Subject: [PATCH 31/41] adding from_intrinsics method --- dimos/msgs/sensor_msgs/CameraInfo.py | 32 +++++++++++++++++++++++ dimos/simulation/sensors/mujoco_camera.py | 18 +++---------- 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/dimos/msgs/sensor_msgs/CameraInfo.py b/dimos/msgs/sensor_msgs/CameraInfo.py index a371475675..a37682f7a5 100644 --- a/dimos/msgs/sensor_msgs/CameraInfo.py +++ b/dimos/msgs/sensor_msgs/CameraInfo.py @@ -90,6 +90,38 @@ def __init__( self.roi_width = 0 self.roi_do_rectify = False + @classmethod + def from_intrinsics( + cls, + fx: float, + fy: float, + cx: float, + cy: float, + width: int, + height: int, + frame_id: str = "", + ) -> CameraInfo: + """Create CameraInfo from pinhole intrinsics (no distortion). + + Args: + fx: Focal length x (pixels) + fy: Focal length y (pixels) + cx: Principal point x (pixels) + cy: Principal point y (pixels) + width: Image width + height: Image height + frame_id: Frame ID + """ + return cls( + height=height, + width=width, + distortion_model="plumb_bob", + D=[0.0, 0.0, 0.0, 0.0, 0.0], + K=[fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0], + P=[fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0], + frame_id=frame_id, + ) + def with_ts(self, ts: float) -> CameraInfo: """Return a copy of this CameraInfo with the given timestamp. diff --git a/dimos/simulation/sensors/mujoco_camera.py b/dimos/simulation/sensors/mujoco_camera.py index 237a6987d4..8ec39a643d 100644 --- a/dimos/simulation/sensors/mujoco_camera.py +++ b/dimos/simulation/sensors/mujoco_camera.py @@ -155,20 +155,10 @@ def _build_camera_info(self) -> None: fovy_rad = math.radians(fovy_deg) fy = h / (2.0 * math.tan(fovy_rad / 2.0)) fx = fy # square pixels - cx = w / 2.0 - cy = h / 2.0 - - K = [fx, 0.0, cx, 0.0, fy, cy, 0.0, 0.0, 1.0] - P = [fx, 0.0, cx, 0.0, 0.0, fy, cy, 0.0, 0.0, 0.0, 1.0, 0.0] - D = [0.0, 0.0, 0.0, 0.0, 0.0] - - self._camera_info_base = CameraInfo( - height=h, - width=w, - distortion_model="plumb_bob", - D=D, - K=K, - P=P, + + self._camera_info_base = CameraInfo.from_intrinsics( + fx=fx, fy=fy, cx=w / 2.0, cy=h / 2.0, + width=w, height=h, frame_id=self._color_optical_frame, ) From 1f3d8fa7f7fcf7e23cd674de52e54e961b05f22a Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 31 Mar 2026 14:48:25 -0700 Subject: [PATCH 32/41] fix: precommit fixes --- dimos/simulation/sensors/mujoco_camera.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/dimos/simulation/sensors/mujoco_camera.py b/dimos/simulation/sensors/mujoco_camera.py index 8ec39a643d..2e5cde393a 100644 --- a/dimos/simulation/sensors/mujoco_camera.py +++ b/dimos/simulation/sensors/mujoco_camera.py @@ -21,9 +21,9 @@ from __future__ import annotations import math +from pathlib import Path import threading import time -from pathlib import Path from typing import Any from pydantic import Field @@ -157,8 +157,12 @@ def _build_camera_info(self) -> None: fx = fy # square pixels self._camera_info_base = CameraInfo.from_intrinsics( - fx=fx, fy=fy, cx=w / 2.0, cy=h / 2.0, - width=w, height=h, + fx=fx, + fy=fy, + cx=w / 2.0, + cy=h / 2.0, + width=w, + height=h, frame_id=self._color_optical_frame, ) From 4e3acddebd10f2794e259b227938978097176076 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Sat, 4 Apr 2026 18:46:01 -0700 Subject: [PATCH 33/41] new arch: using shm to cummunicate bw modules --- dimos/hardware/manipulators/sim/adapter.py | 263 +++++++++++-- .../hardware/manipulators/sim/test_adapter.py | 205 ++++++++++ dimos/manipulation/blueprints.py | 14 +- dimos/robot/all_blueprints.py | 2 +- dimos/simulation/conftest.py | 11 +- dimos/simulation/engines/mujoco_engine.py | 79 ++-- dimos/simulation/engines/mujoco_shm.py | 362 ++++++++++++++++++ .../mujoco_sim_module.py} | 334 ++++++++++------ .../manipulators/sim_manip_interface.py | 221 ----------- .../manipulators/test_sim_adapter.py | 89 ----- .../simulation/sensors/test_mujoco_camera.py | 237 ------------ 11 files changed, 1059 insertions(+), 758 deletions(-) create mode 100644 dimos/hardware/manipulators/sim/test_adapter.py create mode 100644 dimos/simulation/engines/mujoco_shm.py rename dimos/simulation/{sensors/mujoco_camera.py => engines/mujoco_sim_module.py} (51%) delete mode 100644 dimos/simulation/manipulators/sim_manip_interface.py delete mode 100644 dimos/simulation/manipulators/test_sim_adapter.py delete mode 100644 dimos/simulation/sensors/test_mujoco_camera.py diff --git a/dimos/hardware/manipulators/sim/adapter.py b/dimos/hardware/manipulators/sim/adapter.py index 39cf6ffc5e..03c90ba1f5 100644 --- a/dimos/hardware/manipulators/sim/adapter.py +++ b/dimos/hardware/manipulators/sim/adapter.py @@ -12,61 +12,260 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""MuJoCo simulation adapter for ControlCoordinator integration. +"""Shared-memory adapter for MuJoCo-based manipulator simulation. -Thin wrapper around SimManipInterface that plugs into the adapter registry. -Arm joint methods are inherited from SimManipInterface. +This adapter is the hot-path counterpart to :class:`MujocoSimModule`. The sim +module owns the physics engine and publishes joint state to SHM; this adapter +(which lives inside ``ControlCoordinator``'s process) reads from and writes to +the same SHM buffers. + +Why SHM instead of RPC? The coordinator's tick loop calls the adapter +synchronously at 100Hz — ~4 calls per tick (positions, velocities, efforts, +command). LCM RPC latency (~1-5ms per call) would blow the 10ms budget. SHM +reads are sub-microsecond. + +Discovery is keyless: both sides derive the SHM buffer names from the MJCF +path, so no name exchange is required. """ from __future__ import annotations -from pathlib import Path +import logging +import math +import time from typing import TYPE_CHECKING, Any -from dimos.simulation.engines.mujoco_engine import MujocoEngine, get_or_create_engine -from dimos.simulation.manipulators.sim_manip_interface import SimManipInterface +from dimos.hardware.manipulators.spec import ( + ControlMode, + JointLimits, + ManipulatorInfo, +) +from dimos.simulation.engines.mujoco_shm import ( + ManipShmReader, + shm_key_from_path, +) if TYPE_CHECKING: from dimos.hardware.manipulators.registry import AdapterRegistry -class SimMujocoAdapter(SimManipInterface): - """Uses ``address`` as the MJCF XML path (same field real adapters use for IP/port). - If the engine has more joints than ``dof``, the extra joint at index ``dof`` - is treated as the gripper, with ctrl range scaled automatically. +_READY_WAIT_TIMEOUT_S = 60.0 +_READY_WAIT_POLL_S = 0.1 +_ATTACH_RETRY_TIMEOUT_S = 30.0 +_ATTACH_RETRY_POLL_S = 0.2 + + +class ShmMujocoAdapter: + """``ManipulatorAdapter`` that proxies to a ``MujocoSimModule`` via SHM. + + Uses ``address`` (the MJCF XML path) as the discovery key. The sim module + must be running and have signalled ready before ``connect()`` returns. """ def __init__( self, dof: int = 7, address: str | None = None, - headless: bool = True, - engine: MujocoEngine | None = None, + hardware_id: str | None = None, **_: Any, ) -> None: - if engine is None: - if address is None: - raise ValueError("address (MJCF XML path) is required for sim_mujoco adapter") - engine = get_or_create_engine(config_path=Path(address), headless=headless) - - # Detect gripper from engine joints - gripper_idx = None - gripper_kwargs = {} - joint_names = list(engine.joint_names) - if len(joint_names) > dof: - gripper_idx = dof - ctrl_range = engine.get_actuator_ctrl_range(dof) - joint_range = engine.get_joint_range(dof) - if ctrl_range is None or joint_range is None: - raise ValueError(f"Gripper joint at index {dof} missing ctrl/joint range in MJCF") - gripper_kwargs = {"gripper_ctrl_range": ctrl_range, "gripper_joint_range": joint_range} - - super().__init__(engine=engine, dof=dof, gripper_idx=gripper_idx, **gripper_kwargs) + if address is None: + raise ValueError("address (MJCF XML path) is required for sim_mujoco adapter") + self.logger = logging.getLogger(self.__class__.__name__) + self._dof = dof + self._address = address + self._hardware_id = hardware_id + self._shm_key = shm_key_from_path(address) + self._shm: ManipShmReader | None = None + self._connected = False + self._servos_enabled = False + self._control_mode = ControlMode.POSITION + self._error_code = 0 + self._error_message = "" + self._has_gripper = False + + # ---------------- connection --------------------------------- + + def connect(self) -> bool: + try: + # Attach to SHM. If the sim module hasn't created the buffers yet, + # retry for a short window — module startup ordering is not + # strictly guaranteed, only that they're both in the same blueprint. + deadline = time.monotonic() + _ATTACH_RETRY_TIMEOUT_S + while True: + try: + self._shm = ManipShmReader(self._shm_key) + break + except FileNotFoundError: + if time.monotonic() > deadline: + self.logger.error( + f"ShmMujocoAdapter: SHM buffers not found for {self._address} " + f"(key={self._shm_key}) after {_ATTACH_RETRY_TIMEOUT_S}s" + ) + return False + time.sleep(_ATTACH_RETRY_POLL_S) + + # Wait for sim module to signal ready. + deadline = time.monotonic() + _READY_WAIT_TIMEOUT_S + while not self._shm.is_ready(): + if time.monotonic() > deadline: + self.logger.error( + f"ShmMujocoAdapter: sim module not ready after {_READY_WAIT_TIMEOUT_S}s" + ) + return False + time.sleep(_READY_WAIT_POLL_S) + + num_joints = self._shm.num_joints() + self._has_gripper = num_joints > self._dof + self._connected = True + self._servos_enabled = True + self.logger.info( + f"ShmMujocoAdapter connected (dof={self._dof}, gripper={self._has_gripper})" + ) + return True + except Exception as exc: + self.logger.error(f"ShmMujocoAdapter: connect() failed: {exc}") + return False + + def disconnect(self) -> None: + try: + if self._shm is not None: + self._shm.cleanup() + finally: + self._shm = None + self._connected = False + + def is_connected(self) -> bool: + return self._connected and self._shm is not None + + # ---------------- info --------------------------------------- + + def get_info(self) -> ManipulatorInfo: + return ManipulatorInfo( + vendor="Simulation", + model="Simulation", + dof=self._dof, + firmware_version=None, + serial_number=None, + ) + + def get_dof(self) -> int: + return self._dof + + def get_limits(self) -> JointLimits: + lower = [-math.pi] * self._dof + upper = [math.pi] * self._dof + max_vel_rad = math.radians(180.0) + return JointLimits( + position_lower=lower, + position_upper=upper, + velocity_max=[max_vel_rad] * self._dof, + ) + + # ---------------- control mode ------------------------------- + + def set_control_mode(self, mode: ControlMode) -> bool: + self._control_mode = mode + return True + + def get_control_mode(self) -> ControlMode: + return self._control_mode + + # ---------------- read state (SHM -> caller) ----------------- + + def read_joint_positions(self) -> list[float]: + if self._shm is None: + return [0.0] * self._dof + return self._shm.read_positions(self._dof) + + def read_joint_velocities(self) -> list[float]: + if self._shm is None: + return [0.0] * self._dof + return self._shm.read_velocities(self._dof) + + def read_joint_efforts(self) -> list[float]: + if self._shm is None: + return [0.0] * self._dof + return self._shm.read_efforts(self._dof) + + def read_state(self) -> dict[str, int]: + velocities = self.read_joint_velocities() + is_moving = any(abs(v) > 1e-4 for v in velocities) + mode_int = list(ControlMode).index(self._control_mode) + return {"state": 1 if is_moving else 0, "mode": mode_int} + + def read_error(self) -> tuple[int, str]: + return self._error_code, self._error_message + + # ---------------- write commands (caller -> SHM) ------------- + + def write_joint_positions(self, positions: list[float], velocity: float = 1.0) -> bool: + if not self._servos_enabled or self._shm is None: + return False + self._control_mode = ControlMode.POSITION + self._shm.write_position_command(positions[: self._dof]) + return True + + def write_joint_velocities(self, velocities: list[float]) -> bool: + if not self._servos_enabled or self._shm is None: + return False + self._control_mode = ControlMode.VELOCITY + self._shm.write_velocity_command(velocities[: self._dof]) + return True + + def write_joint_efforts(self, efforts: list[float]) -> bool: + # Effort mode not exposed via SHM yet; caller can fall back to position. + return False + + def write_stop(self) -> bool: + # Hold current position. + if self._shm is None: + return False + positions = self._shm.read_positions(self._dof) + self._shm.write_position_command(positions) + return True + + def write_enable(self, enable: bool) -> bool: + self._servos_enabled = enable + return True + + def read_enabled(self) -> bool: + return self._servos_enabled + + def write_clear_errors(self) -> bool: + self._error_code = 0 + self._error_message = "" + return True + + # ---------------- cartesian (not supported) ------------------ + + def read_cartesian_position(self) -> dict[str, float] | None: + return None + + def write_cartesian_position(self, pose: dict[str, float], velocity: float = 1.0) -> bool: + return False + + # ---------------- gripper ------------------------------------ + + def read_gripper_position(self) -> float | None: + if not self._has_gripper or self._shm is None: + return None + return self._shm.read_gripper_position() + + def write_gripper_position(self, position: float) -> bool: + if not self._has_gripper or self._shm is None: + return False + # Raw joint-space position; sim module handles the joint->ctrl scaling. + self._shm.write_gripper_command(position) + return True + + def read_force_torque(self) -> list[float] | None: + return None def register(registry: AdapterRegistry) -> None: """Register this adapter with the registry.""" - registry.register("sim_mujoco", SimMujocoAdapter) + registry.register("sim_mujoco", ShmMujocoAdapter) -__all__ = ["SimMujocoAdapter"] +__all__ = ["ShmMujocoAdapter"] diff --git a/dimos/hardware/manipulators/sim/test_adapter.py b/dimos/hardware/manipulators/sim/test_adapter.py new file mode 100644 index 0000000000..d8ea1fe485 --- /dev/null +++ b/dimos/hardware/manipulators/sim/test_adapter.py @@ -0,0 +1,205 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ShmMujocoAdapter (the SHM-backed ManipulatorAdapter).""" + +from __future__ import annotations + +from unittest.mock import MagicMock +import uuid + +import pytest + +from dimos.hardware.manipulators.sim.adapter import ShmMujocoAdapter, register +from dimos.hardware.manipulators.spec import ControlMode, ManipulatorAdapter +from dimos.simulation.engines.mujoco_shm import ManipShmWriter + +ARM_DOF = 7 + + +@pytest.fixture +def shm_key(): + return f"test_{uuid.uuid4().hex[:10]}" + + +@pytest.fixture +def writer(shm_key, monkeypatch): + """Pretend we're the sim module: create SHM, signal ready. + + We monkey-patch ``shm_key_from_path`` so the adapter under test resolves + to our fixture's key regardless of the address string. + """ + import dimos.hardware.manipulators.sim.adapter as adapter_mod + + monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) + w = ManipShmWriter(shm_key) + w.signal_ready(num_joints=ARM_DOF) + yield w + w.cleanup() + + +@pytest.fixture +def writer_with_gripper(shm_key, monkeypatch): + import dimos.hardware.manipulators.sim.adapter as adapter_mod + + monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) + w = ManipShmWriter(shm_key) + w.signal_ready(num_joints=ARM_DOF + 1) + yield w + w.cleanup() + + +@pytest.fixture +def adapter(writer): + a = ShmMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml") + assert a.connect() is True + yield a + a.disconnect() + + +@pytest.fixture +def adapter_with_gripper(writer_with_gripper): + a = ShmMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml") + assert a.connect() is True + yield a + a.disconnect() + + +class TestProtocolConformance: + def test_implements_manipulator_adapter(self): + a = ShmMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml") + assert isinstance(a, ManipulatorAdapter) + + def test_address_required(self): + with pytest.raises(ValueError, match="address"): + ShmMujocoAdapter(dof=ARM_DOF, address=None) + + def test_register(self): + registry = MagicMock() + register(registry) + registry.register.assert_called_once_with("sim_mujoco", ShmMujocoAdapter) + + +class TestReadState: + def test_read_joint_positions(self, adapter, writer): + writer.write_joint_state( + positions=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], + velocities=[0.0] * 7, + efforts=[0.0] * 7, + ) + assert adapter.read_joint_positions() == [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7] + + def test_read_joint_velocities(self, adapter, writer): + writer.write_joint_state( + positions=[0.0] * 7, + velocities=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], + efforts=[0.0] * 7, + ) + assert adapter.read_joint_velocities() == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] + + def test_read_joint_efforts(self, adapter, writer): + writer.write_joint_state( + positions=[0.0] * 7, + velocities=[0.0] * 7, + efforts=[-1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0], + ) + assert adapter.read_joint_efforts() == [-1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0] + + def test_returns_only_dof_joints(self, adapter_with_gripper, writer_with_gripper): + # Writer publishes 8 joints (7 arm + 1 gripper); adapter should return 7. + writer_with_gripper.write_joint_state( + positions=list(range(8)), + velocities=[0.0] * 8, + efforts=[0.0] * 8, + ) + positions = adapter_with_gripper.read_joint_positions() + assert len(positions) == ARM_DOF + + +class TestWriteCommand: + def test_write_joint_positions(self, adapter, writer): + assert adapter.write_joint_positions([0.1] * 7) is True + cmd = writer.read_position_command(7) + assert cmd is not None + assert cmd.tolist() == pytest.approx([0.1] * 7) + + def test_write_joint_velocities(self, adapter, writer): + assert adapter.write_joint_velocities([0.5] * 7) is True + cmd = writer.read_velocity_command(7) + assert cmd is not None + assert cmd.tolist() == pytest.approx([0.5] * 7) + + def test_write_when_disabled(self, adapter): + adapter.write_enable(False) + assert adapter.write_joint_positions([0.0] * 7) is False + + def test_control_mode_tracked(self, adapter): + adapter.write_joint_positions([0.0] * 7) + assert adapter.get_control_mode() == ControlMode.POSITION + adapter.write_joint_velocities([0.0] * 7) + assert adapter.get_control_mode() == ControlMode.VELOCITY + + +class TestGripper: + def test_gripper_detected(self, adapter_with_gripper): + assert adapter_with_gripper._has_gripper is True + + def test_no_gripper_when_dof_matches(self, adapter): + assert adapter._has_gripper is False + + def test_read_gripper_position(self, adapter_with_gripper, writer_with_gripper): + writer_with_gripper.write_gripper_state(0.33) + assert adapter_with_gripper.read_gripper_position() == pytest.approx(0.33) + + def test_read_gripper_position_no_gripper(self, adapter): + assert adapter.read_gripper_position() is None + + def test_write_gripper_position(self, adapter_with_gripper, writer_with_gripper): + assert adapter_with_gripper.write_gripper_position(0.5) is True + # Gripper command is raw (unscaled) — sim module handles joint->ctrl. + assert writer_with_gripper.read_gripper_command() == pytest.approx(0.5) + + def test_write_gripper_position_no_gripper(self, adapter): + assert adapter.write_gripper_position(0.5) is False + + +class TestConnect: + def test_connect_before_sim_ready_times_out(self, shm_key, monkeypatch): + """If sim module never signals ready, connect() returns False after timeout.""" + import dimos.hardware.manipulators.sim.adapter as adapter_mod + + monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) + # Shrink timeouts so the test runs fast. + monkeypatch.setattr(adapter_mod, "_READY_WAIT_TIMEOUT_S", 0.2) + monkeypatch.setattr(adapter_mod, "_READY_WAIT_POLL_S", 0.02) + + w = ManipShmWriter(shm_key) + try: + # Note: writer exists but signal_ready is NOT called. + a = ShmMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml") + assert a.connect() is False + finally: + w.cleanup() + + def test_connect_waits_for_shm(self, shm_key, monkeypatch): + """If SHM buffers don't exist yet, connect() retries briefly.""" + import dimos.hardware.manipulators.sim.adapter as adapter_mod + + monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) + monkeypatch.setattr(adapter_mod, "_ATTACH_RETRY_TIMEOUT_S", 0.2) + monkeypatch.setattr(adapter_mod, "_ATTACH_RETRY_POLL_S", 0.02) + + a = ShmMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml") + # SHM was never created — attach must time out. + assert a.connect() is False diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index c327ac5c25..37518d1231 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -494,14 +494,14 @@ def _make_piper_config( ) -# Sim perception: MuJoCo camera replaces RealSense for sim-based pick-and-place -# Both the sim adapter and camera resolve the same MujocoEngine via the registry -# (keyed by MJCF path), so they share physics state. -# The engine is created lazily when the adapter connects / camera starts. +# Sim perception: MujocoSimModule owns the MujocoEngine and publishes both +# camera streams (replacing the old MujocoCamera) and joint state via shared +# memory. The sim_xarm7 hardware component resolves to ShmMujocoAdapter, which +# attaches to the same SHM by MJCF path — no globals, no cross-process sharing. from dimos.control.blueprints._hardware import XARM7_SIM_PATH, sim_xarm7 from dimos.control.coordinator import TaskConfig as TaskConfig -from dimos.simulation.sensors.mujoco_camera import MujocoCamera +from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode xarm_perception_sim = autoconnect( @@ -521,8 +521,10 @@ def _make_piper_config( planning_timeout=10.0, enable_viz=True, ), - MujocoCamera.blueprint( + MujocoSimModule.blueprint( address=str(XARM7_SIM_PATH), + headless=False, + dof=7, camera_name="wrist_camera", base_frame_id="link7", ), diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index bd31e9a518..34566206df 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -139,7 +139,7 @@ "mock-b1-connection-module": "dimos.robot.unitree.b1.connection.MockB1ConnectionModule", "module-a": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleA", "module-b": "dimos.robot.unitree.demo_error_on_name_conflicts.ModuleB", - "mujoco-camera": "dimos.simulation.sensors.mujoco_camera.MujocoCamera", + "mujoco-sim-module": "dimos.simulation.engines.mujoco_sim_module.MujocoSimModule", "navigation-module": "dimos.robot.unitree.rosnav.NavigationModule", "navigation-skill-container": "dimos.agents.skills.navigation.NavigationSkillContainer", "object-db-module": "dimos.perception.detection.moduleDB.ObjectDBModule", diff --git a/dimos/simulation/conftest.py b/dimos/simulation/conftest.py index 6a12c3d1cc..270c010c9a 100644 --- a/dimos/simulation/conftest.py +++ b/dimos/simulation/conftest.py @@ -22,7 +22,6 @@ import numpy as np import pytest -from dimos.simulation.engines.mujoco_engine import _clear_registry from dimos.simulation.utils.xml_parser import JointMapping ARM_DOF = 7 @@ -96,21 +95,13 @@ def _patch_mujoco_engine(n_joints: int = ARM_DOF): return patches -@pytest.fixture(autouse=True) -def clean_engine_registry(): - """Clear the engine registry before and after each simulation test.""" - _clear_registry() - yield - _clear_registry() - - @pytest.fixture def patched_mujoco_engine(): """Start MuJoCo engine patches for the test, stop on teardown. Usage: def test_something(patched_mujoco_engine): - engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) + engine = MujocoEngine(config_path=Path("/fake/scene.xml"), headless=True) """ patches = _patch_mujoco_engine(ARM_DOF) for p in patches: diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 1cb61ebf45..cbc29d69e6 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -16,6 +16,7 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path import threading @@ -36,6 +37,9 @@ logger = setup_logger() +# Step hook signature: called with the engine instance inside the sim thread. +StepHook = Callable[["MujocoEngine"], None] + @dataclass class CameraConfig: @@ -71,57 +75,6 @@ class _CameraRendererState: last_render_time: float = 0.0 -_engine_registry: dict[str, MujocoEngine] = {} -_engine_registry_lock = threading.Lock() - - -def get_or_create_engine( - config_path: Path, - headless: bool = True, - cameras: list[CameraConfig] | None = None, -) -> MujocoEngine: - """Return the shared MujocoEngine for *config_path*, creating one if needed. - - If an engine already exists for the resolved path and *cameras* are supplied, - any **new** camera configs are appended (idempotent by name). - """ - key = str(config_path.resolve()) - with _engine_registry_lock: - if key in _engine_registry: - engine = _engine_registry[key] - if engine._headless != headless: - logger.warning( - f"get_or_create_engine: ignoring headless={headless} — " - f"existing engine for '{key}' was created with headless={engine._headless}" - ) - # Merge new camera configs (by name), guarded against sim-thread iteration - if cameras: - with engine._camera_lock: - existing_names = {c.name for c in engine._camera_configs} - for cam in cameras: - if cam.name not in existing_names: - engine._camera_configs.append(cam) - return engine - - engine = MujocoEngine(config_path=config_path, headless=headless, cameras=cameras) - _engine_registry[key] = engine - return engine - - -def unregister_engine(engine: MujocoEngine) -> None: - """Remove an engine from the registry (called on disconnect).""" - with _engine_registry_lock: - keys_to_remove = [k for k, v in _engine_registry.items() if v is engine] - for k in keys_to_remove: - del _engine_registry[k] - - -def _clear_registry() -> None: - """Clear the engine registry (for test teardown only).""" - with _engine_registry_lock: - _engine_registry.clear() - - class MujocoEngine(SimulationEngine): """ MuJoCo simulation engine. @@ -136,8 +89,18 @@ def __init__( config_path: Path, headless: bool, cameras: list[CameraConfig] | None = None, + on_before_step: StepHook | None = None, + on_after_step: StepHook | None = None, ) -> None: super().__init__(config_path=config_path, headless=headless) + # Optional observer callbacks invoked inside the sim thread. + # ``on_before_step`` runs before ``_apply_control`` (use it to pull + # commands into the engine); ``on_after_step`` runs after + # ``_update_joint_state`` (use it to publish joint state to an + # external sink, e.g. shared memory). The engine treats these as + # opaque callables and has no knowledge of what they do. + self._on_before_step: StepHook | None = on_before_step + self._on_after_step: StepHook | None = on_after_step xml_path = self._resolve_xml_path(config_path) self._model = mujoco.MjModel.from_xml_path(str(xml_path)) @@ -265,7 +228,6 @@ def disconnect(self) -> bool: if self._sim_thread and self._sim_thread.is_alive(): self._sim_thread.join(timeout=2.0) self._sim_thread = None - unregister_engine(self) return True except Exception as e: logger.error(f"{self.__class__.__name__}: disconnect() failed: {e}") @@ -341,11 +303,21 @@ def _sim_loop(self) -> None: def _step_once(sync_viewer: bool) -> None: loop_start = time.time() + if self._on_before_step is not None: + try: + self._on_before_step(self) + except Exception as exc: + logger.error(f"on_before_step failed: {exc}") self._apply_control() mujoco.mj_step(self._model, self._data) if sync_viewer: m_viewer.sync() self._update_joint_state() + if self._on_after_step is not None: + try: + self._on_after_step(self) + except Exception as exc: + logger.error(f"on_after_step failed: {exc}") self._render_cameras(loop_start, cam_renderers) elapsed = time.time() - loop_start @@ -511,6 +483,5 @@ def get_camera_fovy(self, camera_name: str) -> float | None: "CameraConfig", "CameraFrame", "MujocoEngine", - "get_or_create_engine", - "unregister_engine", + "StepHook", ] diff --git a/dimos/simulation/engines/mujoco_shm.py b/dimos/simulation/engines/mujoco_shm.py new file mode 100644 index 0000000000..924fa3fe56 --- /dev/null +++ b/dimos/simulation/engines/mujoco_shm.py @@ -0,0 +1,362 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared-memory buffers for sim-manipulator IPC. + +Layout for exchanging joint state and commands between ``MujocoSimModule`` +(which owns the physics engine) and ``ShmMujocoAdapter`` (which plugs into +ControlCoordinator). Modeled after ``dimos.simulation.mujoco.shared_memory`` +(the Go2 SHM pattern). + +Names are deterministic: both sides derive them from the resolved MJCF path, +so no name exchange over RPC is needed. The sim module creates the buffers +and signals ``ready``; the adapter attaches to them by name. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +from multiprocessing import resource_tracker +from multiprocessing.shared_memory import SharedMemory +from pathlib import Path +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# Upper bound on joint count per sim. Arms + gripper are typically <= 10. +MAX_JOINTS = 16 + +_FLOAT_BYTES = 8 # float64 +_INT32_BYTES = 4 + +_joint_array_size = MAX_JOINTS * _FLOAT_BYTES # float64 array + +# Buffer sizes (in bytes). +# seq: one int64 counter per buffer type. +# control: int32 [ready, stop, command_mode, num_joints]. +_shm_sizes = { + "positions": _joint_array_size, + "velocities": _joint_array_size, + "efforts": _joint_array_size, + "position_targets": _joint_array_size, + "velocity_targets": _joint_array_size, + "gripper": 2 * _FLOAT_BYTES, # [gripper_position, gripper_target] + "seq": 8 * _FLOAT_BYTES, # 8 int64 counters + "control": 4 * _INT32_BYTES, # [ready, stop, command_mode, num_joints] +} + +# Sequence counter indices. +SEQ_POSITIONS = 0 +SEQ_VELOCITIES = 1 +SEQ_EFFORTS = 2 +SEQ_POSITION_CMD = 3 +SEQ_VELOCITY_CMD = 4 +SEQ_GRIPPER_STATE = 5 +SEQ_GRIPPER_CMD = 6 + +# Control indices. +CTRL_READY = 0 +CTRL_STOP = 1 +CTRL_COMMAND_MODE = 2 +CTRL_NUM_JOINTS = 3 + +# Command modes. +CMD_MODE_POSITION = 0 +CMD_MODE_VELOCITY = 1 + +_NAME_PREFIX = "dimos_mjmanip" + + +def shm_key_from_path(config_path: Path | str) -> str: + """Derive a deterministic short key from an MJCF path. + + Both sim module and adapter compute the same key from the same path, + so SHM buffer names can be agreed upon without an RPC round-trip. + """ + resolved = str(Path(config_path).expanduser().resolve()) + return hashlib.md5(resolved.encode("utf-8")).hexdigest()[:12] + + +def _buffer_name(key: str, buffer: str) -> str: + return f"{_NAME_PREFIX}_{key}_{buffer}" + + +def _unregister(shm: SharedMemory) -> SharedMemory: + """Detach ``shm`` from ``resource_tracker`` to silence spurious warnings. + + Same technique as ``dimos.simulation.mujoco.shared_memory._unregister``. + """ + try: + resource_tracker.unregister(shm._name, "shared_memory") # type: ignore[attr-defined] + except Exception: + pass + return shm + + +@dataclass(frozen=True) +class ManipShmSet: + """Frozen set of named SharedMemory buffers for manipulator IPC.""" + + positions: SharedMemory + velocities: SharedMemory + efforts: SharedMemory + position_targets: SharedMemory + velocity_targets: SharedMemory + gripper: SharedMemory + seq: SharedMemory + control: SharedMemory + + @classmethod + def create(cls, key: str) -> ManipShmSet: + """Create new SHM buffers with deterministic names derived from *key*. + + If stale buffers from a prior run exist, unlink them first. + """ + buffers: dict[str, SharedMemory] = {} + for buffer_name, size in _shm_sizes.items(): + name = _buffer_name(key, buffer_name) + try: + stale = SharedMemory(name=name) + stale.close() + try: + stale.unlink() + logger.info("ManipShmSet: unlinked stale SHM", name=name) + except FileNotFoundError: + pass + except FileNotFoundError: + pass + buffers[buffer_name] = SharedMemory(create=True, size=size, name=name) + return cls(**buffers) + + @classmethod + def attach(cls, key: str) -> ManipShmSet: + """Attach to existing SHM buffers created by the sim side.""" + buffers: dict[str, SharedMemory] = {} + for buffer_name in _shm_sizes: + name = _buffer_name(key, buffer_name) + buffers[buffer_name] = _unregister(SharedMemory(name=name)) + return cls(**buffers) + + def as_list(self) -> list[SharedMemory]: + return [getattr(self, k) for k in _shm_sizes] + + +class ManipShmWriter: + """Sim-side handle: writes joint state, reads command targets. + + Owned by ``MujocoSimModule``. Creates the SHM buffers on init and + unlinks them on cleanup. + """ + + shm: ManipShmSet + + def __init__(self, key: str) -> None: + self.shm = ManipShmSet.create(key) + self._last_pos_cmd_seq = 0 + self._last_vel_cmd_seq = 0 + self._last_gripper_cmd_seq = 0 + # Zero everything. + for buf in self.shm.as_list(): + np.ndarray((buf.size,), dtype=np.uint8, buffer=buf.buf)[:] = 0 + + # ---------------- joint state (sim -> adapter) ---------------- + + def write_joint_state( + self, + positions: list[float], + velocities: list[float], + efforts: list[float], + ) -> None: + n = min(len(positions), MAX_JOINTS) + pos_arr = self._array(self.shm.positions, MAX_JOINTS, np.float64) + vel_arr = self._array(self.shm.velocities, MAX_JOINTS, np.float64) + eff_arr = self._array(self.shm.efforts, MAX_JOINTS, np.float64) + pos_arr[:n] = positions[:n] + vel_arr[:n] = velocities[:n] + eff_arr[:n] = efforts[:n] + self._increment_seq(SEQ_POSITIONS) + self._increment_seq(SEQ_VELOCITIES) + self._increment_seq(SEQ_EFFORTS) + + def write_gripper_state(self, position: float) -> None: + arr = self._array(self.shm.gripper, 2, np.float64) + arr[0] = position + self._increment_seq(SEQ_GRIPPER_STATE) + + # ---------------- commands (adapter -> sim) ------------------- + + def read_position_command(self, num_joints: int) -> NDArray[np.float64] | None: + """Return a copy of position targets if a new command arrived since last call.""" + seq = self._get_seq(SEQ_POSITION_CMD) + if seq <= self._last_pos_cmd_seq: + return None + self._last_pos_cmd_seq = seq + arr = self._array(self.shm.position_targets, MAX_JOINTS, np.float64) + return arr[:num_joints].copy() + + def read_velocity_command(self, num_joints: int) -> NDArray[np.float64] | None: + seq = self._get_seq(SEQ_VELOCITY_CMD) + if seq <= self._last_vel_cmd_seq: + return None + self._last_vel_cmd_seq = seq + arr = self._array(self.shm.velocity_targets, MAX_JOINTS, np.float64) + return arr[:num_joints].copy() + + def read_gripper_command(self) -> float | None: + seq = self._get_seq(SEQ_GRIPPER_CMD) + if seq <= self._last_gripper_cmd_seq: + return None + self._last_gripper_cmd_seq = seq + arr = self._array(self.shm.gripper, 2, np.float64) + return float(arr[1]) + + def read_command_mode(self) -> int: + return int(self._control()[CTRL_COMMAND_MODE]) + + # ---------------- control / lifecycle ------------------------- + + def signal_ready(self, num_joints: int) -> None: + ctrl = self._control() + ctrl[CTRL_NUM_JOINTS] = num_joints + ctrl[CTRL_READY] = 1 + + def signal_stop(self) -> None: + self._control()[CTRL_STOP] = 1 + + def should_stop(self) -> bool: + return bool(self._control()[CTRL_STOP] == 1) + + def cleanup(self) -> None: + for shm in self.shm.as_list(): + try: + shm.close() + except Exception: + pass + try: + shm.unlink() + except Exception: + pass + + # ---------------- helpers ------------------------------------ + + def _array(self, buf: SharedMemory, n: int, dtype: Any) -> NDArray[Any]: + return np.ndarray((n,), dtype=dtype, buffer=buf.buf) + + def _control(self) -> NDArray[np.int32]: + return np.ndarray((4,), dtype=np.int32, buffer=self.shm.control.buf) + + def _increment_seq(self, index: int) -> None: + seq_arr = np.ndarray((8,), dtype=np.int64, buffer=self.shm.seq.buf) + seq_arr[index] += 1 + + def _get_seq(self, index: int) -> int: + seq_arr = np.ndarray((8,), dtype=np.int64, buffer=self.shm.seq.buf) + return int(seq_arr[index]) + + +class ManipShmReader: + """Adapter-side handle: reads joint state, writes command targets. + + Owned by ``ShmMujocoAdapter``. Attaches to existing buffers created by + the sim module; does not unlink them on cleanup. + """ + + shm: ManipShmSet + + def __init__(self, key: str) -> None: + self.shm = ManipShmSet.attach(key) + + # ---------------- joint state (sim -> adapter) ---------------- + + def read_positions(self, num_joints: int) -> list[float]: + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.positions.buf) + return [float(x) for x in arr[:num_joints]] + + def read_velocities(self, num_joints: int) -> list[float]: + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.velocities.buf) + return [float(x) for x in arr[:num_joints]] + + def read_efforts(self, num_joints: int) -> list[float]: + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.efforts.buf) + return [float(x) for x in arr[:num_joints]] + + def read_gripper_position(self) -> float: + arr = np.ndarray((2,), dtype=np.float64, buffer=self.shm.gripper.buf) + return float(arr[0]) + + # ---------------- commands (adapter -> sim) ------------------- + + def write_position_command(self, positions: list[float]) -> None: + n = min(len(positions), MAX_JOINTS) + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.position_targets.buf) + arr[:n] = positions[:n] + self._set_command_mode(CMD_MODE_POSITION) + self._increment_seq(SEQ_POSITION_CMD) + + def write_velocity_command(self, velocities: list[float]) -> None: + n = min(len(velocities), MAX_JOINTS) + arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.velocity_targets.buf) + arr[:n] = velocities[:n] + self._set_command_mode(CMD_MODE_VELOCITY) + self._increment_seq(SEQ_VELOCITY_CMD) + + def write_gripper_command(self, position: float) -> None: + arr = np.ndarray((2,), dtype=np.float64, buffer=self.shm.gripper.buf) + arr[1] = position + self._increment_seq(SEQ_GRIPPER_CMD) + + # ---------------- control / lifecycle ------------------------- + + def is_ready(self) -> bool: + return bool(self._control()[CTRL_READY] == 1) + + def num_joints(self) -> int: + return int(self._control()[CTRL_NUM_JOINTS]) + + def signal_stop(self) -> None: + self._control()[CTRL_STOP] = 1 + + def cleanup(self) -> None: + for shm in self.shm.as_list(): + try: + shm.close() + except Exception: + pass + + # ---------------- helpers ------------------------------------ + + def _control(self) -> NDArray[np.int32]: + return np.ndarray((4,), dtype=np.int32, buffer=self.shm.control.buf) + + def _set_command_mode(self, mode: int) -> None: + self._control()[CTRL_COMMAND_MODE] = mode + + def _increment_seq(self, index: int) -> None: + seq_arr = np.ndarray((8,), dtype=np.int64, buffer=self.shm.seq.buf) + seq_arr[index] += 1 + + +__all__ = [ + "MAX_JOINTS", + "ManipShmReader", + "ManipShmSet", + "ManipShmWriter", + "shm_key_from_path", +] diff --git a/dimos/simulation/sensors/mujoco_camera.py b/dimos/simulation/engines/mujoco_sim_module.py similarity index 51% rename from dimos/simulation/sensors/mujoco_camera.py rename to dimos/simulation/engines/mujoco_sim_module.py index 2e5cde393a..ce872c2fde 100644 --- a/dimos/simulation/sensors/mujoco_camera.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -12,10 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""MuJoCo simulation camera module. +"""Unified MuJoCo simulation Module. -Drop-in replacement for RealSenseCamera in manipulation blueprints. -Renders RGB + depth from a named camera in a shared MujocoEngine. +Owns a single ``MujocoEngine`` and publishes: +- camera streams (Out ports), replacing ``MujocoCamera`` +- joint state via shared memory, consumed by ``ShmMujocoAdapter`` inside + ``ControlCoordinator`` + +This avoids the prior pattern of sharing engines via a global in-process +registry, which was fragile when ``WorkerManager`` places the adapter and +the camera in different worker processes. """ from __future__ import annotations @@ -33,21 +39,24 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.hardware.sensors.camera.spec import ( - DepthCameraConfig, - DepthCameraHardware, -) +from dimos.hardware.sensors.camera.spec import DepthCameraConfig, DepthCameraHardware from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.simulation.engines.mujoco_engine import ( CameraConfig, CameraFrame, MujocoEngine, - get_or_create_engine, +) +from dimos.simulation.engines.mujoco_shm import ( + CMD_MODE_POSITION, + CMD_MODE_VELOCITY, + ManipShmWriter, + shm_key_from_path, ) from dimos.spec import perception from dimos.utils.logging_config import setup_logger @@ -65,18 +74,25 @@ def _default_identity_transform() -> Transform: ) -class MujocoCameraConfig(ModuleConfig, DepthCameraConfig): - """Configuration for the MuJoCo simulation camera.""" +class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): + """Configuration for the unified MuJoCo simulation module.""" + # MJCF path — also used as the SHM key so the adapter can discover the + # shared memory region without RPC. address: str = "" headless: bool = False + # Number of arm joints (excluding gripper). If the engine exposes more + # joints than ``dof``, the extra joint at index ``dof`` is treated as a + # gripper and its commands get scaled from joint-space into actuator ctrl. + dof: int = 7 + + # Camera config (matches former MujocoCameraConfig). camera_name: str = "wrist_camera" width: int = 640 height: int = 480 fps: int = 15 base_frame_id: str = "link7" base_transform: Transform | None = Field(default_factory=_default_identity_transform) - # MuJoCo renders color+depth from same virtual camera, alignment is a no-op align_depth_to_color: bool = True enable_depth: bool = True enable_pointcloud: bool = False @@ -84,14 +100,17 @@ class MujocoCameraConfig(ModuleConfig, DepthCameraConfig): camera_info_fps: float = 1.0 -class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.DepthCamera): - """Simulated depth camera that renders from a MujocoEngine. - - Implements the same output ports and TF chain as RealSenseCamera so it can - be used as a drop-in replacement in manipulation blueprints. - - The engine is resolved automatically from the registry via ``address`` - (the same MJCF path used by the sim_mujoco adapter). +class MujocoSimModule( + DepthCameraHardware, + Module[MujocoSimModuleConfig], + perception.DepthCamera, +): + """Single Module that owns a MujocoEngine, publishes camera streams, and + exposes joint state/commands to a ``ShmMujocoAdapter`` via shared memory. + + The adapter attaches to the same SHM buffers using the MJCF path as the + discovery key — no RPC, no globals. From ControlCoordinator's perspective + the adapter is an ordinary ``ManipulatorAdapter``; SHM is its transport. """ color_image: Out[Image] @@ -100,15 +119,21 @@ class MujocoCamera(DepthCameraHardware, Module[MujocoCameraConfig], perception.D camera_info: Out[CameraInfo] depth_camera_info: Out[CameraInfo] - default_config = MujocoCameraConfig + default_config = MujocoSimModuleConfig def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._engine: MujocoEngine | None = None + self._shm: ManipShmWriter | None = None + self._gripper_idx: int | None = None + self._gripper_ctrl_range: tuple[float, float] = (0.0, 1.0) + self._gripper_joint_range: tuple[float, float] = (0.0, 1.0) self._stop_event = threading.Event() - self._thread: threading.Thread | None = None + self._publish_thread: threading.Thread | None = None self._camera_info_base: CameraInfo | None = None + # ---------------- frame id helpers --------------------------- + @property def _camera_link(self) -> str: return f"{self.config.camera_name}_link" @@ -129,6 +154,8 @@ def _depth_frame(self) -> str: def _depth_optical_frame(self) -> str: return f"{self.config.camera_name}_depth_optical_frame" + # ---------------- RPC spec methods --------------------------- + @rpc def get_color_camera_info(self) -> CameraInfo | None: return self._camera_info_base @@ -141,108 +168,207 @@ def get_depth_camera_info(self) -> CameraInfo | None: def get_depth_scale(self) -> float: return 1.0 - def _build_camera_info(self) -> None: - """Compute camera intrinsics from the MuJoCo model (pinhole, no distortion).""" - if self._engine is None: - return - fovy_deg = self._engine.get_camera_fovy(self.config.camera_name) - if fovy_deg is None: - logger.error("Camera not found in MJCF", camera_name=self.config.camera_name) - return - - h = self.config.height - w = self.config.width - fovy_rad = math.radians(fovy_deg) - fy = h / (2.0 * math.tan(fovy_rad / 2.0)) - fx = fy # square pixels - - self._camera_info_base = CameraInfo.from_intrinsics( - fx=fx, - fy=fy, - cx=w / 2.0, - cy=h / 2.0, - width=w, - height=h, - frame_id=self._color_optical_frame, - ) + # ---------------- lifecycle ---------------------------------- @rpc def start(self) -> None: - if self._engine is None and self.config.address: - self._engine = get_or_create_engine( - config_path=Path(self.config.address), - headless=self.config.headless, - cameras=[ - CameraConfig( - name=self.config.camera_name, - width=self.config.width, - height=self.config.height, - fps=self.config.fps, - ) - ], - ) - if self._engine is None: - raise RuntimeError( - "MujocoCamera: set address (MJCF path) in config or call set_engine()" + if not self.config.address: + raise RuntimeError("MujocoSimModule: config.address (MJCF path) is required") + + # SHM key — adapter derives the same key from the same MJCF path. + shm_key = shm_key_from_path(self.config.address) + self._shm = ManipShmWriter(shm_key) + + # Build engine with SHM hooks installed. + self._engine = MujocoEngine( + config_path=Path(self.config.address), + headless=self.config.headless, + cameras=[ + CameraConfig( + name=self.config.camera_name, + width=self.config.width, + height=self.config.height, + fps=float(self.config.fps), + ) + ], + on_before_step=self._apply_shm_commands, + on_after_step=self._publish_shm_state, + ) + + # Detect gripper (extra joint beyond dof). + dof = self.config.dof + joint_names = list(self._engine.joint_names) + if len(joint_names) > dof: + ctrl_range = self._engine.get_actuator_ctrl_range(dof) + joint_range = self._engine.get_joint_range(dof) + if ctrl_range is None or joint_range is None: + raise ValueError(f"Gripper joint at index {dof} missing ctrl/joint range in MJCF") + self._gripper_idx = dof + self._gripper_ctrl_range = ctrl_range + self._gripper_joint_range = joint_range + logger.info( + "MujocoSimModule: gripper detected", + idx=dof, + ctrl_range=ctrl_range, + joint_range=joint_range, ) - # Ensure camera config is registered (handles the set_engine() path) - cam_cfg = CameraConfig( - name=self.config.camera_name, - width=self.config.width, - height=self.config.height, - fps=self.config.fps, - ) - with self._engine._camera_lock: - existing_names = {c.name for c in self._engine._camera_configs} - if cam_cfg.name not in existing_names: - self._engine._camera_configs.append(cam_cfg) + # Start physics (sim thread spawned inside engine.connect()). + if not self._engine.connect(): + raise RuntimeError("MujocoSimModule: engine.connect() failed") - logger.info("MujocoCamera engine resolved", connected=self._engine.connected) + # Publish num_joints + ready so the adapter can proceed past its + # ready-wait in connect(). Use the engine's full joint count (arm + + # gripper if present) so the adapter can detect the gripper via + # ``num_joints > dof``. + self._shm.signal_ready(num_joints=len(joint_names)) + # Camera intrinsics. self._build_camera_info() + # Camera publish thread (same pattern as former MujocoCamera). self._stop_event.clear() - self._thread = threading.Thread(target=self._publish_loop, daemon=True) - self._thread.start() + self._publish_thread = threading.Thread( + target=self._publish_loop, daemon=True, name="MujocoSimPublish" + ) + self._publish_thread.start() - # Periodic camera_info publishing + # Periodic camera_info publishing. interval_sec = 1.0 / self.config.camera_info_fps self._disposables.add( rx.interval(interval_sec).subscribe( on_next=lambda _: self._publish_camera_info(), - on_error=lambda e: logger.error("CameraInfo publish error", error=e), + on_error=lambda e: logger.error(f"CameraInfo publish error: {e}"), ) ) - # Optional pointcloud generation + # Optional pointcloud generation. if self.config.enable_pointcloud and self.config.enable_depth: pc_interval = 1.0 / self.config.pointcloud_fps self._disposables.add( backpressure(rx.interval(pc_interval)).subscribe( on_next=lambda _: self._generate_pointcloud(), - on_error=lambda e: logger.error("Pointcloud error", error=e), + on_error=lambda e: logger.error(f"Pointcloud error: {e}"), ) ) + logger.info( + "MujocoSimModule started", + address=self.config.address, + dof=dof, + camera=self.config.camera_name, + shm_key=shm_key, + ) + @rpc def stop(self) -> None: self._stop_event.set() - if self._thread and self._thread.is_alive(): - self._thread.join(timeout=2.0) - self._thread = None + if self._publish_thread and self._publish_thread.is_alive(): + self._publish_thread.join(timeout=2.0) + self._publish_thread = None + if self._engine is not None: + try: + self._engine.disconnect() + except Exception as exc: + logger.error(f"MujocoSimModule: engine.disconnect() failed: {exc}") + self._engine = None + if self._shm is not None: + try: + self._shm.signal_stop() + self._shm.cleanup() + except Exception as exc: + logger.error(f"MujocoSimModule: SHM cleanup failed: {exc}") + self._shm = None self._camera_info_base = None super().stop() + # ---------------- SHM <-> engine hooks ----------------------- + + def _apply_shm_commands(self, engine: MujocoEngine) -> None: + """Pre-step hook: pull command targets from SHM into the engine.""" + shm = self._shm + if shm is None: + return + dof = self.config.dof + + pos_cmd = shm.read_position_command(dof) + if pos_cmd is not None: + engine.write_joint_command(JointState(position=pos_cmd.tolist())) + + vel_cmd = shm.read_velocity_command(dof) + if vel_cmd is not None: + engine.write_joint_command(JointState(velocity=vel_cmd.tolist())) + + # Gripper: adapter writes raw joint-space position; sim module scales + # into actuator ctrl range before setting the target. + if self._gripper_idx is not None: + gripper_cmd = shm.read_gripper_command() + if gripper_cmd is not None: + ctrl_value = self._gripper_joint_to_ctrl(gripper_cmd) + engine.set_position_target(self._gripper_idx, ctrl_value) + + def _publish_shm_state(self, engine: MujocoEngine) -> None: + """Post-step hook: publish joint state to SHM.""" + shm = self._shm + if shm is None: + return + shm.write_joint_state( + positions=engine.joint_positions, + velocities=engine.joint_velocities, + efforts=engine.joint_efforts, + ) + if self._gripper_idx is not None: + positions = engine.joint_positions + if self._gripper_idx < len(positions): + shm.write_gripper_state(positions[self._gripper_idx]) + + def _gripper_joint_to_ctrl(self, joint_position: float) -> float: + """Map joint-space gripper position to actuator control value. + + The high-level API uses ``position ∈ [jlo, jhi]`` where ``jlo`` is + closed and ``jhi`` is open (e.g. xArm: 0.0 closed, 0.85 open). The + MJCF actuator ``ctrl`` range is inverted on xArm-style grippers + (0 open, 255 closed). We map ``jlo → chi`` and ``jhi → clo``. + """ + jlo, jhi = self._gripper_joint_range + clo, chi = self._gripper_ctrl_range + clamped = max(jlo, min(jhi, joint_position)) + if jhi == jlo: + return clo + t = (clamped - jlo) / (jhi - jlo) + return chi - t * (chi - clo) + + # ---------------- camera publishing -------------------------- + + def _build_camera_info(self) -> None: + if self._engine is None: + return + fovy_deg = self._engine.get_camera_fovy(self.config.camera_name) + if fovy_deg is None: + logger.error(f"Camera not found in MJCF: {self.config.camera_name}") + return + h = self.config.height + w = self.config.width + fovy_rad = math.radians(fovy_deg) + fy = h / (2.0 * math.tan(fovy_rad / 2.0)) + fx = fy # square pixels + self._camera_info_base = CameraInfo.from_intrinsics( + fx=fx, + fy=fy, + cx=w / 2.0, + cy=h / 2.0, + width=w, + height=h, + frame_id=self._color_optical_frame, + ) + def _publish_loop(self) -> None: """Poll engine for rendered frames and publish at configured FPS.""" interval = 1.0 / self.config.fps last_timestamp = 0.0 published_count = 0 - logger.info("MujocoCamera publish loop started", camera=self.config.camera_name) - - # Wait for engine to connect (adapter may not have started yet) + # Wait for engine to actually be connected (sim thread may take a tick). deadline = time.monotonic() + 30.0 while ( not self._stop_event.is_set() @@ -250,26 +376,22 @@ def _publish_loop(self) -> None: and not self._engine.connected ): if time.monotonic() > deadline: - logger.error("MujocoCamera: timed out waiting for engine to connect") + logger.error("MujocoSimModule: timed out waiting for engine to connect") return self._stop_event.wait(timeout=0.1) if self._stop_event.is_set(): return - logger.info("MujocoCamera: engine connected, polling for frames") - while not self._stop_event.is_set() and self._engine is not None: try: frame = self._engine.read_camera(self.config.camera_name) if frame is None or frame.timestamp <= last_timestamp: self._stop_event.wait(timeout=interval * 0.5) continue - last_timestamp = frame.timestamp ts = time.time() - # Publish RGB color_img = Image( data=frame.rgb, format=ImageFormat.RGB, @@ -278,7 +400,6 @@ def _publish_loop(self) -> None: ) self.color_image.publish(color_img) - # Publish depth (float32 meters) if self.config.enable_depth: depth_img = Image( data=frame.depth, @@ -288,24 +409,22 @@ def _publish_loop(self) -> None: ) self.depth_image.publish(depth_img) - # Publish TF (world -> camera from MuJoCo pose) self._publish_tf(ts, frame) published_count += 1 if published_count == 1: logger.info( - "MujocoCamera first frame published", + "MujocoSimModule first frame published", rgb_shape=frame.rgb.shape, depth_shape=frame.depth.shape, ) - # Rate limit elapsed = time.time() - ts sleep_time = interval - elapsed if sleep_time > 0: time.sleep(sleep_time) - except Exception as e: - logger.error("MujocoCamera publish error", error=e) + except Exception as exc: + logger.error(f"MujocoSimModule publish error: {exc}") time.sleep(1.0) def _publish_camera_info(self) -> None: @@ -313,7 +432,7 @@ def _publish_camera_info(self) -> None: if base is None: return ts = time.time() - color_info = CameraInfo( + info = CameraInfo( height=base.height, width=base.width, distortion_model=base.distortion_model, @@ -323,25 +442,21 @@ def _publish_camera_info(self) -> None: frame_id=base.frame_id, ts=ts, ) - self.camera_info.publish(color_info) - self.depth_camera_info.publish(color_info) + self.camera_info.publish(info) + self.depth_camera_info.publish(info) - def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: + def _publish_tf(self, ts: float, frame: CameraFrame | None) -> None: if frame is None: return - mj_rot = R.from_matrix(frame.cam_mat.reshape(3, 3)) optical_rot = mj_rot * _RX180 q = optical_rot.as_quat() # xyzw - pos = Vector3( float(frame.cam_pos[0]), float(frame.cam_pos[1]), float(frame.cam_pos[2]), ) rot = Quaternion(float(q[0]), float(q[1]), float(q[2]), float(q[3])) - - # Publish world -> all optical/link frames (all co-located in sim) self.tf.publish( Transform( translation=pos, @@ -367,7 +482,6 @@ def _publish_tf(self, ts: float, frame: CameraFrame | None = None) -> None: ) def _generate_pointcloud(self) -> None: - """Generate pointcloud from latest depth frame (optional, for visualization).""" if self._engine is None or self._camera_info_base is None: return frame = self._engine.read_camera(self.config.camera_name) @@ -394,8 +508,12 @@ def _generate_pointcloud(self) -> None: ) pcd = pcd.voxel_downsample(0.005) self.pointcloud.publish(pcd) - except Exception as e: - logger.error("Pointcloud generation error", error=e) + except Exception as exc: + logger.error(f"Pointcloud generation error: {exc}") + + +# Silence unused-import warning: these are re-exported for test/debug use. +_ = (CMD_MODE_POSITION, CMD_MODE_VELOCITY) -__all__ = ["MujocoCamera", "MujocoCameraConfig"] +__all__ = ["MujocoSimModule", "MujocoSimModuleConfig"] diff --git a/dimos/simulation/manipulators/sim_manip_interface.py b/dimos/simulation/manipulators/sim_manip_interface.py deleted file mode 100644 index abcfcc96f7..0000000000 --- a/dimos/simulation/manipulators/sim_manip_interface.py +++ /dev/null @@ -1,221 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Simulation-agnostic manipulator interface.""" - -from __future__ import annotations - -import logging -import math -from typing import TYPE_CHECKING - -from dimos.hardware.manipulators.spec import ControlMode, JointLimits, ManipulatorInfo -from dimos.msgs.sensor_msgs.JointState import JointState - -if TYPE_CHECKING: - from dimos.simulation.engines.base import SimulationEngine - - -class SimManipInterface: - """Adapter wrapper around a simulation engine to provide a uniform manipulator API.""" - - def __init__( - self, - engine: SimulationEngine, - dof: int | None = None, - gripper_idx: int | None = None, - gripper_ctrl_range: tuple[float, float] = (0.0, 1.0), - gripper_joint_range: tuple[float, float] = (0.0, 1.0), - ) -> None: - self.logger = logging.getLogger(self.__class__.__name__) - self._engine = engine - self._joint_names = list(engine.joint_names) - self._dof = dof if dof is not None else len(self._joint_names) - self._connected = False - self._servos_enabled = False - self._control_mode = ControlMode.POSITION - self._error_code = 0 - self._error_message = "" - self._gripper_idx = gripper_idx - self._gripper_ctrl_range = gripper_ctrl_range - self._gripper_joint_range = gripper_joint_range - - def connect(self) -> bool: - """Connect to the simulation engine.""" - try: - self.logger.info("Connecting to simulation engine...") - if not self._engine.connect(): - self.logger.error("Failed to connect to simulation engine") - return False - if self._engine.connected: - self._connected = True - self._servos_enabled = True - self.logger.info( - "Successfully connected to simulation", - extra={"dof": self._dof}, - ) - return True - self.logger.error("Failed to connect to simulation engine") - return False - except Exception as exc: - self.logger.error(f"Sim connection failed: {exc}") - return False - - def disconnect(self) -> None: - """Disconnect from simulation.""" - try: - self._engine.disconnect() - except Exception as exc: - self.logger.error(f"Sim disconnection failed: {exc}") - finally: - self._connected = False - - def is_connected(self) -> bool: - return bool(self._connected and self._engine.connected) - - def get_info(self) -> ManipulatorInfo: - vendor = "Simulation" - model = "Simulation" - dof = self._dof - return ManipulatorInfo( - vendor=vendor, - model=model, - dof=dof, - firmware_version=None, - serial_number=None, - ) - - def get_dof(self) -> int: - return self._dof - - def get_joint_names(self) -> list[str]: - return list(self._joint_names) - - def get_limits(self) -> JointLimits: - lower = [-math.pi] * self._dof - upper = [math.pi] * self._dof - max_vel_rad = math.radians(180.0) - return JointLimits( - position_lower=lower, - position_upper=upper, - velocity_max=[max_vel_rad] * self._dof, - ) - - def set_control_mode(self, mode: ControlMode) -> bool: - self._control_mode = mode - return True - - def get_control_mode(self) -> ControlMode: - return self._control_mode - - def read_joint_positions(self) -> list[float]: - positions = self._engine.read_joint_positions() - return positions[: self._dof] - - def read_joint_velocities(self) -> list[float]: - velocities = self._engine.read_joint_velocities() - return velocities[: self._dof] - - def read_joint_efforts(self) -> list[float]: - efforts = self._engine.read_joint_efforts() - return efforts[: self._dof] - - def read_state(self) -> dict[str, int]: - velocities = self.read_joint_velocities() - is_moving = any(abs(v) > 1e-4 for v in velocities) - mode_int = list(ControlMode).index(self._control_mode) - return { - "state": 1 if is_moving else 0, - "mode": mode_int, - } - - def read_error(self) -> tuple[int, str]: - return self._error_code, self._error_message - - def write_joint_positions(self, positions: list[float], velocity: float = 1.0) -> bool: - if not self._servos_enabled: - return False - self._control_mode = ControlMode.POSITION - self._engine.write_joint_command(JointState(position=positions[: self._dof])) - return True - - def write_joint_velocities(self, velocities: list[float]) -> bool: - if not self._servos_enabled: - return False - self._control_mode = ControlMode.VELOCITY - self._engine.write_joint_command(JointState(velocity=velocities[: self._dof])) - return True - - def write_joint_efforts(self, efforts: list[float]) -> bool: - if not self._servos_enabled: - return False - self._control_mode = ControlMode.TORQUE - self._engine.write_joint_command(JointState(effort=efforts[: self._dof])) - return True - - def write_stop(self) -> bool: - self._engine.hold_current_position() - return True - - def write_enable(self, enable: bool) -> bool: - self._servos_enabled = enable - return True - - def read_enabled(self) -> bool: - return self._servos_enabled - - def write_clear_errors(self) -> bool: - self._error_code = 0 - self._error_message = "" - return True - - def read_cartesian_position(self) -> dict[str, float] | None: - return None - - def write_cartesian_position( - self, - pose: dict[str, float], - velocity: float = 1.0, - ) -> bool: - _pose = pose - _velocity = velocity - return False - - def read_gripper_position(self) -> float | None: - if self._gripper_idx is None: - return None - positions = self._engine.read_joint_positions() - return positions[self._gripper_idx] - - def write_gripper_position(self, position: float) -> bool: - if self._gripper_idx is None: - return False - jlo, jhi = self._gripper_joint_range - clo, chi = self._gripper_ctrl_range - position = max(jlo, min(jhi, position)) - if jhi != jlo: - t = (position - jlo) / (jhi - jlo) - ctrl_value = clo + t * (chi - clo) - else: - ctrl_value = clo - self._engine.set_position_target(self._gripper_idx, ctrl_value) - return True - - def read_force_torque(self) -> list[float] | None: - return None - - -__all__ = [ - "SimManipInterface", -] diff --git a/dimos/simulation/manipulators/test_sim_adapter.py b/dimos/simulation/manipulators/test_sim_adapter.py deleted file mode 100644 index 1413004fee..0000000000 --- a/dimos/simulation/manipulators/test_sim_adapter.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Tests for SimMujocoAdapter and gripper integration.""" - -from __future__ import annotations - -from unittest.mock import MagicMock, patch - -import pytest - -from dimos.hardware.manipulators.sim.adapter import SimMujocoAdapter, register -from dimos.simulation.conftest import ARM_DOF - - -class TestSimMujocoAdapter: - """Tests for SimMujocoAdapter with and without gripper.""" - - @pytest.fixture - def adapter_with_gripper(self, patched_mujoco_engine_with_gripper): - """SimMujocoAdapter with ARM_DOF arm joints + 1 gripper joint.""" - return SimMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml", headless=True) - - @pytest.fixture - def adapter_no_gripper(self, patched_mujoco_engine): - """SimMujocoAdapter with ARM_DOF arm joints, no gripper.""" - return SimMujocoAdapter(dof=ARM_DOF, address="/fake/scene.xml", headless=True) - - def test_address_required(self, patched_mujoco_engine): - with pytest.raises(ValueError, match="address"): - SimMujocoAdapter(dof=ARM_DOF, address=None) - - def test_gripper_detected(self, adapter_with_gripper): - assert adapter_with_gripper._gripper_idx == ARM_DOF - - def test_no_gripper_when_dof_matches(self, adapter_no_gripper): - assert adapter_no_gripper._gripper_idx is None - - def test_read_gripper_position(self, adapter_with_gripper): - pos = adapter_with_gripper.read_gripper_position() - assert pos is not None - - def test_write_gripper_sets_target(self, adapter_with_gripper): - """Write a gripper position and verify the control target was set.""" - assert adapter_with_gripper.write_gripper_position(0.42) is True - target = adapter_with_gripper._engine._joint_position_targets[ARM_DOF] - assert target != 0.0, "write_gripper_position should update the control target" - - def test_read_gripper_position_no_gripper(self, adapter_no_gripper): - assert adapter_no_gripper.read_gripper_position() is None - - def test_write_gripper_position_no_gripper(self, adapter_no_gripper): - assert adapter_no_gripper.write_gripper_position(0.5) is False - - def test_write_gripper_does_not_clobber_arm(self, adapter_with_gripper): - """Gripper write must not overwrite arm joint targets.""" - engine = adapter_with_gripper._engine - for i in range(ARM_DOF): - engine._joint_position_targets[i] = float(i) + 1.0 - - adapter_with_gripper.write_gripper_position(0.0) - - for i in range(ARM_DOF): - assert engine._joint_position_targets[i] == pytest.approx(float(i) + 1.0) - - def test_read_joint_positions_excludes_gripper(self, adapter_with_gripper): - positions = adapter_with_gripper.read_joint_positions() - assert len(positions) == ARM_DOF - - def test_connect_and_disconnect(self, adapter_with_gripper): - with patch("dimos.simulation.engines.mujoco_engine.mujoco.mj_step"): - assert adapter_with_gripper.connect() is True - adapter_with_gripper.disconnect() - - def test_register(self): - registry = MagicMock() - register(registry) - registry.register.assert_called_once_with("sim_mujoco", SimMujocoAdapter) diff --git a/dimos/simulation/sensors/test_mujoco_camera.py b/dimos/simulation/sensors/test_mujoco_camera.py deleted file mode 100644 index a67f673966..0000000000 --- a/dimos/simulation/sensors/test_mujoco_camera.py +++ /dev/null @@ -1,237 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Tests for MujocoCamera and engine registry.""" - -from __future__ import annotations - -import math -from pathlib import Path -import threading -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest - -from dimos.simulation.engines.mujoco_engine import ( - CameraConfig, - CameraFrame, - MujocoEngine, - _engine_registry, - get_or_create_engine, - unregister_engine, -) -from dimos.simulation.sensors.mujoco_camera import MujocoCamera - -ARM_DOF = 7 - - -def _make_camera_frame( - cam_pos: list[float] | None = None, - cam_mat: list[float] | None = None, -) -> CameraFrame: - """Create a CameraFrame with sensible defaults.""" - return CameraFrame( - rgb=np.zeros((480, 640, 3), dtype=np.uint8), - depth=np.ones((480, 640), dtype=np.float32), - cam_pos=np.array(cam_pos or [1.0, 2.0, 3.0]), - cam_mat=np.array(cam_mat or np.eye(3).flatten()), - fovy=45.0, - timestamp=1.0, - ) - - -def _make_mock_engine(fovy: float = 45.0) -> MagicMock: - mock_engine = MagicMock(spec=MujocoEngine) - mock_engine.get_camera_fovy.return_value = fovy - mock_engine.connected = True - mock_engine._camera_configs = [] - mock_engine._camera_lock = threading.Lock() - mock_engine.read_camera.return_value = _make_camera_frame() - return mock_engine - - -@pytest.fixture -def mock_engine() -> MagicMock: - return _make_mock_engine() - - -@pytest.fixture -def camera_with_mock_engine(mock_engine: MagicMock): - cam = MujocoCamera(camera_name="wrist_camera") - cam._engine = mock_engine - yield cam - cam.stop() - - -class TestEngineRegistry: - def test_creates_new(self, patched_mujoco_engine): - engine = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) - assert engine is not None - assert len(_engine_registry) == 1 - - def test_returns_same_instance(self, patched_mujoco_engine): - e1 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) - e2 = get_or_create_engine(config_path=Path("/fake/scene.xml"), headless=True) - assert e1 is e2 - assert len(_engine_registry) == 1 - - def test_merges_new_cameras(self, patched_mujoco_engine): - e1 = get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_a")], - ) - get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_b")], - ) - names = {c.name for c in e1._camera_configs} - assert names == {"cam_a", "cam_b"} - - def test_deduplicates_cameras(self, patched_mujoco_engine): - get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_a")], - ) - get_or_create_engine( - config_path=Path("/fake/scene.xml"), - cameras=[CameraConfig(name="cam_a")], - ) - engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) - cam_names = [c.name for c in engine._camera_configs] - assert cam_names.count("cam_a") == 1 - - def test_unregister_removes(self, patched_mujoco_engine): - engine = get_or_create_engine(config_path=Path("/fake/scene.xml")) - assert len(_engine_registry) == 1 - unregister_engine(engine) - assert len(_engine_registry) == 0 - - -class TestCameraIntrinsics: - def test_fovy_45(self, camera_with_mock_engine: MujocoCamera): - cam = camera_with_mock_engine - cam._build_camera_info() - info = cam._camera_info_base - assert info is not None - - expected_fy = 480.0 / (2.0 * math.tan(math.radians(45.0) / 2.0)) - # K = [fx, 0, cx, 0, fy, cy, 0, 0, 1] - assert info.K[0] == pytest.approx(expected_fy, abs=0.01) # fx - assert info.K[4] == pytest.approx(expected_fy, abs=0.01) # fy - assert info.K[2] == pytest.approx(320.0) # cx - assert info.K[5] == pytest.approx(240.0) # cy - - def test_fovy_90(self, mock_engine: MagicMock): - mock_engine.get_camera_fovy.return_value = 90.0 - cam = MujocoCamera(camera_name="wrist_camera") - cam._engine = mock_engine - cam._build_camera_info() - info = cam._camera_info_base - assert info is not None - # tan(45°) = 1.0, so fy = 480 / 2 = 240 - assert info.K[0] == pytest.approx(240.0, abs=0.01) - assert info.K[4] == pytest.approx(240.0, abs=0.01) - cam.stop() - - def test_unknown_camera(self, mock_engine: MagicMock): - mock_engine.get_camera_fovy.return_value = None - cam = MujocoCamera(camera_name="nonexistent") - cam._engine = mock_engine - cam._build_camera_info() - assert cam._camera_info_base is None - cam.stop() - - def test_distortion_and_frame_id(self, camera_with_mock_engine: MujocoCamera): - cam = camera_with_mock_engine - cam._build_camera_info() - info = cam._camera_info_base - assert info is not None - assert info.D == [0.0, 0.0, 0.0, 0.0, 0.0] - assert info.distortion_model == "plumb_bob" - assert info.frame_id == "wrist_camera_color_optical_frame" - - -class TestMujocoCameraLifecycle: - def test_start_no_engine_raises(self): - cam = MujocoCamera(camera_name="cam", address="") - try: - with pytest.raises(RuntimeError, match="set address"): - cam.start() - finally: - cam.stop() - - def test_start_creates_thread(self, camera_with_mock_engine: MujocoCamera): - cam = camera_with_mock_engine - cam._engine.connected = False # thread will wait, won't spin - # Patch rx.interval to avoid spawning scheduler threads that leak - with patch("dimos.simulation.sensors.mujoco_camera.rx.interval", return_value=MagicMock()): - cam.start() - assert cam._thread is not None - assert cam._thread.is_alive() - cam.stop() - assert cam._thread is None - - def test_stop_sets_event(self, camera_with_mock_engine: MujocoCamera): - cam = camera_with_mock_engine - cam._engine.connected = False - with patch("dimos.simulation.sensors.mujoco_camera.rx.interval", return_value=MagicMock()): - cam.start() - cam.stop() - assert cam._stop_event.is_set() - assert cam._thread is None - - def test_frame_name_properties(self): - cam = MujocoCamera(camera_name="wrist_camera") - assert cam._camera_link == "wrist_camera_link" - assert cam._color_frame == "wrist_camera_color_frame" - assert cam._color_optical_frame == "wrist_camera_color_optical_frame" - assert cam._depth_frame == "wrist_camera_depth_frame" - assert cam._depth_optical_frame == "wrist_camera_depth_optical_frame" - cam.stop() - - -class TestTFPublishing: - def test_publish_tf_correct_frames(self, camera_with_mock_engine: MujocoCamera): - cam = camera_with_mock_engine - mock_tf = MagicMock() - frame = _make_camera_frame(cam_pos=[1.0, 2.0, 3.0]) - - with patch.object(type(cam), "tf", new_callable=lambda: property(lambda self: mock_tf)): - cam._publish_tf(ts=0.0, frame=frame) - - mock_tf.publish.assert_called_once() - transforms = mock_tf.publish.call_args.args - assert len(transforms) == 3 - - child_ids = {t.child_frame_id for t in transforms} - assert child_ids == { - "wrist_camera_color_optical_frame", - "wrist_camera_depth_optical_frame", - "wrist_camera_link", - } - for t in transforms: - assert t.frame_id == "world" - assert t.translation.x == pytest.approx(1.0) - assert t.translation.y == pytest.approx(2.0) - assert t.translation.z == pytest.approx(3.0) - - def test_publish_tf_none_noop(self, camera_with_mock_engine: MujocoCamera): - cam = camera_with_mock_engine - mock_tf = MagicMock() - - with patch.object(type(cam), "tf", new_callable=lambda: property(lambda self: mock_tf)): - cam._publish_tf(ts=0.0, frame=None) - - mock_tf.publish.assert_not_called() From 69cc28c09434774012efd5f823ceb81b8efc33bd Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Sat, 4 Apr 2026 19:00:15 -0700 Subject: [PATCH 34/41] misc: fstring style logging --- dimos/simulation/engines/mujoco_engine.py | 24 +++++++++++-------- dimos/simulation/engines/mujoco_sim_module.py | 14 +++++------ 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index cbc29d69e6..773426ef59 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -202,7 +202,7 @@ def _update_joint_state(self) -> None: def connect(self) -> bool: try: - logger.info(f"{self.__class__.__name__}: connect()") + logger.info("connect()", cls=self.__class__.__name__) with self._lock: self._connected = True self._stop_event.clear() @@ -216,12 +216,12 @@ def connect(self) -> bool: self._sim_thread.start() return True except Exception as e: - logger.error(f"{self.__class__.__name__}: connect() failed: {e}") + logger.error("connect() failed", cls=self.__class__.__name__, error=str(e)) return False def disconnect(self) -> bool: try: - logger.info(f"{self.__class__.__name__}: disconnect()") + logger.info("disconnect()", cls=self.__class__.__name__) with self._lock: self._connected = False self._stop_event.set() @@ -230,7 +230,7 @@ def disconnect(self) -> bool: self._sim_thread = None return True except Exception as e: - logger.error(f"{self.__class__.__name__}: disconnect() failed: {e}") + logger.error("disconnect() failed", cls=self.__class__.__name__, error=str(e)) return False def _init_new_cameras(self, cam_renderers: dict[str, _CameraRendererState]) -> None: @@ -245,7 +245,7 @@ def _init_new_cameras(self, cam_renderers: dict[str, _CameraRendererState]) -> N continue cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name) if cam_id < 0: - logger.warning(f"Camera '{cfg.name}' not found in MJCF, skipping") + logger.warning("Camera not found in MJCF, skipping", camera_name=cfg.name) continue rgb_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) depth_renderer = mujoco.Renderer(self._model, height=cfg.height, width=cfg.width) @@ -259,7 +259,11 @@ def _init_new_cameras(self, cam_renderers: dict[str, _CameraRendererState]) -> N interval=interval, ) logger.info( - f"Camera '{cfg.name}' renderer created ({cfg.width}x{cfg.height} @ {cfg.fps}fps)" + "Camera renderer created", + camera_name=cfg.name, + width=cfg.width, + height=cfg.height, + fps=cfg.fps, ) def _render_cameras(self, now: float, cam_renderers: dict[str, _CameraRendererState]) -> None: @@ -294,7 +298,7 @@ def _close_cam_renderers(cam_renderers: dict[str, _CameraRendererState]) -> None state.depth_renderer.close() def _sim_loop(self) -> None: - logger.info(f"{self.__class__.__name__}: sim loop started") + logger.info("sim loop started", cls=self.__class__.__name__) dt = 1.0 / self._control_frequency # Camera renderers — created in sim thread (MuJoCo thread-safety). @@ -307,7 +311,7 @@ def _step_once(sync_viewer: bool) -> None: try: self._on_before_step(self) except Exception as exc: - logger.error(f"on_before_step failed: {exc}") + logger.error("on_before_step failed", error=str(exc)) self._apply_control() mujoco.mj_step(self._model, self._data) if sync_viewer: @@ -317,7 +321,7 @@ def _step_once(sync_viewer: bool) -> None: try: self._on_after_step(self) except Exception as exc: - logger.error(f"on_after_step failed: {exc}") + logger.error("on_after_step failed", error=str(exc)) self._render_cameras(loop_start, cam_renderers) elapsed = time.time() - loop_start @@ -336,7 +340,7 @@ def _step_once(sync_viewer: bool) -> None: _step_once(sync_viewer=True) self._close_cam_renderers(cam_renderers) - logger.info(f"{self.__class__.__name__}: sim loop stopped") + logger.info("sim loop stopped", cls=self.__class__.__name__) @property def connected(self) -> bool: diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index ce872c2fde..2a6ed7547c 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -238,7 +238,7 @@ def start(self) -> None: self._disposables.add( rx.interval(interval_sec).subscribe( on_next=lambda _: self._publish_camera_info(), - on_error=lambda e: logger.error(f"CameraInfo publish error: {e}"), + on_error=lambda e: logger.error("CameraInfo publish error", error=str(e)), ) ) @@ -248,7 +248,7 @@ def start(self) -> None: self._disposables.add( backpressure(rx.interval(pc_interval)).subscribe( on_next=lambda _: self._generate_pointcloud(), - on_error=lambda e: logger.error(f"Pointcloud error: {e}"), + on_error=lambda e: logger.error("Pointcloud error", error=str(e)), ) ) @@ -270,14 +270,14 @@ def stop(self) -> None: try: self._engine.disconnect() except Exception as exc: - logger.error(f"MujocoSimModule: engine.disconnect() failed: {exc}") + logger.error("engine.disconnect() failed", error=str(exc)) self._engine = None if self._shm is not None: try: self._shm.signal_stop() self._shm.cleanup() except Exception as exc: - logger.error(f"MujocoSimModule: SHM cleanup failed: {exc}") + logger.error("SHM cleanup failed", error=str(exc)) self._shm = None self._camera_info_base = None super().stop() @@ -345,7 +345,7 @@ def _build_camera_info(self) -> None: return fovy_deg = self._engine.get_camera_fovy(self.config.camera_name) if fovy_deg is None: - logger.error(f"Camera not found in MJCF: {self.config.camera_name}") + logger.error("Camera not found in MJCF", camera_name=self.config.camera_name) return h = self.config.height w = self.config.width @@ -424,7 +424,7 @@ def _publish_loop(self) -> None: if sleep_time > 0: time.sleep(sleep_time) except Exception as exc: - logger.error(f"MujocoSimModule publish error: {exc}") + logger.error("MujocoSimModule publish error", error=str(exc)) time.sleep(1.0) def _publish_camera_info(self) -> None: @@ -509,7 +509,7 @@ def _generate_pointcloud(self) -> None: pcd = pcd.voxel_downsample(0.005) self.pointcloud.publish(pcd) except Exception as exc: - logger.error(f"Pointcloud generation error: {exc}") + logger.error("Pointcloud generation error", error=str(exc)) # Silence unused-import warning: these are re-exported for test/debug use. From da771882c66e8e0e0bde4b9159cee408a778eeb4 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Sat, 4 Apr 2026 19:28:21 -0700 Subject: [PATCH 35/41] except checks --- dimos/simulation/engines/mujoco_shm.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/dimos/simulation/engines/mujoco_shm.py b/dimos/simulation/engines/mujoco_shm.py index 924fa3fe56..9976da3ec4 100644 --- a/dimos/simulation/engines/mujoco_shm.py +++ b/dimos/simulation/engines/mujoco_shm.py @@ -247,12 +247,16 @@ def cleanup(self) -> None: for shm in self.shm.as_list(): try: shm.close() - except Exception: - pass + except FileNotFoundError: + pass # already detached + except OSError as exc: + logger.warning("SHM close failed", name=shm.name, error=str(exc)) try: shm.unlink() - except Exception: - pass + except FileNotFoundError: + pass # already unlinked (e.g. cleanup called twice) + except OSError as exc: + logger.warning("SHM unlink failed", name=shm.name, error=str(exc)) # ---------------- helpers ------------------------------------ @@ -337,8 +341,10 @@ def cleanup(self) -> None: for shm in self.shm.as_list(): try: shm.close() - except Exception: - pass + except FileNotFoundError: + pass # already detached + except OSError as exc: + logger.warning("SHM close failed", name=shm.name, error=str(exc)) # ---------------- helpers ------------------------------------ From 1118704fcdc269b8fe40cb947b1dcef220c81f77 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Sat, 4 Apr 2026 19:42:59 -0700 Subject: [PATCH 36/41] misc: cleanup --- dimos/hardware/manipulators/sim/adapter.py | 91 ++++++----- .../hardware/manipulators/sim/test_adapter.py | 9 +- dimos/simulation/conftest.py | 122 -------------- dimos/simulation/engines/mujoco_shm.py | 8 +- dimos/simulation/engines/mujoco_sim_module.py | 153 +++++++++++------- 5 files changed, 156 insertions(+), 227 deletions(-) delete mode 100644 dimos/simulation/conftest.py diff --git a/dimos/hardware/manipulators/sim/adapter.py b/dimos/hardware/manipulators/sim/adapter.py index 03c90ba1f5..e315b17de7 100644 --- a/dimos/hardware/manipulators/sim/adapter.py +++ b/dimos/hardware/manipulators/sim/adapter.py @@ -30,7 +30,6 @@ from __future__ import annotations -import logging import math import time from typing import TYPE_CHECKING, Any @@ -44,11 +43,14 @@ ManipShmReader, shm_key_from_path, ) +from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from dimos.hardware.manipulators.registry import AdapterRegistry +logger = setup_logger() + _READY_WAIT_TIMEOUT_S = 60.0 _READY_WAIT_POLL_S = 0.1 _ATTACH_RETRY_TIMEOUT_S = 30.0 @@ -71,7 +73,6 @@ def __init__( ) -> None: if address is None: raise ValueError("address (MJCF XML path) is required for sim_mujoco adapter") - self.logger = logging.getLogger(self.__class__.__name__) self._dof = dof self._address = address self._hardware_id = hardware_id @@ -83,49 +84,55 @@ def __init__( self._error_code = 0 self._error_message = "" self._has_gripper = False + self._effort_mode_warned = False # ---------------- connection --------------------------------- def connect(self) -> bool: - try: - # Attach to SHM. If the sim module hasn't created the buffers yet, - # retry for a short window — module startup ordering is not - # strictly guaranteed, only that they're both in the same blueprint. - deadline = time.monotonic() + _ATTACH_RETRY_TIMEOUT_S - while True: - try: - self._shm = ManipShmReader(self._shm_key) - break - except FileNotFoundError: - if time.monotonic() > deadline: - self.logger.error( - f"ShmMujocoAdapter: SHM buffers not found for {self._address} " - f"(key={self._shm_key}) after {_ATTACH_RETRY_TIMEOUT_S}s" - ) - return False - time.sleep(_ATTACH_RETRY_POLL_S) - - # Wait for sim module to signal ready. - deadline = time.monotonic() + _READY_WAIT_TIMEOUT_S - while not self._shm.is_ready(): + # Attach to SHM. If the sim module hasn't created the buffers yet, + # retry for a short window — module startup ordering is not + # strictly guaranteed, only that they're both in the same blueprint. + # Only the ManipShmReader construction is wrapped, since that's the + # only call here with a recoverable failure mode (buffers not yet + # created). Everything below that point either succeeds or raises + # a real bug that shouldn't be silently swallowed. + deadline = time.monotonic() + _ATTACH_RETRY_TIMEOUT_S + while True: + try: + self._shm = ManipShmReader(self._shm_key) + break + except FileNotFoundError: if time.monotonic() > deadline: - self.logger.error( - f"ShmMujocoAdapter: sim module not ready after {_READY_WAIT_TIMEOUT_S}s" + logger.error( + "SHM buffers not found", + address=self._address, + shm_key=self._shm_key, + timeout_s=_ATTACH_RETRY_TIMEOUT_S, ) return False - time.sleep(_READY_WAIT_POLL_S) - - num_joints = self._shm.num_joints() - self._has_gripper = num_joints > self._dof - self._connected = True - self._servos_enabled = True - self.logger.info( - f"ShmMujocoAdapter connected (dof={self._dof}, gripper={self._has_gripper})" - ) - return True - except Exception as exc: - self.logger.error(f"ShmMujocoAdapter: connect() failed: {exc}") - return False + time.sleep(_ATTACH_RETRY_POLL_S) + + # Wait for sim module to signal ready. + deadline = time.monotonic() + _READY_WAIT_TIMEOUT_S + while not self._shm.is_ready(): + if time.monotonic() > deadline: + logger.error( + "sim module not ready", + timeout_s=_READY_WAIT_TIMEOUT_S, + ) + return False + time.sleep(_READY_WAIT_POLL_S) + + num_joints = self._shm.num_joints() + self._has_gripper = num_joints > self._dof + self._connected = True + self._servos_enabled = True + logger.info( + "ShmMujocoAdapter connected", + dof=self._dof, + gripper=self._has_gripper, + ) + return True def disconnect(self) -> None: try: @@ -215,6 +222,14 @@ def write_joint_velocities(self, velocities: list[float]) -> bool: def write_joint_efforts(self, efforts: list[float]) -> bool: # Effort mode not exposed via SHM yet; caller can fall back to position. + # Log once per adapter instance so the caller sees feedback without + # spamming the log each tick. + if not self._effort_mode_warned: + logger.warning( + "write_joint_efforts not supported by sim adapter; ignoring and returning False", + dof=self._dof, + ) + self._effort_mode_warned = True return False def write_stop(self) -> bool: diff --git a/dimos/hardware/manipulators/sim/test_adapter.py b/dimos/hardware/manipulators/sim/test_adapter.py index d8ea1fe485..45dae82964 100644 --- a/dimos/hardware/manipulators/sim/test_adapter.py +++ b/dimos/hardware/manipulators/sim/test_adapter.py @@ -21,6 +21,7 @@ import pytest +import dimos.hardware.manipulators.sim.adapter as adapter_mod from dimos.hardware.manipulators.sim.adapter import ShmMujocoAdapter, register from dimos.hardware.manipulators.spec import ControlMode, ManipulatorAdapter from dimos.simulation.engines.mujoco_shm import ManipShmWriter @@ -40,8 +41,6 @@ def writer(shm_key, monkeypatch): We monkey-patch ``shm_key_from_path`` so the adapter under test resolves to our fixture's key regardless of the address string. """ - import dimos.hardware.manipulators.sim.adapter as adapter_mod - monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) w = ManipShmWriter(shm_key) w.signal_ready(num_joints=ARM_DOF) @@ -51,8 +50,6 @@ def writer(shm_key, monkeypatch): @pytest.fixture def writer_with_gripper(shm_key, monkeypatch): - import dimos.hardware.manipulators.sim.adapter as adapter_mod - monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) w = ManipShmWriter(shm_key) w.signal_ready(num_joints=ARM_DOF + 1) @@ -177,8 +174,6 @@ def test_write_gripper_position_no_gripper(self, adapter): class TestConnect: def test_connect_before_sim_ready_times_out(self, shm_key, monkeypatch): """If sim module never signals ready, connect() returns False after timeout.""" - import dimos.hardware.manipulators.sim.adapter as adapter_mod - monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) # Shrink timeouts so the test runs fast. monkeypatch.setattr(adapter_mod, "_READY_WAIT_TIMEOUT_S", 0.2) @@ -194,8 +189,6 @@ def test_connect_before_sim_ready_times_out(self, shm_key, monkeypatch): def test_connect_waits_for_shm(self, shm_key, monkeypatch): """If SHM buffers don't exist yet, connect() retries briefly.""" - import dimos.hardware.manipulators.sim.adapter as adapter_mod - monkeypatch.setattr(adapter_mod, "shm_key_from_path", lambda _: shm_key) monkeypatch.setattr(adapter_mod, "_ATTACH_RETRY_TIMEOUT_S", 0.2) monkeypatch.setattr(adapter_mod, "_ATTACH_RETRY_POLL_S", 0.02) diff --git a/dimos/simulation/conftest.py b/dimos/simulation/conftest.py deleted file mode 100644 index 270c010c9a..0000000000 --- a/dimos/simulation/conftest.py +++ /dev/null @@ -1,122 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Shared test fixtures for simulation tests.""" - -from __future__ import annotations - -from pathlib import Path -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest - -from dimos.simulation.utils.xml_parser import JointMapping - -ARM_DOF = 7 - - -def _make_joint_mapping(name: str, idx: int) -> JointMapping: - """Create a JointMapping for a simple revolute joint.""" - return JointMapping( - name=name, - joint_id=idx, - actuator_id=idx, - qpos_adr=idx, - dof_adr=idx, - tendon_qpos_adrs=(), - tendon_dof_adrs=(), - ) - - -def _make_gripper_mapping(name: str, idx: int) -> JointMapping: - """Create a JointMapping for a tendon-driven gripper.""" - return JointMapping( - name=name, - joint_id=None, - actuator_id=idx, - qpos_adr=None, - dof_adr=None, - tendon_qpos_adrs=(idx, idx + 1), - tendon_dof_adrs=(idx, idx + 1), - ) - - -def _patch_mujoco_engine(n_joints: int = ARM_DOF): - """Patch only the MuJoCo C-library and filesystem boundaries. - - Mocks ``_resolve_xml_path``, ``MjModel.from_xml_path``, ``MjData``, and - ``build_joint_mappings`` — the rest of ``MujocoEngine.__init__`` runs as-is. - """ - mappings = [_make_joint_mapping(f"joint{i}", i) for i in range(ARM_DOF)] - if n_joints > ARM_DOF: - mappings.append(_make_gripper_mapping(f"joint{ARM_DOF}", ARM_DOF)) - - fake_model = MagicMock() - fake_model.opt.timestep = 0.002 - fake_model.nu = n_joints - fake_model.nq = n_joints - fake_model.njnt = n_joints - fake_model.actuator_ctrlrange = np.array( - [[-6.28, 6.28]] * ARM_DOF + ([[0.0, 255.0]] if n_joints > ARM_DOF else []) - ) - fake_model.jnt_range = np.array( - [[-6.28, 6.28]] * ARM_DOF + ([[0.0, 0.85]] if n_joints > ARM_DOF else []) - ) - fake_model.jnt_qposadr = np.arange(n_joints) - - fake_data = MagicMock() - fake_data.qpos = np.zeros(n_joints + 4) # extra for tendon qpos addresses - fake_data.actuator_length = np.zeros(n_joints) - - patches = [ - patch( - "dimos.simulation.engines.mujoco_engine.MujocoEngine._resolve_xml_path", - return_value=Path("/fake/scene.xml"), - ), - patch( - "dimos.simulation.engines.mujoco_engine.mujoco.MjModel.from_xml_path", - return_value=fake_model, - ), - patch("dimos.simulation.engines.mujoco_engine.mujoco.MjData", return_value=fake_data), - patch("dimos.simulation.engines.mujoco_engine.build_joint_mappings", return_value=mappings), - ] - return patches - - -@pytest.fixture -def patched_mujoco_engine(): - """Start MuJoCo engine patches for the test, stop on teardown. - - Usage: - def test_something(patched_mujoco_engine): - engine = MujocoEngine(config_path=Path("/fake/scene.xml"), headless=True) - """ - patches = _patch_mujoco_engine(ARM_DOF) - for p in patches: - p.start() - yield - for p in patches: - p.stop() - - -@pytest.fixture -def patched_mujoco_engine_with_gripper(): - """Same as patched_mujoco_engine but with ARM_DOF+1 joints (gripper).""" - patches = _patch_mujoco_engine(ARM_DOF + 1) - for p in patches: - p.start() - yield - for p in patches: - p.stop() diff --git a/dimos/simulation/engines/mujoco_shm.py b/dimos/simulation/engines/mujoco_shm.py index 9976da3ec4..64fb5fdbc2 100644 --- a/dimos/simulation/engines/mujoco_shm.py +++ b/dimos/simulation/engines/mujoco_shm.py @@ -127,13 +127,17 @@ class ManipShmSet: def create(cls, key: str) -> ManipShmSet: """Create new SHM buffers with deterministic names derived from *key*. - If stale buffers from a prior run exist, unlink them first. + If stale buffers from a prior run exist, unlink them first. The stale + probe uses ``_unregister`` because that handle is transient — we never + owned it and must not let the tracker think we did. The create-side + handle, by contrast, *must* stay registered with ``resource_tracker`` + so that ``shm.unlink()`` in ``cleanup()`` matches up cleanly. """ buffers: dict[str, SharedMemory] = {} for buffer_name, size in _shm_sizes.items(): name = _buffer_name(key, buffer_name) try: - stale = SharedMemory(name=name) + stale = _unregister(SharedMemory(name=name)) stale.close() try: stale.unlink() diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 2a6ed7547c..48c57f5051 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -53,14 +53,11 @@ MujocoEngine, ) from dimos.simulation.engines.mujoco_shm import ( - CMD_MODE_POSITION, - CMD_MODE_VELOCITY, ManipShmWriter, shm_key_from_path, ) from dimos.spec import perception from dimos.utils.logging_config import setup_logger -from dimos.utils.reactive import backpressure logger = setup_logger() @@ -158,11 +155,17 @@ def _depth_optical_frame(self) -> str: @rpc def get_color_camera_info(self) -> CameraInfo | None: - return self._camera_info_base + # Return a fresh copy — callers must not be able to mutate our + # cached intrinsics, and the color/depth getters must not alias. + if self._camera_info_base is None: + return None + return self._camera_info_base.with_ts(time.time()) @rpc def get_depth_camera_info(self) -> CameraInfo | None: - return self._camera_info_base + if self._camera_info_base is None: + return None + return self._camera_info_base.with_ts(time.time()) @rpc def get_depth_scale(self) -> float: @@ -246,7 +249,7 @@ def start(self) -> None: if self.config.enable_pointcloud and self.config.enable_depth: pc_interval = 1.0 / self.config.pointcloud_fps self._disposables.add( - backpressure(rx.interval(pc_interval)).subscribe( + rx.interval(pc_interval).subscribe( on_next=lambda _: self._generate_pointcloud(), on_error=lambda e: logger.error("Pointcloud error", error=str(e)), ) @@ -266,22 +269,38 @@ def stop(self) -> None: if self._publish_thread and self._publish_thread.is_alive(): self._publish_thread.join(timeout=2.0) self._publish_thread = None + + # Collect teardown failures across both halves (engine + SHM) so that + # a failure in the first doesn't skip the second. Only clear the + # corresponding ``self._*`` reference if teardown actually succeeded + # — silently nulling on failure would let the next ``start()`` run on + # top of a half-torn-down state and mask the real bug. + errors: list[tuple[str, BaseException]] = [] if self._engine is not None: try: self._engine.disconnect() + self._engine = None except Exception as exc: logger.error("engine.disconnect() failed", error=str(exc)) - self._engine = None + errors.append(("engine.disconnect", exc)) if self._shm is not None: try: self._shm.signal_stop() self._shm.cleanup() + self._shm = None except Exception as exc: logger.error("SHM cleanup failed", error=str(exc)) - self._shm = None + errors.append(("shm.cleanup", exc)) + self._camera_info_base = None super().stop() + if errors: + # Surface the first failure so the operator sees it. Subsequent + # ones are already logged above. + op, exc = errors[0] + raise RuntimeError(f"MujocoSimModule.stop() failed during {op}: {exc}") from exc + # ---------------- SHM <-> engine hooks ----------------------- def _apply_shm_commands(self, engine: MujocoEngine) -> None: @@ -363,18 +382,32 @@ def _build_camera_info(self) -> None: ) def _publish_loop(self) -> None: - """Poll engine for rendered frames and publish at configured FPS.""" + """Poll engine for rendered frames and publish at configured FPS. + + Threading note: this loop reads ``self._engine`` without a lock. The + safety invariant is enforced by ``stop()``, which sets + ``self._stop_event`` *before* it joins this thread and only clears + ``self._engine`` *after* the join returns. So while this loop is + live, either ``_stop_event`` is unset (and ``_engine`` is guaranteed + non-None by construction in ``start()``), or ``_stop_event`` is set + and the next iteration will exit. The ``self._engine is not None`` + guards below are belt-and-suspenders, not the real correctness + argument — don't rely on them if you reorder ``stop()``. + """ + # Capture engine locally once up front. This makes the invariant + # explicit to the reader and removes the "is this racing with stop?" + # question from every subsequent access. + engine = self._engine + if engine is None: + return + interval = 1.0 / self.config.fps last_timestamp = 0.0 published_count = 0 # Wait for engine to actually be connected (sim thread may take a tick). deadline = time.monotonic() + 30.0 - while ( - not self._stop_event.is_set() - and self._engine is not None - and not self._engine.connected - ): + while not self._stop_event.is_set() and not engine.connected: if time.monotonic() > deadline: logger.error("MujocoSimModule: timed out waiting for engine to connect") return @@ -383,49 +416,59 @@ def _publish_loop(self) -> None: if self._stop_event.is_set(): return - while not self._stop_event.is_set() and self._engine is not None: + while not self._stop_event.is_set(): + # Only MuJoCo's rendering path (``read_camera``) is known to raise + # transient RuntimeErrors — everything else in this loop is our + # own code and a raise there is a real bug we want to see. We + # catch narrowly around the render call, not the whole body. try: - frame = self._engine.read_camera(self.config.camera_name) - if frame is None or frame.timestamp <= last_timestamp: - self._stop_event.wait(timeout=interval * 0.5) - continue - last_timestamp = frame.timestamp - ts = time.time() - - color_img = Image( - data=frame.rgb, - format=ImageFormat.RGB, + frame = engine.read_camera(self.config.camera_name) + except RuntimeError as exc: + logger.error( + "MuJoCo render failed; stopping publish loop", + camera_name=self.config.camera_name, + error=str(exc), + exc_info=True, + ) + return + + if frame is None or frame.timestamp <= last_timestamp: + self._stop_event.wait(timeout=interval * 0.5) + continue + last_timestamp = frame.timestamp + ts = time.time() + + color_img = Image( + data=frame.rgb, + format=ImageFormat.RGB, + frame_id=self._color_optical_frame, + ts=ts, + ) + self.color_image.publish(color_img) + + if self.config.enable_depth: + depth_img = Image( + data=frame.depth, + format=ImageFormat.DEPTH, frame_id=self._color_optical_frame, ts=ts, ) - self.color_image.publish(color_img) - - if self.config.enable_depth: - depth_img = Image( - data=frame.depth, - format=ImageFormat.DEPTH, - frame_id=self._color_optical_frame, - ts=ts, - ) - self.depth_image.publish(depth_img) - - self._publish_tf(ts, frame) - - published_count += 1 - if published_count == 1: - logger.info( - "MujocoSimModule first frame published", - rgb_shape=frame.rgb.shape, - depth_shape=frame.depth.shape, - ) - - elapsed = time.time() - ts - sleep_time = interval - elapsed - if sleep_time > 0: - time.sleep(sleep_time) - except Exception as exc: - logger.error("MujocoSimModule publish error", error=str(exc)) - time.sleep(1.0) + self.depth_image.publish(depth_img) + + self._publish_tf(ts, frame) + + published_count += 1 + if published_count == 1: + logger.info( + "MujocoSimModule first frame published", + rgb_shape=frame.rgb.shape, + depth_shape=frame.depth.shape, + ) + + elapsed = time.time() - ts + sleep_time = interval - elapsed + if sleep_time > 0: + time.sleep(sleep_time) def _publish_camera_info(self) -> None: base = self._camera_info_base @@ -512,8 +555,4 @@ def _generate_pointcloud(self) -> None: logger.error("Pointcloud generation error", error=str(exc)) -# Silence unused-import warning: these are re-exported for test/debug use. -_ = (CMD_MODE_POSITION, CMD_MODE_VELOCITY) - - __all__ = ["MujocoSimModule", "MujocoSimModuleConfig"] From e09af324516eb789a0650847c021bd917d40b0bb Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 6 Apr 2026 12:33:22 -0700 Subject: [PATCH 37/41] misc: cleanup code + comments --- dimos/hardware/manipulators/sim/adapter.py | 47 ++-------------- dimos/manipulation/blueprints.py | 6 +- dimos/simulation/engines/mujoco_engine.py | 40 +++----------- dimos/simulation/engines/mujoco_shm.py | 13 +---- dimos/simulation/engines/mujoco_sim_module.py | 55 +------------------ 5 files changed, 17 insertions(+), 144 deletions(-) diff --git a/dimos/hardware/manipulators/sim/adapter.py b/dimos/hardware/manipulators/sim/adapter.py index e315b17de7..b1ab342218 100644 --- a/dimos/hardware/manipulators/sim/adapter.py +++ b/dimos/hardware/manipulators/sim/adapter.py @@ -13,19 +13,7 @@ # limitations under the License. """Shared-memory adapter for MuJoCo-based manipulator simulation. - -This adapter is the hot-path counterpart to :class:`MujocoSimModule`. The sim -module owns the physics engine and publishes joint state to SHM; this adapter -(which lives inside ``ControlCoordinator``'s process) reads from and writes to -the same SHM buffers. - -Why SHM instead of RPC? The coordinator's tick loop calls the adapter -synchronously at 100Hz — ~4 calls per tick (positions, velocities, efforts, -command). LCM RPC latency (~1-5ms per call) would blow the 10ms budget. SHM -reads are sub-microsecond. - -Discovery is keyless: both sides derive the SHM buffer names from the MJCF -path, so no name exchange is required. +this adapter reads from and writes to the same SHM buffers. """ from __future__ import annotations @@ -86,16 +74,7 @@ def __init__( self._has_gripper = False self._effort_mode_warned = False - # ---------------- connection --------------------------------- - def connect(self) -> bool: - # Attach to SHM. If the sim module hasn't created the buffers yet, - # retry for a short window — module startup ordering is not - # strictly guaranteed, only that they're both in the same blueprint. - # Only the ManipShmReader construction is wrapped, since that's the - # only call here with a recoverable failure mode (buffers not yet - # created). Everything below that point either succeeds or raises - # a real bug that shouldn't be silently swallowed. deadline = time.monotonic() + _ATTACH_RETRY_TIMEOUT_S while True: try: @@ -116,10 +95,7 @@ def connect(self) -> bool: deadline = time.monotonic() + _READY_WAIT_TIMEOUT_S while not self._shm.is_ready(): if time.monotonic() > deadline: - logger.error( - "sim module not ready", - timeout_s=_READY_WAIT_TIMEOUT_S, - ) + logger.error("sim module not ready", timeout_s=_READY_WAIT_TIMEOUT_S) return False time.sleep(_READY_WAIT_POLL_S) @@ -127,11 +103,7 @@ def connect(self) -> bool: self._has_gripper = num_joints > self._dof self._connected = True self._servos_enabled = True - logger.info( - "ShmMujocoAdapter connected", - dof=self._dof, - gripper=self._has_gripper, - ) + logger.info("ShmMujocoAdapter connected", dof=self._dof, gripper=self._has_gripper) return True def disconnect(self) -> None: @@ -145,8 +117,6 @@ def disconnect(self) -> None: def is_connected(self) -> bool: return self._connected and self._shm is not None - # ---------------- info --------------------------------------- - def get_info(self) -> ManipulatorInfo: return ManipulatorInfo( vendor="Simulation", @@ -169,8 +139,6 @@ def get_limits(self) -> JointLimits: velocity_max=[max_vel_rad] * self._dof, ) - # ---------------- control mode ------------------------------- - def set_control_mode(self, mode: ControlMode) -> bool: self._control_mode = mode return True @@ -178,8 +146,6 @@ def set_control_mode(self, mode: ControlMode) -> bool: def get_control_mode(self) -> ControlMode: return self._control_mode - # ---------------- read state (SHM -> caller) ----------------- - def read_joint_positions(self) -> list[float]: if self._shm is None: return [0.0] * self._dof @@ -204,8 +170,6 @@ def read_state(self) -> dict[str, int]: def read_error(self) -> tuple[int, str]: return self._error_code, self._error_message - # ---------------- write commands (caller -> SHM) ------------- - def write_joint_positions(self, positions: list[float], velocity: float = 1.0) -> bool: if not self._servos_enabled or self._shm is None: return False @@ -222,8 +186,6 @@ def write_joint_velocities(self, velocities: list[float]) -> bool: def write_joint_efforts(self, efforts: list[float]) -> bool: # Effort mode not exposed via SHM yet; caller can fall back to position. - # Log once per adapter instance so the caller sees feedback without - # spamming the log each tick. if not self._effort_mode_warned: logger.warning( "write_joint_efforts not supported by sim adapter; ignoring and returning False", @@ -252,7 +214,7 @@ def write_clear_errors(self) -> bool: self._error_message = "" return True - # ---------------- cartesian (not supported) ------------------ + # ---------------- cartesian (not supported yet) ------------------ def read_cartesian_position(self) -> dict[str, float] | None: return None @@ -270,7 +232,6 @@ def read_gripper_position(self) -> float | None: def write_gripper_position(self, position: float) -> bool: if not self._has_gripper or self._shm is None: return False - # Raw joint-space position; sim module handles the joint->ctrl scaling. self._shm.write_gripper_command(position) return True diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index 37518d1231..7114f92779 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -495,9 +495,9 @@ def _make_piper_config( # Sim perception: MujocoSimModule owns the MujocoEngine and publishes both -# camera streams (replacing the old MujocoCamera) and joint state via shared -# memory. The sim_xarm7 hardware component resolves to ShmMujocoAdapter, which -# attaches to the same SHM by MJCF path — no globals, no cross-process sharing. +# camera streams and joint state via shared memory. +# The sim_xarm7 hardware component resolves to ShmMujocoAdapter, which +# attaches to the same SHM by MJCF path from dimos.control.blueprints._hardware import XARM7_SIM_PATH, sim_xarm7 from dimos.control.coordinator import TaskConfig as TaskConfig diff --git a/dimos/simulation/engines/mujoco_engine.py b/dimos/simulation/engines/mujoco_engine.py index 773426ef59..b46c4bde12 100644 --- a/dimos/simulation/engines/mujoco_engine.py +++ b/dimos/simulation/engines/mujoco_engine.py @@ -43,8 +43,6 @@ @dataclass class CameraConfig: - """Configuration for a camera to be rendered inside the sim loop.""" - name: str width: int = 640 height: int = 480 @@ -53,8 +51,6 @@ class CameraConfig: @dataclass class CameraFrame: - """Thread-safe container for a rendered camera frame.""" - rgb: NDArray[np.uint8] depth: NDArray[np.float32] cam_pos: NDArray[np.float64] @@ -65,8 +61,6 @@ class CameraFrame: @dataclass class _CameraRendererState: - """Mutable state for a single camera renderer (sim-thread only).""" - cfg: CameraConfig cam_id: int rgb_renderer: mujoco.Renderer @@ -93,12 +87,6 @@ def __init__( on_after_step: StepHook | None = None, ) -> None: super().__init__(config_path=config_path, headless=headless) - # Optional observer callbacks invoked inside the sim thread. - # ``on_before_step`` runs before ``_apply_control`` (use it to pull - # commands into the engine); ``on_after_step`` runs after - # ``_update_joint_state`` (use it to publish joint state to an - # external sink, e.g. shared memory). The engine treats these as - # opaque callables and has no knowledge of what they do. self._on_before_step: StepHook | None = on_before_step self._on_after_step: StepHook | None = on_after_step @@ -233,16 +221,10 @@ def disconnect(self) -> bool: logger.error("disconnect() failed", cls=self.__class__.__name__, error=str(e)) return False - def _init_new_cameras(self, cam_renderers: dict[str, _CameraRendererState]) -> None: - """Create renderers for any camera configs not yet initialized. - - Must be called from the sim thread (MuJoCo thread-safety). - """ - with self._camera_lock: - configs_snapshot = list(self._camera_configs) - for cfg in configs_snapshot: - if cfg.name in cam_renderers: - continue + def _init_cameras(self) -> dict[str, _CameraRendererState]: + """Create renderers for all configured cameras""" + cam_renderers: dict[str, _CameraRendererState] = {} + for cfg in self._camera_configs: cam_id = mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_CAMERA, cfg.name) if cam_id < 0: logger.warning("Camera not found in MJCF, skipping", camera_name=cfg.name) @@ -258,17 +240,10 @@ def _init_new_cameras(self, cam_renderers: dict[str, _CameraRendererState]) -> N depth_renderer=depth_renderer, interval=interval, ) - logger.info( - "Camera renderer created", - camera_name=cfg.name, - width=cfg.width, - height=cfg.height, - fps=cfg.fps, - ) + return cam_renderers def _render_cameras(self, now: float, cam_renderers: dict[str, _CameraRendererState]) -> None: """Render all due cameras and store frames. Must be called from sim thread.""" - self._init_new_cameras(cam_renderers) for state in cam_renderers.values(): if now - state.last_render_time < state.interval: continue @@ -301,9 +276,8 @@ def _sim_loop(self) -> None: logger.info("sim loop started", cls=self.__class__.__name__) dt = 1.0 / self._control_frequency - # Camera renderers — created in sim thread (MuJoCo thread-safety). - # Checked each tick so cameras added after connect() are picked up. - cam_renderers: dict[str, _CameraRendererState] = {} + # Camera renderers: created once in the sim thread + cam_renderers = self._init_cameras() def _step_once(sync_viewer: bool) -> None: loop_start = time.time() diff --git a/dimos/simulation/engines/mujoco_shm.py b/dimos/simulation/engines/mujoco_shm.py index 64fb5fdbc2..d46739c1e8 100644 --- a/dimos/simulation/engines/mujoco_shm.py +++ b/dimos/simulation/engines/mujoco_shm.py @@ -42,7 +42,6 @@ # Upper bound on joint count per sim. Arms + gripper are typically <= 10. MAX_JOINTS = 16 - _FLOAT_BYTES = 8 # float64 _INT32_BYTES = 4 @@ -125,14 +124,7 @@ class ManipShmSet: @classmethod def create(cls, key: str) -> ManipShmSet: - """Create new SHM buffers with deterministic names derived from *key*. - - If stale buffers from a prior run exist, unlink them first. The stale - probe uses ``_unregister`` because that handle is transient — we never - owned it and must not let the tracker think we did. The create-side - handle, by contrast, *must* stay registered with ``resource_tracker`` - so that ``shm.unlink()`` in ``cleanup()`` matches up cleanly. - """ + """Create new SHM buffers with deterministic names derived from *key*""" buffers: dict[str, SharedMemory] = {} for buffer_name, size in _shm_sizes.items(): name = _buffer_name(key, buffer_name) @@ -164,7 +156,6 @@ def as_list(self) -> list[SharedMemory]: class ManipShmWriter: """Sim-side handle: writes joint state, reads command targets. - Owned by ``MujocoSimModule``. Creates the SHM buffers on init and unlinks them on cleanup. """ @@ -180,8 +171,6 @@ def __init__(self, key: str) -> None: for buf in self.shm.as_list(): np.ndarray((buf.size,), dtype=np.uint8, buffer=buf.buf)[:] = 0 - # ---------------- joint state (sim -> adapter) ---------------- - def write_joint_state( self, positions: list[float], diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 48c57f5051..beb4afc769 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -74,13 +74,8 @@ def _default_identity_transform() -> Transform: class MujocoSimModuleConfig(ModuleConfig, DepthCameraConfig): """Configuration for the unified MuJoCo simulation module.""" - # MJCF path — also used as the SHM key so the adapter can discover the - # shared memory region without RPC. address: str = "" headless: bool = False - # Number of arm joints (excluding gripper). If the engine exposes more - # joints than ``dof``, the extra joint at index ``dof`` is treated as a - # gripper and its commands get scaled from joint-space into actuator ctrl. dof: int = 7 # Camera config (matches former MujocoCameraConfig). @@ -129,8 +124,6 @@ def __init__(self, **kwargs: Any) -> None: self._publish_thread: threading.Thread | None = None self._camera_info_base: CameraInfo | None = None - # ---------------- frame id helpers --------------------------- - @property def _camera_link(self) -> str: return f"{self.config.camera_name}_link" @@ -151,12 +144,8 @@ def _depth_frame(self) -> str: def _depth_optical_frame(self) -> str: return f"{self.config.camera_name}_depth_optical_frame" - # ---------------- RPC spec methods --------------------------- - @rpc def get_color_camera_info(self) -> CameraInfo | None: - # Return a fresh copy — callers must not be able to mutate our - # cached intrinsics, and the color/depth getters must not alias. if self._camera_info_base is None: return None return self._camera_info_base.with_ts(time.time()) @@ -171,8 +160,6 @@ def get_depth_camera_info(self) -> CameraInfo | None: def get_depth_scale(self) -> float: return 1.0 - # ---------------- lifecycle ---------------------------------- - @rpc def start(self) -> None: if not self.config.address: @@ -220,16 +207,11 @@ def start(self) -> None: if not self._engine.connect(): raise RuntimeError("MujocoSimModule: engine.connect() failed") - # Publish num_joints + ready so the adapter can proceed past its - # ready-wait in connect(). Use the engine's full joint count (arm + - # gripper if present) so the adapter can detect the gripper via - # ``num_joints > dof``. self._shm.signal_ready(num_joints=len(joint_names)) # Camera intrinsics. self._build_camera_info() - # Camera publish thread (same pattern as former MujocoCamera). self._stop_event.clear() self._publish_thread = threading.Thread( target=self._publish_loop, daemon=True, name="MujocoSimPublish" @@ -270,11 +252,6 @@ def stop(self) -> None: self._publish_thread.join(timeout=2.0) self._publish_thread = None - # Collect teardown failures across both halves (engine + SHM) so that - # a failure in the first doesn't skip the second. Only clear the - # corresponding ``self._*`` reference if teardown actually succeeded - # — silently nulling on failure would let the next ``start()`` run on - # top of a half-torn-down state and mask the real bug. errors: list[tuple[str, BaseException]] = [] if self._engine is not None: try: @@ -296,8 +273,6 @@ def stop(self) -> None: super().stop() if errors: - # Surface the first failure so the operator sees it. Subsequent - # ones are already logged above. op, exc = errors[0] raise RuntimeError(f"MujocoSimModule.stop() failed during {op}: {exc}") from exc @@ -318,8 +293,6 @@ def _apply_shm_commands(self, engine: MujocoEngine) -> None: if vel_cmd is not None: engine.write_joint_command(JointState(velocity=vel_cmd.tolist())) - # Gripper: adapter writes raw joint-space position; sim module scales - # into actuator ctrl range before setting the target. if self._gripper_idx is not None: gripper_cmd = shm.read_gripper_command() if gripper_cmd is not None: @@ -342,13 +315,7 @@ def _publish_shm_state(self, engine: MujocoEngine) -> None: shm.write_gripper_state(positions[self._gripper_idx]) def _gripper_joint_to_ctrl(self, joint_position: float) -> float: - """Map joint-space gripper position to actuator control value. - - The high-level API uses ``position ∈ [jlo, jhi]`` where ``jlo`` is - closed and ``jhi`` is open (e.g. xArm: 0.0 closed, 0.85 open). The - MJCF actuator ``ctrl`` range is inverted on xArm-style grippers - (0 open, 255 closed). We map ``jlo → chi`` and ``jhi → clo``. - """ + """Map joint-space gripper position to actuator control value.""" jlo, jhi = self._gripper_joint_range clo, chi = self._gripper_ctrl_range clamped = max(jlo, min(jhi, joint_position)) @@ -382,21 +349,7 @@ def _build_camera_info(self) -> None: ) def _publish_loop(self) -> None: - """Poll engine for rendered frames and publish at configured FPS. - - Threading note: this loop reads ``self._engine`` without a lock. The - safety invariant is enforced by ``stop()``, which sets - ``self._stop_event`` *before* it joins this thread and only clears - ``self._engine`` *after* the join returns. So while this loop is - live, either ``_stop_event`` is unset (and ``_engine`` is guaranteed - non-None by construction in ``start()``), or ``_stop_event`` is set - and the next iteration will exit. The ``self._engine is not None`` - guards below are belt-and-suspenders, not the real correctness - argument — don't rely on them if you reorder ``stop()``. - """ - # Capture engine locally once up front. This makes the invariant - # explicit to the reader and removes the "is this racing with stop?" - # question from every subsequent access. + """Poll engine for rendered frames and publish at configured FPS.""" engine = self._engine if engine is None: return @@ -417,10 +370,6 @@ def _publish_loop(self) -> None: return while not self._stop_event.is_set(): - # Only MuJoCo's rendering path (``read_camera``) is known to raise - # transient RuntimeErrors — everything else in this loop is our - # own code and a raise there is a real bug we want to see. We - # catch narrowly around the render call, not the whole body. try: frame = engine.read_camera(self.config.camera_name) except RuntimeError as exc: From fceeb35f7b88b6c0f0c78c37a80716ccccd9faaa Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 6 Apr 2026 13:08:49 -0700 Subject: [PATCH 38/41] fix: change to match new manip blueprint style --- dimos/manipulation/blueprints.py | 41 +++++++++++++------------------- 1 file changed, 16 insertions(+), 25 deletions(-) diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index a4bb2a8e43..3c7f23f768 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -264,28 +264,26 @@ # Sim perception: MujocoSimModule owns the MujocoEngine and publishes both # camera streams and joint state via shared memory. -# The sim_xarm7 hardware component resolves to ShmMujocoAdapter, which -# attaches to the same SHM by MJCF path +# ShmMujocoAdapter attaches to the same SHM buffers by MJCF path. -from dimos.control.blueprints._hardware import XARM7_SIM_PATH, sim_xarm7 -from dimos.control.coordinator import TaskConfig as TaskConfig +from dimos.robot.catalog.ufactory import XARM7_SIM_PATH from dimos.simulation.engines.mujoco_sim_module import MujocoSimModule from dimos.visualization.rerun.bridge import RerunBridgeModule, _resolve_viewer_mode +_xarm7_sim_cfg = _catalog_xarm7( + name="arm", + adapter_type="sim_mujoco", + address=str(XARM7_SIM_PATH), + add_gripper=True, + pitch=math.radians(45), + tf_extra_links=["link7"], + home_joints=[0.0, 0.0, 0.0, 0.0, 0.0, -0.7, 0.0], + pre_grasp_offset=0.05, +) + xarm_perception_sim = autoconnect( PickAndPlaceModule.blueprint( - robots=[ - _make_xarm7_config( - "arm", - joint_prefix="arm_", - coordinator_task="traj_arm", - add_gripper=True, - gripper_hardware_id="arm", - tf_extra_links=["link7"], - pre_grasp_offset=0.05, - home_joints=[0.0, 0.0, 0.0, 0.0, 0.0, -0.7, 0.0], - ), - ], + robots=[_xarm7_sim_cfg.to_robot_model_config()], planning_timeout=10.0, enable_viz=True, ), @@ -301,15 +299,8 @@ tick_rate=100.0, publish_joint_state=True, joint_state_frame_id="coordinator", - hardware=[sim_xarm7("arm", headless=False, gripper=True)], - tasks=[ - TaskConfig( - name="traj_arm", - type="trajectory", - joint_names=[f"arm_joint{i + 1}" for i in range(7)], - priority=10, - ), - ], + hardware=[_xarm7_sim_cfg.to_hardware_component()], + tasks=[_xarm7_sim_cfg.to_task_config()], ), RerunBridgeModule.blueprint(viewer_mode=_resolve_viewer_mode()), ).transports( From de6d8ee70d2a0646f396a5671d79f1e14061940f Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 6 Apr 2026 13:20:12 -0700 Subject: [PATCH 39/41] tests: rename file --- .../manipulators/sim/{test_adapter.py => test_shm_adapter.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename dimos/hardware/manipulators/sim/{test_adapter.py => test_shm_adapter.py} (100%) diff --git a/dimos/hardware/manipulators/sim/test_adapter.py b/dimos/hardware/manipulators/sim/test_shm_adapter.py similarity index 100% rename from dimos/hardware/manipulators/sim/test_adapter.py rename to dimos/hardware/manipulators/sim/test_shm_adapter.py From 63e4300ea187efa5a60b45c6ebe7db92fd6cba37 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 6 Apr 2026 13:23:41 -0700 Subject: [PATCH 40/41] fix: mypy issues --- dimos/simulation/engines/mujoco_sim_module.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index beb4afc769..1f2d1df020 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -273,8 +273,8 @@ def stop(self) -> None: super().stop() if errors: - op, exc = errors[0] - raise RuntimeError(f"MujocoSimModule.stop() failed during {op}: {exc}") from exc + op, err = errors[0] + raise RuntimeError(f"MujocoSimModule.stop() failed during {op}: {err}") from err # ---------------- SHM <-> engine hooks ----------------------- From 6cef3f422285a90f7f493bae05c2b9d4fbd37189 Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Mon, 6 Apr 2026 17:23:19 -0700 Subject: [PATCH 41/41] fix: fix test fails --- dimos/hardware/manipulators/sim/adapter.py | 4 ---- dimos/simulation/engines/mujoco_shm.py | 14 -------------- dimos/simulation/engines/mujoco_sim_module.py | 4 ---- 3 files changed, 22 deletions(-) diff --git a/dimos/hardware/manipulators/sim/adapter.py b/dimos/hardware/manipulators/sim/adapter.py index b1ab342218..581dfb43a9 100644 --- a/dimos/hardware/manipulators/sim/adapter.py +++ b/dimos/hardware/manipulators/sim/adapter.py @@ -214,16 +214,12 @@ def write_clear_errors(self) -> bool: self._error_message = "" return True - # ---------------- cartesian (not supported yet) ------------------ - def read_cartesian_position(self) -> dict[str, float] | None: return None def write_cartesian_position(self, pose: dict[str, float], velocity: float = 1.0) -> bool: return False - # ---------------- gripper ------------------------------------ - def read_gripper_position(self) -> float | None: if not self._has_gripper or self._shm is None: return None diff --git a/dimos/simulation/engines/mujoco_shm.py b/dimos/simulation/engines/mujoco_shm.py index d46739c1e8..1679a00e42 100644 --- a/dimos/simulation/engines/mujoco_shm.py +++ b/dimos/simulation/engines/mujoco_shm.py @@ -193,8 +193,6 @@ def write_gripper_state(self, position: float) -> None: arr[0] = position self._increment_seq(SEQ_GRIPPER_STATE) - # ---------------- commands (adapter -> sim) ------------------- - def read_position_command(self, num_joints: int) -> NDArray[np.float64] | None: """Return a copy of position targets if a new command arrived since last call.""" seq = self._get_seq(SEQ_POSITION_CMD) @@ -223,8 +221,6 @@ def read_gripper_command(self) -> float | None: def read_command_mode(self) -> int: return int(self._control()[CTRL_COMMAND_MODE]) - # ---------------- control / lifecycle ------------------------- - def signal_ready(self, num_joints: int) -> None: ctrl = self._control() ctrl[CTRL_NUM_JOINTS] = num_joints @@ -251,8 +247,6 @@ def cleanup(self) -> None: except OSError as exc: logger.warning("SHM unlink failed", name=shm.name, error=str(exc)) - # ---------------- helpers ------------------------------------ - def _array(self, buf: SharedMemory, n: int, dtype: Any) -> NDArray[Any]: return np.ndarray((n,), dtype=dtype, buffer=buf.buf) @@ -280,8 +274,6 @@ class ManipShmReader: def __init__(self, key: str) -> None: self.shm = ManipShmSet.attach(key) - # ---------------- joint state (sim -> adapter) ---------------- - def read_positions(self, num_joints: int) -> list[float]: arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.positions.buf) return [float(x) for x in arr[:num_joints]] @@ -298,8 +290,6 @@ def read_gripper_position(self) -> float: arr = np.ndarray((2,), dtype=np.float64, buffer=self.shm.gripper.buf) return float(arr[0]) - # ---------------- commands (adapter -> sim) ------------------- - def write_position_command(self, positions: list[float]) -> None: n = min(len(positions), MAX_JOINTS) arr = np.ndarray((MAX_JOINTS,), dtype=np.float64, buffer=self.shm.position_targets.buf) @@ -319,8 +309,6 @@ def write_gripper_command(self, position: float) -> None: arr[1] = position self._increment_seq(SEQ_GRIPPER_CMD) - # ---------------- control / lifecycle ------------------------- - def is_ready(self) -> bool: return bool(self._control()[CTRL_READY] == 1) @@ -339,8 +327,6 @@ def cleanup(self) -> None: except OSError as exc: logger.warning("SHM close failed", name=shm.name, error=str(exc)) - # ---------------- helpers ------------------------------------ - def _control(self) -> NDArray[np.int32]: return np.ndarray((4,), dtype=np.int32, buffer=self.shm.control.buf) diff --git a/dimos/simulation/engines/mujoco_sim_module.py b/dimos/simulation/engines/mujoco_sim_module.py index 1f2d1df020..9250a2c6ac 100644 --- a/dimos/simulation/engines/mujoco_sim_module.py +++ b/dimos/simulation/engines/mujoco_sim_module.py @@ -276,8 +276,6 @@ def stop(self) -> None: op, err = errors[0] raise RuntimeError(f"MujocoSimModule.stop() failed during {op}: {err}") from err - # ---------------- SHM <-> engine hooks ----------------------- - def _apply_shm_commands(self, engine: MujocoEngine) -> None: """Pre-step hook: pull command targets from SHM into the engine.""" shm = self._shm @@ -324,8 +322,6 @@ def _gripper_joint_to_ctrl(self, joint_position: float) -> float: t = (clamped - jlo) / (jhi - jlo) return chi - t * (chi - clo) - # ---------------- camera publishing -------------------------- - def _build_camera_info(self) -> None: if self._engine is None: return