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


1 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$.

  • $\begingroup$ 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}$. $\endgroup$
    – Matrix001
    Aug 13, 2022 at 7:12
  • $\begingroup$ @Matrix001 Updated. $\endgroup$
    – Uduse
    Aug 13, 2022 at 20:55
  • $\begingroup$ @Matrix001 does it make sense now? $\endgroup$
    – Uduse
    Aug 28, 2022 at 20:01

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .