From 28025e5ebfbb473fd2ae6437293d67cf9cb7da76 Mon Sep 17 00:00:00 2001
From: rachelmoan <moanrachel516@gmail.com>
Date: Wed, 5 Mar 2025 12:56:42 -0600
Subject: [PATCH] Do not use ego frame in MPC, use global frame throughout

---
 guided_mrmp/controllers/multi_mpc.py          | 133 +++--
 .../controllers/multi_path_tracking.py        | 190 +++++--
 .../controllers/multi_path_tracking_db.py     | 509 ++++++++++--------
 guided_mrmp/controllers/optimizer.py          |   9 +-
 guided_mrmp/controllers/path_tracker.py       | 265 +++++++--
 guided_mrmp/controllers/utils.py              | 100 ++--
 guided_mrmp/controllers/vis_helpers.py        | 304 ++++++++++-
 guided_mrmp/simulator.py                      |  17 +-
 8 files changed, 1122 insertions(+), 405 deletions(-)

diff --git a/guided_mrmp/controllers/multi_mpc.py b/guided_mrmp/controllers/multi_mpc.py
index cb64663..febd1a5 100644
--- a/guided_mrmp/controllers/multi_mpc.py
+++ b/guided_mrmp/controllers/multi_mpc.py
@@ -19,12 +19,16 @@ class MultiMPC:
 
         self.circle_obs = circle_obs
         self.rectangle_obs = rectangle_obs
+        self.obstacle_bias = 0.05
 
         state_cost = settings['model_predictive_controller']['Q']  # state error cost
         final_state_cost = settings['model_predictive_controller']['Qf']  # state final error cost
         input_cost = settings['model_predictive_controller']['R']  # input cost
         input_rate_cost = settings['model_predictive_controller']['P']  # input rate of change cost
 
+        self.bf_c = settings['model_predictive_controller']['barrier_function_c']
+        self.bf_max_d = settings['model_predictive_controller']['barrier_function_max_d']
+
         # how far we can look into the future divided by our dt 
         # is the number of control intervals
         self.control_horizon = int(T / DT) 
@@ -89,47 +93,95 @@ class MultiMPC:
             - Cs: List of C matrices for each robot
         """
 
+        print(f"Initial state: {initial_state}")
+
         opti = ca.Opti()
 
-        # Decision variables
+        # Set up state variable
+        # X is of the form [[x1_k0, y1_k0, theta1_k0, x2_k0, y2_k0, theta2_k0, ...], [x1_k1, y1_k1, theta1_k1, x2_k1, y2_k1, theta2_k1, ...], ...]
         X = opti.variable(self.nx*self.num_robots, self.control_horizon + 1)
         pos = X[:self.num_robots*2,:]               # position is the first two values
         x = pos[0::2,:]
         y = pos[1::2,:]
         heading = X[self.num_robots*2:,:]           # heading is the last value
 
+        # control variables
+        # U is of the form [[v1_k0, omega1_k0, v2_k0, omega2_k0, ...], [v1_k1, omega1_k1, v2_k1, omega2_k1, ...], ...]
         U = opti.variable(self.nu*self.num_robots, self.control_horizon)
 
+        # # parameters, these parameters are the reference trajectories of the pose and inputs
+        # opt_u_ref = opti.parameter(N, 2)
+        # opt_x_ref = opti.parameter(N+1, 3)
+
         # Parameters
-        initial_state = ca.MX(initial_state)
+        # initial_state = ca.MX(initial_state)
+
+        
+
+        # create model
+        f = lambda x_, u_: ca.vertcat(*[u_[0]*ca.cos(x_[2]), u_[0]*ca.sin(x_[2]), u_[1]])
+        f_np = lambda x_, u_: np.array([u_[0]*np.cos(x_[2]), u_[0]*np.sin(x_[2]), u_[1]])
+
+        # initial condition
+        for i in range(self.num_robots):
+            opti.subject_to(X[i*3:i*3+3, 0] == initial_state[i])
+
+        # next state constraints
+        for i in range(self.control_horizon):
+            for j in range(self.num_robots):
+                this_X = X[j*3:j*3+3, i]
+                this_U = U[j*2:j*2+2, i]
+                next_state = X[j*3:j*3+3, i+1]
+                x_next = this_X + f(this_X, this_U)*self.dt
+                opti.subject_to(next_state == x_next)
+
+        # add constraints to obstacle
+        # for i in range(self.control_horizon+1):
+        #     for obs in self.circle_obs:
+        #         for j in range(self.num_robots):
+        #             # print("obs", obs)
+        #             temp_constraints_ = ca.sqrt((x[j, i]-obs[0]-self.obstacle_bias)**2+(y[j, i]-obs[1]-self.obstacle_bias)**2)-self.robot_radius-obs[2]
+        #             opti.subject_to(opti.bounded(.75, temp_constraints_, 20.0))
+
+        # bounds constraints
+        v_max = 20.0
+        omega_max = np.pi/2.0
+        for i in range(self.num_robots):
+            opti.subject_to(opti.bounded(-v_max, U[i*2, :], v_max))
+            opti.subject_to(opti.bounded(-omega_max, U[i*2+1, :], omega_max))
+            # opti.subject_to(ca.fabs(U[i*2, 0] - prev_cmd[i][0]) / self.dt <= self.robot_model.max_d_vel)
+            # opti.subject_to(ca.fabs(U[i*2+1, 0] - prev_cmd[i][1]) / self.dt <= self.robot_model.max_d_steer)
+
+            # for k in range(1, self.control_horizon):
+            #     opti.subject_to(ca.fabs(U[i*2, k] - U[i*2, k-1]) / self.dt <= self.robot_model.max_d_vel)
+            #     opti.subject_to(ca.fabs(U[i*2+1, k] - U[i*2+1, k-1]) / self.dt <= self.robot_model.max_d_steer)
 
-        # print(f"target = {target}")
-        # target = target
-        # prev_cmd = ca.MX(prev_cmd)
-        # As = ca.MX(As)
-        # Bs = ca.MX(Bs)
-        # Cs = ca.MX(Cs)
+    
 
-        # Cost function
+        # Set up cost function
         cost = 0
         for k in range(self.control_horizon):
             for i in range(self.num_robots):
                 this_target = [target[i][0][k], target[i][1][k], target[i][2][k]]
+
+
                 # difference between the current state and the target state
-                cost += ca.mtimes([(X[i*3 : i*3 +3, k+1] - this_target).T, self.Q, X[i*3 : i*3 +3, k+1] - this_target])
+                # cost += ca.mtimes([(X[i*3 : i*3 +3, k+1] - this_target).T, self.Q, X[i*3 : i*3 +3, k+1] - this_target])
             
+                state_error = X[i*3 : i*3 +3, k] - this_target
+                cost += ca.mtimes([state_error.T, self.Q, state_error])
 
                 # control effort
                 cost += ca.mtimes([U[i*2:i*2+2, k].T, self.R, U[i*2:i*2+2, k]])
-            if k > 0:
-                # Penalize large changes in control
-                cost += ca.mtimes([(U[i*2:i*2+2, k] - U[i*2:i*2+2, k-1]).T, self.P, U[i*2:i*2+2, k] - U[i*2:i*2+2, k-1]])
+            # if k > 0:
+            #     # Penalize large changes in control
+            #     cost += ca.mtimes([(U[i*2:i*2+2, k] - U[i*2:i*2+2, k-1]).T, self.P, U[i*2:i*2+2, k] - U[i*2:i*2+2, k-1]])
 
 
         # Final state cost
-        for i in range(self.num_robots):
-            final_target = this_target = [target[i][0][-1], target[i][1][-1], target[i][2][-1]]
-            cost += ca.mtimes([(X[i*3 : i*3 +3, -1] - final_target).T, self.Qf, X[i*3 : i*3 +3, -1] - final_target])
+        # for i in range(self.num_robots):
+        #     final_target = this_target = [target[i][0][-1], target[i][1][-1], target[i][2][-1]]
+        #     cost += ca.mtimes([(X[i*3 : i*3 +3, -1] - final_target).T, self.Qf, X[i*3 : i*3 +3, -1] - final_target])
 
         # robot-robot collision cost
         dist_to_other_robots = 0
@@ -139,43 +191,49 @@ class MultiMPC:
                     if r1 != r2:
                         d = ca.sumsqr(pos[2*r1 : 2*r1+1, k] - pos[2*r2 : 2*r2+1, k]) 
                         d = ca.sqrt(d)
-                        dist_to_other_robots += self.apply_quadratic_barrier(6*self.robot_radius, d-self.robot_radius*2, 1)
+                        dist_to_other_robots += self.apply_quadratic_barrier(self.bf_max_d*self.robot_radius, d-self.robot_radius*2, self.bf_c)
 
         # obstacle collision cost
         obstacle_cost = 0
+        # distance = 0
         for k in range(self.control_horizon):
             for i in range(self.num_robots):
+                print(f"Robot {i} at time {k}")
                 for obs in self.circle_obs:
                     d = ca.sumsqr(x[i, k] - obs[0]) + ca.sumsqr(y[i, k] - obs[1])
                     d = ca.sqrt(d)
-                    obstacle_cost += self.apply_quadratic_barrier(6*self.robot_radius, d-self.robot_radius*2, 1)
-
+                    obstacle_cost += self.apply_quadratic_barrier(self.bf_max_d*(self.robot_radius + obs[2]) , d-self.robot_radius - obs[2], self.bf_c)
                 for obs in self.rectangle_obs:
                     d = self.dist_robot_to_rectangle(x[i, k], y[i, k], self.robot_radius, obs)
                     obstacle_cost += self.apply_quadratic_barrier(6*self.robot_radius, d, 1)
 
         opti.minimize(cost + self.robot_robot_collision_weight*dist_to_other_robots + self.obstacle_collision_weight*obstacle_cost)
-
+       
         # Constraints
-        for i in range(self.num_robots):
-            for k in range(self.control_horizon):
-                A = ca.MX(As[i])
-                B = ca.MX(Bs[i])
-                C = ca.MX(Cs[i])
-                opti.subject_to(X[i*3:i*3+3, k+1] == ca.mtimes(A, X[i*3:i*3+3, k]) + ca.mtimes(B, U[i*2:i*2+2, k]) + C)
+        # for i in range(self.num_robots):
+        #     for k in range(self.control_horizon):
+                # A = ca.MX(As[i])
+                # B = ca.MX(Bs[i])
+                # C = ca.MX(Cs[i])
+                # opti.subject_to(X[i*3:i*3+3, k+1] == ca.mtimes(A, X[i*3:i*3+3, k]) + ca.mtimes(B, U[i*2:i*2+2, k]) + C)
 
-        for i in range(self.num_robots):
-            opti.subject_to(X[i*3:i*3+3, 0] == initial_state[i])
+        # pi = [3.14159]*self.num_robots
+        # pi = np.array(pi)
+        # pi = ca.DM(pi)
 
-        for i in range(self.num_robots):
-            opti.subject_to(opti.bounded(-self.robot_model.max_vel, U[i*2, :], self.robot_model.max_vel))
-            opti.subject_to(opti.bounded(-self.robot_model.max_steer_angle, U[i*2+1, :], self.robot_model.max_steer_angle))
-            opti.subject_to(ca.fabs(U[i*2, 0] - prev_cmd[i][0]) / self.dt <= self.robot_model.max_d_vel)
-            opti.subject_to(ca.fabs(U[i*2+1, 0] - prev_cmd[i][1]) / self.dt <= self.robot_model.max_d_steer)
+        # parameters, these parameters are the reference trajectories of the pose and inputs
+        # opt_u_ref = opti.parameter(N, 2)
+        # opt_x_ref = opti.parameter(N+1, 3)
+
+
+        # for k in range(self.control_horizon): # loop over control intervals
+        #     dxdt = vel[:,k] * ca.cos(heading[:,k])
+        #     dydt = vel[:,k] * ca.sin(heading[:,k])
+        #     dthetadt = omega[:,k]
+        #     opti.subject_to(x[:,k+1]==x[:,k] + self.dt*dxdt)
+        #     opti.subject_to(y[:,k+1]==y[:,k] + self.dt*dydt) 
+        #     opti.subject_to(heading[:,k+1]==ca.fmod(heading[:,k] + self.dt*dthetadt, 2*pi))
 
-            for k in range(1, self.control_horizon):
-                opti.subject_to(ca.fabs(U[i*2, k] - U[i*2, k-1]) / self.dt <= self.robot_model.max_d_vel)
-                opti.subject_to(ca.fabs(U[i*2+1, k] - U[i*2+1, k-1]) / self.dt <= self.robot_model.max_d_steer)
 
 
         return {
@@ -186,7 +244,8 @@ class MultiMPC:
             'target': target,
             'prev_cmd': prev_cmd,
             'cost': cost,
-            'dist_to_other_robots': dist_to_other_robots
+            'dist_to_other_robots': dist_to_other_robots,
+            'obs_cost': obstacle_cost
         }
 
     def solve_optimization_problem(self, problem, initial_guesses=None, solver_options=None):
diff --git a/guided_mrmp/controllers/multi_path_tracking.py b/guided_mrmp/controllers/multi_path_tracking.py
index 26e5852..6cdf729 100644
--- a/guided_mrmp/controllers/multi_path_tracking.py
+++ b/guided_mrmp/controllers/multi_path_tracking.py
@@ -3,7 +3,8 @@ import matplotlib.pyplot as plt
 
 from guided_mrmp.controllers.utils import compute_path_from_wp, get_ref_trajectory
 from guided_mrmp.controllers.multi_mpc import MultiMPC
-from guided_mrmp.controllers.mpc import MPC
+from guided_mrmp.utils import Roomba
+from guided_mrmp.utils import Env
 
 class MultiPathTracker:
     def __init__(self, env, initial_positions, dynamics, target_v, T, DT, waypoints, trees, settings):
@@ -28,6 +29,7 @@ class MultiPathTracker:
         self.DT = DT
         self.target_v = target_v
         self.trees = trees
+        self.trajs = None
 
         self.radius = dynamics.radius
 
@@ -84,62 +86,29 @@ class MultiPathTracker:
             if np.linalg.norm(traj1[0:2, i] - traj2[0:2, i]) < 2*threshold:
                 return True
         return False
-    
-    def ego_to_global_roomba(self, state, mpc_out):
-        """
-        Transforms optimized trajectory XY points from ego (robot) reference
-        into global (map) frame.
-
-        Args:
-            mpc_out (numpy array): Optimized trajectory points in ego reference frame.
-
-        Returns:
-            numpy array: Transformed trajectory points in global frame.
-        """
-        # Extract x, y, and theta from the state
-        x = state[0]
-        y = state[1]
-        theta = state[2]
-
-        # Rotation matrix to transform points from ego frame to global frame
-        Rotm = np.array([
-            [np.cos(theta), -np.sin(theta)],
-            [np.sin(theta), np.cos(theta)]
-        ])
-
-        # Initialize the trajectory array (only considering XY points)
-        trajectory = mpc_out[0:2, :]
-
-        # Apply rotation to the trajectory points
-        trajectory = Rotm.dot(trajectory)
-
-        # Translate the points to the robot's position in the global frame
-        trajectory[0, :] += x
-        trajectory[1, :] += y
-
-        return trajectory
-    
-    def advance(self, screen, state, show_plots=False):
+     
+    def advance(self, state, show_plots=False):
         # optimization loop
         # start=time.time()
 
-        # Get Reference_traj -> inputs are in worldframe
         # 1. Get the reference trajectory for each robot
         targets = []
         for i in range(self.num_robots):
             ref, visited_guide_points = get_ref_trajectory(np.array(state[i]), 
-                                                           np.array(self.paths[i]), 
-                                                           self.target_v, 
+                                                           np.array(self.paths[i]),
                                                            self.T, 
                                                            self.DT, 
-                                                           [])
+                                                           self.visited_points_on_guide_paths[i])
             
             self.visited_points_on_guide_paths[i] = visited_guide_points
 
             targets.append(ref)
         # dynamycs w.r.t robot frame
         # curr_state = np.array([0, 0, self.state[2], 0])
-        curr_states = np.zeros((self.num_robots, 3))
+        # curr_states = np.zeros((self.num_robots, 3))
+        curr_states = state
+        curr_states = np.array(curr_states)
+        print(f"curr_states = {curr_states}")
         x_mpc, u_mpc = self.coupled_mpc.step(
             curr_states,
             targets,
@@ -222,7 +191,7 @@ class MultiPathTracker:
             self.h_history[i].append(self.states[i, 2])
         if show_plots: self.plot_sim()
 
-        self.plot_current_world_state()
+        # self.plot_current_world_state()
         
         while 1:
             # check if all robots have reached their goal
@@ -231,7 +200,7 @@ class MultiPathTracker:
                 return np.asarray([self.x_history, self.y_history, self.h_history])
             
             # plot the current state of the robots
-            self.plot_current_world_state()
+            # self.plot_current_world_state()
             
             # get the next control for all robots
             x_mpc, controls = self.advance(self.states)
@@ -242,17 +211,110 @@ class MultiPathTracker:
 
             self.states = next_states
 
+            print(f"x_mpc = {x_mpc}")
             self.states = np.array(self.states)
             for i in range(self.num_robots):
                 self.x_history[i].append(self.states[i, 0])
                 self.y_history[i].append(self.states[i, 1])
                 self.h_history[i].append(self.states[i, 2])
+                
+                self.optimized_trajectories_hist[i] = x_mpc[i*3:i*3+3, :]
+
+                print(f"robot {i} optimized trajectory = {x_mpc[i*3:i*3+3, :]}")
 
             # use the optimizer output to preview the predicted state trajectory
-            # self.optimized_trajectory = self.ego_to_global(x_mpc.value)
-            if show_plots: self.optimized_trajectory = self.ego_to_global_roomba(x_mpc)
+            # if show_plots: self.optimized_trajectory = x_mpc
             if show_plots: self.plot_sim()
 
+
+    def plot_sim(self):
+
+        plt.clf()
+        import matplotlib.cm as cm
+        colors = cm.rainbow(np.linspace(0, 1, self.num_robots))
+
+        grid = plt.GridSpec(2, 3)
+
+        plt.subplot(grid[0:2, 0:2])
+        plt.title("MPC Simulation \n" + "Simulation elapsed time {}s".format(self.sim_time))
+
+        # circle = plt.Circle((5, 0), .5, color='r', fill=False)
+        # plt.gca().add_artist(circle)
+        for obs in self.circle_obs:
+            circle1 = plt.Circle((obs[0], obs[1]), obs[2], color='k', fill=True)
+            plt.gca().add_artist(circle1)
+
+        for i in range(self.num_robots):
+
+            plt.plot(
+                self.paths[i][0, :],
+                self.paths[i][1, :],
+                c="tab:orange",
+                marker=".",
+                label="reference track",
+            )
+
+            plt.plot(
+                self.x_history[i],
+                self.y_history[i],
+                c=colors[i],
+                marker=".",
+                alpha=0.5,
+                label="vehicle trajectory",
+            )
+
+            # print(f"opt hist shape = {self.optimized_trajectories_hist[i].shape}")
+            
+            if len(self.optimized_trajectories_hist[i]) >0:
+                print(f"opt hist x = {self.optimized_trajectories_hist[i]}")
+                plt.plot(
+                    self.optimized_trajectories_hist[i][0],
+                    self.optimized_trajectories_hist[i][1],
+                    c=colors[i],
+                    marker="+",
+                    alpha=0.5,
+                    label="mpc opt trajectory",
+                )
+
+
+            plot_roomba(self.x_history[i][-1], self.y_history[i][-1], self.h_history[i][-1], colors[i], False, self.radius)
+
+
+        plt.ylabel("map y")
+        plt.yticks(np.arange(min(self.paths[0][1, :]) - 1.0, max(self.paths[0][1, :] + 1.0) + 1, 1.0))
+        plt.xlabel("map x")
+        plt.xticks(np.arange(min(self.paths[0][0, :]) - 1.0, max(self.paths[0][0, :] + 1.0) + 1, 1.0))
+        plt.axis("equal")
+        # plt.legend()
+
+        plt.subplot(grid[0, 2])
+        # plt.title("Linear Velocity {} m/s".format(self.v_history[-1]))
+        # plt.plot(self.a_history, c="tab:orange")
+        # locs, _ = plt.xticks()
+        # plt.xticks(locs[1:], locs[1:] * DT)
+        # plt.ylabel("a(t) [m/ss]")
+        # plt.xlabel("t [s]")
+
+        plt.subplot(grid[1, 2])
+        # plt.title("Angular Velocity {} m/s".format(self.w_history[-1]))
+        # plt.plot(np.degrees(self.d_history), c="tab:orange")
+        # plt.ylabel("gamma(t) [deg]")
+        # locs, _ = plt.xticks()
+        # plt.xticks(locs[1:], locs[1:] * DT)
+        # plt.xlabel("t [s]")
+
+        plt.tight_layout()
+
+        ax = plt.gca()
+        ax.set_xlim([0, 10])
+        ax.set_ylim([0, 10])
+
+        plt.show()
+
+        # plt.draw()
+        # plt.pause(0.1)
+
+
 def plot_roomba(x, y, yaw, color, fill, radius):
     """
 
