Lab 9: Policy Trees

This lab moves from evaluating heterogeneity (Lab 8) to making decisions. Policy trees learn simple, interpretable treatment assignment rules from causal forest predictions. You will fit policy trees of different depths, evaluate their performance, and discuss the ethical implications of algorithmic treatment assignment.

Setup

library(causalworkshop)
library(grf)
library(policytree)
library(tidyverse)

Re-fit the causal forest (or re-use from previous labs):

# simulate data
d <- simulate_nzavs_data(n = 5000, seed = 2026)
d0 <- d |> filter(wave == 0)
d1 <- d |> filter(wave == 1)
d2 <- d |> filter(wave == 2)

# construct matrices
covariate_cols <- c(
  "age", "male", "nz_european", "education", "partner", "employed",
  "log_income", "nz_dep", "agreeableness", "conscientiousness",
  "extraversion", "neuroticism", "openness",
  "community_group", "wellbeing"
)

X <- as.matrix(d0[, covariate_cols])
Y <- d2$wellbeing
W <- d1$community_group

cf <- causal_forest(
  X, Y, W,
  num.trees = 1000,
  honesty = TRUE,
  tune.parameters = "all",
  seed = 2026
)

tau_hat <- predict(cf)$predictions

The gamma matrix

A policy tree needs a reward matrix (called the "gamma matrix"). Each row is an individual; each column is an action. The entry gives the expected reward for assigning that individual to that action.

With two actions (treat vs not treat), the gamma matrix has two columns:

# construct gamma matrix
# column 1: reward if not treated (control) = 0 (baseline)
# column 2: reward if treated = predicted treatment effect
gamma_matrix <- cbind(
  control = rep(0, length(tau_hat)),
  treatment = tau_hat
)

head(gamma_matrix)

Fit a depth-1 policy tree

A depth-1 tree makes a single split, dividing the population into two groups based on one covariate:

# subsample for speed (policy tree fitting can be slow on large datasets)
set.seed(2026)
n_sample <- 500
idx <- sample(seq_len(nrow(X)), n_sample)

X_sample <- as.data.frame(X[idx, ])
gamma_sample <- gamma_matrix[idx, ]

# fit depth-1 policy tree
pt_depth1 <- policy_tree(X_sample, gamma_sample, depth = 1)

# print the tree
print(pt_depth1)

Visualise the tree:

plot(pt_depth1)

Fit a depth-2 policy tree

A depth-2 tree makes two sequential splits, creating four groups:

# fit depth-2 policy tree
pt_depth2 <- policy_tree(X_sample, gamma_sample, depth = 2)

# print and plot
print(pt_depth2)
plot(pt_depth2)

Evaluate policies

Predict treatment assignments on the full dataset and compare with random assignment:

# predict actions for full dataset
X_full <- as.data.frame(X)
actions_depth1 <- predict(pt_depth1, X_full)
actions_depth2 <- predict(pt_depth2, X_full)

# compute expected reward under each policy
# action = 1 means control, action = 2 means treatment
reward_depth1 <- ifelse(actions_depth1 == 1, gamma_matrix[, 1], gamma_matrix[, 2])
reward_depth2 <- ifelse(actions_depth2 == 1, gamma_matrix[, 1], gamma_matrix[, 2])
reward_random <- 0.5 * gamma_matrix[, 1] + 0.5 * gamma_matrix[, 2]

# compare policies
policy_comparison <- tibble(
  policy = c("Random assignment", "Depth-1 tree", "Depth-2 tree", "Treat everyone"),
  expected_reward = c(
    mean(reward_random),
    mean(reward_depth1),
    mean(reward_depth2),
    mean(tau_hat)
  ),
  treat_rate = c(
    0.50,
    mean(actions_depth1 == 2),
    mean(actions_depth2 == 2),
    1.00
  )
)

print(policy_comparison |> mutate(across(where(is.numeric), \(x) round(x, 3))))

Interpret the rules in plain language

Read the tree output and translate it into a decision rule:

# what variables and thresholds does the depth-2 tree use?
print(pt_depth2)

# example interpretation (your values will differ):
# "treat individuals with extraversion > 0.3 and baseline wellbeing < 0.1"

Compare policy assignments with actual treatment

How do the policy-recommended assignments compare with who actually received treatment in the data?

# agreement between policy and observed treatment
agreement <- tibble(
  actual = W,
  policy_depth2 = ifelse(actions_depth2 == 2, 1, 0)
) |>
  mutate(agree = actual == policy_depth2)

cat("Agreement rate:", round(mean(agreement$agree), 3), "\n")
cat("Policy treats: ", round(mean(agreement$policy_depth2), 3), "\n")
cat("Actual treated:", round(mean(agreement$actual), 3), "\n")

Ethical considerations

Exercises

1. Different outcome. Fit a policy tree for religious_service on belonging. Do the splitting variables change? What does this suggest about which covariates drive effect modification for different outcomes?

2. Depth-3 tree. Fit a depth = 3 policy tree. Does the expected reward improve substantially over depth-2? Is the tree still interpretable?

3. Discuss override. In one paragraph, describe a scenario where a clinician should override a policy tree recommendation. What information would the clinician have that the model does not?