mirror of
https://github.com/AtsushiSakai/PythonRobotics.git
synced 2026-01-13 16:47:55 -05:00
22cbee4921
Updated all instances of "Ref:" to "Reference" for consistency in both code and documentation. This change improves clarity and aligns with standard terminology practices.
675 lines
24 KiB
Python
675 lines
24 KiB
Python
"""
|
||
|
||
A rocket powered landing with successive convexification
|
||
|
||
author: Sven Niederberger
|
||
Atsushi Sakai
|
||
|
||
Reference:
|
||
- Python implementation of 'Successive Convexification for 6-DoF Mars Rocket Powered Landing with Free-Final-Time' paper
|
||
by Michael Szmuk and Behcet Acıkmese.
|
||
|
||
- EmbersArc/SuccessiveConvexificationFreeFinalTime: Implementation of "Successive Convexification for 6-DoF Mars Rocket Powered Landing with Free-Final-Time" https://github.com/EmbersArc/SuccessiveConvexificationFreeFinalTime
|
||
|
||
"""
|
||
import warnings
|
||
from time import time
|
||
import numpy as np
|
||
from scipy.integrate import odeint
|
||
import cvxpy
|
||
import matplotlib.pyplot as plt
|
||
|
||
# Trajectory points
|
||
K = 50
|
||
|
||
# Max solver iterations
|
||
iterations = 30
|
||
|
||
# Weight constants
|
||
W_SIGMA = 1 # flight time
|
||
W_DELTA = 1e-3 # difference in state/input
|
||
W_DELTA_SIGMA = 1e-1 # difference in flight time
|
||
W_NU = 1e5 # virtual control
|
||
|
||
print(cvxpy.installed_solvers())
|
||
solver = 'ECOS'
|
||
verbose_solver = False
|
||
|
||
show_animation = True
|
||
|
||
|
||
class Rocket_Model_6DoF:
|
||
"""
|
||
A 6 degree of freedom rocket landing problem.
|
||
"""
|
||
|
||
def __init__(self, rng):
|
||
"""
|
||
A large r_scale for a small scale problem will
|
||
ead to numerical problems as parameters become excessively small
|
||
and (it seems) precision is lost in the dynamics.
|
||
"""
|
||
self.n_x = 14
|
||
self.n_u = 3
|
||
|
||
# Mass
|
||
self.m_wet = 3.0 # 30000 kg
|
||
self.m_dry = 2.2 # 22000 kg
|
||
|
||
# Flight time guess
|
||
self.t_f_guess = 10.0 # 10 s
|
||
|
||
# State constraints
|
||
self.r_I_final = np.array((0., 0., 0.))
|
||
self.v_I_final = np.array((-1e-1, 0., 0.))
|
||
self.q_B_I_final = self.euler_to_quat((0, 0, 0))
|
||
self.w_B_final = np.deg2rad(np.array((0., 0., 0.)))
|
||
|
||
self.w_B_max = np.deg2rad(60)
|
||
|
||
# Angles
|
||
max_gimbal = 20
|
||
max_angle = 90
|
||
glidelslope_angle = 20
|
||
|
||
self.tan_delta_max = np.tan(np.deg2rad(max_gimbal))
|
||
self.cos_theta_max = np.cos(np.deg2rad(max_angle))
|
||
self.tan_gamma_gs = np.tan(np.deg2rad(glidelslope_angle))
|
||
|
||
# Thrust limits
|
||
self.T_max = 5.0
|
||
self.T_min = 0.3
|
||
|
||
# Angular moment of inertia
|
||
self.J_B = 1e-2 * np.diag([1., 1., 1.])
|
||
|
||
# Gravity
|
||
self.g_I = np.array((-1, 0., 0.))
|
||
|
||
# Fuel consumption
|
||
self.alpha_m = 0.01
|
||
|
||
# Vector from thrust point to CoM
|
||
self.r_T_B = np.array([-1e-2, 0., 0.])
|
||
|
||
self.set_random_initial_state(rng)
|
||
|
||
self.x_init = np.concatenate(
|
||
((self.m_wet,), self.r_I_init, self.v_I_init, self.q_B_I_init, self.w_B_init))
|
||
self.x_final = np.concatenate(
|
||
((self.m_dry,), self.r_I_final, self.v_I_final, self.q_B_I_final, self.w_B_final))
|
||
|
||
self.r_scale = np.linalg.norm(self.r_I_init)
|
||
self.m_scale = self.m_wet
|
||
|
||
def set_random_initial_state(self, rng):
|
||
if rng is None:
|
||
rng = np.random.default_rng()
|
||
|
||
self.r_I_init = np.array((0., 0., 0.))
|
||
self.r_I_init[0] = rng.uniform(3, 4)
|
||
self.r_I_init[1:3] = rng.uniform(-2, 2, size=2)
|
||
|
||
self.v_I_init = np.array((0., 0., 0.))
|
||
self.v_I_init[0] = rng.uniform(-1, -0.5)
|
||
self.v_I_init[1:3] = rng.uniform(-0.5, -0.2,
|
||
size=2) * self.r_I_init[1:3]
|
||
|
||
self.q_B_I_init = self.euler_to_quat((0,
|
||
rng.uniform(-30, 30),
|
||
rng.uniform(-30, 30)))
|
||
self.w_B_init = np.deg2rad((0,
|
||
rng.uniform(-20, 20),
|
||
rng.uniform(-20, 20)))
|
||
|
||
def f_func(self, x, u):
|
||
m, _, _, _, vx, vy, vz, q0, q1, q2, q3, wx, wy, wz = x[0], x[1], x[
|
||
2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13]
|
||
ux, uy, uz = u[0], u[1], u[2]
|
||
|
||
return np.array([
|
||
[-0.01 * np.sqrt(ux**2 + uy**2 + uz**2)],
|
||
[vx],
|
||
[vy],
|
||
[vz],
|
||
[(-1.0 * m - ux * (2 * q2**2 + 2 * q3**2 - 1) - 2 * uy
|
||
* (q0 * q3 - q1 * q2) + 2 * uz * (q0 * q2 + q1 * q3)) / m],
|
||
[(2 * ux * (q0 * q3 + q1 * q2) - uy * (2 * q1**2
|
||
+ 2 * q3**2 - 1) - 2 * uz * (q0 * q1 - q2 * q3)) / m],
|
||
[(-2 * ux * (q0 * q2 - q1 * q3) + 2 * uy
|
||
* (q0 * q1 + q2 * q3) - uz * (2 * q1**2 + 2 * q2**2 - 1)) / m],
|
||
[-0.5 * q1 * wx - 0.5 * q2 * wy - 0.5 * q3 * wz],
|
||
[0.5 * q0 * wx + 0.5 * q2 * wz - 0.5 * q3 * wy],
|
||
[0.5 * q0 * wy - 0.5 * q1 * wz + 0.5 * q3 * wx],
|
||
[0.5 * q0 * wz + 0.5 * q1 * wy - 0.5 * q2 * wx],
|
||
[0],
|
||
[1.0 * uz],
|
||
[-1.0 * uy]
|
||
])
|
||
|
||
def A_func(self, x, u):
|
||
m, _, _, _, _, _, _, q0, q1, q2, q3, wx, wy, wz = x[0], x[1], x[
|
||
2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13]
|
||
ux, uy, uz = u[0], u[1], u[2]
|
||
|
||
return np.array([
|
||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||
[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
|
||
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
|
||
[(ux * (2 * q2**2 + 2 * q3**2 - 1) + 2 * uy * (q0 * q3 - q1 * q2) - 2 * uz * (q0 * q2 + q1 * q3)) / m**2, 0, 0, 0, 0, 0, 0, 2 * (q2 * uz
|
||
- q3 * uy) / m, 2 * (q2 * uy + q3 * uz) / m, 2 * (q0 * uz + q1 * uy - 2 * q2 * ux) / m, 2 * (-q0 * uy + q1 * uz - 2 * q3 * ux) / m, 0, 0, 0],
|
||
[(-2 * ux * (q0 * q3 + q1 * q2) + uy * (2 * q1**2 + 2 * q3**2 - 1) + 2 * uz * (q0 * q1 - q2 * q3)) / m**2, 0, 0, 0, 0, 0, 0, 2 * (-q1 * uz
|
||
+ q3 * ux) / m, 2 * (-q0 * uz - 2 * q1 * uy + q2 * ux) / m, 2 * (q1 * ux + q3 * uz) / m, 2 * (q0 * ux + q2 * uz - 2 * q3 * uy) / m, 0, 0, 0],
|
||
[(2 * ux * (q0 * q2 - q1 * q3) - 2 * uy * (q0 * q1 + q2 * q3) + uz * (2 * q1**2 + 2 * q2**2 - 1)) / m**2, 0, 0, 0, 0, 0, 0, 2 * (q1 * uy
|
||
- q2 * ux) / m, 2 * (q0 * uy - 2 * q1 * uz + q3 * ux) / m, 2 * (-q0 * ux - 2 * q2 * uz + q3 * uy) / m, 2 * (q1 * ux + q2 * uy) / m, 0, 0, 0],
|
||
[0, 0, 0, 0, 0, 0, 0, 0, -0.5 * wx, -0.5 * wy,
|
||
- 0.5 * wz, -0.5 * q1, -0.5 * q2, -0.5 * q3],
|
||
[0, 0, 0, 0, 0, 0, 0, 0.5 * wx, 0, 0.5 * wz,
|
||
- 0.5 * wy, 0.5 * q0, -0.5 * q3, 0.5 * q2],
|
||
[0, 0, 0, 0, 0, 0, 0, 0.5 * wy, -0.5 * wz, 0,
|
||
0.5 * wx, 0.5 * q3, 0.5 * q0, -0.5 * q1],
|
||
[0, 0, 0, 0, 0, 0, 0, 0.5 * wz, 0.5 * wy,
|
||
- 0.5 * wx, 0, -0.5 * q2, 0.5 * q1, 0.5 * q0],
|
||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
|
||
|
||
def B_func(self, x, u):
|
||
m, _, _, _, _, _, _, q0, q1, q2, q3, _, _, _ = x[0], x[1], x[
|
||
2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13]
|
||
ux, uy, uz = u[0], u[1], u[2]
|
||
|
||
return np.array([
|
||
[-0.01 * ux / np.sqrt(ux**2 + uy**2 + uz**2),
|
||
-0.01 * uy / np.sqrt(ux ** 2 + uy**2 + uz**2),
|
||
-0.01 * uz / np.sqrt(ux**2 + uy**2 + uz**2)],
|
||
[0, 0, 0],
|
||
[0, 0, 0],
|
||
[0, 0, 0],
|
||
[(-2 * q2**2 - 2 * q3**2 + 1) / m, 2
|
||
* (-q0 * q3 + q1 * q2) / m, 2 * (q0 * q2 + q1 * q3) / m],
|
||
[2 * (q0 * q3 + q1 * q2) / m, (-2 * q1**2 - 2
|
||
* q3**2 + 1) / m, 2 * (-q0 * q1 + q2 * q3) / m],
|
||
[2 * (-q0 * q2 + q1 * q3) / m, 2 * (q0 * q1 + q2 * q3)
|
||
/ m, (-2 * q1**2 - 2 * q2**2 + 1) / m],
|
||
[0, 0, 0],
|
||
[0, 0, 0],
|
||
[0, 0, 0],
|
||
[0, 0, 0],
|
||
[0, 0, 0],
|
||
[0, 0, 1.0],
|
||
[0, -1.0, 0]
|
||
])
|
||
|
||
def euler_to_quat(self, a):
|
||
a = np.deg2rad(a)
|
||
|
||
cy = np.cos(a[1] * 0.5)
|
||
sy = np.sin(a[1] * 0.5)
|
||
cr = np.cos(a[0] * 0.5)
|
||
sr = np.sin(a[0] * 0.5)
|
||
cp = np.cos(a[2] * 0.5)
|
||
sp = np.sin(a[2] * 0.5)
|
||
|
||
q = np.zeros(4)
|
||
|
||
q[0] = cy * cr * cp + sy * sr * sp
|
||
q[1] = cy * sr * cp - sy * cr * sp
|
||
q[3] = cy * cr * sp + sy * sr * cp
|
||
q[2] = sy * cr * cp - cy * sr * sp
|
||
|
||
return q
|
||
|
||
def skew(self, v):
|
||
return np.array([
|
||
[0, -v[2], v[1]],
|
||
[v[2], 0, -v[0]],
|
||
[-v[1], v[0], 0]
|
||
])
|
||
|
||
def dir_cosine(self, q):
|
||
return np.array([
|
||
[1 - 2 * (q[2] ** 2 + q[3] ** 2), 2 * (q[1] * q[2]
|
||
+ q[0] * q[3]), 2 * (q[1] * q[3] - q[0] * q[2])],
|
||
[2 * (q[1] * q[2] - q[0] * q[3]), 1 - 2
|
||
* (q[1] ** 2 + q[3] ** 2), 2 * (q[2] * q[3] + q[0] * q[1])],
|
||
[2 * (q[1] * q[3] + q[0] * q[2]), 2 * (q[2] * q[3]
|
||
- q[0] * q[1]), 1 - 2 * (q[1] ** 2 + q[2] ** 2)]
|
||
])
|
||
|
||
def omega(self, w):
|
||
return np.array([
|
||
[0, -w[0], -w[1], -w[2]],
|
||
[w[0], 0, w[2], -w[1]],
|
||
[w[1], -w[2], 0, w[0]],
|
||
[w[2], w[1], -w[0], 0],
|
||
])
|
||
|
||
def initialize_trajectory(self, X, U):
|
||
"""
|
||
Initialize the trajectory with linear approximation.
|
||
"""
|
||
K = X.shape[1]
|
||
|
||
for k in range(K):
|
||
alpha1 = (K - k) / K
|
||
alpha2 = k / K
|
||
|
||
m_k = (alpha1 * self.x_init[0] + alpha2 * self.x_final[0],)
|
||
r_I_k = alpha1 * self.x_init[1:4] + alpha2 * self.x_final[1:4]
|
||
v_I_k = alpha1 * self.x_init[4:7] + alpha2 * self.x_final[4:7]
|
||
q_B_I_k = np.array([1, 0, 0, 0])
|
||
w_B_k = alpha1 * self.x_init[11:14] + alpha2 * self.x_final[11:14]
|
||
|
||
X[:, k] = np.concatenate((m_k, r_I_k, v_I_k, q_B_I_k, w_B_k))
|
||
U[:, k] = m_k * -self.g_I
|
||
|
||
return X, U
|
||
|
||
def get_constraints(self, X_v, U_v, X_last_p, U_last_p):
|
||
"""
|
||
Get model specific constraints.
|
||
|
||
:param X_v: cvx variable for current states
|
||
:param U_v: cvx variable for current inputs
|
||
:param X_last_p: cvx parameter for last states
|
||
:param U_last_p: cvx parameter for last inputs
|
||
:return: A list of cvx constraints
|
||
"""
|
||
# Boundary conditions:
|
||
constraints = [
|
||
X_v[0, 0] == self.x_init[0],
|
||
X_v[1:4, 0] == self.x_init[1:4],
|
||
X_v[4:7, 0] == self.x_init[4:7],
|
||
# X_v[7:11, 0] == self.x_init[7:11], # initial orientation is free
|
||
X_v[11:14, 0] == self.x_init[11:14],
|
||
|
||
# X_[0, -1] final mass is free
|
||
X_v[1:, -1] == self.x_final[1:],
|
||
U_v[1:3, -1] == 0,
|
||
]
|
||
|
||
constraints += [
|
||
# State constraints:
|
||
X_v[0, :] >= self.m_dry, # minimum mass
|
||
cvxpy.norm(X_v[2: 4, :], axis=0) <= X_v[1, :] / \
|
||
self.tan_gamma_gs, # glideslope
|
||
cvxpy.norm(X_v[9:11, :], axis=0) <= np.sqrt(
|
||
(1 - self.cos_theta_max) / 2), # maximum angle
|
||
# maximum angular velocity
|
||
cvxpy.norm(X_v[11: 14, :], axis=0) <= self.w_B_max,
|
||
|
||
# Control constraints:
|
||
cvxpy.norm(U_v[1:3, :], axis=0) <= self.tan_delta_max * \
|
||
U_v[0, :], # gimbal angle constraint
|
||
cvxpy.norm(U_v, axis=0) <= self.T_max, # upper thrust constraint
|
||
]
|
||
|
||
# linearized lower thrust constraint
|
||
rhs = [U_last_p[:, k] / cvxpy.norm(U_last_p[:, k]) @ U_v[:, k]
|
||
for k in range(X_v.shape[1])]
|
||
constraints += [
|
||
self.T_min <= cvxpy.vstack(rhs)
|
||
]
|
||
|
||
return constraints
|
||
|
||
|
||
class Integrator:
|
||
def __init__(self, m, K):
|
||
self.K = K
|
||
self.m = m
|
||
self.n_x = m.n_x
|
||
self.n_u = m.n_u
|
||
|
||
self.A_bar = np.zeros([m.n_x * m.n_x, K - 1])
|
||
self.B_bar = np.zeros([m.n_x * m.n_u, K - 1])
|
||
self.C_bar = np.zeros([m.n_x * m.n_u, K - 1])
|
||
self.S_bar = np.zeros([m.n_x, K - 1])
|
||
self.z_bar = np.zeros([m.n_x, K - 1])
|
||
|
||
# vector indices for flat matrices
|
||
x_end = m.n_x
|
||
A_bar_end = m.n_x * (1 + m.n_x)
|
||
B_bar_end = m.n_x * (1 + m.n_x + m.n_u)
|
||
C_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u)
|
||
S_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 1)
|
||
z_bar_end = m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2)
|
||
self.x_ind = slice(0, x_end)
|
||
self.A_bar_ind = slice(x_end, A_bar_end)
|
||
self.B_bar_ind = slice(A_bar_end, B_bar_end)
|
||
self.C_bar_ind = slice(B_bar_end, C_bar_end)
|
||
self.S_bar_ind = slice(C_bar_end, S_bar_end)
|
||
self.z_bar_ind = slice(S_bar_end, z_bar_end)
|
||
|
||
self.f, self.A, self.B = m.f_func, m.A_func, m.B_func
|
||
|
||
# integration initial condition
|
||
self.V0 = np.zeros((m.n_x * (1 + m.n_x + m.n_u + m.n_u + 2),))
|
||
self.V0[self.A_bar_ind] = np.eye(m.n_x).reshape(-1)
|
||
|
||
self.dt = 1. / (K - 1)
|
||
|
||
def calculate_discretization(self, X, U, sigma):
|
||
"""
|
||
Calculate discretization for given states, inputs and total time.
|
||
|
||
:param X: Matrix of states for all time points
|
||
:param U: Matrix of inputs for all time points
|
||
:param sigma: Total time
|
||
:return: The discretization matrices
|
||
"""
|
||
for k in range(self.K - 1):
|
||
self.V0[self.x_ind] = X[:, k]
|
||
V = np.array(odeint(self._ode_dVdt, self.V0, (0, self.dt),
|
||
args=(U[:, k], U[:, k + 1], sigma))[1, :])
|
||
|
||
# using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
|
||
# flatten matrices in column-major (Fortran) order for CVXPY
|
||
Phi = V[self.A_bar_ind].reshape((self.n_x, self.n_x))
|
||
self.A_bar[:, k] = Phi.flatten(order='F')
|
||
self.B_bar[:, k] = np.matmul(Phi, V[self.B_bar_ind].reshape(
|
||
(self.n_x, self.n_u))).flatten(order='F')
|
||
self.C_bar[:, k] = np.matmul(Phi, V[self.C_bar_ind].reshape(
|
||
(self.n_x, self.n_u))).flatten(order='F')
|
||
self.S_bar[:, k] = np.matmul(Phi, V[self.S_bar_ind])
|
||
self.z_bar[:, k] = np.matmul(Phi, V[self.z_bar_ind])
|
||
|
||
return self.A_bar, self.B_bar, self.C_bar, self.S_bar, self.z_bar
|
||
|
||
def _ode_dVdt(self, V, t, u_t0, u_t1, sigma):
|
||
"""
|
||
ODE function to compute dVdt.
|
||
|
||
:param V: Evaluation state V = [x, Phi_A, B_bar, C_bar, S_bar, z_bar]
|
||
:param t: Evaluation time
|
||
:param u_t0: Input at start of interval
|
||
:param u_t1: Input at end of interval
|
||
:param sigma: Total time
|
||
:return: Derivative at current time and state dVdt
|
||
"""
|
||
alpha = (self.dt - t) / self.dt
|
||
beta = t / self.dt
|
||
x = V[self.x_ind]
|
||
u = u_t0 + beta * (u_t1 - u_t0)
|
||
|
||
# using \Phi_A(\tau_{k+1},\xi) = \Phi_A(\tau_{k+1},\tau_k)\Phi_A(\xi,\tau_k)^{-1}
|
||
# and pre-multiplying with \Phi_A(\tau_{k+1},\tau_k) after integration
|
||
Phi_A_xi = np.linalg.inv(
|
||
V[self.A_bar_ind].reshape((self.n_x, self.n_x)))
|
||
|
||
A_subs = sigma * self.A(x, u)
|
||
B_subs = sigma * self.B(x, u)
|
||
f_subs = self.f(x, u)
|
||
|
||
dVdt = np.zeros_like(V)
|
||
dVdt[self.x_ind] = sigma * f_subs.transpose()
|
||
dVdt[self.A_bar_ind] = np.matmul(
|
||
A_subs, V[self.A_bar_ind].reshape((self.n_x, self.n_x))).reshape(-1)
|
||
dVdt[self.B_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * alpha
|
||
dVdt[self.C_bar_ind] = np.matmul(Phi_A_xi, B_subs).reshape(-1) * beta
|
||
dVdt[self.S_bar_ind] = np.matmul(Phi_A_xi, f_subs).transpose()
|
||
z_t = -np.matmul(A_subs, x) - np.matmul(B_subs, u)
|
||
dVdt[self.z_bar_ind] = np.dot(Phi_A_xi, z_t.T).flatten()
|
||
|
||
return dVdt
|
||
|
||
|
||
class SCProblem:
|
||
"""
|
||
Defines a standard Successive Convexification problem and
|
||
adds the model specific constraints and objectives.
|
||
|
||
:param m: The model object
|
||
:param K: Number of discretization points
|
||
"""
|
||
|
||
def __init__(self, m, K):
|
||
# Variables:
|
||
self.var = dict()
|
||
self.var['X'] = cvxpy.Variable((m.n_x, K))
|
||
self.var['U'] = cvxpy.Variable((m.n_u, K))
|
||
self.var['sigma'] = cvxpy.Variable(nonneg=True)
|
||
self.var['nu'] = cvxpy.Variable((m.n_x, K - 1))
|
||
self.var['delta_norm'] = cvxpy.Variable(nonneg=True)
|
||
self.var['sigma_norm'] = cvxpy.Variable(nonneg=True)
|
||
|
||
# Parameters:
|
||
self.par = dict()
|
||
self.par['A_bar'] = cvxpy.Parameter((m.n_x * m.n_x, K - 1))
|
||
self.par['B_bar'] = cvxpy.Parameter((m.n_x * m.n_u, K - 1))
|
||
self.par['C_bar'] = cvxpy.Parameter((m.n_x * m.n_u, K - 1))
|
||
self.par['S_bar'] = cvxpy.Parameter((m.n_x, K - 1))
|
||
self.par['z_bar'] = cvxpy.Parameter((m.n_x, K - 1))
|
||
|
||
self.par['X_last'] = cvxpy.Parameter((m.n_x, K))
|
||
self.par['U_last'] = cvxpy.Parameter((m.n_u, K))
|
||
self.par['sigma_last'] = cvxpy.Parameter(nonneg=True)
|
||
|
||
self.par['weight_sigma'] = cvxpy.Parameter(nonneg=True)
|
||
self.par['weight_delta'] = cvxpy.Parameter(nonneg=True)
|
||
self.par['weight_delta_sigma'] = cvxpy.Parameter(nonneg=True)
|
||
self.par['weight_nu'] = cvxpy.Parameter(nonneg=True)
|
||
|
||
# Constraints:
|
||
constraints = []
|
||
|
||
# Model:
|
||
constraints += m.get_constraints(
|
||
self.var['X'], self.var['U'], self.par['X_last'], self.par['U_last'])
|
||
|
||
# Dynamics:
|
||
# x_t+1 = A_*x_t+B_*U_t+C_*U_T+1*S_*sigma+zbar+nu
|
||
constraints += [
|
||
self.var['X'][:, k + 1] ==
|
||
cvxpy.reshape(self.par['A_bar'][:, k], (m.n_x, m.n_x), order='F') @
|
||
self.var['X'][:, k] +
|
||
cvxpy.reshape(self.par['B_bar'][:, k], (m.n_x, m.n_u), order='F') @
|
||
self.var['U'][:, k] +
|
||
cvxpy.reshape(self.par['C_bar'][:, k], (m.n_x, m.n_u), order='F') @
|
||
self.var['U'][:, k + 1] +
|
||
self.par['S_bar'][:, k] * self.var['sigma'] +
|
||
self.par['z_bar'][:, k] +
|
||
self.var['nu'][:, k]
|
||
for k in range(K - 1)
|
||
]
|
||
|
||
# Trust regions:
|
||
dx = cvxpy.sum(cvxpy.square(
|
||
self.var['X'] - self.par['X_last']), axis=0)
|
||
du = cvxpy.sum(cvxpy.square(
|
||
self.var['U'] - self.par['U_last']), axis=0)
|
||
ds = self.var['sigma'] - self.par['sigma_last']
|
||
constraints += [cvxpy.norm(dx + du, 1) <= self.var['delta_norm']]
|
||
constraints += [cvxpy.norm(ds, 'inf') <= self.var['sigma_norm']]
|
||
|
||
# Flight time positive:
|
||
constraints += [self.var['sigma'] >= 0.1]
|
||
|
||
# Objective:
|
||
sc_objective = cvxpy.Minimize(
|
||
self.par['weight_sigma'] * self.var['sigma'] +
|
||
self.par['weight_nu'] * cvxpy.norm(self.var['nu'], 'inf') +
|
||
self.par['weight_delta'] * self.var['delta_norm'] +
|
||
self.par['weight_delta_sigma'] * self.var['sigma_norm']
|
||
)
|
||
|
||
objective = sc_objective
|
||
|
||
self.prob = cvxpy.Problem(objective, constraints)
|
||
|
||
def set_parameters(self, **kwargs):
|
||
"""
|
||
All parameters have to be filled before calling solve().
|
||
Takes the following arguments as keywords:
|
||
|
||
A_bar
|
||
B_bar
|
||
C_bar
|
||
S_bar
|
||
z_bar
|
||
X_last
|
||
U_last
|
||
sigma_last
|
||
E
|
||
weight_sigma
|
||
weight_nu
|
||
radius_trust_region
|
||
"""
|
||
|
||
for key in kwargs:
|
||
if key in self.par:
|
||
self.par[key].value = kwargs[key]
|
||
else:
|
||
print(f'Parameter \'{key}\' does not exist.')
|
||
|
||
def get_variable(self, name):
|
||
if name in self.var:
|
||
return self.var[name].value
|
||
else:
|
||
print(f'Variable \'{name}\' does not exist.')
|
||
return None
|
||
|
||
def solve(self, **kwargs):
|
||
error = False
|
||
try:
|
||
with warnings.catch_warnings(): # For User warning from solver
|
||
warnings.simplefilter('ignore')
|
||
self.prob.solve(verbose=verbose_solver,
|
||
solver=solver)
|
||
except cvxpy.SolverError:
|
||
error = True
|
||
|
||
stats = self.prob.solver_stats
|
||
|
||
info = {
|
||
'setup_time': stats.setup_time,
|
||
'solver_time': stats.solve_time,
|
||
'iterations': stats.num_iters,
|
||
'solver_error': error
|
||
}
|
||
|
||
return info
|
||
|
||
|
||
def axis3d_equal(X, Y, Z, ax):
|
||
|
||
max_range = np.array([X.max() - X.min(), Y.max()
|
||
- Y.min(), Z.max() - Z.min()]).max()
|
||
Xb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2,
|
||
- 1:2:2][0].flatten() + 0.5 * (X.max() + X.min())
|
||
Yb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2,
|
||
- 1:2:2][1].flatten() + 0.5 * (Y.max() + Y.min())
|
||
Zb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2,
|
||
- 1:2:2][2].flatten() + 0.5 * (Z.max() + Z.min())
|
||
# Comment or uncomment following both lines to test the fake bounding box:
|
||
for xb, yb, zb in zip(Xb, Yb, Zb):
|
||
ax.plot([xb], [yb], [zb], 'w')
|
||
|
||
|
||
def plot_animation(X, U): # pragma: no cover
|
||
|
||
fig = plt.figure()
|
||
ax = fig.add_subplot(projection='3d')
|
||
# for stopping simulation with the esc key.
|
||
fig.canvas.mpl_connect('key_release_event',
|
||
lambda event: [exit(0) if event.key == 'escape' else None])
|
||
|
||
for k in range(K):
|
||
plt.cla()
|
||
ax.plot(X[2, :], X[3, :], X[1, :]) # trajectory
|
||
ax.scatter3D([0.0], [0.0], [0.0], c="r",
|
||
marker="x") # target landing point
|
||
axis3d_equal(X[2, :], X[3, :], X[1, :], ax)
|
||
|
||
rx, ry, rz = X[1:4, k]
|
||
# vx, vy, vz = X[4:7, k]
|
||
qw, qx, qy, qz = X[7:11, k]
|
||
|
||
CBI = np.array([
|
||
[1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy + qw * qz),
|
||
2 * (qx * qz - qw * qy)],
|
||
[2 * (qx * qy - qw * qz), 1 - 2
|
||
* (qx ** 2 + qz ** 2), 2 * (qy * qz + qw * qx)],
|
||
[2 * (qx * qz + qw * qy), 2 * (qy * qz - qw * qx),
|
||
1 - 2 * (qx ** 2 + qy ** 2)]
|
||
])
|
||
|
||
Fx, Fy, Fz = np.dot(np.transpose(CBI), U[:, k])
|
||
dx, dy, dz = np.dot(np.transpose(CBI), np.array([1., 0., 0.]))
|
||
|
||
# attitude vector
|
||
ax.quiver(ry, rz, rx, dy, dz, dx, length=0.5, linewidth=3.0,
|
||
arrow_length_ratio=0.0, color='black')
|
||
|
||
# thrust vector
|
||
ax.quiver(ry, rz, rx, -Fy, -Fz, -Fx, length=0.1,
|
||
arrow_length_ratio=0.0, color='red')
|
||
|
||
ax.set_title("Rocket powered landing")
|
||
plt.pause(0.5)
|
||
|
||
|
||
def main(rng=None):
|
||
print("start!!")
|
||
m = Rocket_Model_6DoF(rng)
|
||
|
||
# state and input list
|
||
X = np.empty(shape=[m.n_x, K])
|
||
U = np.empty(shape=[m.n_u, K])
|
||
|
||
# INITIALIZATION
|
||
sigma = m.t_f_guess
|
||
X, U = m.initialize_trajectory(X, U)
|
||
|
||
integrator = Integrator(m, K)
|
||
problem = SCProblem(m, K)
|
||
|
||
converged = False
|
||
w_delta = W_DELTA
|
||
for it in range(iterations):
|
||
t0_it = time()
|
||
print('-' * 18 + f' Iteration {str(it + 1).zfill(2)} ' + '-' * 18)
|
||
|
||
A_bar, B_bar, C_bar, S_bar, z_bar = integrator.calculate_discretization(
|
||
X, U, sigma)
|
||
|
||
problem.set_parameters(A_bar=A_bar, B_bar=B_bar, C_bar=C_bar, S_bar=S_bar, z_bar=z_bar,
|
||
X_last=X, U_last=U, sigma_last=sigma,
|
||
weight_sigma=W_SIGMA, weight_nu=W_NU,
|
||
weight_delta=w_delta, weight_delta_sigma=W_DELTA_SIGMA)
|
||
problem.solve()
|
||
|
||
X = problem.get_variable('X')
|
||
U = problem.get_variable('U')
|
||
sigma = problem.get_variable('sigma')
|
||
|
||
delta_norm = problem.get_variable('delta_norm')
|
||
sigma_norm = problem.get_variable('sigma_norm')
|
||
nu_norm = np.linalg.norm(problem.get_variable('nu'), np.inf)
|
||
|
||
print('delta_norm', delta_norm)
|
||
print('sigma_norm', sigma_norm)
|
||
print('nu_norm', nu_norm)
|
||
|
||
if delta_norm < 1e-3 and sigma_norm < 1e-3 and nu_norm < 1e-7:
|
||
converged = True
|
||
|
||
w_delta *= 1.5
|
||
|
||
print('Time for iteration', time() - t0_it, 's')
|
||
|
||
if converged:
|
||
print(f'Converged after {it + 1} iterations.')
|
||
break
|
||
|
||
if show_animation: # pragma: no cover
|
||
plot_animation(X, U)
|
||
|
||
print("done!!")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|