Neural Schrödinger Bridge With Sinkhorn Losses: Application to Data-Driven Minimum Effort Control of Colloidal Self-Assembly

We show that the minimum effort control of colloidal self-assembly (SA) can be naturally formulated in the order-parameter space as a generalized Schrödinger bridge problem (GSBP)—a class of fixed-horizon stochastic optimal control problems that originated in the works of Erwin Schrödinger in the early 1930s. In recent years, this class of problems has seen a resurgence of research activities in the control and machine-learning communities. Different from the existing literature on the theory and computation for such problems, the controlled drift and diffusion coefficients for colloidal SA are typically nonaffine in control and are difficult to obtain from physics-based modeling. We deduce the conditions of optimality for such generalized problems and show that the resulting system of equations is structurally very different from the existing results in a way that standard computational approaches no longer apply. Thus motivated, we propose a data-driven learning and control framework, named “neural Schrödinger bridge,” to solve such generalized Schrödinger bridge problems by innovating on recent advances in neural networks (NNs). We illustrate the effectiveness of the proposed framework using a numerical case study of colloidal SA. We learn the controlled drift and diffusion coefficients as two NNs using molecular dynamics (MD) simulation data and then use these two to train a third network with Sinkhorn losses designed for distributional endpoint constraints, specific for this class of control problems.

I. Introduction

Motivated by feedback control of colloidal self-assembly (SA), this work focuses on learning the solution of the nonlinear stochastic optimal control problems over a given fixed time horizon of the form $$\begin{align*} &\underset { \boldsymbol {u}\in \mathcal {U}}{\inf }\:\mathbb {E}_{\mu ^{ \boldsymbol {u}}}\left [{\int _{0}^{T}\frac {1}{2}\| {u}\left ({t, \boldsymbol {x}}\right)\|_{2}^{2}\: {\mathrm{d}} t}\right] \tag{1a}\\ &\text {subject to}\,{\mathrm{d}}\, \boldsymbol {x} = \boldsymbol {f}\left ({t, \boldsymbol {x}, \boldsymbol {u}}\right) {\mathrm{d}} t + \sqrt {2}\: \boldsymbol {g}\left ({t, \boldsymbol {x}, \boldsymbol {u}}\right) {\mathrm{d}} \boldsymbol {w} \tag{1b}\\ &\hphantom {\text {subject to} } \boldsymbol {x}\left ({t=0}\right)\sim \mu _{0}\;\left ({\text {given}}\right), \; \boldsymbol {x}\left ({t=T}\right)\sim \mu _{T}\;\left ({\text {given}}\right)\tag{1c}\end{align*}$$ where denote the prescribed probability measures over the state space at and , respectively. The constraint in (1b) is a controlled Itô stochastic differential equation (SDE) with the state vector , the control vector , and the standard Wiener process . For the solution to the SDE (1b) to be for colloidal SA systems, the state vector represents suitable order parameters. The drift coefficient is a vector field given by mapping , and the diffusion coefficient is a matrix field given by mapping . For the SDE solutions to be well-posed, we will detail suitable smoothness assumptions on and .