mujoco sim module

dimos/e2e_tests/test_simulation_module.py

@@ -0,0 +1,86 @@
+# 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.
+
+"""End-to-end tests for the simulation module."""
+
+import os
+
+import pytest
+
+from dimos.msgs.sensor_msgs import JointCommand, JointState, RobotState
+
+
+def _positions_within_tolerance(
+    positions: list[float],
+    target: list[float],
+    tolerance: float,
+) -> bool:
+    if len(positions) < len(target):
+        return False
+    return all(abs(positions[i] - target[i]) <= tolerance for i in range(len(target)))
+
+
+@pytest.mark.skipif(bool(os.getenv("CI")), reason="LCM doesn't work in CI.")
+@pytest.mark.e2e
+class TestSimulationModuleE2E:
+    def test_xarm7_joint_state_published(self, lcm_spy, start_blueprint) -> None:
+        joint_state_topic = "/xarm/joint_states#sensor_msgs.JointState"
+        lcm_spy.save_topic(joint_state_topic)
+
+        start_blueprint("simulation-xarm7")
+        lcm_spy.wait_for_saved_topic(joint_state_topic, timeout=15.0)
+
+        with lcm_spy._messages_lock:
+            raw_joint_state = lcm_spy.messages[joint_state_topic][0]
+
+        joint_state = JointState.lcm_decode(raw_joint_state)
+        assert len(joint_state.name) == 8
+        assert len(joint_state.position) == 8
+
+    def test_xarm7_robot_state_published(self, lcm_spy, start_blueprint) -> None:
+        robot_state_topic = "/xarm/robot_state#sensor_msgs.RobotState"
+        lcm_spy.save_topic(robot_state_topic)
+
+        start_blueprint("simulation-xarm7")
+        lcm_spy.wait_for_saved_topic(robot_state_topic, timeout=15.0)
+
+        with lcm_spy._messages_lock:
+            raw_robot_state = lcm_spy.messages[robot_state_topic][0]
+
+        robot_state = RobotState.lcm_decode(raw_robot_state)
+        assert robot_state.mt_able in (0, 1)
+
+    def test_xarm7_joint_command_updates_joint_state(self, lcm_spy, start_blueprint) -> None:
+        joint_state_topic = "/xarm/joint_states#sensor_msgs.JointState"
+        joint_command_topic = "/xarm/joint_position_command#sensor_msgs.JointCommand"
+        lcm_spy.save_topic(joint_state_topic)
+
+        start_blueprint("simulation-xarm7")
+        lcm_spy.wait_for_saved_topic(joint_state_topic, timeout=15.0)
+
+        target_positions = [0.2, -0.2, 0.1, -0.1, 0.15, -0.15, 0.05]
+        lcm_spy.publish(joint_command_topic, JointCommand(positions=target_positions))
+
+        tolerance = 0.03
+        lcm_spy.wait_for_message_result(
+            joint_state_topic,
+            JointState,
+            predicate=lambda msg: _positions_within_tolerance(
+                list(msg.position),
+                target_positions,
+                tolerance,
+            ),
+            fail_message=("joint_state did not reach commanded positions within tolerance"),
+            timeout=10.0,
+        )

dimos/robot/all_blueprints.py

@@ -64,6 +64,7 @@
     "xarm-perception": "dimos.manipulation.manipulation_blueprints:xarm_perception",
     "xarm6-planner-only": "dimos.manipulation.manipulation_blueprints:xarm6_planner_only",
     "xarm7-planner-coordinator": "dimos.manipulation.manipulation_blueprints:xarm7_planner_coordinator",
+    "xarm7-trajectory-sim": "dimos.simulation.sim_blueprints:xarm7_trajectory_sim",
 }
 
 
@@ -103,6 +104,7 @@
     "replanning_a_star_planner": "dimos.navigation.replanning_a_star.module",
     "rerun_bridge": "dimos.visualization.rerun.bridge",
     "ros_nav": "dimos.navigation.rosnav",
+    "simulation": "dimos.simulation.manipulators.sim_module",
     "spatial_memory": "dimos.perception.spatial_perception",
     "speak_skill": "dimos.agents.skills.speak_skill",
     "temporal_memory": "dimos.perception.experimental.temporal_memory.temporal_memory",

dimos/simulation/engines/__init__.py

@@ -0,0 +1,25 @@
+"""Simulation engines for manipulator backends."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from dimos.simulation.engines.base import SimulationEngine
+from dimos.simulation.engines.mujoco_engine import MujocoEngine
+
+EngineType = Literal["mujoco"]
+
+_ENGINES: dict[EngineType, type[SimulationEngine]] = {
+    "mujoco": MujocoEngine,
+}
+
+
+def get_engine(engine_name: EngineType) -> type[SimulationEngine]:
+    return _ENGINES[engine_name]
+
+
+__all__ = [
+    "EngineType",
+    "SimulationEngine",
+    "get_engine",
+]

dimos/simulation/engines/base.py

@@ -0,0 +1,84 @@
+# 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.
+
+"""Base interfaces for simulator engines."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from dimos.msgs.sensor_msgs import JointState
+
+
+class SimulationEngine(ABC):
+    """Abstract base class for a simulator engine instance."""
+
+    def __init__(self, config_path: Path, headless: bool) -> None:
+        self._config_path = config_path
+        self._headless = headless
+
+    @property
+    def config_path(self) -> Path:
+        return self._config_path
+
+    @property
+    def headless(self) -> bool:
+        return self._headless
+
+    @abstractmethod
+    def connect(self) -> bool:
+        """Connect to simulation and start the engine."""
+
+    @abstractmethod
+    def disconnect(self) -> bool:
+        """Disconnect from simulation and stop the engine."""
+
+    @property
+    @abstractmethod
+    def connected(self) -> bool:
+        """Whether the engine is connected."""
+
+    @property
+    @abstractmethod
+    def num_joints(self) -> int:
+        """Number of joints for the loaded robot."""
+
+    @property
+    @abstractmethod
+    def joint_names(self) -> list[str]:
+        """Joint names for the loaded robot."""
+
+    @abstractmethod
+    def read_joint_positions(self) -> list[float]:
+        """Read joint positions in radians."""
+
+    @abstractmethod
+    def read_joint_velocities(self) -> list[float]:
+        """Read joint velocities in rad/s."""
+
+    @abstractmethod
+    def read_joint_efforts(self) -> list[float]:
+        """Read joint efforts in Nm."""
+
+    @abstractmethod
+    def write_joint_command(self, command: JointState) -> None:
+        """Command joints using a JointState message."""
+
+    @abstractmethod
+    def hold_current_position(self) -> None:
+        """Hold current joint positions."""