# Robust Sample Weighting to Facilitate Individualized Treatment Rule Learning
for a Target Population

Rui Chen
Department of Statistics, University of Wisconsin
Jared D. Huling
Department of Biostatistics, University of Minnesota
Guanhua Chen
Department of Biostatistics and Medical Informatics, University of Wisconsin
Menggang Yu
Department of Biostatistics and Medical Informatics, University of Wisconsin
Corresponding authors<EMAIL_ADDRESS>authors:
<EMAIL_ADDRESS>
###### Abstract

Learning individualized treatment rules (ITRs) is an important topic in
precision medicine. Current literature mainly focuses on deriving ITRs from a
single source population. We consider the observational data setting when the
source population differs from a target population of interest. We assume
subject covariates are available from both populations, but treatment and
outcome data are only available from the source population. Although adjusting
for differences between source and target populations can potentially lead to
an improved ITR for the target population, it can substantially increase the
variability in ITR estimation. To address this dilemma, we develop a weighting
framework that aims to tailor an ITR for a given target population and protect
against high variability due to superfluous covariate shift adjustments. Our
method seeks covariate balance over a nonparametric function class
characterized by a reproducing kernel Hilbert space and can improve many ITR
learning methods that rely on weights. We show that the proposed method
encompasses importance weights and the so-called overlap weights as two
extreme cases, allowing for a better bias-variance trade-off in between.
Numerical examples demonstrate that the use of our weighting method can
greatly improve ITR estimation for the target population compared with other
weighting methods.

Keywords: covariate balance, covariate shift, precision medicine, weighting
adjustments

##  1 Introduction

Individualized decision making based on a subject’s characteristics has
received increasing attention in the past decades. It has broad application in
many fields, such as establishing personalized treatment regimes based on
electronic health record data (Wang et al.,, 2016), personalizing ad placement
based on user profiles (Bottou et al.,, 2013), and optimally allocating public
resources based on individual characteristics (Kube et al.,, 2019). The common
goal of these problems is to find an individualized treatment rule (ITR), a
mapping from individual characteristics to a treatment assignment such that
the expected outcome is optimized under the rule. In order for the estimated
ITR to be practically useful, it is usually desirable to limit the search of
the optimal ITR within a prespecified class of restricted rules. Such
restriction may arise from logistic, legal, ethical or political
considerations (Kitagawa and Tetenov,, 2018; Athey and Wager,, 2021), or from
the need for an interpretable ITR to inform better clinical practice; see, for
example, Zhang et al., (2018); Murdoch et al., (2019).

A wide array of statistical methods have been developed to estimate the
optimal ITR using experimental or observational data. One type of these
approaches relies on modeling the outcomes as functions of the covariates and
then obtains the optimal ITR by inverting the regression estimates (Qian and
Murphy,, 2011; Foster et al.,, 2011). However, the success of these methods
highly depends on accuracy of the outcome models. In contrast, weighting-based
ITR learning approaches obviate or reduce the need for fitting outcome models
by directly optimizing the expected ITR outcome or value function (Zhao et
al.,, 2012; Zhou et al.,, 2017; Chen et al.,, 2017), thus making estimated
ITRs more robust to modeling choices. A critical component of this class of
methods is a weighting function that reweights the source data based on the
covariates and treatment information such that the treated and control groups
are comparable. Nearly all such methods use inverse propensity score weights
(IPW) for this task. Reweighting by the IPW makes the treated and control
groups look as though they followed the same covariate distribution as the
whole source sample. As a result, the estimated ITR is tailored for the source
population, which might be suboptimal for a target population that has a
different covariate distribution (Zhao et al.,, 2019; Mo et al.,, 2020).

The distributional difference in covariates between the two populations is
termed as covariate shift in the literature (Sugiyama et al.,, 2007). Such a
scenario is not uncommon, for instance, when we try to generalize the ITR
learned from one health care system to another (Institute of Medicine,, 2013).
Despite the vast literature in estimating the optimal ITR, much less attention
has been paid to its generalizability to a target population. The aim of this
paper is to develop a weighting framework to facilitate generalizability of
optimal ITR from a source population to a target population. We assume subject
covariates are available from both populations but treatment and outcome data
are only available from the source population, where the treatment assignments
are not randomized and thus confounding is also a key concern. In this work,
we address issues related to adjusting for both confounding and covariate
shift and propose a flexible design-based weighting strategy that balances on
a sufficiently large class of covariate functions. Compared with causal
generalization for the average treatment effect (ATE), ITR generalization
poses new challenges due to the need to model and generalize the rules based
on a prespecified class of functions. This puts more stringent requirements on
covariate balance. In addition, the prespecified class of functions may not
contain the unrestricted true optimal ITR and, as we explain below, our
proposed weighting method is designed to mitigate the impact of such
misspecification.

Weighting adjustment by the importance weights has been a commonly used
approach for dealing with covariate shift (Sugiyama et al.,, 2007). In the
context of ITR learning, the source data are reweighted by the density ratio
between the target and source populations on top of the IPW, as suggested by
Kitagawa and Tetenov, (2018, Remark 2.2). Although the importance weights can
ensure unbiased ITR estimation by correcting for covariate shift and
confounding simultaneously, they can contain extreme values especially when
there is insufficient overlap between the treated and control groups or
between the source and target populations, leading to unstable estimation.
Kallus, (2020) proposed to alleviate the instability issue in ITR learning in
this context by retargeting to a population whose density ratio against the
source population is proportional to $\pi(x)(1-\pi(x))$, where $\pi(x)$
denotes the propensity score. This retargeted “hypothetical” population
allocates more emphasis on the covariate region where the treated and control
groups have good overlap. When combined with the IPW, the resulting weights
are essentially the overlap weights of Li et al., (2018). Kallus, (2020)
showed that when the prespecified rule class can well approximate the
unrestricted true optimal ITR, their retargeting approach incurs no bias to
the learning process and can lead to better performing ITRs, and thus can work
well if one expects certain amount of covariate shift in the target
population. However, when the unrestricted true optimal ITR is not contained
in the prespecified rule class, this claim no longer holds because the
retargeted population can be substantially different from the target
population of interest, and thus their restricted optimal ITRs within the
prespecified rule class also differ. Therefore, practitioners interested in
generalizing an ITR to a new population are presented with two extreme
weighting methods: the importance weights that strive for eliminating bias
from confounding and covariate shift but may lead to large variance, and the
overlap weights that emphasize robust and stable estimation but are prone to
bias as they only mitigate confounding. In this paper, we aim to develop a
framework that not only encompasses the importance and the overlap weights,
but also allows one to search a better bias-variance trade-off between them in
a data-driven manner.

Another increasingly recognized issue of these existing weighting methods is
that they are constructed in an indirect way based on the propensity score and
density ratio. In practice these functions need to be estimated and can be
extremely sensitive to model misspecification (Kang et al.,, 2007). Our
proposed method addresses this issue by constructing the weights directly,
bypassing the need for direct modeling of the propensity score or density
ratio. Our approach to doing so is inspired by the recent growing literature
of covariate balancing weights for estimation of the ATE or similar causal
estimands (Hainmueller,, 2012; Imai and Ratkovic,, 2014; Wong and Chan,, 2018;
Hirshberg et al.,, 2019; Huling and Mak,, 2020; Chattopadhyay et al.,, 2020).
Instead of explicitly specifying a functional form for the underlying
propensity score model, covariate balancing approaches directly construct
weights so that they improve the balance between treatment groups over a given
set of covariate functions. The variation of the weights is typically
controlled by minimizing the dispersion of the weights subject to the
specified balance constraints for these covariate functions.

As we will show, proper weights for ITR learning should be able to achieve
balance on a possibly very complex set of covariate functions to cover
candidate assignment rule classes. We start by establishing a covariate
balancing weight strategy that minimizes the imbalance between the treated
group, the control group, and the target population over a nonparametric
function class characterized by a reproducing kernel Hilbert space (RKHS, see
Wahba, (1990)). By carefully inspecting the target population value function,
we observe that different from adjusting for confounding, adjusting for
covariate shift is not essential when the prespecified ITR class contains the
unrestricted true optimal ITR. As we cannot know such containment in practice,
we propose a three-way balancing strategy based on a convex combination of
objectives that seek to induce covariate shift balance and confounding
balance. As our weighting approach emphasizes full distributional balance of
covariates, it is robust to the functional form of confounding and works well
empirically. Despite the fact that our weights are nonparametrically
constructed via balance-seeking objectives, we show that the resulting weights
converge to either the importance weights or the overlap weights under
different hyperparameters when sample size tends to infinity. To guide the
selection of hyperparameters that control the tradeoff in between the
importance weights and overlap weights, we devise a tuning procedure based on
treatment-control balance and predicted outcomes on the target sample and show
the procedure works well in the numerical examples. The proposed weights have
potential for broad application, because they can be incorporated into any
weighting-based ITR learning methods as substitutes for the IPW and facilitate
a better performing ITR for the target population of interest.

##  2 Preliminaries

Suppose we have collected data on $n$ subjects:
$(X_{i},S_{i},S_{i}A_{i},S_{i}Y_{i})$, $i=1,...,n$.
$X_{i}\in\mathcal{X}\subset\mathbb{R}^{p}$ denotes a vector of pre-treatment
covariates which contain confounding factors and treatment effect modifiers.
$S_{i}$ is a population indicator such that $S_{i}=1$ indicates subject $i$ is
from the source population $\mathcal{S}$ and $S_{i}=0$ indicates $i$ is from
the target population $\mathcal{T}$. $A_{i}\in\\{0,1\\}$ is a binary treatment
indicator and $Y_{i}$ is the outcome of interest for which we assume larger
outcomes are preferred without loss of generality. In our setting, $A_{i}$ and
$Y_{i}$ are only observed for observations from the source population
($S_{i}=1$). We will refer to the observations from $\mathcal{S}$ as the
source sample and the rest as the target sample. Let $n_{s}$, $n_{t}$ be the
sample sizes of the source and target samples, respectively, and
$n=n_{s}+n_{t}$ be the total sample size.

