Newer
Older
rachelmoan
committed
import numpy as np
import matplotlib.pyplot as plt
rachelmoan
committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from guided_mrmp.controllers.utils import compute_path_from_wp, get_ref_trajectory
from guided_mrmp.controllers.multi_mpc import MultiMPC
class MultiPathTracker:
def __init__(self, env, initial_positions, dynamics, target_v, T, DT, waypoints, settings, lib_2x3, lib_3x3, lib_2x5):
"""
Initializes the PathTracker object.
Parameters:
- initial_positions: List of the initial positions of the robots [x, y, heading].
- dynamics: The dynamics model of the robots.
- target_v: The target velocity of the robots.
- T: The time horizon for the model predictive control (MPC).
- DT: The time step for the MPC.
- waypoints: A list of waypoints defining the desired path for each robot.
"""
# State of the robot [x,y, heading]
self.env = env
self.states = initial_positions
self.num_robots = len(initial_positions)
self.dynamics = dynamics
self.T = T
self.DT = DT
self.target_v = target_v
self.radius = dynamics.radius
self.update_ref_paths = False
# helper variable to keep track of mpc output
# starting condition is 0,0
self.control = np.zeros((self.num_robots, 2))
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)
# libraries for the discrete solver
self.lib_2x3 = lib_2x3
self.lib_3x3 = lib_3x3
self.lib_2x5 = lib_2x5
# For a circular robot (easy dynamics)
Q = settings['model_predictive_controller']['Q'] # state error cost
Qf = settings['model_predictive_controller']['Qf'] # state final error cost
R = settings['model_predictive_controller']['R'] # input cost
P = settings['model_predictive_controller']['P'] # input rate of change cost
rachelmoan
committed
self.mpc = MultiMPC(self.num_robots, dynamics, T, DT, Q, Qf, R, P, settings['model_predictive_controller'], settings['environment']['circle_obstacles'])
self.circle_obs = settings['environment']['circle_obstacles']
# Path from waypoint interpolation
self.paths = []
for wp in waypoints:
self.paths.append(compute_path_from_wp(wp[0], wp[1], 0.05))
self.visited_points_on_guide_paths = [[]]*self.num_robots
rachelmoan
committed
print(f"paths = {len(self.paths)}")
# Helper variables to keep track of the sim
self.sim_time = 0
self.x_history = [ [] for _ in range(self.num_robots) ]
self.y_history = [ [] for _ in range(self.num_robots) ]
self.v_history = [ [] for _ in range(self.num_robots) ]
self.h_history = [ [] for _ in range(self.num_robots) ]
self.a_history = [ [] for _ in range(self.num_robots) ]
self.d_history = [ [] for _ in range(self.num_robots) ]
self.optimized_trajectories_hist = [ [] for _ in range(self.num_robots) ]
self.optimized_trajectory = None
def trajectories_overlap(self, traj1, traj2, threshold):
"""
Checks if two trajectories overlap. We only care about xy positions.
Args:
traj1 (3xn numpy array): First trajectory. First row is x, second row is y, third row is heading.
traj2 (3xn numpy array): Second trajectory.
threshold (float): Distance threshold to consider a collision.
rachelmoan
committed
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
Returns:
bool: True if trajectories overlap, False otherwise.
"""
for i in range(traj1.shape[1]):
for j in range(traj2.shape[1]):
if np.linalg.norm(traj1[0:2, i] - traj2[0:2, j]) < 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)
rachelmoan
committed
# Translate the points to the robot's position in the global frame
trajectory[0, :] += x
trajectory[1, :] += y
rachelmoan
committed
return trajectory
def get_next_control(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):
targets.append(get_ref_trajectory(np.array(state[i]), np.array(self.paths[i]), self.target_v, self.T, self.DT, len(self.x_history[i])+1))
# dynamycs w.r.t robot frame
# curr_state = np.array([0, 0, self.state[2], 0])
curr_states = np.zeros((self.num_robots, 3))
x_mpc, u_mpc = self.mpc.step(
curr_states,
targets,
self.control
rachelmoan
committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
# only the first one is used to advance the simulation
# self.control[:] = [u_mpc[0, 0], u_mpc[1, 0]]
self.control = []
for i in range(self.num_robots):
self.control.append([u_mpc[i*2, 0], u_mpc[i*2+1, 0]])
return x_mpc, self.control
def done(self):
for i in range(self.num_robots):
# print(f"state = {self.states[i]}")
# print(f"path = {self.paths[i][:, -1]}")
if (np.sqrt((self.states[i][0] - self.paths[i][0, -1]) ** 2 + (self.states[i][1] - self.paths[i][1, -1]) ** 2) > 1):
return False
return True
def plot_current_world_state(self):
"""
Plot the current state of the world.
"""
import matplotlib.pyplot as plt
import matplotlib.cm as cm
# Plot the current state of each robot using the most recent values from
# x_history, y_history, and h_history
colors = cm.rainbow(np.linspace(0, 1, self.num_robots))
for i in range(self.num_robots):
plot_roomba(self.x_history[i][-1], self.y_history[i][-1], self.h_history[i][-1], colors[i], False, self.radius)
# plot the goal of each robot with solid circle
for i in range(self.num_robots):
x, y, theta = self.paths[i][:, -1]
plt.plot(x, y, 'o', color=colors[i])
circle1 = plt.Circle((x, y), self.radius, color=colors[i], fill=False)
plt.gca().add_artist(circle1)
# plot the ref path of each robot
for i in range(self.num_robots):
plt.plot(self.paths[i][0, :], self.paths[i][1, :], '--', color=colors[i])
# set the size of the plot to be 10x10
plt.xlim(0, 10)
plt.ylim(0, 10)
# force equal aspect ratio
plt.gca().set_aspect('equal', adjustable='box')
plt.show()
def run(self, show_plots=False):
"""
Run the path tracker algorithm.
Parameters:
- show_plots (bool): Flag indicating whether to show plots during the simulation. Default is False.
Returns:
- numpy.ndarray: Array containing the history of x, y, and h coordinates.
"""
# Add the initial state to the histories
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])
if show_plots: self.plot_sim()
self.plot_current_world_state()
while 1:
# check if all robots have reached their goal
if self.done():
print("Success! Goal Reached")
return np.asarray([self.x_history, self.y_history, self.h_history])
# plot the current state of the robots
self.plot_current_world_state()
# get the next control for all robots
x_mpc, controls = self.get_next_control(self.states)
next_states = []
for i in range(self.num_robots):
next_states.append(self.dynamics.next_state(self.states[i], controls[i], self.DT))
self.states = next_states
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])
if self.update_ref_paths:
self.update_reference_paths()
self.update_ref_paths = False
# 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.plot_sim()
def plot_roomba(x, y, yaw, color, fill, radius):
"""
Args:
x ():
y ():
yaw ():
"""
fig = plt.gcf()
ax = fig.gca()
if fill: alpha = .3
else: alpha = 1
circle = plt.Circle((x, y), radius, color=color, fill=fill, alpha=alpha)
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__":
# main()