\section{Introduction.}\label{intro}

\section{Introduction}
\label{section:intro}

Personalized decision-making is concerned with determining the best treatment for an individual based on their personal characteristics. This problem arises in a variety of domains, from optimizing online advertisements for different users \citep{li2010contextual} to assigning suitable therapies for patients suffering from a medical condition \citep{flume2007cystic}. In this paper, we aim to design such personalized policies based on data collected in deployment rather than during randomized trials. In such settings, historical treatments are not assigned at random but rather based on a (potentially unknown) policy (e.g., social workers or doctors triaging clients/patients based on their personal characteristics or symptoms). Being able to design personalized treatment assignment policies without running randomized trials \comment{that assign or withhold treatment at random}is extremely important since such trials are often considered unethical, impractical, or may take too long to conduct. Our work is in particular motivated by high-stakes domains such as public health and social interventions, where both interpretability and optimality are critical \citep{rudin2019stop}.

We therefore focus our attention on the design of \emph{optimal prescriptive trees}. A prescriptive tree takes the form of a binary tree. In each branching node of the tree, a binary test is performed on a specific feature. Two branches emanate from each branching node, with each branch representing the outcome of the test. If a datapoint passes (resp.\ fails) the test, it is directed to the left (resp.\ right) branch. A treatment is assigned to all leaf nodes. Thus, each path from root to leaf represents a treatment assignment rule that assigns the same treatment to all datapoints that reach that leaf. The goal of learning prescriptive trees is to select branching decisions and treatment assignments such that the expected value of the outcomes of the treatments on the population are maximized. We say that a prescriptive tree is \emph{optimal} (in-sample, given a training set) if there exists a mathematical proof that no other tree yields better expected outcomes in the population used for training the method. Accordingly, we will say that a tree is \emph{optimal out-of-sample} if there exists a proof that no other tree yields better expected outcomes in the entire population.

\subsection{Problem Statement}
\label{section:problem_statement}

We now formalize the problem we study. In our discussion, we use terminology and concepts from \cite{hernan2019causal}. The goal is to design a personalized policy $\pi : \sets X \rightarrow \sets K$ in the form of a prescriptive tree that maps an individual's characteristics $x \in \sets X \subseteq \reals^F$ to a treatment from a finite set of candidates indexed in the set $\sets{K}$. We note that the set $\sets K$ may also include a ``no treatment'' option. Each individual is characterized by their covariates $X \in \reals^F$ and their \emph{potential outcomes} $Y(k) \in \sets Y \subseteq \reals$ under each treatment $k \in \sets K$\comment{, see e.g., \cite{hernan2019causal}}. The joint distribution of $X$ and $\{Y(k)\}_{k \in \sets K}$ is unknown. However, we have access to $I$ historical observations $\sets D:= \{(X_i, K_i, Y_i)\}_{i \in \mathcal{I}}$ indexed in the set $\sets I :=\{1,\ldots, I\}$, where $X_i$ denote the covariates of the $i$th observation, $K_i$ is the treatment assigned to it, and $Y_i = Y_i(K_i)$ is the \emph{observed outcome}, i.e., the outcome under the treatment received\comment{, see e.g., \cite{hernan2019causal}}. We use the convention that higher values of $Y_i$ are more desirable. Critically, as in all observational experiments, we cannot control the historical treatment assignment policy and the outcomes $Y_i(k)$ for $k \neq K_i$ remain \emph{unobserved}. Formally,  our aim is to learn, from the observational data $\sets D$, a \comment{personalized treatment assignment} policy $\pi$ in the form of a tree of maximum depth $d$ that maximizes the quantity
$$
Q(\pi) \; := \;  \mathbb{E}\left[Y(\pi(X))\right],
$$
where $\mathbb E(\cdot)$ denotes the expectation operator with respect to the joint distribution $\mathbb P$ of $X, K, Y(1), \ldots, Y(k)$. A major challenge in learning the best policy $\pi$ is the fact that we cannot observe the \emph{counterfactual} outcomes $Y_i(k)$, $k\in \sets K, \; k \neq K_i$ that were not received by datapoint~$i$.  Indeed, the missing data make it difficult to identify the best possible treatment for each datapoint. In fact, without further assumptions on the historical policy, it is impossible to even identify the treatment that is the best \emph{in expectation} for any given datapoint.

\subsection{Background on Causal Inference}

In this section, we provide a brief overview of the tools from causal inference that motivate the assumptions we make in our work and our proposed methods.

\subsubsection{Identifiability, Randomized Experiments, and Observational Studies.}
\label{section:randomization}

Learning an optimal prescriptive tree \comment{(or more generally a policy $\pi$)} from the observed data requires that the expected values of the counterfactual outcomes be \emph{identifiable}, i.e., possible to express as a function of the observed data. A sufficient condition for identifiability is for the data to have been collected during a \emph{(marginally) randomized} or \emph{conditionally randomized} experiment, see \cite{hernan2019causal}.

In a randomized experiment, \comment{treatments are assigned at random and }all individuals have the same chance of getting any given treatment independent of their characteristics $x$. \comment{(although treatments need not be equally probable)} Formally, we have $K \independent X$. When treatment assignment is randomized, we say that treatment groups are \emph{exchangeable} because it is irrelevant which group received a particular treatment. Exchangeability means that for all $k$ and~$k' \in \sets K$, those who got treatment~$k$ in the data would have the same outcome distribution under $k'$ as those who got treatment~$k'$. Equivalently, exchangeability means that the counterfactual outcome and the actual treatment are independent, or~$Y(k) \independent K$, for all values of~$k$.

In a conditionally randomized experiment, treatment assignment probabilities depend on the characteristics $x$ of each individual. A conditionally randomized experiment will, in general, not result in exchangeability because treatment groups may differ systemically from one another. However, treatment groups are exchangeable \emph{conditionally on~$X$} since no information other than~$X$ is used to decide the treatment. Thus, conditional randomization ensures \emph{conditional exchangeability}, i.e., $Y(k) \independent K \; | \; X$, for all values of~$k$. We can then infer the outcome distribution under treatment $k'$ for those with covariates $x$ that got treatment $k$ in the data by looking at their counterparts with the same (or similar) $x$ that received treatment $k'$.

As mentioned previously, in high-stakes domains we typically only have access to data from an \emph{observational study}. In such contexts, it is common practice in the literature on causal inference to assume that treatment is assigned at random conditional on $X$, although this may be an approximation and not possible to check in practice. The observational study can then be viewed as a conditionally randomized experiment, provided: \emph{a)} interventions are well defined and are recorded as distinct treatment values in the data; \emph{b)} the conditional probability of receiving a treatment depends only on the covariates $X$; and \emph{c)} the probability of receiving any given treatment conditional on $X$ is positive, i.e., $\mathbb P(K=k|X=x)>0$ almost surely for all $k$. This last condition is termed \emph{positivity} and holds in practice for randomized and conditionally randomized experiments.

We now discuss several methods from the literature for evaluating the performance of a \emph{counterfactual} treatment assignment policy $\pi$ from (conditionally) randomized experiments or, more generally, from observational studies that satisfy the above conditions. In particular, in the following and throughout the paper, we assume that treatments in the data have been assigned according to a (potentially unknown) \emph{logging policy} $\mu$, where $\mu(k,x) := \mathbb P(K=k|X=x)$ and $\mu(k,x)>0$ for all $k \in \sets K$ and $x \in \sets X$.

\subsubsection{Inverse Propensity Weighting\newpv{.}}
\label{section:ipw_intro}

The performance of a counterfactual treatment assignment policy $\pi$ can be evaluated using Inverse Probability Weighting (IPW).

IPW was originally proposed to estimate causal quantities such as expected values of counterfactual outcomes, average treatment effects, risk ratios, etc., see \cite{horvitz1952generalization}. It relies on reweighting the outcome for each individual $i$ in the dataset by the inverse of their \emph{propensity score}, given by $\mu(K_i,X_i)$. Such reweighting has the effect of creating a \emph{pseudo-population} where all individuals in the data are hypothetically given all treatments (thus simulating a randomized treatment assignment). Doing so allows for the estimation of the distribution of the unobserved counterfactual outcomes for all possible values of $X$.

Importantly, IPW can also be used to estimate the performance of a counterfactual policy~$\pi$ by reweighting each individual~$i$ in the data by $\frac{\mathbbm{1}(\pi(X_i) = K_i)}{\mu(K_i, X_i)}$ as illustrated in Figure~\ref{fig:ipw_policy_example}, see also \cite{charles2013counterfactual}. This reweigthing simulates the performance of the policy $\pi$ by leveraging the conditional exchangeability property of the data. Using this estimator, the quantity $Q(\pi)$ is estimated from the data as
\begin{equation}\label{eq:ipw_obj}
\QipwI(\pi) := \frac{1}{I}\sum_{i \in \mathcal I} \frac{\mathbbm{1}(\pi(X_i) = K_i) }{\hat \mu(K_i, X_i)} Y_i,
\end{equation}
where $\hat \mu$ is an estimator of $\mu$, which is obtained for instance using machine learning, by fitting a model to $\{(X_i,K_i)\}_{i \in \sets I}$. If $\mu$ is known or if~$\hat \mu$ converges almost surely to~$\mu$, the IPW estimator of $Q(\pi)$ is unbiased. However, it may suffer from high variance, particularly if some of the propensity scores $\mu(K, X)$ are small, see~\cite{Dudik2011DoublyLearning}.

\begin{figure}[h]
\includegraphics[width=\textwidth]{ipw_policy_example_v4}
\centering
\caption{Figure illustrating the evaluation of the performance of the counterfactual policy~$\pi$ (bottom tree) from data collected using the logging policy~$\mu$ (top tree) using the IPW estimator--notational conventions based on~\cite{hernan2019causal}. The historical data presents 50 instances, each with a single binary covariate $X^1 \in \{0,1\}$. $X^1 = 1$ (resp.\ $0$) indicates that the patient is sick (resp. healthy). There are two treatment options ($\sets K:=\{0,1\}$), where $K=1$ ($0$) is to (not) treat the patient. The potential outcomes are binary ($Y(k) \in \{0,1\}$, $k \in \sets K$). The expected reward under policy $\mu$ is $0.6$, while policy $\pi$ has a $0.82$ probability of getting a positive outcome--in fact, it can be shown that $\pi$ is optimal.}
\label{fig:ipw_policy_example}
\end{figure}

\subsubsection{Direct Method.} \label{section:dm}

An alternative estimator of the performance of the counterfactual policy $\pi$ is the Direct Method (DM)\comment{, see \cite{Dudik2011DoublyLearning}}.

DM is based upon the so-called Regress and Compare (RC) approach, which proceeds in three steps to design a policy that maximizes the expected value of the outcomes \citep{Kallus2018BalancedLearning}. First, it partitions the dataset $\sets{D}$ by treatments. Second, for all $k\in \sets K$, it learns a model $\hat{\nu}_{k}(X)$ of $\mathbb{E}(Y|K=k, X)$ using the subpopulation that was assigned treatment $k$. Third, given any $X\in \reals^F$, it estimates $Y(k)$ as $\hat{\nu}_{k}(X)$ and assigns the best treatment according to $\hat \nu$, that is, $\pi^{\text{RC}}(X) = \argmax_{k \in \sets{K}} \hat{\nu}_{k}(X)$. RC has been used, for example in the online setting by \cite{bastani2020online} and \cite{qian2011performance}. Usually, variants of linear regression are used to get estimators $\hat{\nu}$, see \cite{bastani2020online,goldenshluger2013linear,li2010contextual}.

The RC approach, however, does not allow the evaluation of arbitrary policies nor does it enable the design of policies that belong to an arbitrary class. DM adapts the regress and compare framework to resolve these shortcomings. It proposes to estimate $Q(\pi)$ via
\begin{equation}\label{eq:dm_obj}
\QdmI(\pi) := \frac{1}{I}\sum_{i \in \mathcal I} \hat{\nu}_{\pi(X_i)}(X_i).
\end{equation}
Both RC and DM result in poor performance if $\hat{\nu}_k(X)$, $k\in \sets K$, are biased estimators or make predictions that are far from the true expected outcomes, see~\cite{beygelzimer2009offset}. As pointed out by~\cite{Dudik2011DoublyLearning}, $\hat{\nu}_k(X)$ are constructed without information about~$\pi$ and thus may approximate $\mathbb E(Y|K=k,X)$ poorly in areas that are relevant for $Q(\pi)$. That said, if $\hat \nu_k$ are unbiased for all $k \in \sets K$, the DM estimator of~$Q(\pi)$ will be unbiased.

\subsubsection{Doubly Robust Approach.}
\label{section:robust_intro}

Doubly robust (DR) estimation is a family of techniques that combine two estimators. If either one of the two estimators is accurate, then the doubly robust approach is also accurate, which dampens the errors brought by the individual estimators.

In particular, realizing the drawbacks of IPW and DM, \cite{Dudik2011DoublyLearning} proposed a doubly robust approach that estimates the counterfactual performance of a policy $\pi$ as
\begin{equation}
\label{eq:dr_obj}
\QdrI(\pi) := \frac{1}{I}\sum_{i \in \mathcal I} \left(\hat{\nu}_{\pi(X_i)}(X_i)+ (Y_i - \hat{\nu}_{K_i}(X_i)) \frac{\mathbbm{1}(\pi(X_i) = K_i)}{\hat \mu(K_i, X_i)}\right).
\end{equation}
Intuitively, DR uses IPW to estimate the difference between the true outcome and the prediction made by the DM estimator, $Y_i - \hat{\nu}_{K_i}(X_i)$, thus removing any bias caused by DM, assuming the IPW estimator is unbiased. Provided at least one of $\hat \mu$ or $\hat \nu$ converges almost surely to $\mu$ or $\nu$, respectively, the DR estimator of $Q(\pi)$ will be asymptotically unbiased. Furthermore, in practice, $\QdrI$ usually has a smaller variance compared to $\QipwI$ but a higher variance than $\QdmI$, see~\cite{Dudik2011DoublyLearning}.

\subsection{Related Works}\label{section:related_works}
In this section, we position our work in the fields of operations research, causal inference, and machine learning.

\subsubsection{Estimating Heterogeneous Causal Effects.}
To estimate heterogeneous causal effects, \cite{Athey2016RecursiveEffects} and \cite{wager2018estimation} use recursive partitioning to produce causal trees and causal forests, respectively. Non-tree-based methods can be found in \cite{abrevaya2015estimating}, which use inverse propensity weighting, and \cite{fan2020estimation}, which extends \cite{abrevaya2015estimating} to account for high-dimensional data. \cite{powers2018some} adapt three methods--random forests, boosting, and MARS (Multivariate Adaptive Regression Splines)--to also estimate treatment effects in high-dimensional data. The overall goal of these approaches is to find good estimates of the treatment effects rather than to design treatment assignment policies. While some of these, such as the method from \cite{Athey2016RecursiveEffects}, can be \emph{adapted} to design tree policies, these policies will not be interpretable due to the need to combine the predictive trees associated with many different treatments.

\subsubsection{Mixed-Integer Optimization.}

MIO has been gaining traction as a tool to solve challenging learning problems. It has been used for example to tackle sparse regression problems \citep{wilson2017alamo, atamturk2019rank,bertsimas2020sparse,hazimeh2020sparse,xie2020scalable,gomez2021mixed}, verification of neural networks \citep{fischetti2018deep,khalil2018combinatorial,tjandraatmadja2020convex}\newpv{,} and sparse principal component analysis \citep{dey2018convex, bertsimas2020solving}, among others. More importantly in the context of this paper, MIO methods have been proposed to learn optimal decision trees \citep{bertsimas2017optimal, verwer2019learning, aghaei2019learning, aghaeistrong, elmachtoub2020decision, mivsic2020optimization}.

MIO technology has also been used in causal inference. Existing approaches have mostly focused on matching \citep{bennett2020building,zubizarreta2017optimal}, i.e., to estimate causal effects by pairing each treatment instance with a control instance that has similar covariates \citep{rubin2006matched}. A related approach to matching that utilizes MIO technology is subset selection to achieve \emph{balance} (i.e., similar distributions) between the covariates in the treatment and control groups \citep{nikolaev2013balance}. Finally, \citet{mintz2017behavioral} use MIO to estimate the effects of weight-loss interventions in myopic agents and design optimal policies from the perspective of a healthcare provider.

\subsubsection{Learning Prescriptive Trees.}
\label{sec:lit_prescriptive_trees}

In a recent work closely related to this paper, \cite{Kallus2017RecursiveData} proposes an MIO formulation to learn optimal prescriptive trees that relies on a critical assumption: that trees are sufficiently deep such that the historical treatment assignment and covariates are independent \emph{at each leaf}. Building upon this approach, \citet{Bertsimas2019OptimalTrees} seek to ameliorate it by augmenting the objective function with a term aimed to improve accuracy of the outcome predictions. The paper uses coordinate descent with multiple starts to train the prescriptive trees (as opposed to MIO technology), and empirically verifies that the resulting solutions are indeed high-quality. We provide an in-depth analysis of both methods in Section~\ref{section:prev_mio} and compare to them in our experiments.

\subsection{Contributions}\label{section:contributions}

In this work, we do an in-depth analysis of existing methods to learn prescriptive trees, and propose new MIO formulations for solving such problems. In particular, our key contributions are:
\begin{enumerate}[label=\emph{(\alph*)}]
    \item We analyze the assumption made by the existing methods to learn prescriptive trees discussed in Section~\ref{sec:lit_prescriptive_trees}. We demonstrate, by means of examples, that when the assumption fails to hold \emph{for all feasible trees}, the trees returned as ``optimal'' by these methods may be severely suboptimal. They may even branch on noise covariates that are not predictive of outcomes nor treatment assignments. Moreover, our examples show that this assumption is very strict in the sense that trees deep enough to allow branching on all levels of all covariates (including covariates that are independent of outcomes and treatment assignments) may be needed for this assumption to hold at an optimal tree
    \item Motivated by the limitations of existing approaches, we propose novel MIO formulations for learning optimal prescriptive trees based on tools from causal inference. Contrary to methods from the literature, our approaches enable the design of shallow and/or deep trees, thus allowing the decision-maker to tune the trade-off between prescription accuracy and interpretability. We demonstrate that, under mild conditions,  our methods are asymptotically exact in the sense that they converge to an optimal out-of-sample treatment assignment policy as the number of historical data samples tends to infinity. Based on extensive computational experiments on both synthetic and real data, we demonstrate that our asymptotic guarantees translate to significant out-of-sample performance improvements even in finite samples.
\end{enumerate}

The remainder of this paper is organized as follows. In Section \ref{section:prev_mio}, we present the existing prescriptive tree formulations from the literature and study their performance in the context of two simple synthetic examples. We introduce our proposed MIO formulations for learning optimal prescriptive tree in Section~\ref{section:proposed_formulations}. Finally, we summarize our computational experiments in Section~\ref{section:experiments}. All proofs and detailed computational results are provided in the Electronic Companion.

\section{Existing Prescriptive Tree Formulations}
\label{section:prev_mio}
This section presents the methods introduced by \cite{Kallus2017RecursiveData} and \cite{Bertsimas2019OptimalTrees}--from hereon, we refer to these as K-PT and B-PT, respectively. It further analyzes their advantages and drawbacks using two simple examples.

To describe these two approaches, we work with the notation introduced in~\cite{Kallus2017RecursiveData}, which views trees as producing a partition $\sets X = \bigcup_{n \in \sets{L}} \sets{X}_n$ such that $\sets{X}_n \cap \sets{X}_m = \emptyset$ whenever $n \neq m$, where $\sets L$ collects the set of leaf nodes and $\sets X_n \subseteq \sets X$ represents the set of covariate values that are associated with datapoints landing in leaf $n \in \sets L$. The approaches K-PT and B-PT rely on the following \emph{crucial} assumption about the structure of the prescriptive tree.

\begin{asu}[\cite{Kallus2017RecursiveData}]
    The partition $\{\sets X_n\}_{n \in \sets L}$ induced by the prescription tree is sufficiently fine so that
    $K \independent X \; | \; X \in \sets{X}_n \quad \forall n \in \sets{L}.$
    \label{asu:sufficient_partition}
