Bayes Adaptive Monte Carlo Tree Search for Offline Model-based Reinforcement Learning

## Abstract

Abstract Offline reinforcement learning (RL) is a powerful approach for data-driven decision-making and control. Compared to model-free methods, offline model-based reinforcement learning (MBRL) explicitly learns world models from a static dataset and uses them as surrogate simulators, improving the data efficiency and enabling the learned policy to potentially generalize beyond the dataset support. However, there could be various MDPs that behave identically on the offline dataset and so dealing with the uncertainty about the true MDP can be challenging. In this paper, we propose modeling offline MBRL as a Bayes Adaptive Markov Decision Process (BAMDP), which is a principled framework for addressing model uncertainty. We further introduce a novel Bayes Adaptive Monte-Carlo planning algorithm capable of solving BAMDPs in continuous state and action spaces with stochastic transitions. This planning process is based on Monte Carlo Tree Search and can be integrated into offline MBRL as a policy improvement operator in policy iteration. Our “RL + Search" framework follows in the footsteps of superhuman AIs like AlphaZero, improving on current offline MBRL methods by incorporating more computation input. The proposed algorithm significantly outperforms state-of-the-art model-based and model-free offline RL methods on twelve D4RL MuJoCo benchmark tasks and three target tracking tasks in a challenging, stochastic tokamak control simulator.

## 1 Introduction

The success of reinforcement learning (RL) typically relies on large amounts of interactions with the environment. However, in real-world scenarios, such interactions can be unsafe or costly. As an alternative, offline RL (Levine et al. ()) leverages offline datasets of transitions, collected by a behavior policy, to train a policy that can transfer to an online task. To avoid overestimation of the value function for some (out-of-sample) states in the environment, which can mislead policy learning, model-free offline RL methods (Kumar et al. (); Wu et al. ()) often constrain the learned policy to remain close to the behavior policy or within the support of the offline dataset. However, collecting transitions that comprehensively cover possible task scenarios, or acquiring a large volume of demonstrations from a high-quality behavior policy, can be expensive. This challenge has led to the development of offline model-based reinforcement learning (MBRL) approaches, such as (Lu et al. (); Guo et al. ()). These methods train dynamics models from offline data and optimize policies using imaginary rollouts generated by the models. Notably, the dynamics modeling is independent of the behavior policy, making it possible to achieve high returns even with data collected from a random policy. Furthermore, with careful dynamics modeling and thorough simulation, the learned policy can more effectively handle the environmental stochasticity and generalize to states beyond the support of the offline dataset.

Given a dataset, there may be various potential MDPs that behave identically on the limited set of states and actions, but their dynamics and reward functions could differ, especially on out-of-sample states and actions. This implies that we are dealing with a distribution of possible world models underlying the dataset. A common strategy in offline MBRL is to learn an ensemble of world models and treat them equally. For instance, when determining the next state, a world model is uniformly sampled from the ensemble and generate its prediction. However, different ensemble members may perform better in different regions of the state-action space, making it necessary to adapt the belief over each ensemble member based on the experience accumulated since the start of the episode. The Bayes Adaptive Markov Decision Process (BAMDP, Duff ()) provides a principled framework for modelling such an adaptive process. We show in Section 4.1 that, despite the need for Bayesian posterior updates, BAMDPs can still be efficiently simulated using deep ensembles. BAMCP (Guez et al. ()) is an efficient online planning method for solving BAMDPs. However, BAMCP has several limitations: (1) it relies on a ground-truth world model for planning; (2) it is restricted to discrete state and action spaces; and (3) its outcome is an action choice at a particular state, rather than a policy function. To address these challenges: (1) we apply a reward penalty (defined with the adapting belief) to construct a pessimistic BAMDP, preventing overexploitation of inaccurate world models (learned from the offline dataset); (2) we propose a novel planning algorithm to solve BAMDPs in continuous state and action spaces by extending BAMCP with double progressive widening (Auger et al. ()); and (3) we integrate the planning component as a policy improvement operator within policy iteration RL methods (Sutton & Barto ()), enabling the derivation of a policy suitable for real-time execution from the planning results. Specifically, the planning process is carried out through Monte Carlo Tree Search on a BAMDP. Integrating search with RL allows for the use of significantly more computation input, thereby improving policy learning performance. Grounded in the “scaling law", this paradigm has seen tremendous success in sophisticated policy learning, as demonstrated in (Silver et al. (); Schrittwieser et al. (); AlphaProof & AlphaGeometry ()). Its application to offline MBRL, particularly in continuous control tasks, is a promising area for exploration.

To summarize, the main contributions of this work include:
(1) Introducing BAMDPs to handle model uncertainties in offline MBRL;
(2) Proposing an efficient Bayes Adaptive Monte Carlo Tree Search method for planning in continuous, stochastic BAMDPs;
(3) Developing the first algorithm to successfully integrate Bayesian RL, offline MBRL, and deep search for sophisticated policy learning in continuous control under highly stochastic environments;
(4) Demonstrating the improvements brought by Bayesian RL and deep search across twelve D4RL MuJoCo tasks and three target tracking tasks in a stochastic tokamak control scenario (for nuclear fusion), highlighting the potential of our algorithm to tackle challenging, real-world problems.

## 2 Background

A Markov Decision Process (MDP, Puterman ()) is described as a tuple $\mathcal{M}=\langle\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\gamma\rangle$. $\mathcal{S}$ and $\mathcal{A}$ are the state space and action space, respectively. $\mathcal{P}:\mathcal{S}\times\mathcal{A}\rightarrow\Delta_{\mathcal{S}}$ is the dynamics function and $\mathcal{R}:\mathcal{S}\times\mathcal{A}\rightarrow\Delta_{[0,1]}$ is the reward function, where $\Delta_{\mathcal{X}}$ denotes the set of possible probability distributions on $\mathcal{X}$. $\gamma\in[0,1)$ is a discount factor. The goal of solving an MDP is to find a policy $\pi:\mathcal{S}\rightarrow\Delta_{\mathcal{A}}$ that maximizes the expected return, defined as the cumulative discounted reward over time.

A Bayes Adaptive MDP (BAMDP, Duff ()) can model scenarios where the precise MDP $\mathcal{M}_{\theta}=\langle\mathcal{S},\mathcal{A},\mathcal{P}_{\theta},
\mathcal{R}_{\theta},\gamma\rangle$ is uncertain but is known to follow a prior distribution $b_{0}(\theta)$. During planning, a Bayes-optimal agent would update its belief over the MDP based on experience. Formally, a BAMDP can be described as a tuple $\mathcal{M}^{+}=\langle\mathcal{S}^{+},\mathcal{A},\mathcal{P}^{+},\mathcal{R}
^{+},\gamma\rangle$. $\mathcal{S}^{+}$ denotes the space of information states $(s,b)$, which is a composition of the physical state and the current belief over the MDP. After each transition $(s,a,r,s^{\prime})$, the belief is updated to the corresponding Bayesian posterior: $b^{\prime}(\theta)\propto b(\theta)P((s,a,r,s^{\prime})|\theta)=b(\theta)
\mathcal{P}_{\theta}(s^{\prime}|s,a)\mathcal{R}_{\theta}(r|s,a)$. Accordingly, $\mathcal{P}^{+}$ and $\mathcal{R}^{+}$ can be defined as follows:

$$ $\mathcal{P}^{+}((s^{\prime},b^{\prime\prime})|(s,b),a)=\mathds{1}(b^{\prime \prime}=b^{\prime})\int_{\theta}\mathcal{P}_{\theta}(s^{\prime}|s,a)b(\theta)d \theta,\ \mathcal{R}^{+}((s,b),a)=\int_{\theta}\mathcal{R}_{\theta}(s,a)b( \theta)d\theta$ (1) $$

The Q-function that satisfies the Bellman optimality equations: $(\forall x=(s,b)\in\mathcal{S}^{+},a\in\mathcal{A})$

$$ $Q^{*}(x,a)=\mathcal{R}^{+}(x,a)+\gamma\int_{x^{\prime}}V^{*}(x^{\prime}) \mathcal{P}^{+}(x^{\prime}|x,a)dx^{\prime},\ V^{*}(x^{\prime})=\max_{a}Q^{*}(x ^{\prime},a)$ (2) $$

is the Bayes-optimal Q-function and $\pi^{*}(s,b)=\arg\max_{a}Q^{*}((s,b),a)$ is the Bayes-optimal policy. Actions derived from $\pi^{*}$ are executed in the real MDP and constitute the best course of action for a Bayesian agent with respect to its prior belief $b_{0}$ over the environment (Guez et al. ()). A BAMDP can be cast into a partially observable MDP (POMDP)(^1^11Such a POMDP can be defined as a tuple $\langle\mathcal{S}^{+},\mathcal{A},\mathcal{P}^{+},\mathcal{R}^{+},\mathcal{Z} ,\mathcal{S},\gamma\rangle$, according to its original definition in (Littman ()).) by viewing $\mathcal{S}^{+}$ and $\mathcal{S}$ as the state and observation spaces, respectively, and setting the observation function as $\mathcal{Z}((s^{\prime},b^{\prime}),a,o)=\mathds{1}(s^{\prime}=o)$. ($\mathcal{Z}$ specifies the probability of observing $o$ under the previous action $a$ and current state $(s^{\prime},b^{\prime})$.) As a result, approaches developed for POMDPs can potentially be used to solve BAMDPs.

Bayesian Reinforcement Learning (BRL, Ghavamzadeh et al. ()), as introduced above, is a principled approach to dealing with uncertainty in the world model $\mathcal{M}_{\theta}$ and has two main advantages: (1) Domain knowledge can be injected by defining a proper prior belief; (2) A Bayes Adaptive policy solves the exploration-exploitation dilemma by explicitly including the belief in its state representation and incorporating belief updates into the planning process (Sorg et al. ()). However, Bayes-optimal planning is generally intractable, for which we introduce some approximate methods in the next section.

## 3 Related Works