@@ -273,7 +335,39 @@ def plot_roomba(x, y, yaw, color, fill, radius):
     dy = 1 * np.sin(yaw)
     ax.arrow(x, y, dx, dy, head_width=0.1, head_length=0.1, fc='r', ec='r')
 
+if __name__ == "__main__":
+
+    # Example usage
+
+    file_path = "experiments/settings_files/4robotscircle.yaml"
+    import yaml
+    with open(file_path, 'r') as file:
+        settings = yaml.safe_load(file)
+
+    initial_pos1 = np.array([0.0, 5, 0.0])
+    initial_pos2 = np.array([5.0,0.0, 0.0])
+    dynamics = Roomba(settings)
+    target_vocity = 3.0 # m/s
+    T = 5  # Prediction Horizon [s]
+    DT = 0.2  # discretization step [s]
+    wp1 = [[0, 1,2,3,4,5,6,7,8,9],
+          [5, 5,5,5,5,5,5,5,5,5]]
+    
+    
+    wp2 = [[5, 5,5,5,5,5,5,5,5,5],
+          [0, 1,2,3,4,5,6,7,8,9]]
+
+    # wp1 = [[0, 1,2,3,4,5,6,7,8,9],
+    #       [5, 5,5,5,5,5,5,5,5,5]]
+    
+    
+    # wp2 = [[0, 1,2,3,4,5,6,7,8,9],
+    #       [6, 6,6,6,6,6,6,6,6,6]]
+    
+
+    env = Env([0,10], [0,10], circle_obs=[[5,5,1]], rectangle_obs=[])
 
+    sim = MultiPathTracker(env, [initial_pos1], dynamics, target_vocity, T, DT, [wp1], [], settings)
+    x,y,h = sim.run(show_plots=True)
 
-# if __name__ == "__main__":
-#     main()
+    
\ No newline at end of file
diff --git a/guided_mrmp/controllers/multi_path_tracking_db.py b/guided_mrmp/controllers/multi_path_tracking_db.py
index eaaf148..731eec9 100644
--- a/guided_mrmp/controllers/multi_path_tracking_db.py
+++ b/guided_mrmp/controllers/multi_path_tracking_db.py
@@ -5,7 +5,7 @@ from guided_mrmp.controllers.utils import compute_path_from_wp, get_ref_trajecto
 from guided_mrmp.conflict_resolvers.discrete_resolver import DiscreteResolver
 from guided_mrmp.conflict_resolvers.curve_path import smooth_path, calculate_headings
 from guided_mrmp.utils.helpers import plan_decoupled_path
-from guided_mrmp.controllers.vis_helpers import draw_grid_solution
+from guided_mrmp.controllers.vis_helpers import draw_grid_solution, draw_chosen_subp, draw_temp_goals, draw_grid_world_solution
 
 from guided_mrmp.controllers.multi_mpc import MultiMPC
 from guided_mrmp.controllers.place_grid import place_grid
