Merge branch 'updated-place-grid' into 'main'

Updated place grid

See merge request !2

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
+            - subgoals (list): locations of subgoals for each robot
+            - obstacles (list): locations of circular ('c') and rectangular ('r') obstacles [['c',x,y,r], ['c',x,y,r], ['r',x,y,w,h], ['r',x,y,w,h], ...]
         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
     """
     robot_locations = np.array(robot_locations)
-    N = len(robot_locations)
+    subgoals = np.array(subgoals)
+    num_robots = len(robot_locations)
+    num_obst = len(obstacles)
+    
+    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
     
     # Decision variable: Bottom-left corner of the grid in continuous space
-    origin = cp.Variable(2, name='origin')
+    bottom_left = cp.Variable(2, name='origin')
     
     # 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(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)
+    # Existence of overlap between each obstacle and each robot
+    if num_obst > 0:
+        overlaps = cp.Variable((num_obst, num_robots), boolean=True)
+    
+    # Objective: Minimize the sum of squared distances and number of robot cell / obstacle overlaps
+    if num_obst > 0:
+        alpha = 1
+        cost = cp.sum_squares(robot_locations - cell_centers) + alpha * cp.sum(overlaps)
+    else:
+        cost = cp.sum_squares(robot_locations - cell_centers)
     
     # Constraints
     constraints = []
@@ -37,171 +52,416 @@ 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
-    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
-            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
+    top_right = bottom_left + grid_size * cell_size
+    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)
+
+    # Determine overlaps between obstacles and robots
+    for obst_idx, obst_info in enumerate(obstacles):
+        
+        # Define the obstacle's bounds in grid coordinates
+        if obst_info[0] == 'c':
+            _, cx, cy, r = obst_info
+            x_min = (cx - r - bottom_left[0]) / cell_size
+            x_max = (cx + r - bottom_left[0]) / cell_size
+            y_min = (cy - r - bottom_left[1]) / cell_size
+            y_max = (cy + r - bottom_left[1]) / cell_size
+        elif obst_info[0] == 'r':
+            _, blx, bly, w, h = obst_info
+            x_min = (blx - bottom_left[0]) / cell_size
+            x_max = (blx + w - bottom_left[0]) / cell_size
+            y_min = (bly - bottom_left[1]) / cell_size
+            y_max = (bly + h - bottom_left[1]) / cell_size
+             
+        for i in range(num_robots):
+            # Define temp binary variables for each condition
+            temp_x_min = cp.Variable(boolean=True)
+            temp_x_max = cp.Variable(boolean=True)
+            temp_y_min = cp.Variable(boolean=True)
+            temp_y_max = cp.Variable(boolean=True)
+            
+            # Enforce that robots cannot occupy cells overlapping with obstacles
+            buffer = 0.05
+            constraints.append(grid_indices[i, 0] + 1 + buffer <= x_min + M_ind * (1 - temp_x_min))
+            constraints.append(grid_indices[i, 0] - buffer >= x_max - M_ind * (1 - temp_x_max))
+            constraints.append(grid_indices[i, 1] + 1 + buffer <= y_min + M_ind * (1 - temp_y_min))
+            constraints.append(grid_indices[i, 1] - buffer >= y_max - M_ind * (1 - temp_y_max))
+            
+            # Define boolean variables for overlaps in the x-direction and y-direction
+            temp_x_sep = cp.Variable(boolean=True)
+            temp_y_sep = cp.Variable(boolean=True)
+            constraints.append(temp_x_min + temp_x_max >= 1 - temp_x_sep)
+            constraints.append(temp_y_min + temp_y_max >= 1 - temp_y_sep)
+            
+            # The obstacle and cell overlap if there is overlap in both directions
+            constraints.append(overlaps[obst_idx, i] <= temp_x_sep)
+            constraints.append(overlaps[obst_idx, i] <= temp_y_sep)
+            constraints.append(overlaps[obst_idx, i] >= temp_x_sep + temp_y_sep - 1)                
+
     # Solve the optimization problem
     prob = cp.Problem(cp.Minimize(cost), constraints)
+    
+    solve_start_time = time.time()
     prob.solve(solver=cp.SCIP)
+    solve_end_time = time.time()
+    
+    print("Solve time:", 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
 
 
-def place_grid_with_rotation(robot_locations, cell_size=1, grid_shape=(5, 5), num_angles=18, return_angle_info=False):
+# Not convex
+def two_corner_place_grid(robot_locations, grid_size=5, subgoals=[], obstacles=[]):
     """
