Hessian-augmented Supervised Learning for Hamilton-Jacobi-Bellman PDEs

Technical Summary: Hessian-Augmented Supervised Learning for Hamilton–Jacobi–Bellman PDEs

Problem Statement
The paper addresses the synthesis of optimal feedback control laws for deterministic optimal control problems with nonlinear control-affine dynamics. The standard approach involves solving the Hamilton–Jacobi–Bellman (HJB) equation, a fully nonlinear partial differential equation posed over an $n$-dimensional state space. Classical grid-based discretizations suffer from the "curse of dimensionality," scaling exponentially with the state dimension $n$, which restricts their applicability to low-dimensional systems ($n \lesssim 5$). While various non-standard numerical strategies (e.g., causality-free characteristics, tensor decompositions, and deep learning) have been proposed to mitigate this, the challenge of approximating value functions and recovering feedback laws in high dimensions ($n \gg 5$) remains significant.

Methodology
The authors propose a data-driven framework that approximates the value function $V(x)$ using supervised learning with derivative augmentation. The core methodology consists of three components:

1. Data Generation via PMP and Riccati Equations:
   Instead of solving the HJB PDE directly, the method generates training data by solving the first-order necessary optimality conditions of the Pontryagin Maximum Principle (PMP). For a set of initial conditions, the PMP yields a two-point boundary value problem (TPBVP) describing the optimal state and adjoint trajectories.

   • Zeroth-order: The cost along the optimal trajectory provides the value function $V(x)$.
   • First-order: The adjoint variable $p(t)$ coincides with the gradient $\nabla V(x)$ along the trajectory.
   • Second-order: A matrix Riccati equation, integrated backward along the optimal trajectory, yields the Hessian $\nabla^2 V(x)$.
     This process avoids state-space discretization and scales linearly with the number of initial conditions rather than exponentially with dimension.

2. Sparse Polynomial Regression:
   The value function is approximated by a sparse polynomial model $\tilde{V}(x) = \sum \theta_k \Phi_k(x)$. The basis functions are constructed from tensor-product Legendre polynomials on a hyperbolic cross index set, enriched with low-degree total-degree terms to capture essential multivariate couplings.

3. Derivative-Augmented Least Squares:
   The regression minimizes a weighted least-squares objective that simultaneously penalizes errors in the function values, gradients, and Hessians.

   • Each training sample $x_j$ contributes $1$ equation for the value, $n$ equations for the gradient, and $n(n+1)/2$ equations for the Hessian.
   • This results in a total equation count of $M = N(1 + n + n(n+1)/2)$, where $N$ is the number of samples.
   • To manage computational costs in high dimensions, a "partial Hessian" strategy is introduced, where Hessian equations are enforced only for a fraction $\rho \in (0, 1]$ of the training samples.

Key Contributions

• Extension to Second-Order Information: The paper extends PMP-based data generation to include Hessian information via Riccati equations, systematically developing this direction for optimal control.
• Partial Hessian Strategy: A tunable strategy is introduced to subsample second-order data, balancing the high cost of data generation and regression system size against the benefits of curvature information.
• Basis Interaction: The study investigates the interaction between second-order augmentation and hyperbolic cross polynomial bases enriched with low-degree total-degree terms.
• Systematic Benchmarking: The approach is validated on problems of increasing state dimension ($n=2, 6, 19$), measuring accuracy in discrete Sobolev norms, HJB residuals, and closed-loop control performance.

Results
Numerical experiments demonstrate that second-order augmentation significantly reduces sample complexity compared to zeroth- and first-order methods:

• Analytical Benchmark: For a 12-dimensional test function, achieving a prescribed $L^2$ error required $N=50$ samples with second-order augmentation, whereas first-order regression required $N \approx 350$, and zeroth-order regression failed to reach the target accuracy within the tested range.
• 2D Van der Pol Oscillator: Second-order augmentation produced smooth, accurate reconstructions with $N=25$ samples, whereas lower-order methods exhibited oscillations. In terms of sample efficiency, second-order methods required up to 11.7 times fewer samples than zeroth-order methods to achieve a control error tolerance of 0.1%.
• 6D Satellite Attitude Control: On an enriched basis ($q=1,800$), first-order regression remained underdetermined ($M/q < 1$) for $N \lesssim 250$, while second-order augmentation achieved a well-determined regime at $N=100$ and high accuracy at $N=200$.
• 19D Allen–Cahn Equation: In this high-dimensional regime, the amplification of equations per sample ($n(n+1)/2 = 190$ additional equations) became a structural necessity. First-order regression remained underdetermined at moderate budgets, while second-order augmentation was the only formulation capable of producing a meaningful approximation. Partial Hessian strategies ($\rho=0.5$) achieved near-full accuracy with half the computational cost.

Significance and Claims
The paper claims that the primary motivation for incorporating Hessian data is the reduction of sample complexity. By leveraging the multiplicative gain in the number of linear equations per sample, the method can solve high-dimensional optimal control problems that are intractable for lower-order supervised learning approaches due to underdetermined systems.

The authors emphasize that in high dimensions (e.g., $n=19$), second-order augmentation shifts from being merely an accuracy improvement to a structural requirement for solvability. The method is causality-free, allowing for flexible design choices in sampling and basis selection. The dominant computational cost is the offline generation of the dataset (solving TPBVPs and Riccati equations), while online evaluation of the feedback law is efficient, requiring only basis function evaluations and dot products.

The work concludes that while the method is currently limited to deterministic control problems with smooth dynamics, the framework offers a robust pathway for approximating value functions in dimensions where traditional grid methods fail and lower-order learning methods struggle with sample inefficiency.