Epistemic Monte Carlo Tree Search

Dear reviewer,

Thank you for your valuable time and many detailed comments and suggestions.

1. Thank you for pointing this out, we will clarify the text.

2. This situation is indeed possible, but only when the estimator of the epistemic uncertainty is not working correctly. In the presence of errors in uncertainty estimation we expect EMCTS to have an advantage over methods that do not use planning for exploration, such as the A/MZ+UBE ablation in Figures 2 and 4, because EMCTS averages across multiple estimations of the uncertainty and thus can average out independent errors, while model-free exploration methods (such as the ablation) only have access to one estimate for the uncertainty of the value of the current state $\mathbb V[Q(s_t,a)]$ for each $a$. This is corroborated in Figure 3 in the appendix where the uncertainty estimates of EMCTS (top row) match the ground truth much better than those of the one-prediction UBE estimator (middle row). We touch on this briefly (line 061 in the introduction), and will make sure this is emphasized more strongly.

3. Popular methods such as AlphaZero [1] and MuZero [2] learn one model $\hat m$ rather than a full posterior $\hat M$, because approximating $\hat M$ is hard to intractable in practice. In that case, the sample standard deviation $\frac{1}{N-1} \sum_{i}^N (\nu_{\hat m}^i - \bar \nu)$ yields a measure of how much the backups $\nu^i_{\hat m}$ vary inside the model, because every backup is from the same model $\hat m$. This variance is the non-epistemic part of the uncertainty (i.e. how stochastic are $\hat m$ and the search policy), and is not suited for directing exploration into novel areas of the state-action space. For this reason, we must use other methods such as Equation 11 (see full derivation in Appendix A.2), that estimate the variance of a sum of random variables using the sum of the individual variances.

4. Thank you for pointing out that our explanation wasn't sufficiently clear. We will make sure this is described more clearly in the paper. Please see our address of concerns with this assumption in the general rebuttal.

5. Our intent was to point out that to solve Deep Sea the only way for an RL agent that does not have prior knowledge over the goal transition, such as EMCTS, is by being able to explore novel state actions that are far in the future compared to the current state action, i.e., deep exploration. We will make sure this is phrased clearly in the paper.

6. We include results in a suite of easy-exploration MinAtar problems in the uploaded revision, where the performance of EMCTS is comparable to baseline AlphaZero using the same uncertainty estimator used for SUBLEQ. Note however that in easy-exploration environments the behavior of UCB exploration methods is expected to be from non-existent down to detrimental. For example, in environments such as Cartpole any exploration is detrimental for performance because the goal state is observed immediately. Further, tuning EMCTS' hyperparameters includes the possibility of setting $\beta = 0$, which means that exploration with EMCTS can always reduce to the baseline when random or no exploration performs better than UCB-based exploration. In regards to more complex, large state-action spaced environments such as Chess (although this also applies to image-based Atari and Minatar), it is a property of UCB-based exploration methods such as EMCTS that if the uncertainty estimator is not suited, UCB exploration can induce searching every individual state-action which is not useful when the state-action spaces are large (we discuss this in more detail in Section 2.3), and again, can even be detrimental. Given a suited uncertainty estimator (such as the IO-Hash used by EMCTS for SUBLEQ in Figure 1), the confidence bounds provide relevant information for exploration in the environment. Finding such estimators is an open problem and an active area of research. However, it is orthogonal to the contributions of EMCTS.

[1] Silver, David, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot et al. "A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play." Science 362, no. 6419 (2018): 1140-1144.

[2] Schrittwieser, Julian, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez et al. "Mastering atari, go, chess and shogi by planning with a learned model." Nature 588, no. 7839 (2020): 604-609.

Dear reviewer,

Thank you for your valuable comments and questions.

Some assumptions might make EMCTS inapplicable to real complex problems:

1. Assume that the transition model P is true and in many cases, the transition models need to be estimated, and then the EMCTS will not work: We demonstrate in Figures 2 and 4 that EMCTS works well even when the transition model is learned. Theorem 1 indeed assumes the Q-values use the true transition model. However, as discussed in the common rebuttal, this is addressed in Appendix A.3 and the method proposed in A.3 is evaluated and successful in Figure 4. In Figure 2 we demonstrate that EMCTS works with learned transition models even without modifying the upper bound.

