DIM-1842 drone simulator using MuJoCo simulator and Crazyflie 2

dimos/robot/all_blueprints.py

@@ -52,6 +52,8 @@
     "demo-skill": "dimos.agents.skills.demo_skill:demo_skill",
     "drone-agentic": "dimos.robot.drone.blueprints.agentic.drone_agentic:drone_agentic",
     "drone-basic": "dimos.robot.drone.blueprints.basic.drone_basic:drone_basic",
+    "drone-basic-sim": "dimos.robot.drone.blueprints.basic.drone_basic_sim:drone_basic_sim",
+    "drone-primitive-no-nav": "dimos.robot.drone.blueprints.primitive.drone_primitive_no_nav:drone_primitive_no_nav",
     "dual-xarm6-planner": "dimos.manipulation.blueprints:dual_xarm6_planner",
     "keyboard-teleop-openarm": "dimos.robot.manipulators.openarm.blueprints:keyboard_teleop_openarm",
     "keyboard-teleop-openarm-mock": "dimos.robot.manipulators.openarm.blueprints:keyboard_teleop_openarm_mock",
@@ -128,6 +130,7 @@
     "detection3-d-module": "dimos.perception.detection.module3D.Detection3DModule",
     "drone-camera-module": "dimos.robot.drone.camera_module.DroneCameraModule",
     "drone-connection-module": "dimos.robot.drone.connection_module.DroneConnectionModule",
+    "drone-sim-connection": "dimos.robot.drone.sim.DroneSimConnection",
     "drone-tracking-module": "dimos.robot.drone.drone_tracking_module.DroneTrackingModule",
     "embedding-memory": "dimos.memory.embedding.EmbeddingMemory",
     "emitter-module": "dimos.utils.demo_image_encoding.EmitterModule",

dimos/robot/drone/blueprints/basic/drone_basic_sim.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+
+""" Basic Drone Sim Blueprint with sim connection and visualization"""
+
+from typing import TYPE_CHECKING
+
+from dimos.core.coordination.blueprints import autoconnect
+from dimos.core.global_config import global_config
+from dimos.navigation.replanning_a_star.module import ReplanningAStarPlanner
+from dimos.robot.drone.blueprints.primitive.drone_primitive_no_nav import (
+    drone_primitive_no_nav,
+)
+from dimos.robot.drone.sim import DroneSimConnection
+
+drone_basic_sim = autoconnect(
+    drone_primitive_no_nav,
+    DroneSimConnection.blueprint(),
+    ReplanningAStarPlanner.blueprint(),
+)
+
+__all__ = ["drone_basic_sim"]

