Tuning PID Controllers with Automatic Differentiation

PID controllers are the workhorses of industrial control. They are simple enough to fit on a microcontroller, yet surprisingly effective across a wide range of systems. The catch is that tuning them well takes time and, if you are being honest, a fair amount of intuition. Classic recipes like Ziegler-Nichols are fast but often leave performance on the table. Model-free autotuning approaches work but can require live experiments on the real plant.

Here is a different take: write down a differentiable simulation of your system, define a loss that captures what “good control” means to you, and let gradient descent find the gains. Automatic differentiation handles the calculus; you just specify the objective.

This post walks through exactly that, using JAX¹. The controller structure follows the practical implementation guidelines of Sundström et al.²

The Example: A Second-Order Mass-Spring-Damper

The plant is a linear mass-spring-damper:

m\ddot{x} + c\dot{x} + kx = u(t)

with mass $m = 1\,\text{kg}$ , damping $c = 0.5\,\text{N·s/m}$ , and spring stiffness $k = 2\,\text{N/m}$ . The PID controller drives the position $x$ to a unit step setpoint over a $10\,\text{s}$ simulation.

This system has a natural frequency of $\omega_n = \sqrt{k/m} \approx 1.41\,\text{rad/s}$ and a damping ratio of $\zeta = c/(2\sqrt{mk}) \approx 0.18$ , so the open-loop step response is lightly damped and oscillatory — a good stress test for the tuning method.

The Incremental (Velocity) Form

Rather than computing the control output $u$ directly at each step, the simulation uses the incremental (or velocity) form: only the increment $\Delta u$ is computed, and the output is accumulated as $u[k] = u[k-1] + \Delta u[k]$ .

This structure has several practical advantages over the position form. Actuator saturation and bumpless switching are easier to handle because the accumulator naturally keeps track of where the integrator state is. There is also no integral wind-up from the proportional or derivative paths — only the explicit integral increment $\Delta u_i$ can cause wind-up.

Setpoint weighting

A standard PID controller applies the setpoint equally to all three terms. In practice it is often better to split this: let the proportional and derivative paths see a weighted version of the setpoint, while the integral path always sees the full error.

Two weights $b$ and $c$ parameterise this:

e_p[k] = b\,r - y[k], \qquad e_d[k] = c\,r - y[k], \qquad e[k] = r - y[k]

Setting $b < 1$ reduces the proportional kick when the setpoint steps, giving a smoother response without sacrificing integral action. Setting $c = 0$ makes the derivative act only on measurement changes — eliminating the large derivative spike that would otherwise appear at $t = 0$ when $r$ jumps discontinuously.

Derivative approximation

Rather than using a filtered derivative, the derivative is approximated separately for $r$ and $y$ via raw finite differences:

\dot{r}[k] = \frac{r[k] - r[k-1]}{\Delta t}, \qquad \dot{y}[k] = \frac{y[k] - y[k-1]}{\Delta t}

The second-difference increments $\Delta\dot{r}[k] = \dot{r}[k] - \dot{r}[k-1]$ and $\Delta\dot{y}[k] = \dot{y}[k] - \dot{y}[k-1]$ then enter the derivative increment directly.

The full incremental update at step $k$ is:

\begin{aligned} \Delta u_p[k] &= K_p\bigl(b\,\Delta r[k] - \Delta y[k]\bigr) \\ \Delta u_i[k] &= K_i\,\Delta t\,e[k] \\ \Delta u_d[k] &= K_d\bigl(c\,\Delta\dot{r}[k] - \Delta\dot{y}[k]\bigr) \\ u_{\text{raw}}[k] &= u[k-1] + \Delta u_p[k] + \Delta u_i[k] + \Delta u_d[k] \end{aligned}

where $\Delta r[k] = r[k] - r[k-1]$ and $\Delta y[k] = y[k] - y[k-1]$ .

Here, setting $b = 1$ and $c = 0$ implements derivative-on-measurement. For a constant setpoint $\Delta r = 0$ on every step after the initial one, the proportional increment reduces to $-K_p\,\Delta y$ and only responds to how the output is changing. The integral term $K_i\,\Delta t\,e[k]$ carries the absolute error information needed to eliminate steady-state offset.

The plant is integrated with semi-implicit Euler (velocity updated before position), which is slightly more stable than explicit Euler for oscillatory systems at the same step size $\Delta t = 0.01\,\text{s}$ .

Saturation and Anti-Windup

Rate-limited bounds

