Subgoals added to place_grid without center matching

guided_mrmp/controllers/place_grid.py

 import cvxpy as cp
 import numpy as np
+import time
 
-def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5), return_loss=False):
+def place_grid(robot_locations, cell_size, grid_size=5, subgoals=[], obstacles=[]):
     """
         Place a grid to cover robot locations with alignment to centers.
 
         inputs:
             - robot_locations (list): locations of robots involved in conflict [[x,y], [x,y], ...]
             - cell_size (float): the width of each grid cell in continuous space
-            - grid_shape (tuple): (# of rows, # of columns) of the grid
+            - grid_size (tuple): width of the grid in cells
+            - obstacles (list): locations of circular obstacles [[x,y,r], [x,y,r], ...]
         outputs:
             - origin (tuple): bottom-left corner of the grid in continuous space
             - cell_centers (list): centers of grid cells for each robot (same order as robot_locations)
-            - loss: when return_loss=True, sum of squared differences loss
     """
+    start_time = time.time()
+    
     robot_locations = np.array(robot_locations)
+    subgoals = np.array(subgoals)
+    obstacles = np.array(obstacles)
     N = len(robot_locations)
     
     # Decision variable: Bottom-left corner of the grid in continuous space
-    origin = cp.Variable(2, name='origin')
+    bottom_left = cp.Variable(2, name='origin')
+    # Defin top right for convenience
+    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')
     
     # Calculate cell centers for each robot based on grid indices
     # Reshape origin to (1, 2) for broadcasting
-    cell_centers = cp.reshape(origin, (1, 2), order='C') + grid_indices * cell_size + cell_size / 2
+    cell_centers = cp.reshape(bottom_left, (1, 2), order='C') + grid_indices * cell_size + cell_size / 2
     
     # Objective: Minimize the sum of squared distances
     cost = cp.sum_squares(robot_locations - cell_centers)
@@ -37,102 +44,126 @@ def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5), return_loss=Fals
     constraints.append(grid_indices >= 0)
     
     # Grid indices must fit within grid bounds
-    if grid_shape[0] == grid_shape[1]: # Square grid
-        constraints.append(grid_indices <= grid_shape[0] - 1)
-    else: # Rectangular grid
-        constraints.append(grid_indices[:,0] <= grid_shape[1] - 1)
-        constraints.append(grid_indices[:,1] <= grid_shape[0] - 1)
+    constraints.append(grid_indices <= grid_size - 1)
     
     # No two robots can share a cell
     # Use Big M method to ensure unique grid indices
-    M = max(grid_shape) * 10
+    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):
             # At least one of the two constraints below must be true
-            y1 = cp.Variable(boolean=True)
-            y2 = cp.Variable(boolean=True)
-            constraints.append(y1 + y2 >= 1)
+            xsep = cp.Variable(boolean=True)
+            ysep = cp.Variable(boolean=True)
+            constraints.append(xsep + ysep >= 1)
             
             # Enforces separation by at least 1 in the x direction
-            if robot_locations[i, 0] >= robot_locations[j, 0]:
-                constraints.append(grid_indices[i, 0] - grid_indices[j, 0] + M * (1 - y1) >= 1)
-            else:
-                constraints.append(grid_indices[j, 0] - grid_indices[i, 0] + M * (1 - y1) >= 1)
+            b0 = cp.Variable(boolean=True) # b0 = 0 if robot i's x >= robot j's x, 1 otherwise
+            # b0 = 0
+            constraints.append(robot_locations[j, 0] - robot_locations[i, 0] <= M_cts * b0)
+            constraints.append(grid_indices[i, 0] - grid_indices[j, 0] + M_ind * b0 + M_ind * (1 - xsep) >= 1)
+            # b0 = 1
+            constraints.append(robot_locations[i, 0] - robot_locations[j, 0] <= M_cts * (1 - b0))
+            constraints.append(grid_indices[j, 0] - grid_indices[i, 0] + M_ind * (1 - b0) + M_ind * (1 - xsep) >= 1)
             
             # Enforces separation by at least 1 in the y direction