Offline RL (Levine et al. (); Chen et al. ()) enables an agent to learn control policies from datasets of environment transitions pre-collected by a behavior policy $\mu$, i.e., $\mathcal{D}_{\mu}=\{[(s_{t}^{i},a_{t}^{i},r_{t}^{i})_{t=1}^{T}]_{i=1}^{N}\}$, circumventing the need for potentially expensive or unsafe online interactions. Model-free offline RL methods directly learn value/policy functions from $\mathcal{D}_{\mu}$ but restrict the policy to remain close to $\mu$ or to behave within the support of $\mathcal{D}_{\mu}$. Offline Model-based RL (MBRL) methods, on the other hand, explicitly learn world models $\mathcal{M}_{\theta}$ from $\mathcal{D}_{\mu}$ (through supervised learning) and adopt $\mathcal{M}_{\theta}$ as a surrogate simulator, enabling the learned policy to possibly generalize to states beyond $\mathcal{D}_{\mu}$. Specifically, both planning methods (Argenson & Dulac-Arnold (); Zhan et al. (); Diehl et al. ()), such as MPC (Garcia et al. ()), and RL methods (Yu et al. (); Kidambi et al. (); Lu et al. (); Yu et al. (); Guo et al. ()) can be applied on top of the learned $\mathcal{M}_{\theta}$. However, since $\mathcal{D}_{\mu}$ may not span the entire state-action space, $\mathcal{M}_{\theta}$ is unlikely to be globally accurate. Learning/Planning without any safeguards against such model inaccuracy can yield poor results. In this case, the authors of (Yu et al. (); Kidambi et al. (); Lu et al. ()) propose learning an ensemble of world models, using ensemble-based uncertainty estimations to construct a pessimistic MDP (P-MDP), and learning a near-optimal policy atop it. Ideally, for any policy, the performance in the real environment is lower-bounded by the performance in the corresponding P-MDP (with high probability), thus avoiding being overly optimistic about an inaccurate model. For practical implementation of the P-MDP, Lu et al. () conduct a thorough empirical study to compare a series of heuristics. However, none of these offline MBRL methods have modeled the problem as a BAMDP, even though Bayesian RL/planning provides a principled framework for handling model uncertainty. Additionally, with the learned model $\mathcal{M}_{\theta}$, deep and broad search in the policy learning process is possible, but the offline RL methods introduced above have not attempted such search-based policy improvement.

Monte-Carlo Tree Search (MCTS, Browne et al. ()) has been successfully integrated with RL, as exemplified by AlphaZero (Silver et al. ()) and MuZero (Schrittwieser et al. ()). These methods have achieved superhuman performance in domains requiring highly precise and sophisticated decision-making processes.
AlphaZero relies on given world models, whereas MuZero learns the world model and policy simultaneously by interacting with the environment. Although there have been various extensions of MuZero (Hubert et al. (); Schrittwieser et al. (); Ye et al. (); Danihelka et al. (); Antonoglou et al. (); Oren et al. (); Chen et al. (); Zhao et al. ()), most algorithms are designed for online MBRL. According to Niu et al. (), the applications of MuZero in offline learning, especially for continuous control in highly stochastic environments, which is our focus, still require significant improvement. Although search-based, MuZero learns a dynamics model and a reward model in a latent state space rather than the real state space, does not apply uncertainty estimation of the learned models as in usual offline MBRL methods, and does not adopt the Bayesian RL framework, making it fundamentally different from our algorithm design.

As mentioned in Section 2, Bayes-optimal planning is typically intractable. Approximate methods, such as (Asmuth et al. (); Sorg et al. (); Castro & Precup (); Asmuth & Littman (); Wang et al. (); Fonteneau et al. (); Guez et al. (); Slade et al. ()) have been developed. As a representative work, BAMCP (Guez et al. ()) adopts MCTS for Bayes Adaptive planning and is shown to converge in probability to a near Bayes-optimal policy at the root node of the search tree. However, all these methods cannot be directly applied to large-scale MDPs with continuous state and action spaces. Guez et al. () and Hoang & Vien () attempt to extend the use of BAMCP to scenarios with continuous states and actions using value/policy function approximations, but they replace MCTS with simple Monte Carlo rollouts (Bertsekas & Castañón ()) and do not utilize UCT (Kocsis & Szepesvári ()). Tree search makes branches at each node in the tree, providing a more thorough look-ahead search than simple rollouts, and UCT, as an efficient exploration algorithm, is the basis for the convergence guarantee of BAMCP. Moreover, these planning algorithms are not designed for offline MBRL. How to incorporate search-based planning in the policy improvement stage of RL and how to handle the uncertainty of the learned world model require further exploration, which are shown in the following section.

## 4 Methodology

In this work, we propose a novel offline MBRL algorithm based on Bayes Adaptive MCTS. The core challenge is to design a Bayes Adaptive planning method that is efficient in large stochastic MDPs. In this case, we propose Continuous BAMCP in Section 4.2, which can be applied to continuous control tasks with high dynamics stochasticity. Then, in Section 4.3, we present a search-based policy iteration framework, where the search results are distilled into policy and value networks for policy improvement and policy evaluation, respectively, at each iteration. In this way, we integrate offline MBRL with Bayes Adaptive MCTS. Both components require the use of an ensemble of world models for practical implementation and uncertainty quantification, which is detailed in Section 4.1.

### 4.1 The Key Role of Deep Ensembles

Offline MBRL methods estimate world models $\mathcal{M}_{\theta}$ from a static dataset $\mathcal{D}_{\mu}$, which would inevitably induce epistemic uncertainty about the identity of the real MDP $\mathcal{M}^{*}$. Specifically, there could be various potential MDPs that behave identically on the limited set of states and actions in $\mathcal{D}_{\mu}$, but their dynamics and reward functions may differ, especially on out-of-sample states and actions. Thus, we are actually dealing with a distribution of world models that follow a prior distribution $b_{0}(\theta)\triangleq P(\mathcal{M}_{\theta}|\mathcal{D}_{\mu})$. As introduced in Section 2, Bayesian RL based on BAMDP is a principled framework for handling model uncertainty by explicitly including the belief over the models in its state representation.
Essentially, the belief is updated with experience, providing a measure of how the models’ uncertainty has changed since the beginning of the episode. As a result, the agent can adjust its behavior upon receiving new information that reduces the epistemic uncertainty. Such an adaptive policy is necessary to act optimally in offline RL, as demonstrated in (Ghosh et al. ()).

The idea of deep ensembles (Lakshminarayanan et al. ()) is to train multiple deep neural networks as approximations of a function, each using a different weight initialization and optimized with a different mini-batch sequence. For offline MBRL, we can learn an ensemble of dynamics models $\{\mathcal{P}_{\theta}^{1},\cdots,\mathcal{P}_{\theta}^{K}\}$ and reward models $\{\mathcal{R}_{\theta}^{1},\cdots,\mathcal{R}_{\theta}^{K}\}$(^2^22In some MBRL scenarios, a certain reward function is available, for instance, as defined by domain experts. Otherwise, the reward and dynamics function $(\mathcal{R}_{\theta}^{i},\mathcal{P}_{\theta}^{i})$ are usually trained as a unified probabilistic model $\mathcal{N}(\mu_{\theta}^{i},\sigma_{\theta}^{i})$, since the reward $r$ can be viewed as an element of the next state $s^{\prime}$.) from the dataset $\mathcal{D}_{\mu}$ by minimizing the following supervised learning loss: ($i=1,\cdots,K$)

$$ $\mathcal{L}(\mathcal{P}_{\theta}^{i})=-\mathbb{E}_{(s,a,s^{\prime})\sim \mathcal{D}_{\mu}}\left[\log\mathcal{P}_{\theta}^{i}(s^{\prime}|s,a)\right],\ \mathcal{L}(\mathcal{R}_{\theta}^{i})=-\mathbb{E}_{(s,a,r)\sim\mathcal{D}_{\mu }}\left[\log\mathcal{R}_{\theta}^{i}(r|s,a)\right]$ (3) $$

$\{(\mathcal{P}_{\theta}^{i},\mathcal{R}_{\theta}^{i})_{i=1}^{K}\}$ can be viewed as a set of independent and identically distributed (IID) samples from the prior $P(\mathcal{M}_{\theta}|\mathcal{D}_{\mu})$ and constitute a finite approximation of the space of world models. With such an ensemble, the belief over the world models can be converted to a mass function over a set of $K$ items, where the $i$-th element denotes the probability of being in the MDP $(\mathcal{P}_{\theta}^{i},\mathcal{R}_{\theta}^{i})$. In this case, a reasonable prior distribution is $b_{0}(\theta)=[1/K,\cdots,1/K]$, since these models are IID prior samples. After receiving a transition $(s,a,r,s^{\prime})$, the belief can be updated as follows:

$$ $b^{\prime}(\theta)(i)=x^{i}/\sum_{j=1}^{K}x^{j},\ x^{i}=b(\theta)(i)\mathcal{P }_{\theta}^{i}(s^{\prime}|s,a)\mathcal{R}_{\theta}^{i}(r|s,a)$ (4) $$

This update requires a single inference from each ensemble member, but can be parallelized for computational efficiency. Equation (4) is a practical implementation of the Bayesian posterior update based on deep ensembles, and the simplification of the definition of $b(\theta)$ makes the transitions in Bayesian RL (i.e., Equation (1)) both practical and efficient.

The ensemble can also be used for uncertainty quantification. As aforementioned, our algorithm relies on thorough search on the learned world models. Without any constraints on the search process, the learned policy may overfit to an inaccurate model (by overestimating the expected return) and fail in the true MDP.
Although the agent could adapt its belief and follow more reliable ensemble members in the Bayesian RL framework, there could be regions in the state-action space where none of the members generalize well, as they are all learned from a static offline dataset.
A typical solution is to construct a P-MDP (see Section 3), which lower-bounds the true MDP and discourages the policy from regions where there is large discrepancy between the true and learned world models.
We construct the P-MDP by modifying each reward estimation $r$ into $\tilde{r}$: ($\mu_{\theta}(s,a)=\sum_{i=1}^{K}b(\theta)(i)\mu_{\theta}^{i}(s,a)$)

$$ $\tilde{r}(s,a,r,b(\theta))=r-\lambda\sqrt{\sum_{i=1}^{K}b(\theta)(i)(\sigma_{ \theta}^{i}(s,a)^{2}+\mu_{\theta}^{i}(s,a)^{2})-\mu_{\theta}(s,a)^{2}}$ (5) $$

The reward penalty is weighted by a hyperparameter $\lambda>0$ and corresponds to the standard deviation (std) of the mixture of Gaussian dynamics models, where $\mu_{\theta}^{i}$ and $\sigma_{\theta}^{i}$ are the mean and std from the ensemble member $i$.
This penalty design combines epistemic and aleatoric model uncertainty and has been shown to be successful at capturing errors in predicted dynamics (Lu et al. ())(^3^33In the original literature (Lakshminarayanan et al. (); Lu et al. ()), the ensemble is treated as a uniformly-weighted mixture model, i.e., $b(\theta)(i)=1/K,\ (i=1,\cdots K)$, since their belief is not adaptive. Equation (5) fits into the Bayesian RL framework by adapting the belief $b(\theta)$, which is part of our novelty.).

Figure: Algorithm 1 Continuous BAMCP

### 4.2 Bayes Adaptive MCTS in Continuous State and Action Spaces

