This doesn't look right to me - this should use the equation from the "Doubly Robust" section of https://github.com/py-why/EconML/blob/main/doc/spec/estimation/dr.rst

Mmm, if you mean the right equation should be the one right below "The Doubly Robust approach": Y_{i, t}^{DR} = g_t(X_i, W_i) + \frac{Y_i -g_t(X_i, W_i)}{p_t(X_i, W_i)} \cdot 1{T_i=t}
It's:
Y_DR[i,t] <- g_t(X[i], W[i]) + (Y[i] - g_t(X[i], W[i])) / p_t(X[i], W[i]) * (T[i] == t)
What I put here should be a short format combine with line 16 and line 18 (Where I put the input, weights)?

Yeah, it's a bit tricky to write out - the key is that you're not multiplying by T at the end, you're multiplying by the indicator function selecting the specific case of T, and likewise you're dividing by the probability of that specific treatment, which is a bit awkward to express in pseudo-notation.

I think something like

Ydr(g,p) = g(X,W,T) + (Y - g(X,W,T)) / p_T(X,W)

would make it more obvious that there's only one term being included, it's not really being multiplied by T in a meaningful way, and it expresses the more-than-binary treatment outcome correctly. This does require writing out the arguments to g and p, but I think that's okay since otherwise it's hard to be precise about what's being computed.

Ok updated acoordingly.

This code should basically reflect the code in

where we take Y_pred from the nuisances and we form the doubly-robust estimate by subtracting Y_pred[:,0] from the other Y_pred[:,t] values. Since in the model fitting code the propensities nuisance is used only for adjusting sample_var, which we don't support here, I think you can ignore all of that code. So, I think this should look something more like:

Good suggestion, thanks!