Technical Summary: Constrained Bayesian Experimental Design via Online Planning (COPEx)

Problem Formulation

Bayesian Experimental Design (BED) provides a framework for adaptively selecting optimal experiments to maximize information gain, typically measured by Expected Information Gain (EIG). However, existing state-of-the-art BED methods, particularly amortized policy-based approaches, struggle to adapt to dynamic constraints inherent in real-world deployment. These constraints include:

• Transition Constraints: Physical limitations restricting how designs can evolve (e.g., sensor movement limits requiring smooth trajectories where $\|x_t - x_{t-1}\| \leq \delta$).
• Budget Constraints: Global limits on total resources or costs that vary across instruments and time.
• Feasibility Constraints: Operational restrictions dictating admissible future experiments.

Current amortized methods train policies offline to maximize total EIG under a fixed notion of feasibility. When constraints change at deployment (e.g., a new budget limit or a stricter movement constraint), these policies often fail to adapt without expensive retraining. Naive post-hoc masking of invalid actions forces the policy onto trajectories unseen during training, leading to sub-optimal designs and high posterior uncertainty.

Methodology: COPEx

The authors propose COPEx (Constrained Online Planning for Experimental design), a semi-amortized method that combines offline pre-training with online multi-step lookahead planning. The core innovation is solving the constrained BED problem via a scenario tree search over hypothetical experimental outcomes.

1. Offline Pre-training (Amortization)

To make online planning computationally feasible, COPEx pre-trains two neural networks:

• Amortized Posterior Network ($q_\phi$): A conditional density estimator (specifically a Mixture Density Network based on Transformer Neural Processes) that learns to map datasets $D$ to an approximate posterior $q_\phi(\theta|D)$. This enables rapid belief updates and the generation of "fantasy" outcomes without retraining.
• Amortized Design Policy ($\pi_\psi$): A policy trained offline (using reinforcement learning on unconstrained tasks) to propose high-quality initial designs. This serves as a "warm start" for the online optimizer, biasing the search toward high-utility regions.

2. Online Multi-Step Lookahead Planning

At each time step $t$, COPEx constructs a finite-horizon scenario tree of depth $H$:

• Tree Construction: The root represents the current state $(D_{t-1}, z_t)$, where $z_t$ encodes constraint states (e.g., remaining budget, previous design). The tree alternates between selecting a design $x$ and branching over a set of fantasy observations $\tilde{y}$ sampled from the amortized posterior predictive distribution.
• Objective: The planner optimizes the cumulative EIG over the tree, subject to constraints enforced at every node. The objective is approximated via Monte Carlo estimation over fantasy branches.
• One-Shot Optimization: To avoid nested dynamic programming, COPEx employs a reparameterization trick. Base noise variables are sampled once and fixed, making all fantasy outcomes deterministic functions of the decision variables. This transforms the problem into a single-shot nonlinear optimization over the entire tree structure.
• Initialization: The optimization is initialized using the amortized policy $\pi_\psi$. To handle cases where constraints shift the feasible region significantly, a hybrid initialization strategy is used, combining policy-based proposals with random exploration.

3. Constraint Handling

Constraints are integrated directly into the optimization:

• Transition Constraints: Enforced locally by restricting the feasible set $X(z_t)$ at each node based on the previous design.
• Budget Constraints: Modeled as a state variable $b_t$ that evolves deterministically ($b_{t+1} = b_t - c(x_t)$). The feasible set is defined by the remaining budget.

Key Contributions

1. COPEx Framework: A novel semi-amortized approach that enables non-myopic optimization under dynamic, test-time constraints without requiring retraining.
2. Scenario Tree Planning: A method for constructing and optimizing multi-step lookahead trees using amortized surrogates to bypass the computational intractability of nested expectations in sequential BED.
3. Warm-Start Strategy: The use of a pre-trained amortized policy to initialize the tree optimizer, significantly reducing the computational cost of finding high-quality solutions compared to random initialization.
4. Empirical Validation: Demonstration that COPEx adapts to constraints (transition and budget) where existing amortized and non-amortized baselines fail or perform sub-optimally.

Experimental Results

The authors evaluated COPEx on three tasks:

1. Location Finding (Transition Constraints): In a 2D source localization task with bounded movement constraints ($\|x_t - x_{t-1}\| \leq \delta$), COPEx consistently achieved higher cumulative EIG than baselines (VPCE and ALINE). The performance gap widened as constraints became tighter ($\delta$ decreased). A planning horizon of $H=1$ to $3$ offered the best trade-off between EIG gain and runtime.
2. Constant Elasticity of Substitution (Budget Constraints): In a behavioral economics task with a global budget, COPEx with lookahead ($H=1$) significantly outperformed myopic COPEx and all baselines (VPCE, RL-BOED, ALINE) in terms of EIG accrued before budget exhaustion.
3. Cost-Aware Active Learning: On benchmark functions (Ackley, Branin, Goldstein-Price) with heterogeneous cost landscapes (e.g., "Hazard Center"), COPEx achieved lower Root Mean Squared Error (RMSE) for the same accumulated cost compared to Gaussian Process baselines. Non-myopic planning ($H>0$) provided clear advantages, particularly when the remaining budget was small relative to the planning horizon.

Significance and Claims

The paper claims that COPEx addresses a critical gap in sequential experimental design: the inability of current methods to handle dynamic, path-dependent constraints during deployment.

• Adaptability: Unlike fully amortized policies that are rigid to constraint changes, COPEx performs online planning that explicitly accounts for the current constraint state.
• Efficiency: By leveraging amortized surrogates for posterior updates and EIG estimation, COPEx achieves superior information gain with only a modest computational overhead (seconds-scale runtime) compared to the prohibitive costs of exact non-myopic planning.
• Practicality: The method is applicable to real-world scenarios such as mobile sensing, clinical trials, and material design, where physical movement limits and budget constraints are fundamental to the experimental strategy.

The authors acknowledge limitations, including the accumulation of approximation errors from the amortized posterior over long horizons and the reliance on gradient-based optimization which is best suited for continuous design spaces. They suggest future work in more expressive generative models and extensions to discrete action spaces.