\end{asu}

Assumption~\ref{asu:sufficient_partition} is automatically satisfied in randomized experiments (since $K \independent X$ unconditionally, by definition) or in conditionally randomized experiments provided the partition is ``sufficiently fine'' -- for example, a tree where each $\sets{X}_n$ is a singleton.\comment{. For example, it is satisfied in a conditionally randomized setting for a tree with~$|\sets X|$ leafs that branches on all levels of all features, so that each $\sets{X}_n$ is a singleton.}

\subsection{K-PT}
\label{section:kallus}

Given a partition satisfying Assumption~\ref{asu:sufficient_partition}, \cite{Kallus2017RecursiveData} proves that the \emph{observable} quantity $\mathbb E(Y| K=k, X \in \sets X_n)$ is an unbiased estimator of $\mathbb E(Y(k)|X \in \sets X_n)$. In other words, to estimate the potential outcomes under treatment~$k$ at leaf~$n$, it suffices to average outcomes for those that actually received treatment~$k$ at leaf~$n$ under~$\mathbb P$; this quantity can be estimated from the observed data. This statement is intuitive under Assumption~\ref{asu:sufficient_partition}, which implies that all individuals with covariates in the set $\sets X_n$ have the same chance of getting any given treatment, i.e., they are exchangeable within $\sets X_n$. It follows that, under Assumption~\ref{asu:sufficient_partition}, the outcomes of a policy $\pi$ that assigns treatment $\pi(\sets X_n) \in \sets K$ to all individuals that land in leaf~$n$ can be estimated as
\begin{equation}
\sum_{n \in \sets L} \mathbb P ( X \in \sets X_n) \cdot \mathbb E ( Y | K= \pi(\sets X_n), X \in \sets X_n).
\label{eq:kallus_estimator}
\end{equation}
Accordingly, this quantity can be estimated from data as
$$
\QkptI(\pi) := \sum_{n\in \sets{L}} \frac{\sum_{i \in \sets{I}}\mathbbm{1}[X_i \in \sets{X}_n]}{|\sets I|}\widehat{Y}_n \left(\pi(\sets{X}_n)\right),
$$
where $\widehat{Y}_n(k)$ is the sample average approximation of $\mathbb E(Y| K=k, X \in \sets X_n)$, defined through
\begin{equation}
\widehat{Y}_n(k):=\frac{1}{\sum_{i \in \sets{I}}\mathbbm{1}[K_{i} = k, X_{i} \in \sets{X}_n]}\sum_{i\in \sets{I}}\mathbbm{1}[K_i =k, X_i \in \sets{X}_n] Y_i.
\label{eq:avg_estimate}
\end{equation}

\subsubsection*{Analysis of K-PT.}

\cite{Kallus2017RecursiveData} proved that $\QkptI(\pi)$ is an unbiased estimator of $Q(\pi)$ when Assumption~\ref{asu:sufficient_partition} holds, which is the case in \emph{marginally} randomized settings for example. A salient advantage of K-PT is that, if Assumption 1 \emph{does} hold, it avoids predicting~$\mu$ and thus avoids the bias caused by potentially inaccurate estimates~$\hat \mu$. Unfortunately, though, in \emph{conditionally} randomized settings, decision trees will in general not satisfy Assumption~\ref{asu:sufficient_partition} and the tree obtained by maximizing $\QkptI(\pi)$ may be severely suboptimal for the problem of maximizing $Q(\pi)$. In fact, as we now demonstrate, K-PT may actively choose trees that violate Assumption~\ref{asu:sufficient_partition} even when the complete out-of-sample distribution is available \emph{and} the tree is deep enough to allow branching on all covariates that are predictive of outcomes and treatment assignments.

\begin{example}
Consider the following data distribution. Let the covariate vector $X = (X^1, X^2) \in \{0,1\}^2$, where $X^1$ and $X^2$ are independent and Bernoulli distributed with parameter 0.5. We interpret $X^1$ to denote the severity of a patient's condition, where $X^1 = 0$ (resp.\ $X^1=1$) indicates a healthy (resp.\ sick) patient. $X^2$ is a noise variable, unrelated to both the outcomes and the treatment assignments. There are two possible treatments, where $K=1$ (resp.\ $K=0$) indicates that the patient was treated (resp.\ not treated). Treatments are assigned according to a conditionally randomized experiment, with $\mathbb P(K=0|X^1=0) = 0.9$ and $\mathbb P(K=0|X^1=1) = 0.1$, i.e., 90\% of healthy patients do not receive the treatment, and 90\% of sick patients do. Table~\ref{table:expected} shows the expected values of the potential outcomes in dependence of the covariate values. \comment{To circumvent any finite sample issues, we assume that we have sufficient historical data at our disposal to be able to estimate all quantities in the K-PT estimator of~$Q(\pi)$ without errors; thus we can work with the true values in~\eqref{eq:kallus_estimator}, rather than with their empirical estimates provided in the definition of $\QkptI(\pi)$.}We assume that we have sufficient historical data at our disposal so that we can work with the true values in~\eqref{eq:kallus_estimator} rather than with their empirical estimates provided in the definition of~$\QkptI(\pi)$.

\begin{table}[H]
\centering
 \begin{tabular}{c | c c}
 \hline
 $X^1$ & $\mathbb{E}(Y(0))$ & $\mathbb{E}(Y(1))$ \\
 \hline\hline
 0 & 1 & 0.8 \\
 \hline
 1 & 0 & 0.2 \\
 \hline
\end{tabular}
\caption{Companion table to Example~\ref{ex:example1}.  Expectation of potential outcomes in dependence of the covariate values. A value $X^1 = 0$ (resp.\ $X^1=1$) of the covariate indicates a healthy (resp.\ sick) patient. Larger values of the outcomes are preferred.}
\label{table:expected}
\end{table}

From Table~\ref{table:expected}, it can be seen that the (unique) optimal policy treats sick patients only. In particular, this optimal policy can be modelled by a prescriptive tree of depth one (i.e., a tree with a single branching node and two leafs) that branches on~$X^1$, does not treat any of the healthy patients (left leaf) and treats all sick patients (right leaf). The expected outcome under this policy is given by
$$
\mathbb{E}(Y(0)|X^1=0)\mathbb{P}(X^1=0)+\mathbb{E}(Y(1)|X^1=1)\mathbb{P}(X^1=1)=1\times\frac{1}{2}+0.2\times \frac{1}{2}=0.6.
$$
One may similarly calculate the expected outcomes under a policy that branches on the noise feature~$X^2$. In this case, no matter which treatment is assigned at each leaf, the resulting policies would have a lower expected outcome of 0.5. These candidate policies and the corresponding expected outcomes under each policy are illustrated in the top row of Figure~\ref{fig:kallus_example}.

\begin{figure}[h]
	\includegraphics[width=\textwidth]{kallus_example}
	\centering
	\caption{Companion figure to Example~\ref{ex:example1}. Each row shows the estimates of performance of different prescriptive trees (branch on $X^1$--left-- or $X^2$--right--) under different estimators. In each case, only the best treatment assignments at each leaf (as measured by the corresponding estimator) are shown. In this example, the true optimal policy (top left) is to branch on $X^1$ and treat only those who are sick. Indeed, the expected outcome of the best tree branching on $X^1$ (top left) is $1.0 \cdot 0.5 + 0.2 \cdot 0.5 = 0.6$, which is larger than the expected outcome of the best $X^2$-branching tree (top right), which is $0.5 < 0.6$. However, the method from \cite{Kallus2017RecursiveData} (middle), which is discussed in Section \ref{section:kallus}, will prefer to branch on $X^2$ (noise covariate) and not treat any patient (middle right). Indeed, it estimates the outcomes of the best $X^1$- (resp.\ $X^2$-) branching tree as $0.6$ (resp.\ $0.9$). We note that Assumption~\ref{asu:sufficient_partition} is not satisfied at any of the leaves of the tree deemed as best by this estimator. In this case, the method of \cite{Bertsimas2019OptimalTrees}, see Section \ref{section:bertsimas}, is able to identify the truly optimal tree (bottom left); it correctly branches on~$X^1$ provided~$\theta$ is sufficiently small. Note that this method penalizes the variances corresponding to all treatments, but for simplicity the figure only depicts variances corresponding to the chosen treatment.}
	\label{fig:kallus_example}
\end{figure}