BAMCP (Guez et al. ()) has been successful in solving large-scale BAMDPs, as detailed in Appendix A, but it is limited to scenarios with discrete state and action spaces.
In this subsection, we introduce an online planning method to approximate the Bayes-optimal policy at a decision point $(s,h)$. ($h$ denotes the transition history that ends at $s$.) In particular, this method can be used to solve BAMDPs with continuous states and actions.
Specifically, we adopt double progressive widening (DPW, Couëtoux et al. (); Auger et al. ()), which maintains a finite list of chance nodes to be searched at each decision point and incrementally adds a new chance node to the list based on the visitation counts. For a node $(s,h)$, a new action $a$ is sampled (with the current policy) and added to its children set $C((s,h))$, if $\left\lfloor{N((s,h))^{\alpha}}\right\rfloor\geq|C((s,h))|$, where $\alpha\in(0,1)$ is a hyperparameter that controls the growth rate and $N$ denotes the visitation counts. Otherwise, an action will be sampled from $C((s,h))$ according to the UCT (Kocsis & Szepesvári ()) rule. Similarly, to handle the infinitely many possible transitions, a new next state $s^{\prime}$ is added to the children set $C((s,h),a)$ only if $\left\lfloor{N((s,h),a)^{\beta}}\right\rfloor\geq|C((s,h),a)|$ ($\beta\in(0,1)$). Otherwise, the least visited child in $C((s,h),a)$ will be selected as the next state. With DPW, the sets of possible actions or next states for search are all finite, allowing deep search as in discrete scenarios, and the more promising states and actions (with higher $N$) have more subsequent branches, thereby reducing corresponding estimation uncertainty.
However, directly integrating DPW and BAMCP (i.e., Algorithm 3) likely cannot solve BAMDPs with continuous state and action spaces. As introduced in Appendix A, BAMCP relies on root sampling, which samples dynamics functions (from $b(\theta)$) only at the root node and follows a specific dynamics function throughout a simulation procedure. However, the rationale of root sampling (i.e., Lemma A.1) does not hold when applying DPW(^4^44The last equality of Eq. (6) does not hold, since $\tilde{b}(\theta|has^{\prime})\propto\tilde{b}(\theta|ha)\widetilde{\mathcal{P }}_{\theta}(s^{\prime}|s,a)\neq\tilde{b}(\theta|ha)\mathcal{P}_{\theta}(s^{ \prime}|s,a)$. $\widetilde{\mathcal{P}}_{\theta}(s^{\prime}|s,a)$ represents the state transition distribution when applying DPW, which differs from the true distribution $\mathcal{P}_{\theta}(s^{\prime}|s,a)$, as dictated by the DPW rule.). Also, BAMCP assumes that the reward function is certain rather than a function of $\theta$.

On the other hand, Polynomial Upper Confidence Trees (PUCT, Auger et al. ()), based on DPW, is a provably consistent planning method for solving MDPs with infinite-scale state/action spaces and arbitrarily stochastic transition kernels.
In this case, we propose casting BAMDPs into MDPs (i.e., $\mathcal{M}^{+}$ defined in Section 2) and solving them with PUCT. The pseudo code is shown as Algorithm 1. Ideally, as the number of simulations $E\rightarrow\infty$, PUCT can find a near-optimal solution of $\mathcal{M}^{+}$, which is also a near Bayes-optimal solution for the true environment. Each simulation follows a path that begins at the root node and ends at a leaf node, utilizing progressive widening for sampling both actions and next states, as illustrated in the ActionPW and StatePW procedures. Compared to PUCT, the significant modifications include: (1) replacing $\langle\mathcal{S},\mathcal{P},\mathcal{R}\rangle$ with their extended definitions in BAMDPs, $\langle\mathcal{S}^{+},\mathcal{P}^{+},\mathcal{R}^{+}\rangle$, and (2) applying reward penalties to account for model uncertainty. To be specific, in StatePW, $r$ and $s^{\prime}$ are sampled from a mixture of ensemble members, which is a practical implementation of sampling from $\mathcal{R}^{+}$ and $\mathcal{P}^{+}$ as shown in Eq. (1). After receiving the transition $(s,a,r,s^{\prime})$, the belief vector $b(\theta)$ is updated to $b^{\prime}(\theta)$ following Eq. (4), finishing the transition in $\mathcal{S}^{+}$ from $(s,b(\theta))$ to $(s,b^{\prime}(\theta))$. Moreover, the reward $r$ is adjusted with a penalty term as defined in Eq. (5), which is then used to calculate the return $R$. As aforementioned, applying such a reward penalty can effectively mitigate the issue of model exploitation.

The proposed algorithm can be used to approximate the Bayes-optimal policy at $(s,h)$, which is $\pi_{\text{ret}}(a|(s,h))\propto N((s,h),a),a\in C((s,h))$ (Auger et al. ()). However, our objective is to solve the entire BAMDP offline, eliminating the need for anything beyond simple lookup or inference using the learned policy function during execution. This necessitates a well-learned policy function at each state, but we cannot execute Algorithm 1 at every $(s,h)$ due to the scale of the state space. Therefore, we integrate the planning algorithm into a policy iteration framework, which is introduced in the next subsection. In this case, $\pi$ and $V$ in Algorithm 1 denote the policy and value functions from the previous learning iteration.(^5^55$\pi$ and $V$ are implemented as functions of $(s,h)$ because the states in BAMDPs consist of both $s$ and the corresponding belief $b(\theta)$, with $b(\theta)$ being a function of the history $h$ (according to its recursive updating rule).) As further details, multiple terms (labeled in blue) in Algorithm 1 have alternative design choices across different literatures, which we elaborate on in Appendix B.

### 4.3 The Overall Framework: Search-based Policy Iteration

In this subsection, we present a framework (inspired by MuZero(^6^66The novelty of our algorithm in comparison to MuZero is detailed in Section 3.)) that integrates continuous BAMCP (i.e., Algorithm 1) into policy improvement and policy evaluation. By iteratively running these procedures, we can approach a near Bayes-optimal policy, i.e., $\pi$, that can be directly referred to during execution in the true environment. The pseudo code of the overall framework is shown as Algorithm 2. For efficiency, a learner and a number of actors execute in parallel, reading from and sending data to the replay buffer $\mathcal{D}$ respectively. The actors update their copies of policy and value functions every $E_{l}$ learner steps.

In particular, each actor would interact with the learned world models to sample trajectories $\tau$.(^7^77In Algorithm 2, $\hat{\rho}(s)$ denotes the empirical distribution of initial states in the offline dataset $\mathcal{D}_{\mu}$.) At each time step of the trajectory, a Search procedure (defined in Algorithm 1) is executed using the actor’s current copy of $\pi$ and $V$ for planning at the current decision point. The search result $\pi_{\text{ret}}$ is then used to indicate the action choice, i.e., $a\sim\pi_{\text{ret}}(\cdot|(s,h))$. The subsequent transition process follows a BAMDP, where $r\sim\mathcal{R}^{+}(\cdot|(s,h),a)$, $s^{\prime}\sim\mathcal{P}^{+}(\cdot|(s,h),a)$, and the belief $b({\theta})$ adapts with experience, as described in StatePW. The collected trajectories are used in the learning process, where $\pi$ is trained to mimic the search result $\pi_{\text{ret}}$ by minimizing a cross-entropy loss (i.e., $\mathcal{L}(\pi,\{\tau_{i}\}_{i=1}^{B})$), while $V$ is trained to predict the $n$-step pessimistic return $z$.(^8^88In the definition of $z$, $\tilde{r}^{(0)}=\tilde{r}$ is the pessimistic reward associated with $(s,h)$ and the subscript $(i)$ denotes that the corresponding quantity is collected in $i$ time steps within that trajectory.) As noted in (Hubert et al. ()), $\pi_{\text{ret}}$ improves $\pi$ at each decision point, so repeatedly applying continuous BAMCP to obtain $\pi_{\text{ret}}$ and projecting the search results to the parameter space of $\pi$ (by minimizing $\mathcal{L}(\pi,\{\tau_{i}\}_{i=1}^{B})$) represents the policy improvement process. Meanwhile, value backups to update the Q estimates (i.e., $R\leftarrow V((s^{\prime},hars^{\prime}))$ in Algorithm 1) and the temporal difference learning for the value function $V$ constitute the policy evaluation process. In this way, the search algorithm is used as a powerful improvement operator to iteratively enhance the performance of the learned policy $\pi$.

Figure: Algorithm 2 Search-based Policy Iteration

## 5 Evaluation

To evaluate the effectiveness of each component in our algorithm design, we introduce three variants: (1) BA-MBRL leverages learned world models as surrogate simulators, applying reward penalties to collected transitions and using standard online RL algorithms (e.g., SAC (Haarnoja et al. ())) to learn a policy. While following existing offline MBRL methods, it models the problem as a BAMDP (rather than an MDP), with environment transitions defined by Equations (1) and (4) and the reward penalty by Equation (5), and is designed to evaluate the effectiveness of Bayesian RL. (2) BA-MCTS builds on BA-MBRL by introducing Continuous BAMCP (Algorithm 1) to plan at decision points, rather than inferring directly from the policy, to generate trajectories for downstream SAC, demonstrating the impact of deep search on policy learning. (3) BA-MCTS-SL, described in Algorithm 2, replaces the policy learning algorithm in BA-MCTS from policy gradient methods (as in SAC) with supervised learning (SL), allowing us to compare which approach offers a more efficient policy update mechanism, particularly for continuous control tasks.

**Table 1: Comparisons between the proposed algorithms and SOTA offline model-based RL methods on the D4RL benchmark suite. Each value represents the normalized score, as proposed in (Fu et al. ()), of the policy trained by the corresponding algorithm. These scores are undiscounted returns normalized to approximately range between 0 and 100, where a score of 0 corresponds to a random policy and a score of 100 corresponds to an expert-level policy. For our algorithms, we report the average score of the final ten policy learning epochs and its standard deviation across three random seeds. Results in the last four columns are taken from the original papers (Lu et al. (); Yu et al. (); Kidambi et al. (); Yu et al. ()), respectively.**
| Data Type | Environment | BA-MCTS<br>-SL (ours) | BA-MCTS<br>(ours) | BA-MBRL<br>(ours) | Optimized | COMBO | MOReL | MOPO |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| random | HalfCheetah | 29.20 $\pm$ 2.00 | 36.23 $\pm$ 1.04 | 32.76 $\pm$ 1.16 | 31.7 | 38.8 | 25.6 | 35.4 |
| random | Hopper | 33.83 $\pm$ 0.10 | 31.56 $\pm$ 0.12 | 31.47 $\pm$ 0.03 | 12.1 | 17.9 | 53.6 | 11.7 |
| random | Walker2d | 21.89 $\pm$ 0.07 | 21.59 $\pm$ 0.32 | 21.45 $\pm$ 0.53 | 21.7 | 7.0 | 37.3 | 13.6 |
| medium | HalfCheetah | 70.47 $\pm$ 3.52 | 75.84 $\pm$ 3.81 | 56.54 $\pm$ 5.20 | 45.7 | 54.2 | 42.1 | 42.3 |
| medium | Hopper | 97.75 $\pm$ 7.09 | 96.70 $\pm$ 14.0 | 98.25 $\pm$ 3.42 | 69.3 | 97.2 | 95.4 | 28.0 |
| medium | Walker2d | 82.24 $\pm$ 1.85 | 74.73 $\pm$ 3.25 | 75.41 $\pm$ 4.17 | 79.7 | 81.9 | 77.8 | 17.8 |
| med-replay | HalfCheetah | 61.16 $\pm$ 1.60 | 65.45 $\pm$ 0.81 | 62.50 $\pm$ 0.18 | 58.0 | 55.1 | 40.2 | 53.1 |
| med-replay | Hopper | 106.3 $\pm$ 0.13 | 101.8 $\pm$ 3.46 | 93.91 $\pm$ 4.25 | 90.8 | 89.5 | 93.6 | 67.5 |
| med-replay | Walker2d | 92.13 $\pm$ 5.13 | 95.06 $\pm$ 2.11 | 97.54$\pm$ 1.93 | 65.8 | 56.0 | 49.8 | 39.0 |
| med-expert | HalfCheetah | 80.53 $\pm$ 6.63 | 76.16 $\pm$ 10.3 | 90.52 $\pm$ 4.13 | 104.2 | 90.0 | 53.3 | 63.3 |
| med-expert | Hopper | 112.2 $\pm$ 0.29 | 108.3 $\pm$ 0.22 | 107.8 $\pm$ 0.37 | 105.8 | 111.1 | 108.7 | 23.7 |
| med-expert | Walker2d | 107.7 $\pm$ 0.82 | 110.0 $\pm$ 1.74 | 84.71 $\pm$ 0.87 | 97.1 | 103.3 | 95.6 | 44.6 |
| Average Score | 74.62 | 74.45 | 71.06 | 65.16 | 66.83 | 64.42 | 36.67 |  |