2. Theorem 1 need $Q^\pi$ to be linear: There seems to be a misunderstanding. Note that any value is a sum of discounted expected future rewards and is therefore linear in the rewards. The proof of Theorem 1 (Appendix A.1) only uses this property which is true for every value function. It does not assume that $Q^\pi$ is a generally linear function.

3. The learned models are assumed to be consistent: We do not believe that this is a concern in practice regardless of the complexity of the problem, see our discussion in the above common rebuttal.

To summarize, we believe none of these assumptions prevent EMCTS from being applicable to real complex problems, such as programming, demonstrated on SUBLEQ in Figure 1.

Is it satisfied that the reward is independent in different types of tasks: Does the reviewer refer to the independence assumption made in line 287-288? Please note that we assume the epistemic uncertainty of different rewards to be independent, not the rewards themselves. This assumption is irrespective of the task. It depends only on the class of estimator used (for example, whether $\hat R$ is a DNN or a table), for the reason that it is only assumed with respect to the learned reward function $\hat R$. If this is not the assumption the reviewer refers to, could the reviewer please clarify which assumption is referred to?

How to get the true transition model in typical tasks? In many typical tasks where RL is used in practice (board games, programming and algorithm design, video games, robotics) the real transition model is either known (for example, in single-player games or zero sum 2-player games with self-play), definable by the practitioner (for example, programming and algorithm design) or formulated using physics models in robotics.

Is the accuracy and uncertainty of the learned model dynamically improved in the learning procedure? The accuracy of the learned models (reward, value and or transition, but also the learned value uncertainty estimator $\hat u$) indeed all improve as the agent sees a growing number of training steps and data points, demonstrated also in Figure 3 in the appendix where the estimates improve left to right with respect to the ground truth.

And will this affect the consistency assumption? For the above reason, indeed any possible inconsistency or error in $\hat u$ should further reduce as training progresses. In practice however we see in the right subplots of Figure 4 in the appendix that the agent explores the environment very quickly, demonstrating that the $\hat u$ estimates are sufficiently accurate to result in reliable exploration of novel states even early on in training.

We hope that these answers along with the main rebuttal address the reviewer's concerns and questions. We are happy to clarify anything that is left unclear.

[1] Schrittwieser, Julian, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez et al. "Mastering atari, go, chess and shogi by planning with a learned model." Nature 588, no. 7839 (2020): 604-609.

Dear reviewer,

Thank you for your time, detailed review, comments and questions.

Consistency: Please see our reply in the main rebuttal above.

High dimensional, large or continuous state-action spaces: To effectively explore high dimensional, large or continuous state action spaces, Upper-Confidence-Bound (UCB) based exploration approaches (such as EMCTS) require estimators that estimate the epistemic uncertainty in reward/value of unvisited transitions $(s',a')$ based on a suitable similarity metric between $(s',a')$ and visited transitions $(s,a)$ (see Section 2.3). Developing such estimators for complex domains is indeed a hard problem and remains an active area of research. It is however orthogonal to the contributions of EMCTS. In the presence of suitable estimators EMCTS can be used directly without any computational cost that depends specifically on the dimensionality or the size of the state-action space. This is demonstrated in Figure 1 where EMCTS paired with a suitable estimator, the IO-Hash, can generate SUBLEQ programs that solve specific programming tasks much more efficiently than the variation which uses a less-suited source of uncertainty (the full observation hash), despite the size of the state space ($\approx 16$ million unique states). An example of an uncertainty estimator that is viable in the presence of continuous state/action spaces is RND, with which EMCTS works well, see Figure 2 and Figure 4. For these reasons, we do not see a conflict between sampling-based MCTS approaches such as [1] and the extensions proposed by EMCTS, given an uncertainty estimator that is suited for the problem, similarly to other UCB-based exploration methods. Finally, note that MuZero [2] is very successful in relatively large action spaces (such as in Go) because of the search that combines value and prior policy. The same search properties apply to EMCTS, with a prior-exploration-policy (Equation 10 and line 255). An EMCTS agent that uses a prior-exploration-policy is evaluated in Figure 1, succsfully solving SUBLEQ instances much faster than the baseline.

The sensitivity of EMCTS to $\beta$ is very low in the sparse-reward domain of Deep Sea, which is demonstrated in the right subplot of Figure 2. Specifically, $\beta$ need only be large enough to induce deep exploration. This aligns with our expectations for a UCB-based exploration method in sparse-reward environments, where the main driver for performance is the ability to estimate the uncertainty well enough and use it to consistently explore novel states until the rewarding transition is observed.

