How does Generalized Random Forest calculate the gradient of the score function? The reference is GENERALIZED RANDOM FORESTS by ATHEY, TIBSHIRANI and WAGER (2019). They construct a general algorithm to grow trees and forest for estimation of target parameters that are conditional on a covariate vector $X$.
Notation and Setup
Observable data $O_{i}=(Y_{i},W_{i})$, where $Y$ defines the outcome, $W\in\{0,1\}$ defines the binary treatment, and $X\in\mathbb{R}^{p}$ represents a covariates matrix. In matrix form
$$W=\left(\begin{array}{c}
W_{1}\\
\vdots\\
W_{N}
\end{array}\right)_{N\times1}Y=\left(\begin{array}{c}
Y_{1}\\
\vdots\\
Y_{N}
\end{array}\right)_{N\times1}X=\left(\begin{array}{ccc}
X_{11} & \cdots & X_{1p}\\
\vdots & \ddots\\
X_{N1} &  & X_{Np}
\end{array}\right)_{N\times p}
$$
The paper defines the target parameter $\theta(x)$ (a possible nuisance parameters $\nu(x)$) as solutions to local moment conditions
$$\mathbb{E}\left[\psi_{\theta,\nu}\left(O\right)\mid X=x\right]=0.$$
In the example of conditional treatment effect (CATE) estimation we assume $Y=W_ib_i+e_i$ and conditional independence, and the CATE is defined as $\beta(x)=\mathbb{E}[b_i\mid X_i=x]$. Then a for a given contrast $\xi$, the target parameter $\theta(x)=\xi\beta(x)$ is identified via the score function
\begin{align*}
\psi_{\beta,c}(Y_{i},W_{i})_{\mid X_{i}=x} & =\left[\begin{array}{c}
Y_{i}-\beta(x)W_{i}-c(x)\\
\left(Y_{i}-\beta(x)W_{i}-c(x)\right)W_{i}
\end{array}\right]_{2\times1}
\end{align*}
which defines the two moment conditions in experiments,
\begin{align*}
\mathbb{E}\left[\psi_{\beta,c}\left(Y,W\right)\mid X_{i}=x\right]=0\leftrightarrow\left[\begin{array}{c}
\mathbb{E}\left[Y-\beta(x)W-c(x)\mid X_{i}=x\right]\\
\mathbb{E}\left[\left(Y-\beta(x)W-c(x)\right)W\mid X_{i}=x\right]
\end{array}\right]_{2\times1}=\left[\begin{array}{c}
0\\
0
\end{array}\right]_{2\times1}
\end{align*}
for all $x\in\mathcal{X}$.
A gradient of $\mathbb{E}[\psi\mid X_i=x]$
The generalized random forest performs CART splits on psuedo-outcomes, calculated as follows. Given a parent node $P$, we define the pseudo-outcome
$\rho$ by
$$\rho_{i}=-\xi^{T}A_{p}^{-1}\psi_{\hat{\theta}_{P},\hat{\nu}_{P}}(O_{i}),$$
where $A_{P}$ is a consistent estimate for the gradient of $\mathbb{E}\left[\psi_{b,c}\left(Y,W\right)\mid X_{i}\in P\right]$.
In the CATE estimation this amounts to
\begin{align}
    \rho_i &= A_P^{-1}\left(W_i-\bar{W}_P\right)\left(Y_i-\bar{Y}_P-\left(W_i-\bar{W}_P\right)\hat{\beta}_P\right) \\
    A_P &= \frac{1}{\left| \{i:X_i\in P\} \right|} \left[ \sum_{\{i:X_i\in P\}} \left(W_i-\bar{W}_P\right)\left(W_i-\bar{W}_P\right)^T \right]
\end{align}
My question is: why? Why is $A_p$ the gradient of the local expectation score function $\mathbb{E}[\psi\mid X_i=x]$? How can $A_p$ be defined as a gradient, when the $\mathbb{E}[\psi\mid X_i=x]$ is in fact a matrix of vector?
 A: I believe there is a typo in the paper. Equation (20) should be:
\begin{align}
A_P &= \frac{1}{\left| \{i:X_i\in P\} \right|} \left[ \sum_{\{i:X_i\in P\}} \left(\tilde{W_i}\right)\left(\tilde{W_i}\right)^T \right]
\end{align}
where $\tilde{W_{i}}$ is a $(q + 1) \times 1$ vector that has a 1 in the first position and the vector $W_{i}$ in the next $q$ positions. I explain in detail below.
Gradients and Jacobians
I'll follow the gradient conventions laid out in the Mathematical Appendix of Patterns, Predictions and Actions by Hardt and Recht (2022).

*

*Let $\Phi: \mathbb{R}^{d} \to \mathbb{R}$. The gradient of $\Phi(x)$ with respect to the variable $x$ evaluated at a value $w \in \mathbb{R}^{d}$ is the vector of partial derivatives