@@ -20,9 +20,11 @@ class DiscreteRobot:
 
 class MultiPathTrackerDB(MultiPathTracker):
     def get_temp_starts_and_goals(self, state, grid_origin):
-        # the temporary starts are the current positions of the robots snapped to the grid
-        # based on the continuous space location of the robot, we find the cell in the grid that 
-        # includes that continuous space location using the top left of the grid as a reference point
+        """
+        Return the temporary starts and goals of the robots in the grid space. 
+        The temporary starts are the current positions of the robots snapped to the grid
+        The temporary goals are the points at the end of the robot's predicted trajectory
+        """
 
         # find the minimum and maximum x and y values of the grid
         min_x = grid_origin[0]
@@ -45,7 +47,8 @@ class MultiPathTrackerDB(MultiPathTracker):
         # the temmporary goal is the point at the end of the robot's predicted traj
         temp_goals = []
         for r in range(self.num_robots):
-            traj = self.ego_to_global_roomba(state[r], self.trajs[r])
+            # traj = self.ego_to_global_roomba(state[r], self.trajs[r])
+            traj = self.trajs[r]
             x = traj[0][-1]
             y = traj[1][-1]
             cell_x = min(max(math.floor((x - grid_origin[0]) / self.cell_size), 0), self.grid_size - 1)
@@ -57,7 +60,8 @@ class MultiPathTrackerDB(MultiPathTracker):
     def get_subgoals(self, state, robots_in_conflict):
         subgoals = []
         for idx in robots_in_conflict:
-            traj = self.ego_to_global_roomba(state[idx], self.trajs[idx])
+            # traj = self.ego_to_global_roomba(state[idx], self.trajs[idx])
+            traj = self.trajs[idx]
             x = traj[0][-1]
             y = traj[1][-1]
             subgoals.append((x, y))
@@ -108,27 +112,95 @@ class MultiPathTrackerDB(MultiPathTracker):
                     this_conflict.append(r)
             disc_all_conflicts.append(this_conflict)
 
-        print(f"discrete robots = ")
+        grid_solver1 = DiscreteResolver(disc_conflict, disc_robots, conflict_starts, conflict_goals, disc_all_conflicts,grid, libs[0], libs[1], libs[2])
+        subproblem1 = grid_solver1.find_subproblem()
+
+        if subproblem1 is None:
+            print("Couldn't find a discrete subproblem")
+            draw_grid_solution(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, [], state, grid_origin, grid, conflict, None, None, [])
+            return None
+        
+        print(f"old conflict goals = {conflict_goals}")
+        print(f"discrete solver goals = {grid_solver1.goals}")
+        print(f"Subproblem1 goals = {subproblem1.get_goals()}")
+        print(f"subproblem1 temp goals = {subproblem1.temp_goals}")
+
+        # reassign the temp goals in the subproblem by modify subproblem.temp_goals
+        # the new temp goals should be the point furthest along the robot's guide path that is still 
+        # contatined inside the subproblem, then snapped to the grid
+        top_left = subproblem1.top_left
+        bottom_right = subproblem1.bottom_right
+        max_x = grid_origin[0] + (bottom_right[0]+1)*self.cell_size 
+        min_x = grid_origin[0] + top_left[0]*self.cell_size
+        min_y = grid_origin[1] + top_left[1]*self.cell_size
+        max_y = grid_origin[1] + (bottom_right[1]+1)*self.cell_size
+
+        print(f"min_x = {min_x}, max_x = {max_x}, min_y = {min_y}, max_y = {max_y}")
+
+
+        temp_goals = []
+        for i in range(self.num_robots):
+            guide_path = self.paths[i]
+            # for point in guide_path:
+            xs = guide_path[0]
+            ys = guide_path[1]
+            
+            # iterate over x,y values in reverse
+            for x, y in zip(reversed(xs), reversed(ys)):
+                # find the first point that is inside the subproblem
+                if x >= min_x and x <= max_x and y >= min_y and y <= max_y:
+                    cell_x, cell_y = get_grid_cell(x, y, grid_origin, self.grid_size, self.cell_size)
+                    temp_goals.append([cell_x, cell_y])
+                    break
+
+        
+        new_conflict_goals = [temp_goals[c] for c in conflict]
+
+        # create a "discrete" version of the robot, i.e. the 
+        # format that the discrete solver expects to receive 
+        disc_robots = self.create_discrete_robots(starts, temp_goals)
+
+        disc_conflict = []
         for r in disc_robots:
-            print(f"Robot {r.label}")
+            if r.label in conflict:
+                disc_conflict.append(r)
+
+        disc_all_conflicts = []
+        for c in all_conflicts:
+            this_conflict = []
+            for r in disc_robots:
+                if r.label in c:
+                    this_conflict.append(r)
+            disc_all_conflicts.append(this_conflict)
+
+        temp_goals_continuous = [get_grid_cell_location(temp_goals[c][0], temp_goals[c][1], grid_origin, self.cell_size) for c in conflict]
 
-        print(f"discrete conflict = ")
-        for r in disc_conflict:
-            print(f"Robot {r.label}")
+        grid_solver2 = DiscreteResolver(disc_conflict, disc_robots, conflict_starts, new_conflict_goals, disc_all_conflicts,grid, libs[0], libs[1], libs[2])
+        subproblem2 = grid_solver2.find_subproblem()
         
-        print(f"discrete all conflicts = ")
-        for c in disc_all_conflicts:
-            print(f"Conflict = ")
-            for r in c:
-                print(f"Robot {r.label}")
+        print(f"New Conflict goals = {new_conflict_goals}")
+        print(f"discrete solver goals = {grid_solver2.goals}")
+        print(f"Subproblem goals = {subproblem2.get_goals()}")
+        print(f"subproblem temp goals = {subproblem2.temp_goals}")
 
-        grid_solver = DiscreteResolver(disc_conflict, disc_robots, conflict_starts, conflict_goals, disc_all_conflicts,grid, libs[0], libs[1], libs[2])
-        subproblem = grid_solver.find_subproblem()
+        
+
+        top_left = subproblem2.top_left
+        bottom_right = subproblem2.bottom_right
+        max_x = grid_origin[0] + (bottom_right[0]+1)*self.cell_size 
+        min_x = grid_origin[0] + top_left[0]*self.cell_size
+        min_y = grid_origin[1] + top_left[1]*self.cell_size
+        max_y = grid_origin[1] + (bottom_right[1]+1)*self.cell_size
 
-        if subproblem is None:
+        draw_chosen_subp(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, state, grid_origin, grid, conflict, (min_x, min_y), (max_x, max_y))
+        draw_temp_goals(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, state, grid_origin, grid, conflict, (min_x, min_y), (max_x, max_y), temp_goals_continuous, self.paths)
+        
+        if subproblem2 is None:
             print("Couldn't find a discrete subproblem")
+            draw_grid_solution(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, [], state, grid_origin, grid, conflict)
             return None
-        grid_solution = grid_solver.solve_subproblem(subproblem)
+        grid_solution = grid_solver2.solve_subproblem(subproblem2)
+
 
         if grid_solution is None:
             draw_grid_solution([], state, grid_origin, grid, conflict, 0)
@@ -145,6 +217,10 @@ class MultiPathTrackerDB(MultiPathTracker):
                 new_sol.append(point)
             unpadded_grid_solution.append(new_sol)
 
+        print(f"Grid Solution = {unpadded_grid_solution}")
+        draw_grid_world_solution(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, [], state, grid_origin, grid, c, (min_x, min_y), (max_x, max_y), unpadded_grid_solution)
+
+
         return unpadded_grid_solution
     
     def convert_db_sol_to_continuous(self, state, grid_solution, grid_origin, robots_in_conflict, N, cp_dist, smoothing=0.0):
