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?