根据B站UP代码修改的3D-RRT程序

B站视频连接：机器人路径规划、轨迹优化课程-第五讲-RRT算法原理和代码讲解-哔哩哔哩
原视频提供的代码功能基本都已经完善，只要在原有2维基础上在增加一个维度就能实现3维的，碰撞还是用的原算法，高级的不太会

import copy
import math
import random
import time

import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation as Rot
import numpy as np

show_animation = True
def DrawBall(ax,r,location):
    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, np.pi, 100)
    x = r * np.outer(np.cos(u), np.sin(v))+location[0]
    y = r * np.outer(np.sin(u), np.sin(v))+location[1]
    z = r * np.outer(np.ones(np.size(u)), np.cos(v))+location[2]

    # Plot the surface
    ax.plot_surface(x, y, z)


class RRT:

    def __init__(self, obstacleList, randArea,
                 expandDis=2.0, goalSampleRate=10, maxIter=200):

        self.start = None
        self.goal = None
        self.min_rand = randArea[0]
        self.max_rand = randArea[1]
        self.expand_dis = expandDis
        self.goal_sample_rate = goalSampleRate
        self.max_iter = maxIter
        self.obstacle_list = obstacleList
        self.node_list = None
        self.fig=plt.figure()
        self.ax = self.fig.add_subplot(projection='3d')
        self.Draw_first=True
    def rrt_planning(self, start, goal, animation=True):
        start_time = time.time()
        self.start = Node(start[0], start[1], start[2])
        self.goal = Node(goal[0], goal[1], goal[2])
        self.node_list = [self.start]
        path = None

        for i in range(self.max_iter):
            rnd = self.sample()
            n_ind = self.get_nearest_list_index(self.node_list, rnd)
            nearestNode = self.node_list[n_ind]

            # steer
            thetax = math.atan2(rnd[1] - nearestNode.y, rnd[0] - nearestNode.x)
            lxy = math.sqrt((rnd[1] - nearestNode.y) ** 2 + (rnd[0] - nearestNode.x) ** 2)
            thetaz = math.atan2(rnd[2] - nearestNode.z, lxy)
            newNode = self.get_new_node(thetax, thetaz, n_ind, nearestNode)

            noCollision = self.check_segment_collision(newNode.x, newNode.y, newNode.z, nearestNode.x, nearestNode.y,
                                                       nearestNode.z)
            if noCollision:
                self.node_list.append(newNode)
                if animation:
                    self.draw_graph(newNode, path)

                if self.is_near_goal(newNode):
                    if self.check_segment_collision(newNode.x, newNode.y,newNode.z,
                                                    self.goal.x, self.goal.y,self.goal.z):
                        lastIndex = len(self.node_list) - 1
                        path = self.get_final_course(lastIndex)
                        pathLen = self.get_path_len(path)
                        print("current path length: {}, It costs {} s".format(pathLen, time.time() - start_time))

                        if animation:
                            self.draw_graph(newNode, path)
                        return path

    def rrt_star_planning(self, start, goal, animation=True):
        start_time = time.time()
        self.start = Node(start[0], start[1])
        self.goal = Node(goal[0], goal[1])
        self.node_list = [self.start]
        path = None
        lastPathLength = float('inf')

        for i in range(self.max_iter):
            rnd = self.sample()
            n_ind = self.get_nearest_list_index(self.node_list, rnd)
            nearestNode = self.node_list[n_ind]

            # steer
            theta = math.atan2(rnd[1] - nearestNode.y, rnd[0] - nearestNode.x)
            newNode = self.get_new_node(theta, n_ind, nearestNode)

            noCollision = self.check_segment_collision(newNode.x, newNode.y, nearestNode.x, nearestNode.y)
            if noCollision:
                nearInds = self.find_near_nodes(newNode)
                newNode = self.choose_parent(newNode, nearInds)

                self.node_list.append(newNode)
                self.rewire(newNode, nearInds)

                if animation:
                    self.draw_graph(newNode, path)

                if self.is_near_goal(newNode):
                    if self.check_segment_collision(newNode.x, newNode.y,
                                                    self.goal.x, self.goal.y):
                        lastIndex = len(self.node_list) - 1

                        tempPath = self.get_final_course(lastIndex)
                        tempPathLen = self.get_path_len(tempPath)
                        if lastPathLength > tempPathLen:
                            path = tempPath
                            lastPathLength = tempPathLen
                            print(
                                "current path length: {}, It costs {} s".format(tempPathLen, time.time() - start_time))

        return path

    def informed_rrt_star_planning(self, start, goal, animation=True):
        start_time = time.time()
        self.start = Node(start[0], start[1])
        self.goal = Node(goal[0], goal[1])
        self.node_list = [self.start]
        # max length we expect to find in our 'informed' sample space,
        # starts as infinite
        cBest = float('inf')
        path = None

        # Computing the sampling space
        cMin = math.sqrt(pow(self.start.x - self.goal.x, 2)
                         + pow(self.start.y - self.goal.y, 2))
        xCenter = np.array([[(self.start.x + self.goal.x) / 2.0],
                            [(self.start.y + self.goal.y) / 2.0], [0]])
        a1 = np.array([[(self.goal.x - self.start.x) / cMin],
                       [(self.goal.y - self.start.y) / cMin], [0]])

        e_theta = math.atan2(a1[1], a1[0])

        # 论文方法求旋转矩阵（2选1）
        # first column of identity matrix transposed
        # id1_t = np.array([1.0, 0.0, 0.0]).reshape(1, 3)
        # 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)

        # 直接用二维平面上的公式（2选1）
        C = np.array([[math.cos(e_theta), -math.sin(e_theta), 0],
                      [math.sin(e_theta), math.cos(e_theta), 0],
                      [0, 0, 1]])

        for i in range(self.max_iter):
            # Sample space is defined by cBest
            # cMin is the minimum distance between the start point and the goal
            # xCenter is the midpoint between the start and the goal
            # cBest changes when a new path is found

            rnd = self.informed_sample(cBest, cMin, xCenter, C)
            n_ind = self.get_nearest_list_index(self.node_list, rnd)
            nearestNode = self.node_list[n_ind]

            # steer
            theta = math.atan2(rnd[1] - nearestNode.y, rnd[0] - nearestNode.x)
            newNode = self.get_new_node(theta, n_ind, nearestNode)

            noCollision = self.check_segment_collision(newNode.x, newNode.y, nearestNode.x, nearestNode.y)
            if noCollision:
                nearInds = self.find_near_nodes(newNode)
                newNode = self.choose_parent(newNode, nearInds)

                self.node_list.append(newNode)
                self.rewire(newNode, nearInds)

                if self.is_near_goal(newNode):
                    if self.check_segment_collision(newNode.x, newNode.y,
                                                    self.goal.x, self.goal.y):
                        lastIndex = len(self.node_list) - 1
                        tempPath = self.get_final_course(lastIndex)
                        tempPathLen = self.get_path_len(tempPath)
                        if tempPathLen < cBest:
                            path = tempPath
                            cBest = tempPathLen
                            print(
                                "current path length: {}, It costs {} s".format(tempPathLen, time.time() - start_time))
            if animation:
                self.draw_graph_informed_RRTStar(xCenter=xCenter,
                                                 cBest=cBest, cMin=cMin,
                                                 e_theta=e_theta, rnd=rnd, path=path)

        return path

    def sample(self):
        if random.randint(0, 100) > self.goal_sample_rate:
            rnd = [random.uniform(self.min_rand, self.max_rand), random.uniform(self.min_rand, self.max_rand),
                   random.uniform(self.min_rand, self.max_rand)]
        else:  # goal point sampling
            rnd = [self.goal.x, self.goal.y, self.goal.z]
        return rnd

    def choose_parent(self, newNode, nearInds):
        if len(nearInds) == 0:
            return newNode

        dList = []
        for i in nearInds:
            dx = newNode.x - self.node_list[i].x
            dy = newNode.y - self.node_list[i].y
            d = math.hypot(dx, dy)
            theta = math.atan2(dy, dx)
            if self.check_collision(self.node_list[i], theta, d):
                dList.append(self.node_list[i].cost + d)
            else:
                dList.append(float('inf'))

        minCost = min(dList)
        minInd = nearInds[dList.index(minCost)]

        if minCost == float('inf'):
            print("min cost is inf")
            return newNode

        newNode.cost = minCost
        newNode.parent = minInd

        return newNode

    def find_near_nodes(self, newNode):
        n_node = len(self.node_list)
        r = 50.0 * math.sqrt((math.log(n_node) / n_node))
        d_list = [(node.x - newNode.x) ** 2 + (node.y - newNode.y) ** 2
                  for node in self.node_list]
        near_inds = [d_list.index(i) for i in d_list if i <= r ** 2]
        return near_inds

    def informed_sample(self, cMax, cMin, xCenter, C):
        if cMax < float('inf'):
            r = [cMax / 2.0,
                 math.sqrt(cMax ** 2 - cMin ** 2) / 2.0,
                 math.sqrt(cMax ** 2 - cMin ** 2) / 2.0]
            L = np.diag(r)
            xBall = self.sample_unit_ball()
            rnd = np.dot(np.dot(C, L), xBall) + xCenter
            rnd = [rnd[(0, 0)], rnd[(1, 0)]]
        else:
            rnd = self.sample()

        return rnd

    @staticmethod
    def sample_unit_ball():
        a = random.random()
        b = random.random()

        if b < a:
            a, b = b, a

        sample = (b * math.cos(2 * math.pi * a / b),
                  b * math.sin(2 * math.pi * a / b))
        return np.array([[sample[0]], [sample[1]], [0]])

    @staticmethod
    def get_path_len(path):
        pathLen = 0
        for i in range(1, len(path)):
            node1_x = path[i][0]
            node1_y = path[i][1]
            node2_x = path[i - 1][0]
            node2_y = path[i - 1][1]
            pathLen += math.sqrt((node1_x - node2_x)
                                 ** 2 + (node1_y - node2_y) ** 2)

        return pathLen

    @staticmethod
    def line_cost(node1, node2):
        return math.sqrt((node1.x - node2.x) ** 2 + (node1.y - node2.y) ** 2)

    @staticmethod
    def get_nearest_list_index(nodes, rnd):
        dList = [(node.x - rnd[0]) ** 2
                 + (node.y - rnd[1]) ** 2
                 + (node.z - rnd[2]) ** 2 for node in nodes]
        minIndex = dList.index(min(dList))
        return minIndex

    def get_new_node(self, theta1, theta2, n_ind, nearestNode):
        newNode = copy.deepcopy(nearestNode)

        newNode.x += self.expand_dis * math.cos(theta1)
        newNode.y += self.expand_dis * math.sin(theta1)
        newNode.z += self.expand_dis * math.sin(theta2)

        newNode.cost += self.expand_dis
        newNode.parent = n_ind
        return newNode

    def is_near_goal(self, node):
        d = self.line_cost(node, self.goal)
        if d < self.expand_dis:
            return True
        return False

    def rewire(self, newNode, nearInds):
        n_node = len(self.node_list)
        for i in nearInds:
            nearNode = self.node_list[i]

            d = math.sqrt((nearNode.x - newNode.x) ** 2
                          + (nearNode.y - newNode.y) ** 2)

            s_cost = newNode.cost + d

            if nearNode.cost > s_cost:
                theta = math.atan2(newNode.y - nearNode.y,
                                   newNode.x - nearNode.x)
                if self.check_collision(nearNode, theta, d):
                    nearNode.parent = n_node - 1
                    nearNode.cost = s_cost

    @staticmethod
    def distance_squared_point_to_segment(v, w, p):
        # Return minimum distance between line segment vw and point p
        if np.array_equal(v, w):
            return (p - v).dot(p - v)  # v == w case
        l2 = (w - v).dot(w - v)  # i.e. |w-v|^2 -  avoid a sqrt
        # Consider the line extending the segment,
        # parameterized as v + t (w - v).
        # We find projection of point p onto the line.
        # It falls where t = [(p-v) . (w-v)] / |w-v|^2
        # We clamp t from [0,1] to handle points outside the segment vw.
        t = max(0, min(1, (p - v).dot(w - v) / l2))
        projection = v + t * (w - v)  # Projection falls on the segment
        return (p - projection).dot(p - projection)

    def check_segment_collision(self, x1, y1, z1, x2, y2, z2):
        for (ox, oy, oz, size) in self.obstacle_list:
            dd = self.distance_squared_point_to_segment(
                np.array([x1, y1,z1]),
                np.array([x2, y2,z2]),
                np.array([ox, oy,oz]))
            if dd <= size ** 2+1:
                return False  # collision
        return True

    def check_collision(self, nearNode, theta, d):
        tmpNode = copy.deepcopy(nearNode)
        end_x = tmpNode.x + math.cos(theta) * d
        end_y = tmpNode.y + math.sin(theta) * d
        return self.check_segment_collision(tmpNode.x, tmpNode.y, end_x, end_y)

    def get_final_course(self, lastIndex):
        path = [[self.goal.x, self.goal.y,self.goal.z]]
        while self.node_list[lastIndex].parent is not None:
            node = self.node_list[lastIndex]
            path.append([node.x, node.y,node.z])
            lastIndex = node.parent
        path.append([self.start.x, self.start.y, self.start.z])
        return path

    def draw_graph_informed_RRTStar(self, xCenter=None, cBest=None, cMin=None, e_theta=None, rnd=None, path=None):
        plt.clf()
        # for stopping simulation with the esc key.
        plt.gcf().canvas.mpl_connect(
            'key_release_event',
            lambda event: [exit(0) if event.key == 'escape' else None])
        if rnd is not None:
            plt.plot(rnd[0], rnd[1], "^k")
            if cBest != float('inf'):
                self.plot_ellipse(xCenter, cBest, cMin, e_theta)

        for node in self.node_list:
            if node.parent is not None:
                if node.x or node.y is not None:
                    plt.plot([node.x, self.node_list[node.parent].x], [
                        node.y, self.node_list[node.parent].y], "-g")
        if self.Draw_first:
            for (ox, oy, size) in self.obstacle_list:
                plt.plot(ox, oy, "ok", ms=30 * size)
            self.Draw_first=False
        if path is not None:
            plt.plot([x for (x, y) in path], [y for (x, y) in path], '-r')

        plt.plot(self.start.x, self.start.y, "xr")
        plt.plot(self.goal.x, self.goal.y, "xr")
        plt.axis([-2, 18, -2, 15])
        plt.grid(True)
        plt.pause(0.01)

    @staticmethod
    def plot_ellipse(xCenter, cBest, cMin, e_theta):  # pragma: no cover

        a = math.sqrt(cBest ** 2 - cMin ** 2) / 2.0
        b = cBest / 2.0
        angle = math.pi / 2.0 - e_theta
        cx = xCenter[0]
        cy = xCenter[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_matrix()[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 draw_graph(self, rnd=None, path=None):
        # plt.clf()
        # # for stopping simulation with the esc key.
        # plt.gcf().canvas.mpl_connect(
        #     'key_release_event',
        #     lambda event: [exit(0) if event.key == 'escape' else None])


        if rnd is not None:
            self.ax.plot(rnd.x, rnd.y,rnd.z, "^k")

        for node in self.node_list:
            if node.parent is not None:
                if node.x or node.y is not None:
                    self.ax.plot([node.x, self.node_list[node.parent].x], [
                        node.y, self.node_list[node.parent].y],[node.z, self.node_list[node.parent].z], "-g")
        if self.Draw_first:

            for (ox, oy,oz, size) in self.obstacle_list:
                # self.plot_circle(ox, oy, size)
                DrawBall(self.ax,size,[ox,oy,oz])
            self.Draw_first=False

        self.ax.plot(self.start.x, self.start.y,self.start.z, "xr")
        self.ax.plot(self.goal.x, self.goal.y,self.goal.z, "xr")

        if path is not None:
            self.ax.plot([x for (x, y,z) in path], [y for (x, y,z) in path],[z for (x, y,z) in path], '-r')

        # ax.axis([-2, 18, -2, 15,-2,15])
        self.ax.set_xlim(-2, 18)
        self.ax.set_ylim(-2, 15)
        self.ax.set_zlim(-2, 15)
        self.ax.grid(True)
        plt.pause(1)


class Node:

    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z
        self.cost = 0.0
        self.parent = None


def main():
    print("Start rrt planning")

    # create obstacles
    obstacleList = [
        (3, 3, 3, 1.5),
        (12, 5, 5, 5),
        (3, 9, 9, 4),
        (9, 11, 11, 2),
    ]
    # obstacleList = [(5, 5, 1), (3, 6, 2), (3, 8, 2), (3, 10, 2), (7, 5, 2),
    #                 (9, 5, 2), (8, 10, 1)]

    # Set params
    rrt = RRT(randArea=[-2, 18], obstacleList=obstacleList, maxIter=200)
    path = rrt.rrt_planning(start=[0, 0, 0], goal=[15, 12, 12], animation=show_animation)
    # path = rrt.rrt_star_planning(start=[0, 0], goal=[15, 12], animation=show_animation)
    # path = rrt.informed_rrt_star_planning(start=[0, 0], goal=[15, 12], animation=show_animation)
    print("Done!!")

    if show_animation and path:
        plt.show()


if __name__ == '__main__':
    main()