Does `CausalForestDML` assume linear treatment?

[AlxndrMlk]

Hi!

The documentation here states that CausalForestDML assumes linear treatment.

Is this because CausalForestDML estimates a partially linear model (similar to LinearDML)?

Thank you!

[kbattocchi]

Yes, exactly, all of the DML methods estimate partially linear models. But note that:

1. It is linear in the treatment T, but CausalForestDML estimates a flexible conditional average treatment effect in terms of the features X. That is, we are assuming something like Y=\theta(X) T + f(X,W) where we learn a flexible forest-based \theta during fitting. By contrast, LinearDML learns an effect model that is linear in (a featurization of) the features X.
2. With our most recent release, you can easily pass a treatment featurizer (e.g. CausalForestDML(..., treatment_featurizer=PolynomialFeatures(degree=2, include_bias=False)) to estimate a quadratic instead of linear effect in T). Under the covers, we ultimately still estimate a linear model in each of the columns in the transformed treatment, but this saves you from having to do the featurization yourself.