In practice, the quantities $\mathbb{E}(Y(k))$, $k\in \sets K$, cannot be easily estimated from data since $Y_i(k)$ is only observable for $i$ such that $K_i=k$ and treatments are not assigned at random. Thus, the procedure described in the previous paragraph to select an optimal prescriptive tree is not possible to follow. \comment{Now, l}Let us instead consider the tree of depth one that maximizes the objective function of K-PT given by~\eqref{eq:kallus_estimator}. Suppose that we branch on covariate $X^j$ at the sole branching node of the tree, $j\in \{1,2\}$. Then, at the leaf where $X^j = x^j \in \{0,1\}$, K-PT estimates the counterfactual outcomes $\mathbb E(Y(k)|X^j=x^j)$ as $\mathbb{E}(Y|K=k, \; X^j=x^j)$. However, these estimates can be incorrect. For example,
$$
\mathbb{E}(Y(0)|X^2=0) = \sum_{x^1 \in \{0,1\}} \mathbb{E}(Y(0)|X^1=x^1,X^2=0)\mathbb{P}(X^1=x^1|X^2=0)=1 \times \frac{1}{2} + 0\times \frac{1}{2} = 0.5,
$$
whereas
$$
\mathbb{E}(Y|K=0,X^2=0) = \sum_{x^1 \in \{0,1\}} \mathbb{E}(Y|K=0,X^1=x^1,X^2=0)\mathbb{P}(X^1=x^1|K=0,X^2=0
=0.9.
$$
The difference between these two quantities is intuitive. Individuals that were not treated ($K=0$) in the data are more likely to be healthy (with a large associated outcome). The K-PT estimate $\mathbb{E}(Y|K=0,X^2=0)$ \comment{(K-PT estimate) }thus \emph{overestimates} $\mathbb{E}(Y(0)|X^2=0)$. Similarly, since the vast majority of individuals that were treated ($K=1$) in the data are sick (with low outcomes), the K-PT estimate $\mathbb{E}(Y|K=0,X^2=1)=0.26$ is an \emph{underestimator} for $\mathbb{E}(Y(1)|X^2=0)=0.5$.

The middle row of Figure~\ref{fig:kallus_example} depicts the outcomes of the two candidate trees as estimated by K-PT. When branching on $X^1$ (middle left), Assumption~\ref{asu:sufficient_partition} is satisfied and the estimated outcome of~$0.6$ provided by K-PT is correct. However, K-PT prefers the alternative tree which branches on the noise variable and does not treat any of the patients (middle right); it wrongly estimated an outcome of~$0.9$ while the true expected outcome is only~$0.5$. We conclude that if not \emph{all feasible trees} in the problem satisfy Assumption~\ref{asu:sufficient_partition}, the method from~\cite{Kallus2017RecursiveData} may fail.

\label{ex:example1}
\end{example}

\subsection{B-PT}
\label{section:bertsimas}

\cite{Bertsimas2019OptimalTrees} also observe that K-PT may fail in certain conditions. They posit that failings can be attributed to poor estimates of the outcomes at the leaves. Thus, they augment the K-PT estimator with a regularization term, which penalizes trees that induce a high variance in the outcomes conditional on each treatment at the leaves, i.e.,
$$
\sum_{k \in \sets K} \sum_{n \in \sets L} \mathbb V{\rm{ar}} ( Y | X \in \sets X_n , K = k ).
$$
Specifically, they optimize the quantity
\begin{gather*}
\QbptI(\pi) = \theta \left(\QkptI(\pi)\right) - (1-\theta) \sum_{i \in \sets{I}} \left(Y_i - \widehat{Y}_{\sets{X}_{n(i)}}(K_i) \right)^2,
\end{gather*}
where $n(i) \in \sets L$ denotes the leaf datapoint $i$ belongs to, and $\theta\geq 0$ is a parameter to be tuned.

\subsubsection*{Analysis of B-PT.} Recall that \cite{Kallus2017RecursiveData}
estimates $\mathbb E(Y(k) | X \in \sets X_n)$ through $\mathbb E(Y | K=k, X \in \sets X_n)$ and shows that this estimator is unbiased under Assumption~\ref{asu:sufficient_partition}. By adding a regularization term that seeks to promote trees with low variance $\mathbb V{\rm{ar}}(Y| K=k, X \in \sets X_n)$, \cite{bertsimas2017personalized} improves the statistical properties of the K-PT estimator (lower variance). Therefore, if Assumption~\ref{asu:sufficient_partition} holds (so that $\mathbb{E}(Y|K=K_i, X\in \sets{X}_{n(i)})=\mathbb{E}(Y(K_i)| X\in \sets{X}_{n(i)})$), then B-PT can lead to improved performance over K-PT (for a suitable choice of $\theta$). However, B-PT does not directly address the issues discussed in Example~\ref{ex:example1}. Indeed, when Assumption~\ref{asu:sufficient_partition} does not hold, then $\mathbb{E}(Y|K=K_i, X\in \sets{X}_{n(i)})$ and $\mathbb{E}(Y(K_i)| X\in \sets{X}_{n(i)})$ can be drastically different from one another, and the policies chosen by B-PT may still be suboptimal.

In the specific case of Example~\ref{ex:example1}, \comment{it can be checked that if $\theta$ is sufficiently small, then} B-PT can in fact recover the best tree if $\theta$ is sufficiently small. Nonetheless, as we now show, a small modification of Example~\ref{ex:example1} results in a setting in which B-PT has \emph{additional} incentives to choose the worst tree, which branches on irrelevant features.

\setcounter{example}{0}
\begin{example}[Continued]
We now revisit Example~\ref{ex:example1} and use the method from~\cite{Bertsimas2019OptimalTrees} to identify an optimal tree. Figure~\ref{fig:kallus_example} (bottom) shows the estimated outcomes and variances associated with each candidate tree. For simplicity, we assume that the  potential outcomes in Table~\ref{table:expected} have zero variance. Since the tree that branches on the relevant feature~$X^1$ has smaller variance in this case, it will be chosen by B-PT provided that $\theta$ is sufficiently small. \hfill $\blacksquare$

\end{example}

\begin{example}
\label{ex:example2}
We now consider a variant of Example~\ref{ex:example1} where we add to all potential outcomes the quantity $(1-X^2)$, see Table~\ref{table:expected2}. Note that while $X^2$ now affects the value of the outcomes, it does not affect the value of the difference $\mathbb E(Y(1)-Y(0)|X)$ and is thus \emph{irrelevant} in deciding which treatment is preferable. Therefore, under suitable identifiability conditions, there exists \comment{in this problem }an optimal prescriptive tree of depth one. Figure~\ref{fig:bertsimas_example2} (top) shows the outcomes, for a depth one tree, under the two possible branching decisions if the counterfactuals are known--the best tree branches on the only predictive feature $X^1$. However, when using estimates $\mathbb{E}[Y|K=k,X=x]$ for the counterfactuals (Figure~\ref{fig:bertsimas_example2}, bottom), the tree that branches on the irrelevant feature $X^2$ has both a higher estimated outcome and a smaller variance. In this case B-PT has even more incentive to select a policy based on the irrelevant feature and both K-PT and B-PT will select a suboptimal tree. \hfill $\blacksquare$

\begin{table}[H]
\centering
 \begin{tabular}{c c | c c}
 \hline
 $X^1$ & $X^2$ & $\mathbb{E}(Y(0))$ & $\mathbb{E}(Y(1))$ \\ [0.5ex]
 \hline\hline
 0 & 0 & 2 & 1.8 \\
 \hline
 1 & 0 & 1 & 1.2 \\
 \hline
 0 & 1 & 1 & 0.8 \\
 \hline
 1 & 1 & 0 & 0.2 \\
 \hline
\end{tabular}
\caption{Companion table to Example~\ref{ex:example2}. Expectation of the potential outcomes in dependence of the covariate values. Similar to Example~\ref{ex:example1} (Continued), we assume that the potential outcomes have a variance of zero.}
\label{table:expected2}
\end{table}

\begin{figure}[h]
\includegraphics[width=\textwidth]{bertsimas_example2}
\centering
\caption{Companion figure to Example~\ref{ex:example2}. The layout and interpretation of the figure exactly parallel that in Figure~\ref{fig:kallus_example}. The optimal tree (top left) has the objective $0.5(1.5) + 0.5(0.7) = 1.1$, which is higher than the $X^2$-branching tree (top right), $0.5(1.5) + 0.5(0.5) = 1.0 < 1.1$. However, B-PT will incorrectly branch on $X^2$ no matter the regularization strength $\theta$, since the $X^1$-branching tree (bottom left) has a lower estimated value and also higher variance.}
\label{fig:bertsimas_example2}
\end{figure}
\end{example}

\comment{Continuing from the motivating example in Section \ref{section:example_kallus}, we show that variance reduction might similarly lead to unfavorable results. Consider, again, a decision tree that branches once on $X^1$, and we want to determine the ideal treatment $K \in \{0, 1\}$ for the left leaf (i.e. datapoints with $X^1 = 0$). For $K=0$, the average outcome is $\mathbb{E}(Y|X^1=0, K=0) = 1.5$ and the variance is
\begin{gather*}
\mathbb{E}[(Y-\mathbb{E}(Y))^2|X^1=0, K=0] = 0.5 \cdot (2-1.5)^2 + 0.5 \cdot (1-1.5)^2 = 0.25.
\end{gather*}

The regularized term then becomes $0.5(1.5) - 0.5(0.25) = 0.625$, where 0.5 represents the regularization strength $\theta$. \cite{Bertsimas2019OptimalTrees} cited that 0.5 seems to yield the best results, although this parameter needs to be tuned in practice. We compare this value with $K=1$, where $\mathbb{E}(Y|X^1=0, K=1) = 1.3$ and $\mathbb{E}[(Y-\mathbb{E}(Y))^2|X^1=0, K=0] = 0.25$. The regularization term becomes 0.525, so we assign the more favorable $K=0$. Continue to evaluate all possible permutations to find the optimal tree, which is summarized in Figure \ref{fig:bertsimas_example}. The algorithm still branches on $X^2$, failing to remedy the problem in Section \ref{section:example_kallus}.}

\section{Proposed Formulations}
\label{section:proposed_formulations}

In this section, we present our proposed MIO formulations for learning optimal prescriptive trees from observational data. Without loss of generality, we assume that the covariates are binary, i.e., $\sets X \subseteq \{0,1\}^F$ -- this assumption can easily be relaxed to cater for integer or categorical features, see Remark~1 in~\cite{aghaeistrong}. Importantly, and in sharp contrast with the existing literature on prescriptive trees, we \emph{do not} make Assumption~\ref{asu:sufficient_partition}. Instead, we directly optimize the estimators~\eqref{eq:ipw_obj}, \eqref{eq:dm_obj}, or \eqref{eq:dr_obj} from the causal inference literature. This enables us to design interpretable prescriptive trees that are guaranteed to be optimal out-of-sample as the number of historical samples grows.

As will become clear shortly, our proposed approaches give rise to \emph{weighted} classification trees. We thus adapt the MIO formulations of~\cite{aghaeistrong} for learning optimal classification trees from the predictive to the prescriptive setting. We extend our notation from the problem statement in Section \ref{section:problem_statement} to describe optimal prescriptive trees. The building block of our formulation is a perfect binary tree of depth~$d$, whose nodes are numbered $1$ through $(2^{d+1}-1)$ in the order in which they appear in a breadth-first search. We let $\sets{B} := \{1, \dots, 2^{d}-1\}$ denote the set of branching nodes and $\sets{T} := \{2^d, \dots, 2^{d+1}-1\}$ collect all terminal nodes (see Figure \ref{fig:flow_graph}, left). From the perfect binary tree, we build a flow graph as follows. We connect a source node~$s$ to the root node, and connect all nodes other than $s$ to the sink node $t$ (see Figure \ref{fig:flow_graph}, right). All links in the graph are directed from source to sink and have capacity 1.
\begin{figure}[h]
\includegraphics[width=0.7\textwidth]{flow_graph}
\centering
\caption{A prescriptive tree with depth 2 (left) and its associated flow graph (right).}
\label{fig:flow_graph}
\end{figure}

\subsection{Formulation Based on Inverse Propensity Weighting}
\label{section:ipw_formulation}

Equipped with the flow graph associated with a prescriptive tree, we now formulate an MIO problem to optimize~\eqref{eq:ipw_obj} over all trees of maximum depth $d$. For every branching node $n\in \sets B$ and feature $f\in \sets F$, we let the binary variable~$b_{nf}$ indicate if feature~$f$ is selected for branching at node~$n$. Accordingly, for every node~$n \in \sets{T} \cup \sets{B}$, we let the binary variable~$p_n$ indicate if node~$n$ is a treatment node, in which case a treatment must be assigned to all datapoints that land at that node and no further branching is allowed. For~$k\in \sets K$, we let~$w_{nk}\in \{0,1\}$ equal~1 if and only if (iff) treatment~$k$ is selected at node~$n$. Finally, the ``flow variables'' $z$ capture the flow of datapoints from the source to the sink: $z^{a(n),n}_{i} = 1$ iff the following two conditions are satisfied: \textit{(i)} datapoint~$i$  flows to node~$n$ from its ancestor~$a(n)$; and \textit{(ii)} datapoint~$i$ reaches a treatment node where treatment~$K_i$ was assigned (i.e., $w_{nK_i} = 1$ for some~$n\in \sets B\cup \sets T$). \delag{If datapoint $i$ were to flow to node $n$ from its ancestor $a(n)$, then $z^{a(n),n}_{i} = 1$ and $0$ otherwise. Critically, since flow must be preserved, datapoint $i$ flows to sink $t$ if it is directed to a terminal node that assigns the same treatment as was observed in the data (i.e., $w_{nK_i} = 1$). Otherwise, datapoint $i$ will have a flow value of 0.}

\delag{From the training data, unless the true propensity scores are known, determine a-priori the scores $\mu(K_i, X_i)$, which is learned using a suitable machine learning model.}

The formulation is as follows:
\begin{subequations}
\begin{align}
\text{maximize} \;\; & \displaystyle \sum_{i \in \mathcal I} \sum_{n \in \sets B \cup \sets T } (z^{n,t}_{i}) \frac{Y_i}{\hat\mu(K_i, X_i)} = \sum_{i \in \sets I}z^{s,1}_{i}\frac{Y_i}{\hat\mu(K_i, X_i)}
\label{eq:flow_reg_obj}\\
\text{subject to}\;\;
& \displaystyle \sum_{f \in \sets F}b_{nf} + p_n + \sum_{m \in \sets A(n)}p_m = 1   &\hspace{-5cm}  \forall n \in \sets B \label{eq:flow_reg_branch_or_predict}\\
&  p_n+\sum_{m \in \sets A(n)}p_{m} =1   &  \forall n \in \sets T \label{eq:flow_reg_terminal_leaf}\\
& \displaystyle z^{a(n),n}_{i} =  z^{n,\ell(n)}_{i} + z^{n,r(n)}_{i} + z^{n,t}_{i}  &\hspace{-5cm}  \forall n \in \sets B, i \in \mathcal I \label{eq:flow_reg_conservation_internal}\\
&  \displaystyle z^{a(n),n}_{i} = z^{n,t}_{i} &\hspace{-5cm}   \forall i \in \mathcal I,n \in \sets T \label{eq:flow_reg_conservation_terminal}\\
& \displaystyle z^{s,1}_{i} \leq 1 &\hspace{-5cm} \forall i \in \mathcal I\label{eq:flow_reg_source}\\
&  \displaystyle z^{n,\ell(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i\leq0}b_{nf} &\hspace{-5cm} \forall n \in \sets B, \; i \in \mathcal I \label{eq:flow_reg_branch_left}\\
&  \displaystyle z^{n,r(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i=1}b_{nf}  &\hspace{-5cm} \forall n \in \sets B, \; i \in \mathcal I \label{eq:flow_reg_branch_right}\\
&  \displaystyle z^{n,t}_{i} \leq  w_{nK_i} &\hspace{-5cm} \forall i \in \mathcal I, \; n \in \sets B \cup \sets T \label{eq:flow_reg_sink}\\
&  \displaystyle \sum_{k \in \sets K}w_{nk} = p_n  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T \label{eq:flow_reg_leaf_prediction}\\
&  \displaystyle w_{nk} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B \cup \sets T, \; k \in \sets K \\
&  \displaystyle b_{nf} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B, \; f \in \sets F \\
&  \displaystyle p_{n} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B \cup \sets T \\
&  \displaystyle z^{a(n),n}_{i}, z^{n,t}_{i}\in \{0,1\}  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T, \; i \in \sets I,
\end{align}
\label{eq:flow_reg}
\end{subequations}
where $\sets A(n)$ denotes the set of all ancestors of node $n\in \sets B \cup \sets T$, $r(n)$ (resp.\ $\ell(n)$) is the right (resp.\ left) descendant of $n$, and $\hat \mu$ is as defined in Section~\ref{section:ipw_intro}. The objective \eqref{eq:flow_reg_obj} maximizes the IPW estimator of the counterfactual performance of the prescriptive tree defined through variables $b$, $w$, and $p$, as introduced in \eqref{eq:ipw_obj}. Constraints~\eqref{eq:flow_reg_branch_or_predict} ensure that we branch at all branching nodes where no treatment is assigned, at the node itself or at any of its ancestors. Similarly, constraints~\eqref{eq:flow_reg_terminal_leaf} impose that a treatment must be assigned at all terminal nodes, unless a treatment has already been assigned to one of its ancestors. Constraints~\eqref{eq:flow_reg_conservation_internal} and~\eqref{eq:flow_reg_conservation_terminal} are flow conservation constraints, whereby any datapoint that flows into a node $n$ must either exit to its right or left descendant, or flow directly to the sink. Constraints~\eqref{eq:flow_reg_source} ensure that a flow of at most one can enter the source node for each datapoint. Constraints~\eqref{eq:flow_reg_branch_left} and~\eqref{eq:flow_reg_branch_right} guarantee that datapoints flow according to the branching decisions specified by variables $b$. Constraints~\eqref{eq:flow_reg_sink} state that only datapoints whose assigned treatments match the one observed in the data can flow to the sink. Finally, constraints~\eqref{eq:flow_reg_leaf_prediction} combined with integrality of the $w_{nk}$ variables ensure that if a treatment is assigned at a node, exactly one of the available treatments must be assigned. Note that we can relax integrality on the~$w$ and~$z$ variables and still reach an integral solution.

\subsection{Formulation Based on Doubly Robust Estimation}
\label{section:robust_formulation}

We now formulate an MIO problem to optimize~\eqref{eq:dr_obj} over all trees of maximum depth~$d$. To this end, we adapt the flow graph associated with the decision tree so that both correctly and incorrectly classified datapoints are tracked and also redefine the variables $z$. Specifically, we create $|\sets{K}|$ sink nodes (denoted by $t_k, \forall k \in \sets{K}$) and connect each sink node to all nodes other than the source node~$s$. As before, all arcs are directed from source to sink, see Figure~\ref{fig:flow_graph_multisink}. In this formulation, each datapoint will be directed towards the sink corresponding to its assigned treatment. Therefore,
we redefine variables $z$ as $z^{a(n),n}_{i} = 1$ iff datapoint~$i$ flows to node~$n$ from its ancestor $a(n)$ (regardless of the treatment assigned to the datapoint in the data).

\begin{figure}[h]
\includegraphics[scale=0.35]{images/flow_graph_multisink}
\centering
\caption{A modified flow graph with multiple sink nodes; in this case, there are 2 sinks to reflect 2 treatments.}
\label{fig:flow_graph_multisink}
\end{figure}

The MIO formulation for DR remains largely unchanged from formulation~\eqref{eq:flow_reg}:
\begin{subequations}
\begin{align}
\text{maximize} \;\; & \displaystyle \sum_{i \in \mathcal I} \sum_{k \in \sets K } \sum_{n \in \sets B \cup \sets T}z^{n,t_{k}}_{i} \hat{\nu}_k(X_i) + \displaystyle \sum_{i \in \mathcal I} \sum_{n \in \sets B \cup \sets T} (Y_i -\hat{\nu}_{K_{i}}(X_i)) \frac{{z^{n,t_{K_{i}}}_{i}} }{\hat\mu(K_i, X_i)} \label{eq:robust_obj}\\
\text{subject to}\;\;
& \displaystyle \sum_{f \in \sets F}b_{nf} + p_n + \sum_{m \in \sets A(n)}p_{m} = 1   &\hspace{-5cm}  \forall n \in \sets B \label{eq:robust_branch_or_predict}\\
&  p_n+\sum_{m \in \sets A(n)}p_{m} =1   &  \forall n \in \sets T \label{eq:robust_all_terminal_leaf}\\
& \displaystyle z^{a(n),n}_{i} =  z^{n,\ell(n)}_{i} + z^{n,r(n)}_{i} + \sum_{k \in \sets K}z^{n,t_{k}}_{i}   &\hspace{-5cm}  \forall n \in \sets B,\; i \in \mathcal I \label{eq:robust_conservation_internal}\\
&  \displaystyle z^{a(n),n}_{i} = \sum_{k \in \sets K}z^{n,t_{k}}_{i}  &\hspace{-5cm}   \forall i \in \mathcal I, \; n \in \sets T \label{eq:robust_conservation_terminal}\\
& \displaystyle z^{s,1}_{i} = 1 &\hspace{-5cm} \forall i \in \mathcal I\label{eq:robust_source}\\
&  \displaystyle z^{n,\ell(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i=0}b_{nf} &\hspace{-5cm} \forall n \in \sets B,\; i \in \mathcal I \label{eq:robust_branch_left}\\
&  \displaystyle z^{n,r(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i=1}b_{nf}  &\hspace{-5cm} \forall n \in \sets B, \; i \in \mathcal I \label{eq:robust_branch_right}\\
&  \displaystyle z^{n,t_{k}}_i \leq  w_{nk} &\hspace{-5cm} \forall i \in \mathcal I,\; n \in \sets T,k \in \sets K \label{eq:robust_sink}\\
&  \displaystyle \sum_{k \in \sets K}w_{nk} = p_n  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T \label{eq:robust_leaf_prediction}\\
&  \displaystyle w_{nk} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets T, \; k \in \sets K \\
&  \displaystyle b_{nf} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B, \; f \in \sets F \\
&  \displaystyle p_{n} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B \cup \sets T \\
&  \displaystyle z^{a(n),n}_{i}, z^{n,t_{k}}_{i} \in \{0,1\}  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T, \; i \in \sets I, \; k \in \sets K.
\end{align}
\label{eq:robust}
\end{subequations}
Here, $\hat \mu$ and $\hat \nu_k$ are as defined in Sections~\ref{section:ipw_intro} and~\ref{section:dm}, respectively. The main difference relative to formulation~\eqref{eq:flow_reg} is the objective \eqref{eq:robust_obj}, which now maximizes the doubly robust objective \eqref{eq:dr_obj}. The other difference is in the flow conservation constraints \eqref{eq:robust_conservation_internal} and \eqref{eq:robust_conservation_terminal}, which restrict points to flow to at most one sink node. The last difference is in constraint~\eqref{eq:robust_sink}, which states that datapoints can only flow to the sink associated with their assigned treatment. Similar to formulation~\eqref{eq:flow_reg}, we can relax integrality on the $w$ and $z$ variables and still reach an integral solution.

\subsection{Formulation Based on Direct Method Estimator}
\label{section:dm_formulation}

Formulation~\eqref{eq:robust} can easily be adapted to optimize the direct method objective~\eqref{eq:dm_obj} by simply dropping the second term in the objective~\eqref{eq:robust_obj}.

\subsection{Analysis of IPW, DM, and DR}
\label{section:analysis}

It can be shown that if $\mu$ (resp.\ either $\mu$ or $\nu$) are known, a depth one tree trained according to formulations~\eqref{eq:flow_reg} (resp.\ \eqref{eq:robust}) in Examples~\ref{ex:example1} and~\ref{ex:example2} will be optimal out-of-sample, provided the number of samples in the data is sufficiently large. In this section, we go one step further and demonstrate that, under mild conditions, the tree returned by our optimal MIO formulations~\eqref{eq:flow_reg}, \eqref{eq:robust}, and~\eqref{eq:robust} modified as in Section~\ref{section:dm_formulation} will converge almost surely to an optimal tree (in terms of out-of-sample treatment assignment).

We introduce some notation to help formalize our claim. We let $\Pi_d$ denote the set of tree-based policies of maximum depth~$d$. For any tree based policy $\pi \in \Pi_d$, define
$$\Qipw (\pi) := \mathbb E\left[ \frac{ \mathbb I (K= \pi(X)) Y }{ \mu (K,X) } \right] \quad \text{and} \quad \QipwIN (\pi) := \frac{1}{I} \sum_{i=1}^{I} \frac{ \mathbb I (K_i= \pi(X_i)) Y_i }{ \muhatN (K_i,X_i) },$$
where $\muhatN (K,X)$ is an estimator of $\mu (K,X)$ trained over a sample of cardinality $N$ independent from~$\mathcal D$.
We denote by~$v^\star$ and~$\sets S^\star$, the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} Q(\pi)$, respectively.
Similarly, we denote by~$\vhatINipw$ and~$\shatINipw$, the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} \QipwIN(\pi)$, respectively. Finally, we use the abbreviation ``w.p.1" for ``with probability one".
\begin{proposition}
\label{prop:exact_IPW}
Suppose that: conditional exchangeability holds, i.e., $Y(k) \independent K | X$, for all values of~$k \in \sets K$; the potential outcomes $Y(k)$ are bounded, for all values of~$k \in \sets K$; and, both $\mu(k,x) >0$ and $\muhatN(k,x)>0$ for all $x \in \sets X$, $k\in \sets K$, and $N$ sufficiently large. If $\muhatN$ converges almost surely to~$\mu$, then $\vhatINipw \rightarrow v^\star$ w.p.1 and $\shatINipw \subseteq \sets S^\star$ w.p.1 for $I$ and $N$ large enough.
\end{proposition}

Similar to the IPW case, we now show that under mild conditions, the MIO formulation~\eqref{eq:robust} for the DR estimator returns an optimal out-of-sample tree. We introduce additional notation to help formalize our claim.
For any tree based policy $\pi \in \Pi_d$, define
$$\Qdr (\pi) \; := \; \mathbb E\left[ \hat{\nu}_{\pi(X)}(X) + (Y-\hat{\nu}_{\pi(X)}(X))\frac{ \mathbb I (K= \pi(X)) }{ \hat{\mu} (K,X) } \right]$$
$$\text{and} \quad \QdrIN (\pi)\; := \; \frac{1}{I} \sum_{i=1}^{I} \left(\nuhatN_{\pi(X_i)}(X_i) + (Y-\nuhatN_{\pi(X_i)}(X_i))\frac{ \mathbb I (K_i= \pi(X_i)) }{ \muhatN(K_i,X_i) }\right),$$
where $\nuhatN_K (\cdot)$ is an estimator of~$\nu_K(\cdot)$, trained over a sample of cardinality~$N$ independent from~$\mathcal D$. Let $\hat{\mu}(K,X)$ (resp. $\hat{\nu}_K(X)$) be the limit of $\muhatN (K,X)$ (resp. $\nuhatN_K (X)$) as $N$ goes to infinity.
We denote by~$\vhatINdr$ and~$\shatINdr$, the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} \QdrIN(\pi)$, respectively.

\begin{proposition}
\label{prop:exact_DR}
Suppose that: conditional exchangeability holds, i.e., $Y(k) \independent K | X$, for all values of~$k \in \sets K$; the potential outcomes $Y(k)$ are bounded, for all values of~$k \in \sets K$;  both $\mu(k,x) >0$ and $\muhatN(k,x)>0$ for all $x \in \sets X$, $k\in \sets K$ and $N$ sufficiently large; and both $\hat \mu(k,x)$ and $\hat \nu_k (x)$ exist and are bounded for all $x \in \sets X$, $k\in \sets K$. If either $\muhatN$ converges almost surely to $\mu$ or $\nuhatN_K$ converges almost surely to $\nu_k$, then $\vhatINdr \rightarrow v^\star$ w.p.1 and $\shatINdr \subseteq \sets S^\star$ w.p.1 for $I$ and $N$ large enough.
\end{proposition}

At the end, we show that under mild conditions, the MIO formulation~\eqref{eq:robust} modified as in Section~\ref{section:dm_formulation} for the direct estimator returns an optimal out-of-sample tree. We introduce additional notation to help formalize our claim.
For any tree based policy $\pi \in \Pi_d$, define
$$\Qdm (\pi) \; := \; \mathbb E\left[ \hat{\nu}_{\pi(X)}(X) \right] \quad \text{and} \quad  \QdmIN (\pi)\; := \; \frac{1}{I} \sum_{i=1}^{I} \left(\nuhatN_{\pi(X_i)}(X_i)\right).$$
We denote by~$\vhatINdm$ and~$\shatINdm$ the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} \QdmIN(\pi)$, respectively.

\begin{proposition}
\label{prop:exact_DM}
Suppose that: conditional exchangeability holds, i.e., $Y(k) \independent K | X$, for all values of~$k \in \sets K$; and $\hat \nu_k (x)$ exists, for all $k \in \sets K$. If $\nuhatN_K$ converges almost surely to $\nu_k$, then $\vhatINdm \rightarrow v^\star$ w.p.1 and $\shatINdm \subseteq \sets S^\star$ w.p.1 for $I$ and $N$ large enough.
\end{proposition}

\subsection{Extensions} Our formulations in Sections \ref{section:ipw_formulation}, \ref{section:robust_formulation}, and \ref{section:dm_formulation} have strong modeling power; they can be used to learn trees that interpretable and/or fair, to design randomized policies, and to impose budget constraints relevant for under-resourced settings. We discuss these in the present section. We introduce all these constraints on formulation~\eqref{eq:robust} which tracks which datapoints are assigned to which treatments. We note that formulation~\eqref{eq:robust} can be used to optimize all three objectives (IPW, DR, or DM).

\subsubsection*{Intepretability.} In settings where learning interpretable policies is desired, see e.g., \cite{rudin2019stop} and~\cite{azizi2018designing}, one may limit the number of branching nodes in the tree to $M$ by augmenting formulation~\eqref{eq:robust} with the constraint $$\sum_{n\in \sets B\cup \sets T}p_n \; \leq \; M.$$

\subsubsection*{Budget Constraints.} In low resource settings such as organ allocation \citep{bertsimas2013fairness, zenios2000dynamic}, housing allocation \citep{azizi2018designing}, and security \citep{xu2018strategic}, many treatments/interventions often have limited supply. Tree based policies that satisfy such capacity constraints can be learned by augmenting formulation~\eqref{eq:robust} with constraints
$$ \displaystyle \sum_{n \in \sets B \cup \sets T} \sum_{i \in \sets I} z^{n,t_{k}}_i \; \leq \; |I|C_k \quad\forall k \in \sets K, $$
where $C_k$ denotes the percentage of instances that can be assigned treatment $k$.

\subsubsection*{Fairness in Treatment Assignment.} Our formulations can be used to learn trees where treatments are assigned fairly across groups. Concretely, we will say that a policy satisfies treatment assignment parity if the probability of assigning a particular treatment is equal across protected groups. We let $\sets P$ collect all levels of a (collection of) protected covariates (e.g., levels of race and LGBTQ status). Accordingly, we let $P_i \in \sets P$ represent the value of the protected feature(s) of datapoint~$i$. It is worth noting that the features in $\sets P$ are usually not included in covariates $X$ to make sure that the policy does not make decisions based on protected features. To ensure that the learned policy satisfies treatment assignment parity up to a bias $\delta$, we may augment~\eqref{eq:robust} with the constraint
$$\left| \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p} z^{n, t_k}_i}{|\{i \in \sets I : P_i = p\}|} - \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p'} z^{n, t_k}_i}{|\{i \in \sets I : P_i = p'\}|} \right| \; \leq \; \delta \qquad \forall p, p' \in \sets P: p \neq p', \; k \in \sets K.$$

