mpc_python_learn/mpc_pybullet_demo/mpc_demo_pybullet.py

296 lines
7.3 KiB
Python
Raw Normal View History

2020-04-08 18:43:38 +08:00
import numpy as np
import matplotlib.pyplot as plt
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-07-01 00:21:27 +08:00
2020-04-08 18:43:38 +08:00
import sys
import time
import pathlib
2020-04-08 18:43:38 +08:00
import pybullet as p
import time
2023-10-24 03:58:21 +08:00
# Params
TARGET_VEL = 1.0 # m/s
L = 0.3 # vehicle wheelbase [m]
T = 5 # Prediction Horizon [s]
DT = 0.2 # discretization step [s]
2022-07-21 19:23:08 +08:00
2020-04-08 18:43:38 +08:00
def get_state(robotId):
2023-10-25 21:17:25 +08:00
"""
Args:
robotId ():
Returns:
"""
2020-04-08 18:43:38 +08:00
robPos, robOrn = p.getBasePositionAndOrientation(robotId)
2022-07-21 19:23:08 +08:00
linVel, angVel = p.getBaseVelocity(robotId)
return np.array(
[
robPos[0],
robPos[1],
np.sqrt(linVel[0] ** 2 + linVel[1] ** 2),
p.getEulerFromQuaternion(robOrn)[2],
]
)
2020-04-08 18:43:38 +08:00
2022-07-21 19:23:08 +08:00
def set_ctrl(robotId, currVel, acceleration, steeringAngle):
2023-10-25 21:17:25 +08:00
"""
Args:
robotId ():
currVel ():
acceleration ():
steeringAngle ():
"""
2022-07-21 19:23:08 +08:00
gearRatio = 1.0 / 21
steering = [0, 2]
wheels = [8, 15]
2020-07-06 23:54:22 +08:00
maxForce = 50
2020-04-08 18:43:38 +08:00
2023-10-24 03:58:21 +08:00
targetVelocity = currVel + acceleration * DT
2020-04-08 18:43:38 +08:00
2020-07-06 23:54:22 +08:00
for wheel in wheels:
2022-07-21 19:23:08 +08:00
p.setJointMotorControl2(
robotId,
wheel,
p.VELOCITY_CONTROL,
targetVelocity=targetVelocity / gearRatio,
force=maxForce,
)
2020-07-06 23:54:22 +08:00
for steer in steering:
2022-07-21 19:23:08 +08:00
p.setJointMotorControl2(
robotId, steer, p.POSITION_CONTROL, targetPosition=steeringAngle
)
def plot_results(path, x_history, y_history):
2020-04-08 18:43:38 +08:00
plt.style.use("ggplot")
plt.figure()
plt.title("MPC Tracking Results")
2022-07-21 19:23:08 +08:00
plt.plot(
path[0, :], path[1, :], c="tab:orange", marker=".", label="reference track"
)
plt.plot(
x_history,
y_history,
c="tab:blue",
marker=".",
alpha=0.5,
label="vehicle trajectory",
)
2020-07-01 00:21:27 +08:00
plt.axis("equal")
2020-04-08 18:43:38 +08:00
plt.legend()
plt.show()
2022-07-21 19:23:08 +08:00
2020-04-08 21:01:42 +08:00
def run_sim():
2020-04-08 18:43:38 +08:00
p.connect(p.GUI)
2022-07-21 19:23:08 +08:00
p.resetDebugVisualizerCamera(
cameraDistance=1.0,
cameraYaw=-90,
cameraPitch=-45,
cameraTargetPosition=[-0.1, -0.0, 0.65],
)
2020-04-08 18:43:38 +08:00
2020-07-06 23:54:22 +08:00
p.resetSimulation()
2020-04-08 21:01:42 +08:00
2022-07-21 19:23:08 +08:00
p.setGravity(0, 0, -10)
2020-07-06 23:54:22 +08:00
useRealTimeSim = 1
2022-07-21 19:23:08 +08:00
p.setTimeStep(1.0 / 120.0)
p.setRealTimeSimulation(useRealTimeSim) # either this
2020-07-06 23:54:22 +08:00
file_path = pathlib.Path(__file__).parent.resolve()
plane = p.loadURDF(str(file_path) + "/racecar/plane.urdf")
car = p.loadURDF(
str(file_path) + "/racecar/f10_racecar/racecar_differential.urdf", [0, 0.3, 0.3]
)
2020-07-06 23:54:22 +08:00
for wheel in range(p.getNumJoints(car)):
2022-07-21 19:23:08 +08:00
# print("joint[",wheel,"]=", p.getJointInfo(car,wheel))
p.setJointMotorControl2(
car, wheel, p.VELOCITY_CONTROL, targetVelocity=0, force=0
)
p.getJointInfo(car, wheel)
c = p.createConstraint(
car,
9,
car,
11,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=1, maxForce=10000)
c = p.createConstraint(
car,
10,
car,
13,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=-1, maxForce=10000)
c = p.createConstraint(
car,
9,
car,
13,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=-1, maxForce=10000)
c = p.createConstraint(
car,
16,
car,
18,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=1, maxForce=10000)
c = p.createConstraint(
car,
16,
car,
19,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=-1, maxForce=10000)
c = p.createConstraint(
car,
17,
car,
19,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=-1, maxForce=10000)
c = p.createConstraint(
car,
1,
car,
18,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=-1, gearAuxLink=15, maxForce=10000)
c = p.createConstraint(
car,
3,
car,
19,
jointType=p.JOINT_GEAR,
jointAxis=[0, 1, 0],
parentFramePosition=[0, 0, 0],
childFramePosition=[0, 0, 0],
)
p.changeConstraint(c, gearRatio=-1, gearAuxLink=15, maxForce=10000)
2020-04-08 18:43:38 +08:00
# Interpolated Path to follow given waypoints
2022-07-21 19:23:08 +08:00
path = compute_path_from_wp(
[0, 3, 4, 6, 10, 11, 12, 6, 1, 0],
[0, 0, 2, 4, 3, 3, -1, -6, -2, -2],
2023-10-13 04:36:42 +08:00
0.05,
2022-07-21 19:23:08 +08:00
)
for x_, y_ in zip(path[0, :], path[1, :]):
p.addUserDebugLine([x_, y_, 0], [x_, y_, 0.33], [0, 0, 1])
2020-07-01 00:21:27 +08:00
2023-10-26 00:04:35 +08:00
# starting conditions
2023-11-10 17:40:55 +08:00
control = np.zeros(2)
2022-07-21 19:23:08 +08:00
# Cost Matrices
2023-04-23 23:44:43 +08:00
Q = [20, 20, 10, 20] # state error cost [x,y,v,yaw]
Qf = [30, 30, 30, 30] # state error cost at final timestep [x,y,v,yaw]
R = [10, 10] # input cost [acc ,steer]
P = [10, 10] # input rate of change cost [acc ,steer]
2022-07-21 19:23:08 +08:00
2023-10-24 04:02:41 +08:00
mpc = MPC(VehicleModel(), T, DT, Q, Qf, R, P)
2022-07-21 19:23:08 +08:00
x_history = []
y_history = []
2020-04-08 18:43:38 +08:00
2022-07-21 21:14:14 +08:00
input("\033[92m Press Enter to continue... \033[0m")
2022-07-21 19:23:08 +08:00
while 1:
state = get_state(car)
2020-04-08 18:43:38 +08:00
x_history.append(state[0])
y_history.append(state[1])
2022-07-21 19:23:08 +08:00
# track path in bullet
p.addUserDebugLine(
[state[0], state[1], 0], [state[0], state[1], 0.5], [1, 0, 0]
)
2020-04-08 18:43:38 +08:00
2022-07-21 19:23:08 +08:00
if np.sqrt((state[0] - path[0, -1]) ** 2 + (state[1] - path[1, -1]) ** 2) < 0.2:
2020-04-08 18:43:38 +08:00
print("Success! Goal Reached")
2022-07-21 19:23:08 +08:00
set_ctrl(car, 0, 0, 0)
plot_results(path, x_history, y_history)
input("Press Enter to continue...")
2020-04-08 18:43:38 +08:00
p.disconnect()
return
2023-10-13 04:36:42 +08:00
# Get Reference_traj
# NOTE: inputs are in world frame
2023-10-24 04:02:41 +08:00
target = get_ref_trajectory(state, path, TARGET_VEL, T, DT)
2023-10-13 04:36:42 +08:00
2023-04-23 23:44:43 +08:00
# for MPC base link frame is used:
# so x, y, yaw are 0.0, but speed is the same
2023-10-13 04:36:42 +08:00
ego_state = np.array([0.0, 0.0, state[2], 0.0])
2023-04-23 23:44:43 +08:00
# to account for MPC latency
2023-10-13 04:36:42 +08:00
# simulate one timestep actuation delay
2023-10-24 03:58:21 +08:00
ego_state[0] = ego_state[0] + ego_state[2] * np.cos(ego_state[3]) * DT
ego_state[1] = ego_state[1] + ego_state[2] * np.sin(ego_state[3]) * DT
2023-11-10 17:40:55 +08:00
ego_state[2] = ego_state[2] + control[0] * DT
ego_state[3] = ego_state[3] + control[0] * np.tan(control[1]) / L * DT
2023-10-13 04:36:42 +08:00
# optimization loop
start = time.time()
2022-07-21 19:23:08 +08:00
2023-04-22 00:24:28 +08:00
# MPC step
2023-11-10 17:40:55 +08:00
_, u_mpc = mpc.step(ego_state, target, control, verbose=False)
control[0] = u_mpc.value[0, 0]
control[1] = u_mpc.value[1, 0]
2022-07-21 19:23:08 +08:00
elapsed = time.time() - start
2020-04-08 18:43:38 +08:00
print("CVXPY Optimization Time: {:.4f}s".format(elapsed))
2023-11-10 17:40:55 +08:00
set_ctrl(car, state[2], control[0], control[1])
2022-07-21 19:23:08 +08:00
2023-10-24 03:58:21 +08:00
if DT - elapsed > 0:
time.sleep(DT - elapsed)
2020-04-08 18:43:38 +08:00
2022-07-21 19:23:08 +08:00
if __name__ == "__main__":
2020-04-08 21:01:42 +08:00
run_sim()