dimos/robot/drone/blueprints/primitive/drone_primitive_no_nav.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+
+""" Minimal Drone Stack without navigation. An attempt to create a minimal stack"""
+
+from typing import Any
+
+from dimos_lcm.sensor_msgs import CameraInfo
+
+from dimos.core.coordination.blueprints import autoconnect
+from dimos.core.global_config import global_config
+from dimos.core.transport import LCMTransport
+# from dimos.hardware.sensors.camera.module import CameraModule
+# from dimos.hardware.sensors.camera.webcam import Webcam
+# from dimos.hardware.sensors.camera.zed import compat as zed
+from dimos.mapping.costmapper import CostMapper
+from dimos.mapping.voxels import VoxelGridMapper
+from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped
+# from dimos.msgs.geometry_msgs.Quaternion import Quaternion
+# from dimos.msgs.geometry_msgs.Transform import Transform
+from dimos.msgs.geometry_msgs.Twist import Twist
+# from dimos.msgs.geometry_msgs.Vector3 import Vector3
+# from dimos.msgs.nav_msgs.Odometry import Odometry
+# from dimos.msgs.nav_msgs.Path import Path
+from dimos.msgs.sensor_msgs.Image import Image
+from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2
+# from dimos.msgs.std_msgs.Bool import Bool
+# from dimos.navigation.frontier_exploration.wavefront_frontier_goal_selector import (
+#     WavefrontFrontierExplorer,
+# )
+from dimos.protocol.pubsub.impl.lcmpubsub import LCM
+from dimos.visualization.vis_module import vis_module
+from dimos.web.websocket_vis.websocket_vis_module import WebsocketVisModule
+
+def _convert_camera_info(camera_info: Any) -> Any:
+    return camera_info.to_rerun(
+        image_topic="/world/color_image",
+        optical_frame="camera_optical",
+    )
+
+def _convert_navigation_costmap(grid: Any) -> Any:
+    return grid.to_rerun(
+        colormap="Accent",
+        z_offset=0.015,
+        opacity=0.2,
+        background="#484981",
+    )
+
+def _static_base_link(rr: Any) -> list[Any]:
+    return [
+        rr.Boxes3D(
+            half_sizes=[0.05, 0.05, 0.02],
+            colors=[(0, 255, 127)],
+            fill_mode="MajorWireframe",
+        ),
+        rr.Transform3D(parent_frame="tf#/base_link"),
+    ]
+
+def _drone_rerun_blueprint() -> Any:
+    """Same as the G1 rerun blueprint but with a different base link visualization."""
+    import rerun as rr
+    import rerun.blueprint as rrb
+
+    return rrb.Blueprint(
+        rrb.Horizontal(
+            rrb.Spatial2DView(origin="world/color_image", name="Camera"),
+            rrb.Spatial3DView(
+                origin="world",
+                name="3D",
+                background=rrb.Background(kind="SolidColor", color=[0, 0, 0]),
+                line_grid=rrb.LineGrid3D(
+                    plane=rr.components.Plane3D.XY.with_distance(0.5),
+                ),
+            ),
+            column_shares=[1, 2],
+        ),
+    )
+
+rerun_config = {
+    "blueprint": _drone_rerun_blueprint,
+    "visual_override": {
+        "world/camera_info": _convert_camera_info,
+        "world/navigation_costmap": _convert_navigation_costmap,
+    },
+    "static": {
+        "world/tf/base_link": _static_base_link,
+    }
+}
+
+_with_vis = vis_module(viewer_backend=global_config.viewer, rerun_config=rerun_config)
+
+drone_primitive_no_nav = (
+    autoconnect(
+        _with_vis,
+        VoxelGridMapper.blueprint(),
+        CostMapper.blueprint(),
+        #Visualization
+        WebsocketVisModule.blueprint(),
+    )
+    .global_config(n_workers=4, robot_model="drone", mujoco_room="empty")
+    .transports(
+        {   
+            # Movement command 
+            ("cmd_vel", Twist): LCMTransport("/cmd_vel", Twist),
+            # Odometry output from ROSNavigationModule
+            ("odom", PoseStamped): LCMTransport("/odom", PoseStamped),
+            ("pointcloud", PointCloud2): LCMTransport("/lidar", PointCloud2),
+            ("goal_pose", PoseStamped): LCMTransport("/goal_pose", PoseStamped),
+            # Camera topics
+            ("color_image", Image): LCMTransport("/color_image", Image),
+            ("camera_info", CameraInfo): LCMTransport("/camera_info", CameraInfo),
+        }
+    )
+)
+
+__all__ = ["drone_primitive_no_nav"]

dimos/robot/drone/sim.py

