# Gaussian process policy iteration with additive Schwarz acceleration for forward and inverse HJB and mean field game problems

arXiv: https://arxiv.org/abs/2505.00909 · HTML: https://mlsanity.com/paper/2505.00909/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 71 references that are in the map.
- First posted on arXiv 1 May 2025.

## Builds on (most cited first, 28 of 28)

- [Machine Learning of Linear Differential Equations using Gaussian Processes](https://mlsanity.com/paper/1701.02440/) (2017)
- [Solving and Learning Nonlinear PDEs with Gaussian Processes](https://mlsanity.com/paper/2103.12959/) (2021)
- [In-Context Operator Learning with Data Prompts for Differential Equation Problems](https://mlsanity.com/paper/2304.07993/) (2023)
- [A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems](https://mlsanity.com/paper/1912.01825/) (2019)
- [Numerical Gaussian Processes for Time-dependent and Non-linear Partial Differential Equations](https://mlsanity.com/paper/1703.10230/) (2017)
- [Machine learning approximation algorithms for high-dimensional fully nonlinear partial differential equations and second-order backward stochastic differential equations](https://mlsanity.com/paper/1709.05963/) (2017)
- [Alternating the Population and Control Neural Networks to Solve High-Dimensional Stochastic Mean-Field Games](https://mlsanity.com/paper/2002.10113/) (2020)
- [Convergence Analysis of Machine Learning Algorithms for the Numerical Solution of Mean Field Control and Games: II -- The Finite Horizon Case](https://mlsanity.com/paper/1908.01613/) (2019)
- [Sparse Cholesky Factorization for Solving Nonlinear PDEs via Gaussian Processes](https://mlsanity.com/paper/2304.01294/) (2023)
- [Learning dynamical systems from data: a simple cross-validation perspective](https://mlsanity.com/paper/2007.05074/) (2020)
- [Learning Dynamical Systems from Data: A Simple Cross-Validation Perspective, Part V: Sparse Kernel Flows for 132 Chaotic Dynamical Systems](https://mlsanity.com/paper/2301.10321/) (2023)
- [Proximal methods for stationary Mean Field Games with local couplings](https://mlsanity.com/paper/1608.07701/) (2016)
- [On the implementation of a primal-dual algorithm for second order time-dependent mean field games with local couplings](https://mlsanity.com/paper/1802.07902/) (2018)
- [A mean field game inverse problem](https://mlsanity.com/paper/2007.11551/) (2020)
- [Numerical methods for Mean field Games based on Gaussian Processes and Fourier Features](https://mlsanity.com/paper/2112.05414/) (2021)
- [Splitting methods for a class of non-potential mean field games](https://mlsanity.com/paper/2007.00099/) (2020)
- [Policy iteration method for time-dependent Mean Field Games systems with non-separable Hamiltonians](https://mlsanity.com/paper/2110.02552/) (2021)
- [Mean field control problems for vaccine distribution](https://mlsanity.com/paper/2104.11887/) (2021)
- [Inverse problems for mean field games](https://mlsanity.com/paper/2205.11350/) (2022)
- [A Policy Iteration Method for Inverse Mean Field Games](https://mlsanity.com/paper/2409.06184/) (2024)
- [An Efficient Policy Iteration Algorithm for Dynamic Programming Equations](https://mlsanity.com/paper/1308.2087/) (2013)
- [Learning optimal policies in potential Mean Field Games: Smoothed Policy Iteration algorithms](https://mlsanity.com/paper/2212.04791/) (2022)
- [Deep Policy Iteration for High-Dimensional Mean Field Games](https://mlsanity.com/paper/2310.10827/) (2023)
- [Decoding Mean Field Games from Population and Environment Observations By Gaussian Processes](https://mlsanity.com/paper/2312.06625/) (2023)
- [A Bilevel Optimization Method for Inverse Mean-Field Games](https://mlsanity.com/paper/2401.05539/) (2024)
- [Learning Surrogate Potential Mean Field Games via Gaussian Processes: A Data-Driven Approach to Ill-Posed Inverse Problems](https://mlsanity.com/paper/2502.11506/) (2025)
- [Unique determination of cost functions in a multipopulation mean field game model](https://mlsanity.com/paper/2312.01622/) (2023)
- [Reconstructing a state-independent cost function in a mean-field game model](https://mlsanity.com/paper/2402.09297/) (2024)

## Abstract (arXiv)

In this paper, we propose a Gaussian Process (GP)-based policy iteration framework for addressing both forward and inverse problems in Hamilton--Jacobi--Bellman (HJB) equations and mean field games (MFGs). Policy iteration is formulated as an alternating procedure between evaluating the value function under a fixed control policy and improving the policy. In our approach, we model the unknown fields using GPs within a policy-iteration framework that converts the nonlinear system into a sequence of linear PDE subproblems. Then, leveraging the linear structure, the updates for the value function and, in the MFG setting, the population density admit explicit representer formulas under linear PDE collocation constraints. The policy is subsequently updated pointwise via a Legendre transform step, which involves a low-dimensional maximization over the control variable. This maximization is explicit for standard quadratic costs. For smooth, strictly convex costs, this pointwise maximization is solved through its first-order optimality condition, whereas in constrained or non-smooth cases, it becomes a low-dimensional constrained maximization problem. To improve convergence, we incorporate the additive Schwarz acceleration as a preconditioning step following each policy update. Numerical experiments demonstrate the effectiveness of the Schwarz acceleration in improving computational efficiency.