Epistemic Monte Carlo Tree Search

## Abstract

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.

## 1 Introduction

Many recent successes of reinforcement learning (RL) have been achieved by the model-based algorithm family of AlphaZero/MuZero (A/MZ, Silver et al., ; Schrittwieser et al., ).
A/MZ have outperformed humans in games, in tasks that traditionally relied on intricate human engineering (Mandhane et al., ) and even made real world impact with the design of novel, more efficient algorithms (Fawzi et al., ; Mankowitz et al., ) for day-to-day problems, a task that is often formulated as a challenging sparse reward environment.
When rewards are sparse, it is difficult to learn good policies without employing some form of deep exploration to search for the rewards.
Deep exploration refers to the ability of the agent to direct itself towards novel transitions in the environment irrespective of how far away they are from the current state and promises up to exponential improvement in sample efficiency in sparse reward environments (Osband et al., ).

At the core of A/MZ is the combination of Monte Carlo Tree Search (MCTS, Swiechowski et al., ) with learned models of value and/or environment dynamics.
Learned models introduce epistemic uncertainty, which refers to the uncertainty in the predictions of the model sourced in limited coverage of the state-action space during training (Hüllermeier & Waegeman, ).
Accounting for epistemic uncertainty allows the agent to discern between predictions that are based on evidence (i.e. the learned predictor was trained on this input), or based on generalization (i.e. the learned predictor was not trained on this input) and is useful for many purposes in online and offline RL.
Common uses range from reducing overestimation errors through pessimism in the face of uncertainty (see Kumar et al., ), to directing deep exploration through optimism in the face of uncertainty (see Neu & Pike-Burke, ).
Harnessing search for exploration is a popular approach in similar model based algorithms, such as Dreamer (Sekar et al., ).

