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.utils import Roomba
# Classes
class PathTracker:
def __init__(self, initial_position, dynamics, target_v, T, DT, waypoints, settings):
"""
Initializes the PathTracker object.
Parameters:
- initial_position: The initial position of the robot [x, y, heading].
- dynamics: The dynamics model of the robot.
- target_v: The target velocity of the robot.
- 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.
"""
# State of the robot [x,y, heading]
self.state = initial_position
self.dynamics = dynamics
self.T = T
self.DT = DT
self.target_v = target_v
# helper variable to keep track of mpc output
# starting condition is 0,0
self.control = np.zeros(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)
# 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'])
# Path from waypoint interpolation
self.path = compute_path_from_wp(waypoints[0], waypoints[1], 0.05)
# Helper variables to keep track of the sim
self.sim_time = 0
self.x_history = []
self.y_history = []
self.v_history = []
self.h_history = []
self.a_history = []
self.d_history = []
self.optimized_trajectory = None
# Initialise plot
# plt.style.use("ggplot")
# self.fig = plt.figure()
# plt.ion()
# plt.show()
def ego_to_global(self, mpc_out):
"""
transforms optimized trajectory XY points from ego (robot) reference
into global (map) frame
Args:
mpc_out ():
"""
trajectory = np.zeros((2, self.K))
trajectory[:, :] = mpc_out[0:2, 1:]
Rotm = np.array(
[
[np.cos(self.state[3]), np.sin(self.state[3])],
[-np.sin(self.state[3]), np.cos(self.state[3])],
]
)
trajectory = (trajectory.T.dot(Rotm)).T
trajectory[0, :] += self.state[0]
trajectory[1, :] += self.state[1]
return trajectory
def ego_to_global_roomba(self, 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 = self.state[0]
y = self.state[1]
theta = self.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 get_next_control(self, state, show_plots=False):
# optimization loop
# start=time.time()
# 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)
# dynamycs w.r.t robot frame
# curr_state = np.array([0, 0, self.state[2], 0])
curr_state = np.array([0, 0, 0])
x_mpc, u_mpc = self.mpc.step(
curr_state,
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 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.x_history.append(self.state[0])
self.y_history.append(self.state[1])
self.h_history.append(self.state[2])
if show_plots: self.plot_sim()
while 1:
if (np.sqrt((self.state[0] - self.path[0, -1]) ** 2 + (self.state[1] - self.path[1, -1]) ** 2) < 0.1):
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)
self.state = next_state
self.x_history.append(self.state[0])
self.y_history.append(self.state[1])
self.h_history.append(self.state[2])
# 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_sim(self):
self.sim_time = self.sim_time + self.DT
# self.x_history.append(self.state[0])
# self.y_history.append(self.state[1])
# self.v_history.append(self.control[0])
self.h_history.append(self.state[2])
self.d_history.append(self.control[1])
plt.clf()
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))
plt.plot(
self.path[0, :],
self.path[1, :],
c="tab:orange",
marker=".",
label="reference track",
)
plt.plot(
self.x_history,
self.y_history,
c="tab:blue",
marker=".",
alpha=0.5,
label="vehicle trajectory",
)
if self.optimized_trajectory is not None:
plt.plot(
self.optimized_trajectory[0, :],
self.optimized_trajectory[1, :],
c="tab:green",
marker="+",
alpha=0.5,
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])
plt.ylabel("map y")
plt.yticks(np.arange(min(self.path[1, :]) - 1.0, max(self.path[1, :] + 1.0) + 1, 1.0))
plt.xlabel("map x")
plt.xticks(np.arange(min(self.path[0, :]) - 1.0, max(self.path[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.draw()
plt.pause(0.1)
def plot_roomba(x, y, yaw):
"""
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='b', 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__":
# Example usage
file_path = "settings_files/settings.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])
dynamics = Roomba(settings)
target_vocity = 3.0 # m/s
T = 1 # 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]]
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)