More Efficient Policy Learning via Optimal Retargeting

TL;DR

This paper introduces a retargeting approach to improve policy learning from observational data by changing the population focus, reducing bias, and enhancing policy performance across various applications.

Contribution

It characterizes the optimal retargeting weights and reference policies to minimize estimation variance, providing a practical method to improve policy learning.

Findings

Retargeting can significantly improve policy learning performance.

The method reduces bias without introducing substantial new bias.

Empirical results show improved policies in simulations and real-world case studies.

Abstract

Policy learning can be used to extract individualized treatment regimes from observational data in healthcare, civics, e-commerce, and beyond. One big hurdle to policy learning is a commonplace lack of overlap in the data for different actions, which can lead to unwieldy policy evaluation and poorly performing learned policies. We study a solution to this problem based on retargeting, that is, changing the population on which policies are optimized. We first argue that at the population level, retargeting may induce little to no bias. We then characterize the optimal reference policy and retargeting weights in both binary-action and multi-action settings. We do this in terms of the asymptotic efficient estimation variance of the new learning objective. Extensive empirical results in a simulation study and a case study of personalized job counseling demonstrate that retargeting is a fairly easy way to significantly improve any policy learning procedure applied to observational data.

Figures

Tables

Table 1. Table 1: Average policy values of different policy learning methods applied to the job counseling dataset, with standard errors over replications.

	Average Policy Value (in 1000’s)
Method	Conventional	Retargeted	Improvement
DC	$-3.44\pm 0.005$	—	—
DM	$-3.42\pm 0.005$	$-3.42\pm 0.005$	0%
DR	$1.14\pm 0.006$	$1.93\pm 0.006$	70%
IPW	$2.09\pm 0.006$	$2.65\pm 0.005$	27%

Equations

$$V(\pi )=\mathbb{E}\left[\sum _{a\in \mathcal{A}}\pi (a\mid X)Y(a)\right]=\mathbb{E}\left[\sum _{a\in \mathcal{A}}\pi (a\mid X)\mu (a\mid X)\right].$$

$$V(\pi ;w)=\mathbb{E}\left[w(X)\sum _{a\in \mathcal{A}}\pi (a\mid X)\mu (a\mid X)\right].$$

$$R(\pi ;w,\rho )=\mathbb{E}\left[w(X)\sum _{a\in \mathcal{A}}(\pi (a\mid X)-\rho (a\mid X))\mu (a\mid X)\right].$$

$${\Pi }^{\ast }=\underset{\pi \in \Pi }{\mathrm{argmax}}V(\pi ).$$

$${\Pi }_0^{\ast }=\underset{\pi \in {\Pi }_0}{\mathrm{argmax}}V(\pi ).$$

$${\Pi }_0^{\ast }(w,\rho )=\underset{\pi \in {\Pi }_0}{\mathrm{argmax}}R(\pi ;w,\rho ).$$

$$V_n(\pi )$$

$$R_n(\pi ;w,\rho )$$

$${\Pi }_n^{\ast }=\underset{\pi \in \Pi }{\mathrm{argmax}}{V}_n(\pi ),{\Pi }_{0,n}^{\ast }(w,\rho )=\underset{\pi \in {\Pi }_0}{\mathrm{argmax}}{R}_n(\pi ;w,\rho ).$$

$$\displaystyle\psi(x,a,y;\phi_{0},\mu_{0})=w(x)\Biggl{(}$$

$$\displaystyle\qquad\qquad\qquad+\frac{\pi(a\mid x)-\rho(a\mid x)}{\phi_{0}(a\mid x)}(y-\mu_{0}(a\mid x))\Biggr{)}.$$

$${\hat{R}}_n(\pi ;w,\rho )={\tilde{R}}_n(\pi ;w,\rho )+o_p(1/\sqrt{n}),\text{where}{\tilde{R}}_n(\pi ;w,\rho )=\frac{1}{n}\sum _{i=1}^n\psi (X_i,A_i,Y_i;\varphi ,\mu ).$$

$$n\mathrm{Var}\mathbb{E}\left[{\tilde{R}}_n(\pi ;w,\rho )\mid X_{1:n}\right]$$

$$=\mathrm{Var}\left(w(X)\sum _{a\in \mathcal{A}}(\pi (a\mid X)-\rho (a\mid X))\mu (a\mid X)\right),$$

$$n\mathbb{E}\mathrm{Var}\left({\tilde{R}}_n(\pi ;w,\rho )\mid X_{1:n}\right)$$

$$=\mathbb{E}\left[w^2(X)\sum _{a\in \mathcal{A}}\frac{{\sigma }^2(a\mid X)}{\varphi (a\mid X)}{\left(\pi (a\mid X)-\rho (a\mid X)\right)}^2\right],$$

$$\frac{\sqrt{n}({\hat{R}}_n-R_n(\pi ;w,\rho ))\text{buildrel}d}{⟶\mathcal{N}\left(0,\mathbb{E}\left[w^2(X){\sum }_{a\in \mathcal{A}}\frac{{\sigma }^2(a\mid X)}{\varphi (a\mid X)}{\left(\pi (a\mid X)-\rho (a\mid X)\right)}^2\right]\right).}$$

$$\Omega (w,\rho )$$

$$\underset{\pi ,{\pi }^{\prime }\in {\Pi }_0}{\sup }\left|({\tilde{V}}_n(\pi ;w)-{\tilde{V}}_n({\pi }^{\prime };w))-(V_n(\pi ;w)-V_n({\pi }^{\prime };w))\right|$$

$$\le C\left(\kappa ({\Pi }_0)+1+\sqrt{\log (1/\delta )}\right)\sqrt{\frac{\Omega (w,\rho )}{n}}+o\left(\frac{\log (1/\delta )}{\sqrt{n}}\right)$$

$$\Omega (w,\rho )$$

$${\rho }_0$$

$$({\rho }_0,w_0)$$

$$\Omega (w_0,{\rho }_0)$$

$$R(\pi ;w_0,{\rho }_0)$$

$$w_0(x)$$

$$\Omega (w_0,{\rho }_0)$$

$$R(\pi ;w_0,{\rho }_0)$$

$$\xi (x)$$

$${\rho }_0(a\mid x)$$

$$\kappa (x)$$

$$w_0(x)$$

$${\rho }_0$$

$$({\rho }_0,w_0)$$

$$\Omega (w_0,{\rho }_0)$$

$$R(\pi ;w_0,{\rho }_0)$$

$${\rho }_0(a\mid x)=\frac{1}{2}(1-(m-2)\varphi (a\mid x)),$$

$$w_0(x)\propto {\left(\sum _{a\in \mathcal{A}}\frac{1}{\varphi (a\mid x)}+\frac{m}{2}-1\right)}^{-1}.$$

$$X^{\prime \prime }$$

Full text

Keywords: individualized treatment regimes, overlap, efficient policy learning, optimization.

1 Introduction

Personalized intervention policies, also known as individualized treatment regimes, are having increasing impact in domains such as education (Mandel et al., 2014), healthcare (Bertsimas et al., 2017), and public policy (Kube et al., 2019). In many of these domains exploration is costly, unethical, or otherwise prohibitive, and so it is crucial to learn new policies using existing observational data. The potential for this can be big: in medicine, for example, the widespread adoption of electronic medical records provides very rich and plentiful data that can help drive the future of personalized care (Kosorok & Laber, 2019).

Learning new policies from observational data relies on our ability to evaluate a counterfactual: on average, how would an individual have fared under a treatment different than that given in the past. If we have rich enough covariates we may be able to defend an assumption of unconfounded treatment assignment, in which case the solution is essentially to compare like to like: on average the individual would fare like a similar individual who received the alternative treatment. For this to be possible, of course, a like comparison must be available in the data. The lack of sufficiently similar comparisons is known as limited overlap and it can lead to highly unwieldy policy value estimates and poorly performing learned policies in practice.

In this paper we study one way to address this problem: change the population on which we optimize the policy, which we term retargeting. We demonstrate that in many cases the optimal policy can actually be characterized as the population minimizer of a variety of population objective functions, so it stands to reason that we should use the one that is easiest to estimate. Each objective is centered with respect to a different reference policy and averaged over a different weighted subpopulation. We characterize the estimability of these objectives in terms of their efficiency bounds, that is, the minimal asymptotic variance that a consistent regular estimator can have in estimating them. We show how to optimize these bounds in the binary- and multiple-action cases. We demonstrate these new tools empirically in an extensive simulation study and in an application to personalized job counseling. The results show that retargeting can offer significant benefits and is fairly robust. In sum, retargeting is an easy way to improve any policy learning procedure using observational data where limited overlap is a major concern.

1.1 Problem set up

A policy $\pi$ is a map from the space $\mathcal{X}$ of baseline covariates to a probability distribution over the action space $\mathcal{A}$. Specifically, let $\Delta^{\mathcal{A}}=\left\{p\in\mathbb{R}^{\mathcal{A}}_{+}:\sum_{a\in\mathcal{A}}p(a)=1\right\}$ be the space of probability distributions over actions and let $\Pi=\left(\Delta^{\mathcal{A}}\right)^{\mathcal{X}}$ be the space of all maps $\mathcal{X}\to\Delta^{\mathcal{A}}$. We index an element $\pi\in\Pi$ by letting $\pi(a\mid x)$ be the $a$ coordinate in the image of $x$ under the map $\pi$, i.e., the probability of taking action $a$ when seeing baseline covariates $x$.

The population of decision instances is described by covariates $X\in\mathcal{X}$ and potential outcomes $(Y(a))_{a\in\mathcal{A}}\in\mathbb{R}^{\mathcal{A}}$. That is, in each decision instance, first the covariates $X$ are revealed, then we take some action $a\in\mathcal{A}$, and then we get the reward $Y(a)$. Let $\mu(a\mid X)=\mathbb{E}\left[{Y(a)\mid X}\right]$ and $\sigma^{2}(a\mid X)=\operatorname{Var}(Y(a)\mid X)$. The value of a policy $\pi$ is the average reward we get over the population of instances when we choose our action from the distribution given by $\pi(\cdot\mid X)$:

[TABLE]