For causal interpretation, we use the potential outcome framework (Rosenbaum
and Rubin,, 1983; Imbens and Rubin,, 2015), wherein each subject $i$ is
assumed to have two potential outcomes $Y_{i}(1)$, the outcome if subject $i$
were to receive the treatment, and $Y_{i}(0)$, the outcome if $i$ were not to
receive the treatment. Under the Stable Unit Treatment Value Assumption
(SUTVA), the observed outcome in the source sample can be expressed as
$Y_{i}=Y(A_{i})$. Throughout this paper, we assume SUTVA holds. Under SUTVA we
can associate each observation with a “full” random variate
$(X_{i},S_{i},A_{i},Y_{i}(0),Y_{i}(1))$, which across $i$ are assumed to be
i.i.d. draws from a super-population.

Consider a prespecified ITR class $\mathcal{D}\subset\\{0,1\\}^{\mathcal{X}}$.
For an ITR $d\in\mathcal{D}$, the potential outcome following the treatment
assignment prescribed by $d$ is $Y(d(X))$. The utility of the ITR $d$ can be
measured by the value function, which is defined as the population mean
outcome were all subjects to receive treatment according to $d$. Specifically,
the value function under the source distribution and the target distribution
are

$\mathcal{V}_{\mathcal{S}}(d)=\mathbb{E}[Y(d(X))\mid
S=1]\quad\text{and}\quad\mathcal{V}_{\mathcal{T}}(d)=\mathbb{E}[Y(d(X))\mid
S=0],$

respectively. The objective is then to find the optimal ITR in $\mathcal{D}$
for the target population, i.e. the one that maximizes
$\mathcal{V}_{\mathcal{T}}(d)$ among all rules in the prespecified rule class:

$\operatorname*{arg\,max}_{d\in\mathcal{D}}\;\mathcal{V}_{\mathcal{T}}(d).$
(1)

We assume that the treatment assignment mechanism in the source sample is
determined by a propensity score $\pi(x)=\mathbb{P}(A=1\mid X=x,S=1)$
(Rosenbaum and Rubin,, 1983). If the source sample is from a randomized trial,
then $\pi(x)$ is known. We further denote $\rho(x)=\mathbb{P}(S=1\mid X=x)$
and refer to this as participation probability (Dahabreh et al.,, 2020). To
identify $\mathcal{V}_{\mathcal{T}}(d)$ from the observed data, we impose the
following assumptions throughout the paper.

######  Assumption 1.

(Unconfoundedness of treatment assignment) In the source population,
$(Y(0),Y(1))$ are conditionally independent of $A$ given $X$:
$(Y(0),Y(1))\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}A\mid X,S=1$.

######  Assumption 2.

(Positivity of propensity score) The propensity score of the source population
is bounded away from 0 and 1: for some $c>0$, $c\leq\pi(X)\leq 1-c$ almost
surely.

######  Assumption 3.

(Mean exchangability across populations) The conditional mean of the potential
outcomes given the covariates are equal between the two populations:
$\mathbb{E}[Y(a)\mid X,S=1]=\mathbb{E}[Y(a)\mid X,S=0]$ almost surely for
$a\in\\{0,1\\}$.

######  Assumption 4.

(Positivity of participation probability) The participation probability is
bounded away from 0: $\rho(X)>c$ almost surely for some $c>0$.

The first two are common assumptions in causal inference; and together with
the SUTVA enable identification of causal quantities with respect to the
source population, such as $\mathcal{V}_{\mathcal{S}}(d)$, from the observed
sample. The last two assumptions are adopted from Dahabreh et al., (2020) and
Rudolph and van der Laan, (2017) and allow us to generalize causal estimates
to the target population (Colnet et al.,, 2020). Under these assumptions, the
value function on the target population can be identified with the observed
source sample quantities through

$\displaystyle\mathcal{V}_{\mathcal{T}}(d)$
$\displaystyle=\mathbb{E}\left[w^{*}(A,X)\mathbbm{1}(d(X)=A)Y\mid S=1\right]$
$\displaystyle=\mathbb{E}\left[Aw^{*}(1,X)d(X)Y\mid
S=1\right]+\mathbb{E}\left[(1-A)w^{*}(0,X)(1-d(X))Y\mid S=1\right]$ (2)

where $\mathbbm{1}(\cdot)$ is the indicator function and

$w^{*}(a,x)=\left\\{\frac{a}{\pi(x)}+\frac{1-a}{1-\pi(x)}\right\\}\frac{\mathbb{E}(S)(1-\rho(x))}{(1-\mathbb{E}(S))\rho(x)}.$
(3)

For $a=1$ (or 0), $w^{*}(1,x)$ (or $w^{*}(0,x)$) is proportional to the
density ratio between the target population and the treated group (or control
group). In other words, $w^{*}(a,x)$ acts as the importance weight that
reweights the underlying distribution of the source data to the target
distribution. It can be estimated from data because $\mathbb{E}(S)$ can be
estimated as the empirical average of $S_{i}$ and $\pi(x),\rho(x)$ can be
estimated by probability models for binary outcomes such as logistic
regression. However, since the probability estimates appear in the denominator
of $w^{*}(a,x)$, estimation error and model misspecification of the
probability models can lead to significant error in estimating
$\mathcal{V}_{\mathcal{T}}(d)$ and thus jeopardize the ITR learning process.
In the next section, we propose approaches to directly estimate weights that
are suitable for ITR learning.

##  3 Method

Let $w:\\{0,1\\}\times\mathcal{X}\rightarrow[0,\infty)$ denote a general
weighting function. Let $\mathcal{S}_{a}=\\{i:S_{i}=1,A_{i}=a\\}$ be the index
set of treatment group $a$ among the source sample, $a\in\\{0,1\\}$. By abuse
of notation we let $\mathcal{S}$ also denote $\\{i:S_{i}=1\\}$ and similarly
$\mathcal{T}=\\{i:S_{i}=0\\}$. Following (2) we consider estimators for
$\mathcal{V}_{\mathcal{T}}(d)$ of the form:

$\widehat{\mathcal{V}}_{n}(d;w)=\frac{1}{n_{s}}\sum_{i\in\mathcal{S}_{1}}w_{i}d(X_{i})Y_{i}+\frac{1}{n_{s}}\sum_{i\in\mathcal{S}_{0}}w_{i}(1-d(X_{i}))Y_{i},$
(4)

where the weights $w_{i}$ are autonormalized as

$w_{i}=\begin{cases}\dfrac{w(1,X_{i})}{\tfrac{1}{n_{s}}\sum_{i\in\mathcal{S}_{1}}w(1,X_{i})},&\text{if
}i\in\mathcal{S}_{1},\vspace{6pt}\\\
\dfrac{w(0,X_{i})}{\tfrac{1}{n_{s}}\sum_{i\in\mathcal{S}_{0}}w(0,X_{i})},&\text{if
}i\in\mathcal{S}_{0}.\end{cases}$

The weights are normalized by treatment group, i.e.,
$\sum_{i\in\mathcal{S}_{1}}w_{i}=\sum_{i\in\mathcal{S}_{0}}w_{i}=n_{s}$. Such
weighted estimators are also known as the Hájek estimator (Chattopadhyay et
al.,, 2020). Given the weights, the optimal ITR can be estimated as

$\hat{d}=\operatorname*{arg\,max}_{d\in\mathcal{D}}\widehat{\mathcal{V}}_{n}(d;w).$
(5)

When $w(1,x)=1/\pi(x)$ and $w(0,x)=1/(1-\pi(x))$, (4) becomes the objective
function of outcome weighted learning approach (Zhao et al.,, 2012).

In this section, we will study the requirements on $w(\cdot)$ or $w_{i}$ such
that maximizing $\widehat{\mathcal{V}}_{n}(d;w)$ produces a well-performing
ITR for the target population and will then develop a balancing weight
framework to meet these conditions. We first derive two different balancing
weight methods, one that makes a covariate shift adjustment in addition to
adjusting for confounding and the other that only adjusts for confounding. We
will demonstrate that each of these two approaches is suitable under different
situations, depending on whether or not the prespecified rule class contains
the unrestricted optimal ITR. However, it is not clear a priori which of the
weighting approaches is best suited for a given data-generating mechanism. To
address this, our final approach will be formulated as a hybrid of these two
balancing weighting methods that combines them in a data-driven manner.

###  3.1 General case: balancing to the target population

In this section we focus on the general case where the true optimal ITR is
potentially not contained in the prespecified class $\mathcal{D}$. In this
case we investigate what properties a set of weights should have in order to
reduce error in estimation of the optimal ITR. From the perspective of ITR
evaluation (Kallus,, 2018), a natural idea is to choose $w$ for the source
population so that make $\widehat{\mathcal{V}}_{n}(d;w)$ is as close to
$\mathcal{V}_{\mathcal{T}}(d)$ as possible for all $d\in\mathcal{D}$. Given
such weights, an ITR $d$ that is the maximizer of
$\widehat{\mathcal{V}}_{n}(d;w)$ will also achieve a high value of
$\mathcal{V}_{\mathcal{T}}(d)$.

Define $\mu_{a}(x)=\mathbb{E}[Y(a)\mid X=x]$ and $\varepsilon_{i}\equiv
Y_{i}-\mu_{A_{i}}(X_{i})$. We assume that $\varepsilon_{i}$ are mutually
independent and have uniformly bounded variance, i.e.,
$\mathbb{E}[\varepsilon_{i}^{2}\mid X_{i},A_{i}]\leq\sigma^{2}<\infty$. Let
$\mathcal{V}_{n,\mathcal{T}}(d)=\frac{1}{n_{t}}\sum_{i\in\mathcal{T}}\\{d(X_{i})\mu_{1}(X_{i})+(1-d(X_{i}))\mu_{0}(X_{i})\\}$
be the sample version of $\mathcal{V}_{\mathcal{T}}(d)$, where
$n_{t}=\sum_{i}(1-S_{i})$. Then the estimation error
$\widehat{\mathcal{V}}_{n}(d;w)-\mathcal{V}_{\mathcal{T}}(d)$ admits the
following decomposition:

$\displaystyle\widehat{\mathcal{V}}_{n}(d;w)-\mathcal{V}_{\mathcal{T}}(d)$
$\displaystyle=\frac{1}{n_{s}}\sum_{i\in\mathcal{S}_{1}}w_{i}d(X_{i})(\mu_{1}(X_{i})+\varepsilon_{i})+\frac{1}{n_{s}}\sum_{i\in\mathcal{S}_{0}}w_{i}(1-d(X_{i}))(\mu_{0}(X_{i})+\varepsilon_{i})$
(6)
$\displaystyle\quad-\frac{1}{n_{t}}\sum_{i\in\mathcal{T}}d(X_{i})\mu_{1}(X_{i})-\frac{1}{n_{t}}\sum_{i\in\mathcal{T}}(1-d(X_{i}))\mu_{0}(X_{i})+\big{\\{}\mathcal{V}_{n,\mathcal{T}}(d)-\mathcal{V}_{\mathcal{T}}(d)\big{\\}}$
$\displaystyle=\int_{\mathcal{X}}d(x)\mu_{1}(x)\textnormal{d}\big{[}\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}-\mathbb{P}_{n,\mathcal{T}}\big{]}(x)+\int_{\mathcal{X}}(1-d(x))\mu_{0}(x)\textnormal{d}\big{[}\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)}-\mathbb{P}_{n,\mathcal{T}}\big{]}(x)$
$\displaystyle\quad+\frac{1}{n_{s}}\sum_{i\in\mathcal{S}}\mathbbm{1}(d(X_{i})=A_{i})w_{i}\varepsilon_{i}+\big{\\{}\mathcal{V}_{n,\mathcal{T}}(d)-\mathcal{V}_{\mathcal{T}}(d)\big{\\}}.$

Here
$\mathbb{P}_{n,\mathcal{S}_{a}}^{(w)}=\frac{1}{n_{s}}\sum_{i\in\mathcal{S}_{a}}w_{i}\delta_{X_{i}}$
are the weighted empirical distributions of the covariates of the treated
($a=1$) and control ($a=0$) units in the source sample, and
$\mathbb{P}_{n,\mathcal{T}}$ is the empirical distribution of the covariates
of the target sample.

The decomposition (6) allows a transparent understanding about what properties
$w_{i}^{\prime}s$ should possess to ensure the estimation error of the value
function in the target population to be small. The first term in the
decomposition is the source of estimation bias arising from the imbalance
between $\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}$ and
$\mathbb{P}_{n,\mathcal{T}}$ on the function $d(x)\mu_{1}(x)$. In other words,
depending on the behavior of the unknown quantity $d(x)\mu_{1}(x)$, this term
can contribute to the bias if the weighted distribution of covariates of
treated group in the source sample differs from the empirical distribution of
covariates in the target sample. Similarly, the second term is due to the
imbalance between $\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)}$ and
$\mathbb{P}_{n,\mathcal{T}}$ on $(1-d(x))\mu_{0}(x)$. The third term
demonstrates how variance is inflated by the weighting function $w(a,x)$.
These three terms can be controlled by properly selecting the weighting
function $w(a,x)$. The last term on the right-hand side of (6) is completely
due to sampling error and is not related to $w(a,x)$ in any manner. It
converges to 0 at the rate of $\mathcal{O}(n^{-\frac{1}{2}})$ as long as the
target sample is representative of the target population.

We first consider the first term on the right-hand side of (6). Ideally we
want to eliminate or mitigate the bias by balancing
$\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}$ and $\mathbb{P}_{n,\mathcal{T}}$ on
$d(x)\mu_{1}(x)$. However, $\mu_{1}(x)$ is typically unknown. Even if
$\mu_{1}(x)$ can be well approximated by a single function, multiplication by
$d(x)$ with $d(x)$ varying over the ITR class $\mathcal{D}$ can induce a
highly complex set of functions. Therefore, we consider the maximum imbalance
between $\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}$ and
$\mathbb{P}_{n,\mathcal{T}}$ over a class of flexible functions. Specifically,
we choose this function class to be the unit ball of an RKHS $\mathcal{H}$
associated with a positive definite kernel function $K(x,x^{\prime})$. We will
restrict our attention to universal and bounded kernel functions. This
includes many commonly used kernel functions, such as the Gaussian kernel and
the Matérn kernel. This maximum or worst-case imbalance is also commonly
referred to as the maximum mean discrepancy (MMD) (Gretton et al.,, 2012):

$\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})=\sup_{h\in\mathcal{H},\|h\|_{\mathcal{H}}\leq
1}\int_{\mathcal{X}}h(x)\textnormal{d}\big{[}\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}-\mathbb{P}_{n,\mathcal{T}}\big{]}(x).$
(7)

Although $d(x)\mu_{1}(x)$ is typically not itself a member of the RKHS
$\mathcal{H}$ due to the discrete nature of $d(x)$, (7) still provides a
probabilistic bound for the imbalance between
$\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}$ and $\mathbb{P}_{n,\mathcal{T}}$ on
$d(x)\mu_{1}(x)$. Specifically, if $\mu_{1}\in\mathcal{H}$ and the kernel
function takes the form of
$K(x,x^{\prime})=\theta^{-p}K_{0}(\|x-x^{\prime}\|/\theta)$ then under
suitable regularity conditions provided in the supplementary materials, we can
show that the following

$\int_{\mathcal{X}}d(x)\mu_{1}(x)\textnormal{d}\big{[}\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}-\mathbb{P}_{n,\mathcal{T}}\big{]}(x)\leq\|\mu_{1}\|_{\mathcal{H}}\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})+C_{1}\left(\frac{1}{n\theta^{p}}+\theta^{\gamma/2}\right)$
(8)

holds with probability at least $1-\exp\\{-C_{2}n\theta^{p+\gamma}\\}$, where
$\gamma,C_{1},C_{2}$ are constants not related to $d(x)$. The derivation of
this result is based on Gretton et al., (2012) and is deferred to the
supplementary materials. This observation motivates us to find a set of
appropriate weights by minimizing
$\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})$.
Gretton et al., (2012) showed that when $\mathcal{X}$ is a compact space,
$\textnormal{MMD}_{\mathcal{H}}$ is indeed a metric of the probability
distributions on $\mathcal{X}$. Therefore, by minimizing
$\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})$
with respect to the weights, we are in essence searching for a set of weights
that make the treated group in the source sample close to the target sample
and further due to (8), this minimization can help mitigate one of the key
sources of systematic bias, i.e. the first term in the decomposition (6). By a
similar argument we propose to control the other key source of systematic
bias, the second term of (6), by optimizing
$\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)},\mathbb{P}_{n,\mathcal{T}})$.

Reducing the impact from the third term of (6) requires some element of
control over the magnitude of $w_{i}$. In the literature of covariate
balancing weights (Hainmueller,, 2012; Chan et al.,, 2016; Wang and
Zubizarreta,, 2020), this is usually achieved by minimizing a dispersion
measure of $w_{i}$. We choose to penalize a large
$\frac{1}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2}$ for several reasons.
First, the conditional mean square of the third term of (6) given
$\\{X_{i},S_{i},A_{i}\\}$,
$\frac{1}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2}\mathbbm{1}(d(X_{i})=A_{i})\textnormal{Var}(\varepsilon_{i}\mid
X_{i},S_{i},A_{i})$, is bounded by
$\frac{\sigma^{2}}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2}$. So
$\frac{1}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2}$ directly relates to the
magnitude of this term. Second, for weighted data the effective sample size is
usually approximated as $(\sum_{i}w_{i})^{2}/\sum_{i}w_{i}^{2}$ (Kish,, 1965),
which is equal to $4n_{s}^{2}/\sum_{i\in\mathcal{S}}w_{i}^{2}$ in our case due
to the normalizing constraints. Thus,
$\frac{1}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2}$ has a clear
interpretation as the inverse sample size and has been advocated for
diagnosing the stability of weighting approaches (Chattopadhyay et al.,,
2020).

Putting all the pieces discussed above together, we propose to estimate
weights via the constrained optimization

$\displaystyle\min_{w}$
$\displaystyle\qquad\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})^{2}+\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)},\mathbb{P}_{n,\mathcal{T}})^{2}+\frac{\lambda}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2},$
(9) subject to
$\displaystyle\qquad\sum_{i\in\mathcal{S}_{1}}w_{i}=\sum_{i\in\mathcal{S}_{0}}w_{i}=n_{s},$

where $\lambda>0$ is a hyperparameter controlling the level of penalty. Here
we take square of the $\textnormal{MMD}_{\mathcal{H}}$ terms for computational
convenience, which will be explained in Section 4.1.

###  3.2 Special case with correctly specified rule class: balancing between
the treated and control in the source population

In the previous section, we discussed a strategy to find an appropriate
weighting function by making $\widehat{\mathcal{V}}_{n}(d;w)$ a good estimator
for the value function $\mathcal{V}_{\mathcal{T}}(d)$ without any
consideration for whether the prespecified ITR class $\mathcal{D}$ provided a
reasonable approximation to the true optimal ITR. The resulting optimization
problem minimizes the distance between $\mathbb{P}_{n,\mathcal{S}_{a}}^{(w)}$
and $\mathbb{P}_{n,\mathcal{T}}$ subject to a smoothing penalty. However,
since the objective is to find a good ITR rather than to estimate
$\mathcal{V}_{\mathcal{T}}(d)$, such an intermediate objective might not
always be necessary. In this section, we consider the situation where the
unrestricted true optimal ITR $d^{*}(x)=\mathbbm{1}(\mu_{1}(x)\geq\mu_{0}(x))$
is contained in $\mathcal{D}$. In this case, the optimal ITR does not depend
on the distribution of the covariates, obviating the need for covariate shift
adjustments that make $\mathbb{P}_{n,\mathcal{S}_{a}}^{(w)}$ similar to
$\mathbb{P}_{n,\mathcal{T}}$. As a consequence, we will see that the balancing
requirements can be relaxed, thus allowing for use of a balancing criterion
that yields less variable weights compared with the general case in Section
3.1.