We first evaluate our algorithms on a widely-used continuous control benchmark for offline RL methods – D4RL MuJoCo (Fu et al. ()). The evaluation results for three types of robotic agents, each with offline datasets of four different qualities, are presented in Table 1. (1) Compared to SOTA offline MBRL methods, our algorithms achieve superior performance on nine out of twelve tasks. In terms of average performance, BA-MBRL significantly outperforms the baselines, demonstrating the effectiveness of using BAMDPs to handle model uncertainties in offline MBRL. (2) Both BA-MCTS and BA-MCTS-SL further improve upon BA-MBRL, highlighting the enhancement brought by deep search in policy learning. Notably, we apply Continuous BAMCP to only 10% of states when collecting training trajectories, while for the remaining states, we sample actions directly from the policy, i.e., $a\sim\pi(\cdot|s)$. Increasing the search ratio could further enhance policy performance at the cost of increased computation. (3) BA-MCTS-SL performs similarly to BA-MCTS, validating the effectiveness of both policy update mechanisms. However, BA-MCTS-SL struggles on Walker2d, where a warm-up training phase (using BA-MBRL) is required to establish a better initial policy. On the other hand, the advantage of the SL-based policy update is evident in the training plots of our algorithms in Figure 3, where BA-MCTS-SL exhibits much smoother learning curves compared to the other two algorithms, indicating greater robustness in model selection. (4) We further compare our algorithms with model-free offline policy learning methods, as shown in Appendix D. The performance improvement is even greater than that over model-based methods, highlighting the necessity of model-based learning. Particularly, when data quality is low, merely mimicking or staying close to the behavior policy would result in an underperforming policy. (5) We provide an ablation study in Appendix F to demonstrate the necessity of incorporating the reward penalty in offline MBRL to prevent the overexploitation of inaccurate world models. Additionally, we find that the SL-based policy update is less sensitive to model inaccuracies.

Figure: (a) hc-med-expert
Refer to caption: x1.png

Figure: Figure 2: Evaluation results for the tokamak control tasks. The figure shows the change in episodic returns over training epochs for the proposed algorithms and baselines across three target tracking tasks in the nuclear fusion scenario. Solid lines represent the average performance, while shaded areas indicate the 95% confidence intervals.
Refer to caption: x13.png

For fair comparisons and real-time execution, we do not perform test-time search and adopt the same policy architecture as the baselines, i.e., a feedforward neural network, rather than an RNN that incorporates transition history as input. These alternative designs have the potential to further improve our algorithms. For implementation, we build on the codebase of Optimized (Lu et al. ()), which thoroughly explores design choices in offline MBRL, making minimal changes to the code and hyperparameter settings(^9^99The detailed hyperparameter setups of our algorithms are provided in Appendix C.). Therefore, we believe the performance improvements stem from the Bayesian RL framework and deep search component. Further, both components can be seamlessly integrated with other advancements in offline MBRL, such as more accurate world model learning and improved uncertainty quantification for constructing pessimistic MDPs.

MuZero also applies deep search to MBRL. To evaluate its performance on D4RL MuJoCo tasks, we use the open-source implementation and hyperparameter configurations of Sampled EfficientZero (Ye et al. ()) provided by LightZero (Niu et al. ()). Benchmarking results from LightZero indicate that Sampled EfficientZero, equipped with a Gaussian policy, achieves the best performance on (online) MuJoCo locomotion tasks compared to other MuZero variants. To adapt Sampled EfficientZero for offline learning, we employ the reanalyse technique proposed by (Schrittwieser et al. ()). The evaluation results are presented in Figure 1. For reference, the expert-level episodic returns (corresponding to scores of 100) for HalfCheetah, Hopper, and Walker2d are 12135, 3234.3, and 4592.3, respectively. As shown, the results are significantly worse compared to the performance of offline RL methods listed in Table 1, despite Sampled EfficientZero’s higher computational cost (detailed in Appendix E). Notably, both Sampled EfficientZero and BA-MCTS-SL rely on supervised learning for policy improvement. However, for continuous control tasks, the agent can only sample a finite number of actions at a decision point, and the search result (e.g., $\pi_{\text{ret}}$ in Algorithm 1) is a distribution over this finite set, which could be a poor approximation of the optimal action distribution. Thus, purely mimicking the search result may be less sample-efficient than policy gradient methods, as it fails to account for the continuous nature of the action space. Furthermore, world model learning is the foundation of MBRL and can be particularly challenging in continuous control and offline learning settings, where the state-action space is vast but training data is limited. Sampled EfficientZero integrates model learning and policy training into a single stage, which significantly increases the learning difficulty (compared to BA-MCTS-SL).

Finally, we evaluate our algorithms on three target tracking tasks in tokamak control. The tokamak is one of the most promising confinement devices for achieving controllable nuclear fusion, where the primary challenge lies in confining the plasma, i.e., an ionized gas of hydrogen isotopes, while heating it and increasing its pressure to initiate and sustain fusion reactions (Pironti & Walker ()). Tokamak control involves applying a series of direct actuators (e.g., neutral beam, ECH power, magnetic field) and indirect actuators (e.g., setting targets for the plasma shape and density) to confine the plasma to achieve a desired state or track a given target. This sophisticated physical process is an ideal test bed for our algorithms. Specifically, we use a well-trained data-driven dynamics model provided by Char et al. () as a “ground truth" simulator for the nuclear fusion process during evaluation, and generate a dataset containing 725270 transitions using this model for offline RL. We select a reference shot (i.e., an episode of a fusion process) from DIII-D(^10^1010DIII-D is a tokamak device located in San Diego, California, operated by General Atomics.), and use its trajectories of Ion Rotation, Electron Temperature, and $\beta_{n}$ as targets for three tracking tasks. These are critical quantities in tokamak control, particularly $\beta_{n}$, which serves as an economic indicator of the efficiency of nuclear fusion. The tracking tasks have a 28-dimensional state space and a 14-dimensional action space, both continuous. Moreover, these tasks are highly stochastic, as the underlying dynamics model is a probabilistic neural network and each state transition is a sample from this model. For details on the simulator, and the design of the state/action spaces and reward functions, please refer to Appendix G. We compare our algorithms with SOTA model-free and model-based offline RL methods, specifically CQL and Optimized. The learning performance on the three tracking tasks is shown in Figure 2, where the x-axis and y-axis represent the training epochs and (negative) full-shot tracking errors, respectively. Our algorithms consistently outperform the baselines. Notably, the offline dataset does not include the reference shot or any similar, nearby shots. Therefore, restricting the policy to stay close to the behavior policy, as done in model-free methods, can be problematic. Also, learning dynamics models for MBRL is quite challenging in this nuclear fusion scenario. Our algorithms share the same ensemble of dynamics models with “Optimized" for policy learning, and the comparisons can demonstrate the superiority of Bayesian RL and deep search. Figure 2 has been smoothed for visualization(^11^1111The episodic return is plotted every 10 training epochs, with the y-axis representing the average value of a sliding window of length 5.). We further report the average return over the final 10 training epochs in Table 2, and the conclusions align with those from the D4RL MuJoCo tasks, showing the robustness of our proposed algorithms.

**Table 2: Comparisons between the proposed algorithms and offline RL baselines on the target tracking tasks. For each algorithm, we report the average return of the final ten policy learning epochs and its standard deviation across three different random seeds.**
| Task | BA-MCTS<br>-SL (ours) | BA-MCTS<br>(ours) | BA-MBRL<br>(ours) | CQL | Optimized |
| --- | --- | --- | --- | --- | --- |
| Temperature | -21.16 $\pm$ 5.00 | -23.83 $\pm$ 9.66 | -29.35 $\pm$ 4.72 | -59.62 $\pm$ 1.57 | -83.55 $\pm$ 10.56 |
| Rotation | -14.14 $\pm$ 1.88 | -19.07 $\pm$ 5.85 | -31.33 $\pm$ 11.54 | -85.48 $\pm$ 2.72 | -71.54 $\pm$ 9.88 |
| $\beta_{n}$ | -37.03 $\pm$ 17.98 | -18.93 $\pm$ 1.75 | -23.4 $\pm$ 10.77 | -36.37 $\pm$ 1.17 | -57.84 $\pm$ 10.27 |
| Average | -24.11 | -20.61 | -28.03 | -60.49 | -70.98 |

## 6 Conclusion and Discussions

In this work, we propose framing offline model-based reinforcement learning (MBRL) as a Bayes Adaptive Markov Decision Process (BAMDP) to better address uncertainties in the world models learned from offline datasets. We also introduce a novel planning method for solving BAMDPs in continuous state and action spaces using Monte Carlo Tree Search. This planning process is integrated into a policy iteration framework, enabling the derivation of a policy suitable for real-time execution from the planning results. In our evaluation, we test several variants of our algorithms to separately highlight the effectiveness of Bayesian RL and deep search. Additionally, we compare two different approaches for policy updates (based on the search results) in continuous control tasks: supervised learning and policy gradient methods. Our findings demonstrate that: (1) adapting beliefs over an ensemble of world models based on experience yields more accurate model approximations for MBRL; (2) deep search improves learning performance by incorporating planning and additional computation input; and (3) while supervised-learning-based policy updates result in smoother learning curves, they may struggle in complex continuous control tasks due to their approximation of the continuous action space as a finite set of action samples. For future work, our algorithms can be further improved by integrating advancements in offline MBRL and Bayesian RL, such as Bayesian RL methods that do not rely on deep ensembles, techniques to address sparse rewards in MBRL, and more principled approaches to construct pessimistic MDPs beyond those based on ensemble discrepancy.

