Conditional Average Treatment Effect
Ruoqing Zhu
Last Updated: September 23, 2026
1 Overview
We have introduced the average treatment effect (ATE) estimation in the previous lecture. In this lecture, we will take a step further to study the conditional average treatment effect (CATE) estimation, which aims to estimate the average treatment effect among individuals who share the covariate value \(X=x\). This is different from the individual causal effect \(Y_i(1)-Y_i(0)\), which remains unobserved because we only observe one potential outcome for each individual. CATE estimation is a key component in personalized medicine, where the goal is to recommend treatment based on individual characteristics. Precisely, we are interested in the conditional average treatment effect:
\[ \tau(x) = \E[Y(1) - Y(0) | X = x] \]
and the optimal treatment regime can be inferred from the CATE:
\[ \begin{aligned} \pi_\text{opt}(x) &= \underset{a \in \{0,1\}}{\arg\max} \,\, \E\left[Y(a) | X = x\right] \\ &= \mathbf{1}\{\tau(x) > 0\} \end{aligned} \]
where \(\pi(\cdot)\) is a treatment assignment function that takes the patient’s covariate \(X\) as the input. When a larger outcome is preferred, the rule assigns treatment 1 whenever \(\tau(x)>0\) and treatment 0 otherwise. The best of all such functions, \(\pi_\text{opt}(x)\), is called the optimal treatment regime. In this section, we will discuss CATE estimation, which is often done by regression methods. When these methods are used to estimate the optimal treatment regime, they are often referred to as indirect learning methods. On the other hand, there are also direct learning methods that directly estimate the optimal treatment regime without estimating the CATE first. We will discuss both types of methods in this lecture.
2 The Virtual Twins with Random Forests
The most straightforward idea is to perform regression models on both the treated and untreated groups, and use their difference to infer the treatment decision. Here, we assume consistency, no unmeasured confounding, and overlap. Hence, the approach involves two steps:
- In the first step, we learn the regression functions \(\E[Y | X, A = 1]\) and \(\E[Y | X, A = 0]\).
- In the second step, we estimate the CATE for each covariate value \[ \hat{\tau}(X_i) = \widehat{\E}[Y | X_i, A = 1] - \widehat{\E}[Y | X_i, A = 0] \] and use \(\mathbf{1}\{\hat{\tau}(X_i)>0\}\) as the label for learning an estimated decision rule.
A typical example of this type of approach is called the Virtual Twins model (Foster et al. 2011). It uses random forests to learn the regression functions in the first step, and in the second step, it uses a single tree model to summarize the decision rule and make it interpretable. The two-stage workflow here follows the vignette of the aVirtualTwins R package. We also need the packages for random forests and CART.
Let’s consider a simulated data set derived from the SIDES package. The data in .csv format can be downloaded from our course website. In this data set, 470 patients and 14 variables are collected. The variables are listed below.
Health: Health outcome (larger the better)THERAPY: 1 for active treatment, 0 for the control treatmentTIMFIRST: Time from first sepsis-organ fail to start drugAGE: Patient age in yearsBLLPLAT: Baseline local plateletsblSOFA: Sum of baseline sofa score (cardiovascular, hematology, hepatorenal, and respiration scores)BLLCREAT: Base creatinineORGANNUM: Number of baseline organ failuresPRAPACHE: Pre-infusion apache-ii scoreBLGCS: Base GLASGOW coma scale scoreBLIL6: Baseline serum IL-6 concentrationBLADL: Baseline activity of daily living scoreBLLBILI: Baseline local bilirubinBEST: A treatment label supplied with the course data. You should not use this variable when fitting models!
For each patient, sepsis was observed during their hospital stay. Hence, one of the two treatments (indicated by variable THERAPY) must be chosen to prevent further adverse events. After the treatment, the patient’s health outcome (Health) was measured, with a larger value being the better outcome. The course data include a BEST label, but its provenance is not documented here. We use it only for a descriptive comparison, not as the true optimal treatment.
Run the cell below to load the data set and display the first few rows. It will only take a few seconds to complete. Make sure that you have the working directory setup correctly. You can do this by putting your .rmd file and the .csv file in the same folder, then open your .rmd file to pop up RStudio.
Sepsis <- read.csv("..//..//dataset//Sepsis.csv")
# remove the first column, which is observation ID
Sepsis = Sepsis[, -1]
head(Sepsis)2.1 Step 1: Outcome Regressions
We will fit two random forests to model the outcome Health: one model is based on all patients who received treatment (THERAPY) 1, and the other model is based on all patients who received treatment 0 (corresponding to our previous label \(-1\)). Denote these two models as \(\hat f_1(x)\) and \(\hat f_0(x)\), respectively. Here, we will also use a little trick in random forest, which is the out-of-bag prediction. This is similar to a leave-one-out type of approach and the fitted value of a subject is calculated from the trees that do not use this subject. For more details, you may refer to this lecture note. This mechanism helps prevent overfitting to a certain extent. In the meantime, we should also tune both mtry and nodesize, but that step will be skipped in the demonstration. The following code demonstrates this idea:
library(randomForest)
## randomForest 4.7-1.2
## Type rfNews() to see new features/changes/bug fixes.
# fit model for treatment 0
model0 <- randomForest(Health ~ . - BEST,
data = Sepsis[Sepsis$THERAPY == 0, ],
ntree=2000,
mtry = 5,
nodesize = 5)
# fit model for treatment 1
model1 <- randomForest(Health ~ . - BEST,
data = Sepsis[Sepsis$THERAPY == 1, ],
ntree=2000,
mtry = 5,
nodesize = 5)
# out-of-bag prediction for treatment 0 group with model 0
model0.treat0 = model0$predicted
# out-of-bag prediction for treatment 1 group with model 1
model1.treat1 = model1$predicted
# prediction for treatment 1 group with model 0
model0.treat1 = predict(model0, Sepsis[Sepsis$THERAPY == 1, ])
# prediction for treatment 0 group with model 1
model1.treat0 = predict(model1, Sepsis[Sepsis$THERAPY == 0, ])
# combine predictions together
pred.treat0 = rep(NA, nrow(Sepsis))
pred.treat0[Sepsis$THERAPY == 0] = model0.treat0
pred.treat0[Sepsis$THERAPY == 1] = model0.treat1
pred.treat1 = rep(NA, nrow(Sepsis))
pred.treat1[Sepsis$THERAPY == 0] = model1.treat0
pred.treat1[Sepsis$THERAPY == 1] = model1.treat1
# which is better?
pred.best = (pred.treat1 > pred.treat0)
# agreement with the supplied label
mean(pred.best == Sepsis$BEST)
## [1] 0.8085106The fitted forest rule agrees with the supplied BEST label for 81% of these patients. This in-sample agreement does not establish treatment-rule accuracy or policy value. The rule chooses the treatment with the higher predicted mean outcome, but the twin random forests are difficult to interpret.
2.2 Step 2: A Tree Decision Rule
This second step can be optional, but very useful in practice if a simple decision rule is preferred. We will construct a single-tree model (CART) to approximate the fitted forest rule. The idea is very simple: we will create an artificial label, best.label, which is 1 if pred.treat1 is larger than pred.treat0, and 0 otherwise. Then we will use all of our covariates to model the best.label we learned.
library(rpart)
# fit a decision tree to predict pred.best,
# excluding BEST, Health, and THERAPY from the dataset
best.label = ifelse(pred.treat1 > pred.treat0, "Treatment 1", "Treatment 0")
best.label = as.factor(best.label)
rpart.fit <- rpart(best.label ~ . - BEST - Health - THERAPY, data = Sepsis)
# in the coming cells, we will prune the tree
# start by plotting the cross-validation relative error at different tree sizes
plotcp(rpart.fit)
After some tuning of the tree model, using the cp parameter, we can view the fitted tree model:
# cut the tree
cutted.tree = prune(rpart.fit, cp = 0.04)
library(rpart.plot)
# make a good plot
rpart.plot(cutted.tree)
The fitted tree recommends treatment 1 when \(\text{Age} < 52\). For \(\text{Age} \geq 52\), it recommends treatment 0 if the apache score is less than 32, and treatment 1 otherwise.
3 Direct Learning with Linear Regression
We can of course fit a linear regression model with interactions between the treatment label and the covariates. If we simply do that and use the fitted model to predict the best treatment, we are doing indirect learning.
\[ \E[Y | X, A] = \boldsymbol \beta^T X + \boldsymbol \alpha^T X \cdot A \]
The following code demonstrates this idea using a simulated data set and the Lasso model. Please note that we will be using \(A \in \{-1, 1\}\) as the treatment label since it will benefit our later development.
set.seed(1)
n = 100
p = 200
X = matrix(rnorm(n*p), n, p)
A = rbinom(n, 1, 0.5)*2-1
beta = c(rep(0.2, 10), rep(0, p-10))
alpha = c(rep(0, p-2), 0.75, -0.75)
R = X %*% beta + A * (X %*% alpha) + rnorm(n, sd = 0.5)
testn = 500
testX = matrix(rnorm(testn*p), testn, p)
TRUE_BEST = sign(testX %*% alpha) library(glmnet)
## Loading required package: Matrix
## Loaded glmnet 5.0
# fit two regression models
Apos.fit <- cv.glmnet(x = cbind(X, A)[A == 1,], y = R[A == 1])
Aneg.fit <- cv.glmnet(x = cbind(X, A)[A == -1,], y = R[A == -1])
# compare the models for new data
Apos.pred = predict(Apos.fit, newx = cbind(testX, 1))
Aneg.pred = predict(Aneg.fit, newx = cbind(testX, -1))
# the predicted treatment label
indirect.pred = sign(Apos.pred - Aneg.pred)
# the accuracy
mean(indirect.pred == TRUE_BEST)
## [1] 0.876This seems to be fairly accurate. We got 88% of accuracy, meaning that, for around 88% of the new patients, we suggested the better treatment.
3.1 Case 1: Randomized Trial
But let’s also explore an interesting fact about our previous equation. This method was proposed by Tian et al. (2014). Suppose we do not use \(Y\) as our outcome of interest. Instead, we use \(AY\), the product of treatment label and the outcome. Then, we have
\[ \begin{aligned} \E(AY| X) =& \E( A | X) \cdot \boldsymbol \beta^T X + \E(A^2 | X) \cdot \boldsymbol \alpha^T X.\\ =& \big[ P( A = 1 | X) - P(A = -1 | X) \big] \cdot \boldsymbol \beta^T X + \boldsymbol \alpha^T X \end{aligned} \]
If our study is collected from a balanced randomized trial, i.e., \(\Pr(A = 1 | X) = P(A = -1 | X) = 0.5\), then the first term is zero. And we only have the second term
\[ \E(AY | X) = \boldsymbol \alpha^T X \]
This suggests a very simple idea for directly modeling the best treatment rule using this modified outcome:
# fitting the direct learning model
direct.fit <- cv.glmnet(x = X, y = R*A)
# direct learning predicted label
direct.pred <- sign(predict(direct.fit, testX))
# the accuracy
mean(direct.pred == TRUE_BEST)
## [1] 0.916This time, the accuracy 92% is slightly better than the previous one.
3.2 Case 2: Observational Study
As we can expect, the way to deal with an observational study is to use propensity score weighting. We can modify the previous equation using an inverse propensity weight. Of course, we also assume suitable conditions as we did in previous sections.
\[ \begin{aligned} & \E \left( \frac{A}{\Pr( A | X)}Y \Biggm| X \right) \\ =& \left[ \frac{1}{\Pr( A | X)}\Pr( A = 1 | X) + \frac{-1}{\Pr( A = -1 | X)} \Pr( A = -1 | X) \right] \cdot \bbeta^T X + \\ &\quad \left[ \frac{1^2}{\Pr( A = 1| X)}\Pr( A = 1| X) + \frac{(-1)^2}{\Pr( A = -1| X)} \Pr( A = -1| X) \right] \balpha^T X \\ =& \left[ 1 - 1 \right] \cdot \bbeta^T X + \left[ 1 + 1 \right] \balpha^T X \\ =& 2 \balpha^T X \end{aligned} \]
So the first term can still be canceled out. Then, our job simply involves one more step: estimate the propensity scores \(\Pr( A | X)\), using, e.g. logistic regression or random forests, and then plug them into the previous linear regression estimation using the inverse subject weights \(1/\widehat{\Pr}(A = A_i | X = X_i)\).
4 CATE Estimation with Causal Forests
We now use grf for CATE estimation; its splitting rule will be studied further in the Module 3 generalized random forest lecture. Under consistency, conditional unconfoundedness, and overlap, the CATE is identified by the following orthogonal residual moment:
\[ \mathbb{E}\left[ (A - e(X)) \left\{Y - m(X) - \tau(X)(A - e(X))\right\} \mid X \right] = 0, \]
where \(e(X) = \mathbb{P}(A=1 \mid X)\) is the propensity score and
\[ m(X) = \mathbb{E}[Y \mid X] = e(X) \mathbb{E}[Y \mid X, A=1] + (1 - e(X)) \mathbb{E}[Y \mid X, A=0] \]
is the marginal outcome regression that does not condition on treatment. It can be estimated by directly regressing \(Y\) on \(X\) while ignoring the treatment label. To see why the moment identifies \(\tau(X)\), write \(\mu_a(X)=\mathbb{E}[Y\mid X,A=a]\). Then
\[ \begin{aligned} \mathbb{E}\left[(A-e(X))(Y-m(X))\mid X\right] &=\mathbb{E}\left[(A-e(X))Y\mid X\right] \\ &=e(X)(1-e(X))\left\{\mu_1(X)-\mu_0(X)\right\}. \end{aligned} \]
The first equality uses \(\mathbb{E}[A-e(X)\mid X]=0\), while conditional unconfoundedness gives \(\mu_1(X)-\mu_0(X)=\tau(X)\). Also,
\[ \mathbb{E}\left[(A-e(X))^2\mid X\right]=e(X)(1-e(X)). \]
Thus the residual moment equals zero at the CATE. The grf package estimates \(e(X)\) and \(m(X)\), then fits a forest-weighted local regression of \(Y-\hat m(X)\) on \(A-\hat e(X)\) to estimate \(\tau(x)\).
Three related ideas should be distinguished here. Orthogonalization makes the moment’s first-order behavior insensitive to small nuisance-estimation errors. Cross-fitting fits \(e(X)\) and \(m(X)\) on training folds and evaluates the score on held-out observations; grf uses out-of-bag nuisance predictions to achieve a similar separation. Honesty is a forest construction strategy that uses different observations to choose splits and estimate effects within leaves. The following code uses the grf package with honest tree construction. In the code, \(W\) denotes the same binary treatment indicator as \(A\) in the equations.
library(grf)
# Simulated data
n <- 500
p <- 5
X <- matrix(rnorm(n * p), n, p)
W <- rbinom(n, 1, 0.5)
tau <- 1.5 * X[,1] # the true CATE function
Y <- X[,2] + X[, 3] + W * tau + rnorm(n)
par(mfrow=c(1,2))
# plot data with two regression lines
plot(X[,1], Y, col = ifelse(W==1, "deepskyblue", "darkorange"), pch=16, xlab="X1", ylab="Y")
abline(lm(Y[W==1] ~ X[W==1,1]), col="deepskyblue", lwd=2)
abline(lm(Y[W==0] ~ X[W==0,1]), col="darkorange", lwd=2)
# Fit a causal forest
cf <- causal_forest(X, Y, W,
honesty = TRUE,
honesty.fraction = 0.5)
# Estimate conditional treatment effects
cf.pred <- predict(cf, estimate.variance = TRUE)
tau_hat <- cf.pred$predictions
# plot the estimated CATE against the true CATE
plot(X[,1], tau_hat, pch=16, xlab="X1", ylab="Estimated CATE", col="black")
points(X[,1], tau, col="red", pch=16, cex = 0.5)
legend("bottomright", legend=c("Estimated CATE", "True CATE"),
col=c("black", "red"), pch=16, cex = 1.5)
An interesting view of this method, similar to the virtual twins method, is that its splitting rule aims to focus on treatment effect heterogeneity rather than prognostic variation alone. This can be seen from the variable importance plot below. Here, only the first variable \(X_1\) affects the CATE, while the second and third variables affect the regression function but not the CATE. However, due to the estimation of \(m(X)\), this causal forest is still an indirect learning method.
variable_importance(cf)
## [,1]
## [1,] 0.74158514
## [2,] 0.05624845
## [3,] 0.06119734
## [4,] 0.07624305
## [5,] 0.064726025 The Meta-Learner View
A popular framework of CATE estimation is the so-called meta-learner proposed by Künzel et al. (2019). The idea is to decompose the CATE estimation problem into several standard supervised learning problems, and then combine the results to get the final CATE estimate. Several meta-learners are proposed in their paper, most notably the X-learner. At this point, the ideas of these estimators should be relatively straight forward based on what we have discussed so far. Hence, let’s just look at the construction of the X-learner. And here are its steps:
Fit Outcome models: \[ \hat{\mu}_1(x) = \E[Y \mid X=x, A=1], \quad \hat{\mu}_0(x) = \E[Y \mid X=x, A=0]. \]
Create augmented treatment effect data for each individual:
- For treated individuals (\(A_i=1\)): \[ \tilde{D}_i^{(1)} = Y_i - \hat{\mu}_0(X_i) \]
- For control individuals (\(A_i=0\)): \[ \tilde{D}_i^{(0)} = \hat{\mu}_1(X_i) - Y_i \]
Fit separate CATE models for each group: \[ \hat{\tau}_1(x) = \E[\tilde{D}^{(1)} \mid X=x, A=1], \quad \hat{\tau}_0(x) = \E[\tilde{D}^{(0)} \mid X=x, A=0]. \]
Combine the two estimates using an estimated propensity score \(e(x) = \Pr(A=1 \mid X=x)\): \[ \hat{\tau}_X(x) = e(x)\hat{\tau}_0(x) + (1 - e(x))\hat{\tau}_1(x). \]
The following code implements the X-learner using random forests via the ranger package. Again, we will use a simulated data set where the true CATE is known for evaluation purposes.
library(ranger)
##
## Attaching package: 'ranger'
## The following object is masked from 'package:randomForest':
##
## importance
# 1) Outcome models: mu1(x) = E[Y|A=1,X], mu0(x) = E[Y|A=0,X]
idx1 <- which(W == 1L)
idx0 <- which(W == 0L)
df_t1 <- data.frame(y = Y[idx1], x = X[idx1, , drop = FALSE])
df_t0 <- data.frame(y = Y[idx0], x = X[idx0, , drop = FALSE])
mu1_fit <- ranger(y ~., data = df_t1)
mu0_fit <- ranger(y ~., data = df_t0)
# 2) Impute pseudo-effects
# Treated: D^(1) = Y - mu0(X); Control: D^(0) = mu1(X) - Y
mu0_pred_on_1 <- predict(mu0_fit, data = df_t1)$predictions
mu1_pred_on_0 <- predict(mu1_fit, data = df_t0)$predictions
D1 <- Y[idx1] - mu0_pred_on_1
D0 <- mu1_pred_on_0 - Y[idx0]
# 3) CATE models within each group: tau1(x), tau0(x)
tau1_fit <- ranger(y ~ .,
data = data.frame(x = X[idx1, , drop = FALSE], y = D1))
tau0_fit <- ranger(y ~ .,
data = data.frame(x = X[idx0, , drop = FALSE], y = D0))
# 4.1) Propensity model e(x) = P(A=1|X) using ranger classifier (probability=TRUE)
prop_fit <- ranger(y ~ .,
data = data.frame(x = X, y = as.factor(W)),
probability = TRUE)
e_hat <- prop_fit$predictions[, 2]
# 4.2) Combine: tau_X(x) = e(x)*tau0(x) + (1 - e(x))*tau1(x)
tau1_hat <- predict(tau1_fit, data = data.frame(x = X))$predictions
tau0_hat <- predict(tau0_fit, data = data.frame(x = X))$predictions
tau_hat_X <- e_hat * tau0_hat + (1 - e_hat) * tau1_hat
# 4.3) Evaluation (true CATE known: tau = 1.5*X[,1])
xlearner_mse <- mean((tau_hat_X - tau)^2)
cat(sprintf("X-learner (ranger) in-sample CATE MSE: %.4f\n", xlearner_mse))
## X-learner (ranger) in-sample CATE MSE: 0.2597
# (Optional) Simple diagnostic plot against the driver X1
plot(X[,1], tau_hat_X, pch = 16, xlab = "X1", ylab = "Estimated CATE (X-learner, ranger)")
points(X[,1], tau, col = "red", pch = 16, cex = 0.6)
legend("topleft", legend = c("Estimated CATE", "True CATE"),
col = c("black", "red"), pch = 16, bty = "n")
Following similar ideas, the DR-learner was proposed in Kennedy (2023), based on the doubly robust idea we have introduced before. The goal is to construct a pseudo-outcome whose conditional mean remains \(\tau(x)\) if either the outcome model or the propensity model is correctly specified. It should be noted that the theoretical properties of these meta-learners rely on the concept of Neyman orthogonality, which is a core tool used in modern causal inference with machine learning plug-ins. We will not go into the details of this topic in this course, but interested readers may refer to Chernozhukov et al. (2018) for more information.