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
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(
                        '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(
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?

  • $\begingroup$ The theory has not been worked out for the nonlinear case, and the code reflects that. The discussion you linked says that as well. Can you clarify why this is a problem? The ATE is unlikely to be outside the bounds, and the individual responses can be recoded to the appropriate boundary if they fall outside. $\endgroup$
    – dimitriy
    Feb 16 at 22:23

2 Answers 2


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.

  • 1
    $\begingroup$ Thanks. Do you know if there's a generally recommended method of estimating binary treatment effects then? Of course I could train a classifier and observe the effect of the treatment parameter, but that will suffer from regularization bias, so I'm not sure if that's the best method available. $\endgroup$ Sep 23, 2021 at 8:43
  • 1
    $\begingroup$ The function you used is fine. You can just bound the outputted predictions into [-1,1] manually. You might not want to use xgboost/random-forests/tree-based algorithms since they will lead to bound violations more often than smoother/spline-based estimators. $\endgroup$
    – Lars
    Sep 24, 2021 at 19:34
  • $\begingroup$ How would you recommend to bound the output? @Lars. $\endgroup$
    – jkortner
    Sep 30, 2022 at 22:04

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.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.