diff --git a/assets/policies/go1_policy.onnx b/assets/policies/go1_policy.onnx new file mode 100644 index 0000000000..af52536397 --- /dev/null +++ b/assets/policies/go1_policy.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:386cde6a1eac679e2f2a313ade2b395d9f26905b151ee3b019c2c4163d49b6f2 +size 772916 diff --git a/assets/robots/go1/assets/hfield.png b/assets/robots/go1/assets/hfield.png new file mode 100644 index 0000000000..62af27a2bc Binary files /dev/null and b/assets/robots/go1/assets/hfield.png differ diff --git a/assets/robots/go1/assets/rocky_texture.png b/assets/robots/go1/assets/rocky_texture.png new file mode 100644 index 0000000000..1456b3ff47 Binary files /dev/null and b/assets/robots/go1/assets/rocky_texture.png differ diff --git a/assets/robots/go1/robot.xml b/assets/robots/go1/robot.xml new file mode 100644 index 0000000000..2b1251cbe7 --- /dev/null +++ b/assets/robots/go1/robot.xml @@ -0,0 +1,306 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dimos/robot/unitree_webrtc/mujoco_connection.py b/dimos/robot/unitree_webrtc/mujoco_connection.py new file mode 100644 index 0000000000..277977a8a6 --- /dev/null +++ b/dimos/robot/unitree_webrtc/mujoco_connection.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 + +# Copyright 2025 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. + + +import atexit +import functools +import logging +import threading +import time +from typing import List + +from reactivex import Observable + +from dimos.msgs.geometry_msgs import Vector3 +from dimos.msgs.sensor_msgs import Image + +try: + from dimos.simulation.mujoco.mujoco import MujocoThread +except ImportError: + MujocoThread = None + +LIDAR_FREQUENCY = 10 +ODOM_FREQUENCY = 50 +VIDEO_FREQUENCY = 30 + +logger = logging.getLogger(__name__) + + +class MujocoConnection: + def __init__(self, *args, **kwargs): + if MujocoThread is None: + raise ImportError("'mujoco' is not installed. Use `pip install -e .[sim]`") + self.mujoco_thread = MujocoThread() + self._stream_threads: List[threading.Thread] = [] + self._stop_events: List[threading.Event] = [] + self._is_cleaned_up = False + + # Register cleanup on exit + atexit.register(self.cleanup) + + def start(self): + self.mujoco_thread.start() + + def standup(self): + print("standup supressed") + + def liedown(self): + print("liedown supressed") + + @functools.cache + def lidar_stream(self): + def on_subscribe(observer, scheduler): + if self._is_cleaned_up: + observer.on_completed() + return lambda: None + + stop_event = threading.Event() + self._stop_events.append(stop_event) + + def run(): + try: + while not stop_event.is_set() and not self._is_cleaned_up: + lidar_to_publish = self.mujoco_thread.get_lidar_message() + + if lidar_to_publish: + observer.on_next(lidar_to_publish) + + time.sleep(1 / LIDAR_FREQUENCY) + except Exception as e: + logger.error(f"Lidar stream error: {e}") + finally: + observer.on_completed() + + thread = threading.Thread(target=run, daemon=True) + self._stream_threads.append(thread) + thread.start() + + def dispose(): + stop_event.set() + + return dispose + + return Observable(on_subscribe) + + @functools.cache + def odom_stream(self): + print("odom stream start") + + def on_subscribe(observer, scheduler): + if self._is_cleaned_up: + observer.on_completed() + return lambda: None + + stop_event = threading.Event() + self._stop_events.append(stop_event) + + def run(): + try: + while not stop_event.is_set() and not self._is_cleaned_up: + odom_to_publish = self.mujoco_thread.get_odom_message() + if odom_to_publish: + observer.on_next(odom_to_publish) + + time.sleep(1 / ODOM_FREQUENCY) + except Exception as e: + logger.error(f"Odom stream error: {e}") + finally: + observer.on_completed() + + thread = threading.Thread(target=run, daemon=True) + self._stream_threads.append(thread) + thread.start() + + def dispose(): + stop_event.set() + + return dispose + + return Observable(on_subscribe) + + @functools.cache + def video_stream(self): + print("video stream start") + + def on_subscribe(observer, scheduler): + if self._is_cleaned_up: + observer.on_completed() + return lambda: None + + stop_event = threading.Event() + self._stop_events.append(stop_event) + + def run(): + try: + while not stop_event.is_set() and not self._is_cleaned_up: + with self.mujoco_thread.pixels_lock: + if self.mujoco_thread.shared_pixels is not None: + img = Image.from_numpy(self.mujoco_thread.shared_pixels.copy()) + observer.on_next(img) + time.sleep(1 / VIDEO_FREQUENCY) + except Exception as e: + logger.error(f"Video stream error: {e}") + finally: + observer.on_completed() + + thread = threading.Thread(target=run, daemon=True) + self._stream_threads.append(thread) + thread.start() + + def dispose(): + stop_event.set() + + return dispose + + return Observable(on_subscribe) + + def move(self, vector: Vector3, duration: float = 0.0): + if not self._is_cleaned_up: + self.mujoco_thread.move(vector, duration) + + def stop(self): + """Stop the MuJoCo connection gracefully.""" + self.cleanup() + + def cleanup(self): + """Clean up all resources. Can be called multiple times safely.""" + if self._is_cleaned_up: + return + + logger.debug("Cleaning up MuJoCo connection resources") + self._is_cleaned_up = True + + # Stop all stream threads + for stop_event in self._stop_events: + stop_event.set() + + # Wait for threads to finish + for thread in self._stream_threads: + if thread.is_alive(): + thread.join(timeout=2.0) + if thread.is_alive(): + logger.warning(f"Stream thread {thread.name} did not stop gracefully") + + # Clean up the MuJoCo thread + if hasattr(self, "mujoco_thread") and self.mujoco_thread: + self.mujoco_thread.cleanup() + + # Clear references + self._stream_threads.clear() + self._stop_events.clear() + + # Clear cached methods to prevent memory leaks + if hasattr(self, "lidar_stream"): + self.lidar_stream.cache_clear() + if hasattr(self, "odom_stream"): + self.odom_stream.cache_clear() + if hasattr(self, "video_stream"): + self.video_stream.cache_clear() + + def __del__(self): + """Destructor to ensure cleanup on object deletion.""" + try: + self.cleanup() + except Exception: + pass diff --git a/dimos/robot/unitree_webrtc/type/map.py b/dimos/robot/unitree_webrtc/type/map.py index 78d26427e8..c0a0338de7 100644 --- a/dimos/robot/unitree_webrtc/type/map.py +++ b/dimos/robot/unitree_webrtc/type/map.py @@ -41,11 +41,15 @@ def __init__( voxel_size: float = 0.05, cost_resolution: float = 0.05, global_publish_interval: Optional[float] = None, + min_height: float = 0.15, + max_height: float = 0.6, **kwargs, ): self.voxel_size = voxel_size self.cost_resolution = cost_resolution self.global_publish_interval = global_publish_interval + self.min_height = min_height + self.max_height = max_height super().__init__(**kwargs) @rpc @@ -60,8 +64,8 @@ def publish(_): occupancygrid = OccupancyGrid.from_pointcloud( self.to_lidar_message(), resolution=self.cost_resolution, - min_height=0.15, - max_height=0.6, + min_height=self.min_height, + max_height=self.max_height, ) self.global_costmap.publish(occupancygrid) diff --git a/dimos/robot/unitree_webrtc/unitree_go2.py b/dimos/robot/unitree_webrtc/unitree_go2.py index 6296b38b5c..8e8460bd86 100644 --- a/dimos/robot/unitree_webrtc/unitree_go2.py +++ b/dimos/robot/unitree_webrtc/unitree_go2.py @@ -124,14 +124,14 @@ class ConnectionModule(Module): lidar: Out[LidarMessage] = None video: Out[Image] = None ip: str - playback: bool + connection_type: str = "webrtc" _odom: PoseStamped = None _lidar: LidarMessage = None - def __init__(self, ip: str = None, playback: bool = False, *args, **kwargs): + def __init__(self, ip: str = None, connection_type: str = "webrtc", *args, **kwargs): self.ip = ip - self.playback = playback + self.connection_type = connection_type self.tf = TF() self.connection = None Module.__init__(self, *args, **kwargs) @@ -139,10 +139,18 @@ def __init__(self, ip: str = None, playback: bool = False, *args, **kwargs): @rpc def start(self): """Start the connection and subscribe to sensor streams.""" - if self.playback: - self.connection = FakeRTC(self.ip) - else: - self.connection = UnitreeWebRTCConnection(self.ip) + match self.connection_type: + case "webrtc": + self.connection = UnitreeWebRTCConnection(self.ip) + case "fake": + self.connection = FakeRTC(self.ip) + case "mujoco": + from dimos.robot.unitree_webrtc.mujoco_connection import MujocoConnection + + self.connection = MujocoConnection() + self.connection.start() + case _: + raise ValueError(f"Unknown connection type: {self.connection_type}") # Connect sensor streams to outputs self.connection.lidar_stream().subscribe(self.lidar.publish) @@ -208,7 +216,7 @@ def __init__( output_dir: str = None, websocket_port: int = 7779, skill_library: Optional[SkillLibrary] = None, - playback: bool = False, + connection_type: Optional[str] = "webrtc", ): """Initialize the robot system. @@ -221,7 +229,9 @@ def __init__( """ super().__init__() self.ip = ip - self.playback = playback or (ip is None) # Auto-enable playback if no IP provided + self.connection_type = connection_type or "webrtc" + if ip is None and self.connection_type == "webrtc": + self.connection_type = "fake" # Auto-enable playback if no IP provided self.output_dir = output_dir or os.path.join(os.getcwd(), "assets", "output") self.websocket_port = websocket_port self.lcm = LCM() @@ -291,7 +301,9 @@ def start(self): def _deploy_connection(self): """Deploy and configure the connection module.""" - self.connection = self.dimos.deploy(ConnectionModule, self.ip, playback=self.playback) + self.connection = self.dimos.deploy( + ConnectionModule, self.ip, connection_type=self.connection_type + ) self.connection.lidar.transport = core.LCMTransport("/lidar", LidarMessage) self.connection.odom.transport = core.LCMTransport("/odom", PoseStamped) @@ -300,7 +312,10 @@ def _deploy_connection(self): def _deploy_mapping(self): """Deploy and configure the mapping module.""" - self.mapper = self.dimos.deploy(Map, voxel_size=0.5, global_publish_interval=2.5) + min_height = 0.3 if self.connection_type == "mujoco" else 0.15 + self.mapper = self.dimos.deploy( + Map, voxel_size=0.5, global_publish_interval=2.5, min_height=min_height + ) self.mapper.global_map.transport = core.LCMTransport("/global_map", LidarMessage) self.mapper.global_costmap.transport = core.LCMTransport("/global_costmap", OccupancyGrid) @@ -597,10 +612,11 @@ def navigate_to_object(self, bbox: List[float], distance: float = 0.5, timeout: def main(): """Main entry point.""" ip = os.getenv("ROBOT_IP") + connection_type = os.getenv("CONNECTION_TYPE", "webrtc") pubsub.lcm.autoconf() - robot = UnitreeGo2(ip=ip, websocket_port=7779, playback=False) + robot = UnitreeGo2(ip=ip, websocket_port=7779, connection_type=connection_type) robot.start() try: diff --git a/dimos/simulation/mujoco/depth_camera.py b/dimos/simulation/mujoco/depth_camera.py new file mode 100644 index 0000000000..2920e9a140 --- /dev/null +++ b/dimos/simulation/mujoco/depth_camera.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 + +# Copyright 2025 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. + +import math +import numpy as np +import open3d as o3d + +RANGE_FINDER_MAX_RANGE = 4 + + +def depth_image_to_point_cloud( + depth_image: np.ndarray, + camera_pos: np.ndarray, + camera_mat: np.ndarray, + fov_degrees: float = 120, + min_range: float = 0.0, +) -> np.ndarray: + """ + Convert a depth image from a camera to a 3D point cloud using perspective projection. + + Args: + depth_image: 2D numpy array of depth values in meters + camera_pos: 3D position of camera in world coordinates + camera_mat: 3x3 camera rotation matrix in world coordinates + fov_degrees: Vertical field of view of the camera in degrees + min_range: Minimum distance from camera to include points (meters) + + Returns: + numpy array of 3D points in world coordinates, shape (N, 3) + """ + height, width = depth_image.shape + + # Calculate camera intrinsics similar to StackOverflow approach + fovy = math.radians(fov_degrees) + f = height / (2 * math.tan(fovy / 2)) # focal length in pixels + cx = width / 2 # principal point x + cy = height / 2 # principal point y + + # Create Open3D camera intrinsics + cam_intrinsics = o3d.camera.PinholeCameraIntrinsic(width, height, f, f, cx, cy) + + # Convert numpy depth array to Open3D Image + o3d_depth = o3d.geometry.Image(depth_image.astype(np.float32)) + + # Create point cloud from depth image using Open3D + o3d_cloud = o3d.geometry.PointCloud.create_from_depth_image(o3d_depth, cam_intrinsics) + + # Convert Open3D point cloud to numpy array + camera_points = np.asarray(o3d_cloud.points) + + if camera_points.size == 0: + return np.array([]).reshape(0, 3) + + # Flip y and z axes + camera_points[:, 1] = -camera_points[:, 1] + camera_points[:, 2] = -camera_points[:, 2] + + # Filter points based on depth range (note: z is now negative) + valid_mask = (camera_points[:, 2] < -min_range) & ( + camera_points[:, 2] > -RANGE_FINDER_MAX_RANGE + ) + camera_points = camera_points[valid_mask] + + if camera_points.size == 0: + return np.array([]).reshape(0, 3) + + # Transform to world coordinates + world_points = (camera_mat @ camera_points.T).T + camera_pos + + return world_points diff --git a/dimos/simulation/mujoco/model.py b/dimos/simulation/mujoco/model.py new file mode 100644 index 0000000000..465adbcea7 --- /dev/null +++ b/dimos/simulation/mujoco/model.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 + +# Copyright 2025 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. + + +import mujoco +import numpy as np +from etils import epath +from mujoco_playground._src import mjx_env + + +from dimos.simulation.mujoco.policy import OnnxController +from dimos.simulation.mujoco.types import InputController + +RANGE_FINDER_MAX_RANGE = 4 +RANGE_FINDER_MIN_RANGE = 0.2 +LIDAR_RESOLUTION = 0.05 +VIDEO_FREQUENCY = 30 +DEPTH_CAMERA_FOV = 160 + +_HERE = epath.Path(__file__).parent + + +def get_assets() -> dict[str, bytes]: + assets: dict[str, bytes] = {} + assets_path = _HERE / "../../../assets/robots/go1" + mjx_env.update_assets(assets, assets_path, "*.xml") + mjx_env.update_assets(assets, assets_path / "assets") + path = mjx_env.MENAGERIE_PATH / "unitree_go1" + mjx_env.update_assets(assets, path, "*.xml") + mjx_env.update_assets(assets, path / "assets") + return assets + + +def load_model(input_device: InputController, model=None, data=None): + mujoco.set_mjcb_control(None) + + model = mujoco.MjModel.from_xml_path( + (_HERE / "../../../assets/robots/go1/robot.xml").as_posix(), + assets=get_assets(), + ) + data = mujoco.MjData(model) + + mujoco.mj_resetDataKeyframe(model, data, 0) + + ctrl_dt = 0.02 + sim_dt = 0.004 + n_substeps = int(round(ctrl_dt / sim_dt)) + model.opt.timestep = sim_dt + + policy = OnnxController( + policy_path=(_HERE / "../../../assets/policies/go1_policy.onnx").as_posix(), + default_angles=np.array(model.keyframe("home").qpos[7:]), + n_substeps=n_substeps, + action_scale=0.5, + input_controller=input_device, + ) + + mujoco.set_mjcb_control(policy.get_control) + + return model, data diff --git a/dimos/simulation/mujoco/mujoco.py b/dimos/simulation/mujoco/mujoco.py new file mode 100644 index 0000000000..e7cc096983 --- /dev/null +++ b/dimos/simulation/mujoco/mujoco.py @@ -0,0 +1,376 @@ +#!/usr/bin/env python3 + +# Copyright 2025 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. + + +import atexit +import logging +import threading +import time + +import mujoco +import numpy as np +import open3d as o3d +from mujoco import viewer + + +from dimos.msgs.geometry_msgs import Quaternion, Vector3 +from dimos.robot.unitree_webrtc.type.lidar import LidarMessage +from dimos.robot.unitree_webrtc.type.odometry import Odometry +from dimos.simulation.mujoco.depth_camera import depth_image_to_point_cloud +from dimos.simulation.mujoco.model import load_model + +RANGE_FINDER_MAX_RANGE = 4 +RANGE_FINDER_MIN_RANGE = 0.2 +LIDAR_RESOLUTION = 0.05 +VIDEO_FREQUENCY = 30 +DEPTH_CAMERA_FOV = 160 +STEPS_PER_FRAME = 7 # Adjust this number to control physics speed + +logger = logging.getLogger(__name__) + + +class MujocoThread(threading.Thread): + def __init__(self): + super().__init__(daemon=True) + self.shared_pixels = None + self.pixels_lock = threading.RLock() + self.shared_depth_front = None + self.depth_lock_front = threading.RLock() + self.shared_depth_left = None + self.depth_left_lock = threading.RLock() + self.shared_depth_right = None + self.depth_right_lock = threading.RLock() + self.odom_data = None + self.odom_lock = threading.RLock() + self.lidar_lock = threading.RLock() + self.model = None + self.data = None + self._command = np.zeros(3, dtype=np.float32) + self._command_lock = threading.RLock() + self._is_running = True + self._stop_timer: threading.Timer | None = None + self._viewer = None + self._rgb_renderer = None + self._depth_renderer = None + self._depth_left_renderer = None + self._depth_right_renderer = None + self._cleanup_registered = False + + # Register cleanup on exit + atexit.register(self.cleanup) + + def run(self): + try: + self.run_simulation() + except Exception as e: + logger.error(f"MuJoCo simulation thread error: {e}") + finally: + self._cleanup_resources() + + def run_simulation(self): + self.model, self.data = load_model(self) + + camera_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, "head_camera") + lidar_camera_id = mujoco.mj_name2id( + self.model, mujoco.mjtObj.mjOBJ_CAMERA, "lidar_front_camera" + ) + lidar_left_camera_id = mujoco.mj_name2id( + self.model, mujoco.mjtObj.mjOBJ_CAMERA, "lidar_left_camera" + ) + lidar_right_camera_id = mujoco.mj_name2id( + self.model, mujoco.mjtObj.mjOBJ_CAMERA, "lidar_right_camera" + ) + + with viewer.launch_passive(self.model, self.data) as m_viewer: + self._viewer = m_viewer + window_size = (320, 240) + + # Create separate renderers for RGB and depth + self._rgb_renderer = mujoco.Renderer( + self.model, height=window_size[1], width=window_size[0] + ) + self._depth_renderer = mujoco.Renderer( + self.model, height=window_size[1], width=window_size[0] + ) + # Enable depth rendering only for depth renderer + self._depth_renderer.enable_depth_rendering() + + # Create renderers for left and right depth cameras + self._depth_left_renderer = mujoco.Renderer( + self.model, height=window_size[1], width=window_size[0] + ) + self._depth_left_renderer.enable_depth_rendering() + + self._depth_right_renderer = mujoco.Renderer( + self.model, height=window_size[1], width=window_size[0] + ) + self._depth_right_renderer.enable_depth_rendering() + + scene_option = mujoco.MjvOption() + + while m_viewer.is_running() and self._is_running: + # Step multiple times per frame to speed up physics + for _ in range(STEPS_PER_FRAME): + mujoco.mj_step(self.model, self.data) + + with self.odom_lock: + # base position + pos = self.data.qpos[0:3] + # base orientation + quat = self.data.qpos[3:7] # (w, x, y, z) + self.odom_data = (pos.copy(), quat.copy()) + + # Render regular camera for video (RGB) + self._rgb_renderer.update_scene( + self.data, camera=camera_id, scene_option=scene_option + ) + pixels = self._rgb_renderer.render() + + with self.pixels_lock: + self.shared_pixels = pixels.copy() + + # Render fisheye camera for depth/lidar data + self._depth_renderer.update_scene( + self.data, camera=lidar_camera_id, scene_option=scene_option + ) + # When depth rendering is enabled, render() returns depth as float array in meters + depth = self._depth_renderer.render() + + with self.depth_lock_front: + self.shared_depth_front = depth.copy() + + # Render left depth camera + self._depth_left_renderer.update_scene( + self.data, camera=lidar_left_camera_id, scene_option=scene_option + ) + depth_left = self._depth_left_renderer.render() + + with self.depth_left_lock: + self.shared_depth_left = depth_left.copy() + + # Render right depth camera + self._depth_right_renderer.update_scene( + self.data, camera=lidar_right_camera_id, scene_option=scene_option + ) + depth_right = self._depth_right_renderer.render() + + with self.depth_right_lock: + self.shared_depth_right = depth_right.copy() + + m_viewer.sync() + + def _process_depth_camera(self, camera_name: str, depth_data, depth_lock) -> np.ndarray | None: + """Process a single depth camera and return point cloud points.""" + with depth_lock: + if depth_data is None: + return None + + depth_image = depth_data.copy() + camera_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, camera_name) + if camera_id == -1: + return None + + camera_pos = self.data.cam_xpos[camera_id] + camera_mat = self.data.cam_xmat[camera_id].reshape(3, 3) + points = depth_image_to_point_cloud( + depth_image, + camera_pos, + camera_mat, + fov_degrees=DEPTH_CAMERA_FOV, + min_range=RANGE_FINDER_MIN_RANGE, + ) + return points if points.size > 0 else None + + def get_lidar_message(self) -> LidarMessage | None: + all_points = [] + origin = None + + with self.lidar_lock: + if self.model is not None and self.data is not None: + pos = self.data.qpos[0:3] + origin = Vector3(pos[0], pos[1], pos[2]) + + cameras = [ + ("lidar_front_camera", self.shared_depth_front, self.depth_lock_front), + ("lidar_left_camera", self.shared_depth_left, self.depth_left_lock), + ("lidar_right_camera", self.shared_depth_right, self.depth_right_lock), + ] + + for camera_name, depth_data, depth_lock in cameras: + points = self._process_depth_camera(camera_name, depth_data, depth_lock) + if points is not None: + all_points.append(points) + + # Combine all point clouds + if not all_points: + return None + + combined_points = np.vstack(all_points) + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(combined_points) + + # Apply voxel downsampling to remove overlapping points + pcd = pcd.voxel_down_sample(voxel_size=LIDAR_RESOLUTION) + lidar_to_publish = LidarMessage( + pointcloud=pcd, + ts=time.time(), + origin=origin, + resolution=LIDAR_RESOLUTION, + ) + return lidar_to_publish + + def get_odom_message(self) -> Odometry | None: + with self.odom_lock: + if self.odom_data is None: + return None + pos, quat_wxyz = self.odom_data + + # MuJoCo uses (w, x, y, z) for quaternions. + # ROS and Dimos use (x, y, z, w). + orientation = Quaternion(quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]) + + odom_to_publish = Odometry( + position=Vector3(pos[0], pos[1], pos[2]), + orientation=orientation, + ts=time.time(), + frame_id="world", + ) + return odom_to_publish + + def _stop_move(self): + with self._command_lock: + self._command = np.zeros(3, dtype=np.float32) + self._stop_timer = None + + def move(self, vector: Vector3, duration: float = 0.0): + if self._stop_timer: + self._stop_timer.cancel() + + with self._command_lock: + self._command = np.array([vector.x, vector.y, vector.z], dtype=np.float32) + + if duration > 0: + self._stop_timer = threading.Timer(duration, self._stop_move) + self._stop_timer.daemon = True + self._stop_timer.start() + else: + self._stop_timer = None + + def get_command(self) -> np.ndarray: + with self._command_lock: + return self._command.copy() + + def stop(self): + """Stop the simulation thread gracefully.""" + self._is_running = False + + # Cancel any pending timers + if self._stop_timer: + self._stop_timer.cancel() + self._stop_timer = None + + # Wait for thread to finish + if self.is_alive(): + self.join(timeout=5.0) + if self.is_alive(): + logger.warning("MuJoCo thread did not stop gracefully within timeout") + + def cleanup(self): + """Clean up all resources. Can be called multiple times safely.""" + if self._cleanup_registered: + return + self._cleanup_registered = True + + logger.debug("Cleaning up MuJoCo resources") + self.stop() + self._cleanup_resources() + + def _cleanup_resources(self): + """Internal method to clean up MuJoCo-specific resources.""" + try: + # Cancel any timers + if self._stop_timer: + self._stop_timer.cancel() + self._stop_timer = None + + # Clean up renderers + if self._rgb_renderer is not None: + try: + self._rgb_renderer.close() + except Exception as e: + logger.debug(f"Error closing RGB renderer: {e}") + finally: + self._rgb_renderer = None + + if self._depth_renderer is not None: + try: + self._depth_renderer.close() + except Exception as e: + logger.debug(f"Error closing depth renderer: {e}") + finally: + self._depth_renderer = None + + if self._depth_left_renderer is not None: + try: + self._depth_left_renderer.close() + except Exception as e: + logger.debug(f"Error closing left depth renderer: {e}") + finally: + self._depth_left_renderer = None + + if self._depth_right_renderer is not None: + try: + self._depth_right_renderer.close() + except Exception as e: + logger.debug(f"Error closing right depth renderer: {e}") + finally: + self._depth_right_renderer = None + + # Clear data references + with self.pixels_lock: + self.shared_pixels = None + + with self.depth_lock_front: + self.shared_depth_front = None + + with self.depth_left_lock: + self.shared_depth_left = None + + with self.depth_right_lock: + self.shared_depth_right = None + + with self.odom_lock: + self.odom_data = None + + # Clear model and data + self.model = None + self.data = None + + # Reset MuJoCo control callback + try: + mujoco.set_mjcb_control(None) + except Exception as e: + logger.debug(f"Error resetting MuJoCo control callback: {e}") + + except Exception as e: + logger.error(f"Error during resource cleanup: {e}") + + def __del__(self): + """Destructor to ensure cleanup on object deletion.""" + try: + self.cleanup() + except Exception: + pass diff --git a/dimos/simulation/mujoco/policy.py b/dimos/simulation/mujoco/policy.py new file mode 100644 index 0000000000..2ab78f6c4c --- /dev/null +++ b/dimos/simulation/mujoco/policy.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 + +# Copyright 2025 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. + + +import mujoco +import numpy as np +import onnxruntime as rt + +from dimos.simulation.mujoco.types import InputController + + +class OnnxController: + """ONNX controller for the Go-1 robot.""" + + def __init__( + self, + policy_path: str, + default_angles: np.ndarray, + n_substeps: int, + action_scale: float, + input_controller: InputController, + ): + self._output_names = ["continuous_actions"] + self._policy = rt.InferenceSession(policy_path, providers=["CPUExecutionProvider"]) + + self._action_scale = action_scale + self._default_angles = default_angles + self._last_action = np.zeros_like(default_angles, dtype=np.float32) + + self._counter = 0 + self._n_substeps = n_substeps + self._input_controller = input_controller + + def get_obs(self, model, data) -> np.ndarray: + linvel = data.sensor("local_linvel").data + gyro = data.sensor("gyro").data + imu_xmat = data.site_xmat[model.site("imu").id].reshape(3, 3) + gravity = imu_xmat.T @ np.array([0, 0, -1]) + joint_angles = data.qpos[7:] - self._default_angles + joint_velocities = data.qvel[6:] + obs = np.hstack( + [ + linvel, + gyro, + gravity, + joint_angles, + joint_velocities, + self._last_action, + self._input_controller.get_command(), + ] + ) + return obs.astype(np.float32) + + def get_control(self, model: mujoco.MjModel, data: mujoco.MjData) -> None: + self._counter += 1 + if self._counter % self._n_substeps == 0: + obs = self.get_obs(model, data) + onnx_input = {"obs": obs.reshape(1, -1)} + onnx_pred = self._policy.run(self._output_names, onnx_input)[0][0] + self._last_action = onnx_pred.copy() + data.ctrl[:] = onnx_pred * self._action_scale + self._default_angles diff --git a/dimos/simulation/mujoco/types.py b/dimos/simulation/mujoco/types.py new file mode 100644 index 0000000000..42fd28efd2 --- /dev/null +++ b/dimos/simulation/mujoco/types.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python3 + +# Copyright 2025 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. + + +from typing import Protocol + +import numpy as np + + +class InputController(Protocol): + """A protocol for input devices to control the robot.""" + + def get_command(self) -> np.ndarray: ... + def stop(self) -> None: ... diff --git a/pyproject.toml b/pyproject.toml index 80ad9f0f7b..2b81ebb52f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -158,6 +158,12 @@ dev = [ "textual==3.7.1" ] +sim = [ + # Simulation + "mujoco>=3.3.4", + "playground>=0.0.5", +] + [tool.ruff] line-length = 100 exclude = [