Statistically Efficient Bayesian Sequential Experiment Design via Reinforcement Learning with Cross-Entropy Estimators
Extracted source text · arXiv 2305.18435
\documentclass[nohyperref]{article}
\input{preamble}
\input{math}
\input{acronyms}
\begin{document}
\input{content}
\end{document}
% if you need to pass options to natbib, use, e.g.:
% \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2022
% ready for submission
\usepackage[nonatbib,preprint]{neurips_2023}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
% \usepackage[preprint]{neurips_2022}
% to compile a camera-ready version, add the [final] option, e.g.:
% \usepackage[final]{neurips_2022}
% to avoid loading the natbib package, add option nonatbib:
% \usepackage[nonatbib]{neurips_2022}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage[hidelinks]{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
\usepackage{xcolor} % colors
\usepackage{algorithmic}
\usepackage[ruled, linesnumbered]{algorithm2e}
% Attempt to make hyperref and algorithmic work together better:
\newcommand{\theHalgorithm}{\arabic{algorithm}}
% For theorems and such
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{amsthm}
% if you use cleveref..
\usepackage[capitalize,noabbrev]{cleveref}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% THEOREMS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\theoremstyle{plain}
\newtheorem{theorem}{Theorem}
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{corollary}[theorem]{Corollary}
\theoremstyle{definition}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{assumption}[theorem]{Assumption}
\theoremstyle{remark}
\newtheorem{remark}[theorem]{Remark}
% Todonotes is useful during development; simply uncomment the next line
% and comment out the line below the next line to turn off comments
%\usepackage[disable,textsize=tiny]{todonotes}
\usepackage[textsize=tiny]{todonotes}
\usepackage[utf8]{inputenc}
\usepackage[numbers]{natbib}
\usepackage[acronym,nowarn]{glossaries}
\glsdisablehyper
\makeglossaries
\usepackage{color}
\usepackage{cancel}
\newcommand{\tom}[1]{\textcolor{cyan}{[{\sc tom}] #1}}
\newcommand{\iadine}[1]{\textcolor{magenta}{[{\sc ic}] #1}}
\def\UrlBreaks{\do\/\do-}
% Designs
\newcommand{\opt}{*}
\newcommand{\dopt}{d^\opt}
% Policies
\newcommand{\piopt}{\pi^\opt}
\newcommand{\mathbold}[1]{\ensuremath{\boldsymbol{\mathbf{#1}}}}
% # PROBABILITY
\newcommand{\g}{\,|\,}
% \renewcommand{\gg}{\,\|\,}
\renewcommand{\d}[1]{\ensuremath{\operatorname{d}\!{#1}}}
\newcommand{\nestedmathbold}[1]{{\mathbold{#1}}}
% # BOLD MATHEMATICS
\newcommand{\mba}{\nestedmathbold{a}}
\newcommand{\mbb}{\nestedmathbold{b}}
\newcommand{\mbc}{\nestedmathbold{c}}
\newcommand{\mbd}{\nestedmathbold{d}}
\newcommand{\mbe}{\nestedmathbold{e}}
\newcommand{\mbf}{\nestedmathbold{f}}
\newcommand{\mbg}{\nestedmathbold{g}}
\newcommand{\mbh}{\nestedmathbold{h}}
\newcommand{\mbi}{\nestedmathbold{i}}
\newcommand{\mbj}{\nestedmathbold{j}}
\newcommand{\mbk}{\nestedmathbold{k}}
\newcommand{\mbl}{\nestedmathbold{l}}
\newcommand{\mbm}{\nestedmathbold{m}}
\newcommand{\mbn}{\nestedmathbold{n}}
\newcommand{\mbo}{\nestedmathbold{o}}
\newcommand{\mbp}{\nestedmathbold{p}}
\newcommand{\mbq}{\nestedmathbold{q}}
\newcommand{\mbr}{\nestedmathbold{r}}
\newcommand{\mbs}{\nestedmathbold{s}}
\newcommand{\mbt}{\nestedmathbold{t}}
\newcommand{\mbu}{\nestedmathbold{u}}
\newcommand{\mbv}{\nestedmathbold{v}}
\newcommand{\mbw}{\nestedmathbold{w}}
\newcommand{\mbx}{\nestedmathbold{x}}
\newcommand{\mby}{\nestedmathbold{y}}
\newcommand{\mbz}{\nestedmathbold{z}}
\newcommand{\mbA}{\nestedmathbold{A}}
\newcommand{\mbB}{\nestedmathbold{B}}
\newcommand{\mbC}{\nestedmathbold{C}}
\newcommand{\mbD}{\nestedmathbold{D}}
\newcommand{\mbE}{\nestedmathbold{E}}
\newcommand{\mbF}{\nestedmathbold{F}}
\newcommand{\mbG}{\nestedmathbold{G}}
\newcommand{\mbH}{\nestedmathbold{H}}
\newcommand{\mbI}{\nestedmathbold{I}}
\newcommand{\mbJ}{\nestedmathbold{J}}
\newcommand{\mbK}{\nestedmathbold{K}}
\newcommand{\mbL}{\nestedmathbold{L}}
\newcommand{\mbM}{\nestedmathbold{M}}
\newcommand{\mbN}{\nestedmathbold{N}}
\newcommand{\mbO}{\nestedmathbold{O}}
\newcommand{\mbP}{\nestedmathbold{P}}
\newcommand{\mbQ}{\nestedmathbold{Q}}
\newcommand{\mbR}{\nestedmathbold{R}}
\newcommand{\mbS}{\nestedmathbold{S}}
\newcommand{\mbT}{\nestedmathbold{T}}
\newcommand{\mbU}{\nestedmathbold{U}}
\newcommand{\mbV}{\nestedmathbold{V}}
\newcommand{\mbW}{\nestedmathbold{W}}
\newcommand{\mbX}{\nestedmathbold{X}}
\newcommand{\mbY}{\nestedmathbold{Y}}
\newcommand{\mbZ}{\nestedmathbold{Z}}
\newcommand{\mbalpha}{\nestedmathbold{\alpha}}
\newcommand{\mbbeta}{\nestedmathbold{\beta}}
\newcommand{\mbdelta}{\nestedmathbold{\delta}}
\newcommand{\mbepsilon}{\nestedmathbold{\epsilon}}
\newcommand{\mbchi}{\nestedmathbold{\chi}}
\newcommand{\mbeta}{\nestedmathbold{\eta}}
\newcommand{\mbgamma}{\nestedmathbold{\gamma}}
\newcommand{\mbiota}{\nestedmathbold{\iota}}
\newcommand{\mbkappa}{\nestedmathbold{\kappa}}
\newcommand{\mblambda}{\nestedmathbold{\lambda}}
\newcommand{\mbmu}{\nestedmathbold{\mu}}
\newcommand{\mbnu}{\nestedmathbold{\nu}}
\newcommand{\mbomega}{\nestedmathbold{\omega}}
\newcommand{\mbphi}{\nestedmathbold{\phi}}
\newcommand{\mbpi}{\nestedmathbold{\pi}}
\newcommand{\mbpsi}{\nestedmathbold{\psi}}
\newcommand{\mbrho}{\nestedmathbold{\rho}}
\newcommand{\mbsigma}{\nestedmathbold{\sigma}}
\newcommand{\mbtau}{\nestedmathbold{\tau}}
\newcommand{\mbtheta}{\nestedmathbold{\theta}}
\newcommand{\mbupsilon}{\nestedmathbold{\upsilon}}
\newcommand{\mbvarepsilon}{\nestedmathbold{\varepsilon}}
\newcommand{\mbvarphi}{\nestedmathbold{\varphi}}
\newcommand{\mbvartheta}{\nestedmathbold{\vartheta}}
\newcommand{\mbvarrho}{\nestedmathbold{\varrho}}
\newcommand{\mbxi}{\nestedmathbold{\xi}}
\newcommand{\mbzeta}{\nestedmathbold{\zeta}}
\newcommand{\mbDelta}{\nestedmathbold{\Delta}}
\newcommand{\mbGamma}{\nestedmathbold{\Gamma}}
\newcommand{\mbLambda}{\nestedmathbold{\Lambda}}
\newcommand{\mbOmega}{\nestedmathbold{\Omega}}
\newcommand{\mbPhi}{\nestedmathbold{\Phi}}
\newcommand{\mbPi}{\nestedmathbold{\Pi}}
\newcommand{\mbPsi}{\nestedmathbold{\Psi}}
\newcommand{\mbSigma}{\nestedmathbold{\Sigma}}
\newcommand{\mbTheta}{\nestedmathbold{\Theta}}
\newcommand{\mbUpsilon}{\nestedmathbold{\Upsilon}}
\newcommand{\mbXi}{\nestedmathbold{\Xi}}
\newcommand{\mbzero}{\nestedmathbold{0}}
\newcommand{\mbone}{\nestedmathbold{1}}
\newcommand{\mbtwo}{\nestedmathbold{2}}
\newcommand{\mbthree}{\nestedmathbold{3}}
\newcommand{\mbfour}{\nestedmathbold{4}}
\newcommand{\mbfive}{\nestedmathbold{5}}
\newcommand{\mbsix}{\nestedmathbold{6}}
\newcommand{\mbseven}{\nestedmathbold{7}}
\newcommand{\mbeight}{\nestedmathbold{8}}
\newcommand{\mbnine}{\nestedmathbold{9}}
% # MISCELLANEOUS
\newcommand{\ELBO}{\textsc{elbo}}
\newcommand{\GELBO}{\textsc{gelbo}}
\newcommand{\scH}{\textsc{h}}
\DeclareRobustCommand{\KL}[2]{\ensuremath{\textsc{kl}\left[#1\;\|\;#2\right]}}
\DeclareRobustCommand{\Df}[2]{\ensuremath{\mathcal{D}_f\left[#1\;\|\;#2\right]}}
\newcommand{\diag}{\textrm{diag}}
\newcommand{\supp}{\textrm{supp}}
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\newcommand\indep{\protect\mathpalette{\protect\independenT}{\perp}}
\def\independenT#1#2{\mathrel{\rlap{$#1#2$}\mkern2mu{#1#2}}}
\newcommand{\cA}{\mathcal{A}}
\newcommand{\cB}{\mathcal{B}}
\newcommand{\cD}{\mathcal{D}}
\newcommand{\cE}{\mathcal{E}}
\newcommand{\cF}{\mathcal{F}}
\newcommand{\cG}{\mathcal{G}}
\newcommand{\cH}{\mathcal{H}}
\newcommand{\cI}{\mathcal{I}}
\newcommand{\cL}{\mathcal{L}}
\newcommand{\cM}{\mathcal{M}}
\newcommand{\cN}{\mathcal{N}}
\newcommand{\cP}{\mathcal{P}}
\newcommand{\cQ}{\mathcal{Q}}
\newcommand{\cR}{\mathcal{R}}
\newcommand{\cS}{\mathcal{S}}
\newcommand{\cT}{\mathcal{T}}
\newcommand{\cU}{\mathcal{U}}
\newcommand{\cV}{\mathcal{V}}
\newcommand{\cY}{\mathcal{Y}}
\newcommand{\E}{\mathbb{E}}
\newcommand{\bbH}{\mathbb{H}}
\newcommand{\bbR}{\mathbb{R}}
\newcommand{\bbN}{\mathbb{N}}
\newcommand{\bbC}{\mathbb{C}}
\newcommand{\bbV}{\mathbb{V}}
% GP stuff
\newcommand{\bigO}{\mathcal{O}}
\newcommand{\GP}{{\mathcal{GP}}}
% Distributions
\newcommand{\Poisson}{\text{Poisson}}
\newacronym{DL}{dl}{deep learning}
\newacronym{VI}{VI}{variational inference}
\newacronym{BO}{BO}{Bayesian optimization}
\newacronym{DOE}{DoE}{design of experiments}
\newacronym{SGD}{SGD}{stochastic gradient descent}
\newacronym{EIG}{EIG}{expected information gain}
\newacronym{BOED}{BOED}{Bayesian optimal experimental design}
\newacronym[longplural=Markov decision processes]{MDP}{MDP}{Markov decision process}
\newacronym{RL}{RL}{reinforcement learning}
\newacronym{PCE}{PCE}{prior contrastive estimation}
\newacronym{SMC}{SMC}{sequential Monte Carlo}
\newacronym{SPCE}{sPCE}{sequential prior contrastive estimation}
\newacronym{SCEE}{sCEE}{sequential cross-entropy estimator}
\newacronym{DAD}{DAD}{deep adaptive design}
\newacronym{SED}{SED}{sequential experimental design}
\newacronym{REDQ}{REDQ}{Randomized ensembled double q-learning}
\newacronym{SNMC}{sNMC}{sequential nested monte carlo}
\newacronym{CES}{CES}{constant elasticity of substitution}
\newacronym{hipmdp}{HIP-MDP}{Hidden Parameter MDP}
\title{Cross-Entropy Estimators for Sequential Experiment Design with Reinforcement Learning}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
\author{%
Tom Blau \\
CSIRO's Data61\\
Australia\\
\texttt{tom.blau@data61.csiro.au} \\
\And
Edwin Bonilla \\
CSIRO's Data61\\
Australia\\
\And
Iadine Chades \\
CSIRO's Land and Water\\
Australia\\
\And
Amir Dezfouli \\
CSIRO's Data61\\
Australia\\
}
\maketitle
\begin{abstract}
Reinforcement learning can effectively learn amortised design policies for designing sequences of experiments. However, current methods rely on contrastive estimators of expected information gain, which require an exponential number of contrastive samples to achieve an unbiased estimation. We propose an alternative lower bound estimator, based on the cross-entropy of the joint model distribution and a flexible proposal distribution. This proposal distribution approximates the true posterior of the model parameters given the experimental history and the design policy. Our estimator requires no contrastive samples, can achieve more accurate estimates of high information gains, allows learning of superior design policies, and is compatible with implicit probabilistic models. We assess our algorithm's performance in various tasks, including continuous and discrete designs and explicit and implicit likelihoods.
\end{abstract}
\section{Introduction}
A key challenge in science is to develop predictive models based on experimental observations. As far back as~\citet{lindley1956measure} it has been recognised that experimental designs can be optimised to be maximally informative, under the assumptions of a Bayesian framework. Since then optimal experimental design has been applied to a wide variety of fields with different models and assumptions, including neuroscience~\cite{shababo2013bayesian}, biology~\cite{treloar2022deep}, ecology~\cite{drovandi2014sequential} and causal structure learning~\cite{agrawal2019abcd}.
Under the framework of \gls{BOED}, we have a probabilistic model $p(y \g \theta, d)$ where $d$ is the design (e.g. where to measure), $y$ is the outcome (e.g. the measurement value) and $\theta$ are parameters over which we have a prior belief $p(\theta)$.
The objective is to find the optimal design $d^*$ that maximises the \gls{EIG}, formally:
\begin{align}
\label{eq:first-eig}
EIG(d) & \coloneqq \E_{p(y \g d)} \left[H(p(\theta)) - H(p(\theta \g y, d ))\right],\\
\label{eq:first-optimal-d-eig}
d^* &= \argmax_{d \in \mathcal{D}} EIG(d).
\end{align}
We see that na\"ive computation of the \gls{EIG} requires an expectation over the marginal likelihood $p(y \g d)$ and estimation of the posterior $p(\theta \g y, d)$. Since sampling from $p(y \g d)$ is typically intractable and there is usually no closed-form solution for the posterior, minimising this expression involves estimating a nested expectation numerically, which is challenging~\cite{rainforth2018nesting}. Furthermore, we are often interested in conducting more than one experiment, in which case optimal designs must incorporate the outcomes of previous experiments sequentially~\cite{krause2007nonmyopic}.
In settings where computational or application-specific constraints demand fast deployment times, the state-of-the-art approach to \gls{SED} is to use \textit{amortised} methods, which learn an optimal design policy as a function of the experimental history instead of optimising each design in turn~\cite{blau2022optimizing,foster2021deep,ivanova2021implicit}. Once trained, a policy can be reused to design experiments as many times as desired, thus amortising the cost of training.
However, all amortised methods introduced thus far have the drawback that they rely on maximising contrastive lower bounds of the objective. To achieve an unbiased estimate of the \gls{EIG}, these contrastive bounds require a number of samples that is exponential in its magnitude~\cite{mcallester2020formal, poole2019variational}. Thus their performance degrades in cases where the \gls{EIG} is large.
To address this limitation, we propose a new amortised method, using \gls{RL} and a non-contrastive bound based on the cross-entropy of the joint model distribution and a flexible proposal distribution. This proposal approximates the true posterior of the model parameters given the experimental history and the design policy. Our method does not suffer from exponential sample complexity and is thus able to achieve higher \gls{EIG} than prior art, especially in settings where the information gain of the optimal policy is large. Furthermore, unlike previous amortised methods, our method is generally applicable to continuous and discrete design spaces, non-differentiable likelihoods, and even implicit likelihoods.
Our experiments provide supporting evidence of the benefits of our approach when compared to previous methods in these settings. We provide our code as supplementary material and will publish it upon acceptance.
In the next section we cover the theoretical background necessary to understand our proposed method, which we explain in detail in~\cref{sec:theory,sec:method}. Experiments then follow in~\cref{sec:results} to test the claimed benefits of our approach. Related literature is discussed in~\cref{sec:related} and finally we present conclusions and potential future work in~\cref{sec:conclusion}.
\section{Amortised design of experiments}\label{sec:background}
In \acrfull{BOED} the goal is to identify the parameters of a probabilistic model by sending queries to that model. Let $p(y \g \theta, d)$ be the model of concern, with some prior belief $p(\theta)$ regarding the value of parameters $\theta$.
As described above, an optimal design $d^*$ is one that maximises the \gls{EIG} as given by \cref{eq:first-eig,eq:first-optimal-d-eig}, where computational intractabilities readily appear in the estimation of the marginal likelihood and the posterior distribution.
Furthermore, more challenging than optimising a single experiment is the problem of optimising an entire sequence of experiments $d_{1:T}$ where $T \in \mathbb{N}$ is some fixed budget. The state-of-the-art approach to this problem is to forego optimising the designs directly, and instead optimise a design policy $\pi : \mathcal{H} \rightarrow \mathcal{D}$ that designs experiments conditioned on a history $h_t = (d_i, y_i)_{i=1}^t$. The computational cost of learning such a policy is high, but designing experiments with a trained policy is computationally efficient, requiring only a single forward pass of a neural network. Therefore the training cost is amortised over the lifetime of the policy, and this class of algorithms is known as \textit{amortised} design of experiments. ~\citet{foster2021deep} have shown that the \gls{EIG} of a design policy can be maximised through a lower bound called \gls{SPCE}:
\\
\begin{align}
EIG(\pi, T) \ge sPCE(\pi, T, L) = \E_{p(\theta_{0:L}) p(h_T \g \theta_0, \pi)} \left[ \log \frac{p(h_T \g \theta_0, \pi)}{\frac{1}{L+1} \sum^L_{l=0} p(h_T \g \theta_l, \pi)} \right], \label{eq:spce}
\end{align}
where $\theta_{1:L} \sim p(\theta)$ are \textit{contrastive samples}, giving the bound its name, and it is assumed that the probability of the history factorises into $p(h_T \g \theta, \pi)
= \prod_{t=1}^T p(y_t \g \theta, d_t)$. A corresponding upper bound, sequential nested Monte Carlo (sNMC) estimator, can be achieved by excluding the $l=0$ term from the denominator.
Current amortised methods all use contrastive bounds to optimise the policy. However, such bounds require a number of contrastive samples $L$ that is exponential in the magnitude of the quantity being estimated~\citep{mcallester2020formal}. In other words, if the \gls{EIG} of a policy is large, then computing an accurate contrastive bound is intractable.
\subsection{Reinforcement learning for experiment design}
Problems in reinforcement learning are formulated as \glspl{MDP}, which are tuples $<\mathcal{S}, \mathcal{A}, \mathcal{T}, \mathcal{R}, \rho_0>$. Here $\mathcal{S}$ and $\mathcal{A}$ are state and action spaces, $\mathcal{T} : \mathcal{S} \times \mathcal{A} \rightarrow \cP_s(\cS)$ is the transition dynamics, $\mathcal{R} : \mathcal{S} \times \mathcal{A} \times \mathcal{S} \rightarrow \bbR$ is a reward function and $\rho_0$ is the distribution of initial states.
For a given discount factor $\gamma \in [0,1]$ and time horizon $T \in \bbN$, the solution to the problem is a policy $\pi: \cS \rightarrow \cP_a(\cA)$ mapping states to (distributions over) actions that maximises the expected discounted return $J(\pi) \coloneqq \E_{\mathcal{T}, \pi, \rho_0} \left[ \sum^T_{t=1} \gamma^{t-1} \mathcal{R}(s_{t-1}, a_{t-1}, s_t )\right]$.
\citet{blau2022optimizing} have shown that the problem of learning an experiment design policy can be formulated as a special case of \gls{MDP} called the SED-MDP, which has the form $<\cS, \cA, \Theta, \cT, \cR, \rho_0, P_\Theta>$. Here $\Theta$ are unobserved parameters that are fixed over the course of a rollout, and correspond to the parameters $\theta$ of the likelihood model in design of experiments, where $P_\Theta$ is the prior. The reward function and transition dynamics are parameterised by $\theta$. The key idea is to map experimental designs to policy actions $a_{t-1} = d_t$ and history information to system states:
\\
\begin{align}
\begin{split}
s_t &= (B_{\psi,t}, C_t, y_t), \\
y_t &\sim p(y_t \g d_t, \theta_0) , \\
B_{\psi,t} &= B_{\psi,t-1} + ENC_\psi(d_t, y_t) \text{ and}\\
C_t &= C_{t-1} \odot \left[ p(y_t \g \theta_l, d_t) \right]_{l=0}^L,
\end{split}
\end{align}
\\
where $ENC_\psi(d_t, y_t)$ is a learned encoding of the most recent experiment, $\psi$ are the parameters of the encoding and $\odot$ is the Hadamard product. If the discount is set to $\gamma = 1$ and the reward function is set to $\mathcal{R}(s_{t-1}, a_{t-1}, s_t, \theta) = \log p(y_t \g \theta_0, d_t) - \log (C_t \cdot \mathbf{1}) + \log (C_{t-1} \cdot \mathbf{1})$, then the expected discounted return is identical to the \gls{SPCE} and the design policy can be optimised with any \gls{RL} algorithm we choose.
\section{The sequential cross-entropy estimator}\label{sec:theory}
From Theorem 2 of~\citet{foster2021deep}, we have that the total EIG of a sequence of $T$ experiments, given a design policy, is lower bounded by the \gls{SPCE}. However, since the term $p(h_T \g \theta_0, \pi)$ appears in both the numerator and denominator of \gls{SPCE} (as per~\cref{eq:spce}), the lower bound is itself upper bounded by $\log(L+1)$, where $L$ is the number of contrastive samples. Therefore, to provide an unbiased estimate of $EIG(\pi, T)$, we require an exponential number of contrastive samples $L \ge \exp(EIG(\pi, T)) - 1$.
To resolve this limitation, we will introduce a proposal distribution $q(\theta \g h_T, \pi)$ that approximates the true posterior $p(\theta \g h_T, \pi)$. Using the cross-entropy of the two, we construct the following estimator, which we refer to as the \gls{SCEE}:
\begin{equation}
\label{eq:sceet}
sCEE(\pi, T) \coloneqq \E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right] + H\left[ p(\theta) \right]
\end{equation}
From Jensen's inequality it follows that the cross-entropy of two random variables is a lower bound for the self-entropy of either variable. By extending this to the sequential case, the following theorem shows that the sCEE is a lower bound of the true EIG:
\begin{theorem}\label{thm:bound}
Let $p(y \g \theta, d)$ be a probabilistic model with prior $p(\theta)$. For an arbitrary fixed design policy $\pi$ and sequence length $T$, the \gls{EIG} of using $\pi$ to design $T$ experiments is denoted $EIG(\pi, T)$.
Let $q(\theta \g h_T, \pi)$ be a proposal distribution over parameters $\theta$ conditioned on experimental history $h_T$, and the sCEE bound is
\begin{equation}
sCEE(\pi, T) \coloneqq \E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right] + H\left[ p(\theta) \right]
\end{equation}
we have that
\begin{equation}
sCEE(\pi, T) \leq EIG(\pi, T)
\end{equation}
\end{theorem}
\begin{proof}
A sketch of proof follows here, with the full proof in~\cref{app:bound_proof}. The main idea is to rewrite the \gls{EIG} as an expectation w.r.t. distribution $p(h_T, \theta \g \pi)$, and then show that the difference between \gls{EIG} and \gls{SCEE} is an expectation over KL divergences.
\end{proof}
%
~\Cref{thm:bound} leads to $2$ corollaries that describe the nature of the lower bound (cf.~\cref{app:corollaries}):
\begin{corollary}
The bound is tight if and only if $q(\theta \g h_T, \pi) = p(\theta \g h_T, \pi)$
\end{corollary}
\begin{corollary}
The bias of the \gls{SCEE} estimator is $-\E_{h_T} \left[ \KL{p(\theta \g h_T, \pi)}{q(\theta \g h_T, \pi)} \right]$
\end{corollary}
In other words, the quality of the estimation rests on how well the proposal distribution can match the true posterior, in terms of KL divergence.
\subsection{Proposal parameterization}
To evaluate the \gls{SCEE}, we sample from the joint $p(\theta, h_T \g \pi)$ simply by rolling out the policy. Under mild assumptions, it can be shown that this Monte Carlo estimation approaches the true value of the sCEE at a rate of $O(\frac{1}{\sqrt{n}})$, where $n$ is the number of samples (cf.~\cref{app:convergence_proof}).
We will parameterise the proposal distribution by a conditional normalising flow~\cite{winkler2019learning} with parameters $\kappa$ and, therefore, refer to it using $q_\kappa(\cdot)$. Thus, we can maximise the sCEE w.r.t. $\kappa$ using stochastic gradient descent. Note that we only need to optimise the
negative cross-entropy term $\E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right]$, since the prior entropy is constant. Details of the normalising flows used in our experiments are in~\cref{app:algos}.
\section{Experiment design with sCEE and reinforcement learning}\label{sec:method}
To use the \gls{SCEE} bound in the context of an RL algorithm, we substitute sCEE for sPCE in the reward definition of the SED-MDP:
%
\begin{equation}
\mathcal{R}(s_{t-1}, a_{t-1}, s_t, \theta) = \log q(\theta \g B_{\psi,t}) - \log q(\theta \g B_{\psi,t-1}). \label{eq:reward}
\end{equation}
%
Policy and critic networks $\pi_\phi$ and $\mathcal{C}_\chi$ can be updated following the rules of the RL algorithm of our choice, using mini-batches of either off-policy or on-policy samples. The posterior network $q_\kappa(\cdot)$ can be updated by using the same mini-batches
to maximise the log-likelihood of the observations under our posterior model.
Note that this means rewards are now no longer fixed but depend on $q_\kappa(\cdot)$, and change with every update of $\kappa$. The computational cost thus incurred can be minimised by lazy evaluation~\cite{bloss1988code}: we only update each reward when we are about to use it to update the policy and critic networks of the RL agent. The procedure is summarised in~\cref{alg:scee-rl}.
\begin{algorithm}[t]
\caption{sCEE-RL}
\label{alg:scee-rl}
\begin{algorithmic}
\STATE {\bfseries Input:} $\mathcal{M}$: SED-MDP, $L_\pi$: policy loss function, $L_\mathcal{C}$: critic loss function
\STATE Initialise replay buffer $\cB$
\WHILE{convergence criterion not reached}
\STATE Generate rollouts $(s_{0:T}, a_{0:T}, \theta)^{1:N}$ using $\cM$ and $\pi$ and push to $\cB$.
\STATE Sample mini-batch of size $mb$ from $\cB$
\STATE Compute posterior loss $L_q = -\frac{1}{mb}\sum_{i=1}^{mb}\log q_\kappa(\theta^i \g B_{\psi,t}^i)$
\STATE Take gradient step to minimise $\nabla_\kappa L_q$
\STATE Compute rewards for mini-batch using~\cref{eq:reward}
\STATE Use mini-batch and rewards to compute $L_\pi$ and $L_\mathcal{C}$
\STATE Take gradient step to minimise $\nabla_\phi L_\pi$ and $\nabla_\chi L_\mathcal{C}$
\ENDWHILE
\end{algorithmic}
\end{algorithm}
\subsection{Advantages of sCEE-RL}
Our method based on the sCEE lower bound and RL delivers a number of advantages. \textbf{(i) Better sample complexity}: it does not require the use of contrastive samples, and hence does not suffer from the exponential sample complexity issue of the sPCE bound.
Thus, sCEE can more closely estimate \gls{EIG} when the true quantity is large, although estimation accuracy depends on learning a good
posterior network $q_\kappa(\cdot)$. \textbf{(ii) Applicable to implicit models}:
Furthermore, we see that the sCEE estimator, as defined in \cref{eq:sceet}, only requires sampling of the model distribution and avoids explicit log-likelihood computations $\log p(h_T \g \theta, \pi)$. This means that our method is compatible with implicit likelihood models where the likelihood is a black-box or intractable and, therefore, can only be sampled but not evaluated explicitly.
Interestingly, the sCEE bound is closely related to the sACE bound introduced in the appendices of~\citet{foster2021deep}. We discuss this relationship in~\cref{app:sace}. \textbf{(iii) Suitable for continuous and discrete design spaces}:
Finally, similar to the method proposed in
\citet{blau2022optimizing}, our approach using the sCEE estimator along with reinforcement learning, as described in \cref{alg:scee-rl}, can handle both continuous and discrete design spaces.
\subsection{Simultaneous policy and reward learning}
We propose to learn the design policy $\pi_\phi$ and the proposal distribution $q_\kappa$ from data simultaneously. Since the reward function depends on $q_\kappa$, and the objective function of $q_\kappa$ in turn depends on $\pi_\phi$, this leads to inherent instability, similar to the "deadly triad" that is often observed in value-based reinforcement learning~\citep{van2018deep}. We therefore apply several stabilisation mechanisms to prevent the neural network estimators from diverging.
\textbf{Target posterior network:} similar to the use of target Q-networks as introduced by~\citet{lillicrap2016continuous}, we maintain a primary posterior network $q_\kappa$ and a target network $q_\kappa^\prime$. The primary network $q_\kappa$ is updated using gradient descent in every iteration of the algorithm, but is not used directly to compute rewards. Instead, the target network $q_\kappa^\prime$ is used to compute~\cref{eq:reward}, and its weights are periodically updated to maintain a moving average:
\begin{equation}
\kappa^\prime \gets \kappa^\prime \cdot (1 - \tau) + \kappa \cdot \tau
\end{equation}
where $\tau \in (0,1)$ is a constant controlling the rate of change.
\textbf{Fixed % prior likelihood
initial posterior:} the reward definition of~\cref{eq:reward} assigns each experiment its own (estimated) information gain.
The return of an entire trajectory is a telescoping sum that reduces to $\log q(\theta \g B_{\psi,T}) - \log q(\theta \g B_{\psi, 0})$, and the expected return over infinitely many trajectories recovers the sCEE. Therefore, the component $q(\theta \g B_{\psi,0})$ of the first reward $r_0$ is the only contributor to the prior entropy term $H[p(\theta)]$ of the sCEE. Since this term is constant w.r.t. all networks, we can simply ignore it when training as it does not change the optimal policy. Furthermore, learning the correct estimator for $q(\theta \g B_{\psi,0})$ that maps the null inputs to the prior $p(\theta)$ can be challenging. Therefore instead of learning this mapping for the empty first state, we assigned it a fixed value of $\log q(\theta \g B_{\psi,0}) \coloneqq 0$.
\section{Experimental results}\label{sec:results}
In the following, we evaluate the performance and properties or our proposed method on a number of \gls{SED} problems. We compare our method (RL-sCEE) with a number of baselines, including RL with the sPCE bound (RL-sPCE)~\cite{blau2022optimizing}, Deep Adaptive Design (DAD)~\cite{foster2021deep}, implicit Deep Adaptive Design (iDAD)~\cite{ivanova2021implicit}, and a non-amortised sequential Monte Carlo experiment design approach (SMC-ED)~\cite{moffat2020sequential}. The Pyro~\cite{bingham2018pyro} and normflows~\cite{normflows} frameworks were used to implement our algorithm. For reinforcement learning, we use the Garage framework~\cite{garage} and the REDQ algorithm~\cite{chen2020randomized}. For complete details about algorithms and hyperparameters, see~\cref{app:algos}.
To evaluate the \gls{EIG} we used contrastive estimators with $L=1\textrm{E}8$, a number of contrastive samples that is impractical for learning, but achieves better estimation than sCEE in most problems we investigated. We discuss this limitation further in~\cref{sec:conclusion}. Additional results and ablation studies are included in the appendices.
\subsection{Accuracy of EIG estimators on synthetic data}
Given the theoretical results about sCEE and contrastive bounds, we expect that sCEE should perform well in situations where the \gls{EIG} is large and $q_\kappa(\cdot)$ is easy to learn. To assess this, we evaluate the estimator on $7$ experimental tasks which allow us to know the true \gls{EIG} in closed form. The priors are isotropic Gaussians of the form $\cN(\mu_0, \sigma_0\mathbf{I}_k)$ and the likelihoods are similarly Gaussian with known isotropic covariance $\sigma\mathbf{I}_k$, where $k$ is the number of dimensions. Each task has an experimental budget of $T=10$ experiments. Thus we can manipulate $k, \sigma_0$ and $\sigma$ to create tasks where the \gls{EIG} of the optimal design is known exactly. See~\cref{app:conjugate} for details.
\Cref{table:estimators} enumerates these tasks, alongside the optimal \gls{EIG} and the estimates of sCEE and sPCE with different numbers of contrastive samples. As can be seen from the left-most columns of the table, when the \gls{EIG} is small enough, sPCE can provide a better estimate than sCEE (note that the sPCE at times slightly overestimates the \gls{EIG} due to variance in estimating the expectation with Monte Carlo samples). However, as the \gls{EIG} becomes large relative to $\log(L)$, the underestimation of sPCE becomes more severe, and for the right-most columns all sPCE variants have reached their upper limit. Meanwhile, sCEE consistently provides good estimates regardless of the magnitude of the \gls{EIG}. It should be noted, however, that this is in part because the posterior is easy to learn from data. A more complex posterior, or less training, would worsen the underestimation.
\begin{table}[h]
\caption{Different estimators for EIG of increasing magnitudes in synthetic data problems with conjugate priors. Averages computed over $1000$ samples. $k$ is the number of random variable dimensions, $\sigma_0$ is prior variance, and $\sigma$ is likelihood variance.}
\label{table:estimators}
\begin{center}
\begin{small}
\vskip -0.1in
\begin{tabular}{lccccccc}
\toprule
Method & $k=10$ & $k=10$ & $k=10$ & $k=10$ & $k=10$ & $k=10$ & $k=20$ \\
& $\sigma_0 = 0.5$ & $\sigma_0 = 0.5$ & $\sigma_0 = 1$ & $\sigma_0 = 2$ & $\sigma_0 = 2$ & $\sigma_0 = 4$ & $\sigma_0 = 4$ \\
& $\sigma = 5$ & $\sigma = 1$ & $\sigma = 1$ & $\sigma = 1$ & $\sigma = 0.5$ & $\sigma = 0.5$ & $\sigma = 0.5$ \\
\midrule
True EIG & $3.47$ & $8.96$ & $11.99$ & $15.22$ & $18.57$ & $21.97$ & $43.94$ \\
sCEE & $3.40$ & $\mathbf{8.90}$ & $\mathbf{11.92}$ & $\mathbf{15.07}$ & $\mathbf{18.41}$ & $\mathbf{20.47}$ & $\mathbf{43.89}$ \\
sPCE($L=1\textrm{E}4)$ & $\mathbf{3.45}$ & $7.92$ & $8.95$ & $9.18$ & $9.21$ & $9.21$ & $9.21$ \\
sPCE($L=1\textrm{E}6)$ & $3.48$ & $8.89$ & $11.45$ & $13.18$ & $13.75$ & $13.81$ & $13.81$ \\
sPCE($L=1\textrm{E}8)$ & $3.48$ & $8.97$ & $11.85$ & $14.35$ & $16.71$ & $18.08$ & $18.42$ \\
\bottomrule
\end{tabular}
\end{small}
\end{center}
\vskip -0.1in
\end{table}
\subsection{Continuous designs}
\begin{figure}[ht]
\vskip 0.1in
\begin{center}
\centerline{
\includegraphics[width=0.5\columnwidth]{figures/ces_results.png}
\includegraphics[width=0.5\columnwidth]{figures/source_results.png}
}
\caption{\gls{EIG} for the CES problem (left) and the source location problem (right), estimated using \gls{SPCE} with $L=1\textrm{E}8$. Trendlines are means and shaded regions are standard errors aggregated from $1000$ rollouts. Our method is referred to as RL-sCEE.}
\label{fig:ces_source}
\end{center}
\vskip -0.2in
\end{figure}
\begin{table}[h]
\caption{Lower and upper bounds for the EIG computed using the sPCE and sNMC estimators, respectively. $L=1\textrm{E}8$ contrastive samples were used for the CES and Source Location problems, and $L=1\textrm{E}6$ for the Prey Population problem. Means and standard errors aggregated from $1000$ rollouts.}
\label{table:eig}
\setlength\tabcolsep{4.75 pt}
\begin{center}
\begin{small}
\vskip -0.1in
\begin{tabular}{lcccccc}
\toprule
Method & \multicolumn{2}{c}{CES} & \multicolumn{2}{c}{Source Location} & \multicolumn{2}{c}{Prey Population} \\
& Lower bound & Upper bound & Lower bound & Upper bound & Lower bound & Upper bound \\
\midrule
RL-sCEE & $\mathbf{15.91} \pm 0.10$ & $20.78 \pm 0.43$ & $\mathbf{13.37} \pm 0.07$ & $13.42 \pm 0.08$ & $4.41 \pm 0.05$ & $4.41 \pm 0.05$ \\
RL-sPCE & $14.81 \pm 0.12$ & $15.56 \pm 0.17$ & $11.65 \pm 0.06$ & $12.01 \pm 0.07$ & $4.38 \pm 0.05$ & $4.41 \pm 0.04$ \\
DAD & $10.77 \pm 0.08$ & $13.20 \pm 0.68$ & $11.22 \pm 0.07$ & $11.29 \pm 0.07$ & N/A & N/A \\
iDAD & $9.67 \pm 0.08$ & $10.63 \pm 0.52$ & $10.37 \pm 0.07$ & $10.41 \pm 0.08$ & N/A & N/A \\
SMC-ED & N/A & N/A & N/A & N/A & $\mathbf{4.52} \pm 0.07$ & $4.52 \pm 0.06$ \\
\bottomrule
\end{tabular}
\end{small}
\end{center}
\vskip -0.1in
\end{table}
\begin{figure}[ht]
\vskip 0.1in
\begin{center}
\centerline{
\includegraphics[width=1.2\columnwidth]{figures/ces_post.png}
}
\caption{Example posteriors for the CES problem after $10$ experiments. Histograms are based on $1\textrm{E}5$ samples. Dashed vertical lines indicate the ground truth value of each variable. Middle plot shows the marginals of the $3$ different elements of $\alpha$.}
\label{fig:post_ces}
\end{center}
\vskip -0.2in
\end{figure}
\begin{figure}[ht]
\vskip 0.1in
\begin{center}
\centerline{
\includegraphics[width=\columnwidth]{figures/correlation_source.png}
}
\caption{Histograms of example posteriors for the source location problem after $30$ experiments, showing the joint distributions of the $x$ co-ordinates (left) and $y$ co-ordinates (right) of the $2$ sources. The plots show symmetry with respect to the dashed red line, which is predicted by Bayes' theorem. Insets zoom in on the modes of each posterior. Histograms are based on $1\textrm{E}5$ samples. Black rings denote the ground truth value of each variable.}
\label{fig:post_source}
\end{center}
\vskip -0.2in
\end{figure}
\textbf{Constant elasticity of substitution (CES)}: Moving on to realistic tasks, we evaluate a design problem in behavioural economics where we must estimate the parameters of a Constant Elasticity of Substitution (CES) utility function~\cite{baltas2001utility}. Experimental designs consist of $2$ bags of goods with $3$ goods in each, so that the design space is $\cD = [0,100]^6$. The outcome is the relative preference of a test subject in the range $[0,1]$, as determined by the agent's CES utility function, and the specific values of its parameters $\theta = \{\rho, \alpha, u\}$, with $\rho \in [0,1]$, $ \alpha \in \Delta_3$ and $u > 0$. See~\cref{app:ces} for full details.
The left-most plot of~\cref{fig:ces_source} shows the \gls{EIG} achieved at each point in a sequence of $T=10$ experiments. \gls{EIG} was estimated by the sPCE estimator with $L=1\textrm{E}8$ contrastive samples. Our proposed method, RL-sCEE, performs significantly better than all baselines. Indeed, by the $5^{th}$ experiment, our method already exceeds the performance that the state-of-the-art RL-sPCE baseline attains after $10$ experiments. Considering the sNMC upper bound shown in~\cref{table:eig}, it is clear that the \textit{lower} bound for RL-sCEE is higher than the \textit{upper} bound for RL-sPCE. This is particularly important since the \gls{EIG} for the RL-sCEE policy is approaching the limits of what can be estimated by the lower bound, so that the gap between the true \gls{EIG} and the estimate is potentially large.
Finally, since we trained the conditional normalising flow $q_\kappa$ to predict posteriors, we can now efficiently sample from $q_\kappa(\theta \g h_T)$ for any given history. It is important to note, however, that such posteriors can be expected to be accurate only for histories sampled from the joint $p(\theta, \pi_\phi)$ where $\pi_\phi$ is the policy that was trained alongside $q_\kappa$. A representative posterior is shown in~\cref{fig:post_ces}. The histograms were computed using $1\textrm{E}5$ samples and $500$ bins. The marginal distribution for each random variable is unimodal and highly concentrated. The mode of each marginal is very close to the true value of the corresponding random variable, denoted by a dashed vertical line of the same colour.
\textbf{Source location}: In the next experiment we consider the problem of locating signal sources in space. Here we have $2$ sources with co-ordinates $s_1 = (x_1,y_1); s_2=(x_2,y_2)$, and designs consist of choosing the co-ordinates at which to take a noisy sample of signal magnitude. Signals decay with distance from the source and individual signals always superpose constructively. The full probabilistic model is described in~\cref{app:source}.
We trained policies to design sequences of $T=30$ experiments. The mean and standard error of \gls{EIG} achieved by our proposed method, as well as several baselines, are shown in the right-most plot of~\cref{fig:ces_source}. RL-sCEE outperforms all baselines by a statistically significant margin. As before, ~\cref{table:eig} shows that the lower bound estimate for our proposed method exceeds the upper bound estimate for all other methods.
As with the CES problem, we can examine the posteriors obtained from $q_\kappa(\theta \g h_T)$. This time, however, we plotted the marginals $q(x_1,x_2)$ and $q(y_1, y_2)$, i.e. the $x$ and $y$ co-ordinates of the $2$ signal sources, using $x_1$ (respectively $y_1$) as the X-axis and $x_2$ (respectively $y_2$) as the Y-axis. The sources are exchangeable, i.e. $p(y|d, s_1=(x_1,y_1), s_2=(x_2,y_2)) = p(y|d, s_1=(x_2,y_2), s_2=(x_1,y_1))$. Therefore the marginals of the true Bayesian posterior, $p(x_1,x_2)$, should be symmetric w.r.t. the line $x_2 = x_1$ in the $x_1x_2$-plane. Indeed, the plots in~\cref{fig:post_source} exhibit this symmetry, which has been learned entirely from data, without providing any inductive bias. We include a histogram of the full posterior in~\cref{app:results}.
\textbf{Implicit likelihood}: in both the CES and source location problem, we simulate an implicit likelihood by withholding the explicit likelihood values $p(y \g \theta, d)$ from the RL-sCEE and iDAD agents. One would expect that both agents would perform worse than the RL-sPCE and DAD agents, which do exploit explicit likelihood information. Indeed, the iDAD agent is the least performant of any method considered. However, our proposed RL-sCEE method outperforms all baselines, both in the source location and CES problems, in spite of not having access to explicit likelihoods.
\subsection{Discrete designs}
\begin{figure}[ht]
\vskip 0.1in
\begin{center}
\centerline{
\includegraphics[width=0.5\columnwidth]{figures/prey_results.png}
\includegraphics[width=0.5\columnwidth]{figures/prey_post.png}
}
\caption{LEFT: \gls{EIG} for the prey population problem, estimated using \gls{SPCE} with $L=1\textrm{E}6$. Trendlines are means and shaded regions are standard errors aggregated from $1000$ rollouts (\gls{RL}) or $500$ rollouts (SMC-ED). RIGHT: priors (orange) and posteriors (blue) after $10$ experiments. Histograms computed from $1\textrm{E}5$ samples.}
\label{fig:prey}
\end{center}
\vskip -0.2in
\end{figure}
To evaluate our method in tasks with discrete design spaces, we consider the prey population problem from~\citet{moffat2020sequential}. Designs are the initial population of a prey species, limited to the discrete interval $\cD = {1,2, \dots, 300}$. The outcome is the number of individual who were consumed by predators at the end of a $24$ hour period, based on the attack rate and handling time of the predators. Full details are available in~\cref{app:prey}.
Since DAD and iDAD cannot optimise over discrete design spaces, we added the sequential Monte Carlo design algorithm proposed by~\citet{moffat2020sequential} as a baseline. Note that this method is not amortised, and requires considerable computation time to design each experiment ($>1$ minute per design, whereas amortised policies take milliseconds). As can be seen from~\cref{fig:prey}, RL-sCEE performs similarly to the baselines, in spite of having no access to explicit likelihood information (a circumstances which would cause both baselines to fail) and using orders of magnitude less time to compute designs than the SMC-ED baseline. The numerical estimates in~\cref{table:eig} show a relative difference of \textasciitilde{}$1\%$ in the mean \gls{EIG} estimates, and the standard errors overlap.
Finally, the right hand side of~\cref{fig:prey} shows the posteriors after $10$ experiments, with their corresponding priors. Both marginals place the mode near the true value of the random variable. The handling time is fitted very tightly, whereas the attack rate has a poorer fit. Both variables have the same prior, but the posterior for handling time is so concentrated that the prior is barely visible at the bottom of the plot.
\section{Related work}\label{sec:related}
Considerable work has been done on \gls{BOED}~\cite{chaloner1995bayesian, ryan-doe-review-2016}, and particularly on using machine learning to optimise experimental designs~\cite{rainforth2023modern}. Greedy algorithms have been developed based on variational bounds~\cite{foster2019variational,foster2020unified} or neural network estimates~\cite{kleinegesse2020bayesian} of the \gls{EIG}. In the active learning literature, the BALD~\cite{houlsby2011bayesian} score is equivalent to \gls{EIG}, and can be estimated using Monte Carlo dropout neural networks~\cite{gal2017deep}. Other works attempt a non-greedy approach, i.e. they can sacrifice information gain in the current experiment in exchange for higher information gain in future experiments. Such approaches include n-step look-ahead~\cite{zhao2021uncertainty,yue2020non} or using batch designs as a lower bound for the utility of sequential designs~\cite{jiang2020binoculars}.
~\citet{foster2021deep} were the first to propose an amortised method for sequential experiment design, and showed empirically that the learned policies can exhibit non-myopic behaviour. This was extended to the case of implicit likelihood models by~\citet{ivanova2021implicit}. ~\citet{blau2022optimizing} formulated the \gls{SED} problem as a special \gls{MDP}, and showed that design policies can be learned with \gls{RL} algorithms.
The field of Reinforcement Learning has two specialised frameworks for dealing with unknown parameters that govern the system dynamics: Bayes Adaptive MDPs~\cite{duff2002optimal} and Partially Observable MDPs~\cite{kaelbling1998planning}. In both cases, agents must balance between improving their posterior belief about the MDP and achieving high rewards. The variBAD~\cite{zintgraf2021varibad} algorithm fits a variational autoencoder to observed transitions and rewards, and its latent variables are used as input for the policy. The FORBES~\cite{chen2022flow} algorithm, on the other hand, uses normalising flows to model beliefs over latent variables, transitions and rewards. MAX~\cite{shyam2019model} uses an ensemble model to generate synthetic data and compute information gain, which guides each policy action.
\section{Discussion}\label{sec:conclusion}
In this paper we introduced the sequential Cross-Entropy Estimator (sCEE), a lower bound estimate for the \gls{EIG} of an experiment design policy, as well as a reinforcement learning algorithm (RL-sCEE) that uses it to optimise policies. Experiments show the sCEE is capable of estimating large \glspl{EIG} that are intractable to estimate with contrastive estimators, which are the state of the art. In tasks where \gls{EIG} is large, RL-sCEE significantly outperforms all baselines and learns policies whose lower bound \gls{EIG} estimates exceed the upper bound estimate of the strongest alternative. In tasks where \gls{EIG} is small, RL-sCEE matches the performance of state-of-the-art baselines.
Our proposed method is highly flexible, and works in problems with discrete design spaces and with non-differentiable likelihood models. Furthermore, RL-sCEE does not require explicit likelihood information, and is thus applicable when only implicit likelihood models are available. In spite of this, it outperforms baselines that rely on explicit likelihoods.
Finally, the RL-sCEE learns not just a policy but also a neural estimator of the posterior. Our experiments show that this neural posterior can learn significant structure that we know theoretically should be present in the true Bayesian posterior, such as symmetries or constraints on the support. This structure is learned entirely from data and self-guided experiments, with no inductive bias of any kind.
One limitation of the RL-sCEE approach is its reliance on learning a neural conditional likelihood that closely matches the true posterior. In problems where this is challenging, the estimator bias will be large, and the design policies will be degraded. Indeed our experiments show that with an extremely large number of samples, it is possible for a contrastive estimator to outdo sCEE, although the cost of doing so during training, where millions of estimations are needed, is prohibitive for current hardware. Similarly, it is difficult to use the sCEE post-hoc, for a policy that was not co-trained with the estimator, as this requires learning the neural likelihood from scratch.
Future work can focus on improving the learning of this neural likelihood, which will accelerate the training of design policies as well as post-hoc evaluation.
\bibliographystyle{icml2022}
\bibliography{main.bib}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% APPENDIX
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
\appendix
\onecolumn
\setcounter{theorem}{0}
\section{Proofs}
This appendix enumerates the proofs for the theorems, corollaries and other claims made in the main paper.
\subsection{Proof of~\cref{thm:bound}}\label{app:bound_proof}
Here we prove the main theorem of the paper, which is restated for convenience
\begin{theorem}
Let $p(y \g \theta, d)$ be a probabilistic model with prior $p(\theta)$. For an arbitrary fixed design policy $\pi$ and sequence length $T$, the \gls{EIG} of using $\pi$ to design $T$ experiments is denoted $EIG(\pi, T)$.
Let $q(\theta \g h_T, \pi)$ be a proposal distribution over parameters $\theta$ conditioned on experimental history $h_T$, and the sCEE bound is
\begin{equation}
sCEE(\pi, T) \coloneqq \E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right] + H\left[ p(\theta) \right] \label{eq:scee}
\end{equation}
we have that
\begin{equation}
sCEE(\pi, T) \leq EIG(\pi, T)
\end{equation}
\end{theorem}
\begin{proof}
From Theorem 1 of~\citet{foster2021deep} we have that the \gls{EIG} is:
\begin{equation}
EIG(\pi, T) = \E_{p(h_T, \theta \g \pi)} \left[ \log p(h_T \g \theta, \pi) - \log p(h_T \g \pi) \right]
\end{equation}
This can be rewritten into a more convenient form:
%
\begin{align}
EIG(\pi, T) &= \E_{p(h_T, \theta \g \pi)} \left[ \log\frac{ p(h_T \g \theta, \pi)}{ p(h_T \g \pi)} \right] = \E_{p(h_T, \theta \g \pi)} \left[ \log\frac{ p(h_T, \theta \g \pi)}{ p(h_T \g \pi) p(\theta)} \right] \\
&= \E_{p(h_T, \theta \g \pi)} \left[ \log\frac{ p(\theta \g h_T, \pi) \cancel{p(h_T \g \pi)} }{ p(\theta) \cancel{p(h_T \g \pi)} } \right] = \E_{p(h_T, \theta \g \pi)} \left[ \log\frac{ p(\theta \g h_T, \pi) }{ p(\theta) } \right] \\
&= \E_{p(h_T, \theta \g \pi)} \left[ \log p(\theta \g h_T, \pi) - \log p(\theta) \right] \\
&= \E_{p(h_T, \theta \g \pi)} \left[ \log p(\theta \g h_T, \pi) \right] + H\left[ p(\theta) \right].\label{eq:eig}
\end{align}
%
We proceed to show that sCEE lower bounds this form. Consider the KL divergence between 2 conditional distributions given a fixed value $y$:
%
\begin{equation}
\KL{p(x \g y)}{q(x \g y)} = \E_{p(x \g y)} \left[ \log \frac{p(x \g y)}{q(x \g y)} \right]
\end{equation}
%
If $y$ is not fixed but random we then take an expectation:
%
\begin{align}
\E_{p(y)}\left[ \KL{p(x \g y)}{q(x \g y)} \right] &= \E_{p(x \g y)p(y)} \left[ \log \frac{p(x \g y)}{q(x \g y)} \right] \\
&= \E_{p(x, y)} \left[ \log \frac{p(x \g y)}{q(x \g y)} \right] \\
&= \E_{p(x, y)} \left[ \log p(x \g y) - \log q(x \g y) \right]
\end{align}
%
rearranging the sides gives
%
\begin{align}
\E_{p(x, y)} \left[ \log q(x \g y) \right] &= \E_{p(x, y)} \left[ \log p(x \g y) \right] - \E_{p(y)}\left[ \KL{p(x \g y)}{q(x \g y)} \right] \label{eq:difference}\\
&\leq \E_{p(x, y)} \left[ \log p(x \g y) \right]
\end{align}
%
where the last line exploits the fact that the KL divergence is always non-negative. Plugging in $x = \theta; y = (h_T; \pi)$ yields the lower bound:
\begin{equation}
\E_{p(\theta, h_T, \pi)} \left[ \log q(\theta \g h_T, \pi) \right] \leq \E_{p(\theta, h_T, \pi)} \left[ \log p(\theta \g h_T, \pi) \right]
\end{equation}
For a known policy $\pi$ this becomes:
%
\begin{equation}
\E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right] \leq \E_{p(\theta, h_T \g \pi)} \left[ \log p(\theta \g h_T, \pi) \right]
\end{equation}
%
Adding the prior entropy to both sides yields:
%
\begin{equation}
\E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right] + H\left[ p(\theta) \right] \leq \E_{p(\theta, h_T \g \pi)} \left[ \log p(\theta \g h_T, \pi) \right] + H\left[ p(\theta) \right].
\end{equation}
%
Finally, plugging in \cref{eq:scee,eq:eig} completes the proof:
\begin{equation}
sCEE(\pi, T) \leq EIG(\pi, T)
\end{equation}
\end{proof}
\subsection{Proof of corollaries}\label{app:corollaries}
In the main paper we state $2$ corollaries of the above theorem:
\begin{corollary}
The bound is tight if and only if $p(\theta \g h_T, \pi) = q(\theta \g h_T, \pi)$
\end{corollary}
\begin{corollary}
The bias of the sCEE estimator is $-\E_{h_T} \left[ \KL{p(\theta \g h_T, \pi)}{q(\theta \g h_T, \pi)} \right]$
\end{corollary}
If we subtract the lower bound from the \gls{EIG} we get the difference:
\\
\begin{equation}
\E_{p(\theta, h_T \g \pi)} \left[ \log p(\theta \g h_T, \pi) \right] - \E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right].
\end{equation}
\\
From~\cref{eq:difference} it follows that this difference is
\begin{equation}
-\E_{h_T} \left[ \KL{p(\theta \g h_T, \pi)}{q(\theta \g h_T, \pi)} \right]
\end{equation}
Since the KL divergence is always non-negative, this difference is $0$ and the bound is tight if and only if $\KL{p(\theta \g h_T, \pi)}{q(\theta \g h_T, \pi)} = 0$ for all realisations of $h_T$. This establishes both corollaries.
\subsection{Proof of convergence}\label{app:convergence_proof}
In the main paper we make the claim that a Monte Carlo estimator of the sCEE converges at a rate of $O(\frac{1}{\sqrt{n}})$, where $n$ is the number of MC samples. Since the prior is known, we can rely on standard MC convergence proofs for the prior entropy component. Thus we need only worry about a convergence proof for estimating the cross-entropy component $\E_{p(\theta, h_T \g \pi)} \left[ \log q(\theta \g h_T, \pi) \right]$.
We denote the cross-entropy as $H\left[ p(\theta \g h_T, \pi), q(\theta \g h_T, \pi) \right]$ and the MC estimator as
\begin{equation}
\hat{H}\left[ p(\theta \g h_T, \pi), q(\theta \g h_T, \pi) \right] = \frac{1}{n} \sum_{i=1}^n -\log q(\theta^i \g h_T^i, \pi)
\end{equation}
According to Theorem 5.1 of~\citet{mcallester2020formal}, if there is a minimum log-likelihood $F_{max}$ such that $\log q(\theta \g h_T, \pi) \geq F_{max}$, then with probability at least $1-\delta$ we have that
\begin{equation}
|H\left[ p(\theta \g h_T, \pi), q(\theta \g h_T, \pi) \right] - \hat{H}\left[ p(\theta \g h_T, \pi), q(\theta \g h_T, \pi) \right]| \leq F_{max}\sqrt{\frac{\log(\frac{2}{\delta})}{2n}}
\end{equation}
Thus the MC estimator converges to the true sCEE with high probability at the desired rate of $O(\frac{1}{\sqrt{n}})$.
\section{Relationship between sCEE and sACE}\label{app:sace}
\citet{foster2021deep} propose in the appendices a lower bound EIG estimator that relies on a parameterised proposal distribution that approximates the posterior $p(\theta \g h_T)$. They called this the \textit{sequential Adaptive Constrastive Estimation} (sACE):
\begin{equation}
\E_{p(\theta_0, h_T \g \pi)q(\theta_{1:L}; h_T)} \left[ \log \frac{p(h_T \g \theta_0, \pi)}{\frac{1}{L+1} \sum_{l=0}^L \frac{p(h_T \g \theta_l, \pi) p(\theta_l)}{q(\theta_l;h_T)}} \right]
\end{equation}
This is a contrastive bound where the contrastive samples are distributed according to the proposal distribution $\theta_{1:l} \sim q(\theta \g h_T)$. The construction and proof assume a minimum of $1$ contrastive sample. However, if we set $L=0$ in this expression, the sampling of contrastive samples from $q(\theta \g h_T)$ disappears and we get:
\begin{align}
\E_{p(\theta_0, h_T \g \pi)} \left[ \log \frac{p(h_T \g \theta_0, \pi)}{\frac{1}{0+1} \sum_{l=0}^0 \frac{p(h_T \g \theta_l, \pi) p(\theta_l)}{q(\theta_l;h_T)}} \right] &=
\E_{p(\theta_0, h_T \g \pi)} \left[ \log \frac{p(h_T \g \theta_0, \pi)}{\frac{1}{0+1} \frac{p(h_T \g \theta_0, \pi) p(\theta_0)}{q(\theta_0;h_T)}} \right] \\
&= \E_{p(\theta_0, h_T \g \pi)} \left[ \log \frac{\cancel{p(h_T \g \theta_0, \pi)}}{\cancel{p(h_T \g \theta_0, \pi)}} \frac{q(\theta_0;h_T)}{p(\theta_0)} \right] \\
&= \E_{p(\theta_0, h_T \g \pi)} \left[ \log \frac{q(\theta_0;h_T)}{p(\theta_0)} \right],
\end{align}
which is equivalent to the sCEE. Note that by avoiding the need for contrastive samples, the sCEE gains a considerable computational advantage. In the RL setting, the rewards depend on $q(\theta \g h_T)$ and hence need to be recomputed every time $q$ is updated. With the sACE estimator, this recomputation requires resampling the contrastive samples, increasing the computational effort by a factor of $O(L)$. Indeed, depending on memory constraints, it may not be possible to recompute an entire batch of rewards in a single vectorised operation. With the sCEE, however, reward recomputation requires only a single neural network pass.
In addition to the computational benefits, sCEE has the further advantage that it is compatible with implicit likelihood models, wherease sACE requires explicit models, since it includes the term $p(h_T \g \theta_0, \pi)$ in the numerator.
\section{Normalising flows on the probability simplex}
If we have a random variable with support on the canonical (open) simplex $\Delta_{k-1}$ rather than in $\cR^k$, additional caution is required for fitting a normalising flow to this RV. Since the $k^{th}$ dimension of the RV is fully determined by the first $k-1$ dimensions, the NF is free to fit this dimension with extremely high confidence, leading to an overestimation of log-likelihood of the entire RV.
The fix to this issue is rather involved. First, we exclude the $k^{th}$ dimension as input to the NF. Then, at the penultimate layer of a normalising flow, it implements the diffeomorphism $\cF : \cR^{k-1} \rightarrow \cR^{k-1}$ i.e. the base distribution is a standard Gaussian and the resulting distribution can have support in the entire real space. Now we add a series of bijections that will produce a map $\cG : \cR^{k-1} \rightarrow \Delta_{k-1}$. Note that it is not enough simply to concatenate $1 - \sum_{i=1}^{k-1} \cF(x)_i$ with the intermediate vector $\cF(x)$ because we are not guaranteed that $0 \leq \cF(x)_k \leq 1 \ \forall k$ and that $\sum_{i=1}^{k-1} \cF(x)_i \leq 1$. First we must transform the output to ensure these properties:
\begin{align}
&u = \cF(x)\\
&v_i = \sigma(u_i) \label{eq:open-box}\\
&w_i = v_i\left(1 - \sum_{j=1}^{i-1}w_j \right) \ \forall i\in[1, k-1] \label{eq:low-simplex}\\
&\theta_i = \frac{w_i}{1 - \epsilon}. \label{eq:epsilon}
\end{align}
\Cref{eq:open-box} projects $\cR^{k-1}$ to the semi-open box $[0,1)^{k-1}$. \Cref{eq:low-simplex} projects this box to the $k-1$ dimensional simplex $\mathbf{s} = \{x : \sum_{i=1}^{k-1} x_i < 1 \ \text{and} \ 0 \leq x_k < 1 \ \forall i \in [1,k-1]\}$. This non-canonical simplex is in fact the equivalent of projecting the $k$-dimensional canonical simplex $\Delta_{k-1}$ down to $k-1$ dimensions. The simplex $\mathbf{s}$ can be lifted to $\Delta_{k-1}$ by assigning $w_k = 1 - \sum_{j=1}^{k-1}w_j$. However, we won't include this in the mapping $\cG$ because it makes the Jacobian low-rank and hence the inverse ill-defined. To avoid floating-point errors, each element of the RV actually has to be in the range $[\epsilon, 1-\epsilon]$ where $\epsilon$ is the machine epsilon.~\Cref{eq:epsilon} maps between this space and the actual canonical simplex.
The corresponding log-det-Jacobians are:
\begin{align}
&\sum_{i=1}^{k-1} \log (0.99 \cdot v_i(1-v_i)), \\
&\sum_{i=1}^{k-1} \log (1 - \sum_{j=1}^{i-1}w_j), \\
&(1 - k) \cdot \log(1 - \epsilon).
\end{align}
The inverse $\cG^{-1}$ can be written compactly as:
\begin{equation}
u_i = \sigma^{-1} \left( \frac{(1 - \epsilon) \cdot \theta_i}{1 - \sum_{j=1}^{i-1} (1 - \epsilon) \cdot \theta_j} \right) \quad \forall i \in [1,k-1].
\end{equation}
\section{Experiment details}
This appendix describes the probabilistic models, hyperparameters, and all other details relating to the experiment design problems appearing in the paper.
\subsection{EIG for conjugate priors}\label{app:conjugate}
For an isotropic Gaussian prior $\cN(\mu_0, \sigma_0\mathbf{I_k})$ and Gaussian likelihood with known isotropic covariance $\sigma\mathbf{I_k}$, the posterior after $n$ observations is an isotropic Gaussian with covariance:
\begin{align}
\Sigma_{post} &= (\sigma_0^{-1}\mathbf{I}_k + n\sigma^{-1}\mathbf{I}_k)^{-1} \\
&= (\sigma_0^{-1} + n\sigma^{-1})^{-1} \mathbf{I}_k
\end{align}
The mean of the posterior is unimportant to us as it does not affect the entropy:
\begin{align}
H_{post} &= \frac{k}{2} + \frac{k}{2}\log(2\pi) + \frac{1}{2}\log (|\Sigma_{post}|) \\
&= \frac{k}{2} + \frac{k}{2}\log(2\pi) - \frac{k}{2}\log(\sigma_0^{-1} + n\sigma^{-1}).
\end{align}
Therefore the entropy is independent of the designs and we can compute the entropy of the "optimal" policy by subtracting the posterior entropy from the prior entropy:
\begin{align}
I_n(\pi) &= H[\cN(\mu_0, \sigma_0\mathbf{I_k})] - H_{post} \\
&= \cancel{\frac{k}{2}} + \cancel{\frac{k}{2}\log(2\pi)} + \frac{k}{2}\log (\sigma_0) - \cancel{\frac{k}{2}} - \cancel{\frac{k}{2}\log(2\pi)} + \frac{k}{2}\log(\sigma_0^{-1} + n\sigma^{-1}) \\
&= \frac{k}{2} (\log (\sigma_0) + \log(\sigma_0^{-1} + n\sigma^{-1})) \\
&= \frac{k}{2} \log (1 + n \frac{\sigma_0}{\sigma})
\end{align}
Thus we can create an \gls{EIG} estimation problem with an \gls{EIG} of our choice by setting $k, n, \sigma_0$ and $\sigma$ appropriately. In our experiments sCEE was trained for $10,000$ epochs, and was exposed to $10,000$ data points in each epoch. Each estimator was evaluated using $1,000$ Monte Carlo samples.
\subsection{Constant elasticity of substitution}\label{app:ces}
In this experiment economic agents compare $2$ baskets of goods and give a rating on a sliding scale from $0$ to $1$. Each basket consists of $k$ different goods with different value. The designs are vectors $d = (x, x^\prime)$ where $x, x^\prime \in [0, 100]^k$ are the baskets of goods. The latent parameters of the likelihood and their priors are:
\begin{align}
\rho &\sim Beta(1,1)\\
\alpha &\sim Dirichlet(\mathbf{1}_k)\\
\log u &\sim \cN(1,3).
\end{align}
The probabilistic model is:
\begin{align}
U(x) &= \left(\sum_i x^\rho_i\alpha_i\right)^{1/\rho}\\
\mu_\eta &= (U(x) - U(x^\prime)) u\\
\sigma_\eta &= (1 + ||x - x^\prime||) \tau\cdot u\\
\eta &\sim \cN(\mu_\eta, \sigma^2_\eta)\\
y &= clip(sigmoid(\eta), \epsilon, 1-\epsilon),
\end{align}
In our experiments we used the following hyperparameters:
\begin{table}[h]
\vskip 0.15in
\begin{center}
\begin{small}
\begin{sc}
\begin{tabular}{ll}
\toprule
Parameter & Value \\
\midrule
$k$ & $3$ \\
$\tau$ & $0.005$ \\
$\epsilon$ & $2^{-22}$ \\
\bottomrule
\end{tabular}
\end{sc}
\end{small}
\end{center}
\vskip -0.1in
\end{table}
\subsection{Prey population}\label{app:prey}
In this experiment an initial population of prey animals is left to survive for $\cT$ hours, and we measure the number of individuals consumed by predators at the end of the experiment. The designs are the initial populations $d = N_0 \in {1,2, \dots, 300}$. The latent parameters and priors are:
\begin{align}
\log a &\sim \cN(-1.4, 1.35)\\
\log T_h &\sim \cN(-1.4, 1.35),
\end{align}
where $a$ represents the attack rate and $T_h$ is the handling time.
The population changes over time according to a Holling's Type III model, which is a differential equation:
\begin{equation}
\frac{dN}{d\tau} = - \frac{aN^2}{1 + aT_hN^2}.
\end{equation}
And the population $N_\cT$ is thus the solution of an initial value problem. The probabilistic model is:
\begin{align}
p_\cT &= \frac{d - N_\cT}{d}\\
y &\sim Binom(d, p_\cT).
\end{align}
We used a simulation time of $\cT = 24$ hours.
\subsection{Source location}\label{app:source}
In this experiment there are $n$ sources embedded in $k$-dimensional space that emit independent signals. the designs are the co-ordinates at which to measure signal intensity, and we restrict the space to $d \in [-4, 4]^k$. The total intensity at any given co-ordinate $d$ in the plane is given the sum of individual signals:
\begin{equation}
\mu(\theta, d) = b + \sum_i\frac{1}{m + ||\theta_i - d||^2},
\end{equation}
where $b, m > 0$ are the background and maximum signals, respectively, $||\cdot||^2$ is the squared Euclidean norm, and $\theta_i$ are the co-ordinates of the $i^{th}$ signal source. The probabilistic model is:
\begin{equation}
\theta_i \sim \cN(0, I); \quad \log y \g \theta, d \sim \cN(\log(\mu(\theta, d), \sigma),
\end{equation}
i.e. the prior is unit Gaussian and we observe the log of the total signal intensity with some Gaussian observation noise $\sigma$. The hyperparameters we used are
\begin{table}[h]
\vskip 0.15in
\begin{center}
\begin{small}
\begin{sc}
\begin{tabular}{lr}
\toprule
Parameter & Value \\
\midrule
$n$ & $2$ \\
$k$ & $2$ \\
$b$ & $1\textrm{E}-1$ \\
$m$ & $1\textrm{E}-4$ \\
$\sigma$ & $0.5$ \\
\bottomrule
\end{tabular}
\end{sc}
\end{small}
\end{center}
\vskip -0.1in
\end{table}
\section{Algorithm details}\label{app:algos}
This appendix provides the implementation details for all design of experiment algorithms used in the paper.
\subsection{RL-sCEE}
We used the implementation of REDQ from~\citet{blau2022optimizing} as the basis of our algorithm, although we limited the ensemble size to $N=2$. Normalising flows were implemented using the normflows~\cite{normflows} package, which we extended to create a conditioned version of realNVP~\cite{dinhdensity}. In every experiment we used a normalising flow with $6$ layers, and the parameter map is a $2$-layer neural network with sizes $(128,128)$. Both normalising flows and policies use a permutation invariant representation similar to~\citet{ivanova2021implicit}, including a single self-attention layer with $8$ attention heads.
Additional hyperparameters are listed in the table below, and are largerly derived from~\citet{blau2022optimizing}:
\begin{table}[h]
\vskip 0.15in
\begin{center}
\begin{small}
\begin{sc}
\begin{tabular}{llll}
\toprule
Parameter & Source Location & CES & Prey Population \\
\midrule
training iterations & $1\textrm{E}5$ & $1\textrm{E}5$ & $2\textrm{E}4$ \\
$T$ & $30$ & $10$ & $10$ \\
$\gamma$ & $0.9$ & $0.9$ & $0.95$ \\
$\tau$ & $1\textrm{E}-3$ & $5\textrm{E}-3$ & $1\textrm{E}-2$ \\
policy learning rate & $1\textrm{E}-4$ & $3\textrm{E}-4$ & $1\textrm{E}-4$ \\
critic learning rate & $3\textrm{E}-4$ & $3\textrm{E}-4$ & $1\textrm{E}-3$ \\
buffer size & $1\textrm{E}7$ & $1\textrm{E}7$ & $1\textrm{E}6$ \\
\bottomrule
\end{tabular}
\end{sc}
\end{small}
\end{center}
\vskip -0.1in
\end{table}
\subsection{RL-sPCE}
We used the implementation of~\citet{blau2022optimizing}, which is available at \url{https://github.com/csiro-mlai/rl-boed}. We kept all hyperparameters and network architectures the same, with the exception of adding a self-attention layer to the policy network. This layer is identical to the one described in the previous section. We did not find that adding attention lead to significant change in performance, but included it in order to maintain a fair comparison with the RL-sCEE implementation.
In particular, we used $L=1\textrm{E}5$ contrastive samples for training. Not only is it the value used by~\citet{blau2022optimizing}, but is also pushing the limits of the number of samples that can be used in a reasonable amount of time. Since tens of millions of simulated experiments have to be run to train a single agent, we must leverage vectorisation over multiple sequences of experiments in parallel. Although in the evaluation we used $L=1\textrm{E}8$ samples, this only allows a single experiment at a time to fit in a GPU, and requires multiple seconds per experiment. It would require several years to train a single agent in this manner.
\subsection{DAD and iDAD}
For these baselines we used the implementations of the original papers, which are available at \url{https://github.com/ae-foster/dad} and \url{https://github.com/desi-ivanova/idad}, respectively. We kept the default hyperparameters of those implementations. The only exception is for iDAD on the source location problem, which we found unstable for a sequence of $T=30$ experiments. We therefore used early stopping, and stopped learning at $40k$ iterations instead of the original $100k$.
\subsection{SMC-ED}
We used the implementation made available in \url{https://github.com/csiro-mlai/rl-boed}, which in turn uses the R language implementation of~\cite{moffat2020sequential} and executes it from within a Python script by using the $rpy2$ bindings. The original R code is available at \url{https://github.com/haydenmoffat/sequential_design_for_predator_prey_experiments}.
\section{Hardware details}
SMC-ED experiments were run on a desktop machine with an Intel i7-10610U CPU and no GPU. All other experiments were run in a high-performance computing cluster, using a single node each with $4$ cores of an Intel Xeon E5-2690 CPU and an Nvidia Tesla P100 GPU.
\section{Ablation Study}
To evaluated the stabilisation mechanisms incorporated in the implementation of RL-sCEE, we conduct an ablation study where we remove either the target posterior network, the fixed initial posterior, or both. The results can be seen in~\cref{fig:ablation}, with each variant replicated $10$ times, using common random seeds between different variants (e.g. the blue trendline labeled "0" represents the same random seed in all $4$ plots).
It is clear that the removal of the target network causes significant degradation in performance, with many replications converging to a lower final performance or even peaking early and then decreasing in EIG. On the other hand, the use of a fixed initial posterior doesn't seem to offer a clear advantage over a learned one.
\begin{figure}[ht]
\vskip 0.1in
\begin{center}
\centerline{
\includegraphics[width=0.5\columnwidth]{figures/no_target.png}
\includegraphics[width=0.5\columnwidth]{figures/learned_no_target.png}
}
\centerline{
\includegraphics[width=0.5\columnwidth]{figures/target.png}
\includegraphics[width=0.5\columnwidth]{figures/learned_target.png}
}
\caption{Ablation studies for the stabilisation mechanisms.}
\label{fig:ablation}
\end{center}
\vskip -0.2in
\end{figure}
\section{Additional results}\label{app:results}
\begin{figure}[ht]
\vskip 0.1in
\begin{center}
\centerline{
\includegraphics[width=\columnwidth]{figures/hist.png}
}
\caption{Posterior for the source location problem. Computed from $1\textrm{E}5$ samples. Black rings denote the true co-ordinates of signal sources.}
\label{fig:extra_post}
\end{center}
\vskip -0.2in
\end{figure}
\begin{thebibliography}{39}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Agrawal et~al.(2019)Agrawal, Squires, Yang, Shanmugam, and
Uhler]{agrawal2019abcd}
Agrawal, R., Squires, C., Yang, K., Shanmugam, K., and Uhler, C.
\newblock Abcd-strategy: Budgeted experimental design for targeted causal
structure discovery.
\newblock In \emph{International Conference on Artificial Intelligence and
Statistics}, 2019.
\bibitem[Baltas(2001)]{baltas2001utility}
Baltas, G.
\newblock Utility-consistent brand demand systems with endogenous category
consumption: principles and marketing applications.
\newblock \emph{Decision Sciences}, 32\penalty0 (3):\penalty0 399--422, 2001.
\bibitem[Bingham et~al.(2018)Bingham, Chen, Jankowiak, Obermeyer, Pradhan,
Karaletsos, Singh, Szerlip, Horsfall, and Goodman]{bingham2018pyro}
Bingham, E., Chen, J.~P., Jankowiak, M., Obermeyer, F., Pradhan, N.,
Karaletsos, T., Singh, R., Szerlip, P., Horsfall, P., and Goodman, N.~D.
\newblock {Pyro: Deep Universal Probabilistic Programming}.
\newblock \emph{Journal of Machine Learning Research}, 2018.
\bibitem[Blau et~al.(2022)Blau, Bonilla, Chades, and
Dezfouli]{blau2022optimizing}
Blau, T., Bonilla, E.~V., Chades, I., and Dezfouli, A.
\newblock Optimizing sequential experimental design with deep reinforcement
learning.
\newblock In \emph{International Conference on Machine Learning}, 2022.
\bibitem[Bloss et~al.(1988)Bloss, Hudak, and Young]{bloss1988code}
Bloss, A., Hudak, P., and Young, J.
\newblock Code optimizations for lazy evaluation.
\newblock \emph{Lisp and Symbolic Computation}, 1\penalty0 (2):\penalty0
147--164, 1988.
\bibitem[Chaloner \& Verdinelli(1995)Chaloner and
Verdinelli]{chaloner1995bayesian}
Chaloner, K. and Verdinelli, I.
\newblock Bayesian experimental design: A review.
\newblock \emph{Statistical Science}, pp.\ 273--304, 1995.
\bibitem[Chen et~al.(2021)Chen, Wang, Zhou, and Ross]{chen2020randomized}
Chen, X., Wang, C., Zhou, Z., and Ross, K.~W.
\newblock Randomized ensembled double q-learning: Learning fast without a
model.
\newblock In \emph{International Conference on Learning Representations}, 2021.
\bibitem[Chen et~al.(2022)Chen, Mu, Luo, Li, and Chen]{chen2022flow}
Chen, X., Mu, Y.~M., Luo, P., Li, S., and Chen, J.
\newblock Flow-based recurrent belief state learning for pomdps.
\newblock In \emph{International Conference on Machine Learning}, pp.\
3444--3468. PMLR, 2022.
\bibitem[Dinh et~al.(2017)Dinh, Sohl-Dickstein, and Bengio]{dinhdensity}
Dinh, L., Sohl-Dickstein, J., and Bengio, S.
\newblock Density estimation using real nvp.
\newblock In \emph{International Conference on Learning Representations}, 2017.
\bibitem[Drovandi et~al.(2014)Drovandi, McGree, and
Pettitt]{drovandi2014sequential}
Drovandi, C.~C., McGree, J.~M., and Pettitt, A.~N.
\newblock A sequential monte carlo algorithm to incorporate model uncertainty
in bayesian sequential design.
\newblock \emph{Journal of Computational and Graphical Statistics}, 23\penalty0
(1):\penalty0 3--24, 2014.
\bibitem[Duff(2002)]{duff2002optimal}
Duff, M.~O.
\newblock \emph{Optimal Learning: Computational procedures for Bayes-adaptive
Markov decision processes}.
\newblock University of Massachusetts Amherst, 2002.
\bibitem[Foster et~al.(2019)Foster, Jankowiak, Bingham, Horsfall, Teh,
Rainforth, and Goodman]{foster2019variational}
Foster, A., Jankowiak, M., Bingham, E., Horsfall, P., Teh, Y.~W., Rainforth,
T., and Goodman, N.
\newblock Variational bayesian optimal experimental design.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2019.
\bibitem[Foster et~al.(2020)Foster, Jankowiak, O’Meara, Teh, and
Rainforth]{foster2020unified}
Foster, A., Jankowiak, M., O’Meara, M., Teh, Y.~W., and Rainforth, T.
\newblock A unified stochastic gradient approach to designing bayesian-optimal
experiments.
\newblock In \emph{International Conference on Artificial Intelligence and
Statistics}, 2020.
\bibitem[Foster et~al.(2021)Foster, Ivanova, Malik, and
Rainforth]{foster2021deep}
Foster, A., Ivanova, D.~R., Malik, I., and Rainforth, T.
\newblock Deep adaptive design: Amortizing sequential bayesian experimental
design.
\newblock \emph{International Conference on Machine Learning}, 2021.
\bibitem[Gal et~al.(2017)Gal, Islam, and Ghahramani]{gal2017deep}
Gal, Y., Islam, R., and Ghahramani, Z.
\newblock Deep bayesian active learning with image data.
\newblock In \emph{International Conference on Machine Learning}, 2017.
\bibitem[{Garage Contributors}(2019)]{garage}
{Garage Contributors}.
\newblock Garage: A toolkit for reproducible reinforcement learning research.
\newblock \url{https://github.com/rlworkgroup/garage}, 2019.
\bibitem[Houlsby et~al.(2011)Houlsby, Husz{\'a}r, Ghahramani, and
Lengyel]{houlsby2011bayesian}
Houlsby, N., Husz{\'a}r, F., Ghahramani, Z., and Lengyel, M.
\newblock Bayesian active learning for classification and preference learning.
\newblock \emph{arXiv preprint arXiv:1112.5745}, 2011.
\bibitem[Ivanova et~al.(2021)Ivanova, Foster, Kleinegesse, Gutmann, and
Rainforth]{ivanova2021implicit}
Ivanova, D., Foster, A., Kleinegesse, S., Gutmann, M.~U., and Rainforth, T.
\newblock Implicit deep adaptive design: Policy-based experimental design
without likelihoods.
\newblock In \emph{Advances in Neural Information Processing Systems}, 2021.
\bibitem[Jiang et~al.(2020)Jiang, Chai, Gonzalez, and
Garnett]{jiang2020binoculars}
Jiang, S., Chai, H., Gonzalez, J., and Garnett, R.
\newblock Binoculars for efficient, nonmyopic sequential experimental design.
\newblock In \emph{International Conference on Machine Learning}, 2020.
\bibitem[Kaelbling et~al.(1998)Kaelbling, Littman, and
Cassandra]{kaelbling1998planning}
Kaelbling, L.~P., Littman, M.~L., and Cassandra, A.~R.
\newblock Planning and acting in partially observable stochastic domains.
\newblock \emph{Artificial intelligence}, 101\penalty0 (1-2):\penalty0 99--134,
1998.
\bibitem[Kleinegesse \& Gutmann(2020)Kleinegesse and
Gutmann]{kleinegesse2020bayesian}
Kleinegesse, S. and Gutmann, M.~U.
\newblock Bayesian experimental design for implicit models by mutual
information neural estimation.
\newblock In \emph{International Conference on Machine Learning}, 2020.
\bibitem[Krause \& Guestrin(2007)Krause and Guestrin]{krause2007nonmyopic}
Krause, A. and Guestrin, C.
\newblock Nonmyopic active learning of gaussian processes: an
exploration-exploitation approach.
\newblock In \emph{International Conference on Machine learning}, 2007.
\bibitem[Lillicrap et~al.(2016)Lillicrap, Hunt, Pritzel, Heess, Erez, Tassa,
Silver, and Wierstra]{lillicrap2016continuous}
Lillicrap, T.~P., Hunt, J.~J., Pritzel, A., Heess, N., Erez, T., Tassa, Y.,
Silver, D., and Wierstra, D.
\newblock Continuous control with deep reinforcement learning.
\newblock \emph{International Conference on Learning Representations}, 2016.
\bibitem[Lindley(1956)]{lindley1956measure}
Lindley, D.~V.
\newblock On a measure of the information provided by an experiment.
\newblock \emph{The Annals of Mathematical Statistics}, pp.\ 986--1005, 1956.
\bibitem[McAllester \& Stratos(2020)McAllester and
Stratos]{mcallester2020formal}
McAllester, D. and Stratos, K.
\newblock Formal limitations on the measurement of mutual information.
\newblock In \emph{International Conference on Artificial Intelligence and
Statistics}, 2020.
\bibitem[Moffat et~al.(2020)Moffat, Hainy, Papanikolaou, and
Drovandi]{moffat2020sequential}
Moffat, H., Hainy, M., Papanikolaou, N.~E., and Drovandi, C.
\newblock Sequential experimental design for predator--prey functional response
experiments.
\newblock \emph{Journal of the Royal Society Interface}, 17\penalty0 (166),
2020.
\bibitem[Poole et~al.(2019)Poole, Ozair, Van Den~Oord, Alemi, and
Tucker]{poole2019variational}
Poole, B., Ozair, S., Van Den~Oord, A., Alemi, A., and Tucker, G.
\newblock On variational bounds of mutual information.
\newblock In \emph{International Conference on Machine Learning}, 2019.
\bibitem[Rainforth et~al.(2018)Rainforth, Cornish, Yang, Warrington, and
Wood]{rainforth2018nesting}
Rainforth, T., Cornish, R., Yang, H., Warrington, A., and Wood, F.
\newblock On nesting monte carlo estimators.
\newblock In \emph{International Conference on Machine Learning}, 2018.
\bibitem[Rainforth et~al.(2023)Rainforth, Foster, Ivanova, and
Smith]{rainforth2023modern}
Rainforth, T., Foster, A., Ivanova, D.~R., and Smith, F.~B.
\newblock Modern bayesian experimental design.
\newblock \emph{arXiv preprint arXiv:2302.14545}, 2023.
\bibitem[Ryan et~al.(2016)Ryan, Drovandi, McGree, and
Pettitt]{ryan-doe-review-2016}
Ryan, E.~G., Drovandi, C.~C., McGree, J.~M., and Pettitt, A.~N.
\newblock A review of modern computational algorithms for {B}ayesian optimal
design.
\newblock \emph{International Statistical Review}, 84\penalty0 (1):\penalty0
128--154, 2016.
\newblock \doi{10.1111/insr.12107}.
\newblock URL \url{https://onlinelibrary.wiley.com/doi/abs/10.1111/insr.12107}.
\bibitem[Shababo et~al.(2013)Shababo, Paige, Pakman, and
Paninski]{shababo2013bayesian}
Shababo, B., Paige, B., Pakman, A., and Paninski, L.
\newblock Bayesian inference and online experimental design for mapping neural
microcircuits.
\newblock \emph{Advances in Neural Information Processing Systems}, 2013.
\bibitem[Shyam et~al.(2019)Shyam, Ja{\'s}kowski, and Gomez]{shyam2019model}
Shyam, P., Ja{\'s}kowski, W., and Gomez, F.
\newblock Model-based active exploration.
\newblock In \emph{International Conference on Machine Learning}, 2019.
\bibitem[Stimper et~al.(2023)Stimper, Liu, Campbell, Berenz, Ryll,
Sch{\"o}lkopf, and Hern{\'a}ndez-Lobato]{normflows}
Stimper, V., Liu, D., Campbell, A., Berenz, V., Ryll, L., Sch{\"o}lkopf, B.,
and Hern{\'a}ndez-Lobato, J.~M.
\newblock {N}ormflows: {A} {P}y{T}orch {P}ackage for {N}ormalizing {F}lows.
\newblock \emph{arXiv preprint arXiv:2302.12014}, 2023.
\bibitem[Treloar et~al.(2022)Treloar, Braniff, Ingalls, and
Barnes]{treloar2022deep}
Treloar, N.~J., Braniff, N., Ingalls, B., and Barnes, C.~P.
\newblock Deep reinforcement learning for optimal experimental design in
biology.
\newblock \emph{PLOS Computational Biology}, 18\penalty0 (11), 2022.
\bibitem[Van~Hasselt et~al.(2018)Van~Hasselt, Doron, Strub, Hessel, Sonnerat,
and Modayil]{van2018deep}
Van~Hasselt, H., Doron, Y., Strub, F., Hessel, M., Sonnerat, N., and Modayil,
J.
\newblock Deep reinforcement learning and the deadly triad.
\newblock In \emph{NeurIPS Deep Reinforcement Learning Workshop}, 2018.
\bibitem[Winkler et~al.(2019)Winkler, Worrall, Hoogeboom, and
Welling]{winkler2019learning}
Winkler, C., Worrall, D., Hoogeboom, E., and Welling, M.
\newblock Learning likelihoods with conditional normalizing flows.
\newblock \emph{arXiv preprint arXiv:1912.00042}, 2019.
\bibitem[Yue \& Kontar(2020)Yue and Kontar]{yue2020non}
Yue, X. and Kontar, R.~A.
\newblock Why non-myopic bayesian optimization is promising and how far should
we look-ahead? a study via rollout.
\newblock In \emph{International Conference on Artificial Intelligence and
Statistics}, 2020.
\bibitem[Zhao et~al.(2021)Zhao, Dougherty, Yoon, Alexander, and
Qian]{zhao2021uncertainty}
Zhao, G., Dougherty, E., Yoon, B.-J., Alexander, F., and Qian, X.
\newblock Uncertainty-aware active learning for optimal bayesian classifier.
\newblock In \emph{International Conference on Learning Representations}, 2021.
\bibitem[Zintgraf et~al.(2021)Zintgraf, Schulze, Lu, Feng, Igl, Shiarlis, Gal,
Hofmann, and Whiteson]{zintgraf2021varibad}
Zintgraf, L., Schulze, S., Lu, C., Feng, L., Igl, M., Shiarlis, K., Gal, Y.,
Hofmann, K., and Whiteson, S.
\newblock Varibad: variational bayes-adaptive deep rl via meta-learning.
\newblock \emph{The Journal of Machine Learning Research}, 22\penalty0
(1):\penalty0 13198--13236, 2021.
\end{thebibliography}