@@ -165,7 +241,7 @@ class MultiPathTrackerDB(MultiPathTracker):
         """
         num_robots = len(robots_in_conflict)
         estimated_state = np.zeros((num_robots*3, N+1))
-        final_estimated_state = np.zeros((num_robots*3, 12+N+1))
+        final_estimated_state = np.zeros((num_robots*3, 15+N+1))
 
 
         for i in range(num_robots):
@@ -175,110 +251,100 @@ class MultiPathTrackerDB(MultiPathTracker):
             smoothed_curve, _ = smooth_path(points, N+1, cp_dist, smoothing)
             estimated_state[i*3, :] = smoothed_curve[:, 0]        # x
             estimated_state[i*3 + 1, :] = smoothed_curve[:, 1]    # y
+            # import matplotlib.pyplot as plt
+            # plt.plot(estimated_state[i*3, :], estimated_state[i*3 + 1, :], 'b-', label="Bezier Curve")
+            # plt.plot(points[:, 0], points[:, 1], 'ro', label="Control Points")
+            # plt.legend()
+            # plt.show()
 
             # translate the initial guess so that the first point is at (0,0)
             estimated_state[i*3, :] -= estimated_state[i*3, 1]
             estimated_state[i*3 + 1, :] -= estimated_state[i*3+1, 1]
 
-            smoothed_curve_first_point = get_grid_cell_location(points[0][0], points[0][1], grid_origin, self.cell_size)
+            estimated_state[i*3, :] *= self.cell_size
 
-            # translate the initial guess so that the first point is at the current robot position
-            current_robot_position = state[robots_in_conflict[i]][:2]
-            estimated_state[i*3, :] += smoothed_curve_first_point[0]
-            estimated_state[i*3 + 1, :] += smoothed_curve_first_point[1] 
-            # estimated_state[i*3, :] += current_robot_position[0]
-            # estimated_state[i*3 + 1, :] += current_robot_position[1] 
+            cell_loc = get_grid_cell_location(points[0][0], points[0][1], grid_origin, self.cell_size)
+
+            # translate the initial guess so that the first point is at the first point of the continuous space soln
+            estimated_state[i*3, :] += cell_loc[0]
+            estimated_state[i*3 + 1, :] += cell_loc[1] 
             
             headings = calculate_headings(smoothed_curve)
             headings.append(headings[-1])
 
             estimated_state[i*3 + 2, :] = headings
 
-            # now that we have the continuous space soln, we need to connect the robot to it
-            x_start, y_start = state[robots_in_conflict[i]][:2]
-            x_end, y_end = estimated_state[i*3, 0], estimated_state[i*3 + 1, 0]
+            
 
-            # first, rotate the robot in place to face the first point of the continuous space soln
-            current_heading = state[robots_in_conflict[i]][2]
-            first_heading = np.arctan2(y_end - y_start, x_end - x_start)
+            # vector a should be the vector from the cell center to the robot position
+            current_robot_position = state[robots_in_conflict[i]][:2]
+            current_robot_position = np.array(current_robot_position)
+            cell_loc = np.array(cell_loc)
+            a_vector = current_robot_position - cell_loc
 
-            delta_heading = first_heading - current_heading
+            # vector b should be the vector from the cell center to the 4th point of the estimated state
+            b_vector = np.array([estimated_state[i*3, 3], estimated_state[i*3 + 1, 3]]) - cell_loc
 
-            headings = [current_heading+(delta_heading/4), 
-                        current_heading+(delta_heading/2), 
-                        current_heading+(delta_heading*3/4), 
-                        first_heading]
-            
-            
+            # project a onto b
+            a_dot_b = np.dot(a_vector, b_vector)
+            b_dot_b = np.dot(b_vector, b_vector)
+            proj = a_dot_b / b_dot_b * b_vector
 
-            
+            cp2 = proj
 
-            xs = [x_start, x_start, x_start, x_start]
-            ys = [y_start, y_start, y_start, y_start]
+            cp1 = a_vector + cell_loc
+            cp2 = proj + cell_loc
+            cp3 = b_vector + cell_loc
 
-            # print(f"i = {i}")
-            # for j in range(len(xs)):
-            #     import matplotlib.pyplot as plt
-            #     plot_roomba(xs[j], ys[j], headings[j], 'black', False, self.radius)
-            #     plot_roomba(xs[j], ys[j], first_heading, 'red', False, self.radius)
-            #     plt.plot(estimated_state[i*3, 0], estimated_state[i*3 + 1, 0], 'o-')
-            #     plt.show()
 
-            # interpolate between the current state of the robot and the first point of the continuous space soln
-            num_points = 4
+            # plot the control points, the cell center, and the robot position, b vecotr, and a vector
+            import matplotlib.pyplot as plt
 
-            # Create evenly spaced x-values between the start and end x-values
-            x_interp = np.linspace(x_start, x_end, num_points + 2)[1:-1]
+            plt.plot([cp1[0], cp2[0], cp3[0]], [cp1[1], cp2[1], cp3[1]], 'ro', label="control points")
+            plt.plot(cell_loc[0], cell_loc[1], 'bo', label="cell center")
+            plt.plot(current_robot_position[0], current_robot_position[1], 'go', label="robot position")
+            plt.plot(b_vector[0], b_vector[1], 'yo', label="b vector")
+            plt.plot(a_vector[0], a_vector[1], 'yo', label="a vector")
+            plt.plot(proj[0], proj[1], 'yx', label="proj")
+            plt.legend()
+            plt.show()
 
-            # Calculate corresponding y-values using linear interpolation
-            y_interp = np.interp(x_interp, [x_start, x_end], [y_start, y_end])
 
-            for x,y in zip(x_interp[1:], y_interp[1:]):
-                xs.append(x)
-                ys.append(y)
+            # create a smooth bezier curve using these control points
+            smoothed_curve, _ = smooth_path([cp1, cp2, cp3], 10+1, 0, 0)
+            
+            xs = smoothed_curve[:, 0]
+            ys = smoothed_curve[:, 1]
+            headings = calculate_headings(smoothed_curve)
+            headings.append(headings[-1])
 
-            xs.append(x_end)
-            ys.append(y_end)
+            # now add a few points to the beginning of the curve to orient the heading
+            current_robot_heading = state[robots_in_conflict[i]][2]
+            goal_heading = np.arctan2(ys[1] - ys[0], xs[1] - xs[0])
 
-            for j in range(num_points-1):
-                dx = x_interp[j+1] - x_interp[j]
-                dy = y_interp[j+1] - y_interp[j]
-                headings.append(np.arctan2(dy, dx))
+            diff = (goal_heading - current_robot_heading) % (2*np.pi)
 
-            headings.append(headings[-1])
+            # if diff < 0: diff += 2*np.pi
 
-            # now get the heading between the first two points in the estimated soln
-            dx = estimated_state[i*3, 1] - estimated_state[i*3, 0]
-            dy = estimated_state[i*3 + 1, 1] - estimated_state[i*3 + 1, 0]
-            next_heading = np.arctan2(dy, dx)
+            delta_heading = diff / 4
 
-            last_heading = headings[-1]
+            for j in range(4):
+                xs = np.insert(xs, 0, xs[0])
+                ys = np.insert(ys, 0, ys[0])
+                headings.insert(0, current_robot_heading + delta_heading*(j+1))
+                current_robot_heading += delta_heading
 
-            delta_heading = next_heading - last_heading
 
-            headings.append(last_heading + (delta_heading/4))
-            headings.append(last_heading + (delta_heading/2))
-            headings.append(last_heading + (delta_heading*3/4))
-            headings.append(next_heading)
+            # plot the bezier curve, along with the control points
+            # plt.plot(xs, ys, 'b-', label="Bezier Curve")
+            # plt.plot([cp1[0], cp2[0], cp3[0]], [cp1[1], cp2[1], cp3[1]], 'ro', label="control points")
+            # plt.legend()
+            # plt.show()
 
-            xs.append(xs[-1])
-            xs.append(xs[-1])
-            xs.append(xs[-1])
-            xs.append(xs[-1])
-            ys.append(ys[-1])
-            ys.append(ys[-1])
-            ys.append(ys[-1])
-            ys.append(ys[-1])
 
-            
-            # print(f"i = {i}")
-            # for j in range(len(xs)):
-            #     import matplotlib.pyplot as plt
-            #     plot_roomba(xs[j], ys[j], headings[j], 'black', False, self.radius)
-            #     plt.plot(estimated_state[i*3, :], estimated_state[i*3 + 1, :], 'o-')
-            #     plt.show()
-            
-            
+            # # plot the estimated state
+            # plt.plot(estimated_state[i*3, :], estimated_state[i*3 + 1, :], 'b-', label="Bezier Curve")
+
             
             # Finally, combine the interpolated points with the estimated state
             final_estimated_state[i*3, :] = np.concatenate((xs, estimated_state[i*3, :]))
@@ -308,16 +374,17 @@ class MultiPathTrackerDB(MultiPathTracker):
                            self.env.circle_obs, 
                            self.env.rect_obs)
             
-            curr_state = np.zeros((1, 3))
+            # curr_state = np.zeros((1, 3))
+            curr_state = np.array([state[i]])
             x_mpc, u_mpc = mpc.step(
                 curr_state,
                 [ref],
                 [self.control[i]]
             )
 
-            x_mpc_global = self.ego_to_global_roomba(state[i], x_mpc)
+            # x_mpc_global = self.ego_to_global_roomba(state[i], x_mpc)
             u_next[i] = [u_mpc[0, 0], u_mpc[1, 0]]
-            x_next[i] = x_mpc_global
+            x_next[i] = x_mpc
 
         self.trajs = targets 
         return x_next, u_next
@@ -348,7 +415,8 @@ class MultiPathTrackerDB(MultiPathTracker):
                            self.env.circle_obs, 
                            self.env.rect_obs)
             
-            curr_state = np.zeros((1, 3))
+            # curr_state = np.zeros((1, 3))
+            curr_state = np.array([state])
             x_mpc, u_mpc = mpc.step(
                 curr_state,
                 [ref],
@@ -386,7 +454,7 @@ class MultiPathTrackerDB(MultiPathTracker):
                                             self.env.rect_obs, 
                                             vis=False, 
                                             iter=self.settings["sampling_based_planner"]["num_samples"],
-                                            obstacle_tol=.5)
+                                            obstacle_tol=0)
 
         xs = []
         ys = []
@@ -409,7 +477,6 @@ class MultiPathTrackerDB(MultiPathTracker):
 
         self.visited_points_on_guide_paths[i] = []
         
-
     def advance(self, state, time, libs, show_plots=False):
 
         waiting = False
@@ -418,90 +485,140 @@ class MultiPathTrackerDB(MultiPathTracker):
         u_next = [[] for _ in range(self.num_robots)]
         x_next = [[] for _ in range(self.num_robots)]
 
-        # 1. Get the reference trajectory for each robot
+        # 1. Get the mpc trajectory for each robot
         x_next, u_next = self.get_next_state_control(state, u_next, x_next)
         
         # 2. Check if the targets of any two robots overlap
         all_conflicts = []
         for i in range(self.num_robots):
-            traj1 = self.ego_to_global_roomba(state[i], self.trajs[i])
+            # traj1 = self.ego_to_global_roomba(state[i], self.trajs[i])
+            traj1 = self.trajs[i]
             this_robot_conflicts = [i]
             for j in range(i + 1, self.num_robots):
-                traj2 = self.ego_to_global_roomba(state[j], self.trajs[j])
+                # traj2 = self.ego_to_global_roomba(state[j], self.trajs[j])
+                traj2 = self.trajs[j]
                 if self.trajectories_overlap(traj1, traj2, self.radius):
                     this_robot_conflicts.append(j)
             if len(this_robot_conflicts) > 1:
                 all_conflicts.append(this_robot_conflicts)
 
         for c in all_conflicts:
+            # plot the conflict:
+            # draw_conflict(self.num_robots, self.radius, self.env, self.trajs, state, c)
+
             # 3. If they do collide, then reroute the reference trajectories of these robots
             # Get the positions of robots involved in the conflict
             robot_positions = [ state[i][:2] for i in c ]
 
-            # get the largest x and y difference between the robots in conflict
-            x_diff = max([abs(robot_positions[i][0] - robot_positions[j][0]) for i in range(len(c)) for j in range(len(c))])
-            y_diff = max([abs(robot_positions[i][1] - robot_positions[j][1]) for i in range(len(c)) for j in range(len(c))])
-            diff = max(x_diff, y_diff)
+            # find the  diff in y values of the robots in conflict
+            y_diff = robot_positions[0][1] - robot_positions[1][1]
+
+            
 
             # Put down a local grid
             self.cell_size = self.radius*2
-            self.grid_size = 5
-
-            if diff < 5*self.cell_size:
-                circle_obs = []
-                for obs in self.circle_obs:
-                    circle_obs.append(('c',obs[0], obs[1], obs[2]))
-
-                subgoals = self.get_subgoals(state, c)
-
-                grid_origin, centers = place_grid(robot_positions, 
-                                                  cell_size=self.cell_size, 
-                                                  grid_size=5, 
-                                                  subgoals=subgoals,
-                                                  obstacles=circle_obs)
-                grid_obstacle_map = get_obstacle_map(grid_origin, self.grid_size, self.cell_size, circle_obs) 
-
-                # Solve a discrete version of the problem 
-                # Find a subproblem and solve it
-                grid_solution = self.get_discrete_solution(state, c, all_conflicts, grid_obstacle_map, grid_origin, libs)
-
-                if grid_solution:
-                    # if we found a grid solution, we can use it to reroute the robots
-                    continuous_soln = self.convert_db_sol_to_continuous(state, 
-                                                                        grid_solution,
-                                                                        grid_origin, 
-                                                                        c,
-                                                                        self.settings["database"]["num_intervals"],
-                                                                        self.settings["database"]["control_point_dist"], 
-                                                                        self.settings["database"]["smoothing_sigma"]
-                                                                        )
-                    
-                    if show_plots:
-                        draw_grid_solution(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, continuous_soln, state, grid_origin, grid_obstacle_map, c, time)
-
-                    # for each robot in conflict, reroute its reference trajectory to match the grid solution
-                    import copy
-                    old_paths = copy.deepcopy(self.paths)
-
-                    for robot_idx, i in enumerate(c):
-                        new_ref = continuous_soln[robot_idx*3:robot_idx*3+3, :]
-                        self.reroute_guide_paths(new_ref, old_paths, robot_idx, i)
-
-                    # Now replan the reference trajectories to reflect the new guide paths
-                    for i in c:
-                        ref, visited_guide_points = get_ref_trajectory(np.array(state[i]), 
-                                                            np.array(self.paths[i]),
-                                                            self.T, 
-                                                            self.DT, 
-                                                            self.visited_points_on_guide_paths[i])
+            self.grid_size = 6
+
+            if y_diff > self.cell_size*4:
+                continue
+
+            labeled_circle_obs = []
+            for obs in self.circle_obs:
+                labeled_circle_obs.append(('c',obs[0], obs[1], obs[2]))
+
+            subgoals = self.get_subgoals(state, c)
+
+            grid_origin, centers = place_grid(robot_positions, 
+                                                cell_size=self.cell_size, 
+                                                grid_size=self.grid_size, 
+                                                subgoals=subgoals,
+                                                obstacles=labeled_circle_obs)
+            grid_obstacle_map = get_obstacle_map(grid_origin, self.grid_size, self.cell_size, self.circle_obs) 
+
+            # draw the grid that was placed
+            draw_grid_solution(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, [], state, grid_origin, grid_obstacle_map, c, None, None, [])
+
+            # Solve a discrete version of the problem 
+            # Find a subproblem and solve it
+            grid_solution = self.get_discrete_solution(state, c, all_conflicts, grid_obstacle_map, grid_origin, libs)
+
+            # if there is solution that has a sequence of nodes that look like: 
+            # [x,y], [x+1,y], [x,y] or [x,y], [x-1,y], [x,y]
+            # change it to just be [x,y], [x,y], [x,y]
+            if grid_solution:
+                for i in range(len(grid_solution)):
+                    for j in range(len(grid_solution[i])-2):
+                        if grid_solution[i][j] == grid_solution[i][j+2]:
+                            grid_solution[i][j+1] = grid_solution[i][j]
+
+
+
+            if grid_solution:
+                # if we found a grid solution, we can use it to reroute the robots
+                continuous_soln = self.convert_db_sol_to_continuous(state, 
+                                                                    grid_solution,
+                                                                    grid_origin, 
+                                                                    c,
+                                                                    self.settings["database"]["num_intervals"],
+                                                                    self.settings["database"]["control_point_dist"], 
+                                                                    self.settings["database"]["smoothing_sigma"]
+                                                                    )
                 
-                        self.visited_points_on_guide_paths[i] = visited_guide_points
-                        self.trajs[i] = ref
-
-                    # use MPC to track the new reference trajectories
-                    # include all the robots that were in conflict in the MPC problem
-                    # TODO: might instead want to include all robots within a certain distance
-                    mpc = MultiMPC(len(c), # num robots
+                draw_grid_solution(self.num_robots, self.radius, self.grid_size, self.cell_size, self.env, self.trajs, continuous_soln, state, grid_origin, grid_obstacle_map, c, None, None, grid_solution)
+
+                # for each robot in conflict, reroute its reference trajectory to match the grid solution
+                import copy
+                old_paths = copy.deepcopy(self.paths)
+
+                for robot_idx, i in enumerate(c):
+                    new_ref = continuous_soln[robot_idx*3:robot_idx*3+3, :]
+                    self.reroute_guide_paths(new_ref, old_paths, robot_idx, i)
+
+                # Now replan the reference trajectories to reflect the new guide paths
+                for i in c:
+                    ref, visited_guide_points = get_ref_trajectory(np.array(state[i]), 
+                                                        np.array(self.paths[i]),
+                                                        self.T, 
+                                                        self.DT, 
+                                                        self.visited_points_on_guide_paths[i])
+            
+                    self.visited_points_on_guide_paths[i] = visited_guide_points
+                    self.trajs[i] = ref
+
+                # use MPC to track the new reference trajectories
+                # include all the robots that were in conflict in the MPC problem
+                # TODO: might instead want to include all robots within a certain distance
+                mpc = MultiMPC(len(c), # num robots
+                    self.dynamics, 
+                    self.T, 
+                    self.DT, 
+                    self.settings, 
+                    self.env.circle_obs, 
+                    self.env.rect_obs)
+
+                # curr_states = np.zeros((len(c), 3))
+                curr_states = np.array([state[i] for i in c])
+                these_trajs = [self.trajs[i] for i in c]
+                these_controls = [self.control[i] for i in c]
+                x_mpc, u_mpc = mpc.step(
+                    curr_states,
+                    these_trajs,
+                    these_controls
+                )
+
+                for i, r in enumerate(c):
+                    u_next[r] = [u_mpc[i*2, 0], u_mpc[i*2+1, 0]]
+                    x_next[r] = [x_mpc[i*3, 1], x_mpc[i*3+1, 1], x_mpc[i*3+2, 1]]
+
+            else: # Discrete solver failed, resolve the conflict another way
+                if waiting: 
+                    # Choose all but one robot to "wait" (i.e. controls are 0)
+                    c_wait = c[1:]
+                    all_waiting_robots.extend(c_wait)
+
+                else:
+                    print("Using coupled solver to resolve conflict")
+                    mpc = MultiMPC(self.num_robots, # num robots
                         self.dynamics, 
                         self.T, 
                         self.DT, 
@@ -509,59 +626,30 @@ class MultiPathTrackerDB(MultiPathTracker):
                         self.env.circle_obs, 
                         self.env.rect_obs)
 
-                    curr_states = np.zeros((len(c), 3))
-                    these_trajs = [self.trajs[i] for i in c]
-                    these_controls = [self.control[i] for i in c]
+                    curr_states = np.zeros((self.num_robots, 3))
+                    for i in range(self.num_robots):
+                        curr_states[i] = state[i]
                     x_mpc, u_mpc = mpc.step(
                         curr_states,
-                        these_trajs,
-                        these_controls
+                        self.trajs,
+                        self.control
                     )
 
-                    for i, r in enumerate(c):
-                        u_next[r] = [u_mpc[i*2, 0], u_mpc[i*2+1, 0]]
-                        x_next[r] = [x_mpc[i*3, 1], x_mpc[i*3+1, 1], x_mpc[i*3+2, 1]]
-
-                else: # Discrete solver failed, resolve the conflict another way
-                    if waiting: 
-                        # Choose all but one robot to "wait" (i.e. controls are 0)
-                        c_wait = c[1:]
-                        all_waiting_robots.extend(c_wait)
-
-                    else:
-                        print("Using coupled solver to resolve conflict")
-                        mpc = MultiMPC(self.num_robots, # num robots
-                           self.dynamics, 
-                           self.T, 
-                           self.DT, 
-                           self.settings, 
-                           self.env.circle_obs, 
-                           self.env.rect_obs)
-
-                        curr_states = np.zeros((self.num_robots, 3))
-                        x_mpc, u_mpc = mpc.step(
-                            curr_states,
-                            self.trajs,
-                            self.control
-                        )
-
-                        # for each robot in conflict, reroute its reference trajectory to match the grid solution
-                        import copy
-                        old_paths = copy.deepcopy(self.paths)
-                        for i in c:
-                            u_next[i] = [u_mpc[i*2, 0], u_mpc[i*2+1, 0]]
-                            x_next[i] = [x_mpc[i*3, 1], x_mpc[i*3+1, 1], x_mpc[i*3+2, 1]]
-
-                            # update the guide paths of the robots to reflect the new state
-                            new_state = self.ego_to_global_roomba(state[i], x_mpc)
-                            new_theta = x_mpc[i*3+2, :]
-                            new_ref = np.array([new_state[0], new_state[1], new_theta])
+                    # for each robot in conflict, reroute its reference trajectory to match the grid solution
+                    import copy
+                    old_paths = copy.deepcopy(self.paths)
+                    for i in c:
+                        u_next[i] = [u_mpc[i*2, 0], u_mpc[i*2+1, 0]]
+                        x_next[i] = [x_mpc[i*3, 1], x_mpc[i*3+1, 1], x_mpc[i*3+2, 1]]
 
-                            self.reroute_guide_paths(new_ref, old_paths, i, i)
+                        # update the guide paths of the robots to reflect the new state
+                        # new_state = self.ego_to_global_roomba(state[i], x_mpc)
+                        new_x = x_mpc[i*3, :]
+                        new_y = x_mpc[i*3+1, :]
+                        new_theta = x_mpc[i*3+2, :]
+                        new_ref = np.array([new_x, new_y, new_theta])
 
-            else:
-                print("The robots are too far apart to place a grid")
-                print("Proceeding with the current paths")
+                        self.reroute_guide_paths(new_ref, old_paths, i, i)
 
         if waiting:
             all_waiting_robots = self.find_all_waiting_robots(all_waiting_robots, state, x_next)
@@ -579,7 +667,6 @@ class MultiPathTrackerDB(MultiPathTracker):
         whose next state intersects with the waiting robots' current state. Those robots
         should also be marked as waiting. This will need to be resursive since the robots
         that are marked as waiting will also need to be checked for intersections with others
-    
 
         params:
             - waiting (list): list of robots that have already been marked as waiting
diff --git a/guided_mrmp/controllers/optimizer.py b/guided_mrmp/controllers/optimizer.py
index 965c14d..14408d6 100644
--- a/guided_mrmp/controllers/optimizer.py
+++ b/guided_mrmp/controllers/optimizer.py
@@ -5,8 +5,6 @@ class Optimizer:
     def solve_optimization_problem(self, initial_guesses=None, solver_options=None, lam_g=None):
         opti = self.problem['opti']
 
-        X = self.problem['X']
-        U = self.problem['U']
         
         if initial_guesses:
             for param, value in initial_guesses.items():
@@ -23,9 +21,10 @@ class Optimizer:
 
         def print_intermediates_callback(i):
             # print the current value of the objective function
-            print("Iteration:", i, "Current cost cost:", opti.debug.value(self.problem['cost']))
-            print("Iteration:", i, "Current robot cost:", opti.debug.value(self.problem['dist_to_other_robots']))
-            # print("Iteration:", i, "Current obstacle cost:", opti.debug.value(self.problem['obs_cost']))
+            # print("Iteration:", i, "Current cost cost:", opti.debug.value(self.problem['cost']))
+            # print("Iteration:", i, "Current robot cost:", opti.debug.value(self.problem['dist_to_other_robots']))
+            # print("Iteration:", i, "Current obstacle cost:", opti.debug.value(self.problem['obs_cost'])),
+            print("Iteration:", i, "Current distances:", opti.debug.value(self.problem['distance']))
             # print("Iteration:", i, "Current control cost:", opti.debug.value(self.problem['control_cost']))
             # print("Iteration:", i, "Current time cost:", opti.debug.value(self.problem['time_cost']))
             # print("Iteration:", i, "Current goal cost:", opti.debug.value(self.problem['goal_cost']))
diff --git a/guided_mrmp/controllers/path_tracker.py b/guided_mrmp/controllers/path_tracker.py
index 5f7b987..0ccbc36 100644
--- a/guided_mrmp/controllers/path_tracker.py
+++ b/guided_mrmp/controllers/path_tracker.py
@@ -1,10 +1,162 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
-from guided_mrmp.controllers.utils import compute_path_from_wp, get_ref_trajectory
-from guided_mrmp.controllers.mpc import MPC
+# from guided_mrmp.controllers.utils import compute_path_from_wp, get_ref_trajectory
+from guided_mrmp.controllers.multi_mpc import MultiMPC
 from guided_mrmp.utils import Roomba
 
+from scipy.interpolate import interp1d
+
+
+def compute_path_from_wp(start_xp, start_yp, step=0.1):
+    """
+    params:
+        start_xp (array-like): 1D array of x-positions
+        start_yp (array-like): 1D array of y-positions
+        step (float): interpolation step size 
+
+    output:
+        ndarray of shape (3,N) representing the  path as x,y,heading
+    """
+    final_xp = []
+    final_yp = []
+    delta = step  # [m]
+    for idx in range(len(start_xp) - 1):
+        # find the distance between consecutive waypoints
+        section_len = np.sum(
+            np.sqrt(
+                np.power(np.diff(start_xp[idx : idx + 2]), 2)
+                + np.power(np.diff(start_yp[idx : idx + 2]), 2)
+            )
+        )
+
+        # how many interpolated points are needed to reach the next waypoint
+        interp_range = np.linspace(0, 1, np.floor(section_len / delta).astype(int))  
+
+        # interpolate between waypoints
+        fx = interp1d(np.linspace(0, 1, 2), start_xp[idx : idx + 2], kind=1)
+        fy = interp1d(np.linspace(0, 1, 2), start_yp[idx : idx + 2], kind=1)
+
+        # append the interpolated points to the final path 
+        final_xp = np.append(final_xp, fx(interp_range)[1:])
+        final_yp = np.append(final_yp, fy(interp_range)[1:])
+    dx = np.append(0, np.diff(final_xp))
+    dy = np.append(0, np.diff(final_yp))
+    theta = np.arctan2(dy, dx)
+    print(f"theta = {theta}")
+    return np.vstack((final_xp, final_yp, theta))
+
+
+def fix_angle_reference(angle_ref, angle_init):
+    """
+    Removes jumps greater than 2PI to smooth the heading.
+
+    Args:
+        angle_ref (array-like): Reference angles
+        angle_init (float): Initial angle
+
+    Returns:
+        array-like: Smoothed reference angles
+    """
+    diff_angle = angle_ref - angle_init
+    diff_angle = np.unwrap(diff_angle)
+    return angle_init + diff_angle
+
+
+def get_ref_trajectory(state, path, T, DT, path_visited_points=[]):
+    """
+    Generates a reference trajectory for the Roomba.
+
+    Args:
+        state (array-like): Current state [x, y, theta]
+        path (ndarray): Path points [x, y, theta] in the global frame
+        path_visited_points (array-like): Visited path points [[x, y], [x, y], ...]
+        target_v (float): Desired speed
+        T (float): Control horizon duration
+        DT (float): Control horizon time-step
+
+    Returns:
+        ndarray: Reference trajectory [x_k, y_k, theta_k] in the ego frame
+    """
+    K = T/DT
+    xref = np.zeros((3, K+1))  # Reference trajectory for [x, y, theta]
+
+    # Find the last visited point
+    last_visited_idx = 0 if path_visited_points == [] else path_visited_points[-1]
+
+    # Find the spatially closest point after the last visited point
+    next_ind = last_visited_idx + 1
+    path_visited_points.append(next_ind)
+
+    # print(f"guide path = {path}")
+    
+    if K +1  + next_ind < len(path[0]):
+        # return the next k points on the path
+        xref[0,:] = path[0, next_ind:next_ind+K+1]
+        xref[1,:] = path[1, next_ind:next_ind+K+1]
+        xref[2,:] = path[2, next_ind:next_ind+K+1]
+
+    else:
+        xref[0, 0:len(path[0]) - next_ind] = path[0, next_ind:]
+        xref[1, 0:len(path[0]) - next_ind] = path[1, next_ind:]
+        xref[2, 0:len(path[0]) - next_ind] = path[2, next_ind:]
+
+        # Fill the rest of the trajectory with the last point
+        xref[0, len(path[0]) - next_ind:] = path[0, -1]
+        xref[1, len(path[0]) - next_ind:] = path[1, -1]
+        xref[2, len(path[0]) - next_ind:] = path[2, -1]
+
+    # Transform to ego frame
+    # dx = xref[0, :] - state[0]
+    # dy = xref[1, :] - state[1]
+    # xref[0, :] = dx * np.cos(-state[2]) - dy * np.sin(-state[2])  # X
+    # xref[1, :] = dy * np.cos(-state[2]) + dx * np.sin(-state[2])  # Y
+
+    # if next_ind < len(path[0]):
+    #     xref[2, :] = path[2, next_ind] - state[2]  # Theta
+    # else:
+    #     xref[2, :] = path[2, -1] - state[2]
+    
+
+    # Normalize the angles
+    # xref[2, :] = (xref[2, :] + np.pi) % (2.0 * np.pi) - np.pi
+    # xref[2, :] = fix_angle_reference(xref[2, :], xref[2, 0])
+
+
+    # print(f"x_ref = {xref}")
+    return xref, path_visited_points
+
+
+def desired_command_and_trajectory(x0_:np.array, N_, guide_path, points_visited):
+
+    # Find the last visited point
+    last_visited_idx = 0 if points_visited == [] else points_visited[-1]
+
+
+    next_ind = last_visited_idx + 1
+    points_visited.append(next_ind)
+
+    
+
+
+
+    x_ = np.zeros((3, N_+1))
+    x_[:,0] = x0_
+
+    for i in range(N_):
+        if next_ind+i >= len(guide_path[0]):
+            x_ref_ = guide_path[0, -1]
+            y_ref_ = guide_path[1, -1]
+            theta_ref_ = guide_path[2, -1]
+        else:
+            x_ref_ = guide_path[0, next_ind+i]
+            y_ref_ = guide_path[1, next_ind+i]
+            theta_ref_ = guide_path[2, next_ind+i]
+        x_[:,i+1] = np.array([x_ref_, y_ref_, theta_ref_])
+
+    return x_, points_visited
+
+
 # Classes
 class PathTracker:
     def __init__(self, initial_position, dynamics, target_v, T, DT, waypoints, settings):
@@ -31,23 +183,17 @@ class PathTracker:
 
         self.K = int(T / DT)
 
-        # For a car model 
-        # Q = [20, 20, 10, 20]  # state error cost
-        # Qf = [30, 30, 30, 30]  # state final error cost
-        # R = [10, 10]  # input cost
-        # P = [10, 10]  # input rate of change cost
-        # self.mpc = MPC(VehicleModel(), T, DT, Q, Qf, R, P)
-
-
         # For a circular robot (easy dynamics)
-        Q = [20, 20, 20]  # state error cost
-        Qf = [30, 30, 30]  # state final error cost
-        R = [10, 10]  # input cost
-        P = [10, 10]  # input rate of change cost
-        self.mpc = MPC(dynamics, T, DT, Q, Qf, R, P, settings['model_predictive_controller'])
+        self.mpc = MultiMPC(1,
+                            dynamics, 
+                            T, 
+                            DT, 
+                            settings, 
+                            circle_obs=[[5,0,.25]],
+                            rectangle_obs=[])
 
         # Path from waypoint interpolation
-        self.path = compute_path_from_wp(waypoints[0], waypoints[1], 0.05)
+        self.path = compute_path_from_wp(waypoints[0], waypoints[1], .1)
 
         # Helper variables to keep track of the sim
         self.sim_time = 0
@@ -59,6 +205,8 @@ class PathTracker:
         self.d_history = []
         self.optimized_trajectory = None
 
+        self.path_visited_points = []
+
         # Initialise plot
         # plt.style.use("ggplot")
         # self.fig = plt.figure()
@@ -120,30 +268,51 @@ class PathTracker:
 
         return trajectory
 
+    # def shift(T, t0, x0, u, x_n, f):
+        # f_value = f(x0, u[0])
+        # st = x0 + T*f_value
+        # t = t0 + T
+        # u_end = np.concatenate((u[1:], u[-1:]))
+        # x_n = np.concatenate((x_n[1:], x_n[-1:]))
+        # return t, st, u_end, x_n
+
     def get_next_control(self, state, show_plots=False):
         # optimization loop
         # start=time.time()
 
+        print(f"state = {state}")
+
         # Get Reference_traj -> inputs are in worldframe
-        target = get_ref_trajectory(np.array(state), np.array(self.path), self.target_v, self.T, self.DT,0)
+        # target, self.path_visited_points = get_ref_trajectory(np.array(state), np.array(self.path), self.T, self.DT, self.path_visited_points)
+        target, self.path_visited_points = desired_command_and_trajectory(np.array(state), self.K, np.array(self.path), self.path_visited_points)
+
+        print(f"target = {target}")
+        print(f"points_visited = {self.path_visited_points}")
+
 
         # dynamycs w.r.t robot frame
         # curr_state = np.array([0, 0, self.state[2], 0])
-        curr_state = np.array([0, 0, 0])
+        curr_state = np.array([state])
+        # curr_state = np.array([state])
         x_mpc, u_mpc = self.mpc.step(
             curr_state,
-            target,
-            self.control
+            [target],
+            [self.control]
         )
         
         # only the first one is used to advance the simulation
         self.control[:] = [u_mpc[0, 0], u_mpc[1, 0]]
-        # self.state = self.predict_next_state(
-        #     self.state, [self.control[0], self.control[1]], DT
-        # )
 
         return x_mpc, self.control
 
+    def shift(self, DT, t0, x0, u, x_n, f):
+        f_value = f(x0, u)
+        st = x0 + DT*f_value
+        t = t0 + DT
+        u_end = np.concatenate((u[1:], u[-1:]))
+        x_n = np.concatenate((x_n[1:], x_n[-1:]))
+        return t, st, u_end, x_n 
+
     def run(self, show_plots=False):
         """
         Run the path tracker algorithm.
