Beating the Winner's Curse via Inference-Aware Policy Optimization

Beating the Winner's Curse via Inference-Aware Policy Optimization

\RUNAUTHOR{Bastani et al.}

\RUNTITLE{Inference-Aware Policy Optimization}

\TITLE{Beating the Winner's Curse via \\Inference-Aware Policy Optimization}

\ARTICLEAUTHORS{ \AUTHOR{Hamsa Bastani} \AFF{Operations, Information, and Decisions Department, The Wharton School, \EMAIL{[email removed]}}

\AUTHOR{Osbert Bastani} \AFF{Department of Computer and Information Science, University of Pennsylvania, \EMAIL{[email removed]}}

\AUTHOR{Bryce McLaughlin} \AFF{Wharton Healthcare Analytics Lab, Wharton AI & Analytics Initiative, \EMAIL{[email removed]}}

}

\ABSTRACT{ There has been a surge of recent interest in automatically learning policies to target treatment decisions based on rich individual covariates. In addition, practitioners want confidence that the learned policy has better performance than the incumbent policy according to downstream policy evaluation. However, due to the winner's curse---an issue where the policy optimization procedure exploits prediction errors rather than finding actual improvements---predicted performance improvements are often not substantiated by downstream policy evaluation. To address this challenge, we propose a novel strategy called inference-aware policy optimization, which modifies policy optimization to account for how the policy will be evaluated downstream. Specifically, it optimizes not only for the estimated objective value, but also for the chances that the estimate of the policy's improvement passes a significance test during downstream policy evaluation. We mathematically characterize the Pareto frontier of policies according to the tradeoff of these two goals. Based on our characterization, we design a policy optimization algorithm that estimates the Pareto frontier using machine learning models; then, the decision-maker can select the policy that optimizes their desired tradeoff, after which policy evaluation can be performed on the test set as usual. Finally, we perform simulations to illustrate the effectiveness of our methodology. \footnote{We gratefully acknowledge Katherine Milkman, Dena Gromet, and the Behavior Change for Good Initiative for providing access to the megastudy data used in Section (ref) of our paper. Furthermore, we would like to thank Ron Berman, Dean Eckles, Vishal Gupta, Dean Knox, Daniel Rock, Neha Sharma, Hummy Song, Alp Sungu, Sophie Yu, and members of the Wharton AI & Analytics Initiative for their helpful discussion and feedback. } }

Introduction

Due to rapid improvements in the capabilities of machine learning over the past two decades, there has been a surge of recent interest in systems that automatically make decisions based on the predictions of machine learning models dudik2011doubly,bottou2013counterfactual. These methods have been deployed across a wide variety of application domains, including refugee matching programs bansak_improving_2018,ahani_placement_2021 and personalized healthcare treatment decisions boutilier_improving_2022, lekwijit2024evaluating.

However, existing methodologies for policy learning face a major challenge known as the winner's curse. This issue arises when training a model to approximate the true objective and then optimizing against this model. Intuitively, even if the model is highly accurate on the training data, the optimization procedure can systematically exploit its prediction errors to obtain a policy whose predicted value is substantially higher than its true value. This failure has previously been observed across a variety of contexts, including common value auctions capen1971competitive,thaler1988anomalies and managerial decision-making (where it has also been called post-decision surprise harrison1984decision or the optimizer's curse smith2006optimizer); it is also an instance of regression to the mean harrison1984decision. It is greatly exacerbated by modern policy optimization pipelines, which consider complex policies mapping rich, individual covariates to actions.

Much of the early work on policy evaluation assumes that a policy is fixed prior to evaluation rosenbaum_central_1983, rosenbaum1984reducing,bang2005doubly. Due to increasing interest in policy optimization, recent work has proposed methodologies for correcting the optimistic bias arising from the winner's curse gupta2024debiasing,andrews_inference_2024,zrnic_flexible_2024,xu2025winner. However, they only correct for the evaluation error---if the learned policy is ineffective, they cannot help compute a more effective policy. Furthermore, practitioners often want evidence that the learned policy actually increases performance---e.g., by obtaining a statistically significant estimate of the learned policy's performance improvement over the incumbent policy (which we call the observational policy). Thus, a critical remaining question is how to identify a policy that improves performance according to a rigorous downstream evaluation methodology. Addressing this challenge requires rethinking the entire policy optimization pipeline---our key insight is that policy optimization should account for how the policy will be evaluated downstream, which we call inference-aware policy optimization.

We formalize the inference-aware policy optimization problem and mathematically characterize the space of solutions to this problem. We focus on a simple policy evaluation methodology using inverse propensity weighting (IPW), which evaluates a learned policy by reweighting test outcomes to appear as if they were collected under the learned policy rosenbaum_central_1983; we restrict the policy optimization algorithm to only see training data and test covariates, which suffices to avoid the winner's curse during evaluation. Our goal is to design a policy optimization algorithm that aims to achieve statistically significant performance improvement. Specifically, we aim to compute policies that trade off (1) the expected objective value according to the downstream IPW estimate, and (2) the expected $z$-score of the IPW estimate used to assess whether the learned policy is statistically significantly better than the observational policy. Intuitively, if the learned and observational policies are quite different, then the IPW estimator will downweight a large part of the test dataset when evaluating the proposed policy li_unbiased_2011, thereby effectively reducing sample size and failing to establish a statistically significant performance improvement. Thus, to be conservative, we want to keep the learned policy closer to the random policy used to collect data, avoiding large propensity weights that can result in high variance.

In this setting, we obtain a precise mathematical characterization of the Pareto frontier of optimal solutions---namely, the tradeoff between the expected objective value and $z$-score. Specifically, we create a superset of this frontier by considering a family of convex optimization problems that minimize the variance of the IPW estimator for a fixed expected improvement $\lambda$. By analyzing the Lagrangian of this problem, we obtain a parametric representation of this family which we prune to obtain the Pareto frontier of optimal policies.

Using this characterization, we show that the policy with the largest expected $z$-score is close to the observational policy, making a small but confident improvement. Specifically, it differs from the observational policy for all subjects where an expected improvement can be made, with larger differences when the counterfactual means of treatments differ more and their related variances are small. The step size reaches the point where some treatment option is eliminated for some unit. Moving along the frontier from this initial solution, the tradeoff between expected objective and statistical significance becomes strict---i.e., any strict increase in expected objective requires a strict decrease in statistical significance. Furthermore, we provide a simple closed form for the largest expected $z$-score, effectively allowing for power calculations that determine whether a test dataset has enough useful variation to evaluate any policy's improvement as significant.

Next, we show how our theoretical results can be used to derive our algorithm for inference-aware policy optimization (IAPO). The main barrier is that the solution to our convex program is specialized to assigning treatments to our test units instead of being a generalizable policy that can be applied to any unit. However, we note that the policy assigns treatments based purely on the mean and standard deviation of the predicted counterfactual outcomes; thus, we can straightforwardly devise a generalizable policy from our solution. A full policy optimization and evaluation pipeline consists of the following steps: (1) train a predictive model on the training dataset, (2) compute the Pareto frontier using our characterization and select a policy along this frontier, and (3) evaluate the efficacy of this policy using IPW. We implement this pipeline in multiple simulation studies, demonstrating that it can effectively extract policies that deliver confident improvements in cases where na\"{i}ve policy optimization approaches fail.

Related Work

\paragraph{Policy evaluation.}