Note that the last expectation is only over the variable $X$, where the inside of the expectation is the expected reward within the context of seeing an instance with covariates $X$. The best possible policy value is $V^{*}=\max_{\pi\in\Pi}V(\pi)$.

The data available for learning policies consists of $n$ observations of covariates, action, and reward, $X_{1},A_{1},Y_{1},\dots,X_{n},A_{n},Y_{n}$. We drop subscripts to indicate a generic draw. We assume that the observed reward is the potential reward of the observed action, $Y=Y(A)$, encapsulating non-interference and consistency, also known as stable unit value assumption (SUTVA; Rubin, 1980). We further assume that the data is mean-unconfounded in that $\mu(a\mid X)=\mathbb{E}\left[{Y\mid A=a,X}\right]\,\forall a\in\mathcal{A}$. We let $\phi(a\mid X)=\mathbb{P}\left(A=a\mid X\right)$ be the propensity score.

1.2 Related Literature

Policy learning.

Since an optimal policy is given by $\pi(a^{*}(X)\mid X)=1$ for any $a^{*}(X)\in\operatornamewithlimits{argmax}_{a\in\mathcal{A}}\mu(a\mid X)$, one approach to policy learning is to fit a regression estimate $\hat{\mu}$ and plug it into $a^{*}$ above. Under unconfoundedness, $\hat{\mu}$ can be estimated by regressing the observed outcome $Y$ on $A,X$. This policy is called direct comparison (DC) and in the absence of additional assumption or constraints it enjoys certain minimax guarantees (Hirano & Porter, 2009; Stoye, 2009). However, often we are interested in a more structured policy, but at the same time have to deal with confounding. For example, we might want the policy to not consider ethnoracial features but we do need to consider them in accounting for confounding factors. Additionally, we may want the policy to be simple and interpretable, like a separating hyperplane or a decision tree. In such cases, we want to restrict ourselves to a constrained policy class $\Pi_{0}\subseteq\Pi$. This may also afford us additional efficiency as we are only truly interested in a potentially much lower dimensional object than $\mu$ itself.

In such cases, various methods have been proposed, which estimate $V(\pi)$ and then optimize this estimate over $\pi\in\Pi_{0}$ (we embed any potential regularization implicitly as constraints on $\Pi_{0}$ itself). Estimates of $V(\pi)$ are similar to estimates of averages with missing-at-random data or of average treatment effects under unconfoundedness. One policy learning approach is the direct method (DM), which uses the regression-adjusted estimate $\hat{V}_{n}(\pi)=\frac{1}{n}\sum_{i=1}^{n}\sum_{a\in\mathcal{A}}\pi(a\mid X_{i})\hat{\mu}(a\mid X_{i})$ given a fitted $\hat{\mu}$ (Qian & Murphy, 2011). Another is to use inverse propensity weighting (IPW), which, given a fitted $\hat{\phi}$, estimates $\hat{V}_{n}(\pi)=\frac{1}{n}\sum_{i=1}^{n}\pi(A_{i}\mid X_{i})Y_{i}/\hat{\phi}(A_{i}\mid X_{i})$ (Beygelzimer & Langford, 2009; Li et al., 2011). When $\Pi_{0}$ has bounded capacity, bounds on the regret of the learned policy compared to the a-priori-best policy in $\Pi_{0}$ can be obtained via uniform concentration inequalities on $\hat{V}_{n}(\pi)$ (Kitagawa & Tetenov, 2018). Optimizing such estimates $\hat{V}_{n}(\pi)$ is a form of cost-sensitive classification and can be approached using convex surrogate losses (Zhao et al., 2012). In the absence of unconfoundedness, $V(\pi)$ may be unidentifiable but well-performing policies may still be learned (Kallus & Zhou, 2018a).

To deal with the instability of inverse probability weights, normalization (Hájek estimates) and clipping are often employed (Swaminathan & Joachims, 2015b). Balanced policy learning (Kallus, 2018) uses optimal balance (Kallus, 2016) to deal more directly with this instability. Doubly robust (DR) policy learning (Dudík et al., 2011) uses the augmented IPW (AIPW) estimator (Robins et al., 1994). Using cross-fitted models makes this estimate efficient under mild conditions (Chernozhukov et al., 2018), which has been shown to lead to better regret bounds (Athey & Wager, 2017). IPW and DR can also be generalized to continuous action spaces (Chen et al., 2016; Kallus & Zhou, 2018b).

We further discuss the statistical efficiency of various policy value estimates and its relationship to policy learning in Section 2.2.

Dealing with limited overlap in average treatment effect estimation.

Overlap refers to the extent to which the covariate distributions conditioned on different observed actions overlap with one another in the data, which controls the availability of like comparisons. As richer covariates are used to better support the plausibility of unconfoundedness and to enable more pinpoint personalization of policies, the issue of overlap becomes more dire (D’Amour et al., 2017). Overlap is often quantified by how close to 0 or 1 the propensities $\phi(a\mid X)$ are at any point $X$. Lack of overlap leads to a fundamental difficulty in estimating causal estimands (LaLonde, 1986; Heckman et al., 1997; Dehejia & Wahba, 1999). In the estimation of $\hat{\mu}$ it can be understood as a covariate shift between the distribution where $\hat{\mu}(\cdot\mid a)$ is fit and where it is deployed. This shift means that parametric models will precariously extrapolate from areas of high density to areas of low density, while non-parametric models will have little local information to drive predictions. In IPW and AIPW estimates, the size of the density ratio $\pi(A\mid X)/\phi(A\mid X)$ drives the variance and grows big as overlap diminishes.

In the context of estimating average treatment effects, i.e., $\mathbb{E}\left[{Y(+)-Y(-)}\right]$ or$\mathbb{E}\left[{Y(+)-Y(-)\mid A=+}\right]$ with $\mathcal{A}=\{-,+\}$, various approaches to dealing with this fundamental issue have been investigated. Clipping case weights (Ionides, 2008) or regularizing them (Kallus, 2016; Santacatterina & Bottai, 2018) shrinks the estimate toward the unadjusted one to trade bias for variance. In the context of matching, another approach is to drop units without good matches (Iacus et al., 2011; Rubin, 2010; Cochran & Rubin, 1973). Similarly, one can drop units with extreme propensities (Dehejia & Wahba, 1999; Heckman et al., 1997; Smith & Todd, 2005). Both limit the estimates to a subpopulation with better overlap. Crump et al. (2006) formalize this notion and frame it as choosing a subpopulation so to optimize (part of) the efficiency bound of the corresponding estimand. They show that thresholding the propensity score gives the best subpopulation defined as a subset of $\mathcal{X}$ and they also characterize the best weighted subpopulation, which under homoskedasticity is reweighted by $\phi(+\mid X)\phi(-\mid X)$. Li et al. (2018) consider these latter weights and show that they correspond to certain balancing conditions. All of these lead to retargeted estimated effects that are local to a subpopulation different from the study one.

In the context of policy learning, study of the ramifications of limited overlap and ways to address it has been limited. Swaminathan & Joachims (2015a) regularize to force the learned policy to be close to $\phi$ in order to control the size of density ratios. Moreover, weight clipping (Swaminathan & Joachims, 2015a) and regularization (Kallus, 2018) are generally also used in policy learning. However, the modification of the target population on which policies are optimized has not been previously studied and neither has its optimal choice. We will here show it holds particularly unique promise for policy learning, which may incur little to no asymptotic bias due to retargeting, where a well performing policy is much more important than an unbiased or interpretable population estimate, and where some train-test shift is often inevitable anyway.

2 Centering and Retargeting the Policy Learning Objective

In order to alleviate the issues of lack of overlap we will consider changing the objective that we optimize. Specifically, we will consider reweighting it. For every $w\in\mathbb{R}^{\mathcal{X}}_{++}$, define the $w$-weighted population policy value as

[TABLE]

This retargets policy learning to the subpopulation weighted by $w$. We will seek the weights $w$ for which $V(\pi;w)$ is in some sense most convenient from the point of view of estimation. Because the scale of our objective function does not matter, we will restrict our attention to $w$ satisfying $\mathbb{E}\left[{w(X)}\right]=1$.

But, in fact, the optimizer over $\pi$ of $V(\pi;w)-V(\rho;w)$ remains the same over choice of a fixed reference policy $\rho$, and the same is true of $\hat{V}_{n}(\pi;w)-\hat{V}_{n}(\rho;w)$ for any value estimator. Therefore, it is actually not immediately clear the estimability of which objective we should be considering. We therefore will also seek the best reference policy $\rho$ that makes the centered value the easiest to estimate. Even though centering by a reference policy has no impact on the policy learning, finding the best reference policy allows us to correctly characterize the estimability of different retargeted learning objectives.

Specifically, for every $w\in\mathbb{R}^{\mathcal{X}}_{++}$ and $\rho\in\mathbb{R}^{\mathcal{A}\times\mathcal{X}}$, define

[TABLE]

We interpret this as the additional average value of $\pi$ compared to $\rho$ over the retargeted population where contexts are reweighted by $w(X)$. Letting $\bm{1}(X)=1$ we recognize $R(\pi;\bm{1},\rho)=V(\pi)-V(\rho)$. Note that we do not actually constrain $\rho$ to be a policy, that is, its image need not be nonnegative and sum to one. We therefore also recognize $V(\pi;w)=R(\pi;w,\bm{0})$.

In the next sections we will first consider how this new objective impacts the learning problem and then discuss its efficient estimation from data for a fixed $w,\rho$.

2.1 The Effect of Retargeting on Policy Learning

In policy learning, we seek a policy that optimizes $V(\pi)$. Define the corresponding set of optimal policies as

[TABLE]

Note that $\pi^{*}\in\Pi^{*}$ if and only if $\pi^{*}(a\mid X)(\max_{a^{\prime}\in\mathcal{A}}\mu(a^{\prime}\mid X)-\mu(a\mid X))=0\,\forall a\in\mathcal{A}$ almost surely. That is, $\pi^{*}(\cdot\mid X)$ can only support actions with largest conditional mean reward given $X$.

As discussed in Section 1.2, in practice we only optimize over a restricted policy class, $\Pi_{0}\subseteq\Pi$. Let the best policies in $\Pi_{0}$ be denoted by

