Sequential Bayesian Optimal Experimental Design in Infinite Dimensions via Policy Gradient Reinforcement Learning

TL;DR

This work tackles sequential Bayesian optimal experimental design for PDE-governed inverse problems with infinite-dimensional parameters. It casts SBOED as a finite-horizon MDP and learns an amortized policy via policy-gradient reinforcement learning, aided by a derivative-informed latent attention neural operator surrogate and dual dimension reduction. A Laplace-based D-optimality reward and a low-rank eigenvalue evaluation enable scalable, sensitivity-aware design without repeated MAP solves. On a contaminant source tracking case, the method delivers about 100x speedups over high-fidelity solvers and discovers physically interpretable upstream sensor placement strategies that outperform random designs.

Abstract

Sequential Bayesian optimal experimental design (SBOED) for PDE-governed inverse problems is computationally challenging, especially for infinite-dimensional random field parameters. High-fidelity approaches require repeated forward and adjoint PDE solves inside nested Bayesian inversion and design loops. We formulate SBOED as a finite-horizon Markov decision process and learn an amortized design policy via policy-gradient reinforcement learning (PGRL), enabling online design selection from the experiment history without repeatedly solving an SBOED optimization problem. To make policy training and reward evaluation scalable, we combine dual dimension reduction -- active subspace projection for the parameter and principal component analysis for the state -- with an adjusted derivative-informed latent attention neural operator (LANO) surrogate that predicts both the parameter-to-solution map and its Jacobian. We use a Laplace-based D-optimality reward while noting that, in general, other expected-information-gain utilities such as KL divergence can also be used within the same framework. We further introduce an eigenvalue-based evaluation strategy that uses prior samples as proxies for maximum a posteriori (MAP) points, avoiding repeated MAP solves while retaining accurate information-gain estimates. Numerical experiments on sequential multi-sensor placement for contaminant source tracking demonstrate approximately $100\times$ speedup over high-fidelity finite element methods, improved performance over random sensor placements, and physically interpretable policies that discover an ``upstream'' tracking strategy.

Figures (14)

• Figure 1: (a) The physical domain $\Omega \subset \mathbb{R}^{2}$ featuring two impenetrable buildings (blocks) and the computational mesh (refined twice in our computation). (b) The steady-state velocity field $\vec{v}$ resulting from boundary condition $\vec{g}_{1}$. (c) The velocity field resulting from boundary condition $\vec{g}_{2}$, which induces a reversed flow pattern.
• Figure 2: Top: Four independent samples drawn from the Gaussian prior distribution of the parameter $m$. The selected covariance parameters ($\gamma, \delta$) induce significant spatial variability, resulting in distinct high-value regions across samples. Bottom: The corresponding initial condition fields, obtained by squaring the prior samples ($m^2$).
• Figure 3: Time evolution of the contaminant concentration $u(t, x)$, simulated using the Finite Element Method (FEM) for a parameter $m$ sampled from the prior. The state variable exhibits negligible variation after $t=1.6$, which marks the final sensor placement time step. This steady-state behavior justifies the concentration of experiments in the earlier time intervals.
• Figure 4: (a) The initial positions of the three sensors used for sparse observations. (b) The red dashed rectangles indicate the permissible regions within which each sensor must remain. (c) An illustration of the distance constraints for the sensor at $(0.800, 0.375)$; the bounds for the design vector are derived from these distances to ensure validity across all $N$ experiments.
• Figure 5: Decay of the generalized eigenvalues of active subspace projection for random field parameter dimension reduction (left) and the singular values of PCA for state variable dimension reduction (right).