@@ -0,0 +1,102 @@
+import threading 
+from threading import Thread
+import time
+from typing import Any
+
+from pydantic import Field
+from reactivex.disposable import Disposable
+
+from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT
+from dimos.core.core import rpc
+from dimos.core.module import Module, ModuleConfig
+from dimos.core.stream import In, Out
+from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped
+from dimos.msgs.geometry_msgs.Quaternion import Quaternion
+from dimos.msgs.geometry_msgs.Transform import Transform
+from dimos.msgs.geometry_msgs.Twist import Twist
+from dimos.msgs.geometry_msgs.Vector3 import Vector3
+from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo
+from dimos.msgs.sensor_msgs.Image import Image
+from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2
+
+from dimos.robot.drone.connection_module import DroneConnectionModule
+from dimos.robot.unitree.mujoco_connection import MujocoConnection
+from dimos.robot.unitree.type.odometry import Odometry as SimOdometry
+from dimos.utils.logging_config import setup_logger
+
+logger = setup_logger()
+
+class DroneSimConfig(ModuleConfig):
+    ip: str = Field(default_factory=lambda m: m["g"].robot_ip)
+
+class DroneSimConnection(DroneConnectionModule):
+    config: DroneSimConfig
+    cmd_vel: In[Twist]
+    lidar: Out[PointCloud2]
+    odom: Out[PoseStamped]
+    color_image: Out[Image]
+    camera_info: Out[CameraInfo]
+    connection: MujocoConnection | None = None
+    _camera_info_thread: Thread | None = None
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._stop_event = threading.Event()
+
+    @rpc
+    def start(self) -> None:
+        Module.start(self)
+
+        from dimos.robot.unitree.mujoco_connection import MujocoConnection
+
+        self.connection = MujocoConnection(self.config.g)
+        assert self.connection is not None
+        self.connection.start()
+
+        self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move)))
+        self.register_disposable(self.connection.odom_stream().subscribe(self._publish_sim_odom))
+        self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish))
+        self.register_disposable(self.connection.video_stream().subscribe(self.color_image.publish))
+
+        self._camera_info_thread = Thread(
+            target=self._publish_camera_info_loop,
+            daemon=True,
+        )
+        self._camera_info_thread.start()
+
+    @rpc
+    def stop(self) -> None:
+        self._stop_event.set()
+        assert self.connection is not None
+        self.connection.stop()
+        if self._camera_info_thread and self._camera_info_thread.is_alive():
+            self._camera_info_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT)
+        Module.stop(self)
+
+    def _publish_camera_info_loop(self) -> None:
+        assert self.connection is not None
+        info = self.connection.camera_info_static
+        while not self._stop_event.is_set():
+            self.camera_info.publish(info)
+            self._stop_event.wait(1.0)
+
+    def _publish_sim_odom(self, msg: SimOdometry) -> None:
+        self._publish_tf(
+            PoseStamped(
+                ts=msg.ts,
+                frame_id=msg.frame_id,
+                position=msg.position,
+                orientation=msg.orientation,
+            )
+        )
+
+    @rpc
+    def move(self, twist: Twist, duration: float = 0.0) -> None:
+        assert self.connection is not None
+        self.connection.move(twist, duration)
+
+    @rpc
+    def publish_request(self, topic: str, data: dict[str, Any]) -> dict[Any, Any]:
+        logger.info(f"Publishing request to topic: {topic} with data: {data}")
+        assert self.connection is not None
+        return self.connection.publish_request(topic, data)

dimos/simulation/mujoco/model.py

@@ -27,7 +27,7 @@
 from dimos.mapping.occupancy.extrude_occupancy import generate_mujoco_scene
 from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid
 from dimos.simulation.mujoco.input_controller import InputController
-from dimos.simulation.mujoco.policy import G1OnnxController, Go1OnnxController, OnnxController
+from dimos.simulation.mujoco.policy import G1OnnxController, Go1OnnxController, OnnxController, DroneController
 from dimos.utils.data import get_data
 
 
@@ -52,6 +52,10 @@ def get_assets() -> dict[str, bytes]:
     mjx_env.update_assets(assets, person_dir, "*.obj")
     mjx_env.update_assets(assets, person_dir, "*.png")
 
+    # From: https://github.com/google-deepmind/mujoco_menagerie/tree/main/bitcraze_crazyflie_2
+    mjx_env.update_assets(assets, mjx_env.MENAGERIE_PATH / "bitcraze_crazyflie_2")
+    mjx_env.update_assets(assets, mjx_env.MENAGERIE_PATH / "bitcraze_crazyflie_2" / "assets")
+
     return assets
 
 
