Generative optimal transport via forward-backward HJB matching

2026-04-08

AI Summary

This work addresses the problem of steering an uncontrolled stochastic system to a target distribution implicitly defined by samples, without requiring its explicit form, while minimizing control cost. By establishing a time-reversal duality, the backward optimal control problem is transformed into a forward Hamilton-Jacobi-Bellman (HJB) equation, enabling a novel paradigm that computes the value function and optimal policy solely from forward relaxed trajectories. The approach integrates the Cole-Hopf transformation, Feynman-Kac path integrals, and risk-sensitive control, eliminating the need for backward simulation or prior knowledge of the target distribution. It thereby unifies frameworks from stochastic optimal control, Schrödinger bridges, and nonequilibrium statistical mechanics. Experimental visualizations confirm that the learned value function effectively governs spatial regulation in controlled diffusion processes, aligning closely with theoretical predictions.

Abstract

Controlling the evolution of a many-body stochastic system from a disordered reference state to a structured target ensemble, characterized empirically through samples, arises naturally in non-equilibrium statistical mechanics and stochastic control. The natural relaxation of such a system - driven by diffusion - runs from the structured target toward the disordered reference. The natural question is then: what is the minimum-work stochastic process that reverses this relaxation, given a pathwise cost functional combining spatial penalties and control effort? Computing this optimal process requires knowledge of trajectories that already sample the target ensemble - precisely the object one is trying to construct. We resolve this by establishing a time-reversal duality: the value function governing the hard backward dynamics satisfies an equivalent forward-in-time HJB equation, whose solution can be read off directly from the tractable forward relaxation trajectories. Via the Cole-Hopf transformation and its associated Feynman-Kac representation, this forward potential is computed as a path-space free energy averaged over these forward trajectories - the same relaxation paths that are easy to simulate - without any backward simulation or knowledge of the target beyond samples. The resulting framework provides a physically interpretable description of stochastic transport in terms of path-space free energy, risk-sensitive control, and spatial cost geometry. We illustrate the theory with numerical examples that visualize the learned value function and the induced controlled diffusions, demonstrating how spatial cost fields shape transport geometry analogously to Fermat's Principle in inhomogeneous media. Our results establish a unifying connection between stochastic optimal control, Schrödinger bridge theory, and non-equilibrium statistical mechanics.

Problem

• optimal transport
• stochastic control
• non-equilibrium statistical mechanics
• Schrödinger bridge
• path-space free energy

Innovation

• optimal transport
• stochastic control
• Hamilton-Jacobi-Bellman equation
• time-reversal duality
• path-space free energy