## Appendix A BAMCP

Figure: Algorithm 3 BAMCP

Bayes Adaptive Monte Carlo Planning (BAMCP, Guez et al. ()) is a sample-based online planning method, aiming to find the action $a^{*}$ that approximately maximizes the expected return at a decision point $(s,h)$ under the BAMDP. Its detailed pseudo code is shown as Algorithm 3. BAMCP has demonstrated success in solving BAMDPs with large-scale discrete state and action spaces. Its key algorithmic ideas include: (1) applying MCTS with an efficient exploration strategy – UCT (Kocsis & Szepesvári ()) to the BAMDP in order to simulate the outcomes of different action choices; (2) utilizing root sampling to avoid frequent Bayesian posterior updates. Specifically, the UCT rule is used for selecting actions at non-leaf nodes, i.e., $a\leftarrow\arg\max_{x}Q((s,h),x)+c\sqrt{\frac{\log N((s,h))}{N((s,h),x)}}$, managing the tradeoff between exploration and exploitation. Root sampling refers to sampling the dynamics model only at the root node (i.e., $\theta\sim b(\cdot)$) and not adapting the belief $b(\cdot)$ according to the Bayes rule during the search process, of which the rationality is justified in the following lemma.

###### Lemma A.1 .

For all suffix histories $h^{\prime}$ of $h$, $b(\theta|h^{\prime})=\tilde{b}(\theta|h^{\prime})$. Here, $b(\theta|h^{\prime})$ is the true posterior probability of $\theta$ at the decision point $h^{\prime}$, while $\tilde{b}(\theta|h^{\prime})$ is the probability of experiencing $\theta$ at $h^{\prime}$ when using root sampling.

###### Proof.

This lemma can be proved by induction.

Base case: When $h^{\prime}=h$, $b(\theta|h^{\prime})=\tilde{b}(\theta|h^{\prime})=b(\theta)$.

Step case:

$$ $\displaystyle b(\theta|has^{\prime})$ $\displaystyle=P(has^{\prime}|\theta)P(\theta)/P(has^{\prime})$ (6) $\displaystyle=P(h|\theta)\mathcal{P}_{\theta}(s^{\prime}|s,a)P(\theta)/P(has^{ \prime})$ $\displaystyle=b(\theta|h)P(h)\mathcal{P}_{\theta}(s^{\prime}|s,a)/P(has^{ \prime})$ $\displaystyle=Zb(\theta|h)\mathcal{P}_{\theta}(s^{\prime}|s,a)$ $\displaystyle=Z\tilde{b}(\theta|h)\mathcal{P}_{\theta}(s^{\prime}|s,a)$ $\displaystyle=Z\tilde{b}(\theta|ha)\mathcal{P}_{\theta}(s^{\prime}|s,a)=\tilde {b}(\theta|has^{\prime})$ $$

Here, $Z=1/\int_{\theta}\mathcal{P}_{\theta}(s^{\prime}|s,a)b(\theta|h)d\theta=1/\int
_{\theta}\mathcal{P}_{\theta}(s^{\prime}|s,a)\tilde{b}(\theta|h)d\theta=1/\int
_{\theta}\mathcal{P}_{\theta}(s^{\prime}|s,a)\tilde{b}(\theta|ha)d\theta$ is the normalization constant. The fifth equality in Eq. (6) holds due to the inductive hypothesis. The sixth equality is based on the fact that the choice of $a$ at each node $h$ is made independently of the sample $\theta$. As for the last equality, to experience $\theta$ at $has^{\prime}$, the sample $\theta$ needs to traverse $ha$ (with probability $\tilde{b}(\theta|ha)$) and then the state $s^{\prime}$ needs to be sampled, which is with probability $\mathcal{P}_{\theta}(s^{\prime}|s,a)$, so $\tilde{b}(\theta|has^{\prime})\propto\tilde{b}(\theta|ha)\mathcal{P}_{\theta}(
s^{\prime}|s,a)$.
∎

###### Lemma A.1 .

###### Proof.

## Appendix B Alternative Design Choices for Continuous BAMCP

The terms labeled in blue in Algorithm 1 have alternative design choices. Empirical comparisons among these alternatives are reserved for future work.

$\tilde{Q}((s,h),x)$ in the exploration strategy, i.e., $a\leftarrow\arg\max_{x\in C((s,h))}\tilde{Q}((s,h),x)$, could take various forms. For instance, in (Couëtoux et al. (); Guez et al. (); Lee et al. ()), $\tilde{Q}((s,h),x)=Q((s,h),x)+c\sqrt{\frac{\log N((s,h))}{N((s,h),x)}}$; in PUCT (Auger et al. ()), $\tilde{Q}((s,h),x)=Q((s,h),x)+\sqrt{\frac{N((s,h))^{e(d)}}{N((s,h),x)}}$, where $e(d)$ is a schedule of coefficients related to the search depth $d$; in Sampled MuZero (a variant of MuZero that can be applied in continuous action spaces (Hubert et al. ())), $\tilde{Q}((s,h),x)=Q((s,h),x)+\hat{\pi}(x|(s,h))\frac{\sqrt{N((s,h))}}{1+N((s,
h),x)}\left(c_{1}+\log\left(\frac{N((s,h))+c_{2}+1}{c_{2}}\right)\right)$.
Here, $c$, $c_{1}$, $c_{2}$ are hyperparameters, $\hat{\pi}=\hat{\beta}\pi^{1-1/\tau}$ is a sample policy defined upon the real policy $\pi$. In particular, at each decision point $(s,h)$, Sampled MuZero would sample $M$ actions $\{a_{i}\}$ from the distribution $\pi^{1/\tau}$ and accordingly define $\hat{\beta}(a|(s,h))=\sum_{i}\mathds{1}_{a_{i}=a}/M$, where $\tau>0$ is a temperature hyperparameter. Thus, Sampled MuZero does not adopt progressive widening like ours. Following BAMCP, we adopt the first definition of $\tilde{Q}((s,h),x)$, though it could potentially be improved with other choices. In addition, as an implementation trick (Hamrick et al. ()), the Q estimates are usually normalized into $\bar{Q}\in[0,1]$ before being used to calculate $\tilde{Q}$ as above. The normalized estimates can be computed as $\bar{Q}((s,h),x)=\frac{Q((s,h),x)-Q_{\text{min}}}{Q_{\text{max}}-Q_{\text{min}}}$, where $Q_{\text{max}}$ and $Q_{\text{min}}$ are the maximum and minimum Q values observed in the search tree so far.

As the planning/search result, $\pi_{\text{ret}}$ can take multiple forms. In (Guez et al. (); Sunberg & Kochenderfer (); Lee et al. ()), $\pi_{\text{ret}}((s,h))=\arg\max_{a\in C((s,h))}Q((s,h),a)$; in (Sampled) MuZero, $\pi_{\text{ret}}(a|(s,h))=\frac{N((s,h),a)^{1/\tau}}{\sum_{x\in C((s,h))}N((s,
h),x)^{1/\tau}}$; in ROSMO (a variant of MuZero with improved performance in offline scenarios (Liu et al. ())), $\pi_{\text{ret}}(a|(s,h))\propto\pi(a|(s,h))\exp(Q((s,h),a)-V((s,h)))$. Here, $\tau\in(0,1]$ is a temperature parameter and decays with the training process, ensuring the action selection becomes greedier. We select the second form for $\pi_{\text{ret}}$ in Algorithm 1. This is because (1) as described in PUCT, the returned action should be the most visited one, which is not necessarily the one with the highest Q value, and (2) ROSMO adopts one-step look-ahead rather than deep tree search at each root node, which does not align with our approach.

As for the conditions of double progressive widening, PUCT designs $\alpha$ and $\beta$ to be functions of the search depth $d$, while UCT-DPW (Couëtoux et al. ()) utilizes a different set of conditions: $\left\lceil{K_{a}N((s,h))^{\alpha}}\right\rceil\geq|C((s,h))|$, $\left\lceil{K_{s}N((s,h),a)^{\beta}}\right\rceil\geq|C((s,h),a)|$, where $K_{a},K_{s},\alpha,\beta$ are all constant hyperparameters. When the progressive widening condition for sampling the next state is not satisfied, either the least visited node in $C((s,h),a)$ can be selected (following PUCT), or a random node can be sampled from $C((s,h),a)$ following a distribution proportional to the number of visits (following UCT-DPW). As shown in Algorithm 1, we follow the designs of PUCT, but keep $\alpha$ and $\beta$ as constants for simplicity in hyperparameter fine-tuning.

Finally, the condition for continuing the simulation procedure, i.e., $N((s,h))>1$, could potentially be replaced with $N((s,h),a)>1$ or $N((s^{\prime},hars^{\prime}))>0$. These conditions indicate that the nodes $(s,h)$, $(s,h,a)$, and $(s^{\prime},hars^{\prime})$ have been visited before, respectively. At the end of the simulation procedure, we can either apply rollouts, i.e., simulating a single path until the end of an episode, to estimate the expected value for a leaf node $(s,h)$, or directly use $V((s,h))$ as the estimation. The former approach is widely used in online planning algorithms (Guez et al. (); Sunberg & Kochenderfer (); Lee et al. ()), while the latter is used in iterative frameworks like MuZero.

## Appendix C Key Hyperparameter Setup

**Table 3: Key hyperparameters of the proposed algorithms for each evaluation task. $K$: ensemble size, $\lambda$: reward penalty coefficient, $H$: rollout horizon, $N$: number of training epochs.**
| Data Type | Environment | BA-MBRL | BA-MCTS | BA-MCTS-SL |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $K$ | $\lambda$ | $H$ | $N$ | $K$ | $\lambda$ | $H$ | $N$ | $K$ | $\lambda$ | $H$ | $N$ |  |  |
| random | HalfCheetah | 10 | 7 | 6 | 200 | 10 | 7 | 6 | 800 | 10 | 7 | 6 | 500 |
| random | Hopper | 6 | 50 | 47 | 700 | 6 | 50 | 47 | 800 | 6 | 50 | 47 | 500 |
| random | Walker2d | 10 | 0.5 | 20 | 700 | 10 | 0.5 | 20 | 800 | 10 | 0.5 | 20 | 500 |
| medium | HalfCheetah | 12 | 6 | 6 | 300 | 12 | 6 | 5 | 800 | 12 | 6 | 5 | 500 |
| medium | Hopper | 12 | 40 | 42 | 200 | 12 | 40 | 42 | 800 | 12 | 40 | 42 | 200 |
| medium | Walker2d | 8 | 5 | 20 | 700 | 8 | 5 | 20 | 800 | 8 | 5 | 20 | 500 |
| med-replay | HalfCheetah | 11 | 40 | 10 | 300 | 11 | 40 | 10 | 800 | 11 | 40 | 10 | 500 |
| med-replay | Hopper | 7 | 5 | 5 | 700 | 7 | 5 | 5 | 800 | 7 | 5 | 5 | 500 |
| med-replay | Walker2d | 13 | 2.5 | 47 | 1000 | 13 | 2.5 | 47 | 800 | 13 | 2.5 | 47 | 500 |
| med-expert | HalfCheetah | 7 | 100 | 5 | 1000 | 7 | 100 | 5 | 800 | 7 | 100 | 5 | 1100 |
| med-expert | Hopper | 12 | 40 | 43 | 600 | 12 | 40 | 43 | 800 | 12 | 40 | 43 | 500 |
| med-expert | Walker2d | 6 | 20 | 37 | 400 | 6 | 20 | 37 | 800 | 6 | 20 | 37 | 500 |