Our approach builds on the causal model introduced by rubin_estimating_1974. Specifically, the key challenge with policy evaluation is the need to estimate counterfactual outcomes---i.e., for each treatment unit (e.g., an individual), we only observe the outcome for the treatment that is assigned to that unit by the observational policy (called the factual outcome), not the outcomes for other treatments (called counterfactual outcomes). Then, the policy evaluation problem is to estimate the average treatment effect (ATE) of one treatment compared to another (usually the control). This problem does not directly consider targeted interventions, where the treatment is targeted based on unit covariates instead of being constant across all units. However, when a specific candidate targeting policy is given, then the problem of policy evaluation can be reduced to this one by considering the policy to be the intervention, and evaluating the efficacy of that policy.

We focus on the setting where there are no unobserved confounders; for instance, this is the case in applications where data is collected from a known observational policy such as random assignment. In this setting, there are two strategies that can be applied for policy evaluation. The first (called the direct method in dudik2011doubly) is to train a predictive model to estimate counterfactual outcomes, and use it to evaluate the candidate policy. This approach assumes that the predictive model is accurate, which is typically not true in practice. An alternative strategy is inverse propensity weighting rosenbaum_central_1983, which uses importance weights to correct for the shift from the observational policy to the candidate policy. This approach requires that the inverse propensity weights are known (e.g., this is the case when the observational policy is random assignment); otherwise, they must be estimated. One approach to further improve on importance weighting is doubly robust estimation robins1994estimation, which subtracts a prediction of the treatment outcome from the realized outcome to normalize it, resulting in the augmented IPW (AIPW) estimator. When the predicted outcomes are accurate, this strategy can significantly reduce variance compared to the IPW estimator. Doubly robust estimation is often used in settings where weights must be estimated, but it can reduce variance even when they are known. While we do not explicitly include the doubly robust correction in our main derivation, it can be straightforwardly incorporated into our framework.

\paragraph{Policy learning.}

The problem of learning an optimal treatment policy has long been studied in the multi-armed bandit lai1985asymptotically and reinforcement learning sutton1998reinforcement,murphy2003optimal literatures. While a majority of work in this space focuses on learning in the sequential setting, which seeks to learn through data collected online from the current policy, there has been recent interest in offline learning dudik2011doubly,bottou2013counterfactual (also called batch learning ernst2005tree), where the goal is to learn from historical data (the literature also distinguishes off-policy learning, which seeks to learn from a combination of historical and online data watkins1992q,precup2001off; see levine2020offline for a discussion).

The standard approach to policy learning is to construct an objective based on a policy evaluation technique, and then select the policy that optimizes this objective. Using the direct method for policy evaluation results in a plug-in approach qian_performance_2011, whereas using IPW results in a weight-based approach bottou2013counterfactual,kitagawa_who_2018; doubly robust approaches combining the two can also be used dudik2011doubly,zhang2012robust,zhou_residual_2017. There has also been work on using self-normalization to improve out-of-sample performance NIPS2015_39027dfa,kuzborskij_confident_2021, improving robustness to confounders Kallus_Confounding, and handling multiple observational policies zhan2024policy.

However, these strategies all suffer from the winner's curse---even if policy evaluation is unbiased for a single candidate policy, the objective value of the learned policy is optimistically biased. These approaches avoid this issue by using sample splitting---i.e., optimize the objective on an optimization set, and then perform policy evaluation (e.g., using IPW) on a held-out test set (the policy optimization procedure can use the test covariates; only treatment assignments and outcomes need to be held out). More sophisticated approaches have recently been proposed to avoid the reduction in sample complexity from sample splitting gupta2024debiasing,andrews_inference_2024,zrnic_flexible_2024,xu2025winner. However, all of these approaches can only correct for optimistic bias in policy evaluation; they cannot select a policy that is more likely to achieve statistically significant policy improvement, and often fail to do so.

There are two existing strategies that can be viewed as performing policy optimization to improve statistical significance. The first is treatment pooling, which tries to pool treatment options to reduce the size of the policy class in consideration banerjee_selecting_2025,Ma_adaptive_fusion,ladhania_personalized_2023,saito_potec:_2024; a related approach is identifying promising subgroups to reduce the size of the policy space spiess_finding_2023. This procedure can be viewed as reducing variance in exchange for a small amount of bias; however, it does so in an ad hoc way and can only be applied to small policy classes.

The second is regularization, where the policy optimization objective includes some kind of term to penalize variance in the policy evaluation objective. For instance, kallus_balanced_2018 estimates policy performance using weights that minimize the conditional mean squared error of the policy's performance estimates. Similarly, the Policy Optimizer for Exponential Model (POEM) swaminathan_batch_2015 and its self-normalized variant (Norm-POEM) NIPS2015_39027dfa penalize the variance of importance weighted samples. There has also been work along these lines in reinforcement learning, including conservative $Q$-learning kumar2020conservative and KL-regularized policy optimization nachum2019algaedice. However, these approaches do not connect variance regularization to statistical significance; furthermore, their optimization objectives are non-convex and are not amenable to theoretical analysis. Finally, in concurrent work, chernozhukov_policy_2025 connects variance-regularized policy optimization to statistical inference. However, they do not provide any insights into the structure of the Pareto frontier; instead, they provide valid inference after selecting a policy along the frontier. In contrast, we obtain a closed-form solution to the Pareto frontier, and use it to obtain interesting insights on its structure.

Contributions

Our contributions are three-fold. (1) We formalize the inference-aware policy optimization problem, where the goal is to obtain a policy with statistically significant policy improvement. (2) We perform a theoretical analysis on the tradeoff between expected performance and statistically significant policy improvement---specifically, we compute the entire Pareto frontier of policies trading off expected improvement and statistical significance of improvement. We find a number of interesting insights (summarized below). (3) We propose a novel policy optimization algorithm based on this analysis; our policy optimization algorithm is the first one designed to enable practitioners to obtain statistically significant policy improvements. By selecting a desired policy along the Pareto frontier, practitioners can obtain the best policy at a desired level of confidence.

Now, we summarize some of the key insights from our theoretical analysis:

General Framework

Our goal is to characterize the Pareto frontier of policies that jointly maximize the expected value and expected significance of an inverse propensity weighting (IPW) estimate of the policy improvement over the observational policy used to collect an observational dataset. We begin by formalizing this problem (Section (ref)), and establishing the convex optimization framework we use to solve it (Section (ref)). In Section (ref), we examine the special case of a single treatment (i.e., treatment vs. control group), enabling us to derive closed-form expressions for the Pareto frontier and obtain qualitative insights. Then, in Section (ref), we analyze the general case of multiple treatments; while we can no longer obtain a closed-form representation of the frontier, we nevertheless show that our key insights about its structure continue to hold.

Problem Formulation

Policy evaluation combines the Rubin causal model with inverse propensity weighting (IPW) to provide an unbiased estimate of an (unobserved) policy's performance using data collected by an observational policy. Under the Rubin causal model, we cannot observe the counterfactual outcomes units would realize under treatments other than the one they were assigned by the observational policy, inhibiting our ability to estimate the performance of other policies. IPW forms an unbiased estimate of any policies performance by reweighting the outcomes of the observational policy so they appear as if generated by a different policy. As these estimates can be quite noisy, confidently believing any policy delivers an improvement over the observational policy requires passing a significance test. Thus, when proposing policies, practitioners must find high performing policies which they can show are significantly better than the observational policy.