-        Place a grid to cover robot locations with alignment to centers. 
+        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
-            - sample_size (int): number of evenly spaced angles to sample between 0 and pi/2 radians
+            - grid_size (tuple): width of the grid in cells
+            - obstacles (list): locations of circular obstacles [[x,y,r], [x,y,r], ...]
         outputs:
-            - grid_center (tuple): center 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)
-            - angle_info (dict): when return_angle_info=True, dict of { grid rotation angle : ( grid_center, cell_centers, loss ) for that angle
-    """ 
+    """
+    start_time = time.time()
     
-    # Dictionary of { grid rotation angle : ( grid_center, cell_centers, loss ) for that angle }
-    angle_info = {}
+    robot_locations = np.array(robot_locations)
+    subgoals = np.array(subgoals)
+    obstacles = np.array(obstacles)
+    N = len(robot_locations)
     
-    # The rotation angle from the angles sampled which minimizes loss
-    min_loss_angle = 0
-    min_loss = None 
+    # 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')
     
-    from tqdm import tqdm
+    # 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')
     
-    for angle in tqdm(np.linspace(start=0, stop=np.pi/2, num=num_angles, endpoint=False), desc="Processing..."):
-        rotation_matrix = np.array([[np.cos(angle), -np.sin(angle)],
-                                    [np.sin(angle), np.cos(angle)]])
-        
-        # Rotate robot locations about the origin in continuous space
-        rotated_robot_locations = robot_locations @ rotation_matrix
-        
-        # Run place_grid on transformed robot locations
-        rotated_origin, rotated_cell_centers, loss = place_grid(rotated_robot_locations, cell_size, grid_shape, return_loss=True)
-        rotated_grid_center = rotated_origin + np.array(grid_shape) * cell_size / 2
-        
-        # Undo the rotation transformation for the origin and cell centers 
-        grid_center = rotated_grid_center @ np.linalg.inv(rotation_matrix)
-        cell_centers = rotated_cell_centers @ np.linalg.inv(rotation_matrix)
-        
-        # Check if the loss is the smallest observed so far
-        if min_loss is None or loss < min_loss:
-            min_loss_angle = angle
-            min_loss = loss
+    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')
+    
+    # Calculate cell centers for each robot based on grid indices
+    # 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)
+    
+    # Constraints
+    constraints = []
+    
+    # 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)
+    
+    # Grid indices must fit within grid bounds
+    constraints.append(grid_indices <= grid_size - 1)
+    
+    # No two robots can share a cell
+    # 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):
+            # 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)
             
-        angle_info[angle] = grid_center, cell_centers, loss
+            # 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)
+            
+            # 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()
     
-    if return_angle_info:
-        return *angle_info[min_loss_angle][:2], min_loss_angle, angle_info
-    return *angle_info[min_loss_angle][:2], min_loss_angle
+    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 != "optimal":
+        print("Problem could not be solved to optimality.")
+        return None
+    
+    print("Grid Indices:", grid_indices.value)
+    
+    return bottom_left.value, cell_centers.value
 
-def plot_grid(origin, cell_size, grid_shape):
+
+def plot_grid(bottom_left, top_right, grid_size):
     import matplotlib.pyplot as plt
     
+    bottom_left = np.array(bottom_left)
+    top_right = np.array(top_right)
+    bottom_right = np.array([bottom_left[0] + top_right[0] - bottom_left[1] + top_right[1],
+                             bottom_left[0] - top_right[0] + bottom_left[1] + top_right[1]]) / 2
+    top_left = np.array([bottom_left[0] + top_right[0] + bottom_left[1] - top_right[1],
+                         -bottom_left[0] + top_right[0] + bottom_left[1] + top_right[1]]) / 2
+    
+    x_prime_hat = (bottom_right - bottom_left) / grid_size
+    y_prime_hat = (top_left - bottom_left) / grid_size
+    
     # Draw the grid
-    for i in range(grid_shape[1] + 1):
+    for i in range(grid_size + 1):
         # Draw vertical lines
-        plt.plot([origin[0] + i * cell_size, origin[0] + i * cell_size], 
-                    [origin[1], origin[1] + grid_shape[0] * cell_size], 'k-')
-    for i in range(grid_shape[0] + 1):
+        plt.plot([(bottom_left + i * x_prime_hat)[0], (top_left + i * x_prime_hat)[0]], 
+                 [(bottom_left + i * x_prime_hat)[1], (top_left + i * x_prime_hat)[1]], 'k-')
         # Draw horizontal lines
-        plt.plot([origin[0], origin[0] + grid_shape[1] * cell_size], 
-                    [origin[1] + i * cell_size, origin[1] + i * cell_size], 'k-')
-        
+        plt.plot([(bottom_left + i * y_prime_hat)[0], (bottom_right + i * y_prime_hat)[0]], 
+                 [(bottom_left + i * y_prime_hat)[1], (bottom_right + i * y_prime_hat)[1]], 'k-')
+
+
+def get_locations(low, high, num_robots, robot_radius=0.5, circ_obstacle_radii=[0.5, 1, 0.3], rect_obstacle_wh=[]):
+    
+    from shapely.geometry import Point, Polygon
+    
+    # Sort all entities by size
+    queue = []
+    for i in range(num_robots):
+        queue.append((robot_radius, "robot"))
+        queue.append((robot_radius, "subgoal"))
+    for r in circ_obstacle_radii:
+        queue.append((r, "c_obstacle"))
+    for (w, h) in rect_obstacle_wh:
+        queue.append(((w, h), "r_obstacle"))
+    queue = sorted(queue, key=lambda x: x[0] if x[1] != "r_obstacle" else max(x[0]))
+    
+    robot_locs = []
+    robot_geoms = []
+    subgoal_locs = []
+    subgoal_geoms = []
+    obst_locs = []
+    obst_geoms = []
+    iter = 1
+    while len(queue) > 0:
+        curr_dims, curr_type = queue[-1]
+        curr_loc = np.random.uniform(low, high, 2)
         
-def plot_grid_with_rotation(grid_center, cell_size, grid_shape, rotation_angle):
-    rotation_matrix = np.array([[np.cos(rotation_angle), -np.sin(rotation_angle)],
-                                [np.sin(rotation_angle), np.cos(rotation_angle)]])
-    grid_center = np.array(grid_center)
+        if curr_type == "r_obstacle":
+            w, h = curr_dims
+            rect_corners = np.array([(0, 0), (w, 0), (w, h), (0, h)])
+            curr_geom = Polygon(rect_corners + np.array(curr_loc))
+        else:
+            curr_geom = Point(*curr_loc).buffer(curr_dims)
+
+        valid = True
+        for geom in obst_geoms:
+            if curr_geom.intersects(geom):
+                valid = False
+                break    
+        if curr_type != "subgoal":     
+            for geom in robot_geoms:
+                if curr_geom.intersects(geom):
+                    valid = False
+                    break 
+            for geom in subgoal_geoms:
+                if curr_geom.intersects(geom):
+                    valid = False
+                    break 
+            
+        if valid:
+            match curr_type:
+                case "robot":
+                    robot_locs.append(curr_loc)
+                    robot_geoms.append(curr_geom)
+                case "subgoal":
+                    subgoal_locs.append(curr_loc)
+                    subgoal_geoms.append(curr_geom)
+                case "c_obstacle":
+                    obst_locs.append(['c', *curr_loc, curr_dims])
+                    obst_geoms.append(curr_geom)
+                case "r_obstacle":
+                    obst_locs.append(['r', *curr_loc, *curr_dims])
+                    obst_geoms.append(curr_geom)
+            queue = queue[:-1]
+        iter += 1
     
-    import matplotlib.pyplot as plt
+    return robot_locs, subgoal_locs, obst_locs
+
+
+def main(seed, num_robots):
+    np.random.seed(1)
     
-    # Draw the grid
-    for i in range(grid_shape[1] + 1):
-        # Construct vertical line endpoints
-        line_points = np.array([[(-grid_shape[1] / 2 + i) * cell_size, (-grid_shape[1] / 2 + i) * cell_size],
-                                [-grid_shape[0] / 2, grid_shape[0] / 2]])
-        # Rotate vertical lines and draw
-        rotated_line_points = rotation_matrix @ line_points + grid_center[:, None]
-        plt.plot(rotated_line_points[0], rotated_line_points[1], 'k-')
-    for i in range(grid_shape[1] + 1):
-        # Construct horizontal line endpoints
-        line_points = np.array([[-grid_shape[1] / 2, grid_shape[1] / 2],
-                                [(-grid_shape[0] / 2 + i) * cell_size, (-grid_shape[0] / 2 + i) * cell_size]])
-        # Rotate horizontal lines and draw
-        rotated_line_points = rotation_matrix @ line_points + grid_center[:, None]
-        plt.plot(rotated_line_points[0], rotated_line_points[1], 'k-')
-
-
-def main(allow_rotations, num_robots, num_angles, seed):
     if seed is not None:
         np.random.seed(seed)
     
-    robot_locations = np.random.uniform(low=0, high=5, size=(num_robots, 2))
-    cell_size = 1
-    grid_shape = (5, 5)
+    roomba_radius = 0.5
+    cell_size = 2.5 * roomba_radius
+    grid_size = 5
+    circ_obstacle_radii = [0.8]
+    rect_obstacle_wh = [(2, 1.5), (1.1, 1.1)]
+    
+    robot_locations, subgoals, obstacles = get_locations(low=0, 
+                                                            high=6, 
+                                                            num_robots=num_robots, 
+                                                            robot_radius=roomba_radius,
+                                                            circ_obstacle_radii=circ_obstacle_radii,
+                                                            rect_obstacle_wh=rect_obstacle_wh)
+
+    bottom_left, cell_centers = place_grid(robot_locations=robot_locations, 
+                                            cell_size=cell_size,
+                                            grid_size=grid_size,
+                                            subgoals=subgoals,
+                                            obstacles=obstacles)
+    
+    print("Grid Origin (Bottom-Left Corner):", bottom_left)
+    print("Cell Centers:", cell_centers)
     
     import matplotlib.pyplot as plt
+    import matplotlib.patches as patches
     
-    plt.figure(figsize=(6, 6))
+    fig, ax = plt.subplots()
     
-    if allow_rotations:
-        grid_center, cell_centers, rotation_angle = place_grid_with_rotation(robot_locations, cell_size, grid_shape, num_angles=num_angles)
-        print("Grid Center:", grid_center)
-        # angles = angle_info.keys()
-        # losses = [loss for (_, _, loss) in angle_info.values()]
-        # plt.scatter(angles, losses, c='r')
-        
-        plot_grid_with_rotation(grid_center, cell_size, grid_shape, rotation_angle)
-    else:
-        origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape)
-        print("Grid Origin (Bottom-Left Corner):", origin)
-        
-        plot_grid(origin, cell_size, grid_shape)
+    top_right = np.array(bottom_left) + grid_size * cell_size
+    plot_grid(bottom_left, top_right, grid_size=grid_size)
+    
+    # Plot cell centers
+    cell_centers = np.array(cell_centers)
+    # plt.scatter(cell_centers[:, 0], cell_centers[:, 1], c='r', label='Cell Centers')
+    for center in cell_centers:
+        square = patches.Rectangle(center - cell_size/2, cell_size, cell_size, edgecolor='r', facecolor='r', alpha=0.2, linewidth=2)
+        ax.add_patch(square)
     
     # Plot robot locations
     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=roomba_radius, edgecolor='r', fill=False, linewidth=2)