[TABLE]

We next consider the question of what happens if we instead optimize our alternative objective $R(\pi;w,\rho)$ over $\Pi_{0}$. Let

[TABLE]

First, notice that $R(\pi;w,\rho)$ decomposes as a sum of the weighted value of $\pi$ minus the weighted value of $\rho$. Therefore, as also discussed above, the choice of $\rho$ has no influence on the optimizer of $R(\pi;w,\rho)$ no matter what set of policies one optimizes over, as it only amounts to constant shift in the objective. That is, $\Pi_{0}^{*}(w,\rho)=\Pi_{0}^{*}(w,\rho^{\prime})$. In particular, we immediately see that $\Pi_{0}^{*}(\bm{1},\rho)=\Pi_{0}^{*}$ for any $\rho$.

Next, we consider the effect of weighting. As noted above, an optimal unconstrained policy simply selects the best action in each context. If we reweight the contexts, we will arrive at the same conclusion. But this may not be true of a constrained optimal policy. We next establish that if there exists some policy in $\Pi_{0}$ that is also optimal in an unconstrained way, then neither retargeting nor centering affects policy learning over $\Pi_{0}$ at the population level.

Lemma 2.1**.**

Suppose $\Pi^{*}\cap\Pi_{0}\neq\varnothing$. Then $\Pi_{0}^{*}(w,\rho)=\Pi^{*}\cap\Pi_{0}$ for every $w\in\mathbb{R}^{\mathcal{X}}_{++}$, $\rho\in\mathbb{R}^{\mathcal{A}\times\mathcal{X}}$. In particular, if $\pi$ is a solution to Eq. 3 then $V(\pi)=V^{*}$.

Lemma 2.1 shows that if $\Pi_{0}$ is well-specified, in that it contains one of the true unconstrained optimal policies, then policy learning over the class $\Pi_{0}$ is invariant to retargeting. In other words, there are many objectives that characterize the best policy in $\Pi_{0}$. This, of course, depends on well-specification; we discuss the case of misspecification in Section 5.2.

A result akin to Lemma 2.1 also applies in a given sample. Let

[TABLE]

respectively be the value of $\pi$ on the given sample and the corresponding sample version of $R(\pi;w,\rho)$. Similarly, let

[TABLE]

Lemma 2.2**.**

Suppose $\Pi^{*}\cap\Pi_{0}\neq\varnothing$. Then $\Pi_{0,n}^{*}(w,\rho)=\Pi^{*}_{n}\cap\Pi_{0}$ for every $w\in\mathbb{R}^{\mathcal{X}}_{++}$, $\rho\in\mathbb{R}^{\mathcal{A}\times\mathcal{X}}$ almost surely.

Note that $\Pi^{*}\subseteq\Pi^{*}_{n}$ almost surely. The conclusion from Lemma 2.2 is that if $\Pi_{0}$ is well-specified then optimizing the oracle sample objective, which uses perfect counterfactual predictions to score policies on the given $X$-sample, is equivalent to optimizing its retargeted version.

2.2 Efficient Estimation of Retargeted Policy Value

In the above we introduced a new estimand, $R(\pi;w,\rho)$, and its sample version, $R_{n}(\pi;w,\rho)$. For it to serve as an objective for policy learning, we also need to estimate it. To characterize how easily estimable these are, we use consider the corresponding asymptotic efficiency bound, that is, we focus on the asymptotic estimation variance of any efficient regular estimator for $R(\pi;w,\rho)$.

Lemma 2.3**.**

The efficient influence function for $R(\pi;w,\rho)$ is $\psi(x,a,y;\phi,\mu)-R(\pi;w,\rho)$ where

[TABLE]

Hence every efficient regular estimator for $R(\pi;w,\rho)$ must have the form

[TABLE]

Moreover,

[TABLE]

And the sum of the above is the efficiency bound for $R(\pi;w,\rho)$.

This result follows easily following the same proof arguments as in Hahn (1998); Hirano et al. (2003), using the results of Bickel et al. (1993), and finally by the law of total variance.

With fixed known $w,\rho$, the efficient estimation of $R(\pi;w,\rho)$ is very similar to the efficient estimation of an average from missing-at-random data. Therefore, there are also various immediate examples of efficient estimators for $R(\pi;w,\rho)$ arising from existing literature. When $\mu$ is sufficiently smooth, the sample average of $\psi(X_{i},A_{i},Y_{i};\bm{\infty},\hat{\mu})$, that is, the direct estimator, with estimated outcome regression $\hat{\mu}$ is efficient (Hahn, 1998). When $\phi$ is sufficiently smooth, the sample average of $\psi(X_{i},A_{i},Y_{i};\hat{\phi},\bm{0})$, that is, the IPW estimator, with estimated propensities $\hat{\phi}$ is efficient (Hirano et al., 2003). Alternatively, with less stringent smoothness assumptions, the sample average of $\psi(X_{i},A_{i},Y_{i};\hat{\phi}^{(-i)},\hat{\mu}^{(-i)})$ is efficient when $\hat{\phi}^{(-i)}$, $\hat{\mu}^{(-i)}$ are “cross-fold” estimated with convergence rates $o_{p}(n^{-1/4})$, that is the double machine learning (DML) estimator (Chernozhukov et al., 2018).

Since all of the above efficient estimators rely on plugging in nuisance estimates of $\mu,\phi$ that are shared across policies $\pi$, it easily follows that these estimates also obtain uniform efficiency across $\pi\in\Pi_{0}$ (i.e., the $o_{p}(1/\sqrt{n})$ is uniform in $\pi\in\Pi_{0}$ in Eq. 4) as long as $\Pi_{0}$ has bounded capacity (i.e., has a polynomial covering number in inverse radius, which is implied by, e.g., having finite pseudo- or Vapnik-Chervonenkis dimension, Vapnik, 2000). This follows from a straightforward invocation of standard empirical process theory arguments (see, e.g., Pollard, 1990); we omit the details. For example, (without any retargeting) Athey & Wager (2017) use a careful application of this to obtain tighter regret bounds than are possible with optimizing inefficient policy value estimates.

Lemma 2.3 decomposes the efficiency bound into two terms. The first, Eq. 5, reflects only the variance due to $X$-variation in the population and depends only on $\mu$ and not $\phi$. The second, Eq. 6, reflects the variance at any one $X$-level due to residual noise and, importantly, any lack of overlap, averaged over $X$ values, and depends only on $\phi$ and not $\mu$. This latter term can be understood as the “efficiency bound” of estimating $R_{n}(\pi;w,\rho)$ in the sense that if $\hat{R}_{n}$ is an efficient regular estimator then we must have

[TABLE]

3 Optimal Retargeting

In the previous section, we first argued that there may be many objectives that characterize the optimal policy as their maximizer, both in the sample and in the population, and we then presented results on how efficiently estimable these objectives are. It therefore stands to reason we should pick the objective that is most convenient, or easiest to estimate. We next consider precisely how to do this.

First, we discuss precisely what we will be mean by easiest to estimate. Ideally, if we had access to prefect counterfactual predictions for all actions, we would optimize the sample policy value, $V_{n}(\pi)$, and not suffer from lack of overlap. In view of Lemma 2.2, under correct specification, this is equivalent to optimizing ${R_{n}}(\pi;w,\rho)$ for any $w$. Moreover, in view of Lemma 2.2, the choice of $w$ does not affect this latter optimization, so we need not consider how $w$ affects the variance of ${R_{n}}(\pi;w,\rho)$ itself, only how $w$ affects the variance of estimating ${R_{n}}(\pi;w,\rho)$. In other words, the asymptotic variance of any asymptotically unbiased estimator (e.g., an efficient an estimator) can be decomposed into the sum of the variance of estimating ${R_{n}}(\pi;w,\rho)$ and the variance of ${R_{n}}(\pi;w,\rho)$ itself, and since $w$ does not change the optimizer of ${R_{n}}(\pi;w,\rho)$ (only the optimizer of our estimate thereof), we only consider the variance of estimating ${R_{n}}(\pi;w,\rho)$.

The efficiency bounds from the last section exactly characterize how difficult this estimation is. In particular, any efficient regular estimator will have an asymptotic mean-squared error given by Eq. 6 in estimating ${R_{n}}(\pi;w,\rho)$. We therefore seek the objective ${R_{n}}(\pi;w,\rho)$ requiring minimal such asymptotic error. (See also Section 5.1 for further discussion.)

Still, our variance in estimating ${R_{n}}(\pi;w,\rho)$ also depends on $\pi$. When doing policy learning, we need to estimate our objective for various policies and we want the weights $w$ to be independent of the policy $\pi$ to avoid any bias. Therefore, our criterion for a good objective for policy learning will be whether it is easily estimable uniformly for all policies. Toward that end, we define,

[TABLE]

Our aim will be to seek $w,\rho$ with small $\Omega(w,\rho)$, i.e., that lead to uniformly easily estimable objective functions for policy learning. Given such desirable $w,\rho$, our approach will then be to construct an efficient estimate $\hat{R}_{n}(\pi;w,\rho)$ for $R_{n}(\pi;w,\rho)$ and to minimize it over $\pi\in\Pi_{0}$. We recall that $\rho$ actually has no effect on the policy learning, in the sample or in the population. Nonetheless, the best $\rho$ provides the best centered objective and will facilitate choosing the best retargeting weights $w$.

In particular, if $\Pi_{0}$ is of bounded capacity, then $\Omega(w,\rho_{0})$ controls the uniform deviation between our estimated objective and the ideal finite-sample objective, $V_{n}(\pi;w)$. Focusing on deterministic policies ($\Pi_{0}\subseteq\{0,1\}^{\mathcal{A}\times\mathcal{X}}$), we characterize the capacity of $\Pi_{0}$ in terms of its uniform entropy integral $\kappa(\Pi_{0})=\int_{0}^{1}\sup_{m,x_{1},\dots,x_{m}}\sqrt{\log(N(\epsilon^{2},\Pi_{0},\{x_{1},\dots,x_{m}\}))}d\epsilon$, where $N(\epsilon,\Pi_{0},\{x_{1},\dots,x_{m}\})$ is the smallest number $N$ of policies $\pi_{1},\dots,\pi_{N}\in\Pi_{0}$ such that $\sup_{\pi\in\Pi_{0}}\inf_{j=1,\dots,N}\frac{1}{m}\sum_{i=1}^{m}\sum_{a\in\mathcal{A}}\left|\pi(a\mid x_{i})-\pi(a\mid x_{i})\right|\leq\epsilon$.

