Lab: Causal Forests for Heterogeneous Treatment Effects

Overview

In this lab you will analyze a simulated randomized controlled trial where the treatment effect varies substantially across individuals. Rather than estimating a single average treatment effect (ATE), you will use causal forests to estimate the conditional average treatment effect (CATE) as a function of covariates, identify which variables drive treatment effect heterogeneity, and design an optimal treatment targeting policy.

What you will learn:

• How to estimate heterogeneous treatment effects with causal forests
• How to assess variable importance for treatment effect heterogeneity
• How to evaluate CATE estimates using calibration tests and RATE curves
• How to design and evaluate optimal targeting policies
• How to compare causal forests with simple subgroup analysis

Prerequisites: Familiarity with random forests and basic causal inference. Completion of the OLS and DML tutorial labs is recommended.

Step 1: Simulate an RCT with Heterogeneous Effects

We create an experiment where the treatment effect depends on age and baseline risk.

1library(grf)
2
3set.seed(42)
4n <- 4000
5
6age <- runif(n, 25, 65)
7income <- rlnorm(n, 10.5, 0.6)
8educ <- pmin(pmax(rnorm(n, 14, 3), 8), 22)
9health_score <- pmin(pmax(rnorm(n, 50, 15), 0), 100)
10female <- rbinom(n, 1, 0.5)
11risk <- rnorm(n)
12
13W <- rbinom(n, 1, 0.5)
14
15tau_true <- 5 + 3 * (risk > 0) - 0.1 * (age - 40) + 2 * (educ > 16)
16mu0 <- 50 + 0.5 * age + 0.3 * health_score - 2 * risk + rnorm(n, 0, 5)
17Y <- mu0 + W * tau_true
18
19X <- data.frame(age, income, educ, health_score, female, risk)
20
21cat("True ATE:", mean(tau_true), "\n")
22cat("CATE range:", range(tau_true), "\n")

Expected output:

True ATE: ~6.00
True CATE range: [~0.0, ~12.0]
Treatment rate: ~50%
Covariates that drive heterogeneity: age, risk, educ

Step 2: Estimate the Average Treatment Effect

First, confirm the overall ATE before looking at heterogeneity.

1# Simple difference in means
2ate_simple <- mean(Y[W == 1]) - mean(Y[W == 0])
3
4cat("Diff in means:", ate_simple, "\n")
5
6# Covariate-adjusted
7ols <- lm(Y ~ W + age + income + educ + health_score + female + risk)
8cat("OLS-adjusted:", coef(ols)["W"], "\n")
9cat("True ATE:", mean(tau_true), "\n")

Expected output:

Difference in means: ~6.00
OLS-adjusted ATE:    ~6.00 (SE: ~0.16)
True ATE:            ~6.00

The covariate-adjusted estimator has a much smaller standard error because adjusting for strong predictors of the outcome (age, health_score, risk) reduces residual variance.

Step 3: Estimate CATEs with a Causal Forest

1# Causal forest using grf
2X_mat <- as.matrix(X)
3
4cf <- causal_forest(X_mat, Y, W,
5                   num.trees = 2000,
6                   min.node.size = 5,
7                   seed = 42)
8
9# Predicted CATEs
10tau_hat <- predict(cf)$predictions
11
12# Evaluation
13cat("Correlation:", cor(tau_true, tau_hat), "\n")
14cat("RMSE:", sqrt(mean((tau_true - tau_hat)^2)), "\n")
15cat("Estimated ATE:", mean(tau_hat), "\n")
16
17# ATE with confidence interval from forest
18ate_cf <- average_treatment_effect(cf)
19cat("Forest ATE:", ate_cf[1], " SE:", ate_cf[2], "\n")

Expected output:

Correlation(true CATE, estimated CATE): ~0.85
RMSE:                                   ~1.2
Estimated ATE (mean of CATEs):          ~6.00
True ATE:                                ~6.00

The correlation of ~0.85 indicates the causal forest successfully identifies who benefits more vs. less from treatment, even though individual-level estimates are noisy.

Step 4: Variable Importance and Heterogeneity Drivers

1# Variable importance from grf
2vi <- variable_importance(cf)
3rownames(vi) <- colnames(X)
4vi_sorted <- sort(vi[,1], decreasing = TRUE)
5barplot(vi_sorted, horiz = TRUE, main = "Variable Importance for Heterogeneity",
6      xlab = "Importance")
7
8# CATE by subgroup
9cat("\nCATEs by risk group:\n")
10cat("Low risk:", mean(tau_hat[risk <= 0]), "(true:", mean(tau_true[risk <= 0]), ")\n")
11cat("High risk:", mean(tau_hat[risk > 0]), "(true:", mean(tau_true[risk > 0]), ")\n")

Expected output:

Expected output: CATEs by age group

CATEs decline with age, consistent with the DGP: tau = 5 + 3*(risk > 0) - 0.1*(age - 40) + 2*(educ > 16).

Step 5: Calibration and the RATE Curve

Evaluate whether the CATE estimates actually predict treatment effect heterogeneity.

1# Calibration test from grf
2test_calibration <- test_calibration(cf)
3print(test_calibration)
4
5# RATE curve
6rate <- rank_average_treatment_effect(cf, tau_true)
7plot(rate, main = "RATE Curve")
8
9# Best linear projection
10blp <- best_linear_projection(cf, X_mat)
11print(blp)

Expected output: CATE calibration by quintile

Calibration is good: predicted CATEs closely track both the observed and true effects across quintiles, confirming the forest identifies genuine heterogeneity.

Step 6: Optimal Targeting Policy

Use CATE estimates to design a policy that treats only those who benefit most.

1# Policy: treat top 50%
2threshold <- median(tau_hat)
3policy <- as.integer(tau_hat >= threshold)
4
5cat("Average effect, targeted:", mean(tau_true[policy == 1]), "\n")
6cat("Average effect, not targeted:", mean(tau_true[policy == 0]), "\n")
7cat("Average effect, treating all:", mean(tau_true), "\n")
8
9# Compare with simple subgroup
10simple_policy <- as.integer(risk > 0)
11cat("\nSimple rule benefit:", mean(tau_true[simple_policy == 1]), "\n")
12cat("Forest targeting benefit:", mean(tau_true[policy == 1]), "\n")

Expected output:

=== Optimal Targeting Policy (treat top 50%) ===
Average effect for those targeted:      ~7.8
Average effect for those NOT targeted:  ~4.2
Average effect if treating everyone:    ~6.0

Gain from targeting: ~1.8 per treated individual
Fraction correctly identified (true top 50%): ~80%

Simple rule (risk > 0) benefit:         ~7.5
Forest-based targeting benefit:         ~7.8

The forest-based targeting outperforms the simple rule because it combines information from risk, age, and education simultaneously, while the simple rule uses only one variable.

Exercises

1. Compare ML methods. Estimate CATEs using a causal forest, a T-learner (separate random forests for treated and control), and a linear interaction model. Which has the best RMSE for the true CATEs?

2. Out-of-sample validation. Split the data 50/50. Train the causal forest on the first half and evaluate CATE predictions on the second half using the calibration test.

3. Budget-constrained targeting. Suppose you can only treat 25% of the population. Design the optimal targeting policy and compute the expected average effect.

4. Nonlinear heterogeneity. Modify the DGP so that the treatment effect has a U-shape in age. Does the causal forest capture this pattern?

Summary

In this lab you learned:

• Causal forests estimate the conditional average treatment effect (CATE) as a function of covariates, moving beyond the ATE
• The method uses honest splitting (separate samples for determining splits and estimating effects) to produce valid confidence intervals
• Variable importance reveals which covariates drive treatment effect heterogeneity
• Calibration tests and RATE curves assess whether estimated CATEs are predictive of actual treatment effect heterogeneity
• Optimal targeting policies assign treatment to individuals with the highest predicted CATEs, potentially improving welfare
• CATE estimates from observational data should be validated before informing policy, as heterogeneity may reflect differential selection rather than true effect variation