Doubly Robust ATE Estimate

Problem

Estimate the ATE with the doubly robust formula: mean(mu1 + t*(y - mu1)/e - mu0 - (1-t)*(y - mu0)/(1-e)), clipping propensities to [1e-6, 1 - 1e-6].

Empty input returns 0.0.

Hint

The estimator stays consistent if either the propensity model or the outcome model is correct.

case 1

in: [[1,1,0,0],[1,0,1,0],[0.5,0.5,0.5,0.5],[0,0,0,0],[0,0,0,0]]

expected: 0

case 2

in: [[1,0],[3,2],[0.75,0.5],[2,1],[3,2]]

expected: 0

case 3

in: [[1,1,0],[3,1,0],[0.5,0.5,0.5],[2,0,0],[3,2,2]]

expected: 1

case 4

in: [[],[],[],[],[]]

expected: 0

Starter codePython

def doubly_robust_ate(treatment, outcome, propensity, mu0, mu1):
    # Your code here
    pass

4 test cases · run your code against them in the editor