torch-testbed 0.1.0

Modular PyTorch building blocks for Bayesian experimental design.

Quickstart · Documentation · Citation

TestBED is a PyTorch library for developing, evaluating and applying Bayesian experimental design (BED) methods. Tasks, design objectives, designers and policy networks come as small, composable pieces, so you can write your own training loops on top of them instead of adopting a fixed pipeline. For fair benchmarking, TestBED also ships reproducible training recipes for established methods.

Installation

pip install torch-testbed

TestBED requires Python 3.10+ and PyTorch 2.5+, and the core library depends only on PyTorch. Two optional extras add more: train installs Hydra, which the training methods and checkpoint evaluators use, and full also adds plotting, Pyro and Weights & Biases support.

pip install "torch-testbed[train]"
pip install "torch-testbed[full]"

Quickstart

Here is how to design a source location-finding experiment, choosing each design by optimising the prior contrastive estimation (PCE) bound on the expected information gain:

from testbed.tasks import LocationFinding
from testbed.design_objectives import PCE
from testbed.designers import SequentialDesigner

task = LocationFinding(n_sources=1, dim_theta=2)
pce = PCE(task, n_outer=128, n_inner=128)
designer = SequentialDesigner(task, T=5, objective=pce, lr=5e-3, n_step_opt=200)

theta = task.sample_prior()
rollout = designer.rollout(theta)  # T designs and observations, plus a posterior estimate

The Location Finding tutorial walks through this example step by step and compares it with other designers.

What's included

Tasks. Widely used BED benchmark problems behind a shared interface, including location finding, CES, pharmacokinetics, Michaelis–Menten and temporal discounting. The interface covers continuous and discrete designs and outcomes, as well as likelihoods that depend on the full history.

Designers. PolicyDesigner, SequentialDesigner, StaticDesigner and RandomDesigner offer different strategies for choosing designs. You can customise or replace them as needed.

Design objectives and networks. Common bounds on and estimators of the expected information gain (PCE, sPCE, NMC, sNMC and Barber–Agakov), REINFORCE gradient estimators, and policy networks for amortised design.

Training methods. Reference implementations of state-of-the-art algorithms for learning amortised design policies, with Hydra recipes for most tasks. For example, training a policy with Action-BED on the temporal discounting task takes a single command:

python -m testbed.methods.actionbed --config-name experiments/temporal_discounting/actionbed

The following methods are currently available:

See Configuration recipes to learn how recipes are organised and how to override them.

Documentation

The full documentation lives at torch-testbed.readthedocs.io. It includes guides on using tasks, using designers and implementing custom methods, as well as a complete API reference.

Contributing

Contributions are very welcome! See CONTRIBUTING.md to set up a development install and open a pull request, and please read our Code of Conduct before taking part.

Citation

If you use TestBED in your research, please cite:

@software{testbed,
  title  = {{TestBED}: A {PyTorch} Library for {Bayesian} Experimental Design},
  author = {Kerrigan, Gavin and Hedman, Marcel and Iqbal, Sahel and Rainforth, Tom},
  year   = {2026},
  url    = {https://github.com/testbed-benchmark/testbed},
}