-            if robot_locations[i, 1] >= robot_locations[j, 1]:
-                constraints.append(grid_indices[i, 1] - grid_indices[j, 1] + M * (1 - y2) >= 1)
-            else:
-                constraints.append(grid_indices[j, 1] - grid_indices[i, 1] + M * (1 - y2) >= 1)
+            b1 = cp.Variable(boolean=True) # b1 = 0 if robot i's y >= robot j's y, 1 otherwise
+            # b1 = 0
+            constraints.append(robot_locations[j, 1] - robot_locations[i, 1] <= M_cts * b1)
+            constraints.append(grid_indices[i, 1] - grid_indices[j, 1] + M_ind * b1 + M_ind * (1 - ysep) >= 1)
+            # b1 = 1
+            constraints.append(robot_locations[i, 1] - robot_locations[j, 1] <= M_cts * (1 - b1))
+            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 sg in subgoals:
+        constraints.append(bottom_left <= sg)
+        constraints.append(sg <= top_right)
     
     # Solve the optimization problem
+    prob_init_start_time = time.time()
     prob = cp.Problem(cp.Minimize(cost), constraints)
+    solve_start_time = time.time()
     prob.solve(solver=cp.SCIP)
+    solve_end_time = time.time()
+    
+    print("Time to add vars/constraints:", prob_init_start_time - start_time)
+    print("Time to parse:", solve_start_time - prob_init_start_time)
+    print("Time to solve:", solve_end_time - solve_start_time)
 
-    if prob.status not in ["optimal", "optimal_inaccurate"]:
+    if prob.status != "optimal":
         print("Problem could not be solved to optimality.")
         return None
     
-    if return_loss:
-        return origin.value, cell_centers.value, prob.value
-    return origin.value, cell_centers.value
+    return bottom_left.value, cell_centers.value
 
 
-# This currently does not follow DCP, working on it
-def place_grid_with_rotation(robot_locations, grid_size=5, return_loss=False):
+def two_corner_place_grid(robot_locations, grid_size=5, subgoals=[], obstacles=[]):
     """
-        Place a square grid to cover robot locations with alignment to centers. Allows for rotation and scaling of the grid.
+        Place a grid to cover robot locations with alignment to centers.
 
         inputs:
             - robot_locations (list): locations of robots involved in conflict [[x,y], [x,y], ...]
-            - grid_size (float): the number of cells in each row/column of the grid
+            - cell_size (float): the width of each grid cell in continuous space
+            - grid_size (tuple): width of the grid in cells
+            - obstacles (list): locations of circular obstacles [[x,y,r], [x,y,r], ...]
         outputs:
-            - bottom_left (tuple): bottom-left corner of the grid in continuous space
-            - top_right (tuple): top-right corner of the grid in continuous space
+            - origin (tuple): bottom-left corner of the grid in continuous space
             - cell_centers (list): centers of grid cells for each robot (same order as robot_locations)
-            - loss: when return_loss=True, sum of squared differences loss
     """
+    start_time = time.time()
+    
     robot_locations = np.array(robot_locations)
+    subgoals = np.array(subgoals)
+    obstacles = np.array(obstacles)
     N = len(robot_locations)
     
-    # Decision variables: Bottom-left and top-right corners of the grid in continuous space
+    # Decision variable: Bottom-left corner of the grid in continuous space
     bottom_left = cp.Variable(2, name='bottom_left')
     top_right = cp.Variable(2, name='top_right')
     
+    # Bottom-right and top-left corners of the grid for convenience
+    # bottom_right = 0.5 * cp.hstack([bottom_left[0] + top_right[0] - bottom_left[1] + top_right[1],
+    #                                 bottom_left[0] - top_right[0] + bottom_left[1] + top_right[1]])
+    # top_left = 0.5 * cp.hstack([bottom_left[0] + top_right[0] + bottom_left[1] - top_right[1],
+    #                             -bottom_left[0] + top_right[0] + bottom_left[1] + top_right[1]])
+    bottom_right = cp.Variable(2, name='bottom_right')
+    top_left = cp.Variable(2, name='top_left')
+    
+    grid_x_hat = cp.Variable(2, name='grid_x_hat')
+    grid_y_hat = cp.Variable(2, name='grid_y_hat')
+    
     # Decision variable: Integer grid indices for each robot
     grid_indices = cp.Variable((N, 2), integer=True, name='grid_indices')
