Regression Methods in Causal Inference In the most basic regression methods of causal inference (randomized experiment case), it's known that we can use covariates to predict the observed outcome, i.e. $Y^{obs}$ and the model is
$$
Y^{obs}_i=\alpha+\tau W_i+\beta X+\epsilon_i
$$
in which $\tau$ is ATE and $W_i$ is assignment. And we know that the causality of regression coefficients is guaranteed by randomized experiments. In other words, even if the linear relationship is wrong, we can still get the correct ATE. Regression helps to reduce the variance of ATE.
Therefore, my question is whether we can use another model, such as Tree or Neural Networks to add the covariates, that is
$$
Y^{obs}_i=\tau W_i+f(X)+\epsilon_i
$$
If so, can this model still reduce the variance? How to prove it and how to reduce possible overfitting?

Yes, we can use a non-linear estimator to reduce the variance and get more accurate results. There are many different techniques. To start of, for example, BART (Bayesian Additive Regression Trees) has been found to be an excellent algorithm for such  "out-of-the-box" causal inference tasks, see Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition by Dorie et al. (2017) for a more detailed investigation. In the last 5 to 6 years representation learning-based methods have also blossomed (starting with the work of Johansson et al.  (2016) Learning Representations for Counterfactual Inference) offering often very competitive results too.