In a similar fashion, we will say that a policy satisfies \textit{conditional} treatment assignment parity if the probability of assigning a particular treatment is equal across protected groups, conditional on some legitimate feature(s) that affect the outcome. For instance, in prescribing housing resources to people experiencing homelessness, one may require that all individuals with the same vulnerability have the same likelihood of receiving a certain treatment.  We let $\sets A$ be the set of features indicative of risk, and let $A_i \in \sets A$ be the value of the risk feature of datapoint $i$. Conditional treatment assignment parity is satisfied up to a bias $\delta$ and for all features $a \in \sets A$ by adding the following constraint to \eqref{eq:robust},
$$\left| \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p, A_i = a} z^{n, t_k}_i}{|\{i \in \sets I : P_i = p \cap A_i = a\}|} - \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p', A_i = a}  z^{n, t_k}_i}{|\{i \in \sets I : P_i = p' \cap A_i = a \}|} \right| \; \leq \; \delta  \qquad \forall p, p' \in \sets P: p \neq p', \; k \in \sets K, \; a \in \sets A$$

\subsubsection*{Fairness in Treatment Outcomes}. Our formulations can also be used to ensure fairness in expected outcomes across groups. In general, this requires the average expected outcomes of a protected group $p \in \sets P$ to be above some threshold~$\gamma_p$ $\in \mathbb R$:
$$\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p} \left[(Y_i -\hat{\nu}_{K_{i}}(X_i)) \frac{{z^{n,t_{K_{i}}}_{i}}}{\hat\mu(K_i, X_i)} + \sum_{k \in \sets K} z^{n, t_k}_i \hat{\nu}_k(X_i)\right] \; \geq \; \gamma_p \qquad \forall p \in \sets P.$$
For example, \cite{bertsimas2013fairness} set
$\gamma_p := \sum_{i \in \sets I: P_i = p} Y_i$, i.e., the expected outcomes of all protected groups under the learned policy should be greater than or equal to what was observed in the data. Alternatively, the constraint can be used to impose max-min fairness to protect the outcomes of the groups that are worst off, see \cite{rawls1974some}. In this case, one sets $\gamma_p := \gamma$ for all $p \in \sets P$, where $\gamma$ is the largest value for which the MIO problem remains feasible.

\subsubsection*{Randomized Treatment Assignment Policies.} The addition of fairness and budget constraints may make problem~\eqref{eq:robust} infeasible. In such cases, it may be desirable to design \emph{randomized} policies where datapoints that land at the same leaf are assigned each treatment with a certain probability (rather than all getting the same treatment). This can be achieved by relaxing integrality on the variables~$w$ and~$z$ in formulation~\eqref{eq:robust}. The variable $w_{nk}$ can then be interpreted as the probability of assigning treatment~$k$ to datapoints that fall on node~$n$. Similarly, $z^{n, t_k}_i$ can be viewed as the probability of assigning treatment~$k$ to datapoint~$i$ if it lands at node~$n$ (and 0 otherwise).

\section{Experiments}
\label{section:experiments}

We now evaluate the empirical performance of our proposed formulations in Section~\ref{section:proposed_formulations} in two problem settings: a synthetic setting from the literature and a real setting based on a warfarin dosing dataset. In both cases, we take the viewpoint that data comes from an observational study\comment{where the treatment assignment policy is not known by the researcher, but that the study}, which can be viewed as a conditionally randomized experiment. We benchmark against the approaches of~\cite{Kallus2017RecursiveData} and \cite{Bertsimas2019OptimalTrees} for learning prescriptive trees, see Section~\ref{section:prev_mio}. As before, we refer to these as K-PT and B-PT, respectively. Since \cite{Bertsimas2019OptimalTrees} do not propose an MIO formulation, we adapt the method from \cite{Kallus2017RecursiveData} to implement B-PT. For completeness, we provide the MIO formulations that we have implemented for these two approaches in Appendices~\ref{section:kallus_formulation} and~\ref{section:bertsimas_formulation}, respectively.
We investigate the performance of all approaches as the probability of correct treatment assignment is varied in the historical data

\subsection{Dataset Description}
\label{section:synthetic1}

\subsubsection*{Synthetic Data.} For our experiments on synthetic data, we adapt the data generation process from \cite{Athey2016RecursiveEffects}. In this problem, there are two treatment possibilities indexed in the set $\sets K = \{0, 1\}$. The covariate vector has two independent and identically distributed features, $X=(X^1, X^2)$, where each $X^j \sim \mathcal{N}(0, 1)$. The potential outcome of datapoint~$i$ with covariates~$X_i$ under treatment~$k$ is
\begin{equation}\label{eq:synthetic}
Y_i(k) = \phi(X_i) + \frac{1}{2}(2k-1) \cdot \kappa(X_i) + \epsilon_i,
\end{equation}
where $\phi(x) := \frac{1}{2}x^1 + x^2$ models the mean effect, $\kappa(x):=\frac{1}{2}x^1$ models the treatment effect, and $\epsilon_i \sim \sets N(0,0.1)$ is noise added to the outcome that is independent of the covariates.

To study different settings of observational experiments, we vary, in the data, the probability~$p$ of assigning the treatment that is best \emph{in expectation} for each unit\comment{; we denote this probability by~$p$}. We let $p \in \{ 0.1, 0.25, 0.5 , 0.75, 0.9\}$ -- since there are 2 treatments, 0.5 corresponds to the (marginally) randomized setting, while the other settings correspond to conditionally randomized experiments. For each~$p$, we randomly generate 5 training and test sets, each with 500 and 10,000 datapoints, respectively.

\subsubsection*{Warfarin Dosing.} For our study on real data, we employ a dataset for personalized warfarin dosing for which counterfactuals are available. Warfarin is the most widely used oral anticoagulant agent, but despite its prevalence, determining one's optimal warfarin dosage is difficult because it can vary widely depending on demographic variables, clinical factors, and genetics \citep{international2009estimation}. The publicly available dataset that we \comment{employ}use was collected by the International Warfarin Pharmacogenetics Consortium and published at the Pharmacogenetics and Pharmacogenomics Knowledge Base, see \cite{international2009estimation}. The advantage of using this dataset is that we can model a patient's true outcomes when given varying doses of warfarin, which allows us to evaluate the performance of arbitrary counterfactual policies. \cite{international2009estimation} published a learned affine function $f(x) = \beta x + c$ that determines the optimal warfarin dose based on a patient's age, weight, race, VKORC1 genotype, CYP2c9 genotype, and whether or not the patient is currently taking amiodarone or an enzyme reducer, see equation~\eqref{warfarin_function} in Section~\ref{sec:appendix_warfarin}. We calculate a patient $i$'s optimal dosage using $f(X_i) + \epsilon$, where $\epsilon \sim \mathcal{N}(0, 0.02).$ This dosage is then discretized into three groups (i.e., $|\mathcal K|=3$) using the same convention as \cite{international2009estimation}: $K_{i}^{\text{opt}} = 0$ ($\leq$ 3 mg/day), 1 (between 3 and 7 mg/day), and 2 ($\geq$ 7 mg/day). The observed outcome is $Y_{i}(K_i) = 1$, if $K_i = K_{i}^{\text{opt}}$; and $=0$ otherwise.

To study different settings of observational experiments, we consider three different treatment assignment mechanisms, as follows. For the marginally randomized setting, we assign each treatment with probability $1/3$, and repeat to produce 5 datasets. To simulate data that is based on a more informed policy, we assign treatments based on modified versions of $f$ obtained by perturbing its coefficients so that patients closer to the boundary of each treatment bucket are at higher risk of receiving an incorrect treatment. Specifically, for each coefficient~$a$ of $f$, we generate $a'\sim U(a-a\cdot r, a+a \cdot r)$, where $r$ denotes the range of sampling. We assign patient $i$ the treatment corresponding to this modified function, discretized into three groups as discussed previously. In testing, we found that $r \in [0.05, 0.12]$ was the ideal range to simulate a reasonable treatment policy where the probability of correct treatment assignment is between 0.6 and 0.9. We fix 2 values of~$r$, 0.06 and 0.11, and for each value, we generate 5 sets of randomly sampled coefficients (from hereon, we refer to this variability as a ``realization''). At this point we have~3 experiment designs--each with~5 realizations--making up~15 datasets. Each dataset is then split randomly~5 times into training and test sets of~3,000 and~1,387 instances, respectively, yielding a total of~75 train-test pairs.

\subsection{Experimental Setup}
\delag{Throughout our experiments, we assume that we are in the \emph{observational setting}, where the historical logging policy is unknown, see Section~\ref{section:randomization}.} Since the proposed methods\delag{our proposed formulations} rely on estimating the propensity scores and/or counterfactuals, we discuss the estimation models we use for both datasets. We also discuss the preprocessing steps required.

\subsubsection*{Synthetic Data.} The synthetic data has real covariates but our methods only allow for binary features to split on. We thus discretize each covariate feature into ten buckets corresponding to quantiles from a normal distribution. Since the outcome distribution is known by construction (it is linear), we learn the potential outcomes using both linear regression (LR) and lasso regression with $\alpha = 0.08$ (Lasso), knowing that the latter will be a slightly less accurate predictor. We also use true propensity scores--they are known by construction--as well as two methods to estimate propensity scores: logistic regression (Log) and decision trees (DT). Therefore, there will be six doubly robust methods, each corresponding to a pair of models for IPW and DM. These methods are then compared to K-PT and B-PT. All approaches are allotted a solve time of 1 hour.

\subsubsection*{Warfarin Dosing.} Most of the features are already binary with the exception of age, height, and weight, so these features were split into 5 buckets, where each bucket contains approximately the same number of datapoints. A variety of other preprocessing and data imputation steps were done based on recommendations from \cite{international2009estimation}. A patient's possible outcomes are binary, and in selecting DM's model we found that using a random forest (RF) gave the best validation set accuracy. The propensity scores are learned using decision trees (DT), also shown to have the best validation set performance. These methods are then compared to K-PT and B-PT. All approaches are allotted a solve time of 1 hour.

\subsubsection*{Policy Evaluation.} We evaluate a learned policy~$\pi$ in terms of its out of sample probability of assigning the optimal treatment to a patient (OOSP). In the next section, we compare all methods using this metric.

\subsection{Experimental Results}
In this section, we focus on analyzing the results for trees of depth 1 and 2, i.e., with 2 and 4 leaves, respectively. We decided to focus on these results because K-PT did not scale to deeper trees and resulted in large optimality gaps, making for an unfair comparison. Indeed, from Figure~\ref{fig:performance}, we observe that K-PT had an optimality gap of more than $100\%$ for most instances of trees of depth 3 and 4. Refer to Section~\ref{section:experiments_detail} for a detailed reporting of computational times.

\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/performance_v2}
\centering
\caption{Comparison of computational performance for the synthetic and real data over all methods except for B-PT. B-PT is excluded because our implementation is an MIO adaptation of the local search method described in \cite{Bertsimas2019OptimalTrees}, and so showing either of B-PT (which is slow) or the method of \cite{Bertsimas2019OptimalTrees} (which is fast but does not provide optimality gaps) on the graph would not yield a fair comparison. We only include the models associated with the best statistical performance. K-PT solved 157 instances to optimality within 3,600 seconds while DR solved the same number of instances in 303 seconds, for a 11.9$\times$ speed-up.}
\label{fig:performance}
\end{figure}

\label{section:experimental_results}
\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/athey_v1}
\centering
\caption{Results on the synthetic dataset. The left graph compares the out-of-sample probability of correct treatment assignment of the best models for IPW, DM, and DR with the methods described in \cite{Kallus2017RecursiveData} (K-PT) and \cite{Bertsimas2019OptimalTrees} (B-PT). The right graph shows the same, but only includes the DR approaches where at least one of the combined methods is correct. Both figures are averages from trees of $d \in \{1, 2\}$ with 2 and 4 leaf nodes, respectively.}
\label{fig:athey_v1}
\end{figure}

\subsubsection*{Synthetic Data.}
Our results on the synthetic data are summarized on Figure~\ref{fig:athey_v1} (more details can be found in Section~\ref{section:experiments_detail}). The figure shows the average OOSP for each method in dependence of the probability of correct treatment assignment, $p$. From the figure, it can be seen that the DR methods, with and OOSP of around 75\%, consistently outperform both K-PT and B-PT whose best performance is 70\%. Even when one of the two estimators in DR is inaccurate (i.e., includes Lasso or Log), DR favors the other (better) estimator as expected. Notably, K-PT and B-PT drop to around 50\% OOSP (i.e., similar to random treatment assignment) when the probability of correct treatment assignment in the data is 0.75 and 0.9. Their poor performance in these settings is particularly relevant since most observational data is based on an informed policy that presumably treats most patients correctly. The likely explanation for this drop in performance is that a tree of moderate depth is unlikely to partition the dataset sufficiently fine for Assumption~\ref{asu:sufficient_partition} to hold.

\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/warfarin_box}
\centering
\caption{Out of sample probability of correct treatment assignment over all 5 realizations on the warfarin dataset. The left graph shows the distribution for the randomized experiments. The middle (resp.\ right) figure shows the distribution of results for $r = 0.06$ (resp.\ $0.11$). All figures are averages from trees with $d \in \{1, 2\}$ with 2 and 4 leaf nodes, respectively.}
\label{fig:warfarin_box}
\end{figure}

\subsubsection*{Warfarin Dosing.}
Our results on the warfarin data are summarized on Figures~\ref{fig:warfarin_box} and~\ref{fig:warfarin_bar} (more details can be found in Section~\ref{section:experiments_detail}). Figure~\ref{fig:warfarin_box} plots the distribution of OOSP over the 5 realizations. The left figure shows that all methods have similar performance in the randomized setting, except for B-PT which suffered from optimality gaps. From hereon, we discuss results only from the conditionally randomized experiments ($r \in \{0.06, 0.11\}$), which are more pertinent in our context. In both these settings (middle and right subfigures), both K-PT and B-PT have relatively low variance, but they perform consistently below 30\% OOSP (i.e., worse than random assignment). IPW seems to be the strongest estimator consistently at around 60\%. Both DM and DR have high variances that span from outperforming IPW in some realizations to performing below 20\% in others--we now discuss why this is the case.

\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/warfarin_bar}
\centering
\caption{Out-of-sample probability of correct treatment assignment of the 5 methods on the warfarin dataset with conditionally randomized experiments ($r=0.06, 0.11$). The left graph displays averages for all 5 realizations. The middle figure shows averages over seeds where each treatment class had a sufficient balance of patients getting outcomes~$Y=0$ and~$Y=1$. The right graph shows averages over realizations where treatment $K=2$ had extreme imbalances in the outcome distribution. All figures are averages from trees with $d \in \{1, 2\}$ with~2 and~4 leaf nodes, respectively.}
\label{fig:warfarin_bar}
\end{figure}

Figure~\ref{fig:warfarin_bar} shows the performance of the various methods in the conditionally randomized settings over different subsets of realizations. In some realizations, patients in each treatment group are distributed evenly such that the DM model can predict outcomes accurately, yielding high performance (see Figure~\ref{fig:warfarin_bar}, middle). In other realizations, however, there exist extreme class imbalance in one of the treatments ($K=2$), which results in a poor DM model that predicts all patients would have favorable outcomes from said treatment. An attempt at oversampling was made to remedy the problem, but there were so few samples in the minority class that oversampling could not be done effectively. In these cases, the learned policy is biased toward treating all patients with $K=2$--which is only suitable for a minority of patients--yielding poor results (see Figure~\ref{fig:warfarin_bar}, right).

In Section~\ref{section:intro}, we mentioned that IPW suffers from higher variance, and that DR should be accurate when either the model for IPW or DM is accurate. Our results above do not contradict the theory because, as \cite{Dudik2011DoublyLearning} analyzes, the variance of DR depends on how accurately DM predicts the true outcomes. In this case, DM's prediction is very poor, which explains the high variance for DR relative to IPW.

