Doubly robust learning with binary treatment and outcome

Question

I'm trying to use doubly robust learning to estimate heterogenous treatment effects. My treatments T and outcomes y are both binary. I'm following the example listed under "How do I select the hyperparameters of the first stage models or the final model?" in the econML documentation.

For binary outcomes, the recommendation seems to be to simply wrap classifiers to make them behave like regressors, see for instance this github comment. A working example is provided below. However, the model outputs treatment effects which exceed 1, which doesn't make any sense for binary outcomes:

from econml.dr import DRLearner
import numpy as np
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV

# Wrapper to make classifier behave like a regressor
class RegWrapper:
    def __init__(self, clf):
        self._clf = clf

    def fit(self, X, y):
        self._clf.fit(X, y)
        return self

    def predict(self, X):
        return self._clf.predict_proba(X)[:, 1]

### Simulate data
N = 10000
np.random.seed(0)
X = np.random.uniform(size=(N, 3))  # data affecting outcome and treatment
y = [int(np.random.uniform() < 0.1) for _ in range(N)]  # binary outcomes
T = [int(np.random.uniform() < 0.3) for _ in range(N)]  # binary treatment variables

# Propensity model - predicts T from X (classification)
model_clf = LogisticRegression()

# 'Regression' model - predicts y from X
model_reg = RegWrapper(LogisticRegression())

# Final model - predicts (I think) y_t=1 - y_t=0 from X
model_fin = GridSearchCV(
                estimator=RandomForestRegressor(n_jobs=1),
                param_grid={
                        'max_depth': [3, None],
                        'n_estimators': (10, 50)
                    }, cv=3, n_jobs=1, scoring='neg_mean_squared_error'
                )

# Create and fit doubly robust learner
est = DRLearner(
    model_regression=model_reg,
    model_propensity=model_clf,
    model_final=model_fin,
    cv=2)
est.fit(y, T, X=X, W=None)

# Estimate heterogenous treatment effect for all points
point = est.effect(X, T0=0, T1=1)

print(max(point))  # prints 1.456

Can anyone point me in the right direction to get correct treatment effect estimates?

Answer 1

The DR Learner does not respect known bounds. It uses least-squares regression and therefore can lead to treatment effect predictions outside (-1,1). This is a limitation of the method.

Answer 2

I have come across this issue before and this is a limitation of the package. However, an easy solution is just to do it manually, since it's easy enough. Just fit the nuisance functions using your method of choice (like 2 logistic regressions), and then just compute the doubly robust scores yourself according to the formula.

Then you use your random forest regressor or whatever to fit these scores and estimate the conditional average treatment effect. It is still possible at this step for the algo to produce results outside the range of [-1,1], but it becomes increasingly improbable with larger data size.