Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The authors extend mu-zero to MDPs with stochasticity by adding after-states to the tree and using a VQ-VAE model. They show that this enables them to solve stochastic environments where mu-zero fails.
main_review: Potential Areas of Improvements / Questions:
- One thing I would have liked to see, possibly in the appendix, is confirmation that this model trained on Go-Zero learned to become fully, or almost-fully deterministic. It’s not essential since it clearly works, but I think it’d be interesting! I suspect it doesn’t completely do this as this is the only explanation I am able to come up for why you need a larger sampling budget for stochastic mu-zero?
- To the former point, why do you need a larger sampling budget for stochastic mu-zero in the Go games?
- Would this model also then be applicable to multi-player, imperfect information games? This also does not affect my score nor am I asking for additional experiments on this, I just think such an experiment would be interesting.
- Is the notion of after-states strictly needed? Why is it advantageous to explicitly split the deterministic and stochastic components instead of having each node be stochastic?
- It would be good to include the hyperparam grids that were searched over to get the results in the appendix, or if none was done, to say so. The robustness section is nice but it’s valuable to know how much tuning was involved.

Clarity:
- In the “Model” section, I might have missed it, but it looks like l^p is not defined? Same for l^v.
- Notational question: why is c_{t+k+1} the only thing where the k index is a subscript rather than a superscript?
- In the section “chance outcomes” there seems to be something wrong with the sentence “By using a fixed codebook of one hot vectors, we can simplify the equations of the VQ-VAE 3”
- I may have missed it but in case I didn’t, the appendix should include an exact description of the training details of VQ-VAE, including when different components are frozen, the size of the codebook, etc. Unless I missed something, the paper is not reproducible as is.
- After equation (5), I am not totally clear what is meant by the expression $\sigma_t^k, ,c_{t+k+1}\sim \sigma_t^k$. I assume you mean to say that chance outcomes are drawn from sigma but the construction of the sentence makes it read like you are drawing $\sigma_t^k$ from $\sigma_t^k$.
- I think it would be worthwhile, at least in the appendix, to be clearer about the process by which the chance variables are actually sampled. The description in the “chance outcomes” section “The resulting encoder can also be viewed as a stochastic it 
function of the observation which makes use of the Gumbel softmax reparameterization trick (Jang et al., 2016) with zero temperature during the forward pass and a straight through estimator during the backward”  is a little unclear.
- Given that the VQ-VAE model is a key contribution, I really do think clearer explanation of how the “novel” variant of VQ-VAE works could be given. I don’t have a substantive solution, I just want the authors to be aware that this section was slightly confusing to read. In particular, it might be worth expanding on how the deterministic codes let you still acquire stochasticity.
- Is $l^\sigma$ defined anywhere?
- I think in Fig. 1B you want the sampling symbol rather than the approximately equal symbol?
summary_of_the_review: A very good paper with a couple of areas that need some work in rewriting to make the method clearer.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper aims to extend previous work on value-equivalent MBRL, such as MuZero, to stochastic environments. In contrast to conventional work in MBRL that fit transition models to be consistent with environmental observations, this line of work fits transition models to improve the accuracy / utility of a downstream value / policy. To this end authors consider the MuZero algorithm and advocate for learning a stochastic model with a VQ-VAE, and they modify MCTS so that it can be used with their stochastic model.

​The authors propose Stochastic MuZero. Their algorithm makes use of a stochastic model to predict future values, policies and rewards. The authors suggest utilizing "afterstates": an imaginary state that is the result of taking an action but it is also before the environment responds with an actual state. In 2048, for example, an afterstate could be the state reached after applying a tile moving action but before a number "2" tile appears in a random place.

As illustrated in Figure 1, the stochastic model consists of 5 functions in contrast to 3 functions in MuZero. The notable addition is in incorporating afterstates in these functions which allows for incorporating chance outcomes. There is an afterstate dynamics function that predicts a latent after-state given a state and action. The typical dynamics function would then still predict a next actual state and reward but its input will be a latent afterstate and a chance outcome. There is also an afterstate prediction function for value and a distribution prediction, where the distribution is that of a chance outcome given an afterstate. That distribution could then be used for sampling chance outcomes in inference.

To adapt MCTS to this model, search starts from a state and then proceeds to alternate at every level between afterstates and states by using the corresponding dynamics function to reach each type of state. ​ ​
main_review: I generally liked the ideas in this paper. Though, I believe there is room for improving the related work and empirical results. 

