Epistemic Monte Carlo Tree Search

Authors: Yaniv Oren, Viliam Vadocz, Matthijs T. J. Spaan, Wendelin Boehmer
Primary area: reinforcement learning
Keywords: model based, epistemic uncertainty, exploration, planning, alphazero, muzero

Abstract:
The AlphaZero/MuZero (A/MZ) family of algorithms has achieved remarkable success across various challenging domains by integrating Monte Carlo Tree Search (MCTS) with learned models. Learned models introduce epistemic uncertainty, which is caused by learning from limited data and is useful for exploration in sparse reward environments. MCTS does not account for the propagation of this uncertainty however. To address this, we introduce Epistemic MCTS (EMCTS): a theoretically motivated approach to account for the epistemic uncertainty in search and harness the search for deep exploration. In the challenging sparse-reward task of writing code in the Assembly language SUBLEQ, AZ paired with our method achieves significantly higher sample efficiency over baseline AZ. Search with EMCTS 
solves variations of the commonly used hard-exploration benchmark Deep Sea - which baseline A/MZ are practically unable to solve - much faster than an otherwise equivalent method that does not use search for uncertainty estimation, demonstrating significant benefits from search for epistemic uncertainty estimation.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730717702014

thread_type:
official_review

<thread_object>
creator: reviewer_v4ih

<content>
summary:
The authors propose a novel method, Epistemic - Monte Carlo Tree Search (E-MCTS) which incorporates the uncertainty estimates into MCTS. The authors theoretically motivate their approach as well as provide an experimental study showing the benefit of using their method in hard exploration problems.

strengths:
Overall, the paper is generally well-written (except for a few places), the proposed method is novel and the contribution is significant.

MCTS is a technique which showed great results in relatively simple domains and it still remains a challenge to apply it in hard-to-explore domains. This work addresses this challenge by adding the uncertainty estimation mechanism into MCTS. The paper demonstrates good results in hard-to-explore domains.

weaknesses:
1. In abstract, lines 14-15, it is written "MCTS does not account for the propagation of this uncertainty however.". At this point of the text, it is not clear what it means to propagate the uncertainty and why it is necessary. It might be worth adding a sentence about "propagating uncertainty" before that.

2. The authors mention in many places that the model will be uncertain about the part of the state-action space which was not observed during training, like for example, lines 125-126 "thus their predictions are uncertain outside of the training set". But what about the situation when the model is actually **certain** about the unobserved part of the state-action space, but **wrong** ? I believe this situation is also likely to happen. There is no discussion of it at all in the paper.

3. The authors propose to estimate the square root of the variance $\sqrt{V[Q]}$ as $1/N \sum_{i=1}^{N} \sqrt{V[v]}$ (see, for example eq.11), but this is not how it generally is estimated. Instead, one should estimate  $\sqrt{V[Q]} = \sqrt{1/(N-1) \sum_{i=1}^{N} (v_i - \bar{v})^2}$, where $\bar{v}$ is the empirical average of $v$. Where does this estimate come from? What is the impact of using this estimate rather than a more correct one which I wrote?

4. The authors write "models are consistent" (line 269), but it is not clear what is meant by consistent. They also say in line 270, "V[V[s]] ~V[Vpi[s]]". It is unclear why this must be true. Also, the authors say "Imposing again the assumption that the learned models are consistent" (in lines 276-277). It is unclear what they mean by consistency. Suggestion -- write clearly (as an equaition) the definition of the consistency. Then, argue / discuss the implications of this assumption and why this assumption is true in the first place.

5. In lines 388-390, the authors say "achieve deep exploration", what does it mean? From the experiments, I can only conclude that the method solves these specific tasks, but it is wrong to claim that it "achieves deep exploration", since first, this "achieving" is not defined and second, it is not implied from the experiments that it will hold in general.

6. The authors evaluate their method in two domains. I think in order to increase the significance of their work, the authors should also evaluate on more easy-to-explore domains (i.e., in Atari) to show that their method works consistently well overall, as well as on more hard-to-explore domains (i.e., in Atari) to show that in situations where it should work well, it does it consistently. For example, would it also work well in domains like chess? Will the findings still hold? Will E-MCTS actually be more effective here compared to AZ?

questions:
See weaknesses for the questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730738755346

thread_type:
official_review

<thread_object>
creator: reviewer_yivx

<content>
summary:
This manuscript tries to solve an important issue in Monte Carlo Tree Search (MCTS), that MCTS doesn’t consider the propagation of the uncertainty of learned model. The authors proposed Epistemic Monte Carlo Tree Search (EMCTS), which can account for the epistemic uncertainty in search and harness the search for deep exploration. The proposed method is achieved in three steps: search with learned reward model, planning for exploration with epistemic search, and propagating epistemic uncertainty in search. EMCTS is evaluated in the challenging sparse-reward task of writing code in the Assembly language SUBLEQ and hard-exploration benchmark Deep Sea, and the results show that the proposed framework can solve the quiz faster than AlphaZero/Muzero.

