Bayesian Sequential Optimal Experimental Design for Nonlinear Models Using Policy Gradient Reinforcement Learning
Wanggang Shen, Xun Huan
TL;DR
This work addresses optimally designing a finite sequence of experiments under Bayesian uncertainty by formulating sOED as a finite-horizon POMDP with information-theoretic rewards. It develops a policy-gradient, actor-critic framework (PG-sOED) that parameterizes both the policy and value functions with deep neural networks, enabling efficient gradient-based optimization for continuous states and actions. The authors prove the exact PG expression, propose a Monte Carlo estimator, and implement a nonparametric belief representation that avoids intermediate posterior updates. Numerical results on a linear-Gaussian benchmark and a nonlinear convection–diffusion contaminant-inversion problem demonstrate that PG-sOED matches analytic optima in the benchmark and outperforms batch and greedy designs in challenging, expensive-forward-model scenarios, with significant online-speed advantages. The work provides code and outlines future enhancements for scalability and advanced RL techniques to broaden applicability to high-dimensional and more complex experimental design problems.
Abstract
We present a mathematical framework and computational methods to optimally design a finite number of sequential experiments. We formulate this sequential optimal experimental design (sOED) problem as a finite-horizon partially observable Markov decision process (POMDP) in a Bayesian setting and with information-theoretic utilities. It is built to accommodate continuous random variables, general non-Gaussian posteriors, and expensive nonlinear forward models. sOED then seeks an optimal design policy that incorporates elements of both feedback and lookahead, generalizing the suboptimal batch and greedy designs. We solve for the sOED policy numerically via policy gradient (PG) methods from reinforcement learning, and derive and prove the PG expression for sOED. Adopting an actor-critic approach, we parameterize the policy and value functions using deep neural networks and improve them using gradient estimates produced from simulated episodes of designs and observations. The overall PG-sOED method is validated on a linear-Gaussian benchmark, and its advantages over batch and greedy designs are demonstrated through a contaminant source inversion problem in a convection-diffusion field.