Following the Rubin causal model rubin_estimating_1974, we consider a set of units $n \in [N]$ that each receive a treatment $t \in \{0,\ldots,K\}$, yielding a random real-valued outcome $Y_{n,t}$. We assume that each of these outcomes are generated independently upon treatment assignment without interference between units. We let $\mu_{n,t} = \mathbb{E}\left[Y_{n,t}\right]$ denote the mean and $\sigma^2_{n,t} = \mathrm{Var}(Y_{n,t})$ denote the variance of these counterfactual outcomes. Each unit is assigned a treatment $T^o_n$ according to a known observational policy $\pi^o_n$ generating an observed outcome $Y_n^o = Y_{n,T_n^o}$. Each observational policy $\pi^o_n\in\Delta^K$ is an element of the standard $K$-simplex, and treatment assignments are independent random variables $T^o_n \sim \mathrm{Categorical}(\pi_{n}^o)$. We assume each treatment is assigned to each unit with positive weight (i.e. positivity). We refer to the collection $\mathcal{O} = \{(\pi^o_1,T^o_1,Y^o_1),\ldots,(\pi^o_N,T^o_N,Y^o_N)\}$ as the observational dataset and let $\pi^o_{n,t}$ be the probability unit $n$ is assigned treatment $t$.

We can use $\mathcal{O}$ to compare the performance of an alternate policy $\pi\in (\Delta^K)^N$ to $\pi^o$ using inverse propensity weighting (IPW) rosenbaum_central_1983. Specifically, the difference in expected performance between $\pi$ and $\pi^o$ is given by the sample average treatment effect $$\tau(\pi) = \frac{1}{N}\sum_{n=1}^N \sum_{t=0}^K \mu_{n,t}(\pi_{n,t} - \pi_{n,t}^o).$$ Applying IPW to $\mathcal{O}$ yields an unbiased estimate of $\tau(\pi)$:

is the inverse propensity weight. As $\hat{\tau}(\pi)$ is only an estimate of the improvement of $\pi$ over $\pi^o$, a $z$-test needs to be applied to gain confidence that $\tau(\pi)>0$; we use a $z$-test as we assume knowledge of the variances $\sigma^2$ (otherwise a $t$-test would be used), and we expect to use large sample sizes where the central limit theorem guarantees $\hat{\tau}(\pi)$ is normally distributed. The $z$-score for $\hat{\tau}(\pi)$ is

is the standard error of $\hat{\tau}(\pi)$.\footnote{Typically, the standard error is unknown, so we instead use a sample approximation $s(\pi)\approx\hat{s}(\pi)$ in conjunction with the $t$-score. For simplicity, we assume the standard error is known, so we can use the $z$-score.} The policy is accepted if the $z$-score exceeds some level $Z_{\text{min}}$.

Our goal is to construct a policy which improves upon the observational policy as much as possible while passing a significance test on the observational dataset. To do so, we have access to the counterfactual outcome distributions of $Y_{n,t}$ and the observational policy $\pi^o$, but not the observed treatments or outcomes. Note that for all $\pi$, $Z(\pi)$ should be approximately normal with $\mathrm{Var}(Z(\pi)) = 1$, meaning that to maximize our probability of achieving statistical significance, it suffices to maximize the expected $z$-score. Thus, we formalize our goal as computing the Pareto frontier $\Pi_{\mathrm{Opt}}$ for a multi-objective optimization problem, where the objectives are (1) the expected policy improvement $\tau(\pi)$, and (2) the expected $z$-score $\mathbb{E}[Z(\pi)] = \frac{\tau(\pi)}{s(\pi)}$ for the IPW estimator used in conjunction with the observational dataset $\mathcal{O}$. Mathematically, our goal is to compute

i.e., the Pareto frontier $\Pi_{\text{Opt}}$ consists of exactly the policies $\pi$ that cannot be strictly improved upon in terms of both expected improvement and $z$-score.

Convex Relaxation Framework

This frontier is difficult to characterize directly as one of the objectives (the expected $z$-score) is not convex. We start by showing that a superset of $\Pi_{\mathrm{Opt}}$ is given by a family of convex optimization problems $\Pi_{\Lambda}$. Then, we decompose the first order condition for this family of problems in order to characterize $\Pi_{\Lambda}$, which allows us to characterize $\Pi_{\mathrm{Opt}}$. Proofs for all claims are given in Appendix (ref).

Intuitively, any policy on the frontier $\Pi_{\text{Opt}}$ also minimizes the standard error $s(\pi)$ for some expected improvement $\tau(\pi)$; if this were not this case, then there would exist another policy with the same expected improvement but a higher expected $z$-score, removing it from $\Pi_{\text{Opt}}$. As the standard error is positive, this also holds for minimizing $s^2(\pi)$. We give complete proofs in Appendix (ref).

Proposition (ref) offers a path to solving for $\Pi_{\text{Opt}}$ by first solving for $\Pi_\Lambda$ then removing extraneous solutions to form the Pareto frontier. This is particularly useful as each optimization problem in $\Pi_{\Lambda}$ is convex in terms of the policy $\pi$.

To characterize $\Pi_\Lambda$ we need to solve this convex problem for all values of $\lambda \in [0,\tau_{\text{max}}]$. The first order condition of this optimization problem does not contain $\lambda$, instead giving a form for $\pi_{n,t}$ in terms of dual variables whose optimal values are induced by $\lambda$. The Lagrangian associated with all problems forming $\Pi_{\Lambda}$ is, $$\mathcal{L}(\pi,\zeta,\beta,\kappa) = s^2(\pi) + \zeta\left(\lambda - \tau(\pi)\right) +\sum_{n=1}^N\beta_n \left(1-\sum_{t=0}^K \pi_{n,t}\right) - \sum_{n=1}^N\sum_{t=0}^K\kappa_{n,t}\pi_{n,t},$$ where $\kappa > 0$ is the only dual constraint. To write its first order condition to respect $\pi_{n,t}$ we first need the partial derivatives of $s^2(\pi)$ and $\tau(\pi)$,

Thus, the first order condition with respect to $\pi_{n,t}^*$ for an optimal solution $(\pi^*,\zeta^*,\beta^*,\kappa^*)$ is $$ \frac{2}{N^2} \left[\frac{\mu_{n,t}^2 + \sigma_{n,t}^2}{\pi^o_{n,t}} \left(\pi_{n,t}^* - \pi^o_{n,t}\right) - \mu_{n,t} \sum_{v=0}^K \mu_{n,v} \left(\pi_{n,v}^* - \pi^o_{n,v}\right)\right] + \frac{\zeta^*}{N}\mu_{n,t} - \beta_{n}^* - \kappa_{n,t}^* = 0.$$ We can eliminate $\beta^*_n$ using the primal constraint $\sum_{t=0}^K \pi_{n,t}^* = 1$ and then reorganize this system of equations to solve for $\pi^*$ in terms of $\zeta^*$ and $\kappa^*$,

To convert Proposition (ref) into a characterization of $\Pi_{\Lambda}$, we solve for the values of all dual variables except $\zeta^*$, leaving a one-dimensional parameterization of $\Pi_{\Lambda}$ in terms of $\zeta^*$.

Single Treatment Setting

For intuition, we first consider for the case of comparing a single treatment to a control (i.e., $K=1$). We begin by solving our convex relaxation to obtain a superset of the Pareto frontier (Section (ref)). Then, we characterize the structure of the Pareto frontier based on this result (Section (ref)). Finally, we discuss implications of our characterization for policy learning methodologies (Section (ref)). Section (ref) generalizes our characterization results to the multi-treatment setting.

Superset of the Pareto Frontier

