arXiv:2606.23662v1 [stat.ML] 22 Jun 2026
ACTION -BED: Task-Driven Bayesian Experimental Design with Singly Intractable Objectives
Tom Rossa Department of Statistics University of Oxford [email protected]
Angus Phillips Department of Statistics University of Oxford [email protected]
Tom Rainforth Department of Statistics University of Oxford [email protected]
Abstract Bayesian experimental design (BED) has traditionally been based on maximising expected uncertainty reductions from prior to posterior. A major shortfall of this approach is that it leads to doubly intractable objectives that are difficult to optimise, while customising them to particular downstream tasks of interest can also be difficult. Following first principles decision theory, we demonstrate that BED can alternatively be formulated in terms of an expected future loss (EFL) on downstream actions, providing a simple and naturally task-driven framework. Critically, we then show that all such EFLs can be rearranged into singly intractable objectives that can be jointly optimised with respect to both the design policy and a downstream action policy using stochastic gradients, an approach we refer to as ACTION -BED. This formulation further sidesteps the need for any explicit posterior or marginal likelihood estimation and is naturally implicit, requiring only the ability to sample from the joint model over model parameters and data, and evaluate the downstream loss function. It thus allows design policies to be learned more effectively, efficiently, and simply than existing methods, while providing easy customisation to different downstream tasks and losses.
1
Introduction
Bayesian Experimental Design (BED) [12, 17, 30, 48, 49, 58] is a principled framework for the optimal acquisition of data, in which designs are chosen to maximise the expected uncertainty reduction (EUR) from prior to posterior over unknown values of interest. Typically the uncertainty measure is taken to be the Shannon entropy [61], leading to the classical expected information gain (EIG) objective [48]. In this perspective, data is valued according to how it changes our belief distribution over unknown values, which is theoretically justified by the fact that downstream Bayes-optimal actions only depend on observed data through the updated belief distribution [59]. However, there are two major issues with using this framework in practice. First, defining uncertainties that are customised to particular downstream tasks of interest can be challenging [40], noting that the data gathered is rarely the end goal in itself and is instead used to aid downstream estimation, prediction, or decision-making. In particular, common metrics like Shannon entropy constitute generic uncertainties which fail to reflect that not all information is equally useful for our task of interest.
Preprint.
ED g B es tin oach s i Ex ppr A
Design Policy
Posterior Uncertainty U(p(θ | D))
Action Policy πaψ (D)
rollout
πdϕ Data
Future Loss ℓtask (a, θ) Action
ACTION -BED: Joint policy updates ∇ϕ,ψ ℓtask
Figure 1: ACTION -BED; our task-driven, singly intractable algorithm for policy-based BED. Rather than targeting uncertainty reductions, ACTION -BED jointly optimises a downstream action policy alongside the design policy to minimise the expected future loss using stochastic gradient descent. Second, EURs are generally doubly intractable [58]: without an analytic posterior, the EUR is inevitably a nested expectation [57] as meaningful uncertainty measures are nonlinear [16, 25, 50, 60]. This makes estimating and optimising EURs very challenging as standard estimators cannot be applied and, while a wide array of computational strategies have been proposed to deal with this [2, 21, 22, 26, 36–38], they are typically only applicable to EIGs and tend to suffer from high computational costs and bias. Others have instead looked to learn amortised posterior approximations to circumvent the double intractability [4, 23], often learning this alongside the design/policy optimisation using variational formulations [10, 24, 33, 34, 62]. However, this amortised posterior learning is still itself a very challenging task and inaccuracies in the posterior approximations can undermine the resulting data acquisition. We propose to address these challenges with a fundamental move away from the EUR framework by instead defining the utility of data through the performance of the downstream actions it enables. Specifically, using first principles Bayesian decision theory, we show that the Bayes-optimal data gathering policy can be formulated as the minimiser of the expected future loss (EFL) of downstream Bayes-optimal actions, allowing for the definition of experimental design objectives that are naturally customised to our downstream tasks and losses of interest. While the set of unique objectives indexed by the EFL and EUR frameworks turns out to be equivalent, our critical insight is that the former allows for a simple rearrangement into a singly intractable objective that we jointly optimise with respect to a design policy and a downstream action policy. This results in a much simpler and computationally tractable policy-based framework than existing approaches that we call ACTION -BED. As shown in Figure 1, ACTION -BED avoids the need to directly perform posterior approximations or explicitly estimate any uncertainties: it simply requires the minimisation of an expected loss with respect to the joint policy parameters, which can be done using standard stochastic gradient descent methods. Furthermore, the task-driven EFL objectives it uses are naturally implicit, requiring only reparametrisable simulations from the joint model over parameters and data, and a differentiable loss function. These benefits can substantially simplify the policy training and we empirically find that ACTION -BED often significantly outperforms previous state-of-the-art in terms of EIG values even when targeting different downstream tasks and losses. In summary, we propose a paradigm shift in BED from thinking about uncertainties to future losses. Not only does this allow us to easily customise design policy learning to downstream tasks of interest, it also provides substantial computational and practical benefits, avoiding the double intractability of EURs by jointly optimising a downstream action policy alongside the design policy using our ACTION -BED framework. In turn, this provides significant empirical benefits that allow more informative and targeted data to be collected than existing state-of-the-art approaches.
2
Background
Bayesian experimental design. Bayesian experimental design (BED) [12, 17, 20, 30, 48, 49, 58] is typically defined as choosing designs to maximise the reduction in a particular uncertainty measure between the prior and the posterior. Namely, let p(θ)p(y|θ, ξ) denote the joint generating process over outcome y ∈ Y and variables of interest θ ∈ Θ given experimental design ξ ∈ Ξ. Then the expected uncertainty reduction (EUR) in θ for an uncertainty measure U and design ξ is: EURθ (ξ) = Ep(θ)p(y|θ,ξ) [U[p(θ)] − U [p(θ|y, ξ)]] .
(1)
Typically, U is taken to be the Shannon entropy, U[q(θ)] = Eq(θ) [− log q(θ)], in which case the EUR corresponds to the expected information gain (EIG) in θ, i.e. the mutual information between θ and y. 2
Sequential and policy-based BED. BED is often most impactful in sequential settings where one performs an experiment over T rounds. At each round t a design ξt ∈ Ξ is proposed and outcome yt ∈ Y is observed according to our proposed model p(yt |ξt , θ). Let ht := {(ξ1 , y1 ), . . . , (ξt , yt )} be the experimental history to time t, referred to in generality as the data. Traditional approaches to sequential BED [53] would at each step infer the posterior and select the next design by optimising the EIG, however this process is prohibitively expensive. Instead, policy-based approaches [8, 10, 18, 22, 28, 33, 34, 38, 39, 62, 63] train an adaptive policy network πd ∈ Πd which maps historical data ht onto optimal next designs ξt+1 . The policy is generally trained to maximise (an approximation of) the overall EUR from observing the full experimental history. This results in an acquisition policy which acts non-myopically to maximise the EUR from observing the full dataset and which can be deployed almost instantaneously. Note that the optimal design policy is always deterministic [49] and so we will assume deterministic policies throughout, but our methods can equally be applied with probabilistic design policies if desired. Optimising EIG. The longstanding challenge with EUR objectives such as the EIG is their fundamental double intractability. This arises because the uncertainty measures considered are non-linear functionals of the intractable posterior distribution [16, 25, 50, 60], which are then nested within expectations over the data generating process. Much work has been devoted to developing estimation and optimisation procedures for these objectives, in particular for the EIG [2, 13, 22–24, 26, 27, 31, 32, 36– 39, 42, 43, 57]. Of particular relevance is the family of variational bounds that allow the EIG to be represented as an optimisation over some kind of amortised posterior approximation [23, 24] (which is sometimes indirectly defined through a marginal likelihood approximation). For example, the Barber Agakov (BA) bound is given by [4]: IT (πd ) ≥ L(πd , ψ) = Ep(θ)p(y1:T |θ;πd ) [log qψ (θ|hT ) − log p(θ)]
(2)
where qψ (θ|hT ) is an amortised approximation of the posterior distribution. If there are some variational parameters ψ ∈ Ψ such that qψ (θ|hT ) perfectly captures the true posterior p(θ|hT ) for all possible hT , then maxψ∈Ψ L(πd , ψ) = IT (πd ). This bound then allows a singly intractable joint optimisation of both design policy πd and variational posterior qψ . However, in practice, it is limited by the variational family and difficulty of learning accurate amortised posteriors.
3
Loss-Driven Bayesian Experimental Design
As we have established, the de-facto perspective on BED as maximising EURs leads to objectives which are fundamentally doubly intractable and therefore pose a real computational barrier to their adoption, particularly at scale. Furthermore, variational formulations for circumventing this double intractability rely on learning accurate amortised posterior approximations, which is typically challenging and often unnecessary when the posterior itself is not the downstream goal of the experiment. They are also specific to the EIG, and there is a lack of more general computationally viable frameworks for other uncertainty measures, in particular customising to downstream tasks and losses of interest. Instead we propose an entirely loss-driven approach, inspired by the principle that the value of data should be measured by how well it allows us to complete downstream tasks. Grounding ourselves in first principles Bayesian decision theory (BDT), we show Bayes-optimal design policies take the form of the minimisers of an expected future loss when future actions are made Bayes-optimally. We then show how this can be turned into a practical, computationally tractable BED algorithm, ACTION -BED, which solves the aforementioned challenges and is additionally loss-driven and implicit, only relying on joint simulation of parameters and data under the proposed model. 3.1
Bayesian Experimental Design from First Principles
Under Bayesian coherence [14], a rational agent should plan ahead in a way that assumes that they will act optimally given the information available to them in the future. Using this principle, we can derive a Bayes-optimal design policy from first principles, by starting with terminal actions taken after all data is gathered and then working backwards through our sequential decisions. To derive the optimal downstream action given previously acquired data, we note that the Savage axioms imply the existence, for any rational agent, of a loss function on actions and the state of the world [59], ℓ : A × Θ → R, which encapsulates our preferences over the consequences from the actions we take, and a subjective belief model on the world given observed data p(θ|hT ). 3
Here θ ∈ Θ now denotes whatever aspect of the true world state directly influences our loss function. The Bayes-optimal action, aB , is then the minimiser of the Bayesian expected loss, aB = arg mina∈A Ep(θ|hT ) [ℓ(a, θ)]. This is easily generalised to cases where the loss depends on the data hT itself, but we omit this dependence for ease of notation. For a given belief state and loss function, the minimal achievable Bayesian expected loss, mina∈A Ep(θ|hT ) [ℓ(a, θ)], is now a function of the data we observe. It is analogous to a value function in reinforcement learning (RL) [52]: better data will provide a more certain belief state and, in turn, a better achievable expected downstream loss. Properly formalising this requires the definition of a belief distribution on possible data that will be generated under our design policy, p(hT ; πd ). The existence of this belief distribution is again implied for a rational agent by the Savage axioms, noting that we are uncertain about both hT and θ before the experiment. Further, our beliefs R on θ should not depend on our data gathering policy, so we should have that our “prior” p(θ) = p(hT ; πd )p(θ|hT )dhT is independent of πd . As minimising the loss is our ultimate objective and a rational agent will act optimally once the data is gathered, the Bayes-optimal data gathering policy is then the minimiser of the expected future loss (EFL) when downstream actions are themselves made Bayes-optimally as follows. Definition 1 (Bayes-Optimal Experimental Design Policy). The Bayes-optimal sequential experimental design policy πdB for a given belief model and downstream loss pairing is defined as: B πdB = arg min EFLB (π ) where EFL (π ) = E min E [ℓ(a, θ)] . (3) d d p(hT ;πd ) p(θ|hT ) ℓ ℓ a∈A
πd ∈Πd
We note that similar theoretical EFL formulations of the Bayes-optimal experimental designs have already been established in the classical BED literature, most notably Lindley [49] and Dawid [15]. However, they only considered individual design decisions rather than design policies and, moreover, it has not since been directly used for deriving computationally viable BED approaches. Instead, the BED field has since exclusively relied on the mathematical equivalence of this EFL formulation to EURs (see Section 3.3) when constructing practical methods. 3.2
Singly Intractable and Loss-Driven BED via Joint Policy Optimisation
While Definition 1 provides a precise formulation of the Bayes-optimal data gathering policy as the minimiser of a task-driven EFL, directly solving Equation (3) is very difficult as it involves a nested optimisation due to the dependence of the optimal action aB on the observed data hT . Our key insight is to switch from trying to optimise individual actions to optimising an action policy, πa , which maps from observed data to downstream actions. As we show below, this leads to a singly intractable joint objective for design and action policy training which is compatible with any downstream loss. Let πa ∈ Πa := {π | π : Y T × ΞT → A} be an action policy mapping observed data to downstream actions. Noting that Πa contains all admissible measurable mappings from the data to actions and the optimal action is defined on a data-by-data basis for a given model and loss, under standard measurability assumptions detailed in Appendix A, we can lift the optimisation over individual actions in the EFL outside the expectation to become an optimisation over action policies as follows: EFLB ℓ (πd ) = min Ep(hT ;πd )p(θ|hT ) [ℓ(πa (hT ), θ)] . πa ∈Πa
(4)
While the minimisation over action policies in Equation (4) does not usually have an analytic solution, we can instead think of the EFL under more general downstream action policies before jointly optimising for both πa and πd , which will still recover the Bayes-optimal policies for both actions: (πa , πd )B = arg min EFLℓ (πd , πa ), (πa ,πd )∈Πa ×Πd (5) where EFLℓ (πd , πa ) = Ep(θ)p(y1:T |θ;πd ) [ℓ(πa (hT ), θ)] , and we have rewritten our belief model using p(θ)p(y1:T |θ; πd ) = p(hT ; πd )p(θ|hT ). We refer to approaches that directly train policies using Equation (5) as ACTION -BED. Its seemingly straightforward reinterpretation of experimental design as a joint optimisation problem now has remarkably impactful consequences. Indeed the EFL represents a completely general, singly intractable, naturally implicit and loss-driven objective for optimal experimental design. It never requires the approximation of posterior distributions or evaluation of unknown densities, it simply relies on simulations from the model. As we will see in Section 4, parameterising both action and design policies as deep neural networks leads to a practical end-to-end experimental design algorithm. 4
3.3
Recovering the Expected Uncertainty Reduction Formulation
Traditional EUR objectives can be formalised and reconciled with our BED formulation through the framework of proper scoring rules. A scoring rule is a function s : Pθ × Θ → R which assigns a score s(pθ , θ) for predicting a distribution pθ ∈ Pθ when the realisation of a random variable Θ is θ. The scoring rule is (strictly) proper if it is (strictly) minimised when the predicted distribution is exactly the distribution of the target random variable. Following ideas introduced by [7, 15, 17], we can define general uncertainty measures, known as generalised entropies, using hs [q(θ)] = Eq(θ) [s(q(·), θ)]. We then define an experimental design objective via the expected posterior uncertainty (EPU): EPUs (πd ) = Ep(y1:T ;πd ) [hs [p(θ|hT )]] = Ep(θ)p(y1:T |θ;πd ) [s(pθ (·|hT ), θ)] .
(6)
We now show that via an equivalence between loss-action pairings and proper scoring rules, any optimal BED policy according to Definition 1 can be represented as the minimiser of an EPU for some s, and a minimiser of an EPU for any s corresponds to a Bayes-optimal design policy according to Definition 1 for at least one loss-belief pairing (see Appendix B.1 for proof). Theorem 1 (Equivalence of EFL & EPU). Let s : Pθ × Θ → R be a proper scoring rule on Θ. Then πdB = arg min Ep(θ)p(y1:T |θ;πd ) s pθ (·|hT ), θ (7) πd ∈Πd
is a Bayes-optimal design policy for the belief model p(θ)p(y1:T |θ; πd ) and for any loss ℓ ∈ L, where n o L = ℓ : ℓ π̃a (qθ ), θ = s(qθ , θ), ∀qθ ∈ Pθ , θ ∈ Θ , π̃a (qθ ) = arg min Eθ∼qθ ℓ(a, θ) . (8) a∈A
Furthermore, the set L is never empty, and any lower-bounded loss function ℓ induces a corresponding scoring rule s that is proper, though not necessarily strictly proper. In order to make explicit the connection to EUR objectives, we note that if the scoring rule does not directly depend on πd or y1:T , the expected uncertainty of the prior Ep(y1:T ;πd ) [hs [p(θ)]] = hs [p(θ)] is constant in πd , therefore EUR objectives are equivalent to the EPU objectives considered in Theorem 1. The set of possible EFL and EPU objectives is thus strictly greater than EURs. The EIG itself is equivalent, up to a constant offset, to an EPU objective with the log score s(qθ (·), θ) = − log qθ (θ), for which the optimal action turns out to be the Bayesian posterior itself. Theorem 1 thus reconciles the two BED paradigms as different representations of the same underlying decision problem. The key distinctions are thus first in shifting the burden of problem specification from directly choosing uncertainty measures to specifying an explicit downstream task and loss. This makes it easier to tailor objectives for specific downstream tasks, such as targeting the point estimation of some quantity of interest, or incorporating in real-world decision-making, such as business decisions or choosing medical treatments. It also makes clear how choices of uncertainty measures in the EUR framework implicitly target particular downstream tasks and losses, e.g. EIG targets log loss on probabilistic predictions, expected variance reduction targets point estimation with squared error loss. Second, there are critical practical distinctions between the viewpoints in terms of their computational viability, with the EFL formulation leading to the general purpose singly intractable problem formulation we introduced in Equation (5). Interestingly, this formulation recovers existing variational formulations [4, 23] as a special case when using an EFL that is equivalent to the EIG (as here the optimal action is the posterior, see Appendix B.2). However, it critically generalises to other action spaces and losses, and it turns out the action policies we need to learn for other downstream losses are often much simpler than those required for the EIG. In particular, when the required downstream actions correspond to point estimates or explicit decisions rather than distributions, then the action policy is a simple functional mapping to single values, avoiding the need to implicitly estimate or approximate the posterior. For example, if using squared error loss on an estimation task, then the action policy only needs to learn to map from data to the posterior mean. This can substantially simplify the policy training. Indeed we observe in our experiments (Section 6) that suitably chosen predictive losses indirectly achieve higher EIG values than direct EIG optimisation approaches.
5
4
ACTION -BED
We now show how our ACTION -BED frameRequire: Policies πdϕ , πaψ , initialisations ϕ0 , ψ0 work, based on the joint optimisation of 1: for k = 0, 1, . . . , K − 1 do the EFL with respect to action and design 2: θ1:B , ε1:B 1:T ∼ p(θ)q(ε1:T ) policies as per Equation (5), can be reϕ ψ [ 3: g , g k k ← ∇ϕ,ψ EFL(ϕk , ψk ) (eq. (10)) alised as a practical end-to-end algorithm for ϕ ψ 4: ϕ , ψ task-driven BED. Specifically, building on k+1 k+1 ← U PDATE (ϕk , gk , ψk , gk ) 5: end for the deep adaptive design (DAD) approach 6: return ϕK , ψK of [22, 39], we propose to directly parameterise both policies using deep neural netAlgorithm 1: ACTION -BED works, denoted πdϕ and πaψ , with parameters ϕ ∈ Φ and ψ ∈ Ψ. We then train these networks using end-to-end stochastic gradient descent on the expected future loss for rollouts drawn from the model, as shown in Algorithm 1. Empirically, it is often helpful to initialise ψ ∈ Ψ by pre-training using a fixed random design policy (Appendix E.2), following which our algorithm remains unchanged. For example, if p(y1:T |θ; πdϕ ) is reparameterisable, we can optimise this objective using pathwise gradients [51] obtained via reparameterisation [41]: h i ∇ϕ,ψ EFL(ϕ, ψ) = Ep(θ)q(ε1:T ) ∇ϕ,ψ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ , (9) where we reparameterised hT = hT (θ, ε1:T ; πdϕ ) for ε1:T ∼ q(ε1:T ). See Appendix C for full details of the reparameterisation. We can then estimate these gradients unbiasedly using simple Monte Carlo: 1 [ ∇ ϕ,ψ EFL(ϕ, ψ) = B
B X
∇ϕ,ψ ℓ πaψ (hT (θb , εb1:T ; πdϕ )), θb
(10)
b=1
for θb ∼ p(θ), εb1:T ∼ q(ε1:T ), from which stochastic gradient updates can be performed using any desired optimizer. If the simulator is not reparameterisable, we can resort to the REINFORCE estimator [66] to estimate gradients, provided that we have access to the score function of the data generating process. These requirements are the main limitations of our approach, although we note that a very large class of BED problems satisfy at least one of these assumptions. Details on the architectures used for the policies are provided in Appendix E.
5
Related Work
To the best of our knowledge, there are very limited prior works considering truly task-driven approaches to BED. One approach is Huang et al. [33], who maximise the improvement in expected utility of the optimal downstream action after observing new data. They only consider utility functions on actions living in outcome space Y, meaning their method cannot be applied as a baseline in any of our BED experiments and is instead applied on prediction-oriented active learning settings [65]. Furthermore, they maximise their expected utility objective by approximating the intractable predictive distribution and explicitly solving the nested optimisation problem over downstream actions. In contrast, our approach does not require any distributional approximations and removes the nested optimisation by learning action policies. Short of being truly task-based, recent works [11, 19, 44, 62, 68, 69] consider goal-oriented experimental design, in which designs are sought to maximise the EUR in downstream quantities of interest. While this approach shifts focus from model parameters, it is still fundamentally an EUR objective and therefore inherits the associated computation challenges. Finally, Kerrigan et al. [40] introduces a new perspective on BED, maximising the dependence between y and θ by solving an optimal transport problem with a chosen cost function. The choice of cost function provides a way of tailoring designs to downstream tasks of interest, but cannot be interpreted as minimising a loss over downstream actions in a principled BDT perspective. In alternative contexts, Huang et al. [35] introduces a task-driven framework for active learning, which they solve analytically when the downstream loss is a Bregman divergence. Our approach is instead singly intractable for any downstream loss. In a Bayesian optimisation (BO) setting, Neiswanger et al. [55] frame classical BO acquisition functions in a loss-based perspective via generalised entropies, which they use to derive new acquisition functions for custom BO tasks. In their case, analytic solutions depend on the Gaussian process posterior, specific to the BO setting. 6
Our approach also shares strong connections with RL, since mina∈A Ep(θ|hT ) [ℓ(a, θ)] can be interpreted as a value function giving the best achievable downstream loss given the current belief state. Prior works [5, 8, 9, 62, 63] have adopted an RL approach to BED by defining rewards as the incremental improvement in the above value function. However, this has only been done for the EIG, implicitly choosing the log-loss to define the value function, while the generalisation to any downstream loss has not yet been made. Furthermore, while variational and amortised inference approaches [10, 18, 24, 34, 62] learn amortised variational approximations which can be interpreted as downstream action networks (often within RL frameworks e.g. Shen et al. [62]), the generalisation to arbitrary downstream actions has also not yet been made. Ultimately our approach could directly benefit from inheriting RL ideas, such as reward shaping and credit assignment, in the future.
6
Experiments
We evaluate ACTION -BED across four BED tasks spanning a range of complementary settings. Our experiments illustrate the improved performance and computational efficiency of ACTION -B ED on downstream tasks, due to our task-driven and singly intractable procedure. Baselines and evaluation protocol. We compare against two competitive policy-based BED baselines. DAD [22] serves as the standard EIG-based policy baseline when an explicit likelihood is available, while I DAD [39] provides its likelihood-free counterpart. ALINE [34] is a gold-standard EIG-based approach using a transformer as a joint amortised design and inference network, trained using reinforcement learning with a dense reward. Additionally we include a Random policy which samples designs from a task-specific proposal. In order to evaluate these baselines on downstream actions, we train a design policy using the baseline approach, which is then frozen while a downstream action network is trained to minimise the downstream task loss on trajectories generated by the pre-trained design policy. Thus, the design network of each baseline is trained to maximise EIG while the downstream network is trained to minimise downstream loss by approximating the Bayes-optimal downstream action. Methods are evaluated in terms of downstream loss, design policy EIG bounds [22], and end-to-end wall-clock GPU training time. Training robustness. To confirm that our relative performance gains are not driven by architectural choices, we matched architectures across all baselines and further repeated our experiments using smaller architectures in Appendix G.4. Due to limited computational resources, our main results for the source location finding benchmark are presented for a single training seed. We therefore evaluate training stability using the smaller architectures, with results reported in Appendix G.4.2 and Table 33, showing our downstream performance is stable over training seeds. 6.1
Source Location Finding
The first experiment is inspired by the acoustic energy attenuation model introduced in [64]. The unknown parameter θ ∈ R2×2 specifies the locations of two sources which emit signals that decay with distance. Over T = 30 steps, the design policy sequentially selects locations and observes the resulting aggregate signal intensity. The goal is to learn an adaptive design strategy that enables efficient downstream prediction of source locations. For the downstream predictive Figure 2: Comparison of EIG bounds as a function loss, we consider the mean squared error (MSE) of design policy training time, until convergence. and the Log-MSE, shorthand for an MSE loss augmented with a logarithmic penalty, ℓLog-MSE (θ, θ̂) = ||θ − θ̂||2 + log ||θ − θ̂||2 + ϵ , which increases incentive for precise localisation. Implementation, architectural, and experimental details are provided in Appendix G.1. Results. Table 1 shows the quantitative results on the source location finding task while qualitative design trajectories are included in Appendix G.1.1. We firstly observe the remarkable performance of ACTION -BED under the Log-MSE objective; we achieve both the best downstream Log-MSE error as well as state-of-the-art EIG of the design policy. We believe that this is possible because the Log-MSE remains sensitive to very accurate predictions and continues to reward design policies that identify the sources more precisely, whereas the sPCE contrastive lower bound saturates at 7
Table 1: Source Location Finding. Uncertainties are reported as one standard error over 2048 rollouts. Parentheses indicate the loss used to train the downstream action network. Compute time includes training both the design and downstream action networks. Performance (± 1 s.e.)
Method −2
MSE (10
)↓
Compute
Log-MSE ↓
sPCE ↑
sNMC ↑
Time (min) ↓
Action-BED (MSE) Action-BED (Log)
0.55 ± 0.06 2.8 ± 0.4
−6.83 ± 0.03 −10.10 ± 0.03
11.63 ± 0.04 17.14 ± 0.11
12.84 ± 0.08 21.07 ± 0.45
1171 1184
DAD (MSE) DAD (Log)
5.6 ± 0.6 10.7 ± 0.8
−6.35 ± 0.04 −6.58 ± 0.05
12.94 ± 0.03 12.94 ± 0.03
16.35 ± 0.14 16.35 ± 0.14
1214 1216
ALINE (MSE) ALINE (Log)
0.21 ± 0.03 0.81 ± 0.10
−7.58 ± 0.05 −8.11 ± 0.07
14.83 ± 0.05 14.83 ± 0.05
16.97 ± 0.13 16.97 ± 0.13
2962 2970
Random
22.4 ± 0.8
−3.16 ± 0.04
8.27 ± 0.04
8.49 ± 0.05
31
high EIG values therefore becoming harder to train to. Utilising this informative loss function and sidestepping the challenges of EIG-based training is made directly possible by our approach. Under the MSE loss, ALINE performs best however requires around 2.5 times the compute to achieve such a result. A plausible explanation is that its dense training rewards help identify the latent sources more efficiently, while such guidance is absent in our setting. ACTION -BED significantly outperforms DAD, which uses a similar amount of compute. In terms of computational efficiency, we show in Figure 2 the attained EIG of the design policy versus training time, where ACTION -BED (Log) achieves competitive EIG values faster than competing approaches. Additional ablations and evaluation. We provide several additional analyses in Appendix G.1.2 and G.1.3. We first evaluate the decision-relevant quality of the learned design policies independently of their associated downstream action networks by applying a common posterior-based decision procedure to all methods. This isolates the informativeness of the acquired histories themselves and confirms the objective-dependent ranking observed in the main results. Moreover, it highlights the substantial deployment-time cost avoided by amortised decision-making with ACTION -BED; see Tables 16 and 17. We then investigate the coupling between the jointly learned design policy and downstream action policy in ACTION -BED by evaluating each design and action network crosspairing; see Table 20. Downstream performance degrades substantially under such policy swaps, showing that the action network exploits not only the observed outcomes but also information encoded in the adaptive design trajectory itself. This coupling is strongest for the most accurate policies, whose designs concentrate efficiently near the latent sources. Second, we show for ACTION -BED that retraining the action policy on a fixed pre-trained design policy leads to a mild downstream degradation compared to joint training, suggesting that the joint dynamics beneficially co-adapt the two policies; see Table 19. Finally, we ablate the warmstart used before joint optimisation; see Table 21. Removing it slightly worsens downstream performance, likely because early lowquality action predictions provide a noisier training signal to the design policy, perturbing the joint optimisation dynamics. In this benchmark, warmstarting is therefore useful and incurs only a small computational overhead. 6.2
Dynamical Systems
We next consider stochastic single- and double-pendulum systems [6, 38]. In these tasks, the system evolves over time according to a controlled SDE depending on some unknown physical parameters. The continuous dynamics are discretised using Euler-Maruyama over T = 50 steps with dt = 0.05. At each step an agent influences the dynamics by applying a torque ξt . The goal is to learn an adaptive design policy that enables accurate downstream prediction of the latent physical parameters θ ∈ Θ from the observed trajectories. We evaluate the prediction with several downstream objectives—MSE, log-MSE, and a weighted MSE over the components of θ—the latter requiring the design policy to prioritise parameter regions of particular interest. These tasks pose three main challenges: the datagenerating process is non-exchangeable due to Markovian dynamics, designs must satisfy physical constraints, and longer rollouts induce complex feedback between interventions and observations. Stochastic Pendulum. This task comprises a single pendulum with state yt = (qt , q̇t ) and unknown parameters θ ∈ Θ ⊆ R3 . At each step, the agent applies a bounded torque ξt ∈ [−1, 1] to control the system. Full experimental details are provided in Appendix G.2.1. Table 2 confirms that loss-driven acquisition improves the corresponding downstream objective: each variant performs best under the 8
Table 2: Stochastic Single Pendulum. Uncertainties are reported as one standard error over 2048 rollouts. Compute time includes training both the design and downstream action networks. Performance (± 1 s.e.)
Method
Compute
MSE (10−3 ) ↓
Log-MSE ↓
W. MSE (10−3 ) ↓
sPCE ↑
sNMC ↑
Time (min) ↓
Action-BED (MSE) Action-BED (Log) Action-BED (W. MSE)
6.00 ± 0.14 7.20 ± 0.17 9.80 ± 0.25
−6.60 ± 0.05 −6.75 ± 0.05 −6.07 ± 0.04
4.90 ± 0.13 4.80 ± 0.11 3.50 ± 0.09
3.68 ± 0.04 3.99 ± 0.04 3.33 ± 0.04
3.71 ± 0.04 4.00 ± 0.04 3.36 ± 0.04
240 287 290
DAD (MSE) DAD (Log) DAD (W. MSE)
7.20 ± 0.17 8.10 ± 0.19 7.30 ± 0.17
−6.20 ± 0.04 −6.49 ± 0.05 −6.44 ± 0.05
5.00 ± 0.14 5.20 ± 0.14 4.40 ± 0.12
3.82 ± 0.04 3.82 ± 0.04 3.82 ± 0.04
3.88 ± 0.04 3.88 ± 0.04 3.88 ± 0.04
670 709 671
ALINE (MSE) ALINE (Log) ALINE (W. MSE)
23.50 ± 0.51 27.53 ± 0.59 23.58 ± 0.50
−5.43 ± 0.02 −5.69 ± 0.02 −5.61 ± 0.02
7.59 ± 0.12 9.89 ± 0.17 7.06 ± 0.11
2.23 ± 0.06 2.23 ± 0.06 2.23 ± 0.06
2.24 ± 0.06 2.24 ± 0.06 2.24 ± 0.06
2120 2141 2122
Random
29.0 ± 1.2
−4.25 ± 0.05
12.6 ± 1.0
1.42 ± 0.03
1.50 ± 0.03
19
Table 3: Stochastic Double Pendulum. Uncertainties are reported as one standard error over 2048 rollouts. Compute time includes training both the design and downstream action networks. Performance (± 1 s.e.)
Method
Compute
MSE (10−4 ) ↓
Log-MSE ↓
W. MSE (10−4 ) ↓
sPCE ↑
sNMC ↑
Time (min) ↓
Action-BED (MSE) Action-BED (Log) Action-BED (W. MSE)
1.11 ± 0.03 1.31 ± 0.03 1.14 ± 0.03
−10.21 ± 0.04 −10.32 ± 0.04 −10.28 ± 0.04
1.64 ± 0.04 1.75 ± 0.05 1.51 ± 0.05
10.73 ± 0.05 10.81 ± 0.12 10.65 ± 0.04
12.71 ± 0.19 14.29 ± 0.56 12.10 ± 0.16
501 504 504
DAD (MSE) DAD (Log) DAD (W. MSE)
2.10 ± 0.06 2.43 ± 0.10 2.13 ± 0.06
−9.50 ± 0.03 −9.51 ± 0.03 −9.47 ± 0.03
2.47 ± 0.07 2.79 ± 0.08 2.46 ± 0.07
11.37 ± 0.03 11.37 ± 0.03 11.37 ± 0.03
14.79 ± 0.25 14.79 ± 0.25 14.79 ± 0.25
934 940 933
ALINE (MSE) ALINE (Log) ALINE (W. MSE)
11.21 ± 0.34 19.21 ± 0.51 11.81 ± 0.33
−8.46 ± 0.02 −8.53 ± 0.03 −8.52 ± 0.03
15.12 ± 0.04 23.29 ± 0.07 14.02 ± 0.04
9.28 ± 0.07 9.28 ± 0.07 9.28 ± 0.07
9.78 ± 0.14 9.78 ± 0.14 9.78 ± 0.14
3013 3062 3007
Random
20.0 ± 0.9
−6.84 ± 0.05
48.0 ± 1.3
8.26 ± 0.06
8.41 ± 0.08
58
loss used for training. Interestingly, we observe again that the design policy of ACTION -BED (Log) achieves the highest EIG, outperforming the baselines which directly optimise EIG. We also observe that our method requires substantially less compute than the baselines. Stochastic Double-Pendulum. The double-pendulum task is substantially harder as its dynamics are nonlinear, coupled, and potentially chaotic, making trajectories highly sensitive to designs. The state yt ∈ R4 contains the two joint angles and angular velocities, while θ ∈ Θ ⊆ R4 contains the unknown masses and link lengths. Designs are bounded joint torques ξt = (ξt1 , ξt2 ) ∈ [−4, 4] × [−2, 2]. Details are given in Appendix G.2.2. In this setting, Table 3 again shows that ACTION -BED achieves considerably improved downstream prediction for reduced computational cost. In this task, DAD appears to achieve marginally higher EIG although EIG bounds overlap with ACTION -BED. 6.3
MNIST Masked Classification
Finally, we consider a sequential masked-classification task adapted from the experimental setting in [37]. Each episode begins with an unobserved image θ ∈ R28×28 from MNIST [46], together with its class label z ∈ {0, . . . , 9}. At each of the T = 5 acquisition steps, the policy selects the corner ξ ∈ [1, 28]2 of a local 5 × 5 patch and observes a noisy measurement of that patch, y = Aξ θ + η, with η ∼ N (0, σId ). Here, Aξ θ denotes the masking operator induced by ξ, and applied to the image. The goal is to select the patches most informative for downstream digit classification. This setting is challenging because acquisition is driven by the class label, whereas measurements arise from the latent image. The patch likelihood p(y|θ, ξ) is explicit, but its label-conditioned counterpart p(y|z, ξ) is only implicitly induced by the MNIST class distribution. This problem is also strongly information-constrained, since the downstream policy must infer the class from only a few noisy local measurements in a high-dimensional space. Additional details are provided in Appendix G.3. 9
Table 4: MNIST Masked Classification. Uncertainties are reported as one standard error across evaluation batches. We report cross-entropy and classification accuracy on both the training and test datasets. GPU time includes training of both the design and downstream classification networks. Train dataset Method
Test dataset
Compute
Cross-entropy
Accuracy (%)
Cross-entropy
Accuracy (%)
GPU time (min)
Action-BED
0.033 ± 0.001
98.87 ± 0.04
0.097 ± 0.007
97.42 ± 0.16
9
DAD iDAD ALINE
0.350 ± 0.004 0.099 ± 0.002 0.256 ± 0.008
83.32 ± 0.04 96.49 ± 0.08 92.13 ± 0.08
0.546 ± 0.009 0.164 ± 0.008 0.271 ± 0.01
81.83 ± 0.16 95.10 ± 0.21 91.59 ± 0.10
53 12 357
Random
0.786 ± 0.015
71.78 ± 0.87
0.797 ± 0.015
71.40 ± 0.86
5
ACTION -BED naturally targets this implicit problem by jointly training a design policy and downstream classifier to minimise the cross-entropy loss: min Ep(z)p(θ|z)p(y1:T |θ;πϕ ) ℓCE πaψ (·|hT ), z ϕ,ψ
d
where ℓCE πaψ (·|hT ), z = − log πaψ (z|hT ) is the cross-entropy loss and πaψ (·|hT ) is a probabilistic classifier. Training to this objective only requires simulating z, θ from the empirical MNIST dataset and reparameterising the simulation of p(y1:T |θ; πd ). By contrast, DAD must target EIG in the latent image θ, since the outcomes conditional on class label p(y|z, ξ) do not have explicit likelihood. I DAD is able to target the EIG in the label z, thereby providing a more competitive baseline. ALINE does not naturally apply in this setting due to the hierarchical likelihood and classification target, we therefore adapted the setup of the inference network accordingly. Quantitative and qualitative results are reported respectively in Table 4 and Figure 7. ACTION -BED achieves the best downstream performance across the board, demonstrating the lowest cross-entropy loss and highest classification accuracy amongst all approaches, whilst requiring the least training time. These results demonstrate that ACTION -BED easily maximises downstream performance in implicit models without additional modifications, critic networks, or internal posterior approximations.
7
Conclusion
In this work we introduce ACTION -BED, a practical algorithm for task-driven BED, resulting from the simple principle of minimising an expected future loss over downstream actions. We show that this approach is grounded in first principles decision theory and we demonstrate an equivalence to existing EUR approaches. Our algorithm, which jointly learns a design network and downstream action network, is singly intractable, naturally implicit and directly targets downstream task performance. As a result, we demonstrate improved downstream performance compared to baselines which directly optimise for EIG, while requiring significantly less training effort. Our approach is applicable to BED problems with reparameterisable simulators or access to gradients of the log-likelihood of the observation model, which represents a very broad class of BED problems. Advancing the efficiency of BED methods for downstream task performance has clear societal benefits by enabling better decision making and downstream predictions whenever these actions are informed by experimental data, for instance in medical diagnoses. While experimental data can sometimes be collected with malicious intent, we believe the potential benefits of our approach outweigh the risks of misuse.
Acknowledgements Tom Rossa is supported by the EPSRC CDT in Statistics and Machine Learning (EP/Y034813/1). Angus Phillips is supported by the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1). Tom Rainforth is supported by the EPSRC grant EP/Y037200/1.
10
References [1] Aliprantis, C. D. and Border, K. C. (2006). Infinite Dimensional Analysis: A Hitchhiker’s Guide. Springer, 3 edition. 16 [2] Ao, Z. and Li, J. (2024). On Estimating the Gradient of the Expected Information Gain in Bayesian Experimental Design. Proceedings of the AAAI Conference on Artificial Intelligence, 38(18):20311–20319. Number: 18. 2, 3 [3] Ashman, M., Diaconu, C., Weller, A., and Turner, R. E. (2024). In-context in-context learning with transformer neural processes. In Antorán, J. and Naesseth, C. A., editors, Proceedings of the 6th Symposium on Advances in Approximate Bayesian Inference, volume 253 of Proceedings of Machine Learning Research, pages 1–29. PMLR. 23 [4] Barber, D. and Agakov, F. (2003). Information Maximization in Noisy Channels : A Variational Approach. In Advances in Neural Information Processing Systems, volume 16. MIT Press. 2, 3, 5 [5] Barlas, Y. and Salako, K. (2025). Performance Comparisons of Reinforcement Learning Algorithms for Sequential Experimental Design. Philadelphia, USA. 7 [6] Belousov, B., Abdulsamad, H., Schultheis, M., and Peters, J. (2019). Belief space model predictive control for approximately optimal system identification. In Proceedings of the Multidisciplinary Conference on Reinforcement Learning and Decision Making (RLDM). 8 [7] Bickford Smith, F., Kossen, J., Trollope, E., Van Der Wilk, M., Foster, A., and Rainforth, T. (2025). Rethinking aleatoric and epistemic uncertainty. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 4345–4359. PMLR. 5 [8] Blau, T., Bonilla, E. V., Chades, I., and Dezfouli, A. (2022). Optimizing Sequential Experimental Design with Deep Reinforcement Learning. In Proceedings of the 39th International Conference on Machine Learning, pages 2107–2128. PMLR. 3, 7 [9] Blau, T., Chades, I., Dezfouli, A., Steinberg, D. M., and Bonilla, E. V. (2023). Cross-Entropy Estimators for Sequential Experiment Design with Reinforcement Learning. 7 [10] Bracher, N., Kühmichel, L., Ivanova, D. R., Intes, X., Bürkner, P.-C., and Radev, S. T. (2025). JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference. arXiv:2512.22999 [stat]. 2, 3, 7 [11] Chakraborty, A., Huan, X., and Catanach, T. A. (2024). A Likelihood-Free Approach to Goal-Oriented Bayesian Optimal Experimental Design. 6 [12] Chaloner, K. and Verdinelli, I. (1995). Bayesian Experimental Design: A Review. Statistical Science, 10(3):273–304. 1, 2 [13] Coons, T. E. and Huan, X. (2025). A Multi-fidelity Estimator of the Expected Information Gain for Bayesian Optimal Experimental Design. arXiv:2501.10845 [stat]. 3 [14] Cox, R. T. et al. (1946). Probability, frequency and reasonable expectation. American journal of physics, 14(1):1–13. 3 [15] Dawid, A. P. (1998). Coherent measures of discrepancy, uncertainty and dependence, with applications to bayesian predictive experimental design. Department of Statistical Science, University College London. http://www. ucl. ac. uk/Stats/research/abs94. html, Tech. Rep, 139. 4, 5 [16] Dawid, A. P. (2007). The geometry of proper scoring rules. Annals of the Institute of Statistical Mathematics, 59(1):77–93. 2, 3 [17] DeGroot, M. (1962). Uncertainty, information, and sequential experiments. The Annals of Mathematical Statistics, 33(2):404–419. 1, 2, 5 [18] Dong, J., Jacobsen, C., Khalloufi, M., Akram, M., Liu, W., Duraisamy, K., and Huan, X. (2025). Variational Bayesian optimal experimental design with normalizing flows. Computer Methods in Applied Mechanics and Engineering, 433:117457. 3, 7 11
[19] Filstroff, L., Sundin, I., Mikkola, P., Tiulpin, A., Kylmäoja, J., and Kaski, S. (2024). Targeted Active Learning for Bayesian Decision-Making. Transactions on Machine Learning Research. 6 [20] Fisher, R. (1935). The design of experiments. 2 [21] Fort, G., Gobet, E., and Moulines, E. (2017). Mcmc design-based non-parametric regression for rare event. application to nested risk computations. Monte Carlo Methods and Applications, 23(1):21–42. 2 [22] Foster, A., Ivanova, D. R., Malik, I., and Rainforth, T. (2021). Deep Adaptive Design: Amortizing Sequential Bayesian Experimental Design. In Proceedings of the 38th International Conference on Machine Learning, pages 3384–3395. PMLR. 2, 3, 6, 7, 23, 29, 30, 48 [23] Foster, A., Jankowiak, M., Bingham, E., Horsfall, P., Teh, Y. W., Rainforth, T., and Goodman, N. (2019a). Variational Bayesian Optimal Experimental Design. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc. 2, 3, 5, 18 [24] Foster, A., Jankowiak, M., O’Meara, M., Teh, Y., and Rainforth, T. (2019b). A Unified Stochastic Gradient Approach to Designing Bayesian-Optimal Experiments. 2, 3, 7, 18 [25] Gneiting, T. and Raftery, A. E. (2007). Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477):359–378. 2, 3 [26] Goda, T., Hironaka, T., and Iwamoto, T. (2020). Multilevel Monte Carlo estimation of expected information gains. Stochastic Analysis and Applications, 38(4):581–600. _eprint: https://doi.org/10.1080/07362994.2019.1705168. 2, 3 [27] Goda, T., Hironaka, T., Kitade, W., and Foster, A. (2022). Unbiased MLMC Stochastic GradientBased Optimization of Bayesian Experimental Designs. SIAM Journal on Scientific Computing, 44(1):A286–A311. 3 [28] Hedman, M., Ivanova, D. R., Guan, C., and Rainforth, T. (2025). Step-DAD: Semi-Amortized Policy-Based Bayesian Experimental Design. In Proceedings of the 42nd International Conference on Machine Learning, pages 22904–22923. PMLR. 3 [29] Hochreiter, S. and Schmidhuber, J. (1997). Long short-term memory. Neural Computation, 9(8):1735–1780. 23 [30] Huan, X., Jagalur, J., and Marzouk, Y. (2024). Optimal experimental design: Formulations and computations. Acta Numerica, 33:715–840. 1, 2 [31] Huan, X. and Marzouk, Y. (2012). Gradient-based stochastic optimization methods in Bayesian experimental design. International Journal for Uncertainty Quantification, 4. 3 [32] Huan, X. and Marzouk, Y. M. (2013). Simulation-based optimal Bayesian experimental design for nonlinear systems. J. Comput. Phys., 232(1):288–317. 3 [33] Huang, D., Guo, Y., Acerbi, L., and Kaski, S. (2024). Amortized Bayesian Experimental Design for Decision-Making. 2, 3, 6 [34] Huang, D., Wen, X., Bharti, A., Kaski, S., and Acerbi, L. (2025). ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition. 2, 3, 7 [35] Huang, Z., Smith, F. B., and Rainforth, T. (2026). Loss-Driven Bayesian Active Learning. arXiv:2604.11995 [cs]. 6 [36] Iollo, J., Heinkelé, C., Alliez, P., and Forbes, F. (2024). PASOA- PArticle baSed Bayesian Optimal Adaptive design. In Proceedings of the 41st International Conference on Machine Learning, pages 21020–21046. PMLR. 2, 3 [37] Iollo, J., Heinkelé, C., Alliez, P., and Forbes, F. (2025). Bayesian Experimental Design via Contrastive Diffusions. arXiv:2410.11826 [stat] version: 2. 9, 44 12
[38] Iqbal, S., Corenflos, A., Särkkä, S., and Abdulsamad, H. (2024). Nesting Particle Filters for Experimental Design in Dynamical Systems. In Proceedings of the 41st International Conference on Machine Learning, pages 21047–21068. PMLR. 2, 3, 8, 39 [39] Ivanova, D. R., Foster, A., Kleinegesse, S., Gutmann, M. U., and Rainforth, T. (2021). Implicit Deep Adaptive Design: Policy-Based Experimental Design without Likelihoods. In Advances in Neural Information Processing Systems, volume 34, pages 25785–25798. Curran Associates, Inc. 3, 6, 7, 23, 50 [40] Kerrigan, G., Naesseth, C. A., and Rainforth, T. (2025). A Geometric Approach to Optimal Experimental Design. arXiv:2510.14848 [stat]. 1, 6 [41] Kingma, D. P. and Welling, M. (2013). Auto-Encoding Variational Bayes. arXiv:1312.6114 [stat]. 6 [42] Kleinegesse, S. and Gutmann, M. U. (2019). Efficient Bayesian Experimental Design for Implicit Models. In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, pages 476–485. PMLR. 3 [43] Kleinegesse, S. and Gutmann, M. U. (2020). Bayesian Experimental Design for Implicit Models by Mutual Information Neural Estimation. In Proceedings of the 37th International Conference on Machine Learning, pages 5316–5326. PMLR. 3 [44] Kleinegesse, S. and Gutmann, M. U. (2021). Gradient-based Bayesian Experimental Design for Implicit Models using Mutual Information Lower Bounds. arXiv:2105.04379 [stat]. 6 [45] Kuratowski, K. and Ryll-Nardzewski, C. (1965). A general theorem on selectors. Bull. Acad. Polon. Sci. Sér. Sci. Math. Astronom. Phys, 13(6):397–403. 16 [46] LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. (1998). Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324. 9, 44 [47] Lee, J., Lee, Y., Kim, J., Kosiorek, A. R., Choi, S., and Teh, Y. W. (2019). Set transformer: A framework for attention-based permutation-invariant neural networks. In Proceedings of the 36th International Conference on Machine Learning. 23 [48] Lindley, D. V. (1956). On a Measure of the Information Provided by an Experiment. The Annals of Mathematical Statistics, 27(4):986–1005. 1, 2 [49] Lindley, D. V. (1972). Bayesian statistics. Society for Industrial and Applied Mathematics. tex.eprint: https://epubs.siam.org/doi/pdf/10.1137/1.9781611970654. 1, 2, 3, 4 [50] Lindley, D. V. (1982). Scoring rules and the inevitability of probability. International statistical review/revue internationale de statistique, pages 1–11. 2, 3 [51] Mohamed, S., Rosca, M., Figurnov, M., and Mnih, A. (2020). Monte Carlo Gradient Estimation in Machine Learning. Journal of Machine Learning Research, 21(132):1–62. 6 [52] Murphy, K. (2024). Reinforcement learning: an overview. arXiv preprint arXiv:2412.05265. 4 [53] Myung, J. I., Cavagnaro, D. R., and Pitt, M. A. (2013). A tutorial on adaptive design optimization. Journal of Mathematical Psychology, 57(3):53–67. 3 [54] Neal, R. M. (2011). Mcmc using hamiltonian dynamics. In Brooks, S., Gelman, A., Jones, G. L., and Meng, X.-L., editors, Handbook of Markov Chain Monte Carlo, pages 113–162. Chapman and Hall/CRC. 34 [55] Neiswanger, W., Yu, L., Zhao, S., Meng, C., and Ermon, S. (2022). Generalizing Bayesian optimization with decision-theoretic entropies. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS ’22, pages 21016–21029, Red Hook, NY, USA. Curran Associates Inc. 6 13
[56] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. (2019). Pytorch: An imperative style, high-performance deep learning library. In Wallach, H., Larochelle, H., Beygelzimer, A., d’Alché Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems 32 (NeurIPS), pages 8024–8035. 29 [57] Rainforth, T., Cornish, R., Yang, H., Warrington, A., and Wood, F. (2018). On Nesting Monte Carlo Estimators. In Proceedings of the 35th International Conference on Machine Learning, pages 4267–4276. PMLR. 2, 3 [58] Rainforth, T., Foster, A., Ivanova, D. R., and Smith, F. B. (2024). Modern Bayesian Experimental Design. Statistical Science, 39(1):100–114. 1, 2 [59] Savage, L. J. (1951). The theory of statistical decision, volume 46. Journal of the American Statistical Association. 1, 3 [60] Savage, L. J. (1971). Elicitation of personal probabilities and expectations. Journal of the American Statistical Association, 66(336):783–801. 2, 3 [61] Shannon, C. E. (1948). A Mathematical Theory of Communication. Bell System Technical Journal, 27(3):379–423. _eprint: https://onlinelibrary.wiley.com/doi/pdf/10.1002/j.15387305.1948.tb01338.x. 1 [62] Shen, W., Dong, J., and Huan, X. (2025). Variational sequential optimal experimental design using reinforcement learning. Computer Methods in Applied Mechanics and Engineering, 444:118068. 2, 3, 6, 7 [63] Shen, W. and Huan, X. (2023). Bayesian Sequential Optimal Experimental Design for Nonlinear Models Using Policy Gradient Reinforcement Learning. Computer Methods in Applied Mechanics and Engineering, 416:116304. arXiv:2110.15335 [cs]. 3, 7 [64] Sheng, X. and Hu, Y.-H. (2005). Maximum likelihood multiple-source localization using acoustic energy measurements with wireless sensor networks. IEEE Transactions on Signal Processing, 53(1):44–53. 7 [65] Smith, F. B., Kirsch, A., Farquhar, S., Gal, Y., Foster, A., and Rainforth, T. (2023). PredictionOriented Bayesian Active Learning. In Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, pages 7331–7348. PMLR. 6 [66] Williams, R. J. (1992). Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. Mach. Learn., 8(3-4):229–256. 6 [67] Zaheer, M., Kottur, S., Ravanbakhsh, S., Póczos, B., Salakhutdinov, R., and Smola, A. J. (2017). Deep sets. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems 30 (NeurIPS 2017), pages 3391–3401. Curran Associates, Inc. 23 [68] Zhang, Z., Dong, J., Liu, J., and Huan, X. (2025). Goal-oriented sequential bayesian experimental design for causal learning. 6 [69] Zhong, S., Shen, W., Catanach, T., and Huan, X. (2026). Goal-Oriented Bayesian Optimal Experimental Design for Nonlinear Models using Markov Chain Monte Carlo. SIAM/ASA Journal on Uncertainty Quantification, 14(1):19–47. arXiv:2403.18072 [stat]. 6
14
A
Lifting Bayes-Optimal Actions to Action Policies
In this section, we formalise the equivalence between the nested optimisation over Bayes-optimal downstream actions in Definition 1 and the joint optimisation over design and action policies used in Equation (5). The key point is that, under standard measurability and integrability assumptions, a pointwise optimisation over actions after observing the experimental history can be lifted to an optimisation over measurable action policies. Setup. Let (Ω, F, P) be the underlying probability space, and πd ∈ Πd be a fixed design policy. We assume that hT takes values in a measurable history space (H, H). We denote by ℓ the mapping from pairs (a, θ) ∈ A × Θ to its associated loss, ℓ : A × Θ → R. For a fixed history h ∈ H, we formally define the conditional risk ∀a ∈ A as, L(a, h) := E[ℓ(a, θ) | hT = h] . The Bayes-optimal action after observing h is any solution of aB (h) ∈ arg min L(a, h). a∈A
The corresponding Bayes expected future loss for a design policy πd is EFLB (π ) = E min L(a, h ) . d T p(hT ;πd ) ℓ a∈A
We now show that this nested formulation is equivalent to an optimisation over measurable action policies. Standing assumptions. We work under standard regularity assumptions ensuring that the pointwise Bayes-optimal actions can be selected measurably as functions of the observed history. The history space (H, H) is standard Borel, and the action space A is a compact metric space with Borel σ-algebra B(A). For every πd ∈ Πd , the conditional risk L(a, h) admits a jointly measurable version on A × H, and a 7→ L(a, h) is lower semicontinuous for every h ∈ H. We also assume the relevant expectations are well-defined, for example through an integrable envelope m(hT , θ) satisfying |ℓ(a, θ)| ≤ m(hT , θ) for all a ∈ A,
E[m(hT , θ)] < ∞.
Let
Πmeas := {πa : H → A : πa is H/B(A)-measurable} (11) a denote the set of admissible measurable action policies. The restriction to measurable action policies makes the optimisation probabilistically well-defined. Under the regularity assumptions above, the Bayes-action correspondence admits a measurable selector, so the pointwise Bayes action belongs to . It yields the desired equivalence, as stated in the following proposition. Πmeas a Proposition 1 (Lifting pointwise Bayes actions to action policies). Under the previous assumptions, for every fixed design policy πd ∈ Πd , Ep(hT ;πd ) min L(a, hT ) = min Ep(hT ;πd ) [L(πa (hT ), hT )] . (12) meas πa ∈Πa
a∈A
Equivalently, EFLB ℓ (πd ) =
min
πa ∈Πmeas a
Ep(hT ;πd )p(θ|hT ) [ℓ(πa (hT ), θ)] .
(13)
Proof. Fix πd ∈ Πd . For any measurable action policy πa ∈ Πmeas , we directly have pointwise a min L(a, h) ≤ L(πa (h), h). a∈A
Therefore, Ep(hT ;πd ) min L(a, hT ) ≤ Ep(hT ;πd ) [L(πa (hT ), hT )] . a∈A
Taking the infimum over πa ∈ Πmeas gives a Ep(hT ;πd ) min L(a, hT ) ≤ a∈A
inf
πa ∈Πmeas a
15
Ep(hT ;πd ) [L(πa (hT ), hT )] .
It remains to prove the reverse inequality. Since L(·, h) is lower semicontinuous on the compact action space A, the lower semicontinuous version of Weierstrass’ theorem implies that arg mina∈A L(a, h) is non-empty for every h ∈ H. Hence the argmin correspondence Γ(h) := arg min L(a, h) a∈A
is non-empty and compact-valued, as a closed subset of A. Since L is jointly measurable and lower semicontinuous in the action variable, the measurable maximum theorem from [1] implies that the value function h 7→ min L(a, h) a∈A
is measurable and that Γ is a measurable correspondence. By the Kuratowski–Ryll-Nardzewski measurable selection theorem from [45], there exists a measurable selector πaB : H → A such that πaB (h) ∈ Γ(h) Therefore, ∀h ∈ H
for all h ∈ H.
L(πaB (h), h) = min L(a, h), a∈A
and hence
B Ep(hT ;πd ) L(πa (hT ), hT ) = Ep(hT ;πd ) min L(a, hT ) . a∈A
Thus the infimum over measurable policies is attained by πaB , giving the desired equality. From Bayes-action policies to joint design–action optimisation. Having shown the equality (13), we can directly rewrite it as EFLB ℓ (πd ) =
min
πa ∈Πmeas a
Ep(θ)p(y1:T |θ;πd ) [ℓ(πa (hT ), θ)] .
(14)
Consequently, optimising the Bayes expected future loss over design policies can be written as πdB ∈ arg min EFLB ℓ (πd ) πd ∈Πd
= arg min
min
meas πd ∈Πd πa ∈Πa
Ep(θ)p(y1:T |θ;πd ) [ℓ(πa (hT ), θ)] ,
which ultimately leads to the core objective used by ACTION -BED as described in equation (5). Remark. The compactness and lower semicontinuity assumptions above are standard sufficient conditions ensuring that exact Bayes actions exist and admit measurable selectors, which is enough for the equivalence used in this work. When these assumptions fail, the result can be stated in a weaker but more general form using infima and measurable ε-optimal selectors. Specifically, if for every ε > 0 there exists a measurable policy πaε : H → A satisfying L(πaε (h), h) ≤ inf L(a, h) + ε, a∈A
then
Ep(hT ;πd )
inf L(a, hT ) =
a∈A
inf
πa ∈Πmeas a
Ep(hT ;πd ) [L(πa (hT ), hT )] .
This extension is useful when the action space is non-compact or when the conditional risk is not known to attain its pointwise infimum.
16
B
Equivalence Between BED Viewpoints
In the following Section we expand on the discussion on the equivalence between EUR and EFL perspectives of BED by providing a proof of Theorem 1. In section B.2 we additionally elaborate on the equivalence between ACTION -BED and existing variational approaches in a particular case. B.1
Proof of Theorem 1
We restate the theorem for completeness. Theorem 1 (Equivalence of EFL & EPU). Let s : Pθ × Θ → R be a proper scoring rule on Θ. Then πdB = arg min Ep(θ)p(y1:T |θ;πd ) s pθ (· | hT ), θ (15) πd ∈Πd
is a Bayes-optimal design policy for the belief model p(θ)p(y1:T |θ; πd ) and for any loss ℓ ∈ L, where n o L = ℓ : ℓ π̃a (qθ ), θ = s(qθ , θ), ∀qθ ∈ Pθ , θ ∈ Θ , π̃a (qθ ) = arg min Eθ∼qθ ℓ(a, θ) . (16) a∈A
Furthermore, the set L is never empty, and any lower-bounded loss function ℓ induces a corresponding scoring rule s that is proper, though not necessarily strictly proper. Proof. We start by noting that, for any ℓ ∈ L, Ep(θ)p(y1:T |θ;πd ) [s(pθ (· | y1:T ; πd ), θ)] = Ep(θ)p(y1:T |θ;πd ) [ℓ(π̃a (pθ (· | y1:T ; πd )) , θ)] = Ep(θ)p(y1:T |θ;πd ) ℓ πaB (y1:T , πd ), θ = Ep(y1:T ;πd ) min Ep(θ|y1:T ;πd ) [ℓ(a, θ)] . a∈A
(17) (18) (19)
The first equality follows from the definition of L, while the second follows from the definition of the Bayes action induced by the posterior belief pθ (· | y1:T ; πd ). Therefore, for any ℓ ∈ L, the minimizers of Equation (15) coincide with the Bayes-optimal experimental design policies associated with the downstream loss ℓ, as defined in Definition 1. Hence, πdB is a Bayes-optimal experimental design policy for ℓ under the belief model p(θ)p(y1:T | θ; πd ). Next, we show that for any ℓ, the corresponding scoring rule induced by Equation (16) is proper. By the definition of s, we have min Ep(θ|y1:T ;πd ) [s(qθ , θ)] = min Ep(θ|y1:T ;πd ) [ℓ(π̃a (qθ ), θ)] .
qθ ∈Pθ
qθ ∈Pθ
Since qθ influences the expected loss only through the choice of action π̃a (qθ ), it follows that min Ep(θ|y1:T ;πd ) [ℓ(π̃a (qθ ), θ)] ≥ min Ep(θ|y1:T ;πd ) [ℓ(a, θ)] .
qθ ∈Pθ
a∈A
Finally, using the definition of π̃a , the right-hand side is attained by choosing the true posterior belief qθ = pθ (· | y1:T ; πd ), so that min Ep(θ|y1:T ;πd ) [ℓ(a, θ)] = Ep(θ|y1:T ;πd ) [ℓ(π̃a (pθ (· | y1:T ; πd )) , θ)] a∈A
and therefore = Ep(θ|y1:T ;πd ) [s(pθ (· | y1:T ; πd ), θ)] . Thus, pθ (· | y1:T ; πd ) is a minimizer of Ep(θ|y1:T ;πd ) [s(qθ , θ)] , and so s satisfies the definition of a proper scoring rule on Θ. It need not be strictly proper as the minimiser is not necessarily unique. Finally, we show that L is never empty. Starting from Equation (15), we have πdB = arg min Ep(θ)p(y1:T |θ;πd ) [s(pθ (· | y1:T ; πd ), θ)] πd ∈Πd = arg min Ep(y1:T ;πd ) Ep(θ|y1:T ;πd ) [s(pθ (· | y1:T ; πd ), θ)] . πd ∈Πd
17
(20) (21)
Since s is a proper scoring rule, this is equal to πdB = arg min Ep(y1:T ;πd ) πd ∈Πd
min Ep(θ|y1:T ;πd ) [s(qθ , θ)] .
qθ ∈Pθ
(22)
This is the general form of a Bayes-optimal experimental design policy with action space A = Pθ and downstream loss ℓ(qθ , θ) = s(qθ , θ). Therefore, s is itself an element of L, and consequently L can never be empty. B.2
Recovering Variational Approaches
As noted in Section 3.3, the EIG itself is equivalent, up to a constant offset, to an EPU objective with the log score s(qθ (·), θ) = − log qθ (θ). That is, up to a constant offset, the EIG is equivalent to Ep(θ)p(y1:T |θ;πd ) [− log p(θ|hT )] . We can find an equivalent EFL formulation by choosing ℓ ∈ L according to Theorem 1. The canonical choice of loss function ℓ ∈ L for the log score is the log-loss on actions a(·) ∈ Pθ , that is ℓ(a(·), θ) = − log a(θ). The optimal action under this loss is a standard result: π̃a = arg min Eθ∼p(·|hT ) [− log a(θ)] = arg min KL(p(θ|hT )∥a(θ)) = p(θ|hT ), a(·)∈Pθ
a(·)∈Pθ
i.e. true Bayesian posterior. Therefore ACTION -BED solves the following joint optimisation problem: (πa , πd )B =
arg min
EFLlog-loss (πd , πa ),
P
(πa ,πd )∈Πa θ ×Πd
where
(23)
EFLlog-loss (πd , πa ) = Ep(θ)p(y1:T |θ;πd ) [− log(πa (hT ), θ)] ,
T T Pθ θ and ΠP a = {π | π : Y × Ξ → Pθ }. In practice, we must parameterise Πa , or a subset thereof, in order to pose a tractable parametric optimisation problem. We do this by specifying an amortised posterior approximation policy πaψ : Y T × ΞT → Pθ , defined as:
πaψ (hT ) = qψ (θ|hT ), where ψ ∈ Ψ parametrises an amortised variational family. Thus, ACTION -BED solves the following parametric optimisation: min Ep(θ)p(y1:T |θ;πϕ ) − log πaψ (hT ) = min Ep(θ)p(y1:T |θ;πϕ ) [− log qψ (θ|hT )] . ϕ,ψ
ϕ,ψ
d
d
This objective is equivalent to the variational posterior objectives introduced by Foster et al. [23, 24], who solve the following problem: max Ep(θ)p(y1:T |θ;πϕ ) [log qψ (θ|hT ) − log p(θ)] . ϕ,ψ
d
18
C
Reparameterisation of Simulator and Regularity Conditions
In this Section, we provide details of the reparameterisation function used in our EFL gradient expression Equation (9). We further provide the precise regularity conditions required to differentiate under the integral sign in the reparameterisation trick. We recall the application of the reparameterisation trick as follows: ∇ϕ,ψ EFL(ϕ, ψ) = ∇ϕ,ψ Ep(θ) p(y1:T |θ;πϕ ) ℓ πaψ (hT ), θ d h i = Ep(θ)q(ε1:T ) ∇ϕ,ψ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ ,
(24) (25)
where the simulator p(y1:T | θ; πdϕ ) was reparameterised in order to write hT = hT (θ, ε1:T ; πdϕ ) for ε1:T ∼ q(ε1:T ). We firstly provide the exact form of the reparameterisation function in section C.1 before providing regularity conditions on the loss and reparameterisation functions in section C.2 and consequently on the BED model in section C.3. C.1
Reparametrisation Functions
Recall the data generating process under a given design policy πdϕ is given by p(y1:T | θ; πdϕ ) = QT ϕ t=1 p(yt | θ, ξt = πd (ht−1 ), ht−1 ). Furthermore recall the data hT is defined as hT = {y1:T , ξ1:T } and designs ξt are recovered auto-regressively using the design policy and the previous data ξt = πdϕ (ht−1 ). We can therefore build the reparameterisation of hT recursively as follows. At the first step, we have h1 = {y1 , ξ1 } where y1 ∼ p(· | θ, ξ1 ), ξ1 = πdϕ (∅). Introducing noise variable ε1 ∼ q1 (·), we can directly write the pair (y1 , ξ1 ) as a function of (θ, ϵ1 ; πdϕ ), specifically: (y1 , ξ1 ) = g̃(θ, ϵ1 ; πdϕ ), πdϕ (∅) =: h1 (θ, ϵ1 ; πϕd ). At the second time step, we now have y2 ∼ p(· | θ, ξ2 , h1 ),
ξ2 = πdϕ (h1 ).
Again introducing noise variable ε2 ∼ q2 (·), we can express the pair (y2 , ξ2 ) as a function of (θ, ϵ2 , h1 ; πdϕ ) as follows: (y2 , ξ2 ) = g̃(θ, ϵ2 , h1 ; πdϕ ), πdϕ (h1 ) =: h̃2 (θ, ε2 , h1 ; πdϕ ) By substituting the reparameterisation of the previous step, i.e. h1 = h1 (θ, ϵ1 ; πϕd ) for ε1 ∼ q1 (·) and stacking with h1 , we obtain: h2 = (y1 , ξ1 , y2 , ξ2 ) = (h1 , y2 , ξ2 ) = (h1 , h̃2 ) =: h2 (θ, ϵ1:2 ; πdϕ ). We continue this argument recursively, arriving at a reparamterisation function for the entire experQT imental history hT = hT (θ, ϵ1:T ; πdϕ ) for ε1:T ∼ t=1 qt (·). Note that we may sometimes abuse notation slightly and write this as hT = hT (θ, ε1:T , ϕ), with the dependence on the policy πd implicit. C.2
Conditions for Differentiation Under the Integral Sign
Let (Ω, F, µ) be a measure space on which the random variables θ and ε are defined. Let Φ × Ψ ⊂ Rdϕ × Rdψ be an open set. Equation 25 is justified under the following conditions: 1. The distributions p(θ) and p(ε1:T ) are independent of (ϕ, ψ); 2. For µ-almost every (θ, ε1:T ) ∈ Ω, the mapping (ϕ, ψ) 7→ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ
is F-measurable and integrable: Z ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ dµ(θ, ε1:T ) < ∞. Ω
19
3. For µ-almost every (θ, ε1:T ) ∈ Ω, the mapping is differentiable in (ϕ, ψ) and the partial derivatives ∂ϕ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ , ∂ψ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ exist (which holds under standard smoothness assumptions on hT , πaψ , and ℓ); 4. For each compact K ⊂ Φ × Ψ, there exist integrable functions gϕ : Ω → [0, ∞), gψ : Ω → [0, ∞) such that for all (ϕ, ψ) ∈ K and µ-almost every (θ, ε1:T ) ∈ Ω, Z ϕ ψ ∂ϕ ℓ πa (hT (θ, ε1:T ; πd )), θ ≤ gϕ (ω), and gϕ (θ, ε1:T ) dµ(θ, ε1:T ) < ∞, (26) Ω
and ∂ψ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ
Z ≤ gψ (ω),
gψ (θ, ε1:T ) dµ(θ, ε1:T ) < ∞.
and Ω
(27)
These conditions follow from the differentiation-under-the-integral theorem. The following section explains how these necessary conditions for differentiating under the integral sign can be derived from assumptions on the model. C.3
Conditions on the BED model
Let Φ ⊂ Rdϕ and Ψ ⊂ Rdψ be open parameter sets. In order to have the map (ϕ, ψ) 7→ ℓ(πaψ (hT (θ, ε1:T ; πdϕ )), θ) differentiable, we have to satisfy the following conditions: 1. Differentiable design network: The design-induced history hT : Θ × E × Φ → YT × ΞT (θ, ε1:T , ϕ) 7→ hT (θ, ε1:T ; πdϕ ), is measurable in (θ, ε1:T ) and differentiable in ϕ for almost all (θ, ε1:T ), with ∇ϕ hT (θ, ε1:T , ϕ) existing and being finite. In particular, in all deep learning architectures considered to model hT , this is implemented as a composition of differentiable operations with respect to ϕ (e.g. affine maps and smooth or piecewise-smooth activations), so that it is differentiable almost everywhere. 2. Differentiable prediction network: The downstream prediction network πaψ : HT → Θ is measurable in its input and differentiable in (ψ, hT ), i.e. (ψ, h) 7→ πaψ (h) is differentiable for all (ψ, h), or at least almost everywhere, with Jacobians ∇ψ πaψ (h),
∇h πaψ (h)
existing and finite. This holds in particular when πaψ is a neural network with standard smooth (or piecewise-linear) activations such as tanh, softplus, ReLU or leaky − ReLU, which are differentiable almost everywhere. 3. Differentiable loss: The downstream loss ℓ : Θ × Θ → R, (a, θ) 7→ ℓ(a, θ), is differentiable in its first argument a for all (a, θ) (or almost everywhere), with gradient ∇a ℓ(a, θ) existing and finite (e.g. mean-squared error, cross-entropy, etc.). Under these conditions, for almost all (θ, ε1:T ), the composite map (ϕ, w) 7−→ ℓ πaψ (hT (θ, ε1:T , ϕ)), θ is differentiable, and its gradient can be computed by the chain rule. 20
D
Gradient Computations
D.1
Gradient Computations for Separated Networks
We denote, for a single draw (θ, ε1:T ) ∼ p(θ)q(ε1:T ), L(ϕ, ψ; θ, ε1:T ) := ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ , so that [ ∇ ϕ,ψ EFL(ϕ, ψ) = ∇ϕ,ψ L(ϕ, ψ; θ, ε1:T ). Gradient with respect to ψ. Since hT does not depend on ψ, we have ∇ψ L(ϕ, ψ; θ, ε1:T ) = ∇ψ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ ∇ψ πaψ hT (θ, ε1:T ; πdϕ ) , = ∇a ℓ a, θ ψ a=πa (hT )
(28)
(29) (30)
where ∇a ℓ(a, θ) denotes the gradient of the loss with respect to its first argument. Gradient with respect to ϕ. Here the dependence on ϕ is only through the history hT (θ, ε1:T ; πdϕ ), so the chain rule yields ∇ϕ L(ϕ, ψ; θ, ε1:T ) = ∇ϕ ℓ πaψ (hT (θ, ε1:T ; πdϕ )), θ (31) ϕ ϕ ψ = ∇a ℓ a, θ ∇hT πa hT (θ, ε1:T ; πd ) ∇ϕ hT (θ, ε1:T ; πd ). (32) ψ a=πa (hT )
Joint gradient estimator. Stacking the two components, a single-sample Monte Carlo estimator of the full gradient is thus ∇ϕ L(ϕ, ψ; θ, ε1:T ) [ ∇ . ϕ,ψ EFL(ϕ, ψ) = ∇ψ L(ϕ, ψ; θ, ε1:T ) In practice, automatic differentiation computes precisely these quantities by backpropagating through the computational graph (ϕ, ψ, θ, ε1:T ) 7→ hT (θ, ε1:T ; πdϕ ) 7→ πaψ (hT ) 7→ ℓ(πaψ (hT ), θ). D.2
Gradient Estimator in case of weight sharing
We now make explicit the form of the Monte Carlo gradient estimator in the case where both the acquisition and downstream policies share the same encoder network, with a sum-pooling aggregation step of the encodings and a specific prediction head for both. It is a common choice of architecture for policy-based methods. For a single draw (θ, ε1:T ) ∼ p(θ)q(ε1:T ), we consider the loss L(ϕ1 , ϕ2 , ψ; θ, ε1:T ) := ℓ πaψ (hT (θ, ε1:T ; πdϕ1 ,ϕ2 )), θ , where the downstream prediction network is defined via a shared DeepSets-style encoder: et = Eϕ1 (ξt , yt ) ∈ Rd , R(hT ) =
T X t=1
et =
T X
t = 1, . . . , T,
Eϕ1 (ξt , yt ),
(33) (34)
t=1
πaψ (hT ) = fψ R(hT ) .
(35)
A single-sample Monte Carlo estimator of the gradient of the expected future loss is then ∇\ ϕ1 ,ϕ2 ,ψ EFL(ϕ1 , ϕ2 , ψ) = ∇ϕ1 ,ϕ2 ,ψ L(ϕ1 , ϕ2 , ψ; θ, ε1:T ). Gradient with respect to the prediction head ψ. The dependence on ψ is only through a = fψ (R(hT )), so the chain rule yields ∇ψ L(ϕ1 , ϕ2 , ψ; θ, ε1:T ) = ∇ψ ℓ fψ (R(hT )), θ (36) = ∇a ℓ a, θ ∇ψ fψ R(hT ) . (37) a=fψ (R(hT ))
21
Gradient with respect to the shared encoder parameters ϕ1 . Here, ϕ1 affects the loss both through the encoder outputs et = Eϕ1 (ξt , yt ) and, through the design policy, through the generated history hT (θ, ε1:T ; πdϕ1 ,ϕ2 ). Applying the chain rule gives ∇ϕ1 L(ϕ1 , ϕ2 , ψ; θ, ε1:T ) = ∇a ℓ a, θ ∇R fψ R(hT ) ∇ϕ1 R(hT ). (38) a=fψ (R(hT ))
The term ∇ϕ1 R(hT ) includes both the direct dependence of the encoder Eϕ1 on ϕ1 and the indirect dependence through the history hT (θ, ε1:T ; πdϕ1 ,ϕ2 ). In particular, " T # X ∇ϕ1 R(hT ) = ∇ϕ1 Eϕ1 (ξt , yt ) , (39) t=1
where (ξt , yt )Tt=1 are generated by the design policy πdϕ1 ,ϕ2 . Gradient with respect to the design head ϕ2 . The parameters ϕ2 enter through the design policy that generates the history hT (θ, ε1:T ; πdϕ1 ,ϕ2 ). Denoting this dependence explicitly, the chain rule gives ∇ϕ2 L(ϕ1 , ϕ2 , ψ; θ, ε1:T ) = ∇hT ℓ fψ (R(hT )), θ ∇ϕ2 hT (θ, ε1:T ; πdϕ1 ,ϕ2 ), (40) where the term ∇hT ℓ(·) itself decomposes through the encoder and pooling: ∇hT ℓ fψ (R(hT )), θ = ∇a ℓ a, θ ∇R fψ R(hT ) ∇hT R(hT ). a=fψ (R(hT ))
(41)
In practice, automatic differentiation computes exactly these gradients by backpropagating through the computational graph (ϕ1 , ϕ2 , ψ, θ, ε1:T ) 7→ hT (θ, ε1:T ; πdϕ1 ,ϕ2 ) 7→ {Eϕ1 (ξt , yt )}Tt=1 7→ R(hT ) 7→ fψ (R(hT )) 7→ ℓ(·, θ). Finally, a Monte Carlo estimator of the full gradient is obtained by averaging these single-sample (i) gradients over i.i.d. draws (θ(i) , ε1:T ).
22
E
Network Architecture
Our co-optimization framework, introduced in Section 4, is agnostic to the specific parameterization of the design policy πdϕ and the downstream action policy πaψ . Any differentiable architecture can be used, provided that the resulting computation graph supports pathwise optimization through the reparameterized simulator. Nevertheless, architectural choices play an important role in practice, since they determine how the experimental history is represented and which inductive biases are made available to the policy. In particular, the architecture can encode exchangeability, temporal ordering, Markovian structure, spatial symmetries, or more general dependencies between past designs and outcomes. Design policy. The design policy πdϕ is the central component of the acquisition mechanism. At each experimental step, it maps the current history ht = {(ξs , ys )}ts=1 to the next design, or to a distribution over the designs, ( ξt+1 = πdϕ (ht ) ξt+1 ∼ pϕ (ξ | ht ) =: πdϕ (ht ) In deep adaptive design, this mapping is commonly parameterized through an encoder–aggregator– emitter architecture. First, each observed design–outcome pair (ξs , ys ) is embedded by a local encoder, rs = Eϕ1 (ξs , ys ), or, more generally, by separate design and outcome encoders followed by a fusion module. The resulting sequence or set of pair embeddings (r1 , . . . , rt ) is then compressed into a global history representation, zt = Aϕ2 (r1 , . . . , rt ), where Aϕ2 is a history aggregation module. Finally, an emitter network maps this history state to the next design, ξt+1 = Gϕ3 (zt ), or to a distribution over candidate designs in stochastic acquisition policies. The main decision-making capacity of the design policy is often concentrated in the aggregation module. The encoder produces local representations of individual experiments, whereas the aggregator determines how these pieces of evidence are combined to support the next acquisition decision. Different choices of Aϕ2 encode different assumptions about the structure of the history. For exchangeable probabilistic models of histories, permutation-invariant architectures such as DeepSets [22, 67] are natural: they aggregate pair embeddings through a symmetric pooling operation and therefore make the design decision independent of the ordering of observations. For temporally ordered or dynamically evolving systems, recurrent architectures such as LSTMs [29, 39] can encode the sequential dependence between past interventions and outcomes. Attention-based architectures, including Set Transformers [47] and Transformer Neural Process (TNP)-style encoders [3], provide a more flexible alternative by allowing the policy to learn interactions between all past design–outcome pairs and to condition the next design on the most informative parts of the history. Downstream action network. The downstream action policy maps the final experimental history to a terminal action, aT = πaψ (hT ). This action may represent a point estimate, a classifier output, a decision rule, or any task-specific prediction required after data acquisition. Since the downstream policy is only evaluated after the acquisition phase and needs to extract information contained in the history, πaψ can be implemented with a lighter architecture, or can follow the same encoder–aggregator–emitter structure as the design policy: rs = Eψ1 (ξs , ys ), zT = Aψ2 (r1 , . . . , rT ), aT = Gψ3 (zT ). This shared architectural template is especially useful when the downstream action depends on subtle interactions across the entire history. In such cases, using a rich history encoder for πaψ can improve the quality of the terminal decision and provide stronger gradients to guide the design policy during co-optimization. 23
Enhancing coupling through weight sharing. When the design policy and downstream action policy are built from compatible encoder–aggregator architectures, one can further strengthen their coupling by sharing part of their representations. For example, both policies may use a common local encoder Eω for design–outcome pairs, while retaining separate aggregation and output heads:
ξt+1 = Gϕd
rs = Eω (ξs , ys ), ψ Aϕd (r1 , . . . , rt ) , aT = Gψ a Aa (r1 , . . . , rT ) .
A stronger form of sharing can also be used, where the local encoder and part of the history aggregator are shared, while task-specific heads remain distinct: zt = Aω (Eω (ξ1 , y1 ), . . . , Eω (ξt , yt )), ξt+1 = Gϕd (zt ),
aT = Gψ a (zT ).
This form of weight sharing has several advantages. First, it encourages both the design policy and the downstream predictor to build compatible representations of the experimental history. The design policy is then trained to collect data that are useful under the same representational geometry used by the downstream action network. Second, sharing parameters strengthens the coupling between the two objectives: gradients from the downstream loss directly shape the representation used for acquisition. Third, it reduces the number of free parameters, which can improve statistical efficiency and regularize training, especially when the number of simulated trajectories is limited. E.1
TNP-Like Policy Architecture
In the experiments, we use a modular neural architecture for the adaptive design policy πdϕ . The architecture follows the general encoder–aggregator–emitter decomposition described above, with a Transformer-based module used to summarize the experimental history. Its purpose is to map the information collected so far into the next design decision, while remaining flexible enough to handle different experimental tasks. The architecture can be viewed as a TNP-like design policy where past design–outcome pairs play the role of context observations, and a learned decision token queries this context in order to produce the next design. We detail below the different modules composing the architecture used throughout the experiments. Experiment-specific dimensions and hyperparameters are denoted by variables, whereas numerical values indicate settings shared across all experiments. Encoding design–outcome pairs. Each past design ξt and outcome yt is first mapped to a latent representation. In practice, the design and outcome are encoded separately, eξt = fξ (ξt ),
eyt = fy (yt ),
Table 5: Design encoder architecture.
Table 6: Outcome encoder architecture.
Layer
Description
Dimension
Activation
Layer
Description
Dimension
Activation
Input H1 Output
ξ Fully connected Fully connected
dimξ 64 denc,ξ
– ReLU –
Input H1 Output
y Fully connected Fully connected
dimy 64 denc,y
– ReLU –
and then combined into a single representation of the experiment performed at step t: rt = fpair [eξt ; eyt ] .
Layer
Table 7: Fusion module architecture. Description Dimension Activation
Input H1 Output
[eξ ; ey ] Fully connected Fully connected
denc,y + denc,ξ 64 denc,r
24
– GELU –
This pair representation captures the local information revealed by one design–outcome interaction. Applying the same construction to all previous steps produces a sequence of history tokens Rt = (r1 , . . . , rt ). When the task is temporally structured, these tokens can also be augmented with time information, so that the policy can distinguish early and late observations. When the task is exchangeable, the same representation can instead be interpreted as an unordered set of observed experiments. Aggregating the experimental history. The sequence of pairwise representations (r1 , . . . , rt ) is processed by a Transformer-based history aggregator. Each representation is first projected to a common embedding dimension, then enriched with a learned type embedding indicating that it corresponds to a history token. When time embeddings are enabled, a learned temporal embedding is also added to preserve the order of the acquisition trajectory. A learned decision token is appended at the end of the sequence and is assigned its own type embedding, together with the time index corresponding to the next decision step. The resulting token sequence is normalized and passed through a stack of Transformer encoder layers. The final representation of the decision token is used as the summary of the experimental history. This token acts as a learnable query that attends to all previous design–outcome pairs and extracts the information relevant for choosing the next design. Denoting this representation by ht , the aggregation step can be written as ht = fhist (r1 , . . . , rt ), ht ∈ Rdh . This attention-based aggregation avoids imposing a fixed pooling rule over the experimental history. It allows the policy to model interactions between past observations, account for the temporal order of the acquisitions, and focus on the parts of the trajectory that are most informative for the next experimental decision. Table 8: History aggregator architecture. The module maps the sequence of pair representations (r1 , . . . , rt ) to a single history representation ht . Component
Configuration
Input sequence Input adapter History token embeddings Time embeddings Decision token Transformer encoder Readout
Output shape denc,r
Pair representations (r1 , . . . , rt ), ri ∈ R Linear layer denc,r → denc,h Add learned history type embeddings Add learned time embeddings, enabled Append learned query token with decision type embedding 2 layers, 4 heads, feedforward dimension 128 Final decision-token representation
t × denc,r t × denc,h t × denc,h t × denc,h (t + 1) × denc,h (t + 1) × denc,h denc,h
Emitting the next design. Finally, an emitter network maps the history representation to the next design: ξt+1 = femit (ht ).
Layer
Table 9: Emitter network architecture. Description Dimension Activation
Input H1 Output
History representation ht Fully connected Fully connected
denc,h 64 dimξ
– GELU –
Complete and modular design policy. The complete design policy is obtained by composing the modules. Given the experimental history Ht = {(ξi , yi )}ti=1 , the policy can be written as πϕ (Ht ) = femit ◦ fhist ◦ fpair ◦ (fξ , fy )({(ξi , yi )}ti=1 ). A key advantage of this architecture is its modularity. The task-specific encoders fξ and fy can be adapted to the structure of each experiment, while the fusion, history aggregation, and emission 25
Design encoder eξi = fξ (ξi ) Pair fusion ri = fpair (eξi , eyi )
Experimental history Ht = {(ξi , yi )}ti=1
r1 , . . . , r t
Outcome encoder eyi = fy (yi )
History aggregation Decision-token Transformer ht = fhist (r1 , . . . , rt )
ht
Emitter ξt+1 = femit (ht )
Next design ξt+1
Shared policy modules
Task-specific encoders
Figure 3: Modular architecture of the design policy. Past designs and outcomes are encoded separately, fused into pair representations, aggregated by a decision-token Transformer, and mapped by the emitter to the next design. modules are kept fixed. In low-dimensional settings, the encoders are simple multilayer perceptrons, whereas for structured or high-dimensional observations, fy can be replaced by a convolutional or patch-based encoder. This architecture therefore provides a common backbone across all experiments, both for our task-driven method and for the DAD/iDAD baselines. The same policy network can be trained end-to-end by backpropagating gradients either through the downstream task loss or through the sequential PCE objective. E.2
Downstream warm-up
In practice, before running the joint optimisation described in algorithm 1, we may initialise the action policy with a short downstream warm-up phase. During this phase, the design policy is kept fixed and is not used to generate designs. Instead, design trajectories are sampled from a task-relevant random distribution r(ξ1:T ), and only the downstream-action policy πaψ is trained to minimise the target downstream loss. Given samples θb ∼ p(θ), simulator noise εb1:T ∼ q(ε1:T ), and random b designs ξ1:T ∼ r(ξ1:T ), the warm-up gradient estimator is cψ Ldown (ψ) = 1 ∇ B
B X
b ∇ψ ℓ πaψ (hT (θb , εb1:T ; ξ1:T )), θb .
(42)
b=1
Although the resulting random designs are generally suboptimal, they provide an initial supervised signal for the action policy. Empirically, this warm-up can improve the subsequent joint optimisation: when the design policy starts being trained, the downstream network already maps histories to meaningful actions, yielding more informative gradients for the design policy and leading to a better coupling between acquisition and downstream prediction. Algorithm 2 ACTION -BED with downstream warm-up Require: Policies πdϕ , πaψ , initialisations ϕ0 , ψ0 Require: Warm-up steps Kwarm , joint steps K, random design distribution r(ξ1:T ) Downstream warm-up 1: for k = 0, 1, . . . , Kwarm − 1 do 1:B 2: θ1:B , ε1:B 1:T , ξ1:T ∼ p(θ)q(ε1:T )r(ξ1:T ) ψ cψ Ldown (ψk ) (eq. 42) 3: gk ← ∇ 4: ψk+1 ← U PDATEa (ψk , gkψ ) 5: end for 6: ψ0 ← ψKwarm Joint Action-BED optimisation 7: for k = 0, 1, . . . , K − 1 do 8: θ1:B , ε1:B 1:T ∼ p(θ)q(ε1:T ) ϕ ψ [ 9: gk , gk ← ∇ ϕ,ψ EFL(ϕk , ψk ) (eq. (10)) 10: ϕk+1 , ψk+1 ← U PDATE(ϕk , gkϕ , ψk , gkψ ) 11: end for 12: return ϕK , ψK
26
F
Bayes-optimal action under MSE and permutation-invariant MSE
We provide here a self-contained derivation of the Bayes-optimal action associated with the downstream decision problem considered in the source location finding experiment in Section 6.1. Since the general Bayes-action formulation is introduced in the main text, we only recall the minimal setup needed for the derivation. We then derive the posterior-mean optimality result for the standard squared error loss and extend it to the permutation-invariant MSE used in the experiment. F.1
Bayes-action setup
Let θ ∈ Θ denote the latent parameter and let hT ∈ HT denote the experimental trajectory induced by a fixed data-acquisition policy πdϕ . A downstream action policy is a measurable function πaψ : HT → A, implemented in practice by the downstream-action network. For a loss function ℓ : A × Θ → R+ , and a fixed trajectory hT , we call the Bayes-optimal action the minimizer of the posterior expected loss: a⋆ (hT ) = arg min Ep(θ|hT ) [ℓ(a, θ)] . a∈A
F.2
(43)
Standard MSE loss
We first specialize to the standard squared error loss. We assume that Θ = A ⊆ Rd , which is the case for the source location finding task, and we define ℓMSE (θ̂, θ) = ∥θ̂ − θ∥22 . For a fixed trajectory hT , the conditional risk is h i RMSE (θ̂ | hT ) = Ep(θ|hT ) ∥θ̂ − θ∥22 .
(44)
Expanding the squared norm gives h i RMSE (θ̂ | hT ) = E ∥θ̂∥22 − 2⟨θ̂, θ⟩ + ∥θ∥22 hT D E = ∥θ̂∥22 − 2 θ̂, E[θ | hT ] + E ∥θ∥22 | hT .
(45)
The last term does not depend on θ̂. Differentiating with respect to θ̂ and setting the gradient to zero yields ∇θ̂ RMSE (θ̂ | hT ) = 2θ̂ − 2E[θ | hT ] = 0. Hence, ⋆ θ̂MSE (hT ) = E[θ | hT ]. (46) Thus, under the standard MSE loss, the Bayes-optimal action is the posterior mean. F.3
Permutation-invariant MSE loss
However, in our setting the latent parameter is not naturally ordered. For example, in the source location finding problem with K exchangeable sources, the parameter can be written as θ = (θ1 , . . . , θK ) ∈ RK×d , but the ordering of the sources is arbitrary. In that case, penalizing an estimate θ̂ using the standard MSE would incorrectly distinguish between two estimates that differ only by a permutation of the sources. Let SK denote the symmetric group over K elements. For a permutation σ ∈ SK , write σ · θ = (θσ(1) , . . . , θσ(K) ). The permutation-invariant MSE is defined as K
2 1 X ℓPI-MSE (θ̂, θ) = min θ̂k − θσ(k) . σ∈SK K 2 k=1
27
(47)
Equivalently, identifying permutations with permutation matrices P ∈ PK , this can be written as ℓPI-MSE (θ̂, θ) =
2 1 min θ̂ − P θ , K P ∈PK F
(48)
where ∥ · ∥F denotes the Frobenius norm. For a fixed trajectory hT , the Bayes-optimal action under this loss is therefore 2 1 ⋆ θ̂PI-MSE . (hT ) = arg min Ep(θ|hT ) min θ̂ − P θ K P ∈PK F θ̂∈RK×d
(49)
This is the Fréchet mean of the posterior distribution on the quotient space RK×d /SK , where configurations that differ only by a permutation are identified. To make the optimality condition explicit, define, for a given estimate θ̂ and latent parameter θ, an optimal alignment 2
Pθ̂ (θ) ∈ arg min
P ∈PK
θ̂ − P θ
. F
(50)
Assuming that this optimal alignment is unique almost surely under p(θ | hT ), the conditional risk can be written locally as 2 1 . (51) RPI-MSE (θ̂ | hT ) = Ep(θ|hT ) θ̂ − Pθ̂ (θ)θ K F By the envelope theorem, differentiating the minimized objective with respect to θ̂ gives i h 2 ∇θ̂ RPI-MSE (θ̂ | hT ) = Ep(θ|hT ) θ̂ − Pθ̂ (θ)θ K 2 = θ̂ − Ep(θ|hT ) Pθ̂ (θ)θ . K Therefore, any differentiable Bayes-optimal action satisfies the fixed-point condition i h ⋆ (θ)θ . θ̂PI-MSE (hT ) = Ep(θ|hT ) Pθ̂⋆ (hT ) PI-MSE
(52)
(53)
Equation (53) shows that the Bayes-optimal action under permutation-invariant MSE is the posterior mean after optimally aligning each posterior sample to the estimate itself. The Bayes-optimal action is also not unique in the usual Euclidean sense: if θ̂⋆ is optimal, then P θ̂⋆ is also optimal for any permutation matrix P ∈ PK . Hence, the estimator is identifiable only up to permutation, which is precisely the desired invariance property. F.4
Interpretation for the downstream network
Under the standard MSE loss, minimizing the empirical Bayesian risk trains the downstream-action network to approximate the posterior mean mapping hT 7−→ E[θ | hT ]. Under the permutation-invariant MSE, the downstream network instead learns a representative of the posterior Fréchet mean on the quotient space induced by the permutation symmetry: 2 1 hT 7−→ arg min Ep(θ|hT ) min θ̂ − P θ . P ∈PK K F θ̂ Equivalently, it learns an unordered set-valued estimate of the latent sources. During training and evaluation, the optimal matching in Equation (47) removes the arbitrary label ordering of the sources, ensuring that two estimates that differ only by a permutation incur the same loss.
28
G
Experiment Details
The implementation was carried out in part using the PyTorch library [56], and the code is publicly available. Experiments were conducted on a server equipped with an NVIDIA Titan RTX GPU (24GB, Turing architecture) and 24 CPU cores. Each experiment was run on a single GPU. G.1
Source Location Finding
The source location finding problem is a canonical benchmark in BED, in which the objective is to infer the unknown spatial positions of K signal-emitting sources θ1 , . . . , θK ∈ R2 from a sequence of noisy measurements. Following the experimental setting of [22], at each design step t, an agent selects a sensor location ξt ∈ R2 , and observes a scalar response yt ∈ R corresponding to the aggregation of the signals emitted by the K sources, whose expected magnitudes decay with the squared distance to the sensor according to a known physical attenuation model. A commonly used formulation assumes yt = b +
K X k=1
A + σ εt , m + ∥ξt − θk ∥2
εt ∼ N (0, 1),
(54)
where A > 0 denotes the signal amplitude, b is a background offset, m > 0 is a stabilising constant preventing singularities at zero distance, and σ controls the measurement noise level. Equivalently, denoting the noiseless aggregated intensity by µ(θ, ξ) = b +
K X k=1
A , m + ∥ξ − θk ∥2
the observation model can be written in terms of noisy measurements of the log total intensity, yielding the prior–likelihood specification i.i.d. θk ∼ N (02 , I2 ), log y | θ, ξ ∼ N log µ(θ, ξ), σ , where θ = (θ1 , . . . , θK ) ∈ RK×2 . This formulation clearly shows that the model satisfies the reparameterization assumption C.1. The goal is to localize the sources with the highest possible accuracy: we seek to collect data from which a downstream model can be trained to identify all sources as effectively as possible. In our experimental setting, we consider the case K = 2. The task is challenging because information is spatially structured and highly non-uniform: sensor placements far from the source yield almost uninformative responses, while overly greedy exploitation may lead to premature localisation around suboptimal regions. From an inference perspective, the posterior distribution p(θ | y1:T , ξ1:T ) is typically multi-modal at early stages and becomes sharply concentrated as informative measurements accumulate. Thus, in this setting, amortised methods such as DAD aim to query maximally informative sensors and should generate sequences of designs that are effective for the prediction task. Parameter
Value
Number of sources, K 2 dimξ 2 dimy 1 T 30 Base signal, b 10−1 Max. signal, m 10−4 A 1 Signal noise, σ 0.5 Table 10: Location Finding Parameters Downstream Objective. In this experimental setting, the downstream objective is to predict the source locations. Prediction quality is evaluated using two losses: the mean squared error (MSE) and a logarithmic variant, denoted Log-MSE. For problems with multiple sources, the source labels become exchangeable, meaning that permuting their ordering leaves both the likelihood and the 29
signal observed by the sensor unchanged. Hence, the downstream task is intrinsically permutationinvariant. What matters is the predicted set of source locations, not the arbitrary ordering assigned to its elements. Consequently, the standard MSE loss is not the most relevant evaluation criterion in this setting, since it imposes a fixed correspondence between predicted and true sources and may penalize predictions that are correct up to a permutation. Therefore, we use in this experiment the permutation invariant MSE over the losses, defined as ℓM SE (θ, θ̂) = min
σ∈SK
K X
∥θ̂σ(k) − θk ∥2 ,
(55)
k=1
where SK is the set of permutations of {1, · · · , K} and θ̂ the point prediction of source locations. This optimization problem is significantly more complex than in the single-source case. The expectation involves a minimization over permutations inside the loss, leading to a highly non-linear and generally non-convex objective. This phenomenon poses significant challenges during training: • Label switching and multimodality. The posterior distribution is invariant under permutations of the sources, which typically results in a multimodal posterior with K! symmetric modes. • Failure of the posterior mean. Averaging over these modes leads to a posterior mean that does not correspond to any meaningful configuration of sources, often placing predicted sources between true locations or collapsing them toward the center of mass. As a result, under this permutation-invariant loss, the Bayes estimator is not the usual posterior mean on the labelled parameter space, but rather the posterior Fréchet mean induced by the optimalassignment metric. Equivalently, the Bayes-optimal action is a fixed point where each predicted source is equal to the posterior mean of the true sources optimally matched to it; see Appendix F.3. This loss is adopted because it directly reflects the objective of the experiment, namely recovering the unordered set of source locations. For the Log-MSE loss, however, we avoid this additional complication by training the downstream network to solve a harder labelled prediction problem, with the loss defined as ℓLog-MSE (C(θ), θ) = ||C(θ) − θ̂||2 + log ||C(θ) − θ̂||2 + ϵ (56) where ϵ = 10−6 is used for numerical stability, and C(·) denotes a canonicalisation operator that enforces a deterministic ordering of the sources. Specifically, for a multi-source configuration θ = (θ1 , . . . , θK ), we define C(θ) = θπ(1) , . . . , θπ(K) , π = argsort ∥θk ∥2 , k
so that sources are ordered by their Euclidean distance to the origin. This transformation maps each unordered set of sources to a canonical representative, thereby breaking the permutation symmetry and assigning a stable identity to each source. Consequently, the downstream network is trained to learn a well-defined mapping from trajectories {(ξi , yi )}Ti=1 to ordered source configurations. Experimental Details. In this experiment, we consider three baselines methods: 1. R ANDOM, where designs are naively sampled at random from the marginal prior distribution of a single source location. 2. DAD, the standard deep policy-based method introduced in [22]. It parameterises the adaptive design policy with a neural network and is trained using its native expectedinformation-gain objective. To ensure a fair comparison, we replace the architectural backbone of the original implementation with the same policy network used in our method. 3. ALINE, a recent policy-based method that achieves among the strongest reported performance in the literature. Its effectiveness is driven by a TNP-like policy architecture, a dense trajectory-level reward objective, and an internal posterior-approximation mechanism that amortises posterior inference. We retain the architecture proposed in the original paper, which is already close in spirit to the transformer-based policy network used in our method. To compare methods on this task, each baseline is first trained using its native design objective, derived from an EIG-based criterion, thereby learning an acquisition policy independently of the 30
downstream predictor. Once this policy is fixed, a downstream network is trained on trajectories generated by the learned policy to predict the source locations from the resulting design–observation history. All methods are then evaluated through the predictive accuracy of this downstream network, ensuring that the comparison reflects performance on the final source-localization task rather than on the surrogate design objective alone. The objective of this experiment is threefold. First, we evaluate our method from the perspective of downstream predictive performance. In particular, we assess whether the learned design policy produces informative trajectories that enable the downstream network to accurately solve the sourcelocation point-estimation problem. Second, we evaluate the quality of the learned design policy using the usual lower and upper bounds on the EIG, estimated respectively by sPCE and sNMC.Third, we compare methods in terms of computational cost, measured by the wall-clock training time on a single GPU. We distinguish between the cost of learning the design policy itself and the total training cost, which, for the baselines, also includes the additional training of the downstream prediction network. Architectures. We use the same policy-network architecture for ACTION -BED and DAD, detailed in Appendix E.1. At the start of each rollout, before any design–outcome observations are available, the initial history encoding is parameterised as a learnable vector. The only architectural hyperparameters varied across methods are the encoding dimensions, which are reported in Table 11. Table 11: Policy architecture encoding dimensions. Method DAD ACTION -BED
denc,ξ
denc,y
denc,r
denc,h
32 32
32 32
32 32
64 64
In this task, weight sharing between the design policy and the downstream network did not yield a clear performance improvement without further optimisation. We therefore use a simple fully connected MLP for the downstream network, which maps the design–outcome history {(ξi , yi )}Ti=1 to a point estimate of the source locations. The architecture is detailed in Table 12.
Layer
Table 12: Downstream network architecture. Description Dimension
Input Hidden layer 1 Hidden layer 2 Hidden layer 3 Output
hT Fully connected Fully connected Fully connected θ̂
T × (dimξ + dimy ) = 90 512 256 128 4
Activation – GELU GELU GELU –
Training procedure. During training, all policy networks are optimised from data generated online by rollouts of the corresponding design policy. At each iteration, latent parameters θ are sampled from the prior distribution p(θ), and the policy is rolled out autoregressively over the experimental horizon. The hyperparameters used for the policy networks are reported in Table 13. ALINE is implemented using its native architecture and trained with the hyperparameters reported in the original paper for the same source-location finding task. In particular, ALINE is trained for 100 000 iterations with a joint design policy and posterior approximation network, after an initial burning phase of 20 000 iterations during which only the posterior approximation network is optimised. DAD and Action-BED are implemented using the hyperparameters reported in the same table. For Action-BED, we additionally use a warm-up phase in which the downstream prediction network is trained in isolation to infer the source locations from randomly generated designs. After this warm-up, the design and downstream networks are optimised jointly using a shared optimizer. This pretraining stage provides a more informative learning signal to the design policy at the beginning of joint optimization, improves the coupling between the acquisition policy and the downstream predictor, and empirically leads to convergence toward better optima. Evaluation procedure. After training, we evaluate the downstream action networks associated with each learned design policy, either jointly trained with the acquisition policy in Action-BED, 31
Table 13: Hyperparameters for source location finding. Training and optimization hyperparameters used for the policy networks and the downstream prediction networks. ExpLR denotes an exponential learning-rate scheduler. The R ANDOM downstream network uses the same hyperparameters as DAD. ALINE
DAD Field
Value
Policy training Batch size Contrastive samples L Policy steps Policy seed
1024 2000 150K 42
Policy optimization Optimizer Adam betas Policy LR Scheduler Decay γ Decay period Weight decay Downstream training Downstream steps Downstream seed Downstream LR Downstream sched.
Adam (0.8, 0.998) 10−3 ExpLR 0.95 2000 0 150K 42 5×10−4 ExpLR
Field
Action-BED Value
Policy training Batch size Policy steps Burn-in steps Burn-in training Query pool, burn-in Query pool, train Query proposal Reward discount Policy weight α Policy seed
200 100K 20K Posterior only 30 200 Uniform 1 1 123
Policy optimization Optimizer Adam betas Policy LR Policy sched. Grad. clipping
Adam (0.9, 0.999) 10−3 Cosine ∥∇∥∞ ≤ 1
Downstream training Downstream steps Downstream LR Downstream sched. Downstream decay γ Decay period Downstream seed
Field
Value
Joint training Batch size 2000 Joint steps 150K Warm-up steps 50K Warm-up training Downstream only Warm-up designs Random Seed 42 Optimization Optimizer Adam betas Policy LR Scheduler Decay γ Decay period Weight decay
Adam (0.8, 0.998) 5×10−4 ExpLR 0.95 2000 0
100K 5×10−4 ExpLR 0.95 2000 42
or trained independently after the design policy has been learned for other baselines. We evaluate downstream predictors using the two downstream losses described above. Evaluation rollouts are generated by first sampling latent source parameters θ ∼ p(θ) from the prior and then rolling out the corresponding adaptive design policy autoregressively over the full experimental horizon. The downstream network receives the resulting terminal history hT and outputs a point prediction of the source configuration. All reported downstream losses are computed over 2048 independent evaluation rollouts, and uncertainty is reported as ±1 standard error in Table 1. EIG-bound evaluation. For each trained design policy, we additionally estimate lower and upper bounds on the EIG. This provides a standard, task-agnostic measure of design quality commonly used in the BED literature. It is not the downstream objective considered in our experiment, but DAD and ALINE are explicitly trained with information-theoretic objectives related to EIG maximisation. We recall the definition of these bounds: " # p(hT | θ0 ; πdϕ ) , (57) LsPCE = E log 1 PL ϕ ℓ=0 p(hT | θℓ ; πd ) L+1 " # p(hT | θ0 ; πdϕ ) UsNMC = E log 1 PL , (58) ϕ ℓ=1 p(hT | θ̃ℓ ; πd ) L with θ0 ∼ p(θ), θ1:L ∼ p(θ) are independent Monte Carlo samples used to approximate the marginal likelihood, and L being the number of contrastive samples drawn. For all methods, we use a large number of contrastive samples in both estimators so that the resulting bounds are sufficiently tight for comparison, see Table 14. Reported uncertainties correspond to ±1 standard error over independent rollouts. G.1.1
Qualitative Results
We complement the quantitative evaluation with qualitative trajectory plots of the learned design policies. We roll out each trained policy and visualize the sequence of selected designs over the experimental horizon. This comparison illustrates the characteristic behavior induced by each method: how quickly it explores the design space, whether it concentrates measurements around informative 32
Table 14: Source location finding. Evaluation protocol for design policies. Method
Training objective
Action-BED (MSE) Action-BED (Log) DAD ALINE Random
Downstream loss Downstream loss EIG-oriented objective EIG-oriented objective None
sPCE / sNMC
Downstream losses
L
Batch size
Batch size
2000 200 2000 500 2000
2000 2000 2000 2000 2000
6
10 108 106 5 × 106 105
Notes. L denotes the number of contrastive samples used to estimate the sequential PCE and NMC bounds. The corresponding batch size is the number of rollout histories used for these EIG-bound estimates. Downstream losses are evaluated on independent rollout batches.
(a) DAD
(b) ALINE
(c) ACTION -BED (MSE) (d) ACTION -BED (L OG )
Figure 4: Comparison of four policy rollouts. Each panel shows one experimental design trajectory under the trained design policy.
regions, and how its sampling strategy adapts as observations are collected. These trajectories can provide an interpretable diagnostic of the experimental strategies learned by the different policies. G.1.2
Amortising Downstream Bayes-action Selection
This section details the experimental protocol used to compare trained design policies on the downstream task, in terms of the quality of the Bayes-optimal action induced by the experimental trajectories they generate. This serves two purposes. First, it provides a direct way to compare design policies under the decision-theoretic objective of interest, namely the losses incurred by the action selected after observing the experimental history. Second, it illustrates a key practical advantage of ACTION -BED: since the design policy is trained jointly with an amortised downstream action policy, it directly learns an action network that predict a near-optimal decision, thereby amortising the otherwise costly posterior-inference step. Indeed, for design policies trained on the EIG, the downstream action-selection step is not directly available from training. Therefore, in order to evaluate such policies on downstream performance, one must proceed in two stages: first generate histories using the fixed trained design policy, and then infer or learn the corresponding downstream action from these histories. This is precisely the comparison performed here. In this experiment, we evaluate the posterior Fréchet mean under both the MSE and Log-MSE losses. For the MSE objective, this estimator coincides with the Bayes action; see Appendix F. This provides a useful explicit posterior-based decision rule against which amortised downstream prediction can be compared. For each trained design policy, we compute the posterior Fréchet mean from posterior samples conditioned on the observed trajectory, and report its empirical downstream performance over the evaluation batch under both loss functions. Inference protocol. To compute the explicit posterior Bayes action, we consider three posteriormean estimators. The first one is based on importance sampling with the prior as proposal distribution. For each experimental history hT = (ξ1:T , y1:T ) generated by a fixed design policy, we sample particles θ(1) , . . . , θ(N ) ∼ p(θ) independently from the prior distribution. These particles are then reweighted according to the likelihood of the observed history, (i)
(i)
w (hT ) ∝ p(hT | θ ),
N X i=1
33
w(i) (hT ) = 1.
This yields a weighted empirical approximation of the posterior distribution, N X
p(θ | hT ) ≈
w(i) (hT ) δθ(i) .
i=1
However, because the latent sources are exchangeable, the direct posterior mean is not Bayesoptimal for the permutation-invariant MSE loss considered and may yield a degenerate estimate of the underlying sources by averaging across permutations. To account for this symmetry, we first optimally permuting the sources of each particle. Given a reference ordering θ̄, for instance the (i) (i) current empirical estimate or an initial point estimate, each particle θ(i) = (θ1 , θ2 ) is aligned by solving 2 X 2 (i) σi⋆ ∈ arg min θσ(k) − θ̄k , σ∈S2
k=1
where SK denotes the set of permutations of the latent sources. We then define the aligned particle (i) (i) θe(i) = θσ⋆ (1) , θσ⋆ (2) , i
i
and estimate the MSE Bayes action by the weighted posterior Fréchet mean b aIS (hT ) =
N X
w(i) (hT ) θe(i) .
i=1
This aligned posterior mean is then used as the downstream prediction, and its empirical loss is evaluated on the held-out batch. The numerical parameters used for this importance-sampling estimator are reported in Table ??. The second estimator is based on Hamiltonian Monte Carlo (HMC) [54]. HMC constructs a Markov chain targeting the posterior distribution by simulating Hamiltonian dynamics in an augmented position–momentum space. For each history hT , we use HMC to draw approximate samples from the posterior distribution. In practice, the HMC chain is initialised from a coarse approximation of the optimal action obtained by a lighter version of the prior-proposal importance-sampling procedure described above. This provides a reasonable starting point near a high-posterior-probability region and stabilises the subsequent sampling stage. After burn-in and thinning according to the parameters reported in Table ??, we obtain posterior samples for each observed trajectory. The HMC-based Bayes action is then computed using the same post-processing procedure as for importance sampling. The empirical posterior mean is then computed over the aligned samples and used as the downstream MSE Bayes-action estimate. We finally evaluate the empirical downstream MSE incurred by this prediction over the evaluation batch. The third estimator is based on a Laplace approximation of the posterior distribution. For each history hT , we first compute a maximum-a-posteriori estimate of the latent parameters by optimising the log posterior, θbMAP (hT ) ∈ arg max {log p(hT | θ) + log p(θ)} . θ
In practice, the optimisation is initialised from a coarse particle-based approximation in order to start near a high-posterior-probability region, and multiple restarts are used to improve robustness. The posterior is then locally approximated by a Gaussian distribution centred at the MAP, with precision given by the negative Hessian of the log posterior at this point, −1 p(θ | hT ) ≈ N θbMAP (hT ), (H(hT ) + λridge I) , where H(hT ) denotes the local posterior precision and λridge is a small ridge regularisation parameter added for numerical stability. We then draw samples from this Gaussian approximation using a Cholesky decomposition whenever the regularised precision matrix is positive definite, as described in Table ??. The Laplace-based Bayes action is finally computed using the same post-processing procedure as for importance sampling and HMC: the Gaussian samples are aligned to account for the exchangeability of the latent sources, and the empirical posterior mean over the aligned samples is used as the downstream MSE Bayes-action estimate. We finally evaluate the empirical downstream MSE incurred by this prediction over the evaluation batch. 34
(a) Hamiltonian Monte Carlo Parameter Evaluation batch size HMC samples Burn-in steps Step size ϵ⋆ Leapfrog steps Thinning interval PM init. particles NPM
(b) Laplace approximation
Value
Parameter
2000 5 × 104 104 [10−3 , 10−2 ] 10 1 105
Evaluation batch size MAP optimisation steps Laplace samples Ridge regularisation λridge Factorisation
Value 2000 1000 5 × 104 10−3 Cholesky Decomposition
(c) Prior-proposal IS Parameter
Value
Particles N Evaluation batch size
5 × 106 2000
Table 15: Hyperparameters for posterior Bayes-action estimation. We report the parameters used for the two Bayes-action estimators: prior-proposal IS and HMC. The HMC step size ϵ is adapted separately for each method so as to match a target acceptance rate, while remaining constrained to the interval reported in the table.
Indicative inference cost. Since posterior-action estimators require test-time inference for each observed trajectory, we additionally report, in Table 17, the wall-clock time required to infer actions and evaluate the downstream loss on a batch of 2000 trajectories. These timings were obtained on the evaluation machine used for this posterior-inference study, which differs from the hardware used for training. They should therefore be interpreted as indicative deployment costs rather than hardware-normalised runtime benchmarks. Nevertheless, they highlight the qualitative distinction between amortised action prediction, which only requires a forward pass through the downstream network, and posterior-inference baselines, which require iterative or sampling-based computation at test time. Discussion. The results reported in Table 16 highlight an important distinction between the two downstream losses considered in this work. Although both losses evaluate the quality of a final action inferred from the experimental history, they encode different downstream objectives. Consequently, the corresponding Bayes actions are not the same, and neither are the design policies learned by ACTION -BED. This distinction is crucial when interpreting the posterior-action evaluations of fixed design policies. The experimental trajectories produced by the policy are shaped by the final decision criterion they are meant to serve. ACTION -BED trained with the Log-MSE loss, as shown in Table 1, achieves strong Log-MSE performance, and its learned designs often concentrate tightly around the latent sources. One might therefore expect that, under these highly localised designs, the MSE Bayes action obtained by posterior inference would also be especially accurate, perhaps even outperforming the policies trained with other objectives. However, this is not what we observe in practice. As illustrated in Figure 5, the distribution of rollout-wise MSE losses differs substantially between ACTION -BED (MSE) and ACTION -BED (L OG -MSE). Under the same evaluation conditions, ACTION -BED (L OG -MSE) achieves markedly lower losses for a large fraction of rollouts: in particular, its median MSE is substantially smaller, indicating that most of its downstream predictions are highly accurate. However, this improvement over the bulk of the distribution comes at the cost of a much heavier upper tail. For the largest quantiles, the MSE losses of ACTION -BED (L OG MSE) deteriorate sharply and become much larger than those obtained by ACTION -BED (MSE). Thus, although the Log-MSE-trained policy produces very accurate predictions on the majority of trajectories, it also induces rare rollouts on which the downstream action can fail severely. These high-loss trajectories dominate the average MSE and explain why the overall MSE performance of the Log-MSE-trained policy is worse despite its stronger median behaviour. ACTION -BED (L OG -MSE) learns to generate efficient experimental trajectories and an associated amortised action rule that are well adapted to the geometry of the Log-MSE objective: most predictions are brought extremely close to the true sources, but robustness to rare high-error trajectories is not prioritised in the same way as under the MSE objective. By contrast, the MSE-trained policy learns designs that may be less sharply accurate on the typical rollout but better control the upper 35
Table 16: Deployment-time Bayes-action estimation under fixed acquisition policies. For each pretrained design policy, we compare amortised action networks with posterior-based Bayes-action estimators. This evaluates how well each policy supports downstream decision-making after observing an experimental history. Uncertainties are reported as one standard error over 2048 rollouts. Acquisition
Decision Rule
Performance (± 1 s.e.) MSE (10−2 ) ↓
Log-MSE ↓
ACTION -BED (MSE)
Native network (MSE) Retrained network (MSE) Bayes-act. (IS) Bayes-act. (HMC) Bayes-act. (Laplace)
0.55 ± 0.06 0.86 ± 0.09 0.44 ± 0.03 0.61 ± 0.04 0.46 ± 0.03
-6.83 ± 0.03 -6.93 ± 0.01 -7.25 ± 0.03 -7.11 ± 0.02 -7.33 ± 0.04
ACTION -BED (Log)
Native network (Log) Retrained network (Log) Bayes-act. (IS) Bayes-act. (HMC) Bayes-act. (Laplace)
2.8 ± 0.4 1.7 ± 0.2 1.3 ± 0.6 2.00 ± 0.43 2.31 ± 0.71
-10.10 ± 0.03 -9.97 ± 0.03 -9.37 ± 0.04 -10.12 ± 0.04 -10.32 ± 0.05
DAD
Retrained network Bayes-act. (IS) Bayes-act. (HMC) Bayes-act. (Laplace)
5.6 ± 0.6 3.9 ± 0.2 3.6 ± 0.5 3.5 ± 0.8
-6.58 ± 0.05 -8.06 ± 0.05 -8.11 ± 0.05 -8.05 ± 0.06
ALINE
Retrained network Bayes-act. (IS) Bayes-act. (HMC) Bayes-act. (Laplace)
0.21 ± 0.03 0.16 ± 0.05 0.35 ± 0.21 0.13 ± 0.04
-8.11 ± 0.07 -9.22 ± 0.03 -9.17 ± 0.03 -9.74 ± 0.07
Random
Retrained network Bayes-act. (IS) Bayes-act. (HMC) Bayes-act. (Laplace)
22.4 ± 0.8 6.1 ± 0.7 15.7 ± 2.5 32.4 ± 3.1
-3.16 ± 0.04 -4.36 ± 0.05 -3.45 ± 0.08 -3.28 ± 0.09
Notes. Retrained networks are trained post-hoc on histories generated by the fixed design policy, whereas native networks for ACTION -BED are jointly trained on the downstream training loss. Bayes-action estimators use the fixed acquisition policy and compute the deployment-time action from the posterior induced by the observed history.
Table 17: Indicative test-time cost of posterior-action inference. We report the approximate wall-clock time required to infer actions and evaluate the downstream loss on a batch of 2000 trajectories. Approx. time (seconds) ↓
Action Estimator Amortised downstream network Importance-sampling posterior mean Laplace posterior approximation HMC posterior approximation
≈6 ≈ 7600 ≈ 250 ≈ 3100
Notes. Timings include action inference and loss evaluation. They are measured on a separate evaluation machine and are intended only to indicate the approximate order of magnitude of the deployment cost.
tail of the squared error distribution. This also explains why ACTION -BED trained with Log-MSE remains inferior under the MSE criterion, even when replacing its amortised action network by posterior-inference-based Bayes actions. G.1.3
Additional Results
Training time decomposition. All training-time comparisons are performed on the same GPU, so that reported differences reflect the computational cost of the methods rather than hardware variability. In addition to the main performance metrics, we report in Table 18 a decomposition of the training time into policy and downstream action-network training. This decomposition shows that DAD has 36
(a) ACTION -BED (MSE)
(b) ACTION -BED (L OG -MSE)
Figure 5: Distribution of rollout-wise MSE losses for loss-specialised ACTION -BED policies. We report per-rollout downstream MSE losses sorted from smallest to largest for ACTION -BED trained with the MSE objective and with the Log-MSE objective.
a slightly larger total training time than Action-BED, although its policy network alone is faster to train than the jointly trained Action-BED model. This is likely due to the batch sizes used in this experiment and to the relatively large neural architectures, for which backpropagation represents a substantial fraction of the total cost. Consequently, the computational advantage of the singly intractable objective is partly diluted by the cost of the network updates. ALINE is markedly more expensive: policy training alone already exceeds the full joint-training time of Action-BED, and its downstream stage remains computationally demanding because adaptive rollouts are costly to generate. Due to its simplicity, the random-design baseline is comparatively inexpensive to train. Evolution of downstream performance. Figure 6 reports an additional time-resolved comparison of downstream performance during policy training. For each method, we periodically freeze the design-policy weights and evaluate the downstream losses using the same evaluation protocol as above. For DAD and ALINE, this requires training a downstream action network from scratch at each frozen policy checkpoint, using the same architecture and training hyperparameters as reported in Tables 12 and 13. For Action-BED, we instead evaluate the downstream action network associated with the joint model at the corresponding checkpoint. Importantly, the horizontal axis only accounts for the elapsed policy-training time, or joint-training time for Action-BED, including the potential warm-up time. Thus, it excludes the additional cost of training independent downstream action networks for the baselines. Checkpoints are shown up to convergence. For the joint Action-BED models, both downstream losses are evaluated at every checkpoint, even when the downstream predictor was not explicitly optimised for the loss under consideration. This experiment illustrates how quickly each design policy produces histories that are useful for downstream prediction. Under the permutation-invariant MSE, Action-BED rapidly reaches a regime in which its action network performs well, achieving the best performance during approximately the first six hours of training. At later checkpoints, however, downstream networks trained on ALINE policies achieve stronger performance. Under the log-MSE loss, Action-BED trained with the logMSE objective yields substantial gains very early in training, outperforming the other methods by a large margin. Effect of Joint Training on Policy Co-Adaptation. We further examine whether joint training induces a useful co-adaptation between the design policy network and the downstream action policy network. To isolate this effect, we take a fully trained ACTION -BED design policy, keep it fixed, and retrain a new downstream action network under the same optimisation budget and hyperparameters as in the original joint training procedure; see Tables 12 and 13. We perform this retraining separately with the MSE and Log-MSE objectives. The relevant comparison is the same-loss setting: for each ACTION -BED model, we compare the native jointly trained action network to a post-hoc action network retrained with the same loss used to learn the design policy. As shown in Table 19, this retrained network performs slightly worse than the native joint network, despite being trained on histories generated by an already fully optimised design policy. This suggests that the gain of joint training is due to the optimisation path through which the design and action networks become mutually adapted. In particular, the design policy learns to generate histories that are informative in a form the action network can exploit, while the 37
Figure 6: Comparison of downstream losses as a function of design policy training time, until convergence. Table 18: Training Time Decomposition. Method
Design
Downstream
Total
ACTION -BED (MSE) ACTION -BED (Log)
– –
– –
1171 1184
DAD (MSE) DAD (Log)
879 879
335 337
1214 1216
ALINE (MSE) ALINE (Log)
1780 1780
1182 1190
2962 2970
0
31
31
Random
Notes. Wall-clock GPU times are reported in minutes. For ACTION -BED, design and downstream action networks are trained jointly, so only total training time is reported. Parentheses indicate the loss used to train the downstream action network.
action network simultaneously learns to decode both the observed outcomes and the information implicit in the adaptive sequence of designs. Table 19: Native versus retrained downstream networks for A-BED. Uncertainties are reported as one standard error. Acquisition
Action Network
Performance (± 1 s.e.) MSE (10−2 ) ↓
Log-MSE ↓
ACTION -BED (MSE)
Native network Retrained network
0.55 ± 0.06 0.86 ± 0.09
-6.83 ± 0.03 -6.93 ± 0.01
ACTION -BED (Log)
Native network Retrained network
2.8 ± 0.4 1.7 ± 0.2
−10.10 ± 0.03 -9.97 ± 0.03
Notes. Native networks are obtained directly from joint Action-BED training. Retrained networks are trained post-hoc after freezing the learned acquisition policy, using the same downstream loss. MSE is reported in units of 10−2 .
Cross-Evaluation of Design and Action Policies. We further study the coupling between acquisition and decision-making by cross-evaluating all trained downstream action networks across all design policies in the benchmark. For each trained design policy, we generate histories under that policy and evaluate the predictions produced by action networks trained with every other method with the MSE loss. For ACTION -BED, the action network is the one learned jointly with the design policy; for the other baselines, it is the network trained subsequently on trajectories from the fixed design policy. Table 20 reports normalised performance ratios, measuring, for each fixed design policy, the 38
degradation incurred when replacing its associated action network by an action network trained on another history distribution. This experiment tests whether action policies learn transferable Bayes-optimal action, or instead become specialised to the distribution of histories induced by their own design policy. Design policies that are more accurate and concentrate their measurements tightly around the latent sources induce history distributions on which action networks trained under other acquisition strategies perform substantially worse. In this regime, the downstream predictor must exploit fine-grained, policyspecific patterns in the adaptive design sequence, rather than only aggregate generic information from the observations. Conversely, the action networks paired with these more focused design policies achieve strong performance on their own histories but transfer poorly to histories generated by other policies. This suggests that high-performing methods do not merely collect more information; they also shape the information in a form that is specifically matched to their associated downstream decision rule. Table 20: Normalized cross-policy coupling between acquisition and action networks. Rows correspond to the acquisition policy used at test time, while columns correspond to the downstream action network used for decision-making. Test Acquisition
ACTION -BED (MSE) ACTION -BED (Log) DAD ALINE Random
Downstream Action Network A-BED (MSE)
A-BED (Log)
DAD
ALINE
Random
1.00 17.11 19.48 32.86 11.83
304.55 1.00 1079.98 18.10 1135.97
114.73 16.07 1.00 20.03 4.16
175.64 34.75 16.89 1.00 10.82
13.64 9.21 6.14 131.91 1.00
Notes. Each entry reports a normalized degradation factor. For a test acquisition policy πdi and an eij = R(π i , πaj )/ mink R(π i , π k ). action network πaj , we define R a d d
Impact of Warmstarting in Joint Training. We assess whether warmstarting the downstream action network improves joint ACTION -BED training. Table 21 shows that, for both the MSE and Log-MSE objectives, warmstarting leads to better final performance. This suggests that the benefit is not limited to a faster initial optimisation of the action network. Rather, by providing a more accurate downstream predictor early in training, warmstarting yields a more informative learning signal for the design policy. This improved signal changes the subsequent joint optimisation dynamics and guides the coupled design–action system toward a better final solution. Table 21: Impact of warmstarting in joint Action-BED training. Uncertainties are reported as one standard error. Acquisition
Warmstart
Performance (± 1 s.e.) −2
Log-MSE ↓
sPCE ↑
sNMC ↑
GPU (min) ↓
ACTION -BED (MSE)
Yes No
0.55 ± 0.06 0.87 ± 0.1
−6.83 ± 0.03 -6.82 ± 0.01
11.63 ± 0.04 11.62 ± 0.08
12.84 ± 0.08 12.67 ± 0.15
1171 1159
ACTION -BED (Log)
Yes No
2.8 ± 0.4 3.0 ± 0.3
−10.10 ± 0.03 -9.33 ± 0.03
17.14 ± 0.11 16.48 ± 0.18
21.07 ± 0.45 18.88 ± 0.49
1184 1172
MSE (10
G.2
)↓
Compute Time
Dynamical Systems
We consider sequential design problems in which the underlying simulator is a dynamical system. Unlike exchangeable observation models, each outcome depends on the current system state, so the order in which observations are collected matters. This setting is adapted from [38] and covers both the simple- and double-pendulum experiments considered in this work. More precisely, we view these experiments as an abstraction of dynamical systems governed by stochastic differential equations, detailed below, whose observation model is Markovian: p(yt+1 | yt , ξt+1 , θ), 39
so that the next observation is generated from the previous state or observation yt , the current design ξt+1 , and the physical parameters θ. Hence, the likelihood is sequential and cannot be factorized as with conditionally exchangeable observations. Under these assumptions, the joint distribution of a trajectory can be written as p(hT | θ; πdϕ ) =
T Y
p(yt | ht−1 , θ)
t=1
=
T Y
p(yt | yt−1 , ξt , θ).
t=1
Architectures. For both pendulum tasks, we use the same TNP-like policy architecture for ACTION BED and DAD, as described in Appendix E.1 and with the same encoding dimensions as in Table 11. Since the likelihood is non-exchangeable, we add a time embedding to each history element rt before aggregation. At the beginning of a rollout, when no design–outcome pair has yet been observed, the initial history representation is parameterized as a learnable vector. For both ACTION -BED and DAD, the emitter output is additionally mapped through an affine–tanh transformation, ξt = a tanh(ξ˜t ) + b, which ensures that the selected designs lie in the admissible control domain. The same architectural template is used for the stochastic and double pendulum experiments; only the input and output dimensions can vary, depending on the dimension of the designs and observations. ALINE is permutation-invariant by construction and is therefore not directly tailored to nonexchangeable likelihoods. To make the sequential structure explicit, we augment each context element with a time encoding and include the most recent observation in the context before encoding. This provides the model with the information needed to condition each candidate design on the current dynamical state. R ANDOM samples designs uniformly at random over the valid constrained space. In all cases, the downstream predictor is a fully connected MLP mapping the complete design– outcome history to a point estimate of θ in the relevant space (see Table 22). No parameter sharing is used between the design policy and the downstream network.
Layer
Table 22: Downstream network architecture. Description Dimension Activation
Input Hidden layer 1 Hidden layer 2 Hidden layer 3 Output
hT Fully connected Fully connected Fully connected θ̂
T × (dimξ + dimy ) 512 256 128 dθ
– GELU GELU GELU –
Training procedure. For the dynamical-system experiments, we follow a training protocol closely mirroring that of the source location finding task. All policies are trained from trajectories generated by the corresponding simulator. At each iteration, latent parameters θ are sampled from the prior, the design policy is rolled out autoregressively over the experimental horizon, and each method is trained on its own objective. We compare several downstream objectives in order to assess the accuracy and flexibility of the learned predictive action network. We consider the standard mean-squared error on θ, LMSE (θ̂, θ) = ∥θ̂ − θ∥22 , the logarithmic MSE, Llog −MSE (θ̂, θ) = log ∥θ̂ − θ∥22 + ε , and a weighted MSE, LwMSE (θ̂, θ) =
dθ X j=1
40
wj (θ̂j − θj )2 .
For the stochastic pendulum, we use weights w = (0.1, 1.0, 2.0), whereas for the double pendulum we use w = (1.0, 1.0, 2.0, 3.0). The weighted objective allows different components of θ to be emphasized, illustrating that the downstream target can be adapted to the scientific or operational quantity of interest. Evaluation procedure. In addition to downstream prediction performance, we evaluate each trained policy using sequential PCE and sequential NMC bounds on the EIG, which provide a common information-theoretic measure of design quality across methods. For both dynamical systems and for all downstream losses, prediction metrics are estimated from 2048 rollout batches and reported with ±1 standard error. For the single-pendulum experiment, the sequential PCE and NMC estimates are computed with L = 50K contrastive samples and a batch size of 2000. For the double-pendulum experiment, we use a larger contrastive budget, with L = 200K and a batch size of 1000. We also report the total GPU compute time required for training. G.2.1
Stochastic Pendulum
The first dynamical system considered is a stochastic single pendulum. The simulator describes the trajectory of a pendulum whose state evolves over time according to noisy dynamics. At each step, the design variable corresponds to an external control applied to the system, allowing the policy to influence the subsequent motion of the pendulum and thereby collect informative observations about the latent physical parameters. Task. We consider the problem of actively identifying the physical parameters of a stochastic controlled pendulum. The state is xt = (qt , q̇t ), where qt denotes the angle from the vertical and q̇t the angular velocity. The parameters of interest are the mass and length of the pendulum, (m, l), while the gravitational acceleration and damping coefficient are fixed to g = 9.81, d = 0.1. The initial state is fixed as x0 = (0, 0)⊤ . At each step of the discretized dynamics, with time step ∆t = 0.05, the policy selects a bounded control input ξt ∈ [−1, 1], corresponding to a torque applied to the pendulum. The observation is the state of the system at the next discretization step. The objective is to choose controls that generate trajectories informative about the latent physical parameters. Stochastic Dynamics. In continuous time, the controlled dynamics follow the stochastic differential equation dxt = h(xt , ξt )⊤ θ dt + L dβt , where βt is a Brownian motion, h(xt , ξt ) = (− sin qt , −q̇t , ξt ),
L = (0, 0.1)⊤ .
Equivalently, dqt = q̇t dt, dq̇t = h(xt , ξt )⊤ θ dt + 0.1 dβt . The observations used in the experiment are the discretized states yt = xt = (qt , q̇t ), recorded every ∆t = 0.05. The resulting observation process is sequential: each new observation depends on the previous state, the applied control, and the latent parameter θ. Following the conditionally linear formulation, the unknown parameter is ⊤ 3g 3d 3 θ= , 2, 2 , 2l ml ml which is bijectively related to the physical parameters (m, l). We place the Gaussian prior " # " #! 14.7 0.1 0 0 0 , 0 0.01 0 p(θ) = N . 3.0 0 0 0.1 The design variable ξt represents the external torque applied to the pendulum. 41
Table 23: Hyperparameters for stochastic single pendulum. Training and optimization hyperparameters used for the policy networks and the downstream prediction networks. ExpLR denotes an exponential learning-rate scheduler. The R ANDOM downstream network uses the same hyperparameters as DAD. ALINE
DAD Field
Value
Policy training Batch size Contrastive samples L Policy steps Policy seed
512 2000 50K 42
Policy optimization Optimizer Adam betas Policy LR Scheduler Decay γ Decay period Weight decay Downstream training Downstream steps Downstream seed Downstream LR Downstream sched.
Adam (0.8, 0.998) 10−4 ExpLR 0.96 400 0 50K 42 10−4 ExpLR
Field
Action-BED Value
Policy training Batch size Policy steps Burn-in steps Burn-in training Query pool, burn-in Query pool, train Query proposal Reward discount Policy weight α Policy seed
200 80K 20K Posterior only 50 200 Uniform 1 1 123
Policy optimization Optimizer Adam betas Policy LR Policy sched. Grad. clipping
Adam (0.9, 0.999) 10−3 Cosine ∥∇∥∞ ≤ 1
Downstream training Downstream steps Downstream LR Downstream sched. Downstream decay γ Decay period Downstream seed
Field
Value
Joint training Batch size 512 Joint steps 50K Warm-up steps 10K Warm-up training Downstream only Warm-up designs Random Seed 42 Optimization Optimizer Adam betas Policy LR Scheduler Decay γ Decay period Weight decay
Adam (0.8, 0.998) 10−4 ExpLR 0.96 400 0
50K 10−3 ExpLR 0.95 1000 42
Hyperparameters. The main training hyperparameters are reported in Table 23. ACTION -BED and DAD are trained for 50K gradient steps. For ACTION -BED, we additionally use a 10K-step warm-up phase in which the downstream network is trained on trajectories generated from random designs. ALINE is trained with a total budget of 100K steps, including a 20K-step burning phase during which only the posterior inference network is optimised. G.2.2
Stochastic Double Pendulum
We next consider a controlled stochastic double-pendulum system. Compared with the singlependulum experiment, this benchmark involves higher-dimensional, nonlinear, and coupled dynamics, making the generated trajectories substantially more sensitive to both the applied controls and the latent physical parameters. The state is q xt = t = (q1,t , q2,t , q̇1,t , q̇2,t )⊤ ∈ R4 , q̇t where q1 , q2 are the joint angles and q̇1 , q̇2 their angular velocities. The unknown parameter is θ = (m1 , m2 , l1 , l2 )⊤ , corresponding to the masses and lengths of the two links. As in the single-pendulum setting, observations are the T = 50 discretized states of the system, collected along a trajectory. The design variable is a pair of bounded torques ξt = (ξ1,t , ξ2,t )⊤ ,
ξ1,t ∈ [−4, 4],
ξ2,t ∈ [−2, 2],
applied to the two joints. The downstream task is to infer θ from the resulting controlled trajectory. Stochastic Dynamics. The double-pendulum dynamics follow the standard manipulator form M (q)q̈ + C(q, q̇)q̇ + τg (q) = ξ, where M (q) is the inertia matrix, C(q, q̇) contains the Coriolis and centrifugal terms, and τg (q) is the gravitational torque. We use g = 9.81 and add process noise to the acceleration dynamics. The resulting stochastic differential equation is dqt = q̇t dt, dq̇t = M (qt )−1 τg (qt ) + ξt − C(qt , q̇t )q̇t dt + L dβt , 42
Table 24: Hyperparameters for stochastic double pendulum. ExpLR denotes an exponential learning-rate scheduler. The R ANDOM downstream network uses the same hyperparameters as DAD. ALINE
DAD Field
Value
Policy training Batch size Contrastive samples L Policy steps Policy seed Policy optimization Optimizer Adam betas Policy LR Scheduler Decay γ Decay period Weight decay Downstream training Downstream steps Downstream seed Downstream LR Downstream sched.
512 2000 50K 42 Adam (0.8, 0.998) 10−4 ExpLR 0.96 400 0 50K 42 10−4 ExpLR
Field
Action-BED Value
Policy training Batch size Policy steps Burn-in steps Burn-in training Query pool, burn-in Query pool, train Query proposal Reward discount Policy weight α Policy seed
200 80K 20K Posterior only 50 150 Uniform 1 1 123
Policy optimization Optimizer Adam betas Policy LR Policy sched. Grad. clipping
Adam (0.9, 0.999) 10−3 Cosine ∥∇∥∞ ≤ 1
Downstream training Downstream steps Downstream LR Downstream sched. Downstream decay γ Decay period Downstream seed
Field
Value
Joint training Batch size 512 Joint steps 100K Warm-up steps 20K Warm-up training Downstream only Warm-up designs Random Seed 42 Optimization Optimizer Adam betas Policy LR Scheduler Decay γ Decay period Weight decay
Adam (0.8, 0.998) 5 × 10−4 ExpLR 0.96 400 0
100K 10−3 ExpLR 0.95 1000 42
where βt = (β1,t , β2,t )⊤ is a two-dimensional Brownian motion with independent components and 0.1 0 L= . 0 0.1 The SDE is discretized with the same Euler–Maruyama scheme as in the single-pendulum experiment, with ∆t = 0.05. The observation process is therefore Markovian and non-exchangeable: each new state depends on the previous state, the applied joint torques, and the latent physical parameters. We place a log-normal prior on the physical parameters, 0 0 p(θ) = LogNormal , 0.01 I4 , 0 0 so that masses and lengths remain positive. This task therefore tests whether adaptive design policies can select torque sequences that generate trajectories informative about coupled inertial, gravitational, and velocity-dependent effects. Hyperparameters. The main training hyperparameters are reported in Table 24. For this task, ACTION -BED is trained for 100K gradient steps, preceded by a 20K-step warm-up phase in which only the downstream action network is optimised. ALINE is trained under the same total budget of 100K steps, including a 20K-step burn-in phase for the posterior inference network. Finally, DAD and its associated downstream networks are trained for 50K steps, as no further improvement was observed beyond this point. In all cases, downstream action networks are trained until their predictive performance saturates. G.2.3
Additional Results
We report additional compute-time results for both dynamical-system tasks. Tables 25 and 26 decompose the total wall-clock GPU time into design-policy training and downstream-network training whenever these stages are trained separately. 43
Table 25: Single Pendulum. Training time decomposition. Method
Design
Downstream
Total
ACTION -BED (MSE) ACTION -BED (Log) ACTION -BED (W-MSE)
– – –
– – –
240 287 290
DAD (MSE) DAD (Log) DAD (W-MSE)
463 463 463
207 246 208
670 709 671
ALINE (MSE) ALINE (Log) ALINE (W-MSE)
1520 1520 1520
600 621 602
2120 2141 2122
0
19
19
R ANDOM
Table 26: Double Pendulum. Training time decomposition. Method
Design
Downstream
Total
ACTION -BED (MSE) ACTION -BED (Log) ACTION -BED (W-MSE)
– – –
– – –
501 504 504
DAD (MSE) DAD (Log) DAD (W-MSE)
862 862 862
72 78 71
934 940 933
ALINE (MSE) ALINE (Log) ALINE (W-MSE)
1938 1938 1938
1075 1124 1069
3013 3062 3007
0
58
58
R ANDOM
Notes. Wall-clock GPU times are reported in minutes. For ACTION -BED, the design and downstream action networks are trained jointly, so only total training time is reported. Parentheses indicate the loss used to train the downstream action network.
G.3
Masked MNIST Classification Experiment
We consider a sequential masked-classification task inspired by Iollo et al. [37] and based on the MNIST dataset [46]. The latent variable is an image θ ∈ R28×28 , and the target of interest is its class label z ∈ {0, . . . , 9}. The goal is to identify the digit class from a small number of partial and noisy observations of the image. At each step, the design ξ ∈ [1, 28]2 specifies a continuous spatial location in pixel coordinates, interpreted as the top-left corner of a local patch. The observation is a noisy 5 × 5 patch extracted around this location: y = Aξ θ + η, η ∼ N (0, σ 2 Id ), 2 where d = 5 , Aξ θ denotes the extracted patch, and zero padding is used whenever the patch extends outside the image domain. This illustrates that the observations are reparametrizable, which allows gradients to be propagated through the observation process. Conditionally on the image, the observation model is therefore explicit: p(y | θ, ξ) = N (Aξ θ, σ 2 Id ). A direct implementation of Aξ as a discrete masking operator would be non-differentiable with respect to ξ: small changes in the queried location would not affect the selected pixels until crossing an integer pixel boundary. To obtain a differentiable simulator, we instead use a continuous counterpart of the patch-extraction operator. In practice, the image is queried by grid sampling at the 5 × 5 patch coordinates induced by continuous ξ. When a sampling coordinate is not integer-valued, its intensity is obtained by bilinear interpolation from the four neighboring pixels. Thus, moving the design location smoothly changes the observed patch, making the simulator differentiable with respect to ξ while retaining the interpretation of localized pixel measurements. 44
Over a horizon T = 5, the policy adaptively selects patch locations ξ1:T based on the previously observed patches, with the objective of producing a final history informative for classification. This means that the task has a hierarchical structure. Observations are generated from the latent image θ, but the target of interest is the discrete label z. Consequently, the relevant class-conditional observation model is obtained by marginalizing over images within each class, Z p(y | z, ξ) = p(y | θ, ξ) p(θ | z) dθ. Θ
Although p(y | θ, ξ) is available in closed form, this class-conditional likelihood is not. In practice, the joint prior over images and labels is only available through the empirical dataset, N
p(θ, z) =
1 X δ(θ ,z ) (θ, z). N i=1 i i
The experiment therefore provides a benchmark for sequential design in a high-dimensional, hierarchical, and partially implicit setting. The design problem is challenging because only a small number of localized measurements are allowed, while the latent image is high-dimensional. Informative policies must therefore learn where to observe the image so as to reveal discriminative regions for the downstream classifier. This directly tests whether an acquisition strategy can align its designs with the final prediction task, rather than merely collecting information about the full latent image. G.3.1
Baselines, architectures, and downstream objective.
Protocol. We follow the same experimental protocol as in the previous experiments, with the main difference that the label-level observation model is implicit. As a result, EIG bounds are not reported for this task. Acquisition policies are trained using only images from the training set, while final performance is evaluated and reported on both the training and test sets. All downstream classifiers share the architecture reported in Table 27 and are trained with the cross-entropy loss LCE (w) = −E(z,θ),hT log πaψ (z | hT ) , where z denotes the class label and hT the final acquisition history. For each method, training is monitored on a validation split and stopped by early stopping once validation performance no longer improves. The resulting training budgets and selected hyperparameters are summarised in Tables 28 and 29. To avoid redundancy, these tables report only the hyperparameters that differ from those in Table 13; all omitted entries are kept unchanged. We report the final classification loss and accuracy on both the training and test sets. Table 27: MNIST masked classification downstream architecture. Module Pair encoder
History aggregation
Prediction head
Layer
Description
Dimension / Activation
Input H1 H2 H3 Output
(ξt , yt ) Fully connected Fully connected Fully connected Fully connected
dim(ξ) + dim(y) / – 512 / ReLU 256 / ReLU 128 / ReLU 32 / –
Input Pooling
{R(ξs , ys )}T s=1 Sum or mean pooling
T × 32 / – 32 / –
Input H1 H2 Output
Aggregated history Fully connected Fully connected Class logits
32 / – 128 / ReLU 64 / ReLU K/–
Notes. The downstream classifier uses the same encoder–aggregation structure as the acquisition network, with a larger pair encoder to handle the higher-dimensional masked-image observations. The output dimension K = 10 is the number of classes.
Architectures. The acquisition policies used by DAD, I DAD, and ACTION -BED rely on the same architectural backbone as in Appendix E. Apart from the task-specific dimensions of the design– observation pairs, the encoding dimensions are kept unchanged with respect to Table 11. The only 45
Table 28: MNIST masked classification hyperparameters: DAD and I DAD. Field Batch size Learning rate Policy steps Contrastive samples Downstream steps Downstream LR
DAD
iDAD
512 5 × 10−4 50K L = 500 50K 10−4
512 10−3 10K Exact over labels – –
Notes. For I DAD, no additional downstream classifier is trained: predictions are obtained directly from the critic logits. The contrastive denominator is computed exactly over the discrete label space. The R ANDOM downstream network uses the same hyperparameters as DAD.
Table 29: MNIST masked classification hyperparameters: ACTION -BED and ALINE. Field
ALINE
Action-BED
Batch size Number of queries Number of queries, burn-in Learning rate Total steps Burn-in steps Warm-up steps Joint steps
200 128 128 10−3 90K 10K – –
512 – – 5 × 10−4 10K – 0 10K
Notes. For ALINE, predictions are obtained directly from the classification head added to the posterior inference network.
additional modification is a sigmoid activation at the output of the emitter network, so that the proposed designs lie in the normalized domain [0, 1]2 before being reprojected onto the pixel grid. ACTION -BED jointly trains the acquisition policy and downstream classifier for 10K gradient steps using the downstream cross-entropy objective, with no warm-up phase. DAD is trained on sPCE objective, since it requires an explicit likelihood. This objective is defined with respect to the latent image θ rather than the class label z, therefore, its policy seeks designs that discriminate between images, rather than directly between classes. For both DAD and R ANDOM, a downstream classifier with the same architecture as in ACTION -BED is then trained separately for 50K gradient steps on trajectories generated by the corresponding policy. We also compare against I DAD, which adapts the contrastive InfoNCE objective to the implicit class-conditional setting. It jointly learns a design policy πdiDAD and a critic Uψ (z, hT ) measuring the compatibility between a candidate label and the observed history. Since the latent parameter space is discrete, the nested expectation appearing in the denominator of the InfoNCE objective can be accurately approximated by summing over all possible labels rather than relying on Monte Carlo negative samples. In practice, this amounts to replacing the usual contrastive denominator by a categorical normalization over the label space, " # exp{Uψ (z, hT )} LiDAD (ϕ, ψ) = −Ez,hT log P9 . ′ z ′ =1 exp{Uψ (z , hT )} For ALINE, we keep the native encoder–decoder architecture and modify only the target prediction head. Whereas the original implementation parameterises the posterior by a Gaussian mixture, the masked classification task requires a discrete posterior over labels. We therefore replace the Gaussian-mixture head by a classification head producing categorical logits. The target decoder outputs embeddings ztarget , with one target token per class. Each token embedding is passed through a small MLP to produce one scalar logit, ℓk = fcls (ztarget,k ), k = 1, . . . , 10, and hence a categorical posterior through a softmax. The posterior inference network is trained with the standard cross-entropy loss. This preserves ALINE’s native token-based architecture while adapting its output layer to the discrete label space. 46
(a) DAD
(b) I DAD
(c) ALINE
(d) ACTION -BED
Figure 7: Comparison of four policy rollouts. Each panel shows one experimental design trajectory under the trained design policy.
Table 30: Pair Encoder and Emitter Architectures. Module Pair encoder
Emitter
Layer
Description
Dimension / Activation
Input H1 Output
ξ, y Fully connected Fully connected
3/– 256 / ReLU 16 / –
Input H1 Output
R(ht ) Fully connected ξ
16 / – 2/– 2/–
Similarly, I DAD already includes a critic network which, in this setting, directly scores class labels as a function of the acquisition history. Its outputs can therefore be interpreted as class logits. Consequently, for both I DAD and ALINE, we do not train an additional downstream classifier: predictions are obtained directly from the critic logits for I DAD and from the classification head logits for ALINE. Evaluation. For MNIST masked classification, loss and accuracy are evaluated after training on the full training and test sets. Standard errors are computed across examples: for the cross-entropy loss, we report the empirical standard error of the per-example losses; for accuracy, we use the Bernoulli standard error associated with the empirical success rate. These uncertainty estimates therefore reflect finite-sample evaluation variability. Additional Figure 7 shows representative design trajectories for the different methods.
47
Table 31: Source Location Finding hyperparameters. Hyperparameter Batch size Contrastive samples L Policy steps Downstream steps Warm-up steps Optimizer Adam betas Learning rate LR decay γ LR decay period Weight decay Random seed
DAD
Action-BED
2000 2000 50K 150K – Adam (0.8, 0.998) 5 × 10−4 0.95 2000 0 42
2000 – 150K – 50K Adam (0.8, 0.998) 7 × 10−4 0.95 2000 0 42
Notes. Contrastive samples are used only for the sPCE-based DAD objective. For ACTION -BED, the acquisition policy and downstream action network are optimised with the downstream loss; warm-up steps are also reported.
G.4
Robustness to Network Architecture
This section reports additional results obtained under the same experimental protocols as in the main experiments, but using lighter neural architectures for both acquisition and downstream inference, and smaller computational budgets. The goal is not to improve upon the main quantitative results, but rather to assess whether the qualitative conclusions are robust to reduced model capacity. In particular, we investigate whether the relative behaviour of ACTION -BED and the information-theoretic baselines DAD and I DAD remains consistent. Across tasks, the experimental setting, simulators, evaluation metrics, and train/test protocols are kept identical to those used in the corresponding main experiments. As expected, this generally leads to weaker absolute performance compared with the main results. The precise architecture depends on the structure of the task. For exchangeable observation models, we use a DeepSets-style architecture in which each design–observation pair (ξs , ys ) is encoded independently and the resulting embeddings are aggregated by sum or mean pooling. For sequential and non-exchangeable simulators, such as the stochastic pendulum, the history is instead processed by a recurrent network, since the likelihood depends explicitly on the previous state. For each task, the downstream action policy πaψ is kept identical to that used in the corresponding main experiment. It is implemented as an MLP taking the final history hT as input and returning the estimate or action associated with the chosen downstream loss. G.4.1
Source Location Finding
Training. For the source-location experiment, DAD and ACTION -BED use the same dataacquisition architecture, reported in Table 30. This architecture is the one used in Foster et al. [22] for the corresponding DAD experiment, ensuring that differences in performance are not driven by changes in policy-network capacity. Each design–observation pair (ξ, y) is first processed by a pair encoder, which maps it to a fixed-dimensional embedding. The embeddings associated with the history ht are then aggregated by a permutation-invariant pooling operation, either sum or mean pooling, and passed to an emitter network that outputs the next design ξt+1 . Both the pair encoder and the emitter are fully connected MLPs. The R ANDOM baseline does not use a learned acquisition network: designs are sampled independently from the initial prior distribution p(θ). The corresponding training hyperparameters for the learned policies and downstream action networks are reported in Table 31. Results. The results are reported in Table 32. After training, all policies are evaluated using L = 100,000 contrastive samples and a batch size of 2000 rollouts for the sequential PCE and NMC estimates. Downstream losses are computed on independent batches of 2000 rollouts. Uncertainties are reported as ±1 standard error across rollout histories. We also report the total training time 48
for each method, including both acquisition-policy and downstream-action-network training when applicable. Table 32: Source Location Finding, lightweight architectures. Performance (± 1 s.e.)
Method MSE (10
−2
)↓
Compute
Log-MSE ↓
sPCE ↑
sNMC ↑
Time (min) ↓
Action-BED (MSE) Action-BED (Log)
2.6 ± 0.2 4.2 ± 0.4
−5.59 ± 0.03 −6.39 ± 0.04
9.14 ± 0.04 10.36 ± 0.04
9.95 ± 0.07 12.72 ± 0.03
97 102
DAD (MSE) DAD (Log)
13.6 ± 0.7 15.5 ± 1.2
−3.72 ± 0.04 −4.02 ± 0.04
10.38 ± 0.03 10.38 ± 0.03
12.91 ± 0.03 12.91 ± 0.03
283 285
Random
22.4 ± 0.8
−3.16 ± 0.04
8.27 ± 0.04
8.49 ± 0.05
31
Notes. MSE is reported in units of 10−2 . Uncertainties are reported as ±1 s.e. over 2048 evaluation rollouts. Compute time includes training both the design and downstream action networks.
G.4.2
Training Stability over Training Seeds
To assess the stability of ACTION -BED, we repeated training across multiple random seeds and evaluated the resulting policies under the same protocol. We found that performance remained consistent across runs, with only moderate variation in downstream loss and EIG estimates. Notably, the standard errors across independently trained policies are of a similar order of magnitude to those obtained from rollout-level uncertainty estimates, suggesting that training-seed variability is not substantially larger than evaluation noise. This suggests that the joint optimisation of the design and action policies is not overly sensitive to random initialisation or stochastic training trajectories. In particular, the observed gains over EIG-based baselines in terms of downstream losses are not driven by a single favourable run, but persist across independently trained policies. Table 33: Training stability over random seeds on Source Location Finding. MSE (10−2 ) ↓
Log-MSE ↓
sPCE ↑
sNMC ↑
ACTION -BED (MSE) ACTION -BED (Log)
2.7 ± 0.1 3.9 ± 0.3
−5.54 ± 0.02 −6.27 ± 0.02
9.04 ± 0.04 10.21 ± 0.07
9.79 ± 0.06 12.67 ± 0.09
DAD (MSE) DAD (Log)
14.1 ± 0.3 15.7 ± 0.9
−3.78 ± 0.05 −3.96 ± 0.06
10.36 ± 0.07 10.36 ± 0.07
12.79 ± 0.05 12.79 ± 0.05
Random
21.2 ± 0.4
−3.19 ± 0.08
8.21 ± 0.02
8.41 ± 0.02
Method
Notes. Uncertainties are reported as one standard error over independently trained policies with 15 different random seeds. For each trained policy, metrics are estimated using the same evaluation protocol.
49
Table 34: Acquisition architecture for the pendulum tasks. Module Pair encoder
History aggregator
Emitter
Layer
Description
Dimension / Activation
Input H1 H2 Output
(ξt , xt ) Fully connected Fully connected Fully connected
dim(ξ) + dim(x) / – 256 / ReLU 256 / ReLU 64 / –
Input H1 H2
{R(ξs , xs )}ts=0 LSTM LSTM
(t + 1) × 64 / – 64 / – 64 / –
Input H1 H2 Output
Final recurrent state Fully connected Fully connected Design ξt+1
64 / – 256 / ReLU 256 / ReLU dim(ξ) / tanh
Notes. The design dimension is dim(ξ) = 1 for the stochastic single pendulum and dim(ξ) = 2 for the double pendulum. The tanh output is subsequently rescaled to the task-specific control domain.
G.4.3
Dynamical Systems
Architecture. For both pendulum tasks, the acquisition network must account for the temporal dependence of the simulator. We therefore use the same sequential architecture for the stochastic single and double pendulum, inspired by Ivanova et al. [39]. Each design–state pair (ξt , xt ) is first embedded by an MLP encoder, and the resulting sequence of embeddings is processed by a two-layer LSTM. The final recurrent state is then passed through an emitter MLP, whose output dimension matches the design dimension of the task. The full acquisition architecture is summarised in Table 34. The downstream action network is kept identical to that used in the corresponding main experiments. It is implemented as an MLP taking the full trajectory hT as input and returning the parameter estimate associated with the chosen downstream loss. Unless otherwise stated, all optimisation hyperparameters are also inherited from the main experiments. The only change is the training budget: for the stochastic single pendulum, all policy, downstream, and joint networks are trained for 30K gradient steps; for the double pendulum, all corresponding networks are trained for 50K gradient steps. No downstream-only warm-up phase is used in either task. The results are reported in Tables 35 and 36. They are obtained using the same evaluation protocol as in the main experiments, with downstream losses computed on independent rollout batches and information-theoretic metrics estimated using sequential PCE and NMC bounds. Table 35: Stochastic Single Pendulum. Performance (± 1 s.e.)
Method −2
Log-MSE ↓
sPCE ↑
sNMC ↑
Time (min) ↓
Action-BED (MSE) Action-BED (Log)
1.0 ± 0.09 1.1 ± 0.11
−5.30 ± 0.03 −5.85 ± 0.04
3.35 ± 0.02 2.55 ± 0.01
3.40 ± 0.02 2.66 ± 0.01
45 47
DAD
1.7 ± 0.16
−5.34 ± 0.05
2.63 ± 0.12
2.65 ± 0.11
242
Random
2.90 ± 0.13
−4.15 ± 0.08
1.39 ± 0.07
1.42 ± 0.08
5
MSE (10
)↓
Compute
Notes. For DAD, downstream performance is reported only for the downstream network trained with the corresponding loss. Uncertainties are reported as ±1 s.e over 2048 rollout histories. Compute time includes training both the design and downstream action networks when applicable.
G.4.4
MNIST Masked Classification
For the MNIST masked-classification task, we use the same class of acquisition architectures as in the source-location experiment. In particular, the design policies are implemented with a DeepSets-style backbone: each design–observation pair is encoded independently, the resulting embeddings are aggregated over the acquisition history, and an emitter network outputs the next design. The precise architecture is reported in Table 37. 50
Table 36: Double Pendulum. Performance (± 1 s.e.)
Method MSE (10
−3
)↓
Compute
Log-MSE ↓
sPCE ↑
sNMC ↑
Time (min) ↓
Action-BED (MSE) Action-BED (Log)
0.7 ± 0.05 1.4 ± 0.08
−9.25 ± 0.02 −9.95 ± 0.04
10.40 ± 0.02 9.90 ± 0.03
14.50 ± 0.13 12.50 ± 0.07
276 274
DAD (MSE)
1.0 ± 0.04
−9.32 ± 0.03
10.75 ± 0.03
14.10 ± 0.14
582
Random
2.0 ± 0.09
−6.80 ± 0.05
8.26 ± 0.06
8.41 ± 0.08
58
Notes. For DAD, downstream performance is reported only for the downstream network trained with the corresponding loss.
Table 37: MNIST masked classification acquisition-policy architecture. Module
Layer
Description
Dimension / Activation
Input H1 H2 H3 Output
(ξt , yt ) Fully connected Fully connected Fully connected Fully connected
dim(ξ) + dim(y) / – 256 / ReLU 128 / ReLU 64 / ReLU 16 / –
History aggregation
Input Pooling
{R(ξs , ys )}ts=1 Sum pooling
t × 16 / – 16 / –
Emitter
Input Output
Aggregated history Design ξt+1
16 / – 2 / Sigmoid
Pair encoder
Notes. The design-policy backbone is a DeepSets-style net. A sigmoid activation is used to constrain designs to [0, 1]2 before reprojection onto the pixel grid.
The downstream classifier follows the same encoder–aggregation as summarised in Table 27. For I DAD, we use the same backbone for the design policy but the critic encoder is used as a downstream classifier as in the main experiment. Training hyperparameters for all methods are kept identical to those used in the corresponding main experiment. The corresponding classification results, including cross-entropy loss and accuracy on both training and test sets, are presented in Table 4. Table 38: MNIST masked classification. Performance (± 1 s.e.)
Method
Compute
Train CE ↓
Train Acc. (%) ↑
Test CE ↓
Test Acc. (%) ↑
Time (min) ↓
Action-BED
0.050 ± 0.002
98.24 ± 0.04
0.128 ± 0.006
96.55 ± 0.13
2.5
DAD iDAD
0.398 ± 0.004 0.113 ± 0.003
83.10 ± 0.03 96.28 ± 0.07
0.609 ± 0.010 0.175 ± 0.007
77.75 ± 0.12 94.76 ± 0.19
18 5
Random
0.786 ± 0.015
71.78 ± 0.87
0.797 ± 0.015
71.40 ± 0.86
5
Notes. CE denotes cross-entropy loss and Acc. denotes classification accuracy. Metrics are computed after training on the full training and test datasets. Uncertainties are reported as one standard error across examples. Compute time includes training the acquisition policy and, when applicable, the downstream classifier.
51