+        ax.add_patch(circle)
 
-    # Plot cell centers
-    cell_centers = np.array(cell_centers)
-    plt.scatter(cell_centers[:, 0], cell_centers[:, 1], c='b', label='Cell Centers')
+    # Plot subgoals
+    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)
+
+    from shapely.geometry import Point, Polygon
+    
+    circ_obstacles = np.array([obst[1:] for obst in obstacles if obst[0] == 'c'])
+    rect_obstacles = np.array([obst[1:] for obst in obstacles if obst[0] == 'r'])
+
+    # Plot circular obstacles
+    # plt.scatter(circ_obstacles[:, 0], circ_obstacles[:, 1], c='black', marker='s', label='Obstacles')
+    for (x, y, r) in circ_obstacles:
+        circle = patches.Circle((x, y), radius=r, edgecolor='black', fill=False, linewidth=2)
+        ax.add_patch(circle)
+        obstacle_geom = Point(x, y).buffer(r)
+        
+        for x_idx in range(grid_size):
+            for y_idx in range(grid_size):
+                cell_corners = np.array([(0, 0), (cell_size, 0), (cell_size, cell_size), (0, cell_size)])
+                cell_corners += np.array(bottom_left)
+                cell_corners += np.array([x_idx, y_idx]) * cell_size
+                cell_geom = Polygon(cell_corners)
+                
+                intersection = cell_geom.intersection(obstacle_geom)
+                if cell_geom.intersects(obstacle_geom):
+                    obstacle_cell = patches.Rectangle(cell_corners[0], cell_size, cell_size, edgecolor='black', facecolor='black', alpha=0.2+0.5*intersection.area/(cell_size**2), linewidth=2)
+                    ax.add_patch(obstacle_cell)
     
