Getting started with causalmlr

ATE estimation

The sodium dataset simulates the effect of sodium intake on blood pressure (Luque-Fernandez et al., 2019). Age confounds both the treatment and the outcome, and the true ATE is 1.05.

data(sodium)
head(sodium)
#>   age sodium       bp
#> 1  73      1 149.0206
#> 2  67      0 134.0011
#> 3  69      1 141.6555
#> 4  76      1 153.8871
#> 5  74      1 147.2216
#> 6  60      0 120.3477

The naive difference in means ignores the confounder and is heavily biased:

naive <- ate_naive(sodium, outcome = "bp", treatment = "sodium")
naive
#> ATE estimate - Naive (difference in means) 
#>   Estimate:  5.18
#>   Std. error: 0.1948 
#>   95% CI:    [4.798, 5.561]
#>   N:         10000

Adjusting for age removes (most of) the bias. Inverse propensity weighting (IPW) models the treatment assignment, the doubly robust estimator (AIPW) additionally models the outcome, and double machine learning (DML) uses orthogonalised residuals with cross-fitting:

ipw <- ate_ipw(sodium, outcome = "bp", treatment = "sodium",
               ps_learner = lrn("classif.rpart"))

dr <- ate_dr(sodium, outcome = "bp", treatment = "sodium",
             outcome_learner = lrn("regr.rpart"),
             ps_learner = lrn("classif.rpart"),
             folds = 5, ps_trim = 0.01)

dml <- ate_dml(sodium, outcome = "bp", treatment = "sodium",
               outcome_learner = lrn("regr.rpart"),
               ps_learner = lrn("classif.rpart"),
               folds = 5, ps_trim = 0.01)
dml
#> ATE estimate - Double machine learning 
#>   Estimate:  1.12
#>   Std. error: 0.0429 
#>   95% CI:    [1.036, 1.204]
#>   N:         10000

ps_trim = 0.01 clips the estimated propensity scores to \([0.01, 0.99]\). Tree-based propensity models can output probabilities of exactly 0 or 1, which would make the inverse weights explode - trimming is cheap insurance.

Because the true ATE is known here, we can compare the estimators with the absolute ATE error, eps_ate():

true_ate <- 1.05
sapply(list(naive = naive, ipw = ipw, dr = dr, dml = dml),
       function(est) eps_ate(true_ate, est))
#>      naive        ipw         dr        dml 
#> 4.12969203 1.41636915 0.08537609 0.07033070

CATE estimation with meta-learners

The synth_train / synth_test data follow the data generating process of Künzel et al. (2019). The test set holds only covariates plus the true effect tau, so we can evaluate how well each meta-learner recovers the heterogeneous effects.

data(synth_train)
data(synth_test)
head(synth_train)
#>           x0          x1          x2         x3         x4 t          y
#> 1 -2.1966546  1.35390014 -1.39801445 -1.7138880 -0.6693249 1 12.1009267
#> 2  0.6193071 -0.03483849 -0.03802601  0.1585317  1.0855784 1  0.4292407
#> 3  0.1146665  0.57386661  0.53158533 -1.3231972 -1.3461593 0  1.0622299
#> 4 -0.7829745  1.20219763 -1.04630459 -1.0525001  0.2200478 0  2.4699487
#> 5  0.8731837  1.32112750  0.22711515 -0.5524267  1.0754895 1  8.4183230
#> 6  2.3052394  0.32871172  0.20735839  1.1505775  1.4527570 1  8.5148841

m_s <- s_learner(synth_train, outcome = "y", treatment = "t",
                 learner = lrn("regr.rpart"))

m_t <- t_learner(synth_train, outcome = "y", treatment = "t",
                 learner = lrn("regr.rpart"))

m_x <- x_learner(synth_train, outcome = "y", treatment = "t",
                 learner = lrn("regr.rpart"),
                 ps_learner = lrn("classif.rpart"))

m_dr <- dr_learner(synth_train, outcome = "y", treatment = "t",
                   outcome_learner = lrn("regr.rpart"),
                   ps_learner = lrn("classif.rpart"),
                   ps_trim = 0.01)

m_r <- r_learner(synth_train, outcome = "y", treatment = "t",
                 outcome_learner = lrn("regr.rpart"),
                 ps_learner = lrn("classif.rpart"),
                 ps_trim = 0.01)

The dr_learner() (Kennedy, 2023) is doubly robust: it cross-fits the propensity and outcome nuisances into an AIPW pseudo-outcome and regresses that on the covariates, so it stays consistent if either nuisance model is correct.

The r_learner() (Nie & Wager, 2021) is the heterogeneous-effect version of the partially linear DML estimator ate_dml(): it cross-fits the outcome and propensity residuals and fits a weighted regression that directly minimises the R-Loss (see below). It falls back to an unweighted regression when the second-stage learner does not support observation weights.

predict() returns a vector of CATE estimates, one per row; ate() averages them:

tau_s <- predict(m_s, synth_test)
head(tau_s)
#> [1] 0 0 0 0 0 0
ate(m_s, synth_test)
#> [1] 3.737132

With ground-truth effects available, the standard evaluation metric is PEHE (precision in estimation of heterogeneous effects) - the RMSE between true and predicted CATEs:

c(S = pehe(synth_test$tau, tau_s),
  T = pehe(synth_test$tau, predict(m_t, synth_test)),
  X = pehe(synth_test$tau, predict(m_x, synth_test)),
  DR = pehe(synth_test$tau, predict(m_dr, synth_test)),
  R = pehe(synth_test$tau, predict(m_r, synth_test)))
#>         S         T         X        DR         R 
#> 1.4743224 1.8315454 0.8750192 0.6746229 0.3699063