-
-    # Define bottom-right and top-left corners of the grid
-    bottom_right = (1 / 2) * cp.hstack([
-        bottom_left[0] + top_right[0] - bottom_left[1] + top_right[1],
-        bottom_left[0] - top_right[0] + bottom_left[1] + top_right[1]
-    ])
-    top_left = (1 / 2) * cp.hstack([
-        bottom_left[0] + top_right[0] + bottom_left[1] - top_right[1],
-        -bottom_left[0] + top_right[0] + bottom_left[1] + top_right[1]
-    ])
-    
-    # Define basis vectors for the grid
-    # Vector pointing from **left -> right** on the grid, with length equal to the width of one cell (1 grid index)
-    x_prime_hat = (bottom_right - bottom_left) * (1 / grid_size)
-    # Vector pointing from **bottom -> top** on the grid, with length equal to the width of one cell (1 grid index)
-    y_prime_hat = (top_left - bottom_left) * (1 / grid_size)
     
     # Calculate cell centers for each robot based on grid indices
-    cell_centers = cp.vstack([bottom_left for _ in range(N)]) # Grid origin point
-    cell_centers += cp.vstack([x_prime_hat * (grid_indices[i,0] + 0.5) for i in range(N)]) # Component of cell centers in the x_prime direction
-    cell_centers += cp.vstack([y_prime_hat * (grid_indices[i,1] + 0.5) for i in range(N)]) # Component of cell centers in the y_prime direction
-    print(cell_centers)
+    # Reshape origin to (1, 2) for broadcasting
+    grid_x_offsets = cp.Variable((N, 2), name='grid_x_offsets')
+    grid_y_offsets = cp.Variable((N, 2), name='grid_y_offsets')
+    cell_centers = cp.reshape(bottom_left, (1, 2), order='C') + grid_x_offsets + grid_y_offsets
     
     # Objective: Minimize the sum of squared distances
     cost = cp.sum_squares(robot_locations - cell_centers)
     
-    # Initialize constraints
+    # Constraints
     constraints = []
     
-    # The top right corner of the grid can't be below or to the left of the bottom left corner
+    # Ensure top-right and bottom-left corners are in the right orientation
     constraints.append(top_right >= bottom_left)
     
+    # Fixing bottom-right and top-left corners
+    constraints.append(2 * bottom_right[0] == bottom_left[0] + top_right[0] - bottom_left[1] + top_right[1])
+    constraints.append(2 * bottom_right[1] == bottom_left[0] - top_right[0] + bottom_left[1] + top_right[1])
+    constraints.append(2 * top_left[0] == bottom_left[0] + top_right[0] + bottom_left[1] - top_right[1])
+    constraints.append(2 * top_left[1] == -bottom_left[0] + top_right[0] + bottom_left[1] + top_right[1])
+    
+    # Defining grid_x_hat and grid_y_hat based on corners
+    constraints.append(grid_x_hat == (bottom_right - bottom_left) * (1 / grid_size))
+    constraints.append(grid_y_hat == (top_left - bottom_left) * (1 / grid_size))
+    
+    # Defining offsets in cell centers calculation
+    constraints.append(grid_x_offsets == grid_x_hat * grid_indices)
+    
     # Grid indices must be non-negative
     constraints.append(grid_indices >= 0)
     
@@ -140,52 +171,64 @@ def place_grid_with_rotation(robot_locations, grid_size=5, return_loss=False):
     constraints.append(grid_indices <= grid_size - 1)
     
     # No two robots can share a cell
-    M = 1e6 # Sufficiently large constant for Big-M
+    # 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):
-            x_separated = cp.Variable(boolean=True)
-            y_separated = cp.Variable(boolean=True)
-            
-            # Robot i's coordinate in the x_prime direction
-            robot_i_x_prime = robot_locations[i] @ x_prime_hat
-            # Robot j's coordinate in the x_prime direction
-            robot_j_x_prime = robot_locations[j] @ x_prime_hat
-            
-            b0 = cp.Variable(boolean=True)
-            # When b1 = 1, robot i's x_prime coordinate is greater than robot j's
-            constraints.append(robot_i_x_prime - robot_j_x_prime <= M * b0)  
-            # When b1 = 0, robot j's x_prime coordinate is greater than robot i's
-            constraints.append(robot_j_x_prime - robot_i_x_prime <= M * (1 - b0)) 
-            # Enforces separation by at least 1 between the robots' x indices on the grid
-            constraints.append((2 * b0 - 1) * (grid_indices[i, 0] - grid_indices[j, 0]) + M * (1 - x_separated) >= 1)
+            # At least one of the two constraints below must be true
+            xsep = cp.Variable(boolean=True)
+            ysep = cp.Variable(boolean=True)
+            constraints.append(xsep + ysep >= 1)
             