We begin by characterizing the Pareto frontier $\Pi_{\Lambda}$. The first order condition given in Proposition (ref) is independent across units, thus, for any $\zeta^*$ we can calculate $\pi_n^*(\zeta^*)$ in terms of $\pi_{n}^o$, $\mu_n$, and $\sigma^2_n$. This provides insight into how the dual variable determines optimal solution propensities in terms of the relative quality of the treatment for each unit. We characterize the policies $\pi \in \Pi_{\Lambda}$ in terms of $\pi_{n,1}$, the propensity to give the treatment to unit $n$, noting that $\pi_{n,0} = 1-\pi_{n,1}$. For convenience, define $\tau(\zeta)=\tau(\pi^*(\zeta))$, $s^2(\zeta)=s^2(\pi^*(\zeta))$, and $Z(\zeta)=\mathbb{E}[Z(\pi^*(\zeta))]=\tau(\zeta)/s(\zeta)$. A proof for Theorem (ref) is given in Appendix (ref).

Here, $\tau_n$ is the treatment effect for unit $n$, and $\eta_n$ is the variance of its IPW estimate:\footnote{The second term of $\eta_{n,t}$ involving $\mu_{n,t}$ may seem strange; it arises because the IPW estimator is not invariant to additive constants to the outcome variable. Indeed, normalizing outcomes can reduce variance of IPW, resulting in the augmented IPW (AIPW) estimator robins1994estimation.}

Furthermore, $\Pi_{\Lambda}$ is parameterized by the optimal value of the dual variable $\zeta$. It starts from the observational policy $\pi^*(0)=\pi^o$. As $\zeta$ increases, the treatment probability increases or decreases (depending on the sign of $\tau_n$) continuously at a rate of $N\tau_n/(2\eta_n)$ until it binds to either zero or one; the binding order is determined by the magnitude of this rate. This rate prioritizes units with higher expected treatment effect and lower estimation variance; notably, it is independent of the other units except through the total number of units $N$. The binding value optimizes the expected treatment effect; also, letting $\zeta_{\text{max}}$ denote the smallest $\zeta$ for which which the treatment probabilities are bound for all units, then $\pi^*(\zeta)$ is the expectation maximizing policy for all $\zeta\ge\zeta_{\text{max}}$.

Structure of the Pareto Frontier

Theorem (ref) provides detailed insights into the structure of the Pareto frontier. First, we define

i.e., $\mathcal{N}_u(\zeta)$ is the set of units for which $\pi_{n,t}^*(\zeta)$ is not bound to zero or one (but including the boundary case); we denote its complement by $\mathcal{N}_b(\zeta)=[N]\setminus\mathcal{N}_u(\zeta)$.

The behavior of the $z$-score is only slightly more complicated (though its proof is less trivial):

Before we provide a proof, we first describe the implications of Proposition (ref) & (ref) for the structure of the Pareto frontier, which is illustrated in Figure (ref). First, note that the observational policy is $\pi^o=\pi^*(0)$; thus, we have $\tau(0)=0$ (i.e., observational policy has zero improvement over itself) and $Z(0)$ is undefined since $s^2(0)=0$ (intuitively, we can never “prove” that the observational policy is better than itself). Next, define

Intuitively, $\zeta_{\text{min}}$ is the smallest value of $\zeta$ for which some unit has treatment probability bound to zero or one; analogously, $\zeta_{\text{max}}$ is the smallest value at which all units become bound to zero or one (so $\pi^*(\zeta_{\text{max}})\in\operatorname*{\arg\max}_{\pi}\tau(\pi)$ is the maximizer of the expected treatment effect).

By Proposition (ref), the $z$-score is constant for $\zeta\in(0,\zeta_{\text{min}}]$, and by Proposition (ref), the expected treatment effect is strictly increasing along this interval. Thus, $\zeta_{\text{min}}$ is a strictly better choice than the remainder of the interval, so we prune $(0,\zeta_{\text{min}})$ from the frontier. Now, Propositions (ref) & (ref) say that on $[\zeta_{\text{min}},\zeta_{\text{max}}]$, the $z$-score is strictly decreasing and the expected treatment effect is strictly increasing; thus, all policies $\pi^*(\zeta)$ for $\zeta_{\text{min}}\le\zeta\le\zeta_{\text{max}}$ are part of the frontier. Finally, the policy becomes constant for $\zeta\ge\zeta_{\text{max}}$, so we also prune $(\zeta_{\text{max}},\infty)$ from the frontier. Thus:

Discussion

Here, we discuss several implications of our results for policy learning algorithms targeting statistically significant performance improvements. While our discussion is based on the characterization in a specialized setting, we show that our key insights extend to the general setting in Section (ref); thus, our discussion considers the general case.

\paragraph{Comparing policy learning methodologies.}

Prior work has proposed ad hoc techniques for policy learning targeting statistical significance. They can be grouped into three strategies:

To compare these approaches, we introduce the following standard definitions.

One of our key insights is that policies along the Pareto frontier are largely stochastic.

Now, we consider the effect of treatment aggregation on the policy space. Treatment aggregation typically considers a large number of treatment options, but does not consider targeting---i.e., we have a finite policy class $\Pi=\{\pi^t\}_{t\in\mathcal{T}}$, where $\pi^t_{n,t'}=\mathbbm{1}(t'=t)$\footnote{Targeting can be in principle encoded by having the “treatment” vary depending on the unit, but this is not typical.}. In other words, policy $\pi^t$ deterministically assigns treatment $t$ to all units. Now, merging treatments $t_1,...,t_H\in\mathcal{T}$ is mathematically achieved by creating a new policy $\tilde{\pi}$ that uniformly randomly selects among these treatments---i.e.,

Then, $\pi^{t_1},...,\pi^{t_H}$ are removed from $\Pi$ and $\tilde{\pi}$ is added to $\Pi$. Note that $\tilde{\pi}$ is stochastic, which is consistent with our insight that Pareto optimal policies are stochastic. However, in our analysis in Section (ref), we show that the policy $\pi^*(\zeta_{\text{min}})$ is achieved by increasing $\zeta$ until $\pi_{n,t}^*(\zeta)$ binds to zero for any unit $n\in[N]$ and treatment $t\in\mathcal{T}$; thus, $\pi^*(\zeta_{\text{min}})$ typically randomizes over all treatments. It may be possible to improve treatment aggregation by randomly selecting among all treatments in a non-uniform way rather than merging a few high-value treatments.