@@ -164,9 +333,17 @@ class PathTracker:
                 print("Success! Goal Reached")
                 return np.asarray([self.x_history, self.y_history, self.h_history])
             x_mpc, controls = self.get_next_control(self.state)
-            next_state = self.dynamics.next_state(self.state, [self.control[0], self.control[1]], self.DT)
+            print(f"self.state = {self.state}")
+            print(f"x_mpc = {x_mpc}")
+            # next_state = self.dynamics.next_state(self.state, [self.control[0], self.control[1]], self.DT)
+
+
+            f_np = lambda x_, u_: np.array([u_[0]*np.cos(x_[2]), u_[0]*np.sin(x_[2]), u_[1]])
+
+            t0, current_state, u0, next_states = self.shift(self.DT, 0, self.state, controls, x_mpc, f_np)
 
-            self.state = next_state
+
+            self.state = current_state
 
             self.x_history.append(self.state[0])
             self.y_history.append(self.state[1])
@@ -174,7 +351,7 @@ class PathTracker:
             
             # use the optimizer output to preview the predicted state trajectory
             # self.optimized_trajectory = self.ego_to_global(x_mpc.value)
-            if show_plots: self.optimized_trajectory = self.ego_to_global_roomba(x_mpc)
+            if show_plots: self.optimized_trajectory =  x_mpc
             if show_plots: self.plot_sim()
             
     def plot_sim(self):
