Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A supervised learning scheme for computing Hamilton-Jacobi equation via density coupling

Jianbo Cui, Shu Liu, Haomin Zhou

TL;DR

The paper introduces a density-coupled supervised learning framework to solve high-dimensional first-order Hamilton–Jacobi equations by leveraging Wasserstein Hamiltonian flow. By coupling a continuity equation with the HJ equation, it recasts the problem as a regression using a Bregman divergence, with training data generated from symplectic integration of Hamiltonian ODEs. A rigorous residual bound is established, showing how the L1 residual depends on numerical, sampling, and training errors, and the method can extend the solution beyond caustics via μ_t-weighted momentum. Numerical experiments across separable and non-separable Hamiltonians, including caustic formation and inverted pendulum control with terminal density constraints, validate the approach and demonstrate its ability to capture multi-valued momentum structures and extend solutions beyond classical blow-up times.

Abstract

We propose a supervised learning scheme for the first order Hamilton--Jacobi PDEs in high dimensions. The scheme is designed by using the geometric structure of Wasserstein Hamiltonian flows via a density coupling strategy. It is equivalently posed as a regression problem using the Bregman divergence, which provides the loss function in learning while the data is generated through the particle formulation of Wasserstein Hamiltonian flow. We prove a posterior estimate on $L^1$ residual of the proposed scheme based on the coupling density. Furthermore, the proposed scheme can be used to describe the behaviors of Hamilton--Jacobi PDEs beyond the singularity formations on the support of coupling density. Several numerical examples with different Hamiltonians are provided to support our findings.

A supervised learning scheme for computing Hamilton-Jacobi equation via density coupling

TL;DR

The paper introduces a density-coupled supervised learning framework to solve high-dimensional first-order Hamilton–Jacobi equations by leveraging Wasserstein Hamiltonian flow. By coupling a continuity equation with the HJ equation, it recasts the problem as a regression using a Bregman divergence, with training data generated from symplectic integration of Hamiltonian ODEs. A rigorous residual bound is established, showing how the L1 residual depends on numerical, sampling, and training errors, and the method can extend the solution beyond caustics via μ_t-weighted momentum. Numerical experiments across separable and non-separable Hamiltonians, including caustic formation and inverted pendulum control with terminal density constraints, validate the approach and demonstrate its ability to capture multi-valued momentum structures and extend solutions beyond classical blow-up times.

Abstract

We propose a supervised learning scheme for the first order Hamilton--Jacobi PDEs in high dimensions. The scheme is designed by using the geometric structure of Wasserstein Hamiltonian flows via a density coupling strategy. It is equivalently posed as a regression problem using the Bregman divergence, which provides the loss function in learning while the data is generated through the particle formulation of Wasserstein Hamiltonian flow. We prove a posterior estimate on residual of the proposed scheme based on the coupling density. Furthermore, the proposed scheme can be used to describe the behaviors of Hamilton--Jacobi PDEs beyond the singularity formations on the support of coupling density. Several numerical examples with different Hamiltonians are provided to support our findings.
Paper Structure (33 sections, 7 theorems, 160 equations, 16 figures, 1 table, 1 algorithm)

This paper contains 33 sections, 7 theorems, 160 equations, 16 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Density coupling strategy
  3. Regression problem via the Bregman divergence
  4. Extending solution beyond singularity formation of HJ equation
  5. Connections with Hamiltonian dynamics
  6. Connections with a variational problem
  7. Supervised learning scheme via density coupling
  8. Algorithm
  9. Bound on the residual
  10. Numerical tests
  11. Residual and error bounds
  12. Behavior of neural network solution near caustics
  13. Solving HJ equations
  14. Example with quadratic potential
  15. Swarming particles with LQC of inverted pendulum with terminal distribution constrain
  16. ...and 18 more sections

Key Result

Lemma 2.1

\newlabelthm Legendre and Bregman divergence0 Suppose $f\in\mathcal{C}^2(\mathbb{R}^d)$ is $\alpha$-strongly convex and $L$-strongly smooth ($\alpha, L >0$), i.e., $\alpha I_d \preceq \nabla^2 f(q) \preceq L I_d$ for any $q\in\mathbb R^d$. Then, the Legendre transform $f^*$ of $f$ belongs to $\mat

Figures (16)

  • Figure 1: (Top row) Heat graphs of the residual $\textrm{Res}(x, t)$ of the numerical solution $\psi_\theta$ and the sample points (black) at different time stages $t$. (Bottom row) Heat graphs of the error $\textrm{Err}(x, t)$ of the numerical solution $\psi_\theta$ and the sample points (black) at different time stages $t$.
  • Figure 1: \ref{['subfig1']}: Plots of $\boldsymbol{\eta}^\top\nabla\psi_\theta(z\boldsymbol{\eta})$ with the exact weighted momentum $\partial_z \widehat{f}(z, t)$ on $[-\pi,\pi]$ at $t=1.5$; \ref{['subfig2']}: Plots of the error $|\boldsymbol{\eta}^\top\nabla\psi_\theta(z\boldsymbol{\eta}) - \partial_z \widehat{f}(z, t)|$ on $[-\pi,\pi]$ at $t=1.5$. We test $\psi_\theta$ as ResNets using different activation functions; \ref{['subfig3']}-\ref{['subfig4']}: Same plots at $t=2.5$.
  • Figure 1: Plots of vector fields $\nabla \psi_\theta(\cdot, t)$ (green) with momentums of samples (red) at different time stages on the $5\text{th}-15\text{th}$ plane.
  • Figure 1: Left: Plot of trajectories of the Kepler system for $0\leq t \leq 9$ computed using different numerical integrators; Middle: Training loss (Störmer-Verlet scheme) vs. iteration; Right: Plot of average Hamiltonian $\bar{H}_\theta(t)$ vs $t$ for forward Euler, Störmer-Verlet, and Runge-Kutta(RK4) schemes.
  • Figure 2: Average error versus sample size plots ($\log_2-\log_2$) for 2D and 10D HJ equation (plots with confidence interval ($25\%-75\%$) based on $40$ sets of data)
  • ...and 11 more figures

Theorems & Definitions (15)

  • Definition 2.1: Bregman divergence BREGMAN1967200
  • Lemma 2.1
  • Proposition 2.1
  • Theorem 2.1: Consistency
  • Lemma 2.2
  • Theorem 3.1: Posterior estimation on $L^1$ residual of Hamilton-Jacobi equation
  • Proof 1
  • Lemma 1
  • Proof 2
  • Proof 3: Proof of Lemma 2.1
  • ...and 5 more