diff --git a/guided_mrmp/conflict_resolvers/traj_opt.py b/guided_mrmp/conflict_resolvers/traj_opt.py
index 9dc77a28a4c91ac685275c2c62cedc9b581e4d7b..e4e05b5ce51b09c35bc9ea65a4a28771327a3fcf 100644
--- a/guided_mrmp/conflict_resolvers/traj_opt.py
+++ b/guided_mrmp/conflict_resolvers/traj_opt.py
@@ -96,17 +96,27 @@ class TrajOptMultiRobot():
         # print("Initial pos:", pos[:, 0])
         # print("Initial heading:", heading[:, 0])
 
+        pi = [3.14159*2]*self.num_robots
+        pi = np.array(pi)
+        pi = DM(pi)
+
         for k in range(N): # loop over control intervals
             dxdt = vel[:,k] * cos(heading[:,k])
             dydt = vel[:,k] * sin(heading[:,k])
             dthetadt = omega[:,k]
             opti.subject_to(x[:,k+1]==x[:,k] + dt*dxdt)
             opti.subject_to(y[:,k+1]==y[:,k] + dt*dydt) 
-            opti.subject_to(heading[:,k+1]==heading[:,k] + dt*dthetadt)
+            opti.subject_to(heading[:,k+1]==fmod(heading[:,k] + dt*dthetadt, pi))
+
+        # Calculate the sum of squared differences between consecutive heading angles
+        heading_diff_penalty = 0
+        for k in range(N-1):
+            heading_diff_penalty += sumsqr(heading[:,k+1] - heading[:,k])
 
         opti.minimize(self.rob_dist_weight*dist_to_other_robots 
-                      + self.obs_dist_weight*dist_to_other_obstacles 
-                      + self.time_weight*T)
+                    + self.obs_dist_weight*dist_to_other_obstacles 
+                    + self.time_weight*T
+                    + 5*heading_diff_penalty)
 
 
         # ---- path constraints -----------
@@ -142,7 +152,14 @@ class TrajOptMultiRobot():
 
         # print(f"pos = {opti.debug.value(pos[2:4,:])}")
 
-        return sol,pos
+        # Extract x and y values
+        x_vals = np.array(sol.value(x))
+        y_vals = np.array(sol.value(y))
+
+        # Extract theta values
+        theta_vals = np.array(sol.value(heading))
+
+        return sol,pos, x_vals, y_vals, theta_vals
 
     def plot_paths(self, x_opt, initial_guess):
         fig, ax = plt.subplots()
@@ -182,6 +199,75 @@ class TrajOptMultiRobot():
         plt.show()
 
 
+def plot_sim(x_histories, y_histories, h_histories):
+
+    if len(x_histories) > 20:
+        colors = plt.cm.hsv(np.linspace(0.2, 1.0, len(x_histories)))
+    elif len(x_histories) > 10:
+        colors = plt.cm.tab20(np.linspace(0, 1, len(x_histories)))
+    else:
+        colors = plt.cm.tab10(np.linspace(0, 1, len(x_histories)))
+
+    
+
+
+    longest_traj = max([len(x) for x in x_histories])
+
+    for i in range(longest_traj):
+        plt.clf()
+        for x_history, y_history, h_history, color in zip(x_histories, y_histories, h_histories, colors):
+            
+            print(color)
+
+            plt.plot(
+                x_history[:i],
+                y_history[:i],
+                c=color,
+                marker=".",
+                alpha=0.5,
+                label="vehicle trajectory",
+            )
+
+            if i < len(x_history):
+                plot_roomba(x_history[i-1], y_history[i-1], h_history[i-1], color)
+            else:
+                plot_roomba(x_history[-1], y_history[-1], h_history[-1], color)
+
+        
+        ax = plt.gca()
+        ax.set_xlim([0, 10])
+        ax.set_ylim([0, 10])
+        plt.tight_layout()
+
+        plt.draw()
+        plt.pause(0.2)
+    input()
+
+
+def plot_roomba(x, y, yaw, color):
+    """
+
+    Args:
+        x ():
+        y ():
+        yaw ():
+    """
+    LENGTH = 0.5  # [m]
+    WIDTH = 0.25  # [m]
+    OFFSET = LENGTH  # [m]
+
+    fig = plt.gcf()
+    ax = fig.gca()
+    circle = plt.Circle((x, y), .5, color=color, fill=False)
+    ax.add_patch(circle)
+
+    # Plot direction marker
+    dx = 1 * np.cos(yaw)
+    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__":
 
     # define obstacles
@@ -225,6 +311,9 @@ if __name__ == "__main__":
         # print(f"goal = {goal}")
         initial_guess[i,:] = np.linspace(start[0], goal[0], N+1)
         initial_guess[i+1,:] = np.linspace(start[1], goal[1], N+1)
+        dx = goal[0] - start[0]
+        dy = goal[1] - start[1]
+        initial_guess[i+2,:] = np.arctan2(dy,dx)*np.ones(N+1)
         # initial_guess[i+2,:] = np.linspace(.5, .5, N+1)
         # initial_guess[i+3,:] = np.linspace(.5, .5, N+1)
 
@@ -287,12 +376,22 @@ if __name__ == "__main__":
                                control_weight=control_costs_weight,
                                time_weight=time_weight
                                )
-    sol,pos = solver.solve(N, initial_guess)
+    sol,pos, xs, ys, thetas = solver.solve(N, initial_guess)
     pos_vals = np.array(sol.value(pos))
 
     
     solver.plot_paths(pos_vals, initial_guess)
 
+    print(pos_vals)
+
+
+
+
+    # for r in range(num_robots):
+    #     xs.append(pos_vals[r*2, :])
+    #     ys.append(pos_vals[r*2+1, :])
+    #     thetas.append(pos_vals[num_robots*2 + r, :])
 
+    plot_sim(xs, ys, thetas)