## Related  Work
There are two references that were not mentioned as part of the related work, but which are important in the line of value-equivalent MBRL:
* [Model-Based Reinforcement Learning with Value-Targeted Regression](https://arxiv.org/abs/2006.01107)
* [Value Iteration Networks](https://arxiv.org/abs/1602.02867)

## Experiments
To support their claims, the authors chose two environments that exhibit stochasticity: 2048 and Backgammon, and compared the performance of their algorithm against AlphaZero (where a perfect stochastic simulator is used) and MuZero (where a learned deterministic model is used). A similar methodology was used in an additional experiment, using the deterministic, perfect-information game of Go, and here the question was whether their algorithm's learned stochastic model could match the performance of a learned deterministic model.

In general, I see no issue with the choice of domains and baselines. My concerns are mostly with the way results are evaluated and reported. For instance, Figures 2, 3, and 4 all seem to represent a single trial. This is unacceptable for experiments with random outcomes, and it considerably weakens support for the paper's main claim of high performance in stochastic settings. I would imagine the authors expect their results to hold on average, given the randomness of a stochastic model. 

Some attempt was made to characterize the variability of results in Section 5.4. Each experiment was repeated---under unknown random conditions---three times for a small fraction of the data shown in figures 2, 3, and 4. There are two issues with this methodology:
1. Three samples is insufficient to characterize a distribution's dispersion.
2. The trials should be run with the same amount of data as the reported results. 

To approximate the distribution of performance to a reasonable degree of accuracy, I suggest that a minimum of thirty trials are run for the full length of learning. In addition, the authors should describe exactly the sources of randomness their results represent. 
 
It is also unclear if a hyperparameter search took place for the proposed algorithm. 

While the authors do reference relevant work that has used their domains, a concise description of the games and their sources of stochasticity is missing. This is important to communicate explicitly to the reader, since it is the primary environmental feature that experiments depend on. This information could be included in the appendix if the paper is tight on space.

## Minor points
* I didn't find the pseudocode in the appendix very informative. A simple table with all the considered parameters would be sufficient.

summary_of_the_review: Overall, this is an interesting idea, and I encourage the authors to continue working on this paper to provide sound empirical support.  If the claims were sufficiently supported, the results of this paper would be a significant contribution to MBRL. For now, I am leaning towards a rejection as the current empirical methodology is unsound.

### Final 
Thank you to the authors, for their response. I have read their replies and the other reviews. I will maintain my original decision to reject, because the biggest issue I have with the paper have not been addressed. This related to the empirical methodology.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The submission proposes an algorithm (called Stochastic MuZero) that combines VQ-VAEs with MuZero. Unlike MuZero, Stochastic MuZero can handle settings with stochasticity in a principled way (in terms of value equivalence). The submission shows that Stochastic MuZero can perform comparably to AlphaZero in 2048 and Backgammon. Additionally, it shows that Stochastic MuZero can perform comparably to MuZero in Go in a setting in which its computation budget is twice as large.
main_review: This work can largely be summarized as plugging the ideas from Ozair 2021 into a MuZero training pipeline. Nevertheless, I think it is a valuable contribution to the community, as it addresses a major shortcoming of MuZero and, unlike Ozair 2021, does so in self-play setting without a reconstruction loss.

The paper is nicely written (though see small comments below). My main comments are on the experimental side:

1) Throughout the paper, the submission claims that Stochastic MuZero retains the MuZero’s performance in deterministic settings. However, the paper only demonstrates this in a case in which Stochastic MuZero is given twice the simulation budget of MuZero.  I think that making this claim would require showing that Stochastic MuZero matches the performance of MuZero for equal computation budgets. Maybe Stochastic MuZero’s network architecture could be tuned such that each simulation is only half as costly? Seems like this merits further investigation.

2) One weakness of the paper is that the results on games with stochasticity are not on active (at least in a relative sense) research benchmarks. I think it would be interesting to see Stochastic MuZero benchmarked in Hanabi. While Hanabi is an N-player game, it can easily be turned into a POMDP by fixing N-1 players. This was the approach taken by single-agent SPARTA, for which code has been open sourced: https://github.com/facebookresearch/Hanabi_SPARTA. I think showing Stochastic MuZero’s performance against single-agent SPARTA would be a valuable contribution to the community. This experiment would also move toward addressing reviewer 4qYA's interest (which I also share) in seeing Stochastic MuZero in settings with a larger number of possible outcomes.

3) A third, smaller point, regards the baseline AlphaZero implementation. The appendix gives a number of details about the Stochastic MuZero implementation, but few regarding the AlphaZero implementation. I think it would be beneficial to devote some space to this. For example, did the submission's AlphaZero implementation used Monte Carlo returns, as was done in the AlphaZero paper, or did the submission find it necessary to modify this aspect of AlphaZero to achieve good performance in stochastic settings?

4) A fourth, again smaller point concerns the model learned by Stochastic MuZero. In particular, I would be interested to see some analysis on the extent to which Stochastic MuZero's model performs state abstractions. Ie, does Stochastic MuZero's model look more like the true transition model or more like MuZero's (transition determinization) model? Does Stochastic MuZero end up learning a deterministic model in deterministic environment? How do the shape of Stochastic MuZero's search trees look compared to those of AlphaZero and MuZero?

### Small Comments

