Table of Contents
Fetching ...

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.

Bayesian Sequential Optimal Experimental Design for Nonlinear Models Using Policy Gradient Reinforcement Learning

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.

Paper Structure

This paper contains 28 sections, 2 theorems, 41 equations, 16 figures, 2 tables, 1 algorithm.

Key Result

Theorem 1

Let $U_T(\pi)$ be the sOED expected utility defined in eq:expected_utility subject to the constraints in eq:optimal_policy for a given policy $\pi$ while using the terminal formulation eq:terminal1eq:terminal_info_gN. Let $U_I(\pi)$ be the same except using the incremental formulation eq:incremental

Figures (16)

  • Figure 1: Flowchart of the process involved in a $N$-experiment sOED.
  • Figure 2: Convergence history of PG-sOED.
  • Figure 3: Sample numerical solution of the concentration field $G$ at different time snapshots. The solution is solved in a wider computational domain $[-1,2]^2$ but displayed here in a region of interest $[0,1]^2$. In this case, $\theta=[0.210,0.203,0.05,2]$ and the convection grows over time with $u_x=u_y=10t/0.2$. Hence, isotropic diffusion dominates early on and the plume stretches towards the convection direction with time.
  • Figure 4: Sample comparison of the concentration field $G$ at $t=0.05$ and $t=0.2$ of case 2 using the DNN surrogates (left column) and finite volume (right column). They appear nearly identical.
  • Figure 5: Case 1. Expected utility for one-experiment design at $t=0.32$. The best design locations are at the corners.
  • ...and 11 more figures

Theorems & Definitions (4)

  • Theorem 1
  • Theorem 2
  • proof : Proof of \ref{['prop:terminal_incremental']}
  • proof : Proof of \ref{['thm:PG']}