How does EMCTS handle potential inconsistencies in epistemic uncertainty when the learned models (reward, value, transition) are not fully aligned? This is handled by the definition of $\mathbb V [\hat R]$ (line 216) which guarantees an upper confidence bound over the reward outside of the training distribution. The same idea is used in Equation 22 to train the estimator $\mathbb V [\hat V]$ to approximate a term that guarantees an upper bound over the value. Inside the training distribution, further exploration is not necessary and the epistemic uncertainty can safely drop the zero regardless of the extent to which the models are aligned in practice.

We hope this along with the main rebuttal addresses the reviewer's questions and concerns and are happy to clarify anything still left unclear.

[1] Hubert, Thomas, Julian Schrittwieser, Ioannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. "Learning and planning in complex action spaces." In International Conference on Machine Learning, pp. 4476-4486. PMLR, 2021.

[2] Schrittwieser, Julian, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez et al. "Mastering atari, go, chess and shogi by planning with a learned model." Nature 588, no. 7839 (2020): 604-609.

Dear reviewer,

Thank you for your positive comments and valuable time.

Though the bound doesn't apply to learned dynamics The bound is extended to learned dynamics in Appendix A.3. Please see our reply in the main rebuttal for additional details. We evaluate that this approach works in Figure 4. In Figure 2 we demonstrate that even without this modification, EMCTS with learned transition models works well in practice in Deep Sea.

The practical value of epistemic MCTS remains questionable. In what real-world tasks do we need to run search without knowing the world model? First, note that EMCTS is very well suited for tasks where the world model is known, i.e., when only the value model is learned, such as with AlphaZero [AZ, 1] as well as when it is not known. Many methods that use true world models such as AZ, which is very successful in the real-world (known world model tasks of programming and algorithm design [2,3]), rely on learned value functions which contain epistemic uncertainty. The epistemic uncertainty can be useful to direct the agent towards better novel policies, that are very hard to find using non-directed exploration. We demonstrate this on the task of programming in SUBLEQ in Figure 1, where agent has access to the true world model (i.e., both the transition dynamics as well as the reward dynamics) and EMCTS significantly outperforms baseline AZ which uses the same true model. By reasoning over the uncertainty in the learned value function, EMCTS is able explore the environment for solutions (programs) that the agent has not tried before much more efficiently.

A concrete example of a real world task where search is being used without a known world model is [4], where MuZero (fully learned, latent world model) improves over the engineered-by-hand previous state of the art in the real world task of video compression for YouTube videos. In Figures 2 and 4 we demonstrate that our method works in practice with the learned models of MuZero.

Please let us know if any open questions or unaddressed concerns remain.

[1] Silver, David, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot et al. "A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play." Science 362, no. 6419 (2018): 1140-1144.

[2] Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco J. R. Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, David Silver, Demis Hassabis, and Pushmeet Kohli. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature, 610(7930):47–53, 2022. doi: 10.1038/s41586-022-05172-4.

[3] Daniel J. Mankowitz, Andrea Michi, Anton Zhernov, Marco Gelmi, Marco Selvi, Cosmin Paduraru, Edouard Leurent, Shariq Iqbal, Jean-Baptiste Lespiau, Alex Ahern, Thomas Koppe, Kevin Millikin, Stephen Gaffney, Sophie Elster, Jackson Broshear, Chris Gamble, Kieran Milan, Robert Tung, Minjae Hwang, Taylan Cemgil, Mohammadamin Barekatain, Yujia Li, Amol Mandhane, Thomas Hubert, Julian Schrittwieser, Demis Hassabis, Pushmeet Kohli, Martin Riedmiller, Oriol Vinyals, and David Silver. Faster sorting algorithms discovered using deep reinforcement learning. Nature, 618(7964):257–263, 2023. doi: 10.1038/s41586-023-06004-9.

[4] Mandhane, Amol, Anton Zhernov, Maribeth Rauh, Chenjie Gu, Miaosen Wang, Flora Xue, Wendy Shang et al. "Muzero with self-competition for rate control in vp9 video compression." arXiv preprint arXiv:2202.06626 (2022).

Thank you for providing a rebuttal and for correcting me on learned dynamics. I did not look into the appendix but now I see it does extend to it. My evaluation was already positive. I would maintain my rating.