Hamiltonian Theory and Computation of Optimal Probability Density Control in High Dimensions

TL;DR

This work addresses optimal density control in high-dimensional spaces by developing a Wasserstein-free theoretical framework and a scalable, mesh-free computational method. It derives Pontryagin's Maximum Principle and the Hamilton-Jacobi-Bellman equation directly in $L^2$-type density spaces, avoiding Wasserstein geometry while maintaining rigorous optimality conditions. The proposed algorithm parameterizes both the control and adjoint functions with deep neural networks and represents the evolving density via a particle ensemble solved with neural ODEs, accompanied by convergence guarantees to local optima. Numerical experiments up to dimensionality $d=100$ demonstrate effective density steering in scenarios with particle interactions and obstacles, highlighting practical scalability for complex, high-dimensional systems.

Abstract

We develop a general theoretical framework for optimal probability density control and propose a numerical algorithm that is scalable to solve the control problem in high dimensions. Specifically, we establish the Pontryagin Maximum Principle (PMP) for optimal density control and construct the Hamilton-Jacobi-Bellman (HJB) equation of the value functional through rigorous derivations without any concept from Wasserstein theory. To solve the density control problem numerically, we propose to use reduced-order models, such as deep neural networks (DNNs), to parameterize the control vector-field and the adjoint function, which allows us to tackle problems defined on high-dimensional state spaces. We also prove several convergence properties of the proposed algorithm. Numerical results demonstrate promising performances of our algorithm on a variety of density control problems with obstacles and nonlinear interaction challenges in high dimensions.

Figures (4)

• Figure 1: (Test 1) Visualization of particle movements under $u$ solved by Algorithm \ref{['alg:density-control']} for the density control problem with running reward \ref{['eq:r-collision']}. Particle distributions are show by colored circles at different times $t=0$ (red), $0.25$ (orange), $0.5$ (green), $0.75$ (blue), and $1$ (purple). Top row and bottom row show the results with and without particle collision preventions, i.e., $\gamma=5$ and $\gamma=0$, respectively. The plots from left to right columns show the particle distributions in the $(x_1,x_2)$, $(x_3,x_4)$, $(x_5,x_6)$, and $(x_7,x_8)$ coordinates, respectively.
• Figure 2: (Test 2) Visualization of particle movements under $u$ solved by Algorithm \ref{['alg:density-control']} for the density control problem with cylindrical obstacle \ref{['eq:R-cylindrical-obstacle']}. Particle distributions are show by colored circles at different times $t=0$ (red), $0.25$ (orange), $0.5$ (green), $0.75$ (blue), and $1$ (purple). The plots show the particle distributions in the $(x_1,x_2)$ coordinates for dimensions $d=30$ (left) and $d=100$ (right).
• Figure 3: (Test 3) Visualization of particle movements under $u$ solved by Algorithm \ref{['alg:density-control']} for the density control problem with both wedge obstacles and particle interaction. Particle distributions are show by colored circles at different times $t=0$ (red), $0.25$ (orange), $0.5$ (green), $0.75$ (blue), and $1$ (purple). The plots show the particle distributions in the $(x_1,x_2)$ coordinates for with weight $\gamma=0$ (left) and $\gamma=1$ (right) of interaction term in \ref{['eq:r-collision']}.
• Figure 4: (Left) the cost functional value $I[u_{\theta_k}]$ versus iteration $k$ from 0 to $K=100$; (Right) the ratio $H_t^K/|H_0^K|$ versus time $t$. In both plots, Test 1 with interaction only is shown by the red curves, Test 2 with the cylinder obstacle by blue curves, and Test 3 with both interaction and the squeezing obstacle by green curves.

Theorems & Definitions (20)

• Definition 3.1: Adjoint equation
• Definition 3.2: Hamiltonian functional
• Proposition 3.3: Control Hamiltonian system
• Definition 3.4: Perturbation function
• Lemma 3.5
• proof
• Lemma 3.6
• proof
• Theorem 3.7: Pontryagin maximum principle for optimal density control
• proof