Merge pull request #210 - Cleaner global planner api, faster rendering of real time robot position

Vis/fix2 - cleaner global planner api, faster rendering of real time robot position

Author: spomichter

dimos/robot/global_planner/planner.py

@@ -14,12 +14,12 @@
 
 from dataclasses import dataclass
 from abc import abstractmethod
-from typing import Tuple, Callable, Optional
+from typing import Callable, Optional
 import threading
 
-from dimos.types.vector import VectorLike, to_vector, Vector
 from dimos.types.path import Path
 from dimos.types.costmap import Costmap
+from dimos.types.vector import VectorLike, to_vector, Vector
 from dimos.robot.global_planner.algo import astar
 from dimos.utils.logging_config import setup_logger
 from dimos.web.websocket_vis.helpers import Visualizable
@@ -29,7 +29,7 @@
 
 @dataclass
 class Planner(Visualizable):
-    local_nav: Callable[[Path, Optional[threading.Event]], bool]
+    set_local_nav: Callable[[Path, Optional[threading.Event]], bool]
 
     @abstractmethod
     def plan(self, goal: VectorLike) -> Path: ...
@@ -40,30 +40,23 @@ def set_goal(self, goal: VectorLike, goal_theta: float, stop_event: Optional[thr
         if not path:
             logger.warning("No path found to the goal.")
             return False
+        return self.set_local_nav(path, stop_event=stop_event, goal_theta = goal_theta)
 
-        return self.local_nav(path, goal_theta = goal_theta, stop_event=stop_event)
-
-
+@dataclass
 class AstarPlanner(Planner):
-    def __init__(
-        self,
-        costmap: Callable[[], Costmap],
-        base_link: Callable[[], Tuple[Vector, Vector]],
-        local_nav: Callable[[Vector], bool],
-    ):
-        super().__init__(local_nav)
-        self.base_link = base_link
-        self.costmap = costmap
+    get_costmap: Callable[[], Costmap]
+    get_robot_pos: Callable[[], Vector]
+    set_local_nav: Callable[[Path], bool]
+    conservativism: int = 3
 
     def plan(self, goal: VectorLike) -> Path:
-        [pos, rot] = self.base_link()
-        costmap = self.costmap()
-        smudge = costmap.smudge(iterations=3)
+        pos = self.get_robot_pos()
+        costmap = self.get_costmap().smudge(iterations=self.conservativism)
 
-        self.vis("planner_costmap", smudge)
+        self.vis("planner_costmap", costmap)
         self.vis("target", goal)
 
-        path = astar(smudge, goal, pos)
+        path = astar(costmap, goal, pos)
 
         if path:
             path = path.resample(0.5)

dimos/robot/unitree/unitree_go2.py

@@ -172,13 +172,14 @@ def __init__(
             robot_length=0.6,  # Unitree Go2 length in meters
             max_linear_vel=0.5,
             lookahead_distance=1.0,
-            visualization_size=500  # 500x500 pixel visualization
+            visualization_size=500,  # 500x500 pixel visualization
         )
 
         self.global_planner = AstarPlanner(
-            costmap=lambda: Costmap.from_msg(self.ros_control.topic_latest("map", msg.OccupancyGrid)()),
-            base_link=lambda: self.ros_control.transform_euler("base_link"),
-            local_nav=self.navigate_path_local
+            conservativism=3,  # how close to obstacles robot is allowed to path plan
+            set_local_nav=self.navigate_path_local,
+            get_costmap=self.ros_control.topic_latest("map", Costmap),
+            get_robot_pos=lambda: self.ros_control.transform_euler_pos("base_link"),
         )
 
         # Create the visualization stream at 5Hz
@@ -268,19 +269,19 @@ def navigate_path_local(self, path: Path, timeout: float = 120.0, goal_theta: Op
                 if self.local_planner.is_goal_reached():
                     logger.info("Path traversed successfully.")
                     path_completed = True
-                    break 
+                    break
 
                 # Get planned velocity towards the current waypoint target
                 vel_command = self.local_planner.plan()
-                x_vel = vel_command.get('x_vel', 0.0)
-                angular_vel = vel_command.get('angular_vel', 0.0)
+                x_vel = vel_command.get("x_vel", 0.0)
+                angular_vel = vel_command.get("angular_vel", 0.0)
 
                 # Send velocity command
                 self.ros_control.move_vel_control(x=x_vel, y=0, yaw=angular_vel)
 
                 # Control loop frequency
                 time.sleep(0.1)
-            
+
             if not path_completed:
                 logger.warning(f"Path following timed out after {timeout} seconds before completing the path.")
 
@@ -293,7 +294,7 @@ def navigate_path_local(self, path: Path, timeout: float = 120.0, goal_theta: Op
         finally:
             logger.info("Stopping robot after path navigation attempt.")
             self.ros_control.stop()
-            
+
         return path_completed
 
     def get_skills(self) -> Optional[SkillLibrary]:

tests/test_websocketvis.py

@@ -1,3 +1,4 @@
+import math
 import os
 import time
 import threading
@@ -11,6 +12,7 @@
 from reactivex import operators as ops
 import argparse
 import pickle
+import reactivex as rx
 
 
 def parse_args():
@@ -29,7 +31,7 @@ def main():
     websocket_vis.start()
 
     if args.live:
-        ros_control = UnitreeROSControl(node_name="simple_nav_test", mock_connection=False)
+        ros_control = UnitreeROSControl(node_name="web_nav_test", mock_connection=False)
         robot = UnitreeGo2(ros_control=ros_control, ip=os.getenv("ROBOT_IP"))
         planner = robot.global_planner
 
@@ -40,9 +42,9 @@ def main():
         pickle_path = f"{__file__.rsplit('/', 1)[0]}/mockdata/costmap.pickle"
         print(f"Loading costmap from {pickle_path}")
         planner = AstarPlanner(
-            costmap=lambda: pickle.load(open(pickle_path, "rb")),
-            base_link=lambda: [Vector(6.0, -1.5), Vector(1, 1, 1)],
-            local_nav=lambda x: time.sleep(1) and True,
+            get_costmap=lambda: pickle.load(open(pickle_path, "rb")),
+            get_robot_pos=lambda: Vector(6.0, -1.5),
+            set_local_nav=lambda x: time.sleep(1) and True,
         )
 
     def msg_handler(msgtype, data):
@@ -63,7 +65,18 @@ def threaded_msg_handler(msgtype, data):
     websocket_vis.msg_handler = threaded_msg_handler
 
     print(f"WebSocket server started on port {websocket_vis.port}")
-    planner.plan(Vector(0, 0))
+    print(planner.get_costmap())
+
+    planner.plan(Vector(0, 0))  # plan a path to the origin
+
+    def fakepos():
+        # Simulate a fake vector position change (to test realtime rendering)
+        vec = Vector(math.sin(time.time()) * 2, math.cos(time.time()) * 2, 0)
+        print(vec)
+        return vec
+
+    if not args.live:
+        websocket_vis.connect(rx.interval(0.05).pipe(ops.map(lambda _: ["fakepos", fakepos()])))
 
     try:
         # Keep the server running