\comment{Even when we consider seeds that yield a balance of $Y=0, 1$ in all treatment groups, K-PT and B-PT perform only slightly better (see Figure \ref{fig:warfarin_bar}, middle). With few people getting assigned treatment $K=2$ to begin with, K-PT and B-PT find it easier to branch on features that create homogenous leaves where all instances with $K=2$ experience either $Y=0$ or $Y=1$. Both methods seem to be biased toward splitting the minority treatment group and assigning their associated treatment even though this rule does not generalize well. Similar to the synthetic experiments, one could also increase the tree's depth for better performance because its leaves would satisfy Assumption \ref{asu:sufficient_partition}. However, any method with $d > 2$ \delpv{sought} \newpv{presented} \delpv{significant} \newpv{optimality} gaps\delpv{ to optimality after} \newpv{at} the time limit (see Table~\ref{table:warfarin_table}), \toaddresspv{i do not understand this part about modeling here; also the following seems a little more harsh than is ideal; e.g. words like ``fails completely'' ``such a restrictive assumption''; i recommend toning it down.} not to mention that there are clearly better methods (e.g., IPW) that do not require modelling to satisfy such a restrictive assumption. Further, it is also unclear when Assumption~\ref{asu:sufficient_partition} is satisfied and in cases when the covariate distribution is so dissimilar between treatment groups, this assumption fails completely.}

\comment{One notable finding is that DR performs as poorly as DM in these imbalanced seeds and does not seem to reap the benefits of the IPW estimator as theorized. This is caused by 2 factors: an exceedingly poor DM model and the use of binary outcomes. The second term of DR's objective function (see Section \ref{section:robust_formulation}) includes the re-weighting process using IPW but vanishes when either $\pi(X_i) \neq K_i$ or $Y_i = \hat{\nu}_{K_{i}}(X_{i})$. \toaddresspv{also here the language is a little harsh: ``almost always'' is repeated;} DR is incentivized to assign the minority treatment $k=2$ often because most $i$ will have $\pi(X_i) \neq K_i$. This problem is usually counteracted by a decent $\hat{\nu}$, as we can see from the nearly 70\% OOSP in balanced seeds (Figure \ref{fig:warfarin_bar}, middle). However, with a poor model where $\hat{\nu}_2(X)$ almost always predicts 1, $\hat{\nu}_2(X_i) = Y_i$ almost always when $\pi(X_i) = K_i$ for $K_i = 2$. The residual term also almost never varies with only two possible outcomes $\{0, 1\}$. This leaves no room for the other treatment options to correct DR's policy through the second term. In any case, this experiment has shown that IPW is sometimes a more reliable method, especially when the alternative, DM, is so biased that it adversely affects DR's performance.}

\section{Summary}
We presented MIO formulations to learn personalized treatment assignment policies in the form of prescriptive trees from observational data. We showed that our methods are asymptotically exact, which sets us apart from existing literature on prescriptive trees that require either data to originate from marginally randomized experiments or trees to be very deep to yield correct treatment assignments. Our experiments show that our methods perform consistently better across different experiment designs, in some cases performing~$3 \times$ better than the state-of-the-art.

\ACKNOWLEDGMENT{N.\ Jo acknowledges support from the Epstein Institute at the University of Southern California. P.\ Vayanos and S.\ Aghaei gratefully acknowledge support from the Hilton C.\ Foundation, the Homeless Policy Research Institute, and the Home for Good foundation under the ``C.E.S.\ Triage Tool Research \& Refinement'' grant. P.\ Vayanos is also grateful for the support of the National Science Foundation under CAREER award number 2046230. A.\ G\'omez is funded in part by the National Science Foundation under grant 2006762.}

\bibliographystyle{informs2014}

\section{Introduction.}\label{intro}

\section{Introduction}
\label{section:intro}

Personalized decision-making is concerned with determining the best treatment for an individual based on their personal characteristics. This problem arises in a variety of domains, from optimizing online advertisements for different users \citep{li2010contextual} to assigning suitable therapies for patients suffering from a medical condition \citep{flume2007cystic}. In this paper, we aim to design such personalized policies based on data collected in deployment rather than during randomized trials. In such settings, historical treatments are not assigned at random but rather based on a (potentially unknown) policy (e.g., social workers or doctors triaging clients/patients based on their personal characteristics or symptoms). Being able to design personalized treatment assignment policies without running randomized trials \comment{that assign or withhold treatment at random}is extremely important since such trials are often considered unethical, impractical, or may take too long to conduct. Our work is in particular motivated by high-stakes domains such as public health and social interventions, where both interpretability and optimality are critical \citep{rudin2019stop}.

We therefore focus our attention on the design of \emph{optimal prescriptive trees}. A prescriptive tree takes the form of a binary tree. In each branching node of the tree, a binary test is performed on a specific feature. Two branches emanate from each branching node, with each branch representing the outcome of the test. If a datapoint passes (resp.\ fails) the test, it is directed to the left (resp.\ right) branch. A treatment is assigned to all leaf nodes. Thus, each path from root to leaf represents a treatment assignment rule that assigns the same treatment to all datapoints that reach that leaf. The goal of learning prescriptive trees is to select branching decisions and treatment assignments such that the expected value of the outcomes of the treatments on the population are maximized. We say that a prescriptive tree is \emph{optimal} (in-sample, given a training set) if there exists a mathematical proof that no other tree yields better expected outcomes in the population used for training the method. Accordingly, we will say that a tree is \emph{optimal out-of-sample} if there exists a proof that no other tree yields better expected outcomes in the entire population.

\subsection{Problem Statement}
\label{section:problem_statement}

We now formalize the problem we study. In our discussion, we use terminology and concepts from \cite{hernan2019causal}. The goal is to design a personalized policy $\pi : \sets X \rightarrow \sets K$ in the form of a prescriptive tree that maps an individual's characteristics $x \in \sets X \subseteq \reals^F$ to a treatment from a finite set of candidates indexed in the set $\sets{K}$. We note that the set $\sets K$ may also include a ``no treatment'' option. Each individual is characterized by their covariates $X \in \reals^F$ and their \emph{potential outcomes} $Y(k) \in \sets Y \subseteq \reals$ under each treatment $k \in \sets K$\comment{, see e.g., \cite{hernan2019causal}}. The joint distribution of $X$ and $\{Y(k)\}_{k \in \sets K}$ is unknown. However, we have access to $I$ historical observations $\sets D:= \{(X_i, K_i, Y_i)\}_{i \in \mathcal{I}}$ indexed in the set $\sets I :=\{1,\ldots, I\}$, where $X_i$ denote the covariates of the $i$th observation, $K_i$ is the treatment assigned to it, and $Y_i = Y_i(K_i)$ is the \emph{observed outcome}, i.e., the outcome under the treatment received\comment{, see e.g., \cite{hernan2019causal}}. We use the convention that higher values of $Y_i$ are more desirable. Critically, as in all observational experiments, we cannot control the historical treatment assignment policy and the outcomes $Y_i(k)$ for $k \neq K_i$ remain \emph{unobserved}. Formally,  our aim is to learn, from the observational data $\sets D$, a \comment{personalized treatment assignment} policy $\pi$ in the form of a tree of maximum depth $d$ that maximizes the quantity
$$
Q(\pi) \; := \;  \mathbb{E}\left[Y(\pi(X))\right],
$$
where $\mathbb E(\cdot)$ denotes the expectation operator with respect to the joint distribution $\mathbb P$ of $X, K, Y(1), \ldots, Y(k)$. A major challenge in learning the best policy $\pi$ is the fact that we cannot observe the \emph{counterfactual} outcomes $Y_i(k)$, $k\in \sets K, \; k \neq K_i$ that were not received by datapoint~$i$.  Indeed, the missing data make it difficult to identify the best possible treatment for each datapoint. In fact, without further assumptions on the historical policy, it is impossible to even identify the treatment that is the best \emph{in expectation} for any given datapoint.

\subsection{Background on Causal Inference}

In this section, we provide a brief overview of the tools from causal inference that motivate the assumptions we make in our work and our proposed methods.

\subsubsection{Identifiability, Randomized Experiments, and Observational Studies.}
\label{section:randomization}

Learning an optimal prescriptive tree \comment{(or more generally a policy $\pi$)} from the observed data requires that the expected values of the counterfactual outcomes be \emph{identifiable}, i.e., possible to express as a function of the observed data. A sufficient condition for identifiability is for the data to have been collected during a \emph{(marginally) randomized} or \emph{conditionally randomized} experiment, see \cite{hernan2019causal}.

In a randomized experiment, \comment{treatments are assigned at random and }all individuals have the same chance of getting any given treatment independent of their characteristics $x$. \comment{(although treatments need not be equally probable)} Formally, we have $K \independent X$. When treatment assignment is randomized, we say that treatment groups are \emph{exchangeable} because it is irrelevant which group received a particular treatment. Exchangeability means that for all $k$ and~$k' \in \sets K$, those who got treatment~$k$ in the data would have the same outcome distribution under $k'$ as those who got treatment~$k'$. Equivalently, exchangeability means that the counterfactual outcome and the actual treatment are independent, or~$Y(k) \independent K$, for all values of~$k$.

In a conditionally randomized experiment, treatment assignment probabilities depend on the characteristics $x$ of each individual. A conditionally randomized experiment will, in general, not result in exchangeability because treatment groups may differ systemically from one another. However, treatment groups are exchangeable \emph{conditionally on~$X$} since no information other than~$X$ is used to decide the treatment. Thus, conditional randomization ensures \emph{conditional exchangeability}, i.e., $Y(k) \independent K \; | \; X$, for all values of~$k$. We can then infer the outcome distribution under treatment $k'$ for those with covariates $x$ that got treatment $k$ in the data by looking at their counterparts with the same (or similar) $x$ that received treatment $k'$.

As mentioned previously, in high-stakes domains we typically only have access to data from an \emph{observational study}. In such contexts, it is common practice in the literature on causal inference to assume that treatment is assigned at random conditional on $X$, although this may be an approximation and not possible to check in practice. The observational study can then be viewed as a conditionally randomized experiment, provided: \emph{a)} interventions are well defined and are recorded as distinct treatment values in the data; \emph{b)} the conditional probability of receiving a treatment depends only on the covariates $X$; and \emph{c)} the probability of receiving any given treatment conditional on $X$ is positive, i.e., $\mathbb P(K=k|X=x)>0$ almost surely for all $k$. This last condition is termed \emph{positivity} and holds in practice for randomized and conditionally randomized experiments.

We now discuss several methods from the literature for evaluating the performance of a \emph{counterfactual} treatment assignment policy $\pi$ from (conditionally) randomized experiments or, more generally, from observational studies that satisfy the above conditions. In particular, in the following and throughout the paper, we assume that treatments in the data have been assigned according to a (potentially unknown) \emph{logging policy} $\mu$, where $\mu(k,x) := \mathbb P(K=k|X=x)$ and $\mu(k,x)>0$ for all $k \in \sets K$ and $x \in \sets X$.

\subsubsection{Inverse Propensity Weighting\newpv{.}}
\label{section:ipw_intro}

The performance of a counterfactual treatment assignment policy $\pi$ can be evaluated using Inverse Probability Weighting (IPW).

IPW was originally proposed to estimate causal quantities such as expected values of counterfactual outcomes, average treatment effects, risk ratios, etc., see \cite{horvitz1952generalization}. It relies on reweighting the outcome for each individual $i$ in the dataset by the inverse of their \emph{propensity score}, given by $\mu(K_i,X_i)$. Such reweighting has the effect of creating a \emph{pseudo-population} where all individuals in the data are hypothetically given all treatments (thus simulating a randomized treatment assignment). Doing so allows for the estimation of the distribution of the unobserved counterfactual outcomes for all possible values of $X$.

Importantly, IPW can also be used to estimate the performance of a counterfactual policy~$\pi$ by reweighting each individual~$i$ in the data by $\frac{\mathbbm{1}(\pi(X_i) = K_i)}{\mu(K_i, X_i)}$ as illustrated in Figure~\ref{fig:ipw_policy_example}, see also \cite{charles2013counterfactual}. This reweigthing simulates the performance of the policy $\pi$ by leveraging the conditional exchangeability property of the data. Using this estimator, the quantity $Q(\pi)$ is estimated from the data as
\begin{equation}\label{eq:ipw_obj}
\QipwI(\pi) := \frac{1}{I}\sum_{i \in \mathcal I} \frac{\mathbbm{1}(\pi(X_i) = K_i) }{\hat \mu(K_i, X_i)} Y_i,
\end{equation}
where $\hat \mu$ is an estimator of $\mu$, which is obtained for instance using machine learning, by fitting a model to $\{(X_i,K_i)\}_{i \in \sets I}$. If $\mu$ is known or if~$\hat \mu$ converges almost surely to~$\mu$, the IPW estimator of $Q(\pi)$ is unbiased. However, it may suffer from high variance, particularly if some of the propensity scores $\mu(K, X)$ are small, see~\cite{Dudik2011DoublyLearning}.

\begin{figure}[h]
\includegraphics[width=\textwidth]{ipw_policy_example_v4}
\centering
\caption{Figure illustrating the evaluation of the performance of the counterfactual policy~$\pi$ (bottom tree) from data collected using the logging policy~$\mu$ (top tree) using the IPW estimator--notational conventions based on~\cite{hernan2019causal}. The historical data presents 50 instances, each with a single binary covariate $X^1 \in \{0,1\}$. $X^1 = 1$ (resp.\ $0$) indicates that the patient is sick (resp. healthy). There are two treatment options ($\sets K:=\{0,1\}$), where $K=1$ ($0$) is to (not) treat the patient. The potential outcomes are binary ($Y(k) \in \{0,1\}$, $k \in \sets K$). The expected reward under policy $\mu$ is $0.6$, while policy $\pi$ has a $0.82$ probability of getting a positive outcome--in fact, it can be shown that $\pi$ is optimal.}
\label{fig:ipw_policy_example}
\end{figure}

\subsubsection{Direct Method.} \label{section:dm}

An alternative estimator of the performance of the counterfactual policy $\pi$ is the Direct Method (DM)\comment{, see \cite{Dudik2011DoublyLearning}}.

DM is based upon the so-called Regress and Compare (RC) approach, which proceeds in three steps to design a policy that maximizes the expected value of the outcomes \citep{Kallus2018BalancedLearning}. First, it partitions the dataset $\sets{D}$ by treatments. Second, for all $k\in \sets K$, it learns a model $\hat{\nu}_{k}(X)$ of $\mathbb{E}(Y|K=k, X)$ using the subpopulation that was assigned treatment $k$. Third, given any $X\in \reals^F$, it estimates $Y(k)$ as $\hat{\nu}_{k}(X)$ and assigns the best treatment according to $\hat \nu$, that is, $\pi^{\text{RC}}(X) = \argmax_{k \in \sets{K}} \hat{\nu}_{k}(X)$. RC has been used, for example in the online setting by \cite{bastani2020online} and \cite{qian2011performance}. Usually, variants of linear regression are used to get estimators $\hat{\nu}$, see \cite{bastani2020online,goldenshluger2013linear,li2010contextual}.

The RC approach, however, does not allow the evaluation of arbitrary policies nor does it enable the design of policies that belong to an arbitrary class. DM adapts the regress and compare framework to resolve these shortcomings. It proposes to estimate $Q(\pi)$ via
\begin{equation}\label{eq:dm_obj}
\QdmI(\pi) := \frac{1}{I}\sum_{i \in \mathcal I} \hat{\nu}_{\pi(X_i)}(X_i).
\end{equation}
Both RC and DM result in poor performance if $\hat{\nu}_k(X)$, $k\in \sets K$, are biased estimators or make predictions that are far from the true expected outcomes, see~\cite{beygelzimer2009offset}. As pointed out by~\cite{Dudik2011DoublyLearning}, $\hat{\nu}_k(X)$ are constructed without information about~$\pi$ and thus may approximate $\mathbb E(Y|K=k,X)$ poorly in areas that are relevant for $Q(\pi)$. That said, if $\hat \nu_k$ are unbiased for all $k \in \sets K$, the DM estimator of~$Q(\pi)$ will be unbiased.

\subsubsection{Doubly Robust Approach.}
\label{section:robust_intro}

Doubly robust (DR) estimation is a family of techniques that combine two estimators. If either one of the two estimators is accurate, then the doubly robust approach is also accurate, which dampens the errors brought by the individual estimators.

In particular, realizing the drawbacks of IPW and DM, \cite{Dudik2011DoublyLearning} proposed a doubly robust approach that estimates the counterfactual performance of a policy $\pi$ as
\begin{equation}
\label{eq:dr_obj}
\QdrI(\pi) := \frac{1}{I}\sum_{i \in \mathcal I} \left(\hat{\nu}_{\pi(X_i)}(X_i)+ (Y_i - \hat{\nu}_{K_i}(X_i)) \frac{\mathbbm{1}(\pi(X_i) = K_i)}{\hat \mu(K_i, X_i)}\right).
\end{equation}
Intuitively, DR uses IPW to estimate the difference between the true outcome and the prediction made by the DM estimator, $Y_i - \hat{\nu}_{K_i}(X_i)$, thus removing any bias caused by DM, assuming the IPW estimator is unbiased. Provided at least one of $\hat \mu$ or $\hat \nu$ converges almost surely to $\mu$ or $\nu$, respectively, the DR estimator of $Q(\pi)$ will be asymptotically unbiased. Furthermore, in practice, $\QdrI$ usually has a smaller variance compared to $\QipwI$ but a higher variance than $\QdmI$, see~\cite{Dudik2011DoublyLearning}.

\subsection{Related Works}\label{section:related_works}
In this section, we position our work in the fields of operations research, causal inference, and machine learning.

\subsubsection{Estimating Heterogeneous Causal Effects.}
To estimate heterogeneous causal effects, \cite{Athey2016RecursiveEffects} and \cite{wager2018estimation} use recursive partitioning to produce causal trees and causal forests, respectively. Non-tree-based methods can be found in \cite{abrevaya2015estimating}, which use inverse propensity weighting, and \cite{fan2020estimation}, which extends \cite{abrevaya2015estimating} to account for high-dimensional data. \cite{powers2018some} adapt three methods--random forests, boosting, and MARS (Multivariate Adaptive Regression Splines)--to also estimate treatment effects in high-dimensional data. The overall goal of these approaches is to find good estimates of the treatment effects rather than to design treatment assignment policies. While some of these, such as the method from \cite{Athey2016RecursiveEffects}, can be \emph{adapted} to design tree policies, these policies will not be interpretable due to the need to combine the predictive trees associated with many different treatments.

\subsubsection{Mixed-Integer Optimization.}

MIO has been gaining traction as a tool to solve challenging learning problems. It has been used for example to tackle sparse regression problems \citep{wilson2017alamo, atamturk2019rank,bertsimas2020sparse,hazimeh2020sparse,xie2020scalable,gomez2021mixed}, verification of neural networks \citep{fischetti2018deep,khalil2018combinatorial,tjandraatmadja2020convex}\newpv{,} and sparse principal component analysis \citep{dey2018convex, bertsimas2020solving}, among others. More importantly in the context of this paper, MIO methods have been proposed to learn optimal decision trees \citep{bertsimas2017optimal, verwer2019learning, aghaei2019learning, aghaeistrong, elmachtoub2020decision, mivsic2020optimization}.

MIO technology has also been used in causal inference. Existing approaches have mostly focused on matching \citep{bennett2020building,zubizarreta2017optimal}, i.e., to estimate causal effects by pairing each treatment instance with a control instance that has similar covariates \citep{rubin2006matched}. A related approach to matching that utilizes MIO technology is subset selection to achieve \emph{balance} (i.e., similar distributions) between the covariates in the treatment and control groups \citep{nikolaev2013balance}. Finally, \citet{mintz2017behavioral} use MIO to estimate the effects of weight-loss interventions in myopic agents and design optimal policies from the perspective of a healthcare provider.

\subsubsection{Learning Prescriptive Trees.}
\label{sec:lit_prescriptive_trees}

In a recent work closely related to this paper, \cite{Kallus2017RecursiveData} proposes an MIO formulation to learn optimal prescriptive trees that relies on a critical assumption: that trees are sufficiently deep such that the historical treatment assignment and covariates are independent \emph{at each leaf}. Building upon this approach, \citet{Bertsimas2019OptimalTrees} seek to ameliorate it by augmenting the objective function with a term aimed to improve accuracy of the outcome predictions. The paper uses coordinate descent with multiple starts to train the prescriptive trees (as opposed to MIO technology), and empirically verifies that the resulting solutions are indeed high-quality. We provide an in-depth analysis of both methods in Section~\ref{section:prev_mio} and compare to them in our experiments.

\subsection{Contributions}\label{section:contributions}

In this work, we do an in-depth analysis of existing methods to learn prescriptive trees, and propose new MIO formulations for solving such problems. In particular, our key contributions are:
\begin{enumerate}[label=\emph{(\alph*)}]
    \item We analyze the assumption made by the existing methods to learn prescriptive trees discussed in Section~\ref{sec:lit_prescriptive_trees}. We demonstrate, by means of examples, that when the assumption fails to hold \emph{for all feasible trees}, the trees returned as ``optimal'' by these methods may be severely suboptimal. They may even branch on noise covariates that are not predictive of outcomes nor treatment assignments. Moreover, our examples show that this assumption is very strict in the sense that trees deep enough to allow branching on all levels of all covariates (including covariates that are independent of outcomes and treatment assignments) may be needed for this assumption to hold at an optimal tree
    \item Motivated by the limitations of existing approaches, we propose novel MIO formulations for learning optimal prescriptive trees based on tools from causal inference. Contrary to methods from the literature, our approaches enable the design of shallow and/or deep trees, thus allowing the decision-maker to tune the trade-off between prescription accuracy and interpretability. We demonstrate that, under mild conditions,  our methods are asymptotically exact in the sense that they converge to an optimal out-of-sample treatment assignment policy as the number of historical data samples tends to infinity. Based on extensive computational experiments on both synthetic and real data, we demonstrate that our asymptotic guarantees translate to significant out-of-sample performance improvements even in finite samples.
\end{enumerate}

The remainder of this paper is organized as follows. In Section \ref{section:prev_mio}, we present the existing prescriptive tree formulations from the literature and study their performance in the context of two simple synthetic examples. We introduce our proposed MIO formulations for learning optimal prescriptive tree in Section~\ref{section:proposed_formulations}. Finally, we summarize our computational experiments in Section~\ref{section:experiments}. All proofs and detailed computational results are provided in the Electronic Companion.

\section{Existing Prescriptive Tree Formulations}
\label{section:prev_mio}
This section presents the methods introduced by \cite{Kallus2017RecursiveData} and \cite{Bertsimas2019OptimalTrees}--from hereon, we refer to these as K-PT and B-PT, respectively. It further analyzes their advantages and drawbacks using two simple examples.

To describe these two approaches, we work with the notation introduced in~\cite{Kallus2017RecursiveData}, which views trees as producing a partition $\sets X = \bigcup_{n \in \sets{L}} \sets{X}_n$ such that $\sets{X}_n \cap \sets{X}_m = \emptyset$ whenever $n \neq m$, where $\sets L$ collects the set of leaf nodes and $\sets X_n \subseteq \sets X$ represents the set of covariate values that are associated with datapoints landing in leaf $n \in \sets L$. The approaches K-PT and B-PT rely on the following \emph{crucial} assumption about the structure of the prescriptive tree.

\begin{asu}[\cite{Kallus2017RecursiveData}]
    The partition $\{\sets X_n\}_{n \in \sets L}$ induced by the prescription tree is sufficiently fine so that
    $K \independent X \; | \; X \in \sets{X}_n \quad \forall n \in \sets{L}.$
    \label{asu:sufficient_partition}
\end{asu}

