Added obstacle plotting, starting to integrate into place_grid

guided_mrmp/controllers/place_grid.py

@@ -20,7 +20,8 @@ def place_grid(robot_locations, cell_size, grid_size=5, subgoals=[], obstacles=[
     robot_locations = np.array(robot_locations)
     subgoals = np.array(subgoals)
     obstacles = np.array(obstacles)
-    N = len(robot_locations)
+    num_robots = len(robot_locations)
+    num_obst = len(obstacles)
     
     # Decision variable: Bottom-left corner of the grid in continuous space
     bottom_left = cp.Variable(2, name='origin')
@@ -28,12 +29,21 @@ def place_grid(robot_locations, cell_size, grid_size=5, subgoals=[], obstacles=[
     top_right = bottom_left + grid_size * cell_size
     
     # Decision variable: Integer grid indices for each robot
-    grid_indices = cp.Variable((N, 2), integer=True, name='grid_indices')
+    grid_indices = cp.Variable((num_robots, 2), integer=True, name='grid_indices')
     
     # Calculate cell centers for each robot based on grid indices
     # Reshape origin to (1, 2) for broadcasting
     cell_centers = cp.reshape(bottom_left, (1, 2), order='C') + grid_indices * cell_size + cell_size / 2
     
+    if num_obst:
+        obstacle_cells = cp.Variable((num_obst, grid_size, grid_size), boolean=True, name='obstacle_cells')
+        
+        for (cx, cy, r) in obstacles:
+            for x_idx in range(grid_size):
+                for y_idx in range(grid_size):
+                    p1 = bottom_left + np.array([x_idx, y_idx]) * cell_size
+                    p2 = p1 + np.ones(2) * cell_size  
+    
     # Objective: Minimize the sum of squared distances
     cost = cp.sum_squares(robot_locations - cell_centers)
     
@@ -50,8 +60,8 @@ def place_grid(robot_locations, cell_size, grid_size=5, subgoals=[], obstacles=[
     # Use Big M method to ensure unique grid indices
     M_ind = 10 * grid_size # Big M relative to grid indices
     M_cts = 10 * max(max(robot_locations[:,0]) - min(robot_locations[:,0]), max(robot_locations[:,1]) - min(robot_locations[:,1])) # Big M relative to robot locations
-    for i in range(N):
-        for j in range(i+1, N):
+    for i in range(num_robots):
+        for j in range(i+1, num_robots):
             # At least one of the two constraints below must be true
             xsep = cp.Variable(boolean=True)
             ysep = cp.Variable(boolean=True)
@@ -76,9 +86,13 @@ def place_grid(robot_locations, cell_size, grid_size=5, subgoals=[], obstacles=[
             constraints.append(grid_indices[j, 1] - grid_indices[i, 1] + M_ind * (1 - b1) + M_ind * (1 - ysep) >= 1)
     
     # All robots and subgoals must be within grid bounds
+    for loc in robot_locations:
+        constraints.append(bottom_left <= loc)
+        constraints.append(loc <= top_right)
     for sg in subgoals:
         constraints.append(bottom_left <= sg)
         constraints.append(sg <= top_right)
+
     
     # Solve the optimization problem
     prob_init_start_time = time.time()
@@ -98,6 +112,7 @@ def place_grid(robot_locations, cell_size, grid_size=5, subgoals=[], obstacles=[
     return bottom_left.value, cell_centers.value
 
 
+# Working on making this convex
 def two_corner_place_grid(robot_locations, grid_size=5, subgoals=[], obstacles=[]):
     """
         Place a grid to cover robot locations with alignment to centers.
@@ -254,7 +269,7 @@ def plot_grid(bottom_left, top_right, grid_size):
                  [(bottom_left + i * y_prime_hat)[1], (bottom_right + i * y_prime_hat)[1]], 'k-')
 
 
-def get_roomba_locs(low, high, num_robots, radius=0.5):
+def get_roomba_locs(low, high, num_robots, radius=0.5, obstacles=[]):
     """
     Generates a list of roomba locations within the box bounded by points (low, low), (high, low), (high, high), (low, high).
     The roombas must be separated by at least 2 * radius
@@ -262,10 +277,17 @@ def get_roomba_locs(low, high, num_robots, radius=0.5):
     locs = []
     while len(locs) < num_robots:
         locs.append(np.random.uniform(low, high, 2))
+        invalid = False
+        for (obst_x, obst_y, obst_r) in obstacles:
+            if np.linalg.norm(np.array(locs[-1]) - np.array([obst_x, obst_y])) <= radius + obst_r:
+                invalid = True
+                break         
         for other_loc in locs[:-1]:
             if np.linalg.norm(np.array(locs[-1]) - np.array(other_loc)) <= 2 * radius:
-                locs = locs[:-1]
+                invalid = True
                 break
+        if invalid:
+            locs = locs[:-1]
     return np.array(locs)
 
 
@@ -275,12 +297,15 @@ def main(seed, num_robots, plot, two_corner):
     
     if not two_corner:
         roomba_radius = 0.5
-        robot_locations = get_roomba_locs(low=0, high=6, num_robots=num_robots, radius=roomba_radius)
-        # robot_locations = np.random.uniform(low=0, high=5, size=(num_robots, 2))
         cell_size = 2.5 * roomba_radius
         grid_size = 5
-        # subgoals = np.random.uniform(low=0, high=6, size=(num_robots, 2))
-        subgoals = np.array([[0, 0], [0, 6], [6, 6], [6, 0]])
+        
+        
+        obstacles = np.array([[2, 2, 0.75], [4, 4, 0.5]])
+        # robot_locations = np.random.uniform(low=0, high=5, size=(num_robots, 2))
+        robot_locations = get_roomba_locs(low=0, high=6, num_robots=num_robots, radius=roomba_radius, obstacles=obstacles)
+        # subgoals = np.array([[0, 0], [0, 6], [6, 6], [6, 0]])
+        subgoals = get_roomba_locs(low=0, high=6, num_robots=num_robots, radius=roomba_radius, obstacles=obstacles)
         
         bottom_left, cell_centers = place_grid(robot_locations=robot_locations, 
                                                cell_size=cell_size,
@@ -329,6 +354,12 @@ def main(seed, num_robots, plot, two_corner):
             for (x, y) in subgoals:
                 circle = patches.Circle((x, y), radius=roomba_radius, edgecolor='orange', fill=False, linewidth=2)
                 ax.add_patch(circle)
+    
+            obstacles = np.array(obstacles)
+            plt.scatter(obstacles[:, 0], obstacles[:, 1], c='black', marker='s', label='Obstacles')
+            for (x, y, r) in obstacles:
+                circle = patches.Circle((x, y), radius=r, edgecolor='black', fill=False, linewidth=2)
+                ax.add_patch(circle)
         
         plt.legend(loc='upper left')