@@ -76,9 +80,14 @@ def load_model(
     n_substeps = round(ctrl_dt / sim_dt)
     model.opt.timestep = sim_dt
 
+    if robot == "cf2":
+        keyframe_name = "hover"
+    else: # For unitree_go1 and unitree_g1 -- default 
+        keyframe_name = "home" 
+
     params = {
         "policy_path": (_get_data_dir() / f"{robot}_policy.onnx").as_posix(),
-        "default_angles": np.array(model.keyframe("home").qpos[7:]),
+        "default_angles": np.array(model.keyframe(keyframe_name).qpos[7:]),
         "n_substeps": n_substeps,
         "action_scale": 0.5,
         "input_controller": input_device,
@@ -90,6 +99,8 @@ def load_model(
             policy: OnnxController = Go1OnnxController(**params)
         case "unitree_g1":
             policy = G1OnnxController(**params, drift_compensation=[-0.18, 0.0, -0.09])
+        case "cf2":
+            policy = DroneController(**params)
         case _:
             raise ValueError(f"Unknown robot policy: {robot}")
 

dimos/simulation/mujoco/mujoco_process.py

@@ -63,6 +63,8 @@ def get_command(self) -> NDArray[Any]:
             self._command[0] = linear[0]  # forward/backward
             self._command[1] = linear[1]  # left/right
             self._command[2] = angular[2]  # rotation
+        else:
+            self._command[:] = 0
         result: NDArray[Any] = self._command.copy()
         return result
 
@@ -75,6 +77,8 @@ def _run_simulation(config: GlobalConfig, shm: ShmReader) -> None:
     robot_name = config.robot_model or "unitree_go1"
     if robot_name == "unitree_go2":
         robot_name = "unitree_go1"
+    if robot_name == "drone":
+        robot_name = "cf2"
 
     controller = MockController(shm)
     model, data = load_model(controller, robot=robot_name, scene_xml=load_scene_xml(config))
@@ -87,6 +91,8 @@ def _run_simulation(config: GlobalConfig, shm: ShmReader) -> None:
             z = 0.3
         case "unitree_g1":
             z = 0.8
+        case "cf2":
+            z = 0.5
         case _:
             z = 0
 
@@ -154,14 +160,14 @@ def _run_simulation(config: GlobalConfig, shm: ShmReader) -> None:
             current_time = time.time()
 
             # Video rendering
-            if current_time - last_video_time >= video_interval:
+            if camera_id != -1 and current_time - last_video_time >= video_interval:
                 rgb_renderer.update_scene(data, camera=camera_id, scene_option=scene_option)
                 pixels = rgb_renderer.render()
                 shm.write_video(pixels)
                 last_video_time = current_time
 
             # Lidar/depth rendering
-            if current_time - last_lidar_time >= lidar_interval:
+            if lidar_camera_id != -1 and current_time - last_lidar_time >= lidar_interval:
                 # Render all depth cameras
                 depth_renderer.update_scene(data, camera=lidar_camera_id, scene_option=scene_option)
                 depth_front = depth_renderer.render()

dimos/simulation/mujoco/policy.py

@@ -154,3 +154,28 @@ def get_obs(self, model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray[Any,
     def _post_control_update(self) -> None:
         phase_tp1 = self._phase + self._phase_dt
         self._phase = np.fmod(phase_tp1 + np.pi, 2 * np.pi) - np.pi
+
+_DRONE_HOVER_THRUST = 0.26487 # From cf2.xml file
+_DRONE_MOMENT_SCALE = 0.005
+
+class DroneController():
+    def __init__(
+            self,
+            input_controller: InputController,
+            **kwargs: Any,
+    ) -> None:
+        self._input_controller = input_controller
+
+    def get_obs(self, model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray[Any, Any]:
+        command = self._input_controller.get_command()
+        return command.astype(np.float32)
+
+    def get_control(self, model: mujoco.MjModel, data: mujoco.MjData) -> None:
+        command = self._input_controller.get_command()
+        forward = float(command[0])
+        lateral = float(command[1])
+        yaw = float(command[2])
+        data.ctrl[0] = _DRONE_HOVER_THRUST
+        data.ctrl[1] = lateral * _DRONE_MOMENT_SCALE    # x_moment (roll)
+        data.ctrl[2] = -forward * _DRONE_MOMENT_SCALE   # y_moment (pitch)
+        data.ctrl[3] = -yaw * _DRONE_MOMENT_SCALE        # z_moment (yaw)