strengths:
1.	This manuscript propose a novel and solid framework to introduce and deal with uncertainty of learned model caused by limited data or sparse reward environments, which is not taken into account in the standard MCTS.
2.	EMCTS solves the two-fold objective, i.e., extend MCTS to estimate and propagate the epistemic uncertainty from the uncertain learned model and harness the epistemic uncertainty in the search to achieve deep exploration of the environment, via formulating the epistemic uncertainty, and search based on the propagated UCB.
3.	By giving some assumption, the estimation and propagation of the epistemic uncertainty is defined by a random variable and computed via the variance of reward.
4.	Experiments demonstrate that the proposed EMCTS can solve the harder task in a much smaller number of samples than AZ.

weaknesses:
1.	Some assumptions might make EMCTS inapplicable to real complex problems. They assume that the transition model P is true. Theorem 1 need $Q^{\pi}$ to be linear. The learned models are assumed to be consistent. These is not reasonable in many cases.
2.	In many cases, the transition models need to be estimated, and then the EMCTS will not work.

questions:
1.	This method assumes that the reward is independent when $i \neq j$. Is this satisfied in programming or math task? for example, when $i = j+1$?
2.	How to get the true transition model in typical tasks?
3.	Is accuracy and uncertainty of the learned model is dynamically improved in the learning procedure? And is this will affect the consistent assumption?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731409584273

thread_type:
official_review

<thread_object>
creator: reviewer_acfq

<content>
summary:
This paper proposes Epistemic Monte-Carlo Tree Search (EMTCS) by introducing a method incorporating epistemic uncertainty in traditional MCTS, due to the limited coverage of the data. This paper upper bounds the uncertainty in the MCTS process and designs a novel planning scheme with this estimated uncertainty upper bound. This method makes the agent to explore the environments with sparse reward signals more efficiently. They did experiments on code generation task like SUBLEQ, which involves deep exploration. They also conducted experiments on the Deep Sea Benchmark. Extensive experiments show that EMCTS achieve better sample efficiency and deep exploration compared to baseline methods, particularly in tasks with sparse rewards.

strengths:
1. The paper provides a relatively rigorous theoretical foundation for EMCTS. The authors derive an upper confidence bound for the epistemic uncertain within the learned MCTS model, which theoretically justify the exploration strategy of EMCTS.

2. The experiments demonstrate the effectiveness of EMCTS algorithm in SUBLEQ and Deep Sea task. EMCTS show a better sample efficiency than other baseline methods. Notably, EMCTS demonstrate deep exploration in the presence of the learned value, reward, and transition models as well as both deterministic an stochastic rewards.

3. The motivation of EMCTS is good: to address the epistemic uncertainty in deep exploration tasks and incorporate the uncertainty not just in reward but also in transition dynamics.

weaknesses:
1. The authors assume consistency across the reward and value. This assumption might not hold in dynamic environments or with certain deep neural network structures in practice. To what degree do those assumptions hold in practice? Are there any experiments results showing that those assumptions hold (or nearly hold) in practice? How does EMCTS handle potential inconsistencies in epistemic uncertainty when the learned models (reward, value, transition) are not fully aligned?

2. The upper bound of the epistemic uncertainty only applies to the discrete state-action space and can introduce considerable computational costs if we are in a high dimensional space. Moreover, how does EMTCS handle larger state-action spaces or continuous state-action spaces? Experiments in the paper only show the effectiveness of EMCTS in case of discrete action space, but how can EMCTS (or some of its adapted variants) handle the continuous action space? Moreover, how sensitive is EMCTS to the exploration hyperparameter \beta?

questions:
/
</content>
</thread_object>

</thread>
<thread>
creation date:
1731443796174

thread_type:
official_review

<thread_object>
creator: reviewer_uotu

<content>
summary:
This work studies epistemic MCTS, which is to characterize the propagation of model uncertainty in tree search. This is motivated by tasks where the model dynamics are not known in advanced, but need to be learned instead. It could be also useful to host exploration methods.

The algorithm is inspired by theoretical view of linear Q, where the true optimal Q is bounded with high probability by the Q estimate plus the variance term $O\sqrt{V}$. Though the bound doesn't apply to learned dynamics, it works in the similar way as UCB, where the "inequality" provides intuition of the algorithm design. In the search step, the epistemic UCT then uses a similar bonus term in $O\sqrt{V}$ that is applied to the Q estimate.

This variance term is also propagated through the backup operator, as the authors intended to in their motivation of this work. Through an independency assumption, the variance could be backed up by variance terms of the nodes on the trajectory. With this backup, epistemic UCT could be applied to obtain epistemic MCTS.

The algorithm is tested on Subleq, which is a programming task, and Deep sea, which is a gridworld task. Combining with AlphaZero/MuZero, the algorithm performs quite well on the baseline tasks.

strengths:
1. The new algorithm solves MCTS with a learned model, with deep exploration and variance propagation.
2. The model is supported by theoretical insights.
3. The model performs quite well on the baseline tasks.

