Wasserstein Proximal Algorithms for the Schrödinger Bridge Problem: Density Control With Nonlinear Drift

In this article, we study the Schrödinger bridge problem (SBP) with nonlinear prior dynamics. In control-theoretic language, this is a problem of minimum effort steering of a given joint state probability density function (PDF) to another over a finite-time horizon, subject to a controlled stochastic differential evolution of the state vector. As such, it can be seen as a stochastic optimal control problem in continuous time with endpoint density constraints—A topic that originated in the physics literature in 1930s, and in the recent years, has garnered burgeoning interest in the systems-control community. For generic nonlinear drift, we reduce the SBP to solving a system of forward and backward Kolmogorov partial differential equations (PDEs) that are coupled through the boundary conditions, with unknowns being the “Schrödinger factors”—so named since their product at any time yields the optimal controlled joint state PDF at that time. We show that if the drift is a gradient vector field, or is of mixed conservative–dissipative nature, then it is possible to transform these PDEs into a pair of initial value problems (IVPs) involving the same forward Kolmogorov operator. Combined with a recently proposed fixed point recursion that is contractive in the Hilbert metric, this opens up the possibility to numerically solve the SBPs in these cases by computing the Schrödinger factors via a single IVP solver for the corresponding (uncontrolled) forward Kolmogorov PDE. The flows generated by such forward Kolmogorov PDEs, for the two aforementioned types of drift, in turn, enjoy gradient descent structures on the manifold of joint PDFs with respect to suitable distance functionals. We employ a proximal algorithm developed in our prior work that exploits this geometric viewpoint, to solve these IVPs and compute the Schrödinger factors via weighted scattered point cloud evolution in the state space. We provide the algorithmic details and illustrate the proposed framework of solving the SBPs with nonlinear prior dynamics by numerical examples.