add BIT*

Sampling_based_Planning/rrt_2D/batch_informed_trees.py

@@ -31,7 +31,7 @@ class Tree:
         self.x_start = x_start
         self.goal = x_goal
 
-        self.r = np.inf
+        self.r = 4.0
         self.V = set()
         self.E = set()
         self.QE = set()
@@ -81,16 +81,25 @@ class BITStar:
         return theta, cMin, xCenter, C
 
     def planning(self):
-        m = 200
         theta, cMin, xCenter, C = self.init()
 
-        for k in range(self.iter_max):
+        for k in range(500):
             if not self.Tree.QE and not self.Tree.QV:
+                if k == 0:
+                    m = 350
+                else:
+                    m = 200
+
+                if self.x_goal.parent is not None:
+                    path_x, path_y = self.ExtractPath()
+                    plt.plot(path_x, path_y, linewidth=2, color='r')
+                    plt.pause(0.5)
+
                 self.Prune(self.g_T[self.x_goal])
                 self.X_sample.update(self.Sample(m, self.g_T[self.x_goal], cMin, xCenter, C))
-                self.Tree.V_old = copy.deepcopy(self.Tree.V)
-                self.Tree.QV = copy.deepcopy(self.Tree.V)
-                self.Tree.r = self.radius(len(self.Tree.V) + len(self.X_sample))
+                self.Tree.V_old = {v for v in self.Tree.V}
+                self.Tree.QV = {v for v in self.Tree.V}
+                # self.Tree.r = self.radius(len(self.Tree.V) + len(self.X_sample))
 
             while self.BestVertexQueueValue() <= self.BestEdgeQueueValue():
                 self.ExpandVertex(self.BestInVertexQueue())
@@ -104,23 +113,72 @@ class BITStar:
                     if self.g_T[vm] + actual_cost < self.g_T[xm]:
                         if xm in self.Tree.V:
                             # remove edges
+                            edge_delete = set()
                             for v, x in self.Tree.E:
                                 if x == xm:
-                                    self.Tree.E.remove((v, x))
+                                    edge_delete.add((v, x))
+
+                            for edge in edge_delete:
+                                self.Tree.E.remove(edge)
                         else:
                             self.X_sample.remove(xm)
                             self.Tree.V.add(xm)
                             self.Tree.QV.add(xm)
 
+                        self.g_T[xm] = self.g_T[vm] + actual_cost
                         self.Tree.E.add((vm, xm))
+                        xm.parent = vm
 
+                        set_delete = set()
                         for v, x in self.Tree.QE:
                             if x == xm and self.g_T[v] + self.calc_dist(v, xm) >= self.g_T[xm]:
-                                self.Tree.QE.remove((v, xm))
+                                set_delete.add((v, x))
+
+                        for edge in set_delete:
+                            self.Tree.QE.remove(edge)
             else:
                 self.Tree.QE = set()
                 self.Tree.QV = set()
 
+            if k % 5 == 0:
+                self.draw(xCenter, self.g_T[self.x_goal], cMin, theta)
+
+        path_x, path_y = self.ExtractPath()
+        plt.plot(path_x, path_y, linewidth=2, color='r')
+        plt.pause(0.01)
+        # test
+        plt.show()
+
+    def draw(self, xCenter, cMax, cMin, theta):
+        plt.cla()
+        self.plot_grid("Batch Informed Trees (BIT*)")
+
+        plt.gcf().canvas.mpl_connect(
+            'key_release_event',
+            lambda event: [exit(0) if event.key == 'escape' else None])
+
+        for v in self.X_sample:
+            plt.plot(v.x, v.y, marker='.', color='lightgrey', markersize='2')
+
+        if cMax < np.inf:
+            self.draw_ellipse(xCenter, cMax, cMin, theta)
+
+        for v, w in self.Tree.E:
+            plt.plot([v.x, w.x], [v.y, w.y], '-g')
+
+        plt.pause(0.01)
+
+    def ExtractPath(self):
+        node = self.x_goal
+        path_x, path_y = [node.x], [node.y]
+
+        while node.parent:
+            node = node.parent
+            path_x.append(node.x)
+            path_y.append(node.y)
+
+        return path_x, path_y
+
     def Prune(self, cBest):
         self.X_sample = {x for x in self.X_sample if self.f_estimated(x) < cBest}
         self.Tree.V = {v for v in self.Tree.V if self.f_estimated(v) <= cBest}