To illustrate the key source of error in estimating the optimal ITR when
$\mathcal{D}$ is well-specified, we investigate conditions under which the
maximizer of $\widehat{\mathcal{V}}_{n}(d;w)$ will be close to $d^{*}$. A
sufficient condition of such is to ensure that
$\widehat{\mathcal{V}}_{n}(d,w)$ is only be greater than
$\widehat{\mathcal{V}}_{n}(d^{*},w)$ when $d$ is close to $d^{*}$. In other
words, $\widehat{\mathcal{V}}_{n}(d^{*};w)-\widehat{\mathcal{V}}_{n}(d;w)$
must be positive whenever $d$ differs substantially from $d^{*}$. Let
$m(x)=\\{\mu_{0}(x)+\mu_{1}(x)\\}/2$ denote the main effect function and
$\tau(x)=\mu_{1}(x)-\mu_{0}(x)$ denote the conditional average treatment
effect (CATE) function. Let us consider the following decomposition

$\displaystyle\widehat{\mathcal{V}}_{n}(d^{*};w)-\widehat{\mathcal{V}}_{n}(d;w)$
$\displaystyle=\frac{1}{n_{s}}\sum_{i\in\mathcal{S}_{1}}w_{i}\\{d^{*}(X_{i})-d(X_{i})\\}\left\\{m(X_{i})+\frac{\tau(X_{i})}{2}+\varepsilon_{i}\right\\}$
(10)
$\displaystyle\quad-\frac{1}{n_{s}}\sum_{i\in\mathcal{S}_{0}}w_{i}\\{d^{*}(X_{i})-d(X_{i})\\}\left\\{m(X_{i})-\frac{\tau(X_{i})}{2}+\varepsilon_{i}\right\\}$
$\displaystyle=\frac{1}{2}\int_{\mathcal{X}}\\{d^{*}(x)-d(x)\\}\tau(x)\textnormal{d}\big{[}\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}+\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)}\big{]}(x)$
$\displaystyle\quad+\int_{\mathcal{X}}\\{d^{*}(x)-d(x)\\}m(x)\textnormal{d}\big{[}\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}-\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)}\big{]}(x)$
$\displaystyle\quad+\frac{1}{n_{s}}\sum_{i\in\mathcal{S}}(2A_{i}-1)\\{d^{*}(X_{i})-d(X_{i})\\}w_{i}\varepsilon_{i}.$

Since $d(x)$ and $d^{*}(x)$ only take value in $\\{0,1\\}$,
$d^{*}(x)-d(x)\in\\{0,\pm 1\\}$. When $\mathcal{D}$ contains the unrestricted
optimal ITR, $d^{*}(x)$ can be expressed as $\mathbbm{1}(\tau(x)\geq 0)$, so
$\tau(x)$ is non-negative when $d^{*}(x)-d(x)=1$ and negative when
$d^{*}(x)-d(x)=-1$. Consequently, the first term on the right-hand side of
(10) is always non-negative, and it will take a large positive value if $d(x)$
is substantially different from $d^{*}(x)$, especially when they disagree on
the region in the covariate space where $\tau(x)$ is of large magnitude. On
the contrary, the signs of the second and the last terms in (10) are
uncertain. If the magnitude of these two terms can be controlled, then
$\widehat{\mathcal{V}}_{n}(d^{*};w)-\widehat{\mathcal{V}}_{n}(d;w)$ can only
be positive when $d$ is sufficiently close to $d^{*}$.

The second term arises from the imbalance between
$\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}$ and
$\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)}$ on $\\{d^{*}(x)-d(x)\\}m(x)$. Using
similar arguments as in the previous section, we propose to control this
quantity by optimizing
$\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)})$.
The last term in the decomposition always has mean zero and involves the
variability of the weights; so to mitigate this noise term, we again use a
penalty term $\frac{\lambda}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2}$.
Hence, we obtain the following optimization problem

$\displaystyle\min_{w}$
$\displaystyle\qquad\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)})^{2}+\frac{\lambda}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2},$
(11) subject to
$\displaystyle\qquad\sum_{i\in\mathcal{S}_{1}}w_{i}=\sum_{i\in\mathcal{S}_{0}}w_{i}=n_{s}.$

The balancing objective in (11) is to make
$\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)}$ and
$\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)}$ close, but does not explicitly enforce
them towards any specific distribution. In Section 3.4 we formally establish
that the resulting reweighted distribution is one that focuses on the region
in the covariate space where the two treatment groups have good overlap,
similar to the notion of overlap weights (Li et al.,, 2018).

Since $\textnormal{MMD}_{\mathcal{H}}$ is a metric of probability
distributions, it follows from the triangle inequality that
$\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)})\leq\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})+\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)},\mathbb{P}_{n,\mathcal{T}})$.
This implies

$\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)})^{2}\leq
2\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})^{2}+2\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)},\mathbb{P}_{n,\mathcal{T}})^{2}.$

Thus, the balancing objective in (9) is inherently more stringent than that in
(11). As a result, (11) will generally produce weights with less variability
thus translating to a larger effective sample size. Balancing the weighted
distribution of all confounders between the treated and control groups is
indeed the minimal requirement in ITR estimation and cannot be further relaxed
unless data were generated from a randomized controlled trial or no
confounding is present. Without such a balance, the ITR obtained from
optimizing $\mathcal{V}_{n,\mathcal{T}}(d)$ might recommend a subject to take
a treatment not because it corresponds to a better potential outcome but
because more people with similar characteristics received the treatment in the
observed data.

###  3.3 Back to the general case: a three-way balancing approach

The objective (9) works for general situations but imposes stronger balancing
requirements, which may result in more extreme weights, causing the ITR
learning process to focus too heavily on a few observations with large
weights. On the other hand, the objective (11) only requires minimal balancing
requirements, but the resulting weights are only likely to yield an optimal
ITR when $\mathcal{D}$ contains the unrestricted true optimal ITR. Yet, in
practice, it is unclear how well $\mathcal{D}$ is specified and one may be
unwilling to wager that $\mathcal{D}$ contains the unrestricted true optimal
ITR and thus only balance the treated and control groups to each other. To
achieve some of the benefits of each approach, we propose to combine these two
balancing strategies by optimizing the following convex combination of them

$\displaystyle\min_{w}$
$\displaystyle\qquad\alpha\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{T}})^{2}+\alpha\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)},\mathbb{P}_{n,\mathcal{T}})^{2}$
(12)
$\displaystyle\qquad+(1-\alpha)\textnormal{MMD}_{\mathcal{H}}(\mathbb{P}_{n,\mathcal{S}_{1}}^{(w)},\mathbb{P}_{n,\mathcal{S}_{0}}^{(w)})^{2}+\frac{\lambda}{n_{s}^{2}}\sum_{i\in\mathcal{S}}w_{i}^{2},$
subject to
$\displaystyle\qquad\sum_{i\in\mathcal{S}_{1}}w_{i}=\sum_{i\in\mathcal{S}_{0}}w_{i}=n_{s}.$

Here $\alpha\in[0,1]$ is a hyperparameter controlling the extent to which the
weighted treated/control group should resemble the target sample. We present a
strategy for choosing both $\alpha$ and $\lambda$ in Section 4.2.

###  3.4 Asymptotic properties

The following theorem characterizes the convergence limit of the proposed
weights. Here we impose two conditions on the RKHS: universality and that the
unit ball is $\mathbb{P}$-Donsker. These conditions hold for most commonly
used kernels, such as the Gaussian kernel and Matérn kernel (Sriperumbudur et
al.,, 2016).

######  Theorem 1.

Suppose $\\{\hat{w_{i}},i\in\mathcal{S}\\}$ are the weights that solve (12).
Assume the RKHS is universal, and the unit ball in it $\mathcal{H}$ is
$\mathbb{P}$-Donsker. Assume $\pi(x),\rho(x)$ are continuous functions on
$\mathcal{X}$. If $0<\alpha\leq 1$, then

$\frac{1}{n_{s}}\sum_{i\in\mathcal{S}}\left(\hat{w}_{i}-w^{*}(A_{i},X_{i})\right)^{2}\overset{p}{\longrightarrow}0,$

where $w^{*}(a,x)$ is defined in (3). If $\alpha=0$,

$\frac{1}{n_{s}}\sum_{i\in\mathcal{S}}\left(\hat{w}_{i}-w^{\dagger}(A_{i},X_{i})\right)^{2}\overset{p}{\longrightarrow}0.$

Here

$w^{\dagger}(a,x)=\frac{\pi(X)(1-\pi(X))}{C_{\pi}}\left\\{\frac{a}{\pi(x)}+\frac{1-a}{1-\pi(x)}\right\\},$

where $C_{\pi}=\mathbb{E}[\pi(X)(1-\pi(X))\mid S=1]$ is a normalizing
constant.

The proofs of this theorem and the other theoretical results below are
deferred to the supplementary materials. The convergence results of the
weights straightforwardly imply the convergence limit of
$\widehat{\mathcal{V}}_{n}(d;\hat{w})$, the objective function for learning
the optimal ITR:

######  Corollary 1.1.

Suppose the conditions in Theorem 1 hold and the potential outcomes are square
integrable. Then, for any treatment rule $d$, when $\alpha>0$,
$\widehat{\mathcal{V}}_{n}(d;\hat{w})\overset{p}{\longrightarrow}\mathcal{V}_{\mathcal{T}}(d)$;
when $\alpha=0$,
$\widehat{\mathcal{V}}_{n}(d;\hat{w})\overset{p}{\longrightarrow}{C_{\pi}}^{-1}\mathbb{E}[Y(d)\pi(X)(1-\pi(X))\mid
S=1]$.