weaknesses:
1. The practical value of epistemic MCTS remains questionable. In what real-world tasks do we need to run search without knowing the world model?

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1731752882051

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear reviewers,

Thank you for your valuable time and many comments and questions.

**Learned transition models:** Reviewers yivx and Uotu correctly point out that Theorem 1 requires the transition model $ P $ to be known in advance, and raise the question whether EMCTS is applicable with learned transition models, such as in MuZero [1]. 

EMCTS is indeed applicable with learned transition models, which are discussed in Section 3.4.

In Appendix A.3 we propose a theoretically motivated method to construct an extended upper bound on estimates of $\mathbb V[\hat Q(s,a)]$ that accounts for possible uncertainty in the transition model. 
In Figure 4 (bottom subplots) we evaluate an agent that uses this method and a learned transition model, which successfully solves a challenging instance of Deep Sea.
In Figure 2 (left subplot) we find that in practice the unmodified upper bound already works well even with a learned transition model, demonstrated with the E-MZ agent, which uses a learned transition model and the unmodified upper confidence bound.
We will make sure it is emphasized in the paper that EMCTS *can* be used in practice in the presence of both learned and true transition models by referring to Appendix A.3 and Section 3.4 directly after Theorem 1, as well as describing the E-MZ agent in Figure 4 in more detail in the appendix.

**Consistency assumption:** The consistency assumption was raised as a concern by reviewers acfQ, yivx and v4iH.
Let us first clarify the consistency assumption required by EMCTS (line 269-270).
EMCTS requires that the learned value function $\hat V$ and reward function $\hat R$ are trained on the same data.
This is the popular choice in practice by standard methods such as MuZero [1].
In addition, EMCTS assumes that the epistemic uncertainty over the learned value $\mathbb V[\hat V]$ (Equation 6) is consistent with the uncertainty with the model used by MCTS $ \mathbb V[V^\pi_{\hat M}] $. That is, that they the are approximately the same $ \mathbb V[\hat V] \approx \mathbb V[V^\pi_{\hat M}] $ (line 270).

We believe that this assumption is reasonable in principle for two main reasons:
First, since $\hat V$ and $\hat R$ are trained on the same dataset the uncertainty in $\mathbb V[\hat V]$ analytically coincides with the uncertainty in $\mathbb V[\hat R]$ (Equation 5).
Second, the estimator $\hat u $ based on the uncertainty Bellman equation [2], estimates $ \hat u \approx \mathbb V[\hat V]$ using estimates of $ \mathbb V[\hat R] $ (Equation 6), and for that reason $\hat u$ is trained to operate as an estimate that is consistent with both $ \mathbb V[V_{\hat M}^\pi], \mathbb V[\hat V]$.

In practice, inside the training distribution, the TD-learning based training of $\hat u$ is well established and is empirically observed to produce useful estimates of the epistemic uncertainty in practice (for example, [3, 4]).
Outside of the training distribution, this is addressed in Appendix D.6, by explicitly upper-bounding predictions where $\hat u$ is evaluated outside of the training distribution.
Our experiments corroborate the effectiveness of this approach, and demonstrate robustness of EMCTS to possible errors in uncertainty estimation: 
Figures 1, 2 and 4 demonstrate that EMCTS with $\hat u$ work well across multiple environments, sources of uncertainty, different architectures and variations of learned models.
Finally, in Figure 3 we visualize the uncertainty estimated by EMCTS compared to the uncertainty estimated by one sample of $\hat u$, where EMCTS matches the ground truth epistemic uncertainty much better.

We will make sure this is described more clearly in the paper.

We have uploaded a new version with the discussed changes (marked in green) and additional results on easy-exploration environments requested by reviewer v4iH (Figure 5 in the new version).
We hope that this discussion along with the individual replies addresses the reviewers' concerns, and we are looking forward to further discussion with the reviewers.

[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.

[2] O’Donoghue, Brendan, Ian Osband, Remi Munos, and Volodymyr Mnih. "The uncertainty Bellman equation and exploration." In International Conference on Machine Learning, pp. 3836-3845. 2018.

[3] Moritz Akiya Zanger, Wendelin Boehmer, and Matthijs T. J. Spaan. Diverse projection ensembles for distributional reinforcement learning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.

[4] Luis, Carlos E., Alessandro G. Bottero, Julia Vinogradska, Felix Berkenkamp, and Jan Peters. "Model-based uncertainty in value functions." In International Conference on Artificial Intelligence and Statistics, pp. 8029-8052. PMLR, 2023.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The work studies epistemic MCTS, in particular the propagation of model uncertainty in the tree search.
All reviewers appreciated the solid theoretical contribution as well as the algorithmic innovation; both are recognized as being substantial. 
The experiments in hard to explore environment show the advantage of the proposed methods.
The reviewers have some comments, which can be taken as an opportunity to clarify certain aspects of the work, and they can be addressed in the writing of an updated version.