How to get linear approximation of causal random forest?

[josh-ashkinaze]

Context
I have a multi-arm RCT. I want to run a causal random forest, and then afterwards make results more interpretable by getting a linear approximation of the results. Is there a better way than just running OLS on CATEs with robust SEs?

import numpy as np
import statsmodels.api as sm
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from econml.dml import CausalForestDML

np.random.seed(42)

####################
# Generate data
####################

n = 2000
n_arms = 3
n_mods = 4

X = np.random.randn(n, n_mods)
W = np.random.randn(n, 2)
T = np.random.choice(n_arms + 1, size=n)

cate_true = {
    1:  1.5*X[:,0] - 0.8*X[:,1],
    2:  0.5*X[:,1] + 1.0*X[:,2],
    3: -1.0*X[:,0] + 0.3*X[:,3],
}
Y = np.array([cate_true[t][i] if t > 0 else 0 for i, t in enumerate(T)]) \
    + X.sum(1) + np.random.randn(n)

####################
# Fit causal forest
####################

est = CausalForestDML(
    model_y=RandomForestRegressor(n_estimators=200, random_state=42),
    model_t=RandomForestClassifier(n_estimators=200, random_state=42),
    discrete_treatment=True,
    n_estimators=500,
    random_state=42,
)
est.fit(Y, T, X=X, W=W)

####################
# Regress CATEs on moderators
####################

for arm in range(1, n_arms + 1):
    cate_hat = est.effect(X, T0=0, T1=arm)
    Xc = sm.add_constant(X)
    ols = sm.OLS(cate_hat, Xc).fit(cov_type='HC3')
    coef_str = ", ".join([f"X{j+1}={ols.params[j+1]:.3f} (p={ols.pvalues[j+1]:.3f})" for j in range(n_mods)])
    print(f"\nArm {arm}: {coef_str}")

[kbattocchi]

That seems reasonable to me; perhaps also consider SingleTreeCateInterpreter which learns an interpretable shallow decision tree.