210 lines
6.6 KiB
Python
210 lines
6.6 KiB
Python
"""
|
|
Physics-embedded CW dynamics for spacecraft terminal rendezvous.
|
|
|
|
Implements:
|
|
- Exact ZOH discretization of CW equations (Phi, Gamma closed-form)
|
|
- Continuous-time CW dynamics (A_cw, B)
|
|
- CW drift acceleration a_cw(x)
|
|
- Discrete LQR nominal controller via DARE
|
|
- Kalman filter for state estimation
|
|
- Domain randomization utilities
|
|
|
|
Reference: Section 2 of the paper.
|
|
"""
|
|
|
|
import numpy as np
|
|
from scipy.linalg import solve_discrete_are
|
|
|
|
|
|
class CWDynamics:
|
|
"""
|
|
Clohessy-Wiltshire relative dynamics with exact zero-order-hold discretization.
|
|
|
|
State: x = [r_x, r_y, r_z, v_x, v_y, v_z]^T
|
|
Coordinate: x radial-out, y along-track, z normal
|
|
|
|
Parameters
|
|
----------
|
|
n : float
|
|
Mean orbital angular velocity (rad/s).
|
|
dt : float
|
|
Sampling period (s).
|
|
d_bar : float
|
|
Disturbance acceleration bound (m/s^2).
|
|
"""
|
|
|
|
def __init__(self, n=None, dt=1.0, d_bar=0.0,
|
|
mu=3.986004418e14, R0=42164000.0):
|
|
if n is None:
|
|
n = np.sqrt(mu / R0 ** 3)
|
|
self.n = n
|
|
self.dt = dt
|
|
self.d_bar = d_bar
|
|
self.mu = mu
|
|
self.R0 = R0
|
|
|
|
self.A_cw = self._build_A_cw(n)
|
|
self.B_ct = self._build_B()
|
|
self.Phi, self.Gamma = self._build_exact_discrete(n, dt)
|
|
|
|
@staticmethod
|
|
def _build_A_cw(n):
|
|
"""Continuous-time A matrix of CW equations."""
|
|
A = np.zeros((6, 6))
|
|
A[0:3, 3:6] = np.eye(3)
|
|
A[3, 0] = 3.0 * n ** 2
|
|
A[5, 2] = -n ** 2
|
|
A[3, 4] = 2.0 * n
|
|
A[4, 3] = -2.0 * n
|
|
return A
|
|
|
|
@staticmethod
|
|
def _build_B():
|
|
B = np.zeros((6, 3))
|
|
B[3:6, :] = np.eye(3)
|
|
return B
|
|
|
|
@staticmethod
|
|
def _build_exact_discrete(n, dt):
|
|
"""
|
|
Exact ZOH discretization — closed-form Phi and Gamma.
|
|
See Section 2.2 of the paper.
|
|
"""
|
|
nu = n * dt
|
|
c = np.cos(nu)
|
|
s = np.sin(nu)
|
|
|
|
Phi = np.array([
|
|
[4.0 - 3.0 * c, 0.0, 0.0,
|
|
s / n, 2.0 * (1.0 - c) / n, 0.0],
|
|
[6.0 * (s - nu), 1.0, 0.0,
|
|
-2.0 * (1.0 - c) / n, (4.0 * s - 3.0 * nu) / n, 0.0],
|
|
[0.0, 0.0, c,
|
|
0.0, 0.0, s / n],
|
|
[3.0 * n * s, 0.0, 0.0,
|
|
c, 2.0 * s, 0.0],
|
|
[-6.0 * n * (1.0 - c), 0.0, 0.0,
|
|
-2.0 * s, 4.0 * c - 3.0, 0.0],
|
|
[0.0, 0.0, -n * s,
|
|
0.0, 0.0, c],
|
|
])
|
|
|
|
Gamma = np.array([
|
|
[(1.0 - c) / n ** 2,
|
|
2.0 * (nu - s) / n ** 2, 0.0],
|
|
[-2.0 * (nu - s) / n ** 2,
|
|
(4.0 * (1.0 - c) - 1.5 * nu ** 2) / n ** 2, 0.0],
|
|
[0.0, 0.0, (1.0 - c) / n ** 2],
|
|
[s / n, 2.0 * (1.0 - c) / n, 0.0],
|
|
[-2.0 * (1.0 - c) / n, (4.0 * s - 3.0 * nu) / n, 0.0],
|
|
[0.0, 0.0, s / n],
|
|
])
|
|
|
|
return Phi, Gamma
|
|
|
|
def acw(self, x):
|
|
"""CW drift acceleration: a_cw(x) from Section 3.4."""
|
|
n = self.n
|
|
return np.array([
|
|
3.0 * n ** 2 * x[0] + 2.0 * n * x[4],
|
|
-2.0 * n * x[3],
|
|
-n ** 2 * x[2],
|
|
])
|
|
|
|
def step(self, x, u, disturbance=None):
|
|
"""
|
|
Exact ZOH discrete step: x_{k+1} = Phi x_k + Gamma u_k + w_k
|
|
"""
|
|
x_next = self.Phi @ x + self.Gamma @ u
|
|
if disturbance is not None:
|
|
x_next += self.Gamma @ disturbance
|
|
return x_next
|
|
|
|
def rollout(self, x0, u_sequence):
|
|
"""Nominal multi-step rollout (no disturbance)."""
|
|
L = len(u_sequence)
|
|
traj = np.zeros((L + 1, 6))
|
|
traj[0] = x0
|
|
for k in range(L):
|
|
traj[k + 1] = self.Phi @ traj[k] + self.Gamma @ u_sequence[k]
|
|
return traj
|
|
|
|
def sensitivity(self, ell):
|
|
"""Proposition 1: d x_{k+ell} / d u_k = Phi^{ell-1} Gamma"""
|
|
return np.linalg.matrix_power(self.Phi, ell - 1) @ self.Gamma
|
|
|
|
def sample_disturbance(self):
|
|
"""Sample a bounded disturbance for domain randomization."""
|
|
d = np.random.uniform(-1, 1, size=3)
|
|
d = d / (np.linalg.norm(d) + 1e-8) * self.d_bar * np.random.uniform()
|
|
return d
|
|
|
|
def update_params(self, n=None, dt=None, d_bar=None):
|
|
if n is not None:
|
|
self.n = n
|
|
self.A_cw = self._build_A_cw(n)
|
|
if dt is not None:
|
|
self.dt = dt
|
|
if d_bar is not None:
|
|
self.d_bar = d_bar
|
|
self.Phi, self.Gamma = self._build_exact_discrete(self.n, self.dt)
|
|
|
|
|
|
class NominalLQR:
|
|
"""
|
|
Discrete LQR nominal controller: u_nom = -K (x_hat - x_h)
|
|
Solves DARE for (Phi, Gamma, Q_lqr, R_lqr). See Section 3.1.
|
|
"""
|
|
|
|
def __init__(self, dynamics: CWDynamics, Q_lqr=None, R_lqr=None):
|
|
self.dynamics = dynamics
|
|
self.Q_lqr = Q_lqr if Q_lqr is not None else np.diag(
|
|
[1.0, 1.0, 1.0, 0.1, 0.1, 0.1])
|
|
self.R_lqr = R_lqr if R_lqr is not None else np.eye(3) * 10.0
|
|
self.K = self._solve(dynamics.Phi, dynamics.Gamma,
|
|
self.Q_lqr, self.R_lqr)
|
|
|
|
@staticmethod
|
|
def _solve(Phi, Gamma, Q, R):
|
|
P = solve_discrete_are(Phi, Gamma, Q, R)
|
|
return np.linalg.solve(R + Gamma.T @ P @ Gamma, Gamma.T @ P @ Phi)
|
|
|
|
def compute(self, x_hat, x_h):
|
|
return -self.K @ (x_hat - x_h)
|
|
|
|
def update(self, dynamics: CWDynamics):
|
|
self.dynamics = dynamics
|
|
self.K = self._solve(dynamics.Phi, dynamics.Gamma,
|
|
self.Q_lqr, self.R_lqr)
|
|
|
|
|
|
class KalmanFilter:
|
|
"""
|
|
Discrete Kalman filter: o_k = H x_k + nu_k. See Section 2.3.
|
|
"""
|
|
|
|
def __init__(self, dynamics: CWDynamics, H=None, Sigma_o=None, Q_proc=None):
|
|
self.dynamics = dynamics
|
|
self.H = H if H is not None else np.eye(6)
|
|
self.Sigma_o = Sigma_o if Sigma_o is not None else np.eye(6) * 1.0
|
|
self.Q_proc = Q_proc if Q_proc is not None else np.eye(6) * 0.01
|
|
self.x_hat = np.zeros(6)
|
|
self.P = np.eye(6) * 100.0
|
|
|
|
def predict(self, u):
|
|
Phi = self.dynamics.Phi
|
|
Gamma = self.dynamics.Gamma
|
|
self.x_hat = Phi @ self.x_hat + Gamma @ u
|
|
self.P = Phi @ self.P @ Phi.T + self.Q_proc
|
|
|
|
def update(self, o):
|
|
H = self.H
|
|
S = H @ self.P @ H.T + self.Sigma_o
|
|
K_kal = self.P @ H.T @ np.linalg.inv(S)
|
|
self.x_hat = self.x_hat + K_kal @ (o - H @ self.x_hat)
|
|
self.P = (np.eye(6) - K_kal @ H) @ self.P
|
|
|
|
def reset(self, x0, P0=None):
|
|
self.x_hat = x0.copy()
|
|
self.P = P0 if P0 is not None else np.eye(6) * 100.0
|