@@ -192,6 +369,11 @@ class PathTracker:
         plt.subplot(grid[0:2, 0:2])
         plt.title("MPC Simulation \n" + "Simulation elapsed time {}s".format(self.sim_time))
 
+        # plot a circle at 5,5 with radius .5
+        circle = plt.Circle((5, 0), .5, color='r', fill=False)
+        plt.gca().add_artist(circle)
+
+
         plt.plot(
             self.path[0, :],
             self.path[1, :],
@@ -219,20 +401,7 @@ class PathTracker:
                 label="mpc opt trajectory",
             )
 
-        # plt.plot(self.x_history[-1], self.y_history[-1], c='tab:blue',
-        #                                                  marker=".",
-        #                                                  markersize=12,
-        #                                                  label="vehicle position")
-        # plt.arrow(self.x_history[-1],
-        #           self.y_history[-1],
-        #           np.cos(self.h_history[-1]),
-        #           np.sin(self.h_history[-1]),
-        #           color='tab:blue',
-        #           width=0.2,
-        #           head_length=0.5,
-        #           label="heading")
-
-        # plot_car(self.x_history[-1], self.y_history[-1], self.h_history[-1])
+
         plot_roomba(self.x_history[-1], self.y_history[-1], self.h_history[-1])
 
 
@@ -265,6 +434,8 @@ class PathTracker:
         ax.set_xlim([0, 10])
         ax.set_ylim([0, 10])
 
+        # plt.show()
+
         plt.draw()
         plt.pause(0.1)
 
@@ -294,17 +465,19 @@ if __name__ == "__main__":
 
     # Example usage
 
-    file_path = "settings_files/settings.yaml"
+    file_path = "experiments/settings_files/4robotscircle.yaml"
     import yaml
     with open(file_path, 'r') as file:
         settings = yaml.safe_load(file)
 
-    initial_pos = np.array([0.0, 0.5, 0.0, 0.0])
+    initial_pos = np.array([0.0, 0.5, 0.0])
     dynamics = Roomba(settings)
     target_vocity = 3.0 # m/s
-    T = 1  # Prediction Horizon [s]
+    T = 5  # Prediction Horizon [s]
     DT = 0.2  # discretization step [s]
-    wp = [[0, 3, 4, 6, 10, 12, 13, 13, 6, 1, 0],
-          [0, 0, 2, 4, 3, 3, -1, -2, -6, -2, -2]]
+    wp = [[0, 1,2,3,4,5,6,7,8,9],
+          [0, 0,0,0,0,0,0,0,0,0]]
     sim = PathTracker(initial_position=initial_pos, dynamics=dynamics, target_v=target_vocity, T=T, DT=DT, waypoints=wp, settings=settings)
-    x,y,h = sim.run(show_plots=True)
\ No newline at end of file
+    x,y,h = sim.run(show_plots=True)
+
+    
\ No newline at end of file
diff --git a/guided_mrmp/controllers/utils.py b/guided_mrmp/controllers/utils.py
index a73a289..1ead009 100644
--- a/guided_mrmp/controllers/utils.py
+++ b/guided_mrmp/controllers/utils.py
@@ -87,7 +87,7 @@ def fix_angle_reference(angle_ref, angle_init):
     diff_angle = np.unwrap(diff_angle)
     return angle_init + diff_angle
 
-def get_ref_trajectory(state, path, T, DT, path_visited_points=[]):
+def get_ref_trajectory(state, path, T, DT, points_visited=[]):
     """
     Generates a reference trajectory for the Roomba.
 
@@ -102,46 +102,32 @@ def get_ref_trajectory(state, path, T, DT, path_visited_points=[]):
     Returns:
         ndarray: Reference trajectory [x_k, y_k, theta_k] in the ego frame
     """
-    K = int(T / DT)
-    xref = np.zeros((3, K))  # Reference trajectory for [x, y, theta]
+    K = int(T/DT)
 
     # Find the last visited point
-    last_visited_idx = 0 if path_visited_points == [] else path_visited_points[-1]
+    last_visited_idx = 0 if points_visited == [] else points_visited[-1]
+
 
