Causal Inference IX(d) — Policy Learning (R companion)¶
grf + policytree — doubly-robust policy trees, the Athey-Wager way¶
R's policytree (Sverdrup, Kanodia, Zhou, Athey & Wager) is the reference implementation, built to consume grf doubly-robust scores. This companion reproduces the Python notebook: with a treatment cost, treating everyone destroys value; the targeting curve locates the optimal treated fraction; and a shallow policy_tree learns an interpretable rule that nearly matches the oracle. Same simulation (known $\tau(x)$), so the oracle and regret are computable.
1. Setup — treat-all destroys value¶
Heterogeneous effect $\tau(x)=2.5x_1-1$ (negative to positive), treatment cost $c=0.4$ above the average effect. A policy's value gain (over treating no one) is $\mathbb E[\pi(X)(\tau(X)-c)]$. Treating everyone is negative here; the oracle targeted policy ($\tau>c$) is strongly positive.
suppressMessages({library(grf); library(policytree)})
gen<-function(seed,n=4000){set.seed(seed); X<-matrix(runif(n*5),n,5); tau<-2.5*X[,1]-1
W<-rbinom(n,1,plogis(0.6*X[,2])); Y<-X[,3]+sin(3*X[,4])+W*tau+rnorm(n,0,0.5); list(X=X,W=W,Y=Y,tau=tau)}
tr<-gen(0); te<-gen(1); cost<-0.4
val<-function(pol,tau) mean(pol*(tau-cost))
oracle<-val(as.integer(te$tau>cost),te$tau); treatall<-val(rep(1,length(te$tau)),te$tau)
cat(sprintf("true ATE = %.2f; cost c = %.1f\n", mean(te$tau), cost))
cat(sprintf("treat EVERYONE: value %+.3f (NEGATIVE) treat NONE: 0.000\n", treatall))
cat(sprintf("ORACLE (tau>c): value %+.3f, treats %.0f%%\n", oracle, 100*mean(te$tau>cost)))
Warning message: "package 'grf' was built under R version 4.6.1"
Warning message: "package 'policytree' was built under R version 4.6.1"
true ATE = 0.24; cost c = 0.4
treat EVERYONE: value -0.164 (NEGATIVE) treat NONE: 0.000
ORACLE (tau>c): value +0.244, treats 43%
2. The targeting (Qini) curve¶
Rank test units by the causal_forest CATE estimate, treat the top fraction, and plot cumulative value gain. It peaks at the optimal treated fraction; the area over the random-targeting line is the Qini coefficient.
cf<-causal_forest(tr$X, tr$Y, tr$W)
th<-predict(cf, te$X)$predictions
ord<-order(th, decreasing=TRUE); n<-length(th); fr<-(1:n)/n
gain<-cumsum((te$tau-cost)[ord])/n; rand<-fr*mean(te$tau-cost)
qini<-sum((gain-rand))*mean(diff(fr))
pk<-max(gain); pkf<-fr[which.max(gain)]
cat(sprintf("targeting peak value %.3f at fraction %.2f (oracle %.3f); Qini coefficient %.3f\n", pk, pkf, oracle, qini))
options(repr.plot.width=8.5, repr.plot.height=5)
plot(fr, gain, type="l", col="#2f855a", lwd=2.5, xlab="fraction treated (highest CATE first)", ylab="cumulative value gain", main=sprintf("Targeting / Qini curve (Qini = %.3f)", qini))
lines(fr, rand, col="grey", lwd=2, lty=2); abline(h=0, col="black", lwd=.6); abline(v=pkf, col="#2b6cb0", lty=3)
points(pkf, pk, col="#c53030", pch=19, cex=1.3)
legend("bottomleft", c("targeting by CATE","random targeting",sprintf("optimal frac %.2f",pkf)), col=c("#2f855a","grey","#2b6cb0"), lwd=2, lty=c(1,2,3), bty="n")
targeting peak value 0.243 at fraction 0.44 (oracle 0.244); Qini coefficient 0.211
3. The policy tree — policytree::policy_tree on doubly-robust scores¶
double_robust_scores(cf) returns the DR reward of each treatment; we form the cost-adjusted reward matrix (value of treating $=\hat\Gamma_{\text{treat}}-\hat\Gamma_{\text{control}}-c$, value of not treating $=0$) and fit a depth-2 policy_tree. It learns a shallow, interpretable rule maximizing estimated value — here nearly matching the oracle while remaining auditable. The tree's split reveals the covariate driving who benefits.
dr<-double_robust_scores(cf)
Gamma<-cbind(`do not treat`=0, treat=dr[,2]-dr[,1]-cost)
pt<-policy_tree(tr$X, Gamma, depth=2)
rec<-predict(pt, te$X)-1 # action index {1,2} -> {0,1}
cat(sprintf("policy_tree (depth 2): value %+.3f, treats %.0f%% (oracle %+.3f)\n", val(rec,te$tau), 100*mean(rec), oracle))
cat(sprintf("threshold policy (CATE>c): value %+.3f\n\n", val(as.integer(th>cost),te$tau)))
cat("Learned policy tree:\n"); print(pt)
options(repr.plot.width=8, repr.plot.height=4)
vals<-c(`treat all`=treatall, `threshold`=val(as.integer(th>cost),te$tau), `policy tree`=val(rec,te$tau), oracle=oracle)
bp<-barplot(vals, col=c("#c53030","#2b6cb0","#2f855a","black"), ylab="policy value gain", main="Interpretable policy tree nearly matches the oracle")
abline(h=0); text(bp, vals+ifelse(vals>=0,0.008,-0.02), sprintf("%+.2f",vals))
policy_tree (depth 2): value +0.241, treats 42% (oracle +0.244)
threshold policy (CATE>c): value +0.242
Learned policy tree:
policy_tree object
Tree depth: 2
Actions: 1: do not treat 2: treat
Variable splits:
(1) split_variable: X2 split_value: 0.504759
(2) split_variable: X1 split_value: 0.620587
(4) * action: 1
(5) * action: 2
(3) split_variable: X1 split_value: 0.521261
(6) * action: 1
(7) * action: 2
4. Summary¶
grf + policytree reproduced the decision layer: with a treatment cost, treating everyone destroyed value, the targeting (Qini) curve located the optimal treated fraction, and a depth-2 policy_tree learned an interpretable, doubly-robust rule that nearly matched the oracle. This is the standard modern workflow — estimate CATE with causal_forest, form doubly-robust scores, and learn a deployable policy with policy_tree.
The takeaway for applied work: a CATE model is only useful once turned into a policy that is evaluated honestly — against treat-all/treat-none and the budget, via a targeting curve and out-of-sample policy value. policy_tree gives a rule a program can actually adopt and audit. Cross-links: the decision layer atop the Causal Forest and meta-learners (subsection 9); doubly-robust scores from the AIPW/DML estimators; the Qini curve as the causal analogue of AUC/uplift evaluation. This completes the depth of the Heterogeneous-Effects subsection.