We first focus on $\tilde{V}_{n}(\pi;w):=\tilde{R}_{n}(\pi;w,\bm{0})$ as in Eq. 4.

Lemma 3.1**.**

Suppose that $\phi(A\mid X)$ is bounded away from zero and $Y$ is bounded. Let $w\in\mathbb{R}^{\mathcal{X}}_{++}$ with $\|w\|_{\infty}<\infty$ be given. Then there exists a universal constant $C$ such that, for any $\delta\in(0,1/2)$ and $\rho\in\mathbb{R}^{\mathcal{A}\times X}$, with probability at least $1-\delta$,

[TABLE]

Recall that an efficient estimator must have $\left|\hat{V}_{n}(\pi;w)-\tilde{V}_{n}(\pi;w)\right|=o_{p}(1/\sqrt{n})$. If $\Pi_{0}$ is of bounded capacity, then we can usually obtain uniformly efficient estimators, $\sup_{\pi\in\Pi_{0}}\left|\hat{V}_{n}(\pi;w)-\tilde{V}_{n}(\pi;w)\right|=o_{p}(1/\sqrt{n})$ (see, e.g., Zhou et al., 2018, Lemma 3), so that Lemma 3.1 similarly applies to $\hat{V}_{n}(\pi;w,\rho)$. (See also the discussion in Sections 5.3 and 5.4.)

Before proceeding to discuss how to optimize our uniform efficiency objective, $\Omega(w,\rho)$, we first reformulate it in more digestible terms.

Lemma 3.2**.**

[TABLE]

3.1 The Binary-Action Case

We first the consider the optimal choice of $w,\rho$ in the special case with just two actions, $\mathcal{A}=\{-,+\}$. We will generalize this to multiple actions in the next section, but we study this simple case first in separation because of the simple, elegant solution that arises and the connections to other work.

When $\mathcal{A}=\{-,+\}$, for $\pi\in\Pi$ we use the notation $\pi(x)=\pi(+\mid x)-\pi(-\mid x)=\break 2\pi(+\mid x)-1$. In particular, deterministic policies $\pi$ have $\pi(x)\in\{-1,+1\}$. We similarly apply this notation to the propensity score $\phi$ (which is also in $\Pi$). In particular, note that ${1\pm\phi(x)}=2\phi(\pm\mid x)$ and that $4\phi(+\mid x)\phi(-\mid x)=1-\phi^{2}(x)$.

Lemma 3.3**.**

Let $\rho_{0}(+\mid x)=\rho_{0}(-\mid x)=\frac{1}{2}$ and $w_{0}(x)\propto\left(\frac{\sigma^{2}(+\mid x)}{1+\phi(x)}+\frac{\sigma^{2}(-\mid x)}{1-\phi(x)}\right)^{-1}$ such that $\mathbb{E}\left[{w_{0}(X)}\right]=1$. Then,

[TABLE]

Lemma 3.3 shows that the best reference policy is complete randomization. Note that we did not constrain $\rho$ to be a policy, that is, to be in $\Pi$; it is such at optimality. Additionally, surprisingly, $\rho_{0}$ does not depend on $\phi$. Surprisingly this also coincides with existing practice: because $V(\pi)-V(\rho_{0})$ is conveniently the $\pi(X)$-weighted average treatment effect, this form of centering is already often used in policy learning with binary actions (Beygelzimer & Langford, 2009; Athey & Wager, 2017). Nonetheless, just centering by a reference policy does not actually change the policy learning optimization problem, whether in the sample or the population. For our purpose, the importance of centering is in obtaining the optimal retargeting weights, which does change the optimization problem (but not the population learning target per Lemma 2.1).

Because the optimal reference policy is complete randomization and because centering by complete randomization means that $R(\pi;w,\rho_{0})$ is in fact a $w(X)\pi(X)$-weighted average treatment effect, the optimal retargeting weights we obtain are very similar to the familiar overlap weights of Crump et al. (2006); Li et al. (2018). However, as long as $\Pi_{0}$ remains well-specified, we are not changing the learning target, which remains the optimal policy. Moreover, when we generalize to multiple actions, the connection to weighted average treatment effects will break, but we will still be able to obtain optimal retargeting weights.

Note that if $\sigma^{2}(\pm\mid x)=\sigma^{2}$ is almost surely constant, then the results of Lemma 3.3 simplify to:

[TABLE]

Compare this to no retargeting: $\Omega(\bm{1},\rho_{0})=\frac{\sigma^{2}}{2}\mathbb{E}\left[{(1-\phi^{2}(X))^{-1}}\right]$. Optimality of course guarantees that $\Omega(w_{0},\rho_{0})\leq\Omega(\bm{1},\rho_{0})$, but we can also see this now directly from Jensen’s inequality, where we would have equality if and only if $\phi(X)$ were almost surely constant.

Finally, note that $w_{0}$ is of course not generally known. In the homoskedastic case, for example, because it depends on the unknown propensity score, the efficiency bound for $R(\pi;w_{0},\rho_{0})$ is different from simply plugging in $w_{0}(x)\propto 1-\phi^{2}(X)$ into the sum of Eqs. 5 and 6. However, the additional inflation in the efficiency bound is solely in the former term, i.e., the $X$-variation term (and also only decreases with less overlap). That is, for any efficient regular estimator $\hat{R}_{n}$, the asymptotic variance of $\sqrt{n}\left(\hat{R}_{n}-R_{n}(\pi;w_{0},\rho_{0})\right)$ remains as in Eq. 6 with this $w_{0}$ plugged in. Since we focus solely on the this latter variance, we again ignore the inflation in the efficiency bound due to unknown $w_{0}$ in our consideration of optimality of retargeting weights.

3.2 The Multiple-Action Case

We next treat the more general case with any number of actions, $\mathcal{A}=\{1,\dots,m\}$.

Lemma 3.4**.**

Let

[TABLE]

Then

[TABLE]

Note that if $m=2$ then $\xi(x)=0$ and we immediately recover the binary-action result of Lemma 3.3, which showed that the optimal retargeting weights for policy learning coincided with the overlap weights of Crump et al. (2006); Li et al. (2018) for average effect estimation. For $m>3$, the optimal reference policy and retargeting weights we obtain above are wholly novel and do not resemble anything previously used for either policy learning or average effect estimation. Indeed, because there is no clear choice for a reference policy with multiple actions, as there is for binary actions, previous studies on multiple-action policy learning use no reference-policy centering (Kallus, 2018, 2017; Zhou et al., 2018). Note, in particular, that unlike the binary-action case, here the optimal reference policy depends on $\phi$. Further note that, again, we did not constrain $\rho$ to actually be a policy, that is, to be in $\Pi$, but it is such at optimality. That is, the $\rho_{0}(a\mid x)$ above is nonnegative and sums to $1$ over $a\in\mathcal{A}$ for each $x$.

Note that if $\sigma^{2}(a\mid x)=\sigma^{2}$ is constant across $a$ and almost every $x$, then the results of Lemma 3.4 simplify to:

[TABLE]

In this case it becomes clear that the weights $w_{0}$ penalize any $x$ at which some propensities are small and therefore focuses on areas with good overlap, generalizing a scalarized notion of overlap to multiple treatments in a rigorous and newly relevant way.

We can also compare the efficiency bound objective to no retargeting. It can be shown that $\Omega(\bm{1},\rho_{0})=\frac{1}{4}\mathbb{E}\left[{\kappa(X)}\right]$. Optimality guarantees $\Omega(w_{0},\rho_{0})\leq\Omega(\bm{1},\rho_{0})$, but Jensen’s inequality reveals that the inequality is strict as long as $\kappa(X)$ is not almost surely constant. In the homoskedastic case, $\kappa(X)=\sigma^{2}\left(\sum_{a\in\mathcal{A}}\phi^{-1}(a\mid X)+(m-2)\right)$, which is almost surely constant if and only if $\sum_{a\in\mathcal{A}}\phi^{-1}(a\mid X)$ is almost surely constant. The more it varies, the greater the benefit from retargeting.

4 Empirical Results

We next study the added value of retargeting empirically. We first consider a simulation study where we can vary various knobs and then a case study based on a multiarm job counseling experiment. Replication code is available at

https://github.com/CausalML/RetargetedPolicyLearning.

4.1 Simulation Study for Binary Actions

For our simulation study we focus on the following data generating process with two actions and two covariates, parameterized by $q_{1},q_{2},\beta\geq 0$:

[TABLE]

In the above, $\Phi$ is used to denote the cumulative distribution function of the standard normal. Only $X$ is returned as the covariate data, not $X^{\prime\prime},X^{\prime}$. The parameter $q_{1}$ therefore changes the distribution of $X$. We use it to control the level of covariate shift between training and test by using different values for the two datasets. The parameter $\beta$ controls the amount of overlap in the data. When $\beta=0$, there is perfect overlap, in the sense that at every $X$, the propensity for each action is equal. As $\beta$ grows overlap becomes minimal, in the sense that at every $X$, the sum of inverse propensities action grows as $\exp(\Theta(\beta^{2}))$, using Landau notation. The parameter $q_{2}$ controls the misspecification of a linear policy. Because $Y(+)-Y(-)$ is linear in $X^{\prime}$, when $q_{2}=1$ a linear policy in $X$ is well-specified. However, when $q_{2}\neq 1$ a linear policy in $X$ is misspecified since $X$ is a nonlinear transform of $X^{\prime}$.

Specifically we will consider scenarios where both training and test data are drawn with $q_{2}=1$ (“well-specified”) and with $q_{2}=1/2$ (“misspecified”). We will always take $q_{1}=1$ for training, but we will consider drawing the test data either with $q_{1}=1$ (“stationary”), $q_{1}=2$ (“shift inward”), or $q_{1}=1/2$ (“shift outward”). Out-of-sample policy values $V(\pi)$ are computed using the known $\mu$ on a test set of 100,000 draws. Regret of a policy $\pi$ is computed as $\frac{1}{n}\sum_{i=1}^{n}\left(\max_{a\in\mathcal{A}}\mu(X_{i}\mid a)-\sum_{a\in\mathcal{A}}\pi(a\mid X_{i})\mu(X_{i}\mid a)\right)$ on the test set.

