-
Notifications
You must be signed in to change notification settings - Fork 3
/
quad_opt.py
353 lines (257 loc) · 14.4 KB
/
quad_opt.py
1
2
3
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
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
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
import numpy as np
import casadi as cs
import os
from acados_template import AcadosOcp, AcadosOcpSolver, AcadosModel
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3
from matplotlib import animation
from quad import Quadrotor3D
from utils.utils import skew_symmetric, quaternion_to_euler, unit_quat, v_dot_q, quaternion_inverse
import pdb
class quad_optimizer:
def __init__(self, quad, t_horizon=1, n_nodes=100, gpe=None):
self.n_nodes = n_nodes
self.t_horizon = t_horizon
self.gpe = gpe
#self.optimization_dt = optimization_dt
self.optimization_dt = self.t_horizon/self.n_nodes
#self.t_horizon = self.n_nodes*self.optimization_dt # look-ahead time
self.quad = quad # quad is needed to create the casadi model using quad parameters
self.dynamics = self.setup_casadi_model()
#self.dynamics = self.f_dict['f_dyn']# + self.f_dict['f_augment']
self.x_dot = cs.MX.sym('x_dot', self.dynamics(x=self.x, u=self.u)['f'].shape) # x_dot has the same dimensions as the dynamics function output
#print(dynamics(x=x, u=u)['x_dot'])
#optimization_dt = 5e-2 # dt between predicted states
#n_nodes = 100 # number of predicted states
#self.discrete_dynamics =
self.terminal_cost = 1
self.acados_model = AcadosModel()
self.acados_model.name = 'quad_OCP'
self.acados_model.x = self.x
self.acados_model.u = self.u
#acados_model.x_dot = x_dot
#breakpoint()
self.acados_model.f_expl_expr = self.dynamics(x=self.x, u=self.u)['f'] # x_dot = f
self.acados_model.f_impl_expr = self.x_dot - self.dynamics(x=self.x, u=self.u)['f'] # 0 = f - x_dot
#print(acados_model)
self.acados_ocp = AcadosOcp()
self.acados_source_path = os.environ['ACADOS_SOURCE_DIR']
self.acados_ocp.acados_include_path = self.acados_source_path + '/include'
self.acados_ocp.acados_lib_path = self.acados_source_path + '/lib'
self.acados_ocp.model = self.acados_model
self.acados_ocp.dims.N = self.n_nodes # prediction horizon
self.acados_ocp.solver_options.tf = self.t_horizon # look ahead time
self.acados_ocp.cost.cost_type = 'LINEAR_LS' # weigths times states (as opposed to a nonlinear relationship)
self.acados_ocp.cost.cost_type_e = 'LINEAR_LS' # end state cost
self.nx = self.acados_model.x.size()[0]
self.nu = self.acados_model.u.size()[0]
self.ny = self.nx + self.nu # y is x and u concatenated for compactness of the loss function
## Optimization costs
self.acados_ocp.cost.Vx = np.zeros((self.ny, self.nx)) # raise the dim of x to the dim of y
self.acados_ocp.cost.Vx[:self.nx, :self.nx] = np.eye(self.nx) # weight only x
self.acados_ocp.cost.Vx_e = np.eye(self.nx) # end x cost
self.acados_ocp.cost.Vu = np.zeros((self.ny, self.nu)) # raise the dim of u to the dim of y
self.acados_ocp.cost.Vu[-self.nu:, -self.nu:] = np.eye(self.nu) # weight only u
# x cost (dim 12)
#Original
#q_cost = np.array([10, 10, 10, 0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05]) # euler angles
#q_cost = np.array([0, 1000, 1, 0, 0, 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) # euler angles
q_cost = np.array([10, 10, 10] + [0.1, 0.1, 0.1] + [0.05, 0.05, 0.05] + [0.05, 0.05, 0.05]) # ORIGINAL
#q_cost = np.array([100, 100, 100] + [0.0, 0.0, 0.0] + [0.0, 0.0, 0.0] + [0.05, 0.05, 0.05]) # MPC doesnt care about linear or angular velocity
# add one more weigth to account for the 4 quaternions instead of 3 EA
q_diagonal = np.concatenate((q_cost[:3], np.mean(q_cost[3:6])[np.newaxis], q_cost[3:]))
r_cost = np.array([0.1, 0.1, 0.1, 0.1]) # u cost (dim 4)
#r_cost = np.array([1000, 1000, 1000, 1000]) # u cost (dim 4)
self.acados_ocp.cost.W = np.diag(np.concatenate((q_diagonal, r_cost))) # error costs
self.acados_ocp.cost.W_e = np.diag(q_diagonal) * self.terminal_cost # end error cost
# reference trajectory (will be overwritten later, this is just for dimensions)
x_ref = np.zeros(self.nx)
self.acados_ocp.cost.yref = np.concatenate((x_ref, np.array([0.0, 0.0, 0.0, 0.0])))
#print(acados_ocp.cost.yref.shape)
self.acados_ocp.cost.yref_e = x_ref # end node reference
self.acados_ocp.constraints.x0 = x_ref
# u constraints
self.acados_ocp.constraints.lbu = np.array([0] * self.nu)
self.acados_ocp.constraints.ubu = np.array([1] * self.nu)
self.acados_ocp.constraints.idxbu = np.array([0,1,2,3]) # Indices of bounds on u
## Solver options
self.acados_ocp.solver_options.qp_solver = 'FULL_CONDENSING_HPIPM'
self.acados_ocp.solver_options.hessian_approx = 'GAUSS_NEWTON'
self.acados_ocp.solver_options.integrator_type = 'ERK'
self.acados_ocp.solver_options.print_level = 0
self.acados_ocp.solver_options.nlp_solver_type = 'SQP_RTI'
# Compile acados OCP solver if necessary
json_file = '_acados_ocp.json'
self.acados_ocp_solver = AcadosOcpSolver(self.acados_ocp, json_file=json_file)
print(f'Optimizer MPC lookahead time={self.t_horizon}s, Number of timesteps in prediction horizon={self.n_nodes}')
def setup_casadi_model(self):
"""
Creates the appropriate symbolic variables as class fields and returns the dynamic equation of the quad model
"""
self.p = cs.MX.sym('p',3) # Position
self.q = cs.MX.sym('q',4) # Quaternion
self.v = cs.MX.sym('v',3) # Velocity
self.r = cs.MX.sym('w',3) # Angle rate
self.x = cs.vertcat(self.p,self.q,self.v,self.r) # the complete state
# Control variables
u1 = cs.MX.sym('u1')
u2 = cs.MX.sym('u2')
u3 = cs.MX.sym('u3')
u4 = cs.MX.sym('u4')
self.u = cs.vertcat(u1,u2,u3,u4) # complete control
# d position
f_p = self.v # position dynamicss
# d quaternion
f_q = 1.0 / 2.0 * cs.mtimes(skew_symmetric(self.r),self.q) # quaternion dynamics
# d velocity
f_thrust = self.u * self.quad.max_thrust
a_thrust = cs.vertcat(0, 0, f_thrust[0] + f_thrust[1] + f_thrust[2] + f_thrust[3]) / self.quad.mass
self.g = cs.vertcat(0,0,9.81)
f_v = v_dot_q(a_thrust, self.q) - self.g # velocity dynamics
# d rate
y_f = cs.MX(self.quad.y_f)
x_f = cs.MX(self.quad.x_f)
c_f = cs.MX(self.quad.z_l_tau)
# rate dynamics
f_r = cs.vertcat(
(cs.mtimes(f_thrust.T, y_f) + (self.quad.J[1] - self.quad.J[2]) * self.r[1] * self.r[2]) / self.quad.J[0],
(-cs.mtimes(f_thrust.T, x_f) + (self.quad.J[2] - self.quad.J[0]) * self.r[2] * self.r[0]) / self.quad.J[1],
(cs.mtimes(f_thrust.T, c_f) + (self.quad.J[0] - self.quad.J[1]) * self.r[0] * self.r[1]) / self.quad.J[2])
# concatenated dynamics
f_dyn = cs.vertcat(f_p, f_q, f_v, f_r)
if self.gpe is not None:
# Transform to body frame because thats what the gpe were trained on
v_body = v_dot_q(self.v, quaternion_inverse(self.q))
#x_body = cs.vertcat(self.x[0:7], v_body, self.x[10:])
# Symbolic prediction
gp_means = self.gpe.predict(v_body.T).T
# Transform prediction back to world frame because thats what the simulator uses
gp_means = v_dot_q(gp_means, self.q)
# 13 x 3 matrix that maps the 3 gpe means to dvx, dvy, dvz
B_x = np.concatenate([np.zeros((3,3)), np.zeros((4,3)), np.diag([1,1,1]), np.zeros((3,3))], axis=0)
f_augment = cs.mtimes(B_x, gp_means)
# Dynamics correction using learned GP
f_corrected = f_dyn + f_augment
# Dynamics corrected using the gpe
return cs.Function('x_dot', [self.x,self.u], [f_corrected], ['x','u'], ['f'])
# Dynamics w/o the gpe augmentation
return cs.Function('x_dot', [self.x,self.u], [f_dyn], ['x','u'], ['f'])
# this does not do anything, I want to use the acados ocp solver to make predictions, not the internal quad object, that exists only to get the model parameters
def set_quad_state(self, x):
self.quad.set_state(x)
def set_reference_state(self, x_target=None, u_target=None):
"""
Sets the state x_target as the reference state over the whole optimization horizon (quad_opt.n_nodes)
:param: x_target: np.array of shape 1 x nx
"""
if u_target is None:
u_target = np.ones((self.nu,))*0.16 # hover
if x_target is None:
x_target = np.array([0,0,0, 1,0,0,0, 0,0,0, 0,0,0])
self.yref = np.empty((self.n_nodes, self.ny)) # prepare memory, N x ny
for j in range(self.n_nodes):
self.yref[j,:] = np.concatenate((x_target, u_target)) # load desired trajectory into yref
#print(self.yref[j,:])
self.acados_ocp_solver.set(j, "yref", self.yref[j,:]) # supply desired trajectory to ocp
# end point of the trajectory has no desired u
self.yref_N = x_target
self.acados_ocp_solver.set(self.n_nodes, "yref", self.yref_N) # dimension nx
#############!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
return self.yref, self.yref_N
def set_reference_trajectory(self, x_trajectory, u_trajectory=None):
"""
Passes x_trajectory to acados_ocp_solver. Acados_ocp_solver then tries to find the optimal control
:param: x_trajectory: x_trajectory has to have the same length as quad_opt.n_node
:param: u_trajectory: u_trajectory has to have the same length as quad_opt.n_node
:param: oversampling: if x_trajectory has smaller sampling time then quad_opt.optimization_time, then undersample the signal to match the sampling time
"""
#breakpoint()
if u_trajectory is None:
u_trajectory = np.ones((self.n_nodes, 4))*0.16 # hover
self.yref = np.empty((self.n_nodes, self.ny)) # prepare memory, N x ny
for j in range(self.n_nodes):
#print(j)
self.yref[j,:] = np.concatenate((x_trajectory[j,:], u_trajectory[j,:])) # load desired trajectory into yref
#print(yref[j,:])
self.acados_ocp_solver.set(j, "yref", self.yref[j,:]) # supply desired trajectory to ocp
# end point of the trajectory has no desired u
self.yref_N = x_trajectory[-1,:]
self.acados_ocp_solver.set(self.n_nodes, "yref", self.yref_N) # dimension nx
return self.yref, self.yref_N
def run_optimization(self, x_init=None):
"""
Performs the acados_ocp_solver.solve() method with the current parameters and returns the solution
"""
if x_init is None:
x_init = np.array([0,0,0, 1,0,0,0, 0,0,0, 0,0,0])
# set initial conditions
self.acados_ocp_solver.set(0, 'lbx', x_init) # not sure what these two lines do. Set the lower and upper bound for x to the same value?
self.acados_ocp_solver.set(0, 'ubx', x_init)
#self.acados_ocp_solver.set(0, 'x', x_init)
# Solve OCP
self.acados_ocp_solver.solve()
self.acados_ocp_solver.solve()
#self.acados_ocp_solver.solve()
#print(x_init)
# preallocate memory
w_opt_acados = np.ndarray((self.n_nodes, self.nu))
x_opt_acados = np.ndarray((self.n_nodes + 1, self.nx))
x_opt_acados[0, :] = self.acados_ocp_solver.get(0, "x") # should be x_init?
# write down the optimal solution
for i in range(self.n_nodes):
w_opt_acados[i, :] = self.acados_ocp_solver.get(i, "u") # optimal control
x_opt_acados[i + 1, :] = self.acados_ocp_solver.get(i + 1, "x") # state under optimal control
#w_opt_acados = np.reshape(w_opt_acados, (-1))
return x_opt_acados, w_opt_acados, self.acados_ocp_solver.get_stats('time_tot'), self.acados_ocp_solver.get_cost()
def discrete_dynamics(self, x, u, dt, body_frame=False):
#
"""
Fixed step Runge-Kutta 4 of the dynamic equation in self.dynamics
:param: x: State 1 x nx to discretize around
:param: u: Conrol 1 x nu to discretize around
:param: dt: Time step for dicretization
"""
k1 = self.dynamics(x=x, u=u)['f']
k2 = self.dynamics(x=x + dt / 2 * k1, u=u)['f']
k3 = self.dynamics(x=x + dt / 2 * k2, u=u)['f']
k4 = self.dynamics(x=x + dt * k3, u=u)['f']
x_out = x + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
#print(x_out)
if not body_frame:
return x_out
else:
# velocity is transformed to bodyframe
v_b = v_dot_q(x_out[7:10], quaternion_inverse(x_out[3:7]))
#print(x_out)
x_out = np.array([x_out[0], x_out[1], x_out[2], x_out[3], x_out[4], x_out[5], x_out[6],
v_b[0], v_b[1], v_b[2], x_out[10], x_out[11], x_out[12]])
#print(x_out)
return x_out
'''
def linearized_quad_dynamics():
"""
Jacobian J matrix of the linearized dynamics specified in the function quad_dynamics. J[i, j] corresponds to
the partial derivative of f_i(x) wrt x(j).
:return: a CasADi symbolic function that calculates the 13 x 13 Jacobian matrix of the linearized simplified
quadrotor dynamics
"""
jac = cs.MX(state_dim, state_dim)
# Position derivatives
jac[0:3, 7:10] = cs.diag(cs.MX.ones(3))
# Angle derivatives
jac[3:7, 3:7] = skew_symmetric(r) / 2
jac[3, 10:] = 1 / 2 * cs.horzcat(-q[1], -q[2], -q[3])
jac[4, 10:] = 1 / 2 * cs.horzcat(q[0], -q[3], q[2])
jac[5, 10:] = 1 / 2 * cs.horzcat(q[3], q[0], -q[1])
jac[6, 10:] = 1 / 2 * cs.horzcat(-q[2], q[1], q[0])
# Velocity derivatives
a_u = (u[0] + u[1] + u[2] + u[3]) * quad.max_thrust / quad.mass
jac[7, 3:7] = 2 * cs.horzcat(a_u * q[2], a_u * q[3], a_u * q[0], a_u * q[1])
jac[8, 3:7] = 2 * cs.horzcat(-a_u * q[1], -a_u * q[0], a_u * q[3], a_u * q[2])
jac[9, 3:7] = 2 * cs.horzcat(0, -2 * a_u * q[1], -2 * a_u * q[1], 0)
# Rate derivatives
jac[10, 10:] = (quad.J[1] - quad.J[2]) / quad.J[0] * cs.horzcat(0, r[2], r[1])
jac[11, 10:] = (quad.J[2] - quad.J[0]) / quad.J[1] * cs.horzcat(r[2], 0, r[0])
jac[12, 10:] = (quad.J[0] - quad.J[1]) / quad.J[2] * cs.horzcat(r[1], r[0], 0)
return cs.Function('J', [x, u], [jac])
'''