$$ \nabla_{x} \Phi(w) =\left[\begin{array}{c}
\frac{\partial \Phi(w)}{\partial x_1}\\
\vdots\\
\frac{\partial \Phi(w}{\partial x_d}
\end{array}\right]_{d\times1}.$$

*

*Let $\Phi: \mathbb{R}^{n} \to \mathbb{R}^{m}$ a multivariate mapping. The Jacobian of $\Phi(x)$ with respect to the variable x evaluated at a value $w \in \mathbb{R}^{d}$ is the $m \times n$ matrix

$$ D_{x}\Phi(w) = \left[ \frac{\partial\Phi_{i}(w)}{\partial x_{j}} \right]_{(i=1,...,m),(j=1,...,n)} $$
Estimating conditional average partial effects
We observe samples $(X_{i},Y_{i},W_{i}) \in \mathcal{X} \times \mathbb{R} \times \mathbb{R}^{q}$ and our conditional parameters of interest are $\beta:\mathcal{X} \to \mathbb{R}^{q}$ and $c:\mathcal{X} \to \mathbb{R}$. The scoring function $\Psi: \mathbb{R}^{q+1} \mapsto \mathbb{R}^{q+1}$ is
\begin{align} 
\Psi_{\beta(x),c(x)}(Y_{i},W_{i})^{\color{green}{[(q+1) \times 1]}} &= (Y_{i}^{\color{green}{[1 \times 1]}} - {W_{i}^{T}}^{\color{green}{[1 \times q]}}\beta(x)^{\color{green}{[q \times 1]}}-c(x)^{\color{green}{[1 \times 1]}})^{\color{green}{[1 \times 1]}}{[\quad 1^{\color{green}{[1 \times 1]}} \quad {W_{i}^{T}}^{\color{green}{[1 \times q]}} \quad ]^{T}} ^{\color{green}{[(q+1) \times 1]}} \\
 &= (Y_{i} - {W_{i}^{T}}\beta(x)-c(x))^{\color{green}{[1 \times 1]}} \left[\begin{array}{c} 1 \\ W_{i1} \\ W_{i2} \\ \vdots \\ W_{iq} \end{array}\right]^{\color{green}{[(q+1) \times 1]}}
\end{align}
In the strictest sense, given this setup, differentiation of $\Psi$ is only well defined for its arguments $(Y_{i},W_{i})$. However, we are interested in relating changes in the conditional parameters $(\beta(x),c(x))$ to changes in the scoring function value, so we will be taking partial derivatives with respect to the parameters rather than the arguments. It might help to think about a modified mapping $\tilde{\Psi}:\mathbb{R}^{q+1} \mapsto \mathbb{R}^{q+1}$ of the form $\tilde{\Psi}_{Y_{i},W_{i}}(\beta(x),c(x))$ where $(Y_{i},W_{i})$ are parameters and $(\beta(x),c(x))$ are arguments.
Note that just like in the case of $\Psi$, $\tilde{\Psi}$ is a multivariate mapping and so speaking about its gradient makes no sense. I think that a good part of the OP confusion has to do with the fact that the object $\nabla_{(\beta(x),c(x))} \Psi_{\beta(x),c(x)}(Y_{i},W_{i})$ is described as a "gradient" in the paper. According to the convention described above, it is more accurately described as the Jacobian matrix $D_{(\beta(x),c(x))} \tilde{\Psi}_{Y_{i},W_{i}}(\beta(x),c(x))$ of dimensions $(q+1) \times (q+1)$.
With that out of the way, we can focus on calculating the Jacobian:
\begin{align}
D_{(\beta(x),c(x))} \tilde{\Psi}_{Y_{i},W_{i}}(\beta(x),c(x)) &= 
\begin{bmatrix}
\frac{\partial(Y_{i} - {W_{i}^{T}}\beta(x)-c(x))}{\partial c(x)} & \frac{\partial(Y_{i} - {W_{i}^{T}}\beta(x)-c(x))}{\partial\beta(x)_{1}} & 
\ldots & \frac{\partial(Y_{i} - {W_{i}^{T}}\beta(x)-c(x))}{\partial \beta(x)_{q}} \\
\frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{i1})}{\partial c(x)} & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{i1})}{\partial\beta(x)_{1}} & 
\ldots & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{i1})}{\partial \beta(x)_{q}} \\
\vdots & \vdots & \ddots & \vdots \\
\frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{iq})}{\partial c(x)} & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{iq})}{\partial\beta(x)_{1}} & 
\ldots & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{iq})}{\partial \beta(x)_{q}}
\end{bmatrix} \\
&= \begin{bmatrix}
1 & W_{i1} & \ldots & W_{iq} \\
W_{i1} & W_{i1}^{2} & \ldots & W_{iq}W_{i1} \\
\vdots & \vdots & \ddots & \vdots \\
W_{iq} & W_{i1}W_{iq} & \ldots & W_{iq}^{2}
\end{bmatrix}\\
&= [1 \quad W_{i}] [ 1 \quad W_{i}]^{T}\\
&= \tilde{W_{i}} \tilde{W_{i}}^{T}
\end{align}
Plugging this into equation (7) in the paper gives the result stated at the beginning of this answer. I don't see where is the extra term $\bar{W_{P}}$ mentioned in the original formulation of equation (20) coming from.