MCTS, however, was designed for search with the true dynamics model and without a value model, and as a result, does not account for epistemic uncertainty introduced from learning the models.
For this reason, A/MZ cannot harness MCTS for deep exploration, nor benefit from the epistemic uncertainty associated with the predictions of the search tree in other ways.
In this work we aim to address both, with three main contributions.
Practical and theoretically motivated methods to
(i) harness MCTS and epistemic uncertainty for upper-confidence-bound-based deep exploration (Jin et al., ) and
(ii) propagate the epistemic uncertainty from learned models of value and/or dynamics during search, which we call Epistemic MCTS (EMCTS).
(iii) A parallelized implementation in JAX (Bradbury et al., ) of EMCTS paired with an AZ agent and an environment implementing the Assembly language subleq
(Mazonka & Kolodin, )(^1^11 [sic-1](https://jaredkrinke.itch.io/sic-1) is an open source game demonstrating subleq.).
We find that the propagation of epistemic uncertainty in EMCTS is very similar to that of value in MCTS.
As search with MCTS improves the value estimates at the root, we hypothesize that search with EMCTS similarly improves the epistemic uncertainty estimates at the root, resulting in more accurate UCBs and more sample efficient exploration compared to a method that does not rely on search but is otherwise equivalent.

We evaluate EMCTS in the challenging, similar to real-world applications and sparse-reward task of programming in subleq, as well as in the commonly used hard-exploration benchmark Deep Sea (Osband et al., ).
Our method is able to find correct programs for a harder programming task in a much smaller number of samples than the AZ baseline.
In the Deep Sea benchmark, our method demonstrates deep exploration by solving stochastic and deterministic reward variations of the task, both of which baseline A/MZ is unable to solve in a reasonable number of samples.
In addition, EMCTS significantly outperforms an ablation that does not rely on search for epistemic uncertainty estimation but is otherwise equivalent, demonstrating significant advantages from search for uncertainty estimation.

## 2 Background

In RL, an agent learns a behavior policy through interactions with an environment.
The environment is represented by a Markov Decision Process (MDP, Bellman, ) $\mathcal{M}=\langle\mathcal{S},\mathcal{A},\rho,\mathcal{R},P,\gamma\rangle$, where $\mathcal{S}$ is a set of states, $\mathcal{A}$ a set of actions, $\rho$ the initial state distribution,
$\mathcal{R}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ a bounded possibly stochastic reward function, and $P$ is a transition distribution such that $P(s^{\prime}|s,a)$ specifies the probability of transitioning from state $s$ to state $s^{\prime}$ after executing action $a$.
The objective $J_{\pi}$ of the agent is to find a policy $\pi(a|s)$, specifying the probability of selecting action $a$ in state $s$, that maximizes the expected discounted return, also called value $V^{\pi}$, from the starting state distribution $\rho$:

$$ $\displaystyle J_{\pi}=\mathbb{E}[V^{\pi}(s_{0})|\!\!\begin{array}[]{c} \scriptstyle s_{0}\sim\rho\end{array}\!\!\,]=\mathbb{E}\Big{[}{\textstyle\sum \limits_{t=0}^{H-1}}\gamma^{t}\mathcal{R}(s_{t},a_{t})\Big{|}\!\!\begin{array} []{c}\scriptstyle s_{0}\sim\rho,s_{t+1}\sim P(s_{t},a_{t}),a_{t}\sim\pi(s_{t}) \end{array}\!\!\!\Big{]}.$ (3) $$

The discount factor $0<\gamma<1$ is used in infinite-horizon MDPs,
i.e. $H=\infty$,
to guarantee that the values remain bounded, and is commonly used in RL for training stability.
A state-action Q-value function is also often used:
$Q^{\pi}(s,a)=\mathbb{E}[\mathcal{R}(s,a)+\gamma V^{\pi}(s^{\prime})|\!\!\begin
{array}[]{c}\scriptstyle s^{\prime}\sim P(s,a)\end{array}\!\!\!]$.
We denote the value of the optimal policy $\pi^{*}$ with $V^{*}(s)=\max_{\pi}V^{\pi}(s),\forall s\in\mathcal{S}$.
In offline RL, the agent must maximize this objective given a static dataset.
In model-based RL (MBRL) the agent uses a model of the dynamics of the environment $(P,\mathcal{R})$ to optimize its policy, often through planning (Moerland et al., ).
The dynamics are either learned from interactions (e.g., in MZ, Schrittwieser et al., ) or provided (e.g., in AZ, Silver et al., ).
In Deep MBRL the agent utilizes deep neural networks (DNN, Goodfellow et al., ) to approximate any of the value, policy, reward and transition functions.

### 2.1 Monte Carlo Tree Search

The MCTS algorithm constructs a tree with the current state $s_{t}$ at its root to estimate the objective: $\operatorname*{arg\,max}_{a}\max_{\pi}Q^{\pi}(s_{t},a)$ (Browne et al., ), by iteratively performing selection, expansion, simulation and backup.
At each iteration $i$ of the algorithm a trajectory in the tree is selected using a tree search policy such as UCT (Kocsis & Szepesvári, ):

$$ $\displaystyle a^{i}_{k}\;=\;UCT_{i}(s_{k})\;=\;\operatorname*{arg\,max}_{a\in A }q^{i}(s_{k},a)+C_{UCT}{\textstyle\sqrt{\frac{2\log(\sum_{a^{\prime}}\!N(s_{k} ,a^{\prime}))}{N(s_{k},a)}}},$ (4) $$

where $N(s_{k},a)$ denotes the number of times action $a$ has been selected in node $s_{k}$, $C_{UCT}>0$ trades off exploration of new nodes in the tree with maximizing observed return and $q^{i}(s_{k},a)$ is the averaged return observed for this state action up to step $i$.
Modern algorithms (such as A/MZ) use instead variations of PUCT (Rosin, ):

$$ $\displaystyle a^{i}_{k}\;=\;PUCT_{i}(s_{k})\;=\;\operatorname*{arg\,max}_{a\in A }q^{i}(s_{k},a)+\pi(a|s_{k})\,C_{PUCT}{\textstyle\frac{\sqrt{\sum_{a^{\prime}} \!N(s_{k},a^{\prime})}}{1+N(s_{k},a)}},$ (5) $$

with some learned prior-policy $\pi$.
When selection step $i$ arrives at a leaf $s^{i}_{T}$ the node is expanded with a value estimate $v(s^{i}_{T})$.
MCTS propagates the return $\nu^{i}$ of planning step $i$ back along the planning trajectory
$\{s^{i}_{j},a^{i}_{j}\}_{j=0}^{T}$:

$$ $\displaystyle\nu^{i}(s_{k},a^{i}_{k})=\sum_{j=0}^{T-1}\gamma^{j}r(s^{i}_{k+j}, a^{i}_{k+j})+\gamma^{T}v(s^{i}_{T}),\quad q^{i}(s_{k},a)=\frac{1}{N(s_{k},a)} \sum_{j=1}^{i}\nu^{j}(s_{k},a)$ (6) $$

where $\textstyle{a^{i}_{k+j}=P/UCT_{i}(s^{i}_{k+j})}$ and $\textstyle{s^{i}_{k+j+1}=f(s_{k+j}^{i},a^{i}_{k+j})}$
for a deterministic transition function(^2^22 For notational simplicity we assume here a deterministic transition function $f$, so that nodes $s^{i}_{k+j}$ correspond to individual states. It is also possible to use a stochastic transition model $P$, where nodes correspond to distributions of states, which are sampled by the selection step. The above equations remain the same, but the state corresponding to $s^{i}_{k+j}$ is effectively a random variable.) $f:\mathcal{S}\times\mathcal{A}\to\mathcal{S}$ and a mean-reward function $r(s,a)=\mathbb{E}[\mathcal{R}(s,a)]$.

In the original MCTS, the dynamics model $m=(f,r)$ is provided to the agent and assumed to be correct for the true MDP $\mathcal{M}$ and the value estimate $v(s^{i}_{T})$ is computed using Monte-Carlo (MC) rollouts with $m$.
In AZ the value function (Silver et al., ) and in MZ all functions $f,r,v$, are learned from interactions with the environment
and thus their predictions are uncertain outside of the training set.
The prior-policy $\pi$ is trained with cross-entropy loss on targets extracted from the root of the tree.
With Reanalyze (Schrittwieser et al., ), which uses search to generate fresh value and policy targets from off-policy data, A/MZ are able to learn off-policy.

### 2.2 Quantifying Uncertainty in Deep Reinforcement Learning

Quantifying uncertainty in deep learning is an active field of research (see Hüllermeier & Waegeman, ; Lockwood & Si, ).
In this work, we take the common approach for quantifying epistemic uncertainty as the variance in a random variable that approximates predictions that are consistent with a dataset of observed interactions $\mathcal{D}$. We model a learned function $\hat{r}$, trained to approximate a true function $r$ on data $\mathcal{D}$, $\hat{r}(s,a)\approx r(s,a)$, as a random variable with respect to the data: $\hat{r}(s,a)|\mathcal{D}:=\hat{R}(s,a)$.
We assume unbiased approximation
$\mathbb{E}[\hat{R}(s,a)]:=r(s,a)$,
and use the variance $\mathbb{V}\!\,[\hat{R}(s,a)]$ to quantify the epistemic uncertainty.
In discrete state-action spaces, the epistemic uncertainty in a reward or a transition can be estimated as the novelty of a state or state-action pair using visitation counting.
Exact counting, which is generally intractable in large state-action spaces, can be replaced by Hash-based counting (Tang et al., ), or methods that focus on direct estimation of the novelty of state-action pairs, such Random Network Distillation (RND, Burda et al., ).

In contrast, epistemic uncertainty in a value prediction $\mathbb{V}\!\,[\hat{V}(s)]$ that is trained with TD-based targets, requires reasoning about uncertainty in the value-targets, as well as whether $\hat{V}$ was trained on $s$ at all.
We follow the popular approach by Strens () of defining epistemic uncertainty in the value function as the variance of a Bayesian posterior of the Q-values of a policy conditioned on the data the agent has collected, as follows:

$$ $\displaystyle\mathbb{V}\!\,[\hat{Q}^{\pi}(s,a)]=\mathbb{V}\!\,\Big{[}\hat{R}(s ,a)+\gamma{\textstyle\sum\limits_{s^{\prime},a^{\prime}}}\pi(a^{\prime}|s^{ \prime})\,\hat{P}(s^{\prime}|s,a)\,\hat{Q}^{\pi}(s^{\prime},a^{\prime})\Big{]}\,.$ (7) $$

Note that the epistemic uncertainty about rewards $\mathbb{V}\!\,[\hat{R}]$ and transitions $\mathbb{V}\!\,[\hat{P}]$ are induced by the data $\mathcal{D}$ the value function has been trained on, not by the agent’s knowledge of the environment. Values predicted by learned value functions can therefore be epistemically uncertain, even if a perfect model of the environment is known to the agent.
In their work on the Uncertainty Bellman Equation (UBE), O’Donoghue et al. () propose to approximate an upper bound $u^{\pi}(s_{t},a_{t})\geq\mathbb{V}\!\,[\hat{Q}^{\pi}(s_{t},a_{t})]$.
$u^{\pi}(s_{t},a_{t})$ can be learned with (possibly $n$-step) TD targets in a similar manner to value learning from local uncertainties $\eta(s,a)$:

$$ $u^{\pi}(s_{t},a_{t}):=\eta(s_{t},a_{t})+\gamma^{2}{\textstyle\sum_{a^{\prime}} }\pi(a^{\prime}|s_{t+1})\,u^{\pi}(s_{t+1},a^{\prime})\leq\eta(s_{t},a_{t})+ \gamma^{2}\max_{a^{\prime}}u^{\pi}(s_{t+1},a^{\prime}).$ (8) $$

The local uncertainty $\eta$ can be derived from $\mathbb{V}\!\,[\hat{R}]$ and $\mathbb{V}\!\,[\hat{P}]$.
Note that the above inequality yields an upper bound on the epistemic uncertainty of any policy $\pi$ with training data $\mathcal{D}$.

### 2.3 Deep Exploration with Upper Confidence Bounds

A popular approach to harness epistemic uncertainty for deep exploration is that of optimism in the face of uncertainty (Lattimore & Szepesvari, ), often formalized using an upper confidence bound (UCB) on the true value of the optimal policy (Azar et al., ; Jin et al., ).
Acting greedily with respect to the maximum UCB guarantees exploration of the environment, as long as the UCB tightens with more samples from the environment.
The efficiency of UCB exploration depends on the epistemic uncertainty estimator.
In environments with small state-action spaces, epistemic uncertainty estimators that predict maximum uncertainty for unvisited states can be used to guarantee exploration of every state-action pair.
To explore continuous or large state-action domains in practical numbers of samples, using an UCB requires uncertainty estimators that estimate the epistemic uncertainty in unvisited states $s^{\prime}$ based on a similarity metric between $s^{\prime}$ and visited states $s\in\mathcal{D}$ (Jin et al., ), such as imposed by RND, for example.

## 3 Deep Exploration with Epistemic MCTS

To find an optimal policy, an RL agent must explore the environment until all necessary information is gathered.
When the agent uses search for action selection in the environment, it can search for exploration, exploitation, or a trade-off between the two.
When a learned model $\hat{m}$ is used in search, the agent faces a problem: the values $q^{i}_{\hat{m}}$ in the search tree converge to $Q_{\hat{m}}^{*_{\hat{m}}}$, the $Q$ value in the learned model of the policy that is optimal in that model $\pi^{*}_{\hat{m}}$.
In areas where the learned model is inaccurate, the search may lead to arbitrarily bad actions with respect to expected return in the true environment.
From the perspective of exploration on the other hand, this presents an opportunity: by estimating the epistemic uncertainty the agent can identify the areas where the model is uncertain due to insufficient interactions, and use the uncertainty to direct exploration into these areas.

Our objective is then two-fold: To 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.
To achieve this, we take the following steps:
(i) Formulate the learned model $\hat{m}$ as a random variable and use it to construct a UCB on $Q^{*}$ (Section 3.1).
(ii) Propose search policies to track the maximum UCB exploration objective (Section 3.2).
(iii) Propagate the epistemic uncertainty through search, such that the epistemic uncertainty $\mathbb{V}\!\,[q_{\hat{m}}^{i}(s,a)]$ in nodes’ value predictions $q_{\hat{m}}^{i}(s,a)$ can be estimated (Section 3.3).
Search with learned transition models introduces additional challenges for uncertainty propagation, which we address in Section 3.4.

### 3.1 Search with a Learned Reward Model

For simplicity, we begin by formulating the problem of search with learned models only in terms of a learned reward model and later extend the setup to learned value and transition models.
Consider the following setting: the agent has access to a dataset $\mathcal{D}=\{(s_{i},a_{i},r_{i},s_{i+1})\,|\,0\leq i<N\}$ of transitions and rewards,
as well as to the true (possibly stochastic) transition model $P$.
The reward function is bounded with: $\textstyle{\max_{s,a\in\mathcal{S}\times A}|\mathcal{R}(s,a)|\leq r_{max}}$.
We define the uncertain model as a random variable $\hat{M}=(P,\hat{R})$.
The uncertain mean-reward function is defined as a random variable $\hat{R}$ in the standard manner for defining an epistemically uncertain model (see Section 2.2), such that
$\mathbb{E}[\hat{R}(s,a)]=r(s,a)$ the mean reward, and
$\hat{R}(s,a)=\frac{1}{|(s,a,\cdot,\cdot)\in D|}\sum_{(s,a,r,\cdot)\in D}r
\approx r(s,a)$, as the empirical mean of observed rewards.
The epistemic uncertainty in reward prediction is defined as the variance $\mathbb{V}\!\,[\hat{R}(s,a)]$.
We note that $\hat{R}(s,a)$ is a bounded random variable over the interval
$[-r_{max},r_{max}]$.
To facilitate constructing an upper confidence bound on $Q^{*}$ we
define $\mathbb{V}\!\,[\hat{R}(s,a)]:=r_{max}^{2},\,\forall(s,a)\notin\mathcal{D}$. That is, for unobserved transitions, the maximum variance possible for a bounded random variable (Popoviciu, ).
The value with respect to the Markov chain induced by the random variable $\hat{M}$ is then itself a random variable:

$$ $\displaystyle Q^{\pi}_{\hat{M}}(s,a):=\mathbb{E}_{\pi,P}\Big{[}\sum_{i=0}^{ \infty}\gamma^{i}r_{i}\,\Big{|}\,a_{i}\sim\pi(s_{i}),s_{i+1}\sim P(s_{i},a_{i} ),r_{i}=\hat{R}(s_{i},a_{i}),s_{0}=s,a_{0}=a\Big{]},$ (9) $$

where the expectation is with respect to $\pi$ and $P$ and not $\hat{R}$.
Using this definition for the value in the model $Q^{\pi}_{\hat{M}}$ and its variance $\mathbb{V}\!\,[Q^{\pi}_{\hat{M}}]$ we can construct an upper confidence bound on $Q^{*}$ for deep exploration.

###### Theorem 1

For $\hat{M}$, $\mathcal{M}$, $Q^{*}$, $Q^{\pi}_{\hat{M}}$ defined as above and $\delta\in(0,1]$:

$$ $\displaystyle P\bigg{(}Q^{*}(s,a)\leq\max_{\pi}\Big{(}Q_{\hat{M}}^{\pi}(s,a)+{ \textstyle\frac{1}{\sqrt{\delta}}}\sqrt{\mathbb{V}\!\,[Q_{\hat{M}}^{\pi}(s,a)] }\Big{)}\bigg{)}\,\,\geq\,\,1-\delta\,,\quad\forall(s,a)\in\mathcal{S}\times \mathcal{A}.$ (10) $$

Theorem 1 relies on the linearity of $Q^{\pi}$ in the reward function and for that reason does not immediately extend to learned transition models.
The proof is provided in Appendix A.1.
This UCB can be tracked by MCTS using epistemic search policies, which we propose next.

###### Theorem 1

### 3.2 Planning for Exploration with Epistemic Search

To harness search for deep exploration through the approximation of the UCB in Equation 10, we propose Epistemic P/UCT (EP/UCT, changes are marked in blue):

$$ $\displaystyle a_{\text{{\color[rgb]{0,0,1}E}UCT}}\;=\;\operatorname*{arg\,max} _{a\in A}{\color[rgb]{0,0,1}q^{\beta}_{\hat{M}}(s_{k},a)}+C_{\text{UCT}}{ \textstyle\sqrt{\frac{2\log(\sum_{a^{\prime}}\!N(s_{k},a^{\prime}))}{N(s_{k},a )}}},$ (11) $\displaystyle a_{\text{{\color[rgb]{0,0,1}E}PUCT}}\;=\;\operatorname*{arg\,max }_{a\in A}{\color[rgb]{0,0,1}q^{\beta}_{\hat{M}}(s_{k},a)}+\pi(a|s_{k})\,C_{ \text{PUCT}}{\textstyle\frac{\sqrt{\sum_{a^{\prime}}\!N(s_{k},a^{\prime})}}{1+ N(s_{k},a)}},$ (12) $$

where we define:

$$ $\displaystyle{\color[rgb]{0,0,1}q^{\beta}_{\hat{M}}(s_{k},a)}:=\underbrace{ \textstyle\frac{1}{N(s_{k},a)}\sum_{i=0}^{N(s_{k},a)}\nu_{\hat{M}}^{i}(s_{k},a )}_{:=q_{\hat{M}}(s_{k},a)}+{\color[rgb]{0,0,1}\beta\underbrace{\textstyle \frac{1}{N(s_{k},a)}\sum_{i=0}^{N(s_{k},a)}\sqrt{\mathbb{V}\!\,[\nu_{\hat{M}}^ {i}(s_{k},a)]}}_{:=\sigma_{q_{\hat{M}}}(s_{k},a)\geq\sqrt{\mathbb{V}\!\,[q_{ \hat{M}}(s_{k},a)]}}}.$ (13) $$

The hyper-parameter $\beta\geq 0$ can be tuned per task, or chosen to guarantee an upper confidence bound with specific confidence $1-\delta$, and
$\textstyle{\nu_{\hat{M}}^{i}(s_{k},a)}$ is
as defined in Equation 6 for a specific model $\textstyle{\hat{M}}$.
We suppress the dependence on planning step $i$ in the notation of $\textstyle{q^{\beta}_{\hat{M}}(s_{k},a)}$ for simplicity.
To maintain the property of PUCT, which assumes $q_{\hat{M}}$ is between 0 and 1 (Rosin, ), one can use the Q normalization approach proposed by Schrittwieser et al. () to normalize the $q^{\beta}$ scores.
More modern variations of PUCT such as those proposed by Danihelka et al. () can be used by replacing $q_{\hat{M}}$ with $q^{\beta}_{\hat{M}}$ in the search objective.
We note that in order to properly harness search policies that use a prior-policy for deep exploration, such as EPUCT, the agent must learn an exploration-prior-policy $\pi_{e}$ that is trained with an exploratory objective.
Otherwise, the prior policy may direct the search away from the exploratory planning objective.
To enable EP/UCT, in the following section we propose methodology to estimate the epistemic uncertainty in backups and upper bound $\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]$.

### 3.3 Propagating Epistemic Uncertainty in Search

To estimate $\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]$ we need to achieve the following: (i) extend the problem setup to account for the learned value model, (ii) compute or upper bound the epistemic uncertainty in one backup $\mathbb{V}\!\,[\nu_{\hat{M}}^{i}(s_{k},a)]$, and (iii) compute or upper bound $\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]$ using $\mathbb{V}\!\,[\nu_{\hat{M}}^{i}(s_{k},a)]$.

