In the MuZero paper, how does backprop in the MCTS account for the immediate reward from each edge?

Question

On page 12 of this paper: Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model, it describes how MCTS works for the MuZero algorithm. It states in equation 4 that during the 'backup' after a simulation, the mean value (Q) for every edge in the simulation path is updated by:

$$Q\left(s^{k-1},\ a^k\right){\colon=}\frac{N\left(s^{k-1},\ a^k\right)Q\left(s^{k-1},\ a^k\right)+G^k}{N\left(s^{k-1},\ a^k\right)+1}$$

where $G^k$ is the return from depth $k$ and onwards in the simulation.

However, I don't see how this equation accounts for the immediate reward attained by transitioning from state $s^{k-1}$ to $s^{k}$.

Since $Q\left(s^{k-1},\ a\right)$ is used to determine what action to select from state $s^{k-1}$, shouldn't the backprop use the return from k-1 ($G^{k-1}$), to update the mean value instead of $G^k$?

Answer 1

Let's look at the first simulation where $k = 1$. We have $G^0 = r_1 + \gamma v^0$ and $G^1 = v^1$, and we can pick one of them as the update target. The difference is that $G^1$ is fully conditioned on the action $a_1$ we just took, while $G^0$ is only partially so. We choose $G^1$ for the same reason that we do the search: we assume the evaluations of child nodes conditioned on more actions is better than the evaluation of the root alone.

EDIT: I believe the expectation should match so that $\mathbb{E}[G^n] = \mathbb{E}[r_{n+1} + \gamma G^{n+1}]$. I looked deeper into it and I agree it's confusing. I think what's happening is that in the paper they mixed two different definitions of $Q$: either $Q(s, a) = V(s')$ or $Q(s, a) = R(s, a) + V(s')$. Their implementations (based on MuZeor's pseudo-code and MCTX) suggest that $Q(s, a) = V(s')$ because child nodes store values without considering the last received reward. However, the reward plays a part in the child selection phase, so a better way to write their selection formula is

$$a^{k}=\arg \max_{a} R(s, a) + Q(s, a)+P(s, a) \frac{\sqrt{\sum_{b} N(s, b)}}{1+N(s, a)}\left[c_{1}+\log \left(\frac{\sum_{b} N(s, b)+c_{2}+1}{c_{2}}\right)\right]$$ and this reflects what they are doing in the code.

On the other hand, what you suggested make perfect sense too if we adopt the other definition of $Q$.