Assumption~\ref{asu:sufficient_partition} is automatically satisfied in randomized experiments (since $K \independent X$ unconditionally, by definition) or in conditionally randomized experiments provided the partition is ``sufficiently fine'' -- for example, a tree where each $\sets{X}_n$ is a singleton.\comment{. For example, it is satisfied in a conditionally randomized setting for a tree with~$|\sets X|$ leafs that branches on all levels of all features, so that each $\sets{X}_n$ is a singleton.}

\subsection{K-PT}
\label{section:kallus}

Given a partition satisfying Assumption~\ref{asu:sufficient_partition}, \cite{Kallus2017RecursiveData} proves that the \emph{observable} quantity $\mathbb E(Y| K=k, X \in \sets X_n)$ is an unbiased estimator of $\mathbb E(Y(k)|X \in \sets X_n)$. In other words, to estimate the potential outcomes under treatment~$k$ at leaf~$n$, it suffices to average outcomes for those that actually received treatment~$k$ at leaf~$n$ under~$\mathbb P$; this quantity can be estimated from the observed data. This statement is intuitive under Assumption~\ref{asu:sufficient_partition}, which implies that all individuals with covariates in the set $\sets X_n$ have the same chance of getting any given treatment, i.e., they are exchangeable within $\sets X_n$. It follows that, under Assumption~\ref{asu:sufficient_partition}, the outcomes of a policy $\pi$ that assigns treatment $\pi(\sets X_n) \in \sets K$ to all individuals that land in leaf~$n$ can be estimated as
\begin{equation}
\sum_{n \in \sets L} \mathbb P ( X \in \sets X_n) \cdot \mathbb E ( Y | K= \pi(\sets X_n), X \in \sets X_n).
\label{eq:kallus_estimator}
\end{equation}
Accordingly, this quantity can be estimated from data as
$$
\QkptI(\pi) := \sum_{n\in \sets{L}} \frac{\sum_{i \in \sets{I}}\mathbbm{1}[X_i \in \sets{X}_n]}{|\sets I|}\widehat{Y}_n \left(\pi(\sets{X}_n)\right),
$$
where $\widehat{Y}_n(k)$ is the sample average approximation of $\mathbb E(Y| K=k, X \in \sets X_n)$, defined through
\begin{equation}
\widehat{Y}_n(k):=\frac{1}{\sum_{i \in \sets{I}}\mathbbm{1}[K_{i} = k, X_{i} \in \sets{X}_n]}\sum_{i\in \sets{I}}\mathbbm{1}[K_i =k, X_i \in \sets{X}_n] Y_i.
\label{eq:avg_estimate}
\end{equation}

\subsubsection*{Analysis of K-PT.}

\cite{Kallus2017RecursiveData} proved that $\QkptI(\pi)$ is an unbiased estimator of $Q(\pi)$ when Assumption~\ref{asu:sufficient_partition} holds, which is the case in \emph{marginally} randomized settings for example. A salient advantage of K-PT is that, if Assumption 1 \emph{does} hold, it avoids predicting~$\mu$ and thus avoids the bias caused by potentially inaccurate estimates~$\hat \mu$. Unfortunately, though, in \emph{conditionally} randomized settings, decision trees will in general not satisfy Assumption~\ref{asu:sufficient_partition} and the tree obtained by maximizing $\QkptI(\pi)$ may be severely suboptimal for the problem of maximizing $Q(\pi)$. In fact, as we now demonstrate, K-PT may actively choose trees that violate Assumption~\ref{asu:sufficient_partition} even when the complete out-of-sample distribution is available \emph{and} the tree is deep enough to allow branching on all covariates that are predictive of outcomes and treatment assignments.

\begin{example}
Consider the following data distribution. Let the covariate vector $X = (X^1, X^2) \in \{0,1\}^2$, where $X^1$ and $X^2$ are independent and Bernoulli distributed with parameter 0.5. We interpret $X^1$ to denote the severity of a patient's condition, where $X^1 = 0$ (resp.\ $X^1=1$) indicates a healthy (resp.\ sick) patient. $X^2$ is a noise variable, unrelated to both the outcomes and the treatment assignments. There are two possible treatments, where $K=1$ (resp.\ $K=0$) indicates that the patient was treated (resp.\ not treated). Treatments are assigned according to a conditionally randomized experiment, with $\mathbb P(K=0|X^1=0) = 0.9$ and $\mathbb P(K=0|X^1=1) = 0.1$, i.e., 90\% of healthy patients do not receive the treatment, and 90\% of sick patients do. Table~\ref{table:expected} shows the expected values of the potential outcomes in dependence of the covariate values. \comment{To circumvent any finite sample issues, we assume that we have sufficient historical data at our disposal to be able to estimate all quantities in the K-PT estimator of~$Q(\pi)$ without errors; thus we can work with the true values in~\eqref{eq:kallus_estimator}, rather than with their empirical estimates provided in the definition of $\QkptI(\pi)$.}We assume that we have sufficient historical data at our disposal so that we can work with the true values in~\eqref{eq:kallus_estimator} rather than with their empirical estimates provided in the definition of~$\QkptI(\pi)$.

\begin{table}[H]
\centering
 \begin{tabular}{c | c c}
 \hline
 $X^1$ & $\mathbb{E}(Y(0))$ & $\mathbb{E}(Y(1))$ \\
 \hline\hline
 0 & 1 & 0.8 \\
 \hline
 1 & 0 & 0.2 \\
 \hline
\end{tabular}
\caption{Companion table to Example~\ref{ex:example1}.  Expectation of potential outcomes in dependence of the covariate values. A value $X^1 = 0$ (resp.\ $X^1=1$) of the covariate indicates a healthy (resp.\ sick) patient. Larger values of the outcomes are preferred.}
\label{table:expected}
\end{table}

From Table~\ref{table:expected}, it can be seen that the (unique) optimal policy treats sick patients only. In particular, this optimal policy can be modelled by a prescriptive tree of depth one (i.e., a tree with a single branching node and two leafs) that branches on~$X^1$, does not treat any of the healthy patients (left leaf) and treats all sick patients (right leaf). The expected outcome under this policy is given by
$$
\mathbb{E}(Y(0)|X^1=0)\mathbb{P}(X^1=0)+\mathbb{E}(Y(1)|X^1=1)\mathbb{P}(X^1=1)=1\times\frac{1}{2}+0.2\times \frac{1}{2}=0.6.
$$
One may similarly calculate the expected outcomes under a policy that branches on the noise feature~$X^2$. In this case, no matter which treatment is assigned at each leaf, the resulting policies would have a lower expected outcome of 0.5. These candidate policies and the corresponding expected outcomes under each policy are illustrated in the top row of Figure~\ref{fig:kallus_example}.

\begin{figure}[h]
	\includegraphics[width=\textwidth]{kallus_example}
	\centering
	\caption{Companion figure to Example~\ref{ex:example1}. Each row shows the estimates of performance of different prescriptive trees (branch on $X^1$--left-- or $X^2$--right--) under different estimators. In each case, only the best treatment assignments at each leaf (as measured by the corresponding estimator) are shown. In this example, the true optimal policy (top left) is to branch on $X^1$ and treat only those who are sick. Indeed, the expected outcome of the best tree branching on $X^1$ (top left) is $1.0 \cdot 0.5 + 0.2 \cdot 0.5 = 0.6$, which is larger than the expected outcome of the best $X^2$-branching tree (top right), which is $0.5 < 0.6$. However, the method from \cite{Kallus2017RecursiveData} (middle), which is discussed in Section \ref{section:kallus}, will prefer to branch on $X^2$ (noise covariate) and not treat any patient (middle right). Indeed, it estimates the outcomes of the best $X^1$- (resp.\ $X^2$-) branching tree as $0.6$ (resp.\ $0.9$). We note that Assumption~\ref{asu:sufficient_partition} is not satisfied at any of the leaves of the tree deemed as best by this estimator. In this case, the method of \cite{Bertsimas2019OptimalTrees}, see Section \ref{section:bertsimas}, is able to identify the truly optimal tree (bottom left); it correctly branches on~$X^1$ provided~$\theta$ is sufficiently small. Note that this method penalizes the variances corresponding to all treatments, but for simplicity the figure only depicts variances corresponding to the chosen treatment.}
	\label{fig:kallus_example}
\end{figure}

