On the Convergence of Jacobian-Free Backpropagation for Optimal Control Problems with Implicit Hamiltonians

TL;DR

This work addresses learning semi-global feedback controllers for high-dimensional optimal control problems with implicit Hamiltonians, where closed-form Hamiltonian maximizers are unavailable. It provides the first convergence guarantees for Jacobian-Free Backpropagation (JFB) used within biased stochastic gradient descent, showing convergence to stationary points of the expected objective $\mathbb{E}_x[J_x(\theta)]$ under minibatch training and standard smoothness/contractivity assumptions. The authors prove key alignment and descent lemmas, derive a finite-sample descent bound, and establish convergence in expectation and in probability to critical points, with extensions to neighborhoods under weaker assumptions. Empirically, they verify the theoretical assumptions in practice and demonstrate that JFB scales to high-dimensional, multi-agent control problems (e.g., 100 agents) where implicit differentiation is memory-prohibitive, achieving comparable performance with substantially greater efficiency compared to AD or CVXPYLayers. Overall, the paper provides both theoretical justification and practical evidence that JFB is a principled, scalable approach for learning value-function–based feedback controllers when the Hamiltonian is implicit.

Abstract

Optimal feedback control with implicit Hamiltonians poses a fundamental challenge for learning-based value function methods due to the absence of closed-form optimal control laws. Recent work~\cite{gelphman2025end} introduced an implicit deep learning approach using Jacobian-Free Backpropagation (JFB) to address this setting, but only established sample-wise descent guarantees. In this paper, we establish convergence guarantees for JFB in the stochastic minibatch setting, showing that the resulting updates converge to stationary points of the expected optimal control objective. We further demonstrate scalability on substantially higher-dimensional problems, including multi-agent optimal consumption and swarm-based quadrotor and bicycle control. Together, our results provide both theoretical justification and empirical evidence for using JFB in high-dimensional optimal control with implicit Hamiltonians.

Figures (6)

• Figure 1: Numerical verification on quadrotor experiment. (Left) Maximum singular value of the Jacobian of the operator $T_\theta$, evaluated on the worst-case sample within each minibatch, illustrating that the operator remains contractive throughout training. (Right) Smallest and largest singular values of $M_{\theta}$, computed batch-wise over all time steps for the single quadrotor experiment, confirming the full row-rank property of $M_{\theta}$ and the boundedness of $\sigma_{\max}(M_{\theta})$.
• Figure 2: Angle between $\mathbb{E}_x[\nabla_{\theta}J_x]$ and $\mathbb{E}_x[d_{x}^{JFB}]$ (Left) plotted alongside loss vs. epoch (Right). The values plotted are the largest batch-wise angles. $\mathbb{E}_x[\nabla_{\theta}J_x]$ is computed using AD. The dashed lines are at the angles $0$ and $\frac{\pi}{2}$.
• Figure 3: Comparison of JFB, automatic differentiation (AD), and CVXPYLayersagrawal2019differentiable (Implicit Differentiation) for training the value function for a quadrotor across four metrics. (Top Left) Loss versus training epochs. (Top Right) Loss plotted against cumulative runtime in minutes. (Bottom Left) Loss plotted against cumulative work units, with one work unit being one evaluation of $\frac{\partial T_{\theta}}{\partial \theta}$, which is equivalent to backpropagation through one application of $T_\theta$. (Bottom Right) Maximum GPU memory usage per training epoch.
• Figure 4: Comparison of JFB, automatic differentiation (AD), and CVXPYLayers (Implicit Differentiation) for training the value function for 6 quadrotors across three metrics. (Top Left) Loss versus training epochs. (Top Right) Loss plotted against cumulative runtime in minutes. (Bottom Left) Loss plotted against cumulative work units. (Bottom Right) Maximum GPU memory usage per training epoch.
• Figure 5: Trajectories of trained quadrotors using JFB (left panel), CVXPYLayers (middle panel), and AD (right panel).

Theorems & Definitions (19)

• Definition 4.2: Objective and Gradient and JFB
• Lemma 4.7
• Lemma 4.8
• Lemma 4.9
• Theorem 4.10
• Theorem 4.11
• Corollary 4.12
• proof
• proof
• proof