Bayesian experimental design (BED) is a principled framework for intelligent data acquisition with applications across a wide range of scientific and industrial domains. However, the development and adoption of BED remains hindered by a lack of standardised software and benchmarks. Existing methods are typically implemented in bespoke codebases and evaluated on a small subset of tasks, making methods difficult to apply, compare, or extend. To address this, we present TestBED, a general-purpose, open-source toolkit for Bayesian experimental design. TestBED decomposes the BED process into modular, reusable components spanning tasks, objectives, training mechanisms, policy architectures, and evaluation procedures, together with scalable implementations of core primitives such as expected information gain estimators. We provide a wide range of these components, which can be freely composed, allowing existing methods to be reproduced within a common framework, while also enabling previously unexplored configurations and new methodologies to be easily implemented and evaluated. Using TestBED, we further conduct a large-scale benchmark of contemporary BED methods, including controlled studies of under-appreciated but surprisingly critical methodological choices, such as the importance of time embeddings in design policy architectures. By providing a modular, extensible software foundation and a reproducible benchmarking suite, we hope that TestBED will accelerate methodological research in BED and lower the barrier to its real-world application.
Figures & tables
Figure 1: Left: TestBED workflow. A Task supplies the model and design space, a Designer produces a Rollout of designs and outcomes, and an Objective scores it. A fitting algorithm updates the designer’s parameters. Hooks augment the rollout loop. Right: a complete DAD-style training loop ( Foster et al., 2021 ) , with each line shaded by the component it uses (imports omitted).
Figure 2: TestBED evaluates every method on every task, and enables new methodological contributions. Columns show tasks (left) and design policy architectures (right). Grey: evaluated in prior work. Light green: an existing method, first evaluated in TestBED . Green ( ★ ): new methods, namely REINFORCE estimators extending ActionBED and JointVarPost to discrete outcomes.
Figure 4: Expanded view of the Rollout block of Figure 1 . Each timestep alternates designer and task calls, with hooks interleaved. All calls receive the current rollout state. After the final step and hook, the completed rollout is returned, including any extra fields produced by the hooks, e.g., actions.
Figure 5: Benchmark (mean ± standard error across five seeds) of BED methods under a shared policy architecture and wall-clock training budget. We report EIG upper bounds ( L=109,N=2,048 ), together with downstream MSE (negated so higher is better). The horizontal black line indicates performance of a random policy.
Attention
DeepSets
+ time
− time
+ time
− time
Method
Lower
Upper
Lower
Upper
Lower
Upper
Lower
Upper
DAD
14.334±0.223
14.359±0.228
5.654±1.061
5.654±1.061
15.423±0.265
15.587±0.307
12.051±0.028
12.052±0.028
RL-BOED
12.188±0.022
12.190±0.022
12.121±0.034
12.122±0.034
12.017±0.009
12.018±0.009
12.058±0.016
12.059±0.016
ActionBED
18.563±0.068
20.034±0.114
9.418±0.496
9.418±0.497
15.818±0.548
16.144±0.632
9.579±0.091
9.579±0.091
Table 1: Architecture ablation on the location finding task. Entries report the mean ± standard error across five seeds of the sPCE and sNMC bounds of the EIG ( L=109,N=2048 ).
L
Lower
Upper
Average number of gradient steps
256
12.568±0.084
12.571±0.085
28,355.8
512
12.896±0.111
12.898±0.111
28,778.2
1,024
13.683±0.184
13.691±0.186
27,448.0
2,048
13.982±0.125
13.992±0.127
24,954.6
4,096
14.615±0.175
14.637±0.179
20,571.2
8,192
14.500±0.288
14.521±0.291
15,369.8
Table 2: Effect of the number of contrastive samples L used for training DAD on location finding. Entries report mean ± standard error across five training seeds. All policies are evaluated using 2,048 outer samples and 109 contrastive samples.
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
Value
Gradient estimator
Pathwise (REINFORCE with leave-one-out baseline on CES, Temporal Discounting)
Contrastive samples, L
2000
Outer samples, B
2048 ( 1024 on Double Pendulum, Cart Pole due to memory)
Initial learning rate
10−4
Betas
(0.9,0.999)
Gradient clipping
1.0 (none on Temporal Discounting)
Appendix
Table 3: Training hyperparameters for DAD. Values apply to all eight benchmark tasks with exceptions explicitly called out.
Parameter
Value
Gradient estimator
Pathwise (REINFORCE with leave-one-out baseline on CES, Temporal Discounting)
Outer samples, B
2048
Posterior pretraining steps
0
Mixture components
10
Policy learning rate
10−5 , cosine decay to 5×10−6
Policy betas
(0.9,0.999)
Appendix
Table 4: Training hyperparameters for JointVarPost. Values apply to all eight benchmark tasks with exceptions explicitly called out. The policy and the posterior networks are optimised jointly from scratch, so there is no pretraining phase unlike variants within Phillips et al. (2026) .
Figure 6: The variational posterior network used by JointVarPost. Each history pair is embedded, a time embedding is added, and a learned CLS token is appended as an extra position; the resulting T+1 tokens are normalised and then pass through the pre-norm transformer blocks, each applying x↦x+MHA(LN(x)) and then x↦x+MLP(LN(x)) . The CLS position is read out and passed to the mixture head, whose three linear heads parametrise a m -component Gaussian mixture over unconstrained coordinates z , which a task-specific bijection pushes onto the support of θ . Note, each history token t carries the time embedding of t while the CLS token carries the embedding of index t , the position it occupies at the end of the history so that the readout knows how many experiments have been run. Additional details are provided in Table 5 .
Component
Value
Token encoder
MLP with widths [dX+dY,256,128] and GELU after each hidden layer (no activation on the output)
Time embedding
Learned, dimension 128
Transformer blocks
3 blocks, width 128 , 4 heads
Pooling
Learned CLS token
Trunk
Linear 128→128 followed by GELU
Output heads
Three linear heads from the trunk: component means ( mD ), diagonal log standard deviations ( mD ), mixing logits ( m )
Appendix
Table 5: Architecture of the variational posterior network used within the JointVarPost method. The encoder mirrors the shared design network of Section 5 ; the mixture readout replaces its design head. D denotes the dimension of the space the mixture is defined over, which is dΘ where θ is unconstrained and the dimension of the unconstrained coordinates otherwise.
Parameter
Value
RL algorithm
Soft Actor-Critic with twin critics and a learned temperature
DeepSets encoder and readout MLP with hidden widths [128,128] , embedding dimension 64 , as in the original implementation
Policy head
Gaussian squashed by tanh into the design box, with logσ∈[−20,0] and initial σ=3−1/2 before squashing
Reward
Dense (telescoping) sPCE reward with L=100000 contrastive samples
Episodes per iteration
100 (parallel environments)
Appendix
Table 6: Training setup and hyperparameters for RL-BOED. Values apply to all eight benchmark tasks with exceptions explicitly mentioned. The bounded design space and outcome normalisation range used for each task are given in Table 7 .
Task
Design box
Outcome range
Location Finding †
[−4,4]2
[−3,10]
Temporal Discounting †
D1∈[0.5,200] , R0∈[1,99]
[0,1]
CES †
[0.01,100]6
[0,1]
Pharmacokinetic
[0,24]
[−2,70]
Michaelis-Menten
[0,400]
[−10,210]
Stochastic Pendulum
[−1,1]
[−10,10]
Appendix
Table 7: Task-specific RL-BOED settings. Design box is the bounded design space used to train and evaluate RL-BOED ( Section D.5 ), and equals the task’s own design space except where marked † : on Location Finding the task is unbounded, on Temporal Discounting it is D1≥0 , R0∈[0,100] , and on CES it is [0,100]6 . Outcome range is the reference range used to rescale outcomes before they enter the actor and critics; outcomes outside it are not clipped.
Parameter
Value
Network
Shared transformer spine for design and inference networks; embedding dimension 128 , feedforward dimension 384 , 4 heads, 3 layers, no dropout, 10 -component Gaussian mixture posterior head
Warm-up iterations
20,000 (inference network only, with designs drawn uniformly from the design space)
Outer samples per iteration
200 rollouts
Candidate pool, training
200 designs drawn afresh for each rollout
Candidate pool, evaluation
2,000 designs, with the policy acting deterministically
Optimiser
AdamW, weight decay 0.01
Appendix
Table 8: Training hyperparameters for ALINE. Values apply to all eight benchmark tasks with exceptions explicitly mentioned. Network sizes are given in Table 10 , and the bounded design space and input rescaling used for each task in Table 9 .
Task
Design box
Design (c;s)
Outcome (c;s)
Location Finding †
[−4,4]2
—
—
Temporal Discounting †
[0.5,200]×[1,99]
(100.25,50);(99.75,49)
0.5;0.5
CES
[0,100]6
50;50
0.5;0.5
Pharmacokinetic
[0,24]
—
—
Michaelis-Menten
[0,400]
200;200
100;100
Stochastic Pendulum
[−1,1]
—
0;(3,8)
Appendix
Table 9: Task-specific ALINE settings. Design box is the bounded design space used to train and evaluate ALINE ( Section D.6 ), and equals the task’s own design space except where marked † ; as with RL-BOED, the Location Finding and Temporal Discounting tasks are modified. Designs and outcomes are rescaled elementwise as (x−c)/s before being passed to the network; the table gives the centre and scale as (c;s) . Outcomes are not clipped. A single value applies to every coordinate, and a dash denotes no rescaling.
Original
Benchmark
Embedding dimension
32
128
Feedforward dimension
128
384
Parameters (total)
97.5K
1.15M
Parameters (policy)
51.4K
647K
Appendix
Table 10: Network sizes of ALINE in the original paper, which we also use for the reproduction in Table 11 , and in our benchmark. Both use 4 attention heads, 3 transformer layers and a 10-component Gaussian mixture posterior head. The design policy comprises every parameter except the posterior head. Parameter counts are for location finding; the counts for CES are within 1% of these.
Location Finding ( T=30 )
Constant Elasticity of Substitution ( T=10 )
Original
8.91±0.04
13.50±0.15
Ours
8.92±0.02
13.50±0.06
Appendix
Table 11: EIG lower bound (sPCE) of ALINE on the location finding and CES tasks of Huang et al. (2025) . Original : numbers from Table 2 of Huang et al. (2025) . Ours : our reimplementation, trained on the same task specifications and network sizes.
Component
Value
Action network
MLP with widths [T(dX+dY),512,256,128,dΘ] and GELU after each hidden layer (no activation on the output)
Action net. parameters
175 k to 319 k depending on the task
Gradient estimator
Pathwise; REINFORCE with a leave-one-out baseline on Temporal Discounting
Action warm-up
25,000 steps training on designs sampled at random from the design space
Batch size
2,000 rollouts per step
Optimiser
Adam, learning rate 5×10−4 , β=(0.8,0.998)
Appendix
Table 12: Hyperparameters for ActionBED method, shared across all tasks unless stated otherwise.
Figure 7: The DeepSets design policy network of the architecture ablation ( Table 1 ), following the location finding network of Foster et al. (2021) . Each history pair is embedded by a shared per-step MLP and the embeddings are sum-pooled into a 16 -dimensional summary, which a decoder maps to an unconstrained design x~t+1 before the design space of the task maps it onto the feasible set ( Section A.1 ). An empty history is encoded by a learned null embedding, which is distinct from Foster et al. (2021) who fix it to zero. The ablation switches the time embedding (dashed) off, leaving every other component unchanged. When on, a learned 16 -dimensional embedding of s is concatenated onto history pair s before the per-step MLP, so it passes through the nonlinearity and reordering the history changes the pooled summary; when off, the encoder is permutation invariant. On location finding the network has 11.4 k trainable parameters with the time embedding and 5.2 k without.
Figure 8: The attentive design policy network used throughout the benchmark and in the architecture ablation ( Table 1 ). The same encoder is used by the posterior network of Figure 6 : each history pair is embedded, a time embedding is added, a learned CLS token is appended, and the t+1 tokens are normalised and passed through three pre-norm transformer blocks. The CLS position is read out and decoded into an unconstrained design x~t+1 , which the design space of the task maps onto the feasible set ( Section A.1 ). The ablation switches the time embedding (dashed) on and off, leaving every other component unchanged. When on, history token s carries the embedding of s and the CLS token the embedding of t ; when off, the encoder attends over an unordered set of tokens and is permutation invariant. On location finding the network has 663 k trainable parameters with the time embedding and 646 k without.
Location Finding
Pharmacokinetic
Method
Lower
Upper
Method
Lower
Upper
Random
8.300±0.019
8.300±0.019
Random
3.370±0.018
3.370±0.018
Static
8.867±0.014
8.867±0.013
Static
4.053±0.018
4.053±0.018
DAD
14.334±0.223
14.359±0.228
DAD
4.079±0.008
4.079±0.008
JointVarPost ( m=10 )
17.848±0.198
18.775±0.298
JointVarPost ( m=10 )
4.041±0.014
4.041±0.014
JointVarPost ( m=1 )
17.895±0.165
18.879±0.270
JointVarPost ( m=1 )
3.721±0.028
3.721±0.028
Appendix
Table 13: Lower and upper bounds on the EIG (mean ± standard error over 5 seeds), estimated using L=109,N=2048 . All policy-based methods are trained for a fixed 12-hour wall-clock budget, and use an attention-based design policy architecture.
Location Finding
Pharmacokinetic
Method
MSE
CMSE
Method
MSE
CMSE
Random
0.328±0.002
−2.522±0.012
Random
0.940±0.005
−3.961±0.013
Static
0.090±0.001
−5.151±0.001
Static
0.624±0.021
−4.967±0.060
DAD
0.068±0.015
−6.460±0.167
DAD
0.734±0.046
−4.882±0.070
JointVarPost ( m=10 )
0.030±0.001
−8.283±0.125
JointVarPost ( m=10 )
0.523±0.006
−5.311±0.019
JointVarPost ( m=1 )
0.021±0.003
−8.592±0.099
JointVarPost ( m=1 )
0.375±0.006
−5.415±0.021
Appendix
Table 14: Losses attained by fitting an MLP on top of the frozen designer (mean ± standard error over five seeds), fit by minimising the MSE between the output of the MLP and θ . CES and temporal discounting predict logθ instead of θ .
Variational family
Lower Bound
Upper Bound
Unconstrained
4.327±0.081
4.328±0.081
Transformed ( exp on k )
4.561±0.057
4.561±0.057
Appendix
Table 15: JointVarPost on Temporal Discounting with two variational families. Unconstrained places the Gaussian mixture directly on θ=(k,α)∈R2 . Transformed places it on (logk,α) and pushes it through an exponential map on k , leaving α unchanged. Both give valid Barber-Agakov bounds, but the transformed family puts no mass on k≤0 and so gives a tighter training signal. The rows differ only in the variational family: both use the gated DeepSets design network ( Table 16 ), 1,024 outer samples, a fixed 200 k training steps rather than a time budget, and are evaluated with N=2,048 and L=106 . Mean ± standard error over five seeds.
Component
Value
Token encoder
MLP on the unconstrained design x~s alone, widths [2,256,256,2×16] , Softplus after each hidden layer
Outcome gating
The two 16 -dimensional heads are mixed by the binary outcome, yshs(1)+(1−ys)hs(0)
Time embedding
None
Pooling
Sum over s≤t ; learned null encoding for the empty history
Emitter
MLP with widths [16,256,256,2] , Softplus after each hidden layer
Output
Unconstrained design (x~d,x~r) , mapped to (D1,R0) by the design space
Appendix
Table 16: Architecture of the gated DeepSets design network used in the variational posterior ablation ( Table 15 ). It follows the Temporal Discounting network of previous literature ( Foster et al., 2021 ; Hedman et al., 2025 ) ; the only difference is that the null encoding is learned rather than fixed to zero.
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.
Bayesian experimental design (BED) is a principled framework for data-efficient design of sequential experiments. However, existing BED methods are unable to adapt to dynamic constraints inherent in real-world tasks due to budget limitations, varying costs, or physical constraints that restrict how designs evolve over time. In this paper, we introduce a novel approach to BED that enables constrained optimization of experimental designs by combining offline pre-training of an amortized policy and a posterior network with online multi-step lookahead planning using scenario trees. We empirically demonstrate that our method yields substantially more informative design sequences than existing methods across a range of constrained BED tasks, while incurring only a modest additional computational overhead.
Yujia Guo, Daolang Huang, Xinyu Zhang +3
ELLIS Institute Finland · Department of Computer Science, Aalto University, Finland · Department of Computer Science, University of Manchester, UK
Standard Bayesian experimental design (BED) methods assume the experimenter alone will act on the data, choosing the best action given their own beliefs and goals. In practice, data are often acted on by external parties whose procedures the experimenter cannot change and may not know, particularly when the data are released publicly. To extend BED to such settings, we introduce usage-aware experimental design, a framework in which an adaptive design policy is trained to minimise the usage risk: the experimenter's expected loss when one or more external rules act on the collected data. These rules may be non-Bayesian or black-box, and when several may apply, we either average over them or take the worst case. We show that, unlike in standard BED, more informative data can be harmful under such rules, so the usage risk cannot in general be expressed as an expected uncertainty reduction; our framework thus strictly generalises standard BED. We then use it to tackle a range of problems that are difficult to address with existing approaches.