Source-linked AI summary
Optimizing Sequential Experimental Design with Deep Reinforcement Learning
Tom Blau, Edwin V. Bonilla, Iadine Chades, Amir Dezfouli
TL;DR
Sequential Bayesian experimental design is computationally challenging, and amortized approaches are limited in exploration, model differentiability, and design-space type. The paper formulates policy optimization as an MDP and solves it with deep reinforcement learning. Its approach performs strongly across continuous and discrete designs, including black-box probabilistic models.
Problem
Non-myopic Bayesian design is computationally challenging, while amortized methods have limited exploration and require differentiable models and continuous design spaces.
Source: §1 · pp. 1–2Method
The paper formulates sequential experimental design as an MDP and learns the design policy using reinforcement learning.
Source: §1 · p. 2Results
Takeaways & Limitations
The MDP formulation allows off-the-shelf reinforcement-learning algorithms to be applied to sequential experimental design problems.
Source: §6 · p. 9In the discrete-design evaluation, the RL method does not outperform the myopic baselines.
Source: §4.2 · p. 7
Abstract
from arXiv · showhide
Bayesian approaches developed to solve the optimal design of sequential experiments are mathematically elegant but computationally challenging. Recently, techniques using amortization have been proposed to make these Bayesian approaches practical, by training a parameterized policy that proposes designs efficiently at deployment time. However, these methods may not sufficiently explore the design space, require access to a differentiable probabilistic model and can only optimize over continuous design spaces. Here, we address these limitations by showing that the problem of optimizing policies can be reduced to solving a Markov decision process (MDP). We solve the equivalent MDP with modern deep reinforcement learning techniques. Our experiments show that our approach is also computationally efficient at deployment time and exhibits state-of-the-art performance on both continuous and discrete design spaces, even when the probabilistic model is a black box.
1. Introduction
Bayesian optimal experimental design seeks informative experiments, but non-myopic sequential optimization is computationally intractable and amortized methods have important constraints. The paper formulates sequential design as an MDP and uses reinforcement learning to target these limitations.
1. Introduction
- BOED frames experiment selection as maximizing expected information gain, motivated by the cost and time required to acquire experimental data.Source: §1 · p. 1
- Non-myopic design with a budget of T experiments requires nested expectations whose depth grows as O(T), making expectation-based approaches intractable.The stated convergence rate is O(n^(-1/(T+1))).Source: §1 · p. 1
- Amortized methods optimize a parameterized policy mapping experimental histories to the next design, rather than optimizing designs directly.They optimize a lower-bound estimate of non-myopic expected information gain with respect to policy parameters.Source: §1 · p. 1
- These amortized approaches require a differentiable probabilistic model and continuous design space, excluding discrete designs and models without available gradients.Source: §1 · p. 2
- The proposed deep-RL approach achieves state-of-the-art expected information gain on continuous and discrete designs, including when the probabilistic model is a black box.The authors also report that correct MDP formulation is critical; naive conversions perform significantly worse.Source: §1 · p. 2
2. Background
BOED selects designs by maximizing expected information gain, which can be estimated with contrastive lower bounds. Sequential policies avoid repeated posterior estimation, while reinforcement learning provides a general framework for optimizing policies over states and actions.
2. Background
- BOED chooses a design to maximize expected information gain, and the PCE lower bound provides a numerically estimable route when EIG is not available in closed form.The PCE estimator uses prior samples and contrastive samples to reweight experiment outcomes.Source: §2 · p. 2
2.1. Sequential Experimental Design
- Iterative sequential design updates the posterior after each experiment but is costly and myopic, optimizing only the next experiment's EIG.Source: §2.1 · p. 2
- A policy maps experimental history to the next design; DAD optimizes this policy using the SPCE lower bound and amortizes computation for fast deployment.The required expectations are estimated by sampling model parameters and histories; the paper proposes reducing policy optimization to solving an MDP for broader exploration and support for discrete designs and black-box models.Source: §2.1 · p. 2
2.2. Reinforcement Learning
- An MDP specifies states, actions, transition dynamics, rewards, and initial states, and reinforcement learning seeks a policy maximizing expected discounted return.Source: §2.2 · p. 3
3. RL for Sequential Experiment Design
The section formulates sequential experimental design as a Markov decision process whose optimized return matches the design objective. It then develops a compact state and reward and describes the advantages of reinforcement learning for design-policy optimization.
3. RL for Sequential Experiment Design
- To avoid repeated posterior inference, the formulation uses g(θ, ht) as reward; because this requires θ in the state, where it is unobservable at test time and unaffected by actions, it adopts a Hidden Parameter MDP.Source: §3 · p. 3
3.1. Hidden Parameter MDPs
- The proposed HIP-MDP formulation makes sequential experimental design a policy-optimization problem, with a theorem establishing that its solution optimizes the corresponding design problem.The parameter is drawn from the prior and held fixed during an episode; a terminal reward avoids counting an experiment’s contribution repeatedly.Source: §3.1 · p. 3
3.2. The SED MDP
- The history summary is updated incrementally, while the proposed reward estimates each experiment’s marginal contribution to cumulative expected information gain.Rewards are generally nonzero at each timestep, sum to the cumulative objective for a trajectory or its prefix, and can be computed from the latest likelihoods and outcome; full histories are still needed in the training replay buffer.Source: §3.2 · pp. 3–4
- The refined HIP-MDP uses history summaries and likelihoods in its state, and Theorem 2 establishes that its solution also optimizes the corresponding sequential design problem.The policy is stochastic, mapping the summary to a distribution over designs.Source: §3.2 · p. 4
3.3. Advantages of the RL Formulation
- Reinforcement-learning agents can be trained offline and their learned design policies used for rapid online experiment design.The section presents the MDP formulation as a recipe for deploying RL algorithms for sequential design.Source: §3.3 · p. 4
3.3.1. DISCRETE-DESIGN SPACES AND BLACK-BOX LIKELIHOODS
- Unlike DAD, reinforcement learning can optimize over discrete design spaces and does not require a differentiable reward; this includes likelihoods that are nondifferentiable with respect to design.The paper gives a model containing an argmax as an example where RL can be used but DAD cannot.Source: §3.3.1 · p. 4
3.3.2. EXPLORATION VS EXPLOITATION
- DAD is described as a pure-exploitation algorithm that may explore the design space insufficiently and become trapped in a poor local optimum.The authors contrast this with RL features such as policy noise and regularization, which support exploration.Source: §3.3.2 · pp. 4–5
4. Experimental Results
The experiments compare RL with baseline methods on continuous and discrete design problems, including a non-differentiable likelihood setting. RL performs strongly across these settings, while its comparison with myopic baselines varies by problem.
4.1. Continuous Design Space
- The continuous-design comparisons include DAD, VPCE, and random, while DAD is inapplicable to discrete designs and VPCE can be adapted by evaluating every possible design.The VPCE modification is tractable because it is myopic; the analogous DAD adjustment would require computing SPCE for |A|^T designs.Source: §4.1 · p. 5, §4.2 · pp. 6–7
4.1.1. SOURCE LOCATION FINDING
- On source location, RL outperforms all baselines at every timestep; at t = 30, its lower-bound improvement over DAD exceeds the standard error by an order of magnitude.Its upper bound is comparable to DAD’s and significantly better than the other baselines; both amortized methods are trained to maximize the lower bound.Source: §4.1.1 · p. 5, Table 1 · p. 5, Figure 1 · p. 6
4.1.2. CONSTANT ELASTICITY OF SUBSTITUTION
- On CES, RL outperforms all baselines over 10 experiments, and its lower bound at t = 10 exceeds DAD’s upper bound by several standard errors.The censored sigmoid maps much of the design space to boundary outcomes, creating a difficult exploration problem; RL’s stochastic policy explores more broadly than DAD’s deterministic policy.Source: Figure 2 · p. 6, §4.1.2 · p. 6, Table 2 · p. 7
4.2. Discrete Design Space
- On the prey population problem, RL achieves results similar to SMC and VPCE, while all three methods outperform random; SMC’s small advantage over RL is within the error bars.The reported lower-bound gap between SMC and RL is smaller than the standard error, with a relative effect size on the order of 1%.Source: Figure 4 · p. 7, §4.2 · p. 7, Table 3 · p. 8
4.3. Computational Benefits of the RL Appproach at Deployment
- At deployment, RL proposes prey-problem designs orders of magnitude faster than SMC and VPCE, and has comparable deployment time to DAD on continuous problems.For the prey problem, RL is close to the random baseline’s deployment time.Source: §4.3 · pp. 7–8, Table 4 · p. 8
4.4. Non-differentiable Likelihood
- With a non-differentiable CES likelihood, RL performs nearly identically to the differentiable-likelihood results and outperforms the gradient-free baselines from the first experiment.DAD and VPCE are absent because they require likelihood gradients; PCE-BO replaces VPCE in this setting.Source: Figure 5 · p. 8, §4.4 · p. 8
5. Related Work
Prior work spans myopic Bayesian design, information-gain methods in active learning and reinforcement learning, and amortized approaches. The closest experimental-design precedent uses problem-specific dynamic programming with explicit posterior inference on a discretized grid.
5. Related Work
- Classical Bayesian design optimizes immediate expected information gain using a range of estimators, while some methods exploit properties of particular model classes.The cited approaches include Monte Carlo, variational, neural, and other estimation methods, alongside methods for linear, Gaussian-process, and polynomial models.Source: §5 · p. 8
- Active-learning work includes non-myopic look-ahead and compressed joint-experiment methods, while reinforcement-learning work has used information gain for exploration or auxiliary objectives.The paper distinguishes these prior directions from applying reinforcement learning to solve experimental-design problems.Source: §5 · p. 9
- The closest prior experimental-design approach uses specific dynamic programming and requires explicit posterior inference approximated on a discretized grid.This paper instead applies generic reinforcement-learning algorithms to experimental design.Source: §5 · p. 9
6. Discussion
The paper formulates sequential experimental design as an MDP and learns amortized design policies with off-the-shelf reinforcement learning. Results span continuous, discrete, and non-differentiable settings, while reward factorization is critical to performance.
6. Discussion
- The method transforms sequential experimental design into an MDP so off-the-shelf RL algorithms can learn an amortized policy for each problem pair.The paper describes the resulting approach as being in the spirit of deep adaptive design.Source: §6 · p. 9
- In continuous design spaces, RL significantly outperforms conventional non-amortized methods and DAD; it is comparable to the best non-amortized baselines in discrete spaces and performs best when the model is non-differentiable.Amortized baselines do not work for discrete designs because they depend on backpropagating gradients through the design.Source: §6 · p. 9
- Ablations show that correct reward factorization is critical; using the SPCE objective naively as the reward does not achieve the reported results.The paper frames this as a necessary part of the RL formulation rather than a plug-in use of the SPCE objective.Source: §6 · p. 9
- The MDP formulation may let researchers apply RL techniques for challenges such as sparse rewards or high dimensionality to experimental-design problems.The paper says these techniques may help with particular problems and can be applied by transforming the design problem to an MDP.Source: §6 · p. 9
A. Theorem Proofs
The appendix provides proofs for the paper’s theorems.
A. Theorem Proofs
- The appendix contains proofs for the paper’s theorems.Source: A. Theorem Proofs · p. 12
A.1. Proof of Theorem 1
Theorem 1 establishes an equivalent HIP-MDP formulation for a sequential experiment optimization problem. Under its specified conditions and terminal reward, the expected return equals the corresponding SPCE objective.
A.1. Proof of Theorem 1
- Theorem 1 constructs an equivalent HIP-MDP for a sequential experiment optimization problem using histories as states and designs as actions.Its setup specifies the initial empty history, prior parameter distribution, likelihood transitions, and γ = 1.Source: §A.1 · p. 12
- The theorem sets rewards to zero before the final step and uses g(θ, h_T) as the terminal reward.Source: §A.1 · p. 12
- The proof shows that the HIP-MDP’s expected return equals sPCE by matching the distributions and terminal terms in the two expectations.It concludes that J(π) = sPCE(π, L, T).Source: §A.1 · p. 12
A.2. Proof of Theorem 2
Theorem 2 establishes an equivalent HIP-MDP reward formulation for sequential experiment optimization. Under its conditions and γ = 1, expected return equals SPCE.
A.2. Proof of Theorem 2
- Theorem 2 defines a HIP-MDP reward as the log-likelihood increment log p(yt | θ0, dt) − log(Ct · 1) + log(Ct−1 · 1).The theorem’s expected-return equivalence is stated for γ = 1.Source: §A.2 · p. 13
- The derivation uses the assumption that the history likelihood factorizes over experiments and telescopes the log-probability terms.It also uses that π(dt | ht−1) is independent of θ.Source: §A.2 · p. 14
- The proof concludes that the policy objective J(π) equals sPCE(π, L, T) under the theorem’s conditions.It establishes equality by matching the terms inside expectations over the same distributions.Source: §A.2 · pp. 13–15
A.3. Q-Learning from Priors
Theorem 3 shows that, under specified deterministic state and action mappings, a Q-function estimator can be learned from trajectories sampled using the prior and a policy. The proof derives a recursive Q-function relation without requiring posterior inference for sampling.
A.3. Q-Learning from Priors
- Under deterministic state and action mappings, Theorem 3 permits learning a Q-function estimator from trajectories sampled from p(θ, hT | π).The theorem states this for an arbitrary policy π and transitions driven by observations from the model.Source: §A.3 · p. 15
- The proof derives a Bellman recursion for Q by relating the state-action-parameter value function Uπ to the state-action value function.It uses the Bayesian identity connecting p(θ | hi)p(yi+1 | di+1, θ) with p(θ | hi+1)p(yi+1 | hi, di+1).Source: §A.3 · pp. 15–16
- Rather than sampling from difficult marginal-history and posterior distributions to evaluate Qπ, the estimator can use prior draws and policy rollouts.The authors state that the required samples can be generated by drawing θ from the prior and rolling out the policy.Source: §A.3 · p. 17
B. Experiment Details
This section provides the full descriptions of the experimental design problems considered in the paper.
B. Experiment Details
- The section’s stated scope is the full description of the paper’s experimental design problems.Source: B. Experiment Details · p. 17
B.1. Source Location
The experiments cover source-location, consumer-preference, and prey-population design problems, alongside algorithm and implementation details. In a one-dimensional source-location test, the non-myopic agent sacrifices early information gain for higher final information gain; learning curves also favor the factorized reward formulation.
B.1. Source Location
- The source-location task places two signal sources in a 2D plane, uses a unit-Gaussian prior and Gaussian noise on log intensity, and restricts designs to [−4, 4]2.The observed signal is the log of total intensity at the chosen coordinate.Source: §B.1 · p. 17
B.2. Constant Elasticity of Substitution
- The CES task asks participants to rate comparisons between two goods baskets, with each basket represented by a three-dimensional vector in [0, 100]3.Source: §B.2 · p. 17
B.3. Prey Population
- The prey-population task varies initial population N0 from 1 to 300, then measures predator consumption after 24 hours.Population dynamics are modeled by a differential equation whose solution gives the population at the experiment’s end.Source: §B.3 · p. 18
C.1. RL
- The RL implementation uses REDQ, with an ensemble of critics and a random subset for Q estimation, and task-specific policy outputs for continuous or discrete designs.The policy uses Tanh-Gaussian distributions for source location and CES, and a Gumbel-Softmax distribution for prey population.Source: §C.1 · pp. 18–19
C.2. DAD; C.3. SMC; C.4. VPCE
- The benchmark implementations use the official DAD code, Moffat et al.’s SMC implementation, and VPCE’s iterative PCE-bound optimization with variational inference.PCE-BO uses a Matérn 5/2 kernel, lengthscale 10, and four Bayesian-optimization steps.Source: §C.2 · p. 19, §C.3 · p. 19, §C.4 · p. 19
D. Hardware details
- Except for SMC, experiments ran on a single Tesla P100 GPU and four Intel Xeon E5-2690 CPU cores; SMC ran without GPU acceleration.Source: D. Hardware details · p. 20
E. Non-Myopic Solutions
- After the second experiment, the non-myopic agent significantly outperforms both the myopic agent and the optimal myopic strategy, despite doing worse after the first experiment.Figure 6 plots mean and standard error of expected information gain after each experiment, aggregated over 1e5 rollouts.Source: E. Non-Myopic Solutions · p. 20, Figure 6 · p. 20
F. Learning Curves
- The factorized reward formulation yields faster performance improvement and higher converged rewards, particularly in source location with its 30-experiment budget.The learning curves aggregate mean and standard errors across 10 replications; their rewards use fewer contrastive samples than final evaluation.Source: F. Learning Curves · p. 21