Skip to content

Commit 7a02483

Browse files
ishaanamahajanishaanamahajan
committed
added actual matrices and autograd compatibility
1 parent 8739f5e commit 7a02483

File tree

2 files changed

+108
-26
lines changed

2 files changed

+108
-26
lines changed

‎examples/quadrotor_hover_code_generation.py‎

Lines changed: 69 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,50 @@
11
import tinympc
22
import numpy as np
3+
from autograd import jacobian
4+
import autograd.numpy as anp
35

46
# Toggle switch for adaptive rho
57
ENABLE_ADAPTIVE_RHO = True # Set to True to enable adaptive rho
68

79
# Quadrotor system matrices (12 states, 4 inputs)
8-
A = np.eye(12) # Identity matrix for simplicity
9-
B = np.zeros((12, 4))
10-
# Fill in control effectiveness
11-
B[0:3, 0:4] = 0.01 * np.ones((3, 4)) # Position control
12-
B[3:6, 0:4] = 0.05 * np.ones((3, 4)) # Velocity control
13-
B[6:9, 0:4] = 0.02 * np.ones((3, 4)) # Attitude control
14-
B[9:12, 0:4] = 0.1 * np.ones((3, 4)) # Angular velocity control
15-
16-
# Cost matrices
17-
Q = np.diag([10.0, 10.0, 10.0, 1.0, 1.0, 1.0, 5.0, 5.0, 5.0, 0.1, 0.1, 0.1])
18-
R = np.diag([0.1, 0.1, 0.1, 0.1])
10+
rho_value = 5.0
11+
Adyn = np.array([
12+
1.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0245250, 0.0000000, 0.0500000, 0.0000000, 0.0000000, 0.0000000, 0.0002044, 0.0000000,
13+
0.0000000, 1.0000000, 0.0000000, -0.0245250, 0.0000000, 0.0000000, 0.0000000, 0.0500000, 0.0000000, -0.0002044, 0.0000000, 0.0000000,
14+
0.0000000, 0.0000000, 1.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0500000, 0.0000000, 0.0000000, 0.0000000,
15+
0.0000000, 0.0000000, 0.0000000, 1.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0250000, 0.0000000, 0.0000000,
16+
0.0000000, 0.0000000, 0.0000000, 0.0000000, 1.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0250000, 0.0000000,
17+
0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 1.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0250000,
18+
0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.9810000, 0.0000000, 1.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0122625, 0.0000000,
19+
0.0000000, 0.0000000, 0.0000000, -0.9810000, 0.0000000, 0.0000000, 0.0000000, 1.0000000, 0.0000000, -0.0122625, 0.0000000, 0.0000000,
20+
0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 1.0000000, 0.0000000, 0.0000000, 0.0000000,
21+
0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 1.0000000, 0.0000000, 0.0000000,
22+
0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 1.0000000, 0.0000000,
23+
0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 1.0000000
24+
]).reshape(12, 12)
25+
26+
# Input/control matrix
27+
Bdyn = np.array([
28+
-0.0007069, 0.0007773, 0.0007091, -0.0007795,
29+
0.0007034, 0.0007747, -0.0007042, -0.0007739,
30+
0.0052554, 0.0052554, 0.0052554, 0.0052554,
31+
-0.1720966, -0.1895213, 0.1722891, 0.1893288,
32+
-0.1729419, 0.1901740, 0.1734809, -0.1907131,
33+
0.0123423, -0.0045148, -0.0174024, 0.0095748,
34+
-0.0565520, 0.0621869, 0.0567283, -0.0623632,
35+
0.0562756, 0.0619735, -0.0563386, -0.0619105,
36+
0.2102143, 0.2102143, 0.2102143, 0.2102143,
37+
-13.7677303, -15.1617018, 13.7831318, 15.1463003,
38+
-13.8353509, 15.2139209, 13.8784751, -15.2570451,
39+
0.9873856, -0.3611820, -1.3921880, 0.7659845
40+
]).reshape(12, 4)
41+
42+
43+
Q_diag = np.array([100.0000000, 100.0000000, 100.0000000, 4.0000000, 4.0000000, 400.0000000,
44+
4.0000000, 4.0000000, 4.0000000, 2.0408163, 2.0408163, 4.0000000])
45+
R_diag = np.array([4.0000000, 4.0000000, 4.0000000, 4.0000000])
46+
Q = np.diag(Q_diag)
47+
R = np.diag(R_diag)
1948

2049
N = 20
2150

@@ -25,7 +54,7 @@
2554
u_max = np.ones(4) * 2.0
2655