Search with a Learned Value Model
A learned value model $\hat{V}$ introduces an additional source of epistemic uncertainty into the search, $\mathbb{V}\!\,[\hat{V}(s)]$.
Most often, all learned models are trained on the same dataset with similar losses and network architecture, and for that reason, it is reasonable to assume that the models are consistent.
That is, if $\hat{M}=(f,\hat{R},\hat{V})$ then $\mathbb{V}\!\,[\hat{V}(s)]\approx\mathbb{V}\!\,[V_{\hat{M}}^{\pi}(s)]$.
When the reward model is not learned, such as in AZ, the value model still has epistemic uncertainty as it is trained from interactions.
For this reason $\mathbb{V}\!\,[\hat{V}(s)]$ coincides with the definition of value uncertainty in model-free literature (7).
This definition accounts for uncertainty in transitions that are unobserved in the environment, regardless of whether the transitions are known to the planning model, capturing the uncertainty in a value model that is trained from interactions.
Imposing again the assumption that the learned models are consistent, more generally we have instead $\mathbb{V}\!\,[\hat{V}(s)]\geq\mathbb{V}\!\,[V_{\hat{M}}^{\pi}(s)]$ and we can reason about uncertainty in the MCTS tree both when the reward model is learned (e.g. MZ) as well as when it is not (e.g. AZ).
In the following we assume that both $\mathbb{V}\!\,[\hat{R}(s,a)]$ and $\mathbb{V}\!\,[\hat{V}(s)]$ can be estimated by the agent.
In practice we use approximators such as RND and counting methods for the former, and UBE for the latter.

Epistemic Uncertainty of the Backup
The epistemic uncertainty in one backup step $\mathbb{V}\!\,[\nu_{\hat{M}}^{i}(s_{k},a)]$ starting at node $s_{k}$ and choosing action $a$ can be formulated as follows:

$$ $\displaystyle\mathbb{V}\!\,[\nu_{\hat{M}}^{i}(s_{k},a)]=\sum_{j=0}^{T-k-1} \gamma^{2j}\mathbb{V}\!\,[\hat{R}(s^{i}_{k+j},a^{i}_{k+j})]+\gamma^{2(T-k)} \mathbb{V}\!\,[\hat{V}(s^{i}_{T})],$ (14) $$

under the assumption that $\hat{R}(s_{k+i},a)$ is independent from $\hat{R}(s_{k+j},a)\,,\forall i\neq j,$ and $\hat{V}$ is independent from $\hat{R}$ conditional on the data.
We note that this assumption might be reasonable for some estimators (i.e., a table), but is not generally assumed to be true for DNNs.
To avoid this assumption, we can instead upper bound the standard deviation of correlated backups using
the sum of standard deviations
(see Appendix A.2).

Epistemic Uncertainty of Node Values
Since A/MZ use the same model $\hat{M}$ throughout planning, the returns predicted for different backups cannot be assumed to be de-correlated.
We propose to upper bound the epistemic uncertainty instead:

$$ $\displaystyle\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]\leq\sigma^{2}_{q_{\hat{M}}}( s_{k},a):=\bigg{(}\frac{1}{N(s_{k},a)}\sum_{i=0}^{N(s_{k},a)}\sqrt{\mathbb{V} \!\,[\nu^{i}(s_{k},a)]}\bigg{)}^{\!\!2},$ (15) $$

where $N(s_{k},a)$ is the number of visitations to action $a$ at node $s_{k}$.
We provide a complete derivation in Appendix A.2.
Equation 15 completes the Epistemic MCTS (EMCTS) algorithm for learned value and/or reward models.
See Algorithm 1 for pseudo code of EMCTS with EUCT, where we suppress dependence on the model $\hat{M}$ for notation simplicity.
Extensions to MCTS are marked in blue.
In AZ, with a true reward model, EMCTS will use $\mathbb{V}\!\,[\hat{R}(s,a)]=0,\,\forall(s,a)\in\mathcal{S}\times\mathcal{A}$, but still use $\mathbb{V}\!\,[\hat{V}(s)]\geq 0$.

To conclude this section we note that while in this work we motivate and later evaluate EMCTS from the perspective of deep exploration, EMCTS is not limited to this use case.
Our method introduces to MCTS-based algorithms such as A/MZ a novel capability to estimate the epistemic uncertainty in the value predictions during and post search, which may be attractive for different purposes.
Noted example are:
(i) reducing over estimation errors with lower-bound value targets of the form $q_{\hat{M}}(s,a)-\beta\sqrt{\mathbb{V}\!\,[q_{\hat{M}}(s,a)]}$. This approach is popular in model free AC methods implicitly (Fujimoto et al., ; Haarnoja et al., ) and explicitly (Ciosek et al., ), as well as in model-based RL (Zhou et al., ) and is generally very common in offline RL (Kumar et al., ; Ghasemipour et al., ). This makes EMCTS an espeically attractive candidate to enhance A/MZ’s Reanalyze for offline-RL or off-policy target generation (Schrittwieser et al., ).
We propose to track the objective $q_{\hat{M}}(s,a)-\beta\sqrt{\mathbb{V}\!\,[q_{\hat{M}}(s,a)]}$ in search, i.e. to use EP/UCT but with $\beta<0$.
(ii) Weighting value and policy losses by the estimate of the epistemic uncertainty in the value which was successful in online as well as offline RL (Lee et al., ; Wu et al., ).
The uncertainty of the value of the root of EMCTS can be used for this purpose.

Figure: Algorithm 1 EMCTS with EUCT. Requires $f,r,v^{\pi}$ and uncertainty estimators $\mathbb{V}\!\,[\hat{V}],\mathbb{V}\!\,[\hat{R}]$

### 3.4 Search with a Learned Transition Model

Learned transition models introduce several challenges from the perspective of EMCTS:
(i) estimating epistemic uncertainty in the possibly-abstracted planning space,
(ii) propagating the uncertainty forward during search as future transitions become less certain, in addition to backwards, and (iii) propagating it in such a way that maintains the UCB constructed in Theorem 1.
Multiple methods to overcome (i) have been successful in previous works, see Henaff (); Sekar et al. ().
In Appendix A.3 we propose an approach to overcome (ii) and (iii) as well as discuss the challenges in more detail.
In practice however, in Section 5 we include experiments where a MZ agent fitted with a reliable uncertainty estimator for $\mathbb{V}\!\,[\hat{R}(s,a)]$ successfully demonstrates deep exploration without accounting for (ii) and (iii).

## 4 Related Work

Search for exploration was used successfully in a number of previous works, see
Jaksch et al. (),
Sun et al. (), Hester & Stone (), Shyam et al. (), Henaff (),
Sekar et al. (), Lambert et al. () and
Luis et al. ().
We add to this line of work EMCTS:
designed for MCTS (and planning trees in general), practical and based in theory.
Tesauro et al. () develop a Bayesian approach for aleatoric uncertainty propagation in MCTS.
POMCP (Silver & Veness, ), POMCPOW (Sunberg & Kochenderfer, ) and BOMCP (Mern et al., ) extend MCTS to POMDPs with a probabilistic Bayesian belief state at the nodes using a probabilistic model, while Stochastic MuZero (Antonoglou et al., ) extends MuZero to the stochastic transitions setting by replacing the deterministic transition function with a Vector Quantised Variational AutoEncoder (van den Oord et al., ).
In these works, epistemic uncertainty is not distinguished or used for exploration.
Latent disagreement ensembles (Lakshminarayanan et al., ; Ramesh et al., ) offer a popular alternative to RND and counts.
Wasserstein Temporal Difference (WTD, Metelli et al., ) offers an alternative to UBE (O’Donoghue et al., ) for estimating value uncertainty, using Wasserstein Barycenters (Agueh & Carlier, ) to update a posterior over $Q$ functions in place of a standard Bayesian update.
While UBE was criticized by Janz et al. () for being insufficient for deep exploration with posterior-sampling based RL (PSRL, Osband et al., ), we note that the same shortcomings do not apply when UBE is used for UCB-based exploration.