In practice, the quantities $\mathbb{E}(Y(k))$, $k\in \sets K$, cannot be easily estimated from data since $Y_i(k)$ is only observable for $i$ such that $K_i=k$ and treatments are not assigned at random. Thus, the procedure described in the previous paragraph to select an optimal prescriptive tree is not possible to follow. \comment{Now, l}Let us instead consider the tree of depth one that maximizes the objective function of K-PT given by~\eqref{eq:kallus_estimator}. Suppose that we branch on covariate $X^j$ at the sole branching node of the tree, $j\in \{1,2\}$. Then, at the leaf where $X^j = x^j \in \{0,1\}$, K-PT estimates the counterfactual outcomes $\mathbb E(Y(k)|X^j=x^j)$ as $\mathbb{E}(Y|K=k, \; X^j=x^j)$. However, these estimates can be incorrect. For example,
$$
\mathbb{E}(Y(0)|X^2=0) = \sum_{x^1 \in \{0,1\}} \mathbb{E}(Y(0)|X^1=x^1,X^2=0)\mathbb{P}(X^1=x^1|X^2=0)=1 \times \frac{1}{2} + 0\times \frac{1}{2} = 0.5,
$$
whereas
$$
\mathbb{E}(Y|K=0,X^2=0) = \sum_{x^1 \in \{0,1\}} \mathbb{E}(Y|K=0,X^1=x^1,X^2=0)\mathbb{P}(X^1=x^1|K=0,X^2=0
=0.9.
$$
The difference between these two quantities is intuitive. Individuals that were not treated ($K=0$) in the data are more likely to be healthy (with a large associated outcome). The K-PT estimate $\mathbb{E}(Y|K=0,X^2=0)$ \comment{(K-PT estimate) }thus \emph{overestimates} $\mathbb{E}(Y(0)|X^2=0)$. Similarly, since the vast majority of individuals that were treated ($K=1$) in the data are sick (with low outcomes), the K-PT estimate $\mathbb{E}(Y|K=0,X^2=1)=0.26$ is an \emph{underestimator} for $\mathbb{E}(Y(1)|X^2=0)=0.5$.

The middle row of Figure~\ref{fig:kallus_example} depicts the outcomes of the two candidate trees as estimated by K-PT. When branching on $X^1$ (middle left), Assumption~\ref{asu:sufficient_partition} is satisfied and the estimated outcome of~$0.6$ provided by K-PT is correct. However, K-PT prefers the alternative tree which branches on the noise variable and does not treat any of the patients (middle right); it wrongly estimated an outcome of~$0.9$ while the true expected outcome is only~$0.5$. We conclude that if not \emph{all feasible trees} in the problem satisfy Assumption~\ref{asu:sufficient_partition}, the method from~\cite{Kallus2017RecursiveData} may fail.

\label{ex:example1}
\end{example}

\subsection{B-PT}
\label{section:bertsimas}

\cite{Bertsimas2019OptimalTrees} also observe that K-PT may fail in certain conditions. They posit that failings can be attributed to poor estimates of the outcomes at the leaves. Thus, they augment the K-PT estimator with a regularization term, which penalizes trees that induce a high variance in the outcomes conditional on each treatment at the leaves, i.e.,
$$
\sum_{k \in \sets K} \sum_{n \in \sets L} \mathbb V{\rm{ar}} ( Y | X \in \sets X_n , K = k ).
$$
Specifically, they optimize the quantity
\begin{gather*}
\QbptI(\pi) = \theta \left(\QkptI(\pi)\right) - (1-\theta) \sum_{i \in \sets{I}} \left(Y_i - \widehat{Y}_{\sets{X}_{n(i)}}(K_i) \right)^2,
\end{gather*}
where $n(i) \in \sets L$ denotes the leaf datapoint $i$ belongs to, and $\theta\geq 0$ is a parameter to be tuned.

\subsubsection*{Analysis of B-PT.} Recall that \cite{Kallus2017RecursiveData}
estimates $\mathbb E(Y(k) | X \in \sets X_n)$ through $\mathbb E(Y | K=k, X \in \sets X_n)$ and shows that this estimator is unbiased under Assumption~\ref{asu:sufficient_partition}. By adding a regularization term that seeks to promote trees with low variance $\mathbb V{\rm{ar}}(Y| K=k, X \in \sets X_n)$, \cite{bertsimas2017personalized} improves the statistical properties of the K-PT estimator (lower variance). Therefore, if Assumption~\ref{asu:sufficient_partition} holds (so that $\mathbb{E}(Y|K=K_i, X\in \sets{X}_{n(i)})=\mathbb{E}(Y(K_i)| X\in \sets{X}_{n(i)})$), then B-PT can lead to improved performance over K-PT (for a suitable choice of $\theta$). However, B-PT does not directly address the issues discussed in Example~\ref{ex:example1}. Indeed, when Assumption~\ref{asu:sufficient_partition} does not hold, then $\mathbb{E}(Y|K=K_i, X\in \sets{X}_{n(i)})$ and $\mathbb{E}(Y(K_i)| X\in \sets{X}_{n(i)})$ can be drastically different from one another, and the policies chosen by B-PT may still be suboptimal.

In the specific case of Example~\ref{ex:example1}, \comment{it can be checked that if $\theta$ is sufficiently small, then} B-PT can in fact recover the best tree if $\theta$ is sufficiently small. Nonetheless, as we now show, a small modification of Example~\ref{ex:example1} results in a setting in which B-PT has \emph{additional} incentives to choose the worst tree, which branches on irrelevant features.

\setcounter{example}{0}
\begin{example}[Continued]
We now revisit Example~\ref{ex:example1} and use the method from~\cite{Bertsimas2019OptimalTrees} to identify an optimal tree. Figure~\ref{fig:kallus_example} (bottom) shows the estimated outcomes and variances associated with each candidate tree. For simplicity, we assume that the  potential outcomes in Table~\ref{table:expected} have zero variance. Since the tree that branches on the relevant feature~$X^1$ has smaller variance in this case, it will be chosen by B-PT provided that $\theta$ is sufficiently small. \hfill $\blacksquare$

\end{example}

\begin{example}
\label{ex:example2}
We now consider a variant of Example~\ref{ex:example1} where we add to all potential outcomes the quantity $(1-X^2)$, see Table~\ref{table:expected2}. Note that while $X^2$ now affects the value of the outcomes, it does not affect the value of the difference $\mathbb E(Y(1)-Y(0)|X)$ and is thus \emph{irrelevant} in deciding which treatment is preferable. Therefore, under suitable identifiability conditions, there exists \comment{in this problem }an optimal prescriptive tree of depth one. Figure~\ref{fig:bertsimas_example2} (top) shows the outcomes, for a depth one tree, under the two possible branching decisions if the counterfactuals are known--the best tree branches on the only predictive feature $X^1$. However, when using estimates $\mathbb{E}[Y|K=k,X=x]$ for the counterfactuals (Figure~\ref{fig:bertsimas_example2}, bottom), the tree that branches on the irrelevant feature $X^2$ has both a higher estimated outcome and a smaller variance. In this case B-PT has even more incentive to select a policy based on the irrelevant feature and both K-PT and B-PT will select a suboptimal tree. \hfill $\blacksquare$

\begin{table}[H]
\centering
 \begin{tabular}{c c | c c}
 \hline
 $X^1$ & $X^2$ & $\mathbb{E}(Y(0))$ & $\mathbb{E}(Y(1))$ \\ [0.5ex]
 \hline\hline
 0 & 0 & 2 & 1.8 \\
 \hline
 1 & 0 & 1 & 1.2 \\
 \hline
 0 & 1 & 1 & 0.8 \\
 \hline
 1 & 1 & 0 & 0.2 \\
 \hline
\end{tabular}
\caption{Companion table to Example~\ref{ex:example2}. Expectation of the potential outcomes in dependence of the covariate values. Similar to Example~\ref{ex:example1} (Continued), we assume that the potential outcomes have a variance of zero.}
\label{table:expected2}
\end{table}

\begin{figure}[h]
\includegraphics[width=\textwidth]{bertsimas_example2}
\centering
\caption{Companion figure to Example~\ref{ex:example2}. The layout and interpretation of the figure exactly parallel that in Figure~\ref{fig:kallus_example}. The optimal tree (top left) has the objective $0.5(1.5) + 0.5(0.7) = 1.1$, which is higher than the $X^2$-branching tree (top right), $0.5(1.5) + 0.5(0.5) = 1.0 < 1.1$. However, B-PT will incorrectly branch on $X^2$ no matter the regularization strength $\theta$, since the $X^1$-branching tree (bottom left) has a lower estimated value and also higher variance.}
\label{fig:bertsimas_example2}
\end{figure}
\end{example}

\comment{Continuing from the motivating example in Section \ref{section:example_kallus}, we show that variance reduction might similarly lead to unfavorable results. Consider, again, a decision tree that branches once on $X^1$, and we want to determine the ideal treatment $K \in \{0, 1\}$ for the left leaf (i.e. datapoints with $X^1 = 0$). For $K=0$, the average outcome is $\mathbb{E}(Y|X^1=0, K=0) = 1.5$ and the variance is
\begin{gather*}
\mathbb{E}[(Y-\mathbb{E}(Y))^2|X^1=0, K=0] = 0.5 \cdot (2-1.5)^2 + 0.5 \cdot (1-1.5)^2 = 0.25.
\end{gather*}

The regularized term then becomes $0.5(1.5) - 0.5(0.25) = 0.625$, where 0.5 represents the regularization strength $\theta$. \cite{Bertsimas2019OptimalTrees} cited that 0.5 seems to yield the best results, although this parameter needs to be tuned in practice. We compare this value with $K=1$, where $\mathbb{E}(Y|X^1=0, K=1) = 1.3$ and $\mathbb{E}[(Y-\mathbb{E}(Y))^2|X^1=0, K=0] = 0.25$. The regularization term becomes 0.525, so we assign the more favorable $K=0$. Continue to evaluate all possible permutations to find the optimal tree, which is summarized in Figure \ref{fig:bertsimas_example}. The algorithm still branches on $X^2$, failing to remedy the problem in Section \ref{section:example_kallus}.}

\section{Proposed Formulations}
\label{section:proposed_formulations}

In this section, we present our proposed MIO formulations for learning optimal prescriptive trees from observational data. Without loss of generality, we assume that the covariates are binary, i.e., $\sets X \subseteq \{0,1\}^F$ -- this assumption can easily be relaxed to cater for integer or categorical features, see Remark~1 in~\cite{aghaeistrong}. Importantly, and in sharp contrast with the existing literature on prescriptive trees, we \emph{do not} make Assumption~\ref{asu:sufficient_partition}. Instead, we directly optimize the estimators~\eqref{eq:ipw_obj}, \eqref{eq:dm_obj}, or \eqref{eq:dr_obj} from the causal inference literature. This enables us to design interpretable prescriptive trees that are guaranteed to be optimal out-of-sample as the number of historical samples grows.

As will become clear shortly, our proposed approaches give rise to \emph{weighted} classification trees. We thus adapt the MIO formulations of~\cite{aghaeistrong} for learning optimal classification trees from the predictive to the prescriptive setting. We extend our notation from the problem statement in Section \ref{section:problem_statement} to describe optimal prescriptive trees. The building block of our formulation is a perfect binary tree of depth~$d$, whose nodes are numbered $1$ through $(2^{d+1}-1)$ in the order in which they appear in a breadth-first search. We let $\sets{B} := \{1, \dots, 2^{d}-1\}$ denote the set of branching nodes and $\sets{T} := \{2^d, \dots, 2^{d+1}-1\}$ collect all terminal nodes (see Figure \ref{fig:flow_graph}, left). From the perfect binary tree, we build a flow graph as follows. We connect a source node~$s$ to the root node, and connect all nodes other than $s$ to the sink node $t$ (see Figure \ref{fig:flow_graph}, right). All links in the graph are directed from source to sink and have capacity 1.
\begin{figure}[h]
\includegraphics[width=0.7\textwidth]{flow_graph}
\centering
\caption{A prescriptive tree with depth 2 (left) and its associated flow graph (right).}
\label{fig:flow_graph}
\end{figure}

\subsection{Formulation Based on Inverse Propensity Weighting}
\label{section:ipw_formulation}

Equipped with the flow graph associated with a prescriptive tree, we now formulate an MIO problem to optimize~\eqref{eq:ipw_obj} over all trees of maximum depth $d$. For every branching node $n\in \sets B$ and feature $f\in \sets F$, we let the binary variable~$b_{nf}$ indicate if feature~$f$ is selected for branching at node~$n$. Accordingly, for every node~$n \in \sets{T} \cup \sets{B}$, we let the binary variable~$p_n$ indicate if node~$n$ is a treatment node, in which case a treatment must be assigned to all datapoints that land at that node and no further branching is allowed. For~$k\in \sets K$, we let~$w_{nk}\in \{0,1\}$ equal~1 if and only if (iff) treatment~$k$ is selected at node~$n$. Finally, the ``flow variables'' $z$ capture the flow of datapoints from the source to the sink: $z^{a(n),n}_{i} = 1$ iff the following two conditions are satisfied: \textit{(i)} datapoint~$i$  flows to node~$n$ from its ancestor~$a(n)$; and \textit{(ii)} datapoint~$i$ reaches a treatment node where treatment~$K_i$ was assigned (i.e., $w_{nK_i} = 1$ for some~$n\in \sets B\cup \sets T$). \delag{If datapoint $i$ were to flow to node $n$ from its ancestor $a(n)$, then $z^{a(n),n}_{i} = 1$ and $0$ otherwise. Critically, since flow must be preserved, datapoint $i$ flows to sink $t$ if it is directed to a terminal node that assigns the same treatment as was observed in the data (i.e., $w_{nK_i} = 1$). Otherwise, datapoint $i$ will have a flow value of 0.}

\delag{From the training data, unless the true propensity scores are known, determine a-priori the scores $\mu(K_i, X_i)$, which is learned using a suitable machine learning model.}

The formulation is as follows:
\begin{subequations}
\begin{align}
\text{maximize} \;\; & \displaystyle \sum_{i \in \mathcal I} \sum_{n \in \sets B \cup \sets T } (z^{n,t}_{i}) \frac{Y_i}{\hat\mu(K_i, X_i)} = \sum_{i \in \sets I}z^{s,1}_{i}\frac{Y_i}{\hat\mu(K_i, X_i)}
\label{eq:flow_reg_obj}\\
\text{subject to}\;\;
& \displaystyle \sum_{f \in \sets F}b_{nf} + p_n + \sum_{m \in \sets A(n)}p_m = 1   &\hspace{-5cm}  \forall n \in \sets B \label{eq:flow_reg_branch_or_predict}\\
&  p_n+\sum_{m \in \sets A(n)}p_{m} =1   &  \forall n \in \sets T \label{eq:flow_reg_terminal_leaf}\\
& \displaystyle z^{a(n),n}_{i} =  z^{n,\ell(n)}_{i} + z^{n,r(n)}_{i} + z^{n,t}_{i}  &\hspace{-5cm}  \forall n \in \sets B, i \in \mathcal I \label{eq:flow_reg_conservation_internal}\\
&  \displaystyle z^{a(n),n}_{i} = z^{n,t}_{i} &\hspace{-5cm}   \forall i \in \mathcal I,n \in \sets T \label{eq:flow_reg_conservation_terminal}\\
& \displaystyle z^{s,1}_{i} \leq 1 &\hspace{-5cm} \forall i \in \mathcal I\label{eq:flow_reg_source}\\
&  \displaystyle z^{n,\ell(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i\leq0}b_{nf} &\hspace{-5cm} \forall n \in \sets B, \; i \in \mathcal I \label{eq:flow_reg_branch_left}\\
&  \displaystyle z^{n,r(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i=1}b_{nf}  &\hspace{-5cm} \forall n \in \sets B, \; i \in \mathcal I \label{eq:flow_reg_branch_right}\\
&  \displaystyle z^{n,t}_{i} \leq  w_{nK_i} &\hspace{-5cm} \forall i \in \mathcal I, \; n \in \sets B \cup \sets T \label{eq:flow_reg_sink}\\
&  \displaystyle \sum_{k \in \sets K}w_{nk} = p_n  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T \label{eq:flow_reg_leaf_prediction}\\
&  \displaystyle w_{nk} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B \cup \sets T, \; k \in \sets K \\
&  \displaystyle b_{nf} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B, \; f \in \sets F \\
&  \displaystyle p_{n} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B \cup \sets T \\
&  \displaystyle z^{a(n),n}_{i}, z^{n,t}_{i}\in \{0,1\}  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T, \; i \in \sets I,
\end{align}
\label{eq:flow_reg}
\end{subequations}
where $\sets A(n)$ denotes the set of all ancestors of node $n\in \sets B \cup \sets T$, $r(n)$ (resp.\ $\ell(n)$) is the right (resp.\ left) descendant of $n$, and $\hat \mu$ is as defined in Section~\ref{section:ipw_intro}. The objective \eqref{eq:flow_reg_obj} maximizes the IPW estimator of the counterfactual performance of the prescriptive tree defined through variables $b$, $w$, and $p$, as introduced in \eqref{eq:ipw_obj}. Constraints~\eqref{eq:flow_reg_branch_or_predict} ensure that we branch at all branching nodes where no treatment is assigned, at the node itself or at any of its ancestors. Similarly, constraints~\eqref{eq:flow_reg_terminal_leaf} impose that a treatment must be assigned at all terminal nodes, unless a treatment has already been assigned to one of its ancestors. Constraints~\eqref{eq:flow_reg_conservation_internal} and~\eqref{eq:flow_reg_conservation_terminal} are flow conservation constraints, whereby any datapoint that flows into a node $n$ must either exit to its right or left descendant, or flow directly to the sink. Constraints~\eqref{eq:flow_reg_source} ensure that a flow of at most one can enter the source node for each datapoint. Constraints~\eqref{eq:flow_reg_branch_left} and~\eqref{eq:flow_reg_branch_right} guarantee that datapoints flow according to the branching decisions specified by variables $b$. Constraints~\eqref{eq:flow_reg_sink} state that only datapoints whose assigned treatments match the one observed in the data can flow to the sink. Finally, constraints~\eqref{eq:flow_reg_leaf_prediction} combined with integrality of the $w_{nk}$ variables ensure that if a treatment is assigned at a node, exactly one of the available treatments must be assigned. Note that we can relax integrality on the~$w$ and~$z$ variables and still reach an integral solution.

\subsection{Formulation Based on Doubly Robust Estimation}
\label{section:robust_formulation}

We now formulate an MIO problem to optimize~\eqref{eq:dr_obj} over all trees of maximum depth~$d$. To this end, we adapt the flow graph associated with the decision tree so that both correctly and incorrectly classified datapoints are tracked and also redefine the variables $z$. Specifically, we create $|\sets{K}|$ sink nodes (denoted by $t_k, \forall k \in \sets{K}$) and connect each sink node to all nodes other than the source node~$s$. As before, all arcs are directed from source to sink, see Figure~\ref{fig:flow_graph_multisink}. In this formulation, each datapoint will be directed towards the sink corresponding to its assigned treatment. Therefore,
we redefine variables $z$ as $z^{a(n),n}_{i} = 1$ iff datapoint~$i$ flows to node~$n$ from its ancestor $a(n)$ (regardless of the treatment assigned to the datapoint in the data).

\begin{figure}[h]
\includegraphics[scale=0.35]{images/flow_graph_multisink}
\centering
\caption{A modified flow graph with multiple sink nodes; in this case, there are 2 sinks to reflect 2 treatments.}
\label{fig:flow_graph_multisink}
\end{figure}

The MIO formulation for DR remains largely unchanged from formulation~\eqref{eq:flow_reg}:
\begin{subequations}
\begin{align}
\text{maximize} \;\; & \displaystyle \sum_{i \in \mathcal I} \sum_{k \in \sets K } \sum_{n \in \sets B \cup \sets T}z^{n,t_{k}}_{i} \hat{\nu}_k(X_i) + \displaystyle \sum_{i \in \mathcal I} \sum_{n \in \sets B \cup \sets T} (Y_i -\hat{\nu}_{K_{i}}(X_i)) \frac{{z^{n,t_{K_{i}}}_{i}} }{\hat\mu(K_i, X_i)} \label{eq:robust_obj}\\
\text{subject to}\;\;
& \displaystyle \sum_{f \in \sets F}b_{nf} + p_n + \sum_{m \in \sets A(n)}p_{m} = 1   &\hspace{-5cm}  \forall n \in \sets B \label{eq:robust_branch_or_predict}\\
&  p_n+\sum_{m \in \sets A(n)}p_{m} =1   &  \forall n \in \sets T \label{eq:robust_all_terminal_leaf}\\
& \displaystyle z^{a(n),n}_{i} =  z^{n,\ell(n)}_{i} + z^{n,r(n)}_{i} + \sum_{k \in \sets K}z^{n,t_{k}}_{i}   &\hspace{-5cm}  \forall n \in \sets B,\; i \in \mathcal I \label{eq:robust_conservation_internal}\\
&  \displaystyle z^{a(n),n}_{i} = \sum_{k \in \sets K}z^{n,t_{k}}_{i}  &\hspace{-5cm}   \forall i \in \mathcal I, \; n \in \sets T \label{eq:robust_conservation_terminal}\\
& \displaystyle z^{s,1}_{i} = 1 &\hspace{-5cm} \forall i \in \mathcal I\label{eq:robust_source}\\
&  \displaystyle z^{n,\ell(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i=0}b_{nf} &\hspace{-5cm} \forall n \in \sets B,\; i \in \mathcal I \label{eq:robust_branch_left}\\
&  \displaystyle z^{n,r(n)}_{i}\leq \sum_{f \in \sets F: x_{f}^i=1}b_{nf}  &\hspace{-5cm} \forall n \in \sets B, \; i \in \mathcal I \label{eq:robust_branch_right}\\
&  \displaystyle z^{n,t_{k}}_i \leq  w_{nk} &\hspace{-5cm} \forall i \in \mathcal I,\; n \in \sets T,k \in \sets K \label{eq:robust_sink}\\
&  \displaystyle \sum_{k \in \sets K}w_{nk} = p_n  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T \label{eq:robust_leaf_prediction}\\
&  \displaystyle w_{nk} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets T, \; k \in \sets K \\
&  \displaystyle b_{nf} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B, \; f \in \sets F \\
&  \displaystyle p_{n} \in \{0,1\}  &\hspace{-5cm}   \forall n \in \sets B \cup \sets T \\
&  \displaystyle z^{a(n),n}_{i}, z^{n,t_{k}}_{i} \in \{0,1\}  &\hspace{-5cm}  \forall n \in \sets B \cup \sets T, \; i \in \sets I, \; k \in \sets K.
\end{align}
\label{eq:robust}
\end{subequations}
Here, $\hat \mu$ and $\hat \nu_k$ are as defined in Sections~\ref{section:ipw_intro} and~\ref{section:dm}, respectively. The main difference relative to formulation~\eqref{eq:flow_reg} is the objective \eqref{eq:robust_obj}, which now maximizes the doubly robust objective \eqref{eq:dr_obj}. The other difference is in the flow conservation constraints \eqref{eq:robust_conservation_internal} and \eqref{eq:robust_conservation_terminal}, which restrict points to flow to at most one sink node. The last difference is in constraint~\eqref{eq:robust_sink}, which states that datapoints can only flow to the sink associated with their assigned treatment. Similar to formulation~\eqref{eq:flow_reg}, we can relax integrality on the $w$ and $z$ variables and still reach an integral solution.

\subsection{Formulation Based on Direct Method Estimator}
\label{section:dm_formulation}

Formulation~\eqref{eq:robust} can easily be adapted to optimize the direct method objective~\eqref{eq:dm_obj} by simply dropping the second term in the objective~\eqref{eq:robust_obj}.

\subsection{Analysis of IPW, DM, and DR}
\label{section:analysis}

It can be shown that if $\mu$ (resp.\ either $\mu$ or $\nu$) are known, a depth one tree trained according to formulations~\eqref{eq:flow_reg} (resp.\ \eqref{eq:robust}) in Examples~\ref{ex:example1} and~\ref{ex:example2} will be optimal out-of-sample, provided the number of samples in the data is sufficiently large. In this section, we go one step further and demonstrate that, under mild conditions, the tree returned by our optimal MIO formulations~\eqref{eq:flow_reg}, \eqref{eq:robust}, and~\eqref{eq:robust} modified as in Section~\ref{section:dm_formulation} will converge almost surely to an optimal tree (in terms of out-of-sample treatment assignment).

We introduce some notation to help formalize our claim. We let $\Pi_d$ denote the set of tree-based policies of maximum depth~$d$. For any tree based policy $\pi \in \Pi_d$, define
$$\Qipw (\pi) := \mathbb E\left[ \frac{ \mathbb I (K= \pi(X)) Y }{ \mu (K,X) } \right] \quad \text{and} \quad \QipwIN (\pi) := \frac{1}{I} \sum_{i=1}^{I} \frac{ \mathbb I (K_i= \pi(X_i)) Y_i }{ \muhatN (K_i,X_i) },$$
where $\muhatN (K,X)$ is an estimator of $\mu (K,X)$ trained over a sample of cardinality $N$ independent from~$\mathcal D$.
We denote by~$v^\star$ and~$\sets S^\star$, the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} Q(\pi)$, respectively.
Similarly, we denote by~$\vhatINipw$ and~$\shatINipw$, the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} \QipwIN(\pi)$, respectively. Finally, we use the abbreviation ``w.p.1" for ``with probability one".
\begin{proposition}
\label{prop:exact_IPW}
Suppose that: conditional exchangeability holds, i.e., $Y(k) \independent K | X$, for all values of~$k \in \sets K$; the potential outcomes $Y(k)$ are bounded, for all values of~$k \in \sets K$; and, both $\mu(k,x) >0$ and $\muhatN(k,x)>0$ for all $x \in \sets X$, $k\in \sets K$, and $N$ sufficiently large. If $\muhatN$ converges almost surely to~$\mu$, then $\vhatINipw \rightarrow v^\star$ w.p.1 and $\shatINipw \subseteq \sets S^\star$ w.p.1 for $I$ and $N$ large enough.
\end{proposition}

Similar to the IPW case, we now show that under mild conditions, the MIO formulation~\eqref{eq:robust} for the DR estimator returns an optimal out-of-sample tree. We introduce additional notation to help formalize our claim.
For any tree based policy $\pi \in \Pi_d$, define
$$\Qdr (\pi) \; := \; \mathbb E\left[ \hat{\nu}_{\pi(X)}(X) + (Y-\hat{\nu}_{\pi(X)}(X))\frac{ \mathbb I (K= \pi(X)) }{ \hat{\mu} (K,X) } \right]$$
$$\text{and} \quad \QdrIN (\pi)\; := \; \frac{1}{I} \sum_{i=1}^{I} \left(\nuhatN_{\pi(X_i)}(X_i) + (Y-\nuhatN_{\pi(X_i)}(X_i))\frac{ \mathbb I (K_i= \pi(X_i)) }{ \muhatN(K_i,X_i) }\right),$$
where $\nuhatN_K (\cdot)$ is an estimator of~$\nu_K(\cdot)$, trained over a sample of cardinality~$N$ independent from~$\mathcal D$. Let $\hat{\mu}(K,X)$ (resp. $\hat{\nu}_K(X)$) be the limit of $\muhatN (K,X)$ (resp. $\nuhatN_K (X)$) as $N$ goes to infinity.
We denote by~$\vhatINdr$ and~$\shatINdr$, the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} \QdrIN(\pi)$, respectively.

\begin{proposition}
\label{prop:exact_DR}
Suppose that: conditional exchangeability holds, i.e., $Y(k) \independent K | X$, for all values of~$k \in \sets K$; the potential outcomes $Y(k)$ are bounded, for all values of~$k \in \sets K$;  both $\mu(k,x) >0$ and $\muhatN(k,x)>0$ for all $x \in \sets X$, $k\in \sets K$ and $N$ sufficiently large; and both $\hat \mu(k,x)$ and $\hat \nu_k (x)$ exist and are bounded for all $x \in \sets X$, $k\in \sets K$. If either $\muhatN$ converges almost surely to $\mu$ or $\nuhatN_K$ converges almost surely to $\nu_k$, then $\vhatINdr \rightarrow v^\star$ w.p.1 and $\shatINdr \subseteq \sets S^\star$ w.p.1 for $I$ and $N$ large enough.
\end{proposition}

At the end, we show that under mild conditions, the MIO formulation~\eqref{eq:robust} modified as in Section~\ref{section:dm_formulation} for the direct estimator returns an optimal out-of-sample tree. We introduce additional notation to help formalize our claim.
For any tree based policy $\pi \in \Pi_d$, define
$$\Qdm (\pi) \; := \; \mathbb E\left[ \hat{\nu}_{\pi(X)}(X) \right] \quad \text{and} \quad  \QdmIN (\pi)\; := \; \frac{1}{I} \sum_{i=1}^{I} \left(\nuhatN_{\pi(X_i)}(X_i)\right).$$
We denote by~$\vhatINdm$ and~$\shatINdm$ the optimal value and the set of optimal solutions of~$\max_{\pi \in \Pi_d} \QdmIN(\pi)$, respectively.

\begin{proposition}
\label{prop:exact_DM}
Suppose that: conditional exchangeability holds, i.e., $Y(k) \independent K | X$, for all values of~$k \in \sets K$; and $\hat \nu_k (x)$ exists, for all $k \in \sets K$. If $\nuhatN_K$ converges almost surely to $\nu_k$, then $\vhatINdm \rightarrow v^\star$ w.p.1 and $\shatINdm \subseteq \sets S^\star$ w.p.1 for $I$ and $N$ large enough.
\end{proposition}

\subsection{Extensions} Our formulations in Sections \ref{section:ipw_formulation}, \ref{section:robust_formulation}, and \ref{section:dm_formulation} have strong modeling power; they can be used to learn trees that interpretable and/or fair, to design randomized policies, and to impose budget constraints relevant for under-resourced settings. We discuss these in the present section. We introduce all these constraints on formulation~\eqref{eq:robust} which tracks which datapoints are assigned to which treatments. We note that formulation~\eqref{eq:robust} can be used to optimize all three objectives (IPW, DR, or DM).

\subsubsection*{Intepretability.} In settings where learning interpretable policies is desired, see e.g., \cite{rudin2019stop} and~\cite{azizi2018designing}, one may limit the number of branching nodes in the tree to $M$ by augmenting formulation~\eqref{eq:robust} with the constraint $$\sum_{n\in \sets B\cup \sets T}p_n \; \leq \; M.$$

\subsubsection*{Budget Constraints.} In low resource settings such as organ allocation \citep{bertsimas2013fairness, zenios2000dynamic}, housing allocation \citep{azizi2018designing}, and security \citep{xu2018strategic}, many treatments/interventions often have limited supply. Tree based policies that satisfy such capacity constraints can be learned by augmenting formulation~\eqref{eq:robust} with constraints
$$ \displaystyle \sum_{n \in \sets B \cup \sets T} \sum_{i \in \sets I} z^{n,t_{k}}_i \; \leq \; |I|C_k \quad\forall k \in \sets K, $$
where $C_k$ denotes the percentage of instances that can be assigned treatment $k$.

\subsubsection*{Fairness in Treatment Assignment.} Our formulations can be used to learn trees where treatments are assigned fairly across groups. Concretely, we will say that a policy satisfies treatment assignment parity if the probability of assigning a particular treatment is equal across protected groups. We let $\sets P$ collect all levels of a (collection of) protected covariates (e.g., levels of race and LGBTQ status). Accordingly, we let $P_i \in \sets P$ represent the value of the protected feature(s) of datapoint~$i$. It is worth noting that the features in $\sets P$ are usually not included in covariates $X$ to make sure that the policy does not make decisions based on protected features. To ensure that the learned policy satisfies treatment assignment parity up to a bias $\delta$, we may augment~\eqref{eq:robust} with the constraint
$$\left| \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p} z^{n, t_k}_i}{|\{i \in \sets I : P_i = p\}|} - \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p'} z^{n, t_k}_i}{|\{i \in \sets I : P_i = p'\}|} \right| \; \leq \; \delta \qquad \forall p, p' \in \sets P: p \neq p', \; k \in \sets K.$$

