skip to main content

Title: Wasserstein Proximal Algorithms for the Schr\"{o}dinger Bridge Problem: Density Control with Nonlinear Drift
We study the Schr{\"o}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. 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\"{o}dinger factors". 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. We employ a proximal algorithm developed in our prior work to solve these IVPs and compute the Schr\"{o}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.
Award ID(s):
Publication Date:
Journal Name:
IEEE Transactions on Automatic Control
Page Range or eLocation-ID:
1 to 1
Sponsoring Org:
National Science Foundation
More Like this
  1. We develop a new computational framework to solve the partial differential equations (PDEs) governing the flow of the joint probability density functions (PDFs) in continuous-time stochastic nonlinear systems. The need for computing the transient joint PDFs subject to prior dynamics arises in uncertainty propagation, nonlinear filtering and stochastic control. Our methodology breaks away from the traditional approach of spatial discretization or function approximation – both of which, in general, suffer from the “curse-of-dimensionality”. In the proposed framework, we discretize time but not the state space. We solve infinite dimensional proximal recursions in the manifold of joint PDFs, which in the small time-step limit, is theoretically equivalent to solving the underlying transport PDEs. The resulting computation has the geometric interpretation of gradient flow of certain free energy functional with respect to the Wasserstein metric arising from the theory of optimal mass transport. We show that dualization along with an entropic regularization, leads to a cone-preserving fixed point recursion that is proved to be contractive in Thompson metric. A block co-ordinate iteration scheme is proposed to solve the resulting nonlinear recursions with guaranteed convergence. This approach enables remarkably fast computation for non-parametric transient joint PDF propagation. Numerical examples and various extensions aremore »provided to illustrate the scope and efficacy of the proposed approach.« less
  2. Abstract We study Bayesian data assimilation (filtering) for time-evolution Partial differential equations (PDEs), for which the underlying forward problem may be very unstable or ill-posed. Such PDEs, which include the Navier–Stokes equations of fluid dynamics, are characterized by a high sensitivity of solutions to perturbations of the initial data, a lack of rigorous global well-posedness results as well as possible non-convergence of numerical approximations. Under very mild and readily verifiable general hypotheses on the forward solution operator of such PDEs, we prove that the posterior measure expressing the solution of the Bayesian filtering problem is stable with respect to perturbations of the noisy measurements, and we provide quantitative estimates on the convergence of approximate Bayesian filtering distributions computed from numerical approximations. For the Navier–Stokes equations, our results imply uniform stability of the filtering problem even at arbitrarily small viscosity, when the underlying forward problem may become ill-posed, as well as the compactness of numerical approximants in a suitable metric on time-parametrized probability measures.
  3. We develop algorithms to automate discovery of stochastic dynamical system models from noisy, vector-valued time series. By discovery, we mean learning both a nonlinear drift vector field and a diagonal diffusion matrix for an Itô stochastic differential equation in Rd . We parameterize the vector field using tensor products of Hermite polynomials, enabling the model to capture highly nonlinear and/or coupled dynamics. We solve the resulting estimation problem using expectation maximization (EM). This involves two steps. We augment the data via diffusion bridge sampling, with the goal of producing time series observed at a higher frequency than the original data. With this augmented data, the resulting expected log likelihood maximization problem reduces to a least squares problem. We provide an open-source implementation of this algorithm. Through experiments on systems with dimensions one through eight, we show that this EM approach enables accurate estimation for multiple time series with possibly irregular observation times. We study how the EM method performs as a function of the amount of data augmentation, as well as the volume and noisiness of the data.
  4. Abstract Combining the classical theory of optimal transport with modern operator splitting techniques, we develop a new numerical method for nonlinear, nonlocal partial differential equations, arising in models of porous media, materials science, and biological swarming. Our method proceeds as follows: first, we discretize in time, either via the classical JKO scheme or via a novel Crank–Nicolson-type method we introduce. Next, we use the Benamou–Brenier dynamical characterization of the Wasserstein distance to reduce computing the solution of the discrete time equations to solving fully discrete minimization problems, with strictly convex objective functions and linear constraints. Third, we compute the minimizers by applying a recently introduced, provably convergent primal dual splitting scheme for three operators (Yan in J Sci Comput 1–20, 2018). By leveraging the PDEs’ underlying variational structure, our method overcomes stability issues present in previous numerical work built on explicit time discretizations, which suffer due to the equations’ strong nonlinearities and degeneracies. Our method is also naturally positivity and mass preserving and, in the case of the JKO scheme, energy decreasing. We prove that minimizers of the fully discrete problem converge to minimizers of the spatially continuous, discrete time problem as the spatial discretization is refined. We conclude withmore »simulations of nonlinear PDEs and Wasserstein geodesics in one and two dimensions that illustrate the key properties of our approach, including higher-order convergence our novel Crank–Nicolson-type method, when compared to the classical JKO method.« less
  5. We are concerned with free boundary problems arising from the analysis of multidimensional transonic shock waves for the Euler equations in compressible fluid dynamics. In this expository paper, we survey some recent developments in the analysis of multidimensional transonic shock waves and corresponding free boundary problems for the compressible Euler equations and related nonlinear partial differential equations (PDEs) of mixed type. The nonlinear PDEs under our analysis include the steady Euler equations for potential flow, the steady full Euler equations, the unsteady Euler equations for potential flow, and related nonlinear PDEs of mixed elliptic–hyperbolic type. The transonic shock problems include the problem of steady transonic flow past solid wedges, the von Neumann problem for shock reflection–diffraction, and the Prandtl–Meyer problem for unsteady supersonic flow onto solid wedges. We first show how these longstanding multidimensional transonic shock problems can be formulated as free boundary problems for the compressible Euler equations and related nonlinear PDEs of mixed type. Then we present an effective nonlinear method and related ideas and techniques to solve these free boundary problems. The method, ideas, and techniques should be useful to analyze other longstanding and newly emerging free boundary problems for nonlinear PDEs.