diff --git a/dimos/robot/frontier_exploration/wavefront_frontier_goal_selector.py b/dimos/robot/frontier_exploration/wavefront_frontier_goal_selector.py index 76f2ddbb0a..0e86462090 100644 --- a/dimos/robot/frontier_exploration/wavefront_frontier_goal_selector.py +++ b/dimos/robot/frontier_exploration/wavefront_frontier_goal_selector.py @@ -657,7 +657,9 @@ def explore(self, stop_event: Optional[threading.Event] = None) -> bool: # Navigate to the frontier logger.info(f"Navigating to frontier at {next_goal}") - navigation_successful = self.set_goal(next_goal, stop_event=stop_event) + navigation_successful = self.set_goal( + next_goal, + ) if not navigation_successful: logger.warning("Failed to navigate to frontier, continuing exploration") diff --git a/dimos/robot/global_planner/planner.py b/dimos/robot/global_planner/planner.py index 55eea616a0..9a1cea91b3 100644 --- a/dimos/robot/global_planner/planner.py +++ b/dimos/robot/global_planner/planner.py @@ -38,19 +38,19 @@ def __init__(self): Module.__init__(self) Visualizable.__init__(self) - # def set_goal( - # self, - # goal: VectorLike, - # goal_theta: Optional[float] = None, - # stop_event: Optional[threading.Event] = None, - # ): - # path = self.plan(goal) - # if not path: - # logger.warning("No path found to the goal.") - # return False + @rpc + def set_goal( + self, + goal: VectorLike, + goal_theta: Optional[float] = None, + stop_event: Optional[threading.Event] = None, + ): + path = self.plan(goal) + if not path: + logger.warning("No path found to the goal.") + return False - # print("pathing success", path) - # return self.set_local_nav(path, stop_event=stop_event, goal_theta=goal_theta) + print("pathing success", path) class AstarPlanner(Planner): @@ -59,6 +59,7 @@ class AstarPlanner(Planner): get_costmap: Callable[[], Costmap] get_robot_pos: Callable[[], Vector3] + set_local_nav: Callable[[Path, Optional[threading.Event], Optional[float]], bool] conservativism: int = 8 @@ -66,10 +67,12 @@ def __init__( self, get_costmap: Callable[[], Costmap], get_robot_pos: Callable[[], Vector3], + set_local_nav: Callable[[Path, Optional[threading.Event], Optional[float]], bool] = None, ): super().__init__() self.get_costmap = get_costmap self.get_robot_pos = get_robot_pos + self.set_local_nav = set_local_nav @rpc def start(self): @@ -80,9 +83,10 @@ def plan(self, goal: VectorLike) -> Path: goal = to_vector(goal).to_2d() pos = self.get_robot_pos() print("current pos", pos) - costmap = self.get_costmap().smudge() + costmap = self.get_costmap() print("current costmap", costmap) + self.vis("target", goal) print("ASTAR ", costmap, goal, pos) @@ -92,5 +96,7 @@ def plan(self, goal: VectorLike) -> Path: path = path.resample(0.1) self.vis("a*", path) self.path.publish(path) + if hasattr(self, "set_local_nav"): + self.set_local_nav(path) return path logger.warning("No path found to the goal.") diff --git a/dimos/robot/local_planner/__init__.py b/dimos/robot/local_planner/__init__.py index 472b58dcd2..e69de29bb2 100644 --- a/dimos/robot/local_planner/__init__.py +++ b/dimos/robot/local_planner/__init__.py @@ -1,7 +0,0 @@ -from dimos.robot.local_planner.local_planner import ( - BaseLocalPlanner, - navigate_to_goal_local, - navigate_path_local, -) - -from dimos.robot.local_planner.vfh_local_planner import VFHPurePursuitPlanner diff --git a/dimos/robot/local_planner/local_planner.py b/dimos/robot/local_planner/local_planner.py index 286ee94f2b..3c9633a67c 100644 --- a/dimos/robot/local_planner/local_planner.py +++ b/dimos/robot/local_planner/local_planner.py @@ -14,28 +14,34 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import math -import numpy as np -from typing import Dict, Tuple, Optional, Callable, Any +import threading +import time from abc import ABC, abstractmethod +from collections import deque +from typing import Any, Callable, Dict, Optional, Tuple + import cv2 +import numpy as np +import reactivex as rx from reactivex import Observable +from reactivex import operators as ops from reactivex.subject import Subject -import threading -import time -import logging -from collections import deque -from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import normalize_angle, distance_angle_to_goal_xy -from dimos.types.vector import VectorLike, Vector, to_tuple -from dimos.types.path import Path +from dimos.core import In, Module, Out, rpc +from dimos.msgs.geometry_msgs import PoseStamped, Vector3 +from dimos.robot.unitree_webrtc.type.odometry import Odometry from dimos.types.costmap import Costmap +from dimos.types.path import Path +from dimos.types.vector import Vector, VectorLike, to_tuple +from dimos.utils.logging_config import setup_logger +from dimos.utils.transform_utils import distance_angle_to_goal_xy, normalize_angle logger = setup_logger("dimos.robot.unitree.local_planner", level=logging.DEBUG) -class BaseLocalPlanner(ABC): +class BaseLocalPlanner(Module, ABC): """ Abstract base class for local planners that handle obstacle avoidance and path following. @@ -63,11 +69,13 @@ class BaseLocalPlanner(ABC): If provided, this will be used to generate a path to the goal before local planning. """ + odom: In[PoseStamped] = None + movecmd: Out[Vector3] = None + latest_odom: PoseStamped = None + def __init__( self, get_costmap: Callable[[], Optional[Costmap]], - get_robot_pose: Callable[[], Any], - move: Callable[[Vector], None], safety_threshold: float = 0.5, max_linear_vel: float = 0.8, max_angular_vel: float = 1.0, @@ -83,9 +91,9 @@ def __init__( global_planner_plan: Optional[Callable[[VectorLike], Optional[Any]]] = None, ): # Control frequency in Hz # Store callables for robot interactions + Module.__init__(self) + self.get_costmap = get_costmap - self.get_robot_pose = get_robot_pose - self.move = move # Store parameters self.safety_threshold = safety_threshold @@ -149,15 +157,25 @@ def __init__( self._costmap = None self._update_frequency = 10.0 # Hz - how often to update cached data self._update_timer = None + + async def start(self): + """Start the local planner's periodic updates and any other initialization.""" self._start_periodic_updates() + def setodom(odom: Odometry): + self.latest_odom = odom + + self.odom.subscribe(setodom) + # self.get_move_stream(frequency=20.0).subscribe(self.movecmd.publish) + def _start_periodic_updates(self): self._update_timer = threading.Thread(target=self._periodic_update, daemon=True) self._update_timer.start() def _periodic_update(self): while True: - self._robot_pose = self.get_robot_pose() + self._robot_pose = self._format_robot_pose() + # print("robot pose", self._robot_pose) self._costmap = self.get_costmap() time.sleep(1.0 / self._update_frequency) @@ -193,7 +211,7 @@ def reset(self): logger.info("Local planner state has been reset") - def _get_robot_pose(self) -> Tuple[Tuple[float, float], float]: + def _format_robot_pose(self) -> Tuple[Tuple[float, float], float]: """ Get the current robot position and orientation. @@ -204,8 +222,10 @@ def _get_robot_pose(self) -> Tuple[Tuple[float, float], float]: """ if self._robot_pose is None: return ((0.0, 0.0), 0.0) # Fallback if not yet initialized - pos, rot = self._robot_pose.pos, self._robot_pose.rot - return (pos.x, pos.y), rot.z + + pos = self.latest_odom.position + euler = self.latest_odom.orientation.to_euler() + return (pos.x, pos.y), euler.z def _get_costmap(self): """Get cached costmap data.""" @@ -352,7 +372,7 @@ def _distance_to_position(self, target_position: Tuple[float, float]) -> float: Returns: Distance in meters """ - robot_pos, _ = self._get_robot_pose() + robot_pos, _ = self._format_robot_pose() return np.linalg.norm( [target_position[0] - robot_pos[0], target_position[1] - robot_pos[1]] ) @@ -411,7 +431,7 @@ def plan(self) -> Dict[str, float]: return {"x_vel": 0.0, "angular_vel": 0.0} # Get current robot pose - robot_pos, robot_theta = self._get_robot_pose() + robot_pos, robot_theta = self._format_robot_pose() robot_pos_np = np.array(robot_pos) # Check if close to final waypoint @@ -485,7 +505,7 @@ def _rotate_to_goal_orientation(self) -> Dict[str, float]: Dict[str, float]: Velocity commands with zero linear velocity """ # Get current robot orientation - _, robot_theta = self._get_robot_pose() + _, robot_theta = self._format_robot_pose() # Calculate the angle difference angle_diff = normalize_angle(self.goal_theta - robot_theta) @@ -512,7 +532,7 @@ def _is_goal_orientation_reached(self) -> bool: return True # No orientation goal set # Get current robot orientation - _, robot_theta = self._get_robot_pose() + _, robot_theta = self._format_robot_pose() # Calculate the angle difference and normalize angle_diff = abs(normalize_angle(self.goal_theta - robot_theta)) @@ -712,7 +732,7 @@ def adjust_goal_to_valid_position( Returns: Tuple[float, float]: A valid goal position, or the original goal if already valid """ - [pos, rot] = self._get_robot_pose() + [pos, rot] = self._format_robot_pose() robot_x, robot_y = pos[0], pos[1] @@ -810,7 +830,7 @@ def check_if_stuck(self) -> bool: current_time = time.time() # Get current robot position - [pos, _] = self._get_robot_pose() + [pos, _] = self._format_robot_pose() current_position = (pos[0], pos[1], current_time) # If we're already in recovery, don't add movements to history (they're intentional) @@ -987,179 +1007,181 @@ def execute_recovery_behavior(self) -> Dict[str, float]: return {"x_vel": 0.0, "angular_vel": 0.0} + @rpc + def navigate_to_goal_local( + self, + goal_xy_robot: Tuple[float, float], + goal_theta: Optional[float] = None, + distance: float = 0.0, + timeout: float = 60.0, + stop_event: Optional[threading.Event] = None, + ) -> bool: + """ + Navigates the robot to a goal specified in the robot's local frame + using the local planner. -def navigate_to_goal_local( - robot, - goal_xy_robot: Tuple[float, float], - goal_theta: Optional[float] = None, - distance: float = 0.0, - timeout: float = 60.0, - stop_event: Optional[threading.Event] = None, -) -> bool: - """ - Navigates the robot to a goal specified in the robot's local frame - using the local planner. - - Args: - robot: Robot instance to control - goal_xy_robot: Tuple (x, y) representing the goal position relative - to the robot's current position and orientation. - distance: Desired distance to maintain from the goal in meters. - If non-zero, the robot will stop this far away from the goal. - timeout: Maximum time (in seconds) allowed to reach the goal. - stop_event: Optional threading.Event to signal when navigation should stop - - Returns: - bool: True if the goal was reached within the timeout, False otherwise. - """ - logger.info( - f"Starting navigation to local goal {goal_xy_robot} with distance {distance}m and timeout {timeout}s." - ) - - robot.local_planner.reset() - - goal_x, goal_y = goal_xy_robot - - # Calculate goal orientation to face the target - if goal_theta is None: - goal_theta = np.arctan2(goal_y, goal_x) - - # If distance is non-zero, adjust the goal to stop at the desired distance - if distance > 0: - # Calculate magnitude of the goal vector - goal_distance = np.sqrt(goal_x**2 + goal_y**2) - - # Only adjust if goal is further than the desired distance - if goal_distance > distance: - goal_x, goal_y = distance_angle_to_goal_xy(goal_distance - distance, goal_theta) - - # Set the goal in the robot's frame with orientation to face the original target - robot.local_planner.set_goal((goal_x, goal_y), is_relative=True, goal_theta=goal_theta) - - # Get control period from robot's local planner for consistent timing - control_period = 1.0 / robot.local_planner.control_frequency - - start_time = time.time() - goal_reached = False - - try: - while time.time() - start_time < timeout and not (stop_event and stop_event.is_set()): - # Check if goal has been reached - if robot.local_planner.is_goal_reached(): - logger.info("Goal reached successfully.") - goal_reached = True - break - - # Check if navigation failed flag is set - if robot.local_planner.navigation_failed: - logger.error("Navigation aborted due to repeated recovery failures.") - goal_reached = False - break + Args: + robot: Robot instance to control + goal_xy_robot: Tuple (x, y) representing the goal position relative + to the robot's current position and orientation. + distance: Desired distance to maintain from the goal in meters. + If non-zero, the robot will stop this far away from the goal. + timeout: Maximum time (in seconds) allowed to reach the goal. + stop_event: Optional threading.Event to signal when navigation should stop - # Get planned velocity towards the goal - vel_command = robot.local_planner.plan() - x_vel = vel_command.get("x_vel", 0.0) - angular_vel = vel_command.get("angular_vel", 0.0) + Returns: + bool: True if the goal was reached within the timeout, False otherwise. + """ + logger.info( + f"Starting navigation to local goal {goal_xy_robot} with distance {distance}m and timeout {timeout}s." + ) - # Send velocity command - robot.local_planner.move(Vector(x_vel, 0, angular_vel)) + self.reset() - # Control loop frequency - use robot's control frequency - time.sleep(control_period) + goal_x, goal_y = goal_xy_robot - if not goal_reached: - logger.warning(f"Navigation timed out after {timeout} seconds before reaching goal.") + # Calculate goal orientation to face the target + if goal_theta is None: + goal_theta = np.arctan2(goal_y, goal_x) - except KeyboardInterrupt: - logger.info("Navigation to local goal interrupted by user.") - goal_reached = False # Consider interruption as failure - except Exception as e: - logger.error(f"Error during navigation to local goal: {e}") - goal_reached = False # Consider error as failure - finally: - logger.info("Stopping robot after navigation attempt.") - robot.local_planner.move(Vector(0, 0, 0)) # Stop the robot + # If distance is non-zero, adjust the goal to stop at the desired distance + if distance > 0: + # Calculate magnitude of the goal vector + goal_distance = np.sqrt(goal_x**2 + goal_y**2) - return goal_reached + # Only adjust if goal is further than the desired distance + if goal_distance > distance: + goal_x, goal_y = distance_angle_to_goal_xy(goal_distance - distance, goal_theta) + # Set the goal in the robot's frame with orientation to face the original target + self.set_goal((goal_x, goal_y), is_relative=True, goal_theta=goal_theta) -def navigate_path_local( - robot, - path: Path, - timeout: float = 120.0, - goal_theta: Optional[float] = None, - stop_event: Optional[threading.Event] = None, -) -> bool: - """ - Navigates the robot along a path of waypoints using the waypoint following capability - of the local planner. + # Get control period from robot's local planner for consistent timing + control_period = 1.0 / self.control_frequency - Args: - robot: Robot instance to control - path: Path object containing waypoints in absolute frame - timeout: Maximum time (in seconds) allowed to follow the complete path - goal_theta: Optional final orientation in radians - stop_event: Optional threading.Event to signal when navigation should stop - - Returns: - bool: True if the entire path was successfully followed, False otherwise - """ - logger.info( - f"Starting navigation along path with {len(path)} waypoints and timeout {timeout}s." - ) + start_time = time.time() + goal_reached = False - robot.local_planner.reset() + try: + while time.time() - start_time < timeout and not (stop_event and stop_event.is_set()): + # Check if goal has been reached + if self.is_goal_reached(): + logger.info("Goal reached successfully.") + goal_reached = True + break + + # Check if navigation failed flag is set + if self.navigation_failed: + logger.error("Navigation aborted due to repeated recovery failures.") + goal_reached = False + break + + # Get planned velocity towards the goal + vel_command = self.plan() + x_vel = vel_command.get("x_vel", 0.0) + angular_vel = vel_command.get("angular_vel", 0.0) + + # Send velocity command + self.movecmd.publish(Vector3(x_vel, 0, angular_vel)) + + # Control loop frequency - use robot's control frequency + time.sleep(control_period) + + if not goal_reached: + logger.warning( + f"Navigation timed out after {timeout} seconds before reaching goal." + ) - # Set the path in the local planner - robot.local_planner.set_goal_waypoints(path, goal_theta=goal_theta) + except KeyboardInterrupt: + logger.info("Navigation to local goal interrupted by user.") + goal_reached = False # Consider interruption as failure + except Exception as e: + logger.error(f"Error during navigation to local goal: {e}") + goal_reached = False # Consider error as failure + finally: + logger.info("Stopping robot after navigation attempt.") + self.movecmd.publish(Vector3(0, 0, 0)) # Stop the robot - # Get control period from robot's local planner for consistent timing - control_period = 1.0 / robot.local_planner.control_frequency + return goal_reached - start_time = time.time() - path_completed = False + @rpc + def navigate_path_local( + self, + path: Path, + timeout: float = 120.0, + goal_theta: Optional[float] = None, + stop_event: Optional[threading.Event] = None, + ) -> bool: + """ + Navigates the robot along a path of waypoints using the waypoint following capability + of the local planner. - try: - while time.time() - start_time < timeout and not (stop_event and stop_event.is_set()): - # Check if the entire path has been traversed - if robot.local_planner.is_goal_reached(): - logger.info("Path traversed successfully.") - path_completed = True - break + Args: + robot: Robot instance to control + path: Path object containing waypoints in absolute frame + timeout: Maximum time (in seconds) allowed to follow the complete path + goal_theta: Optional final orientation in radians + stop_event: Optional threading.Event to signal when navigation should stop - # Check if navigation failed flag is set - if robot.local_planner.navigation_failed: - logger.error("Navigation aborted due to repeated recovery failures.") - path_completed = False - break + Returns: + bool: True if the entire path was successfully followed, False otherwise + """ + logger.info( + f"Starting navigation along path with {len(path)} waypoints and timeout {timeout}s." + ) - # Get planned velocity towards the current waypoint target - vel_command = robot.local_planner.plan() - x_vel = vel_command.get("x_vel", 0.0) - angular_vel = vel_command.get("angular_vel", 0.0) + self.reset() + print() + # Set the path in the local planner + self.set_goal_waypoints(path, goal_theta=goal_theta) - # Send velocity command - robot.local_planner.move(Vector(x_vel, 0, angular_vel)) + # Get control period from robot's local planner for consistent timing + control_period = 1.0 / self.control_frequency - # Control loop frequency - use robot's control frequency - time.sleep(control_period) + start_time = time.time() + path_completed = False - if not path_completed: - logger.warning( - f"Path following timed out after {timeout} seconds before completing the path." - ) + try: + while time.time() - start_time < timeout and not (stop_event and stop_event.is_set()): + # Check if the entire path has been traversed + if self.is_goal_reached(): + logger.info("Path traversed successfully.") + path_completed = True + break + + # Check if navigation failed flag is set + if self.navigation_failed: + logger.error("Navigation aborted due to repeated recovery failures.") + path_completed = False + break + + # Get planned velocity towards the current waypoint target + vel_command = self.plan() + x_vel = vel_command.get("x_vel", 0.0) + angular_vel = vel_command.get("angular_vel", 0.0) + + # Send velocity command + self.movecmd.publish(Vector3(x_vel, 0, angular_vel)) + + # Control loop frequency - use robot's control frequency + time.sleep(control_period) + + if not path_completed: + logger.warning( + f"Path following timed out after {timeout} seconds before completing the path." + ) - except KeyboardInterrupt: - logger.info("Path navigation interrupted by user.") - path_completed = False - except Exception as e: - logger.error(f"Error during path navigation: {e}") - path_completed = False - finally: - logger.info("Stopping robot after path navigation attempt.") - robot.local_planner.move(Vector(0, 0, 0)) # Stop the robot + except KeyboardInterrupt: + logger.info("Path navigation interrupted by user.") + path_completed = False + except Exception as e: + logger.error(f"Error during path navigation: {e}") + path_completed = False + finally: + logger.info("Stopping robot after path navigation attempt.") + self.movecmd.publish(Vector3(0, 0, 0)) # Stop the robot - return path_completed + return path_completed def visualize_local_planner_state( diff --git a/dimos/robot/local_planner/vfh_local_planner.py b/dimos/robot/local_planner/vfh_local_planner.py index f97701e5a5..5945f8bd00 100644 --- a/dimos/robot/local_planner/vfh_local_planner.py +++ b/dimos/robot/local_planner/vfh_local_planner.py @@ -38,8 +38,6 @@ class VFHPurePursuitPlanner(BaseLocalPlanner): def __init__( self, get_costmap: Callable[[], Optional[Costmap]], - get_robot_pose: Callable[[], Any], - move: Callable[[Vector], None], safety_threshold: float = 0.8, histogram_bins: int = 144, max_linear_vel: float = 0.8, @@ -80,8 +78,6 @@ def __init__( # Initialize base class super().__init__( get_costmap=get_costmap, - get_robot_pose=get_robot_pose, - move=move, safety_threshold=safety_threshold, max_linear_vel=max_linear_vel, max_angular_vel=max_angular_vel, @@ -129,7 +125,7 @@ def _compute_velocity_commands(self) -> Dict[str, float]: logger.warning("No costmap available for planning") return {"x_vel": 0.0, "angular_vel": 0.0} - robot_pos, robot_theta = self._get_robot_pose() + robot_pos, robot_theta = self._format_robot_pose() robot_x, robot_y = robot_pos robot_pose = (robot_x, robot_y, robot_theta) @@ -346,7 +342,7 @@ def check_collision(self, selected_direction: float, safety_threshold: float = 1 if costmap is None: return False # No costmap available - robot_pos, robot_theta = self._get_robot_pose() + robot_pos, robot_theta = self._format_robot_pose() robot_x, robot_y = robot_pos # Direction in world frame @@ -382,7 +378,7 @@ def update_visualization(self) -> np.ndarray: if costmap is None: raise ValueError("Costmap is None") - robot_pos, robot_theta = self._get_robot_pose() + robot_pos, robot_theta = self._format_robot_pose() robot_x, robot_y = robot_pos robot_pose = (robot_x, robot_y, robot_theta) diff --git a/dimos/robot/unitree_webrtc/connection.py b/dimos/robot/unitree_webrtc/connection.py index 86fe5f6a85..29080d1ebb 100644 --- a/dimos/robot/unitree_webrtc/connection.py +++ b/dimos/robot/unitree_webrtc/connection.py @@ -42,7 +42,7 @@ VideoMessage: TypeAlias = np.ndarray[tuple[int, int, Literal[3]], np.uint8] -class WebRTCRobot(ConnectionInterface): +class UnitreeWebRTCConnection(ConnectionInterface): def __init__(self, ip: str, mode: str = "ai"): self.ip = ip self.mode = mode diff --git a/dimos/robot/unitree_webrtc/multiprocess/example_usage.py b/dimos/robot/unitree_webrtc/multiprocess/example_usage.py new file mode 100644 index 0000000000..6e11998ff5 --- /dev/null +++ b/dimos/robot/unitree_webrtc/multiprocess/example_usage.py @@ -0,0 +1,156 @@ +#!/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. + +"""Example usage of UnitreeGo2Light and UnitreeGo2Heavy classes.""" + +import asyncio +import os +from dimos.robot.unitree_webrtc.multiprocess.unitree_go2 import UnitreeGo2Light +from dimos.robot.unitree_webrtc.multiprocess.unitree_go2_heavy import UnitreeGo2Heavy +from dimos.utils.reactive import backpressure +from dimos.perception.object_detection_stream import ObjectDetectionStream +from dimos.web.robot_web_interface import RobotWebInterface +from dimos.robot.unitree_webrtc.unitree_skills import MyUnitreeSkills +from dimos.skills.skills import SkillLibrary, AbstractRobotSkill +from dimos.utils.reactive import backpressure +import reactivex as rx +import reactivex.operators as ops +from dimos.stream.audio.pipelines import tts, stt +import threading +from dimos.agents.claude_agent import ClaudeAgent + + +async def run_light_robot(): + """Example of running the lightweight robot without GPU modules.""" + ip = os.getenv("ROBOT_IP") + + robot = UnitreeGo2Light(ip) + + await robot.start() + + pose = robot.get_pose() + print(f"Robot position: {pose['position']}") + print(f"Robot rotation: {pose['rotation']}") + + from dimos.types.vector import Vector + + robot.move(Vector(0.5, 0, 0), duration=2.0) + + robot.explore() + + while True: + await asyncio.sleep(1) + + +async def run_heavy_robot(): + """Example of running the heavy robot with full GPU modules.""" + ip = os.getenv("ROBOT_IP") + + # Create heavy robot instance with all features + robot = UnitreeGo2Heavy(ip=ip, new_memory=True, enable_perception=True) + + # Start the robot + await robot.start() + + if robot.spatial_memory: + print("Spatial memory initialized") + + skills = robot.get_skills() + print(f"Available skills: {[skill.__class__.__name__ for skill in skills]}") + + from dimos.types.robot_capabilities import RobotCapability + + if robot.has_capability(RobotCapability.VISION): + print("Robot has vision capability") + + if robot.person_tracking_stream: + print("Person tracking enabled") + + # Start exploration with spatial memory recording + print(robot.spatial_memory.query_by_text("kitchen")) + + # robot.frontier_explorer.explore() + + # Create a subject for agent responses + agent_response_subject = rx.subject.Subject() + agent_response_stream = agent_response_subject.pipe(ops.share()) + audio_subject = rx.subject.Subject() + + video_stream = robot.get_video_stream() # WebRTC doesn't use ROS video stream + + # # Initialize ObjectDetectionStream with robot + # object_detector = ObjectDetectionStream( + # camera_intrinsics=robot.camera_intrinsics, + # get_pose=robot.get_pose, + # video_stream=video_stream, + # draw_masks=True, + # ) + + # # Create visualization stream for web interface + # viz_stream = backpressure(object_detector.get_stream()).pipe( + # ops.share(), + # ops.map(lambda x: x["viz_frame"] if x is not None else None), + # ops.filter(lambda x: x is not None), + # ) + + streams = { + "unitree_video": robot.get_video_stream(), # Changed from get_ros_video_stream to get_video_stream for WebRTC + # "object_detection": viz_stream, # Uncommented object detection + } + text_streams = { + "agent_responses": agent_response_stream, + } + + web_interface = RobotWebInterface( + port=5555, text_streams=text_streams, audio_subject=audio_subject, **streams + ) + + stt_node = stt() + stt_node.consume_audio(audio_subject.pipe(ops.share())) + + agent = ClaudeAgent( + dev_name="test_agent", + # input_query_stream=stt_node.emit_text(), + input_query_stream=web_interface.query_stream, + skills=robot.get_skills(), + system_query="You are a helpful robot.", + model_name="claude-3-5-haiku-latest", + thinking_budget_tokens=0, + max_output_tokens_per_request=8192, + # model_name="llama-4-scout-17b-16e-instruct", + ) + agent.get_response_observable().subscribe(lambda x: agent_response_subject.on_next(x)) + + # Start web interface in a separate thread to avoid blocking + web_thread = threading.Thread(target=web_interface.run) + web_thread.daemon = True + web_thread.start() + + # Keep running + while True: + await asyncio.sleep(1) + + +if __name__ == "__main__": + use_heavy = True + + if use_heavy: + print("Running UnitreeGo2Heavy with GPU modules...") + asyncio.run(run_heavy_robot()) + else: + print("Running UnitreeGo2Light without GPU modules...") + asyncio.run(run_light_robot()) diff --git a/dimos/robot/unitree_webrtc/multiprocess/individual_node_example.py b/dimos/robot/unitree_webrtc/multiprocess/individual_node_example.py index 56bf50bf49..3bcd9af9ff 100644 --- a/dimos/robot/unitree_webrtc/multiprocess/individual_node_example.py +++ b/dimos/robot/unitree_webrtc/multiprocess/individual_node_example.py @@ -25,7 +25,6 @@ from dimos.msgs.sensor_msgs import Image from dimos.protocol import pubsub from dimos.robot.global_planner import AstarPlanner -from dimos.robot.unitree_webrtc.connection import VideoMessage, WebRTCRobot from dimos.robot.unitree_webrtc.type.lidar import LidarMessage from dimos.robot.unitree_webrtc.type.map import Map from dimos.robot.unitree_webrtc.type.odometry import Odometry diff --git a/dimos/robot/unitree_webrtc/multiprocess/test_actors.py b/dimos/robot/unitree_webrtc/multiprocess/test_actors.py index 7585e746cc..c6679d3262 100644 --- a/dimos/robot/unitree_webrtc/multiprocess/test_actors.py +++ b/dimos/robot/unitree_webrtc/multiprocess/test_actors.py @@ -26,7 +26,6 @@ from dimos.protocol import pubsub from dimos.robot.global_planner import AstarPlanner from dimos.robot.local_planner.simple import SimplePlanner -from dimos.robot.unitree_webrtc.connection import VideoMessage, WebRTCRobot from dimos.robot.unitree_webrtc.type.lidar import LidarMessage from dimos.robot.unitree_webrtc.type.map import Map from dimos.robot.unitree_webrtc.type.map import Map as Mapper diff --git a/dimos/robot/unitree_webrtc/multiprocess/unitree_go2.py b/dimos/robot/unitree_webrtc/multiprocess/unitree_go2.py index bf20449e62..9e82d52c9d 100644 --- a/dimos/robot/unitree_webrtc/multiprocess/unitree_go2.py +++ b/dimos/robot/unitree_webrtc/multiprocess/unitree_go2.py @@ -13,16 +13,12 @@ # limitations under the License. import asyncio -import contextvars import functools import threading import time from typing import Callable -from dask.distributed import get_client, get_worker -from distributed import get_worker from reactivex import operators as ops -from reactivex.scheduler import ThreadPoolScheduler import dimos.core.colors as colors from dimos import core @@ -32,20 +28,45 @@ from dimos.protocol import pubsub from dimos.robot.foxglove_bridge import FoxgloveBridge from dimos.robot.global_planner import AstarPlanner -from dimos.robot.local_planner.simple import SimplePlanner -from dimos.robot.unitree_webrtc.connection import VideoMessage, WebRTCRobot +from dimos.robot.local_planner.vfh_local_planner import VFHPurePursuitPlanner +from dimos.robot.unitree_webrtc.connection import VideoMessage, UnitreeWebRTCConnection from dimos.robot.unitree_webrtc.type.lidar import LidarMessage from dimos.robot.unitree_webrtc.type.map import Map from dimos.robot.unitree_webrtc.type.odometry import Odometry from dimos.types.costmap import Costmap from dimos.types.vector import Vector from dimos.utils.data import get_data -from dimos.utils.reactive import backpressure, getter_streaming +from dimos.utils.reactive import getter_streaming from dimos.utils.testing import TimedSensorReplay - - -# can be swapped in for WebRTCRobot -class FakeRTC(WebRTCRobot): +from dimos.robot.frontier_exploration.wavefront_frontier_goal_selector import ( + WavefrontFrontierExplorer, +) +import os +import logging +import warnings +from dimos.utils.logging_config import setup_logger +from reactivex import Observable + +logger = setup_logger("dimos.robot.unitree_webrtc.multiprocess.unitree_go2", level=logging.INFO) + +# Configure logging levels +os.environ["DIMOS_LOG_LEVEL"] = "WARNING" + +# Suppress verbose loggers +logging.getLogger("aiortc.codecs.h264").setLevel(logging.ERROR) +logging.getLogger("lcm_foxglove_bridge").setLevel(logging.ERROR) +logging.getLogger("websockets.server").setLevel(logging.ERROR) +logging.getLogger("FoxgloveServer").setLevel(logging.ERROR) +logging.getLogger("asyncio").setLevel(logging.ERROR) +logging.getLogger("root").setLevel(logging.WARNING) + +# Suppress warnings +warnings.filterwarnings("ignore", message="coroutine.*was never awaited") +warnings.filterwarnings("ignore", message="H264Decoder.*failed to decode") + + +# can be swapped in for UnitreeWebRTCConnection +class FakeRTC(UnitreeWebRTCConnection): def __init__(self, *args, **kwargs): # ensures we download msgs from lfs store data = get_data("unitree_office_walk") @@ -80,11 +101,7 @@ def move(self, vector: Vector): print("move supressed", vector) -class RealRTC(WebRTCRobot): ... - - -# inherit RealRTC instead of FakeRTC to run the real robot -class ConnectionModule(FakeRTC, Module): +class ConnectionModule(UnitreeWebRTCConnection, Module): movecmd: In[Vector3] = None odom: Out[Vector3] = None lidar: Out[LidarMessage] = None @@ -104,12 +121,18 @@ def __init__(self, ip: str, *args, **kwargs): @rpc def start(self): + # Initialize the parent WebRTC connection super().__init__(self.ip) - # ensure that LFS data is available + + # Connect sensor streams to LCM outputs self.lidar_stream().subscribe(self.lidar.publish) self.odom_stream().subscribe(self.odom.publish) self.video_stream().subscribe(self.video.publish) + + # Connect LCM input to robot movement commands self.movecmd.subscribe(self.move) + + # Set up streaming getters for latest sensor data self._odom = getter_streaming(self.odom_stream()) self._lidar = getter_streaming(self.lidar_stream()) @@ -140,85 +163,219 @@ def plancmd(): thread.start() -async def run(ip): - dimos = core.start(4) - connection = dimos.deploy(ConnectionModule, ip) - - # This enables LCM transport - # Ensures system multicast, udp sizes are auto-adjusted if needed - pubsub.lcm.autoconf() - - connection.lidar.transport = core.LCMTransport("/lidar", LidarMessage) - connection.odom.transport = core.LCMTransport("/odom", PoseStamped) - connection.video.transport = core.LCMTransport("/video", Image) - - mapper = dimos.deploy(Map, voxel_size=0.5, global_publish_interval=2.5) - - mapper.global_map.transport = core.LCMTransport("/global_map", LidarMessage) - - local_planner = dimos.deploy( - SimplePlanner, - get_costmap=connection.get_local_costmap, - ) - - global_planner = dimos.deploy( - AstarPlanner, - get_costmap=mapper.costmap, - get_robot_pos=connection.get_pos, - ) - - global_planner.path.transport = core.pLCMTransport("/global_path") - - local_planner.path.connect(global_planner.path) - local_planner.odom.connect(connection.odom) - - local_planner.movecmd.transport = core.LCMTransport("/move", Vector3) - connection.movecmd.connect(local_planner.movecmd) - - ctrl = dimos.deploy(ControlModule) - - mapper.lidar.connect(connection.lidar) - - ctrl.plancmd.transport = core.LCMTransport("/global_target", Vector3) +class UnitreeGo2Light: + def __init__(self, ip: str): + self.ip = ip + self.dimos = None + self.connection = None + self.mapper = None + self.local_planner = None + self.global_planner = None + self.frontier_explorer = None + self.foxglove_bridge = None + self.ctrl = None + + async def start(self): + self.dimos = core.start(4) + + # Connection Module - Robot sensor data interface via WebRTC =================== + self.connection = self.dimos.deploy(ConnectionModule, self.ip) + + # This enables LCM transport + # Ensures system multicast, udp sizes are auto-adjusted if needed + pubsub.lcm.autoconf() + + # Configure ConnectionModule LCM transport outputs for sensor data streams + # OUTPUT: LiDAR point cloud data to /lidar topic + self.connection.lidar.transport = core.LCMTransport("/lidar", LidarMessage) + # OUTPUT: Robot odometry/pose data to /odom topic + self.connection.odom.transport = core.LCMTransport("/odom", PoseStamped) + # OUTPUT: Camera video frames to /video topic + self.connection.video.transport = core.LCMTransport("/video", Image) + # ====================================================================== + + # Map Module - Point cloud accumulation and costmap generation ========= + self.mapper = self.dimos.deploy(Map, voxel_size=0.5, global_publish_interval=2.5) + + # OUTPUT: Accumulated point cloud map to /global_map topic + self.mapper.global_map.transport = core.LCMTransport("/global_map", LidarMessage) + + # Connect ConnectionModule OUTPUT lidar to Map INPUT lidar for point cloud accumulation + self.mapper.lidar.connect(self.connection.lidar) + # ==================================================================== + + # Local planner Module, LCM transport & connection ================ + self.local_planner = self.dimos.deploy( + VFHPurePursuitPlanner, + get_costmap=self.connection.get_local_costmap, + ) + + # Connects odometry LCM stream to BaseLocalPlanner odometry input + self.local_planner.odom.connect(self.connection.odom) + + # Configures BaseLocalPlanner movecmd output to /move LCM topic + self.local_planner.movecmd.transport = core.LCMTransport("/move", Vector3) + + # Connects connection.movecmd input to local_planner.movecmd output + self.connection.movecmd.connect(self.local_planner.movecmd) + # =================================================================== + + # Global Planner Module =============== + self.global_planner = self.dimos.deploy( + AstarPlanner, + get_costmap=self.mapper.costmap, + get_robot_pos=self.connection.get_pos, + set_local_nav=self.local_planner.navigate_path_local, + ) + + # Configure AstarPlanner OUTPUT path: Out[Path] to /global_path LCM topic + self.global_planner.path.transport = core.pLCMTransport("/global_path") + # ====================================== + + # Global Planner Control Module =========================== + # Debug module that sends (0,0,0) goal after 4 second delay + self.ctrl = self.dimos.deploy(ControlModule) + + # Configure ControlModule OUTPUT to publish goal coordinates to /global_target + self.ctrl.plancmd.transport = core.LCMTransport("/global_target", Vector3) + + # Connect ControlModule OUTPUT to AstarPlanner INPUT - triggers A* planning when goal received + self.global_planner.target.connect(self.ctrl.plancmd) + # ========================================== + + # Visualization ============================ + self.foxglove_bridge = FoxgloveBridge() + # ========================================== + + self.frontier_explorer = WavefrontFrontierExplorer( + set_goal=self.global_planner.set_goal, + get_costmap=self.mapper.costmap, + get_robot_pos=self.connection.get_pos, + ) + + # Prints full module IO + print("\n") + for module in [ + self.connection, + self.mapper, + self.local_planner, + self.global_planner, + self.ctrl, + ]: + print(module.io().result(), "\n") + + # Start modules ============================= + self.mapper.start() + self.connection.start() + self.local_planner.start() + self.global_planner.start() + self.foxglove_bridge.start() + # self.ctrl.start() # DEBUG + + await asyncio.sleep(2) + print("querying system") + print(self.mapper.costmap()) + logger.info("UnitreeGo2Light initialized and started") + + def get_pose(self) -> dict: + """Get the current pose (position and rotation) of the robot. + + Returns: + Dictionary containing: + - position: Vector (x, y, z) + - rotation: Vector (roll, pitch, yaw) in radians + """ + if not self.connection: + raise RuntimeError("Connection not established. Call start() first.") + odom = self.connection.get_odom() + position = Vector(odom.x, odom.y, odom.z) + rotation = Vector(odom.roll, odom.pitch, odom.yaw) + return {"position": position, "rotation": rotation} + + def move(self, velocity: Vector, duration: float = 0.0) -> bool: + """Move the robot using velocity commands. + + Args: + velocity: Velocity vector [x, y, yaw] + duration: Duration to apply command (seconds) + + Returns: + bool: True if movement succeeded + """ + if not self.connection: + raise RuntimeError("Connection not established. Call start() first.") + self.connection.move(Vector3(velocity.x, velocity.y, velocity.z)) + if duration > 0: + time.sleep(duration) + self.connection.move(Vector3(0, 0, 0)) # Stop + return True + + def explore(self, stop_event=None) -> bool: + """Start autonomous frontier exploration. + + Args: + stop_event: Optional threading.Event to signal when exploration should stop + + Returns: + bool: True if exploration completed successfully + """ + if not self.frontier_explorer: + raise RuntimeError("Frontier explorer not initialized. Call start() first.") + return self.frontier_explorer.explore(stop_event=stop_event) - global_planner.target.connect(ctrl.plancmd) + def standup(self): + """Make the robot stand up.""" + if self.connection and hasattr(self.connection, "standup"): + return self.connection.standup() - foxglove_bridge = FoxgloveBridge() + def liedown(self): + """Make the robot lie down.""" + if self.connection and hasattr(self.connection, "liedown"): + return self.connection.liedown() - # we review the structure - print("\n") - for module in [connection, mapper, local_planner, global_planner, ctrl]: - print(module.io().result(), "\n") + @property + def costmap(self): + """Access to the costmap for navigation.""" + if not self.mapper: + raise RuntimeError("Mapper not initialized. Call start() first.") + return self.mapper.costmap - print(colors.green("starting mapper")) - mapper.start() + def get_video_stream(self, fps: int = 30) -> Observable: + """Get the video stream with rate limiting and processing. - print(colors.green("starting connection")) - connection.start() + Args: + fps: Frames per second for rate limiting - print(colors.green("local planner start")) - local_planner.start() + Returns: + Observable stream of video frames + """ + # Import required modules for LCM subscription + from dimos.protocol.pubsub.lcmpubsub import LCM, Topic + from dimos.msgs.sensor_msgs import Image + from reactivex import create + from reactivex.disposable import Disposable - print(colors.green("starting global planner")) - global_planner.start() + lcm_instance = LCM() + lcm_instance.start() - print(colors.green("starting foxglove bridge")) - foxglove_bridge.start() + topic = Topic("/video", Image) - print(colors.green("starting ctrl")) - ctrl.start() + def subscribe(observer, scheduler=None): + unsubscribe_fn = lcm_instance.subscribe(topic, lambda msg, _: observer.on_next(msg)) - print(colors.red("READY")) + return Disposable(unsubscribe_fn) - await asyncio.sleep(2) - print("querying system") - print(mapper.costmap()) - while True: - await asyncio.sleep(1) + return create(subscribe).pipe( + ops.map( + lambda img: img.data if hasattr(img, "data") else img + ), # Convert Image message to numpy array + ops.sample(1.0 / fps), + ) if __name__ == "__main__": import os - asyncio.run(run(os.getenv("ROBOT_IP"))) + robot = UnitreeGo2Light(os.getenv("ROBOT_IP")) + asyncio.run(robot.start()) # asyncio.run(run("192.168.9.140")) diff --git a/dimos/robot/unitree_webrtc/multiprocess/unitree_go2_heavy.py b/dimos/robot/unitree_webrtc/multiprocess/unitree_go2_heavy.py new file mode 100644 index 0000000000..b38f476faa --- /dev/null +++ b/dimos/robot/unitree_webrtc/multiprocess/unitree_go2_heavy.py @@ -0,0 +1,229 @@ +# 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. + +"""Heavy version of Unitree Go2 with GPU-required modules.""" + +import os +import asyncio +from typing import Optional, List +import numpy as np +from reactivex import Observable, operators as ops +from reactivex.disposable import CompositeDisposable +from reactivex.scheduler import ThreadPoolScheduler + +from dimos.robot.unitree_webrtc.multiprocess.unitree_go2 import UnitreeGo2Light +from dimos.perception.spatial_perception import SpatialMemory +from dimos.perception.person_tracker import PersonTrackingStream +from dimos.perception.object_tracker import ObjectTrackingStream +from dimos.skills.skills import SkillLibrary, AbstractRobotSkill +from dimos.robot.unitree_webrtc.unitree_skills import MyUnitreeSkills +from dimos.types.robot_capabilities import RobotCapability +from dimos.types.vector import Vector +from dimos.utils.logging_config import setup_logger +from dimos.utils.threadpool import get_scheduler + +logger = setup_logger("dimos.robot.unitree_webrtc.multiprocess.unitree_go2_heavy") + + +class UnitreeGo2Heavy(UnitreeGo2Light): + """Heavy version of Unitree Go2 with additional GPU-required modules. + + This class extends UnitreeGo2Light with: + - Spatial memory with ChromaDB + - Person tracking stream + - Object tracking stream + - Skill library integration + - Full perception capabilities + """ + + def __init__( + self, + ip: str, + output_dir: str = os.path.join(os.getcwd(), "assets", "output"), + skill_library: Optional[SkillLibrary] = None, + robot_capabilities: Optional[List[RobotCapability]] = None, + spatial_memory_collection: str = "spatial_memory", + new_memory: bool = True, + enable_perception: bool = True, + pool_scheduler: Optional[ThreadPoolScheduler] = None, + ): + """Initialize Unitree Go2 Heavy with full capabilities. + + Args: + ip: IP address of the robot + output_dir: Directory for output files + skill_library: Skill library instance + robot_capabilities: List of robot capabilities + spatial_memory_collection: Collection name for spatial memory + new_memory: Whether to create new spatial memory + enable_perception: Whether to enable perception streams + pool_scheduler: Thread pool scheduler for async operations + """ + super().__init__(ip) + + self.output_dir = output_dir + self.enable_perception = enable_perception + self.disposables = CompositeDisposable() + self.pool_scheduler = pool_scheduler if pool_scheduler else get_scheduler() + + # Initialize capabilities + self.capabilities = robot_capabilities or [ + RobotCapability.LOCOMOTION, + RobotCapability.VISION, + RobotCapability.AUDIO, + ] + + # Create output directory + os.makedirs(self.output_dir, exist_ok=True) + logger.info(f"Robot outputs will be saved to: {self.output_dir}") + + # Initialize memory directories + self.memory_dir = os.path.join(self.output_dir, "memory") + os.makedirs(self.memory_dir, exist_ok=True) + + # Initialize spatial memory properties + self.spatial_memory_dir = os.path.join(self.memory_dir, "spatial_memory") + self.spatial_memory_collection = spatial_memory_collection + self.db_path = os.path.join(self.spatial_memory_dir, "chromadb_data") + self.visual_memory_path = os.path.join(self.spatial_memory_dir, "visual_memory.pkl") + + # Create spatial memory directory + os.makedirs(self.spatial_memory_dir, exist_ok=True) + os.makedirs(self.db_path, exist_ok=True) + + # Camera configuration for Unitree Go2 + self.camera_intrinsics = [819.553492, 820.646595, 625.284099, 336.808987] + self.camera_pitch = np.deg2rad(0) # negative for downward pitch + self.camera_height = 0.44 # meters + + # Initialize skill library + if skill_library is None: + skill_library = MyUnitreeSkills() + self.skill_library = skill_library + + # Initialize spatial memory (will be created after connection is established) + self._spatial_memory = None + self._video_stream = None + + # Tracking streams (initialized after start) + self.person_tracker = None + self.object_tracker = None + self.person_tracking_stream = None + self.object_tracking_stream = None + + async def start(self): + """Start the robot modules and initialize heavy components.""" + # First start the lightweight components + await super().start() + + await asyncio.sleep(0.5) + + # Now we have connection publishing to LCM, initialize video stream + self._video_stream = self.get_video_stream(fps=10) # Lower FPS for processing + + # Initialize spatial memory if perception is enabled + if self.enable_perception and self._video_stream is not None: + self._spatial_memory = SpatialMemory( + collection_name=self.spatial_memory_collection, + db_path=self.db_path, + visual_memory_path=self.visual_memory_path, + new_memory=True, # Always create new for now + output_dir=self.spatial_memory_dir, + video_stream=self._video_stream, + get_pose=self.get_pose, + ) + logger.info("Spatial memory initialized") + + # Initialize person and object tracking + self.person_tracker = PersonTrackingStream( + camera_intrinsics=self.camera_intrinsics, + camera_pitch=self.camera_pitch, + camera_height=self.camera_height, + ) + self.object_tracker = ObjectTrackingStream( + camera_intrinsics=self.camera_intrinsics, + camera_pitch=self.camera_pitch, + camera_height=self.camera_height, + ) + + # Create tracking streams + self.person_tracking_stream = self.person_tracker.create_stream(self._video_stream) + self.object_tracking_stream = self.object_tracker.create_stream(self._video_stream) + + logger.info("Person and object tracking initialized") + else: + logger.info("Perception disabled or video stream unavailable") + + # Initialize skills with robot reference + if self.skill_library is not None: + for skill in self.skill_library: + if isinstance(skill, AbstractRobotSkill): + self.skill_library.create_instance(skill.__name__, robot=self) + if isinstance(self.skill_library, MyUnitreeSkills): + self.skill_library._robot = self + self.skill_library.init() + self.skill_library.initialize_skills() + + logger.info("UnitreeGo2Heavy initialized with all modules") + + @property + def spatial_memory(self) -> Optional[SpatialMemory]: + """Get the robot's spatial memory. + + Returns: + SpatialMemory instance or None if perception is disabled + """ + return self._spatial_memory + + @property + def video_stream(self) -> Optional[Observable]: + """Get the robot's video stream. + + Returns: + Observable video stream or None if not available + """ + return self._video_stream + + def get_skills(self): + """Get the robot's skill library. + + Returns: + The robot's skill library for adding/managing skills + """ + return self.skill_library + + def has_capability(self, capability: RobotCapability) -> bool: + """Check if the robot has a specific capability. + + Args: + capability: The capability to check for + + Returns: + bool: True if the robot has the capability + """ + return capability in self.capabilities + + def cleanup(self): + """Clean up resources used by the robot.""" + # Dispose of reactive resources + if self.disposables: + self.disposables.dispose() + + # Clean up tracking streams + if self.person_tracker: + self.person_tracker = None + if self.object_tracker: + self.object_tracker = None + + logger.info("UnitreeGo2Heavy cleanup completed") diff --git a/dimos/robot/unitree_webrtc/unitree_go2.py b/dimos/robot/unitree_webrtc/unitree_go2.py index 94676bfffc..805c8efb28 100644 --- a/dimos/robot/unitree_webrtc/unitree_go2.py +++ b/dimos/robot/unitree_webrtc/unitree_go2.py @@ -18,7 +18,7 @@ import os from dimos.robot.robot import Robot from dimos.robot.unitree_webrtc.type.map import Map -from dimos.robot.unitree_webrtc.connection import WebRTCRobot +from dimos.robot.unitree_webrtc.connection import UnitreeWebRTCConnection from dimos.robot.global_planner.planner import AstarPlanner from dimos.utils.reactive import getter_streaming from dimos.skills.skills import AbstractRobotSkill, SkillLibrary @@ -66,7 +66,7 @@ def __init__( enable_perception: Whether to enable perception streams and spatial memory """ # Create WebRTC connection interface - self.webrtc_connection = WebRTCRobot( + self.webrtc_connection = UnitreeWebRTCConnection( ip=ip, mode=mode, )