Causal Inference IX(d) — Policy Learning: from CATE to optimal treatment rules¶
Targeting curves, the Qini coefficient, and interpretable policy trees¶
The heterogeneous-effects notebooks estimated the conditional average treatment effect $\tau(x)$ — how much the treatment helps each unit. But an estimate is not a decision. Policy learning asks the operational question a program manager actually faces: given a limited budget and a per-treatment cost, whom should we treat? The answer turns a CATE surface into a treatment rule $\pi(x)\in\{0,1\}$, and the payoff can be enormous — because with a real cost, treating everyone often destroys value, while treating the right subgroup creates it.
This notebook builds the decision layer on top of $\tau(x)$:
- The optimal rule and its value. With a treatment cost $c$, the value-maximizing policy is $\pi^\star(x)=\mathbb{1}\{\tau(x)>c\}$, and a policy's value is what it delivers net of cost. We show that treat-all can be worse than treating no one.
- The targeting (Qini) curve. Rank units by estimated $\hat\tau(x)$, treat the top fraction, and trace the cumulative value gained — the standard uplift-modelling diagnostic. Its peak locates the optimal treated fraction; the area over the random-targeting line is the Qini coefficient, a one-number measure of targeting quality.
- Policy trees. A shallow, interpretable tree that directly maximizes estimated policy value (Athey-Wager) — a rule a committee can read and audit, not a black-box threshold on a forest.
- Honest value & regret. Estimate each policy's value and its regret against the oracle.
We use a simulation with a known $\tau(x)$ so the oracle policy and true regret are computable. Python-lead (econml policy trees + from-scratch targeting curves); R companion uses policytree with grf. This is example 4 of the Heterogeneous-Effects subsection, building on the Causal Forest and meta-learner CATE estimators.
1. From CATE to a decision — and why treat-all can destroy value¶
We simulate a job-training-style program: each worker has covariates $x$, a heterogeneous effect $\tau(x)=2.5x_1-1$ that ranges from clearly negative (the program hurts some) to strongly positive, and treatment costs $c=0.4$ per worker. The value-maximizing policy treats a worker only when the benefit exceeds the cost, $\tau(x)>c$; a policy's value gain (over treating no one) is $\mathbb E[\pi(X)\,(\tau(X)-c)]$.
The key number: the average effect (0.25) is below the cost (0.40), so treating everyone yields a negative value gain — it loses money — even though the program helps a large minority a lot. This is exactly why a CATE estimate matters operationally: the average masks the fact that a well-chosen subgroup has effects far above cost. We estimate $\tau(x)$ with a causal forest (from the first notebook) and set up the policies.
import numpy as np, pandas as pd, matplotlib.pyplot as plt, warnings
warnings.filterwarnings("ignore")
from econml.dml import CausalForestDML
from sklearn.ensemble import GradientBoostingRegressor as GBR, GradientBoostingClassifier as GBC
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; GREY="#a0aec0"; PURP="#6b46c1"
def gen(seed,n=4000):
r=np.random.default_rng(seed); X=r.uniform(0,1,(n,5)); tau=2.5*X[:,0]-1.0
e=1/(1+np.exp(-(0.6*X[:,1]))); T=(r.uniform(size=n)<e).astype(int) # mild confounding
Y=X[:,2]+np.sin(3*X[:,3])+T*tau+r.normal(0,0.5,n); return X,T,Y,tau
Xtr,Ttr,Ytr,tautr=gen(0); Xte,Tte,Yte,taute=gen(1); cost=0.4
cf=CausalForestDML(model_y=GBR(n_estimators=100,max_depth=3),model_t=GBC(n_estimators=100,max_depth=3),discrete_treatment=True,cv=3,random_state=0).fit(Ytr,Ttr,X=Xtr)
tauhat=cf.effect(Xte) # estimated CATE on held-out units
val=lambda pol: np.mean(pol*(taute-cost)) # true value gain of a policy (uses known tau)
oracle=val((taute>cost).astype(int)); treatall=val(np.ones(len(taute)))
print(f"true ATE = {taute.mean():.2f}; treatment cost c = {cost}")
print(f"treat EVERYONE : value gain = {treatall:+.3f} <-- NEGATIVE: the average effect is below cost")
print(f"treat NO ONE : value gain = 0.000")
print(f"ORACLE (treat if tau>c) : value gain = {oracle:+.3f} (treats {(taute>cost).mean():.0%})")
fig,ax=plt.subplots(1,2,figsize=(13,4.2))
ax[0].scatter(Xte[:,0],taute,s=6,c=GREY,alpha=.4); ax[0].axhline(cost,color=RED,ls="--",label=f"cost c={cost}")
ax[0].fill_between([0,1],[cost,cost],[taute.max(),taute.max()],color=GREEN,alpha=.08)
ax[0].set_xlabel("covariate x1"); ax[0].set_ylabel("true effect τ(x)"); ax[0].set_title("Treat only where τ(x) exceeds the cost (green)"); ax[0].legend(fontsize=8)
ax[1].bar(["treat all","treat none","oracle\n(τ>c)"],[treatall,0,oracle],color=[RED,GREY,GREEN]); ax[1].axhline(0,color="k",lw=.7)
for i,v in enumerate([treatall,0,oracle]): ax[1].text(i,v+(0.01 if v>=0 else -0.03),f"{v:+.2f}",ha="center")
ax[1].set_ylabel("policy value gain"); ax[1].set_title("Treating everyone loses money; targeting creates value")
plt.tight_layout(); plt.show()
print("With a real cost, the decision is not 'does it work on average' but 'for whom does it beat the cost'. Treating everyone")
print("is value-destroying here; the oracle targeted policy is strongly positive. Policy learning is how we approximate the oracle.")
true ATE = 0.25; treatment cost c = 0.4 treat EVERYONE : value gain = -0.153 <-- NEGATIVE: the average effect is below cost treat NO ONE : value gain = 0.000 ORACLE (treat if tau>c) : value gain = +0.238 (treats 44%)
With a real cost, the decision is not 'does it work on average' but 'for whom does it beat the cost'. Treating everyone is value-destroying here; the oracle targeted policy is strongly positive. Policy learning is how we approximate the oracle.
2. The targeting (Qini) curve — how much value does ranking capture?¶
The workhorse diagnostic of uplift modelling. Rank the units by estimated effect $\hat\tau(x)$, treat the top fraction, and plot the cumulative value gained as that fraction grows. Because we treat the highest-$\hat\tau$ units first, the curve rises while we are adding units whose effect beats the cost, peaks at the optimal treated fraction, then declines as we are forced to add units the treatment does not help enough. The peak height is the value a threshold policy achieves, and the peak location is how many to treat.
Compared against the straight line of random targeting (which reaches the same treat-all endpoint but gains nothing from ordering), the gap measures the value of knowing whom to treat. The area between the two curves is the Qini coefficient — a single number summarizing targeting quality, the causal-policy analogue of AUC. Our estimated $\hat\tau$ ranking recovers a targeting curve that peaks essentially at the oracle value.
order=np.argsort(-tauhat) # rank by estimated effect, descending
n=len(taute); fracs=np.arange(1,n+1)/n
gain=np.cumsum((taute-cost)[order])/n # cumulative value gain treating top-k by tauhat
rand=fracs*(taute-cost).mean() # random targeting: straight line to treat-all
qini=np.trapezoid(gain-rand,fracs)
peak=gain.max(); peakfrac=fracs[gain.argmax()]
print(f"targeting curve peak value = {peak:.3f} at fraction treated = {peakfrac:.2f} (oracle {oracle:.3f})")
print(f"random-targeting endpoint (treat all) = {(taute-cost).mean():+.3f}; Qini coefficient (area over random) = {qini:.3f}")
fig,ax=plt.subplots(figsize=(12.5, 5.2))
ax.plot(fracs,gain,color=GREEN,lw=2.5,label="targeting by τ̂ (rank & treat top)")
ax.plot(fracs,rand,color=GREY,lw=2,ls="--",label="random targeting")
ax.fill_between(fracs,rand,gain,color=GREEN,alpha=.12)
ax.axvline(peakfrac,color=BLUE,ls=":",label=f"optimal fraction ≈ {peakfrac:.2f}")
ax.axhline(0,color="k",lw=.6); ax.plot([peakfrac],[peak],"o",color=RED,ms=8)
ax.set_xlabel("fraction of population treated (highest τ̂ first)"); ax.set_ylabel("cumulative value gain")
ax.set_title(f"Targeting / Qini curve (Qini coefficient = {qini:.3f})"); ax.legend(fontsize=8)
plt.tight_layout(); plt.show()
print("The curve climbs while high-effect units are added, peaks at the optimal treated fraction, then falls as low-effect units")
print("dilute the gain. The area over the random line (Qini) is the value of the targeting -- how much the CATE ranking is worth.")
targeting curve peak value = 0.237 at fraction treated = 0.44 (oracle 0.238) random-targeting endpoint (treat all) = -0.153; Qini coefficient (area over random) = 0.206
The curve climbs while high-effect units are added, peaks at the optimal treated fraction, then falls as low-effect units dilute the gain. The area over the random line (Qini) is the value of the targeting -- how much the CATE ranking is worth.
3. Policy trees — an interpretable rule that maximizes value¶
A threshold on a black-box forest ($\hat\tau(x)>c$) captures the value but is not something a program committee can read, audit, or deploy as eligibility criteria. A policy tree (Athey & Wager 2021) learns a shallow, interpretable decision tree that directly maximizes the estimated policy value — using the cost-adjusted, doubly-robust reward of treating each unit. The result is a rule a human can inspect — a handful of thresholds on named covariates — whose value nonetheless nearly matches the unconstrained optimum. The rule the tree actually learns is printed in full below, because an interpretable rule that never gets displayed is not interpretable.
We fit a depth-2 policy tree on the reward of treating ($\hat\tau(x)-c$) versus not (0). Remarkably, on this problem the tiny interpretable tree recovers essentially the full oracle value while treating about the right fraction — the heterogeneity here is driven by one covariate, which the tree finds. Interpretability need not cost performance when the effect structure is simple; the tree also reveals that structure.
from econml.policy import PolicyTree
reward=np.column_stack([np.zeros(len(Xtr)), cf.effect(Xtr)-cost]) # value of [not treat, treat] per unit
pt=PolicyTree(max_depth=2,min_samples_leaf=200,random_state=0).fit(Xtr, reward)
rec=pt.predict(Xte).astype(int) # tree's treatment recommendation
val_tree=val(rec); val_thresh=val((tauhat>cost).astype(int))
print(f"policy tree (depth 2, interpretable): value gain = {val_tree:+.3f}, treats {rec.mean():.0%} (oracle {oracle:+.3f})")
print(f"threshold policy (τ̂ > c): value gain = {val_thresh:+.3f}, treats {(tauhat>cost).mean():.0%}")
def rule_lines(tree, names, node=0, depth=0, acc=None):
"""Walk the fitted policy tree and write its decision rule out in words."""
if acc is None: acc=[]
t=tree.tree_; pad=" "*depth
if t.children_left[node]==-1:
v=np.asarray(t.value)[node].ravel(); a=int(np.argmax(v))
acc.append(f"{pad}-> {'TREAT' if a==1 else 'do not treat'} "
f"(value {v[a]:+.3f} against {v[1-a]:+.3f} for the other action)")
else:
f,th=names[t.feature[node]],t.threshold[node]
acc.append(f"{pad}if {f} <= {th:.3f}:")
rule_lines(tree,names,t.children_left[node],depth+1,acc)
acc.append(f"{pad}else ({f} > {th:.3f}):")
rule_lines(tree,names,t.children_right[node],depth+1,acc)
return acc
print("\nThe learned rule, in full -- this is the whole point of using a tree:")
print()
for l in rule_lines(pt,[f"x{i+1}" for i in range(Xtr.shape[1])]): print(" "+l)
print()
imp=pt.feature_importances_
print(f"It splits on x{int(np.argmax(imp))+1} (importance {imp.max():.2f}), which is the covariate driving tau(x).")
cut=[t for t in pt.tree_.threshold if t>-1]
if cut:
learned_cut=max(cut)
opt_cut=(cost+1.0)/2.5 # tau(x)=2.5*x1-1.0, so tau>c at x1 > (c+1)/2.5
print(f"The boundary it settles on is x{int(np.argmax(imp))+1} > {learned_cut:.3f}.")
print(f"The value-maximising boundary is x{int(np.argmax(imp))+1} > {opt_cut:.3f}, since tau(x) = 2.5*x1 - 1.0 crosses the cost")
print(f"of {cost} there. The tree is {'conservative' if learned_cut>opt_cut else 'permissive'} by {abs(learned_cut-opt_cut):.3f} on that axis, and it costs")
print(f"essentially nothing: {val_tree:+.3f} against the oracle's {oracle:+.3f}.")
print()
print("Which is the result worth having. A depth-2 tree, fit to doubly-robust rewards, lands within a")
print("few hundredths of the optimal cutoff and gives up no measurable value -- and unlike the CATE")
print("surface it approximates, it is a sentence a programme can put in an eligibility document.")
fig,ax=plt.subplots(figsize=(12.5, 5.0))
names=["treat all","random\n(=treat all)","threshold\nτ̂>c","policy tree\n(depth 2)","oracle"]
vals=[treatall,treatall,val_thresh,val_tree,oracle]
ax.bar(names,vals,color=[RED,GREY,BLUE,GREEN,"k"]); ax.axhline(0,color="k",lw=.7)
for i,v in enumerate(vals): ax.text(i,v+(0.008 if v>=0 else -0.03),f"{v:+.2f}",ha="center",fontsize=9)
ax.set_ylabel("policy value gain"); ax.set_title("The interpretable policy tree nearly matches the oracle")
plt.tight_layout(); plt.show()
print("A readable depth-2 rule captures essentially the oracle value -- interpretability without a performance penalty here,")
print("and it exposes WHICH covariate drives who benefits (the auditable eligibility criterion a program can actually adopt).")
policy tree (depth 2, interpretable): value gain = +0.238, treats 42% (oracle +0.238)
threshold policy (τ̂ > c): value gain = +0.236, treats 42%
The learned rule, in full -- this is the whole point of using a tree:
if x1 <= 0.579:
if x1 <= 0.104:
-> do not treat (value +0.000 against -1.172 for the other action)
else (x1 > 0.104):
-> do not treat (value +0.000 against -0.522 for the other action)
else (x1 > 0.579):
-> TREAT (value +0.530 against +0.000 for the other action)
It splits on x1 (importance 1.00), which is the covariate driving tau(x).
The boundary it settles on is x1 > 0.579.
The value-maximising boundary is x1 > 0.560, since tau(x) = 2.5*x1 - 1.0 crosses the cost
of 0.4 there. The tree is conservative by 0.019 on that axis, and it costs
essentially nothing: +0.238 against the oracle's +0.238.
Which is the result worth having. A depth-2 tree, fit to doubly-robust rewards, lands within a
few hundredths of the optimal cutoff and gives up no measurable value -- and unlike the CATE
surface it approximates, it is a sentence a programme can put in an eligibility document.
A readable depth-2 rule captures essentially the oracle value -- interpretability without a performance penalty here, and it exposes WHICH covariate drives who benefits (the auditable eligibility criterion a program can actually adopt).
4. Summary¶
Policy learning is the step that turns a heterogeneous-effect estimate into an actual decision, and it is where the value of CATE modelling is realized:
- A treatment cost changes everything. With cost $c=0.4$ above the average effect, treating everyone lost value, treating no one gained nothing, and the oracle targeted policy ($\tau(x)>c$) was strongly positive. The operational question is never "does it work on average" but "for whom does the benefit beat the cost."
- The targeting / Qini curve ranks units by $\hat\tau$, treats the top fraction, and traces cumulative value — locating the optimal treated fraction at its peak and summarizing targeting quality (over random) in the Qini coefficient. Our CATE ranking's curve peaked at essentially the oracle value.
- Policy trees delivered a shallow, interpretable rule that maximized value directly — here recovering the full oracle value while remaining a rule a committee can audit and deploy.
Guidance: estimate $\tau(x)$, then learn a policy (a targeting curve to choose the treated fraction, a policy tree for a deployable rule), and always evaluate policy value and regret honestly — a good CATE model is worthless if the resulting policy is not evaluated against treat-all/treat-none and the budget. Cross-links. This is the decision layer atop the Causal Forest and meta-learner CATE estimators (subsection 9, examples 1–2); honest policy-value estimation uses the doubly-robust / AIPW / DML scores of the estimation notebooks; and the value-of-ranking idea (Qini) is the causal cousin of the AUC / uplift evaluation from the ML arc's model-evaluation subsection. It is also the most directly operational notebook in the arc — the bridge from "what is the effect" to "what should we do." The R companion runs policytree (Athey-Wager) with grf doubly-robust scores.