Constrained Bayesian Experimental Design via Online Planning

Abstract: 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.

• The paper introduces COPEx, which integrates offline amortized inference with online scenario tree optimization to handle dynamic constraints in experimental design.
• It models sequential experiments as a dynamic program over evolving constraint states, using lookahead optimization to maximize expected information gain.
• Experiments across diverse tasks show that COPEx outperforms state-of-the-art methods in efficiency, accuracy, and cost-awareness.

Constrained Bayesian Experimental Design via Online Planning

Motivation and Problem Statement

Bayesian Experimental Design (BED) harnesses information-theoretic criteria (most notably, expected information gain—EIG) for sequential experiment selection under model uncertainty. However, conventional BED approaches are agnostic to dynamic real-world constraints, such as variable costs, operational budgets, or physical feasibility restrictions. These constraints materially alter the design space and dynamics—as evidenced in scenarios like mobile sensing, drug discovery, and behavioral experiments—driving the need for methods that explicitly model adaptive feasibility during deployment.

Figure 1: COPEx produces more informative posterior concentration compared to constraint-naive amortized policies when consecutive designs are restricted by transition constraints.

COPEx: Semi-Amortized Online Planning Approach

The paper introduces COPEx, a constrained online planning framework for BED based on multi-step online planning over scenario trees, leveraging offline pretraining of amortized inference and policy modules. The approach tackles the unsolved challenge of non-myopic constrained experiment design: the utility of each experiment depends on subsequent feasible paths, and constraints can be dynamic or design-dependent.

COPEx consists of two foundational amortized modules:

• An inference network that maps histories to fast, differentiable posterior distributions (Mixture Density Network / Transformer Neural Process).
• A design policy, pretrained on unconstrained BED, used to initialize scenario tree optimization and bias search toward informative regions.

The online component constructs a finite-horizon scenario tree (depth $H$), simulates fantasy outcomes using the amortized posterior to enable rapid belief updates and sample downstream experiment effects, and jointly optimizes all decision variables in the tree under evolving constraint states.

Figure 2: COPEx architecture, combining amortized posterior network and design policy for efficient multi-step scenario tree planning.

Methodological Contributions

COPEx formalizes constrained BED as a dynamic program over both experiment history and constraint state, where admissible designs at each step are determined by evolving feasibility and budget. The policy is computed via lookahead optimization over scenario trees, efficiently simulated with amortized inference. The approach achieves the following:

• Handles arbitrary constraint forms, including local transition limits and global budgets, by recursing constraint state transitions in the scenario tree.
• Produces deterministic policies adaptable at test time, unlike prior amortized approaches which require retraining for each constraint configuration.
• Enables fast, differentiable EIG surrogates using adaptive contrastive objectives with the amortized posterior, replacing computationally expensive exact methods.
• Supports hybrid initialization: warm-starting optimization with the amortized policy, augmented by random exploration to mitigate policy shift under heavy constraints.
• Efficient vectorized evaluation and reparameterization trick allow scalable optimization over large scenario trees, despite the exponential growth in branching factor.

Experimental Validation

COPEx is comprehensively benchmarked against state-of-the-art BED and active learning baselines on three tasks: location finding under transition constraints, behavioral utility elicitation with budgeted change, and cost-aware active learning across complex landscapes.

Location Finding (Transition-Constrained Exploration):

COPEx demonstrates superior posterior concentration and cumulative EIG under strict step-size limits, outperforming both variational and policy-based baselines. The benefit compounds as feasibility narrows, and the computational overhead remains modest for moderate horizons. Warm-start from the amortized policy sharpens performance compared to random tree initialization.

CES (Budgeted Behavioral Elicitation):

Under global change budgets to avoid cognitive overload, COPEx with $H=1$ achieves the highest information gain prior to budget exhaustion, outstripping amortized and reinforcement learning baselines. Deep planning (large $H$) can amplify bias from the amortized surrogate, highlighting the error propagation tradeoff.

Cost-Aware Active Learning:

COPEx shows consistently lower RMSE per unit cost than GP-based acquisition policies on diverse function/cost pairs—demonstrating robust adaptation to spatially varying costs and constraints. Its non-myopic variants exploit future budget/constraint structure more effectively, yielding improved convergence especially in low-budget regimes.

Figure 3: COPEx achieves lower error than GP baselines for the same exploration cost across benchmark functions and heterogeneous cost landscapes.

Figure 4: COPEx and GP baselines under transition and budget constraints. Non-myopic variants yield significant improvements in small-budget regimes.

Figure 5: COPEx sequential query trajectory under tight transition constraints, illustrating cost-aware, smooth design progression.

Numerical Results and Claims

• COPEx consistently achieves higher cumulative EIG, lower predictive error, and improved posterior utility under constraints than unconstrained amortized policies and myopic variational methods.
• The advantage is greatest as feasibility tightens (smaller transition constraints or lower budgets).
• The runtime overhead is modest ($\approx$ seconds per design) for horizons up to $H=3$ and branch factors $m=2$; diminishing returns are observed for larger $H$ due to surrogate bias propagation.
• COPEx outperforms GP-based acquisition methods in cost-sensitive settings, establishing its efficacy for offline-to-online adaptation in active learning.

Theoretical and Practical Implications

COPEx bridges the gap between unconstrained amortized design and constraint-aware non-myopic planning. It provides a tractable and flexible solution for dynamic constraints, enabling deployment of BED policies across environments where costs, budgets, and physical feasibility evolve unpredictably. The scenario-tree formalism unifies local and global constraints, and semi-amortized planning accommodates arbitrary design-dependent or state-dependent feasibility.

On the theoretical side, COPEx establishes an error decomposition for planning under surrogate posteriors, offering quantitative guidance on horizon-depth vs. amortization error accumulation. The modular architecture allows for future integration with more expressive generative inference models (e.g., diffusion-based surrogates), advanced planning algorithms (e.g., DESPOT, parallel rollout), and discrete action space optimizers.

Future Directions

Several opportunities arise for further advancement:

• Reducing amortization bias with flow-matching or diffusion surrogates [wildberger2024flow].
• Hybridizing scenario-tree planning with advanced rollout and parallel search [somani2013despot, cai2021hyp].
• Extending to mixed discrete-continuous action spaces (e.g., integer program solvers [bonami2011algorithms, le2011algorithm]).
• Addressing model misspecification via robust Bayesian criteria [huang2023learning, forster2025improving, tang2025generalization].

Conclusion

COPEx provides an efficient, constraint-adaptive, non-myopic BED framework, combining offline amortization with online scenario-tree planning for diverse real-world design scenarios. It demonstrates substantial improvements in information-efficient sequential experiment design under a broad range of constraints, forming a practical bridge between amortized experimental design and dynamic, adaptive deployment (2605.26990).