It would be more generous to also cite Libratus (not just DeepStack) for tree-based planning algorithms for card games.

“However, learning a model in isolation from its use during planning has proven to be problematic in complex environments”
What about Dreamer v2?

VQ-VAE 3 — make clear this is referring to an equation

code ct+k+1 = one hot (arg maxi (e(o i ≤t+k+1))) produced 
The = part of the <= on this line is hard to see because of its proximity to the capital V on the line below.

It took me a minute to figure out how the encoder was being used. Might be worth discussing it at a greater length in the training section.

In Figure 1, it took me a minute to figure out that the pictures of the 2048 game were not pictures of a VQ-VAE embedding. Would be better to use backgammon or something else less misinterpretable.
summary_of_the_review: I think that, as a matter of correctness, my 1st point needs to be addressed for the submission to merit acceptance. Simply qualifying the claim would suffice but it would obviously be better if the submission were actually able to substantiate the claim that Stochastic MuZero can match MuZero's performance with an equal computation budget. I also think that the 3rd point should be addressed, toward making the submission's experiments more reproducible. While I think the paper would merit acceptance even without addressing my 2nd and 4th concerns, I will be disappointed if the authors choose not to perform these experiments, as I think they would provide significant value to the community.
review_title: null

final_score: rating: 10: strong accept, should be highlighted at the conference
review_content: summary_of_the_paper: The paper proposes an extension of MuZero to stochastic environments.

The stochasticity of the environment is handled by using afterstates, as_t, and chance outcomes, c_t. This decomposes the modelling of the stochastic environment dynamics into a deterministic model s_{t+1}, r_{t+1} = M(as_t, c_t) and modelling chance outcomes p(c_t | as_t). The chance outcomes are modeled as a discrete categorical variable (1 of M) and learned using a VQ-VAE like setup. 

The paper shows how the proposed model achieves ~SOTA on two stochastic environments: 2048 and backgammon, and retains SOTA performance on a single non-stochastic environment, Go, although, using twice the computational budget.
main_review: Strengths: 

 - This paper represents a significant contribution to reinforcement learning, by extending the state-of-the-art MuZero to stochastic or partially observed environments.
- The paper is clearly written and clearly explains the methods it builds on.
- The paper is evaluated on both stochastic and non-stochastic environments.

Weakness:
 - I would have liked a discussion of optimality. Expectimax is known to be optimal, but computationally infeasible in many applications. How does this work compare to expectimax and what are the approximations?
 - I would also have liked a discussion of limitations. What are the limitations of the proposed method? For instance, modeling discrete chance outcomes seems to limit this to environments with discrete randomness, e.g. dice rolls, cars, etc. What about environments with continuous stochasticity? Also, how large are the random outcome spaces of the environments tried, and how well does the one-hot chance outcomes scale with larger random outcome spaces? How well does it scale to partially observed environments, like e.g. poker?
summary_of_the_review: The paper proposes an extension of MuZero to stochastic environments, and achieves ~SOTA results on 2 stochastic environments, 2048 and backgammon, while retaining SOTA performance on a hard non-stochastic environment, Go. The proposed method is a significant step towards a strong, universally applicable, reinforcement learning algorithm.
review_title: null

Meta-review: The paper extends MuZero to stochastic (but observable) MDPs. To represent stochastic dynamics, it splits transitions into two parts: a deterministic transition to an afterstate (incorporating all observations and actions up to the current time), followed by a stochastic outcome (accounting for new randomness that follows the last action). The transition to an afterstate is similar in spirit to ordinary MuZero's dynamics model; the stochastic outcome is learned by a VQ-VAE. At planning time, MuZero retains the MCTS lookahead from ordinary MuZero. Stochastic MuZero achieves impressive results: e.g., it maintains the original MuZero's strong performance on the deterministic game of Go, while improving on MuZero significantly (and achieving superhuman performance) on the stochastic game of backgammon.

This is a strong paper overall: it presents a convincing and successful extension of the already-influential MuZero work, along with large-scale computational experiments confirming the utility of the approach. There are nonetheless a few weaknesses: first, compared to the original AlphaZero and MuZero work, it is perhaps less surprising that the given approach is successful, since it is more closely related to prior work. Second, due to the large-scale computational infrastructure needed, it is only possible to run some of the experiments once. This is not in itself a problem, but care needs to be taken in interpreting the results of such single-run experiments: e.g., any figures that show results of single-run experiments should have a clear warning label, and any statements such as "stochastic MuZero performs better than original MuZero" should be tempered with a caveat about how reliable these conclusions are likely to be. Section 5.4 (which runs shorter experiments using three random seeds each) makes a start at evaluating reliability, but (a) the headline results in previous sections do not contain any caveats or pointers to 5.4, and (b) 5.4 should explicitly acknowledge that it cannot hope to detect even quite-common failure cases with so few seeds.
Decision: Accept (Spotlight)