Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
22,598 characters · 11 sections · 11 citation commands
Bridging the Gap between Empirical Welfare Maximization and Conditional Average Treatment Effect Estimation in Policy Learning
Decision making about treatment choice is a central objective in causal inference Manski2002treatmentchoice. In this study, to recommend a treatment for an individual with covariates $X$, we aim to train a policy function that maps covariates to treatment recommendations using observational data. This policy learning problem has been widely studied in economics, statistics, and machine learning Swaminathan2015counterfactualrisk,Swaminathan2015batchlearning,Kitagawa2018whoshould,Athey2021policylearning,Zhou2023offlinemultiaction.
There are two main approaches in policy learning. The first is empirical welfare maximization (EWM) or equivalently, counterfactual risk minimization, where we estimate the welfare for each candidate policy and select the one that maximizes the estimated welfare Swaminathan2015counterfactualrisk,Swaminathan2015batchlearning,Kitagawa2018whoshould. The second is the plug-in approach, where we estimate the conditional average treatment effect (CATE) and treat whenever the estimated effect is nonnegative. Existing studies have argued that the EWM approach is more preferable, since it directly targets the policy objective rather than relying on an intermediate regression step. For instance, from a theoretical viewpoint, Kitagawa2018whoshould derive regret bounds for EWM with VC-type policy classes, yielding $1/\sqrt{n}$-type rates for the minimax regret. Note that such (minimax) optimal rates depend on the type of worst-case scenarios considered, as discussed in Audibert2007fastlearning. When the underlying expected outcomes belong to a H\"{o}lder class, the minimax optimal rate is typically characterized by the nonparametric rate $n^{\beta/(2\beta + d)}$. However, if we apply VC-type bounds, we cannot achieve this rate, since the VC dimension of the H\"{o}lder class is infinite.
We reconsider the distinction between the EWM and plug-in approaches and aim to bridge the gap between them. Our key message is that the two are two faces of the same underlying optimization: EWM can be reformulated as least squares on the CATE with certain restrictions on the regression models. This observation has two important practical implications. First, it conceptually unifies the literature by showing that EWM and plug-in are equivalent under restricted regression models. Second, it allows us to avoid the NP-hard combinatorial optimization that arises in the EWM approach.
\paragraph{Contributions.} Here, we summarize our contributions:
We mostly follow Kitagawa2018whoshould while allowing for both deterministic ($0$–$1$) and randomized ($[0,1]$) policies.
\paragraph{Observations.} Let the sample size be $n$. We observe i.i.d.\ draws $Z_i=(Y_i,D_i,X_i)$, $i=1,\dots,n$, where $X_i\in{\mathcal{X}}\subset{\mathbb{R}}^{d_{\mathcal{X}}}$ are pre-treatment covariates, $D_i\in\{0,1\}$ is a binary treatment indicator, and $Y_i\in{\mathbb{R}}$ is the observed outcome.
\paragraph{Potential outcomes.} Under the Neyman–Rubin model, each unit has potential outcomes $(Y_{0,i},Y_{1,i})$ and $Y_i=D_i Y_{1,i}+(1-D_i)Y_{0,i}$. We assume unconfoundedness, $(Y_{0,i},Y_{1,i})\perp\perp D_i \mid X_i$, and define $P$ as the population distribution of $(Y_{0},Y_{1},D,X)$.
\paragraph{Policy.} Let $\Pi$ be a class of policies $\pi:{\mathcal{X}}\to[0,1]$. In Kitagawa2018whoshould, a common deterministic class is $\pi(X)=\mathbbm{1}\{X\in G\}$ for $G\subset{\mathcal{X}}$. We analyze both $0$–$1$ and $[0,1]$ policies.
\paragraph{Welfare.} The (utilitarian) social welfare of $\pi$ is \[ W(\pi)\coloneqq E_P\bigl[Y_1\pi(X)+Y_0(1-\pi(X))\bigr]. \] This reduces to $W(G)=E_P[Y_1\mathbbm{1}\{X\in G\}+Y_0\mathbbm{1}\{X\notin G\}]$ when $\pi=\mathbbm{1}\{X\in G\}$.
\paragraph{Goal and regret.} Given data, we construct $\widehat\pi$ and assess it via regret \[ \mathrm{Regret}\bigl(\pi^*_\Pi,\widehat\pi\bigr)\coloneqq W(\pi^*_\Pi)-W(\widehat\pi), \qquad \pi^*_\Pi \in \operatorname*{arg\,max}_{\pi\in\Pi} W(\pi). \] Let $\tau(x)\coloneqq E[Y_1-Y_0\mid X=x]$ denote the CATE. The first-best policy over all measurable policies is \[ \pi^*_{\mathrm{FB}}(x)=\mathbbm{1}\{\tau(x)\ge 0\}. \] If $\Pi$ is unrestricted, then $\pi^*_\Pi=\pi^*_{\mathrm{FB}}$; otherwise, $\pi^*_\Pi$ is the second-best policy within $\Pi$.
Let $e(x)=P(D=1\mid X=x)$ be the propensity score and $m_d(x)=E[Y_d\mid X=x]$ the conditional mean outcomes. We assume that there exists a constant $0< \epsilon< 1/2$ independent of $n$ such that $\epsilon < e(X) < 1 - \epsilon$ almost surely and that the variables $X$ and $Y$ are bounded.
Given an estimator $\widehat W(\pi)$ of $W(\pi)$, we train a policy as \[ \widehat\pi_{\mathrm{EWM}} \in \operatorname*{arg\,max}_{\pi\in\Pi} \widehat W(\pi). \] A basic estimator of the welfare $W(\pi)$ is the IPW estimator, defined as \[ \widehat W_n^{\mathrm{IPW}}(\pi) \coloneqq \frac{1}{n}\sum_{i=1}^n\left(\frac{Y_i D_i}{e(X_i)}\pi(X_i)+\frac{Y_i(1-D_i)}{1-e(X_i)}\big(1-\pi(X_i)\big)\right). \]
When $e(\cdot)$ is unknown, we use an estimate of $\widehat{e}(\cdot)$ instead of its true value. To remove estimation bias at a fast rate, we can employ the augmented IPW (AIPW), also called the doubly robust (DR), welfare estimator, defined as
where $\widehat{e}(x)$ is an estimator of the propensity score $e(x)$, and $\widehat{m}_d(x)$ ($d\in \{1, 0\}$) is an estimator of the conditional expected outcome $m_d(x)$ Athey2021policylearning. If the estimators of the nuisance parameters satisfy the Donsker condition or are constructed via sample splitting Klaassen1987consistentestimation, also called cross-fitting Chernozhukov2018doubledebiased, and meet mild convergence rate conditions, then the bias from estimation error vanishes at a faster rate than $1/\sqrt{n}$. We can estimate the bias-correction weight by using Riesz regression or density-ratio estimation, as proposed in Chernozhukov2024automaticdebiased,Chernozhukov2022debiasedmachine,Kato2025directbias,Kato2025directdebiased.
The plug-in approach first estimates the CATE $\tau(X)$. We can estimate $\tau(X)$ by estimating the conditional expected outcomes $m_d(X) = E_P[Y_d\mid X]$. Let $\widehat{\tau}(X)$ and $\widehat{m}_d(X)$ be estimators of $\tau(X)$ and $m_d(X)$, respectively. For example, we can estimate $m_d(X)$ by regressing $Y_i$ on $X_i$ using only data with $D_i = d$; that is, we estimate $m_d$ as \[\widehat{m}_d \coloneqq \operatorname*{arg\,min}_{f\in{\mathcal{F}}} \sum^n_{i=1}\mathbbm{1}[D_i = d]\big(Y_i - f(X_i)\big)^2,\] where ${\mathcal{F}}$ is a class of regression functions $f\colon {\mathcal{X}} \to {\mathbb{R}}$, and construct $\widehat{\tau}$ as $\widehat{\tau} = \widehat{m}_1 - \widehat{m}_0$. As another example, we can estimate $\tau(X)$ directly by regressing a variable that is (asymptotically) unbiased for $Y_{1, i} - Y_{0, i}$ on covariates $X_i$. For instance, using an IPW estimator for the target variable, we estimate $\tau$ as \[\widehat{\tau} \coloneqq \operatorname*{arg\,min}_{f\in{\mathcal{F}}} \sum^n_{i=1}\left(\frac{Y_i D_i}{e(X_i)} - \frac{Y_i (1 - D_i)}{1 - e(X_i)} - f(X_i)\right)^2,\]
Using an estimator $\widehat{m}_d(X)$, we construct a policy as
We now formalize the core equivalence. Throughout this section, let ${\mathcal{G}}_\Pi\coloneqq\{g=2\pi-1:\pi\in\Pi\}$.
We first consider the case where $\pi$ is a $0$-$1$-valued policy; that is, $\pi \colon {\mathcal{X}} \to \{0, 1\}$, where $\pi(X) = 1$ (resp.\ $\pi(X) = 0)$ means that the policy recommends treatment $1$ (resp.\ treatment $0$) for given covariates $X$.
\paragraph{Oracle problem.} We show the equivalence between the EWM and plug-in approaches when we can observe $Y_1$ and $Y_0$ directly, ignoring the counterfactual nature. Recall that the welfare maximization problem can be written as
In this case, we consider the following regression problem:
where ${\mathcal{G}}_\Pi \coloneqq \{g = (2\pi - 1) \colon \pi \in \Pi\}$. By definition, $g$ is a function such that $g\colon {\mathcal{X}} \to \{-1, 1\}$.
We now show the equivalence between EWM and least squares. Recall that we defined the optimal policy $\pi^*$ as
Let $g^*$ be the optimal predictor defined as
Then, the following theorem holds.
The proof is straightforward and shown below.
This theorem implies that the EWM approach is equivalent to least squares, where we regress $Y_1 - Y_0$ using a function $g \colon {\mathcal{X}} \to \{-1, 1\}$.
\paragraph{Empirical version.} We can similarly show the empirical version of this equivalence. Recall that given true $e(x)$, the EWM approach trains a policy $\widehat{\pi}_n$ as
Let $\widehat{g}_n$ be the predictor for the CATE defined as
Then, the following holds:
The proof follows the same logic as Theorem (ref) and is omitted for brevity.
We next consider the case where $\pi$ is a $[0, 1]$-valued function; that is, $\pi\colon {\mathcal{X}} \to [0, 1]$. Let $\Pi$ be a policy class that contains such policies, $\pi\colon {\mathcal{X}} \to [0, 1]$.
\paragraph{Oracle problem} We consider the equivalence between the following two problems: regularized EWM and modified least squares. We define the regularized EWM as follows:
where $\lambda > 0$ is a regularization coefficient. Note that if $\pi(X) = 1$, then $2\pi(X) - 1 = 1$; if $\pi(X) = 0$, then $2\pi(X) - 1 = -1$. This regularization term helps prevent the policy from overfitting to the observations by discouraging extreme values of $1$ or $0$.
We also define a corresponding least squares problem. Given a regularization coefficient $\lambda > 0$, we define the mean squared error between $\frac{1}{\sqrt{\lambda}} \big(Y_1 - Y_0\big)$ and $\sqrt{\lambda} g(X)$ as
Let $\widetilde{\pi}(\lambda)$ be an optimal policy defined as
Let $\widetilde{g}(\lambda)$ be an optimal predictor defined as
We now show the equivalence between these EWM and least squares formulations.
The proof is straightforward and shown below.
\paragraph{Empirical version} Similarly, we can show the empirical version of the equivalence result. In empirical analysis, we can only observe either $Y_1$ or $Y_0$ based on the treatment $D$. Therefore, we need to estimate $(Y_1-Y_0)$ as a pseudo-outcome.
Given known $e(\cdot)$, we define a policy empirically trained with regularized EWM as
We also define a trained predictor $\widehat{g}_n$ for the CATE as
Then, the following theorem holds. Since the proof is almost the same as that of the above theorem, we omit it.
Maximizing $\widehat{W}(\pi)$ over $\pi\in\Pi$ is, in general, a combinatorial problem (e.g., empirical set selection) and is NP–hard for many natural classes $\Pi$. The equivalence established above replaces this discrete search with empirical risk minimization under a squared loss on a pseudo–outcome, followed by the simple back–mapping $g\mapsto \pi=(g+1)/2$. This yields a training pipeline that is substantially easier to optimize than direct $0$–$1$ search.
The learning problem reduces to least squares regression of a pseudo–outcome, such as an IPW estimator, on a predictor $g\in\mathcal G_\Pi=\{2\pi-1:\pi\in\Pi\}$. This reformulation enables the use of standard continuous optimization tools (closed–form solvers, SGD, or second–order methods).
If the model class for $g$ is convex, then the least–squares is a convex optimization problem. In these cases, global optima are attained by off–the–shelf solvers. When $\Pi$ is discrete/nonconvex (e.g., trees, rule lists, or set indicators), the induced class $\mathcal G_\Pi$ is also nonconvex, so the objective need not be convex. Even then, the squared–loss formulation remains markedly easier to handle than maximizing a combinatorial welfare objective, because it replaces $0$–$1$ structure with a smooth surrogate amenable to gradient–based methods.
Regardless of the specific function class, the least–squares view avoids the NP–hard policy search that arises in direct EWM and provides a stable, scalable ERM workflow with standard regularization and tuning.
From the decision-making, EWM and plug-in policies are often presented as distinct. From a decision-making perspective, the former optimizes welfare directly, whereas the latter bypasses the estimation of treatment effects as an intermediate step. We showed that this distinction collapses under a simple reparameterization: maximizing estimated welfare over a policy class is equivalent to least squares on a CATE target with predictors constrained to $g=2\pi-1$. This equivalence clarifies the relationship between the two traditions, allows regression guarantees to inform policy learning, and replaces combinatorial search with an empirical risk minimization routine that is typically easier to optimize in practice, while the underlying combinatorial hardness of exact policy search generally remains.
The same viewpoint extends beyond binary treatments. For multiple treatments, one can work with one-vs-all pseudo-outcomes or a simplex-valued predictor and a squared-loss on a vector CATE. Budget or fairness constraints can be incorporated through penalties in the least-squares objective. We hope this unification simplifies both analysis and implementation of policy learning in applied work, and we anticipate that remaining gaps in theoretical results can be narrowed under common regularity conditions.