In Table 3, we list the key hyperparameters of the proposed algorithms. For each task, an ensemble of $K$ dynamics and reward models is trained using the provided offline dataset. These learned models are then utilized as a simulator to train a control policy using off-the-shelf RL methods, such as SAC. The policy is trained for $N$ epochs. At each epoch, $50000H$ transitions are sampled by interacting with the simulator, followed by $1000$ RL training iterations. In particular, $50000$ states are randomly sampled from the offline dataset, with each state followed by a rollout lasting $H$ time steps. To mitigate overestimation, a reward penalty based on the discrepancy among the ensemble members is applied with a coefficient $\lambda$, as shown in Equation (5). The setups for $K$, $\lambda$, and $H$ are almost the same across the three algorithms and primarily inherited from the baseline – “Optimized" (Lu et al. ()), to make sure the improvements are brought by the Bayesian RL and deep search components.

The policy is evaluated on the ground truth environment for 10 episodes at the end of each training epoch. We report the average scores across the final 10 training epochs of our algorithms in Tables 1 and 5. It is important to note that increasing the number of training epochs $N$ does not necessarily lead to better policy performance, since the training is based on learned dynamics and reward models rather than the ground truth.
According to (Lu et al. ()) and our experiments, the hyperparameters listed above can significantly influence the performance of model-based RL. Adjusting these hyperparameters could either enhance or impair the learning performance of our algorithms. We also suspect that the performance of the baselines listed in Tables 1 and 5, which are from their original papers, could be further improved by fine-tuning the relevant hyperparameters.

**Table 4: Important hyperparameters used in the search process.**
| Data Type | Environment | BA-MCTS | BA-MCTS-SL |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\rho$ | $\alpha$ | $c$ | $\eta$ | $n_{s}$ | $n_{a}$ | $\rho$ | $\alpha$ | $c$ | $\eta$ | $n_{s}$ | $n_{a}$ | $N_{SL}$ | $N_{P}$ |  |  |
| random | HalfCheetah | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 5 | 0 |
| random | Hopper | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 5 | 0 |
| random | Walker2d | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 5 | 100 |
| medium | HalfCheetah | 0.1 | 0.5 | 1.0 | 0.1 | 5 | 10 | 0.1 | 0.5 | 1.0 | 0.1 | 5 | 10 | 20 | 0 |
| medium | Hopper | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 0.1 | 0.5 | 1.0 | 0.3 | 1 | 20 | 5 | 0 |
| medium | Walker2d | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 0.1 | 0.5 | 2.5 | 0.3 | 1 | 20 | 5 | 100 |
| med-replay | HalfCheetah | 0.1 | 0.8 | 1.0 | 0.3 | 5 | 10 | 0.1 | 0.8 | 2.5 | 0.3 | 1 | 20 | 5 | 0 |
| med-replay | Hopper | 0.1 | 0.8 | 1.0 | 0.1 | 1 | 20 | 0.1 | 0.8 | 1.0 | 0.1 | 1 | 20 | 15 | 0 |
| med-replay | Walker2d | 0.1 | 0.8 | 2.5 | 0.3 | 1 | 20 | 0.1 | 0.8 | 2.5 | 0.3 | 1 | 20 | 5 | 200 |
| med-expert | HalfCheetah | 0.1 | 0.8 | 1.0 | 0.3 | 5 | 10 | 0.1 | 0.8 | 2.5 | 0.3 | 1 | 20 | 5 | 0 |
| med-expert | Hopper | 0.1 | 0.8 | 1.0 | 0.3 | 1 | 20 | 0.1 | 0.8 | 1.0 | 0.3 | 1 | 20 | 5 | 0 |
| med-expert | Walker2d | 0.1 | 0.8 | 2.5 | 0.3 | 1 | 20 | 0.1 | 0.8 | 2.5 | 0.3 | 1 | 20 | 15 | 100 |

Figure: (a) hc-med-expert
Refer to caption: x14.png

BA-MCTS and BA-MCTS-SL utilize Bayes Adaptive Monte Carlo Tree Search to collect samples for offline model-based RL. Instead of performing a tree search at every state, we randomly select a proportion (i.e., $\rho$) of states from the available 50000 states at each rollout time step as root nodes for tree search. For the remaining states, actions are sampled directly from the policy, i.e., $a\sim\pi(\cdot|s)$. The tree search procedure is detailed in Algorithm 1, with the number of MCTS iterations, $E$, set to 50. Increasing $\rho$ and $E$ can potentially enhance performance, but it will also linearly increase the computational cost. Table 4 outlines the key hyperparameters related to the search process for each algorithm and task. (1) As described in Algorithm 1, the parameters $\alpha$ and $\beta$ control the rate of double progressive widening. ($\beta$ is set as 0.5 across all tasks.) To encourage deeper search, we limit the number of actions sampled from a state under $n_{a}$ and the number of next states sampled from an action under $n_{s}$, respectively. Action selection follows the UCT rule, as discussed in Appendix A, where $c>0$ balances the exploration and exploitation. Additionally, inspired by the success of MuZero in enhancing exploration, we introduce Dirichlet noise $x_{d}$ at the root nodes, where actions are sampled from a mixture of distributions: $a\sim\eta x_{d}+(1-\eta)\pi(\cdot|s)$ and $\eta$ controls the mixture rate. Notably, for $(c,\eta,n_{a},n_{s})$, we explore the set of possible combinations: $\{(2.5,0.3,20,1),(1.0,0.3,20,1),(1.0,0.1,20,1),(1.0,0.3,10,5),(1.0,0.1,10,5)\}$ during hyperparameter fine-tuning. We believe there are likely more optimal search settings yet to be discovered. (2) In BA-MCTS-SL, policy improvement is achieved through supervised learning. We find that, rather than learning solely from samples collected within the current epoch, incorporating a buffer of samples from the past $N_{SL}$ epochs helps to stabilize the learning process. Also, in the Walker2d environment, BA-MCTS-SL requires a warm-up training phase of $N_{P}$ epochs using BA-MBRL, allowing the initial policy to generate effective signals for supervised learning.

## Appendix D Comparisons with Model-free Methods on D4RL MuJoCo

**Table 5: Comparisons between the proposed algorithms and model-free offline policy learning methods on the D4RL benchmark suite. Each value represents the normalized score, as proposed in (Fu et al. ()), of the policy trained by the corresponding algorithm. These scores are undiscounted returns normalized to approximately range between 0 and 100, where a score of 0 corresponds to a random policy and a score of 100 corresponds to an expert-level policy. For our algorithms, we report the average score of the final ten policy learning epochs and its standard deviation across three different random seeds.**
| Data Type | Environment | BA-MCTS<br>-SL (ours) | BA-MCTS<br>(ours) | BA-MBRL<br>(ours) | CQL | BEAR | BRAC-v | SAC | BC |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| random | HalfCheetah | 29.20 $\pm$ 2.00 | 36.23 $\pm$ 1.04 | 32.76 $\pm$ 1.16 | 35.4 | 25.1 | 31.2 | 30.5 | 2.1 |
| random | Hopper | 33.83 $\pm$ 0.10 | 31.56 $\pm$ 0.12 | 31.47 $\pm$ 0.03 | 10.8 | 11.4 | 12.2 | 11.3 | 1.6 |
| random | Walker2d | 21.89 $\pm$ 0.07 | 21.59 $\pm$ 0.32 | 21.45 $\pm$ 0.53 | 7.0 | 7.3 | 1.9 | 4.1 | 9.8 |
| medium | HalfCheetah | 70.47 $\pm$ 3.52 | 75.84 $\pm$ 3.81 | 56.54 $\pm$ 5.20 | 44.4 | 41.7 | 46.3 | -4.3 | 36.1 |
| medium | Hopper | 97.75 $\pm$ 7.09 | 96.70 $\pm$ 14.0 | 98.25 $\pm$ 3.42 | 86.6 | 52.1 | 31.1 | 0.8 | 29.0 |
| medium | Walker2d | 82.24 $\pm$ 1.85 | 74.73 $\pm$ 3.25 | 75.41 $\pm$ 4.17 | 74.5 | 59.1 | 81.1 | 0.9 | 6.6 |
| med-replay | HalfCheetah | 61.16 $\pm$ 1.60 | 65.45 $\pm$ 0.81 | 62.50 $\pm$ 0.18 | 46.2 | 38.6 | 47.7 | -2.4 | 38.4 |
| med-replay | Hopper | 106.3 $\pm$ 0.13 | 101.8 $\pm$ 3.46 | 93.91 $\pm$ 4.25 | 48.6 | 33.7 | 0.6 | 3.5 | 11.8 |
| med-replay | Walker2d | 92.13 $\pm$ 5.13 | 95.06 $\pm$ 2.11 | 97.54$\pm$ 1.93 | 32.6 | 19.2 | 0.9 | 1.9 | 11.3 |
| med-expert | HalfCheetah | 80.53 $\pm$ 6.63 | 76.16 $\pm$ 10.3 | 90.52 $\pm$ 4.13 | 62.4 | 53.4 | 41.9 | 1.8 | 35.8 |
| med-expert | Hopper | 112.2 $\pm$ 0.29 | 108.3 $\pm$ 0.22 | 107.8 $\pm$ 0.37 | 111 | 96.3 | 0.8 | 1.6 | 111.9 |
| med-expert | Walker2d | 107.7 $\pm$ 0.82 | 110.0 $\pm$ 1.74 | 84.71 $\pm$ 0.87 | 98.7 | 40.1 | 81.6 | -0.1 | 6.4 |
| Average Score | 74.62 | 74.45 | 71.06 | 54.85 | 39.83 | 31.44 | 4.13 | 25.07 |  |

As a complement to Table 1, we compare our algorithms with a series of model-free offline policy learning (Chen et al. ()) methods. We include SOTA model-free offline RL methods: CQL (Kumar et al. ()), BEAR (Kumar et al. ()), and BRAC-v (Wu et al. ()). Additionally, we show the performance of directly applying SAC or Behavioral Cloning (BC, Chen et al. ()) to the provided offline dataset in the last two columns. The mean performance of the baselines are taken from related works (Yu et al. (); Kidambi et al. (); Fu et al. ()). Our algorithms show significantly better performance, demonstrating the necessity of model-based learning in these environments. The training plots of our proposed algorithms in each environment is further detailed in Figure 3.

## Appendix E Computation Cost of Sampled EfficientZero on D4RL MuJoCo