We will consider a variety of approaches for learning a deterministic linear policy from $n$ observations from the above data generating process. All of them will take the form of first computing some numbers $\Gamma_{1},\dots,\Gamma_{n}$ and then optimizing

[TABLE]

In order to avoid any additional layer of complexity that may obscure the conclusions of the numerical results, we solve the above exactly by using grid search rather than use a convex surrogate classification loss. Specifically, we conduct a grid search on $\theta$, sort the data along $\cos(\theta)X_{i1}+\sin(\theta)X_{i2}$, and consider $b$ being every midpoint along with $-\infty,\infty$.

Given a training dataset, we first estimate $\hat{\phi},\hat{\mu}$ using gradient boosting machines, as implemented in the R package gbm. Given these estimates, we consider the following policy learning methods:

Unweighted: $\Gamma_{i}=A_{i}Y_{i}$. 2. 2.

IPW: $\Gamma_{i}=Y_{i}/(A_{i}+\hat{\phi}(X_{i}))$. 3. 3.

Direct: $\Gamma_{i}=\hat{\mu}(+\mid X_{i})-\hat{\mu}(-\mid X_{i})$. 4. 4.

DR: $\Gamma_{i}=\hat{\mu}(+\mid X_{i})-\hat{\mu}(-\mid X_{i})+(Y_{i}-\hat{\mu}(A_{i}\mid X_{i}))/(A_{i}+\hat{\phi}(X_{i}))$. 5. 5.

Retargeted: for each of the above, we modify $\Gamma_{i}\leftarrow(1-\hat{\phi}^{2}(X_{i}))\Gamma_{i}$.

For numerical stability, we normalize all of these by $\frac{1}{n}\sum_{i=1}^{n}\left|\Gamma_{i}\right|$ before solving the above optimization problem.

Figure 1 shows the average regret as we vary the training dataset size $n$ in the various scenarios (we run 5000 replicates for each of 20 $n$ values evenly spaced on the logarithmic scale of the plot axis and report averages and standard errors within each $n$ value). In the well-specified scenarios we see that the retargeting significantly improves policy learning, achieving an order of magnitude lower regret (note log scale). In the misspecified scenarios, retargeting still significantly improves performance in moderate samples. In very large samples with misspecification, the improvement is diluted due to hitting an approximation error limit. In particular, at the rightmost end of the plot, we see the regret of retargeted policy learning reaching a floor and slightly overtaken by the non-retargeted DR method, which flattens out as well due to the misspecification. The retargeted versions of other methods still uniformly improve on their non-retargeted versions in this regime. Figure 2 shows the regret as we vary the level of overlap from perfect, $\beta=0$, to perfectly bad, $\beta\to\infty$ (we run 5000 replicates for each of 30 $\beta$ values, evenly spaced on the logarithmic scale of the plot axis). We first note that while the unweighted method is indeed consistent when $\beta=0$, it succumbs to confounding bias quickly as $\beta$ grows. The non-retargeted methods account for the confounding but succumb to lack of overlap for slightly larger $\beta$. The retargeted methods account both for confounding and lack of overlap and, while they of course suffer too from the worsening overlap, they consistently perform much better than their non-retargeted versions. (Note that for Fig. 1 we chose a moderate level of overlap, $\beta=3.5$, where the effect is most nuanced. Figure 2 shows that the effect would be uniformly lopsided in favor of retargeting when overlap is particularly bad, e.g., $\beta=10$.)

The conclusion from these experiments is that retargeting offers an easy way to significantly improve policy learning in the absence of good overlap, and that it is generally robust to both correct specification, which underlied Lemma 2.1 and our motivation for retargeting, and to covariate shifts.

4.2 Case Study: Multi-Action Personalized Job Counseling

We next consider an application to a dataset derived from a large scale experiment comparing different programs offered to unemployed individuals in France (Behaghel et al., 2014). The experiment compared three arms: a control arm where individuals receive the standard services provided by the Public Employment Services, a public treatment arm where individuals receive an intensive counseling program run by a public agency, and a private treatment arm with a similar program run by a private agency. The hypothetical application we consider is the learning of a personalized intervention policy to efficiently allocate the resources of the intensive treatments to maximize the number of unemployed individuals who reenter employment within six months, minus costs. (The original study focused instead on the differences of the public and private programs, their cost effectiveness, and the incentive problems in contracting the private provider.)

Some individuals assigned to the two treatment arms refused the additional program and reverted to control. Nonetheless, we restrict our attention to intent-to-treat interventions; that is, we consider the actions of our policy as the offer (or no offer) of an intensive program of either type. This make sense since our personalized interventions may not force individuals into a program, only allocate access to the resource. There are of course important fairness considerations in offering differential access to public resources. Kallus & Zhou (2019) study how to assess disparate impact metrics for personalized interventions despite fundamental unidentifiability issues. They also study this same dataset. Here, for the sake of focus, we only consider the social welfare of policies and ignore fairness or equity considerations. Furthermore, to make our policies focus on heterogeneous effects, we set the costs of each arm to be equal to their within-arm average outcome in the original data. That is, the outcome we consider is equal whether one reentered employment within 6 months, minus the average number of individuals who entered employment within 6 months in that arm. The covariates we consider personalizing on are: statistical risk of long-term unemployment, whether individual is seeking full-time employment, whether individual lives in sensitive suburban area, whether individual has a college education, the number of years of experience in the desired job, and the nature of the desired job (e.g., technician, skilled clerical worker, etc.).

To evaluate the different policy learning methods, we run 1440 replications of the following procedure. Out of the 43977 individuals in the study, we remove a random subset of 8795 (20%) for training; the remaining individuals will be used for evaluation. We then introduce some confounding into the training dataset. We consider the following three binary variables: whether individual has 1–5 years experience in the desired job, whether they seek a skilled blue collar job, and whether their statistical risk of long-term unemployment is medium. After studentizing each variable, we segment the data by the tertiles of their sum. In the first tertile, we drop each individual in the control arm with probability $1/4$ and in the private and public treatment arms with probability $7/8$. In the second and third tertiles, we repeat this, exchanging the probability for the control arm with that for the private and public treatment arms, respectively. This now serves as our training data.

We will consider a variety of approaches for learning a linear policy on this data. All of them will take the form of first computing some numbers $\Gamma\in\mathbb{R}^{n\times\mathcal{A}}$ and then optimizing

[TABLE]

Again, to avoid the additional complexity of relaxations that may obscure the results, we solve the above as a mixed-integer optimization problem using Gurobi (we put a time limit of 5 minutes on each call).

We next estimate $\hat{\phi},\hat{\mu}$ using gradient boosting machines, as implemented in the R package gbm. Given these estimates, we consider the following policy learning methods:

IPW: $\Gamma_{ia}=\mathbb{I}\left[A_{i}=a\right]Y_{i}/\hat{\phi}(a\mid X_{i})$. 2. 2.

Direct: $\Gamma_{ia}=\hat{\mu}(a\mid X_{i})$. 3. 3.

