{'fraction_non_alphanumeric': 0.04196336773161723, 'fraction_numerical': 0.015460188111648452, 'mean_word_length': 4.585282858358146, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Efficient Real-world Testing of Causal Decision Making via Bayesian Experimental Design for Contextual Optimisation

The real-world testing of decisions made using causal machine learning models is an essential prerequisite for their successful application. We focus on evaluating and improving contextual treatment assignment decisions: these are personalised treatments applied to e.g. customers, each with their own contextual information, with the aim of maximising a reward. In this paper we introduce a model-agnostic framework for gathering data to evaluate and improve contextual decision making through Bayesian Experimental Design. Specifically, our method is used for the data-efficient evaluation of the regret of past treatment assignments. Unlike approaches such as A/B testing, our method avoids assigning treatments that are known to be highly sub-optimal, whilst engaging in some exploration to gather pertinent information. We achieve this by introducing an information-based design objective, which we optimise end-to-end. Our method applies to discrete and continuous treatments. Comparing our information-theoretic approach to baselines in several simulation studies demonstrates the superior performance of our proposed approach.


Introduction
Machine learning models, particularly causal machine learning models, can be used to make real-world decisions (Joyce, 1999;Geffner et al., 2022). We consider the problem of collecting experimental data to evaluate such models and improve decisions made in future. For example, suppose a company has developed a model describing the relationship between customer engagement and revenue. This model will be used to offer personalised promotions to each customer in order to maximise some reward, e.g. revenue. Before deploying it to all of its customers, the company may wish to collect experimental data to validate the model.
To formalise this, we consider the contextual optimisation problem in which a treatment or treatments must be assigned to each context to receive a reward. We focus on the experimental design problem of finding treatments that allow us to evaluate the regret of past actions and to improve decisions made in the future. Our aim, then, is to choose actions to test the model; we are not directly optimising rewards in the experimental phase. In line with our real-world scenario, we focus on designing large batch experiments, in which we must select a large number of treatments for different contexts before receiving any feedback.
A standard experimental design procedure is random treatment assignment, which corresponds to A/B testing when treatment is binary (Fisher, 1936;Kohavi and Longbotham, 2017). However, this naive approach has significant drawbacks, as it may involve applying many sub-optimal actions, meaning that the opportunity cost of experimentation becomes high and experimental resources are wasted applying treatments that are known a priori to be sub-optimal. On the other hand, applying the best treatments based on existing beliefs-a pure exploitation strategy-may also be uninformative for model testing as it fails to account for uncertainty and the need to gather data where uncertainty is high.
To design an experiment in a more principled way, we turn to the framework of Bayesian experimental design (BED) (Lindley, 1956;Chaloner and Verdinelli, 1995;Foster et al., 2019). In contrast to standard BED, which maximises the information gained about all model parameters, we define the objective for experimental design as the expected information gained about the maximum reward obtainable for a set of evaluation contexts (max-value EIG). This directly addresses the problem of gaining data to evaluate the regret of past decisions, because, given a past outcome, the unknown component of regret is how much better the reward could have been. The formulation is also relevant to the problem of gathering data (exploration) to obtain better rewards in later interactions (exploitation).
As a model-based framework, BED requires the specification of a Bayesian model (Gelman et al., 2013) relating contexts, treatments and rewards. Since the model makes predictions about interventional distributions, causality plays a central role in selecting a correct model, particularly when the model is pre-trained with observational data. One approach is to rely on assumptions about the causal graph (Pearl, 2009;Sharma et al., 2021). Alternatively, the model can incorporate Bayesian uncertainty about the causal structure (Heckerman et al., 1999;Annadani et al., 2021;Geffner et al., 2022). In this paper, we propose a model-agnostic approach to experimental design that accommodates this full range of Bayesian models; we only require differentiability of samples from the model.
To find the optimal experimental design for a given Bayesian model, we first show that the InfoNCE bound (van den Oord et al., 2018) can be used to estimate the max-value EIG objective. We then optimise the testing treatment assignment (experimental design) by gradients (Foster et al., 2020;Kleinegesse and Gutmann, 2020). Our method is applicable to multiple treatments, both continuous and discrete, using carefully tuned Gumbel-Softmax (Maddison et al., 2016;Jang et al., 2016) for the latter case. We show experimentally that our approach is effective on a number of synthetic models, learning intuitively correct large batch designs that outperform standard baselines by significant margins.

Problem set-up
We are interested in the efficient evaluation of the regret of past treatment assignment decisions. We formalise this by considering a contextual optimisation problem in which treatment, a, is applied in context, c, to receive reward y. For example, in the customer engagement scenario, c would be the information we have about the customer that cannot be directly acted upon, such as age, income and historical data of past interactions, a would denote a particular choice of promotions that we offer to the customer, and the reward y would be the future revenue generated by the customer, net of promotion costs.
Expected regret We are interested in gathering data to evaluate the expected regret of the past action a performed in a context of interest c * . This is defined as where p(·|do(a), c * ) denotes the interventional distribution when action a is applied to context c * and m(c * ) = max a E[y|do(a ), c * ] denotes the best obtainable outcome for context c * . Suppose that a has been performed historically to c * , and we wish evaluate the regret of this treatment. The primary obstacle to evaluating this regret in the real world is that m(c * ) is not known, i.e. we do not know how much better the outcome could have been if a different treatment had been applied. We therefore wish to obtain experimental data, possibly from different contexts, that will help us to efficiently infer m(c * ). Gathering data to learn m(c * ) also aids in choosing the best treatment in future interactions with c * .
Bayesian model In general, we also do not know the true interventional distribution p(y|do(a), c). We therefore model the relationship between a, c and y is a Bayesian manner, by introducing a Bayesian parameter ψ with a prior p(ψ). Different samples of ψ corresponds to different hypotheses p(y|do(a), c, ψ), whilst marginalising over ψ summarises the total Bayesian uncertainty in y|do(a), c. Figure 1 illustrates an example with a binary treatment.
The role of causality Causal considerations are central to building a correct model of the interventional distribution p(y|do(a), c), particularly when the prior p(ψ) is fitted by conditioning on past observational data. By making certain assumptions about the causal graph and the (non-)existence of unobserved confounders, one can reduce the problem to that of learning specific functions (Pearl, 2009;Sharma et al., 2021). These assumptions are often made implicitly in the optimisation literature (Bareinboim et al., 2015). Through suitable assumptions, any Bayesian model that predicts y from a, c can be used as a model for p(y|do(a), c), so a causal treatment of the problem does not restrict the class of models that can be considered. Alternatively, the Bayesian parameter ψ can incorporate uncertainty in the causal graph (Heckerman et al., 1999;Annadani et al., 2021;Geffner et al., 2022). Our approach to experimental design is concordant with either of these modelling choices.

Bayesian Experimental Design for Contextual Optimisation
Whilst standard BED would design experiments to maximise the information gathered about the Bayesian parameter ψ, our goal is to efficiently evaluate the regret (1). This translates to estimating the maximum achievable rewards for a set of D * evaluation contexts C * = c * 1 , . . . , c * D * , denoted m * = m(c * 1 ), . . . , m(c * D * ). In the customer engagement example, C * represents a set of customers for which we are aiming to infer the best possible rewards m * . To learn about m * , we are allowed to perform a batch of D experiments at contexts C = c 1 , . . . , c D . In the customer engagement example, C represents the set of customers that will participate in the real-world test. The experimental design problem is to select the treatments A = a 1 , . . . , a D to apply to C so that the experimental outcomes y = y 1 , . . . , y D The dark (resp. light) shaded area shows our uncertainty about y|do(a), c arising from uncertainty in ψ, measured by one (resp. two) st. dev. from the mean. We want to estimate m * = m(c * ), the unknown maximum achievable reward at the evaluation context c * . We perform three experiments at c 1 , c 2 , c 3 (dotted orange lines). UCB chooses the reward-maximising treatments in each experimental context, bringing no information about m * , as Treatment 2 is a priori suboptimal at c * . Random design selects Treatments 1 and 2 with equal probability; our method selects treatments whose outcomes will be most informative about m * given the prior.
that we get from the real-world test are maximally informative about m * . More formally, we design A to maximize expected information gain between y and m * . We refer to I(A; C, C * ) as the max-value EIG. The objective is doubly intractable (Rainforth et al., 2018) meaning that its optimisation, or even evaluation with a fixed A, is a major challenge. The likelihoods in the expectation are implicit (not available analytically) making the problem more difficult.

Training stage: Objective estimation and optimization
Recently, gradient-based approaches for optimising EIG have been developed for experimental design for parameter learning (Foster et al., 2020;Kleinegesse and Gutmann, 2020) as opposed to contextual optimisation, which is the focus of this work. Inspired by likelihood-free mutual information estimators, we show the InfoNCE bound (van den Oord et al., 2018) can be adapted to the contextual optimisation setting we consider here. Concretely, is a lower bound on I(A; C, C * ) for any L ≥ 1 and any critic function U . We provide a proof of this claim in Appendix A. To learn a suitable critic U , we use a neural network with trainable parameters φ. We then optimise the lower bound L(A, U ; C, C * , L) simultaneously with respect to φ and A, thereby improving the tightness of the bound and optimising the experimental design together. When the actions A are continuous, we require that a sample y of p(y|do(a), c, ψ) is differentiable with respect to a. In this case, a pathwise gradient estimator (Mohamed et al., 2020) for ∇ A,φ L can be readily computed and we apply standard stochastic gradient ascent (SGA) (Robbins and Monro, 1951) to optimise our objective (3).
Discrete action space Previous gradient-based BED work has focused on fully differentiable models. Here we propose a practical way to deal with discrete designs. Rather than learning the treatments A directly, we introduce a stochastic treatment policy π α with trainable parameters α representing the probabilities of selecting each treatment. During training we use a Gumbel-Softmax (Maddison et al., 2016;Jang et al., 2016) relaxed form of A and update α by gradients. Further detail is provided in Appendix A.

Deployment stage: Real-world testing
Once we optimise the objective (3), we perform the batch of experiments A to obtain real-world outcomes y and estimate p(ψ|D)-the posterior of our Bayesian model, given the experimental data D = (y, A, C). We then calculate the posterior p(m * |C * , D), and this it to estimate the regret (1) of a set of past treatments. Importantly, through our design of A, the data contained in D should lead to the most accurate estimate of m * , which translates to an accurate estimate of the regret and an accurate evaluation metric for the original decisions. Algorithm 1 in the Appendix gives a complete description of our method.

Related Work
The most closely related objective to our max-value EIG (2) is Max-value Entropy Search (MES, Wang and Jegelka, 2017). It was proposed as a computationally efficient alternative to Entropy Search (ES, Hennig and Schuler, 2012) and Predictive Entropy Search (PES, Hernández-Lobato et al., 2014). Our objective differs from MES in two essential ways: firstly, MES is applicable to Gaussian processes (GPs), whilst our objective is model-agnostic; secondly, we focus on contextual optimisation, rather than finding a single maximiser. For a more comprehensive discussion of related work, see Appendix B. We test the efficacy of our method on several synthetic experiments. In all synthetic experiments, we assume the causal graph is as shown in Figure 2, but note that our method is not restricted to this graph structure. We compare our method against random treatment assignment and the Upper Confidence Bound (UCB) (Auer, 2002) algorithm. Our key evaluation metrics are EIG and accuracy of inferring m * . For the former we evaluate (3) with the trained critic and learnt designs. For the latter, we sample a realisatioñ ψ andm * from the prior, treating it as the ground truth environment in the Deployment stage of our framework ( § 3.2). We report the mean squared error (MSE) between our estimate m * and the ground truthm * . We report three additional metrics-average regret, MSE of estimatingψ, and deviation from the true optimal treatment. Full details about the models, training, and further results are given in Appendix C.   Discrete treatments We consider a model with four treatment options, two of which are a priori sub-optimal; the other two have the same mean but Treatment 1 has higher variance than Treatment 2 ( Figure 3). Intuitively, the optimal strategy would be to A/B test only the top 2 treatments, which is exactly what our method has learnt. The random method is wasteful as it queries the sub-optimal treatments, while UCB 1 only ever queries Treatment 1. As a result, our method is much more accurate at estimating the true m * .

Continuous treatments
We test our method on a continuous treatment problem, where we design a batch of 40 experiments to learn about the max-values at 39 evaluation contexts. As Table 1 shows, our method outperforms the baselines on all metrics. Although our method is optimised to accurately estimatem * , it is also significantly better at estimatingψ, which in turn translates to choosing better treatments in the future, as the lower MSE(A) indicates.

Discussion and Conclusion
By combining ideas from Bayesian experimental design, contextual bandits and Bayesian optimisation, we proposed a method for real-world model testing by regret evaluation of past context-dependent treatment assignments. Our method also applies to gathering data to improve future decision making, in fact, we saw that gaining information for regret evaluation and for model improvement are really two sides of the same coin. To compute large batch designs efficiently, we cast the problem in terms of implicit likelihood BED, resulting in a model agnostic approach that is able to handle both discrete and continuous designs. Future work may introduce the ability to handle constraints, and explore the scaling of our method to high dimensional treatments and contexts, both continuous and discrete. Our method could also be extended to apply to sequential, adaptive experimentation. Algorithm 1: Data-efficient regret evaluation Input: Bayesian simulator p(y, m * |do(a), c, c * , ψ), batch size D of experimental contexts C, batch size D * of evaluation contexts C * , initial A, initial U φ , number of contrastive samples L ≥ 1 Output: Optimal treatments A for experimental contexts C to be tested Training stage: while Computational training budget not exceeded do Sample ψ ∼ p(ψ) Sample y ∼ p(y|C, do(A), ψ) and m * ∼ p(m * |C * , ψ) Compute L(A, U φ ; L) (3) and update the parameters (A, φ) using a SGA end Deployment stage: Run a batch of experiments A to obtain real-world observations y and estimate a posterior p(ψ|D), D = {y, A, C} . Use the updated model to get an estimate of m * and evaluate the regret (1) associated with some (past) treatment assignment in contexts C * .

Appendix A. Method
For completeness we formalise the claim made in Section 3.1 in Proposition 1, and provide a proof, following standard arguments. The bound is tight for the optimal critic U * (y, m * ) = log p(y|m * , do(A), C) + c(y) as L → ∞, where c(y) is an arbitrary function depending only on the outcomes y.
Finally, substituting U * (y, m * ) = log p(y|m * , do ( which is monotonically increasing in L and tight in the limit as L → ∞ since (See e.g. Foster et al., 2020, Theorem 1 for more detail).
Discrete action space Given K ≥ 2 possible treatments, we learn a (non-deterministic) policy π with parameters α where the parameter τ is called the temperature, which we anneal during optimisationstarting the optimisation at high temperatures, when gradients exhibit low variance, and gradually reducing it towards 0, a which stage gradients are high variance, but we (should) have already learnt good treatments. We optimise the parameters α and those of the critic network φ jointly with SGA. Once the policy is trained, the optimal design for experiment d in the batch is a d = arg max(π d,1 , . . . , π d,K ).

Appendix B. Additional Related Work
Our setting is related to contextual Bayesian optimisation; some of the more closely related methods include Profile Expected Improvement (PEI, Ginsbourger et al., 2014), Multi-task Thompson Sampling (MTS, Char et al., 2019) and conditional Bayesian optimization (ConBO, Pearce et al., 2020). All of these methods are, however, restricted to GPs and the criteria they use to choose optimal designs are not information-based. Contextual bandits is another broad framework that our work is related to. An extensive line of research is focused on online linear bandits and discrete actions chosen using (variations of) UCB, Thompson sampling or -greedy strategy (Auer, 2002;Chu et al., 2011;Agrawal and Goyal, 2013;Han et al., 2020). Krause and Ong (2011) instead model the reward as a GP defined over the context-action space and develop CGP-UCB. More recently, Zanette et al. (2021) proposed designing a batch of experiments offline to collect a good dataset from which to learn a policy. Although this is similar to our setting, our approach is not restricted to linear rewards and uses an information-theoretic design objective.
Similar variational EIG objectives have been used in implicit likelihood BED methods for parameter learning, but not for contextual optimisation. Gradient-based methods for large batch experimentation include SG-BOED (Foster et al., 2020) and MINEBED (Kleinegesse and Gutmann, 2020), while the policy-based iDAD (Ivanova et al., 2021) applies to batch and adaptive settings. Our ability to handle discrete designs is another important distinction of our method.

C.1 Training details
We implement all experiments in Pyro (Bingham et al., 2018). All experiments baselines ran for 50K gradient steps, using a batch size of 2048. We used the Adam optimiser (Kingma and Ba, 2014) with initial learning rate 0.001 and exponential learning rate annealing with coefficient 0.96 applied every 1000 steps. We used a separable critic architecture (Poole et al., 2019) with simple MLP encoders with ReLU activations and 32 output units.
For the discrete treatment example: we added batch norm to the critic architecture, which helped to stabilise the optimisation. We had one hidden layer of size 512. Additionally, for the Gumbel-Softmax policy, we started with a temperature τ = 2.0 and hard=False constraint. We applied temperature annealing every 10K steps with a factor 0.5. We switch to hard=True in the last 10K steps of training.
For the continuous treatment example: We used MLPs with hidden layers of sizes [design dimension × 2, 412, 256] and 32 output units.
Note: In order to evaluate the EIG of various baselines, we train a critic network for each one of them with the same hyperparameters as above.

C.2 Posterior inference details
After completing the training stage of our method (Algorithm 1), we need to deploy the learnt optimal designs in the real world in order to obtain rewards y. This experimental data is then used to fit a posterior p(ψ|D).   There are many ways to do the posterior inference and quality of the results will crucially depend on the accuracy of the fitted posteriors. In both of our examples and for all baselines we use use Pyro's self-normalised importance sampling (SNIS). Samples from this posterior are used for the evaluation metrics.
We validate the accuracy of estimated posteriors by running various sanity checks, including diagnostic plots such as Figure 4, showing the standard deviation of our posterior mean estimate (measure of how uncertainty about the parameter) and L 2 error to the true parameter. The red line shows the rolling mean over 200 points of the latter, and the gray band-the 2 standard deviations. For this plot we used the continuous treatments example with D = 20 experimental contexts.

C.3 Evaluation metrics details
As discussed in the main text, we evaluate how well we can estimate m * by sampling a ground truthψ from the prior and obtaining a corresponding ground truthm * . We approximate the max-values m * empirically using 2000 posterior samples of ψ We similarly estimate ψ using 2000 posterior samples. We define the optimal treatment under the posterior model to be average (with respect to that posterior) optimal treatment when treatments are continuous or UCB 0 when treatments are discrete. Finally the regret is computed as the average difference between the true max value (from the true environment and the true optimal treatment) and the one obtained by applying the estimated optimal treatment. We used 4000 (resp. 2000) true environment realisation for the continuous (resp. discrete) example.

C.4 Discrete treatments example
Model We first give details about the toy model we consider in Figure 3. Each of the four treatments a = 1, 2, 3, 4 is a random function with two parameters ψ k = (ψ k,1 , ψ k,2 ) with the  and reward (outcome) likelihoods: γ = (ψ a,1 + ψ a,2 + 18)/2 (20) β = (ψ a,2 − γ + 9)/3 Intuition about the parameterisation: The first component of each ψ i defines the mean reward at context c = −3, while the second one defines the mean reward at context c = 3. The reward is then the quadratic equation that passes through those points and leading coefficient equal to −1.
Experimental and evaluation contexts We use experimental and evaluation contexts of the same sizes. The experimental context, c is an equally spaced grid of size 10 between −3 and −1. We set the evaluation context c * = −c. Figure 3 in the main text visually illustrates this: the x-axis of the points in each plot are the experimental contests, while the dashed gray lines are the evaluation contexts. Table 2 shows all the evaluation metrics for the discrete treatment example from the main text. Our method achieves substantially higher EIG and lower MSE of estimating the max-rewards. On all other metrics all methods perform similarly. This is to be expected, since Treatments 1 and 2 have exactly the same means and due to the way the model was parameterised (by the value of the quadratic at contexts 3 and -3), the probability of the optimal treatment being 1 or 2 is exactly 50% (the hit rate all baselines achieve). Note that UCB 1 and UCB 2 achieve statistically identical results, which is expected given they select the same designs.

Further results
Training stability We perform our method with the same hyperparameters but different training seeds and report the mean and standard error in Table 3.  Table 4: Continuous treatments example: 40D design training stability. Mean and standard error are reported across 6 different training seeds.

C.5 Continuous treatment example
Model For the continuous treatment example we use the following model: Likelihood: y|c, a, ψ ∼ N (f (ψ, a, c), σ 2 ), where f (ψ, a, c) = exp − a − g(ψ, c) 2 h(ψ, c) g(ψ, c) = ψ 0 + ψ 1 c + ψ 2 c 2 h(ψ, c) = ψ 3 Experimental and evaluation contexts The experimental context, c is an equally spaced grid of size D = 40 (or 20 or 60 in Further Results below) between −3.5 and and 3.5. The evaluation context c * is of size D * = D − 1 and consists of the midpoints of the experimental context (see Figure 5 for an illustration). Baselines Since we have a continuous treatment, for the random baseline we consider sampling designs at random from N (0, 0.2), N (0, 1) or N (0, 2), which we denote by Random 0 .2, Random 1 and Random 2 , respectively.
Training stability We perform our method with the same hyperparameters but different training seeds and report the mean and standard error in Table 4.

Further results
We report the results of the same experiment, but with a smaller and larger batch sizes of experimental and evaluation contexts. Table 5 shows results for an experimental batch size of 20 contexts to learn about 19 evaluation contexts, while Figure 5 visually illustrates the model and the designs. Finally, Table 6 shows results for an experimental batch size of 60 contexts to learn about 59 evaluation contexts.