X-Learner Effect

Problem

Return the X-learner CATE at x as a propensity-weighted blend of the two arm-specific effect estimates: g*tau1 + (1-g)*tau0, where g is the propensity score at x.

Hint

Weight the estimates by how likely each arm is at x.

Test cases

case 1

in: [4,2,0.25]

expected: 2.5

case 2

in: [3,5,0.5]

expected: 4

case 3

in: [1,1,0.9]

expected: 1

case 4

in: [5,2,0]

expected: 2

case 5

in: [5,2,1]

expected: 5

Starter codePython

def x_learner_effect(tau1, tau0, propensity):
    # Your code here
    pass

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