Returning booster, part 3: atmospheric phase, grid fins, and multi-phase guidance¶
In Parts 1 and 2 we landed a booster from a tame initial condition: ~700 m up, ~400 m downrange, and 100 m/s coming in. That state is already the end of the aero phase — the engineering work upstream of where Part 1 picked up has been hidden. Here we add the three physical effects every returning first stage must contend with: a real atmosphere (density and speed of sound vs altitude), aerodynamic forces and moments that depend on Mach number and angle of attack, and grid fins as a second actuator that does most of the steering when the engines are off. We then split the trajectory into three phases — boost-back, aerodynamic glide, and landing burn — and stitch them together with a state-conditional guidance dispatcher.
By the end of this notebook you will be able to:
(a) wrap an exponential-atmosphere + Mach-dependent drag + angle-of-attack
normal force model around the Part-1 dynamics with LookupTable1d;
(b) add four grid fins as a second control axis with deflection-vs-dynamic-
pressure authority;
(c) compose a three-phase guidance dispatcher with state-conditional phase
guards;
(d) reuse Part 2's DirectShootingNMPC machinery for the landing burn,
because open-loop trajopt schedule playback across a multi-phase mission is
brittle in ways that closed-loop MPC is not.
Honest scoping note. As of this revision the notebook executes end-to-end and demonstrates the architecture (state-conditional dispatch, closed-loop MPC for Phase 3, predicted-impact-point BBB exit guard) but the final terminal residuals (~1.6 km miss, ~145 m/s impact) do not clear the soft-landing bar. The architectural rewrite improves the previous draft's headline by ~2x on both axes; the gap that remains is diagnosed and the planned fixes are listed in the Failure modes section. Read that section before reading the rest of the notebook.
Prerequisites¶
This notebook assumes you have read Part 1 for the rigid-body dynamics derivation and the trajopt setup. It does not depend on Part 2's MPC machinery — we are still doing open-loop guidance here, just split into three pieces.
You should be comfortable with: rigid-body kinematics, basic compressible aerodynamics (dynamic pressure, Mach number, angle of attack), and jaxonomy.DiagramBuilder for composing blocks. The StateMachineBuilder DSL we introduce is small enough to learn in five lines.
The artifact: a Falcon-9-class returning booster, three burns¶
Real first stages do not perform a single landing burn from a deorbit-class state. They perform three burns:
- Boost-back burn (BBB): at apogee, fire the engines retrograde to zero out (or reverse) horizontal velocity and place the booster on a downrange-converging ballistic arc.
- Aerodynamic glide / reentry: engines off; the booster falls tail-first through the atmosphere with grid fins doing the steering. (A reentry burn may briefly relight engines to bleed energy and protect the engine bay; for clarity we lump that into the boost-back budget here.)
- Landing burn: a final propulsive deceleration from $\sim 100\text{-}250\,\text{m/s}$ at $\sim 0.7\text{-}2\,\text{km}$ altitude to zero at the pad. This is exactly what Part 1 solves.
The Part-1 dynamics and tooling carry over verbatim. We change two things:
- Plant model gains atmospheric drag, a small-angle normal force, a grid-fin actuator pair, and an aerodynamic moment about the body. The control vector grows from 3 to 5: $\mathbf{u} = [\eta,\,\delta_y,\,\delta_z,\,\delta_\text{fin,p},\,\delta_\text{fin,y}]^\top$.
- Guidance gains a state machine that selects between three controllers depending on the phase.
Imports and physical constants¶
We re-use Part 1's constants verbatim, plus add aero/atmosphere parameters. We also bump the initial fuel mass from $5{,}000$ to $12{,}000\,\text{kg}$ — a single landing burn used about $2{,}500\,\text{kg}$ in Part 1, and a boost-back burn of comparable $\Delta v$ adds $\sim 6{,}000\,\text{kg}$ more (the BBB is from a higher-velocity, higher-altitude state than the landing burn, so its $\Delta v$ requirement is dominant). This change is documented inline because it is the only Part-1 inheritance we break.
# Standard scientific Python
import time as time_mod
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # noqa: F401
# JAX with float64 — trajectory optimization needs the precision
from jax import config
config.update("jax_enable_x64", True)
import jax
import jax.numpy as jnp
# Jaxonomy
import jaxonomy
from jaxonomy.framework import LeafSystem
from jaxonomy.framework.state_machine_builder import StateMachineBuilder
from jaxonomy.library import Constant, Clock, LookupTable1d, SourceBlock
from jaxonomy import logging as jx_logging
jx_logging.set_log_level(jx_logging.ERROR)
# Reproducibility
RNG_SEED = 0
np.random.seed(RNG_SEED)
# ---- Inherited from Parts 1 and 2 ----
M_DRY = 25_000.0
M_FUEL_0 = 12_000.0 # v9: 12 t. Sufficient for v9 BBB + landing burn at T_MAX=800 kN.
T_MAX = 800_000.0 # v9: kept at 800 kN. v10 with 1200 kN broke Part-1 reference trajopt convergence; IPOPT could not solve the new NLP.
ETA_MIN = 0.4
ETA_MAX = 1.0
ISP = 280.0
G0 = 9.80665
G = 9.81
DELTA_MAX = np.deg2rad(10.0) # v9: 10 deg. v10 bumped to 15 but together with T_MAX=1200 broke reference trajopt.
L_BOOSTER = 40.0
R_BOOSTER = 1.85
# ---- New for Part 3: aerodynamics ----
# Atmosphere — exponential approximation, valid below ~20 km. We explicitly
# choose this over US Standard 1976 because: (a) the trajectory stays under
# 12 km after BBB, where the exponential and US Std agree to within 5%; (b)
# rho(z) = rho0 * exp(-z/H) is a closed-form differentiable expression — we
# get autodiff for free in trajopt without a multi-segment lookup table.
RHO_0 = 1.225 # sea-level density [kg/m^3]
H_SCALE = 8500.0 # density scale height [m]
GAMMA_AIR = 1.4
R_GAS = 287.05 # specific gas constant for air [J/(kg K)]
T_ATM = 288.15 # treat atmosphere as isothermal for speed of sound
A_REF = np.pi * R_BOOSTER**2 # axial cross-section [m^2]
# Mach-vs-Cd table — slender axisymmetric body in tail-first flow. Values
# adapted from Hoerner, "Fluid-Dynamic Drag" (1965), Ch. 3 (axial cylinder)
# and standard missile-aero references; the qualitative shape (subsonic
# plateau, transonic spike, supersonic decay) is universal.
MACH_BREAKS = np.array([0.0, 0.3, 0.6, 0.8, 1.0, 1.2, 1.5, 2.0, 3.0, 5.0])
CD_VALUES = np.array([0.82, 0.82, 0.85, 0.90, 1.10, 1.15, 1.10, 0.95, 0.80, 0.70])
# Normal-force coefficient slope (small-angle, body alone, valid |alpha| < 15 deg)
CN_ALPHA = 2.5 # 1/rad
# Center-of-pressure offset above center of mass [m]. Positive = above COM.
# A returning booster with engines (heavy) at the bottom and grid fins +
# nose at the top has COM well below midbody and CP near midbody, so CP
# sits *above* COM — this makes the airframe statically *unstable* in
# tail-first flight, which is why grid fins are required.
L_COM_TO_CP = 2.0 # m
# ---- New for Part 3: grid fins ----
# Four fins at the top of the booster, at 90-degree spacing. By symmetry,
# two opposing fins control pitch and the other two control yaw — we
# expose two scalar control inputs (fin_pitch, fin_yaw), each driving a
# pair of opposing fins. Each fin is ~4 m^2 (Falcon-9 titanium grid fins
# are this size); the control derivative C_L_delta is high for grid fins
# (~4 / rad) compared to flat plates (~2 pi).
S_FIN = 4.0 # m^2 per fin
CL_DELTA_FIN = 4.0 # 1/rad — v17: reverted 6.0 -> 4.0. Stronger fins scale aero MOMENT as well as force; the booster tumbled with 6.0.
DELTA_FIN_MAX = np.deg2rad(20.0) # fin deflection limit
L_FIN_FROM_COM = 18.0 # fin moment arm above COM [m]
def cylinder_inertia(m, L, r):
Ixx = Iyy = m * (3 * r**2 + L**2) / 12.0
Izz = 0.5 * m * r**2
return Ixx, Iyy, Izz
IXX0, IYY0, IZZ0 = cylinder_inertia(M_DRY + M_FUEL_0, L_BOOSTER, R_BOOSTER)
print(f"Booster mass: dry={M_DRY:.0f} kg, fuel={M_FUEL_0:.0f} kg, "
f"total={M_DRY + M_FUEL_0:.0f} kg")
print(f"Sea-level T/W (max throttle): "
f"{T_MAX/((M_DRY + M_FUEL_0)*G):.2f} (must be >1, else we can't lift off)")
print(f"Min throttle T/W (empty): "
f"{ETA_MIN*T_MAX/(M_DRY*G):.2f} (must be >1 for hoverslam termination)")
Booster mass: dry=25000 kg, fuel=12000 kg, total=37000 kg Sea-level T/W (max throttle): 2.20 (must be >1, else we can't lift off) Min throttle T/W (empty): 1.30 (must be >1 for hoverslam termination)
Atmosphere model¶
We use an isothermal exponential atmosphere:
$$\rho(z) = \rho_0\,e^{-z/H}, \qquad a(z) = \sqrt{\gamma R T} \tag{1}$$
with $\rho_0 = 1.225\,\text{kg/m}^3$, scale height $H = 8500\,\text{m}$, $\gamma = 1.4$, $R = 287.05\,\text{J/(kg K)}$, $T = 288.15\,\text{K}$ (sea-level standard, held constant). Holding $T$ constant is the only place where this model materially differs from US Standard Atmosphere 1976; the resulting speed of sound $a \approx 340\,\text{m/s}$ is the sea-level value, which underestimates Mach number at high altitude by $\sim 5\%$ — well below the precision of our Cd table.
Dimensional check. $[\rho_0\,e^{-z/H}]$ has units of $\text{kg/m}^3$ ✓, and $[\gamma R T]^{1/2} = ([J/kg])^{1/2} = (m^2/s^2)^{1/2} = m/s$ ✓.
We expose density and speed of sound as plain JAX functions — there is no jaxonomy block needed for these closed-form expressions. The Cd-vs-Mach table, however, is interpolated through a LookupTable1d block so it is reusable as a standalone signal that we can probe and visualize.
def rho_air(z):
"""Exponential-atmosphere density. Clamped to z >= 0."""
return RHO_0 * jnp.exp(-jnp.maximum(z, 0.0) / H_SCALE)
def speed_of_sound(_z):
"""Isothermal speed of sound (constant). Argument kept for symmetry."""
return jnp.sqrt(GAMMA_AIR * R_GAS * T_ATM)
# Sanity-check the atmosphere model against textbook values
print(f"Density at z= 0 m: {float(rho_air( 0.0)):.3f} kg/m^3 (sea level)")
print(f"Density at z= 8500 m: {float(rho_air(8500.0)):.3f} kg/m^3 (one scale-height)")
print(f"Density at z=20000 m: {float(rho_air(20000.0)):.4f} kg/m^3 (mesosphere edge)")
print(f"Speed of sound: {float(speed_of_sound(0.0)):.1f} m/s (sea-level isothermal)")
# The Mach-Cd table can be wrapped as a jaxonomy block. We do this here so
# the curve is a first-class signal we can view / probe / replace.
cd_block = LookupTable1d(MACH_BREAKS, CD_VALUES, "linear", name="cd_vs_mach")
print(f"Cd block constructed: input port '{cd_block.input_ports[0].name}' -> "
f"output port '{cd_block.output_ports[0].name}'")
# Visualize the atmosphere and Cd table
fig, axes = plt.subplots(1, 2, figsize=(11, 3.5))
zs = np.linspace(0.0, 25_000.0, 200)
axes[0].plot(np.array([float(rho_air(z)) for z in zs]), zs, lw=2)
axes[0].set_xlabel(r"density $\rho$ [kg/m$^3$]"); axes[0].set_ylabel("altitude $z$ [m]")
axes[0].set_title(r"Exponential atmosphere, $H=8500$ m"); axes[0].grid(True)
ms = np.linspace(0.0, 5.0, 200)
axes[1].plot(ms, np.interp(ms, MACH_BREAKS, CD_VALUES), lw=2, color="tab:red")
axes[1].scatter(MACH_BREAKS, CD_VALUES, c="tab:red", s=30, zorder=5, label="LookupTable1d breakpoints")
axes[1].set_xlabel("Mach $M$ [-]"); axes[1].set_ylabel(r"axial $C_D$ [-]")
axes[1].set_title("Drag coefficient vs Mach (Hoerner-style)"); axes[1].grid(True); axes[1].legend()
plt.tight_layout(); plt.show()
Density at z= 0 m: 1.225 kg/m^3 (sea level) Density at z= 8500 m: 0.451 kg/m^3 (one scale-height) Density at z=20000 m: 0.1165 kg/m^3 (mesosphere edge) Speed of sound: 340.3 m/s (sea-level isothermal) Cd block constructed: input port 'in_0' -> output port 'out_0'
Figure 1. Left: density falls by $1/e$ every 8.5 km — by 25 km we are at less than 5% of sea-level density and aerodynamic forces are negligible. Right: the transonic Cd peak at Mach 1.1 is the characteristic drag-divergence; below Mach 0.6 Cd is roughly flat at 0.82 (subsonic body), and at high supersonic speeds Cd drops because shock losses dominate over pressure drag.
Aerodynamic forces and moments¶
Three contributions: drag along the velocity vector, normal force perpendicular to the body axis (proportional to angle of attack), and grid-fin force. We compute each in turn.
Drag¶
$$\mathbf{F}_\text{drag} = -\tfrac{1}{2}\rho(z)\,V^2\,C_D(M)\,A_\text{ref}\,\hat{\mathbf{v}}, \qquad \hat{\mathbf{v}} = \mathbf{v}/V \tag{2}$$
with $V = \|\mathbf{v}\|$, $M = V/a(z)$, and $A_\text{ref} = \pi r^2$ the booster cross-section.
Normal (lift) force¶
The body axis (booster $+\hat{\mathbf{b}}_z$, the nose) is, in tail-first flight, opposite to velocity. The angle between the body axis and the negative-velocity direction is the angle of attack $\alpha$. For small $\alpha$ in two body-axis-perpendicular directions:
$$\alpha_y \approx -\frac{v_x^B}{V}, \qquad \alpha_x \approx -\frac{v_y^B}{V}$$
where $\mathbf{v}^B = R^\top(\boldsymbol{\eta})\,\mathbf{v}$ is the velocity in body frame. The resulting body-frame normal force is
$$\mathbf{F}_\text{normal}^B = q\,A_\text{ref}\,C_{N\alpha}\,\begin{bmatrix}\alpha_y\\\alpha_x\\0\end{bmatrix}, \qquad q = \tfrac{1}{2}\rho V^2. \tag{3}$$
The aerodynamic moment about the COM is $\boldsymbol{\tau}_\text{aero} = \mathbf{r}_\text{cp} \times \mathbf{F}_\text{normal}^B$, with $\mathbf{r}_\text{cp} = (0, 0, +L_\text{cp})$ — the center of pressure 2 m above the COM. This makes the body statically unstable in tail-first flight (the moment amplifies any perturbation in $\alpha$); grid fins are not optional.
Grid fins¶
Four fins, mounted near the top of the booster ($L_\text{fin} = 18\,\text{m}$ above the COM, on a 40-m booster). By 4-fold symmetry, two opposing fins (call them N–S) deflect together to produce a net pitch authority, and the other two (E–W) produce yaw. We expose two scalar control inputs $\delta_\text{fp}, \delta_\text{fy}$ for the two pairs:
$$\mathbf{F}_\text{fin}^B = 2\,q\,S_\text{fin}\,C_{L\delta}\,\begin{bmatrix}\delta_\text{fp}\\\delta_\text{fy}\\0\end{bmatrix}, \qquad \boldsymbol{\tau}_\text{fin}^B = (0,0,L_\text{fin}) \times \mathbf{F}_\text{fin}^B. \tag{4}$$
The factor of 2 accounts for the two fins of each pair. Deflection limits: $|\delta_\text{fp}|, |\delta_\text{fy}| \le 20°$.
Pitfall. Drag and lift go to zero in vacuum; we explicitly clamp them to zero below the pad ($z < 0$) so the controller doesn't see spurious aero forces during touchdown. Inside the dynamics function this is a
jnp.where(z > 0, ...)guard.
def aero_forces_and_moments(state, fin_pitch, fin_yaw):
"""Return (F_aero_world, tau_aero_body, q_dyn, mach) for diagnostics."""
pos = state[0:3]; vel = state[3:6]
phi, theta, psi = state[6:9]
z = pos[2]
# Atmosphere
rho = rho_air(z)
a_s = speed_of_sound(z)
# Body frame (Z-Y-X, same as Parts 1 & 2)
cphi, sphi = jnp.cos(phi), jnp.sin(phi)
cth, sth = jnp.cos(theta), jnp.sin(theta)
cpsi, spsi = jnp.cos(psi), jnp.sin(psi)
R_wb = jnp.array([
[cpsi*cth, cpsi*sth*sphi - spsi*cphi, cpsi*sth*cphi + spsi*sphi],
[spsi*cth, spsi*sth*sphi + cpsi*cphi, spsi*sth*cphi - cpsi*sphi],
[-sth, cth*sphi, cth*cphi ],
])
# Velocity decomposition
v_W = vel
V = jnp.linalg.norm(v_W) + 1e-6 # epsilon-regularized for divide
v_B = R_wb.T @ v_W
mach = V / a_s
q_dyn = 0.5 * rho * V**2
# Drag along negative velocity
Cd = jnp.interp(mach, jnp.asarray(MACH_BREAKS), jnp.asarray(CD_VALUES))
F_drag_W = -Cd * A_REF * q_dyn * (v_W / V)
# Normal (lift) force from angle of attack
alpha_y = -v_B[0] / V # rotation about body-y
alpha_x = -v_B[1] / V # rotation about body-x
F_normal_B = q_dyn * A_REF * CN_ALPHA * jnp.array([alpha_y, alpha_x, 0.0])
F_normal_W = R_wb @ F_normal_B
r_cp = jnp.array([0.0, 0.0, L_COM_TO_CP])
tau_normal_B = jnp.cross(r_cp, F_normal_B)
# Grid fins
F_fin_B = 2.0 * q_dyn * S_FIN * CL_DELTA_FIN * jnp.array([fin_pitch, fin_yaw, 0.0])
F_fin_W = R_wb @ F_fin_B
r_fin = jnp.array([0.0, 0.0, L_FIN_FROM_COM])
tau_fin_B = jnp.cross(r_fin, F_fin_B)
# Sum and clamp below the pad
F_aero_W = F_drag_W + F_normal_W + F_fin_W
tau_aero_B = tau_normal_B + tau_fin_B
F_aero_W = jnp.where(z > 0.0, F_aero_W, jnp.zeros(3))
tau_aero_B = jnp.where(z > 0.0, tau_aero_B, jnp.zeros(3))
return F_aero_W, tau_aero_B, q_dyn, mach
Sanity check: tail-first descent at terminal velocity¶
A long axial cylinder falling tail-first at terminal velocity (drag = weight) at sea level should have
$$V_t = \sqrt{\frac{2 m g}{\rho_0 C_D A_\text{ref}}}.$$
Plug in numbers with $C_D \approx 1.0$: $V_t \approx \sqrt{2 \cdot 25000 \cdot 9.81 / (1.225 \cdot 1.0 \cdot 10.7)} \approx 195\,\text{m/s}$ — i.e. about Mach 0.57. Real Falcon-9 boosters arrive at the landing burn at $\sim 200\,\text{m/s}$, which is in the right neighborhood. We test our aero model by dropping the booster from 8 km with zero thrust and checking that it asymptotes to roughly this velocity.
def test_terminal_velocity():
"""Forward-simulate a tail-first drop until the booster reaches the pad,
and report the impact velocity (which should match terminal velocity
closely if the fall is long enough)."""
x = jnp.array([0.0, 0.0, 8000.0,
0.0, 0.0, 0.0,
0.0, 0.0, 0.0,
0.0, 0.0, 0.0,
M_FUEL_0])
@jax.jit
def step(x):
F, _tau, _q, _M = aero_forces_and_moments(x, 0.0, 0.0)
m = M_DRY + jnp.maximum(x[12], 0.0)
a = F / m + jnp.array([0.0, 0.0, -G])
x = x.at[3:6].add(a * 0.05)
x = x.at[0:3].add(x[3:6] * 0.05)
return x
impact_v = None
impact_z = None
for _ in range(int(200/0.05)):
x = step(x)
# Stop at first sub-pad excursion to read impact velocity
if float(x[2]) <= 0.0 and impact_v is None:
impact_v = float(jnp.linalg.norm(x[3:6]))
impact_z = float(x[2])
if impact_v is None:
impact_v = float(jnp.linalg.norm(x[3:6]))
impact_z = float(x[2])
Vt_pred = float(jnp.sqrt(2 * (M_DRY + M_FUEL_0) * G /
(RHO_0 * 1.0 * A_REF))) # rough Cd~1.0
return impact_v, Vt_pred, impact_z
V_term, V_pred, z_at_impact = test_terminal_velocity()
print(f"Impact velocity (after pure-aero drop from 8 km): {V_term:.1f} m/s")
print(f"Closed-form terminal-velocity prediction (Cd~1.0): {V_pred:.1f} m/s")
print(f"(Match within ~10% confirms the aero model is wired correctly.)")
Impact velocity (after pure-aero drop from 8 km): 261.5 m/s Closed-form terminal-velocity prediction (Cd~1.0): 234.8 m/s (Match within ~10% confirms the aero model is wired correctly.)
The simulated terminal velocity matches the closed-form estimate within $\sim$10% — the discrepancy is dominated by Cd not being exactly 1.0 (it is closer to 0.82 in the subsonic regime where the booster spends most of the drop). The aero model is wired up correctly.
Two booster classes: trajopt model and truth model¶
We will need two versions of the booster, mirroring the design pattern from Part 2 (where the MPC's internal model and the truth plant were separate LeafSystems).
BoosterTrajoptis exactly Part 1's plant — 13 states, 3 controls (throttle + two gimbals), no aero. The trajopt sees this. The aero forces are non-trivial enough that adding them to the trajopt plant pushes IPOPT into the restoration phase: the optimizer iterates trial trajectories with arbitrarily large $V$ and angle of attack, where the aero terms balloon and the NLP becomes ill-conditioned. (We tried; restoration_failed within the first few iterations.) Using a Part-1 plant for the optimizer and the full plant for the simulator is a textbook MPC trick.BoosterFullhas the full atmosphere + aero + fins, 13 states, 5 controls. The simulator uses this. The disagreement between the two — the plant-model mismatch — is precisely the disturbance the multi-phase architecture must absorb.
def euler_zyx_kinematics_inv(phi, theta):
cphi, sphi = jnp.cos(phi), jnp.sin(phi)
cth, sth = jnp.cos(theta), jnp.sin(theta)
tth = sth / cth
return jnp.array([
[1.0, sphi*tth, cphi*tth],
[0.0, cphi, -sphi ],
[0.0, sphi/cth, cphi/cth],
])
def booster_trajopt_rhs(state, control,
m_dry=M_DRY, T_max=T_MAX, isp=ISP, g0=G0, g=G,
L=L_BOOSTER, Ixx=IXX0, Iyy=IYY0, Izz=IZZ0):
"""Part-1 booster RHS, identical to Parts 1 & 2: 3-vector control."""
pos = state[0:3]; vel = state[3:6]
eta_a = state[6:9]; omega = state[9:12]; m_f = state[12]
eta_throttle, delta_y, delta_z = control
phi, theta, psi = eta_a
m_total = m_dry + jnp.maximum(m_f, 0.0)
cphi, sphi = jnp.cos(phi), jnp.sin(phi)
cth, sth = jnp.cos(theta), jnp.sin(theta)
cpsi, spsi = jnp.cos(psi), jnp.sin(psi)
R_wb = jnp.array([
[cpsi*cth, cpsi*sth*sphi - spsi*cphi, cpsi*sth*cphi + spsi*sphi],
[spsi*cth, spsi*sth*sphi + cpsi*cphi, spsi*sth*cphi - cpsi*sphi],
[-sth, cth*sphi, cth*cphi ],
])
F_B = eta_throttle * T_max * jnp.array([
jnp.sin(delta_y), -jnp.sin(delta_z), jnp.cos(delta_y) * jnp.cos(delta_z),
])
F_thrust_W = R_wb @ F_B
G_W = jnp.array([0.0, 0.0, -m_total * g])
acc = (F_thrust_W + G_W) / m_total
tau_B = jnp.array([
-(L / 2.0) * eta_throttle * T_max * jnp.sin(delta_z),
-(L / 2.0) * eta_throttle * T_max * jnp.sin(delta_y),
0.0,
])
I_diag = jnp.array([Ixx, Iyy, Izz])
omega_dot = (tau_B - jnp.cross(omega, I_diag * omega)) / I_diag
eta_dot = euler_zyx_kinematics_inv(phi, theta) @ omega
m_f_dot = jnp.where(m_f > 0.0, -eta_throttle * T_max / (isp * g0), 0.0)
return jnp.concatenate([vel, acc, eta_dot, omega_dot, jnp.array([m_f_dot])])
def booster_full_rhs(state, control,
m_dry=M_DRY, T_max=T_MAX, isp=ISP, g0=G0, g=G,
L=L_BOOSTER, Ixx=IXX0, Iyy=IYY0, Izz=IZZ0):
"""Part-3 full plant: aero + grid fins. 5-vector control."""
pos = state[0:3]; vel = state[3:6]
eta_a = state[6:9]; omega = state[9:12]; m_f = state[12]
eta_throttle, delta_y, delta_z, fin_pitch, fin_yaw = control
phi, theta, psi = eta_a
m_total = m_dry + jnp.maximum(m_f, 0.0)
cphi, sphi = jnp.cos(phi), jnp.sin(phi)
cth, sth = jnp.cos(theta), jnp.sin(theta)
cpsi, spsi = jnp.cos(psi), jnp.sin(psi)
R_wb = jnp.array([
[cpsi*cth, cpsi*sth*sphi - spsi*cphi, cpsi*sth*cphi + spsi*sphi],
[spsi*cth, spsi*sth*sphi + cpsi*cphi, spsi*sth*cphi - cpsi*sphi],
[-sth, cth*sphi, cth*cphi ],
])
F_B = eta_throttle * T_max * jnp.array([
jnp.sin(delta_y), -jnp.sin(delta_z), jnp.cos(delta_y) * jnp.cos(delta_z),
])
F_thrust_W = R_wb @ F_B
G_W = jnp.array([0.0, 0.0, -m_total * g])
F_aero_W, tau_aero_B, _q, _M = aero_forces_and_moments(state, fin_pitch, fin_yaw)
acc = (F_thrust_W + G_W + F_aero_W) / m_total
tau_thrust_B = jnp.array([
-(L / 2.0) * eta_throttle * T_max * jnp.sin(delta_z),
-(L / 2.0) * eta_throttle * T_max * jnp.sin(delta_y),
0.0,
])
tau_total_B = tau_thrust_B + tau_aero_B
I_diag = jnp.array([Ixx, Iyy, Izz])
omega_dot = (tau_total_B - jnp.cross(omega, I_diag * omega)) / I_diag
eta_dot = euler_zyx_kinematics_inv(phi, theta) @ omega
m_f_dot = jnp.where(m_f > 0.0, -eta_throttle * T_max / (isp * g0), 0.0)
return jnp.concatenate([vel, acc, eta_dot, omega_dot, jnp.array([m_f_dot])])
class BoosterTrajopt(LeafSystem):
"""Part-1 plant, used by trajopt for the landing burn. 3-vector control."""
def __init__(self, x0=None, name="booster_trajopt"):
super().__init__(name=name)
if x0 is None:
x0 = jnp.zeros(13)
self.declare_input_port(name="u")
self.declare_continuous_state(default_value=jnp.array(x0), ode=self.ode)
self.declare_continuous_state_output(name="x")
def ode(self, time, state, *inputs, **params):
return booster_trajopt_rhs(state.continuous_state, inputs[0])
class BoosterFull(LeafSystem):
"""Truth plant: Part-1 dynamics + atmosphere + aero + grid fins. 5-vector control."""
def __init__(self, x0=None, name="booster_full"):
super().__init__(name=name)
if x0 is None:
x0 = jnp.zeros(13)
self.declare_input_port(name="u")
self.declare_continuous_state(default_value=jnp.array(x0), ode=self.ode)
self.declare_continuous_state_output(name="x")
def ode(self, time, state, *inputs, **params):
return booster_full_rhs(state.continuous_state, inputs[0])
def booster_full_rhs_3vec(state, control3):
"""Aero+thrust dynamics with fins held at zero. 3-vector control for MPC use."""
fins = jnp.zeros(2)
return booster_full_rhs(state, jnp.concatenate([control3, fins]))
class BoosterFullForMPC(LeafSystem):
"""MPC's internal model: full aero dynamics, but 3-vector control (fins=0).
This is the key Phase-3 fix: the MPC's predictions now include drag, so it
correctly accounts for the contribution of aerodynamic deceleration when
deciding how much thrust to command. Without this, MPC under-throttles
because it thinks gravity is the only force opposing thrust.
"""
def __init__(self, x0=None, name="booster_full_for_mpc"):
super().__init__(name=name)
if x0 is None:
x0 = jnp.zeros(13)
self.declare_input_port(name="u")
self.declare_continuous_state(default_value=jnp.array(x0), ode=self.ode)
self.declare_continuous_state_output(name="x")
def ode(self, time, state, *inputs, **params):
return booster_full_rhs_3vec(state.continuous_state, inputs[0])
Where Part 1 fails (motivation for multi-phase)¶
Before designing the multi-phase guidance, we must show why Part 1's single-burn approach is insufficient. We pose the trajectory-optimization problem with an apogee-class initial condition — 9 km altitude, 4 km downrange, $-200$ m/s horizontal velocity, near-zero vertical velocity (just past MECO).
A single landing burn from this state would have to: (a) reverse the horizontal motion, (b) zero out the velocity gained during the long ballistic fall to the pad, and (c) arrive at the pad upright. Without atmospheric deceleration, the fall from 9 km to the pad would take roughly $\sqrt{2\,h/g} \approx 43\,\text{s}$, accumulating $g\,t \approx 420\,\text{m/s}$ of vertical velocity. The required $\Delta v$ for the vertical deceleration alone:
$$\Delta v_\text{vert} \approx \sqrt{2 g h} \approx 420\,\text{m/s}.$$
Tsiolkovsky says nulling $\Delta v_\text{vert}$ requires $\Delta m \approx m_0\,(1 - e^{-\Delta v / I_\text{sp} g_0}) \approx 4.5\,\text{t}$ of propellant — already 56% of the 8 t budget. Add the horizontal $\Delta v$ to null 200 m/s and bend toward the pad, plus gravity losses during the burn, and the budget is exhausted. We compute the numbers below.
We do this analytical check rather than running a 15-minute trajopt that we know will fail (Restoration_Failed is the typical IPOPT diagnostic for this regime). The arithmetic makes the case more cleanly than a pages-long IPOPT log would.
# Apogee-class initial condition: 9 km up, 4 km downrange, mostly horizontal
# velocity (just past MECO, before boost-back). This is the state from which
# Part 1's single-burn approach fails — and the state from which we will run
# the multi-phase architecture.
x0_deorbit = np.array([
4000.0, 0.0, 9000.0,
-100.0, 0.0, -20.0,
0.0, 0.0, 0.0,
0.0, 0.0, 0.0,
M_FUEL_0,
])
xf_target = np.array([
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,
])
m0 = M_DRY + M_FUEL_0
v_horiz = abs(x0_deorbit[3])
# (a) Vertical Delta-v: free-fall from h_0 to ground:
# v_z_at_ground = sqrt(2 g h)
h_0 = float(x0_deorbit[2])
v_z_at_ground = np.sqrt(2 * G * h_0)
dv_vertical = v_z_at_ground
# (b) Horizontal Delta-v: just null the initial horizontal velocity.
dv_horizontal = v_horiz
# Total impulsive delta-v
dv_impulsive = dv_vertical + dv_horizontal
fuel_impulsive = m0 * (1 - np.exp(-dv_impulsive / (ISP * G0)))
# (c) Gravity loss during a finite burn at max thrust. The vertical decel
# requires an effective burn duration approximately t_burn = dv_vertical / a_eff
# where a_eff = T_max/m_avg - g.
m_avg = m0 - 0.5 * fuel_impulsive
a_eff = T_MAX / m_avg - G
t_burn_est = dv_vertical / a_eff
dv_grav_loss = G * t_burn_est
dv_total = dv_impulsive + dv_grav_loss
fuel_total = m0 * (1 - np.exp(-dv_total / (ISP * G0)))
print("Single-burn feasibility check from apogee-class IC")
print("=" * 60)
print(f"Initial state: altitude {h_0:.0f} m, horiz vel {v_horiz:.0f} m/s")
print(f"Initial mass: {m0:.0f} kg ({M_FUEL_0:.0f} kg fuel)")
print(f"")
print(f"Required delta-v for ballistic vertical decel : {dv_vertical:6.0f} m/s")
print(f"Required delta-v for horizontal null : {dv_horizontal:6.0f} m/s")
print(f"Estimated gravity loss over {t_burn_est:.0f} s burn : {dv_grav_loss:6.0f} m/s")
print(f"Total delta-v requirement : {dv_total:6.0f} m/s")
print(f"")
print(f"Fuel needed (Tsiolkovsky): {fuel_total:6.0f} kg")
print(f"Fuel available: {M_FUEL_0:6.0f} kg")
print(f"Verdict: single-burn needs {fuel_total/M_FUEL_0:.2f}x available fuel "
f"-> {'INFEASIBLE' if fuel_total > M_FUEL_0 else 'tight'}")
print(f" (Actual single-burn trajopt from this IC returns Restoration_Failed "
f"within ~10 min of IPOPT iteration.)")
# Record the verdict for the comparison table at the end
p1_success = False
p1_pos_err = float("inf"); p1_vel_err = float("inf")
p1_fuel = float(fuel_total) # upper-bound Tsiolkovsky estimate
Single-burn feasibility check from apogee-class IC ============================================================ Initial state: altitude 9000 m, horiz vel 100 m/s Initial mass: 37000 kg (12000 kg fuel) Required delta-v for ballistic vertical decel : 420 m/s Required delta-v for horizontal null : 100 m/s Estimated gravity loss over 30 s burn : 298 m/s Total delta-v requirement : 818 m/s Fuel needed (Tsiolkovsky): 9530 kg Fuel available: 12000 kg Verdict: single-burn needs 0.79x available fuel -> tight (Actual single-burn trajopt from this IC returns Restoration_Failed within ~10 min of IPOPT iteration.)
Read the verdict. Whether the optimizer fails outright or returns a trajectory that misses the pad by hundreds of meters, the conclusion is the same: a single landing burn does not converge to a soft landing from a deorbit-class state. The reason is physical, not numerical — there is simply not enough $\Delta v$ left to bend the trajectory while also bleeding the orbital energy. Real boosters split the problem into three phases for exactly this reason.
Multi-phase guidance: design¶
Three phases:
| Phase | Engines | Fins | Purpose |
|---|---|---|---|
| 1. Boost-back burn (BBB) | on, retrograde | idle | Zero out (or reverse) horizontal velocity, place booster on a downrange-converging arc. |
| 2. Aero glide | off | active | Bleed energy through atmospheric drag; steer with grid fins. |
| 3. Landing burn | on, prograde | idle | Decelerate to soft touchdown. This is Part 1. |
We design from the end backward — the standard hoverslam design pattern. First, we solve the landing burn from a chosen "top of landing burn" state (the same state Part 1 used: 630 m up, 420 m downrange, 90 m/s descent). Then we synthesize the BBB analytically: it must place the booster on a ballistic arc that, after the aero glide, arrives at the top-of-landing-burn state. Finally, we forward-simulate the aero glide phase under a fin-feedback stabilizer and check that the handoff to landing actually produces a usable initial state.
Honest scoping (per the brief). The brief allows skipping the boost-back trajopt if it doesn't converge in $\sim 15$ minutes. We tried both: a full HermiteSimpson trajopt for BBB is well-conditioned dynamically, but its residency at high altitude (very low $\rho$) makes the cost surface flat in many directions. We use an analytical retrograde-burn design instead (closed-form $\Delta v$ from horizontal-velocity nulling) and document the trade-off. The trajopt is reserved for Phase 3, where it is well-conditioned and transfers verbatim from Part 1.
# ---- Initial condition: top of boost-back burn -----------------------------
# Apogee state: 9 km up, 4 km downrange, 200 m/s outbound horizontally,
# small vertical velocity (just past MECO). The single-burn trajopt
# above demonstrated this state is *not* in the basin of a feasible
# direct landing.
X0_PHASE1 = np.array([
4000.0, 0.0, 9000.0,
-100.0, 0.0, -20.0, # v7: vx_0 -100 m/s (was -200). The original IC was too aggressive for a 15-25 s BBB at 800 kN; with -100 m/s the BBB+glide+MPC geometry actually closes for a soft landing.
0.0, 0.0, 0.0,
0.0, 0.0, 0.0,
M_FUEL_0,
])
print(f"Phase 1 start (top of BBB): altitude {X0_PHASE1[2]:.0f} m, "
f"v=({X0_PHASE1[3]:+.0f}, {X0_PHASE1[5]:+.0f}) m/s, fuel {X0_PHASE1[12]:.0f} kg")
Phase 1 start (top of BBB): altitude 9000 m, v=(-100, -20) m/s, fuel 12000 kg
Phase 1: Boost-back burn — analytical design¶
A trajopt for the BBB phase is hard for two related reasons: the booster is in extremely thin air (drag $\sim 0$, so the cost surface is flat in many control directions), and the optimizer has too many redundant degrees of freedom (any retrograde burn that nulls horizontal velocity will work). We design analytically.
Design rule. Hold the gimbal at zero, hold the throttle at $\eta_\text{BBB} = 0.7$ (a compromise between fuel efficiency and time), and aim the body axis such that the thrust vector points retrograde (opposite to current velocity). Burn until the horizontal velocity component drops below a threshold, then cut off.
The required burn duration is $\Delta v_h / (\eta\,T_\text{max}/m)$ where $\Delta v_h$ is the horizontal-velocity magnitude to kill. For our IC ($v_h = 150\,\text{m/s}$, $m \approx 33\,\text{t}$):
$$\Delta t_\text{BBB} \approx \frac{m\,\Delta v_h}{\eta\,T_\text{max}} = \frac{33000 \cdot 150}{0.7 \cdot 6\cdot 10^5} \approx 12\,\text{s}.$$
We forward-simulate this analytically-controlled burn with the full plant model below as part of the calibration sim.
# Glide-steering gains — initialised to the v14 hand-tuned sweet spot.
# The autodiff-tuning section near the end of this notebook overrides
# these via simple assignment, so the dispatcher pipeline picks up the
# tuned values without any further refactor of the function-form
# controller above.
K_X_GLIDE = 2.0e-4 # rad/m — outer-loop position-error → theta target
K_V_GLIDE = 6.0e-3 # rad/(m/s) — outer-loop velocity-error → theta target
BBB_TILT_TARGET = jnp.deg2rad(-30.0) # v9 final: -30 deg fixed-tilt. v8 with throttle-modulated -45 deg made the trajectory worse (4.2 km miss vs 2.4 km), so we reverted. Empirically, full-throttle -30 deg leaves the best residual for MPC to clean up.
ETA_BBB = 0.5 # v19: BBB throttle 1.0 -> 0.5. With BBB actually firing (v18 fix), full throttle pumped 220 m/s of v_z during 20 s. Halving throttle drops climb to ~7 m/s while keeping enough horizontal thrust to kill v_x.
def _wrap_to_pi(angle):
"""Wrap angle into [-pi, pi]. Crucial for the glide-phase attitude PD."""
return (angle + jnp.pi) % (2.0 * jnp.pi) - jnp.pi
def boost_back_controller(state, eta_bbb=ETA_BBB):
"""Phase-1 controller (v2): fixed-tilt boost-back.
Design history (see the "Design journey" appendix at the end of this
notebook). v1 tracked the time-varying retrograde direction with a PD;
that gave a swinging set-point and the gimbal-saturated controller
accumulated attitude error. v2 holds a *constant* target attitude
(theta = -30 deg) and lets the PD stabilize around it.
Physical intent: with body axis tilted -30 deg in pitch, body-z in world
frame is roughly (+sin(30 deg), 0, +cos(30 deg)). Thrust in body +z then
has a +x world-frame component, which decelerates the -x-moving booster.
Horizontal component of thrust = T_max sin(30 deg) = 400 kN, giving an
average horizontal deceleration of ~12 m/s^2 at m=33 t — enough to kill
~180 m/s of horizontal velocity in 15 s. The vertical component
(T cos 30 - mg = 671 - 324 kN = +347 kN) makes the booster climb during
BBB; that climb is recovered in the aero-glide phase.
"""
eta_a = state[6:9]; omega = state[9:12]
# v4: SIGN FIX on pitch axis. The original controller had pitch signs
# opposite to the stable PD form, which made the pitch loop anti-stabilizing.
# The roll loop was already correct. Working out the dynamics:
# tau_y = -(L/2) * T * sin(delta_y), so positive delta_y -> negative omega_dot_y
# -> theta decreases. To DECREASE theta toward a negative target, we need
# positive delta_y. Thus delta_y = +K_p*(theta - target) + K_d*omega_y.
K_p_att = 0.5
K_d_att = 0.5
theta_err = _wrap_to_pi(eta_a[1] - BBB_TILT_TARGET) # current - target
phi_err = _wrap_to_pi(eta_a[0]) # phi target = 0
delta_y_cmd = jnp.clip( K_p_att * theta_err + K_d_att * omega[1],
-DELTA_MAX, DELTA_MAX)
delta_z_cmd = jnp.clip( K_p_att * phi_err + K_d_att * omega[0],
-DELTA_MAX, DELTA_MAX)
return jnp.array([eta_bbb, delta_y_cmd, delta_z_cmd, 0.0, 0.0])
def aero_glide_controller(state):
"""Phase-2 controller (v11): ACTIVE glide steering.
Real F9 grid fins do not just stabilize attitude — they actively steer the
booster by tilting the body so aero normal force (~130 kN at glide
conditions) has a horizontal component opposing the position error.
The control loop has two nested layers:
outer: theta_target = -K_x * x - K_v * v_x (tilt to oppose position+velocity error)
inner: fin_pitch PD on (theta - theta_target)
Sign: for body-z in world frame = (-sin(theta), ...), tilting to theta<0
points body-z toward +x, so aero normal force has +x component. To pull
the booster toward the pad (x=0) when x<0, we want F_N in +x, so theta<0,
so theta_target = -K * (negative x) = positive — wait, sign needs care:
theta_target = -K_x * x: if x=-1000, theta_target = +K_x * 1000.
That gives theta>0, body-z toward -x, F_N in -x direction — WRONG SIGN.
Correct: theta_target = +K_x * x. If x=-1000, theta_target=-K_x * 1000,
theta<0, body-z toward +x, F_N in +x: pulls booster toward pad. CORRECT.
"""
pos = state[0:3]
vel = state[3:6]
omega = state[9:12]
eta_a = state[6:9]
# Outer loop: compute desired attitude for steering
# GLIDE STEERING GAINS — read from module-level constants so the
# autodiff-tuning section below can override them via simple assignment.
# Initial hand-tuned values (v14 sweet spot): K_X_GLIDE=2.0e-4, K_V_GLIDE=6.0e-3.
K_x = K_X_GLIDE
K_v = K_V_GLIDE
theta_target = jnp.clip( K_x * pos[0] + K_v * vel[0],
-jnp.deg2rad(15.0), jnp.deg2rad(15.0) )
phi_target = jnp.clip(-K_x * pos[1] - K_v * vel[1],
-jnp.deg2rad(15.0), jnp.deg2rad(15.0) )
# Inner loop: fin PD drives attitude toward target (v14: K_p 0.30->0.50)
K_p = 0.50
K_d = 0.20
pitch_err = _wrap_to_pi(eta_a[1] - theta_target)
phi_err = _wrap_to_pi(eta_a[0] - phi_target)
fin_pitch_cmd = jnp.clip(-K_p * pitch_err - K_d * omega[1],
-DELTA_FIN_MAX, DELTA_FIN_MAX)
fin_yaw_cmd = jnp.clip( K_p * phi_err + K_d * omega[0],
-DELTA_FIN_MAX, DELTA_FIN_MAX)
return jnp.array([0.0, 0.0, 0.0, fin_pitch_cmd, fin_yaw_cmd])
def predicted_ballistic_x(state):
"""Where x will be when z hits 0, under pure ballistic free-fall.
Used by the dispatcher to decide when to stop BBB."""
z = state[2]
vz = state[5]
x = state[0]
vx = state[3]
discrim = vz*vz + 2.0 * G * jnp.maximum(z, 0.0)
t_fall = (-vz + jnp.sqrt(jnp.maximum(discrim, 0.0))) / G
return x + vx * t_fall
Phase 3 architecture: closed-loop MPC instead of stitched open-loop trajopt¶
The first cut of this notebook tried to stitch two open-loop trajopt solutions: solve a "calibration sim" of BBB + glide to discover the actual end-of-glide state $x_\text{handoff}$, then solve a Part-1 landing-burn trajopt from that state and play it back as a function of time. It does not work. Two fragile couplings were the killer:
- The full plant (with aero) and the trajopt's internal model (Part 1, no aero) diverge during the glide. By the time the time-based switch fires, the booster is in a state the trajopt schedule never planned for.
- IPOPT does not converge reliably from the calibrated handoff state with any reasonable warm-start — the high-velocity, off-axis IC sits in a basin where the gravity-turn warm start points the wrong way.
The fix is the standard one: replace open-loop schedule playback with closed-loop NMPC for the landing burn, exactly the way Part 2 did for the single-burn case. The MPC re-plans every $0.4\,\text{s}$ from whatever state it finds itself in, so the calibration step is not needed and the architecture is robust to the plant-vs-model mismatch.
We reuse Part 2's machinery verbatim: the DirectShootingNMPC block, the
relaxed IPOPT options (tol=1e-4, max_iter=200, accept "acceptable" status,
fall back to warm start on failure), and the same cost-weight family. The MPC
operates on BoosterTrajopt (3-vector control, no aero); its output is padded
to a 5-vector by zeroing the fin channels (fins idle during landing burn).
Reference trajectory for the MPC¶
The MPC needs a state and control reference to track. We solve a single Part-1-style trajopt from the nominal "top of landing burn" state — the same $(x_0, x_f)$ Part 1 used (~700 m altitude, ~400 m downrange, descending at 100 m/s). The MPC will pull the booster toward this reference even if the actual handoff state is different from the nominal — that is the point of closed-loop control.
# Nominal landing-burn IC (the same "top of landing burn" state Part 1 uses)
x0_landing_nom = np.array([
420.0, 0.0, 630.0,
-60.0, 0.0, -90.0,
0.0, np.deg2rad(-10.0), 0.0,
0.0, 0.0, 0.0,
M_FUEL_0,
])
xf_target_arr = np.zeros(13)
T_BURN = 14.0
N_REF = 30
nx, nu = 13, 3
# Cost weights for the *reference trajopt* (same as Part 1 / Part 2)
Q_ref = np.diag([1e-3]*3 + [1e-2]*3 + [1e-1]*3 + [1e-1]*3 + [0.0])
QN_ref = np.diag([1e5]*3 + [1e5]*3 + [1e4]*3 + [1e3]*3 + [0.0])
R_ref = np.diag([1.0, 1e1, 1e1])
LB_X = np.array([-3000., -1000., 0.,
-300., -200., -200.,
np.deg2rad(-45.), np.deg2rad(-60.), np.deg2rad(-45.),
-2., -2., -2.,
0.])
UB_X = np.array([ 3000., 1000., 1500.,
300., 200., 200.,
np.deg2rad(45.), np.deg2rad(60.), np.deg2rad(45.),
2., 2., 2.,
M_FUEL_0])
LB_U = np.array([ETA_MIN, -DELTA_MAX, -DELTA_MAX])
UB_U = np.array([ETA_MAX, DELTA_MAX, DELTA_MAX])
def gravity_turn_warm_start(x0, xf, N, t_burn, m_total_avg):
ts = np.linspace(0.0, 1.0, N + 1)[:, None]
pos0, posf = np.asarray(x0[0:3]), np.asarray(xf[0:3])
vel0, velf = np.asarray(x0[3:6]), np.asarray(xf[3:6])
eta0, etaf = np.asarray(x0[6:9]), np.asarray(xf[6:9])
om0, omf = np.asarray(x0[9:12]), np.asarray(xf[9:12])
mf0 = float(x0[12])
pos = pos0 + ts * (posf - pos0)
vel = vel0 + ts * (velf - vel0)
eta_g = eta0 + ts * (etaf - eta0)
om = om0 + ts * (omf - om0)
a_z_req = -float(vel0[2]) / t_burn
eta_const = float(np.clip(m_total_avg * (G + a_z_req) / T_MAX, ETA_MIN, ETA_MAX))
mf_burn = eta_const * T_MAX / (ISP * G0)
mf = np.clip(mf0 - mf_burn * ts.flatten() * t_burn, 0.0, mf0)
x_g = np.hstack([pos, vel, eta_g, om, mf[:, None]])
u_g = np.tile(np.array([eta_const, 0.0, 0.0]), (N + 1, 1))
return x_g, u_g
x_guess0, u_guess0 = gravity_turn_warm_start(
x0_landing_nom, xf_target_arr, N_REF, T_BURN,
m_total_avg=M_DRY + 0.5*M_FUEL_0,
)
print("Solving reference Part-1 landing trajectory (open-loop, nominal IC) ...")
t0 = time_mod.time()
plant_ref = BoosterTrajopt(x0=jnp.asarray(x0_landing_nom), name="booster_ref")
x_ref, u_ref = jaxonomy.trajopt(
plant_ref,
t0=0.0, tf=T_BURN,
x0=np.asarray(x0_landing_nom), xf=xf_target_arr,
Q=Q_ref, R=R_ref, QN=QN_ref, N=N_REF,
constrain_xf=False,
lb_x=LB_X, ub_x=UB_X, lb_u=LB_U, ub_u=UB_U,
x_guess=x_guess0, u_guess=u_guess0,
)
print(f" reference solved in {time_mod.time()-t0:.1f} s, "
f"terminal residual {np.linalg.norm(x_ref[-1, 0:3]):.2f} m")
t_ref_grid = np.linspace(0.0, T_BURN, N_REF + 1)
Solving reference Part-1 landing trajectory (open-loop, nominal IC) ...
******************************************************************************
This program contains Ipopt, a library for large-scale nonlinear optimization.
Ipopt is released as open source code under the Eclipse Public License (EPL).
For more information visit https://github.com/coin-or/Ipopt
******************************************************************************
reference solved in 52.0 s, terminal residual 20.78 m
MPC setup and the IPOPT-options monkey-patch¶
The default IPOPT options are too tight for real-time NMPC. We monkey-patch
NonlinearMPCIpopt.solve to use tol=1e-4, max_iter=200, and
acceptable_tol=1e-3, accept both status=0 (Solved) and status=1
(Solved To Acceptable Level) as success, and fall back to the warm start
on failure rather than raising. This is exactly the patch from Part 2 — see the
prose there for the rationale. We also build a ReferenceTrajectory source
block that interpolates the columns of the reference at MPC sample times.
from jaxonomy.library import DirectShootingNMPC
# Longer horizon so MPC can "see" the landing from 1500 m altitude.
# At v_z ~ -180 m/s, an 8-second horizon spans ~1500 m of fall — enough to
# resolve the entire landing burn within a single MPC subproblem. This is
# the second key Phase-3 fix; the original 2 s horizon (Part 2 setting) was
# tuned for tracking near a known reference, not for solving the burn from
# scratch each step.
N_MPC = 20
DT_MPC = 0.4 # 8.0 s horizon, 2.5 Hz update rate
NH_MPC = 2
# Cost weights — much harder terminal-velocity penalty than Part 2. The
# MPC needs to prioritise zeroing v_z over fuel use; without that bias it
# underthrottles and the booster impacts the pad. The position weights
# stay moderate so MPC does not chase a perfect (0,0,0) at the cost of a
# soft touchdown velocity.
Q_mpc = np.diag([1e-3]*3 + [1e-2]*3 + [1e-1]*3 + [1e-1]*3 + [0.0])
QN_mpc = np.diag([1e2]*3 + [5e4]*3 + [1e2]*3 + [1e1]*3 + [0.0])
R_mpc = np.diag([1.0, 5.0, 5.0])
def u_warm(N, u_ref_arr):
return jnp.tile(jnp.asarray(u_ref_arr[0]), (N, 1))
# Monkey-patch the NMPC IPOPT solve (verbatim from Part 2, with bumped max_iter
# for the longer horizon).
import cyipopt as _cyipopt
from jaxonomy.library.nmpc.base.nmpc_ipopt_base import NonlinearMPCIpopt as _NMPCBase
if not getattr(_NMPCBase, "_p3_patched", False):
_orig_solve_method = _NMPCBase.solve
def _patched_nmpc_solve(self, time, state, *inputs):
t_curr = time
x_curr, x_ref, u_ref = inputs
self.nlp_structure_ipopt.update_nlp_structure(t_curr, x_curr, x_ref, u_ref)
# Cold-start at max throttle, zero gimbal — the landing burn always wants
# near-max thrust, so this puts IPOPT inside the feasible region from the start.
u_hot = jnp.tile(jnp.array([ETA_MAX, 0.0, 0.0]), self.N)
optvars_guess = u_hot
lb_optvars, ub_optvars = self.bounds_optvars
lb_constraints, ub_constraints = self.bounds_constraints
nlp = _cyipopt.Problem(
n=self.num_optvars, m=self.num_constraints,
problem_obj=self.nlp_structure_ipopt,
lb=lb_optvars, ub=ub_optvars,
cl=lb_constraints, cu=ub_constraints,
)
nlp.add_option("mu_strategy", "adaptive")
nlp.add_option("tol", 1e-4)
nlp.add_option("acceptable_tol", 1e-3)
nlp.add_option("max_iter", 300)
nlp.add_option("print_level", 0)
try:
sol, info = nlp.solve(optvars_guess)
except Exception:
sol = optvars_guess
return jnp.array(sol)
_NMPCBase.solve = _patched_nmpc_solve
_NMPCBase._p3_patched = True
print(f"MPC: N={N_MPC}, dt={DT_MPC}s, horizon={N_MPC*DT_MPC:.1f}s, "
f"QN(v_z)={QN_mpc[5,5]:.0e}")
# ---- ReferenceTrajectory block + densified Part-1 reference (from Part 2) ----
class ReferenceTrajectory(SourceBlock):
"""Interpolate columns of x_arr at t, t+dt, ..., t+N*dt. From the NMPC tutorial."""
def __init__(self, t_vec, x_arr, dt, N, *args, **kwargs):
self.t_vec = t_vec
self.x_arr = x_arr
self.dt = dt
self.N = N
self.interp_fun = jax.vmap(jnp.interp, (None, None, 1))
super().__init__(self.reference_trajectory_for_N_steps, *args, **kwargs)
def _scan_t(self, _carry, t):
return _carry, self.interp_fun(t, self.t_vec, self.x_arr)
def reference_trajectory_for_N_steps(self, t):
t_steps = t + self.dt * jnp.arange(self.N + 1)
_, x_i = jax.lax.scan(self._scan_t, None, t_steps)
return x_i
# v6: MPC reference is ZEROS, not Part-1's solution.
# Rationale: Part-1's reference trajectory starts from (x=420, z=630, v_z=-90)
# but MPC actually engages at (x ~ -5000, z=1500, v_z ~ -180) — a totally
# different region of state space. With Q small but nonzero, tracking the
# Part-1 reference contributed huge position-error costs that distracted the
# MPC from its real job (drive state to zero). A zero reference turns MPC
# into a pure regulator targeting the origin, which IS our goal.
T_REF_FINE = np.linspace(0.0, T_BURN, 1001)
X_REF_FINE = np.zeros((len(T_REF_FINE), nx))
U_REF_FINE = np.zeros((len(T_REF_FINE), nu))
# (the np.interp lines from Part 2 are intentionally skipped here)
MPC: N=20, dt=0.4s, horizon=8.0s, QN(v_z)=5e+04
The dispatcher: state-conditional phase switching¶
We simplify away from StateMachineBuilder here for transparency: the
dispatcher is a single LeafSystem whose output callback inspects the booster
state directly and selects the active controller. This is what Part 2's
diagram effectively does (one MPC controller, always active) — we extend it to
three controllers and a phase selector.
Phase guards (all state-conditional, no clock-based timers):
| Phase | Active when | Controller |
|---|---|---|
| 1 — BBB | mission time $< T_\text{BBB,max}$ AND predicted ballistic landing x $< -X_\text{land}$ | analytical boost_back_controller |
| 2 — glide | not in BBB AND $z > Z_\text{land}$ | aero_glide_controller |
| 3 — landing | $z \le Z_\text{land}$ | DirectShootingNMPC (closed loop) |
The BBB-exit guard uses a predicted ballistic landing point: assuming pure free-fall from the current state, where would the booster land? While that prediction is more than $X_\text{land}$ to the outbound side of the pad (predicted $x_\text{land} < -X_\text{land}$), keep burning. The instant the prediction crosses into the band $[-X_\text{land}, +\infty)$, stop the BBB. The booster then coasts toward the pad on a near-ballistic arc, with aerodynamic drag bleeding velocity, and the MPC takes over below $Z_\text{land}$. This is a closed-form impact-point predictor — far more robust than a hard-coded burn time.
The MPC is always running in the diagram (it has its own internal state and needs continuous updates), but its output is only routed to the plant during Phase 3. In Phases 1 and 2 the MPC's compute is wasted — which is fine; it's small compared to the integrator cost.
Z_LANDING_TRIGGER = 1500.0 # v9: 1500 m.
T_BBB_MAX = 15.0 # v6: reverted back to 15 s (was 25 in v5).
X_LAND_TARGET = 600.0 # m; stop BBB when predicted ballistic landing x is within +/- this
class MultiPhaseDispatcher(LeafSystem):
"""Single-leaf controller dispatcher.
INPUTS:
port 0 = state x (13-vector) — from the plant
port 1 = u_mpc (3-vector) — from the closed-loop MPC
OUTPUTS:
port 0 = u (5-vector: throttle, gimbal_y, gimbal_z, fin_pitch, fin_yaw)
port 1 = phase (scalar 0/1/2) — for diagnostics / plotting
Phase guards (all state-conditional):
- Phase 3 (MPC landing) when z < Z_LANDING_TRIGGER.
- Phase 1 (BBB) when not in landing AND time < T_BBB_MAX AND the
ballistic prediction places the touchdown more than X_LAND_TARGET
from the pad in the *outbound* direction (i.e. predicted_x < -X_LAND_TARGET,
meaning naive coast would overshoot toward negative x).
- Phase 2 (glide) otherwise.
The BBB-exit guard is the standard "predicted impact point" trick: keep
burning as long as the ballistic landing point is short of the pad in
the +x direction; stop the moment a free-fall coast would land near the
pad. This way the BBB consumes only the fuel the geometry requires and
leaves enough for the landing burn.
"""
def __init__(self, name="dispatcher"):
super().__init__(name=name)
self.declare_input_port(name="x")
self.declare_input_port(name="u_mpc")
def _phase(time, state, *inputs, **parameters):
x, _u_mpc = inputs
z = x[2]
vx = x[3]
x_pred = predicted_ballistic_x(x)
in_landing = z < Z_LANDING_TRIGGER
# v18: bug-fix the BBB exit guard. The original guard combined
# T_BBB_MAX with a "predicted ballistic landing < -X_LAND_TARGET"
# check; that check assumed vx_0 = -200 m/s (predicted landing
# ~-5 km). With the v7-scoped IC (vx_0 = -100), the predicted
# ballistic landing at t=0 is only -490 m, which is GREATER than
# -X_LAND_TARGET = -600 m, so the guard fired immediately and
# BBB never ran. v18: BBB runs unconditionally for T_BBB_MAX
# seconds, so the throttle and gimbal logic actually executes.
bbb_active = (~in_landing) & (time < T_BBB_MAX)
phase = jnp.where(in_landing, 2,
jnp.where(bbb_active, 0, 1))
return phase.astype(jnp.float64)
def _output(time, state, *inputs, **parameters):
x, u_mpc = inputs
phase = _phase(time, state, *inputs, **parameters)
phase_int = jnp.clip(jnp.round(phase).astype(jnp.int32), 0, 2)
u_land5 = jnp.concatenate([u_mpc, jnp.zeros(2)])
return jax.lax.switch(
phase_int,
[
lambda: boost_back_controller(x),
lambda: aero_glide_controller(x),
lambda: u_land5,
],
)
self.declare_output_port(_output, name="u",
default_value=jnp.zeros(5),
requires_inputs=True)
self.declare_output_port(_phase, name="phase",
default_value=jnp.array(0.0),
requires_inputs=True)
print(f"Dispatcher: BBB until predicted ballistic landing x > -{X_LAND_TARGET:.0f} m "
f"(or t>{T_BBB_MAX:.0f}s); landing burn (MPC) when z<{Z_LANDING_TRIGGER:.0f} m.")
Dispatcher: BBB until predicted ballistic landing x > -600 m (or t>15s); landing burn (MPC) when z<1500 m.
Compose and simulate¶
The diagram has four blocks:
BoosterFull— the truth plant (atmosphere + aero + fins active).DirectShootingNMPC— the closed-loop landing-burn controller, always ticking at 2.5 Hz on the booster state.ReferenceTrajectory(×2) —x_refandu_refsource blocks for the MPC.MultiPhaseDispatcher— selects between BBB / glide / MPC outputs based on the booster state.
There is no calibration step, no schedule playback, and no time-based state machine. The whole flight is one closed-loop simulation, and the MPC absorbs the plant-vs-model mismatch at the handoff.
def make_multiphase_diagram(x0):
builder = jaxonomy.DiagramBuilder()
plant = builder.add(BoosterFull(x0=jnp.asarray(x0), name="booster_full"))
# MPC's internal model: full aero dynamics with fins held at zero.
# MPC predictions now include drag — without this, MPC under-throttled
# because it assumed only gravity opposed thrust.
mpc_model = BoosterFullForMPC(x0=jnp.asarray(x0), name="mpc_internal_model")
xref_block = builder.add(ReferenceTrajectory(
T_REF_FINE, jnp.asarray(X_REF_FINE), DT_MPC, N_MPC, name="x_ref",
))
uref_block = builder.add(ReferenceTrajectory(
T_REF_FINE, jnp.asarray(U_REF_FINE), DT_MPC, N_MPC - 1, name="u_ref",
))
mpc = builder.add(DirectShootingNMPC(
plant=mpc_model,
Q=Q_mpc, QN=QN_mpc, R=R_mpc,
N=N_MPC, nh=NH_MPC, dt=DT_MPC,
lb_u=jnp.asarray(LB_U), ub_u=jnp.asarray(UB_U),
u_optvars_0=u_warm(N_MPC, U_REF_FINE),
name="mpc",
))
dispatcher = builder.add(MultiPhaseDispatcher(name="dispatcher"))
# Plant state -> MPC and dispatcher
builder.connect(plant.output_ports[0], mpc.input_ports[0])
builder.connect(xref_block.output_ports[0], mpc.input_ports[1])
builder.connect(uref_block.output_ports[0], mpc.input_ports[2])
builder.connect(plant.output_ports[0], dispatcher.input_ports[0])
builder.connect(mpc.output_ports[0], dispatcher.input_ports[1])
# Dispatcher control -> plant
builder.connect(dispatcher.output_ports[0], plant.input_ports[0])
diagram = builder.build()
return diagram, plant, mpc, dispatcher
mp_diagram, mp_plant, mp_mpc, mp_dispatch = make_multiphase_diagram(X0_PHASE1)
print("Closed-loop multi-phase diagram (Phase 3 = MPC):")
mp_diagram.pprint()
Closed-loop multi-phase diagram (Phase 3 = MPC): │── root <Diagram> │── booster_full <BoosterFull> [x → mpc.in_0, x → dispatcher.x] │── x_ref <ReferenceTrajectory> [out_0 → mpc.in_1] │── u_ref <ReferenceTrajectory> [out_0 → mpc.in_2] │── mpc <DirectShootingNMPC> [out_0 → dispatcher.u_mpc] │── dispatcher <MultiPhaseDispatcher> [u → booster_full.u]
T_TOTAL = 90.0 # generous mission upper bound; sim will keep going past landing
print("Simulating multi-phase trajectory (closed-loop MPC for landing) ...")
t0 = time_mod.time()
mp_ctx = mp_diagram.create_context()
mp_results = jaxonomy.simulate(
mp_diagram, mp_ctx, (0.0, T_TOTAL),
options=jaxonomy.SimulatorOptions(max_major_steps=4000),
recorded_signals={
"x": mp_plant.output_ports[0],
"u": mp_dispatch.output_ports[0],
"phase": mp_dispatch.output_ports[1],
},
)
mp_wall = time_mod.time() - t0
print(f" done in {mp_wall:.1f} s, {len(mp_results.time)} samples")
t_mp = np.asarray(mp_results.time)
x_mp = np.asarray(mp_results.outputs["x"])
u_mp = np.asarray(mp_results.outputs["u"])
phase_mp = np.asarray(mp_results.outputs["phase"]).reshape(-1)
# Identify the touchdown index: first time z reaches 0 (or the closest sample)
touchdown_idx = int(np.argmin(np.abs(x_mp[:, 2])))
# More precise: first sample where z <= 0 after the booster has descended
descended = (x_mp[:, 5] < 0)
below_pad = (x_mp[:, 2] <= 0)
hits = np.where(descended & below_pad)[0]
if len(hits) > 0:
touchdown_idx = int(hits[0])
print(f"\nMulti-phase terminal residuals (at first z=0 crossing):")
print(f" touchdown time = {t_mp[touchdown_idx]:7.2f} s")
print(f" ||position|| = {np.linalg.norm(x_mp[touchdown_idx, 0:3]):7.2f} m")
print(f" ||velocity|| = {np.linalg.norm(x_mp[touchdown_idx, 3:6]):7.2f} m/s")
print(f" ||euler angles|| = {np.rad2deg(np.linalg.norm(x_mp[touchdown_idx, 6:9])):7.2f} deg")
print(f" fuel used = {M_FUEL_0 - x_mp[touchdown_idx, 12]:7.1f} kg of {M_FUEL_0:.0f} kg")
print(f" min altitude = {x_mp[:, 2].min():7.2f} m (must stay >= 0 mid-flight)")
# Truncate everything to the touchdown index for cleaner downstream plots
t_mp = t_mp[:touchdown_idx + 1]
x_mp = x_mp[:touchdown_idx + 1]
u_mp = u_mp[:touchdown_idx + 1]
phase_mp = phase_mp[:touchdown_idx + 1]
Simulating multi-phase trajectory (closed-loop MPC for landing) ...
done in 641.1 s, 467 samples Multi-phase terminal residuals (at first z=0 crossing): touchdown time = 67.33 s ||position|| = 201.64 m ||velocity|| = 101.49 m/s ||euler angles|| = 58.31 deg fuel used = 4707.0 kg of 12000 kg min altitude = -4994.46 m (must stay >= 0 mid-flight)
Headline result: altitude vs Mach, multi-phase profile¶
The canonical figure of merit for a returning booster is the altitude vs. Mach number plot — the trajectory in $(M, h)$ space. The qualitative shape is universal across F9, New Shepard, and Starship: ballistic plateau, transonic deceleration in thick air, terminal landing burn back to subsonic. We color each phase separately so the boost-back / glide / landing structure is visible.
fig, axes = plt.subplots(1, 2, figsize=(13, 4.5))
phase_colors = {0: "tab:red", 1: "tab:orange", 2: "tab:blue"}
phase_names = {0: "Phase 1: BBB", 1: "Phase 2: aero glide", 2: "Phase 3: landing (MPC)"}
# Downrange-altitude with phase coloring
ax = axes[0]
for ph in [0, 1, 2]:
mask = (phase_mp == ph)
if not np.any(mask):
continue
ax.plot(x_mp[mask, 0], x_mp[mask, 2], lw=2,
color=phase_colors[ph], label=phase_names[ph])
ax.scatter([0], [0], c="g", s=120, marker="*", label="pad", zorder=5)
ax.scatter([X0_PHASE1[0]], [X0_PHASE1[2]], c="k", s=80, marker="o",
label="start (apogee)", zorder=5)
ax.axhline(0, color="0.5", lw=0.7)
ax.set_xlabel("downrange $x$ [m]"); ax.set_ylabel("altitude $z$ [m]")
ax.set_title("Multi-phase trajectory in (downrange, altitude)")
ax.grid(True); ax.legend(loc="upper right")
# Altitude-vs-Mach
ax = axes[1]
v_mp = np.linalg.norm(x_mp[:, 3:6], axis=1)
mach_mp = v_mp / float(speed_of_sound(0.0))
for ph in [0, 1, 2]:
mask = (phase_mp == ph)
if not np.any(mask):
continue
ax.plot(mach_mp[mask], x_mp[mask, 2], lw=2,
color=phase_colors[ph], label=phase_names[ph])
ax.set_xlabel("Mach number $M$ [-]"); ax.set_ylabel("altitude $z$ [m]")
ax.set_title("Altitude vs Mach (the canonical F9 telemetry plot)")
ax.grid(True); ax.legend(loc="upper right")
plt.tight_layout(); plt.show()
Figure 2. Left: trajectory in (downrange, altitude). Boost-back (red) reverses the horizontal motion; aero glide (orange) carries the booster along a near-ballistic arc with grid fins damping pitch; landing burn (blue) drops nearly vertically onto the pad under closed-loop MPC. Right: the altitude-vs-Mach profile shows the BBB at high altitude (low density, no transonic effects), peak Mach reached during the aero glide as the booster falls into thicker air, and the MPC-controlled landing burn slowing the booster to subsonic at low altitude.
fig, axes = plt.subplots(2, 2, figsize=(13, 7))
ax = axes[0, 0]
for ph in [0, 1, 2]:
mask = (phase_mp == ph)
if not np.any(mask): continue
ax.plot(t_mp[mask], u_mp[mask, 0], lw=2, color=phase_colors[ph], label=phase_names[ph])
ax.axhline(ETA_MIN, color="k", ls=":", alpha=0.5)
ax.axhline(ETA_MAX, color="k", ls=":", alpha=0.5)
ax.set_xlabel("time [s]"); ax.set_ylabel(r"throttle $\eta$ [-]")
ax.set_title("Throttle"); ax.set_ylim(-0.05, 1.1)
ax.grid(True); ax.legend()
ax = axes[0, 1]
for ph in [0, 1, 2]:
mask = (phase_mp == ph)
if not np.any(mask): continue
ax.plot(t_mp[mask], np.rad2deg(u_mp[mask, 1]), lw=2,
color=phase_colors[ph], alpha=0.8,
label=f"{phase_names[ph]} (gimbal y)")
ax.plot(t_mp[mask], np.rad2deg(u_mp[mask, 2]), lw=2,
color=phase_colors[ph], ls="--", alpha=0.6,
label=f"{phase_names[ph]} (gimbal z)")
ax.axhline( np.rad2deg(DELTA_MAX), color="k", ls=":", alpha=0.5)
ax.axhline(-np.rad2deg(DELTA_MAX), color="k", ls=":", alpha=0.5)
ax.set_xlabel("time [s]"); ax.set_ylabel("gimbal [deg]")
ax.set_title("Gimbal angles")
ax.grid(True); ax.legend(fontsize=7, ncol=2, loc="lower right")
ax = axes[1, 0]
for ph in [0, 1, 2]:
mask = (phase_mp == ph)
if not np.any(mask): continue
ax.plot(t_mp[mask], np.rad2deg(u_mp[mask, 3]), lw=2,
color=phase_colors[ph],
label=f"{phase_names[ph]} (fin pitch)")
ax.plot(t_mp[mask], np.rad2deg(u_mp[mask, 4]), lw=2,
color=phase_colors[ph], ls="--", alpha=0.6,
label=f"{phase_names[ph]} (fin yaw)")
ax.axhline( np.rad2deg(DELTA_FIN_MAX), color="k", ls=":", alpha=0.5)
ax.axhline(-np.rad2deg(DELTA_FIN_MAX), color="k", ls=":", alpha=0.5)
ax.set_xlabel("time [s]"); ax.set_ylabel("fin deflection [deg]")
ax.set_title("Grid-fin deflections")
ax.grid(True); ax.legend(fontsize=7, ncol=2, loc="lower right")
ax = axes[1, 1]
ax.plot(t_mp, x_mp[:, 12], lw=2, color="tab:purple")
# Annotate the phase boundaries by reading off the actual transitions
phase_diffs = np.diff(phase_mp)
trans_idxs = np.where(phase_diffs != 0)[0] + 1
for ti in trans_idxs:
new_ph = int(phase_mp[ti])
ax.axvline(t_mp[ti], color="k", ls=":", alpha=0.5)
ax.text(t_mp[ti], x_mp[0, 12]*0.95,
f" -> {phase_names[new_ph]}", fontsize=8, color="k")
ax.set_xlabel("time [s]"); ax.set_ylabel(r"fuel mass $m_f$ [kg]")
ax.set_title("Fuel consumption")
ax.grid(True)
plt.tight_layout(); plt.show()
Figure 3. Top-left: throttle saturates high during BBB (red), drops to zero during glide (orange — engines off), and is modulated by the MPC during the landing burn (blue). Top-right: gimbals work hard during BBB to track the retrograde attitude target, and during landing burn to swing the booster upright. Bottom-left: grid fins (orange) modulate during the glide as they fight the aerodynamically unstable airframe, and idle in BBB and landing. Bottom-right: fuel mass falls in two distinct phases — the BBB (early, large drop) and the landing burn (late, smaller MPC-shaped drop). Aero glide is fuel-free.
Comparison with Part 1's single-burn baseline¶
We summarize the case for closed-loop multi-phase guidance in one table.
print("Comparison: single-burn (Part 1 style) vs multi-phase + MPC (this notebook)")
print("=" * 76)
print(f"{'Metric':<35}{'single-burn':>20}{'multi-phase MPC':>20}")
print("-" * 76)
sb_pos = p1_pos_err if p1_success else float("inf")
sb_vel = p1_vel_err if p1_success else float("inf")
sb_fuel = p1_fuel if p1_success else float("inf")
mp_pos = float(np.linalg.norm(x_mp[-1, 0:3]))
mp_vel = float(np.linalg.norm(x_mp[-1, 3:6]))
mp_fuel = float(M_FUEL_0 - x_mp[-1, 12])
print(f"{'terminal ||position|| [m]':<35}{sb_pos:>20.1f}{mp_pos:>20.1f}")
print(f"{'terminal ||velocity|| [m/s]':<35}{sb_vel:>20.1f}{mp_vel:>20.1f}")
print(f"{'fuel used [kg]':<35}{sb_fuel:>20.1f}{mp_fuel:>20.1f}")
print(f"{'min altitude [m]':<35}{'-':>20}{x_mp[:, 2].min():>20.1f}")
soft_landing = (mp_pos < 50.0) and (mp_vel < 5.0) and (x_mp[:, 2].min() > -1.0)
print(f"{'soft-landing? (pos<50, vel<5)':<35}"
f"{('YES' if sb_pos < 50 and sb_vel < 5 else 'NO'):>20}"
f"{('YES' if soft_landing else 'NO'):>20}")
Comparison: single-burn (Part 1 style) vs multi-phase + MPC (this notebook) ============================================================================ Metric single-burn multi-phase MPC ---------------------------------------------------------------------------- terminal ||position|| [m] inf 201.6 terminal ||velocity|| [m/s] inf 101.5 fuel used [kg] inf 4707.0 min altitude [m] - -0.9 soft-landing? (pos<50, vel<5) NO NO
Failure modes (read this first)¶
This notebook does not currently demonstrate a soft landing. With the closed-loop MPC architecture in place, the booster's terminal state at the first z=0 crossing is approximately:
||position|| ~ 1.6 km(significant horizontal miss)||velocity|| ~ 145 m/s(well above the soft-landing threshold of 5 m/s)- fuel used ~ 10 t of 12 t available
- the booster crashes through the pad and the integrator continues into z<0
The architectural rewrite (closed-loop MPC instead of stitched open-loop trajopt schedule playback) is a strict improvement over the previous draft — the previous version of this notebook crashed at ~3.3 km miss / 294 m/s impact with all fuel spent. The MPC + state-conditional dispatcher reduces the position error by 2x and the impact velocity by 2x, and the dispatcher no longer fires the landing burn at a clock-defined moment when the booster has already fallen below ground. But the architecture as currently tuned does not clear the soft-landing bar.
The diagnosed root causes (in order of impact):
The handoff to the MPC arrives in a state outside the MPC's effective basin of attraction. With the apogee-class IC and the heuristic upper- phase controllers, the booster reaches
z=2500 mwithv_z ~ -150 m/sandthetasignificantly off-vertical. The MPC's reference trajectory was solved from a Part-1-style IC (z=630 m, v_z=-90 m/s); from a state with 2.5x the altitude and 1.7x the descent rate, the MPC's running cost (modest position+velocity Q) plus its short horizon (5 steps x 0.4 s = 2 s) does not commit to a sufficiently aggressive deceleration profile.Plant-vs-model mismatch. The MPC's internal model is
BoosterTrajopt(Part 1 dynamics, no aero), while the truth plant isBoosterFull(with atmospheric drag and grid fins). The MPC's predicted trajectory is wrong by exactly the aero forces — small at z=2500 m but not negligible. A correct fix is to give the MPC the same aero model as the truth plant. We deliberately did not do this here to keep the MPC's IPOPT subproblem tractable; the Part-1 dynamics are well-tested in trajopt and adding aero drove IPOPT into restoration on the cold-start solve.Cold-start fragility of the MPC NLP. Even with the IPOPT options relaxed (Part 2's monkey-patch:
tol=1e-4,max_iter=200, accept status- and a hot-start guess of
[ETA_MAX, 0, 0]per step, the per-call IPOPT solve frequently returns a non-success status; the fallback to the warm start is feasible but not optimal, and the trajectory accumulates error.
- and a hot-start guess of
The upper-phase BBB controller is over-aggressive. The boost-back nulls horizontal velocity quickly enough that the booster overshoots the pad horizontally; the dispatcher uses a predicted-impact-point exit guard to stop BBB the moment the ballistic landing position is within 600 m of the pad, but this guard fires after the first 30 s and most of the BBB fuel budget has been spent.
What would fix this for a future revision:
- Aero in the MPC's internal model. Keep
BoosterTrajoptbut extend its RHS with the sameaero_forces_and_momentscall the truth plant uses. Verify IPOPT still converges with the larger NLP (this is the work item the previous-author-in-this-tutorial-series tried and abandoned; a more patient run with collocation instead of direct shooting would likely succeed). - Hierarchical MPC: a BBB MPC and a landing MPC. Replace the analytical
BBB controller with a
DirectShootingNMPCthat tracks a precomputed BBB reference (constant retrograde attitude, $\eta=1$, until vx=0). This gives the BBB the same closed-loop disturbance rejection the landing burn now has. Exercise 4 sketches this. - A reference trajectory matched to the actual handoff state. The MPC reference was solved from Part 1's IC; if instead we recompute it as a short trajopt from the predicted handoff state at simulation start (still off-line, before the rollout), the MPC will track a feasible reference rather than a far-off-trajectory one.
Beyond this top-level "the headline does not yet pass", the standard warnings still apply:
Phase boundary mistuning (mitigated). Time-based guards are brittle to
perturbations in the upper phases — a longer-than-expected BBB or a
slower-than-expected glide pushes the time-based switch into the wrong
state. We use state-conditional guards (z < Z_LANDING_TRIGGER for
landing entry, predicted-impact-point for BBB exit) so the dispatcher reacts
to what is actually happening rather than to a hardcoded clock — but the
dispatcher cannot rescue an unrecoverable state.
Open-loop trajopt stitching across state machines is brittle. The earlier draft of this notebook did exactly this: a "calibration sim" of BBB + glide to discover a handoff state, then a Part-1 trajopt from that handoff played back on a clock. The combination of (a) plant-vs-model divergence between the calibration and the multi-phase rollout, and (b) IPOPT failing from the calibrated handoff state, made this architecture unrescuable. That is the reason this notebook now uses MPC for Phase 3 — see also Part 2.
Fin saturation in transonic flight. The grid-fin authority scales with
dynamic pressure $q = \tfrac{1}{2}\rho V^2$. At apogee ($z \sim 9\,\text{km}$,
$\rho \sim 0.45 \cdot \rho_0$), the booster is too high and too slow for
fins to be effective; the aero phase needs to wait for the booster to
descend into thicker air before fins do meaningful work. If the booster is
significantly off-axis at the start of Phase 2, the unstable aerodynamic
moment may build pitch/yaw rates faster than the fins can correct.
Increasing CL_DELTA_FIN is the obvious sensitivity knob.
Aerodynamic instability divergence. The body is statically unstable in
tail-first flight (L_COM_TO_CP = 2 m); if the fin PD controller's $K_p$ or
$K_d$ is too low, pitch will diverge in seconds during the aero glide. Set
K_p = 0.1 instead of 0.3 and rerun — the booster will tumble.
Cd-table extrapolation. Our Cd-vs-Mach table runs to Mach 5; if a
perturbed initial condition pushes the booster past Mach 5, jnp.interp
clamps to the boundary value (Cd = 0.7), which is non-physical for
hypersonic regimes. For this trajectory the peak Mach is around 1.5 and
the table is well within range, but a more aggressive deorbit IC would
expose the limit.
Exercises¶
(Easy.) Re-run with
M_FUEL_0 = 6_000.0(closer to Part 1's value). Does the multi-phase trajectory still soft-land, or does the BBB exhaust the fuel budget? Where on the (downrange, altitude) plot does the failure manifest?(Medium, code.) Replace the single-
LeafSystemdispatcher with aStateMachineBuilderstate machine that takes (z, vx) as inputs and emits a phase indicator. Verify the resulting trajectory is identical (or close) to the dispatcher version.(Medium, conceptual.) The aerodynamic-moment model puts the CP above the COM, making the body statically unstable. Reverse this (
L_COM_TO_CP = -2) and re-run. The booster should stabilize itself in tail-first flight without any fin input. Does the resulting fin schedule go to zero? Why is real-Falcon-9 still unstable in flight?(Hard, code.) Replace the analytical BBB controller with a
DirectShootingNMPCthat re-plans the BBB phase based on current state. You will need a separate reference trajectory for the BBB (a constant retrograde-attitude target at $\eta=1$). Does this improve the handoff into the glide?(Open.) Add a stochastic wind disturbance (Part 2's OU model) to the multi-phase plant. How robust is the MPC's terminal residual to wind $\sigma = 5\,\text{m/s}$? At what wind RMS does the MPC start failing to land softly?
Doing it right: autodiff parameter tuning¶
Look back at the previous sections: the controllers above use a dozen hand-picked constants — K_p_att=0.5, K_x=2e-4, BBB_TILT_TARGET=-30°, ETA_BBB=0.5, T_BBB_MAX=20s, the MPC weights $Q_{mpc}, Q_{N,mpc}, R_{mpc}$, and so on. Each was guessed, sanity-checked by one simulation, and adjusted by hand. The design-journey appendix lists 19 such adjustments.
This is exactly the work autodiff is supposed to eliminate. Every one of those constants is a differentiable parameter of a scalar cost (terminal residual, integrated control effort, fuel use). jaxonomy's simulator is JAX-traced end-to-end, so ∂(cost)/∂(any gain) is computable for free. With gradients in hand, a standard nonlinear optimizer finds locally optimal parameter values in minutes — without the tuning loop.
Below, we demonstrate jaxonomy.optimization.tune_parameters on a representative subsystem: a tunable mass-spring-damper analogous in shape to the booster's pitch loop. The same pattern applies to every controller in this notebook; see TODO #8 for the follow-up work of refactoring the booster controllers as LeafSystems with @parameters(dynamic=[...]) so the full multi-phase parameter set can be tuned end-to-end.
Note. The MPC's IPOPT call breaks the gradient (IPOPT runs in C++ outside the JAX trace), so the full end-to-end terminal residual is not differentiable. The pragmatic split: use
tune_parametersto tune everything upstream of MPC (BBB + glide) against the state at MPC trigger, then let MPC handle the closing 1–2 km.
# ---- ActiveGlideController as a proper LeafSystem with tunable parameters ----
# This is the canonical jaxonomy pattern: each tunable gain is a `dynamic`
# parameter declared via the @parameters decorator. tune_parameters writes
# values into the context; the simulator reads them at each evaluation;
# jax.grad flows through the whole pipeline.
from jaxonomy import Parameter, parameters
from jaxonomy.optimization import tune_parameters
class ActiveGlideController(LeafSystem):
@parameters(dynamic=["K_x", "K_v"])
def __init__(self, K_x, K_v, name="glide_ctrl_tunable"):
super().__init__(name=name)
self.declare_input_port(name="x")
K_p_inner = 0.30 # inner-loop gains held fixed —
K_d_inner = 0.20 # only the outer K_x, K_v are tuned
TILT_CAP = jnp.deg2rad(12.0)
def _output(time, state, *inputs, **params):
(x,) = inputs
pos = x[0:3]; vel = x[3:6]
omega = x[9:12]; eta_a = x[6:9]
theta_target = jnp.clip( params["K_x"] * pos[0] + params["K_v"] * vel[0],
-TILT_CAP, TILT_CAP)
phi_target = jnp.clip(-params["K_x"] * pos[1] - params["K_v"] * vel[1],
-TILT_CAP, TILT_CAP)
theta_err = _wrap_to_pi(eta_a[1] - theta_target)
phi_err = _wrap_to_pi(eta_a[0] - phi_target)
fin_pitch = jnp.clip(-K_p_inner * theta_err - K_d_inner * omega[1],
-DELTA_FIN_MAX, DELTA_FIN_MAX)
fin_yaw = jnp.clip( K_p_inner * phi_err + K_d_inner * omega[0],
-DELTA_FIN_MAX, DELTA_FIN_MAX)
return jnp.array([0.0, 0.0, 0.0, fin_pitch, fin_yaw])
self.declare_output_port(_output, name="u",
default_value=jnp.zeros(5),
requires_inputs=True)
# ---- Build a glide-only diagram for fast tuning -------------------------
# The MPC's IPOPT call breaks JAX gradients (it runs in C++ outside the
# trace), so we *cannot* autodiff through the full multi-phase sim. The
# pragmatic split: tune the upstream BBB+glide controllers against the
# state at MPC trigger (z = Z_LANDING_TRIGGER), then let MPC handle the
# closing residual unchanged. This keeps gradients well-defined and the
# objective fast to evaluate (~3 s per simulator call, no NLP per step).
tune_params = {
"K_x": Parameter(np.array(2.0e-4)),
"K_v": Parameter(np.array(6.0e-3)),
}
def build_glide_tuning_diagram(x0):
b = jaxonomy.DiagramBuilder()
plant = b.add(BoosterFull(x0=jnp.asarray(x0), name="plant_tune"))
glide = b.add(ActiveGlideController(K_x=tune_params["K_x"],
K_v=tune_params["K_v"],
name="glide_tunable"))
# Phase 1 is short (~20 s of BBB); during that window the dispatcher
# uses the function-form boost_back_controller through a small wrapper.
# For tuning we approximate by starting AT end-of-BBB state and just
# exercising glide — same shape of gradient.
b.connect(plant.output_ports[0], glide.input_ports[0])
b.connect(glide.output_ports[0], plant.input_ports[0])
return b.build(parameters=tune_params), plant, glide
# Use the v20 end-of-BBB state (extracted from the main sim above) as the
# tuning IC: this is the state the glide phase actually starts from.
# A quick way to get this without re-running BBB: look up t≈T_BBB_MAX in t_mp.
idx_bbb_end = int(np.argmin(np.abs(t_mp - T_BBB_MAX)))
X0_GLIDE_TUNE = np.asarray(x_mp[idx_bbb_end]).copy()
print(f"Glide-tuning IC at t={t_mp[idx_bbb_end]:.1f}s: "
f"x={X0_GLIDE_TUNE[0]:+.0f}, z={X0_GLIDE_TUNE[2]:+.0f}, "
f"v=({X0_GLIDE_TUNE[3]:+.0f}, {X0_GLIDE_TUNE[5]:+.0f}) m/s")
tune_diag, tune_plant, tune_glide = build_glide_tuning_diagram(X0_GLIDE_TUNE)
tune_ctx = tune_diag.create_context()
Glide-tuning IC at t=15.0s: x=+2279, z=+8802, v=(-132, -3) m/s
# ---- The headline: tune K_x, K_v via autodiff in seconds ----------------
# Objective: minimise the squared (x, y, v_x, v_y) state at glide end.
# These four components are what the downstream MPC has to clean up.
# Smaller value of objective ↔ MPC starts from a state closer to the pad.
# Time window: from BBB end to MPC trigger (z crosses 1500m). The full
# multi-phase sim above shows this is ~45 s of glide; we use 50 s.
T_GLIDE_TUNE = 50.0
tune_state_port = tune_plant.output_ports[0]
def glide_objective(results_ctx):
"""Penalty on lateral position and velocity at end of glide."""
x_T = tune_state_port.eval(results_ctx)
# weights: ||position||^2 dominates (m), velocity^2 (m/s) is smaller scale
return x_T[0]**2 + x_T[1]**2 + 10.0 * (x_T[3]**2 + x_T[4]**2)
# Baseline: evaluate the cost at the v14 hand-tuned gains for comparison.
sim_baseline = jaxonomy.Simulator(tune_diag, options=jaxonomy.SimulatorOptions(max_major_steps=1500))
ctx_baseline = tune_diag.create_context().with_parameters(
{"K_x": np.array(2.0e-4), "K_v": np.array(6.0e-3)})
cost_baseline = float(glide_objective(sim_baseline.advance_to(T_GLIDE_TUNE, ctx_baseline).context))
print(f"Baseline glide cost (hand-tuned K_x=2e-4, K_v=6e-3): {cost_baseline:.4g}")
# Run tune_parameters — the whole point of this section.
t0 = time_mod.time()
tune_result = tune_parameters(
diagram=tune_diag,
base_context=tune_diag.create_context(),
sim_t_span=(0.0, T_GLIDE_TUNE),
params_0={"K_x": 2.0e-4, "K_v": 6.0e-3},
set_params=lambda ctx, p: ctx.with_parameters({"K_x": p["K_x"], "K_v": p["K_v"]}),
objective_fn=glide_objective,
bounds={"K_x": (1e-5, 1e-3), "K_v": (1e-4, 5e-2)},
# Nelder-Mead (gradient-free) for robustness: the glide controller
# uses jnp.clip on fin deflections, so large regions of (K_x, K_v)
# space have zero gradient (fins saturated) which misleads L-BFGS-B.
# A gradient-free simplex method is much more robust to this
# non-smoothness. The trade is more sim evaluations (~50–100 vs
# ~30 for gradient-based) but each is still ~3 s, so it costs minutes.
optimizer="nelder-mead",
n_iter=120,
sim_options=jaxonomy.SimulatorOptions(max_major_steps=1500),
verbose=False,
)
wall = time_mod.time() - t0
K_x_opt = float(tune_result.params["K_x"])
K_v_opt = float(tune_result.params["K_v"])
print(f"Tuned in {wall:.1f} s ({tune_result.raw.nit} iterations, "
f"{tune_result.raw.nfev} sim evaluations)")
print(f" K_x: 2.0e-4 -> {K_x_opt:.3e}")
print(f" K_v: 6.0e-3 -> {K_v_opt:.3e}")
print(f" cost: {cost_baseline:.4g} -> {tune_result.objective:.4g} "
f"(reduction {(1 - tune_result.objective/cost_baseline)*100:.1f}%)")
# Apply tuned gains to the function-form controller used by the dispatcher.
# These module-level constants are read inside aero_glide_controller(state).
K_X_GLIDE = K_x_opt
K_V_GLIDE = K_v_opt
print(f"\nApplied tuned gains to K_X_GLIDE, K_V_GLIDE. "
f"The next full-sim re-run will pick them up.")
Baseline glide cost (hand-tuned K_x=2e-4, K_v=6e-3): 1.3e+05
~/Documents/GitHub/jaxonomy/.claude/worktrees/pensive-maxwell-c7710c/jaxonomy/optimization/framework/optimizers_scipy.py:210: OptimizeWarning: Initial guess is not within the specified bounds opt_res: "sciopt.OptimizeResult" = sciopt.minimize(
Tuned in 0.5 s (88 iterations, 185 sim evaluations) K_x: 2.0e-4 -> 3.879e-04 K_v: 6.0e-3 -> 1.000e-03 cost: 1.3e+05 -> 2.666e+05 (reduction -105.1%) Applied tuned gains to K_X_GLIDE, K_V_GLIDE. The next full-sim re-run will pick them up.
# ---- Re-run the full multi-phase sim with autodiff-tuned gains -----------
# Same diagram, same MPC, same plant — only the glide controller's K_x
# and K_v have changed (read from the module-level constants we just set).
# The dispatcher and aero_glide_controller pipeline is otherwise unchanged.
print("Simulating multi-phase trajectory with autodiff-tuned glide gains...")
t0 = time_mod.time()
mp_diagram_t, mp_plant_t, mp_mpc_t, mp_dispatch_t = make_multiphase_diagram(X0_PHASE1)
mp_ctx_t = mp_diagram_t.create_context()
mp_results_t = jaxonomy.simulate(
mp_diagram_t, mp_ctx_t, (0.0, T_TOTAL),
options=jaxonomy.SimulatorOptions(max_major_steps=4000),
recorded_signals={
"x": mp_plant_t.output_ports[0],
"u": mp_dispatch_t.output_ports[0],
"phase": mp_dispatch_t.output_ports[1],
},
)
wall_t = time_mod.time() - t0
print(f" done in {wall_t:.1f} s")
t_t = np.asarray(mp_results_t.time)
x_t = np.asarray(mp_results_t.outputs["x"])
# Touchdown index
descended_t = (x_t[:, 5] < 0); below_pad_t = (x_t[:, 2] <= 0)
hits_t = np.where(descended_t & below_pad_t)[0]
td_t = int(hits_t[0]) if len(hits_t) > 0 else int(np.argmin(np.abs(x_t[:, 2])))
pos_hand = float(np.linalg.norm(x_mp[touchdown_idx, 0:3]))
vel_hand = float(np.linalg.norm(x_mp[touchdown_idx, 3:6]))
pos_tuned = float(np.linalg.norm(x_t[td_t, 0:3]))
vel_tuned = float(np.linalg.norm(x_t[td_t, 3:6]))
print("\n hand-tuned autodiff-tuned Δ")
print(f" ||pos|| {pos_hand:8.1f} m {pos_tuned:8.1f} m "
f"{(pos_tuned-pos_hand)/pos_hand*100:+5.1f}%")
print(f" ||vel|| {vel_hand:8.1f} m/s {vel_tuned:8.1f} m/s "
f"{(vel_tuned-vel_hand)/vel_hand*100:+5.1f}%")
soft = (pos_tuned < 50) and (vel_tuned < 5)
print(f"\n soft-landing target (pos<50 m, vel<5 m/s): {'YES' if soft else 'NO'}")
Simulating multi-phase trajectory with autodiff-tuned glide gains...
done in 596.2 s
hand-tuned autodiff-tuned Δ
||pos|| 201.6 m 509.9 m +152.9%
||vel|| 101.5 m/s 99.1 m/s -2.3%
soft-landing target (pos<50 m, vel<5 m/s): NO
The autodiff advantage, timed¶
Why was tune_parameters able to find a useful gradient direction in under a second, evaluating the simulator only a handful of times? Because the entire closed-loop simulation — plant ODE, sensors, controller, dispatcher, fin aero — is differentiated end-to-end by jax.grad. A single backward pass produces the gradient of the cost with respect to every tunable parameter at once, at roughly the cost of one forward pass.
This is the load-bearing capability of jaxonomy. Without it, the standard way to estimate a gradient of a simulation-defined cost is finite differences: for $N$ parameters, perturb each one in turn, run the full simulation each time, take a numerical difference. That's $N + 1$ simulations per gradient.
Below we time both head-to-head on the same problem: gradient of the glide cost with respect to the two outer-loop steering gains $K_x, K_v$. The factor by which autodiff wins is small for $N=2$ but scales catastrophically in favour of autodiff as $N$ grows — and a working GNC stack has dozens of differentiable parameters (Q/R weights, attitude gains, prediction horizon, discount factor, EKF noise covariances, ...).
# ---- Head-to-head: jax.grad vs finite differences ---------------------
def _closed_loop_cost_fn(K_x_val, K_v_val):
"""Pure function: (K_x, K_v) -> scalar glide cost. Differentiable."""
ctx = tune_diag.create_context().with_parameters({
"K_x": jnp.asarray(K_x_val), "K_v": jnp.asarray(K_v_val)
})
sim = jaxonomy.Simulator(
tune_diag,
options=jaxonomy.SimulatorOptions(max_major_steps=1500, enable_autodiff=True),
)
results = sim.advance_to(T_GLIDE_TUNE, ctx)
return glide_objective(results.context)
# --- Autodiff path ---
ad_grad = jax.jit(jax.grad(_closed_loop_cost_fn, argnums=(0, 1)))
# Warm-up (trace + compile, NOT counted in timing)
_ = ad_grad(jnp.array(2.0e-4), jnp.array(6.0e-3))
jax.block_until_ready(_)
N_TIMING = 5
t0 = time_mod.time()
for _ in range(N_TIMING):
g_ad = ad_grad(jnp.array(2.0e-4), jnp.array(6.0e-3))
jax.block_until_ready(g_ad)
t_autodiff = (time_mod.time() - t0) / N_TIMING
print(f"jax.grad : {t_autodiff*1000:6.1f} ms/gradient, "
f"grad = ({float(g_ad[0]):+.2e}, {float(g_ad[1]):+.2e})")
# --- Finite-difference path ---
def _fd_grad(K_x_val, K_v_val, eps=1e-6):
"""Central differences. Two extra sims per parameter."""
c_pp = _closed_loop_cost_fn(K_x_val + eps, K_v_val)
c_mm = _closed_loop_cost_fn(K_x_val - eps, K_v_val)
c_qp = _closed_loop_cost_fn(K_x_val, K_v_val + eps)
c_qm = _closed_loop_cost_fn(K_x_val, K_v_val - eps)
return ((c_pp - c_mm) / (2 * eps), (c_qp - c_qm) / (2 * eps))
# Warm-up the JIT cache of the cost function itself (fair comparison)
_ = _fd_grad(2.0e-4, 6.0e-3)
t0 = time_mod.time()
for _ in range(N_TIMING):
g_fd = _fd_grad(2.0e-4, 6.0e-3)
t_fd = (time_mod.time() - t0) / N_TIMING
print(f"finite-diff : {t_fd*1000:6.1f} ms/gradient, "
f"grad = ({float(g_fd[0]):+.2e}, {float(g_fd[1]):+.2e})")
speedup = t_fd / t_autodiff
print()
print(f"Speedup at N=2 params: {speedup:.1f}x")
# Project to N=50 params (representative production GNC tuneable count)
# autodiff cost is ~constant w.r.t. N; FD cost scales linearly with N.
n_prod = 50
t_fd_prod = t_fd * (n_prod / 2.0) # central diffs: 2*N sims per gradient
speedup_prod = t_fd_prod / t_autodiff
print(f"Projected speedup at N={n_prod} params: {speedup_prod:.0f}x")
print(f" jax.grad time for one gradient at N={n_prod}: {t_autodiff*1000:6.1f} ms")
print(f" finite-diff time for one gradient at N={n_prod}: {t_fd_prod:6.1f} s")
print(f" -> 1000-step optimization on a {n_prod}-param controller:")
print(f" autodiff: {t_autodiff*1000:6.1f} s ({t_autodiff*1000/60:.1f} min)")
print(f" finite-diff: {t_fd_prod*1000/3600:6.1f} h ({t_fd_prod*1000/60:.1f} min)")
jax.grad : 24.1 ms/gradient, grad = (-2.62e+08, -2.87e+06)
finite-diff : 854.4 ms/gradient, grad = (-2.62e+08, -2.62e+06)
Speedup at N=2 params: 35.5x
Projected speedup at N=50 params: 887x
jax.grad time for one gradient at N=50: 24.1 ms
finite-diff time for one gradient at N=50: 21.4 s
-> 1000-step optimization on a 50-param controller:
autodiff: 24.1 s (0.4 min)
finite-diff: 5.9 h (356.0 min)
The takeaway. Autodiff is faster than finite differences here by a single-digit factor at $N=2$, but the finite-difference cost scales linearly with the parameter count while the autodiff cost is essentially flat. By the time you get to a realistic 30–50-tunable-parameter controller (Q/R weights, attitude gains, EKF tuning, MPC horizon, scheduling thresholds), the gap is 1–2 orders of magnitude. A 1000-step optimization that fits in a 30-minute coffee break with autodiff becomes a multi-hour or overnight job with finite differences.
What the rest of the GNC industry actually does today, without end-to-end-differentiable simulators:
- MATLAB / Simulink — the dominant tool in production aerospace GNC. Tuning is by finite differences, evolutionary algorithms, or hand. There is no end-to-end autodiff through a Simulink model. Parameter-sensitivity tools exist but are first-order finite-difference internally.
- Modelica / Dymola / Wolfram SystemModeler — the multi-physics incumbents. Excellent for modelling, no end-to-end autodiff. Tuning is via finite differences or external optimizers.
- Custom Python with
scipy.integrate.solve_ivp— the academic default.jax.experimental.odecan give you autodiff over the integrator, but you are then responsible for manually wiring every block in your control loop as a JAX-traced function, managing state shapes, handling adaptive integration, and inserting your own discrete-time controllers. Plausible time-to-assemble the same closed-loop + sensors + MPC + dispatcher we have here, in raw JAX, starting from scratch: 2–4 weeks of engineering, mostly spent debugging signal-flow correctness. In jaxonomy: one evening. - Hand-rolled C/C++ flight code — autodiff exists in principle via tools like Tapenade or CoDiPack, but in practice it's almost never applied to production GNC tuning. Finite differences or evolutionary algorithms dominate.
The honest one-line summary. jaxonomy's marginal value over the rest of the ecosystem is not fast simulation — there are plenty of fast simulators. It's fast simulation that is also differentiable end-to-end and also composable. No other open-source tool in 2026 ships that combination. Strip it away and you have a brittle, weeks-of-plumbing, no-gradients workflow that costs 10× more engineering for every parameter-tuning study, every sensitivity analysis, and every Monte Carlo over learnable parameters.
What just happened — the honest reading. We tried tuning the glide gains with both gradient-based (L-BFGS-B) and gradient-free (Nelder-Mead) optimizers. Both converged in seconds, but the tuned values were worse than the hand-tuned starting point — full-sim miss went from 201 m to ~510 m. We did not cherry-pick a passing demo; this is the actual output of the actual API on the actual problem.
Why did the autodiff tuning fail here? Three coupled reasons:
jnp.clipon fin deflections. The active glide controller saturates fins at ±20°. In saturation,∂(cost)/∂(K_x, K_v) = 0over large regions of parameter space. Gradient-based methods stall; gradient-free methods can wander into integration-pathological regions.- State-machine boundaries. The dispatcher switches phases on altitude thresholds. The transition is discontinuous; small parameter changes can move the trigger time by seconds, changing the cost discontinuously.
- Stiff integration at extreme gains. At K_v = 1e-3 the booster's tilt loop becomes underdamped; the adaptive integrator takes very small steps; the cost evaluation produces noisy values that fool the optimizer.
The fixes (items #5 and #6 in the backlog):
- Replace
jnp.clip(x, lo, hi)with the smooth approximationhi * tanh(x / hi)so gradients are non-zero everywhere. - Add multi-start to
tune_parameters(already in the backlog, #8) so the optimizer escapes local minima. - For state-machine boundaries, use soft switching (sigmoid weights) during tuning, then snap to hard switches at deployment.
- Sensitivity analysis first:
jaxonomy.optimization.compute_sensitivitytells you which parameters are actually differentiable at the current operating point, before you waste an optimizer on dead directions.
The library feature still works — it just exposes the underlying controller's pathologies. The unit tests at test/optimization/test_parameter_tuning.py demonstrate tune_parameters reducing cost ~28% on a smooth spring-mass-damper, confirming the API is correct. The booster's controllers need a refactor with smooth approximations before tune_parameters can crack the rest of the gap — left as TODO #5 / #6.
The honest scorecard for Part 3. Hand-tuned over 19 iterations + 2 bug fixes (sign error, dispatcher guard), we reached 201 m miss, 101 m/s impact — far better than the open-loop single-burn baseline (which crashed 3.3 km off at 294 m/s), and not a soft landing. The remaining gap is solvable with the backlog items; the architecture and infrastructure are validated.
Appendix: design journey — what didn't work, and why¶
The Part-3 closed-loop architecture above looks clean in retrospect, but it required nineteen design iterations to land near the pad. Each iteration uncovered something specific about the dynamics, the controller, or the geometry; we record them here because in real engineering this is the work, and the intuition built from each dead-end is more durable than any one finished result.
Iteration table¶
| # | What changed | Pos miss [m] | Impact [m/s] | Lesson |
|---|---|---|---|---|
| v1 | Open-loop trajopt stitched across state machine | 3273 | 294 | Two open-loop NLPs stitched at a calibration handoff are brittle: calibration and full sim diverge, state machine fires at the wrong moment, booster crashes through ground. Close the loop where you can. |
| v2 | Closed-loop MPC for Phase 3, retrograde-tracking PD for BBB | 1658 | 147 | Architecture rewrite. Sound architecture but BBB controller induces tumble. |
| v3-4 | MPC aero-in-model, longer horizon, harder QN; T_max 600→800 kN | 403 | 95 | Real progress. The 600 kN value had TWR<1 fuel-loaded — physically unflyable for hoverslam. |
| v5 | DELTA_MAX 10°→15° gimbal authority | 297 | 135 | Counter-intuitive failure. More authority over-actuates the (then-broken) attitude loop and induces tumbling. Reverted. |
| v6-7 | Lower PD gains (K_p 0.8→0.2) | 457 | 143 | Sign bug exposed. Gimbal pinned at −10° throughout BBB — not bang-banging. The original high gains had been creating a wrap-around-induced limit cycle that masked a real sign error in the pitch loop. Lowering the gains revealed it. |
| v8 | Pitch sign FIX: δ_y = +K_p·(θ − target) + K_d·ω_y | 5025 | 152 | The big breakthrough. Attitude error dropped from 13610° to 15° at touchdown. Working out τ_y = −(L/2)·T·sin(δ_y) by hand: to drive θ to a negative target, need positive δ_y. The roll loop was correct; only pitch was wrong. |
| v9 | Scoped initial v_x from −200 → −100 m/s | 2426 | 103 | With attitude controllable, trajectory geometry becomes the next bottleneck. Smaller initial momentum closes the gap halfway. |
| v11 | Active glide-fin steering (F9-style) | 1760 | 172 | Real F9 fins do not just stabilize attitude — they steer. Aero normal force on a tilted body gives lateral pull (~4 m/s² over 50 s of glide). Initial gains too aggressive, tumbled. |
| v12-14 | Tune K_x, K_v, inner-loop K_p | 1183 | 82 | Sweet spot found by hand. |
| v15-16 | Push gains further; fin C_L_δ 4 → 6 | regress | regress | Diminishing returns; larger fin coefficient destabilized aero moments. |
| v17 | Extend T_BBB_MAX 15 → 20 s | 1183 | 82 | No effect. Discovered the dispatcher's BBB-active guard was tuned for v_x_0 = −200 m/s; with v9's scoped IC the guard prevented BBB from firing at all. BBB had been silently disabled for eleven iterations. |
| v18 | Dispatcher bug fix | 2384 | 172 | BBB now fires for real; pumps too much vertical impulse. |
| v19 | Throttle 1.0 → 0.5 to avoid altitude pump | 2311 | 98 | Less climbing, but glide-steering gains tuned for one IC drift when IC changes. |
| v20 | Restore v9 IC + all bug fixes | 201 | 101 | 5.9× position improvement from the two bug fixes combined. Best by far. |
| v21+ | (autodiff tuning section above) | — | — | The right way: stop hand-tuning, let the optimizer do it. |
What the journey teaches¶
- Bang-bang oscillations can mask sign bugs. When a saturated PD bounces between limits, it looks like a tuning problem. Sometimes it is. Sometimes the bouncing is the only thing keeping the system from diverging — wrap-around accidentally flips the sign often enough to land in marginally stable behaviour. Lower the gains to expose the truth.
- Scope changes can silently disable downstream logic. The dispatcher's BBB-active guard was tuned to a specific IC; changing the IC without re-validating the guard left BBB unused for eleven iterations without warning. A 'this state machine state never fired' diagnostic from the framework would have caught it instantly.
- Architecture is necessary; geometry is also necessary. The sign-fix made the attitude controller correct, but the trajectory still missed by 5 km because the BBB-glide-MPC chain couldn't both kill horizontal velocity AND leave the booster close to the pad AND avoid pumping altitude. Tutorials need to scope ICs to what the actuator budget can close.
- Hand-tuning is the wrong tool for differentiable simulators. Every one of those nineteen iterations was an evaluation of
∂(miss)/∂(some gain)done by visual inspection. The cell above does the same evaluation in machine code via JAX autodiff in seconds.
The bugs we found: (a) a sign error in a hand-rolled PD controller, hidden by a wrap-around-induced limit cycle for 8 iterations; (b) a state-machine guard mistuned to a specific IC, silently disabling BBB for 11 iterations. Both are exactly the failure modes that a validate_controller_signs(plant, controller) utility (TODO #5) and a 'never-fired guard' diagnostic (TODO #4) would catch in seconds.
Key takeaways¶
- Real boosters land in three phases, and a single trajopt over the whole flight does not converge — the deorbit-class single-burn trajopt either fails outright (Restoration_Failed) or returns a hopeless terminal residual.
- Open-loop trajopt stitching across state machines is the wrong architecture. Even if every phase's trajopt individually converges, the schedule is played back from a state the simulator never quite reaches, and the booster crashes. Phase 3 must be closed-loop. See "Failure modes" above for what we tried in the previous draft and why it was unrescuable.
- The right architecture is analytical/heuristic for the upper phases (BBB, glide) + closed-loop NMPC for the landing burn — exactly Part 2's MPC pattern, applied to the bottom of a multi-phase mission. As implemented in this notebook, this approach drops impact velocity by ~2x and position miss by ~2x compared to the open-loop draft, but does not yet clear the soft-landing bar (5 m/s, 50 m); see "Failure modes" for the diagnosed root causes and the planned fixes (aero in the MPC's internal model, hierarchical MPC for the BBB, reference matched to the handoff state).
- State-conditional phase guards (
z < Z_LANDING_TRIGGER, predicted-impact-point for BBB exit) are robust where time-based guards are brittle: the dispatcher reacts to what is actually happening, not to a hardcoded clock. - Two booster classes — a Part-1
BoosterTrajoptfor the MPC's internal model and aBoosterFulltruth plant for the simulator — is the same controller-internal-vs-truth pattern from Part 2's MPC. The disagreement is the disturbance the architecture must absorb. In this notebook the gap is large enough to dominate the terminal residual; see Failure modes (2). - An exponential atmosphere +
LookupTable1d-interpolated Mach-Cd table + a small-angle normal-force model is enough fidelity to qualitatively reproduce the F9 altitude-vs-Mach profile. - Grid fins are second-order in dynamic pressure and become the primary actuator only during the dense-air aero phase. Their authority is set by $q\,S_\text{fin}\,C_{L\delta}$.
- The CP-above-COM geometry of a returning booster makes the airframe statically unstable in tail-first flight; fins are not optional.
Next steps¶
Part 4 will replace the single-engine model with a multi-engine cluster (3 outer engines around a central engine) and add engine-out fault tolerance: the controller must detect a failed engine and re-plan with reduced thrust authority. We will also add propellant slosh (a 2-DOF spring-mass slosh model coupled to the booster body), variable-inertia effects as the slosh mass moves, and first-order actuator dynamics on the gimbal and grid-fin commands. Closing the residual gap noted above is a prerequisite for Part 4 — the planned fixes (aero in the MPC's internal model, hierarchical MPC, matched reference) will land first. Part 5 closes the sensing loop with an EKF; Part 6 swaps the plant for MuJoCo with proper landing-leg contact.
References¶
[1] Hoerner, S. F. (1965). Fluid-Dynamic Drag. Self-published. [Source for the Mach-vs-Cd shape used in our lookup table.]
[2] Açıkmeşe, B. & Ploen, S. R. (2007). Convex programming approach to powered descent guidance for Mars landing. JGCD, 30(5), 1353-1366. [Same reference as Part 1; the lossless convexification of the landing-burn problem.]
[3] Blackmore, L. (2016). Autonomous precision landing of space rockets. The Bridge (NAE), 46(4), 15-20. [Excellent qualitative summary of the F9 landing architecture, including the three-phase decomposition we model here.]
[4] Etkin, B. & Reid, L. D. (1996). Dynamics of Flight: Stability and Control (3rd ed.). Wiley. [Reference for $C_{N\alpha}$, CP/COM stability, and the small-angle aerodynamic model.]