---
title: "Getting started with causalmlr"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Getting started with causalmlr}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
```

`causalmlr` estimates **average treatment effects (ATE)** and **conditional
average treatment effects (CATE)** from observational data using machine
learning. All estimators are built on top of the
[mlr3](https://mlr3.mlr-org.com/) ecosystem, so any mlr3 regression or
classification learner can be plugged in as a nuisance model.

```{r setup, message = FALSE}
library(causalmlr)
library(mlr3)

# silence mlr3's info logging
lgr::get_logger("mlr3")$set_threshold("warn")
set.seed(2025)
```

We use simple `rpart` decision trees throughout this vignette because they
ship with base R. In practice you will usually get better results with
stronger learners, e.g. `lrn("regr.ranger")` and `lrn("classif.ranger")`
from `mlr3learners`.

## 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**.

```{r ate-data}
data(sodium)
head(sodium)
```

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

```{r ate-naive}
naive <- ate_naive(sodium, outcome = "bp", treatment = "sodium")
naive
```

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:

```{r ate-adjusted}
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
```

`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()`:

```{r ate-eval}
true_ate <- 1.05
sapply(list(naive = naive, ipw = ipw, dr = dr, dml = dml),
       function(est) eps_ate(true_ate, est))
```

## 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.

```{r cate-data}
data(synth_train)
data(synth_test)
head(synth_train)
```

```{r cate-fit}
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:

```{r cate-predict}
tau_s <- predict(m_s, synth_test)
head(tau_s)
ate(m_s, synth_test)
```

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

```{r cate-eval}
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)))
```

## 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:

```{r cate-ci}
set.seed(1)
ci <- cate_ci(m_t, synth_test, train_data = synth_train, n_boot = 100)
head(ci)
```

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:

```{r rloss}
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)

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)
```

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

```{r rloss-best}
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))
```

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:

```{r mlr3tuning, eval = FALSE}
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

| Dataset | Type | Ground truth | Task |
|---|---|---|---|
| `sodium` | synthetic | ATE = 1.05 | ATE estimation |
| `synth_train` / `synth_test` | synthetic | `tau` | CATE estimation |
| `ihdp_train` / `ihdp_test` | semi-synthetic | `tau` | CATE estimation |
| `jobs_train` / `jobs_test` | real + experiment flag | `e` (randomised subsample) | policy evaluation |
| `abalone`, `housing` | real | - | regression practice |
| `diabetes`, `spirals` | real / synthetic | - | classification practice |

## 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*.