Unit aggregation approaches also assume a finite policy class $\Pi$, where in this case each policy $\pi\in\Pi$ is a targeted policy; then, it identifies similar units $n_1,...,n_H\in[N]$ and imposes the constraint that $\pi_{n_h}=\pi_{n_{h'}}$ for all $h,h'\in[H]$. A typical strategy is to simply remove policies that do not satisfy this constraint from $\Pi$; thus, if $\Pi$ initially only contains deterministic policies, then unit aggregation never introduces any stochasticity into $\Pi$. As a consequence, we might expect these approaches to be more effective if $\Pi$ initially contains stochastic policies.

\paragraph{Statistical power.}

So far, we have been assuming the observational dataset is fixed; however, a common question in practice is how to collect the sample $\mathcal{O}$ to achieve a target level of significance. To address this question, we need to understand $Z(\zeta_{\text{min}})$, which is the best possible $z$-score that can be achieved along the Pareto frontier. In fact, we have an explicit formula for $Z(\zeta_{\text{min}})$:

This result follows from the proof of Proposition (ref). Given assumptions about $\mu_{n,0}$, $\mu_{n,1}$, and $\sigma_n$, it can be used to calculate the number of samples needed for a target level of statistical significance. It also quantifies the value of each unit toward statistical significance, which might help identify promising test subjects when recruitment is costly kouvelis_clinical_2017,tian_optimal_2023. For instance, suppose we further assume the following: (1) $\pi^o$ is an RCT policy (i.e., $\pi^o_{n,t}=1/2$ for all $n\in[N]$ and $t\in\mathcal{T}$), (2) the outcome variances are independent of treatment (i.e., $\sigma_{n,t}=\sigma_n$ for some $\sigma_n$), and (3) outcomes are centered (i.e., $\mu_{n,1}=-\mu_{n,0}=\mu_n$ for some $\mu_n$; for instance, the augmented IPW estimator aims to center outcomes). Then, we have

which is exactly the signal-to-noise ratio (SNR) for unit $n$. For instance, note that

From this result, we observe that the marginal value of increasing the SNR $Z_n$ of unit $n$ only depends on the current $z$-score $Z(\zeta_{\text{min}})$. Intuitively, the marginal value decreases proportional to $Z(\zeta_{\text{min}})$. Thus, it is always better to allocate resources to recruiting units with higher SNR, but the overall efficacy of increasing SNR has diminishing returns.

Finally, we can compare $\pi^*(\zeta_{\text{min}})$ to $\pi^*(\zeta_{\text{max}})$, which is given by

While $Z(\zeta_{\text{max}})$ can be complicated in general, it simplifies when $\pi^o$ is uniform.

This result follows from the proof of Proposition (ref). Under conditions (1), (2), and (3) above,

In particular, in this case, it is no longer true that $Z(\zeta_{\text{max}})$ can be expressed as in terms of the SNRs of individual units. In general, $Z(\zeta_{\text{min}})\ge Z(\zeta_{\text{max}})$; this inequality follows from Proposition (ref), but it also follows directly from Sedrakyan's inequality sedrakyan2018useful. It is tight when $\mu_n=\mu$ and $\sigma_n=\sigma$ are constant; indeed, in this case, it is easy to see that $\pi^*(\zeta_{\text{min}})=\pi^*(\zeta_{\text{max}})$. To understand when there is a gap, consider the following:

Multiple Treatment Setting

Next, we consider the general case of multiple treatments. Mirroring our analysis in the single-treatment analysis, we begin by computing the solutions to our optimization problem, which form a superset of the Pareto frontier (Section (ref)). Then, we establish the structure of the Pareto frontier, which is essentially identical to the structure in Figure (ref) for the single-treatment setting (Section (ref)). Finally, we provide results on the structure of policies along the Pareto frontier (Section (ref)) and the best possible statistical power (Section (ref)). Proofs for all claims are given in Appendix (ref).

Superset of the Pareto Frontier

The lack of a well defined treatment effect when comparing multiple treatments complicates the multiarm characterization of $\Pi_{\Lambda}$ ($K>1$) relative to the single-treatment setting. The policies in $\Pi_{\Lambda}$ can still be parameterized by the dual variable which uniquely generates them $\zeta$ and still form a path of policies from $\pi^o$ to the expectation maximizing policy. However, for each unit $n$, the rate of change for each non-zero propensity $\pi^*_{n,t}(\zeta)$ is impacted by which of the unit's treatments receive zero propensity. These inactive treatments do not impact the weighted average of treatments which act as a control for each treatment's relative impact, causing $\pi^*_{n,t}(\zeta)$ to change in a piece-wise linear manner as $\zeta$ increases. To overcome this issue, our general characterization of $\Pi_{\Lambda}$ first determines the value of $\zeta$ at which each propensity $\pi^*_{n,t}(\zeta)$ first hits $0$, and then determines the non-zero propensities in terms of the set of zero propensity treatments. Doing so allows us to first identify all of the knot points in the piecewise linear evolution of $\pi^*(\zeta)$, and then describe $\pi^*(\zeta)$ in terms of its locally linear form. First, we have:

Proposition (ref) ensures that for each $\omega^*$ which generates a solution, there exists a unique set $\mathcal{Z}_n(\omega^*) = \{t \mid \pi^*_{n,t}(\zeta^*) = 0\}$ of treatments which receive no assignments for unit $n$, and that $\mathcal{Z}_n(\omega')\subseteq\mathcal{Z}_n(\omega'')$ if $\omega'' \geq \omega' > 0$ both generate policies in $\Pi_{\Lambda}$. This provides enough structure to uniquely solve for $\kappa^*$ in Proposition (ref) for any $\omega^* \in \mathbb{R}_+$, generating the full solution to $\Pi_{\Lambda}$ (as $\omega^* < 0$ induces $\lambda < 0$), which can still be uniquely parameterized in terms of $\zeta^* \geq 0$.

To compute $\mathcal{Z}_n(\zeta)$ for each unit, we iteratively compute which treatment's propensity will first hit zero for each unit, remove that treatment and continue. An added benefit to the computation and analysis of $\pi^*(\zeta)$ is that each unit's treatment propensities can be calculated independently. We see that as $\zeta$ increases, $\pi^*(\zeta)$ moves from the observational policy to the expectation maximizing policy by iteratively eliminating suboptimal treatment arms for each unit.

The non-zero propensities of $\pi^*(\zeta)$ are determined by a term which redistributes zero propensity treatments and a linear term similar to the single-treatment characterization. This term replaces $\tau_n$ with a multi-treatment version which compares the expected outcome of unit $n$ in response to treatment $t$ to a weighted average of the expected outcome of the other active treatments. The variance of the conditional treatment effect is then replaced by $Q_n^{-1} \frac{\pi^o_{n,t}}{\mu_{n,t}^2 + \sigma^2_{n,t}}$. The treatment specific term here, which also determines how zero propensity treatments are redistributed as well as the weighted mean each treatment is measured against, is once again related to IPW's evaluation: $$\frac{\mu_{n,t}^2 + \sigma_{n,t}^2}{\pi^o_{n,t}} = \mathbb{E}\left[\left(\mathbbm{1}\left[T_n^o = t\right] \frac{Y_n^o}{\pi^o_{n,T_n^o}}\right)^2\right] = \mathrm{Var}\left(\mathbbm{1}\left[T_n^o = t\right] \frac{Y_n^o}{\pi^o_{n,T_n^o}}\right) + \mathbb{E}\left[\mathbbm{1}\left[T_n^o = t\right] \frac{Y_n^o}{\pi^o_{n,T_n^o}}\right]^2.$$ Thus, as $\zeta$ increases, $\pi^*(\zeta)$ evolves similarly to the single-treatment case; it compares each treatment against the other active treatments, and then changes each unit's policy to increase the policies impact while introducing minimal variance to an estimate of the policy's improvement.

Structure of the Pareto Frontier

As in the single-treatment setting, the expected improvement is strictly increasing as $\pi^*(\zeta)$ moves from the observational policy toward the expectation-maximizing policy, while the expected $z$-score is constant for small $\zeta$, and then strictly decreasing after any single unit has a treatment reach zero propensity. Define $\tau(\zeta)=\tau(\pi^*(\zeta))$, $s^2(\zeta)=s^2(\pi^*(\zeta))$, and $Z(\zeta)=\mathbb{E}[Z(\pi^*(\zeta))]=\frac{\tau(\zeta)}{s^2(\zeta)}$; also,

In other words, $\zeta_{\text{max}}$ is the largest $\zeta$ for which a treatment propensity first reaches $0$. For $\zeta\geq \zeta_{\text{max}}$ we know $\pi^*(\zeta)$ recovers the expectation maximizing policy; up until that point, $\pi^*(\zeta)$ increases the variance of its policy evaluation to improve its expectation as $\zeta$ increases. Similarly, $\zeta_{\text{min}}$ is the smallest $\zeta$ for which some unit is never assigned one of the treatments.

This result mirrors Corollary (ref) in the single-treatment case. The main difference is that only a single unit-treatment pair needs to bind to zero at $\zeta_{\text{min}}$; thus, all units may have stochastic treatments.

Structure of Policies on the Pareto Frontier

Next, we analyze how policies on the Pareto frontier assign treatments that trade off statistically significance vs. expected improvement. At a high level, as $\zeta$ increases, for each unit, policies along the frontier increase the propensity assigned to treatments with expected value above some unit-specific threshold, while decreasing the propensity assigned to the remaining treatments. In addition, a treatment's propensity will remain closer to that of the observational policy if that treatment's counterfactual outcome is often far from $0$. Thus policies on the frontier react conservatively to variance, choosing to adjust the propensities of treatments leading to certain outcomes sooner than uncertain ones.

Our first result shows that policies along the frontier decrease treatment propensities below the observational policy in order of their expected impact.

This result follows a simple logic: for any policy that does not obey this ordering, we can select a pair of offending treatments and contract their propensities towards the observational policy to both increase expected performance and decrease the standard error of the IPW estimator. As such we can understand policies on the frontier as setting a reservation value for each unit; increasing the frequency of treatments in the subgroup that outperform this reservation value while decreasing the frequency of treatments which underperform the reservation value.

However, we also show that policies along the frontier adjust treatments whose counterfactual outcomes are far from 0 more conservatively, meaning a treatment whose average impact is worse than an alternate may be assigned relatively more.

Note that under Proposition (ref), the conditions of Proposition (ref) can only occur if the frontier policy increases both treatments' propensities above that of the observational policy, or decreases both below that of the observational policy. Thus condition (1) says that if the frontier policy increases both treatments, it may increase the lower mean treatment relatively more if its average outcome is closer to zero, while condition (2) says that if the frontier policy decreases both treatments, it may decrease the lower mean treatment relatively less if its average outcome is closer to zero. Re-ranking treatments in this way reduces the variance of the IPW estimator as the estimator scales observed outcomes in the test set with a magnitude that increases with the difference between the frontier policy and the observational policy. If outcomes are far from zero, this scaling can introduce excessive variance that destroys the potential of the IPW estimator to achieve significance, leading policies on the Pareto frontier to adjust treatments with counterfactual outcomes close to zero more than other treatments.

Statistical Power

Finally, we compute the maximum expected z-score any policy can achieve, which is the analog of Corollary (ref) in the general setting.

Intuitively, each unit's contribution increases with the treatment effect for that unit and decreases with the variance in its treatment effect and with the magnitude of the mean outcomes.

Integration to a Policy Optimization Pipeline

Practitioners face two key challenges before they can use Theorem (ref) to propose a new policy. First, they do not know the true counterfactual outcome distribution and thus must form estimates of the counterfactual means $\hat{\mu}$ and variances $\hat{\sigma}^2$ to estimate $\hat{\Pi}_\Lambda$. Second, to avoid multiple testing concerns and streamline their decision process, practitioners need to select a single policy from $\hat{\Pi}_\Lambda$ to propose and evaluate on $\mathcal{O}$. We suggest practitioners use supervised machine learning algorithms on available training data to form counterfactual estimates and offer a method for selecting $\zeta$ to generate $\hat{\pi}^*(\zeta)$ in terms of the desired minimum $z$-score $Z_{\text{min}}$, and the desired improvement $\lambda$.

Estimating the Frontier Using Machine Learning

Forming $\Pi_{\Lambda}$ via Theorem (ref) requires knowledge about the observational dataset $\mathcal{O}$, namely the observational policy $\pi^o$ as well as the mean $\mu$ and variance $\sigma^2$ of the counterfactual outcomes. While practitioners often know $\pi^o$, the policy which generated $\mathcal{O}$, $\mu$ and $\sigma^2$ are fundamentally unknown when optimizing our policy. We suggest first learning estimates of the mean $\hat{\mu}$ and variance $\hat{\sigma}^2$ of the counterfactual outcomes in $\mathcal{O}$ using the practitioner's choice of machine learning algorithm and an available training dataset $\mathcal{T}$. Our approach can also be applied when $\pi^o$ is unknown by additionally estimating $\hat{\pi}^o$ from the training dataset.

To form the estimates $\hat{\mu}$ and $\hat{\sigma}^2$ we assume access to a training dataset $\mathcal{T} = \{(X_1,T_1,Y_1),\ldots,(X_M,T_M,Y_M)\}$ where each unit's treatment and outcome are accompanied by a set of individual covariates $X_m$. We assume that these covariates are similarly available in the observational dataset, expanding its definition to $\mathcal{O} = \{(\pi^o_1,X^o_1,T^o_1,Y^o_1),\ldots,(\pi^o_N,X^o_n,T^o_N,Y^o_N)\}$. The policy which generates the treatments for $\mathcal{T}$ need not be the same as $\pi^o$ unless the observational policy needs to be estimated from $\mathcal{T}$. However treatment assignment in $\mathcal{T}$ should follow standard causal inference assumptions such as positivity and selection on observables.

This training dataset can be used to form counterfactual outcome estimates $\hat{\mu}(x,t)$ and $\hat{\sigma}^2(x,t)$ in terms of a unit's covariates and any treatment. These estimates can be formed by running any supervised machine learning algorithm or regression $f$ on $\mathcal{T}$. If counterfactual variance estimates are difficult to extract from $f$ directly, they can be estimated from the squared residuals of $\mathcal{T}$, $\{(X_1,T_1,(Y_1 - \hat{\mu}(X_1,T_1))^2),\ldots,(X_M,T_M,(Y_M - \hat{\mu}(X_M,T_M))^2)\}$. In our simulations, we choose $f$ to be a set of $K$ honest regression forests, one for each treatment athey_generalized_2019. This approach allows for $f$ to form unbiased estimates non-parametrically when the covariate space is rich, without allowing cross-treatment pooling. Moreover, bag-of-bootstrap approaches allow forests to provide variance estimates directly, eliminating the need for prediction of a separate residual model Lu_varPred.

Once the counterfactual mean $\hat{\mu}$ and variance $\hat{\sigma}^2$ estimators are formed, we plug them into Theorem (ref) with $\mu_{n,t} \leftarrow \hat{\mu}(X^o_n,t)$ and $\sigma^2_{n,t} \leftarrow \hat{\sigma}^2(X^o_n,t)$ to generate an estimate of $\hat{\Pi}_\Lambda$, which contains the Pareto frontier for the counterfactual outcome distribution described by the estimators $\hat{\mu}$ and $\hat{\sigma}^2$. It almost certainly underperforms the true frontier of policies related to $\mu$ and $\sigma^2$. We note that as long as $\hat{\mu}$ and $\hat{\sigma}^2$ are reasonably close to the true parameters $\mu$ and $\sigma$, $\hat{\Pi}_\Lambda$ will extend from the observational policy in a similar direction as $\Pi_{\Lambda}$, ensuring similar policy suggestions for low values of $\zeta$.

The policies trained and evaluated by this pipeline depend on both the training dataset $\mathcal{T}$ and the observed test covariates $X^o$. A relevant concern is whether the IPW estimator still provides an unbiased estimate of a policies performance given this information. Under the assumptions of the Rubin Causal model we show that the IPW estimator is still unbiased for policies which depend on a training set and test covariates, $\pi(\mathcal{T},X^o)$.

Thus, while the learned policy itself may vary if the training dataset or test covariates change, IPW still provides an unbiased estimate of the learned policy on the observational dataset.

Selecting a Policy from the Optimal Frontier

In addition to generating an estimated frontier, practitioners need to select a policy from the frontier to implement. This problem can be reduced to choosing the value of $\zeta$ which generates $\hat{\pi}^*(\zeta) \in \hat{\Pi}_{\Lambda}$. This can be chosen via guess and check or some form of cross-validation, but this still requires some general idea of $\zeta$'s magnitude and may lead to inaccurate evaluation due to concerns with multiple hypothesis testing shaffer_multiple_1995. Instead, we propose to set $\zeta$ using the economic intuition its has as a dual variable.

In the optimization problems defining $\Pi_{\Lambda}$, $\zeta$ is the dual variable associated with the constraint $\tau(\pi) = \lambda$. Thus, $\zeta$ represents the marginal decrease that the squared standard error $s^2(\pi)$ can attain if $\lambda$ is decreased. By selecting $\zeta$, we set a limit on how much any policy change can increase the squared standard error in terms of its expected improvement. All expected policy improvements which can be made while introducing less standard error then this limit are included in $\pi^*(\zeta)$. This intuition offers a conservative way to set $\zeta$ for a practitioner hoping to achieve an improvement of $\lambda$ with a $z$-score of at least $Z_{\text{min}}$. From the proof of Proposition (ref), we have that if $\tau(\pi^*(\zeta)) = \lambda$, then $s^2(\pi^*(\zeta)) \leq \frac{\zeta}{2} \lambda$ due to the limit set by the dual variable $\zeta$ and the concavity of $\tau(\pi^*(\zeta))$ in $\zeta$. As a result, we have $$\mathbb{E}[Z(\pi^*(\zeta))] = \frac{\tau(\pi^*(\zeta))}{s(\pi^*(\zeta))} \geq \sqrt{\frac{2\lambda}{\zeta}}.$$ Thus, we can ensure $\mathbb{E}[Z(\pi^*(\zeta))] \geq Z_{\text{min}}$ by setting $$\zeta < \frac{2\lambda}{Z_{\text{min}}^2}.$$ Lower $\zeta$ will raise the probability of achieving a significant result, but lower the expected improvement. While the limit on $s^2(\pi^*(\zeta))$ is not tight if $\pi^*(\zeta) \in \Pi_{\text{Opt}}$, we recommend practitioners use this bound as our lack of knowledge with respect to the true counterfactual outcome distribution only allows construction of the inferior $\hat{\pi}^*(\zeta)$. Using this value of $\zeta$ ensures that $\hat{\pi}^*(\zeta)$ only attempts to make improvements that introduce allowable variance for the given significance level, limiting the concern that noise will drown out significant policy gains.

Simulations

We perform simulations exploring our method's ability to extract policy improvements when na\"{i}ve techniques fail. We apply our method to data from a megastudy encouraging flu vaccination using text message reminders milkman_680000-person_2022 and find a personalized policy that increases the vaccination rate by 4% while the na\"{i}ve approach fails to extract a significantly improved policy (Section (ref)). Then, we simulate data based on the refugee relocation study of bansak_improving_2018, showing our method can be effectively utilized in settings with additional policy constraints (Section (ref)). In Appendix (ref) we provide additional simulations in a stylized environment in to explore how the accuracy of the prediction model given to our method impacts the performance of learned policies and compare our method to other policy optimization techniques.

Case Study: Vaccine Nudge Megastudy

Behavioral health interventions, such as text message reminders, offer a nearly costless mechanism for practitioners to improve health outcomes such as vaccination rates chapman_opting_2010,milkman_using_2011,regan_randomized_2017---many people intend to get vaccinated, but fail to follow through sheeran_intentionbehavior_2002. Effective reminders can motivate individuals to act on their intentions, improving vaccination rates with minimal cost. We demonstrate that IAPO can be used to learn provably effective policies for targeting these kinds of nudges. We use data from milkman_680000-person_2022, who performed a megastudy (i.e., a large-scale randomized control trial) testing the efficacy of 22 randomly chosen text messages and a control condition at encouraging 680,000 patients at a large US retail pharmacy to receive the flu vaccine during the 2019-2020 flu season.

The data includes the vaccination outcome of each patient (we use vaccinations before October 31, 2019) and their treatment assignment. In addition, it includes patient covariates related to their locations and health records, including (i) age, (ii) gender, (iii) zip code demographics, (iv) past vaccinations, and (v) hypertension. These covariates offer the opportunity to target the choice of text messages---different patients may respond to different motivations, offering the potential for large gains from a personalized treatment. We split the study subjects into training and test samples, train a set of honest random forests to predict each treatment's counterfactual outcomes, use the forests to estimate the Pareto frontier of policies via IAPO, then estimate their performance on the test samples using IPW. After removing samples with missing covariates we maintain 474,667 samples for training and 203,429 samples for testing. In total 25.6% of the patients in our (combined training and test) sample received a vaccination by our cutoff.

Figure (ref) shows both the performance improvement the honest random forest model believes that policies on the estimated frontier can achieve (left) and the evaluation of those policies by IPW (right). It believes that na\"{i}vely choosing the “winning” treatment for each patient should yield statistically significant 20% increase in vaccination rate. However, when properly evaluated using IPW, this policy yields an improvement of 3.5%, and is not statistically significantly better. In contrast, IAPO learns a policy with a similar 3.5% increase in vaccination rate, which is statistically significantly better than the baseline policy at the 95% confidence level. Note that IAPO also suffers from the winner's curse, since it is highly optimistic about performance along the frontier; however, it provides the regularization needed to obtain improvements that hold up under rigorous evaluation.

Table (ref) describes how policies learned by IAPO achieve significant improvements when na\"{i}ve optimization of our model fails. We focus on two policies from the frontier; 95% confident refers to a policy which achieves a similar improvement estimate to the na\"{i}ve approach, but passes a significance test with 95% confidence while 99% confident refers to a policy which achieves half the improvement but passes a significance test with 99% confidence. All policies essentially eliminate the control treatment, but are highly personalized, with the most frequently used treatment given to less than 9% of patients. About half of the units are assigned a deterministic treatment by the 95% confident policy, while the 99% confident policy use stochastic treatments assignments for all units. This increases the overlap of both policies compared to the na\"{i}ve policy, helping them achieve lower standard errors on their improvement estimates and allowing them to pass significance tests at higher confidence levels. The 95% confident policy shows that the na\"{i}ve approach can be improved while only making half of units receive nondeterministic treatments while the 99% confident policy shows that much larger gains in significance require much greater overlap with the observational policy, assigning no treatments deterministically.

Case Study: Refugee Relocation Simulation

In the United States, refugee resettlement agencies are responsible for choosing where to settle incoming refugees. Their goal is to choose an assignment policy $\pi$ that places refugees in available locations in a welfare maximizing way, often measured by outcomes such as achieving employment within one year of settling. Assignment policies must not overload locations' capacities (due to limited resources available to settle refugees at each location). We perform a simulation study to evaluate the efficacy of IAPO at policy learning when varying the best possible performance gains from personalization.

Specifically, we simulate a population of refugees matching the descriptive statistics of the population studied by bansak_improving_2018. We generate 33,782 samples with characteristics of (i) location assigned, (ii) case status (free/fixed assignment), (iii) age, (iv) nationality, (v) education level, and (vi) arrival date. Each sample was assigned characteristics independently at random with probabilities matching the descriptive statistics given in Table S1 of bansak_improving_2018. Age and nationality information were reported at a coarsened level, so we drew those characteristics uniformly at random from coarsened buckets. We refer to characteristics (iii)-(vi) of sample $n$ as $X_n$ and index different characteristics using subscript $j$. We use case status to determine whether samples are used for model training (24,327 samples) or testing (9,455 samples).

We simulate an employment outcome for each sample $n$ independently according to a bernoulli distribution whose mean $\mu_{n,t}$ depends on the sample's characteristics $X_n$ and assigned location $t$. We standardize the characteristics $X_n$ to $X_n^\dagger$ by subtracting the average of each characteristic and dividing by that characteristic's standard deviation across all samples. Then, we generate a fixed causal effect for each covariate $\alpha_j$ and each covariate-location interaction $\beta^{(t)}_j$; these interactions are what enable gains from personalization. We determine $\mu_{n,t}$ using a logistic function whose input is determined by a linear interaction of $\alpha$ and $\beta^{(t)}$ with $X_n^\dagger$:

The parameters $c_{unit}$ and $c_{intx}$ allow us to control the magnitude of the causal effects we introduce while $c_{tot}$ allows us to change the average employment outcome. Increasing $c_{unit}$ increases the heterogeneity in refugees' baseline employment probabilities independent of location while increasing $c_{intx}$ increases the heterogeneity in the impact each location has on each refugee's employment probability. Table (ref) gives the parameters used to generate 3 datasets with varying magnitudes of covariate-location interaction effects. For each, we measure cross-unit heterogeneity as the standard deviation across samples employment probability (averaged over locations) and within-unit heterogeneity as the average standard deviation of samples employment probability across locations:

Our simulated refugee datasets maintain a consistent average outcome and cross-unit heterogeneity while varying the impact locations can have on a refugee's employment.

One challenge applying IAPO to this setting is the presence of resource constraints that limit the set of feasible refugee assignment policies. Each location $t$ can accept a maximum number of refugees $d_t$ given their available resources. We consider a policy feasible if, in expectation, it assigns fewer than $d_t$ refugees to location $t$, as represented by the constraint $$\sum_{n=1}^N \pi_{n,t} \leq d_t \ \forall \ t.$$ Meeting the assignment constraints themselves can then be included as part of the randomization process or can be ignored if the constraints are approximate or somewhat flexible. For our simulations, we set each locations capacity $d_t$ as $5/4$ of the number of freely assigned refugees (our test samples) settled in location $t$. This choice allows proposed policies to slightly increase the number of refugees sent to any given location while requiring the majority of gains to come from identifying refugee-location synergy.

We integrate resource constraints into the IAPO pipeline as linear constraints to the convex optimization problems defining $\Pi_{\Lambda}$. Let $D =\left \{\pi\in(\Delta^K)^N\mid \sum_{n=1}^N \pi_n \leq d \right\}$ be the set of feasible policies according to our capacities $d$. Then, our Pareto frontier is $$\Pi_{\text{Opt}} =\left\{ \pi \in D\biggm\vert \not\exists \ \pi'\ \in D \text{ s.t. } \tau(\pi') \geq\tau(\pi)\wedge\frac{\tau(\pi')}{s(\pi')} \geq \frac{\tau(\pi)}{s(\pi)}\wedge(\tau(\pi'),s(\pi')) \neq (\tau(\pi),s(\pi))\right\},$$ which can be relaxed to the family of convex optimization problems given by $$\Pi_{\Lambda} = \bigcup_{\lambda \in [0,\tau_{\text{max}}]} \operatorname*{\arg\min}_{\{\pi \in (\Delta^K)^N\mid\tau(\pi) = \lambda, \sum_{n=1}^N \pi_n \leq d\}} s^2(\pi).$$ Any policy constraint which can be represented as a valid convex programming constraint in terms of the policy's propensities can be incorporated into IAPO in the same way as resource constraints.

Instead of characterizing the resource constrained $\Pi_{\Lambda}$, we approximate it using convex optimization software. Specifically, we implement the convex program defining $\Pi_{\Lambda}$ in CVXPY and solve it using the Splitting Conic Solver (SCS) while varying the value of $\lambda$ to solve for $\Pi_{\Lambda}$ point-wise. We use a convergence tolerance of $10^{-6}$ to achieve precise results; however, as our objective approaches zero for policies close to the observational policy, we only solve for policies with sufficiently high expected improvement $\lambda$ to ensure stability of the solver's estimate.

We train a set of honest random forest regressions on each of our three refugee datasets to estimate the Pareto frontier of policies for each one. Figure (ref) reports IPW performance estimates using held-out test sets, as well as the predicted and actual performance of the corresponding expectation-maximizing policies. For all three datasets, the winner's curse causes the model to believe it can obtain massive employment gains regardless of the actual potential to improve employment. In contrast, when useful within-unit heterogeneity exists, IAPO obtains gains with much higher confidence and at least as much improvement as the na\"{i}ve approach. In more detail, when policy improvement is impossible (“No Interaction Effect”), no policy along the frontier obtains significant policy improvement, which is expected. For the “Small Interaction Effect” dataset, the na\"{i}ve approach fails to obtain statistically significant policy improvement, whereas IAPO obtains policy improvements up to the 99% confidence level while maintaining similar gains as the na\"{i}ve approach. For the “Medium Interaction Effect” dataset, the na\"{i}ve approach just barely obtains significant improvements at the 95% confidence level; in contrast, IAPO readily obtains similar improvements at a much higher 99% confidence level.

Conclusion

We have proposed a novel policy optimization strategy that aims to identify the full Pareto frontier of policies trading off expected performance with provably significant performance improvements. Focusing on the specific setting where policy evaluation is performed using IPW, we have mathematically characterized the Pareto frontier. Our characterization reveals that focusing on statistical significance prioritizes elimination of bad treatment assignments over selecting the best treatment assignments. Intuitively, rather than try and identify optimal treatments, a decision-maker focuses on the problem of eliminating bad treatments, which can be simpler while yielding confident improvements. Based on our theoretical analysis, we devise a policy optimization algorithm using a plug-in approach. Practitioners first use their choice of supervised machine learning algorithm to train a counterfactual prediction model on the observational data, and then use its predictions to compute the Pareto frontier of optimal policies. Our simulations show that our plug-in approach can extract significant policy improvements, while other common policy optimization approaches cannot, especially when the underlying causal model is complex.

We leave a number of important directions for future work. First, while we have performed some preliminary analysis of how prediction errors affect estimation of the Pareto frontier, significantly more work is required to understand the impact. Our work also assumes the importance weights are known, whereas they may need to be estimated; understanding how IPW estimation error affects our procedure is another important direction. One key benefit of our convex formulation is that additional constraints (e.g., budget constraints) can be incorporated; it may be interesting to explore the impact of these constraints on the characterization of the Pareto frontier. Finally, we have focused on IPW; exploring how to perform policy optimization in conjunction with other policy evaluation methodologies (e.g., randomized controlled trials, instrumental variables, etc.) deserves similar attention.