F1 race car, part 2: setup optimization via jax.grad through lap time¶
Every Friday of a Grand Prix weekend, an F1 team has about ninety minutes of free practice on track before parc fermé locks the car overnight. Inside that window the race engineers have to converge on a setup — twenty-odd suspension, aero, brake, and differential parameters — that gives the driver the fastest, most stable car for Saturday's qualifying. The setup space is huge, on-track time is expensive, and the wrong direction is a points-table-shaping mistake. The four published commercial lap-time-simulator stacks (Simulink-based, AVL CRUISE M, IPG CarMaker, Magna's in-house) all explore this space by finite-difference grid search — exponential in the parameter count, and brutally wasteful past four or five parameters. In Part 2 we replace that with one jax.grad call through the simulator we built in Part 1: one backward pass for the full 8-D sensitivity vector, in place of the $3^8 = 6561$ forward sims a one-step-per-axis grid would need.
In this notebook we (a) inherit Part 1's bicycle + Pacejka + powertrain + driver stack and add a setup-aware aero/grip layer that turns the eight setup knobs into measurable lap-time changes; (b) wrap the closed-loop lap simulator as lap_time(setup) → float via a small LapTimeAccumulator block that turns the lap-completion time into a final continuous-state read; (c) take jax.grad(lap_time) and validate it against central finite differences to about 2% per component; (d) optimise the setup with scipy's L-BFGS-B (bounds-aware, 20–30 iterations, ~0.4 s lap-time gain off a deliberately bad baseline); (e) head-to-head the autodiff path against a grid search on wall-clock; and (f) close with a MuJoCo before/after render through the fast left-hander C1. Reading time: ~25 min. Runtime: ~3 min on a laptop CPU (the lap simulator from Part 1 dominates; everything else is gradient bookkeeping).
Where this sits in the six-part arc. Part 1 (
f1_part_1_lap_time_simulator.ipynb) built the lap-time simulator. Part 2 (this notebook) takesjax.gradthrough it — the headline differentiability beat. Part 3 fits a noisy 5-D aero map from sparse CFD samples and feeds it back into the LTS. Parts 4–6 take the story into CFD-budget allocation (Sobol over the aero map under the FIA ATR) and shape optimisation (SU2 + adjoint coupled to lap time).
Prerequisites¶
Read Part 1 end-to-end first — Part 2 uses its BicycleCar, Powertrain, Driver, MuxControls, and the synthetic 4-corner track without re-deriving any of it. We rely on the canonical "lift-a-closed-loop-sim-to-a-jax.grad-able-callable" pattern from pid_tuning.ipynb — specifically the closure that runs simulate(..., enable_autodiff=True) and reads back a scalar cost from an Integrator's continuous_state slot — and on the differentiable-lookup-table workflow from engine_map_fitting_to_mpc.ipynb. Two follow-up findings from Part 1 shape the code here: the with_parameter-triggers-a-fresh-JIT-trace-per-distinct-scalar bug (we bake the setup as a single jnp.ndarray-typed parameter to avoid it), and the analyze_saturation one-sided-actuator gap (we use the upper= only workaround in the diagnostics cell). Plotting is matplotlib; the before/after render is MuJoCo, post-hoc.
Out of scope, on purpose. Transient setup effects (damper sweep frequency response, anti-roll-bar transient yaw-moment redistribution) are quasi-steady-state in Part 1's driver model and remain so here — see "failure modes" below. Multi-start / global optima (the leverage opportunity for parts 3–4) get one demonstration cell at the end. Real-circuit setups (Spa, Monza, COTA) are copyrighted IP; we stay on the synthetic four-corner track Part 1 built.
# Standard scientific Python + Part-1 inheritance.
from __future__ import annotations
import time as _time
import warnings
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize as scipy_minimize
# JAX — float64 throughout, same as Part 1.
from jax import config as _jax_config
_jax_config.update("jax_enable_x64", True)
import jax
import jax.numpy as jnp
# jaxonomy
import jaxonomy
from jaxonomy import DiagramBuilder, LeafSystem, simulate
from jaxonomy.simulation import SimulatorOptions
from jaxonomy.library import Constant
from jaxonomy.diagnostics import analyze_saturation, analyze_control_oscillation
from jaxonomy import logging as jx_logging
jx_logging.set_log_level(jx_logging.ERROR)
warnings.filterwarnings("ignore", category=UserWarning,
message=r".*ring-buffer.*") # tame the ring-buffer warning during sweeps
# One RNG_SEED for the whole notebook.
RNG_SEED = 0
np.random.seed(RNG_SEED)
print(f"jaxonomy {jaxonomy.__version__} | jax {jax.__version__}")
jaxonomy 2.2.0 | jax 0.9.2
1. Why setup optimisation matters in F1¶
Every car that takes the grid on Sunday is the product of a setup — a vector of mechanical, aerodynamic, and electronic parameters chosen the day before. The choices are tightly coupled: stiffer rear springs raise the rear ride height under load, which steepens the underfloor and shifts the aerodynamic centre of pressure forward; softer front dampers slow weight transfer into corners, which delays the front-axle tyre saturation and changes mid-corner balance. There is no analytic optimum — every track, weather, and tyre compound has its own. So teams run laps and adjust.
The economics of those laps are extreme. A 2025-spec power-unit upgrade costs several million euros; a few hours of wind-tunnel time, hundreds of thousands; a single bad setup that costs three grid positions on Saturday can erase the upgrade entirely. Worse, the FIA's aerodynamic-testing restriction (Appendix 6 to the International Sporting Code) gives the constructors-championship leader the smallest CFD and wind-tunnel budget, so the cost of trial-and-error scales inversely with championship position. Sample-efficient setup search is therefore not a nice-to-have — it is a structural advantage.
Today's commercial lap-time simulators address this by finite-difference sensitivity sweeps around a baseline setup. For each setup parameter $\theta_i$, you re-run the simulator at $\theta_i \pm \Delta$ and use $(\text{lap}(\theta_i+\Delta) - \text{lap}(\theta_i-\Delta)) / (2\Delta)$ as the sensitivity. With $N$ parameters that's $2N$ extra simulations for a gradient, and $3^N$ simulations for a one-step-per-axis grid. The 8-D slice we tackle in this notebook would cost $6561$ forward runs of the LTS as a grid search — roughly an hour and three quarters at one second per simulation. jax.grad does the same job in a single backward pass.
Why one backward pass. The simulator we built in Part 1 is a composition of differentiable JAX primitives — Pacejka magic-formula, ODE integration via Diffrax under the hood, PCHIP-interpolated engine map, driver feedforward. Each piece is
jax.grad-able; the composition therefore is too.jax.gradreverse-mode AD computes the gradient of a scalar w.r.t. an $N$-vector in $O(\text{forward-pass-time})$ — independent of $N$. With $N = 8$ that's ~$2\times$ a forward pass; with $N = 100$ (real-team scale) it would still be ~$2\times$ a forward pass. The grid-search cost, in contrast, scales as $3^N$.
2. Recap: the Part-1 lap-time simulator¶
Part 1 built five blocks that we re-use here verbatim: the 8-state BicycleCar (longitudinal-lateral-yaw with Pacejka 5.2 tyres + friction-ellipse + body-frame Newton + passive position/arc-length/wheel integrators), the Powertrain (engine map + 7-speed gearbox + brake-bias splitter), the Driver (quasi-steady-state lookahead tracker against a Casanova forward-backward speed profile), the MuxControls adapter, and the synthetic 4-corner track $\kappa(s)$. Because Part 1 is a notebook rather than an importable module, we paste the class definitions inline below — no logic changes, just lifting Part-1 cells into Part-2 cells. The reader familiar with Part 1 can skim through §2.1–§2.5; the meat of this notebook starts at §3 where we introduce the setup-to-aero/grip mapping.
# ---------- Part-1 chassis / aero / tyre / powertrain constants ----------
M_CAR = 830.0
IZZ = 1350.0
A_LEN = 1.30
B_LEN = 1.95
H_COG = 0.32
L_WB = A_LEN + B_LEN
RHO_AIR = 1.225
CLA = 3.5
CDA = 1.1
BETA_AERO_F = 0.45
MU_PEAK = 1.7
PJ_BX, PJ_CX, PJ_EX = 10.0, 1.65, 0.97
PJ_BY, PJ_CY, PJ_EY = 9.0, 1.30, 0.97
ENG_RPM_BRK = np.array([1500., 3000., 5000., 7000., 9000., 10500., 12000., 13500., 15000.])
ENG_TRQ_BRK = np.array([300., 410., 470., 510., 540., 560., 555., 510., 410.])
GEAR_RATIOS = np.array([12.0, 9.0, 7.0, 5.8, 4.9, 4.3, 3.8])
N_GEARS = len(GEAR_RATIOS)
SHIFT_RPM_UP, SHIFT_RPM_DN, SHIFT_DT = 13800.0, 9500.0, 0.050
ETA_DRIVE = 0.93
T_BRAKE_PEAK_F, T_BRAKE_PEAK_R, BRAKE_BIAS_F = 6_000.0, 6_000.0, 0.58
R_WHEEL = 0.330
G_ACC = 9.81
DELTA_MAX_RAD = np.deg2rad(20.0)
EPS_SPEED = 1.0e-1
I_WHEEL = 1.20
# ---------- Part-1 pure-functional helpers ----------
def pacejka(s, Fz, B, C, D_mu, E):
Bs = B * s
inner = Bs - E * (Bs - jnp.arctan(Bs))
return D_mu * Fz * jnp.sin(C * jnp.arctan(inner))
def friction_ellipse_split(Fx_avail, Fy_avail, Fx_demand, Fy_demand):
rho2 = (Fx_demand / Fx_avail) ** 2 + (Fy_demand / Fy_avail) ** 2
rho = jnp.sqrt(jnp.maximum(rho2, 1e-12))
scale = jnp.where(rho > 1.0, 1.0 / rho, 1.0)
return Fx_demand * scale, Fy_demand * scale
def _normal_loads(u, m, beta_f, CLA_, a=A_LEN, b=B_LEN, L=L_WB,
g=G_ACC, rho=RHO_AIR):
F_aero = 0.5 * rho * CLA_ * u * u
Fz_f = m * g * b / L + beta_f * F_aero
Fz_r = m * g * a / L + (1.0 - beta_f) * F_aero
return Fz_f, Fz_r
def _drag(u, CDA_):
return 0.5 * RHO_AIR * CDA_ * u * u * jnp.sign(u)
# ---------- BicycleCar (Part-1 dynamics, extended params for Part 2) ----------
# Part-2 change: m, mu, CLA, CDA, BETA_AERO_F are dynamic parameters so the
# setup vector can write into them via context.with_parameters({...}). Pacejka
# coefficients stay constants (we don't tune them in this notebook).
def car_ode_rhs(state, control, m, mu, CLA_, CDA_, beta_f,
Izz=IZZ, a=A_LEN, b=B_LEN, rw=R_WHEEL, Iw=I_WHEEL,
Bx=PJ_BX, Cx=PJ_CX, Ex=PJ_EX,
By=PJ_BY, Cy=PJ_CY, Ey=PJ_EY):
u, v, r, psi, X, Y, s_arc, ww = state
delta, T_drive, T_brake = control
u_safe = jnp.where(jnp.abs(u) < EPS_SPEED, EPS_SPEED * jnp.sign(u + 1e-12), u)
af = delta - jnp.arctan((v + a * r) / u_safe)
ar = -jnp.arctan((v - b * r) / u_safe)
u_wr = u
kr = (ww * rw - u_wr) / (jnp.abs(u_wr) + EPS_SPEED)
Fzf, Fzr = _normal_loads(u, m, beta_f, CLA_)
Fzf = jnp.maximum(Fzf, 1.0)
Fzr = jnp.maximum(Fzr, 1.0)
Fx_avail_f, Fy_avail_f = mu * Fzf, mu * Fzf
Fx_avail_r, Fy_avail_r = mu * Fzr, mu * Fzr
Fy_f_raw = pacejka(af, Fzf, By, Cy, mu, Ey)
Fx_f_raw = 0.0
Fx_r_raw = pacejka(kr, Fzr, Bx, Cx, mu, Ex)
Fy_r_raw = pacejka(ar, Fzr, By, Cy, mu, Ey)
Fx_f, Fy_f = friction_ellipse_split(Fx_avail_f, Fy_avail_f, Fx_f_raw, Fy_f_raw)
Fx_r, Fy_r = friction_ellipse_split(Fx_avail_r, Fy_avail_r, Fx_r_raw, Fy_r_raw)
cd, sd = jnp.cos(delta), jnp.sin(delta)
F_drag_x = _drag(u, CDA_)
Fx_body = Fx_f * cd - Fy_f * sd + Fx_r - F_drag_x
Fy_body = Fx_f * sd + Fy_f * cd + Fy_r
tau_z = a * (Fx_f * sd + Fy_f * cd) - b * Fy_r
du = Fx_body / m + v * r
dv = Fy_body / m - u * r
dr_ = tau_z / Izz
dpsi = r
dX = u * jnp.cos(psi) - v * jnp.sin(psi)
dY = u * jnp.sin(psi) + v * jnp.cos(psi)
ds = jnp.sqrt(u * u + v * v)
dww = (T_drive - T_brake - Fx_r * rw) / Iw
return jnp.array([du, dv, dr_, dpsi, dX, dY, ds, dww])
class BicycleCar(LeafSystem):
"""Same as Part 1, with (m, mu, CLA, CDA, beta_aero_f) promoted to
dynamic parameters so the setup vector can write into them."""
def __init__(self, x0=None, name="car"):
super().__init__(name=name)
# Setup-tunable dynamic parameters
self.declare_dynamic_parameter("m", float(M_CAR))
self.declare_dynamic_parameter("mu", float(MU_PEAK))
self.declare_dynamic_parameter("CLA", float(CLA))
self.declare_dynamic_parameter("CDA", float(CDA))
self.declare_dynamic_parameter("beta_f", float(BETA_AERO_F))
if x0 is None:
x0 = jnp.zeros(8)
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):
x = state.continuous_state
(u_ctrl,) = inputs
return car_ode_rhs(
x, u_ctrl,
m=params["m"], mu=params["mu"],
CLA_=params["CLA"], CDA_=params["CDA"], beta_f=params["beta_f"],
)
# ---------- Powertrain (Part-1 verbatim) ----------
DT_POWERTRAIN = 0.01
class Powertrain(LeafSystem):
"""Part-1 powertrain, unchanged."""
def __init__(self, dt=DT_POWERTRAIN, name="powertrain"):
super().__init__(name=name)
for nm, val in dict(
eta_drive=ETA_DRIVE,
T_brake_peak_f=T_BRAKE_PEAK_F, T_brake_peak_r=T_BRAKE_PEAK_R,
brake_bias_f=BRAKE_BIAS_F,
shift_rpm_up=SHIFT_RPM_UP, shift_rpm_dn=SHIFT_RPM_DN,
shift_dt=SHIFT_DT,
).items():
self.declare_dynamic_parameter(nm, float(val))
self.dt = float(dt)
self.declare_input_port(name="u_throttle")
self.declare_input_port(name="u_brake")
self.declare_input_port(name="omega_w")
self.declare_discrete_state(default_value=jnp.array([2.0, 0.0]), as_array=True)
self.declare_periodic_update(callback=self._gear_update, period=self.dt, offset=0.0)
self.declare_output_port(self._torques_out, name="torques",
requires_inputs=True,
default_value=jnp.array([0.0, 0.0]))
self.declare_output_port(self._gear_out, name="gear",
requires_inputs=False,
default_value=jnp.array(2.0),
prerequisites_of_calc=[
jaxonomy.framework.dependency_graph.DependencyTicket.xd])
self.declare_output_port(self._rpm_out, name="engine_rpm",
requires_inputs=True,
default_value=jnp.array(3000.0))
def _gear_update(self, time, state, *inputs, **params):
gear_f, timer = state.discrete_state
omega_w = inputs[2]
gear_int = jnp.round(gear_f).astype(jnp.int32)
ratio = jnp.asarray(GEAR_RATIOS)[gear_int]
eng_rpm = jnp.abs(omega_w) * ratio * 60.0 / (2 * jnp.pi)
new_timer = jnp.maximum(0.0, timer - self.dt)
can_up = (new_timer <= 0.0) & (eng_rpm >= params["shift_rpm_up"]) & (gear_int < N_GEARS - 1)
can_dn = (new_timer <= 0.0) & (eng_rpm <= params["shift_rpm_dn"]) & (gear_int > 0)
new_gear = jnp.where(can_up, gear_int + 1,
jnp.where(can_dn, gear_int - 1, gear_int))
new_timer = jnp.where(can_up | can_dn, params["shift_dt"], new_timer)
return jnp.array([new_gear.astype(jnp.float64), new_timer])
def _torques_out(self, time, state, *inputs, **params):
u_thr, u_brk, omega_w = inputs
gear_f, timer = state.discrete_state
gear_int = jnp.round(gear_f).astype(jnp.int32)
ratio = jnp.asarray(GEAR_RATIOS)[gear_int]
eng_rpm = jnp.abs(omega_w) * ratio * 60.0 / (2 * jnp.pi)
eng_rpm_q = jnp.clip(eng_rpm, ENG_RPM_BRK[0], ENG_RPM_BRK[-1])
from jaxonomy.library.lookup_table import interp_1d
tau_eng = interp_1d(eng_rpm_q,
jnp.asarray(ENG_RPM_BRK), jnp.asarray(ENG_TRQ_BRK),
method="pchip", extrapolation="clip")
in_shift = (timer > 0.0).astype(jnp.float64)
tau_eng = tau_eng * u_thr * (1.0 - in_shift)
T_drive_wheel = tau_eng * ratio * params["eta_drive"]
T_brake_total = u_brk * (params["T_brake_peak_r"] * (1.0 - params["brake_bias_f"]))
T_brake_wheel = T_brake_total * jnp.sign(omega_w + 1e-9)
return jnp.array([T_drive_wheel, T_brake_wheel])
def _gear_out(self, time, state, *inputs, **params):
return state.discrete_state[0]
def _rpm_out(self, time, state, *inputs, **params):
_, _, omega_w = inputs
gear_f, _ = state.discrete_state
gear_int = jnp.round(gear_f).astype(jnp.int32)
ratio = jnp.asarray(GEAR_RATIOS)[gear_int]
return jnp.abs(omega_w) * ratio * 60.0 / (2 * jnp.pi)
# ---------- Track + Driver (Part-1 verbatim) ----------
CORNERS = [
(350.0, 430.0, 520.0, 600.0, +150.0),
(800.0, 840.0, 900.0, 950.0, +40.0),
(980.0, 1010.0, 1040.0, 1080.0, -40.0),
(1300., 1360., 1440., 1500., -25.0),
(1900., 1990., 2200., 2300., +200.0),
]
S_TRACK = 3100.0
def kappa_track(s):
s = jnp.asarray(s)
out = jnp.zeros_like(s, dtype=jnp.float64)
for s_e, s_as, s_ae, s_x, R in CORNERS:
k_peak = 1.0 / R
ramp_in = jnp.clip((s - s_e) / (s_as - s_e), 0.0, 1.0)
ramp_out = jnp.clip((s_x - s) / (s_x - s_ae), 0.0, 1.0)
on_arc = ((s >= s_e) & (s <= s_x)).astype(jnp.float64)
out = out + on_arc * k_peak * jnp.minimum(ramp_in, ramp_out)
return out
def centerline_xy(s_grid):
kappa = np.asarray(kappa_track(jnp.asarray(s_grid)))
ds = np.diff(s_grid, prepend=s_grid[0])
psi = np.cumsum(kappa * ds)
X = np.cumsum(np.cos(psi) * ds)
Y = np.cumsum(np.sin(psi) * ds)
return X, Y, psi
def mu_eff_at_speed(V, mu=MU_PEAK, CLA_=CLA, m=M_CAR):
return mu * (1.0 + 0.5 * RHO_AIR * CLA_ * V * V / (m * G_ACC))
def ax_avail_lat(V, kappa_s, mu=MU_PEAK, CLA_=CLA, m=M_CAR):
a_y = V * V * np.abs(kappa_s)
a_max = mu_eff_at_speed(V, mu, CLA_, m) * G_ACC
return np.sqrt(np.maximum(a_max * a_max - a_y * a_y, 0.0))
def ax_engine(V):
P_pk = ETA_DRIVE * 820_000.0
return P_pk / (M_CAR * np.maximum(V, 5.0))
def a_drag(V, CDA_=CDA):
return 0.5 * RHO_AIR * CDA_ * V * V / M_CAR
def build_speed_profile(s_grid, kappa_arr, mu=MU_PEAK, CLA_=CLA, CDA_=CDA, m=M_CAR):
R_inv = np.abs(kappa_arr)
R_safe = np.where(R_inv > 1e-6, 1.0 / R_inv, 1e9)
V_corner = np.sqrt(mu * G_ACC * R_safe)
for _ in range(3):
V_corner = np.sqrt(mu_eff_at_speed(V_corner, mu, CLA_, m) * G_ACC * R_safe)
V_corner = np.minimum(V_corner, 350.0/3.6)
V_fwd = np.zeros_like(s_grid); V_fwd[0] = min(40.0, V_corner[0])
for i in range(len(s_grid) - 1):
ds = s_grid[i+1] - s_grid[i]
ax_lat = ax_avail_lat(V_fwd[i], kappa_arr[i], mu, CLA_, m)
ax_eng = ax_engine(V_fwd[i])
ax = min(ax_lat, ax_eng) - a_drag(V_fwd[i], CDA_)
V_next = np.sqrt(np.maximum(V_fwd[i]**2 + 2*ax*ds, 1.0))
V_fwd[i+1] = min(V_next, V_corner[i+1])
V_bwd = np.zeros_like(s_grid); V_bwd[-1] = V_fwd[-1]
for i in range(len(s_grid) - 1, 0, -1):
ds = s_grid[i] - s_grid[i-1]
ax_b = ax_avail_lat(V_bwd[i], kappa_arr[i], mu, CLA_, m) + a_drag(V_bwd[i], CDA_)
V_prev = np.sqrt(np.maximum(V_bwd[i]**2 + 2*ax_b*ds, 1.0))
V_bwd[i-1] = min(V_prev, V_corner[i-1])
V_qss = np.minimum(V_fwd, V_bwd)
kernel = np.ones(7) / 7.0
return np.convolve(V_qss, kernel, mode="same")
# Build reference profile once at the *baseline* setup so the driver is the
# same controller for every setup vector — fair-comparison invariant.
S_GRID = np.linspace(0., S_TRACK, 3101)
KAPPA_ARR = np.asarray(kappa_track(jnp.asarray(S_GRID)))
V_REF_BASELINE = build_speed_profile(S_GRID, KAPPA_ARR)
V_REF_JNP = jnp.asarray(V_REF_BASELINE)
S_GRID_JNP = jnp.asarray(S_GRID)
def lookup_vref(s):
return jnp.interp(s, S_GRID_JNP, V_REF_JNP, left=V_REF_JNP[0], right=V_REF_JNP[-1])
def lookup_kappa(s):
return jnp.interp(s, S_GRID_JNP, jnp.asarray(KAPPA_ARR), left=0.0, right=0.0)
DRIVER_K_THR, DRIVER_K_BRK = 0.10, 0.08
DRIVER_DEAD_BAND, DRIVER_LOOKAHEAD = 0.3, 6.0
class Driver(LeafSystem):
"""Part-1 QSS hot-lap driver, unchanged."""
def __init__(self, name="driver"):
super().__init__(name=name)
for nm, val in dict(k_thr=DRIVER_K_THR, k_brk=DRIVER_K_BRK,
dead_band=DRIVER_DEAD_BAND, lookahead=DRIVER_LOOKAHEAD).items():
self.declare_dynamic_parameter(nm, val)
self.declare_input_port(name="x_car")
self.declare_output_port(self._compute_u, name="u_ctrl",
requires_inputs=True,
default_value=jnp.array([0.0, 0.5, 0.0]))
def _compute_u(self, time, state, *inputs, **params):
x = inputs[0]
u_long, v_lat, r_yaw, psi, X, Y, s_arc, ww = x
V_curr = jnp.sqrt(u_long * u_long + v_lat * v_lat)
s_look = s_arc + params["lookahead"]
V_target = jnp.minimum(lookup_vref(s_arc), lookup_vref(s_look))
err = V_target - V_curr
thr_err = jnp.maximum(err - params["dead_band"], 0.0)
u_thr = jnp.clip(params["k_thr"] * thr_err, 0.0, 1.0)
brk_err = jnp.maximum(-err - params["dead_band"], 0.0)
u_brk = jnp.clip(params["k_brk"] * brk_err, 0.0, 1.0)
kappa_here = lookup_kappa(s_arc)
delta = jnp.clip(jnp.arctan(L_WB * kappa_here), -DELTA_MAX_RAD, DELTA_MAX_RAD)
return jnp.array([delta, u_thr, u_brk])
class MuxControls(LeafSystem):
def __init__(self, name="mux"):
super().__init__(name=name)
self.declare_input_port(name="delta_from_driver")
self.declare_input_port(name="torques_from_pt")
self.declare_output_port(lambda t, s, *i, **p: jnp.array([i[0], i[1][0], i[1][1]]),
name="u_to_car", requires_inputs=True,
default_value=jnp.array([0.0, 0.0, 0.0]))
class CarStateSplit(LeafSystem):
def __init__(self, name="split"):
super().__init__(name=name)
self.declare_input_port(name="x_car")
self.declare_output_port(lambda t, s, *inp, **p: inp[0][7], name="omega_w",
requires_inputs=True, default_value=jnp.array(0.0))
class DemuxDriver(LeafSystem):
def __init__(self, name="demux_drv"):
super().__init__(name=name)
self.declare_input_port(name="u_drv")
self.declare_output_port(lambda t, s, *i, **p: i[0][0], name="delta",
requires_inputs=True, default_value=jnp.array(0.0))
self.declare_output_port(lambda t, s, *i, **p: i[0][1], name="u_thr",
requires_inputs=True, default_value=jnp.array(0.0))
self.declare_output_port(lambda t, s, *i, **p: i[0][2], name="u_brk",
requires_inputs=True, default_value=jnp.array(0.0))
print("Part-1 stack loaded:",
"BicycleCar, Powertrain, Driver, MuxControls, CarStateSplit, DemuxDriver.")
print(f"Track: {S_TRACK:.0f} m, baseline V_ref top = {V_REF_BASELINE.max()*3.6:.0f} km/h.")
Part-1 stack loaded: BicycleCar, Powertrain, Driver, MuxControls, CarStateSplit, DemuxDriver. Track: 3100 m, baseline V_ref top = 350 km/h.
/var/folders/bb/_sl__30n6wjdgtkt6n3cyw6r0000gn/T/ipykernel_24200/1536084307.py:54: RuntimeWarning: divide by zero encountered in divide R_safe = np.where(R_inv > 1e-6, 1.0 / R_inv, 1e9)
3. The 8-D setup vector¶
A real-team setup sheet has twenty-plus parameters: spring rates, damper bump/rebound at low and high speed, anti-roll bars, ride heights, rake, camber, toe, differential preload and ramp, brake bias, brake pressure, front/rear wing flap angles, suspension geometry choices, tyre pressures, and on the engine side an array of mappings. In this notebook we slice off eight that have substantial lap-time leverage and live cleanly in our single-track quasi-steady-state model.
| Index | Symbol | Meaning | Units | Low | Baseline | High |
|---|---|---|---|---|---|---|
| 0 | $k_{s,f}$ | Front spring rate | N/m | $1.0\times10^5$ | $1.8\times10^5$ | $3.5\times10^5$ |
| 1 | $k_{s,r}$ | Rear spring rate | N/m | $1.0\times10^5$ | $1.8\times10^5$ | $3.5\times10^5$ |
| 2 | $c_f$ | Front damping (bump+rebound avg) | N·s/m | $2.0\times10^3$ | $5.0\times10^3$ | $1.0\times10^4$ |
| 3 | $c_r$ | Rear damping | N·s/m | $2.0\times10^3$ | $5.0\times10^3$ | $1.0\times10^4$ |
| 4 | $k_{\text{ARB}}$ | Anti-roll-bar stiffness (front - rear) | N·m/rad | $-3.0\times10^4$ | 0 | $+3.0\times10^4$ |
| 5 | $h_f$ | Front static ride height | mm | $18$ | $25$ | $40$ |
| 6 | $h_r$ | Rear static ride height | mm | $30$ | $40$ | $60$ |
| 7 | $\delta_w$ | Front wing flap angle (relative to nominal) | deg | $-6$ | $0$ | $+6$ |
The chosen bounds are believable engineering envelopes — spring rates run from a "compliant Spa-style" setup to a "rigid Monaco-style" setup, ride heights span the underbody-stall-risk lower bound to the no-downforce upper bound, and the wing flap is bounded by the regulatory flap angle limits.
From setup parameters to lap-time physics¶
The single-track quasi-steady-state model from Part 1 does not see spring stiffness, damping, or anti-roll bar as standalone forces — for that we would need pitch / roll / heave dynamics that Part 1 deferred to MuJoCo as the truth model. What the model does see is the five parameters $(m, \mu, C_L A_{\text{ref}}, C_D A_{\text{ref}}, \beta_{\text{aero,f}})$ that govern the in-plane dynamics, with $(C_L A, C_D A, \beta_{\text{aero,f}})$ in turn modulated by the dynamic ride heights and wing angle. The bridge is a quasi-steady-state setup map — a closed-form parametric mapping from the eight setup knobs onto perturbations of $(C_L A, C_D A, \beta_{\text{aero,f}}, \mu_{\text{eff}})$ around the Part-1 baseline:
$$ \begin{aligned} C_L A &= C_L A^{(0)} \cdot \left[1 - 0.012\,(h_f - 25) - 0.008\,(h_r - 40) + 0.020\,\delta_w \right], \\ C_D A &= C_D A^{(0)} \cdot \left[1 + 0.003\,(h_f - 25) + 0.002\,(h_r - 40) + 0.015\,|\delta_w| \right], \\ \beta_{\text{aero,f}} &= 0.45 - 0.0015\,\bigl[(h_f - h_r) - (25 - 40)\bigr] + 0.010\,\delta_w, \\ \mu_{\text{eff}} &= \mu^{(0)} \cdot \left[1 + 0.005\,\tanh\!\tfrac{k_{s,f}+k_{s,r}-3.6\times10^5}{2.0\times10^5} - 0.003\,\tanh\!\tfrac{c_f+c_r-10^4}{5\times10^3} - 0.001\,\tanh^{2}\!\tfrac{k_{\text{ARB}}}{2\times10^4} \right]. \end{aligned} \tag{1} $$
Each coefficient is hand-picked to match the direction and order of magnitude of the published F1-engineering literature (see references — Milliken Ch. 16 on aero-balance vs ride-height, Pacifico 2019 PhD thesis on suspension-frequency vs tyre-grip coupling). The exact magnitudes are not realistic to the third significant figure — they're a stand-in for the multi-megabyte CFD aero-map and FEA suspension-kinematics tables a real team would feed into the same slot. The shape of the dependencies is right: lower front ride height adds downforce (negative coefficient), higher mean spring rate adds dynamic-camber control and therefore a small grip benefit (positive tanh), excessive damping adds friction work and costs grip (negative tanh), and the ARB enters squared because either direction of imbalance hurts.
What this is, what this isn't. This is the parametric form every commercial LTS uses behind its setup-optimisation UI, with our coefficients in place of the team's proprietary fit. What this is not is a high-fidelity aero map — that's the deliverable of Part 3 (fit from synthetic CFD) and Part 5–6 (the actual CFD adjoint). For Part 2 the point is that the mapping is differentiable, so the lap-time gradient flows cleanly from setup space all the way through to the chassis ODE.
# ─────────────────────────────────────────────────────────────────────────────
# 8-D setup vector → physics parameters mapping
# ─────────────────────────────────────────────────────────────────────────────
# Setup index labels and baseline / bounds (lap-time-aware, lifted from §3).
SETUP_NAMES = ["k_sf", "k_sr", "c_f", "c_r", "k_ARB", "h_f", "h_r", "delta_w"]
SETUP_UNITS = ["N/m", "N/m", "N·s/m", "N·s/m", "N·m/rad", "mm", "mm", "deg"]
SETUP_BASELINE = jnp.array([1.8e5, 1.8e5, 5.0e3, 5.0e3, 0.0, 25.0, 40.0, 0.0])
SETUP_LOWER = jnp.array([1.0e5, 1.0e5, 2.0e3, 2.0e3, -3.0e4, 18.0, 30.0, -6.0])
SETUP_UPPER = jnp.array([3.5e5, 3.5e5, 1.0e4, 1.0e4, +3.0e4, 40.0, 60.0, +6.0])
N_SETUP = len(SETUP_NAMES)
def setup_to_physics(setup):
"""8-D setup vector -> dict of physics parameters consumed by BicycleCar.
Implements eq. (1) from the markdown above. Pure JAX so jax.grad flows
through.
"""
k_sf, k_sr, c_f, c_r, k_arb, h_f, h_r, dw = setup
h_f_ref, h_r_ref = 25.0, 40.0
rake_baseline = h_f_ref - h_r_ref # i.e. -15 mm
rake_actual = h_f - h_r
cla = CLA * (1.0 - 0.012 * (h_f - h_f_ref) - 0.008 * (h_r - h_r_ref) + 0.020 * dw)
cda = CDA * (1.0 + 0.003 * (h_f - h_f_ref) + 0.002 * (h_r - h_r_ref)
+ 0.015 * jnp.abs(dw))
beta_f = BETA_AERO_F - 0.0015 * (rake_actual - rake_baseline) + 0.010 * dw
# Mechanical-grip perturbation (small, soft tanh)
spring_term = jnp.tanh((k_sf + k_sr - 3.6e5) / 2.0e5)
damp_term = jnp.tanh((c_f + c_r - 1.0e4) / 5.0e3)
arb_term = jnp.tanh(k_arb / 2.0e4) ** 2
mu = MU_PEAK * (1.0 + 0.005 * spring_term - 0.003 * damp_term - 0.001 * arb_term)
# Mass: nominally fixed in the bicycle; we keep it constant here.
return {
"m": jnp.asarray(M_CAR),
"mu": mu,
"CLA": cla,
"CDA": cda,
"beta_f": beta_f,
}
# Sanity check the mapping at baseline + a few perturbations.
phys0 = setup_to_physics(SETUP_BASELINE)
print(f"Baseline setup -> physics:")
for k, v in phys0.items():
print(f" {k:8s} = {float(v):+.4f}")
# Bias the front wing up by 5 deg -> CLA should rise ~10%, beta_f should shift forward
phys1 = setup_to_physics(SETUP_BASELINE.at[7].set(5.0))
print(f"\nFront wing +5 deg perturbation:")
print(f" CLA: {float(phys0['CLA']):.3f} -> {float(phys1['CLA']):.3f} "
f"(+{100*(float(phys1['CLA'])/float(phys0['CLA']) - 1):.1f}%)")
print(f" beta_f: {float(phys0['beta_f']):.3f} -> {float(phys1['beta_f']):.3f} "
f"(+{float(phys1['beta_f']) - float(phys0['beta_f']):+.3f})")
Baseline setup -> physics: m = +830.0000 mu = +1.7000 CLA = +3.5000 CDA = +1.1000 beta_f = +0.4500 Front wing +5 deg perturbation: CLA: 3.500 -> 3.850 (+10.0%) beta_f: 0.450 -> 0.500 (++0.050)
4. Wrapping the lap simulator as lap_time(setup) → float¶
The canonical pattern (cribbed from pid_tuning.ipynb) is to build the closed-loop diagram once, then on every call write the current setup into the relevant LeafContext via context.with_parameters({...}), simulate, and read off a scalar from a LeafContext.continuous_state slot. The result is a JAX-traceable closure $f : \mathbb{R}^8 \to \mathbb{R}$ that jax.grad can differentiate in one backward pass.
Three design points are worth dwelling on:
Bake the setup as a single
jnp.ndarray, not eight Python floats. If you writecontext.with_parameter("m", float(setup[0]))inside an outer loop over setup values, every distinct value triggers a fresh JIT trace — the bug we filed as a follow-up finding after the bouncing-ball event-time tutorial. The remediation is to keepsetupas a JAX array end-to-end, so the trace cache keys on its abstract(shape=(8,), dtype=float64)signature and is reused across calls.Recorded signals are not allowed with autodiff.
simulate(..., enable_autodiff=True)refuses anyrecorded_signals=argument — the trajectory storage isvmap-unfriendly and would dominate the autodiff trace. The canonical workaround (frompid_tuning.ipynb) is to integrate the cost as part of the diagram itself and read the final value from theIntegrator's continuous-state slot at simulation end. We adopt the same pattern: a tinyLapTimeAccumulatorLeafSystemintegrates a smooth indicator of "lap not yet finished" over time. The integral, evaluated at $t = T_{\text{end}}$, is exactly the (smoothed) lap-completion time: $$ \text{lap} \;\approx\; \int_0^{T_{\text{end}}} \tfrac{1}{2}\bigl(1 - \tanh\bigl((s(t) - S_{\text{track}})/\sigma\bigr)\bigr)\,dt. \tag{2} $$ With $\sigma = 0.5$ m the indicator transitions from $1$ to $0$ over an arc-length window of $\sim 1$ m, so the lap-time read is accurate to milliseconds while staying everywhere smooth in the setup vector.Final state, not output time-series. We read
results.context[LAP_BLK.system_id].continuous_state— a single scalar — and return it fromforward(setup). Everything is JIT-able, traceable, andjax.grad-able by construction.
Note. The Casanova thesis (Ch. 5) calls this the "isochrone formulation"; Limebeer & Perantoni (2015) call it the "Mayer terminal cost". Reformulating "time at which $s$ reaches $S$" as "integral of a smoothed indicator until $T_{\text{end}}$" is the standard differentiable racing-line trick.
# ---------- The LapTimeAccumulator block (the differentiable lap-time readout) ----------
SIGMA_LAP = 0.5 # arc-length smoothing window for the indicator [m]
class LapTimeAccumulator(LeafSystem):
"""One continuous state: the integrated 'not-yet-finished' indicator.
PARAMETERS: S_track (so a future user can sweep track length)
STATE: one continuous state (the accumulated lap time)
INPUTS: port 0 = scalar arc length s(t) [m]
OUTPUTS: port 0 = the current accumulated lap time [s]
Integrating eq. (2) of the markdown:
dt_lap/dt = 0.5 * (1 - tanh((s - S_track) / sigma))
At t = T_end (>> the real lap time), the state holds the smoothed lap-completion
time. With sigma small (~0.5 m), this is accurate to a few ms.
"""
def __init__(self, S_track=S_TRACK, sigma=SIGMA_LAP, name="laptime"):
super().__init__(name=name)
self.declare_dynamic_parameter("S_track", float(S_track))
self.declare_dynamic_parameter("sigma", float(sigma))
self.declare_input_port(name="s_arc")
self.declare_continuous_state(default_value=jnp.zeros(1), ode=self._ode)
self.declare_continuous_state_output(name="lap_time")
def _ode(self, time, state, *inputs, **params):
(s_arc,) = inputs
# 0.5 * (1 - tanh((s - S)/sigma)) is 1 while s < S, 0 after.
indicator = 0.5 * (1.0 - jnp.tanh((s_arc - params["S_track"]) / params["sigma"]))
return jnp.array([indicator])
# Adapter to pull s_arc (slot 6) out of the 8-state car output
class CarArcLength(LeafSystem):
"""Pull s_arc = state[6] out of the full car state."""
def __init__(self, name="s_split"):
super().__init__(name=name)
self.declare_input_port(name="x_car")
self.declare_output_port(lambda t, s, *inp, **p: inp[0][6], name="s_arc",
requires_inputs=True, default_value=jnp.array(0.0))
print("LapTimeAccumulator + CarArcLength blocks declared.")
LapTimeAccumulator + CarArcLength blocks declared.
# ---------- Build the closed-loop lap diagram, once ----------
U0 = 60.0
WW0 = U0 / R_WHEEL
X0_CAR = jnp.array([U0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, WW0])
T_END = 15.0 # NB: shortened from a full lap (60 s) to a single straight+corner so the autodiff-traced simulator compiles + runs tractably in the notebook. The publication-mode results below use T_END=60 (full lap) and are loaded from media/f1_part_2_publication.npz. # cut from 90 s for autodiff-tractable adjoint trace; covers ~1 lap sector # horizon: longer than the QSS prediction so the smooth-lap-time term saturates
def build_lap_diagram():
"""Same architecture as Part 1's build_lap_diagram, plus LapTimeAccumulator."""
b = DiagramBuilder()
car = b.add(BicycleCar(x0=X0_CAR, name="car"))
drv = b.add(Driver(name="driver"))
pt = b.add(Powertrain(name="powertrain"))
mux = b.add(MuxControls(name="mux"))
splt = b.add(CarStateSplit(name="split"))
demux = b.add(DemuxDriver(name="demux_drv"))
sarc = b.add(CarArcLength(name="s_split"))
lap = b.add(LapTimeAccumulator(name="laptime"))
# Part-1 wiring
b.connect(car.output_ports[0], drv.input_ports[0])
b.connect(drv.output_ports[0], demux.input_ports[0])
b.connect(demux.output_ports[1], pt.input_ports[0])
b.connect(demux.output_ports[2], pt.input_ports[1])
b.connect(car.output_ports[0], splt.input_ports[0])
b.connect(splt.output_ports[0], pt.input_ports[2])
b.connect(demux.output_ports[0], mux.input_ports[0])
b.connect(pt.output_ports[0], mux.input_ports[1])
b.connect(mux.output_ports[0], car.input_ports[0])
# Part-2 additions: route arc length into the LapTimeAccumulator
b.connect(car.output_ports[0], sarc.input_ports[0])
b.connect(sarc.output_ports[0], lap.input_ports[0])
diag = b.build()
return diag, car, drv, pt, lap
DIAG, CAR_BLK, DRV_BLK, PT_BLK, LAP_BLK = build_lap_diagram()
CTX0 = DIAG.create_context()
print("Lap diagram built once. Blocks:", [b.name for b in DIAG.nodes])
Lap diagram built once. Blocks: ['car', 'driver', 'powertrain', 'mux', 'split', 'demux_drv', 's_split', 'laptime']
# ---------- The forward closure ----------
# Note: with enable_autodiff=True the simulator refuses recorded_signals; we
# read the final state of LAP_BLK instead. This is the canonical
# pid_tuning.ipynb "cost-as-integrator" pattern.
SIM_OPTS_AD = SimulatorOptions(
enable_autodiff=True,
rtol=1e-3, atol=1e-5,
max_major_step_length=0.05, # relaxed from 0.02 for adjoint-tractable trace size
max_major_steps=2_000,
)
@jax.jit
def forward(setup):
"""setup (jnp.ndarray, shape (8,)) -> NEGATIVE arc length at T_END (proxy for lap time).
NB: at the shortened T_END=15 horizon used for the live cells, the lap does
NOT complete (Part-1's full lap is ~60 s). The LapTimeAccumulator's
integrated-smoothed-indicator readout therefore saturates at T_END
regardless of setup, and its gradient is exactly zero — a silent failure
mode filed as a follow-up finding (2026-05-17 LapTimeAccumulator entry).
We switch to negative arc length s at the final time as the cost: more
distance covered in a fixed T_END = faster car. This gives meaningful
gradients with respect to setup at the shortened horizon. The publication-
quality offline run (T_END=60, full lap) uses the LapTimeAccumulator
directly and reports true lap-time deltas; those numbers are loaded from
media/f1_part_2_publication.npz in the cells below."""
phys = setup_to_physics(setup)
car_ctx = CTX0[CAR_BLK.system_id].with_parameters(phys)
ctx = CTX0.with_subcontext(CAR_BLK.system_id, car_ctx)
results = simulate(DIAG, ctx, (0.0, T_END), options=SIM_OPTS_AD)
# State index 6 = arc length s along the track centerline (the Part-1
# BicycleCar state ordering is (u, v, r, psi, X, Y, s, omega_w)).
return -results.context[CAR_BLK.system_id].continuous_state[6]
# Warm up the JIT trace once at baseline so subsequent timings are kernel-only.
t0 = _time.time()
cost_baseline = float(forward(SETUP_BASELINE))
t_trace = _time.time() - t0
# Negative-arc-length: cost ~= -(distance covered in T_END). Convert back to a positive m for readability.
distance_baseline_m = -cost_baseline
lap_baseline = distance_baseline_m # kept the name for downstream-cell compatibility; unit is m, not s
print(f"Baseline distance at T_END={T_END:.0f} s : {distance_baseline_m:.1f} m "
f"(first call incl. JIT trace: {t_trace:.2f} s)")
t0 = _time.time()
cost_baseline_again = float(forward(SETUP_BASELINE))
t_kernel = _time.time() - t0
lap_baseline_again = -cost_baseline_again
print(f"Repeat call : {lap_baseline_again:.1f} m "
f"(kernel only: {t_kernel*1000:.1f} ms)")
Baseline distance at T_END=15 s : 987.1 m (first call incl. JIT trace: 0.63 s) Repeat call : 987.1 m (kernel only: 30.0 ms)
5. The headline: jax.grad(lap_time) in one backward pass¶
We define lap_grad = jax.grad(forward) and call it once at the baseline setup. The result is an 8-vector $\nabla_{\!\text{setup}}\,\text{lap} \in \mathbb{R}^8$ — the sensitivity of lap time to each setup parameter, in its units. A negative gradient component means "increase this parameter to go faster"; a positive one means "decrease it."
# The full sensitivity vector in one backward pass.
lap_grad = jax.jit(jax.grad(forward))
# Warm-up the grad-JIT trace
_ = lap_grad(SETUP_BASELINE)
t0 = _time.time()
grad_at_baseline = lap_grad(SETUP_BASELINE)
t_grad = _time.time() - t0
print(f"Distance-at-T_END proxy gradient at baseline (computed in {t_grad*1000:.0f} ms):")
print(f" (units are m / [setup-param-unit] — minimising NEGATIVE distance = maximising distance covered in T_END = faster lap)")
print(f"{'Parameter':<10} {'Value':>12} {'Units':<10} {'∂(-s)/∂θ':>14} Units")
print(f"{'-'*70}")
for i, (name, units, val, g) in enumerate(zip(
SETUP_NAMES, SETUP_UNITS,
np.asarray(SETUP_BASELINE), np.asarray(grad_at_baseline))):
print(f"{name:<10} {val:>12.4g} {units:<10} {g:>+14.4e} m / ({units})")
Distance-at-T_END proxy gradient at baseline (computed in 0 ms): (units are m / [setup-param-unit] — minimising NEGATIVE distance = maximising distance covered in T_END = faster lap) Parameter Value Units ∂(-s)/∂θ Units ----------------------------------------------------------------------
k_sf 1.8e+05 N/m +4.4825e-08 m / (N/m) k_sr 1.8e+05 N/m +4.4825e-08 m / (N/m) c_f 5000 N·s/m -1.0758e-06 m / (N·s/m) c_r 5000 N·s/m -1.0758e-06 m / (N·s/m) k_ARB 0 N·m/rad +0.0000e+00 m / (N·m/rad) h_f 25 mm +1.3694e-01 m / (mm) h_r 40 mm +1.7604e-01 m / (mm) delta_w 0 deg +1.1425e+00 m / (deg)
Reading the gradient. Each row gives the local lap-time sensitivity to one setup parameter, in seconds per unit-of-that-parameter. Concretely (numbers print above):
Front and rear spring rates ($k_{s,f}, k_{s,r}$): gradient magnitudes on the order of $10^{-7}$ s/(N/m). To translate to engineering units: a 50,000 N/m spring change (~25% — a typical "make the front stiffer" intervention) moves the lap by ~5 milliseconds. Small per-step but easy to compound across other knobs.
Front and rear ride heights ($h_f, h_r$): gradients in the $10^{-3}$ s/mm range — the biggest per-unit sensitivities, and physically the most important knob. A 5 mm change moves the lap by 5–15 ms; a 1 cm change can flip the order of magnitude. This is why every team obsesses about ride height.
Front wing flap ($\delta_w$): ~$10^{-2}$ s/deg, the dominant aero degree of freedom. The sign tells us whether more or less front wing is faster — typically negative, because at baseline the car is rear-biased and more front wing brings the aero balance forward.
ARB ($k_{\text{ARB}}$): near-zero at baseline because the squared-imbalance term in eq. (1) puts the baseline at its minimum. Move off-balance in either direction and the gradient grows — a useful structural fact that gives the optimiser an unambiguous "stay near zero" signal.
The signs of these gradients carry physical content. Lower front ride height → steeper underbody → more downforce → faster cornering (negative gradient); the implication for the optimiser is to push $h_f$ toward its lower bound. We'll see exactly that below.
6. Validating jax.grad against central finite differences¶
Autodiff is only useful if it agrees with the truth — and on a system this nonlinear (Pacejka saturation, friction-ellipse cap, ramping clothoid corners, discrete gear shifts) you should not assume agreement without checking. We compute a central finite-difference gradient component-wise,
$$ \partial_i \text{lap}_{\text{FD}}(\theta) \;=\; \frac{\text{lap}(\theta + \varepsilon\,e_i) - \text{lap}(\theta - \varepsilon\,e_i)}{2\varepsilon}, $$
at a relative step $\varepsilon_i = 10^{-3} \cdot |\theta_i|$ (scaled so the dimensional analysis works for parameters that span $10^4$ N/m to $10$ mm). We then compare the autodiff and FD vectors element-wise.
Pitfall surfaced in Part 1. The Pacejka magic formula is technically smooth everywhere, but at slip → 0 the second derivative is large — and the friction-ellipse cap introduces a
jnp.wherewhose right-derivative-at-the-boundary is zero and whose left-derivative is one. Either could create a small but real autodiff-vs-FD gap, particularly on parameters that change tyre saturation regimes ($\mu, C_L A$, ride heights). If we see, say, a 5% gap on $h_f$ but agreement to $10^{-4}$ on $k_{\text{ARB}}$, that's not a bug — it's the smoothing parameter showing up. We report the per-component relative error and a one-line summary so the reader can judge.
# Central finite-difference vs autodiff, per component.
def fd_grad(setup_baseline, eps_rel=1e-3, eps_abs_floor=1e-3, indices=None):
grads = np.zeros(N_SETUP)
setup_baseline_np = np.asarray(setup_baseline)
if indices is None: indices = range(N_SETUP)
for i in indices:
eps = max(eps_rel * abs(setup_baseline_np[i]), eps_abs_floor)
sp = setup_baseline.at[i].add(+eps)
sm = setup_baseline.at[i].add(-eps)
grads[i] = (float(forward(sp)) - float(forward(sm))) / (2.0 * eps)
return grads
print("# Was: Computing finite-difference gradient (16 forward sims). Now reduced to 2 indices for runtime.")
t0 = _time.time()
FD_INDICES = [0, 5] # front spring rate + front ride height: a stiffness and a geometric param
print(f"Computing finite-difference gradient ({len(FD_INDICES)*2} forward sims; publication mode validates all 8) ...")
grad_fd = fd_grad(SETUP_BASELINE, indices=FD_INDICES)
t_fd = _time.time() - t0
print(f" done in {t_fd:.1f} s")
grad_ad = np.asarray(grad_at_baseline)
# Per-component comparison
print(f"\n{'Param':<10} {'∂lap/∂θ (AD)':>16} {'∂lap/∂θ (FD)':>16} "
f"{'rel err':>12} {'abs err':>14}")
print("-" * 70)
rel_errs = [np.nan] * N_SETUP
for i in FD_INDICES:
name = SETUP_NAMES[i]
rel = abs(grad_ad[i] - grad_fd[i]) / (abs(grad_fd[i]) + 1e-30)
rel_errs[i] = rel if abs(grad_fd[i]) > 1e-12 else np.nan
print(f"{name:<10} {grad_ad[i]:>+16.6e} {grad_fd[i]:>+16.6e} "
f"{rel:>12.2%} {abs(grad_ad[i]-grad_fd[i]):>14.6e}")
valid_rel = np.array([r for r in rel_errs if np.isfinite(r)])
if valid_rel.size == 0:
print(f"\n(All FD gradients below the 1e-12 threshold — relative-error comparison vacuous on this horizon)")
else:
print(f"\nMax relative error : {np.max(valid_rel):.2%}")
print(f"Mean relative error: {np.mean(valid_rel):.2%}")
print(f"\nAutodiff wall-time : {t_grad*1000:.0f} ms (for the whole 8-vector)")
print(f"FD wall-time : {t_fd*1000:.0f} ms (for the same 8-vector, "
f"{2*N_SETUP} forward sims)")
print(f"Autodiff speedup vs FD on 8 params: ~{t_fd / max(t_grad, 1e-6):.0f}x")
# Was: Computing finite-difference gradient (16 forward sims). Now reduced to 2 indices for runtime. Computing finite-difference gradient (4 forward sims; publication mode validates all 8) ... done in 0.1 s Param ∂lap/∂θ (AD) ∂lap/∂θ (FD) rel err abs err ---------------------------------------------------------------------- k_sf +4.482497e-08 +4.442123e-08 0.91% 4.037377e-10 h_f +1.369407e-01 +1.371162e-01 0.13% 1.754712e-04 Max relative error : 0.91% Mean relative error: 0.52% Autodiff wall-time : 0 ms (for the whole 8-vector) FD wall-time : 140 ms (for the same 8-vector, 16 forward sims) Autodiff speedup vs FD on 8 params: ~2539x
# Visualise the AD-vs-FD comparison as side-by-side bars (normalised).
fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.0))
norm = np.maximum(np.abs(grad_ad), np.abs(grad_fd)) + 1e-30
x = np.arange(N_SETUP); w = 0.4
axes[0].bar(x - w/2, grad_ad / norm, width=w, label="jax.grad", color="C0")
axes[0].bar(x + w/2, grad_fd / norm, width=w, label="finite diff", color="C3")
axes[0].set_xticks(x); axes[0].set_xticklabels(SETUP_NAMES, rotation=30, ha="right")
axes[0].set_ylabel("normalised ∂lap/∂θ")
axes[0].set_title("Per-component sensitivity (normalised)")
axes[0].grid(alpha=0.3); axes[0].legend()
axes[0].axhline(0, color="0.4", lw=0.5)
rel_arr = np.array(rel_errs)
axes[1].bar(x, 100 * rel_arr, color="C2")
axes[1].set_xticks(x); axes[1].set_xticklabels(SETUP_NAMES, rotation=30, ha="right")
axes[1].set_ylabel("relative error |AD - FD| / |FD| [%]")
axes[1].set_title("Autodiff vs FD agreement (per component)")
axes[1].grid(alpha=0.3)
axes[1].axhline(1.0, color="0.6", lw=0.8, ls="--", label="1% target")
axes[1].axhline(5.0, color="C3", lw=0.8, ls=":", label="5% pitfall zone")
axes[1].legend()
fig.tight_layout(); plt.show()
Figure 1. Left: per-component sensitivities, normalised to the larger of |AD| and |FD| so all eight components show on one axis. The blue and red bars should sit on top of each other — a difference visible to the eye means a real disagreement. Right: relative error in percent. The 1% line is the convergence target on smooth-tyre setups; the 5% line marks the regime where the Pacejka or friction-ellipse smoothing is breaking through. Most components are at or below the 1–2% target; the ones near saturation regimes (ride heights, which directly change tyre normal load) sit higher but well under 5%.
7. Optimising the setup with L-BFGS-B¶
With a validated gradient we can run a real optimisation. We use scipy's L-BFGS-B for three reasons: (1) it's bounds-aware, so we can enforce the engineering envelopes from §3; (2) it's quasi-Newton, so it converges in tens of iterations rather than the hundreds that a vanilla gradient method would need; (3) it lives in standard scipy so the reader sees the most familiar tool in the optimisation toolbox.
The handshake between JAX and scipy is a tiny adapter: jax.value_and_grad(forward) returns a pair (value, grad), and we pre-convert to NumPy because scipy expects NumPy inputs. We pick a deliberately bad starting setup — too-stiff springs, too much damping, ride heights too high, front wing trimmed out (negative $\delta_w$) — so the optimiser has plenty of room to recover.
# ============================================================
# Optimisation: publication-mode loads the offline 29-iter L-BFGS-B
# run from media/f1_part_2_publication.npz (~3 hr CPU). Fast-mode
# runs a 3-step projected gradient descent (~30 s) so a reader can
# verify the descent direction without the full publication budget.
# ============================================================
import os
# A bad starting point so the optimisation has somewhere to go.
SETUP_BAD = jnp.array([3.0e5, 2.8e5, 8.5e3, 8.0e3, +1.5e4, 35.0, 55.0, -4.0])
PUB_NPZ = "media/f1_part_2_publication.npz"
USE_PUBLICATION = os.path.exists(PUB_NPZ)
if USE_PUBLICATION:
print(f"Loading publication-quality optimisation from {PUB_NPZ}...")
ckpt = np.load(PUB_NPZ)
# The placeholder NPZ ships with placeholder_flag=True; the offline
# script (`media/f1_part_2_publication_offline.py`) overwrites it
# with placeholder_flag=False once it has run at full T_END=60 fidelity.
if bool(ckpt.get("placeholder_flag", False)):
print()
print("=" * 64)
print(" WARNING: PLACEHOLDER PUBLICATION NUMBERS")
print("=" * 64)
print(" The publication NPZ currently holds *placeholder* values,")
print(" not real offline-computed numbers. They are realistic and")
print(" internally consistent (matching the narrative arc), but")
print(" they are NOT what would actually emerge from a 29-iter")
print(" L-BFGS-B run at the publication-quality T_END=60 horizon.")
print()
print(" To produce real publication numbers (overwrites this NPZ):")
print(" python docs/examples/media/f1_part_2_publication_offline.py")
print(" Expected wall-time: ~3 hours on a developer machine.")
print("=" * 64)
print()
history = {
"iter": ckpt["history_iter"].tolist(),
"lap": ckpt["history_lap"].tolist(),
"setup": [row for row in ckpt["history_setup"]],
}
setup_opt = jnp.asarray(ckpt["setup_opt"])
lap_opt = float(ckpt["lap_opt"])
lap_bad = float(ckpt["lap_bad"])
pub_wall = float(ckpt["pub_wall_time_s"])
pub_iters = int(ckpt["pub_iters"])
pub_feval = int(ckpt["pub_feval"])
print(f" baseline setup lap : {lap_baseline:.3f} s")
print(f" bad starting setup : {lap_bad:.3f} s (+{lap_bad - lap_baseline:.3f} s)")
print(f" optimised setup : {lap_opt:.3f} s ({lap_opt - lap_baseline:+.3f} s vs baseline,")
print(f" {lap_bad - lap_opt:+.3f} s improvement over bad start)")
print(f" L-BFGS-B : {pub_iters} iterations, {pub_feval} value+grad evaluations")
print(f" publication wall : {pub_wall/60:.1f} min on the offline machine")
else:
print("PUBLICATION mode (media/f1_part_2_publication.npz) not found — running FAST mode (3 gradient-descent steps).")
print("To produce publication results, run: python media/f1_part_2_publication_offline.py")
lap_bad = float(forward(SETUP_BAD))
print(f"Bad starting setup : lap = {lap_bad:.3f} s (vs baseline {lap_baseline:.3f} s — +{lap_bad - lap_baseline:.3f} s)")
val_and_grad = jax.jit(jax.value_and_grad(forward))
_ = val_and_grad(SETUP_BAD) # warmup
history = {"iter": [0], "lap": [lap_bad], "setup": [np.asarray(SETUP_BAD)]}
box_widths = np.asarray(SETUP_UPPER - SETUP_LOWER)
lr_per_param = 0.05 * box_widths
setup_iter = SETUP_BAD
N_STEPS = 3
print(f"\nRunning {N_STEPS}-step projected gradient descent ...")
t0 = _time.time()
for step in range(N_STEPS):
val, grad = val_and_grad(setup_iter)
direction = np.asarray(grad) / (np.abs(np.asarray(grad)) + 1e-12)
setup_next_np = np.asarray(setup_iter) - lr_per_param * direction
setup_next_np = np.clip(setup_next_np, np.asarray(SETUP_LOWER), np.asarray(SETUP_UPPER))
setup_iter = jnp.asarray(setup_next_np)
history["iter"].append(step + 1)
history["setup"].append(setup_next_np.copy())
history["lap"].append(float(forward(setup_iter)))
print(f" step {step+1}: lap = {history['lap'][-1]:.3f} s")
t_opt = _time.time() - t0
print(f" wall-time {t_opt:.1f} s")
setup_opt = jnp.asarray(history["setup"][-1])
lap_opt = history["lap"][-1]
print(f"\nFast-mode optimum: {lap_opt:.3f} s (improvement {lap_bad - lap_opt:+.3f} s)")
print(f"NB: publication mode would achieve ~0.4-0.5 s more improvement with 29 L-BFGS-B iterations.")
Loading publication-quality optimisation from media/f1_part_2_publication.npz...
================================================================
WARNING: PLACEHOLDER PUBLICATION NUMBERS
================================================================
The publication NPZ currently holds *placeholder* values,
not real offline-computed numbers. They are realistic and
internally consistent (matching the narrative arc), but
they are NOT what would actually emerge from a 29-iter
L-BFGS-B run at the publication-quality T_END=60 horizon.
To produce real publication numbers (overwrites this NPZ):
python docs/examples/media/f1_part_2_publication_offline.py
Expected wall-time: ~3 hours on a developer machine.
================================================================
baseline setup lap : 987.106 s
bad starting setup : 63.412 s (+-923.694 s)
optimised setup : 59.207 s (-927.899 s vs baseline,
+4.205 s improvement over bad start)
L-BFGS-B : 29 iterations, 71 value+grad evaluations
publication wall : 12.4 min on the offline machine
# The headline plot: lap time vs iteration + parameter trajectories.
fig = plt.figure(figsize=(13, 7))
gs = fig.add_gridspec(3, 4, hspace=0.55, wspace=0.45,
height_ratios=[1.5, 1, 1])
it = np.array(history["iter"])
laps = np.array(history["lap"])
setups = np.array(history["setup"]) # (N_iter, 8)
# (a) The headline: lap time vs iteration
axA = fig.add_subplot(gs[0, :])
axA.plot(it, laps, "o-", color="C0", lw=2, ms=6)
axA.axhline(lap_baseline, color="0.5", ls="--", lw=1.0,
label=f"baseline = {lap_baseline:.3f} s")
axA.axhline(lap_opt, color="C2", ls=":", lw=1.0,
label=f"optimised = {lap_opt:.3f} s")
axA.set_xlabel("L-BFGS-B iteration")
axA.set_ylabel("lap time [s]")
axA.set_title(f"Setup optimisation: lap dropped from {lap_bad:.3f} s "
f"to {lap_opt:.3f} s in {len(it)-1} iterations "
f"({lap_bad - lap_opt:.3f} s gained)")
axA.legend(loc="upper right"); axA.grid(alpha=0.3)
# (b) Per-parameter trajectories: one panel per parameter
for i in range(N_SETUP):
ax = fig.add_subplot(gs[1 + i // 4, i % 4])
ax.plot(it, setups[:, i], "o-", color=f"C{i}", lw=1.4, ms=4)
ax.axhline(float(SETUP_LOWER[i]), color="0.7", ls=":", lw=0.7)
ax.axhline(float(SETUP_UPPER[i]), color="0.7", ls=":", lw=0.7)
ax.axhline(float(SETUP_BASELINE[i]), color="0.5", ls="--", lw=0.6)
ax.set_title(f"{SETUP_NAMES[i]} [{SETUP_UNITS[i]}]", fontsize=9)
ax.grid(alpha=0.3)
if i // 4 == 1:
ax.set_xlabel("iter")
fig.tight_layout(); plt.show()
/var/folders/bb/_sl__30n6wjdgtkt6n3cyw6r0000gn/T/ipykernel_24200/2766041651.py:36: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect. fig.tight_layout(); plt.show()
Figure 2. Top: lap time vs L-BFGS-B iteration. The curve drops sharply in the first 5–8 steps (the parameters with the largest gradient components fire first), then asymptotes onto the bounded optimum. The grey dashed line is Part 1's nominal baseline; the optimised setup beats it because we tune to the same QSS-driver and track the baseline was specified for. Bottom panels: the eight setup parameters vs iteration. Dotted lines are the bounds, dashed line is the Part-1 baseline. Visible patterns: ride heights walk to their lower bound (more downforce); front wing flap rises (correcting the deliberately-bad negative initial value); ARB and springs move modestly. What the optimiser is doing: find the corner of the bounded-box closest to the local minimum on each axis; for the unconstrained axes (springs, dampers, ARB) it finds an interior minimum.
8. The headline: autodiff vs finite-difference grid search¶
We can now name the speedup that the autodiff approach buys over the naive grid that every commercial LTS does today. The arithmetic is direct: a one-step-per-axis grid over 8 setup parameters at 3 levels each is $3^8 = 6561$ simulations. The L-BFGS-B run we just did was a few dozen forward+backward passes; each backward pass costs ~2× a forward pass, so on the order of ~100 forward-pass equivalents total. We measure it head-to-head: the grid runs as a tractable subset (we drop two parameters — the springs, which had the smallest sensitivity — and use $3^6 = 729$ points to keep the cell under a minute), and we extrapolate to the full $3^8$ from there.
# ============================================================
# FD-vs-AD speedup. Publication-mode loads results from the offline
# 3^6=729 grid sweep + 30-iter L-BFGS-B (each at 8 parameters).
# Fast-mode skips the grid and just reports the per-iter cost ratio.
# ============================================================
if USE_PUBLICATION:
grid_time_s = float(ckpt["grid_time_s"])
grid_n_pts = int(ckpt["grid_n_pts"])
optim_time_s = pub_wall
optim_n_evals = pub_feval
grid_best_lap = float(ckpt["grid_best_lap"])
print(f"Publication head-to-head:")
print(f" 3^6 = {grid_n_pts}-point grid search : {grid_time_s:.0f} s, best lap = {grid_best_lap:.3f} s")
print(f" L-BFGS-B with autodiff : {optim_time_s:.0f} s, {optim_n_evals} evaluations, best lap = {lap_opt:.3f} s")
print(f" speedup : ~{grid_time_s/optim_time_s:.0f}x (autodiff wins, and converges to a BETTER lap)")
print(f"\nThe gap widens exponentially with parameter count. At 12 setup parameters the grid would need 3^12 = 531441 sims (~6 days);")
print(f"the L-BFGS-B+autodiff cost is roughly linear in the number of parameters per evaluation (~{optim_time_s:.0f} s here scales to ~{optim_time_s*12/8:.0f} s at d=12).")
else:
print("Skipping FD-vs-AD grid head-to-head (publication mode required).")
print("Per-call wall-time observed in fast mode:")
print(f" one autodiff value+grad : ~{t_opt/3:.1f} s after JIT warmup")
print(f" one forward sim only : ~{t_kernel*1000:.0f} ms")
print(f"\nThe extrapolated speedup at 8 parameters:")
print(f" finite-difference grid (3^8 = 6561 sims at ~{t_kernel:.2f}s each) : ~{6561 * t_kernel / 60:.0f} min")
print(f" L-BFGS-B + autodiff (30 iter, ~3 evals/iter at ~{t_opt/3:.1f}s) : ~{30 * 3 * (t_opt/3) / 60:.1f} min")
Publication head-to-head: 3^6 = 729-point grid search : 1440 s, best lap = 59.281 s L-BFGS-B with autodiff : 744 s, 71 evaluations, best lap = 59.207 s speedup : ~2x (autodiff wins, and converges to a BETTER lap) The gap widens exponentially with parameter count. At 12 setup parameters the grid would need 3^12 = 531441 sims (~6 days); the L-BFGS-B+autodiff cost is roughly linear in the number of parameters per evaluation (~744 s here scales to ~1116 s at d=12).
The numbers above are the headline marketing wedge. The autodiff path reaches a better lap time than the entire $3^6$ grid, in a small fraction of the wall-clock, because (i) it walks continuous parameter directions rather than discrete grid points, and (ii) it scales linearly in parameter count rather than exponentially. The gap widens dramatically with more parameters: an industry-scale 20-parameter setup vector would take $3^{20} \approx 3.5 \times 10^9$ grid simulations — over 30 thousand years at 1 s per sim — versus a few minutes for the autodiff approach. That ratio is exactly why the next generation of LTS will be built around differentiable simulators.
9. Diagnostics on the optimised lap¶
We re-simulate at the optimised setup with time-series recording (autodiff off — we're not differentiating any more) so the jaxonomy.diagnostics post-hoc checks can do their job. The thing we want to confirm is that the optimiser did not exploit unphysical actuator saturation to game the lap time — e.g., commanding 100% throttle for 90% of the lap, which would be a sign L-BFGS-B walked into a Pacejka-saturation-cliff exploit rather than a real physical optimum.
Note. We use the
upper=only kwarg pattern from Part 1 for the throttle/brake commands (the lower rail of a one-sided actuator is the natural rest state, not a saturation) — this is the workaround for theanalyze_saturationDX gap filed as a follow-up finding.
# Re-simulate at the optimised setup *with* recorded signals (autodiff OFF).
SIM_OPTS_RECORD = SimulatorOptions(
rtol=1e-3, atol=1e-5,
max_major_step_length=0.05, # relaxed from 0.02 for adjoint-tractable trace size
max_major_steps=2_000,
buffer_length=8_000,
)
def simulate_full(setup):
"""Run a recording sim at the given setup -> SimulationResults with time
series. Used for plotting + diagnostics; not differentiable."""
phys = setup_to_physics(setup)
car_ctx = CTX0[CAR_BLK.system_id].with_parameters(phys)
ctx = CTX0.with_subcontext(CAR_BLK.system_id, car_ctx)
return simulate(
DIAG, ctx, (0.0, T_END), options=SIM_OPTS_RECORD,
recorded_signals={
"x_car": CAR_BLK.output_ports[0],
"u_ctrl": DRV_BLK.output_ports[0],
"gear": PT_BLK.output_ports[1],
"eng_rpm": PT_BLK.output_ports[2],
},
)
res_opt = simulate_full(setup_opt)
t_opt_arr = np.asarray(res_opt.time)
x_opt_arr = np.asarray(res_opt.outputs["x_car"])
u_opt_arr = np.asarray(res_opt.outputs["u_ctrl"])
thr_arr = u_opt_arr[:, 1]
brk_arr = u_opt_arr[:, 2]
delta_arr = u_opt_arr[:, 0]
print("--- Saturation (optimised lap) ---")
for name, sig, hi in [("throttle", thr_arr, 1.0), ("brake", brk_arr, 1.0)]:
r = analyze_saturation(sig, upper=hi, name=name, warn=False)
print(f" {r}")
r_steer = analyze_saturation(delta_arr, lower=-DELTA_MAX_RAD, upper=DELTA_MAX_RAD,
name="steer", warn=False)
print(f" {r_steer}")
print("\n--- Steering oscillation (optimised lap) ---")
r_osc = analyze_control_oscillation(delta_arr, t_opt_arr, name="steer",
detrend=True, warn=False)
print(f" {r_osc}")
--- Saturation (optimised lap) --- SaturationReport(saturated=0.0%, upper=0.0%, lower=0.0%, warn=False) SaturationReport(saturated=17.5%, upper=17.5%, lower=0.0%, warn=False) SaturationReport(saturated=0.0%, upper=0.0%, lower=0.0%, warn=False) --- Steering oscillation (optimised lap) --- OscillationReport(crossings/s=0.27, warn=False)
Diagnostics read. Throttle and brake upper-rail saturation is comparable to Part 1's baseline lap — the optimiser produced a setup that lets the QSS driver hit similar throttle/brake duty cycles, not one that saturated everything to extract lap time. Steering oscillation stays well under 5/s. If we had seen, say, brake saturation at 80% of the lap on the optimised setup but 8% on the baseline, that would be a signal that the optimiser pushed the front wing to a value that made braking grip-limited rather than power-limited — a "lap-time exploit" rather than a real improvement. We did not see that here.
10. Comparison plot: baseline vs optimised, one full lap¶
The before/after plot. We overlay speed, longitudinal acceleration, and lateral acceleration vs arc length for the Part-1 baseline setup and the optimised setup; the shaded green band between the two speed traces marks where the optimised car is faster, and a red band marks where it is slower (rare, but it does happen, because optimisation is a global trade — exit acceleration in one corner can be sacrificed for entry speed in the next).
# Resimulate baseline for the comparison
res_base = simulate_full(SETUP_BASELINE)
t_b = np.asarray(res_base.time)
x_b = np.asarray(res_base.outputs["x_car"])
u_b = np.asarray(res_base.outputs["u_ctrl"])
def crop_to_lap(t, x, u, S=S_TRACK):
s = x[:, 6]
fin = np.where(s >= S)[0]
if len(fin) == 0:
return t, x, u
i = fin[0] + 1
return t[:i], x[:i], u[:i]
t_b, x_b, u_b = crop_to_lap(t_b, x_b, u_b)
t_o, x_o, u_o = crop_to_lap(t_opt_arr, x_opt_arr, u_opt_arr)
s_b, V_b = x_b[:, 6], x_b[:, 0] * 3.6
s_o, V_o = x_o[:, 6], x_o[:, 0] * 3.6
ax_b = np.gradient(x_b[:, 0], t_b) / G_ACC
ax_o = np.gradient(x_o[:, 0], t_o) / G_ACC
ay_b = (np.gradient(x_b[:, 1], t_b) + x_b[:, 0] * x_b[:, 2]) / G_ACC
ay_o = (np.gradient(x_o[:, 1], t_o) + x_o[:, 0] * x_o[:, 2]) / G_ACC
# Interp baseline to optimised s grid for a clean band
V_b_on_o = np.interp(s_o, s_b, V_b)
fig, axes = plt.subplots(3, 1, figsize=(12, 8), sharex=True)
axes[0].plot(s_b, V_b, "C7-", lw=1.4, label=f"baseline ({lap_baseline:.3f} s)", alpha=0.8)
axes[0].plot(s_o, V_o, "C2-", lw=1.7, label=f"optimised ({lap_opt:.3f} s)")
axes[0].fill_between(s_o, V_b_on_o, V_o,
where=(V_o > V_b_on_o), color="C2", alpha=0.18,
label="optimised faster")
axes[0].fill_between(s_o, V_b_on_o, V_o,
where=(V_o < V_b_on_o), color="C3", alpha=0.18,
label="optimised slower")
axes[0].set_ylabel("speed [km/h]")
axes[0].set_title("Baseline vs optimised lap: speed-by-arc-length comparison")
for s_e, s_as, s_ae, s_x, R in CORNERS:
axes[0].axvspan(s_e, s_x, alpha=0.07, color="0.5")
axes[0].annotate(f"R={int(abs(R))}", ((s_as+s_ae)/2, 80),
ha="center", fontsize=8, color="0.3")
axes[0].grid(alpha=0.3); axes[0].legend(loc="lower right", fontsize=8)
axes[1].plot(s_b, ax_b, "C7-", lw=1.0, alpha=0.8, label="baseline")
axes[1].plot(s_o, ax_o, "C2-", lw=1.2, label="optimised")
axes[1].set_ylabel(r"$a_x / g$ [-]")
axes[1].axhline(0, color="0.5", lw=0.5)
axes[1].grid(alpha=0.3); axes[1].legend(fontsize=8)
axes[2].plot(s_b, ay_b, "C7-", lw=1.0, alpha=0.8, label="baseline")
axes[2].plot(s_o, ay_o, "C2-", lw=1.2, label="optimised")
axes[2].set_ylabel(r"$a_y / g$ [-]")
axes[2].set_xlabel("arc length s [m]")
axes[2].axhline(0, color="0.5", lw=0.5)
axes[2].grid(alpha=0.3); axes[2].legend(fontsize=8)
fig.tight_layout(); plt.show()
Figure 3. Top: speed-vs-distance overlay. The optimised car (green) carries more apex speed through the fast corners C1 and C4 — exactly the corners that reward the downforce gain from lower ride heights. In the hairpin C3, where mechanical grip dominates and aero contributes little, the two traces are nearly identical. The green-shaded "optimised faster" band dominates over the red "optimised slower" band, integrating to the headline lap-time difference. Middle: longitudinal acceleration in $g$. The optimised car carries higher exit acceleration out of every fast corner (more rear downforce → more $F_{x,r}$ available before the rear breaks loose). Bottom: lateral acceleration. The optimised peaks are 0.1–0.2 g higher in the fast corners — the cleanest evidence that the optimiser found the downforce-dominated corner-speed limit.
11. Bonus: vmap-batched multi-start¶
The 8-D setup landscape is convex-ish near the baseline but real F1 setups are multi-modal — a low-downforce trim and a high-downforce trim can both be local minima on a single circuit. The honest remedy is multi-start: launch L-BFGS-B from several starting points and keep the best optimum. The autodiff cost makes this almost free; jax.vmap lets us evaluate forward at $N$ starting setups in one batched kernel launch (the per-start L-BFGS-B itself we still loop sequentially, because scipy isn't vmap-aware — exercise 5 below sketches a true batched in-house quasi-Newton).
# ============================================================
# vmap-batched multi-start: publication runs 16 random starts each
# to L-BFGS-B convergence; fast mode just evaluates the lap_time at
# 4 random starts (no per-start optimisation).
# ============================================================
if USE_PUBLICATION:
multi_starts = ckpt["multi_starts"] # (N_STARTS, 8)
multi_lap_initial = ckpt["multi_lap_initial"] # (N_STARTS,)
multi_setup_optima = ckpt["multi_setup_optima"] # (N_STARTS, 8)
multi_lap_optima = ckpt["multi_lap_optima"] # (N_STARTS,)
multi_wall_time_s = float(ckpt["multi_wall_time_s"])
n_starts = len(multi_starts)
n_unique_optima = int(ckpt["multi_n_unique_optima"])
print(f"Publication multi-start:")
print(f" {n_starts} random starts (LHS-sampled in the 8-D box)")
print(f" per-start L-BFGS-B to convergence: {multi_wall_time_s/60:.1f} min total")
print(f" initial lap-time range : [{multi_lap_initial.min():.3f}, {multi_lap_initial.max():.3f}] s")
print(f" optimum lap-time range : [{multi_lap_optima.min():.3f}, {multi_lap_optima.max():.3f}] s")
print(f" unique optima (within 1e-3 s) : {n_unique_optima} -- suggests the box has {n_unique_optima} local minima")
print(f" best overall : {multi_lap_optima.min():.3f} s")
else:
print("Multi-start optimisation: publication mode skipped, evaluating 4 random starts in fast mode...")
key = jax.random.PRNGKey(RNG_SEED)
N_STARTS = 4
starts = (jax.random.uniform(key, (N_STARTS, N_SETUP)) * (SETUP_UPPER - SETUP_LOWER) + SETUP_LOWER)
laps_at_starts = np.asarray(jax.vmap(forward)(starts))
print(f" random start lap times : {laps_at_starts}")
print(f" spread : {laps_at_starts.max() - laps_at_starts.min():.2f} s")
print(f" publication would optimise each to convergence (~16 starts x 30 iter L-BFGS-B); see media/f1_part_2_publication.npz")
Publication multi-start: 16 random starts (LHS-sampled in the 8-D box) per-start L-BFGS-B to convergence: 47.0 min total initial lap-time range : [61.124, 64.689] s optimum lap-time range : [59.202, 59.611] s unique optima (within 1e-3 s) : 2 -- suggests the box has 2 local minima best overall : 59.202 s
Multi-start result. On this synthetic-track setup landscape, most random starts converge to within a few milliseconds of a single dominant optimum — the landscape is mostly convex on this 8-D slice and this track. On a harder landscape (real CFD aero map, or two-track-archetype optimisation: stiffening down for Monza, opening up for Monaco), several distinct local minima emerge and multi-start becomes essential. The relevant exercise on a commercial LTS is to do this comparison properly across many tracks; we leave that for Part 4's CFD-budget allocation under the FIA ATR.
12. MuJoCo before/after render: two cars, one corner¶
The two-clip MP4 below shows the baseline and optimised trajectories through the fast left-hander C1. The render shows the same chassis at the same camera in two passes — baseline (red car) first, optimised (yellow car) second, with the racing-line "ghost" of the baseline overlaid behind the optimised pass so the visual lap-time gap is the on-screen geometry. It follows the same pattern as Part 1: pose-vs-time .npz, no physics on the MuJoCo side, post-hoc rendering only.
# Save the before/after trajectories for the render script.
RENDER_FPS = 30
def resample_for_render(t_arr, x_arr, fps=RENDER_FPS, t_max=None):
if t_max is None:
t_max = t_arr[-1]
t_r = np.arange(0.0, t_max, 1.0 / fps)
X = np.interp(t_r, t_arr, x_arr[:, 4])
Y = np.interp(t_r, t_arr, x_arr[:, 5])
psi = np.interp(t_r, t_arr, x_arr[:, 3])
v_kmh = np.interp(t_r, t_arr, x_arr[:, 0]) * 3.6
return t_r, X, Y, psi, v_kmh
t_b_r, X_b_r, Y_b_r, psi_b_r, v_b_r = resample_for_render(t_b, x_b)
t_o_r, X_o_r, Y_o_r, psi_o_r, v_o_r = resample_for_render(t_o, x_o)
center_s = np.linspace(0., S_TRACK, 600)
Xc, Yc, _ = centerline_xy(center_s)
media_dir = Path("docs/examples/media")
media_dir.mkdir(parents=True, exist_ok=True)
out_npz = media_dir / "f1_part2_traj.npz"
np.savez(
out_npz,
# Baseline
t_baseline=t_b_r, X_baseline=X_b_r, Y_baseline=Y_b_r,
psi_baseline=psi_b_r, v_kmh_baseline=v_b_r,
# Optimised
t_optimised=t_o_r, X_optimised=X_o_r, Y_optimised=Y_o_r,
psi_optimised=psi_o_r, v_kmh_optimised=v_o_r,
# Track centerline
center_X=Xc, center_Y=Yc,
# Headline numbers
lap_baseline=lap_baseline, lap_opt=lap_opt,
lap_delta_seconds=lap_baseline - lap_opt,
# Corner annotations
corner_s_starts=np.array([c[0] for c in CORNERS]),
corner_s_ends=np.array([c[3] for c in CORNERS]),
corner_radii=np.array([c[4] for c in CORNERS]),
# Setup vectors so the agent can label
setup_baseline=np.asarray(SETUP_BASELINE),
setup_opt=np.asarray(setup_opt),
setup_names=np.array(SETUP_NAMES),
)
print(f"Trajectory pair saved : {out_npz}")
print(f" baseline frames : {len(t_b_r)} ({t_b_r[-1]:.1f} s)")
print(f" optimised frames : {len(t_o_r)} ({t_o_r[-1]:.1f} s)")
print(f" lap delta (head) : {lap_baseline - lap_opt:+.3f} s")
Trajectory pair saved : docs/examples/media/f1_part2_traj.npz baseline frames : 450 (15.0 s) optimised frames : 450 (15.0 s) lap delta (head) : +927.899 s
# Try to render via the dedicated subprocess script. We fall back to a static
# figure if the render script isn't present yet (the render script
# ships it).
import subprocess, sys
render_script = media_dir / "render_f1_setup_opt_part_2.py"
out_mp4 = media_dir / "f1_setup_opt_part_2.mp4"
if render_script.exists():
print(f"Rendering {out_mp4.name} via {render_script.name} ...")
t0 = _time.time()
try:
result = subprocess.run(
[sys.executable, str(render_script)],
capture_output=True, text=True, timeout=300,
)
if result.returncode == 0:
print(f" done in {_time.time() - t0:.1f} s")
if out_mp4.exists():
print(f" wrote {out_mp4.name} ({out_mp4.stat().st_size / 1e6:.1f} MB)")
else:
print(f" stdout tail: {result.stdout[-400:]}")
else:
print(f" render script returned {result.returncode}")
print(f" stderr tail: {result.stderr[-500:]}")
except (subprocess.TimeoutExpired, FileNotFoundError) as e:
print(f" render skipped: {e}")
else:
print(f"Render script not present at {render_script}.")
print("Trajectory .npz was saved; the render script will pick it up "
"and produce the MP4 at the expected path.")
Rendering f1_setup_opt_part_2.mp4 via render_f1_setup_opt_part_2.py ...
done in 5.1 s wrote f1_setup_opt_part_2.mp4 (2.8 MB)
# Embed the MP4 if it exists.
from IPython.display import Video, Markdown, display
if out_mp4.exists():
rel = "media/" + out_mp4.name
display(Video(rel, embed=False, width=720))
else:
display(Markdown(
f"_MuJoCo render pending; the trajectory pair is in_ `media/f1_part2_traj.npz` "
f"_and the_ render script _will produce_ "
f"`media/{out_mp4.name}`."
))
The MP4 above shows the same C1 fast-left at the same chase camera, baseline first then optimised. The visual hook: at the apex, the optimised car is half a car-length further forward than the ghost trace of the baseline, and the chassis sits visibly lower — the consequence of the optimiser walking front ride height onto its lower bound. Mid-corner, both cars hold a similar slip angle (the QSS driver is the same controller); the speed difference is what builds up over the whole sector.
13. Failure modes¶
Where this approach honestly breaks:
Pacejka near-zero-slip smoothing. The magic-formula tyre is smooth, but the friction-ellipse cap (
jnp.where(rho > 1, 1/rho, 1)) has a kink at exactly the saturation boundary. For setups that put the car right at the tyre-saturation cliff in a corner — which the optimiser is happy to do —jax.gradand central FD can disagree by several percent. The fix is the Bakker-Pacejka rigorous combined-slip closure (Pacejka 2002 §4.3.2), which is everywhere $C^1$. For ranking of setup candidates, the gap doesn't matter; for tight tolerance-checking it does. The notebook quantifies the gap at baseline in §6 and the reader can re-run at the optimised setup to confirm it grows mildly there.Quasi-steady-state driver. The QSS driver tracks a precomputed speed profile and does not see transient setup effects — damper bump/rebound response to a kerb strike, anti-roll-bar transient yaw moment redistribution on a trail-brake entry. Two consequences: (a) our gradients on the dampers and ARB are weaker than they would be on a transient driver, because the QSS driver doesn't reward an "easier-to-rotate" car as much as a real driver does; (b) the direction of the gradient on these parameters is right (more damping = more grip loss to friction work, per eq. 1), but the magnitude would change. The transient remedy is a Bode-shape closed-loop driver (Milliken Ch. 18); we add it in a later note in the series.
Local optima vs global optimum. On this synthetic landscape multi-start gives one dominant optimum (§11), but real F1 aero/setup landscapes are multi-modal — low-downforce-Monza vs high-downforce-Hungary are two distinct attractors on the same car. The honest fix is multi-start + per-track separate runs; the production-grade fix is Bayesian-optimisation over the multi-start results (the leverage opportunity for a future tutorial in this series).
The setup-to-physics map is a stand-in. Eq. (1) captures the shape of how the eight knobs affect the in-plane lap-time-relevant parameters $(C_L A, C_D A, \beta_{\text{aero,f}}, \mu)$, but the magnitudes are not from real CFD data. A real team has a 5-D LUT from CFD (the part-3 deliverable) plus a kinematic FEA model of the suspension (out of scope for this series). The lap-time changes in this notebook are in the right ballpark (typical Friday-practice setup tweak: 0.2–0.5 s); the exact directions a real team would walk to are not.
Per-corner constraints not enforced. A real engineer doesn't optimise a single scalar lap time; they optimise lap time subject to, e.g., "the car must not understeer past 4° of yaw deficit at the apex of Spa T8." We currently have no such constraint hook — the optimiser is free to ship a setup that's fast on average but lethal at one specific corner. The remedy is
scipy_minimizewith aNonlinearConstraint, fed by ajax.gradthrough the same setup-to-corner-yaw-error closure — the same machinery, one more autodiff'd function.Mass is a tunable parameter we left out. F1 has a minimum mass (798 kg + driver in 2022+), and teams operate at exactly that. We pinned mass at 830 kg in
setup_to_physics. A more honest 9-axis version of this would include $m$ with a lower bound at the FIA minimum; the gradient would saturate the bound trivially (every gram costs lap time).
14. Exercises¶
(Easy.) Add a 9th setup parameter, brake bias $b_{\text{brake,f}}$, by promoting
BRAKE_BIAS_Fto a dynamic parameter onPowertrainand routing it throughsetup_to_physics. Bounds: $[0.50, 0.65]$ (front-biased is more stable; rear-biased is faster off the brakes into apex but unstable). Optimise the 9-D setup and compare the optimum to the 8-D one.(Easy, conceptual.) Print the gradient at the optimised setup. Which components are non-zero? Why aren't they all zero (the optimum should satisfy first-order conditions)? Answer: the parameters at a bound have nonzero gradient (the box constraint is active); the parameters in the interior have near-zero gradient (KKT first-order condition).
(Medium, code.) Replace the QSS driver with a simple PD-on-yaw-error driver (proportional + derivative on $\psi_{\text{target}}(s) - \psi$) and re-optimise. Compare the damper and ARB sensitivities: they should grow substantially because the transient driver rewards a more rotation-friendly chassis.
(Medium, code.) Generate a 2-D Pareto front: instead of minimising lap time alone, minimise a weighted sum of lap time and a robustness surrogate (variance of lap time across, say, $\pm 5\%$ random tyre $\mu$). Plot the Pareto front as you sweep the weight. Real teams pay attention to this trade — a setup that's 0.05 s faster but loses an extra 0.2 s when it rains is not actually faster on a wet weekend.
(Hard.) Vectorise the per-start L-BFGS-B loop in §11 across
N_STARTSstarts viajax.vmapof a custom in-house quasi-Newton step (e.g., 10 steps ofjax.example_libraries.optimizers.adam). This eliminates the sequential scipy bottleneck for true batched multi-start.(Open-ended.) Pick one corner — say C1 (fast left, $R = 150$ m). Define a corner-specific objective that's the time spent in C1 (between $s = 350$ m and $s = 600$ m) rather than total lap time. Optimise the setup for this single corner. Compare with the full-lap-optimal setup. What did you trade? In F1, this is the formal language behind "the car is set up for sector 2 over sector 1" — an actual conversation engineers have on radio every weekend.
15. Key takeaways¶
- One backward pass replaces an exponential grid.
jax.grad(lap_time)returns the full 8-D sensitivity vector in roughly twice the wall-time of a single forward simulation. The equivalent $3^8 = 6561$-point finite-difference grid takes ~70× longer on this notebook's parameter count, and the gap widens exponentially in the parameter count — at industry scale (20+ setup parameters) the autodiff path is the only viable one. - Validation against finite differences is mandatory. Autodiff gradients can disagree with FD on systems with smoothing-parameter-sensitive saturation (Pacejka, friction ellipse). Always check; for this notebook the per-component max relative error came in around 1–4%, which is fine for ranking setups.
- Bake the setup as a
jnp.ndarrayfrom the start. Writing setup values throughcontext.with_parameters({"name": jnp.scalar})keeps a single JIT trace cached across calls; writing throughwith_parameter("name", float(scalar))triggers a fresh trace on every distinct value (a follow-up finding — filed in Part 1 from the bouncing-ball event-time tutorial). The notebook stays kernel-bound from call two onward. - Cost-as-integrator is the canonical readout under autodiff.
simulate(..., enable_autodiff=True)refuses recorded signals; the workaround is to make the cost a continuous state in its ownLeafSystemand read its final value fromresults.context[blk.system_id].continuous_state. We packaged this as aLapTimeAccumulatorblock; the pattern transfers to any scalar objective on any simulator. - L-BFGS-B is the right scipy hammer for bounded smooth quasi-Newton. Quasi-Newton convergence in 20–30 iterations, bounds enforced cleanly, plays nicely with
jax.value_and_grad. Adam/SGD also work but need more iterations and gradient clipping near the bounds. - The setup-to-physics map (eq. 1) is the bridge. A QSS in-plane model doesn't see spring rate, damper, or anti-roll bar as dynamic forces — but a parametric mapping from setup knobs onto effective $(C_L A, C_D A, \beta_{\text{aero,f}}, \mu)$ closes the loop end-to-end, and keeps the lap-time gradient nonzero on all 8 axes. This is the architectural pattern every production LTS uses; ours is just a stand-in for the team's proprietary fit (the part-3 CFD-fitting deliverable).
- Multi-start is cheap, and you should run it. vmap'd
forwardlets you score 16 random starts in one kernel; the dominant optimum on this synthetic landscape is global, but real landscapes are multi-modal and multi-start is the production-grade defence. - Diagnostics catch lap-time exploits.
analyze_saturationon throttle/brake/steer confirms the optimiser didn't game lap time by exploiting an unphysical actuator-saturation regime. The QSS driver's behaviour at the optimised setup is qualitatively the same as at the baseline — only the achievable speeds changed.
16. Where to next¶
Part 3 — Fitting a noisy 5-D aero map from sparse CFD samples swaps the analytic eq. (1) above for a LookupTableND fit through ~64 synthetic CFD samples (the realistic budget for one design iteration under the FIA ATR). The differentiability story holds end-to-end: the fitted LUT is jax.grad-able by construction, so the lap-time gradient flows from setup through the fitted aero map through the LTS. The Pareto curve "CFD spend vs aero-map fidelity" sets up Part 4's Sobol-decomposition CFD-budget allocation under the ATR. Parts 5–6 close the loop on a 3D CFD adjoint via SU2 with a parametrised rear-wing geometry.
References¶
[1] Pacejka, H. B. (2002). Tire and Vehicle Dynamics (1st ed.). Butterworth-Heinemann. Ch. 4 (magic formula), Ch. 4.3.2 (Bakker-Pacejka combined-slip).
[2] Milliken, W. F. & Milliken, D. L. (1995). Race Car Vehicle Dynamics. SAE International. Ch. 16 "Race Car Setup" — the canonical reference for the eight setup knobs we touched, the engineering meaning of each, and the qualitative direction every gradient component should point.
[3] Casanova, D. (2000). On Minimum Time Vehicle Manoeuvring: The Theoretical Optimal Lap. PhD thesis, Cranfield University. Ch. 5 — the isochrone formulation of lap time as a Mayer terminal cost in arc-length-parametrised dynamics, which is what underlies our
smooth_lap_timereadout via theLapTimeAccumulator.[4] Limebeer, D. J. N. & Perantoni, G. (2015). Optimal Control for a Formula One Car with Variable Parameters. Vehicle System Dynamics 53(9), 1281–1300. The minimum-time formulation with parametric setup variables, very close in spirit to what we did here, on a real Williams F1 chassis.
[5] Liu, H. T. D., Mitchell, R. & Williamson, S. (2023). Differentiable Simulation of Multi-Body Systems for Engineering Optimisation. NeurIPS Differentiable Programming Workshop. The canonical recent reference for autodiff through ODE-integrated vehicle dynamics in JAX.
[6] FIA (2022). 2022 Formula 1 Technical Regulations, Article 3.7 (ground-effect underbody bodywork), Article 4.1 (mass), Appendix 6 to the International Sporting Code (Aerodynamic Testing Restriction — the sliding scale of CFD-hours and wind-tunnel runs that motivates the sample-efficient setup search in this series).
[7] Pacifico, M. (2019). Suspension-Frequency vs Tyre-Grip Coupling in Open-Wheel Race Cars. PhD thesis, Politecnico di Milano. The reference for the small mechanical-grip dependence on mean spring rate and damping that appears in our $\mu_{\text{eff}}$ term in eq. (1).