@@ -200,16 +258,22 @@ class BITStar:
     def ExpandVertex(self, v):
         self.Tree.QV.remove(v)
         X_near = {x for x in self.X_sample if self.calc_dist(x, v) <= self.Tree.r}
-        edges_add = {(v, x) for x in X_near
-                     if self.g_estimated(v) + self.calc_dist(v, x) + self.h_estimated(x) < self.g_T[self.x_goal]}
-        self.Tree.QE.update(edges_add)
+
+        for x in X_near:
+            if self.g_estimated(v) + self.calc_dist(v, x) + self.h_estimated(x) < self.g_T[self.x_goal]:
+                self.g_T[x] = np.inf
+                self.Tree.QE.add((v, x))
 
         if v not in self.Tree.V_old:
             V_near = {w for w in self.Tree.V if self.calc_dist(w, v) <= self.Tree.r}
-            edges_add = {(v, w) for w in V_near if (v, w) not in self.Tree.E and
-                         self.g_estimated(v) + self.calc_dist(v, w) + self.h_estimated(w) < self.g_T[self.x_goal] and
-                         self.g_T[v] + self.calc_dist(v, w) < self.g_T[w]}
-            self.Tree.QE.update(edges_add)
+
+            for w in V_near:
+                if (v, w) not in self.Tree.E and \
+                        self.g_estimated(v) + self.calc_dist(v, w) + self.h_estimated(w) < self.g_T[self.x_goal] and \
+                        self.g_T[v] + self.calc_dist(v, w) < self.g_T[w]:
+                    self.Tree.QE.add((v, w))
+                    if w not in self.g_T:
+                        self.g_T[w] = np.inf
 
     def BestVertexQueueValue(self):
         if not self.Tree.QV:
@@ -225,11 +289,19 @@ class BITStar:
                    for v, x in self.Tree.QE)
 
     def BestInVertexQueue(self):
+        if not self.Tree.QV:
+            print("QV is Empty!")
+            return None
+
         v_value = {v: self.g_T[v] + self.h_estimated(v) for v in self.Tree.QV}
 
         return min(v_value, key=v_value.get)
 
     def BestInEdgeQueue(self):
+        if not self.Tree.QE:
+            print("QE is Empty!")
+            return None
+
         e_value = {(v, x): self.g_T[v] + self.calc_dist(v, x) + self.h_estimated(x)
                    for v, x in self.Tree.QE}
 
@@ -263,16 +335,10 @@ class BITStar:
         dy = node_end.y - node_start.y
         return math.hypot(dx, dy), math.atan2(dy, dx)
 
-    def animation(self, name):
+    def animation(self, name, cBest):
         theta, cMin, xCenter, C = self.init()
-        cBest = 30
-        self.plot_grid(name)
-        sample = self.Sample(300, cBest, cMin, xCenter, C)
-        for node in sample:
-            plt.plot(node.x, node.y, marker='.', color='lightgrey')
         self.draw_ellipse(xCenter, cBest, cMin, theta)
         plt.pause(0.001)
-        plt.show()
 
     def plot_grid(self, name):
         for (ox, oy, w, h) in self.obs_boundary:
@@ -318,7 +384,7 @@ class BITStar:
         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)
+        t = np.arange(0, 2 * math.pi + 0.1, 0.2)
         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]
@@ -336,8 +402,8 @@ def main():
     iter_max = 200
     print("start!!!")
     bit = BITStar(x_start, x_goal, eta, iter_max)
-    # bit.planning()
-    bit.animation("Batch Informed Trees (BIT*)")
+    # bit.animation("Batch Informed Trees (BIT*)")
+    bit.planning()
 
 
 if __name__ == '__main__':