Random generation of obstacle/robot/subgoals for testcases

guided_mrmp/controllers/place_grid.py

@@ -122,7 +122,7 @@ def place_grid(robot_locations, cell_size, grid_size=5, subgoals=[], obstacles=[
     prob_init_start_time = time.time()
     prob = cp.Problem(cp.Minimize(cost), constraints)
     solve_start_time = time.time()
-    prob.solve(solver=cp.SCIP, verbose=True)
+    prob.solve(solver=cp.SCIP)
     solve_end_time = time.time()
     
     print("Time to add vars/constraints:", prob_init_start_time - start_time)
@@ -296,26 +296,57 @@ 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, obstacles=[]):
+def get_locations(low, high, num_robots, robot_radius=0.5, obstacle_radii=[0.5, 1, 0.3], max_iter=500):
     """
     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))
-        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:
-                invalid = True
-                break
-        if invalid:
-            locs = locs[:-1]
-    return np.array(locs)
+    num_obst = len(obstacle_radii)
+    
+    queue = []
+    for i in range(num_robots):
+        queue.append((robot_radius, "robot"))
+        queue.append((robot_radius, "subgoal"))
+    for i in range(num_obst):
+        queue.append((obstacle_radii[i], "obstacle"))
+    queue = sorted(queue, key=lambda x: x[0])
+    
+    robot_locs = []
+    subgoal_locs = []
+    obst_locs = []
+    while len(queue) > 0:
+        curr_r, curr_type = queue[-1]
+        curr_loc = np.random.uniform(low, high, 2)
+        # match type:
+        #     case "robot":
+        #         robot_locs.append()
+        #     case "subgoal":
+        #     case "obstacle":
+        valid = True
+        for obst in obst_locs:
+            if np.linalg.norm(curr_loc - obst[:2]) <= curr_r + obst[2]:
+                valid = False
+                break    
+        if curr_type != "subgoal":     
+            for robot in robot_locs:
+                if np.linalg.norm(curr_loc - robot) <= curr_r + robot_radius:
+                    valid = False
+                    break
+            for subgoal in subgoal_locs:
+                if np.linalg.norm(curr_loc - subgoal) <= curr_r + robot_radius:
+                    valid = False
+                    break
+            
+        if valid:
+            match curr_type:
+                case "robot":
+                    robot_locs.append(curr_loc)
+                case "subgoal":
+                    subgoal_locs.append(curr_loc)
+                case "obstacle":
+                    obst_locs.append(np.array([*curr_loc, curr_r]))
+            queue = queue[:-1]
+    return robot_locs, subgoal_locs, obst_locs
 
 
 def main(seed, num_robots, plot, two_corner):
@@ -328,18 +359,14 @@ def main(seed, num_robots, plot, two_corner):
         roomba_radius = 0.5
         cell_size = 2.5 * roomba_radius
         grid_size = 5
+        obstacle_radii = [1/3, 1, 0.5]
         
-        
-        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,
-        #                                        grid_size=grid_size,
-        #                                        subgoals=subgoals)
+        robot_locations, subgoals, obstacles = get_locations(low=0, 
+                                                             high=5, 
+                                                             num_robots=3, 
+                                                             robot_radius=roomba_radius,
+                                                             obstacle_radii=obstacle_radii)
+
         bottom_left, cell_centers = place_grid(robot_locations=robot_locations, 
                                                cell_size=cell_size,
                                                grid_size=grid_size,