PaperlayerView exact revision on arXiv

Source-linked AI summary

Policy-Based Bayesian Experimental Design for Non-Differentiable Implicit Models

Vincent Lim, Ellen Novoseller, Jeffrey Ichnowski, Huang Huang, Ken Goldberg

arXiv:2203.04272v1cs.LGcs.AIstat.ME

TL;DR

Sequential experimental design can be difficult when likelihoods are intractable and simulators are non-differentiable, while many approaches also require online computation. RL-DAD learns a policy using reinforcement learning and likelihood-free information bounds; it is competitive with baselines on three benchmarks. Its primary limitation is the training cost and deep RL’s sample inefficiency and hyperparameter brittleness.

  • Problem

    Sequential BOED for implicit models can require costly online information optimization, while existing policy methods may require an analytic likelihood or differentiable simulator.

    Source: §1 · p. 1

  • Method

    RL-DAD formulates BOED as a partially observed Markov decision process and uses reinforcement learning with likelihood-free information lower bounds to learn a design policy.

    Source: §1 · pp. 1–2, Transition and Observation Emission… · p. 5

  • Results

    RL-DAD is competitive with baselines on three benchmarks and approaches policy-based methods requiring analytic likelihoods or differentiable simulators.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    The learned policy provides a sequential BOED approach compatible with black-box simulators without requiring closed-form likelihoods or simulator differentiability.

    Source: §1 · p. 1, §6 · p. 8

    The primary limitation is initial training time; deep RL also brings sample inefficiency and brittleness to hyperparameters and implementation.

    Source: §6 · p. 8

Abstract

from arXiv · show

For applications in healthcare, physics, energy, robotics, and many other fields, designing maximally informative experiments is valuable, particularly when experiments are expensive, time-consuming, or pose safety hazards. While existing approaches can sequentially design experiments based on prior observation history, many of these methods do not extend to implicit models, where simulation is possible but computing the likelihood is intractable. Furthermore, they often require either significant online computation during deployment or a differentiable simulation system. We introduce Reinforcement Learning for Deep Adaptive Design (RL-DAD), a method for simulation-based optimal experimental design for non-differentiable implicit models. RL-DAD extends prior work in policy-based Bayesian Optimal Experimental Design (BOED) by reformulating it as a Markov Decision Process with a reward function based on likelihood-free information lower bounds, which is used to learn a policy via deep reinforcement learning. The learned design policy maps prior histories to experiment designs offline and can be quickly deployed during online execution. We evaluate RL-DAD and find that it performs competitively with baselines on three benchmarks.

1. Introduction

Sequential Bayesian optimal experimental design (BOED) can require expensive online information optimization and may rely on likelihoods or differentiable simulators. RL-DAD learns a design policy with reinforcement learning for non-differentiable implicit models, approaching stricter-assumption baselines across three benchmarks.

1. Introduction

  • Sequential BOED selects experiments from prior outcomes, but conventional methods can be myopic and require costly online mutual-information optimization and posterior updates.The passage describes mutual-information maximization as doubly intractable.Source: §1 · p. 1
  • RL-DAD frames BOED as a partially observed Markov decision process and uses reinforcement learning with likelihood-free information lower bounds to learn a design policy.The policy maps experiment histories to designs and requires only a network forward pass for online decisions.Source: §1 · pp. 1–2
  • RL-DAD targets implicit models without closed-form likelihoods and simulators that cannot be assumed differentiable, extending policy-based BOED beyond prior methods’ assumptions.The introduction identifies robotics, medicine, particle physics, and protein folding as domains where simulator differentiability cannot be assumed.Source: §1 · p. 1
  • Across three benchmarks, RL-DAD approaches policy-based BOED methods requiring analytic likelihoods or differentiable simulators, while also showing promise where both assumptions fail.The reported comparisons concern three benchmarks and include high-dimensional settings.Source: §1 · p. 2
  • The paper also contributes a BOED-suitable policy network architecture.The introduction lists this as a separate contribution from the algorithm and benchmark results.Source: §1 · p. 2

2. Related Work

Related work spans static and sequential BOED, including methods that optimize information estimates online and policy-based methods that amortize design decisions. RL-based BOED methods had not covered general settings without either closed-form likelihoods or differentiable simulators.

Information gain estimation

  • BOED methods estimate mutual information between unknown model parameters and experiment outcomes using variational or contrastive techniques.Examples include MINE and SMILE, as well as sPCE and InfoNCE.Source: §2 · p. 2

Static BOED

  • Prior static BOED work considers likelihood-based and implicit models, using mutual-information estimators such as SMILE and MINE, respectively.For implicit models, the cited estimators are typically optimized using evolutionary methods or Bayesian optimization.Source: §2 · p. 2

Sequential BOED

  • Sequential BOED methods often choose each design by greedy information-gain optimization and perform time-intensive mutual-information estimation between experiments.The passage contrasts this myopic selection with multi-step planning.Source: §2 · p. 2

