mpc_python_learn/mpc_pybullet_demo/mpc_demo_nosim.py

282 lines
7.9 KiB
Python
Raw Normal View History

2020-01-13 20:25:26 +08:00
#! /usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
2020-01-13 20:25:26 +08:00
2023-10-24 02:27:46 +08:00
from cvxpy_mpc.utils import compute_path_from_wp, get_ref_trajectory
2023-10-24 03:58:21 +08:00
from cvxpy_mpc import MPC, VehicleModel
2020-01-13 20:25:26 +08:00
import sys
2023-10-24 03:58:21 +08:00
2020-01-14 21:37:28 +08:00
# Robot Starting position
SIM_START_X = 0.0
SIM_START_Y = 0.5
SIM_START_V = 0.0
SIM_START_H = 0.0
2020-01-13 22:33:29 +08:00
2020-07-01 00:21:27 +08:00
# Params
2023-10-24 03:58:21 +08:00
TARGET_VEL = 1.0 # m/s
T = 5 # Prediction Horizon [s]
DT = 0.2 # discretization step [s]
L = 0.3 # vehicle wheelbase [m]
2020-01-13 20:25:26 +08:00
2023-10-21 02:20:11 +08:00
# Classes
class MPCSim:
2020-01-13 20:25:26 +08:00
def __init__(self):
2023-10-24 16:19:22 +08:00
# State of the robot [x,y,v, heading]
self.state = np.array([SIM_START_X, SIM_START_Y, SIM_START_V, SIM_START_H])
2021-07-08 19:54:48 +08:00
2023-10-24 16:19:22 +08:00
# helper variable to keep track of mpc output
2023-10-26 00:04:35 +08:00
# starting condition is 0,0
2023-11-10 17:40:55 +08:00
self.control = np.zeros(2)
2023-10-24 16:19:22 +08:00
2023-10-24 03:58:21 +08:00
self.K = int(T / DT)
2021-07-08 19:54:48 +08:00
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
2023-10-24 03:58:21 +08:00
self.mpc = MPC(VehicleModel(), T, DT, Q, Qf, R, P)
2020-01-13 20:25:26 +08:00
2023-10-24 16:19:22 +08:00
# Path from waypoint interpolation
self.path = compute_path_from_wp(
[0, 3, 4, 6, 10, 12, 13, 13, 6, 1, 0],
[0, 0, 2, 4, 3, 3, -1, -2, -6, -2, -2],
0.05,
)
2020-01-13 20:25:26 +08:00
2023-10-24 16:19:22 +08:00
# 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 = []
2023-11-10 17:40:55 +08:00
self.optimized_trajectory = None
# Initialise plot
2020-01-13 20:25:26 +08:00
plt.style.use("ggplot")
self.fig = plt.figure()
plt.ion()
plt.show()
2023-11-10 17:40:55 +08:00
def ego_to_global(self, mpc_out):
"""
2023-11-10 17:40:55 +08:00
transforms optimized trajectory XY points from ego(car) reference
into global(map) frame
2023-10-25 21:17:25 +08:00
Args:
mpc_out ():
"""
2023-11-10 17:40:55 +08:00
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])],
]
)
2023-11-10 17:40:55 +08:00
trajectory = (trajectory.T.dot(Rotm)).T
trajectory[0, :] += self.state[0]
trajectory[1, :] += self.state[1]
return trajectory
2020-05-05 00:36:23 +08:00
2020-01-13 20:25:26 +08:00
def run(self):
"""
[TODO:summary]
[TODO:description]
"""
self.plot_sim()
input("Press Enter to continue...")
2020-01-13 20:25:26 +08:00
while 1:
if (
np.sqrt(
(self.state[0] - self.path[0, -1]) ** 2
+ (self.state[1] - self.path[1, -1]) ** 2
)
< 0.5
):
2021-07-08 19:54:48 +08:00
print("Success! Goal Reached")
input("Press Enter to continue...")
return
# optimization loop
# start=time.time()
2023-10-24 16:19:22 +08:00
# Get Reference_traj -> inputs are in worldframe
2023-10-24 04:02:41 +08:00
target = get_ref_trajectory(self.state, self.path, TARGET_VEL, T, DT)
2023-10-24 16:19:22 +08:00
# dynamycs w.r.t robot frame
curr_state = np.array([0, 0, self.state[2], 0])
2023-10-24 02:27:46 +08:00
x_mpc, u_mpc = self.mpc.step(
curr_state,
target,
2023-11-10 17:40:55 +08:00
self.control,
verbose=False,
)
2021-07-08 19:54:48 +08:00
# print("CVXPY Optimization Time: {:.4f}s".format(time.time()-start))
2023-10-24 16:19:22 +08:00
# only the first one is used to advance the simulation
2023-10-26 00:04:35 +08:00
2023-11-10 17:40:55 +08:00
self.control[:] = [u_mpc.value[0, 0], u_mpc.value[1, 0]]
2023-10-28 17:47:35 +08:00
self.state = self.predict_next_state(
2023-11-10 17:40:55 +08:00
self.state, [self.control[0], self.control[1]], DT
2023-10-28 17:47:35 +08:00
)
2023-10-24 16:19:22 +08:00
2023-10-28 17:47:35 +08:00
# use the optimizer output to preview the predicted state trajectory
2023-11-10 17:40:55 +08:00
self.optimized_trajectory = self.ego_to_global(x_mpc.value)
2021-07-08 19:54:48 +08:00
self.plot_sim()
2022-07-21 19:23:08 +08:00
2023-10-28 17:47:35 +08:00
def predict_next_state(self, state, u, dt):
def kinematics_model(x, t, u):
dxdt = x[2] * np.cos(x[3])
dydt = x[2] * np.sin(x[3])
dvdt = u[0]
2023-10-24 03:58:21 +08:00
dthetadt = x[2] * np.tan(u[1]) / L
2023-10-24 02:27:46 +08:00
dqdt = [dxdt, dydt, dvdt, dthetadt]
return dqdt
# solve ODE
2023-10-24 03:58:21 +08:00
tspan = [0, dt]
2023-10-26 00:04:35 +08:00
new_state = odeint(kinematics_model, state, tspan, args=(u[:],))[1]
return new_state
2020-01-13 20:25:26 +08:00
def plot_sim(self):
2023-10-24 03:58:21 +08:00
self.sim_time = self.sim_time + DT
2020-01-13 22:33:29 +08:00
self.x_history.append(self.state[0])
self.y_history.append(self.state[1])
2020-07-01 22:59:13 +08:00
self.v_history.append(self.state[2])
self.h_history.append(self.state[3])
2023-11-10 17:40:55 +08:00
self.a_history.append(self.control[0])
self.d_history.append(self.control[1])
2020-01-13 22:33:29 +08:00
2020-01-13 20:25:26 +08:00
plt.clf()
2020-01-13 22:33:29 +08:00
grid = plt.GridSpec(2, 3)
2020-01-13 20:25:26 +08:00
2020-01-14 21:37:28 +08:00
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",
)
2020-03-04 20:32:29 +08:00
2023-11-10 17:40:55 +08:00
if self.optimized_trajectory is not None:
plt.plot(
2023-11-10 17:40:55 +08:00
self.optimized_trajectory[0, :],
self.optimized_trajectory[1, :],
c="tab:green",
marker="+",
alpha=0.5,
label="mpc opt trajectory",
)
2020-05-05 00:36:23 +08:00
2020-03-04 20:32:29 +08:00
# 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])
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)
)
2020-07-01 00:21:27 +08:00
plt.axis("equal")
# plt.legend()
2020-01-13 22:33:29 +08:00
plt.subplot(grid[0, 2])
# plt.title("Linear Velocity {} m/s".format(self.v_history[-1]))
plt.plot(self.a_history, c="tab:orange")
2020-03-04 20:32:29 +08:00
locs, _ = plt.xticks()
2023-10-24 03:58:21 +08:00
plt.xticks(locs[1:], locs[1:] * DT)
plt.ylabel("a(t) [m/ss]")
plt.xlabel("t [s]")
2020-01-13 22:33:29 +08:00
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]")
2020-03-04 20:32:29 +08:00
locs, _ = plt.xticks()
2023-10-24 03:58:21 +08:00
plt.xticks(locs[1:], locs[1:] * DT)
plt.xlabel("t [s]")
2020-01-13 22:33:29 +08:00
plt.tight_layout()
2020-01-13 20:25:26 +08:00
plt.draw()
plt.pause(0.1)
2020-03-04 20:32:29 +08:00
def plot_car(x, y, yaw):
"""
2023-10-25 21:17:25 +08:00
Args:
x ():
y ():
yaw ():
"""
LENGTH = 0.5 # [m]
WIDTH = 0.25 # [m]
2020-07-01 22:59:13 +08:00
OFFSET = LENGTH # [m]
2020-03-04 20:32:29 +08:00
outline = np.array(
[
[-OFFSET, (LENGTH - OFFSET), (LENGTH - OFFSET), -OFFSET, -OFFSET],
[WIDTH / 2, WIDTH / 2, -WIDTH / 2, -WIDTH / 2, WIDTH / 2],
]
)
2020-03-04 20:32:29 +08:00
Rotm = np.array([[np.cos(yaw), np.sin(yaw)], [-np.sin(yaw), np.cos(yaw)]])
2020-03-04 20:32:29 +08:00
outline = (outline.T.dot(Rotm)).T
outline[0, :] += x
outline[1, :] += y
plt.plot(
np.array(outline[0, :]).flatten(), np.array(outline[1, :]).flatten(), "tab:blue"
)
2020-03-04 20:32:29 +08:00
2020-01-13 20:25:26 +08:00
def do_sim():
2021-07-08 19:54:48 +08:00
sim = MPCSim()
2020-01-13 20:25:26 +08:00
try:
sim.run()
except Exception as e:
sys.exit(e)
if __name__ == "__main__":
2020-01-13 20:25:26 +08:00
do_sim()