On the numerical integration of the Fokker-Planck equation driven by a mechanical force and the Bismut-Elworthy-Li formula

TL;DR

This work develops and validates numerical methods for two key PDEs in stochastic optimal control: the Fokker-Planck equation with a time-dependent mechanical potential and the Hamilton-Jacobi-Bellman equation for the value function. It combines a Girsanov-based Monte Carlo approach to FP with backward diffusion and a Bismut-Elworthy-Li gradient representation to compute optimal protocols, including extensions to degenerate diffusion. The paper demonstrates analytic and numerical examples, including Schrödinger bridge problems and a machine-learning-inspired prototype that solves for optimal drifts without spatial discretization, and confirms consistency with existing iterative methods. The methods offer scalable, high-dimensional alternatives to traditional spatial discretization, with demonstrated speedups and applicability to nanoscale thermodynamics and diffusion-model-based learning tasks.

Abstract

Optimal control theory aims to find an optimal protocol to steer a system between assigned boundary conditions while minimizing a given cost functional in finite time. Equations arising from these types of problems are often non-linear and difficult to solve numerically. In this note, we describe numerical methods of integration for two partial differential equations that commonly arise in optimal control theory: the Fokker-Planck equation driven by a mechanical potential for which we use Girsanov theorem; and the Hamilton-Jacobi-Bellman, or dynamic programming, equation for which we find the gradient of its solution using the Bismut-Elworthy-Li formula. The computation of the gradient is necessary to specify the optimal protocol. Finally, we give an example application of the numerical techniques to solving an optimal control problem without spacial discretization using machine learning.

Figures (6)

• Figure 1: Solution of a Fokker-Planck equation driven by a mechanical potential \ref{['ud:FP']} computed using Monte Carlo integration via Girsanov formula (dashed blue line). We use $\partial_q U(q) = 2\,q^3$. The initial condition is $\mathrm{p}_{t_{\iota}}(q)=\frac{1}{\sqrt{2\pi}}\exp(-q^2/2)$ at $t_{\iota}=0$. The Girsanov method is compared with an implementation of the "proximal gradient descent" method described in CaHa20, shown in orange. For the proximal gradient descent, we use $10^{4}$ samples from the initial distribution and $\gamma=0.05$ as the regularisation parameter, see CaHa20. Both methods simulate trajectories of the auxilliary stochastic process \ref{['od:main']} by the Euler-Maruyama scheme with step size $h=10^{-3}$. For the Girsanov theorem approach, we evolve $10^{3}$ trajectories from $10^{4}$ initial points in the interval $[-6,6]$. Resulting distributions are smoothed by convolution with a box filter. We use $t_{\iota}=0$ and $\mu=\beta=1$. The expected equilibrium state of the distribution is shown by the shaded area in the final panel at $t=0.75$. In our implementation, the Monte Carlo method of integration is roughly three orders of magnitude faster than the proximal gradient descent. Accompanying code for all figures can be found in the link in the Data Availability statement.
• Figure 2: Solution of a Fokker-Planck driven by a time-dependent mechanical potential computed using Monte Carlo integration via Girsanov formula (dashed blue line). The optimal protocol $U$ and reference solution (orange line) is computed using an iterative method CaHa2022. For the Girsanov theorem approach, we evolve $M = 10\,000$ trajectories from $500$ initial points in the interval $[-3,3]$ with a time step of $h=0.005$ by the Euler-Maruyama scheme. The reference solution uses the iteration method of CaHa2022, where integration of the equations \ref{['num:ivpA']} and \ref{['num:ivpB']} is also computed as a numerical average of Monte Carlo sampled trajectories, using $5\,000$ initial points from the interval $[-6,6]$. Ten total iterations are performed. Final distributions are normalized and smoothed by a convolution with a box filter. We use $t_{\iota}=0$ and $t_{f} = 0.2$ and $\mu=\beta=1$. Assigned boundary conditions (shaded blue area in first and final panels) are given by \ref{['eq:boundary_conditions_initial']} with $U_{\iota}(q) = \frac{1}{4}(q-1)^4$ and \ref{['eq:boundary_conditions_final']} with $U_{f}(q) = \frac{1}{4}(q^2-1)^2$.
• Figure 3: Solution of a Fokker-Planck equation driven by a non-linear mechanical underdamped diffusion computed by Monte Carlo integration. Panels (a)-(f) show the following: Center: the joint distribution for the momentum and position; Top: the marginal distribution of the momentum; Left: marginal distribution of the position. The optimal protocol $U$ used in the integration and the reference solutions for the marginal densities (orange) are estimated from a perturbative expansion around the overdamped limit SaBaMG2024. For the integration, we evolve $M = 10\,000$ trajectories from a set of $2601$ equally spaced points from the interval $[-5,5]\times [-5,5]$. We use a time step size $h=0.025$ and integrate over trajectories of \ref{['ud:FP']} using an Euler-Maruyama discretization. We use $t_{\iota}=0$, $t_{f} = 5$, $\beta=25$ and $\tau=m=1$. The assigned initial condition is given by \ref{['UD:initial_boundarycond']} with $U_{\iota}(q) = \frac{1}{4}(q-1)^4$ and final condition \ref{['UD:final_boundarycond']} with $U_{f}(q) = \frac{1}{4}(q^2-1)^2$.
• Figure 4: Gradient of the value function, i.e. the gradient of the solution to the Hamilton-Jacobi-Bellman equation \ref{['eq:hjb_coupled']}, computed using the Bismut-Elworthy-Li formula (BEL) (dashed blue line) described in Algorithm \ref{['alg_dvbel_over']}. We sample $10\ 000$ trajectories of the stochastic process \ref{['od:FP']} from $500$ initial points in the interval $[-3,3]$, discretized by the Euler-Maruyama scheme with time step size $h=0.005$ and compute the BEL weights along the trajectories. The optimal control protocol $U$ and reference solution (orange) used is computed by the iteration as in Fig. \ref{['fig:overdamped_fpk']}. Numerical parameters and boundary conditions are as in Fig. \ref{['fig:overdamped_fpk']}. We use $t_{\iota} = 0$, $t_{f} =0.2$ and $\mu=\beta=1$.
• Figure 5: The gradient of the optimal control potential minimizing the Kullback-Leibler divergence \ref{['eq:kl_divergence']} in the underdamped dynamics. We compute the stationarity condition \ref{['num_ud_stationarity_condition']} at $t=0$ using the gradient of the solution of the Hamilton-Jacobi-Bellman equation \ref{['BEL:dp']} using the Bismut-Elworthy-Li formula (Monte Carlo w. BEL) (blue line) in Algorithm \ref{['alg_dvbel_under']}. The optimal control protocol $U$ and terminal condition $\varphi$ of \ref{['BEL:dp']} are found using numerical integration of the system of equations described in Section IV of SaBaMG2024, using a fourth order co-location method from the DifferentialEquations.jl library differentialequationsjl. We use Gaussian boundary conditions: the initial and final position and momentum means are set as zero; the initial and final cross-correlation is zero; the initial variances are set to $1$; the final position variance is $1.7$, and the final momentum variance is $1$. We sample $10\ 000$ independent trajectories of the stochastic process \ref{['ud:sde']} started from $500$ sample points in the interval $[-5,5]\times[-5,5]$ using an Euler-Maruyama discretization with time-step $h=0.01$. We use $t_{\iota} = 0$, $t_{f} =1$ and $\beta=\tau = m = 1$.

Theorems & Definitions (5)

• Remark
• Proposition 2.1
• Remark
• Proposition 3.1
• Remark