Add RRT*

examples/example_rrt.py

@@ -37,8 +37,10 @@
     options.max_connection_dist = 0.25
     options.goal_biasing_probability = 0.15
     options.max_planning_time = 10.0
-    options.rrt_connect = True
-    options.bidirectional_rrt = True
+    options.rrt_connect = False
+    options.bidirectional_rrt = False
+    options.rrt_star = False
+    options.max_rewire_dist = 3.0
 
     planner = RRTPlanner(model, collision_model)
     path = planner.plan(q_start, q_end, options=options)

src/pyroboplan/core/utils.py

@@ -84,6 +84,26 @@ def configuration_distance(q_start, q_end):
     return np.linalg.norm(q_end - q_start)
 
 
+def get_path_length(q_path):
+    """
+    Returns the configuration distance of a path.
+
+    Parameters
+    ----------
+        q_path : list[array-like]
+            A list of joint configurations describing a path.
+
+    Returns
+    -------
+        float
+            The total configuration distance of the entire path.
+    """
+    total_distance = 0.0
+    for idx in range(1, len(q_path)):
+        total_distance += configuration_distance(q_path[idx - 1], q_path[idx])
+    return total_distance
+
+
 def get_random_state(model, padding=0.0):
     """
     Returns a random state that is within the model's joint limits.

src/pyroboplan/planning/graph.py

@@ -90,6 +90,27 @@ def add_edge(self, nodeA, nodeB):
         nodeB.neighbors.add(nodeA)
         return edge
 
+    def remove_edge(self, edge):
+        """
+        Attempts to remove an edge from a graph, if it exists.
+
+        Parameters
+        ----------
+            edge : `pyroboplan.planning.graph.Edge`
+                The edge to remove from the graph.
+
+        Returns
+        -------
+            bool
+                True if the edge was successfully removed, else False.
+
+        """
+        if edge not in self.edges:
+            return False
+
+        self.edges.remove(edge)
+        return True
+
     def get_nearest_node(self, q):
         """
         Gets the nearest node to a specified robot configuration.

src/pyroboplan/planning/rrt.py

@@ -34,6 +34,18 @@ class RRTPlannerOptions:
     Otherwise, only grows a tree from the start node.
     """
 
+    rrt_star = False
+    """
+    If true, enables the RRT* algorithm.
+    This in turn will use the `max_rewire_dist` parameter.
+    """
+
+    max_rewire_dist = np.inf
+    """
+    Maximum angular distance, in radians, to consider rewiring nodes for RRT*.
+    If set to `np.inf`, all nodes in the trees will be considering for rewiring.
+    """
+
     max_planning_time = 10.0
     """ Maximum planning time, in seconds. """
 
@@ -49,6 +61,7 @@ class RRTPlanner:
     Some good resources:
       * Original RRT paper: https://msl.cs.illinois.edu/~lavalle/papers/Lav98c.pdf
       * RRTConnect paper: https://www.cs.cmu.edu/afs/cs/academic/class/15494-s14/readings/kuffner_icra2000.pdf
+      * RRT* and PRM* paper: https://arxiv.org/abs/1105.1186
     """
 
     def __init__(self, model, collision_model):
@@ -89,9 +102,9 @@ def plan(self, q_start, q_goal, options=RRTPlannerOptions()):
         t_start = time.time()
         self.options = options
 
-        start_node = Node(q_start, parent=None)
+        start_node = Node(q_start, parent=None, cost=0.0)
         self.start_tree.add_node(start_node)
-        goal_node = Node(q_goal, parent=None)
+        goal_node = Node(q_goal, parent=None, cost=0.0)
         self.goal_tree.add_node(goal_node)
 
         goal_found = False
@@ -139,27 +152,32 @@ def plan(self, q_start, q_goal, options=RRTPlannerOptions()):
 
             # Only if extend/connect succeeded, add the new node to the tree.
             if q_new is not None:
-                new_node = Node(q_new, parent=nearest_node)
-                tree.add_node(new_node)
-                tree.add_edge(nearest_node, new_node)
+                new_node = self.add_node_to_tree(tree, q_new, nearest_node, options)
                 if start_tree_phase:
                     latest_start_tree_node = new_node
                 else:
                     latest_goal_tree_node = new_node
 
                 # Check if latest node connects directly to the other tree.
-                # If so, planning is complete.
+                # If so, add it to the tree and mark planning as complete.
                 nearest_node_in_other_tree = other_tree.get_nearest_node(new_node.q)
                 path_to_other_tree = discretize_joint_space_path(
-                    q_new, nearest_node_in_other_tree.q, self.options.max_angle_step
+                    new_node.q,
+                    nearest_node_in_other_tree.q,
+                    self.options.max_angle_step,
                 )
                 if not check_collisions_along_path(
                     self.model, self.collision_model, path_to_other_tree
                 ):
+                    new_node = self.add_node_to_tree(
+                        tree, nearest_node_in_other_tree.q, new_node, options
+                    )
                     if start_tree_phase:
+                        latest_start_tree_node = new_node
                         latest_goal_tree_node = nearest_node_in_other_tree
                     else:
                         latest_start_tree_node = nearest_node_in_other_tree
+                        latest_goal_tree_node = new_node
                     goal_found = True
 
                 # Switch to the other tree next iteration, if bidirectional mode is enabled.
@@ -264,6 +282,44 @@ def extract_path_from_trees(self, start_tree_final_node, goal_tree_final_node):
 
         return path
 
+    def add_node_to_tree(self, tree, q_new, parent_node, options=RRTPlannerOptions()):
+        """
+        Add a new node to the tree. If the RRT* algorithm is enabled, will also rewire.
+        """
+        # Add the new node to the tree
+        new_node = Node(q_new, parent=parent_node)
+        tree.add_node(new_node)
+        edge = tree.add_edge(parent_node, new_node)
+        new_node.cost = parent_node.cost + edge.cost
+
+        # If RRT* is enable it, rewire that node in the tree.
+        if options.rrt_star:
+            min_cost = new_node.cost
+            for other_node in tree.nodes:
+                # Do not consider trivial nodes.
+                if other_node == new_node or other_node == parent_node:
+                    continue
+                # Do not consider nodes farther than the configured rewire distance,
+                new_distance = configuration_distance(other_node.q, q_new)
+                if new_distance > options.max_rewire_dist:
+                    continue
+                # Rewire if this new connections would be of lower cost and is collision free.
+                new_cost = other_node.cost + new_distance
+                if new_cost < min_cost:
+                    new_path = discretize_joint_space_path(
+                        q_new, other_node.q, options.max_angle_step
+                    )
+                    if not check_collisions_along_path(
+                        self.model, self.collision_model, new_path
+                    ):
+                        new_node.parent = other_node
+                        new_node.cost = new_cost
+                        tree.remove_edge(edge)
+                        edge = tree.add_edge(other_node, new_node)
+                        min_cost = new_cost
+
+        return new_node
+
     def visualize(
         self,
         visualizer,

test/core/test_utils.py

@@ -8,6 +8,7 @@
     check_within_limits,
     configuration_distance,
     extract_cartesian_poses,
+    get_path_length,
     get_random_state,
     get_random_transform,
 )
@@ -24,6 +25,16 @@ def test_configuration_distance():
     assert configuration_distance(q_start, q_end) == pytest.approx(0.5)
 
 
+def test_get_path_length():
+    q_path = [
+        np.array([0.0, 0.0, 0.0, 0.0, 0.0]),
+        np.array([0.3, 0.0, -0.4, 0.0, 0.0]),
+        np.array([0.3, 1.0, -0.4, 0.0, 0.0]),
+        np.array([0.3, 1.0, -0.4, 0.0, 10.0]),
+    ]
+    assert get_path_length(q_path) == pytest.approx(11.5)
+
+
 def test_get_random_state():
     model, _, _ = load_models()
     for _ in range(10):