In a similar fashion, we will say that a policy satisfies \textit{conditional} treatment assignment parity if the probability of assigning a particular treatment is equal across protected groups, conditional on some legitimate feature(s) that affect the outcome. For instance, in prescribing housing resources to people experiencing homelessness, one may require that all individuals with the same vulnerability have the same likelihood of receiving a certain treatment.  We let $\sets A$ be the set of features indicative of risk, and let $A_i \in \sets A$ be the value of the risk feature of datapoint $i$. Conditional treatment assignment parity is satisfied up to a bias $\delta$ and for all features $a \in \sets A$ by adding the following constraint to \eqref{eq:robust},
$$\left| \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p, A_i = a} z^{n, t_k}_i}{|\{i \in \sets I : P_i = p \cap A_i = a\}|} - \frac{\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p', A_i = a}  z^{n, t_k}_i}{|\{i \in \sets I : P_i = p' \cap A_i = a \}|} \right| \; \leq \; \delta  \qquad \forall p, p' \in \sets P: p \neq p', \; k \in \sets K, \; a \in \sets A$$

\subsubsection*{Fairness in Treatment Outcomes}. Our formulations can also be used to ensure fairness in expected outcomes across groups. In general, this requires the average expected outcomes of a protected group $p \in \sets P$ to be above some threshold~$\gamma_p$ $\in \mathbb R$:
$$\displaystyle \sum_{n \in \sets B \cup \sets T} \displaystyle \sum_{i \in \sets I : P_i = p} \left[(Y_i -\hat{\nu}_{K_{i}}(X_i)) \frac{{z^{n,t_{K_{i}}}_{i}}}{\hat\mu(K_i, X_i)} + \sum_{k \in \sets K} z^{n, t_k}_i \hat{\nu}_k(X_i)\right] \; \geq \; \gamma_p \qquad \forall p \in \sets P.$$
For example, \cite{bertsimas2013fairness} set
$\gamma_p := \sum_{i \in \sets I: P_i = p} Y_i$, i.e., the expected outcomes of all protected groups under the learned policy should be greater than or equal to what was observed in the data. Alternatively, the constraint can be used to impose max-min fairness to protect the outcomes of the groups that are worst off, see \cite{rawls1974some}. In this case, one sets $\gamma_p := \gamma$ for all $p \in \sets P$, where $\gamma$ is the largest value for which the MIO problem remains feasible.

\subsubsection*{Randomized Treatment Assignment Policies.} The addition of fairness and budget constraints may make problem~\eqref{eq:robust} infeasible. In such cases, it may be desirable to design \emph{randomized} policies where datapoints that land at the same leaf are assigned each treatment with a certain probability (rather than all getting the same treatment). This can be achieved by relaxing integrality on the variables~$w$ and~$z$ in formulation~\eqref{eq:robust}. The variable $w_{nk}$ can then be interpreted as the probability of assigning treatment~$k$ to datapoints that fall on node~$n$. Similarly, $z^{n, t_k}_i$ can be viewed as the probability of assigning treatment~$k$ to datapoint~$i$ if it lands at node~$n$ (and 0 otherwise).

\section{Experiments}
\label{section:experiments}

We now evaluate the empirical performance of our proposed formulations in Section~\ref{section:proposed_formulations} in two problem settings: a synthetic setting from the literature and a real setting based on a warfarin dosing dataset. In both cases, we take the viewpoint that data comes from an observational study\comment{where the treatment assignment policy is not known by the researcher, but that the study}, which can be viewed as a conditionally randomized experiment. We benchmark against the approaches of~\cite{Kallus2017RecursiveData} and \cite{Bertsimas2019OptimalTrees} for learning prescriptive trees, see Section~\ref{section:prev_mio}. As before, we refer to these as K-PT and B-PT, respectively. Since \cite{Bertsimas2019OptimalTrees} do not propose an MIO formulation, we adapt the method from \cite{Kallus2017RecursiveData} to implement B-PT. For completeness, we provide the MIO formulations that we have implemented for these two approaches in Appendices~\ref{section:kallus_formulation} and~\ref{section:bertsimas_formulation}, respectively.
We investigate the performance of all approaches as the probability of correct treatment assignment is varied in the historical data

\subsection{Dataset Description}
\label{section:synthetic1}

\subsubsection*{Synthetic Data.} For our experiments on synthetic data, we adapt the data generation process from \cite{Athey2016RecursiveEffects}. In this problem, there are two treatment possibilities indexed in the set $\sets K = \{0, 1\}$. The covariate vector has two independent and identically distributed features, $X=(X^1, X^2)$, where each $X^j \sim \mathcal{N}(0, 1)$. The potential outcome of datapoint~$i$ with covariates~$X_i$ under treatment~$k$ is
\begin{equation}\label{eq:synthetic}
Y_i(k) = \phi(X_i) + \frac{1}{2}(2k-1) \cdot \kappa(X_i) + \epsilon_i,
\end{equation}
where $\phi(x) := \frac{1}{2}x^1 + x^2$ models the mean effect, $\kappa(x):=\frac{1}{2}x^1$ models the treatment effect, and $\epsilon_i \sim \sets N(0,0.1)$ is noise added to the outcome that is independent of the covariates.

To study different settings of observational experiments, we vary, in the data, the probability~$p$ of assigning the treatment that is best \emph{in expectation} for each unit\comment{; we denote this probability by~$p$}. We let $p \in \{ 0.1, 0.25, 0.5 , 0.75, 0.9\}$ -- since there are 2 treatments, 0.5 corresponds to the (marginally) randomized setting, while the other settings correspond to conditionally randomized experiments. For each~$p$, we randomly generate 5 training and test sets, each with 500 and 10,000 datapoints, respectively.

\subsubsection*{Warfarin Dosing.} For our study on real data, we employ a dataset for personalized warfarin dosing for which counterfactuals are available. Warfarin is the most widely used oral anticoagulant agent, but despite its prevalence, determining one's optimal warfarin dosage is difficult because it can vary widely depending on demographic variables, clinical factors, and genetics \citep{international2009estimation}. The publicly available dataset that we \comment{employ}use was collected by the International Warfarin Pharmacogenetics Consortium and published at the Pharmacogenetics and Pharmacogenomics Knowledge Base, see \cite{international2009estimation}. The advantage of using this dataset is that we can model a patient's true outcomes when given varying doses of warfarin, which allows us to evaluate the performance of arbitrary counterfactual policies. \cite{international2009estimation} published a learned affine function $f(x) = \beta x + c$ that determines the optimal warfarin dose based on a patient's age, weight, race, VKORC1 genotype, CYP2c9 genotype, and whether or not the patient is currently taking amiodarone or an enzyme reducer, see equation~\eqref{warfarin_function} in Section~\ref{sec:appendix_warfarin}. We calculate a patient $i$'s optimal dosage using $f(X_i) + \epsilon$, where $\epsilon \sim \mathcal{N}(0, 0.02).$ This dosage is then discretized into three groups (i.e., $|\mathcal K|=3$) using the same convention as \cite{international2009estimation}: $K_{i}^{\text{opt}} = 0$ ($\leq$ 3 mg/day), 1 (between 3 and 7 mg/day), and 2 ($\geq$ 7 mg/day). The observed outcome is $Y_{i}(K_i) = 1$, if $K_i = K_{i}^{\text{opt}}$; and $=0$ otherwise.

To study different settings of observational experiments, we consider three different treatment assignment mechanisms, as follows. For the marginally randomized setting, we assign each treatment with probability $1/3$, and repeat to produce 5 datasets. To simulate data that is based on a more informed policy, we assign treatments based on modified versions of $f$ obtained by perturbing its coefficients so that patients closer to the boundary of each treatment bucket are at higher risk of receiving an incorrect treatment. Specifically, for each coefficient~$a$ of $f$, we generate $a'\sim U(a-a\cdot r, a+a \cdot r)$, where $r$ denotes the range of sampling. We assign patient $i$ the treatment corresponding to this modified function, discretized into three groups as discussed previously. In testing, we found that $r \in [0.05, 0.12]$ was the ideal range to simulate a reasonable treatment policy where the probability of correct treatment assignment is between 0.6 and 0.9. We fix 2 values of~$r$, 0.06 and 0.11, and for each value, we generate 5 sets of randomly sampled coefficients (from hereon, we refer to this variability as a ``realization''). At this point we have~3 experiment designs--each with~5 realizations--making up~15 datasets. Each dataset is then split randomly~5 times into training and test sets of~3,000 and~1,387 instances, respectively, yielding a total of~75 train-test pairs.

\subsection{Experimental Setup}
\delag{Throughout our experiments, we assume that we are in the \emph{observational setting}, where the historical logging policy is unknown, see Section~\ref{section:randomization}.} Since the proposed methods\delag{our proposed formulations} rely on estimating the propensity scores and/or counterfactuals, we discuss the estimation models we use for both datasets. We also discuss the preprocessing steps required.

\subsubsection*{Synthetic Data.} The synthetic data has real covariates but our methods only allow for binary features to split on. We thus discretize each covariate feature into ten buckets corresponding to quantiles from a normal distribution. Since the outcome distribution is known by construction (it is linear), we learn the potential outcomes using both linear regression (LR) and lasso regression with $\alpha = 0.08$ (Lasso), knowing that the latter will be a slightly less accurate predictor. We also use true propensity scores--they are known by construction--as well as two methods to estimate propensity scores: logistic regression (Log) and decision trees (DT). Therefore, there will be six doubly robust methods, each corresponding to a pair of models for IPW and DM. These methods are then compared to K-PT and B-PT. All approaches are allotted a solve time of 1 hour.

\subsubsection*{Warfarin Dosing.} Most of the features are already binary with the exception of age, height, and weight, so these features were split into 5 buckets, where each bucket contains approximately the same number of datapoints. A variety of other preprocessing and data imputation steps were done based on recommendations from \cite{international2009estimation}. A patient's possible outcomes are binary, and in selecting DM's model we found that using a random forest (RF) gave the best validation set accuracy. The propensity scores are learned using decision trees (DT), also shown to have the best validation set performance. These methods are then compared to K-PT and B-PT. All approaches are allotted a solve time of 1 hour.

\subsubsection*{Policy Evaluation.} We evaluate a learned policy~$\pi$ in terms of its out of sample probability of assigning the optimal treatment to a patient (OOSP). In the next section, we compare all methods using this metric.

\subsection{Experimental Results}
In this section, we focus on analyzing the results for trees of depth 1 and 2, i.e., with 2 and 4 leaves, respectively. We decided to focus on these results because K-PT did not scale to deeper trees and resulted in large optimality gaps, making for an unfair comparison. Indeed, from Figure~\ref{fig:performance}, we observe that K-PT had an optimality gap of more than $100\%$ for most instances of trees of depth 3 and 4. Refer to Section~\ref{section:experiments_detail} for a detailed reporting of computational times.

\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/performance_v2}
\centering
\caption{Comparison of computational performance for the synthetic and real data over all methods except for B-PT. B-PT is excluded because our implementation is an MIO adaptation of the local search method described in \cite{Bertsimas2019OptimalTrees}, and so showing either of B-PT (which is slow) or the method of \cite{Bertsimas2019OptimalTrees} (which is fast but does not provide optimality gaps) on the graph would not yield a fair comparison. We only include the models associated with the best statistical performance. K-PT solved 157 instances to optimality within 3,600 seconds while DR solved the same number of instances in 303 seconds, for a 11.9$\times$ speed-up.}
\label{fig:performance}
\end{figure}

\label{section:experimental_results}
\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/athey_v1}
\centering
\caption{Results on the synthetic dataset. The left graph compares the out-of-sample probability of correct treatment assignment of the best models for IPW, DM, and DR with the methods described in \cite{Kallus2017RecursiveData} (K-PT) and \cite{Bertsimas2019OptimalTrees} (B-PT). The right graph shows the same, but only includes the DR approaches where at least one of the combined methods is correct. Both figures are averages from trees of $d \in \{1, 2\}$ with 2 and 4 leaf nodes, respectively.}
\label{fig:athey_v1}
\end{figure}

\subsubsection*{Synthetic Data.}
Our results on the synthetic data are summarized on Figure~\ref{fig:athey_v1} (more details can be found in Section~\ref{section:experiments_detail}). The figure shows the average OOSP for each method in dependence of the probability of correct treatment assignment, $p$. From the figure, it can be seen that the DR methods, with and OOSP of around 75\%, consistently outperform both K-PT and B-PT whose best performance is 70\%. Even when one of the two estimators in DR is inaccurate (i.e., includes Lasso or Log), DR favors the other (better) estimator as expected. Notably, K-PT and B-PT drop to around 50\% OOSP (i.e., similar to random treatment assignment) when the probability of correct treatment assignment in the data is 0.75 and 0.9. Their poor performance in these settings is particularly relevant since most observational data is based on an informed policy that presumably treats most patients correctly. The likely explanation for this drop in performance is that a tree of moderate depth is unlikely to partition the dataset sufficiently fine for Assumption~\ref{asu:sufficient_partition} to hold.

\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/warfarin_box}
\centering
\caption{Out of sample probability of correct treatment assignment over all 5 realizations on the warfarin dataset. The left graph shows the distribution for the randomized experiments. The middle (resp.\ right) figure shows the distribution of results for $r = 0.06$ (resp.\ $0.11$). All figures are averages from trees with $d \in \{1, 2\}$ with 2 and 4 leaf nodes, respectively.}
\label{fig:warfarin_box}
\end{figure}

\subsubsection*{Warfarin Dosing.}
Our results on the warfarin data are summarized on Figures~\ref{fig:warfarin_box} and~\ref{fig:warfarin_bar} (more details can be found in Section~\ref{section:experiments_detail}). Figure~\ref{fig:warfarin_box} plots the distribution of OOSP over the 5 realizations. The left figure shows that all methods have similar performance in the randomized setting, except for B-PT which suffered from optimality gaps. From hereon, we discuss results only from the conditionally randomized experiments ($r \in \{0.06, 0.11\}$), which are more pertinent in our context. In both these settings (middle and right subfigures), both K-PT and B-PT have relatively low variance, but they perform consistently below 30\% OOSP (i.e., worse than random assignment). IPW seems to be the strongest estimator consistently at around 60\%. Both DM and DR have high variances that span from outperforming IPW in some realizations to performing below 20\% in others--we now discuss why this is the case.

\begin{figure}[h]
\includegraphics[width=0.85\textwidth]{images/warfarin_bar}
\centering
\caption{Out-of-sample probability of correct treatment assignment of the 5 methods on the warfarin dataset with conditionally randomized experiments ($r=0.06, 0.11$). The left graph displays averages for all 5 realizations. The middle figure shows averages over seeds where each treatment class had a sufficient balance of patients getting outcomes~$Y=0$ and~$Y=1$. The right graph shows averages over realizations where treatment $K=2$ had extreme imbalances in the outcome distribution. All figures are averages from trees with $d \in \{1, 2\}$ with~2 and~4 leaf nodes, respectively.}
\label{fig:warfarin_bar}
\end{figure}

Figure~\ref{fig:warfarin_bar} shows the performance of the various methods in the conditionally randomized settings over different subsets of realizations. In some realizations, patients in each treatment group are distributed evenly such that the DM model can predict outcomes accurately, yielding high performance (see Figure~\ref{fig:warfarin_bar}, middle). In other realizations, however, there exist extreme class imbalance in one of the treatments ($K=2$), which results in a poor DM model that predicts all patients would have favorable outcomes from said treatment. An attempt at oversampling was made to remedy the problem, but there were so few samples in the minority class that oversampling could not be done effectively. In these cases, the learned policy is biased toward treating all patients with $K=2$--which is only suitable for a minority of patients--yielding poor results (see Figure~\ref{fig:warfarin_bar}, right).

In Section~\ref{section:intro}, we mentioned that IPW suffers from higher variance, and that DR should be accurate when either the model for IPW or DM is accurate. Our results above do not contradict the theory because, as \cite{Dudik2011DoublyLearning} analyzes, the variance of DR depends on how accurately DM predicts the true outcomes. In this case, DM's prediction is very poor, which explains the high variance for DR relative to IPW.

\comment{Even when we consider seeds that yield a balance of $Y=0, 1$ in all treatment groups, K-PT and B-PT perform only slightly better (see Figure \ref{fig:warfarin_bar}, middle). With few people getting assigned treatment $K=2$ to begin with, K-PT and B-PT find it easier to branch on features that create homogenous leaves where all instances with $K=2$ experience either $Y=0$ or $Y=1$. Both methods seem to be biased toward splitting the minority treatment group and assigning their associated treatment even though this rule does not generalize well. Similar to the synthetic experiments, one could also increase the tree's depth for better performance because its leaves would satisfy Assumption \ref{asu:sufficient_partition}. However, any method with $d > 2$ \delpv{sought} \newpv{presented} \delpv{significant} \newpv{optimality} gaps\delpv{ to optimality after} \newpv{at} the time limit (see Table~\ref{table:warfarin_table}), \toaddresspv{i do not understand this part about modeling here; also the following seems a little more harsh than is ideal; e.g. words like ``fails completely'' ``such a restrictive assumption''; i recommend toning it down.} not to mention that there are clearly better methods (e.g., IPW) that do not require modelling to satisfy such a restrictive assumption. Further, it is also unclear when Assumption~\ref{asu:sufficient_partition} is satisfied and in cases when the covariate distribution is so dissimilar between treatment groups, this assumption fails completely.}

\comment{One notable finding is that DR performs as poorly as DM in these imbalanced seeds and does not seem to reap the benefits of the IPW estimator as theorized. This is caused by 2 factors: an exceedingly poor DM model and the use of binary outcomes. The second term of DR's objective function (see Section \ref{section:robust_formulation}) includes the re-weighting process using IPW but vanishes when either $\pi(X_i) \neq K_i$ or $Y_i = \hat{\nu}_{K_{i}}(X_{i})$. \toaddresspv{also here the language is a little harsh: ``almost always'' is repeated;} DR is incentivized to assign the minority treatment $k=2$ often because most $i$ will have $\pi(X_i) \neq K_i$. This problem is usually counteracted by a decent $\hat{\nu}$, as we can see from the nearly 70\% OOSP in balanced seeds (Figure \ref{fig:warfarin_bar}, middle). However, with a poor model where $\hat{\nu}_2(X)$ almost always predicts 1, $\hat{\nu}_2(X_i) = Y_i$ almost always when $\pi(X_i) = K_i$ for $K_i = 2$. The residual term also almost never varies with only two possible outcomes $\{0, 1\}$. This leaves no room for the other treatment options to correct DR's policy through the second term. In any case, this experiment has shown that IPW is sometimes a more reliable method, especially when the alternative, DM, is so biased that it adversely affects DR's performance.}

\section{Summary}
We presented MIO formulations to learn personalized treatment assignment policies in the form of prescriptive trees from observational data. We showed that our methods are asymptotically exact, which sets us apart from existing literature on prescriptive trees that require either data to originate from marginally randomized experiments or trees to be very deep to yield correct treatment assignments. Our experiments show that our methods perform consistently better across different experiment designs, in some cases performing~$3 \times$ better than the state-of-the-art.

\ACKNOWLEDGMENT{N.\ Jo acknowledges support from the Epstein Institute at the University of Southern California. P.\ Vayanos and S.\ Aghaei gratefully acknowledge support from the Hilton C.\ Foundation, the Homeless Policy Research Institute, and the Home for Good foundation under the ``C.E.S.\ Triage Tool Research \& Refinement'' grant. P.\ Vayanos is also grateful for the support of the National Science Foundation under CAREER award number 2046230. A.\ G\'omez is funded in part by the National Science Foundation under grant 2006762.}

\bibliographystyle{informs2014}