add informed rrt*

Sampling-based Planning/rrt_2D/informed_rrt_star.py

@@ -9,6 +9,7 @@ import math
 import random
 import numpy as np
 import matplotlib.pyplot as plt
+from scipy.spatial.transform import Rotation as Rot
 import matplotlib.patches as patches
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) +
@@ -17,7 +18,6 @@ sys.path.append(os.path.dirname(os.path.abspath(__file__)) +
 from rrt_2D import env
 from rrt_2D import plotting
 from rrt_2D import utils
-from rrt_2D import queue
 
 
 class Node:
@@ -41,7 +41,8 @@ class IRrtStar:
         self.plotting = plotting.Plotting(x_start, x_goal)
         self.utils = utils.Utils()
 
-        # self.fig, self.ax = plt.subplots()
+        self.fig, self.ax = plt.subplots()
+        self.delta = self.utils.delta
         self.x_range = self.env.x_range
         self.y_range = self.env.y_range
         self.obs_circle = self.env.obs_circle
@@ -54,25 +55,15 @@ class IRrtStar:
 
     def planning(self):
         c_best = np.inf
-        c_min = self.Line(self.x_start, self.x_goal)
+        c_min, theta = self.get_distance_and_angle(self.x_start, self.x_goal)
+        C = self.RotationToWorldFrame(self.x_start, self.x_goal, c_min)
         x_center = np.array([[(self.x_start.x + self.x_goal.x) / 2.0],
                              [(self.x_start.y + self.x_goal.y) / 2.0], [0.0]])
-        a1 = np.array([[(self.x_goal.x - self.x_start.x) / c_min],
-                       [(self.x_goal.y - self.x_start.y) / c_min], [0.0]])
-        e_theta = math.atan2(a1[1], a1[0])
-        id1_t = np.array([[1.0, 0.0, 0.0]])
-        M = a1 @ id1_t
-        U, S, Vh = np.linalg.svd(M, True, True)
-        C = np.dot(np.dot(U, np.diag(
-            [1.0, 1.0, np.linalg.det(U) * np.linalg.det(np.transpose(Vh))])), Vh)
 
         for k in range(self.iter_max):
-            if k % 500 == 0:
+            if k % 50 == 0:
                 print(k)
 
-            if self.X_soln:
-                c_best = min([self.Cost(x) for x in self.X_soln])
-
             x_rand = self.Sample(self.x_start, self.x_goal, c_best, x_center, C)
             x_nearest = self.Nearest(self.V, x_rand)
             x_new = self.Steer(x_nearest, x_rand)
@@ -80,8 +71,7 @@ class IRrtStar:
             if x_new and not self.utils.is_collision(x_nearest, x_new):
                 self.V.append(x_new)
                 X_near = self.Near(self.V, x_new)
-                x_min = x_nearest
-                c_min = self.Cost(x_min) + self.Line(x_nearest, x_new)
+                c_min = self.Cost(x_new)
 
                 for x_near in X_near:
                     c_new = self.Cost(x_near) + self.Line(x_near, x_new)
@@ -97,9 +87,96 @@ class IRrtStar:
 
                 if self.InGoalRegion(x_new):
                     self.X_soln.add(x_new)
+                    c_best = min(c_best, self.Cost(x_new) + self.Line(x_new, self.x_goal))
+
+            self.animation(x_center=x_center, c_best=c_best, c_min=c_min, theta=theta)
+
+        self.animation(x_center=x_center, c_best=c_best, c_min=c_min, theta=theta)
+        # plt.plot([x for x, _ in self.path], [y for _, y in self.path], '-r')
+        plt.pause(0.01)
+        plt.show()
+
+    def animation(self, x_center=None, c_best=None, c_min=None, theta=None):
+        plt.cla()
+        self.plot_grid("Informed rrt*")
+        plt.gcf().canvas.mpl_connect(
+            'key_release_event',
+            lambda event: [exit(0) if event.key == 'escape' else None])
+
+        if c_best != np.inf:
+            self.draw_ellipse(x_center, c_best, c_min, theta)
+
+        for node in self.V:
+            if node.parent:
+                plt.plot([node.x, node.parent.x], [node.y, node.parent.y], "-g")
+
+        self.plot_grid("Informed rrt*")
+        plt.pause(0.01)
+
+    def plot_grid(self, name):
+
+        for (ox, oy, w, h) in self.obs_boundary:
+            self.ax.add_patch(
+                patches.Rectangle(
+                    (ox, oy), w, h,
+                    edgecolor='black',
+                    facecolor='black',
+                    fill=True
+                )
+            )
+
+        for (ox, oy, w, h) in self.obs_rectangle:
+            self.ax.add_patch(
+                patches.Rectangle(
+                    (ox, oy), w, h,
+                    edgecolor='black',
+                    facecolor='gray',
+                    fill=True
+                )
+            )
+
+        for (ox, oy, r) in self.obs_circle:
+            self.ax.add_patch(
+                patches.Circle(
+                    (ox, oy), r,
+                    edgecolor='black',
+                    facecolor='gray',
+                    fill=True
+                )
+            )
+
+        plt.plot(self.x_start.x, self.x_start.y, "bs", linewidth=3)
+        plt.plot(self.x_goal.x, self.x_goal.y, "gs", linewidth=3)
+
+        plt.title(name)
+        plt.axis("equal")
 
-        path = self.ExtractPath(self.V[-1])
-        self.plotting.animation(self.V, path, "Informed rrt*")
+    @staticmethod
+    def draw_ellipse(x_center, c_best, c_min, theta):
+        a = math.sqrt(c_best ** 2 - c_min ** 2) / 2.0
+        b = c_best / 2.0
+        angle = math.pi / 2.0 - theta
+        cx = x_center[0]
+        cy = x_center[1]
+        t = np.arange(0, 2 * math.pi + 0.1, 0.1)
+        x = [a * math.cos(it) for it in t]
+        y = [b * math.sin(it) for it in t]
+        rot = Rot.from_euler('z', -angle).as_dcm()[0:2, 0:2]
+        fx = rot @ np.array([x, y])
+        px = np.array(fx[0, :] + cx).flatten()
+        py = np.array(fx[1, :] + cy).flatten()
+        plt.plot(cx, cy, "xc")
+        plt.plot(px, py, "--c")
+
+    def RotationToWorldFrame(self, x_start, x_goal, L):
+        a1 = np.array([[(self.x_goal.x - self.x_start.x) / L],
+                       [(self.x_goal.y - self.x_start.y) / L], [0.0]])
+        e1 = np.array([[1.0], [0.0], [0.0]])
+        M = a1 @ e1.T
+        U, _, V_T = np.linalg.svd(M, True, True)
+        C = U @ np.diag([1.0, 1.0, np.linalg.det(U) * np.linalg.det(V_T.T)]) @ V_T
+
+        return C
 
     def ExtractPath(self, node):
         path = [[self.x_goal.x, self.x_goal.y]]
@@ -129,10 +206,10 @@ class IRrtStar:
 
     def Near(self, nodelist, node):
         n = len(nodelist) + 1
-        r = 2 * self.search_radius * math.sqrt((math.log(n) / n))
+        r = 50 * math.sqrt((math.log(n) / n))
 
-        dist_table = [math.hypot(nd.x - node.x, nd.y - node.y) for nd in nodelist]
-        X_near = [nodelist[ind] for ind in range(len(dist_table)) if dist_table[ind] <= r and
+        dist_table = [(nd.x - node.x) ** 2 + (nd.y - node.y) ** 2 for nd in nodelist]
+        X_near = [nodelist[ind] for ind in range(len(dist_table)) if dist_table[ind] <= r ** 2 and
                   not self.utils.is_collision(node, nodelist[ind])]
 
         return X_near
@@ -145,7 +222,12 @@ class IRrtStar:
                  math.sqrt(c_max ** 2 - c_min ** 2) / 2.0]
             L = np.diag(r)
             x_ball = self.SampleUnitNBall()
-            x_rand = np.dot(np.dot(C, L), x_ball) + x_center
+            # x_rand = C @ L @ x_ball + x_center
+            while True:
+                x_rand = C @ L @ x_ball + x_center
+                if self.x_range[0] + self.delta <= x_rand[0] <= self.x_range[1] + self.delta:
+                    break
+            x_rand = Node((x_rand[0], x_rand[1]))
         else:
             x_rand = self.SampleFreeSpace()
 
@@ -155,26 +237,22 @@ class IRrtStar:
         delta = self.utils.delta
 
         if np.random.random() > self.goal_sample_rate:
-            return [np.random.uniform(self.x_range[0] + delta, self.x_range[1] - delta),
-                    np.random.uniform(self.y_range[0] + delta, self.y_range[1] - delta)]
+            return Node((np.random.uniform(self.x_range[0] + delta, self.x_range[1] - delta),
+                         np.random.uniform(self.y_range[0] + delta, self.y_range[1] - delta)))
 
-        return [self.x_goal.x, self.x_goal.y]
+        return self.x_goal
 
     @staticmethod
     def SampleUnitNBall():
-        x, y = random.random(), random.random()
-
-        if y < x:
-            x, y = y, x
-
-        sample = np.array([[y * math.cos(2 * math.pi * x / y)],
-                           [y * math.sin(2 * math.pi * x / y)], [0.0]])
+        theta, r = random.uniform(0.0, 2*math.pi), random.random()
+        x = r * math.cos(theta)
+        y = r * math.sin(theta)
 
-        return sample
+        return np.array([[x], [y], [0.0]])
 
     @staticmethod
     def Nearest(nodelist, n):
-        return nodelist[int(np.argmin([math.hypot(nd.x - n.x, nd.y - n.y)
+        return nodelist[int(np.argmin([(nd.x - n.x) ** 2 + (nd.y - n.y) ** 2
                                        for nd in nodelist]))]
 
     @staticmethod
@@ -203,7 +281,7 @@ def main():
     x_start = (2, 2)  # Starting node
     x_goal = (49, 24)  # Goal node
 
-    rrt_star = IRrtStar(x_start, x_goal, 10, 0.10, 20, 4000)
+    rrt_star = IRrtStar(x_start, x_goal, 5.0, 0.10, 20, 200)
     rrt_star.planning()