Adding trajectory optimization for multiple robots

LocalPlanners/TrajOpt.py

+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import Circle, Rectangle
+from casadi import *
+
+class TrajOptMultiRobot():
+    def __init__(self, num_robots, robot_radius, starts, goals, circle_obstacles, rectangle_obstacles,
+                 rob_dist_weight, obs_dist_weight, control_weight, time_weight):
+        self.num_robots = num_robots
+        self.starts = starts
+        self.goals = goals
+        self.circle_obs = circle_obstacles
+        self.rect_obs = rectangle_obstacles
+        self.rob_dist_weight = rob_dist_weight
+        self.obs_dist_weight = obs_dist_weight
+        self.control_weight =control_weight
+        self.time_weight = time_weight
+        self.robot_radius = MX(robot_radius)
+
+    def dist(self, robot_position, circle):
+        """
+        Returns the distance between a robot and a circle
+
+        params:
+            robot_position [x,y]
+            circle [x,y,radius]
+        """
+        return norm_2(robot_position - transpose(circle[:2])) - circle[2] - self.robot_radius
+
+    def apply_quadratic_barrier(self, d_max, d, c):
+        """
+        Applies a quadratic barrier to some given distance. The quadratic barrier 
+        is a soft barrier function. We are using it for now to avoid any issues with
+        invalid initial solutions, which hard barrier functions cannot handle. 
+
+        params:
+            d (float): 
+            c (float): controls the steepness of curve. 
+                       higher c --> gets more expensive faster as you move toward obs
+            d_max (float): 
+        """
+        return c*fmax(0, d_max-d)**2
+
+    def solve(self, num_control_intervals):
+
+        N = num_control_intervals
+        opti = Opti() # Optimization problem
+
+        # ---- decision variables --------- #
+        X = opti.variable(self.num_robots*4, N+1)   # state trajectory (x,y,v_x,v_y)
+        pos = X[:self.num_robots*2,:]               # position is the first two values
+        vel = X[self.num_robots*2:,:]               # velocity is the last two values
+
+        
+        circle_obs = DM(self.circle_obs)            # make the obstacles casadi objects 
+        
+        U = opti.variable(self.num_robots*2, N)     # control trajectory (a_x, a_y)
+        T = opti.variable()                         # final time
+
+
+        # sum up the cost of distance to obstacles
+        # TODO:: Include rectangular obstacles
+        dist_to_other_obstacles = 0
+        for r in range(self.num_robots):
+            for k in range(N):
+                for c in range(circle_obs.shape[0]):
+                    circle = circle_obs[c, :]
+                    d = self.dist(pos[2*r : 2*(r+1), k], circle)
+                    dist_to_other_obstacles += self.apply_quadratic_barrier(1, d, 1)
+
+        dist_to_other_robots = 0
+        for k in range(N):
+            for r1 in range(self.num_robots):
+                for r2 in range(self.num_robots):
+                    if r1 != r2:
+                        # print(f"\n{r1} position1 = {pos[2*r1 : 2*(r1+1), k]}")
+                        # print(f"{r2} position2 = {pos[2*r2 : 2*(r2+1), k]}")
+
+                        # note: using norm 2 here gives an invalid num detected error. 
+                        # Must be the sqrt causing an issue
+                        # d = norm_2(pos[2*r1 : 2*(r1+1), k] - pos[2*r2 : 2*(r2+1), k]) - 2*self.robot_radius
+                        d = sumsqr(pos[2*r1 : 2*(r1+1), k] - pos[2*r2 : 2*(r2+1), k]) - 2*self.robot_radius
+                        dist_to_other_robots += self.apply_quadratic_barrier(5, d, 2)
+                        
+        # control costs
+        control_costs = 0
+        dt = T/N # length of a control interval
+        for k in range(N): # loop over control intervals
+            opti.subject_to(pos[:,k+1]==pos[:,k] + dt*vel[:,k]) 
+            opti.subject_to(vel[:,k+1]==vel[:,k] + dt*U[:,k])
+        
+
+        opti.minimize(dist_robots_weight*dist_to_other_robots 
+                      + dist_obstacles_weight*dist_to_other_obstacles 
+                      + control_costs_weight* control_costs 
+                      + time_weight*T)
+
+        # ---- path constraints -----------
+        for k in range(N):
+            for r in range(self.num_robots):
+                opti.subject_to(sumsqr(vel[2*r:2*(r+1),k]) <= 0.2**2)
+                opti.subject_to(sumsqr(U[2*r:2*(r+1),k]) <= 0.1**2)
+
+        opti.subject_to(opti.bounded(0,pos,10))
+        # opti.subject_to(opti.bounded(-0.2,vel,0.2)) 
+        # opti.subject_to(opti.bounded(-.1,U,.1)) # control is limited
+
+        # ---- boundary conditions --------
+        for r in range(self.num_robots):
+            opti.subject_to(vel[2*r : 2*(r+1), 0]==[0,0]) 
+            opti.subject_to(pos[2*r : 2*(r+1), 0]==self.starts[r])
+            opti.subject_to(pos[2*r : 2*(r+1), -1]==self.goals[r])
+
+        # ---- misc. constraints  ----------
+        opti.subject_to(opti.bounded(0,T,100))
+
+        # ---- initial values for solver ---
+        opti.set_initial(T, 20)
+
+        # ---- solve NLP              ------
+        opti.solver("ipopt") # set numerical backend
+        sol = opti.solve()   # actual solve
+
+        return sol,pos
+
+    def plot_paths(self, x_opt):
+        fig, ax = plt.subplots()
+
+        # Plot obstacles
+        for obstacle in self.circle_obs:
+            # if len(obstacle) == 2:  # Circle
+            ax.add_patch(Circle(obstacle, obstacle[2], color='red'))
+            # elif len(obstacle) == 4:  # Rectangle
+            #     ax.add_patch(Rectangle((obstacle[0], obstacle[1]), obstacle[2], obstacle[3], color='red'))
+
+        if self.num_robots > 20:
+            colors = plt.cm.hsv(np.linspace(0.2, 1.0, self.num_robots))
+        elif self.num_robots > 10:
+            colors = plt.cm.tab20(np.linspace(0, 1, self.num_robots))
+        else:
+            colors = plt.cm.tab10(np.linspace(0, 1, self.num_robots))
+
+        # Plot robot paths
+        for r,color in zip(range(self.num_robots),colors):
+            ax.plot(x_opt[r*2, :], x_opt[r*2+1, :], label=f'Robot {r+1}', color=color)
+            ax.scatter(x_opt[r*2, :], x_opt[r*2+1, :], color=color, s=10 )
+            ax.scatter(self.starts[r][0], self.starts[r][1], s=85,color=color)
+            ax.scatter(self.goals[r][0], self.goals[r][1], s=85,facecolors='none', edgecolors=color)
+
+        ax.set_xlabel('X')
+        ax.set_ylabel('Y')
+        ax.legend()
+        ax.set_aspect('equal', 'box')
+
+        plt.ylim(0,10)
+        plt.xlim(0,10)
+        plt.title('Robot Paths')
+        plt.grid(False)
+        plt.show()
+
+
+if __name__ == "__main__":
+
+    # define obstacles
+    circle_obs = np.array([[5,5,1],
+                           [7,7,1],
+                           [3,3,1]])
+
+    rectangle_obs = np.array([])
+
+    #  define all the robots' starts and goals
+    robot_starts = [[1,6],[9,5],[2,2],[1,3]]
+    robot_goals = [[9,6],[1,5],[8,8],[7,3]]
+
+    # weights for the cost function
+    dist_robots_weight = 1.0
+    dist_obstacles_weight = 1.7
+    control_costs_weight = 1.0
+    time_weight = 1.0
+
+    # other params
+    num_robots = 4
+    rob_radius = 0.25
+
+    solver = TrajOptMultiRobot(num_robots=num_robots, 
+                               robot_radius=rob_radius,
+                               starts=robot_starts, 
+                               goals=robot_goals, 
+                               circle_obstacles=circle_obs, 
+                               rectangle_obstacles=rectangle_obs,
+                               rob_dist_weight=dist_robots_weight,
+                               obs_dist_weight=dist_obstacles_weight,
+                               control_weight=control_costs_weight,
+                               time_weight=time_weight
+                               )
+    sol,pos = solver.solve(20)
+    pos_vals = np.array(sol.value(pos))
+
+    
+    solver.plot_paths(pos_vals)
+
+
+
+