Rather than applying fixed saturation limits $[u_{\min},\, u_{\max}]$ directly, the active bounds at each step are tightened by a rate constraint on the previous output $u_s[k-1]$ :

\hat{u}_{\min}[k] = \max\!\bigl(u_{\min},\; u_s[k-1] + \Delta t\,\dot{u}_{\min}\bigr), \quad \hat{u}_{\max}[k] = \min\!\bigl(u_{\max},\; u_s[k-1] + \Delta t\,\dot{u}_{\max}\bigr)

This limits how fast the actuator output can change between steps, independently of the absolute limits. In the current simulation the rate limits are set large enough to be inactive, but the structure is in place for the general case.

Anti-windup by soft pull-back

After computing $u_{\text{raw}}$ , a soft correction pulls the signal back toward the feasible region before the hard clip:

u_s = \text{clip}(u_{\text{raw}},\; \hat{u}_{\min},\; \hat{u}_{\max})

u \leftarrow u_{\text{raw}} - \frac{\Delta t}{T_t}\bigl(u_{\text{raw}} - u_s\bigr)

u[k] = \text{clip}(u,\; \hat{u}_{\min},\; \hat{u}_{\max})

When the signal is within bounds $u_s = u_{\text{raw}}$ and the correction vanishes. When saturated, the correction partially unwinds the accumulator at a rate governed by the time constant $T_t$ . The state carried forward is $u$ (after soft correction, before the final hard clip), so the integrator is reset smoothly rather than hard-clamped. The hard clip on $u[k]$ ensures the plant never sees an out-of-range command.

The Loss Function

Rather than simply minimizing error relative to a step input (which can result in overshoot as the optimizer pushes for the fastest possible rise time), we use a Model-Reference Loss. We define a desired critically damped second-order response:

y_{\text{ref}}(t) = 1 - (1 + \omega_0 t) e^{-\omega_0 t}

with $\omega_0 = 2.5\,\text{rad/s}$ . This reference trajectory has zero overshoot and a settling time of approximately $2.5\,\text{s}$ — a perfect template for high-performance, smooth tracking.

The objective minimizes the tracking error relative to this reference, plus a tiny control effort regulariser:

\mathcal{L} = \sum_{n=0}^{N} \bigl(x_n - y_{\text{ref},n}\bigr)^2\,\Delta t + \lambda\sum_n u_n^2\,\Delta t

with $\lambda = 10^{-5}$ .

Model Reference Tracking. By penalizing deviations from $y_{\text{ref}}(t)$ , the optimizer is guided to find gains that yield zero overshoot and follow a smooth, physical trajectory.

Effort regularizes the gains with a very small coefficient, preventing them from becoming unnecessarily large while allowing the controller to saturate at its physical limit ( $U_{\text{max}} = 20\,\text{N}$ ) to match MATLAB tuning conditions.

Why JAX?

JAX’s jit + grad pipeline means the gradient computation costs roughly the same as two forward passes — that is the magic of reverse-mode AD. With jax.lax.scan you can unroll the simulation loop without materialising every intermediate in Python, keeping compilation time and memory under control even for long horizons.

There are no finite-difference approximations and no symbolic manipulations. The gradients are numerically exact (up to floating point) regardless of the complexity of the integrator or the loss. Crucially, jnp.clip is differentiable in JAX: its subgradient is zero in the saturated region and one elsewhere, so the optimizer receives a correct signal even when the saturation bounds are active.

Optimisation Setup

The optimised parameters are $[K_p,\, K_i,\, K_d]$ , initialised at $[1, 0, 0]$ . The setpoint weights $b = 1$ and $c = 0$ are fixed constants. Setting $c=0$ implements derivative-on-measurement, ensuring that setpoint steps do not create derivative spikes (derivative kick).

Adam runs for 3000 iterations with a cosine-decaying learning rate starting at $0.1$ . The decay allows the optimizer to make large early updates and then freeze the gains as it approaches convergence.

Results

Step Response Comparison

Step response comparison

Orange is the initial controller ( $K_p = 1,\, K_i = 0,\, K_d = 0$ ). With the incremental form and no integral action, the controller accumulates a small control output on the first step (proportional kick from the setpoint step) and then only reacts to changes in the output — there is no mechanism to correct the residual position error once the output stops moving, so the response oscillates slowly without converging to the setpoint.

