TestBED: An Open-Source Toolkit for Bayesian Experimental Design

Gavin Kerrigan; Marcel Hedman; Sahel Iqbal; Tom Rainforth

stat.MLNov 4, 2026

TestBED: An Open-Source Toolkit for Bayesian Experimental Design

Authors: Gavin Kerrigan, Marcel Hedman, Sahel Iqbal, Tom Rainforth

Organizations: Department of Statistics, University of Oxford

Abstract

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

Panel of Figure 1 — 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).

Panel of Figure 2 — 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 ( $\bigstar$ ): new methods, namely REINFORCE estimators extending ActionBED and JointVarPost to discrete outcomes.

Panel of Figure 4 — 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

\pm

standard error across five seeds) of BED methods under a shared policy architecture and wall-clock training budget. We report EIG upper bounds (

L=10^{9},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{\pm\scriptstyle 0.223}$	$14.359{\pm\scriptstyle 0.228}$	$5.654{\pm\scriptstyle 1.061}$	$5.654{\pm\scriptstyle 1.061}$	$15.423{\pm\scriptstyle 0.265}$	$15.587{\pm\scriptstyle 0.307}$	$12.051{\pm\scriptstyle 0.028}$	$12.052{\pm\scriptstyle 0.028}$
RL-BOED	$12.188{\pm\scriptstyle 0.022}$	$12.190{\pm\scriptstyle 0.022}$	$12.121{\pm\scriptstyle 0.034}$	$12.122{\pm\scriptstyle 0.034}$	$12.017{\pm\scriptstyle 0.009}$	$12.018{\pm\scriptstyle 0.009}$	$12.058{\pm\scriptstyle 0.016}$	$12.059{\pm\scriptstyle 0.016}$
ActionBED	$18.563{\pm\scriptstyle 0.068}$	$20.034{\pm\scriptstyle 0.114}$	$9.418{\pm\scriptstyle 0.496}$	$9.418{\pm\scriptstyle 0.497}$	$15.818{\pm\scriptstyle 0.548}$	$16.144{\pm\scriptstyle 0.632}$	$9.579{\pm\scriptstyle 0.091}$	$9.579{\pm\scriptstyle 0.091}$

Table 1: Architecture ablation on the location finding task. Entries report the mean

\pm

standard error across five seeds of the sPCE and sNMC bounds of the EIG (

L=10^{9},N=2\,048

$L$	Lower	Upper	Average number of gradient steps
256	$12.568\pm 0.084$	$12.571\pm 0.085$	$28{,}355.8$
512	$12.896\pm 0.111$	$12.898\pm 0.111$	$28{,}778.2$
1,024	$13.683\pm 0.184$	$13.691\pm 0.186$	$27{,}448.0$
2,048	$13.982\pm 0.125$	$13.992\pm 0.127$	$24{,}954.6$
4,096	$14.615\pm 0.175$	$14.637\pm 0.179$	$20{,}571.2$
8,192	$14.500\pm 0.288$	$14.521\pm 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

\pm

standard error across five training seeds. All policies are evaluated using

2{,}048

outer samples and

10^{9}

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$	$2\,000$
Outer samples, $B$	$2\,048$ ( $1\,024$ 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$	$2\,048$
Posterior pretraining steps	$0$
Mixture components	$10$
Policy learning rate	$10^{-5}$ , cosine decay to $5\times 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) .

Panel of Figure 6 — 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\mapsto x+\mathrm{MHA}(\mathrm{LN}(x))$ and then $x\mapsto x+\mathrm{MLP}(\mathrm{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 $\theta$ . 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 $[d_{\mathcal{X}}+d_{\mathcal{Y}},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\to 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_{\Theta}

where

\theta

is unconstrained and the dimension of the unconstrained coordinates otherwise.

Parameter	Value
RL algorithm	Soft Actor-Critic with twin critics and a learned temperature
Actor network	Shared attention architecture ( Appendix E ): encoder width $256$ , embedding dimension $128$ , $4$ heads, $3$ blocks
Critic networks	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\sigma\in[-20,0]$ and initial $\sigma=3^{-1/2}$ before squashing
Reward	Dense (telescoping) sPCE reward with $L=100\,000$ 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 †	$D_{1}\in[0.5,200]$ , $R_{0}\in[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

D_{1}\geq 0

R_{0}\in[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]\times[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\pm 0.04$	$13.50\pm 0.15$
Ours	$8.92\pm 0.02$	$13.50\pm 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(d_{\mathcal{X}}+d_{\mathcal{Y}}),512,256,128,d_{\Theta}]$ 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\times 10^{-4}$ , $\beta=(0.8,0.998)$

Appendix

Table 12: Hyperparameters for ActionBED method, shared across all tasks unless stated otherwise.

Panel of Figure 7 — 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 $\tilde{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.

Panel of Figure 8 — 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 $\tilde{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{\pm\scriptstyle 0.019}$	$8.300{\pm\scriptstyle 0.019}$	Random	$3.370{\pm\scriptstyle 0.018}$	$3.370{\pm\scriptstyle 0.018}$
Static	$8.867{\pm\scriptstyle 0.014}$	$8.867{\pm\scriptstyle 0.013}$	Static	$4.053{\pm\scriptstyle 0.018}$	$4.053{\pm\scriptstyle 0.018}$
DAD	$14.334{\pm\scriptstyle 0.223}$	$14.359{\pm\scriptstyle 0.228}$	DAD	$4.079{\pm\scriptstyle 0.008}$	$4.079{\pm\scriptstyle 0.008}$
JointVarPost ( $m=10$ )	$17.848{\pm\scriptstyle 0.198}$	$18.775{\pm\scriptstyle 0.298}$	JointVarPost ( $m=10$ )	$4.041{\pm\scriptstyle 0.014}$	$4.041{\pm\scriptstyle 0.014}$
JointVarPost ( $m=1$ )	$17.895{\pm\scriptstyle 0.165}$	$18.879{\pm\scriptstyle 0.270}$	JointVarPost ( $m=1$ )	$3.721{\pm\scriptstyle 0.028}$	$3.721{\pm\scriptstyle 0.028}$

Appendix

Table 13: Lower and upper bounds on the EIG (mean

\pm

standard error over 5 seeds), estimated using

L=10^{9},N=2\,048

. 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{\pm\scriptstyle 0.002}$	$-2.522{\pm\scriptstyle 0.012}$	Random	$0.940{\pm\scriptstyle 0.005}$	$-3.961{\pm\scriptstyle 0.013}$
Static	$0.090{\pm\scriptstyle 0.001}$	$-5.151{\pm\scriptstyle 0.001}$	Static	$0.624{\pm\scriptstyle 0.021}$	$-4.967{\pm\scriptstyle 0.060}$
DAD	$0.068{\pm\scriptstyle 0.015}$	$-6.460{\pm\scriptstyle 0.167}$	DAD	$0.734{\pm\scriptstyle 0.046}$	$-4.882{\pm\scriptstyle 0.070}$
JointVarPost ( $m=10$ )	$0.030{\pm\scriptstyle 0.001}$	$-8.283{\pm\scriptstyle 0.125}$	JointVarPost ( $m=10$ )	$0.523{\pm\scriptstyle 0.006}$	$-5.311{\pm\scriptstyle 0.019}$
JointVarPost ( $m=1$ )	$0.021{\pm\scriptstyle 0.003}$	$-8.592{\pm\scriptstyle 0.099}$	JointVarPost ( $m=1$ )	$0.375{\pm\scriptstyle 0.006}$	$-5.415{\pm\scriptstyle 0.021}$

Appendix

Table 14: Losses attained by fitting an MLP on top of the frozen designer (mean

\pm

standard error over five seeds), fit by minimising the MSE between the output of the MLP and

\theta

. CES and temporal discounting predict

\log\theta

instead of

\theta

Variational family	Lower Bound	Upper Bound
Unconstrained	$4.327{\pm\scriptstyle 0.081}$	$4.328{\pm\scriptstyle 0.081}$
Transformed ( $\exp$ on $k$ )	$4.561{\pm\scriptstyle 0.057}$	$4.561{\pm\scriptstyle 0.057}$

Appendix

Table 15: JointVarPost on Temporal Discounting with two variational families. Unconstrained places the Gaussian mixture directly on

\theta=(k,\alpha)\in\mathbb{R}^{2}

. Transformed places it on

(\log k,\alpha)

and pushes it through an exponential map on

k

, leaving

\alpha

unchanged. Both give valid Barber-Agakov bounds, but the transformed family puts no mass on

k\leq 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=10^{6}

. Mean

\pm

standard error over five seeds.

Component	Value
Token encoder	MLP on the unconstrained design $\tilde{x}_{s}$ alone, widths $[2,256,256,2\times 16]$ , Softplus after each hidden layer
Outcome gating	The two $16$ -dimensional heads are mixed by the binary outcome, $y_{s}h^{(1)}_{s}+(1-y_{s})h^{(0)}_{s}$
Time embedding	None
Pooling	Sum over $s\leq t$ ; learned null encoding for the empty history
Emitter	MLP with widths $[16,256,256,2]$ , Softplus after each hidden layer
Output	Unconstrained design $(\tilde{x}_{d},\tilde{x}_{r})$ , mapped to $(D_{1},R_{0})$ 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

Bayesian experimental design uses probabilistic models to choose measurements that improve learning or support better decisions.

See all 68 papers

Action-BED: Task-Driven Bayesian Experimental Design with Singly Intractable Objectives
Jun 22, 2026Tom Rossa, Angus Phillips, Tom RainforthBayesian Decision Theory Bayesian Experimental Design
Department of Statistics, University of Oxford
Constrained Bayesian Experimental Design via Online Planning
May 26, 2026Yujia Guo, Daolang Huang, Xinyu Zhang +3Bayesian Experimental Design Sequential Experimental Design
ELLIS Institute Finland · Department of Computer Science, Aalto University, Finland · Department of Computer Science, University of Manchester, UK
Designing Experiments for External Decision Makers
Oct 15, 2026Sahel Iqbal, Tom RainforthBayesian Experimental Design Decision-Making under Uncertainty
University of Oxford

Jun 22, 2026J/K move · Enter open · S save

Action-BED: Task-Driven Bayesian Experimental Design with Singly Intractable Objectives

Tom Rossa, Angus Phillips, Tom Rainforth

Department of Statistics, University of Oxford

Bayesian Decision Theory Bayesian Experimental Design Decision-Making under Uncertainty

Open paper AskPDF arXiv

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.

Abstract

Bayesian Experimental Design

Explore similar work

Action-BED: Task-Driven Bayesian Experimental Design with Singly Intractable Objectives

Constrained Bayesian Experimental Design via Online Planning

Designing Experiments for External Decision Makers