-            # Robot i's coordinate in the y_prime direction
-            robot_i_y_prime = robot_locations[i] @ y_prime_hat
-            # Robot j's coordinate in the y_prime direction
-            robot_j_y_prime = robot_locations[j] @ y_prime_hat
+            # Enforces separation by at least 1 in the x direction
+            b0 = cp.Variable(boolean=True) # b0 = 0 if robot i's x >= robot j's x, 1 otherwise
+            # b0 = 0
+            constraints.append(robot_locations[j, 0] - robot_locations[i, 0] <= M_cts * b0)
+            constraints.append(grid_indices[i, 0] - grid_indices[j, 0] + M_ind * b0 + M_ind * (1 - xsep) >= 1)
+            # b0 = 1
+            constraints.append(robot_locations[i, 0] - robot_locations[j, 0] <= M_cts * (1 - b0))
+            constraints.append(grid_indices[j, 0] - grid_indices[i, 0] + M_ind * (1 - b0) + M_ind * (1 - xsep) >= 1)
             
-            b1 = cp.Variable(boolean=True)
-            # When b1 = 1, robot i's y_prime coordinate is greater than robot j's
-            constraints.append(robot_i_y_prime - robot_j_y_prime <= M * b1)  
-            # When b1 = 0, robot j's y_prime coordinate is greater than robot i's
-            constraints.append(robot_j_y_prime - robot_i_y_prime <= M * (1 - b1)) 
-            # Enforces separation by at least 1 between the robots' y indices on the grid
-            constraints.append((2 * b1 - 1) * (grid_indices[i, 1] - grid_indices[j, 1]) + M * (1 - y_separated) >= 1)
-                
-            # Robots must be separated in at least one of the x, y directions
-            constraints.append(x_separated + y_separated >= 1)
+            # Enforces separation by at least 1 in the y direction
+            b1 = cp.Variable(boolean=True) # b1 = 0 if robot i's y >= robot j's y, 1 otherwise
+            # b1 = 0
+            constraints.append(robot_locations[j, 1] - robot_locations[i, 1] <= M_cts * b1)
+            constraints.append(grid_indices[i, 1] - grid_indices[j, 1] + M_ind * b1 + M_ind * (1 - ysep) >= 1)
+            # b1 = 1
+            constraints.append(robot_locations[i, 1] - robot_locations[j, 1] <= M_cts * (1 - b1))
+            constraints.append(grid_indices[j, 1] - grid_indices[i, 1] + M_ind * (1 - b1) + M_ind * (1 - ysep) >= 1)
     
     # Solve the optimization problem
+    prob_init_start_time = time.time()
     prob = cp.Problem(cp.Minimize(cost), constraints)
+    solve_start_time = time.time()
     prob.solve(solver=cp.SCIP)
+    solve_end_time = time.time()
+    
+    print("Time to add vars/constraints:", prob_init_start_time - start_time)
+    print("Time to parse:", solve_start_time - prob_init_start_time)
+    print("Time to solve:", solve_end_time - solve_start_time)
 
-    if prob.status not in ["optimal", "optimal_inaccurate"]:
+    if prob.status != "optimal":
         print("Problem could not be solved to optimality.")
         return None
     
-    if return_loss:
-        return bottom_left.value, top_right.value, cell_centers.value, prob.value
-    return bottom_left.value, top_right.value, cell_centers.value
+    print("Grid Indices:", grid_indices.value)
+    
+    return bottom_left.value, cell_centers.value
+
+
+def mccormick_envelope(w, x, xl, xu, y, yl, yu):
+    """
+    Generates McCormick envelope constraints
+    """
+    mec = []
+    mec.append(w >= xl*y + x*yl - xl*yl)
+    mec.append(w >= xu*y + x*yu - xu*yu)
+    mec.append(w <= xu*y + x*yl - xu*yl)
+    mec.append(w >= x*yu + xl*y - xl*yu)
+    return mec
 
 
 def plot_grid(bottom_left, top_right, grid_size):