Here, we report the training time of Sampled EfficientZero on the D4RL MuJoCo tasks in Table 6. This set of experiments is run on a server with 40 Intel(R) Xeon(R) Gold 5215 CPUs and 4 Tesla V100-SXM2-32GB GPUs. Despite the substantial computational input, no corresponding performance improvement is observed. We welcome researchers to use our released code for further hyperparameter fine-tuning and to improve the results we have provided.

**Table 6: Training time (in hours) of Sampled EfficientZero for each evaluation task. The mean and standard deviation from repeated experiments are reported.**
| hc-med-expert | hc-med-replay | hc-medium | hc-random |
| --- | --- | --- | --- |
| 55.4 $\pm$ 1.3 | 54.8 $\pm$ 1.8 | 54.7 $\pm$ 1.9 | 54.8 $\pm$ 2.6 |
| hp-med-expert | hp-med-replay | hp-medium | hp-random |
| 66.1 $\pm$ 0.9 | 126.6 $\pm$ 12.0 | 70.5 $\pm$ 1.0 | 153.7 $\pm$ 19.4 |
| wk-med-expert | wk-med-replay | wk-medium | wk-random |
| 60.8 $\pm$ 1.8 | 134.9 $\pm$ 4.2 | 63.7 $\pm$ 1.5 | 175.7 $\pm$ 28.2 |

## Appendix F Ablation Study on the Reward Penalty

**Table 7: Comparison of the proposed algorithms with their corresponding versions without the reward penalty (i.e., $\lambda=0$). The definitions of the values in this table are consistent with those in Table 5.**
| Data Type | Environment | BA-MCTS<br>-SL | BA-MCTS | BA-MBRL | BA-MCTS<br>-SL ($\lambda=0$) | BA-MCTS<br>($\lambda=0$) | BA-MBRL<br>($\lambda=0$) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| random | HalfCheetah | 29.20 $\pm$ 2.00 | 36.23 $\pm$ 1.04 | 32.76 $\pm$ 1.16 | 34.80 $\pm$ 1.39 | 38.78 $\pm$ 1.65 | 39.64 $\pm$ 2.86 |
| random | Hopper | 33.83 $\pm$ 0.10 | 31.56 $\pm$ 0.12 | 31.47 $\pm$ 0.03 | 9.16 $\pm$ 0.16 | 7.44 $\pm$ 0.14 | 6.97 $\pm$ 0.07 |
| random | Walker2d | 21.89 $\pm$ 0.07 | 21.59 $\pm$ 0.32 | 21.45 $\pm$ 0.53 | 17.53 $\pm$ 6.16 | 21.53 $\pm$ 0.42 | 21.41 $\pm$ 0.64 |
| medium | HalfCheetah | 70.47 $\pm$ 3.52 | 75.84 $\pm$ 3.81 | 56.54 $\pm$ 5.20 | 61.64 $\pm$ 4.58 | 60.84 $\pm$ 2.00 | 41.49 $\pm$ 2.29 |
| medium | Hopper | 97.75 $\pm$ 7.09 | 96.70 $\pm$ 14.0 | 98.25 $\pm$ 3.42 | 102.8 $\pm$ 2.29 | 104.4 $\pm$ 1.88 | 93.68 $\pm$ 11.4 |
| medium | Walker2d | 82.24 $\pm$ 1.85 | 74.73 $\pm$ 3.25 | 75.41 $\pm$ 4.17 | 82.61 $\pm$ 0.86 | 57.01 $\pm$ 7.24 | 57.97 $\pm$ 15.4 |
| med-replay | HalfCheetah | 61.16 $\pm$ 1.60 | 65.45 $\pm$ 0.81 | 62.50 $\pm$ 0.18 | 42.10 $\pm$ 2.85 | 36.65 $\pm$ 2.39 | 44.03 $\pm$ 7.35 |
| med-replay | Hopper | 106.3 $\pm$ 0.13 | 101.8 $\pm$ 3.46 | 93.91 $\pm$ 4.25 | 107.9 $\pm$ 0.07 | 84.11 $\pm$ 2.97 | 91.81 $\pm$ 11.5 |
| med-replay | Walker2d | 92.13 $\pm$ 5.13 | 95.06 $\pm$ 2.11 | 97.54$\pm$ 1.93 | 88.61 $\pm$ 5.21 | 97.33 $\pm$ 3.51 | 98.19 $\pm$ 1.23 |
| med-expert | HalfCheetah | 80.53 $\pm$ 6.63 | 76.16 $\pm$ 10.3 | 90.52 $\pm$ 4.13 | 51.76 $\pm$ 5.31 | 26.60 $\pm$ 1.46 | 29.88 $\pm$ 2.28 |
| med-expert | Hopper | 112.2 $\pm$ 0.29 | 108.3 $\pm$ 0.22 | 107.8 $\pm$ 0.37 | 106.8 $\pm$ 6.34 | 81.76 $\pm$ 6.45 | 86.79 $\pm$ 18.7 |
| med-expert | Walker2d | 107.7 $\pm$ 0.82 | 110.0 $\pm$ 1.74 | 84.71 $\pm$ 0.87 | 110.8 $\pm$ 1.72 | 110.2 $\pm$ 0.91 | 53.35 $\pm$ 38.0 |
| Average Score | 74.62 | 74.45 | 71.06 | 68.04 | 60.55 | 55.43 |  |

To demonstrate the necessity of incorporating the reward penalty in offline MBRL, we conduct an ablation study by setting $\lambda$ in Eq. (5) to 0, resulting in ablated versions of our proposed three algorithms. The results are presented in Table 7. First, the average performance of the algorithms with the reward penalty is consistently better, demonstrating the importance of using reward penalties in offline MBRL to prevent the overexploitation of the learned world models (which can be inaccurate). Second, the supervised-learning-based algorithm (i.e., BA-MCTS-SL ($\lambda=0$)) is less affected by the absence of the reward penalty, compared to the policy-gradient-based methods. Notably, BA-MCTS-SL ($\lambda=0$) and BA-MCTS-SL achieve comparable performance in the Hopper and Walker2d tasks. This shows an additional advantage of BA-MCTS-SL – its reduced sensitivity to model inaccuracies. Lastly, there are instances where superior performance is achieved with $\lambda$ set to 0. For example, BA-MCTS-SL ($\lambda=0$) performs better than BA-MCTS-SL in 5 out of 12 tasks. This suggests that the performance of our algorithms in Tables 1 and 5 could be further improved by adjusting hyperparameters such as $\lambda$(^12^1212As mentioned in Section 5, we retain most of the hyperparameter settings from Optimized (Lu et al. ()) to ensure that the performance improvements are attributed to our algorithm design.).

## Appendix G Details of the Tokamak Control Tasks

**Table 8: The state and action spaces of the tokamak control tasks.**
| STATE SPACE |  |
| --- | --- |
| Scalar States | $\beta_{n}$, Internal Inductance, Line Averaged Density,<br>Loop Voltage, Stored Energy |
| Profile States | Electron Density, Electron Temperature, Pressure,<br>Safety Factor, Ion Temperature, Ion Rotation |
| ACTION SPACE |  |
| Targets | Current Target, Density Target |
| Shape Variables | Elongation, Top Triangularity, Bottom Triangularity, Minor Radius,<br>Radius and Vertical Locations of the Plasma Center |
| Direct Actuators | Power Injected, Torque Injected, Total Deuterium Gas Injection,<br>Total ECH Power, Magnitude and Sign of the Toroidal Magnetic Field |

Nuclear fusion is a promising energy source to meet the world’s growing demand. It involves fusing the nuclei of two light atoms, such as hydrogen, to form a heavier nucleus, typically helium, releasing energy in the process. The primary challenge of fusion is confining a plasma, i.e., an ionized gas of hydrogen isotopes, while heating it and increasing its pressure to initiate and sustain fusion reactions. The tokamak is one of the most promising confinement devices. It uses magnetic fields acting on hydrogen atoms that have been ionized (given a charge) so that the magnetic fields can exert a force on the moving particles (Pironti & Walker ()).

Char et al. () trained a deep recurrent network as a dynamics model for the DIII-D tokamak, a device located in San Diego, California, and operated by General Atomics, using a large dataset of operational data from that device. A typical shot (i.e., episode) on DIII-D lasts around 6-8 seconds, consisting of a one-second ramp-up phase, a multi-second flat-top phase, and a one-second ramp-down phase. The DIII-D also features several real-time and post-shot diagnostics that measure the magnetic equilibrium and plasma parameters with high temporal resolution. The authors demonstrate that the learned model predicts these measurements for entire shots with remarkable accuracy. Thus, we use this model as a “ground truth" simulator for tokamak control tasks. Specifically, we generate a dataset of 725270 transitions for offline RL and evaluate the learned policy using this data-driven simulator.

The state and action spaces for the tokamak control tasks are outlined in Table 8. For detailed physical explanations of their components, please refer to (Abbate et al. (); Char et al. (); Ariola et al. ()). The state space consists of five scalar values and six profiles which are discretized measurements of physical quantities along the minor radius of the toroid. After applying principal component analysis (Maćkiewicz & Ratajczak ()), the pressure profile is reduced to two dimensions, while the other profiles are reduced to four dimensions each. In total, the state space comprises 27 dimensions. The action space includes direct control actuators for neutral beam power, torque, gas, ECH power, current, and magnetic field, as well as target values for plasma density and plasma shape, which are managed through a lower-level control module. Altogether, the action space consists of 14 dimensions. While for certain tasks, it is possible to prune the state and action spaces to reduce the learning complexity, we have chosen not to apply any domain-specific knowledge in these evaluations for general RL algorithms. We reserve the domain-specific applications of our algorithms, which would require more domain knowledge and engineering efforts, as an important future work.

