| """Fixed-wing 6-DoF flight dynamics, batched torch: the same equations serve the plant and the oracle world model. |
| |
| Internally the textbook frames of Beard & McLain (ch. 3-4): world NED, body FRD, state |
| s = [p_ned(3), v_body uvw(3), q wxyz body->NED(4), w_frd pqr(3)] (B, 13) |
| and the action a = [throttle 0..1, aileron, elevator, rudder (rad)] (B, 4). |
| At the boundary everything is converted to dronepilot's shared 18-dim layout (world ENU z up, body FLU): |
| x = [p, v world, R columns (body -> world), w body] (B, 18) |
| so a plane is read by the same Pose / tasks / aura as every quad. |
| |
| Forces: lift / drag with a flat-plate blend past stall (alpha0, M), side force, gravity, a simple electric |
| propeller (Beard & McLain 1st ed.). Moments: roll, pitch, yaw from the stability and control derivatives. |
| Integration: RK4 at a fixed substep, quaternion renormalized every substep. |
| """ |
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass, field |
|
|
| import torch |
|
|
| T_NE = torch.tensor([[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, -1.0]]) |
| F_FL = torch.diag(torch.tensor([1.0, -1.0, -1.0])) |
|
|
|
|
| @dataclass |
| class Airframe: |
| mass: float |
| J: torch.Tensor |
| S: float |
| b: float |
| c: float |
| rho: float |
| g: float |
| aero: dict |
| prop: dict |
| surface_max: float |
| extra: dict = field(default_factory=dict) |
|
|
| @classmethod |
| def from_cfg(cls, cfg: dict, dtype=torch.float32) -> "Airframe": |
| I = cfg["inertia"] |
| J = torch.tensor([[I["Jx"], 0.0, -I["Jxz"]], [0.0, I["Jy"], 0.0], [-I["Jxz"], 0.0, I["Jz"]]], dtype=dtype) |
| g = cfg["geometry"] |
| return cls(float(cfg["mass"]), J, g["S"], g["b"], g["c"], cfg["rho"], cfg["gravity"], dict(cfg["aero"]), |
| dict(cfg["propulsion"]), cfg["limits"]["surface_rad"]) |
|
|
| @property |
| def AR(self) -> float: |
| return self.b**2 / self.S |
|
|
|
|
| |
| def quat_to_R(q: torch.Tensor) -> torch.Tensor: |
| w, x, y, z = (q / q.norm(dim=-1, keepdim=True)).unbind(-1) |
| return torch.stack([ |
| torch.stack([1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)], -1), |
| torch.stack([2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)], -1), |
| torch.stack([2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], -1)], -2) |
|
|
|
|
| def R_to_quat(R: torch.Tensor) -> torch.Tensor: |
| """(..., 3, 3) -> (..., 4) wxyz, robust branch on the largest diagonal term.""" |
| m = R |
| tr = m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2] |
| cands = torch.stack([tr, m[..., 0, 0], m[..., 1, 1], m[..., 2, 2]], -1) |
| k = cands.argmax(-1) |
| out = torch.empty(R.shape[:-2] + (4,), dtype=R.dtype) |
| s0 = torch.sqrt((1 + tr).clamp(min=1e-12)) * 2 |
| q0 = torch.stack([0.25 * s0, (m[..., 2, 1] - m[..., 1, 2]) / s0, (m[..., 0, 2] - m[..., 2, 0]) / s0, |
| (m[..., 1, 0] - m[..., 0, 1]) / s0], -1) |
| s1 = torch.sqrt((1 + m[..., 0, 0] - m[..., 1, 1] - m[..., 2, 2]).clamp(min=1e-12)) * 2 |
| q1 = torch.stack([(m[..., 2, 1] - m[..., 1, 2]) / s1, 0.25 * s1, (m[..., 0, 1] + m[..., 1, 0]) / s1, |
| (m[..., 0, 2] + m[..., 2, 0]) / s1], -1) |
| s2 = torch.sqrt((1 + m[..., 1, 1] - m[..., 0, 0] - m[..., 2, 2]).clamp(min=1e-12)) * 2 |
| q2 = torch.stack([(m[..., 0, 2] - m[..., 2, 0]) / s2, (m[..., 0, 1] + m[..., 1, 0]) / s2, 0.25 * s2, |
| (m[..., 1, 2] + m[..., 2, 1]) / s2], -1) |
| s3 = torch.sqrt((1 + m[..., 2, 2] - m[..., 0, 0] - m[..., 1, 1]).clamp(min=1e-12)) * 2 |
| q3 = torch.stack([(m[..., 1, 0] - m[..., 0, 1]) / s3, (m[..., 0, 2] + m[..., 2, 0]) / s3, |
| (m[..., 1, 2] + m[..., 2, 1]) / s3, 0.25 * s3], -1) |
| for i, qi in enumerate((q0, q1, q2, q3)): |
| sel = k == i |
| out[sel] = qi[sel] |
| out = out * torch.where(out[..., :1] < 0, -1.0, 1.0) |
| return out / out.norm(dim=-1, keepdim=True) |
|
|
|
|
| |
| def to_shared(s: torch.Tensor) -> torch.Tensor: |
| """Internal (B, 13) NED/FRD -> shared (B, 18) ENU/FLU.""" |
| T, F = T_NE.to(s.dtype), F_FL.to(s.dtype) |
| R_nf = quat_to_R(s[..., 6:10]) |
| R_ef = T @ R_nf @ F |
| p = s[..., 0:3] @ T.T |
| v = (R_nf @ s[..., 3:6, None])[..., 0] @ T.T |
| w = s[..., 10:13] @ F.T |
| return torch.cat([p, v, R_ef[..., :, 0], R_ef[..., :, 1], R_ef[..., :, 2], w], -1) |
|
|
|
|
| def from_shared(x: torch.Tensor) -> torch.Tensor: |
| """Shared (B, 18) ENU/FLU -> internal (B, 13) NED/FRD.""" |
| T, F = T_NE.to(x.dtype), F_FL.to(x.dtype) |
| R_ef = torch.stack([x[..., 6:9], x[..., 9:12], x[..., 12:15]], -1) |
| R_nf = T @ R_ef @ F |
| p = x[..., 0:3] @ T.T |
| v_ned = x[..., 3:6] @ T.T |
| v_body = (R_nf.transpose(-1, -2) @ v_ned[..., None])[..., 0] |
| w = x[..., 15:18] @ F.T |
| return torch.cat([p, v_body, R_to_quat(R_nf), w], -1) |
|
|
|
|
| |
| def air_data(s: torch.Tensor, wind_ned: torch.Tensor | None): |
| R = quat_to_R(s[..., 6:10]) |
| v_r = s[..., 3:6] |
| if wind_ned is not None: |
| v_r = v_r - (R.transpose(-1, -2) @ wind_ned[..., None])[..., 0] |
| u, v, w = v_r.unbind(-1) |
| Va = v_r.norm(dim=-1).clamp(min=0.5) |
| alpha = torch.atan2(w, u) |
| beta = torch.asin((v / Va).clamp(-1.0, 1.0)) |
| return R, Va, alpha, beta |
|
|
|
|
| def forces_moments(af: Airframe, s: torch.Tensor, a: torch.Tensor, wind_ned=None, mass: float | None = None): |
| """Body FRD force (B, 3) and moment (B, 3). mass: the true mass if a payload changed it (gravity).""" |
| A = af.aero |
| R, Va, alpha, beta = air_data(s, wind_ned) |
| p, q, r = s[..., 10:13].unbind(-1) |
| dt, da, de, dr = a.unbind(-1) |
| qbar = 0.5 * af.rho * Va**2 * af.S |
| |
| M, a0 = A["M"], A["alpha0"] |
| e1, e2 = torch.exp(-M * (alpha - a0)), torch.exp(M * (alpha + a0)) |
| sigma = (1 + e1 + e2) / ((1 + e1) * (1 + e2)) |
| CL_lin = A["C_L_0"] + A["C_L_alpha"] * alpha |
| CL = (1 - sigma) * CL_lin + sigma * 2 * torch.sign(alpha) * torch.sin(alpha) ** 2 * torch.cos(alpha) |
| CD = A["C_D_0"] + CL_lin**2 / (math.pi * A["e"] * af.AR) |
| c2V, b2V = af.c / (2 * Va), af.b / (2 * Va) |
| F_lift = qbar * (CL + A["C_L_q"] * c2V * q + A["C_L_delta_e"] * de) |
| F_drag = qbar * (CD + A["C_D_q"] * c2V * q + A["C_D_delta_e"] * de.abs()) |
| ca, sa = torch.cos(alpha), torch.sin(alpha) |
| fx = -ca * F_drag + sa * F_lift |
| fz = -sa * F_drag - ca * F_lift |
| fy = qbar * (A["C_Y_0"] + A["C_Y_beta"] * beta + A["C_Y_p"] * b2V * p + A["C_Y_r"] * b2V * r |
| + A["C_Y_delta_a"] * da + A["C_Y_delta_r"] * dr) |
| |
| P = af.prop |
| thrust = (0.5 * af.rho * P["S_prop"] * P["C_prop"] * ((P["k_motor"] * dt) ** 2 - Va**2)).clamp(min=0.0) |
| fx = fx + thrust |
| |
| m = af.mass if mass is None else mass |
| g_body = (R.transpose(-1, -2) @ torch.tensor([0.0, 0.0, m * af.g], dtype=s.dtype)[:, None])[..., 0] |
| F = torch.stack([fx, fy, fz], -1) + g_body |
| l = qbar * af.b * (A["C_ell_0"] + A["C_ell_beta"] * beta + A["C_ell_p"] * b2V * p + A["C_ell_r"] * b2V * r |
| + A["C_ell_delta_a"] * da + A["C_ell_delta_r"] * dr) |
| pitch = qbar * af.c * (A["C_m_0"] + A["C_m_alpha"] * alpha + A["C_m_q"] * c2V * q + A["C_m_delta_e"] * de) |
| n = qbar * af.b * (A["C_n_0"] + A["C_n_beta"] * beta + A["C_n_p"] * b2V * p + A["C_n_r"] * b2V * r |
| + A["C_n_delta_a"] * da + A["C_n_delta_r"] * dr) |
| return F, torch.stack([l, pitch, n], -1) |
|
|
|
|
| def derivative(af: Airframe, s: torch.Tensor, a: torch.Tensor, wind_ned=None, |
| mass: float | None = None, res: torch.Tensor | None = None, |
| payload: tuple | None = None) -> torch.Tensor: |
| """res (B, 6): a learned residual [linear acc (body FRD, m/s^2), angular acc (body FRD, rad/s^2)] added to the |
| physics (fixed_wing/learned.py). None = pure physics. |
| payload (m_p, r_frd): a point mass m_p (kg, already counted in `mass`) at r_frd (m, body FRD, from the CG): its |
| weight's moment about the CG. Its inertia is the caller's (Airframe.J with the parallel-axis term).""" |
| m = af.mass if mass is None else mass |
| F, Mo = forces_moments(af, s, a, wind_ned, mass) |
| q = s[..., 6:10] |
| v = s[..., 3:6] |
| w = s[..., 10:13] |
| R = quat_to_R(q) |
| p_dot = (R @ v[..., None])[..., 0] |
| v_dot = -torch.cross(w, v, dim=-1) + F / m |
| qw, qx, qy, qz = q.unbind(-1) |
| p, qq, r = w.unbind(-1) |
| q_dot = 0.5 * torch.stack([-qx * p - qy * qq - qz * r, qw * p + qy * r - qz * qq, |
| qw * qq + qz * p - qx * r, qw * r + qx * qq - qy * p], -1) |
| J = af.J.to(s.dtype) |
| Jw = w @ J.T |
| w_dot = torch.linalg.solve(J, (Mo - torch.cross(w, Jw, dim=-1)).unsqueeze(-1)).squeeze(-1) |
| if payload is not None: |
| m_p, r = payload |
| f_p = (R.transpose(-1, -2) @ torch.tensor([0.0, 0.0, m_p * af.g], dtype=s.dtype)[:, None])[..., 0] |
| tau = torch.cross(torch.as_tensor(r, dtype=s.dtype).expand_as(f_p), f_p, dim=-1) |
| w_dot = w_dot + torch.linalg.solve(J, tau.unsqueeze(-1)).squeeze(-1) |
| if res is not None: |
| v_dot = v_dot + res[..., 0:3] |
| w_dot = w_dot + res[..., 3:6] |
| return torch.cat([p_dot, v_dot, q_dot, w_dot], -1) |
|
|
|
|
| def rk4(af: Airframe, s: torch.Tensor, a: torch.Tensor, h: float, wind_ned=None, mass=None, res=None, |
| payload=None) -> torch.Tensor: |
| k1 = derivative(af, s, a, wind_ned, mass, res, payload) |
| k2 = derivative(af, s + 0.5 * h * k1, a, wind_ned, mass, res, payload) |
| k3 = derivative(af, s + 0.5 * h * k2, a, wind_ned, mass, res, payload) |
| k4 = derivative(af, s + h * k3, a, wind_ned, mass, res, payload) |
| out = s + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4) |
| q = out[..., 6:10] |
| return torch.cat([out[..., :6], q / q.norm(dim=-1, keepdim=True), out[..., 10:]], -1) |
|
|