2756
# Setup with adaptive rho based on toggle
28-
prob.setup(A, B, Q, R, N, rho=1.0, max_iter=100, u_min=u_min, u_max=u_max,
57+
prob.setup(Adyn, Bdyn, Q, R, N, rho=rho_value, max_iter=100, u_min=u_min, u_max=u_max,
2958
adaptive_rho=1 if ENABLE_ADAPTIVE_RHO else 0)
3059

3160
if ENABLE_ADAPTIVE_RHO:
@@ -34,20 +63,34 @@
3463
# First compute the cache terms (this will compute K, P, etc.)
3564
Kinf, Pinf, Quu_inv, AmBKt = prob.compute_cache_terms()
3665

37-
# Compute sensitivity matrices using autograd
38-
# For now, we'll use small perturbations to approximate derivatives
39-
eps = 1e-4
40-
rho = 1.0
41-
42-
# Compute perturbed cache terms
43-
prob.setup(A, B, Q, R, N, rho=rho + eps, max_iter=100, u_min=u_min, u_max=u_max, adaptive_rho=1)
44-
Kinf_p, Pinf_p, Quu_inv_p, AmBKt_p = prob.compute_cache_terms()
45-
46-
# Compute derivatives with respect to rho using finite differences
47-
dK = (Kinf_p - Kinf) / eps
48-
dP = (Pinf_p - Pinf) / eps
49-
dC1 = (Quu_inv_p - Quu_inv) / eps # dC1 is derivative of Quu_inv
50-
dC2 = (AmBKt_p - AmBKt) / eps # dC2 is derivative of AmBKt
66+
# Compute derivatives with respect to rho via Autograd's Jacobian
67+
def lqr_direct(rho):
68+
R_rho = anp.array(R) + rho * anp.eye(4)
69+
Q_rho = anp.array(Q) + rho * anp.eye(12)
70+
P = Q_rho
71+
for _ in range(5000):
72+
K = anp.linalg.solve(
73+
R_rho + Bdyn.T @ P @ Bdyn + 1e-8*anp.eye(4),
74+
Bdyn.T @ P @ Adyn
75+
)
76+
P = Q_rho + Adyn.T @ P @ (Adyn - Bdyn @ K)
77+
# Final gain and cache matrices
78+
K = anp.linalg.solve(
79+
R_rho + Bdyn.T @ P @ Bdyn + 1e-8*anp.eye(4),
80+
Bdyn.T @ P @ Adyn
81+
)
82+
C1 = anp.linalg.inv(R_rho + Bdyn.T @ P @ Bdyn)
83+
C2 = (Adyn - Bdyn @ K).T
84+
return anp.concatenate([K.flatten(), P.flatten(), C1.flatten(), C2.flatten()])
85+
86+
derivs = jacobian(lqr_direct)(rho_value)
87+
# split into four blocks and reshape
88+
sizes = [4*12, 12*12, 4*4, 12*12]
89+
parts = np.split(np.array(derivs), np.cumsum(sizes)[:-1])
90+
dK = parts[0].reshape(4, 12)
91+
dP = parts[1].reshape(12, 12)
92+
dC1 = parts[2].reshape(4, 4)
93+
dC2 = parts[3].reshape(12, 12)
5194

5295
# Generate code with sensitivity matrices
5396
prob.codegen_with_sensitivity("out", dK, dP, dC1, dC2, verbose=1)

‎src/tinympc/interface.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,3 +363,42 @@ def compute_cache_terms(self):
363363
print(f"C1={np.linalg.norm(Quu_inv):.6f}, C2={np.linalg.norm(AmBKt):.6f}")
364364

365365
return Kinf, Pinf, Quu_inv, AmBKt
366+
367+
def compute_sensitivity_autograd(self):
368+
"""Compute dK, dP, dC1, dC2 with respect to rho using Autograd's jacobian."""
369+
# Local imports to avoid hard dependency unless this method is called
370+
from autograd import jacobian
371+
import autograd.numpy as anp
372+
373+
# Define the vectorized LQR solution mapping rho -> [K,P,C1,C2].flatten()
374+
def lqr_flat(rho):
375+
R_rho = anp.array(self.R) + rho * anp.eye(self.nu)
376+
Q_rho = anp.array(self.Q) + rho * anp.eye(self.nx)
377+
P = Q_rho
378+
for _ in range(5000):
379+
K = anp.linalg.solve(
380+
R_rho + self.B.T @ P @ self.B + 1e-8 * anp.eye(self.nu),
381+
self.B.T @ P @ self.A
382+
)
383+
P = Q_rho + self.A.T @ P @ (self.A - self.B @ K)
384+
K = anp.linalg.solve(
385+
R_rho + self.B.T @ P @ self.B + 1e-8 * anp.eye(self.nu),
386+
self.B.T @ P @ self.A
387+
)
388+
C1 = anp.linalg.inv(R_rho + self.B.T @ P @ self.B)
389+
C2 = (self.A - self.B @ K).T
390+
return anp.concatenate([K.flatten(), P.flatten(), C1.flatten(), C2.flatten()])
391+
392+
# Compute the Jacobian w.r.t. rho
393+
jac = jacobian(lqr_flat)
394+
vec = jac(self.rho)
395+
396+
# Split derivative vector into four blocks
397+
m, n = self.nu, self.nx
398+
sizes = [m * n, n * n, m * m, n * n]
399+
parts = np.split(np.array(vec), np.cumsum(sizes)[:-1])
400+
dK = parts[0].reshape(m, n)
401+
dP = parts[1].reshape(n, n)
402+
dC1 = parts[2].reshape(m, m)
403+
dC2 = parts[3].reshape(n, n)
404+
return dK, dP, dC1, dC2

0 commit comments

Comments
 (0)