Operator Splitting, Policy Iteration, and Machine Learning for Stochastic Optimal Control

Abstract: We propose a splitting approach to solve the second-order Hamilton--Jacobi equation, reducing it to a heat step and a purely first-order step. The latter is implemented using a gradient value policy iteration algorithm, enabling efficient characteristic-based machine learning methods. We establish convergence rates for the splitting method. In particular, the $L<sup>\infty$ error is bounded below by $\mathcal{O}(h)$ and above by $\mathcal{O}(h<sup>{1/7})$ for Lipschitz initial data; this improves to $\mathcal{O}(h<sup>{1/5})$ for semiconcave data and to $\mathcal{O}(h<sup>{1/3})$ for $C<sup>2$ data. We also prove an upper $L<sup>1$ error estimate of order $\mathcal{O}(h<sup>{1/2})$ in the periodic setting, where $h$ is the splitting step. For the first-order step, we provide a weighted $L<sup>2$ error analysis that shows exponential convergence. Each iteration solves linear characteristic equations and learns the value function by minimizing a weighted value gradient loss. The approach yields stable and accurate numerical results.

• The paper develops a heat–Hamilton–Jacobi operator-splitting method with proven error bounds from O((Tεh)^{1/7}) for Lipschitz data to O((Tεh)^{1/3}) for C² data, plus an O(√(εh)) periodic L¹ estimate.
• The paper proves geometric convergence of value-gradient policy iteration in a weighted space-time L² norm for sufficiently short horizons, with a freely chosen exponential weight rather than a discount-factor restriction.
• The paper combines characteristic-based transport solvers with neural networks or radial basis functions, achieving HJB residuals mostly below 0.02 in experiments up to dimension 32 using only 12–20 trajectories and 30 policy iterations.

Overview and motivation

This paper addresses the numerical solution of the second-order Hamilton–Jacobi–Bellman (HJB) equation

$u_t + H(x,Du) = \varepsilon \Delta u$

on $\mathbb{R}^d$, with $H$ convex and coercive and $u_0 \in W^{1,\infty}(\mathbb{R}^d)$. The parameter $\varepsilon$ interpolates between deterministic ($\varepsilon = 0$) and stochastic ($\varepsilon > 0$) optimal control. The authors' central idea is an operator splitting scheme that decomposes each time step into a pure heat step (trivially computable via the heat kernel) followed by a purely first-order Hamilton–Jacobi step. The first-order step is then solved by a value-gradient policy iteration algorithm (PI-$\lambda$), which reduces each iteration to linear transport equations along characteristics and is amenable to machine learning implementation.

The paper makes two distinct contributions: a quantitative error analysis of the splitting scheme — which the authors state is the first result of its kind for this type of Trotter–Kato product in the viscosity setting — and a convergence proof for the PI-$\lambda$ iteration on the first-order subproblem in a weighted space-time $L^2$ norm. Numerical experiments on quadratic control problems up to dimension $\mathbb{R}^d$0 support the approach.

The splitting scheme and its error analysis

Writing $\mathbb{R}^d$1 for the heat semigroup and $\mathbb{R}^d$2 for the solution operator of the first-order equation $\mathbb{R}^d$3, the splitting approximation is defined by iterating

$\mathbb{R}^d$4

The error analysis rests on commutator estimates between the heat operator and the Hamiltonian. Exploiting convexity of $\mathbb{R}^d$5 through its Legendre transform, the authors prove a linear lower bound on the commutator $\mathbb{R}^d$6 of order $\mathbb{R}^d$7. For the upper bound, they introduce a regularized comparison function built from vanishing viscosity approximations of the first-order step, splitting the total error as $\mathbb{R}^d$8; the second term is controlled by classical semiconcavity-based vanishing viscosity estimates at the optimal rate $\mathbb{R}^d$9.

A key technical ingredient is Proposition 2.4, which shows that Lipschitz bounds, semiconcavity constants, and uniform lower Hessian bounds are preserved uniformly across all splitting iterations. The authors emphasize that the most delicate part is the uniform lower bound on $H$0, since the heat operator accumulates errors across iterations; this estimate relies crucially on the exact linear rate of the lower commutator bound — any weaker rate would not suffice.

The main quantitative results are as follows:

Initial data	Upper bound on $H$1
Lipschitz	$H$2
Lipschitz + semiconcave	$H$3
$H$4	$H$5

In all cases there is a matching linear lower bound $H$6, so the scheme converges from below at rate $H$7. The rates improve monotonically with regularity of $H$8, obtained by balancing mollification, viscosity, and splitting parameters ($H$9, $u_0 \in W^{1,\infty}(\mathbb{R}^d)$0 in the general case). In the periodic setting with semiconcave data, the authors additionally prove an averaged $u_0 \in W^{1,\infty}(\mathbb{R}^d)$1 estimate of order $u_0 \in W^{1,\infty}(\mathbb{R}^d)$2, using BV bounds on $u_0 \in W^{1,\infty}(\mathbb{R}^d)$3 that remain uniformly bounded throughout the iteration. They note that analogous weighted $u_0 \in W^{1,\infty}(\mathbb{R}^d)$4 estimates should follow by the same techniques but defer the details to future work.

An important caveat: these rates are for the splitting error only, measured against the exact solution of the viscous HJB equation, and depend on the product $u_0 \in W^{1,\infty}(\mathbb{R}^d)$5; the analysis assumes the structural conditions $u_0 \in W^{1,\infty}(\mathbb{R}^d)$6–$u_0 \in W^{1,\infty}(\mathbb{R}^d)$7, including uniform convexity of $u_0 \in W^{1,\infty}(\mathbb{R}^d)$8 in the momentum variable and bounded second derivatives.