test/planning/test_graph.py

@@ -90,6 +90,24 @@ def test_add_edges():
     assert edgeAC.cost == pytest.approx(np.linalg.norm([6.0, 6.0, 6.0]))
 
 
+def test_remove_edges():
+    nodeA = Node([1.0, 2.0, 3.0])
+    nodeB = Node([4.0, 5.0, 6.0])
+    graph = Graph()
+    graph.add_node(nodeA)
+    graph.add_node(nodeB)
+    edgeAB = graph.add_edge(nodeA, nodeB)
+    assert len(graph.edges) == 1
+
+    # Remove a valid edge.
+    assert graph.remove_edge(edgeAB)
+    assert len(graph.edges) == 0
+
+    # Remove an edge that was already removed.
+    assert not graph.remove_edge(edgeAB)
+    assert len(graph.edges) == 0
+
+
 def test_get_nearest_node():
     nodeA = Node([1.0, 2.0, 3.0])
     nodeB = Node([4.0, 5.0, 6.0])

test/planning/test_rrt.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from pyroboplan.core.utils import get_path_length
 from pyroboplan.planning.rrt import RRTPlanner, RRTPlannerOptions
 from pyroboplan.models.panda import (
     load_models,
@@ -41,7 +42,7 @@ def test_plan_vanilla_rrt():
     planner = RRTPlanner(model, collision_model)
     path = planner.plan(q_start, q_goal)
 
-    # The path must exist and have more than
+    # The path must exist and have more than the start and goal nodes.
     assert path is not None
     assert len(path) > 2
     assert np.all(path[0] == q_start)
@@ -59,7 +60,7 @@ def test_plan_rrt_connect():
     add_self_collisions(collision_model)
     add_object_collisions(collision_model, visual_model)
 
-    # Plan with RRTConnect and bidirectional RRT
+    # Plan with RRTConnect and bidirectional RRT.
     q_start = np.array([0.0, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0])
     q_goal = np.array([1.57, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0])
 
@@ -70,7 +71,7 @@ def test_plan_rrt_connect():
     planner = RRTPlanner(model, collision_model)
     path = planner.plan(q_start, q_goal, options=options)
 
-    # The path must exist and have more than
+    # The path must exist and have more than the start and goal nodes.
     assert path is not None
     assert len(path) > 2
     assert np.all(path[0] == q_start)
@@ -81,3 +82,29 @@ def test_plan_rrt_connect():
     assert len(planner.start_tree.edges) > 1
     assert len(planner.goal_tree.nodes) > 1
     assert len(planner.goal_tree.edges) > 1
+
+
+def test_plan_rrt_star():
+    model, collision_model, visual_model = load_models()
+    add_self_collisions(collision_model)
+    add_object_collisions(collision_model, visual_model)
+
+    q_start = np.array([0.0, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0])
+    q_goal = np.array([1.57, 1.57, 0.0, 0.0, 1.57, 1.57, 0.0, 0.0, 0.0])
+
+    # First, plan with regular RRT.
+    options = RRTPlannerOptions()
+    options.rrt_star = False
+    options.bidirectional_rrt = True
+    planner = RRTPlanner(model, collision_model)
+    path_original = planner.plan(q_start, q_goal, options=options)
+
+    # Then, plan with RRT* enabled (use maximum rewire distance for effect).
+    options.rrt_star = True
+    options.max_rewire_dist = np.inf
+    path_star = planner.plan(q_start, q_goal, options=options)
+
+    # The RRT* path must be shorter or of equal length, though it also took longer to plan.
+    assert path_original is not None
+    assert path_star is not None
+    assert get_path_length(path_star) <= get_path_length(path_original)