@@ -211,13 +254,52 @@ 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 main(seed, num_robots, plot):
+def get_roomba_locs(low, high, num_robots, radius=0.5):
+    """
+    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
+    """
+    locs = []
+    while len(locs) < num_robots:
+        locs.append(np.random.uniform(low, high, 2))
+        for other_loc in locs[:-1]:
+            if np.linalg.norm(np.array(locs[-1]) - np.array(other_loc)) <= 2 * radius:
+                locs = locs[:-1]
+                break
+    return np.array(locs)
+
+
+def main(seed, num_robots, plot, two_corner):
     if seed is not None:
         np.random.seed(seed)
     
-    robot_locations = np.random.uniform(low=0, high=7.5, size=(num_robots, 2))
-    cell_size = 1.5
-    grid_shape = (5, 5)
+    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]])
+        
+        bottom_left, cell_centers = place_grid(robot_locations=robot_locations, 
+                                               cell_size=cell_size,
+                                               grid_size=grid_size,
+                                               subgoals=subgoals)
+        
+        print("Grid Origin (Bottom-Left Corner):", bottom_left)
+        print("Cell Centers:", cell_centers)
+        
+        top_right = np.array(bottom_left) + grid_size * cell_size
+    else:
+        grid_size = 5
+        robot_locations = np.random.uniform(low=0, high=5, size=(num_robots, 2))
+        print("Robot Locations:", robot_locations)
+        
+        bottom_left, top_right, grid_indices = two_corner_place_grid(robot_locations, grid_size)
+        print("Grid Bottom-Left Corner:", bottom_left)
+        print("Grid Top-Right Corner:", top_right)
+        print("Grid Indices:", grid_indices)
     
     if plot:
         import matplotlib.pyplot as plt
@@ -225,12 +307,7 @@ def main(seed, num_robots, plot):
         
         fig, ax = plt.subplots()
         
-        origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape)
-        print("Grid Origin (Bottom-Left Corner):", origin)
-        
-        
-        top_right = np.array(origin) + np.array(grid_shape) * cell_size
-        plot_grid(origin, top_right, grid_size=5)
+        plot_grid(bottom_left, top_right, grid_size=grid_size)
         
         # Plot cell centers
         cell_centers = np.array(cell_centers)
@@ -243,13 +320,18 @@ def main(seed, num_robots, plot):
         robot_locations = np.array(robot_locations)
         plt.scatter(robot_locations[:, 0], robot_locations[:, 1], c='r', label='Robot Locations')
         for (x, y) in robot_locations:
-            circle = patches.Circle((x, y), radius=0.5, edgecolor='r', fill=False, linewidth=2)
+            circle = patches.Circle((x, y), radius=roomba_radius, edgecolor='r', fill=False, linewidth=2)
             ax.add_patch(circle)
+            
+        if not two_corner:
+            subgoals = np.array(subgoals)
+            plt.scatter(subgoals[:, 0], subgoals[:, 1], c='orange', marker='^', label='Subgoals')
+            for (x, y) in subgoals:
+                circle = patches.Circle((x, y), radius=roomba_radius, edgecolor='orange', fill=False, linewidth=2)
+                ax.add_patch(circle)
         
         plt.legend(loc='upper left')
         
-        # ax.set_xlim(0, 5)
-        # ax.set_ylim(0, 5)   
         ax.set_aspect('equal')
 
         plt.show()
@@ -266,13 +348,16 @@ if __name__ == "__main__":
     parser.add_argument(
         "--num_robots", 
         type=int, 
-        default=2
+        default=3
     )
     parser.add_argument(
         "--plot", 
-        type=bool, 
-        default=False
+        action='store_true'
+    )
+    parser.add_argument(
+        "--two_corner", 
+        action='store_true'
     )
     args = parser.parse_args()
 
-    main(args.seed, args.num_robots, args.plot)
+    main(args.seed, args.num_robots, args.plot, args.two_corner)
\ No newline at end of file