## 5 Experiments

The task of writing code and finding new algorithms, where AZ has recently made world real-world impact (Fawzi et al., ; Mankowitz et al., ), is natural to formulate using sparse-reward environments where the actions of the agent are operations and reward is received when a correct program is completed.
Such environments represent a hard exploration challenge: the state space is often exponential $|\mathcal{A}|^{L}$ in the actions $\mathcal{A}$ and the maximum length of the program $L$,
and while the number of possible solutions is generally unknown, it is in most cases very small compared to the number of possible sequences of operations.
To evaluate the contribution of EMCTS in such real-world tasks we conduct experiments in the one-instruction Assembly programming language subleq (Section 5.1).

To verify that EMCTS demonstrates deep exploration and benefits from search, we conduct experiments in the commonly used hard-exploration benchmark Deep Sea (Osband et al., ) (Section 5.2).
Specifically, we are interested in the following:
RQ I
Does EMCTS paired with an epistemic search policy achieve deep exploration with both AZ as well as MZ and in the presence of stochastic as well as deterministic rewards?
RQ II
Is there benefit in search with EMCTS for deep exploration, compared to an otherwise equivalent approach that does not use search, and if so, is the benefit retained in the presence of the learned transition model of MZ?
We also conduct an ablation study on the exploration parameter $\beta$ to verify that the agent can learn stably from the possibly very-off-policy exploratory data provided by deep exploration.

### 5.1 subleq Experiments

subleq is a Turing-complete (excepting for finite memory) one-instruction programming language (Mazonka & Kolodin, ).
The subleq operation receives three memory addresses as inputs
A, B, C, and executes
mem[A] = mem[A] - mem[B]. If
mem[A] <= 0, the execution branches to address
C (In subleq the program is saved in the same memory it modifies).
Because there is only one instruction, writing code in subleq summarizes to writing a sequence of memory addresses.
We model the action space with $\mathcal{A}=\{0,\dots,N-1\}$, for $N$ the size of the memory.
The agent is rewarded with $1.0$ when the sequence of addresses specifies a program that solves the task, evaluated on a set of test cases.
The reward is otherwise zero.
The observation space constitutes of an example input and corresponding desired output, the program written so far, and the state of the input and output after executing the program on the example input.
We limit the memory size to 16 and the maximum program length to $10$.
In this setup, without prior knowledge, the state space is of size $\leq 16^{10}\approx 1$ trillion unique states.
We present results on two tasks: An easier task of negating and outputting the input, for positive inputs (Negate Positives), and a harder task of implementing the Identity Function.
An implementation of Negate Positives in subleq of length 2 is known (when an address is not specified, the default value in memory is 0 0), which means finding a solution without prior knowledge requires searching on the order of $16^{2}$ unique states.
To the Identity Function, however, an implementation of length $6$ is known, which, under the same assumptions, requires searching on the order of $16^{6}\approx 16$ million unique states.
We provide a more detailed introduction to subleq, a description of the environment and the tasks in Appendix C.
The results are presented in Figure 1.

Figure: Figure 1: Sample efficiency. Left: the easy subleq Negate Positives task. Right: the harder Identity Function task. Mean of 15 seeds, two standard errors.
Refer to caption: x1.png

We compare two variations of EMCTS with AZ (E-AZ) to the AZ baseline.
Both variations of E-AZ are able to solve the harder task in a much smaller number of samples than AZ.
To estimate $\mathbb{V}\!\,[\hat{R}(\cdot,\cdot,s^{\prime})]$ one variation of E-AZ uses a hash based visit count that takes as input the complete state $s^{\prime}$, allowing the agent to, in principle, avoid searching the same state twice.
The other E-AZ variation hashes only part of the state: the example input-output, before and after execution (IO hash), directing the agent to search actions that have an effect on the output of the program.
As expected from a UCB-based method, E-AZ benefits from the more-appropriate uncertainty estimator and finds a correct program much earlier with the IO hash.

All agents use Danihelka et al. ()’s approach that combines Gumbel noise with Sequential Halving at the root and a modern variation of PUCT at non-root nodes.
For E-AZ variations, in online search (when selecting actions in the environment) the tree search policies use $q^{\beta}_{\hat{M}}$ (13) in place of $q_{\hat{M}}$ and an exploration-prior-policy $\pi_{e}$ (See Section 3.2).
To select actions during evaluation, as well as when generating targets with Reanalyze, the search uses EMCTS to propagate the uncertainty, but uses the exploitation prior policy $\pi$ and $q_{\hat{M}}$ in the search objective, in the standard manner of MCTS search.
The policy $\pi_{e}(s)$ is trained with cross entropy loss to fit the softmax across actions over $q^{\beta}(s,\cdot)$ at the root.

### 5.2 Deep Sea Experiments

The Deep Sea environment is structured as a grid, where at each time step the agent chooses between two actions, goes right or left and one row down.
There is a reward $r_{goal}=1.0$ at the bottom right corner of the grid, and the agent starts at the top left corner, and thus the probability of randomly finding the unique optimal trajectory decays exponentially with the size of the grid.
Every transition along the optimal trajectory receives a negative reward that is negligible in comparison to the goal reward, but is otherwise the only reward the agent sees, actively discouraging exploration along the optimal path.
The action mappings are randomized such that the effect of the same action is different in different states, preventing the agent from generalizing across actions.
To evaluate EMCTS in the presence of deterministic transitions and stochastic rewards that more generally align with the assumptions in Section 3, we include a custom stochastic-reward variation of Deep Sea where the goal reward $r_{goal}\sim\mathcal{N}(1,1)$, and a reward $r_{mislead}\sim\mathcal{N}(0,1)$ is given at the bottom left corner.
The agent must explore both transitions a sufficient number of times to correctly identify the larger mean reward.
Results
are presented in Figure 2.

Figure: Figure 2: Left: Scaling to growing Deep Sea sizes, 5 seeds per point. Within the training steps allotted only 2 seeds of MZ+UBE were able to solve size 40, and none size 50, both marked with an X. Middle: Stochastic-reward Deep Sea 50x50, 10 seeds. Right: The effect of the exploration parameter in Deep Sea 30x30, 3 seeds per point. Mean and standard error.
Refer to caption: x2.png

The left subplot shows sample complexity measured as number of interactions until 90% regret is reached against the size of the Deep Sea environment to demonstrate that the scaling is sub-exponential as expected from deep-exploration methods.
The middle subplot presents expected return in evaluation episodes in the stochastic-reward variation of Deep Sea.
In the right subplot, the sensitivity to the exploration parameter $\beta$ is studied and the ability of the agent to learn stably from off-policy data.
In Appendix B we include additional results, including a visualization of the uncertainty estimated by EMCTS.

To answer the research questions the following agents based in A/MZ are compared:
(i) E-A/MZ (purple, dashed for AZ and dotted for MZ), our method. The agents use EMCTS with counts/RND to estimate reward and transition uncertainty and UBE to estimate value uncertainty, search with EUCT and act with the action with the most visitations at the root.
(ii) Baseline AZ (orange) explores by sampling actions proportionally to visitations at the MCTS root.
AZ is included for reference, to demonstrate that indeed Deep Sea cannot be solved in reasonable time with random exploration, and is not able to solve any of the environments in the alloted training steps, with the exception of one seed in the smallest environment size, and that only because the initial replay buffer by chance already contained a trajectory that reached the goal.
We include the exponentially-scaling performance expected of random-exploration based methods such as AZ, for reference.
(iii) Last, A/MZ+UBE (red) acts with a similar UCB-exploration objective to that of EP/UCT, but does not search with the uncertainty: $\textstyle{a_{t}\;=\;\operatorname*{arg\,max}_{a}q_{\hat{M}}(s_{0},a_{t})+
\beta\sqrt{\mathbb{V}\!\,[\hat{R}(s_{0},a_{t})]+\gamma^{2}u(f(s_{0},a_{t}))}}$ (see Appendix D.9).
AZ-based agents are given access to the true transition model.
The value model is learned, as well as the reward model to investigate the behavior under the most general setting considered in Section 3.3.
All A/MZ agents are trained with targets generated by Reanalyze.
For reference performance on Deep Sea, we include Bootstrapped DQN (BDQN, Osband et al., ), a popular model-free, non-search based deep exploration approach that relies on an ensemble to drive deep exploration directly.
For full implementation details see Appendix D.

As expected for deep exploration methods, agents that are informed with respect to epistemic uncertainty (E-A/MZ, A/MZ+UBE and BDQN) demonstrate sample efficiency that scales sub-exponentially with environment size (Figure 2, left).
In addition, E-AZ is able to solve Deep Sea in the presence of stochastic rewards (Figure 2, middle).
This answers RQ I: EMCTS successfully demonstrates deep exploration in the presence of learned value, reward and transition models as well as both deterministic and stochastic rewards.
E-A/MZ (Figure 2, left and middle plots, purple) demonstrate significant improvement in sample efficiency over the equivalent agents that do not use search (red), in deterministic and stochastic reward environments and even with the learned transition model of MZ (dotted line).
This answers RQ II: EMCTS demonstrates benefits from search for uncertainty estimation and exploration which is retained in the presence of MZ’s learned transition dynamics’ model.
The low regret in evaluation with exponentially increasing values of $\beta$ (Figure 2, right) demonstrates that the agent can stably learn the optimal policy even in the presence of off-policy exploratory data.

## 6 Conclusions

In this work we present EMCTS, a novel, practical and theoretically motivated method to incorporate the epistemic uncertainty from learned models into MCTS, as well as harness the search for deep exploration.
AZ paired with EMCTS (E-AZ) achieves significantly higher sample efficiency in the sparse-reward, challenging task of programming in the Assembly language subleq, compared to baseline AZ.
In the popular hard-exploration benchmark Deep Sea, E-A/MZ demonstrate deep exploration by solving variations of the task which cannot be solved by baseline A/MZ at all in a reasonable amount of samples.
EMCTS’ search demonstrates significantly improved epistemic uncertainty estimation through more sample efficient exploration over an otherwise equivalent method that does not use search.
With EMCTS, A/MZ are much better equipped for sparse-reward and hard-exploration environments, which come up in realistic settings such as algorithm design, where AZ has already made significant advances.
By making A/MZ uncertainty aware, EMCTS is also promising for settings that require improved reliability in the face of the unknown such as offline RL and off-policy target generation.

## Acknowledgements

We would like to thank
Marco Loog,
Frank van der Meulen,
Itamar Sher,
Moritz Zanger,
Pascal van der Vaart, Joery de Vries
&
Jinke He
for many fruitful discussions and helpful comments.
We acknowledge the use of computational resources of the DelftBlue supercomputer, provided by Delft High Performance Computing Centre (https://www.tudelft.nl/dhpc) as well as the Delft AI Cluster (DAIC) at TU Delft (RRID: [SCR_025091](https://scicrunch.org/resolver/RRID:scr_025091)).
This work was partially supported by the EU Horizon 2020 programme under grant number 964505 (Epistemic AI),
and partially funded by the Dutch Research Council (NWO) project Reliable Out-of-Distribution Generalization in Deep Reinforcement Learning with project number OCENW.M.21.234.

## Appendix A Additional Results

### A.1 Proof of Theorem 1

By the definition of the model $\mathbb{E}[\hat{M}]=m$, where $m$ is the true model of the environment.
Therefore, $Q^{*}_{m}=Q^{*}_{\mathbb{E}[\hat{M}]}=\mathbb{E}_{\hat{M}}[Q_{\hat{M}}^{*}]$, by linearity of the value in the reward function.
By Chebyshev’s inequality,
$P\Big{(}\mathbb{E}_{\hat{M}}[Q_{\hat{M}}^{*}]\leq Q_{\hat{M}}^{*}+\frac{1}{
\sqrt{\delta}}\sqrt{\mathbb{V}\!\,[Q_{\hat{M}}^{*}]}\Big{)}\geq 1-\delta$.
The inequality
$Q_{\hat{M}}^{*}+\frac{1}{\sqrt{\delta}}\sqrt{\mathbb{V}\!\,[Q_{\hat{M}}^{*}]}
\leq\max_{\pi}Q_{\hat{M}}^{\pi}+\frac{1}{\sqrt{\delta}}\sqrt{\mathbb{V}\!\,[Q_
{\hat{M}}^{\pi}]}$ is true since the right-hand side is the maximum possible policy in a set that contains $\pi^{*}$.

Putting the entire set of inequalities together, we arrive at:

$$ $\displaystyle P\Big{(}Q^{*}=\mathbb{E}_{\hat{M}}[Q_{\hat{M}}^{*}]\leq Q_{\hat{ M}}^{*}+{\textstyle\frac{1}{\sqrt{\delta}}}\sqrt{\mathbb{V}\!\,[Q_{\hat{M}}^{* }]}\leq\max_{\pi}Q_{\hat{M}}^{\pi}+{\textstyle\frac{1}{\sqrt{\delta}}}\sqrt{ \mathbb{V}\!\,[Q_{\hat{M}}^{\pi}]}\Big{)}\geq 1-\delta,$ (16) $\displaystyle\text{and thus:}\quad P\Big{(}Q^{*}\leq\max_{\pi}Q_{\hat{M}}^{\pi }+{\textstyle\frac{1}{\sqrt{\delta}}}\sqrt{\mathbb{V}\!\,[Q_{\hat{M}}^{\pi}]} \Big{)}\geq 1-\delta,$ (17) $$

QED Theorem 1.

### A.2 Derivation of the Upper Bound on 𝕍 ⁢ [ q M ^ ⁢ ( s k , a ) ] 𝕍 delimited-[] subscript 𝑞 ^ 𝑀 subscript 𝑠 𝑘 𝑎 \mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)] blackboard_V [ italic_q start_POSTSUBSCRIPT over^ start_ARG italic_M end_ARG end_POSTSUBSCRIPT ( italic_s start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT , italic_a ) ]

By definition,

$$ $\displaystyle\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]={\textstyle\frac{1}{n^{2}}} \mathbb{V}\!\,\Big{[}{\textstyle\sum\limits_{\scriptscriptstyle i=1}^{ \scriptscriptstyle n}}\nu^{i}(s_{k},a)\Big{]}={\textstyle\frac{1}{n^{2}}}{ \textstyle\sum\limits_{\scriptscriptstyle i,j=1}^{\scriptscriptstyle n}}\text{ Cov}\big{(}\nu^{i}(s_{k},a),\nu^{j}(s_{k},a)\big{)}\,.$ (18) $$

In A/MZ it is standard to learn one set of deterministic models $\hat{V},\hat{R}$ and use them throughout planning, and therefore $\nu^{i},\nu^{j}$ cannot be assumed to be independent $\forall i\neq j$.
Had they been independent, one would compute the variance of the sum directly with the sum of the variances.
Instead, we use the inequality for covariances with known variances (Hössjer & Sjölander, ):
$\text{Cov}[X,Y]~{}\leq~{}\sqrt{\mathbb{V}\!\,[X]\mathbb{V}\!\,[Y]}$ to upper bound the variance $\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]$:

$$ $\displaystyle\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]={\textstyle\frac{1}{n^{2}}} \mathbb{V}\!\,\Big{[}{\textstyle\sum\limits_{\scriptscriptstyle i=1}^{ \scriptscriptstyle n}}\nu^{i}(s_{k},a)\Big{]}={\textstyle\frac{1}{n^{2}}}{ \textstyle\sum\limits_{\scriptscriptstyle i,j=1}^{\scriptscriptstyle n}}\text{ Cov}\big{(}\nu^{i}(s_{k},a),\nu^{j}(s_{k},a)\big{)}$ (19) $\displaystyle\leq{\textstyle\frac{1}{n^{2}}}{\textstyle\sum\limits_{ \scriptscriptstyle i,j=1}^{\scriptscriptstyle n}}\sqrt{\mathbb{V}\!\,[\nu^{i}( s_{k},a)]\,\mathbb{V}\!\,[\nu^{j}(s_{k},a)]}\;=\;\Big{(}{\textstyle\frac{1}{n} }{\textstyle\sum\limits_{\scriptscriptstyle i=1}^{\scriptscriptstyle n}}\sqrt{ \mathbb{V}\!\,[\nu^{i}(s_{k},a)]}\Big{)}^{\!2}.$ (20) $$

In other words, the standard deviation of the averaged backup $\sqrt{\mathbb{V}\!\,[q_{\hat{M}}(s_{k},a)]}$ is upper bounded with the averaged standard deviation across backups ${\textstyle\frac{1}{n}}{\textstyle\sum\limits_{\scriptscriptstyle i=1}^{
\scriptscriptstyle n}}\sqrt{\mathbb{V}\!\,[\nu^{i}(s_{k},a)]}$.
The first inequality is due to the inequality for covariances with known variances,
and the following equality is a sum of squares.

### A.3 EMCTS with learned transition dynamics

When the transition function $\hat{f}$ is learned, we must formulate it as a random variable $\hat{F}$ as well.
We note that in this case $S^{\prime}\sim\hat{F}(S,a)$ is a random variable that depends on the distribution of the previous state $S$ propagated through the uncertain transition function $\hat{F}$, and similarly the distribution of $S^{\prime}$ will propagate through the reward and value and future transition predictions $\hat{R}(S^{\prime},a),\hat{V}_{\hat{M}}^{\pi}(S^{\prime}),\hat{F}(S^{\prime},a)$.
In other words, in order to estimate $\mathbb{V}\!\,[Q^{\pi}_{\hat{M}}]$ we need to account for propagation of the variance through the Markov chain, which is an expensive and non-trivial process.
In addition, as the value is non-linear in the transition dynamics $f$, Theorem 1 does not apply and Equation 10 is not an upper bound with a specific probability.
Finally, to extend EMCTS to MZ fully, we must account for the fact that the dynamics model learned by MZ is not operating in the true state space of the environment, but in a value-equivalent abstraction.
To extend EMCTS to the learned dynamics case, we will assume that it is possible to estimate epistemic uncertainty in the abstracted state space $\hat{\mathcal{S}}$ in a meaningful way.
This problem can be circumvented by driving additional losses through the learned model that incentivize distinguishing between unique states in latent space (Henaff, ), or by learning an auxiliary dynamics model to distinguish between novel and observed starting-states-and-action-sequences, which has been used successfully by Sekar et al. ().
We will therefore use the notation of states $s$ or $S$ in the tree, whether they are in the true state space of the environment $\mathcal{S}$ or the abstracted space of MZ $\mathcal{\hat{S}}$.
In addition, since standard MZ plans with a deterministic transition model, we extend EMCTS to the setting where either the underlying transition dynamics $f$ are deterministic, or the abstracted deterministic transition function is sufficiently meaningful.
For simplicity, let us further assume that the state space $\mathcal{S}$ is continuous, and the starting state distribution $\rho$ is over a finite domain.

To circumvent the problem of the propagation of state uncertainty through the Markov chain, we propose a cheap and simple maximally-optimistic alternative upper-bound approximation for $\mathbb{V}\!\,[Q^{\pi}_{\hat{M}}]$.
We note that: (i) $\hat{F},\hat{R}$ are trained on the same data and thus when the epistemic variance over one of the two is maximal the other can be expected to be maximal as well. (ii) We have assumed a deterministic true transition function $f$ and thus the variance in $\hat{F}$ can be modelled as maximal on the unknown and zero on the known.
(iii) If the uncertainty in the prediction of any state $\mathbb{V}\!\,[\hat{F}(S_{k},a)]\geq 0$, we can expect the uncertainty in all future predictions $S_{j}$ along this trajectory $j>k$ to be associated with maximal uncertainty.
We use these observations to formulate the following simple approximation:

$$ $\displaystyle\mathbb{V}\!\,[\hat{R}(S_{k},a)]\approx\mathbb{V}\!\,[\hat{R}(s_{ k},a)],\,\,\mathbb{V}\!\,[\hat{V}_{\hat{M}}^{\pi}(S_{k})]\approx\mathbb{V}\!\, [\hat{V}_{\hat{M}}^{\pi}(s_{k})],\,\,\forall(s,a)\in\mathcal{D}$ (21) $\displaystyle\mathbb{V}\!\,[\hat{R}(S_{j},a)]\approx r^{2}_{max},\,\,\mathbb{V }\!\,[\hat{V}_{\hat{M}}^{\pi}(S_{j})]\approx({\textstyle\frac{1}{1-\gamma}}r_{ max})^{2},\,\,\text{where}\,\,j\geq k\,\,\text{and}\,\,\forall(s,a)\notin \mathcal{D}$ (22) $$

That is, we propose to ignore any uncertainty in transition along zero-uncertainty trajectories, and with the first uncertain transition, to assume all uncertainties are maximal for all predictions in the rest of the trajectory.
To identify whether $(s,a)\notin\mathcal{D}$ we can directly use $\mathbb{V}\!\,[\hat{R}(s,a)]$, where $\mathbb{V}\!\,[\hat{R}(s,a)]\approx r^{2}_{max}$ indicates $(s,a)\notin\mathcal{D}$.
This will guarantee that (given that the mechanism to identify unobserved transitions is sufficiently reliable) the agent remains sufficiently optimistic and the planning objective of EP/UCT remains an upper bound on $Q^{*}$.
We note that in practice, while this approach is theoretically sound, the results of the experiments in Section 5.2, in the presence of reliable transition-uncertainty, motivate that it is not necessary and the regular setup of EMCTS is sufficient for deep exploration even in the presence of a learned, value-equivalent-abstraction-based transition model, given that $\hat{R}(s,a)$ is a sufficiently reliable estimator of novelty in the environment.

## Appendix B Additional Experiments

We include an example comparison of the uncertainty estimated by EMCTS to that of the uncertainty estimator used by EMCTS, the UBE network head, in Figure 3.
The E-AZ agent used in this experiment uses the true transition model but a learned reward as well as value functions, matching the setup of Section 3.1.
The Deep Sea environment is presented as a grid, where states are the cells including and below the diagonal.
Bold blue in the grids in the bottom row (inverse counts) indicates unvisited states.
By averaging across multiple predictions in search, the uncertainty estimated by EMCTS (top row) is much more varied, more accurately associating less or more uncertainty with states that lead into observed / unobserved trajectories respectively, compared to the UBE predictions (middle row).
Most importantly from the perspective of exploration, EMCTS associates larger uncertainty with states that lead into unobserved directions much more consistently than the single predictions of UBE for each state (for an easily visible example, t$=2000$, top of the diagonal).

Figure: Figure 3: Heat maps over states in DeepSea 40 by 40 at different times (columns) during an example training run of EMCTS with an AZ transition model. Upper row: value uncertainty at the EMCTS root node. Middle row: single prediction of UBE at each state. Lower row: inverse visitation counts as reliable local uncertainty, where score of 2.0 represents unvisited.
Refer to caption: x3.png

We include in addition detailed evaluation curves for E-A/MZ, A/MZ and A/MZ+UBE in Figure 4, as well as rate-of-exploration (number of unique states encountered per interaction with the environment).
An additional variation of MZ is included, where instead of value-equivalent abstraction the model is trained with a reconstruction loss.
For implementation details see Appendix D.3.
In all figures, E-A/MZ outperforms the other agents both in rate of exploration as well as in their ability to learn to reach the goal.
The most interesting behavior is perhaps that of the reconstruction based model, that does not search the environment significantly faster, and yet learns much faster to reach the goal (bottom row in Figure 4).
We hypothesize that due to the search, the uncertainty estimated by the agent is more reliable, resulting in identifying the correct action that leads into novel states, more times in a row.
Just barely visible on the right-hand plot, one can see that indeed the purple curve remains the highest for quite a while, before all others curves match it, despite all curves being very close throughout most of the training.

Figure: Figure 4: Deep Sea 40x40, mean and standard error for 20 seeds. Rows: Different transition models. Left: episodic return in evaluation vs. environment steps. Right: exploration rate (number of discovered states vs. environment steps).
Refer to caption: x4.png

We include a table evaluating interactions-to-goal on Deep Sea 40x40 for all agents.
The results demonstrate that even when the learned dynamics model is not designed for planning (anchored model, third block, Table 1), EMCTS is able to find the goal much faster.

**Table 1: Number of environment steps until the first visitation to the goal transition.**
| Novelty Estimator | Exploration | Average steps to goal transition for<br>seeds that discovered goal $\pm$ STD | % seeds that<br>discovered goal |
| --- | --- | --- | --- |
| RND | E-AZ | $\bm{10539}\pm 9006$ | 94% of 35 seeds |
| AZ+UBE | $22801\pm 7514$ | 91% of 35 seeds |  |
| AZ | - | 0% of 20 seeds |  |
| Counts | E-MZ | $\bm{14339}\pm 6845$ | 100% of 23 seeds |
| MZ+UBE | $29945\pm 8113$ | 57% of 21 seeds |  |
| MZ | - | 0% of 20 seeds |  |
| Reconstruction<br>Model (RND) | E-MZ | $\bm{15241}\pm 3236$ | 95% of 20 seeds |
| MZ+UBE | $22497\pm 6645$ | 85% of 20 seeds |  |
| MZ | - | 0% of 20 seeds |  |

## Appendix C Subleq

What follows is a formal characterization of subleq as used in our experiments, and example programs for the two tasks we presented.

For subleq-$N$, we have a *memory* made up of $N$ *words* $(w_{0},w_{1},...,w_{N-1})$ where each word $w_{i}$ is an integer from 0 0 to $N-1$ inclusive (we write $w_{i}\in[0,N-1]$). We also have an input sequence $(v_{0},v_{1},...)$, $v_{i}\in[0,N-1]$, of variable finite length depending on the particular task to be solved. Tasks correspond to different algorithms we want the agent to implement. Lastly, there is an output buffer $(u_{0},u_{1},...)$, $u_{i}\in[0,N-1]$ which begins empty, but can be extended with output values during execution of a subleq program.

We refer to $w_{i}$ as the word at *address* $i$. Additionally, we give names to some specific addresses: $\texttt{@IN}=N-3$, $\texttt{@OUT}=N-2$, $\texttt{@HALT}=N-1$. The program is stored contiguously in memory starting at address 0 0. Execution begins from address 0 0 as well, meaning the first instruction to execute is $(w_{0},w_{1},w_{2})$.

When executing some instruction at address $i$, we look at the three words $(w_{i},w_{i+1},w_{i+2})$, and modify the memory by subtracting the value at address $w_{i+1}$ from the word at address $w_{i}$, i.e. $w_{w_{i}}\leftarrow w_{w_{i}}-w_{w_{i+1}}$. We say that this instruction *reads* from $w_{i}$ and $w_{i+1}$, and *writes* to $w_{i}$. We also say that reading from $w_{i}$ *returns* $w_{w_{i}}$. All operations are done modulo $N$ so that the words remain in the range $[0,N-1]$. If the result before modulo was less or equal to zero, execution continues from address $w_{i+2}$ (we say we *jump* to $w_{i+2}$), meaning the next instruction would be $(w_{w_{i+2}},w_{w_{i+2}+1},w_{w_{i+2}+2})$. If the result before modulo was strictly positive, then the next instruction is the next three words, i.e. $(w_{i+3},w_{i+4},w_{i+5})$.

Interacting with the addresses @IN, @OUT, and @HALT has different behaviour than normal execution. Reading from the address @IN will instead read the next number from the input sequence. The first read from @IN will return $v_{0}$, the next read returns $v_{1}$, etc. Reading from the input sequence when there is no next value will return 0 0 instead. Writing to @IN is ignored.

Reading from @OUT always returns a 0 0, but writing to @OUT will write to the output buffer. We say that we *output* the value. Each value that we output is added to the output buffer, so after the first output $u_{0}$, the buffer looks like $(u_{0})$, after the second output $u_{1}$, it is $(u_{0},u_{1})$, etc.

The output buffer is compared against the desired output for the task to determine whether the task was solved correctly. The first incorrect output terminates the program and results in a failure of the task. If the output buffer ever equals the desired task output, the program terminates and the task is solved successfully. Finally, if the program tries jumping to an address where the instruction would overlap @HALT, such as $(w_{\texttt{@IN}},w_{\texttt{@OUT}},w_{\texttt{@HALT}})$, the program terminates.

When using SUBLEQ-$N$ as a reinforcement learning environment, the agent writes a program one word at a time. It gets to observe the current state of the memory before execution (it is just the program it has written up to that point followed by zeroes), an example input and desired output pair for the given task, and the resulting state of the input sequence and the output buffer after executing the currently written program.

At each state, the agent has $N$ actions $=\{0,1,...,N-1\}$, each corresponding to writing that number into the next location in memory. The memory is initially filled with zeroes. After an action is chosen, the number is added to memory, and the currently written program is executed to determine the states of the input and output on the example test as well as other (hidden) test cases. If all test cases succeed, the episode terminates after one more (irrelevant) action. On the other hand, if the agent reaches the end of writable memory ($N-3$ actions), the episode terminates unsuccessfully.

The Negate Positives task expects a program which receives positive integers and outputs that input, expect negated. This task are very simple, since it can be done with a single instruction: Reading from @IN will subtract the next input from whichever address to which we write. If we write to @OUT, which reads as 0 0, we are effectively writing $0-v_{i}=-v_{i}$. For Negate Positives, this is exactly what is required, so the only remaining challenge is how to loop back to the start. Luckily, since we know the input is always positive, the output will be negative (or zero), so we always jump. Thus, if we choose the jump address to be 0 0, we can also loop to the start in the same instruction.

$(\texttt{@OUT},\texttt{@IN},0,...)$ in memory forms a solution to the task, and since the memory is initialized with zeroes, it only requires the agent to write $2$ words, meaning only $2$ actions.

In code, a solution to Negate Positives looks like this:

Which the agent would write as:

In the Identity task, we require a program which outputs the input unchanged. This is slightly more difficult that Negate Positives, because an instruction in subleq always subtracts, so we need at least two instruction to undo the negation. This can be done by storing the return of @IN in some auxiliary word, and then writing that value to output. Let us denote the address at which we store the value temporarily as @x. The program that solves Identity could begin $(\texttt{@x},\texttt{@IN},3,\texttt{@OUT},\texttt{@x},...)$. Note the $3$ at address $2$ is needed to progress to the second instruction regardless of whether the input is positive or negative. This program would work for the first input, but we must produce a program which loops. Thus, we still need to erase the value stored at address @x (to prepare it for the next write), and we must loop to the start. These two things can be done in the same instruction. We arrive at a program like $(\texttt{@x},\texttt{@IN},3,\texttt{@OUT},\texttt{@x},6,\texttt{@x},\texttt{@x
},0,...)$. The last instruction $(\texttt{@x},\texttt{@x},0)$ clears @x and jumps to address zero (since $w_{\texttt{@x}}\leftarrow w_{\texttt{@x}}-w_{\texttt{@x}}=0$ which $\leq 0$). Again, we required the constant $6$ at address $5$ to unconditionally continue to the next address. In general, an agent would need $8$ actions to solve Identity, where it needs to pick @IN, @OUT, $3$, and $6$ specifically, and it needs to make sure that all of @x are the same. We say in general, because there is actually a clever solution when we choose $@x=0$. In that case, the program becomes $(0,\texttt{@IN},3,\texttt{@OUT},0,6,0,0,0,...)$ which only requires the agent to write $6$ words, since the memory is initially filled with zeroes.

In code, a general solution to Identity looks like this (left in human syntax, right as the agent writes, the 9s can be any constant $\geq 9$, $<N$):

While the shortest (known) solution (in terms of non-zero bytes written) looks like this:

## Appendix D Implementation Details for Deep Sea

Our implementation for the Deep Sea agents is based in the framework of Ye et al. ().
Below, we detail the implementation details unique to E-A/MZ and A/MZ+UBE in Deep Sea.

### D.1 Targets

Value, policy and reward targets were all computed as in MZ Schrittwieser et al. ().
UBE targets were computed in an n-step manner:

$$ $\displaystyle u_{target}(s_{t})=\sum_{i=0}^{n-1}\gamma^{2i}\eta(s_{t+i},a_{t+i })+\gamma^{2n}\mathbb{V}\!\,[\hat{V}_{\hat{M}}^{\pi}(s_{t+n})]$ (23) $$

Where $\eta$ is the RND / exact or hash count-based novelty estimators.
To guarantee that the UBE estimates remains sufficiently optimistic, the value-uncertainty bootstrap $\mathbb{V}\!\,[\hat{V}_{\hat{M}}^{\pi}(s_{t+n})]$ was computed in one of two ways:

- 1.
Root targets:
The uncertainty at the root of the EMCTS tree at state $s_{t+n}$. Specifically, this tree planned with the uncertainty maximizing objective: $a_{\text{max-uncertainty}}\;=\;\operatorname*{arg\,max}_{a\in A}\sigma_{q_{
\hat{M}}}(s_{k},a)+\pi(a|s_{k})\,C_{\text{PUCT}}{\textstyle\frac{\sqrt{\sum_{a
^{\prime}}\!N(s_{k},a^{\prime})}}{1+N(s_{k},a)}},$.
- 2.
Non-root targets:
$\displaystyle\mathbb{V}\!\,[\hat{V}_{\hat{M}}^{\pi}(s_{t+n})]=\max_{a}\eta(s_{
t+n},a_{t+n})+\gamma^{2}u(s_{t+n+1}),$
(24)
and $u(s_{t+n+1})=\max\big{(}u(s_{t+n+1}),\frac{1}{1-\gamma^{2}}\sigma_{\hat{R}}^{2
}(s_{t+n},a_{t+n})\big{)}$

### D.2 Losses

The original MZ algorithm uses three loss functions $L_{r},L_{v},L_{\pi}$ for the reward, value and policy, respectively.
The gradients from the losses $\mathcal{L}_{r},\mathcal{L}_{v},\mathcal{L}_{\pi}$ propagate through the learned transition model $f$ and are the only learning signal that is used to train the model.

For the anchored model (see Section 5) we use an additional reconstruction loss:

$$ $\displaystyle\mathcal{L}_{re}:={\textstyle\frac{1}{|\mathcal{B}|}}{\textstyle \sum\limits_{\scriptscriptstyle t\in\mathcal{B}}}{\textstyle\sum\limits_{ \scriptscriptstyle k=0}^{\scriptscriptstyle l-1}}||\hat{s}_{t}^{k}-s_{t+k}||^{2}$ $$

where $\mathcal{B}\equiv\{s_{t},a_{t},r_{t},s_{t+1},a_{t+1},\ldots,s_{t+l}\}_{t\in
\mathcal{B}}$ is a training batch containing $b$ trajectories of length $l$ sampled from different episodes.
To estimate value-uncertainty at the leaves, we train a UBE function $u$ with a UBE loss $\mathcal{L}_{u}$:

$$ $\displaystyle\mathcal{L}_{u}:={\textstyle\frac{1}{|\mathcal{B}|}}{\textstyle \sum\limits_{\scriptscriptstyle t\in\mathcal{B}}}{\textstyle\sum\limits_{ \scriptscriptstyle k=0}^{\scriptscriptstyle l-1}}\phi(u^{\text{target}}_{t+k}) ^{T}\log\hat{u}_{t}^{k}$ $$

The final loss is computed as:

$$ $\displaystyle\mathcal{L}:=\lambda_{r}\mathcal{L}_{r}+\lambda_{v}\mathcal{L}_{v }+\lambda_{\pi}\mathcal{L}_{\pi}+\lambda_{u}\mathcal{L}_{u}$ $$

Where the coefficients $\lambda_{r},\lambda_{v},\lambda_{\pi},\lambda_{u}$ are used to weigh the relative effects the individual components of the loss have on the learned transition model $f$.
When $\mathcal{L}_{re}$ was used (the anchored model in Section 5), the model parameters of $f$ were affected only by $L_{re}$, through a second backwards pass.

### D.3 Different Dynamics Models

