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

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$$?

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$$.
• Shouldnt $G^{n}\ =r_{n+1}+\gamma G^{n+1}$? You still havent explained how the immediate reward is accounted for. In the first simulation if you go with $G^{1}$, you disregard $r_{1}$. Aug 13, 2022 at 7:12