The limiting weights under both cases have the form of
$\\{{\pi(x)}^{-1}a+(1-\pi(x))^{-1}(1-a)\\}h(x).$ They balance the weighted
distributions of the covariates between the treated and control groups, and
retarget the distribution to a population whose density ratio to the source
population is $h(x)$. When $\alpha>0$, $h(x)\propto(1-\rho(x))/\rho(x)$ is
exactly the density ratio between the target population and the source
population, so the corresponding weights are the classical importance weights.
For $\alpha=0$, $h(x)\propto\pi(x)(1-\pi(x))$. With this choice we put more
focus on the covariate region where the treated and control groups have good
overlap, i.e. the region where $\pi(x)$ is close to 0.5. Li et al., (2018)
also derived weights of this form from the perspective of minimizing the
asymptotic efficient estimation variance of the weighted average treatment
effect, and they referred to $w^{\dagger}$ as the “overlap weights”. Kallus,
(2020) advocated using the overlap weights for estimating the optimal ITR with
observational data, especially when there is limited overlap between the
treated and control groups, because the overlap weights avoid putting too much
attention to a few observations and thus reduce estimation variability. In our
setting of learning optimal ITR for a target population, the overlap weights
are well-justified when $\mathcal{D}$ is well specified and our nonparametric
balancing weights converge to the overlap weights purely by means of
optimizing within-population imbalance.

Although the weights converge to the same limit for all $\alpha>0$, under
finite samples, varying $\alpha$ from 1 to 0 provides a smooth transition from
the importance weights to the overlap weights. The importance weights aim for
minimizing the distributional difference between the source sample and the
target population, while the overlap weights put more emphasis on improving
stability and reducing any imbalance between treated and control units.
Therefore, the proposed method enables us to search for a better trade-off in
the middle ground, calibrating the estimation towards the target population
while retaining small variance and mitigating imbalance due to selection bias
into the treatment.

Another advantage of our method compared to the importance weights or overlap
weights is that by targeting distributional imbalance directly, it avoids
explicit modeling of either the propensity score $\pi(x)$ and participation
probability $\rho(x)$. As discussed before, modeling may result in the failure
of weighting methods due to model misspecification or estimation error. While
the proposed weights are free of model specification issues and have more
stable values, Theorem 1 guarantees that they converge to interpretable
limits. Further, since the proposed method is entirely formulated upon
covariate balancing, it enjoys better finite sample balance properties, and
this generally leads to more precise treatment comparison in practice
(Chattopadhyay et al.,, 2020). As far as we know, our work is the first to
show that overlap weights can be estimated by constructing weights that purely
target imbalance of covariate distributions.

##  4 Computation

###  4.1 Finite sample representation

Since $\mathcal{H}$ is an infinite-dimensional function space, direct
evaluation and optimization of (12) may appear intractable. The following
lemma is based on Lemma 6 of Gretton et al., (2012) and provides a finite
sample representation for the $\textnormal{MMD}_{\mathcal{H}}$ between two
(weighted) empirical distributions.

######  Lemma 1.

Let $P=\sum_{i\in\mathcal{I}_{P}}p_{i}\delta_{x_{i}}$ and
$Q=\sum_{j\in\mathcal{I}_{Q}}q_{j}\delta_{x_{j}}$ be two discrete probability
distributions on $\mathcal{X}$. Here $\mathcal{I}_{P}$ and $\mathcal{I}_{Q}$
are two finite index sets, and $\\{p_{i},i\in\mathcal{I}_{P}\\}$ and
$\\{q_{j},j\in\mathcal{I}_{Q}\\}$ are point mass that satisfy
$\sum_{i\in\mathcal{I}_{P}}p_{i}=\sum_{i\in\mathcal{I}_{Q}}q_{i}=1$. Then

$\displaystyle\textnormal{MMD}_{\mathcal{H}}(P,Q)^{2}=\sum_{i\in\mathcal{I}_{P}}\sum_{i^{\prime}\in\mathcal{I}_{P}}p_{i}p_{i^{\prime}}K(x_{i},x_{i^{\prime}})+\sum_{j\in\mathcal{I}_{Q}}\sum_{j^{\prime}\in\mathcal{I}_{Q}}q_{j}q_{j^{\prime}}K(x_{j},x_{j^{\prime}})-2\sum_{i\in\mathcal{I}_{P}}\sum_{j\in\mathcal{I}_{Q}}p_{i}q_{j}K(x_{i},x_{j}).$

We can apply Lemma 1 to evaluate the $\textnormal{MMD}_{\mathcal{H}}$ terms in
(12). Specifically, denote
$K_{\mathcal{S}_{a},\mathcal{S}_{a^{\prime}}}=\left(K(x_{i},x_{j})\right)_{i\in\mathcal{S}_{a},j\in\mathcal{S}_{a^{\prime}}}$
for $a,a^{\prime}\in\\{0,1\\}$, and
$K_{\mathcal{S}_{a},\mathcal{T}}=\left(K(x_{i},x_{j})\right)_{i\in\mathcal{S}_{a},j\in\mathcal{T}}$.
Let $w$ be a vector of length $n_{s}$ by stacking
$(w_{i})_{i\in\mathcal{S}_{1}}$ and $(w_{i})_{i\in\mathcal{S}_{0}}$. By Lemma
1, we can write each of the $\textnormal{MMD}_{\mathcal{H}}$ terms in (12) as
a quadratic function of $w$. Then the optimization objective of (12) can be
rearranged as

$w^{\mathsf{T}}\Sigma w-2b^{\mathsf{T}}w+c$

with

$\Sigma=\frac{1}{n_{s}^{2}}\begin{bmatrix}K_{\mathcal{S}_{1},\mathcal{S}_{1}}+\lambda
I&(\alpha-1)K_{\mathcal{S}_{1},\mathcal{S}_{0}}\\\
(\alpha-1)K_{\mathcal{S}_{0},\mathcal{S}_{1}}&K_{\mathcal{S}_{0},\mathcal{S}_{0}}+\lambda
I\end{bmatrix},~{}b=\frac{\alpha}{n_{s}n_{t}}\begin{bmatrix}K_{\mathcal{S}_{1},\mathcal{T}}\mathbf{1}\\\
K_{\mathcal{S}_{0},\mathcal{T}}\mathbf{1}\end{bmatrix},~{}c=\frac{2\alpha}{n_{t}^{2}}\sum_{i\in\mathcal{T}}\sum_{j\in\mathcal{T}}K(x_{i},x_{j}).$

Here $\mathbf{1}$ is a vector of 1’s. The matrix $\Sigma$ is positive
definite, and the constraints in (12) are linear. So the optimization problem
is convex and its global optima exists. We can solve this problem using
standard quadratic programming tools.

###  4.2 Hyperparameter selection

There are two hyperparameters controlling the result of (12). The
hyperparameter $\alpha$ controls the extent to which the weighted source
sample should resemble the target sample, and $\lambda$ regularizes the
smoothness of the weights. In this section we propose a heuristic procedure to
select the hyperparameters by grid search. Before discussing the details and
motivation of the procedure, we first outline the procedure:

Input: A grid of values for $\alpha$: $\\{\alpha_{j}\\}_{j=1}^{J}$ where
$0=\alpha_{1}<\alpha_{2}<\cdots<\alpha_{J}=1$ and a grid for $\lambda$:
$\\{\lambda_{l}\\}_{l=1}^{L}$ where
$0<\lambda_{1}<\lambda_{2}<\cdots<\lambda_{L}$

for  _each $\alpha_{j}\in\\{\alpha_{j}\\}_{j=1}^{J}$_

for  _each $\lambda_{l}\in\\{\lambda_{l}\\}_{l=1}^{L}$_

Optimize (12) to obtain weights $\hat{w}$ ;

Record the resulting treatment-control imbalance using the MMD criterion

end for

Given $\alpha_{j}$, choose the $\lambda_{l}$ (denoted by
$\lambda_{\alpha_{j}}$) that results in the minimum treatment-control
imbalance and evaluate the resulting ITR’s value function on the target
sample;

end for

Output the pair $(\alpha_{j},\lambda_{\alpha_{j}})$ that results in an ITR
with the largest value

Algorithm 1 Tuning parameter selection procedure

For each $\alpha$ in a grid over $[0,1]$, we first select $\lambda$ based on
treatment-control balance. As discussed in Section 3.2, if the distribution of
confounders in the treated and control groups is not well balanced, then the
ITR may be biased towards the treatment assignment in the observed data. For
each $\lambda$, we calculate the weights under $(\alpha,\lambda)$, and
evaluate the $\textnormal{MMD}_{\mathcal{H}}$ between the weighted treated
group and the weighted control group. To obtain an “honest” evaluation, we
compute this quantity on subsamples of the source sample instead of the entire
source sample, and use the average result computed from many subsamples as the
selection criterion. With a right penalization level, the weights are robust
to sampling variation and lead to small imbalance on the subsamples. In this
way, the degree of weight penalization $\lambda$ and treatment-control balance
are not necessarily competing directly. In our implementation, the size of
subsample is set as 80% of the source sample, and the computation is done on
50 subsamples. Then the $\lambda$ that yields the smallest treatment-control
$\textnormal{MMD}_{\mathcal{H}}$ will be selected for the given $\alpha$. We
denote this selected value as $\lambda_{\alpha}$.

The selection of $\alpha$ is primarily based on an assessment of the
performance of the corresponding ITR as measured by the value function. Let
$\hat{w}_{\alpha}$ be the optimal weights under $(\alpha,\lambda_{\alpha})$.
We first solve for $\hat{d}_{\alpha}$ by maximizing
$\widehat{\mathcal{V}}_{n}(d;\hat{w}_{\alpha})$, and then evaluate the value
function of $\hat{d}_{\alpha}$. When $\alpha$ is small, the weighted treated
and control groups do not resemble the target population, so
$\widehat{\mathcal{V}}_{n}(d;\hat{w}_{\alpha})$ is not necessarily a good
estimator for $\mathcal{V}_{\mathcal{T}}(d)$. Hence, we use a regression-based
approach to evaluate the value function: the source sample is used to fit the
potential outcome models and impute the potential outcomes for the target,
then a plug-in estimate of $\mathcal{V}_{\mathcal{T}}(\hat{d}_{\alpha})$ based
on these imputations is computed. Then we select the $\alpha$ that yields the
largest $\mathcal{V}_{\mathcal{T}}(\hat{d}_{\alpha})$.