-    # Find the spatially closest point after the last visited point
     next_ind = last_visited_idx + 1
-    path_visited_points.append(next_ind)
+    points_visited.append(next_ind)
+
+
+    x_ = np.zeros((3, K+1))
+    x_[:,0] = state
+
+    for i in range(K):
+        if next_ind+i >= len(path[0]):
+            x_ref_ = path[0, -1]
+            y_ref_ = path[1, -1]
+            theta_ref_ = path[2, -1]
+        else:
+            x_ref_ = path[0, next_ind+i]
+            y_ref_ = path[1, next_ind+i]
+            theta_ref_ = path[2, next_ind+i]
+        x_[:,i+1] = np.array([x_ref_, y_ref_, theta_ref_])
+
+    return x_, points_visited
 
-    # K = min(K, len(path[0]) - next_ind)
-    
-    if K + next_ind < len(path[0]):
-        # return the next k points on the path
-        xref[0,:] = path[0, next_ind:next_ind+K]
-        xref[1,:] = path[1, next_ind:next_ind+K]
-        xref[2,:] = path[2, next_ind:next_ind+K]
-
-    else:
-        xref[0, 0:len(path[0]) - next_ind] = path[0, next_ind:]
-        xref[1, 0:len(path[0]) - next_ind] = path[1, next_ind:]
-        xref[2, 0:len(path[0]) - next_ind] = path[2, next_ind:]
-
-        # Fill the rest of the trajectory with the last point
-        xref[0, len(path[0]) - next_ind:] = path[0, -1]
-        xref[1, len(path[0]) - next_ind:] = path[1, -1]
-        xref[2, len(path[0]) - next_ind:] = path[2, -1]
-
-    # Transform to ego frame
-    dx = xref[0, :] - state[0]
-    dy = xref[1, :] - state[1]
-    xref[0, :] = dx * np.cos(-state[2]) - dy * np.sin(-state[2])  # X
-    xref[1, :] = dy * np.cos(-state[2]) + dx * np.sin(-state[2])  # Y
-    xref[2, :] = path[2, next_ind] - state[2]  # Theta
-
-    # Normalize the angles
-    xref[2, :] = (xref[2, :] + np.pi) % (2.0 * np.pi) - np.pi
-    xref[2, :] = fix_angle_reference(xref[2, :], xref[2, 0])
-
-    return xref, path_visited_points
 
 def ego_to_global_roomba(state, mpc_out):
     """
@@ -184,19 +170,47 @@ def get_obstacle_map(grid_origin, grid_size, cell_size, circle_obs):
     # create a grid of size grid_size x grid_size
     grid = np.zeros((grid_size, grid_size))
 
-    for i in range(grid_size):
-        for j in range(grid_size):
-            for obs in circle_obs:
-                # collision check this square and the obstacle circle using shapely
-                circle = Point(obs[0], obs[1]).buffer(obs[2])
+    # for i in range(grid_size):
+    #     for j in range(grid_size):
+    #         for obs in circle_obs:
+    #             # collision check this square and the obstacle circle using shapely
+    #             circle = Point(obs[0], obs[1]).buffer(obs[2])
+    #             cell_corners = np.array([(0, 0), (cell_size, 0), (cell_size, cell_size), (0, cell_size)])
+    #             cell_corners += np.array(grid_origin)
+    #             cell_corners += np.array([i, j]) * cell_size
+    #             cell_geom = Polygon(cell_corners)
+    #             overlap_area = circle.intersection(cell_geom).area
+    #             print(f"overlap = {overlap_area}, threshold = {cell_size**2 / 3}")
+    #             if overlap_area >= (cell_size**2 / 3 ):
+    #                 grid[i][j] = 1
+
+
+    for obs in circle_obs: 
+        # find the cell with greatest overlap with the circle
+        circle = Point(obs[0], obs[1]).buffer(obs[2])
+        max_i, max_j = None, None
+        max = 0 
+        for i in range(grid_size):
+            for j in range(grid_size):
                 cell_corners = np.array([(0, 0), (cell_size, 0), (cell_size, cell_size), (0, cell_size)])
                 cell_corners += np.array(grid_origin)
                 cell_corners += np.array([i, j]) * cell_size
                 cell_geom = Polygon(cell_corners)
                 overlap_area = circle.intersection(cell_geom).area
-                if overlap_area >= ((cell_size**2 / 4) - .1):
+                if overlap_area >= (cell_size**2 / 3 ):
                     grid[i][j] = 1
-
+                if overlap_area >= max:
+                    max = overlap_area
+                    max_i = i
+                    max_j = j
+        max_ij_cell_corners = np.array([(0, 0), (cell_size, 0), (cell_size, cell_size), (0, cell_size)])
+        max_ij_cell_corners += np.array(grid_origin)
+        max_ij_cell_corners += np.array([max_i, max_j]) * cell_size
+        max_ij_cell_geom = Polygon(max_ij_cell_corners)
+        overlap_area = circle.intersection(max_ij_cell_geom).area
+        if max_i is not None and overlap_area> .1:
+            grid[max_i][max_j] = 1
+                
             # TODO: add rectangle obstacles
 
     return grid
diff --git a/guided_mrmp/controllers/vis_helpers.py b/guided_mrmp/controllers/vis_helpers.py
index 421d90b..324573e 100644
--- a/guided_mrmp/controllers/vis_helpers.py
+++ b/guided_mrmp/controllers/vis_helpers.py
@@ -75,7 +75,7 @@ def draw_grid(num_robots, radius, grid_size, cell_size, env, trajs, state, grid_
     for i in range(grid_size):
         for j in range(grid_size):
             if obstacle_map[i][j]:
-                x, y = get_grid_cell_location(i, j, grid_origin)
+                x, y = get_grid_cell_location(i, j, grid_origin, cell_size)
                 # create a square with the top left corner at (x - cell_size/2, y - cell_size/2)
                 square = plt.Rectangle((x - cell_size/2, y - cell_size/2), cell_size, cell_size, color='r', alpha=0.3)
                 plt.gca().add_artist(square)
@@ -102,7 +102,82 @@ def draw_grid(num_robots, radius, grid_size, cell_size, env, trajs, state, grid_
 
     plt.show()
 
-def draw_grid_solution(num_robots, radius, grid_size, cell_size, env,trajs, initial_guess, state, grid_origin, obstacle_map, robots_in_conflict, time):
+def draw_chosen_subp(num_robots, radius, grid_size, cell_size, env,trajs, state, grid_origin, obstacle_map, robots_in_conflict, subp_bottom_left, subp_top_right):
+    """
+    params:
+        - state (list): list of robot states
+        - grid_origin (tuple): top left corner of the grid
+        - obstacle_map (bool array): the obstacle map of the grid
+        - num_robots (int): number of robots that are in conflict
+    """
+    circle_obs = env.circle_obs
+    rect_obs = env.rect_obs
+    x_range = env.boundary[0]
+    y_range = env.boundary[1]
+    
+    # draw the whole environment with the local grid drawn on top
+    colors = cm.rainbow(np.linspace(0, 1, num_robots))
+
+    for i in range(num_robots):
+        plot_roomba(state[i][0], state[i][1], state[i][2], colors[i], False, radius)
+
+    # draw the horizontal and vertical lines of the grid
+    for i in range(grid_size + 1):
+        # Draw vertical lines
+        plt.plot([grid_origin[0] + i * cell_size, grid_origin[0] + i * cell_size], 
+                    [grid_origin[1], grid_origin[1] + grid_size * cell_size], 'k-')
+        # Draw horizontal lines
+        plt.plot([grid_origin[0], grid_origin[0] + grid_size * cell_size], 
+                    [grid_origin[1] +  i * cell_size, grid_origin[1] + i * cell_size], 'k-')
+
+    # draw a blue rectangle with the bottom left corner at subp_bottom_left and top right corner at subp_top_right
+    rect = plt.Rectangle(subp_bottom_left, subp_top_right[0] - subp_bottom_left[0], subp_top_right[1] - subp_bottom_left[1], color='b', fill=False, linewidth=4)
+    plt.gca().add_artist(rect)
+
+
+    # draw the obstacles
+    for obs in circle_obs:
+        circle = plt.Circle((obs[0], obs[1]), obs[2], color='red', fill=False)
+        plt.gca().add_artist(circle)
+
+    for obs in rect_obs:
+        rect = plt.Rectangle((obs[0], obs[1]), obs[2], obs[3], color='k', fill=True)
+        plt.gca().add_artist(rect)
+
+    # if a cell is true in the obstacle map, that cell is an obstacle.
+    # fill that cell with a translucent red color
+    for i in range(grid_size):
+        for j in range(grid_size):
+            if obstacle_map[i][j]:
+                x, y = get_grid_cell_location(i, j, grid_origin, cell_size)
+                # create a square with the top left corner at (x - cell_size/2, y - cell_size/2)
+                square = plt.Rectangle((x - cell_size/2, y - cell_size/2), cell_size, cell_size, color='r', alpha=0.3)
+                plt.gca().add_artist(square)
+
+
+    # plot the robots' continuous space subgoals
+    # for idx in robots_in_conflict:
+    #     traj = ego_to_global_roomba(state[idx], trajs[idx])
+    #     x = traj[0][-1]
+    #     y = traj[1][-1]
+    #     plt.plot(x, y, '^', color=colors[idx])
+    #     circle1 = plt.Circle((x, y), radius, color=colors[idx], fill=False)
+    #     plt.gca().add_artist(circle1)
+
+    # set the size of the plot to be 10x10
+    plt.xlim(x_range[0], x_range[1])
+    plt.xlim(y_range[0], y_range[1])
+
+    # force equal aspect ratio
+    plt.gca().set_aspect('equal', adjustable='box')
+
+    # title
+    plt.title("Discrete Solution")
+
+    # plt.savefig(f"example/sim_00{time}5.png")
+    plt.show()
+
+def draw_temp_goals(num_robots, radius, grid_size, cell_size, env,trajs, state, grid_origin, obstacle_map, robots_in_conflict, subp_bottom_left, subp_top_right, temp_goals, guide_paths):
     """
     params:
         - initial_guess (dict): the initial guess for the optimization problem