+    # Plot rectangular obstacles
+    for (x, y, w, h) in rect_obstacles:
+        rectangle = patches.Rectangle((x, y), w, h, edgecolor='black', fill=False, linewidth=2)
+        ax.add_patch(rectangle)
+        obstacle_geom = Polygon([(x, y), (x+w, y), (x+w, y+h), (x, y+h)])
+        
+        for x_idx in range(grid_size):
+            for y_idx in range(grid_size):
+                cell_corners = np.array([(0, 0), (cell_size, 0), (cell_size, cell_size), (0, cell_size)])
+                cell_corners += np.array(bottom_left)
+                cell_corners += np.array([x_idx, y_idx]) * cell_size
+                cell_geom = Polygon(cell_corners)
+                
+                intersection = cell_geom.intersection(obstacle_geom)
+                if cell_geom.intersects(obstacle_geom):
+                    obstacle_cell = patches.Rectangle(cell_corners[0], cell_size, cell_size, edgecolor='black', facecolor='black', alpha=0.2+0.5*intersection.area/(cell_size**2), linewidth=2)
+                    ax.add_patch(obstacle_cell)
+                        
     plt.legend(loc='upper left')
+    ax.set_aspect('equal')
 
     plt.show()
 
@@ -210,25 +470,15 @@ if __name__ == "__main__":
     
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--allow_rotations", 
-        type=bool, 
-        default=False
-    )
-    parser.add_argument(
-        "--num_robots", 
-        type=int, 
-        default=2
-    )
-    parser.add_argument(
-        "--num_angles", 
+        "--seed", 
         type=int, 
-        default=18
+        default=None
     )
     parser.add_argument(
-        "--seed", 
+        "--num_robots", 
         type=int, 
-        default=None
+        default=3
     )
     args = parser.parse_args()
 
-    main(args.allow_rotations, args.num_robots, args.num_angles, args.seed)
\ No newline at end of file
+    main(args.seed, args.num_robots)
\ No newline at end of file