This hyperparameter tuning procedure involves solving the quadratic
programming (12) under various choices of $(\alpha,\lambda)$. To speed up
computation, when solving (12) under a new pair of hyperparameters, we use the
solution under the previous hyperparameters as a warm start. This is possible
when using the OSQP package (Stellato et al.,, 2020). The empirical results in
the following sections suggest that this tuning procedure works quite well in
practice and is able to select suitable hyperparameters.

##  5 Simulation studies

###  5.1 Setup

In this section, we conduct simulation studies to evaluate the performance of
the proposed method in finite sample settings. In our simulation setup, we
generate covariates $X_{i}$ independently from a uniform distribution on
$[-2,2]^{4}$. Given the covariates $X_{i}$, the population indicator $S_{i}$
is then generated from a Bernoulli distribution with success probability
$\rho(X_{i})$, where $\rho(x)=G(x_{2}-1.2x_{1})$, $G(z)=0.8\Phi(z)+0.1$ and
$\Phi(z)$ is the standard normal cumulative probability function. Data with
$S_{i}=1$ and $S_{i}=0$ constitute the source sample and target sample
respectively. For the source sample, we further simulate the treatment
assignment and outcome as follows. To test the performance of the proposed
method under a wide range of settings, we consider three propensity score
models for treatment assignment:

  * (1)

$\pi(x)=G(0.5x_{1}+0.3x_{2}-0.3)$; (linear assignment)

  * (2)

$\pi(x)=G(1.6x_{1}+1.3x_{2}-0.8)$; (linear assignment, bad overlap)

  * (3)

$\pi(x)=G(0.4x_{1}^{2}+0.4x_{2}^{2}+0.5x_{1}x_{2}-0.4x_{1}+0.4x_{2}-0.9)$.
(nonlinear assignment)

In the first setting the propensity score, after a monotone transformation,
linearly depends on the covariates, and thus the logistic regression can
provide a decent fit even though the link function is not perfectly specified.
The second setting is similar, but the large coefficients induce a substantial
difference between the covariate distributions of the treated and control
groups. The third setting includes quadratic terms and an interaction term,
and thus any model that only considers main effects may be insufficient for
controlling for confounding.

The observed outcome is generated as
$Y_{i}=m(X_{i})+(A_{i}-0.5)\tau(X_{i})+\varepsilon_{i}$, where
$m(x)=\Phi(-0.6x_{1}-0.6x_{2}+0.2x_{3}+0.5)+0.5$ and
$\varepsilon_{i}\mathbin{\overset{i.i.d.}{\kern 0.0pt\sim}}N(0,0.5^{2})$. The
CATE function is set as
$\tau(x)=\kappa\tau_{\text{NL}}(x)+(1-\kappa)\tau_{\text{L}}(x)$, where

$\displaystyle\tau_{\text{NL}}(x)$
$\displaystyle=\Phi(1.5x_{2}+0.8x_{1}-0.4(x_{1}-x_{2})^{2}-0.3)-0.07(x_{1}-x_{2})^{2},$
$\displaystyle\tau_{\text{L}}(x)$ $\displaystyle=\Phi(0.4x_{2}+0.6x_{1})-0.5.$

Here, $\tau(x)$ is linear combination of two parts: a nonlinear part
$\tau_{\text{NL}}(x)$ and a linear part $\tau_{\text{L}}(x)$. The parameter
$\kappa$ controls the degree of nonlinearity of $\tau(x)$. We consider two
choices of $\kappa$: 0 and 0.4, the former indicating a linear setting and the
latter a nonlinear setting. Figure 1 visualizes $\tau(x)$ under these two
choices. We choose $\mathcal{D}$ to be the linear rule class, i.e. ITRs of the
form $d(x)=\mathbbm{1}(x^{\mathsf{T}}\beta\geq 0)$. So when $\kappa=0$,
$\mathcal{D}$ contains the unrestricted true optimal ITR, and thus the optimal
ITRs for the source population and the target population are the same;
otherwise, the optimal linear ITR for the target population is different from
that for the source population, which necessitates adjustment for covariate
shift from source to target.

Figure 1: Contour plot of the CATE function $\tau(x)$ and the lines are the
optimal linear decision boundaries, with the solid line indicating the
boundary for the optimal linear ITR for the target population and the dashed
line indicating the boundary for the optimal linear ITR on the source
population.

We implement the proposed method using the Gaussian kernel with bandwidth
chosen via the median heuristic (Gretton et al.,, 2012). To find the linear
rule that maximizes $\widehat{\mathcal{V}}_{n}(d,w)$, we re-cast the
optimization problem as a weighted classification problem. Note that
$\widehat{\mathcal{V}}_{n}(d,w)$ can be expressed as

$\frac{1}{n_{s}}\sum_{i\in\mathcal{S}}\max\\{w_{i}Y_{i},0\\}-\frac{1}{n_{s}}\sum_{i\in\mathcal{S}}|w_{i}Y_{i}|\mathbbm{1}\\{d(X_{i})\neq
A_{i}\textnormal{sign}(w_{i}Y_{i})\\},$

so maximizing $\widehat{\mathcal{V}}_{n}(d,w)$ is equivalent to minimizing
$\frac{1}{n_{s}}\sum_{i\in\mathcal{S}}|w_{i}Y_{i}|\mathbbm{1}\\{d(X_{i})\neq
A_{i}\textnormal{sign}(w_{i}Y_{i})\\}$, which is a weighted binary
classification problem that uses $A_{i}\textnormal{sign}(w_{i}Y_{i})$ as the
labels and $|w_{i}Y_{i}|$ as the weights. We substitute the 0-1 loss with the
logistic loss and solve the optimization approximately with weighted logistic
regression (Xu et al.,, 2015). The predictive modeling tasks used for value
function estimation in the tuning procedure are implemented with xgboost (Chen
and Guestrin,, 2016).

For comparison, we consider other weighting methods, including the overlap
weights, the importance weights and the IPW (labeled as Overlap, Importance
and IPW, respectively). To compute these weights, we use logistic regression
to fit the propensity score and participation probability. So the propensity
score model is nearly correctly specified in the first two treatment
assignment scenarios but is misspecified in the last treatment assignment
scenario. The oracle counterparts that use the underlying true probabilities
to construct weights are also included. We further include entropy balancing
(Hainmueller,, 2012) analogues of the importance weights and IPW in the
comparison, labeled as ebal_t and ebal_s respectively: ebal_t finds the
weights of minimum entropy such that covariate moments of the reweighted
treated/control group are equal to those of the target sample, and ebal_s is
defined similarly but with covariate moments given by the whole source sample.
We note that among the comparator methods only the importance weights and
ebal_t use the covariate information from the target sample. After obtaining
the weights from these methods, we apply the same weighted classification
procedure described above to maximize $\widehat{\mathcal{V}}_{n}(d,w)$ and get
the corresponding optimal ITRs, which are then compared to the one estimated
with our weights. In the simulation, the total sample size is set to be
$n=1600$.

###  5.2 Results

The performance of all methods is measured in terms of two criteria: the
regret of the learned ITR compared to the optimal linear ITR, and the
classification accuracy using the treatment prescription from the optimal
linear ITR as the true “label”. An independent test sample of size $10^{5}$ is
generated from the target population to evaluate these quantities. We also use
this test sample to search for the optimal linear ITR. Since in our settings
only $X_{1}$ and $X_{2}$ are effect modifiers, the optimal linear ITR has the
form of $d(x)=\mathbbm{1}(\cos(\theta)x_{1}+\sin(\theta)x_{2}+b\geq 0)$, then
we find the optimal $\theta$ and $b$ using the grid search procedure described
in Kallus, (2020). The boxplots in Figure 2 summarize the results under
different scenarios based on 500 independent simulation runs.

Figure 2: Simulation results. Top: regret compared to the optimal linear ITR.
Bottom: classification accuracy with true label given by the optimal linear
ITR. Methods labeled with (true) compute the weights using the true values of
$\pi(x)$ and $\rho(x)$ instead of the estimated ones.

In the first column of Figure 2, the underlying assignment mechanisms and the
participation probabilities are linear in the covariates after a monotonic
transformation, and the treated and control groups have sufficient overlap, so
these settings constitute the easiest design across our simulation settings.
As we can see, even for this simple setting, whether calibrating to the target
sample can yield improvement in learning the ITR for the target population
does depend on the underlying outcome model. When the rule class contains the
unrestricted true optimal ITR, weighting the data to the target population
helps only little (if any) in achieving a larger value function. However, it
will make the estimated rule more vulnerable to random noise. In this case,
the overlap weights attain the best performance and our proposed method
achieves a very similar result to the overlap weights, as the tuning parameter
selection procedure allows it to yield weights that mimic the overlap weights.
On the other hand, when the difference in covariate distributions from source
to target leads to a different optimal ITR, it is necessary to tailor the
weights to the target population. As a result, the overlap weights yield a
suboptimal ITR, while our method adaptively chooses the calibration level
$\alpha$ to incorporate covariate shift adjustments and achieves the best
performance among methods both in terms of regret and accuracy. We also notice
that the importance weights based on modeling $\pi(x)$ and $\rho(x)$ can
occasionally achieve great performance, but the result is highly unstable with
some replications yielding poor performance likely as a result of large
weights.

