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

TL;DR

This paper proves convergence guarantees for Jacobian-Free Backpropagation in high-dimensional optimal control problems with implicit Hamiltonians, supported by empirical results on complex multi-agent systems.

Contribution

It provides the first convergence guarantees for JFB in stochastic minibatch settings and demonstrates scalability to high-dimensional control tasks.

Findings

JFB converges to stationary points in stochastic minibatch settings.

Scalability demonstrated on multi-agent control problems.

Empirical results support theoretical convergence guarantees.

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.