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

For Question 1:

As discussed in Section 5.2, MuZero and its variants also utilize MCTS in model-based RL. Specifically, both MuZero and BA-MCTS-SL (ours) rely on supervised learning for policy improvement. Therefore, MuZero and its variants are included as baselines in our comparisons. According to Niu et al. (2023), Sampled EfficientZero (Ye et al., 2021) demonstrates the best performance among MuZero variants on (online) MuJoCo locomotion tasks. Thus, we evaluate Sampled EfficientZero on the set of D4RL MuJoCo tasks, as shown in Figure 3 (which was originally Figure 1). The results indicate that our proposed algorithms outperform Sampled EfficientZero while incurring significantly lower computational costs, as detailed in Appendix E.

For Question 2:

As mentioned in Section 5.3, the tracking tasks in Tokamak control 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. Such a stochastic testing environment poses significant challenges for all algorithms. Notably, while our algorithms and Optimized use the same ensemble of world models, Optimized performs training directly on the ensemble without dynamically adapting the belief, leading to inferior performance due to the discrepancy between the learned and ground-truth world models. To ensure fair comparisons, we set a fixed number of training epochs for all algorithms. $\beta_n$ tracking is relatively easier (since its tracking target has only 1 dimension while the targets in Temperature and Rotation control have 4 dimensions), so we use fewer learning epochs for $\beta_n$ tracking.

For Question 3:

As noted in the first paragraph of Page 9, our implementation is based on Optimized (Lu et al., 2022), with minimal changes to its codebase and hyperparameter settings. Specifically, Optimized utilized Bayesian Optimization to determine key hyperparameters, including $\lambda$ and the size of the ensemble, and we inherited these hyperparameter setups. While further fine-tuning could potentially improve the performance of our algorithms, maintaining minimal changes in implementation ensures that the performance improvements are attributable to our algorithm design. A detailed hyperparameter setup is provided in Appendix C.

In general, either Bayesian Optimization or grid search can be employed to determine these parameters. For example, $\lambda$ can be selected from several values within the range [0, 100], and the size of the ensemble can vary between 5 and 20.

References:

[1] Auger, David, Adrien Couetoux, and Olivier Teytaud. "Continuous upper confidence trees with polynomial exploration–consistency." ECML PKDD, 2013.

[2] Niu, Yazhe, et al. "Lightzero: A unified benchmark for monte carlo tree search in general sequential decision scenarios." Neural Information Processing Systems, 2023.

[3] Ye, Weirui, et al. "Mastering atari games with limited data." Neural Information Processing Systems, 2021.

[4] Lu, Cong, et al. “Revisiting Design Choices in Offline Model Based Reinforcement Learning.” International Conference on Learning Representations, 2022.

For Weakness 1:

In the revised manuscript, we have restructured the related work section (Section 3) by adding paragraph titles. Additionally, we provide a more structured comparison with key related works, including MuZero, Bayesian approaches to offline RL, and algorithms that leverage model-based search to enhance the efficiency of actor-critic training. These modifications are highlighted in blue for ease of reference.

Our algorithm fundamentally differs from MuZero and its variants. (1) MuZero integrates model learning and policy training into a single stage, and the learned world models are defined on a latent state space rather than the real state space. In contrast, our algorithm and the baselines, such as Optimized and MOPO, first learn an ensemble of world models (until their convergence), and then use them as a surrogate simulator for policy learning. (2) MuZero learns a single latent model as the world model and does not involve a mechanism to handle the uncertainty of dynamics/reward predictions. In contrast, our algorithm learns an ensemble of world models, dynamically adapts beliefs for each ensemble member, and constructs a pessimistic MDP to mitigate over-estimation. While estimating uncertainty in the latent space is possible, MuZero does not implement such mechanisms. (3) MuZero does not utilize double progressive widening (DPW, Auger et al. (2013)) or bayes-adaptive planning in its MCTS process, as we do in Algorithm 1. DPW can significantly enhance planning performance in environments with continuous state/action spaces and stochastic dynamics.

For Weakness 2:

As noted in Section 4.1, the practical implementation of a BAMDP relies on the use of deep ensembles, because the ensemble ​​$\{(P_{\theta}^i, R_{\theta}^i)_{i=1}^K\}$ constitute a finite approximation of the space of world models. Also, to construct a pessimistic-MDP, we introduce a reward penalty defined based on the discrepancy in predictions by the deep ensembles, as defined in Eq. (5). Both of these key designs are dependent on deep ensembles, which is why the section is titled “The Key Role of Deep Ensembles”.

As mentioned in Footnote 4, the states in BAMDPs consist of both $s$ and the corresponding belief $b(\theta)$. $b(\theta)$ is a function of, and implied by, the history $h$, since $b(\theta)$ is updated in a Bayesian manner based on the history, as shown in Eq. (4).

Thank you for your valuable feedback. We have refined Section 4.2 by reorganizing the content and adding paragraph titles. Additionally, we have clarified the design motivations and emphasized our novel contributions. Please check the modifications highlighted in blue. In the last paragraph, the introduction of alternative design choices and comparisons between previous methods and ours involve extensive details, so we have to move this content to Appendix B.

For Weakness 3:

In Table 1, the performance gap between BA-MBRL and Optimized is solely attributable to the use of samples to update the belief of ensembles.

As mentioned in the first paragraph of Section 5.1, BA-MBRL follows existing offline MBRL methods such as Optimized, MOPO, and MOReL, by learning an ensemble of world models, applying reward penalties to construct a pessimistic MDP, and using the pessimistic MDP as a surrogate simulator for RL training. Particularly, BA-MBRL and Optimized both use the standard deviation of the predictions across all ensemble members (i.e., Eq. (5)) as the reward penalty design. Thus, the ONLY difference between BA-MBRL and Optimized lies in whether the belief over each ensemble member (i.e., $b(\theta)$) is adapted according to Eq. (4). The benefit of this belief adaptation is clearly demonstrated in Table 1.

In Appendix H of the revised manuscript, we show that employing a Bayes-adaptive ensemble, instead of a uniform ensemble, improves the prediction likelihood for the provided offline trajectories and reduces the prediction errors in imaginary rollouts. As illustrative examples of this performance improvement, Figure 4 tracks belief adaptation during an offline rollout and several imaginary rollouts.

For Weakness 1 & 2:

Thank you for your valuable feedback. We have revised the related work section of our paper, specifically the last paragraph of Section 3, to incorporate research on Bayesian approaches to offline RL and algorithms leveraging model-based search outcomes to enhance the efficiency of actor-critic training. We agree that including these prior works adds clarity to our contributions.

For Weakness 3:

We have improved the clarity of the experiment section (i.e., Section 5). Specifically, we restructured it into three subsections, clearly outlined the six main research questions (in the first paragraph of Section 5), and explicitly linked each result to the corresponding research question to emphasize their significance.

For Weakness 4:

As mentioned in the first paragraph of Page 9, our implementation is based on Optimized (Lu et al. (2022)), with minimal changes to its codebase and hyperparameter settings. In particular, the difference between BA-MBRL and Optimized is whether to adapt the belief over each ensemble member (i.e., $b(\theta)$) according to Eq. (4); the difference between BA-MBRL and BA-MCTS is whether to use MCTS to collect the training samples for downstream SAC; and the difference between BA-MCTS and BA-MCTS-SL is whether to supervised learning to train the policy function. These strictly controlled experiments demonstrate that the performance improvements are driven by the Bayesian RL framework and the deep search component, which is our main experimental objective. 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.

In Appendix I of the revised manuscript, we provide comparisons with RAMBO (Rigter et al., 2022), MAPLE (Chen et al., 2021), and CBOP (Jeong et al., 2023), demonstrating that our algorithm achieves state-of-the-art performance.

For Question 1:

To mitigate error accumulation, we did not do full trajectory rollouts. As mentioned in the first paragraph of Appendix C, at each learning epoch, 50000 states are randomly sampled from the offline dataset, with each state followed by a rollout lasting $H$ time steps. $H$ is no more than 50, according to Table 3. Suppose a state $s_t$ is sampled from the offline dataset. The prior distribution $b_t(\theta)$ at $s_t$ is computed as follows. The prior distribution at the initial state $s_0$ of the trajectory containing $s_t$ is a uniform distribution $b_0(\theta)$. This distribution is then adapted along the offline trajectory, according to Eq. (4), until $b_t(\theta)$ is obtained. Consequently, most belief adaptations are based on the offline trajectories (which are collected from the true environment), with additional adaptations conducted during imaginary rollouts using the learned models. We have updated Algorithm 2 and Section 4.3, with the modifications highlighted in blue, to clarify this issue.

Systematic biases may occur in regions with low data coverage. To address this, we construct a pessimistic MDP to discourage the agent from exploring these regions or overestimating the planning results within them. Also, in Appendix H, we show that employing a Bayes-adaptive ensemble, instead of a uniform ensemble, significantly reduces the prediction errors in imaginary rollouts

For Question 2:

We provided a discussion on the design choices of $\pi_{\text{ret}}$ in the third paragraph of Appendix B. Specifically, $\pi_{\text{ret}}(a|(s, h)) \propto N((s, h), a)^{1/\tau},\ a \in C((s, h))$. Here, $(s, h)$ is a decision node, $C((s, h))$ is the set of sampled actions at $(s, h)$, $N((s, h), a)$ is the number of times $a$ has been sampled at $(s, h)$, $\tau > 0$ is the temperature parameter which decreases with the training process. This design follows MuZero and is practical since it’s defined based on a finite set of actions sampled at $(s, h)$.

For Question 3:

As noted in the first paragraph of Page 9, for fair comparisons, we adopt the same policy architecture as the baselines, i.e., a feedforward neural network, rather than an RNN that incorporates transition history as input. We agree that involving the transition history $h$ as part of the input and adopting an RNN as the policy network could further improve the empirical performance. However, we choose to apply minimal changes to the codebase of existing offline MBRL algorithms like Optimized. This ensures that the performance improvements reported in Table 1 are attributable to dynamically adapting the belief over ensemble members or using MCTS for enhanced planning.

Although $h$ is not directly used as an input to $\pi$, it is utilized to update the belief distribution $b(\theta)$, which influences both the acting and planning phases. The policy $\pi(a|s)$ is trained based on the planning results and can therefore be viewed as a distilled function of $\pi(a|s, h)$. In this sense, the policy implicitly accounts for $h$ through belief adaptations.

In Appendix I, we demonstrate that BA-MCTS outperforms MAPLE (Chen et al., 2021) on the D4RL MuJoCo benchmark.

References:

[1] Lu et al. (2022): "Revisiting design choices in offline model based reinforcement learning." ICLR.

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

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

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

For Weakness 1:

In Table 1, the performance gap between BA-MBRL and Optimized is solely attributable to the use of BAMDP.

As mentioned in the first paragraph of Section 5.1, BA-MBRL follows existing offline MBRL methods such as Optimized, MOPO, and MOReL, by learning an ensemble of world models, applying reward penalties to construct a pessimistic MDP, and using the pessimistic MDP as a surrogate simulator for RL training. Particularly, BA-MBRL and Optimized both use the standard deviation of the predictions across all ensemble members (i.e., Eq. (5)) as the reward penalty design. Thus, the ONLY difference between BA-MBRL and Optimized lies in whether the belief over each ensemble member (i.e., $b(\theta)$) is adapted according to Eq. (4). The benefit of this belief adaptation is clearly demonstrated in Table 1.

In Appendix H of the revised manuscript, we show that employing a Bayes-adaptive ensemble, instead of a uniform ensemble, improves the prediction likelihood for the provided offline trajectories and reduces the prediction errors in imaginary rollouts. As illustrative examples of this performance improvement, Figure 4 tracks belief adaptations during an offline rollout and several imaginary rollouts.

For Weakness 2:

The use of MCTS provides both performance gains (as shown in Table 1 and Figure 2) and additional computational cost. However, this extra computation is incurred only during the offline training stage. To ensure real-time execution, MCTS is not used during deployment. Furthermore, leveraging parallel computation frameworks can significantly accelerate planning with MCTS (as in AlphaZero) and reduce the offline training time.

Thank you for this suggestion. In Appendix E of the revised manuscript, we provide a detailed comparison of the computational costs of our proposed algorithms and Sampled EfficientZero.

For Weakness 3:

Sampled EfficientZero and BA-MCTS-SL both adopt MCTS and supervised-learning-based policy improvement. The performance difference between them could be attributed to several factors: (1) Sampled EfficientZero integrates model learning and policy training into a single stage, which significantly increases the learning difficulty (compared to BA-MCTS-SL). (2) Sampled EfficientZero relies on a single latent model as the world model and lacks a mechanism to address uncertainty in dynamics and reward predictions. In contrast, our algorithm learns an ensemble of world models, dynamically adapts beliefs for each ensemble member, and constructs a pessimistic MDP to mitigate over-estimation. (3) Sampled EfficientZero does not incorporate double progressive widening (Auger et al., 2013) in its planning process, which could enhance planning performance in environments with continuous state/action spaces and stochastic dynamics.

For Weakness 4:

As mentioned in Line 422, BA-MCTS-SL struggles on Walker2d, where a warm-up training phase (using BA-MBRL) is required to establish a better initial policy. Specifically, for Walker2d tasks, during the first 100 out of 500 training steps, we use BA-MBRL to warm up the policy before switching to BA-MCTS-SL for further improvement. Without this warm-up phase, the policy would fail. This observation supports the statement in Line 472: “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.”

For Weakness 5:

Sorry for the confusion. In $(s, h)$, $h$ denotes the transition history ending at state $s$, which is a sequence of states, actions, and rewards. After taking an action $a$ at $s$, the agent receives a reward $r$ and transitions to the next state $s’$. The history is then updated to $h’=hars’$. This clarification has been added to the fourth paragraph of Section 4.2 in the revised manuscript.

For Weakness 6:

The update from $b(\theta)$ to $b’(\theta)$ in Eq. (4) follows the Bayesian rule: $b’(\theta)(i) \propto b(\theta)(i) P_{\theta}^{i}(s' | s, a) R_\theta^i(r | s, a)$, where $i=1, \cdots, K$ and $K$ is the size of the ensemble. $b’(\theta)$ and $b(\theta)$ represent the posterior and prior distributions, respectively, while $P_\theta^i(s' | s, a) R_\theta^i(r | s, a)$ is the likelihood of the transition $(s, a, r, s’)$ under ensemble member $i$ (i.e., ($P_\theta^i, R_\theta^i$)). We have clarified this in the second paragraph of Section 4.1.

For Weakness 7:

We are including the figures as EPS files and are happy to provide clearer versions of the images for the camera-ready version.

References:

[1] Auger, David, Adrien Couetoux, and Olivier Teytaud. "Continuous upper confidence trees with polynomial exploration–consistency." ECML PKDD, 2013.

Thank you for your detailed feedback.

• The warm-up stage was applied only to the Walker2d domain.

• Yes, $h$ represents a memory state, updated as $h' = f(h, a, r, s')$, where $f$ denotes a concatenation/RNN operator. We will clarify this in the revised version.

• Figure 3 consists of 12 subfigures, each rendered as an EPS file. We will refine this figure in the upcoming revision.

For Question 1:

The purpose of search/planning is to simulate the acting process by looking ahead. Therefore, the search and acting processes should be based on the same transition dynamics, that is, utilizing the same ensemble and the same Bayesian adaptive process.

As discussed in Section 4.2, the search algorithm is designed to address a planning problem within a BAMDP (defined in Section 2). According to Eqs. (1) and (4), state and reward transitions should be based on predictions from all ensemble members, with the belief updated at each time step based on the transition outcomes.

For Question 2:

Thank you for the question. We need to adjust our analysis of the performance difference between Sampled EfficientZero and BA-MCTS-SL as follows:

Sampled EfficientZero and BA-MCTS-SL both adopt MCTS and supervised-learning-based policy improvement. The performance difference between them could be attributed to several factors: (1) Sampled EfficientZero integrates model learning and policy training into a single stage, which significantly increases the learning difficulty (compared to BA-MCTS-SL). (2) Sampled EfficientZero relies on a single latent model as the world model and lacks a mechanism to address uncertainty in dynamics and reward predictions. In contrast, our algorithm learns an ensemble of world models, dynamically adapts beliefs for each ensemble member, and constructs a pessimistic MDP to mitigate over-estimation. (3) Sampled EfficientZero does not incorporate double progressive widening (Auger et al., 2013) in its planning process, which could enhance planning performance in environments with continuous state/action spaces and stochastic dynamics.

For Question 3:

As mentioned in Appendix C, at each learning epoch, $50000H$ transitions are sampled by interacting with the learned models, 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. Thus, the training is based on multiple imaginary trajectory segments rather than full trajectories. (Modifications have been made in Section 4.3 for clarification.) To mitigate overestimation caused by model inaccuracies, a reward penalty based on the discrepancy among ensemble members is applied, as defined in Eq. (5). This training protocol is adapted from Optimized and MOPO, which serve as our baselines. We follow the same process to ensure fairness in the evaluation.

Offline variants of MuZero (Schrittwieser et al., 2021) train the policy network exclusively on states from the offline dataset and the planning results derived from those states. This training approach can be somewhat restrictive, particularly when the offline dataset fails to provide adequate coverage of the continuous state space.

References:

[1] Auger, David, Adrien Couetoux, and Olivier Teytaud. "Continuous upper confidence trees with polynomial exploration–consistency." ECML PKDD, 2013.

[2] Schrittwieser, Julian, et al. "Online and offline reinforcement learning by planning with a learned model." Neural Information Processing Systems, 2021.

Regarding Weakness 1:

To the best of our knowledge, there is no other offline model-based RL algorithm using BAMDP.

Also, we’d like to point out that the improvements stem from both the use of BAMDP and the incorporation of MCTS. In Table 1, the comparison between BA-MBRL and Optimized highlights the benefit of using BAMDP, as it is the only difference in their algorithm designs. Further, compared to BA-MBRL, BA-MCTS introduces a planning process with MCTS, which results in significant performance improvements, as shown in Table 1.

Regarding Weakness 2:

In the offline learning setting, we don’t have access to the true world model, so we can only generate rollouts based on learned world models. To mitigate error accumulation, we did not do full trajectory rollouts. As mentioned in the first paragraph of Appendix C, at each learning epoch, 50000 states are randomly sampled from the offline dataset, with each state followed by a rollout lasting $H$ time steps. $H$ is no more than 50, according to Table 3. Suppose a state $s_t$ is sampled from the offline dataset. The prior distribution $b_t(\theta)$ at $s_t$ is computed as follows. The prior distribution at the initial state $s_0$ of the trajectory containing $s_t$ is a uniform distribution $b_0(\theta)$. This distribution is then adapted along the offline trajectory, according to Eq. (4), until $b_t(\theta)$ is obtained. Consequently, most belief adaptations are based on the offline trajectories (which are collected from the true environment), with additional adaptations conducted during imaginary rollouts using the learned models. We have updated Algorithm 2 and Section 4.3, with the modifications highlighted in blue, to clarify this issue.

For your example, we think it’s less likely that the belief would mistakenly concentrate on a wrong model. According to Section 4.1, the belief starts with a uniform distribution and its update is based on either the offline trajectories or imaginary rollouts generated by the predictions of all ensemble members. Thus, for a belief to concentrate on a wrong model, it needs more ensemble members to agree on a wrong prediction than a correct prediction, which is less likely to happen in regions with good data coverage. While, for predictions in out-of-sample regions, it could be problematic but, as mentioned in the last paragraph of Section 4.1, we construct a pessimistic MDP to discourage the agent from exploring those regions or overestimating the planning results there.