rrt & rrt*

Author: omigeft

src/python_motion_planning/common/env/map/grid.py

@@ -5,7 +5,7 @@
 @update: 2025.9.5
 """
 from itertools import product
-from typing import Iterable, Union, Tuple, Callable, List
+from typing import Iterable, Union, Tuple, Callable, List, Dict
 import time
 
 import numpy as np
@@ -235,7 +235,7 @@ def shape(self) -> tuple:
     def esdf(self) -> np.ndarray:
         return self._esdf
 
-    def map_to_world(self, point: tuple) -> tuple:
+    def map_to_world(self, point: Tuple[int, ...]) -> tuple:
         """
         Convert map coordinates to world coordinates.
         
@@ -250,7 +250,7 @@ def map_to_world(self, point: tuple) -> tuple:
 
         return tuple((x + 0.5) * self.resolution + float(self.bounds[i, 0]) for i, x in enumerate(point))
 
-    def world_to_map(self, point: tuple) -> tuple:
+    def world_to_map(self, point: Tuple[float, ...]) -> tuple:
         """
         Convert world coordinates to map coordinates.
         
@@ -265,7 +265,7 @@ def world_to_map(self, point: tuple) -> tuple:
 
         return tuple(round((x - float(self.bounds[i, 0])) * (1.0 / self.resolution) - 0.5) for i, x in enumerate(point))
 
-    def get_distance(self, p1: tuple, p2: tuple) -> float:
+    def get_distance(self, p1: Tuple[int, int], p2: Tuple[int, int]) -> float:
         """
         Get the distance between two points.
 
@@ -278,7 +278,7 @@ def get_distance(self, p1: tuple, p2: tuple) -> float:
         """
         return Geometry.dist(p1, p2, type='Euclidean')
 
-    def within_bounds(self, point: tuple) -> bool:
+    def within_bounds(self, point: Tuple[int, ...]) -> bool:
         """
         Check if a point is within the bounds of the grid map.
         
@@ -300,7 +300,7 @@ def within_bounds(self, point: tuple) -> bool:
                 return False
         return True
 