Policy-based BOED

  • Policy-based BOED learns a history-to-design policy upfront, reducing online optimization at deployment despite significant training computation.Prior approaches include dynamic programming and gradient ascent, with different likelihood or simulator assumptions.Source: §2 · p. 2

RL in BOED

  • RL-based sequential BOED methods cited here assume closed-form likelihoods, while one is restricted to Bayesian function optimization rather than general BOED.A policy-based implicit-model method instead requires a differentiable simulator.Source: §2 · p. 3

3. Preliminaries

The preliminaries define sequential BOED as adaptive experiment selection to learn unknown model parameters, with total information gain as the objective. They review implicit-model and policy-based approaches, their assumptions, and the reinforcement-learning framework used by the paper.

3.1. The Sequential BOED Problem Statement

  • Over a fixed horizon, sequential BOED selects experiment designs and observes outcomes to maximize mutual information between the parameter and observation sequences.The agent uses its history of design-observation pairs to adapt the next experiment.Source: §3.1 · p. 3
  • The simulator may generate observations from the experiment, parameters, and prior history, allowing conditionally dependent experiments without assuming simulator differentiability.The design, observation, and parameter spaces are assumed known.Source: §3.1 · p. 3

3.2. Background: Sequential BOED with design policies

  • For implicit models, conventional sequential BOED alternates likelihood-free posterior estimation and information-gain maximization, a myopic procedure that can yield suboptimal sequences.The conventional objective is the expected reduction in entropy between the current and next intermediate posterior.Source: §3.2 · p. 3
  • Policy-based BOED learns designs from histories to maximize total expected information gain, but its sPCE and InfoNCE bounds require, respectively, tractable likelihoods or differentiable simulation.The passage notes a stochastic-policy extension and a learned critic for InfoNCE.Source: §3.2 · pp. 3–4

3.3. Background: Reinforcement Learning

  • In the episodic partially observed Markov decision process framework, an agent seeks policy parameters maximizing expected cumulative reward over a fixed horizon.States are unobserved; the agent acts on observations and receives rewards and subsequent observations.Source: §3.3 · p. 4
  • The paper uses TD3, an off-policy policy-gradient algorithm that learns a Q-function alongside a neural-network policy and adds Gaussian exploration noise during training.The Q-function is learned via Bellman backups.Source: §3.3 · p. 4

4. The RL-DAD Algorithm

RL-DAD casts sequential BOED as an MDP and learns a history-to-design policy using information-based rewards. Its dense reward preserves each trajectory’s total reward while providing significant empirical gains.

4. The RL-DAD Algorithm

  • RL-DAD frames BOED as an MDP in which an RL agent maps experiment histories to the next design, first using a sparse terminal reward and then a dense alternative.The sparse reward is zero at intermediate steps and reflects total information gained at the final step.Source: §4 · p. 4

4.1. Casting BOED as an MDP

  • The MDP state includes the experiment history and hidden ground-truth and contrastive parameter samples; actions are experiment designs, and observations reveal the history.The parameter samples are drawn from the prior and fixed during each trajectory.Source: §4.1 · p. 5, Transition and Observation Emission… · p. 5

Transition and Observation Emission Probabilities

  • Transitions append the new design and outcome to the history, with transition probabilities determined by the observation under the ground-truth parameter.The next history contains the previous history and the new design–observation pair.Source: Transition and Observation Emission… · p. 5
  • The sparse reward approximates trajectory information gain with an sPCE bound, then an InfoNCE approximation supplies a reward without requiring a tractable likelihood.The expected value of g is bounded above by the mutual information between the trajectory history and parameters.Source: Transition and Observation Emission… · p. 5

4.2. Dense Reward Formulation

  • Dense rewards preserve each trajectory’s total reward and provide significant empirical gains over the sparse formulation.Each dense reward is the change in the information-bound estimate between consecutive histories.Source: §4.2 · p. 6

4.3. Network Architectures

  • For variable-length histories, the critic separately encodes parameters and history, while history encoding uses sum-pooling with attention under conditional independence or an LSTM otherwise.The policy and Q networks use concatenation with zero-padding; this significantly outperforms a more complex sequence-encoding architecture.Source: §4.3 · p. 6

5. Experiments

The experiments evaluate RL-DAD across Location Finding, SIR, and Cartpole, using random and MINEBED-BO baselines and stronger-assumption methods where applicable. RL-DAD outperforms general baselines and is competitive with stronger-assumption methods in several settings, while sparse rewards often fail to converge.

5. Experiments

  • The study tests three design environments and compares RL-DAD with random and MINEBED-BO, adding iDAD or DAD where their assumptions apply.Location Finding uses an sNMC upper bound; SIR and Cartpole use the InfoNCE lower bound, which may be biased because it relies on a learned critic.Source: §5 · p. 6, Table 2 · p. 8