We select a reference shot from DIII-D, which spans 251 time steps, and use its trajectories of Ion Rotation, Electron Temperature, and $\beta_{n}$ as targets for three tracking tasks. Specifically, $\beta_{n}$ is the normalized ratio between plasma pressure and magnetic pressure, a key quantity serving as a rough economic indicator of efficiency. Since the tracking targets vary over time, we include the time step as part of the policy input. The reward function for each task is defined as the negative squared tracking error of the corresponding component (i.e., temperature, rotation, or $\beta_{n}$) at each time step, and the reward is normalized by the episode horizon (i.e., 251 time steps). Notably, for policy learning, the reward function is provided rather than learned from the offline dataset as in D4RL tasks; and the dataset does not include the reference shot or any nearby, similar shots.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. I am not sure about the necessity of adjusting the belief not only in the acting process but also in the searching process. Could the author explain the different influences of adjusting the model at the two stages?
2. The authors indicate that sez does not work well due to the difficulty of end to end model training, but sez also use a different way of choosing actions nodes from DPW. Can the authors explain why they do not consider it as a main factor?
3. I am curious about why the authors choose to update the policy in a trajectory way. MuZero just sample a series of transitions from the buffer, which can avoid the interaction in learned models. Could the authors explain the reason behind this design?
soundness: 2
strengths: 1. The authors describe the entire method in a clear way using pseudocode in Algorithms 1 and 2, and the insight of separating the model training from end-to-end process is explained in experiment section by comparing to sampled EfficientZero (sez). 
2. The experimental results indicate the effectiveness of the BA model training. Although the effectiveness of certain designs like MCTS search and supervised policy learning is not significant according to Table 1, I am glad to see the authors report the results honestly.
summary: This paper propose to use a Bayesian adaptive method to train the model instead of end to end training approach used by MuZero. As indicate by the authors, this can modify the difficulty of model training in complex scenarios like continuous action space. Specifically, the method can be described in three main designs:1) it adapts DPW to handle continuous action and state spaces. 2) it use BA to refine the model during interactions and search process 3) It adapts MuZero’s supervised way to perform policy improvement. The experiment section shows methods using BA generally have a better performance than baselines, but the advantage of certain design seems not significant.
weaknesses: 1. It seems the main improvement comes from BA model training according to Table 1. Thus, I wonder if there are other model-based RL methods using BA and deep ensembles to train the models, and why the authors did not include them in the baseline.
2. In my opinion, the BA process adjust the belief of the models according to the sample results of current belief instead of the interaction with true environment, which may lead to accumulated error. For example, if the belief mistakenly concentrates on a wrong model and the sample results are likely to come from the wrong model, then future belief are likely to be more inclined to this wrong model. Could the authors provide some insights that why this not happen according to the experiment results?

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: * Belief updates during training: are the beliefs updated using imaginary rollouts sampled from the model ensemble? If so, could this introduce systematic biases? The paper mentions applying a reward penalty within a pessimistic MDP framework, but it would be useful to assess whether this is sufficient for accurate belief updates under the BAMDP setup.

* Representation of returned policy in Algorithm 1: how is the policy $\pi_{\mathrm{ret}}$, returned by the SEARCH procedure, represented in practice? Further details would clarify its applicability to real-time policy learning.

* Clarification on state representation: footnote 5 states that the policy takes the states consisting of both $s$ and $h$, but line 416 describes differently. Could the authors clarify this discrepancy? Additionally, MAPLE (Chen et al., 2021) employs an RNN-based policy network, which aligns well with the BAMDP framework. Using MAPLE as a baseline could enhance consistency with the proposed architecture.
soundness: 3
strengths: * Practical MCTS implementation for offline MBRL formulated as BAMDP: the paper introduces a practical implementation of the MCTS algorithm specifically adapted for continuous control tasks under the BAMDP framework. Though a straightforward approach, this integration is effective and brings continuous MCTS to offline MBRL with clear potential.

* Well-executed ablative studies: the paper provides ablation studies that dissect the individual contributions of the algorithm’s components, demonstrating the benefits of the BAMDP formulation, MCTS planning, and the search-guided policy learning approach.
summary: This paper addresses model uncertainty in offline MBRL by framing it as a BAMDP, allowing for adaptive belief updates over multiple potential world models. The authors propose a novel continuous MCTS method to solve BAMDPs in continuous and stochastic settings, integrating this planning approach as a policy improvement step in a policy iteration framework. The approach, which combines Bayesian RL with search, is shown to improve performance across a range of continuous control tasks.
weaknesses: * Insufficient discussion of related work in Bayesian RL and offline MBRL: the related work section could better connect with existing research on modeling offline RL as a Bayes Adaptive MDP or epistemic POMDP. Ghosh et al. (2022) introduced a model-free approach to handling model uncertainty by training offline RL policies to adapt to belief changes, while Dorfman et al. (2021) framed offline meta-RL as a Bayesian RL problem, leveraging a belief-augmented MDP. Discussing these connections would position the current work more clearly within the broader landscape of Bayesian treatments of offline RL.

* Lack of connection to prior policy learning with model-based search: using model-based search outcomes for policy learning within an actor-critic framework was introduced by Feinberg et al. (2018) and later applied to offline MBRL in Jeong et al. (2023). Although these methods do not use MCTS for rollouts, they share the concept of leveraging generated trajectories to improve action-value estimation for policy learning. Discussing these earlier works would clarify the paper’s contributions in relation to prior approaches.

* Clarity issues in the experiment section: the structure of the experiment section could be improved to clarify the main research questions and key takeaways. At present, the content lacks focus on specific questions, making it challenging to follow the significance of each result. Organizing the discussion around clearly defined research questions could improve readability and focus.

* Omission of recent offline MBRL baselines: the empirical evaluations rely on older methods (e.g., MOPO and MOReL), making it difficult to assess claims about state-of-the-art performance. Comparisons with more recent baselines, such as RAMBO (Rigter et al., 2022), MAPLE (Chen et al., 2021), and CBOP (Jeong et al., 2023), would strengthen the experimental claims and provide a more comprehensive evaluation of the proposed method’s effectiveness.

__References:__

- Ghosh et al. (2022): "Offline RL policies should be trained to be adaptive."

- Dorfman et al. (2021): "Offline meta reinforcement learning – identifiability challenges and effective data collection strategies."

- Feinberg et al. (2018): "Model-Based Value Expansion for Efficient Model-Free Reinforcement Learning."

- Jeong et al. (2023): "Conservative Bayesian Model-Based Value Expansion for Offline Policy Optimization."

- Rigter et al. (2022): "Rambo-rl: Robust adversarial model-based offline reinforcement learning."

- Chen et al. (2021): "Offline model-based adaptable policy learning."

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See weaknesses above.
soundness: 3
strengths: The paper investigates a promising direction of leveraging the Bayesian RL framework combined with MCTS in offline RL. The paper is well-motivated and the performance on the D4RL benchmark outperforms baselines.
summary: This paper models offline model-based reinforcement learning as a Bayes Adaptive MDP, which enables the agent to leverage the model uncertainty during learning. By using MCTS with double progressive widening, the agent achieves good performance on the D4RL benchmark and three target-tracking tasks.
weaknesses: Experiments:
- A more thorough analysis of the proposed method's performance is necessary. The observed performance improvement appears to result primarily from two factors: (1) the use of BAMDP, and (2) planning with MCTS. To verify this, it would be helpful to see an experiment where only BAMDP is applied, keeping other elements unchanged. While the BA-MBRL variant seems intended to isolate the effects of BAMDP, the paper does not clearly explain how this differs from other models like MOPO or MoReL. For instance, it’s unclear if the performance gap is solely due to BAMDP. Additionally, tracking the evolution of $b(\theta)$ during updates would strengthen the claims about the role of BAMDP.
- A comparison of runtime across methods would be valuable, especially to understand the impact of using BAMDP and MCTS on computational efficiency. Providing runtime metrics could help illustrate any trade-offs between performance gains and computational cost. (The paper only reports the runtime of Sampled EfficientZero.)
- The inclusion of Figure 1 is somewhat ambiguous. The authors seem to suggest that the proposed method outperforms Sampled EfficientZero. However, a more detailed analysis is needed. For example, what factors contribute to the performance difference with BA-MCTS-SL? Is the advantage primarily due to using the Bayesian RL framework?
-  In Line 472, the paper states, "both Sampled EfficientZero and BA-MCTS-SL rely on supervised learning... Thus, purely mimicking the search result may be less sample-efficient than policy gradient methods." However, in Line 406, "BA-MCTS-SL performs similarly to MA-MCTS, validating the effectiveness of both policy update mechanisms." This conclusion seems inconsistent. 

Writing: 
- Notation: what is hars'. It appears many times but I haven't found the definition. For example, in Algorithm1, the inputs of SIMULATE are defined as ($s, h$), $b(\theta)$, $d$ (Line 225), while in Line 231, the inputs of SIMULATE are $(s', hars'), b', d-1$.
- The update of $b(\theta)$ in Eq 4. should be clarifed. 
- The figures should be added via PDF or SVG, otherwise, the figure can be blurry.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: - what is the major purpose of having Fig.1? 
- What is the explanation of the lack of the convergence in Figure 2 for optimized method? Why in the last subfigure all the curves seem to struggle  to converge?
- How to select $\lambda$ in the reward and the number of the ensembles in practice?
soundness: 3
strengths: - Experiments. This work considers both the D4RL MuJoCo and challenging stochastic target tracking tasks. The experiments also include an ablation study where each component of the proposed algorithms is verified, i.e., BA-MBRL, BA-MCTS, BA-MCTS-SL. The results are compared with many SOTA methods such as COMBO, MoReL, MOPO.
- Problem Formulation. I find the adaptive ensemble mechanism is indeed a crucial problem in the ensemble-based method, and the proposed method using the BAMDP as a backbone to address this issue is novel. The algorithms are well-motivated.  I do believe that this paper has good contribution on the offline model based RL community.
summary: Model-based offline reinforcement learning is a promising approach to improve the data efficiency and the generalization capability beyond the training dataset. To address the uncertainty issue when fitting the model with a static dataset, the conventional approach is to use an ensemble of world models, from which the uncertainty can be estimated. This work focuses on the adaptive ensemble mechanism where the belief over each model in the ensemble will be updated when the new samples are experienced by the algorithms. The proposed Bayes Adaptive Monte-Carlo planning algorithm has two key contributions: cast the offline MBRL as Bayes Adaptive MDP (BAMDP), based on which a Monte Carlo tree search method is used for planning. The experiments on D4RL MuJoCo and the three target tracking tasks show that the proposed method has better performance compared with state-of-the-art model-based and model-free offline RL methods.
weaknesses: - The comparison with related work is unclear. In particular, MuZero is a major  comparison object to motivate this study, while the statement in the paper is not clear. In Section 3, the paper claims that MuZero does not apply uncertainty estimation of the learned model, which has no direct relation with using latent models (i.e., estimating uncertainty in the latent space is possible). Please provide a more structured comparison of key related works to highlight the specific gaps.
- Method section is confusing. Section 4.1 titled “ key role of deep ensemble,” while the contents are more about introducing BAMDP and reward design to incorporate the pessimistic-MDP. In Section 4.2, the newly introduced $(s,h)$ as a decision point seems out of nowhere, what is the connection with $(s,b)$ as introduced in BAMDP? The introduction of the algorithm in Section 4.2 is very hard to follow, and a paragraph title can help a lot. For instance, in the last paragraph where introducing the Bayes-optimal policy,  the author should directly introduce the design choice and then compare it with the previous methods. 
- Lack of comparison of the uncertainty evaluation. One of the key motivations of the proposed method is to use the samples to update the belief of ensembles. In the experiments, the related results on the benefits of such adaptation for ensemble-based methods are lacking (e.g., set the adaptive parameter to be zero).

Meta-review: This paper propose to use Bayesian adaptive method for offline RL. Main drawback of the paper include: (1) a lack of theoretical justification of why a Bayesian framework helps with mitigating the out-of-distribution generalization problem in offline RL, and how one should pick the prior in practice. (2) the clarity of the paper can be improved. (3) in the experimental evaluation, more modern and SOTA baseline should be included.

Decision: Reject