Confidence intervals for predicted CATEs

predict() returns point estimates only. Unlike the ATE estimators - which carry a closed-form standard error and a confint() method - there is no analytic standard error for \(\hat\tau(x)\) when the meta-learner wraps an arbitrary machine-learning base learner. cate_ci() instead uses the nonparametric bootstrap: it resamples the training rows, refits the entire pipeline (nuisance models and cross-fitting included), and re-predicts on the supplied data n_boot times, then forms a pointwise interval from the bootstrap distribution at each row. It works the same way for all five meta-learners, so you pass the fitted model and the original training data:

set.seed(1)
ci <- cate_ci(m_t, synth_test, train_data = synth_train, n_boot = 100)
head(ci)
#>     estimate        se       lower     upper
#> 1 -1.2920574 1.3449870 -2.85874442 1.9112984
#> 2 -0.2280423 2.0505628 -4.08871768 2.7944808
#> 3  1.5599339 0.9771083 -0.09834918 3.4752596
#> 4 -0.4125864 1.4170369 -2.34349639 3.1777231
#> 5  0.1015885 1.4697517 -2.22515306 3.4997011
#> 6 -2.9193186 1.4292442 -4.94250806 0.5162985

The result has one row per row of newdata, with the point estimate, the bootstrap standard error se, and the interval limits lower/upper. The default type = "percentile" uses the empirical quantiles of the bootstrap CATEs; type = "normal" centres a symmetric interval on the point estimate instead.

Two caveats. First, refitting the whole pipeline n_boot times is expensive - 200 replicates (the default) means 200 refits - so budget accordingly and set the seed for reproducibility. Second, bootstrap coverage for CATEs estimated with flexible learners is only approximate and can be anti-conservative; read these intervals as a measure of estimation stability rather than exact frequentist coverage.

Model selection without ground truth: the R-Loss

Real observational data has no tau column, so PEHE cannot be computed. The R-Loss (Nie & Wager, 2021) is an observable score that only needs two nuisance models - the conditional outcome mean \(m(x) = E[Y \mid X]\) and the propensity score \(e(x) = P(T = 1 \mid X)\):

\[\tau\text{-risk}_R = \frac{1}{n} \sum_i \Big[ \big(Y^{(i)} - \hat m(X^{(i)})\big) - \big(T^{(i)} - \hat e(X^{(i)})\big)\,\hat\tau(X^{(i)}) \Big]^2\]

A typical workflow: hold out a validation set, fit the nuisance models on it once (cross-fitted by default), then score any number of candidate CATE models. Here we tune the maxdepth hyperparameter of an S-Learner:

idx <- sample(nrow(synth_train), 0.7 * nrow(synth_train))
train <- synth_train[idx, ]
valid <- synth_train[-idx, ]

nuis <- rloss_nuisance(valid, outcome = "y", treatment = "t",
                       outcome_learner = lrn("regr.rpart"),
                       ps_learner = lrn("classif.rpart"),
                       folds = 5)
#> Warning: Extreme propensity scores detected (near 0 or 1); estimates may be
#> unstable. Consider setting `ps_trim`.

scores <- sapply(1:5, function(md) {
  m <- s_learner(train, outcome = "y", treatment = "t",
                 learner = lrn("regr.rpart", maxdepth = md))
  r_loss(nuis, predict(m, valid))
})
data.frame(maxdepth = 1:5, r_loss = scores)
#>   maxdepth   r_loss
#> 1        1 9.467719
#> 2        2 5.423057
#> 3        3 5.733959
#> 4        4 5.683873
#> 5        5 5.683873

The candidate with the lowest R-Loss is then refit on the full training data:

best_md <- which.min(scores)
m_best <- s_learner(synth_train, outcome = "y", treatment = "t",
                    learner = lrn("regr.rpart", maxdepth = best_md))
pehe(synth_test$tau, predict(m_best, synth_test))
#> [1] 0.06171868

Note that the same recipe compares different estimators too (S vs T vs X vs DR, or different base learners) - the R-Loss only sees the CATE predictions.

Tuning nuisance models directly with mlr3tuning

Because every estimator accepts plain mlr3 learners, the alternative “direct” tuning strategy - tuning each nuisance model on its own supervised objective - composes naturally with mlr3tuning. An AutoTuner is itself a learner, so it can be passed anywhere a learner is expected:

library(mlr3tuning)

at <- auto_tuner(
  tuner = tnr("grid_search", resolution = 5),
  learner = lrn("regr.rpart", maxdepth = to_tune(1, 10)),
  resampling = rsmp("cv", folds = 5),
  measure = msr("regr.mse")
)

# the base learner is tuned internally when the S-Learner fits it
m_tuned <- s_learner(synth_train, outcome = "y", treatment = "t",
                     learner = at)

Keep in mind that the direct strategy optimises a predictive objective (e.g. MSE), which is not the same as optimising causal quality - the R-Loss workflow above targets the causal estimates themselves.

Included datasets

References

• Chernozhukov, V. et al. (2018). Double/debiased machine learning for treatment and structural parameters. The Econometrics Journal.
• Hill, J. L. (2011). Bayesian nonparametric modeling for causal inference. JCGS.
• Kennedy, E. H. (2023). Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics.
• Künzel, S. R. et al. (2019). Metalearners for estimating heterogeneous treatment effects using machine learning. PNAS.
• Machlanski, D., Samothrakis, S., & Clarke, P. (2023). Hyperparameter tuning and model evaluation in causal effect estimation. arXiv.
• Nie, X., & Wager, S. (2021). Quasi-oracle estimation of heterogeneous treatment effects. Biometrika.