Files
PathPlanning/Sampling_based_Planning/rrt_2D/batch_informed_trees.py
T
2020-08-03 11:11:34 -07:00

403 lines
13 KiB
Python

"""
Batch Informed Trees (BIT*)
@author: huiming zhou
"""
import os
import sys
import math
import random
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from scipy.spatial.transform import Rotation as Rot
sys.path.append(os.path.dirname(os.path.abspath(__file__)) +
"/../../Sampling_based_Planning/")
from Sampling_based_Planning.rrt_2D import env, plotting, utils
class Node:
def __init__(self, x, y):
self.x = x
self.y = y
self.parent = None
class Tree:
def __init__(self, x_start, x_goal):
self.x_start = x_start
self.goal = x_goal
self.r = 4.0
self.V = set()
self.E = set()
self.QE = set()
self.QV = set()
self.V_old = set()
class BITStar:
def __init__(self, x_start, x_goal, eta, iter_max):
self.x_start = Node(x_start[0], x_start[1])
self.x_goal = Node(x_goal[0], x_goal[1])
self.eta = eta
self.iter_max = iter_max
self.env = env.Env()
self.plotting = plotting.Plotting(x_start, x_goal)
self.utils = utils.Utils()
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
self.obs_rectangle = self.env.obs_rectangle
self.obs_boundary = self.env.obs_boundary
self.Tree = Tree(self.x_start, self.x_goal)
self.X_sample = set()
self.g_T = dict()
def init(self):
self.Tree.V.add(self.x_start)
self.X_sample.add(self.x_goal)
self.g_T[self.x_start] = 0.0
self.g_T[self.x_goal] = np.inf
cMin, theta = self.calc_dist_and_angle(self.x_start, self.x_goal)
C = self.RotationToWorldFrame(self.x_start, self.x_goal, cMin)
xCenter = np.array([[(self.x_start.x + self.x_goal.x) / 2.0],
[(self.x_start.y + self.x_goal.y) / 2.0], [0.0]])
return theta, cMin, xCenter, C
def planning(self):
theta, cMin, xCenter, C = self.init()
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 = {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())
vm, xm = self.BestInEdgeQueue()
self.Tree.QE.remove((vm, xm))
if self.g_T[vm] + self.calc_dist(vm, xm) + self.h_estimated(xm) < self.g_T[self.x_goal]:
actual_cost = self.cost(vm, xm)
if self.g_estimated(vm) + actual_cost + self.h_estimated(xm) < self.g_T[self.x_goal]:
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:
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]:
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.animation(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)
plt.show()
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}
self.Tree.E = {(v, w) for v, w in self.Tree.E
if self.f_estimated(v) <= cBest and self.f_estimated(w) <= cBest}
self.X_sample.update({v for v in self.Tree.V if self.g_T[v] == np.inf})
self.Tree.V = {v for v in self.Tree.V if self.g_T[v] < np.inf}
def cost(self, start, end):
if self.utils.is_collision(start, end):
return np.inf
return self.calc_dist(start, end)
def f_estimated(self, node):
return self.g_estimated(node) + self.h_estimated(node)
def g_estimated(self, node):
return self.calc_dist(self.x_start, node)
def h_estimated(self, node):
return self.calc_dist(node, self.x_goal)
def Sample(self, m, cMax, cMin, xCenter, C):
if cMax < np.inf:
return self.SampleEllipsoid(m, cMax, cMin, xCenter, C)
else:
return self.SampleFreeSpace(m)
def SampleEllipsoid(self, m, cMax, cMin, xCenter, C):
r = [cMax / 2.0,
math.sqrt(cMax ** 2 - cMin ** 2) / 2.0,
math.sqrt(cMax ** 2 - cMin ** 2) / 2.0]
L = np.diag(r)
ind = 0
delta = self.delta
Sample = set()
while ind < m:
xBall = self.SampleUnitNBall()
x_rand = np.dot(np.dot(C, L), xBall) + xCenter
node = Node(x_rand[(0, 0)], x_rand[(1, 0)])
in_obs = self.utils.is_inside_obs(node)
in_x_range = self.x_range[0] + delta <= node.x <= self.x_range[1] - delta
in_y_range = self.y_range[0] + delta <= node.y <= self.y_range[1] - delta
if not in_obs and in_x_range and in_y_range:
Sample.add(node)
ind += 1
return Sample
def SampleFreeSpace(self, m):
delta = self.utils.delta
Sample = set()
ind = 0
while ind < m:
node = Node(random.uniform(self.x_range[0] + delta, self.x_range[1] - delta),
random.uniform(self.y_range[0] + delta, self.y_range[1] - delta))
if self.utils.is_inside_obs(node):
continue
else:
Sample.add(node)
ind += 1
return Sample
def radius(self, q):
cBest = self.g_T[self.x_goal]
lambda_X = len([1 for v in self.Tree.V if self.f_estimated(v) <= cBest])
radius = 2 * self.eta * (1.5 * lambda_X / math.pi * math.log(q) / q) ** 0.5
return radius
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}
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}
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:
return np.inf
return min(self.g_T[v] + self.h_estimated(v) for v in self.Tree.QV)
def BestEdgeQueueValue(self):
if not self.Tree.QE:
return np.inf
return min(self.g_T[v] + self.calc_dist(v, x) + self.h_estimated(x)
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}
return min(e_value, key=e_value.get)
@staticmethod
def SampleUnitNBall():
while True:
x, y = random.uniform(-1, 1), random.uniform(-1, 1)
if x ** 2 + y ** 2 < 1:
return np.array([[x], [y], [0.0]])
@staticmethod
def RotationToWorldFrame(x_start, x_goal, L):
a1 = np.array([[(x_goal.x - x_start.x) / L],
[(x_goal.y - 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
@staticmethod
def calc_dist(start, end):
return math.hypot(start.x - end.x, start.y - end.y)
@staticmethod
def calc_dist_and_angle(node_start, node_end):
dx = node_end.x - node_start.x
dy = node_end.y - node_start.y
return math.hypot(dx, dy), math.atan2(dy, dx)
def animation(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.001)
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, "rs", linewidth=3)
plt.title(name)
plt.axis("equal")
@staticmethod
def draw_ellipse(x_center, c_best, dist, theta):
a = math.sqrt(c_best ** 2 - dist ** 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.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]
fx = rot @ np.array([x, y])
px = np.array(fx[0, :] + cx).flatten()
py = np.array(fx[1, :] + cy).flatten()
plt.plot(cx, cy, marker='.', color='darkorange')
plt.plot(px, py, linestyle='--', color='darkorange', linewidth=2)
def main():
x_start = (18, 8) # Starting node
x_goal = (37, 18) # Goal node
eta = 2
iter_max = 200
print("start!!!")
bit = BITStar(x_start, x_goal, eta, iter_max)
# bit.animation("Batch Informed Trees (BIT*)")
bit.planning()
if __name__ == '__main__':
main()