Blue is the optimised controller ( $K_p \approx 5.28,\, K_i \approx 2.43,\, K_d \approx 4.21$ ). The step response rises smoothly, reaching the setpoint at around $2.2\,\text{s}$ with zero overshoot, matching the desired critically damped response. Damping from the derivative term $K_d$ prevents transient oscillations, while the integral action $K_i$ guarantees zero steady-state error.

Loss Curve

The loss drops by over two orders of magnitude over 3000 iterations.

Loss over iterations

The steepest descent happens in the first ~200 iterations when the learning rate is near its peak value. After roughly iteration 500 the curve flattens, indicating convergence. The cosine schedule ensures the gains freeze smoothly rather than continuing to drift.

Gain Trajectories

PID gain trajectories

$K_p$ and $K_d$ grow first, establishing the necessary damping and proportional action to track the rise of the reference model. $K_i$ ramps up smoothly to eliminate the steady-state tracking error. All three curves flatten as the cosine schedule drives the learning rate toward zero.

Tracking Error

Tracking error

Orange (initial, $K_i = 0$ ): the error never settles to zero due to the lack of integral action.

Blue (optimised): the tracking error relative to the reference model is driven to zero, resulting in a perfect match between the plant output and the desired critically damped response.

Control Effort

Control effort

At $t = 0$ , the setpoint step from $0$ to $1$ creates a proportional kick of $K_p \cdot b \cdot 1 \approx 5.28\,\text{N}$ . Because $c = 0$ , the derivative term acts only on measurement changes, completely eliminating the setpoint-related derivative kick at $t=0$ . The control signal drops smoothly to $\sim 0.7\,\text{N}$ , then rises to settle at $2.0\,\text{N}$ (the spring force required to hold the mass at $x = 1\,\text{m}$ ), with no actuator saturation.

Practical Considerations

Setpoint weights as design parameters. $b$ and $c$ control the trade-off between setpoint tracking aggressiveness and smooth transients. $b < 1$ softens the proportional response to reference steps. $c = 0$ turns the derivative into pure derivative-on-measurement, removing any setpoint-related derivative kick entirely. These are typically tuned by the engineer rather than optimised, because gradient descent can satisfy them in unintended ways (e.g. by exploiting the interaction between $b$ and $K_p$ ).

Anti-windup time constant. $T_t$ governs how quickly the integrator is unwound when the actuator saturates. A small $T_t$ (aggressive reset) can introduce instability if the plant is slow to respond. A value in the range $\sqrt{K_p/K_i}$ to the closed-loop settling time is a common practical starting point. Here $T_t = T_{\text{end}} = 10\,\text{s}$ as a loose default.

Simulation fidelity. Gradient-based tuning is only as good as your model. If the simulation drifts significantly from the real plant, the optimal gains may not transfer. Sim-to-real gap is the main risk.

Numerical stability. Long rollout horizons can cause gradient magnitudes to grow or shrink exponentially — the same instability that plagues RNN training. For stiff systems or long horizons, consider a smaller step size, a more stable integrator (e.g. RK4), or gradient clipping.

Local minima. The closed-loop simulation loss is generally non-convex in the gains. In practice the basin of attraction is large for common plant types and reasonable initial guesses, but running from several starting points builds confidence.

Takeaways

You can treat PID gain tuning as a standard gradient-based optimisation problem by differentiating through a closed-loop simulation.
JAX makes this straightforward: write the rollout as jax.lax.scan, define the loss, call jax.grad.
The incremental form simplifies saturation and anti-windup handling: only the integral increment accumulates wind-up, and the accumulator state is transparent.
Setpoint weighting ( $b$ , $c$ ) decouples transient aggressiveness from steady-state performance. Fixing these as design constants and optimising only the gains avoids the optimizer finding degenerate solutions.
The soft anti-windup pull-back keeps the integrator state smooth when the actuator saturates, without the discontinuity of hard clamping the increment.
Gradient descent discovers control structure. Starting from a P-only controller, the optimiser finds on its own that integral action is needed — because without it the time-weighted loss is unavoidably large.

The Python script used to generate all plots is available if you want to run the experiments yourself.

J. Bradbury et al., “JAX: Composable transformations of Python+NumPy programs,” version 0.3.13, 2018. [Online]. Available: http://github.com/jax-ml/jax ↩
E. Sundström, M. Bauer, J. L. Guzmán, T. Hägglund, K. Soltesz, “A Practical Guide to PID Controller Implementation,” arXiv:2604.15918 [eess.SY], 2026. https://arxiv.org/abs/2604.15918 ↩