We describe the three transition models used in 5 in more detail.
The AlphaZero dynamics model is a true model of the dynamics of the environment, in the true state space of the environment.
When planning with this model local uncertainty is estimated with RND and value-uncertainty is estimated with UBE.
The MuZero model is a value-equivalent model in latent space. $g,f$ are learned by the agent during training from the value, policy, reward and UBE losses.
When planning with this model local uncertainty is estimated with state-visitation-counts (see D.5 and value-uncertainty is estimated with UBE.
The anchored-MuZero transition model trained only to predict the true transition dynamics of the environment through a reconstruction loss $L_{re}^{k}$ (see Appendix D.2).
When planning with this model local uncertainty is estimated with RND and value-uncertainty is estimated with UBE.

### D.4 Planning with Random Network Distillation Based Epistemic Uncertainty

Many popular novelty-estimators in deep RL
(such as RND, or even deep-ensembles Osband et al., )
do not directly provide a reliable variance estimate.
Many deep ensemble methods for example rely on ensemble disagreement (Sekar et al., ), but do not assume that the variance in the ensemble approximates the variance of a distribution, but rather should be high on unknown and low on known inputs.
To circumvent this limitation of standard and popular uncertainty estimation methods we estimate the uncertainty in latent state $\bm{\Sigma}_{k}$ and the uncertainty in the predictions based on latent state $\sigma_{R}$ together into one score.
More specifically, we use $\bm{\Sigma}_{k}=0$ and $\sigma_{V}(\hat{s}_{k})=\max\Big{(}L_{rnd}(\hat{s}_{k-1},a_{k-1}),u\big{(}f(
\hat{s}_{k-1},a_{k-1})\big{)}\Big{)}$.
This choice is sufficient for EMCTS to significantly improve over a comparable non-planning deep exploration baselines, see Section 5.

When the planning is done with a true model, the agent has access to the true states $s_{t+k}$ and can use RND to evaluate transition uncertainty over the state action pair $(s_{t+k},a_{t+k})$ directly.
When the planning is done with the anchored model, the latent states outputted by the transition model $\hat{s}^{k}_{t}$ approximate the true states $s_{t+k}$ which allows us to use RND over $(\hat{s}^{k}_{t},a_{t+k})$.
In both cases, RND is trained only over the observed transitions $(s_{t+k},a_{t+k})$, not latent state representations $(\hat{s}^{k}_{t},a_{t+k})$, to achieve the objective of yielding large RND prediction errors the further the latent state prediction $\hat{s}^{k}_{t}$ is from observed state $s_{t+k}$.

### D.5 Planning with Visitation-Counts Based Epistemic Uncertainty

When planning with the abstracted model, we provide the agent with access to two additional mechanisms that are used only for local uncertainty estimation: the true model $F(s_{t},a_{t})$ of the environment and a state-action visitation counter $C(s_{t},a_{t})$.
During planning, the true transition model follows the planning decisions $a_{t:t+k}$ and keeps track of the true state $s_{t+k}$.
When the agent evaluates the local uncertainty with transition $(\hat{s}^{k}_{t},a_{t+k})$ the true model provides the matching true state $s_{t+k}$ to the visitation counter, which produces the local uncertainty based on the following formula:

$$ $\displaystyle\eta(s_{t+k},a_{t+k})=\frac{1}{C(s_{t+k},a_{t+k})+\epsilon}$ $$

Where $0<\epsilon\leq 1$ is a constant and $C(s_{t+k},a_{t+k})$ counts the number of times the state action pair $(s_{t+k},a_{t+k})$ has been observed in the environment.
This allows us to evaluate the abstracted-model agent in the presence of a reliable source of local uncertainty.
The leaf-value uncertainty $u(\hat{s}^{k}_{t})$ (which is the dominating factor in visited areas of the state space, as $\sigma^{2}(s_{t+k},a_{t+k})\to 0$ quickly in observed transitions) relies entirely on the learned UBE function $u$ which operates directly on latent states $\hat{s}^{k}_{t}$.

### D.6 Using UBE to Estimate Value-Uncertainty at the Leaves

It is essential for exploration that the epistemic uncertainty prediction is reliably high in unobserved areas of the state action space.
For this reason, a learned function $\hat{u}\approx u$ may not be sufficient to detect that a state $s_{t}$ has not been previously observed.
Instead, we use the following: $\textstyle{\max\big{(}u(s_{t}),\frac{1}{1-\gamma^{2}}\eta(s_{t},\pi(s_{t}))
\big{)}}$.
If the uncertainty $\eta(s_{t},\pi(s_{t}))$ for the transition $(s_{t},\pi(s_{t}))$ is high, the uncertainty will be estimated as high regardless of the UBE prediction $u(s_{t})$, and otherwise, either $u(s_{t})$ is high or both are negligible.

### D.7 Separating Exploration from Exploitation

Acting in the environment with a dedicated exploration policy can be expected to result in samples that are very off-exploitation-policy.
Learning from very off-policy data is known to cause instability in training even in off-policy agents.
To mitigate that, the EMCTS and only-UBE agents (see section 5) alternate between two types of training episodes: exploratory episodes that follow an exploration policy throughout the episode (such as a policy generated by EMCTS with an exploratory planning objective), and exploitatory episodes that follow the standard MuZero exploitation policy throughout the episode.
This enables us to provide the agent with quality exploitation targets to evaluate and train the value and policy functions reliably, while also providing a large amount of exploratory samples that explore the environment much more effectively and are more likely to efficiently search for high-reward interactions.

In practice, rather than alternate between exploration and exploitation episodes we run a certain number of episodes in parallel, a certain portion of which are exploitatory and the rest are exploratory.
In our experiments the ratio was $50/50$.
To avoid learning a a separate prior-policy that may not be necessary in environments with small actions space, we set the policy prediction $\pi(s_{k})$ (see Equation 5) to uniform over all actions, for all $s_{k}$ during exploration episodes.
Dirichlet noise was not used to drive exploration in MCTS with the UBE and EMCTS agents, as any little amount of stochasticity in the policy can prevent the agent from reliably completing the one optimal trajectory.

### D.8 Environment Adaptation

To maintain the exploration difficulty of Deep Sea while reducing numerical challenges, we amplify the goal reward from $1$ to $10$ when deterministic rewards are used.
To further simplify model learning with the anchored model, the representation function $g$ that was used for the anchored model transforms the observations from 2 dimensional $(N,N)$ one-hot representations to 1 dimensional $(2N)$ representations where the first $N$ entries are a 1-hot vector representing the row and following $N$ entries are a 1-hot vector representing the column.
From this perspective, we can view the $\mathcal{L}_{re}$ loss that was used to train the anchored model as a consistency loss between the representation and the state prediction rather than a reconstruction loss.
The loss itself is the same loss specified in Appendix D.2.

### D.9 The A/MZ+UBE Agent

The A/MZ+UBE ablation agent uses MCTS to evaluate the $q$ value of actions in the same manner as A/MZ, and explores by taking the action $a_{t}$ that maximizes the combination of the Q-values approximated by MCTS, local uncertainty $\mathbb{V}\!\,[\hat{R}]$ and UBE head $u$:

$$ $\displaystyle a_{t}\;=\;\operatorname*{arg\,max}_{a}q_{\hat{M}}(s_{0},a_{t})+ \beta\sqrt{\mathbb{V}\!\,[\hat{R}(s_{0},a_{t})]+\gamma^{2}u(f(s_{0},a_{t}))}.$ (25) $$

$q_{\hat{M}}(s_{0},a_{t})$ are the values at the root of the regular MCTS tree after search.
The main difference between A/MZ+UBE and E-A/MZ is that in E-A/MZ the uncertainty $\sqrt{\mathbb{V}\!\,[\hat{R}(s_{0},a_{t})]+\gamma^{2}u(f(s_{0},a_{t}))}$ estimated takes into account estimates from multiple different future trajectories, in the manner MCTS estimates the values $q_{\hat{M}}$.

## Appendix E Network Architecture & Hyperparameters

### E.1 Hyperparameter Search

Due to the large number of hyperparameters in the MuZero framework, our optimization process consisted of manual modifications to the hyperparameters used by Ye et al. () for Deep Sea and Koyamada et al. () for subleq with the objective of achieving learning stability on the target environment with the simplest network architectures possible.
Two exceptions to this statement are the RND network architecture and scale, and the exploration parameter $\beta$.

The RND architecture was designed with the objective of reliably achieving small RND predictions over observed state-action pairs and large predictions over unobserved state-action pairs.
The RND scale was tuned with the objective of achieving local uncertainty measures for unobserved state-action pairs that are significantly larger than the minimum reward of Deep Sea.

The $\beta$ parameter in Deep Sea was tuned with the objective that the EMCTS and only-UBE agents will prioritize exploration of the environment over exploitation until the entire environment has been searched, and was tuned separately for every model.

For subleq we chose $\beta=1$.
We did not experiment with additional values of $\beta$. However, first, the values, UBE prediction, and state-novelty predictions are all bounded $\leq 1$, such that $\beta$ need not account in this case to the possibly arbitrary scales of UBE / the novelty estimator.
Second, to make the most out of Jax’s naturally parallelized setup, each parallel episode explores with a different $\beta_{i}\leq\beta$, evenly spaced from 0 0 to $\beta$.

### E.2 Network Architecture

The functions $f,g,r,v,u,\pi,\psi,\psi^{\prime}$ used fully connected DNNs of varying sizes.
The sizes of the hidden layers and output layers are specified in Table 2 for Deep Sea.
For subleq the FC network architecture constituted value, UBE, exploration policy $\pi_{e}$ and exploitation policy $\pi$ networks.
All networks used two hidden layers of size 256 with ReLu activations between hidden layers.
The value head used tanh activation on the last layer, the UBE head used a $0.5(tanh(x)+1)$ activation to bound the ube prediction between 0 and 1. The policy heads did not use any activation layers.

**Table 2: Network architecture hyperparameters, Deep Sea**
|  | True Model |  |
| --- | --- | --- |
| Function | Hidden Layers Sizes | Output Layer Size |
| f | - | - |
| g | - | - |
| r | [256, 256] | 21 |
| v | [256, 256] | 21 |
| u | [256, 256] | 21 |
| $\pi$ | [256, 256] | 2 |
|  | Anchored Model |  |
| Function | Hidden Layers Sizes | Output Layer Size |
| f | [1024, 1024, 1024] | 80 |
| g | - | - |
| r | [256, 256] | 21 |
| v | [256, 256] | 21 |
| u | [256, 256] | 21 |
| $\pi$ | [256, 256] | 2 |
|  | Abstracted Model |  |
| Function | Hidden Layers Sizes | Output Layer Size |
| f | [1024, 1024, 1024] | 100 |
| g | [512, 512] | 100 |
| r | [128, 128] | 21 |
| v | [128, 128] | 21 |
| u | [128, 128, 128] | 21 |
| $\pi$ | [128, 128] | 2 |
|  | RND network architecture |  |
| Function | Hidden Layers Sizes | Output Layer Size |
| $\psi$ | [1024, 1024] | 512 |
| $\psi^{\prime}$ | [512] | 512 |

### E.3 Deep Sea Hyperparameter Configuration

We detail the full set of hyperparameters in Tables 3 and 4 for Deep Sea.
For the BDQN baseline, we used the default implementation in [https://github.com/deepmind/bsuite](https://github.com/deepmind/bsuite), with ensemble size of 10 and matching batch size to EMCTS: number of unroll steps times batch size $5\cdot 256=1230$.
Otherwise, the default hyper parameters were used.

**Table 3: Hyperparameters used in the Deep Sea experiments**
| Parameter | Setting | Comment |
| --- | --- | --- |
| Stacked Observations | 1 |  |
| $\gamma$ | 0.995 |  |
| Number of simulations in MCTS | 50 |  |
| Dirichlet noise ratio ($\xi$) | 0.3 |  |
| Root exploration fraction | 0 |  |
| Batch size | 256 |  |
| Learning rate | 0.0005 |  |
| Optimizer | Adam (Kingma & Ba, ) |  |
| Unroll steps $l$ | 5 |  |
| Value target TD steps ($n_{v}$) | 5 |  |
| UBE target TD steps ($n_{u}$) | 1 |  |
| value support size | 21 |  |
| UBE support size | 21 |  |
| Reward support size | 21 |  |
| Reanalyzed policy ratio | 0.99 | See (Ye et al., ) |
| Prioritized sampling from the replay | True | See (Schrittwieser et al., )<br>Appendix G |
| Priority exponent ($\alpha$) | 0.6 | See (Schrittwieser et al., )<br>Appendix G |
| Priority correction ($\beta_{p}$) | $0.4\to 1$ | See (Schrittwieser et al., )<br>Appendix G |
| Evaluation episodes | 8 |  |
| Min replay size for sampling | 300 |  |
| Self-play network updating inerval | 5 |  |
| Target network updating interval | 10 |  |

**Table 4: Specific for models and agents**
| Parameter | Setting |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| True Model | Abstracted Model | Anchored Model |  |  |  |  |  |  |  |
| EMCTS | UBE | Uninf. | EMCTS | UBE | Uninf. | EMCTS | UBE | Uninf. |  |
| Training steps /<br>environment<br>interactions | 45K | 45K | 45K | 35K | 35K | 35K | 45K | 45K | 45K |
| Reward loss<br>weight $\lambda_{r}$ | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 |
| Value-loss<br>weight $\lambda_{v}$ | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 |
| Policy-loss<br>weight $\lambda_{\pi}$ | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 |
| UBE-loss<br>weight $\lambda_{u}$ | 0.125 | 0.125 | - | 0.25 | 0.25 | - | 0.125 | 0.125 | - |
| RND scale | 1.0 | 1.0 | - | - | - | - | 0.001 | 0.001 | - |
| Root based<br>targets | False | False | False | True | True | True | False | False | False |
| Disabled<br>policy<br>in exploration | True | True | False | True | True | False | True | True | False |
| Number of<br>parallel<br>episodes | 2 | 2 | 2 | 2 | 2 | 2 | 2 | 2 | 2 |
| Out of are<br>exploration<br>episodes | 1 | 1 | - | 1 | 1 | - | 1 | 1 | - |
| Exploration<br>coefficient $\beta$ | 10 | 10 | - | 1 | 1 | - | 10 | 10 | - |
| Dirichlet noise<br>magnitude $\rho$ | 0 | 0 | 0.25 | 0 | 0 | 0.25 | 0 | 0 | 0.25 |

### E.4 SUBLEQ Hyperparameter Configuration

The subleq E-/AZ agents are implemented in Jax, based in the implementation of Koyamada et al. ().
MCTS parameters used the defaults provided by DeepMind et al. ().
Detailed hyperparameter configuration below:

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: N/A
soundness: 3
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.
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.
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?

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: /
soundness: 3
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.
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.
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?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
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?
soundness: 3
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.
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.
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.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: See weaknesses for the questions.
soundness: 3
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.
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.
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?

Meta-review: 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.

Decision: Accept (Poster)