diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index b88159dd8e..f90ae66afb 100644 --- a/dimos/robot/all_blueprints.py +++ b/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", diff --git a/dimos/robot/drone/blueprints/basic/drone_basic_sim.py b/dimos/robot/drone/blueprints/basic/drone_basic_sim.py new file mode 100644 index 0000000000..3834e133b5 --- /dev/null +++ b/dimos/robot/drone/blueprints/basic/drone_basic_sim.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 + +"""Basic Drone Sim Blueprint with sim connection and visualization""" + +from dimos.core.coordination.blueprints import autoconnect +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"] diff --git a/dimos/robot/drone/blueprints/primitive/drone_primitive_no_nav.py b/dimos/robot/drone/blueprints/primitive/drone_primitive_no_nav.py new file mode 100644 index 0000000000..a5870f3981 --- /dev/null +++ b/dimos/robot/drone/blueprints/primitive/drone_primitive_no_nav.py @@ -0,0 +1,123 @@ +#!/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.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"] diff --git a/dimos/robot/drone/sim.py b/dimos/robot/drone/sim.py new file mode 100644 index 0000000000..0e65a1c255 --- /dev/null +++ b/dimos/robot/drone/sim.py @@ -0,0 +1,99 @@ +import threading +from threading import Thread +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.Twist import Twist +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) diff --git a/dimos/simulation/mujoco/model.py b/dimos/simulation/mujoco/model.py index bc309b7307..d8d14537a9 100644 --- a/dimos/simulation/mujoco/model.py +++ b/dimos/simulation/mujoco/model.py @@ -27,7 +27,12 @@ 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 ( + DroneController, + G1OnnxController, + Go1OnnxController, + OnnxController, +) from dimos.utils.data import get_data @@ -52,6 +57,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 +85,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 +104,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}") diff --git a/dimos/simulation/mujoco/mujoco_process.py b/dimos/simulation/mujoco/mujoco_process.py index 8ed4f20da2..6e705ecd00 100755 --- a/dimos/simulation/mujoco/mujoco_process.py +++ b/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() diff --git a/dimos/simulation/mujoco/policy.py b/dimos/simulation/mujoco/policy.py index 0a792baf1a..ac97091125 100644 --- a/dimos/simulation/mujoco/policy.py +++ b/dimos/simulation/mujoco/policy.py @@ -154,3 +154,30 @@ 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)