DR: $\Gamma_{ia}=\hat{\mu}(a\mid X_{i})+\mathbb{I}\left[A_{i}=a\right](Y_{i}-\hat{\mu}(a\mid X_{i})/\hat{\phi}(a\mid X_{i})$. 4. 4.

Retargeted: for each of the above, we modify $\Gamma_{ia}\leftarrow(\sum_{a\in\mathcal{A}}\hat{\phi}^{-1}(a\mid X_{i})+\frac{1}{2})^{-1}\Gamma_{ia}$.

Given a learned policy, we evaluate it on the held-out test set using a Horvitz-Thompson estimator (note that the randomized study had varying randomization probabilities; we use these in the evaluation). We report the results in Table 1, showing both average value as well as standard errors over replications. We find that DC and DM perform particular badly suggesting that the outcome model was a bad fit, which may also contribute to DR’s worse performance than IPW. Nonetheless, for both DR and IPW, retargeting offers a clear and substantial improvement in policy value.

5 Discussion

5.1 The Choice of Efficiency Bound to Optimize

In our method, we focused on minimizing the efficient asymptotic error of estimating the retargeted sample regret, ${R_{n}}(\pi;w,\rho)$. Notably, this asymptotic error was only part of the efficient asymptotic error of estimating the retargeted population regret, ${R}(\pi;w,\rho)$, which also included the term in Eq. 5. In this section we offer further discussion of this choice.

In certain senses, ${R_{n}}(\pi;w,\rho)$ is the best finite-sample objective we can hope for. It completely eliminates any issues of residual noise and, more importantly, any issue of lack of overlap in the data. Beyond residual noise and overlap issues, any additional error may arise only from the sampling uncertainty of the finite dataset of $X$ values. Recall that ${R_{n}}(\pi;w,\rho)$ is measurable with respect to just $X_{1:n}$. Indeed, the remaining term in the efficiency bound for ${R}(\pi;w,\rho)$, Eq. 5, is the $X$-variance of $R_{n}(\pi;w,\rho)$. This term of the efficiency bound is driven solely by heterogeneity in $X$.

Our primary justification for omitting this term from consideration is that, in view of Lemma 2.2, under correct specification, the choice of $w$ does not affect the optimizer of $R_{n}(\pi;w,\rho)$ and therefore it does not matter which $w$ makes $R_{n}(\pi;w,\rho)$ best estimate the population quantity $R(\pi;w,\rho)$, the result is always the same. What matters is estimating $R_{n}(\pi;w,\rho)$. We therefore focused on Eq. 6, the efficiency bound for this latter estimation.

Under misspecification this argument may not hold, but this is true to the same extent that retargeting may be inappropriate to begin with. We discuss the issue of misspecification further in Section 5.2 below. Note that since our data arrives from the uniformly-weighted population, using a non-constant $w$ would usually (but not necessarily) inflates Eq. 5 (in particular, the minimum such term is given by the non-constant $w\propto\left(\sum_{a\in\mathcal{A}}(\pi(a\mid X)-\rho(a\mid X))\mu(a\mid X)\right)^{-1}$). At any rate, minimizing the $X$-heterogeneity term would have us focus on populations with very little heterogeneity, while heterogeneity is exactly what we focus on and exploit in policy learning.

Finally, we point out that if we know on which $X$-population we wish to evaluate regret, then the efficiency bound for the population quantity is exactly Eq. 6, which we focused on, without an $X$-heterogeneity term, as the following shows.

Lemma 5.1**.**

Suppose the measure $\nu$ of $X$ is known. Then the efficient influence function for $R(\pi;w,\rho)$ is $\psi^{\prime}(x,a,y;\phi,\mu)-R(\pi;w,\rho)$ where

[TABLE]

And, the efficiency bound for $R(\pi;w,\rho)$ in this case is given by Eq. 6.

If we imagine instantiating the above with $\nu$ being the empirical distribution, then, in a sense, this is another way of stating that the “efficiency bound” for estimating ${R_{n}}(\pi;w,\rho)$ is given by Eq. 6. (Note that if $\nu$ is the empirical distribution then the sample average of $\psi(X,A,Y;\phi_{0},\mu_{0})$ is the same as that of $\psi^{\prime}(X,A,Y;\phi_{0},\mu_{0})$.) More generally, this shows that if we know the population distribution of $X$, then the efficiency bound for estimating the population regret, $R(\pi;w,\rho)$, is simply Eq. 6. (Note that in this case the structure of an efficient estimator changes to depend on $\nu$ as in $\psi^{\prime}(X,A,Y;\phi_{0},\mu_{0})$ above.)

5.2 Retargeting, Misspecification, and Train-Test Covariate Shift

In Section 2.1, we argued that if $\Pi_{0}$ is well-specified then we may arbitrarily change the distribution of covariates, and retargeted policy learning will still give the unconstrained best policy. If $\Pi_{0}$ is misspecified, however, retargeted policy learning will instead give the policy in $\Pi_{0}$ that is best on average for the particular retargeted population, which may be different from the best policy in $\Pi_{0}$ that is best for the original training population. In particular, if we use the optimal retargeting weights then we will get the policy in $\Pi_{0}$ that is best on the regions where no action has a small propensity $\phi(a\mid X)$, what Li et al. (2018) call “marginal” units in the binary case. Under correct specification, this is the same as the optimal policy, while under misspecification, this policy depends on the treatment mechanism.

However, in practice, one can inevitably expect some amount of covariate shift between the training data and the test data where the policy will be deployed. Usually, we have training data from a particular study or from the subset of records with sufficient information, we learn a policy on this data, and we intend to deploy it on a larger scale. And usually the exact test distribution is not known a priori (if it is, we can target it; see, e.g., Zhao et al., 2019 and Lemma 5.1). In light of this, it may seem overly dogmatic to insist on getting the best $\Pi_{0}$ policy for the training distribution of covariates for fear of misspecification, especially when it can have very bad overlap and misspecification potentially mild. Therefore, retargeting may still be preferable even under misspecification. Indeed, we saw in Section 4.1 that even under misspecification, with or without any covariate shift, retargeting offered better performance in all reasonably sized datasets. We can also attenuate misspecification effects by making our policy class more flexible or by regularizing retargeting weights as in Appendix A. In the end one must consider and compare the level of misspecification, the level of inevitable (and unknown) train-test covariate shift, and the level of overlap. Given all these interacting uncertain factors, we think it is both safest and most effective to just focus on learning a policy on a subpopulation with good overlap, where the policy learning procedure is actually reliable and robust. That is, we think retargeting is generally a safe bet and should be applied generally.

5.3 On the Interplay of Retargeting and Regret

Our approach to retargeting was motivated by the observation that there are several characterizations of the optimal policy as the maximizer of different objectives. We therefore sought the most convenient characterization, with the most efficiently estimable objective. It remains an open question, however, how to characterize the efficiency of learning the optimal policy itself as the parameter of interest, for example when assuming a parametrized form of the policy class such as linear.

Following a uniform convergence argument, we could characterize the regret on the same retargeted distribution used for training by retargeted policy learning.

Lemma 5.2**.**

Under the same conditions as Lemma 3.1, there exists a universal constant $C$ such that, for any $\delta\in(0,1/2)$ and $\rho_{0}\in\mathbb{R}^{\mathcal{A}\times X}$, with probability at least $1-\delta$,

[TABLE]

where $M(x)=\max_{a\in\mathcal{A}}\mu(a\mid x)-\min_{a\in\mathcal{A}}\mu(a\mid x)$.

If we have a uniformly efficient estimator, $\sup_{\pi\in\Pi_{0}}\left|\hat{V}_{n}(\pi;w)-\tilde{V}_{n}(\pi;w)\right|=o_{p}(1/\sqrt{n})$, then by combining Lemmas 3.1 and 5.2 and instantiating them at a maximizer of $V(\pi;w)$ and at any $\hat{\pi}^{w}_{n}\in\operatornamewithlimits{argmax}_{\pi\in\Pi_{0}}\hat{V}_{n}(\pi;w)$, we obtain

[TABLE]

This shows that we may obtain better regret on the retargeted population. However, this is limited to that population, as could be expected as the result holds regardless of correct specification. While under correct specification we have $\max_{\pi\in\Pi_{0}}V(\pi;w^{\prime})-V(\pi^{\prime};w^{\prime})\leq(\operatorname*{ess\,sup}w^{\prime}(x)/w(x))(\max_{\pi\in\Pi_{0}}V(\pi;w)-V(\pi^{\prime};w))$ for any $w,w^{\prime}\in\mathbb{R}^{\mathcal{X}}_{+}+$, this is still too loose to characterize the regret on the original population ($w^{\prime}=\bm{1}$) and washes out the regret improvements we actually see in practice due to retargeted policy learning. (While in view of the previous section, if there were a covariate shift, then this might be helpful as the ratio between the test distribution and training need not be smaller than that between the test and the retargeted population.) The bound also depends on the second moment of $V_{n}(\pi;w)$, while the set of maximizers of the latter is the same for all $w$, which was our motivation for ignoring this term in our efficiency bound objective. (Although, if the distribution of $X$ were known, then via Lemma 5.1 we could also obtain Eq. 8 without the second variance term.)

Instead, perhaps a better lens through which to understand the regret improvements due to retargeted policy learning is that we learn the underlying optimal policy parameters faster, and can therefore generalize better to any population, including the original one. As stated above, it remains an open question outside our scope how to characterize this in generality. Below we give a primitive treatment of this for the case of a finite class of policies to illustrate the idea.

Lemma 5.3**.**

Let $p_{\pi}=\mathbb{P}\left(\pi(a\mid X)(\max_{a^{\prime}\in\mathcal{A}}\mu(a^{\prime}\mid X)-\mu(a\mid X))=0\;\forall a\in\mathcal{A}\right)$ and $p=\sup_{\pi\in\Pi_{0}:p_{\pi}\neq 1}p_{\pi}$. Suppose $\Pi_{0}\cap\Pi^{*}\neq\varnothing$, $w\in\mathbb{R}^{\mathcal{X}}_{++}$. Define $\Pi_{0,n}^{*}(w)=\operatornamewithlimits{argmax}_{\pi\in\Pi_{0}}V_{n}(\pi;w)$. Then,

[TABLE]

Now consider any $\hat{\Pi}_{n}$ and notice that

[TABLE]

The latter probability is exponentially decaying in $n$ if $\left|\Pi_{0}\right|<\infty$ and in a manner independent of $w$, per Lemma 5.3. The former probability has to do with how well we are approximating the problem of optimizing $V_{n}(\pi;w)$, as characterized by, e.g., Lemma 3.1. In particular, focusing on $\hat{\Pi}_{n}=\operatornamewithlimits{argmax}_{\pi\in\Pi_{0}}\tilde{V}_{n}(\pi;w)$, we have:

Lemma 5.4**.**

Suppose that $\phi(A\mid X)$ is bounded away from zero and $Y$ is bounded. Then there exist $\alpha>0,\,\beta<1$ such that for any $n\geq 1,w\in\mathbb{R}^{\mathcal{X}}_{++},\rho\in\mathbb{R}^{\mathcal{A}\times\mathcal{X}}$,

[TABLE]

where $\gamma_{n}(w)=\max_{\pi\in\Pi_{0}}V_{n}(\pi;w)-\max_{\pi\in\Pi_{0}\backslash\Pi_{0,n}^{*}(w)}V_{n}(\pi;w)$.

Thus, we see that the probability we successfully pick out an optimal policy in $\Pi^{*}$ improves when we make $\Omega(w,\rho)$ small. However, it also depends on the separation that a particular $w$ induces on the empirical value of different policies, $V_{n}(\pi;w)$. Indeed some $w$’s may make it small and some large. However, a priori there is no reason to believe that this separation is smaller for the original population ($w=\bm{1}$) than for any other retargeted population, such as the optimal one ($w=w_{0}$), and therefore we expect to do better when $\Omega(w,\rho)$ is small. (This is unlike the variance of $V_{n}(\pi;w)$, which we do expect to be inflated; see Section 5.1.) This illustrates the regret benefits of optimizing a retargeted objective with small $\Omega(w,\rho)$, but a general characterization remains open.

Much of the complexity here arises from the non-regularity of policy learning (Laber et al., 2014). The benefit of retargeting is more straightforward for regular estimands. Consider estimating the linear projection of some $Z$ onto $X$, i.e., $\beta^{*}$ solving $\mathbb{E}[(Z-\beta^{T}X)X]=0$. Without additional information, ordinary least squares is the best one can do. Consider, for example, doing instead weighted least squares weighted by $1/\operatorname{Var}(Z\mid X)$. Generally, this identifies a different projection from $\beta^{*}$. If, however, $\mathbb{E}[Z\mid X]$ were linear, then it identifies the same $\beta^{*}$ and it increases estimation efficiency, which in turn also translates to lower squared-error risk due its bowl-shaped loss (Van der Vaart, 1998, Lemma 8.5). This is immediately applicable to personalized causal inference. If $\mathcal{A}=\{-,+\}$ and $Z=\psi(X,A,Y;\phi,\mu)$ ($\psi$ as in Lemma 2.3), then $\mathbb{E}[(Z-\beta^{T}X)X]=0$ is the efficient version of the (unobservable) estimating equation $\mathbb{E}[(Y(+)-Y(-)-\beta^{T}X)X]=0$ (Tsiatis, 2007). If $\mathbb{E}[Y(+)-Y(-)\mid X]$, known as the conditional average treatment effect (CATE), were linear, then we would be better off reweighting using $1/\operatorname{Var}(Z\mid X)\propto\left(\frac{\sigma^{2}(+\mid x)}{1+\phi(x)}+\frac{\sigma^{2}(-\mid x)}{1-\phi(x)}\right)^{-1}$, which matches $w_{0}$ in Lemma 3.3. Recently, Bennett & Kallus (2020) considered efficient estimation of the policy parameters that solve smooth convex relaxations of the policy value, which are regular.

As an aside, note that estimating CATE for the purposes of personalization suffers from worse missepcification concerns: not only does a misspecified CATE not transport to new populations, it does not even transport to the training population as it optimizes the wrong loss. (Nonetheless, if we seek the best unrestricted policy in all $\Pi$, learning a nonparametric CATE may be inevitable; cf. Stoye, 2009)

5.4 Local Uniform Efficiency Objectives

For retargeting, we sought the weights and reference policy that minimized the uniform control of efficiency bounds over all $\pi\in\Pi$, i.e., $\Omega(w,\rho)$. But in the end, we only need to estimate the objective for $\pi\in\Pi_{0}$. Therefore, one possible extension of our approach is to instead minimize $\Omega_{0}(w,\rho)=\sup_{\pi\in\Pi_{0}}\mathbb{E}\left[{w^{2}(X)\sum_{a\in\mathcal{A}}\frac{\sigma^{2}(a\mid X)}{\phi(a\mid X)}\left(\pi(a\mid X)-\rho(a\mid X)\right)^{2}}\right]$. In particular, Lemma 3.1 continues to hold with $\Omega(w,\rho)$ replaced by this $\Omega_{0}(w,\rho)$.

In the case of two actions, ${\mathcal{A}}=\{-,+\}$, one can see that, as long as $\Pi_{0}$ contains at least one deterministic policy, then $\Omega(w,\rho_{0})=\Omega_{0}(w,\rho_{0})$ are the same, where $\rho_{0}(\pm\mid x)=1/2$. So, if one is content to use $\rho_{0}(\pm\mid x)=1/2$ as the reference policy then the optimal retargeting weights remain the same. However, unlike the case of $\Omega(w,\rho)$, the optimal $\rho$ in minimizing $\Omega_{0}(w,\rho)$ may not actually be $\rho_{0}$ and may generally depend on $\Pi_{0}$. Finding the $\Pi_{0}$-specific optimal reference policy becomes a difficult non-convex optimization problem. This extension, therefore, while potentially appealing, may be very difficult to both implement and analyze.

6 Conclusion

We studied one countermeasure to make policy learning more robust to limited overlap by simply changing the population on which we optimize the policy, which we termed retargeting. We developed our approach in terms of choosing from among the many equivalent characterizations of the optimal policy the one that is most convenient from the perspective of asymptotic estimation efficiency. In policy learning, one is simply interested in an effective, well-performing policy; not some unbiased estimate with a clean population interpretation, as in the case of treatment effect estimation. This makes policy learning particularly apt for retargeting. When the policy class is well-specified retargeting incurs no bias, but even if the policy class is misspecified, retargeting can lead to more robust policy learning and, in the end, better performance. We therefore believe that in practice, retargeting should be used in any policy learning procedure applied to observational data, where overlap will inevitably pose a serious issue. Since retargeting is quite simple to apply, this offers an easy fix for any policy learning algorithm.

Acknowledgments

The author thanks the anonymous reviewers and associate editor for the constructive inputs. This material is based upon work supported by the National Science Foundation under Grant No. 1846210.

Appendix A Bias Regularization

In this supplemental section we consider interpolating between optimal retargeting weights and uniform weights to protect against cases of misspecification. In Section 3, motivated by the fact that the asymptotic limit of policy learning is invariant under correct specification, we sought to use retargeting weights that lead to the most easily estimable objective and, in particular, attenuate the issue of overlap. However, when the policy class is misspecified, we may be concerned that we may bias our policy learning too much.

In this section we consider seeking retargeting weights that balance bias and asymptotic efficient variance. Specifically, since only retargeting interacts with misspecification, and centering by a reference policy does not, we focus on the bias in the reweighted policy value, i.e.,

[TABLE]

This bias depends on the policy, $\pi$, and the outcome regression function, $\mu$. As we did with the asymptotic efficient variance objective, we will consider a worse-case version of the bias. Suppose $\mu(a\mid X)\in L_{2}$ is square-integrable for each $a\in\mathcal{A}$. Then we have that $\mu_{\max}(X)\in L_{2}$ is also square-integrable where $\mu_{\max}(X)=\max_{a\in\mathcal{A}}\left|\mu(a\mid X)\right|$. Therefore, we consider

[TABLE]

Note that $B(w;\lambda)=\lambda B(w;1)$ is clearly homogeneous in $\lambda$. So we focus on just $B(w)=B(w;1)$.

Lemma A.1**.**

[TABLE]

Note that the form in Lemma A.1 would still follow if we defined bias in the centered values, $R(\pi;w,\rho)$, instead and changed the definition of $\mu_{\max}$ appropriately. The key insight is that, appropriately defined, the worst-case bias will generally take the form of the $L_{2}$ distance between $w$ and uniform weights.

To balance bias and asymptotic efficient variance, we consider choosing weights to minimize a new combined criterion:

[TABLE]

Lemma A.2**.**

Let $\kappa$ be as in Lemma 3.4. Let

[TABLE]

such that $\mathbb{E}\left[{w_{0}(X)}\right]=1$. Then $w_{0}\in\operatornamewithlimits{argmin}_{w\;:\;\mathbb{E}\left[{w(X)}\right]=1}\mathcal{E}^{2}_{\lambda}(w)$.

In other words, Lemma A.2 shows that the optimal weights that balance bias and asymptotic efficient variance are given by simply padding the weights given the previous sections. Essentially, this shrinks the optimal retargeting weights toward uniform weights.

In the case of two actions, $\mathcal{A}=\{-,+\}$, and homoskedastic noise, $\sigma^{2}(\pm\mid x)=\sigma^{2}$, this can be simplified to

[TABLE]

where $c=\frac{\sigma^{2}}{2\lambda^{2}}$ is a transformed weight padding parameter.

A.1 Empirical Results

To study the impact of bias regularization, we repeat the experiential setup of Section 4.1 and consider policy learning methods given by the following scores for varying $c$,

[TABLE]

Figure 3 shows the results when running 50000 replicates for each of 33 $c$ values, evenly spaced on the logarithmic scale of the plot axis, with $n=2000$ and $\beta=3.5$. In the well-specified scenarios, full retargeting ($c\to\infty$) is best while in the misspecified scenarios, full retargeting is not bad but some regularization may offer minimal improvements.

These minimal improvements are given by the best regularization parameter. In practice, $c$ must be chosen, such as using cross validation (where, e.g., out-of-sample estimates of policy value are estimated using some efficient estimator). Therefore these improvements may be too minimal to consider, given that choosing the regularization parameter $c$ is non-trivial. Indeed, as discussed in Section 5.2, optimal retargeting may be both safest and most effective even under misspecification. Nonetheless, bias regularization may provide another knob to control retargeting that may be useful in some applications.

Appendix B Omitted Proofs

Proof of Lemma 2.1.

Since $\Pi$ allows for unconstrained choice of $\pi(\cdot\mid X)$ across values of $X$, we must have

[TABLE]

Let $\pi\in\Pi^{*}\cap\Pi_{0}$ and $\pi^{\prime}\in\Pi_{0}$. Then, since $\pi\in\Pi^{*}$, by the above we must have $R(\pi;w,\rho)\leq R(\pi^{\prime};w,\rho)$. Since $\pi^{\prime}\in\Pi_{0}$ was arbitrary, we have $\pi\in\Pi_{0}^{*}$.

Let $\pi\in\Pi_{0}^{*}$ and $\pi^{\prime}\in\Pi$. By assumption, there exists $\pi^{\prime\prime}\in\Pi^{*}\cap\Pi_{0}\neq\varnothing$. Because $\pi^{\prime\prime}\in\Pi_{0}$, we have $R(\pi;w,\rho)\leq R(\pi^{\prime\prime};w,\rho)$. Because $\pi^{\prime\prime}\in\Pi^{*}$ and by the above, we have $R(\pi^{\prime\prime};w,\rho)\leq R(\pi^{\prime};w,\rho)$. Hence, $R(\pi;w,\rho)\leq R(\pi^{\prime};w,\rho)$. Since $\pi^{\prime}\in\Pi$ was arbitrary, we have $\pi\in\Pi^{*}$. Moreover, by definition, $\pi\in\Pi_{0}$. Hence, $\pi\in\Pi^{*}\cap\Pi_{0}$. ∎

Proof of Lemma 2.2.

We first note that $\Pi^{*}\subseteq\Pi^{*}_{n}$ almost surely because any $\pi^{*}\in\Pi^{*}$ must be point-wise optimal at almost every $x$. Therefore, $\Pi^{*}\cap\Pi_{0}\neq\varnothing$ implies that $\Pi_{n}^{*}\cap\Pi_{0}\neq\varnothing$. The proof then proceeds exactly as that of Lemma 2.1. ∎

Proof of Lemma 3.1.

The quantity to be bounded is:

[TABLE]

Define $\pi(X_{i})=(\pi(a\mid X_{i}))_{a\in\mathcal{A}}\in\mathbb{R}^{\mathcal{A}}$ and $\Gamma_{i}\in\mathbb{R}^{\mathcal{A}}$ such that

[TABLE]

With this notation, notice that we have $\tilde{V}_{n}(\pi;w)-V_{n}(\pi;w)=\frac{1}{n}\sum_{i=1}^{n}\left<\Gamma_{i},\pi(X_{i})\right>$. Moreover, notice that $\mathbb{E}\left[{\Gamma_{i,a}\mid X_{i}}\right]=0$ whence $\mathbb{E}\left<\Gamma_{i},\pi(X_{i})\right>=0$.

By boundedness assumptions, there is some $B>0$ such that $\left|\left<\Gamma_{i},\pi(X_{i})-\pi^{\prime}(X_{i})\right>\right|\leq\left|{\phi^{-1}(A_{i}\mid X_{i})}(Y_{i}-\mu(A_{i}\mid X_{i}))\right|<B$. Then, the final statement of Theorem 2.3 of Bousquet (2002) (simplified using $\sqrt{x+y}\leq\sqrt{x}+\sqrt{y}$ for $x,y\geq 0$) yields that with probability at least $1-\delta$,

[TABLE]

Note that

[TABLE]

We next proceed to bound $\mathbb{E}\mathfrak{D}_{n}$.

Define the Rademacher complexity

[TABLE]

Letting $(\Gamma^{\prime}_{i},X^{\prime}_{i})\sim(\Gamma_{i},X_{i})$ be an iid copy (ghost sample), a standard symmetrization argument yields

[TABLE]

Finally, Theorem 1 of Zhou et al. (2018) yields

[TABLE]

Combining with Eqs. 9 and 10 and $\mathbb{E}\mathfrak{D}_{n}\leq 2\mathfrak{R}_{n}$ completes the proof. ∎

Proof of Lemma 3.2.

Since $\pi\in\Pi$ is unconstrained we have

[TABLE]

Consider the inner optimization for every $X$. Note that it is the sup of a convex function in $p$ over a convex region. Therefore, an optimal solution occurs at an extreme point. Hence, using $1^{2}=1$,

[TABLE]

Taking the maximum over $p$ picks out the largest coefficient, giving the result. ∎

Proof of Lemma 3.3.

From Lemma 3.2, since $\phi(\pm\mid x)=(1\pm\phi(x))/2$, we have

[TABLE]

Fix $w$ and consider $\min_{\rho}\Omega(w,\rho)$. Since $\rho$ is unconstrained, $\rho(\cdot\mid X)$ minimizes the inside of the above expectation for each $X$. Fix $x$. Let $\beta_{\pm}=\frac{\sigma^{2}(\pm\mid x)}{1\pm\phi(x)}\geq 0$. Then $\rho(\cdot\mid x)$ solves $\min_{\rho(\pm\mid x)}f(\rho(+\mid x),\rho(-\mid x))$, where

[TABLE]

We will show that $z_{\pm}=\frac{1}{2}$ is an optimal solution. If either $\beta_{+}=0$ or $\beta_{-}=0$ then $z_{-}=z_{+}=\frac{1}{2}$ is clearly an optimizer. Suppose $\beta_{\pm}>0$. $f(z_{+},z_{-})$ is convex and is therefore optimized at critical points. To enumerate the critical points, first consider the case where $\beta_{+}(1-2z_{+})>\beta_{-}(1-2z_{-})$. Then $f(z_{+},z_{-})$ is differentiable and a critical point must satisfy $\partial_{z_{+}}f(z_{+},z_{-})=2(\beta_{+}z_{+}-\beta_{+})=0\implies z_{+}=1$, $\partial_{z_{-}}f(z_{+},z_{-})=2\beta_{-}z_{-}=0\implies z_{-}=0$, which contradicts $\beta_{+}(1-2z_{+})>\beta_{-}(1-2z_{-})$, so there are no such critical points. The same holds for the case $\beta_{+}(1-2z_{+})<\beta_{-}(1-2z_{-})$. Consider the case $\beta_{+}(1-2z_{+})=\beta_{-}(1-2z_{-})$. Then $z_{+}=\frac{1}{2}+\frac{\beta_{-}}{\beta_{+}}\left(z_{-}-\frac{1}{2}\right)$ and $f(z_{+},z_{-})=\frac{\beta_{-}+\beta_{+}}{4\beta_{+}}(\beta_{+}+\beta_{-}(2z_{-}-1)^{2})$, which is optimized at $z_{-}=\frac{1}{2}$, leading to $z_{+}=\frac{1}{2}$. So $\rho_{0}(\cdot\mid x)=\frac{1}{2}$ is optimal.

For $\rho=\rho_{0}$, we have

[TABLE]

We next optimize this over $w$ subject to $\mathbb{E}w(X)=1$. Let $\kappa(x)=\left(\frac{\sigma^{2}(+\mid X)}{1+\phi(X)}+\frac{\sigma^{2}(-\mid X)}{1-\phi(X)}\right)$. Since the above is convex in $w$, we must have that the optimal $w$ satisfies

[TABLE]

Therefore the optimal $w$ satisfies $w(x)\propto\kappa^{-1}(x)$. ∎

Proof of Lemma 3.4.

From Lemma 3.2,

[TABLE]

Therefore, since $\rho$ is unconstrained $\rho(\cdot\mid x)$ solves $\min_{\rho(\cdot\mid x)}f((\rho(a\mid x))_{a\in\mathcal{A}})$ where

[TABLE]

Fix $x$. Then we can write this convex optimization problem as:

[TABLE]

The Lagrangian of this problem, for a dual variable $\lambda\in\mathbb{R}^{\mathcal{A}}_{+}$, is

[TABLE]

The KKT conditions for $(z,y,\lambda)$ to be a primal-dual optimal are that

[TABLE]

Now consider the solution given by

[TABLE]

It can be directly validated that this solution satisfies all of the above KKT conditions via plugging in the solution and some algebra. Since the problem is convex we conclude this is an optimal solution.

Next, note that for the above solution we have that the objective, which is optimal, is

[TABLE]

and therefore

[TABLE]

Repeating the argument from the proof of Lemma 3.3, we have that $w_{0}\propto\kappa^{-1}(x)$ is optimal. ∎

Proof of Lemma A.1.

Since the sign of $\mu(a\mid x)$ is completely unconstrained across $a,x$, we have

[TABLE]

which yields the result by the self-duality of $L_{2}$. ∎

Proof of Lemma A.2.

From Lemma 3.4, we have that $\inf_{\rho}\Omega(w,\rho)=\frac{1}{4}\mathbb{E}\left[{w^{2}(X)\kappa(X)}\right]$. Therefore,

[TABLE]

Since this is convex in $w$, we must have that the optimal $w$ satisfies

[TABLE]

Therefore, the optimal $w$ satisfies $w(x)\propto(\kappa(x)+4\lambda^{2})^{-1}$. ∎

Proof of Lemma 5.2.

The proof is the same as that of Lemma 3.1 with $\Gamma_{i,a}=\mu(a\mid X_{i})$. ∎

Proof of Lemma 5.3.

Since $\Pi_{0}\cap\Pi^{*}\neq\varnothing$, we have $\max_{\pi\in\Pi_{0}}V_{n}(\pi;w)=\frac{1}{n}\sum_{i=1}^{n}w(X_{i})\max_{a\in\mathcal{A}}\mu(a\mid X_{i})$. Therefore, given $w>0$,

[TABLE]

the probability of which is precisely $p_{\pi}^{n}$. Recall that $p_{\pi}=1$ if and only if $\pi\in\Pi^{*}$, whence by union bound over $\Pi^{*}\cap\Pi_{0}=\{\pi\in\Pi_{0}:p_{\pi}\neq 1\}$ we have the result. ∎

Proof of Lemma 5.4.

First note that if $\sigma^{2}(A\mid X)=0$ almost surely the result is trivial since then $\tilde{R}_{n}(\pi;w)=R_{n}(\pi;w)$ for all $w$. So suppose $\sigma^{2}(A\mid X)>0$ with positive probability. Then $\Omega_{\min}=\inf_{w}\Omega(w,\rho)>0$.

Recall that $\tilde{R}_{n}(\pi;w,\rho)$ has the exact same minimizers as $\tilde{V}_{n}(\pi;w)$ and similarly that $R_{n}(\pi;w,\rho)$ has the exact same minimizers as $V_{n}(\pi;w)$. We therefore proceed to consider the statement about the set of maximizers of the former of each.

Recall that $\tilde{R}_{n}(\pi;w,\rho)-R_{n}(\pi;w,\rho)=\frac{1}{n}\sum_{i=1}^{n}w(X_{i})\frac{\pi(A_{i}\mid X_{i})-\rho(A_{i}\mid X_{i})}{\phi(A_{i}\mid X_{i})}\left(Y_{i}-\mu(A_{i}\mid X_{i})\right)$. Let $s^{2}_{i}(\pi)=w^{2}(X_{i})\sum_{a\in\mathcal{A}}\frac{\sigma^{2}(a\mid X_{i})}{\phi(a\mid X_{i})}\left(\pi(a\mid X_{i})-\rho(a\mid X_{i})\right)^{2}$. Then by boundedness, for some $B>0$, Bernstein’s inequality yields

[TABLE]

Note that $\gamma_{n}(w)=\min_{\pi\in\Pi_{0}}R_{n}(\pi;w,\rho)-\min_{\pi\in\Pi_{0}\backslash\Pi_{0,n}^{*}(w)}R_{n}(\pi;w,\rho)$. Hence, by union bound,

[TABLE]

Let $s^{2}(\pi)=\mathbb{E}s^{2}_{i}(\pi)$. Boundedness and Hoeffding’s inequality gives that$\mathbb{P}\left(\frac{1}{n}\sum_{i=1}^{n}s^{2}_{i}(\pi)>s^{2}(\pi)+\Omega_{\min}\right)\leq\exp(-\eta n\Omega_{\min}^{2})$ for some $\eta$. By union bound, with probability at least $1-\left|\Pi_{0}\right|\exp(-n\eta\Omega_{\min}^{2})$, we have $\sup_{\pi\in\Pi_{0}}\frac{1}{n}\sum_{i=1}^{n}s^{2}_{i}(\pi)\leq\sup_{\pi\in\Pi_{0}}s^{2}(\pi)+\Omega_{\min}\leq 2\Omega(w,\rho)$. Choosing $\alpha=(8\Omega_{\min}/B)^{2}$ yields the result. ∎