The second and third columns of Figure 2 compare the performance under more
challenging scenarios. When there is a substantial difference between the
covariate distributions of the treated and control groups, the performance of
the overlap weights is similar to the previous cases in that it yields
excellent performance but only for correctly specified rule classes; however,
this benefit is attenuated in the nonlinear confounding setting when the model
for the propensity score is more severely misspecified. The other modeling-
based weights appear to be vulnerable to insufficient treatment-control
overlap, the misspecification of the propensity score model, or the wrong
choice of moments to balance. The proposed method achieves the best results
among all the competing methods, especially when covariate shift needs to be
adjusted for. Lastly, when the assignment mechanism is nonlinear, our method
is able to retain satisfactory performance due to its nonparametric nature and
ability to balance full covariate distributions not just prespecified
covariate moments. The relative benefit of the proposed method is the greatest
under the nonlinear assignment mechanism settings.

##  6 Application: transthoracic echocardiography and mortality in sepsis

In this section, we illustrate the proposed method by studying individualized
treatment recommendation of the use of transthoracic echocardiography (TTEC)
for intensive care unit (ICU) patients with sepsis. Feng et al., (2018)
studied the average treatment effect of TTEC on improving 28-day mortality and
suggested the use of TTEC is associated with a significant improvement.
However, It has been argued that TTEC use in surgical ICU is not cost-
effective due to a high failure rate (Cook et al.,, 2002). Hence, studying
potential heterogeneous treatment effects of TTEC and identifying a subgroup
of patients who can benefit from it is clinically relevant.

We use the same dataset in Feng et al., (2018), which is derived from the
MIMIC-III database (Johnson et al.,, 2016). This is an observational dataset
consisting of 6361 ICU patients. Among the patients, 51.3% had TTEC performed
during or in the period less than 24 hours before their ICU admission. We use
28-day survival as the outcome. We use 17 baseline variables in the analysis:
age (range 18 - 91), gender, weight; severity at admission, which is measured
by the simplified acute physiology score (SAPS), the sequential organ failure
assessment (SOFA) score, and the Elixhauser comorbidity score; comorbidity
indicators, including congestive heart failure, atrial fibrillation,
respiratory failure, malignant tumor; vital signs, including mean arterial
pressure, heart rate and temperature; laboratory results, including platelet
count, partial pressure of oxygen, lactate, and blood urea nitrogen. The
distributions of lab results are right skewed, so we apply log transformations
on these variables. All the continuous variables are then standardized for
further analysis. We impute any missing values using MissForest (Stekhoven and
Bühlmann,, 2012), which is a flexible non-parametric missing value imputation
approach.

In order to evaluate ITR learning results and study the influence of covariate
shift and confounding, we use the following procedure to construct the
training data and test data which induces different levels of covariate shift
and confounding. First, we randomly sample 40% of the data without
replacement; this part will be used to construct the source sample and the
remaining represents the target population. The sampling probability is
proportional to a function of the covariates, $f(X_{i})$. We consider two
choices of $f$: (a) $f(x)=0.5$, which means each individual is sampled with
equal probability and (b) $f(x)=G(\text{age}-\text{weight}+0.6\times
I(\text{gender}=\text{female})-0.4)$, so the source and target populations
have different covariate distributions. Here $G(\cdot)$ is the function
defined in Section 5. Then among the first part of the data, we randomly
select (without replacement) 50% of the TTEC patients and 50% of the non-TTEC
patients to form the source sample; however, the treated units are selected
with probability proportional to $g(X_{i})$ while the control units are with
probability proportional to $0.5-g(X_{i})$. This step induces additional
confounding to the source data based on $g$. Three choices of $g$ are
considered: (a) $g(x)=0.5$, so there is no extra confounding; (b)
$g(x)=G(\text{SAPS}+0.8\times\text{SOFA}-0.9\times\text{Elixhauser})$, which
introduces extra confounding determined by a linear combination of the
severity scores; (c)
$g(x)=G(0.3\times\text{SAPS}^{2}+0.2\times\text{SOFA}^{2}+0.2\times\text{Elixhauser}^{2}+0.4\times\text{SAPS}\times\text{SOFA}-0.3\times\text{SAPS}+0.4\times\text{SOFA}+0.2\times\text{Elixhauser}-0.6)$,
which leads a confounding mechanism that involves an interaction term. Among
the second part of the data, 1/3 of the patients are randomly sampled with
equal probability, and their covariate information is used as the target
sample, and the remaining data is used as the test data. So the target sample
and the test data have the same covariate distribution.

By the aforementioned procedure, we preserve the covariate-outcome
relationship in the real data, but use various sampling designs to induce
different levels of covariate shift and confounding. There are 6 scenarios;
under each one, we run 400 replications of the sampling procedure. For each
replication, we then apply the proposed method to estimate the optimal ITR for
the target population. The competing methods in Section 5 are also applied,
except the oracle ones using the true probabilities as they are unknown. In
addition, we include a trivial treatment rule that assigns all patients to
TTEC (labeled as treat_all) to help calibrate the performances of the methods
relative to the average effect of TTEC. The evaluation metric is a modified
version of the value function:
$\mathcal{V}_{\mathcal{T},2}(d)=\mathbb{E}[Y(d(X))-Y(1-d(X))\mid S=0]$. For
comparison purpose, $\mathcal{V}_{\mathcal{T},2}(d)$ is equivalent to
$\mathcal{V}_{\mathcal{T}}(d)$ as
$\mathcal{V}_{\mathcal{T},2}(d)=2\mathcal{V}_{\mathcal{T}}(d)-\mathbb{E}[Y(1)+Y(0)\mid
S=0]$. We use $\mathcal{V}_{\mathcal{T},2}(d)$ because it is equal to the ATE
when $d$ is the trivial treat-all ITR. We estimate
$\mathcal{V}_{\mathcal{T},2}(d)$ on the test data using the conventional
inverse probability estimator. The result is reported in Figure 3.

Figure 3: Modified version of value function in the real data application.
treat_all is a trivial ITR that assigns all patients to treatment and the
result equals the ATE. The other methods are the same as those in Figure 2.

The modified value function of the trivial treat-all ITR is around 6% (i.e. a
6% improvement in 28-day survival rate over no TTEC), which is close to the
ATE estimate in Feng et al., (2018) (5%). The value achieved by the best
performing ITRs is around 8%. This suggests the presence of heterogeneity of
treatment effect, thus personalizing decisions of whether to use TTEC can
potentially lead to an improved overall survival rate. The performance of all
the methods are close when there is no covariate shift and extra confounding.

When the source population and the target population have different covariate
distributions (second row of Figure 3), the methods that calibrate to the
target sample, ebal_t and Importance, perform worse than the others. This
suggests that a linear rule is likely sufficient and that calibrating the
covariates to a specific distribution may be unnecessary; in this case, the
overlap weights is a desirable option. The advantage of overlap weights
becomes more noticeable when the treated and control groups have less overlap
(second column of Figure 3). In this case the proposed method also achieves
good performance because it can adaptively approximate the overlap weights by
selecting a small $\alpha$ value. From the third column of Figure 3 we can see
that the entropy balancing and probability modeling methods are subject to the
choice of balancing functions or model specification. In contrast, the
proposed method is free of such specification issues. Further, we can see that
even with covariate shift and extra confounding with an interaction, our
proposed method does nearly as well as when there is no covariate shift and no
extra confounding.

##  7 Discussion

In this work, we have proposed a weighting framework for ITR learning for the
two-population setting where the target population has different covariate
distribution than the source population. The proposed method seeks covariate
functional balancing over a rich RKHS. The resulting weights adjust for
confounding and covariate shift simultaneously. It is always essential to
achieve treatment-control balance; however, correcting for covariate shift
does not always lead to an improved ITR, especially when the prespecified
class of ITRs contains the unrestricted optimal ITR. Moreover, covariate shift
adjustments may introduce additional variability into the learning process.
Therefore, the proposed method allows for different degrees of covariate shift
adjustment while maintaining control over the discrepancy between the treated
and control groups. Our theoretical results indicate that the proposed
framework encompasses both the importance weights and the overlap weights in
the limiting sense. Thus, by carefully selecting the hyperparameters, it can
achieve a better bias-variance trade-off for finite sample settings than these
two extreme cases. We have devised a regression-based procedure for tuning the
hyperparameters selection and demonstrated the proposed method’s efficiency
and robustness via various numerical examples. A key benefit of our method is
that it directly targets covariate balance and does not require probability
modeling. As such, our method can achieve better finite sample covariate
balance, which generally leads to improved performance for learning ITR in
practice.

In the methodological development we have focused on the objective function
(2), which is based on the outcome weighted learning approach (Zhao et al.,,
2012). The proposed weights, however, can be used in any other weighting-based
ITR learning approaches that utilize the IPW, such as residual weighted
learning of Zhou et al., (2017) and the weighting method of Chen et al.,
(2017). By leveraging the observed subject level covariates information from
the target population, the proposed weights enhance these ITR learning methods
for the generalization setting and can lead to a better ITR for the target
population.

Our proposed method is developed for a two-population setting where the target
population differs from the source, but it is directly applicable to learning
ITR for the source population. To handle optimal ITR estimation for the source
population, we can simply substitute $\mathbb{P}_{n,\mathcal{T}}$ in (12) with
$\mathbb{P}_{n,\mathcal{S}}$, the empirical covariate distribution of the
source sample. Similar to the two-population setting, while treated and
control groups should be weighted to have similar distributions, this weighted
distribution does not necessarily need to resemble the source population,
especially when the ITR class is correctly specified (Kallus,, 2020). In this
case, our method provides a full spectrum from the inverse propensity score
weights to the overlap weights and is able to adaptively select
hyperparameters to provide a suitable mixture of these two extremes.

Although our framework is formulated upon functional covariate balancing, it
can also be extended to incorporate other distribution metrics. In (12) we can
replace $\textnormal{MMD}_{\mathcal{H}}$ with other metrics of probability
distributions, such as the energy distance (Székely and Rizzo,, 2013) and
Wasserstein distance (Villani,, 2008). The resulting weights can also balance
the overall distribution of the covariates. However, the results in Theorem 1
are derived based on the properties of $\textnormal{MMD}_{\mathcal{H}}$, so
the theoretical properties under other choices of distribution metrics would
require further investigation.