Exponential convergence of value-gradient policy iteration

For the first-order step, the Hamiltonian arises from a deterministic control problem with dynamics $u_0 \in W^{1,\infty}(\mathbb{R}^d)$9 and running cost $\varepsilon$0. Rather than iterating on the value function directly, the PI-$\varepsilon$1 algorithm iterates on the costate $\varepsilon$2: given a policy $\varepsilon$3, one solves a linear PDE for $\varepsilon$4 along characteristics, then updates the policy via pointwise maximization. Because the equations for the components of $\varepsilon$5 decouple, they can be solved in parallel.

Under assumptions $\varepsilon$6–$\varepsilon$7 — affine control dependence of $\varepsilon$8, strong convexity of $\varepsilon$9 in $\varepsilon = 0$0, and growth/regularity conditions — the authors establish uniform bounds on $\varepsilon = 0$1, $\varepsilon = 0$2, and $\varepsilon = 0$3 provided the horizon satisfies $\varepsilon = 0$4 for explicit constants $\varepsilon = 0$5. The main convergence theorem states that the weighted error

$\varepsilon = 0$6

satisfies $\varepsilon = 0$7 whenever $\varepsilon = 0$8 and $\varepsilon = 0$9, i.e., geometric (exponential) convergence. The proof proceeds by energy estimates on the difference of successive costate equations, with the weight $\varepsilon > 0$0 controlling boundary terms from integration by parts under the linear growth of $\varepsilon > 0$1.

Two features distinguish this from prior work: the analysis covers the time-dependent problem (extending the time-independent results of Bensoussan, Han, Yam, and Zhou), and the weight parameter $\varepsilon > 0$2 can be chosen freely rather than being tied to a discount factor in the equation. This removes the restriction in earlier work that the discount factor be sufficiently large. The smallness condition on $\varepsilon > 0$3 remains a genuine limitation inherited from the Grönwall-type arguments.

Machine learning implementation and numerics

Each policy iteration requires solving stationary transport equations of the form $\varepsilon > 0$4. Along characteristics, these admit explicit integral representation formulas for both $\varepsilon > 0$5 and its gradient, generating labeled training data along sampled trajectories. A parametric model $\varepsilon > 0$6 (radial basis functions or neural networks) is trained by minimizing a loss that combines mean-squared value error and gradient error with mixing weight $\varepsilon > 0$7, gradients computed via automatic differentiation and optimized with ADAM.

Experiments use the quadratic control problem $\varepsilon > 0$8, $\varepsilon > 0$9, with Hamiltonian $\lambda$0 and initial data $\lambda$1 on $\lambda$2. Three configurations were tested: $\lambda$3 with $\lambda$4 (no splitting needed), $\lambda$5 with $\lambda$6, and $\lambda$7 with $\lambda$8. Accuracy is measured by the mean HJB residual over $\lambda$9 sample points. Reported residuals range roughly between $\lambda$0 and $\lambda$1, with most configurations below $\lambda$2, achieved with remarkably few resources: 12–20 characteristic trajectories, 1000 ADAM steps per iteration, and 30 policy iterations. Notably, accuracy does not degrade substantially when the number of trajectories is reduced, suggesting robustness of the method to limited data coverage. One anomalous entry (residual $\lambda$3 for $\lambda$4, $\lambda$5 in the $\lambda$6 test with random $\lambda$7) indicates occasional instability, though the paper does not analyze when such failures occur.

Limitations and open questions

Several restrictions qualify the results. The $\lambda$8 splitting rates require uniform convexity of $\lambda$9 in $L^2$0 and bounded second derivatives ($L^2$1), stronger than what well-posedness alone demands; the authors acknowledge imposing these for technical convenience. The exponential convergence of PI-$L^2$2 holds only for sufficiently short horizons $L^2$3, with explicit but possibly restrictive constants, and for Hamiltonians with the specific affine-in-control structure of assumptions $L^2$4–$L^2$5; extension to longer horizons or more general dynamics is open. The periodic $L^2$6 estimate requires both periodicity of $L^2$7 and semiconcavity of $L^2$8, and the proposed weighted $L^2$9 analogue is stated without proof. Finally, the gap between the proven upper bounds (e.g., $\mathbb{R}^d$00 for merely Lipschitz data) and the linear lower bound leaves open whether the true splitting convergence rate is in fact $\mathbb{R}^d$01 under weaker regularity — the authors' own machinery suggests the lower bound may be sharp, but no optimality result is established. The numerical section reports residuals rather than errors against known exact solutions, so the practical accuracy of the full pipeline relative to theory is not quantified.

Conclusion

The paper provides the first quantitative convergence rates for a Trotter–Kato-type splitting of viscous Hamilton–Jacobi equations, with rates improving from $\mathbb{R}^d$02 to $\mathbb{R}^d$03 as the initial datum gains regularity, together with an averaged $\mathbb{R}^d$04 rate of $\mathbb{R}^d$05 in the periodic setting. Its second contribution extends value-gradient policy iteration to time-dependent first-order problems with exponential convergence in a freely chosen weighted norm, removing the large-discount restriction of earlier analyses. The combination — splitting into a trivial heat step and a characteristic-based, parallelizable first-order step amenable to supervised learning — yields a practical algorithm whose numerical performance on quadratic control problems up to dimension 32 is stable even with sparse trajectory sampling. The main open issues are the sharpness of the splitting rates, relaxation of the short-horizon assumption in the policy iteration analysis, and rigorous treatment of weighted estimates on unbounded domains.