diff --git a/guided_mrmp/controllers/place_grid.py b/guided_mrmp/controllers/place_grid.py index 81d5b0ba84c17c1721c59311f8a5310569ee09a2..192ad8f0c865ee0e1139c819d095cc9002dbfddc 100644 --- a/guided_mrmp/controllers/place_grid.py +++ b/guided_mrmp/controllers/place_grid.py @@ -1,7 +1,7 @@ import cvxpy as cp import numpy as np -def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5)): +def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5), return_loss=False): """ Place a grid to cover robot locations with alignment to centers. @@ -11,6 +11,8 @@ def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5)): - grid_shape (tuple): (# of rows, # of columns) of the grid 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) @@ -70,17 +72,80 @@ def place_grid(robot_locations, cell_size=1, grid_shape=(5, 5)): if prob.status not in ["optimal", "optimal_inaccurate"]: 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 + +def place_grid_with_rotation(robot_locations, cell_size=1, grid_shape=(5, 5), num_angles=18, return_angle_info=False): + """ + Place a grid to cover robot locations with alignment to centers. Allows rotations of the grid, but only samples + + 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 + - num_angles (int): number of evenly spaced angles to sample between 0 and 90 degrees + outputs: + - grid_center (tuple): center 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 + """ + + # Dictionary of { grid rotation angle : ( grid_center, cell_centers, loss ) for that angle } + angle_info = {} + + # The rotation angle from the angles sampled which minimizes loss + min_loss_angle = 0 + min_loss = None + + for angle in np.linspace(start=0, stop=np.pi/2, num=num_angles, endpoint=False): + print(angle) + rotation_matrix = np.array([[np.cos(angle), -np.sin(angle)], + [np.sin(angle), np.cos(angle)]]) -def main(): - np.random.seed(54) + # 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 + + angle_info[angle] = grid_center, cell_centers, loss + + if return_angle_info: + return angle_info[min_loss_angle][0], angle_info[min_loss_angle][1], angle_info + return angle_info[min_loss_angle][0], angle_info[min_loss_angle][1] + +def main(rotations): + np.random.seed(52) robot_locations = np.random.uniform(low=0, high=5, size=(5, 2)) cell_size = 1 grid_shape = (5, 5) - origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape) - print("Grid Origin (Bottom-Left Corner):", origin) - print(cell_centers) + import matplotlib.pyplot as plt + + if rotations: + grid_center, cell_centers, angle_info = place_grid_with_rotation(robot_locations, cell_size, grid_shape, num_angles=30, return_angle_info=True) + # angles = angle_info.keys() + # losses = [loss for (_, _, loss) in angle_info.values()] + + # plt.scatter(angles, losses, c='r') + + # plt.show() + else: + origin, cell_centers = place_grid(robot_locations, cell_size, grid_shape) + print("Grid Origin (Bottom-Left Corner):", origin) + print(cell_centers) import matplotlib.pyplot as plt @@ -112,4 +177,14 @@ def main(): plt.show() if __name__ == "__main__": - main() \ No newline at end of file + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--rotations", + type=bool, + required=True + ) + args = parser.parse_args() + + main(args.rotations) \ No newline at end of file