## References

  * Athey and Wager, (2021) Athey, S. and Wager, S. (2021).  Policy learning with observational data.  Econometrica, 89(1):133–161.
  * Bottou et al., (2013) Bottou, L., Peters, J., Quiñonero-Candela, J., Charles, D. X., Chickering, D. M., Portugaly, E., Ray, D., Simard, P., and Snelson, E. (2013).  Counterfactual reasoning and learning systems: The example of computational advertising.  The Journal of Machine Learning Research, 14(1):3207–3260.
  * Chan et al., (2016) Chan, K. C. G., Yam, S. C. P., and Zhang, Z. (2016).  Globally efficient non-parametric inference of average treatment effects by empirical balancing calibration weighting.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 78(3):673–700.
  * Chattopadhyay et al., (2020) Chattopadhyay, A., Hase, C. H., and Zubizarreta, J. R. (2020).  Balancing vs modeling approaches to weighting in practice.  Statistics in Medicine, 39(24):3227–3254.
  * Chen et al., (2017) Chen, S., Tian, L., Cai, T., and Yu, M. (2017).  A general statistical framework for subgroup identification and comparative treatment scoring.  Biometrics, 73(4):1199–1209.
  * Chen and Guestrin, (2016) Chen, T. and Guestrin, C. (2016).  Xgboost: A scalable tree boosting system.  In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 785–794.
  * Colnet et al., (2020) Colnet, B., Mayer, I., Chen, G., Dieng, A., Li, R., Varoquaux, G., Vert, J.-P., Josse, J., and Yang, S. (2020).  Causal inference methods for combining randomized trials and observational studies: a review.  arXiv preprint arXiv:2011.08047.
  * Cook et al., (2002) Cook, C. H., Praba, A. C., Beery, P. R., and Martin, L. C. (2002).  Transthoracic echocardiography is not cost-effective in critically ill surgical patients.  Journal of Trauma and Acute Care Surgery, 52(2):280–284.
  * Dahabreh et al., (2020) Dahabreh, I. J., Robertson, S. E., Steingrimsson, J. A., Stuart, E. A., and Hernan, M. A. (2020).  Extending inferences from a randomized trial to a new target population.  Statistics in Medicine, 39(14):1999–2014.
  * Feng et al., (2018) Feng, M., McSparron, J. I., Kien, D. T., Stone, D. J., Roberts, D. H., Schwartzstein, R. M., Vieillard-Baron, A., and Celi, L. A. (2018).  Transthoracic echocardiography and mortality in sepsis: analysis of the mimic-iii database.  Intensive Care Medicine, 44(6):884–892.
  * Foster et al., (2011) Foster, J. C., Taylor, J. M., and Ruberg, S. J. (2011).  Subgroup identification from randomized clinical trial data.  Statistics in Medicine, 30(24):2867–2880.
  * Gretton et al., (2012) Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. (2012).  A kernel two-sample test.  The Journal of Machine Learning Research, 13(1):723–773.
  * Hainmueller, (2012) Hainmueller, J. (2012).  Entropy balancing for causal effects: A multivariate reweighting method to produce balanced samples in observational studies.  Political Analysis, 20(1):25–46.
  * Hirshberg et al., (2019) Hirshberg, D. A., Maleki, A., and Zubizarreta, J. (2019).  Minimax linear estimation of the retargeted mean.  arXiv preprint arXiv:1901.10296.
  * Huling and Mak, (2020) Huling, J. D. and Mak, S. (2020).  Energy balancing of covariate distributions.  arXiv preprint arXiv:2004.13962.
  * Imai and Ratkovic, (2014) Imai, K. and Ratkovic, M. (2014).  Covariate balancing propensity score.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 76(1):243–263.
  * Imbens and Rubin, (2015) Imbens, G. W. and Rubin, D. B. (2015).  Causal Inference for Statistics, Social, and Biomedical Sciences: An Introduction.  Cambridge University Press.
  * Institute of Medicine, (2013) Institute of Medicine (2013).  Observational Studies in a Learning Health System: Workshop Summary.  The National Academies Press, Washington, DC.
  * Johnson et al., (2016) Johnson, A. E., Pollard, T. J., Shen, L., Li-Wei, H. L., Feng, M., Ghassemi, M., Moody, B., Szolovits, P., Celi, L. A., and Mark, R. G. (2016).  Mimic-iii, a freely accessible critical care database.  Scientific Data, 3(1):1–9.
  * Kallus, (2018) Kallus, N. (2018).  Balanced policy evaluation and learning.  In Advances in Neural Information Processing Systems, pages 8895–8906.
  * Kallus, (2020) Kallus, N. (2020).  More efficient policy learning via optimal retargeting.  Journal of the American Statistical Association, pages 1–13.
  * Kang et al., (2007) Kang, J. D., Schafer, J. L., et al. (2007).  Demystifying double robustness: A comparison of alternative strategies for estimating a population mean from incomplete data.  Statistical Science, 22(4):523–539.
  * Kish, (1965) Kish, L. (1965).  Survey sampling.  Wiley, New York.
  * Kitagawa and Tetenov, (2018) Kitagawa, T. and Tetenov, A. (2018).  Who should be treated? empirical welfare maximization methods for treatment choice.  Econometrica, 86(2):591–616.
  * Kube et al., (2019) Kube, A., Das, S., and Fowler, P. J. (2019).  Allocating interventions based on predicted outcomes: A case study on homelessness services.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 622–629.
  * Li et al., (2018) Li, F., Morgan, K. L., and Zaslavsky, A. M. (2018).  Balancing covariates via propensity score weighting.  Journal of the American Statistical Association, 113(521):390–400.
  * Mo et al., (2020) Mo, W., Qi, Z., and Liu, Y. (2020).  Learning optimal distributionally robust individualized treatment rules.  Journal of the American Statistical Association, pages 1–16.
  * Murdoch et al., (2019) Murdoch, W. J., Singh, C., Kumbier, K., Abbasi-Asl, R., and Yu, B. (2019).  Definitions, methods, and applications in interpretable machine learning.  Proceedings of the National Academy of Sciences, 116(44):22071–22080.
  * Qian and Murphy, (2011) Qian, M. and Murphy, S. A. (2011).  Performance guarantees for individualized treatment rules.  The Annals of Statistics, 39(2):1180.
  * Rosenbaum and Rubin, (1983) Rosenbaum, P. R. and Rubin, D. B. (1983).  The central role of the propensity score in observational studies for causal effects.  Biometrika, 70(1):41–55.
  * Rudolph and van der Laan, (2017) Rudolph, K. E. and van der Laan, M. J. (2017).  Robust estimation of encouragement-design intervention effects transported across sites.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 79(5):1509.
  * Sriperumbudur et al., (2016) Sriperumbudur, B. et al. (2016).  On the optimal estimation of probability measures in weak and strong topologies.  Bernoulli, 22(3):1839–1893.
  * Stekhoven and Bühlmann, (2012) Stekhoven, D. J. and Bühlmann, P. (2012).  Missforest—non-parametric missing value imputation for mixed-type data.  Bioinformatics, 28(1):112–118.
  * Stellato et al., (2020) Stellato, B., Banjac, G., Goulart, P., Bemporad, A., and Boyd, S. (2020).  OSQP: an operator splitting solver for quadratic programs.  Mathematical Programming Computation, 12(4):637–672.
  * Sugiyama et al., (2007) Sugiyama, M., Krauledat, M., and MÃžller, K.-R. (2007).  Covariate shift adaptation by importance weighted cross validation.  Journal of Machine Learning Research, 8(May):985–1005.
  * Székely and Rizzo, (2013) Székely, G. J. and Rizzo, M. L. (2013).  Energy statistics: A class of statistics based on distances.  Journal of Statistical Planning and Inference, 143(8):1249–1272.
  * Villani, (2008) Villani, C. (2008).  Optimal transport: old and new, volume 338.  Springer Science & Business Media.
  * Wahba, (1990) Wahba, G. (1990).  Spline models for observational data.  SIAM, Philadelphia.
  * Wang et al., (2016) Wang, Y., Wu, P., Liu, Y., Weng, C., and Zeng, D. (2016).  Learning optimal individualized treatment rules from electronic health record data.  In 2016 IEEE International Conference on Healthcare Informatics (ICHI), pages 65–71. IEEE.
  * Wang and Zubizarreta, (2020) Wang, Y. and Zubizarreta, J. R. (2020).  Minimal dispersion approximately balancing weights: asymptotic properties and practical considerations.  Biometrika, 107(1):93–105.
  * Wong and Chan, (2018) Wong, R. K. and Chan, K. C. G. (2018).  Kernel-based covariate functional balancing for observational studies.  Biometrika, 105(1):199–213.
  * Xu et al., (2015) Xu, Y., Yu, M., Zhao, Y.-Q., Li, Q., Wang, S., and Shao, J. (2015).  Regularized outcome weighted subgroup identification for differential treatment effects.  Biometrics, 71(3):645–653.
  * Zhang et al., (2018) Zhang, Y., Laber, E. B., Davidian, M., and Tsiatis, A. A. (2018).  Interpretable dynamic treatment regimes.  Journal of the American Statistical Association, 113(524):1541–1549.
  * Zhao et al., (2012) Zhao, Y., Zeng, D., Rush, A. J., and Kosorok, M. R. (2012).  Estimating individualized treatment rules using outcome weighted learning.  Journal of the American Statistical Association, 107(499):1106–1118.
  * Zhao et al., (2019) Zhao, Y.-Q., Zeng, D., Tangen, C. M., and LeBlanc, M. L. (2019).  Robustifying trial-derived optimal treatment rules for a target population.  Electronic Journal of Statistics, 13(1):1717.
  * Zhou et al., (2017) Zhou, X., Mayer-Hamblett, N., Khan, U., and Kosorok, M. R. (2017).  Residual weighted learning for estimating individualized treatment rules.  Journal of the American Statistical Association, 112(517):169–187.