From f74dcd93c1a04e5ea11e446416cf40e00c1a270d Mon Sep 17 00:00:00 2001 From: Ruthwik Date: Tue, 24 Mar 2026 11:41:42 -0700 Subject: [PATCH 01/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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, )