5.1. Location Finding

  • In Location Finding, RL-DAD significantly outperforms random and MINEBED-BO, becoming competitive with iDAD at N = 10 and 15 and approaching DAD at N = 15.A measurable performance gap from iDAD and DAD remains; the authors hypothesize deep RL is limiting, since an analytic-likelihood policy also failed to match them.Source: §5.1 · p. 7, Table 1 · p. 7

5.1.1. ABLATION STUDY: SPARSE REWARDS

  • Sparse-reward agents often fail to converge, whereas dense-reward agents quickly become competitive with agents trained using the sPCE-based reward.The two reward signals have the same sum over each trajectory, but sparse-reward training can collapse after initial random exploration.Source: §5.1.1 · p. 7, Figure 1 · p. 7

5.2. SIR Model

  • In the SIR model, RL-DAD outperforms random and MINEBED-BO and remains competitive with iDAD; posterior estimates are generally consistent with ground truth.The model estimates infection and recovery parameters from sequential measurements of infected individuals.Source: §5.2 · pp. 7–8, Table 2 · p. 8

5.3. Cartpole

  • In Cartpole, RL-DAD's posterior estimates are generally tight around ground truth, and iDAD is unavailable because Isaac Gym provides no gradient paths.The unknown parameters are friction and pole mass; observations are limited to the end of each timestep.Source: §5.3 · p. 8, Table 2 · p. 8

6. Discussion

RL-DAD is presented as a general BOED approach compatible with black-box environment simulators, with promising, baseline-competitive results on three benchmarks. Its main stated limitation is the training burden.

  • On three benchmark tasks, RL-DAD shows promising performance and is competitive with baselines.Source: §6 · p. 8
  • The primary limitation is the initial training time required.Source: §6 · p. 8
  • Deep RL also brings sample inefficiency and brittleness to hyperparameters and implementation; future work targets training cost and simulated-sample requirements.Source: §6 · pp. 8–9

A.1. Decomposing the Total Expected Information Gain as a Sum of Marginals

The appendix decomposes total expected information gain into marginal terms and derives its policy-level expression, including the stochastic-policy case.

A.1. Decomposing the Total Expected Information Gain as a Sum of Marginals

  • The chain rule decomposes total expected information gain into marginal information terms conditioned on prior designs and outcomes.Source: §A.1 · p. 11

A.2. Expected Total Information Gain for a Policy

  • For a policy, total expected information gain is expressed using the expected log likelihood ratio of histories generated with and without the parameter.Source: §A.2 · pp. 11–12
  • The policy-level result applies to stochastic policies, whose designs are sampled from the policy, and its information-gain expression can be interpreted as mutual information.Source: §A.2 · p. 11

A.3. InfoNCE Bound for Stochastic Design Policies

This appendix establishes the InfoNCE lower bound for policy-based BOED and explains its extension to stochastic design policies.

  • InfoNCE lower-bounds total information gain for implicit policy-based BOED.Source: §A.3 · p. 13
  • With the optimal critic, the bound recovers the sPCE information-gain bound and is tight as the number of contrastive samples L approaches infinity.Source: §A.3 · p. 13
  • The InfoNCE result applies to stochastic policies as well as deterministic policies.Source: §A.3 · p. 13

A.4. Showing the Markovian Property for the MDP Underlying the BOED POMDP

The appendix shows that the BOED POMDP’s underlying states support Markovian transitions and rewards because each state contains the preceding history.

  • The appendix establishes that the BOED POMDP’s underlying states yield Markovian transition dynamics and rewards.Source: §A.4 · p. 13
  • The Markov property requires the next state and reward to depend on the current state and action, rather than the full trajectory.Source: §A.4 · p. 13
  • Because the current state is the full history, it contains prior states and actions, making it sufficient for the Markov property.Source: §A.4 · p. 13

A.5. The Dense Reward Signal Approximates the Marginal Expected Information Gain Values

The dense reward trajectory sum matches the sparse terminal reward, while each dense reward approaches the marginal expected information gain in a stated limit.

A.5. The Dense Reward Signal Approximates the Marginal Expected Information Gain Values

  • Summing dense rewards over a trajectory gives the sparse terminal reward because the sum telescopes and g(h_0, U_ψ; L) is defined as zero.The remaining term, g(h_T, U_ψ; L), is the reward assigned at the last step under the sparse formulation.Source: §A.5 · p. 13
  • The InfoNCE bound used to connect the reward and mutual information becomes tight as the critic approaches optimality and L→∞.The proof invokes the mutual-information chain rule.Source: §A.5 · pp. 13–14
  • In the limit of an optimal critic and L→∞, each dense reward’s expectation approaches the marginal expected information gain.That gain is I(θ; h_t+1 | π, h_t).Source: §A.5 · p. 14

B. Experiment Details

The experiment details specify a twin-critic training setup and report that a discount factor below one stabilized training.

B. Experiment Details

Related papers

Loading 2203.04272v1…