Machine Learning for Causal Inference with Panel Data: Possible to combine ML estimators with additive/linear terms to derive diff-in-diff estimator?

My question is motivated by the following. First consider the non-panel case, where we have two groups, the treated group ($g=t$) and the comparison group ($g=c$), and are trying to estimate an average treatment effect. Let $\hat\mu_g(X)$ be any good ML learner (e.g. gradient boosted tree, random forest, neural network) of the outcome $Y$ given covariates $X$ for units in group $g$. Letting $i=1,...,n$ index units, one approach to estimating the average treatment effect is the following: $$\tag{1}\hat{ATE}=\frac{1}{n}\sum_{i=1}^n[\hat\mu_t(\textbf{x}_i)-\hat\mu_c(\textbf{x}_i)]$$ (I know that some authors will say that you need to use doubly-robust methods here that reweight residuals by the inverse propensity score or some other distance metric, but personally I am wary of propensity score reweighting approaches so I would prefer to focus on just the regression-based estimator.)

So far so good. But now let's consider the panel data or longitudinal data situation, where each unit is observed in two time periods, and the treated group only receives treatment in the second time period. Let $p\in\{0,1\}$ index the pre-treatment and post-treatment periods respectively. At first glance, a natural way to extend estimator $(1)$ to the panel data context would be to develop four ML estimators: $\hat\mu_{t0}, \hat\mu_{t1}, \hat\mu_{c0}, \hat\mu_{c1}$ that predict $Y$ given $X$ for the treated and comparison groups in the pre-treatment and post-treatment periods respectively, and then apply the following difference-in-differences-esque estimator: $$\tag{2}\hat{ATE}=\frac{1}{n}\sum_{i=1}^n[(\hat\mu_{t1}(\textbf{x}_{i1})-\hat\mu_{t0}(\textbf{x}_{i1}))-(\hat\mu_{c1}(\textbf{x}_{i1})-\hat\mu_{c0}(\textbf{x}_{i1}))]$$ However, estimator $(2)$ is not satisfactory because it allows the extent to which a change in a covariate predicts a change in the outcome ($\hat{\partial{Y}/\partial{\textbf{x}_k}}$) to differ between the treated group and the comparison group in the pre-treatment period even when they have the same $X$ values. Yet there is no reason for this to be the case; I want to allow for different average levels (or intercepts) between the treated and comparison groups in the pre-treatment period, but for a given set of covariate values $X=\textbf{x}$ I want to constrain the predicted impact $\hat{\partial{Y}/\partial{\textbf{x}_k}}$ to be the same for the treated and comparison groups in the pre-treatment period. One thought is to pool the treated and comparison group observations in the pre-period and estimate an ML learner $\hat\mu_0(X,g)$ that simply treats group status $g$ just like any other covariate; but this still does not solve the issue, because with a flexible ML learner there is nothing to prevent interactions between $X$ and $g$.

So finally this brings me to my question. What I want to do is pool the pre-treatment observations from the treatment and comparison groups to estimate $\hat{\partial{Y}/\partial{\textbf{x}_k}}$ using a flexible ML learner that only allows for an average intercept difference between groups in the pre-treamtent periods, i.e. predict $Y$ in the pre-treatment period using an estimator of the form: $$\tag{3}\hat\mu_0(X)+\hat\delta_0\cdot{g}$$ Is it possible to combine a generic ML learner (e.g. gradient boosted tree, random forest, neural network) with a simple linear/additive term to estimate $(3)$?

Given an estimator $(3)$, the average treatment effect could then be estimated as follows: $$\tag{4}\hat{ATE}=\frac{1}{n}\sum_{i=1}^n[(\hat\mu_{t1}(\textbf{x}_{i1})-(\hat\mu_{0}(\textbf{x}_{i1})+\hat\delta_0))-(\hat\mu_{c1}(\textbf{x}_{i1})-\hat\mu_{0}(\textbf{x}_{i1}))]$$ which simplifies to: $$\tag{4'}\hat{ATE}=\frac{1}{n}\sum_{i=1}^n[\hat\mu_{t1}(\textbf{x}_{i1})-\hat\delta_0-\hat\mu_{c1}(\textbf{x}_{i1})]$$

If $(3)$ is estimable, then two natural follow-up questions are:

• Could this be extended to the situation where you have multiple pre-treatment period observations and want to allow for separate time trends between the treated and comparison groups in the pre-treatment periods -- either linear time trends or more flexible approaches such as natural cubic splines;

• and secondly, could this approach be extended to non-linear outcomes, including binary and multi-class outcomes, and outcomes that are continuous but lower-bounded at zero with a non-trivial mass of observations at zero.

Lastly, I would I assume that this approach, if feasible, could be easily extended to estimate heterogeneous conditional average treatment effects for units with any particular range of covariate values by subsetting the $\hat{ATE}$ estimator to only average over units with covariate values in the particular range.