-    def is_expandable(self, point: tuple, src_point: tuple = None) -> bool:
+    def is_expandable(self, point: Tuple[int, ...], src_point: Tuple[int, ...] = None) -> bool:
         """
         Check if a point is expandable.
         
@@ -363,7 +363,7 @@ def get_neighbors(self,
 
         return filtered_neighbors
 
-    def line_of_sight(self, p1: tuple, p2: tuple) -> list:
+    def line_of_sight(self, p1: Tuple[int, ...], p2: Tuple[int, ...]) -> list:
         """
         N-dimensional line of sight (Bresenham's line algorithm)
         
@@ -414,7 +414,7 @@ def line_of_sight(self, p1: tuple, p2: tuple) -> list:
 
         return result
 
-    def in_collision(self, p1: tuple, p2: tuple) -> bool:
+    def in_collision(self, p1: Tuple[int, ...], p2: Tuple[int, ...]) -> bool:
         """
         Check if the line of sight between two points is in collision.
         
@@ -498,17 +498,17 @@ def inflate_obstacles(self, radius: float = 1.0) -> None:
                     self.type_map[i, j] = TYPES.INFLATION
         self.inflation_radius = radius
 
-    def fill_expands(self, expands: List[Node]) -> None:
+    def fill_expands(self, expands: Dict[Tuple[int, int], Node]) -> None:
         """
         Fill the expands in the map.
         
         Args:
             expands: List of expands.
         """
-        for expand in expands:
-            if self.type_map[expand.current] != TYPES.FREE:
+        for expand in expands.keys():
+            if self.type_map[expand] != TYPES.FREE:
                 continue
-            self.type_map[expand.current] = TYPES.EXPAND
+            self.type_map[expand] = TYPES.EXPAND
 
     def update_esdf(self) -> None:
         """
@@ -527,6 +527,30 @@ def update_esdf(self) -> None:
         self._esdf = dist_outside.astype(np.float32)
         self._esdf[obstacle_mask] = -dist_inside[obstacle_mask]
 
+    def path_map_to_world(self, path: List[Tuple[int, int]]) -> List[Tuple[float, float]]:
+        """
+        Convert path from map coordinates to world coordinates
+
+        Args:
+            path: a list of map coordinates
+        
+        Returns:
+            path: a list of world coordinates
+        """
+        return [self.map_to_world(p) for p in path]
+
+    def path_world_to_map(self, path: List[Tuple[float, float]]) -> List[Tuple[int, int]]:
+        """
+        Convert path from world coordinates to map coordinates
+
+        Args:
+            path: a list of world coordinates
+        
+        Returns:
+            path: a list of map coordinates
+        """
+        return [self.world_to_map(p) for p in path]
+
     def _precompute_offsets(self):
         # Generate all possible offsets (-1, 0, +1) in each dimension
         self._diagonal_offsets = np.array(np.meshgrid(*[[-1, 0, 1]]*self.dim), dtype=self.dtype).T.reshape(-1, self.dim)

src/python_motion_planning/common/visualizer/visualizer.py

@@ -4,7 +4,7 @@
 @author: Yang Haodong, Wu Maojia
 @update: 2025.9.20
 """
-from typing import Union, Dict
+from typing import Union, Dict, List, Tuple
 from collections import namedtuple
 import time
 
@@ -16,7 +16,7 @@
 import matplotlib.patheffects as path_effects
 
 from python_motion_planning.controller import BaseController
-from python_motion_planning.common.env import TYPES, ToySimulator, Grid, CircularRobot
+from python_motion_planning.common.env import TYPES, ToySimulator, Grid, CircularRobot, Node
 
 class Visualizer:
     def __init__(self, fig_name: str = ""):
@@ -39,6 +39,10 @@ def __init__(self, fig_name: str = ""):
         self.cmap = mcolors.ListedColormap([info for info in self.cmap_dict.values()])
         self.norm = mcolors.BoundaryNorm([i for i in range(self.cmap.N + 1)], self.cmap.N)
         self.grid_map = None
+        self.dim = None
+
+    def __del__(self):
+        self.close()
 
     def plot_grid_map(self, grid_map: Grid, equal: bool = False, alpha_3d: float = 0.1,
                         show_esdf: bool = False, alpha_esdf: float = 0.5) -> None:
@@ -52,6 +56,8 @@ def plot_grid_map(self, grid_map: Grid, equal: bool = False, alpha_3d: float = 0
             show_esdf: Whether to show esdf.
             alpha_esdf: Alpha of esdf.
         '''
+        self.grid_map = grid_map
+        self.dim = grid_map.dim
         if grid_map.dim == 2:
             plt.imshow(
                 np.transpose(grid_map.type_map.array), 
@@ -117,7 +123,9 @@ def plot_grid_map(self, grid_map: Grid, equal: bool = False, alpha_3d: float = 0
         else:
             raise NotImplementedError(f"Grid map with dim={grid_map.dim} not supported.")
 
-    def plot_path(self, path: list, style: str = "-", color: str = "#13ae00", label: str = None, linewidth: float = 2, marker: str = None) -> None:
+    def plot_path(self, path: List[Union[Tuple[int, ...], Tuple[float, ...]]], 
+                    style: str = "-", color: str = "#13ae00", label: str = None, 
+                    linewidth: float = 2, marker: str = None, map_frame: bool = True) -> None:
         '''
         Plot path-like information.
         The meaning of parameters are similar to matplotlib.pyplot.plot (https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.plot.html).
@@ -129,20 +137,80 @@ def plot_path(self, path: list, style: str = "-", color: str = "#13ae00", label:
             label: label of path
             linewidth: linewidth of path
             marker: marker of path
+            map_frame: whether path is in map frame or not (world frame)
         '''
+        if map_frame:
+            path = [self.grid_map.map_to_world(point) for point in path]
+
         path = np.array(path)
-        if len(path.shape) < 2:
-            return
-        if path.shape[1] == 2:
-            plt.plot(path[:, 0], path[:, 1], style, lw=linewidth, color=color, label=label, marker=marker)
-        elif path.shape[1] == 3:
+        
+        if self.dim == 2:
+            self.ax.plot(path[:, 0], path[:, 1], style, lw=linewidth, color=color, label=label, marker=marker)
+        elif self.dim == 3:
             self.ax.plot(path[:, 0], path[:, 1], path[:, 2], style, lw=linewidth, color=color, label=label, marker=marker)
         else:
-            raise ValueError("Path dimension not supported")
+            raise ValueError("Dimension not supported")
 
         if label:
             self.ax.legend()
 
+    def plot_expand_tree(self, expand_tree: Dict[Union[Tuple[int, ...], Tuple[float, ...]], Node], 
+                        node_color: str = "C5", 
+                        edge_color: str = "C6", 
+                        node_size: float = 10, 
+                        linewidth: float = 1.0, 
+                        node_alpha: float = 1.0,
+                        edge_alpha: float = 1.0,
+                        connect_to_parent: bool = True,
+                        map_frame: bool = True) -> None:
+        """
+        Visualize an expand tree (e.g. RRT).
+        
+        Args:
+            expand_tree: Dict mapping coordinate tuple -> Node (world frame).
+            node_color: Color of the nodes.
+            edge_color: Color of the edges (parent -> child).
+            node_size: Size of node markers.
+            linewidth: Line width of edges.
+            connect_to_parent: Whether to draw parent-child connections.
+            map_frame: whether path is in map frame or not (world frame)
+        """
+        if self.dim == 2:
+            for coord, node in expand_tree.items():
+                current = node.current
+                if map_frame:
+                    current = self.grid_map.map_to_world(current)
+
+                self.ax.scatter(current[0], current[1],
+                                c=node_color, s=node_size, zorder=3, alpha=node_alpha)
+                if connect_to_parent and node.parent is not None:
+                    parent = node.parent
+                    if map_frame:
+                        parent = self.grid_map.map_to_world(parent)
+                    self.ax.plot([parent[0], current[0]],
+                                [parent[1], current[1]],
+                                color=edge_color, linewidth=linewidth, zorder=2, alpha=edge_alpha)
+
+        elif self.dim == 3:
+            for coord, node in expand_tree.items():
+                current = node.current
+                if map_frame:
+                    current = self.grid_map.map_to_world(current)
+
+                self.ax.scatter(current[0], current[1], current[2],
+                                c=node_color, s=node_size, zorder=3, alpha=node_alpha)
+                if connect_to_parent and node.parent is not None:
+                    parent = node.parent
+                    if map_frame:
+                        parent = self.grid_map.map_to_world(parent)
+                    self.ax.plot([parent[0], current[0]],
+                                [parent[1], current[1]],
+                                [parent[2], current[2]],
+                                color=edge_color, linewidth=linewidth, zorder=2, alpha=edge_alpha)
+
+        else:
+            raise ValueError("Dimension must be 2 or 3")
+
     def plot_circular_robot(self, robot: CircularRobot, axis_equal: bool = True) -> None:
         patch = plt.Circle(tuple(robot.pos), robot.radius, 
             color=robot.color, alpha=robot.alpha, fill=robot.fill, linewidth=robot.linewidth, linestyle=robot.linestyle)
@@ -165,7 +233,8 @@ def plot_circular_robot(self, robot: CircularRobot, axis_equal: bool = True) ->
         else:
             return patch, text
 
-    def render_toy_simulator(self, env: ToySimulator, controllers: Dict[str, BaseController], steps: int = 1000, interval: int = 50,
+    def render_toy_simulator(self, env: ToySimulator, controllers: Dict[str, BaseController],
+            steps: int = 1000, interval: int = 50,
             show_traj: bool = True, traj_kwargs: dict = {"linestyle": '-', "alpha": 0.7, "linewidth": 1.5},
             show_env_info: bool = False, rtf_limit: float = 1.0, grid_kwargs: dict = {},
             show_pred_traj: bool = True) -> None:
@@ -281,6 +350,9 @@ def show(self):
 
     def legend(self):
         plt.legend()
+    
+    def close(self):
+        plt.close()