@@ -145,22 +220,34 @@ def draw_grid_solution(num_robots, radius, grid_size, cell_size, env,trajs, init
     for i in range(grid_size):
         for j in range(grid_size):
             if obstacle_map[i][j]:
-                x, y = get_grid_cell_location(i, j, grid_origin)
+                x, y = get_grid_cell_location(i, j, grid_origin, cell_size)
                 # create a square with the top left corner at (x - cell_size/2, y - cell_size/2)
                 square = plt.Rectangle((x - cell_size/2, y - cell_size/2), cell_size, cell_size, color='r', alpha=0.3)
                 plt.gca().add_artist(square)
 
 
-    # for i in range(num_robots):
-    if len(initial_guess) > 0:
-        for robot_idx, i  in enumerate(robots_in_conflict):
-            x = initial_guess[robot_idx*3, :]
-            y = initial_guess[robot_idx*3 + 1, :]
-            plt.plot(x, y, 'x', color=colors[i])
+
+    if subp_bottom_left is not None and subp_top_right is not None:
+        # draw a blue rectangle with the bottom left corner at subp_bottom_left and top right corner at subp_top_right
+        rect = plt.Rectangle(subp_bottom_left, subp_top_right[0] - subp_bottom_left[0], subp_top_right[1] - subp_bottom_left[1], color='b', fill=False, linewidth=4)
+        plt.gca().add_artist(rect)
+
+    for idx, r in enumerate(robots_in_conflict):
+        path = guide_paths[r]
+        temp_goal = temp_goals[idx]
+        color = colors[r]
+        # plot guide path with a dashed line
+        plt.plot(path[0], path[1], '--', color=color)
+
+        # plot the temp goal with a solid circle
+        plt.plot(temp_goal[0], temp_goal[1], 'o', color=color)
+    
+
 
     # plot the robots' continuous space subgoals
     for idx in robots_in_conflict:
-        traj = ego_to_global_roomba(state[idx], trajs[idx])
+        # traj = ego_to_global_roomba(state[idx], trajs[idx])
+        traj = trajs[idx]
         x = traj[0][-1]
         y = traj[1][-1]
         plt.plot(x, y, '^', color=colors[idx])
@@ -180,6 +267,203 @@ def draw_grid_solution(num_robots, radius, grid_size, cell_size, env,trajs, init
     # plt.savefig(f"example/sim_00{time}5.png")
     plt.show()
 
+def plot_background(num_robots, radius, grid_size, cell_size, env,trajs, initial_guess, state, grid_origin, obstacle_map, robots_in_conflict, subp_bottom_left, subp_top_right, grid_points):
+    circle_obs = env.circle_obs
+    rect_obs = env.rect_obs
+    x_range = env.boundary[0]
+    y_range = env.boundary[1]
+    plt.xlim(x_range[0], x_range[1])
+    plt.xlim(y_range[0], y_range[1])
+    
+    # draw the whole environment with the local grid drawn on top
+    colors = cm.rainbow(np.linspace(0, 1, num_robots))
+
+    for i in range(num_robots):
+        plot_roomba(state[i][0], state[i][1], state[i][2], colors[i], False, radius)
+
+    # draw the horizontal and vertical lines of the grid
+    for i in range(grid_size + 1):
+        # Draw vertical lines
+        plt.plot([grid_origin[0] + i * cell_size, grid_origin[0] + i * cell_size], 
+                    [grid_origin[1], grid_origin[1] + grid_size * cell_size], 'k-')
+        # Draw horizontal lines
+        plt.plot([grid_origin[0], grid_origin[0] + grid_size * cell_size], 
+                    [grid_origin[1] +  i * cell_size, grid_origin[1] + i * cell_size], 'k-')
+
+    # draw the obstacles
+    for obs in circle_obs:
+        circle = plt.Circle((obs[0], obs[1]), obs[2], color='red', fill=False)
+        plt.gca().add_artist(circle)
+
+    for obs in rect_obs:
+        rect = plt.Rectangle((obs[0], obs[1]), obs[2], obs[3], color='k', fill=True)
+        plt.gca().add_artist(rect)
+
+    # if a cell is true in the obstacle map, that cell is an obstacle.
+    # fill that cell with a translucent red color
+    for i in range(grid_size):
+        for j in range(grid_size):
+            if obstacle_map[i][j]:
+                x, y = get_grid_cell_location(i, j, grid_origin, cell_size)
+                # create a square with the top left corner at (x - cell_size/2, y - cell_size/2)
+                square = plt.Rectangle((x - cell_size/2, y - cell_size/2), cell_size, cell_size, color='r', alpha=0.3)
+                plt.gca().add_artist(square)
+
+
+
+def draw_grid_solution(num_robots, radius, grid_size, cell_size, env,trajs, initial_guess, state, grid_origin, obstacle_map, robots_in_conflict, subp_bottom_left, subp_top_right, grid_points):
+    """
+    params:
+        - initial_guess (dict): the initial guess for the optimization problem
+        - state (list): list of robot states
+        - grid_origin (tuple): top left corner of the grid
+        - obstacle_map (bool array): the obstacle map of the grid
+        - num_robots (int): number of robots that are in conflict
+    """
+    circle_obs = env.circle_obs
+    rect_obs = env.rect_obs
+    x_range = env.boundary[0]
+    y_range = env.boundary[1]
+    
+    # draw the whole environment with the local grid drawn on top
+    colors = cm.rainbow(np.linspace(0, 1, num_robots))
+
+    for i in range(num_robots):
+        plot_roomba(state[i][0], state[i][1], state[i][2], colors[i], False, radius)
+
+    # draw the horizontal and vertical lines of the grid
+    for i in range(grid_size + 1):
+        # Draw vertical lines
+        plt.plot([grid_origin[0] + i * cell_size, grid_origin[0] + i * cell_size], 
+                    [grid_origin[1], grid_origin[1] + grid_size * cell_size], 'k-')
+        # Draw horizontal lines
+        plt.plot([grid_origin[0], grid_origin[0] + grid_size * cell_size], 
+                    [grid_origin[1] +  i * cell_size, grid_origin[1] + i * cell_size], 'k-')
+
+    # draw the obstacles
+    for obs in circle_obs:
+        circle = plt.Circle((obs[0], obs[1]), obs[2], color='red', fill=False)
+        plt.gca().add_artist(circle)
+
+    for obs in rect_obs:
+        rect = plt.Rectangle((obs[0], obs[1]), obs[2], obs[3], color='k', fill=True)
+        plt.gca().add_artist(rect)
+
+    # if a cell is true in the obstacle map, that cell is an obstacle.
+    # fill that cell with a translucent red color
+    for i in range(grid_size):
+        for j in range(grid_size):
+            if obstacle_map[i][j]:
+                x, y = get_grid_cell_location(i, j, grid_origin, cell_size)
+                # create a square with the top left corner at (x - cell_size/2, y - cell_size/2)
+                square = plt.Rectangle((x - cell_size/2, y - cell_size/2), cell_size, cell_size, color='r', alpha=0.3)
+                plt.gca().add_artist(square)
+
+
+    # for i in range(num_robots):
+    if len(initial_guess) > 0:
+        for robot_idx, i  in enumerate(robots_in_conflict):
+            x = initial_guess[robot_idx*3, :]
+            y = initial_guess[robot_idx*3 + 1, :]
+            plt.plot(x, y, 'x', color=colors[i])
+
+    if len(grid_points) > 0:
+        for i in range(len(grid_points)):
+            for point in grid_points[i]:
+                x, y = grid_origin + np.array(point)*cell_size + cell_size/2
+                plt.plot(x, y, color=colors[i], marker='x', markersize=10)
+
+    if subp_bottom_left is not None and subp_top_right is not None:
+        # draw a blue rectangle with the bottom left corner at subp_bottom_left and top right corner at subp_top_right
+        rect = plt.Rectangle(subp_bottom_left, subp_top_right[0] - subp_bottom_left[0], subp_top_right[1] - subp_bottom_left[1], color='b', fill=False, linewidth=4)
+        plt.gca().add_artist(rect)
+
+    # plot the robots' continuous space subgoals
+    # for idx in robots_in_conflict:
+    #     traj = ego_to_global_roomba(state[idx], trajs[idx])
+    #     x = traj[0][-1]
+    #     y = traj[1][-1]
+    #     plt.plot(x, y, '^', color=colors[idx])
+    #     circle1 = plt.Circle((x, y), radius, color=colors[idx], fill=False)
+    #     plt.gca().add_artist(circle1)
+
+    # set the size of the plot to be 10x10
+    plt.xlim(x_range[0], x_range[1])
+    plt.xlim(y_range[0], y_range[1])
+
+    # force equal aspect ratio
+    plt.gca().set_aspect('equal', adjustable='box')
+
+    # title
+    plt.title("Discrete Solution")
+
+    # plt.savefig(f"example/sim_00{time}5.png")
+    plt.show()
+
+def draw_grid_world_solution(num_robots, radius, grid_size, cell_size, env,trajs, initial_guess, state, grid_origin, obstacle_map, robots_in_conflict, subp_bottom_left, subp_top_right, grid_points):
+    """
+    params:
+        - initial_guess (dict): the initial guess for the optimization problem
+        - state (list): list of robot states
+        - grid_origin (tuple): top left corner of the grid
+        - obstacle_map (bool array): the obstacle map of the grid
+        - num_robots (int): number of robots that are in conflict
+    """
+    # plot_background(num_robots, radius, grid_size, cell_size, env,trajs, initial_guess, state, grid_origin, obstacle_map, robots_in_conflict, subp_bottom_left, subp_top_right, grid_points)
+
+    colors = cm.rainbow(np.linspace(0, 1, num_robots))
+
+    # if len(grid_points) > 0:
+    #     for i in range(len(grid_points)):
+    #         color = colors[robots_in_conflict[i]]
+    #         for point in grid_points[i]:
+    #             x, y = grid_origin + np.array(point)*cell_size + cell_size/2
+    #             plt.plot(x, y, color=color, marker='x', markersize=10)
+
+    # get the length of the longest grid solution
+    max_len = max([len(grid_points[i]) for i in range(len(grid_points))])
+
+    for i in range(max_len):
+        plot_background(num_robots, radius, grid_size, cell_size, env,trajs, initial_guess, state, grid_origin, obstacle_map, robots_in_conflict, subp_bottom_left, subp_top_right, grid_points)
+
+        if subp_bottom_left is not None and subp_top_right is not None:
+            print()
+            # draw a blue rectangle with the bottom left corner at subp_bottom_left and top right corner at subp_top_right
+            rect = plt.Rectangle(subp_bottom_left, subp_top_right[0] - subp_bottom_left[0], subp_top_right[1] - subp_bottom_left[1], color='b', fill=False, linewidth=4)
+            plt.gca().add_artist(rect)
+
+        for j in range(len(grid_points)):
+            if i < len(grid_points[j]):
+                color = colors[robots_in_conflict[j]]
+                point = grid_points[j][i]
+                x, y = grid_origin + np.array(point)*cell_size + cell_size/2
+                plt.plot(x, y, color=color, marker='x', markersize=10)
+            else:
+                color = colors[robots_in_conflict[j]]
+                point = grid_points[j][-1]
+                x, y = grid_origin + np.array(point)*cell_size + cell_size/2
+                plt.plot(x, y, color=color, marker='x', markersize=10)
+
+    # plot the robots' continuous space subgoals
+    # for idx in robots_in_conflict:
+    #     traj = ego_to_global_roomba(state[idx], trajs[idx])
+    #     x = traj[0][-1]
+    #     y = traj[1][-1]
+    #     plt.plot(x, y, '^', color=colors[idx])
+    #     circle1 = plt.Circle((x, y), radius, color=colors[idx], fill=False)
+    #     plt.gca().add_artist(circle1)
+
+    
+
+        # force equal aspect ratio
+        plt.gca().set_aspect('equal', adjustable='box')
+
+        # title
+        plt.title("Discrete Solution")
+
+        # plt.savefig(f"example/sim_00{time}5.png")
+        plt.show()
+
 def plot_trajs(num_robots, paths, env, state, traj1, traj2, radius):
     """
     Plot the trajectories of two robots.
diff --git a/guided_mrmp/simulator.py b/guided_mrmp/simulator.py
index 41aa1c9..651b3b4 100644
--- a/guided_mrmp/simulator.py
+++ b/guided_mrmp/simulator.py
@@ -128,6 +128,13 @@ class Simulator:
         for i in range(self.num_robots):
             plt.plot(self.policy.paths[i][0, :], self.policy.paths[i][1, :], '--', color=colors[i])
 
+        # if self.policy.trajs:
+        #     for i in range(self.num_robots):
+        #         traj = self.policy.ego_to_global_roomba([self.x_history[i][-1], self.y_history[i][-1], self.h_history[i][-1]], self.policy.trajs[i])
+        #         plt.plot(traj[0, :], traj[1, :], '^-',color=colors[i])
+        #         # plot a roomba at the end of the trajectory
+        #         self.plot_roomba(traj[0, -1], traj[1, -1], self.h_history[i][-1], "k", False, self.policy.radius)
+
         x_range = self.env.boundary[0]
         y_range = self.env.boundary[1]
         plt.xlim(x_range[0], x_range[1])
@@ -140,10 +147,10 @@ class Simulator:
 
         plt.tight_layout()
 
-
-        # plt.show()
-        plt.draw()
-        plt.pause(0.3)
+        plt.axis('off')
+        plt.show()
+        # plt.draw()
+        # plt.pause(0.3)
         # plt.savefig(f'figures/sim_00{round(self.time,2)}.png')
         plt.clf()
 
@@ -192,5 +199,5 @@ class Simulator:
             traj = self.optimized_trajectories_hist[i][-1]
             if len(traj) == 0:
                 continue
-            traj = self.policy.ego_to_global_roomba([self.x_history[i][-1], self.y_history[i][-1], self.h_history[i][-1]], traj)
+            # traj = self.policy.ego_to_global_roomba([self.x_history[i][-1], self.y_history[i][-1], self.h_history[i][-1]], traj)
             plt.plot(traj[0, :], traj[1, :], '+-',color=colors[i])
-- 
GitLab