Bayes Adaptive Monte Carlo Tree Search for Offline Model-based Reinforcement Learning
Authors: Jiayu Chen
Le Xu
Wentse Chen
Jeff Schneider
Venue: ICLR 2026 Poster (2026)
Keywords: Offline Reinforcement Learning
Monte-Carlo Tree Search
Primary area: reinforcement learning

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 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 propose 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 offline RL methods on twelve D4RL MuJoCo tasks and three challenging, stochastic tokamak control tasks.

BAYES ADAPTIVE MONTE CARLO TREE SEARCH FOR OFFLINE MODEL-BASED REINFORCEMENT LEARNING
Jiayu Chen
Le Xu Tsinghua University Beijing, China
Wentse Chen & Jeff Schneider Carnegie Mellon University Pitsburgh, PA 15213, USA
The University of Hong Kong INFIFORCE Intelligent Technology Hong Kong SAR j ilayuc@hku.o .k
ABSTRACT
Offline reinforcement learning (RL) is 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 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 dealing with the uncertainty about the true MDP can be challenging. In this paper, we propose modeling offline MBRL as a Bayes Adaptive Markov De- cision Process (BAMDP), which is a principled framework for addressing model uncertainty. We further propose a novel Bayes Adaptive Monte Carto 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 policy iteration. Our "RL. + Search" framework follows in the footsteps of su- perhuman AIs like AlphaZero, improving on current offline MBRL methods by incorporating more computation input. The proposed algorithm significantly out- performs state-of oottth---tt-tttt offline RL methods on twelve D4RL MuJoCo tasks and three challenging, stochastic tokamak control tasks.
INTRODUCTION
The success of RL typically relies on large amounts of interactions with the environment. However, in world scenarios, such interactions can be unsafe or costly. As an alternative, offline RL.
(Levine et al.] 2020) leverages offline datasets of transitions, collected by behavior policy, to train a policy. To avoid overestimation of the expected return for out-of-distribution states, which can mislead policy learning, model-free offline RL methods (Kumar et al., 2020 Wu et al., 2019) often constrain the learned policy to remain close to the behavior policy. However, acquiring large volume of demonstrations from 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 al. (2022). These methods train dynamics models from offline data and optimize policies using imaginary rollouts simulated by the models. A key advantage of this approach is that the dynamics modeling process is decoupled from the behavior policy's 's objectives. This makes it possible, in principle, to learn effective policies even from sub-optima or exploratory data, so long as the dataset provides adequate coverage of the relevant state-action space.
Multiple MDPs can exhibit identical behaviors on an offline dataset of states and actions, yet their underlying dynamics and reward functions may differ, particularly on out-of-sample regions. This implies that we are dealing with a distribution of possible world models underlying 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 would be uniformly sampled from the ensemble to 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 based on the experience. The Bayes Adaptive Markov Decision Process (BAMDP)
Duff 2002) provides principled framework for modeling such adaptations, and BAMCP (Guez

---

et al. 2013) is an efficient online planning method for solving BAMDPs. However, BAMCP has several limitations: (1) it relies on a ground-trutt world model for plamning; (2) it is restricted to discrete state and action spaces; and (3) its outcome is an action choice particular state, rather than policy function. To address these challenges: (1) we learn an ensemble of world models from the offline dataset and construct a pessimistic BAMDP as an estimation of the true MDP; (2) propose a novel planning algorithm to solve BAMDPs in continuous state and action spaces based on double progressive widening (Auger et al., 2013); and (3) we distill the planning outcomes into a policy function via RL, enabling real-time execution. Notably, we provide theoretical proof establishing the consistency of our planner in continuous, Bayes-adaptive MDP settings.
Our main contributions include: (1) Firstly using BAMDPs to handle model uncertainties in of- fline MBRL: (2) Proposing theoretically-grone grounded Continuous BAMCP for planning in continuous, stochastic BAMDPs; (3) Developing the first algorithm to successfully integrate Bayesian RI., of- fline MBRL, and deep search for data-driven continuous control; (4) Demonstrating the state-of-the- art performance twelve D4RL MuJoCo tasks and three target tracking tasks in tokamak control (for nuclear fusion).
BACKGROUND
An MDP (Puterman] 2014) can be described as a tuple M = (S, and A are the state space and action space, respectively. P S x A + As is the dynamies function and R. S x A - 1o. is the reward function, where Ax denotes the set of probability distributions on X. y E [0, 1) is a discount factor A Bayes Adaptive MDP (BAMDP) (Duff) 2002) can model scenarios where the precise MDP Ma = (S,A, Po. Ro,7) 7) is uncertain but is known to follow a prior distribution (0). During planning. Bayes-optimal agent would update its belief over the MDP based on its experi- ence. Formally, a BAMDP can be described as a tuple M = A, , 7) S+ denotes the space of information states (s, b), which is composition of the physical state and the current belief over the MDP. After each transition (s, a, r, s'), the belief is updated to the corresponding Bayesian posterior: (0) x ((0)P((()))))) a, s') s') (( R. (r(s,a), Accordingly, P and R + are defined as follows:
-1) / Po(o'lo, ())0)00
B+ ((*,
=1b==1))
R ((s.b), = (o((sa))((0))
The Q-function that satisfies the Bellman optimality equations (i.e., Eq. (2) is the Bayes-optimal Q- function and n (s,s, b) = arg max. Q' ((x, b), a) is the Bayes-optimal policy. Actions derived from nt are executed in the true MDP and constitute the best course of actions for Bayesian agent with a prior belief by over the underlying MDP (Guez et al. 2014). Bayesian RL (Ghavamzadeh et al. 2015) seeks to compute the optimal policy, TT, and provides principled framework for handling uncertainty in the world model Mg. However, the requirement for Bayesian posterior updates (i.c., Eq. ( often renders it computationally intractable, and discuss several approximate methods in Section 3
Q' a) = R (s, a) +7 [ V* (('la) a)dz'
(2)
= = max (" a), Vr (s, st, aEA
3 RELATED WORKS
Offline model-based RL: Offline RL (Chen et al. 2024) enables an agent to learn control poli- cies from datasets of environment transitions pre-collected by a behavior policy μ4, i.c., D,, = [[(s, at, 1 Offline Model-based RL (MBRL) methods explicitly learn world models Mg from D,, and adopt Mg surrogate simulator, enabling the learned policy to possibly generalize
to states beyond D,,. Both planning methods (Argenson & Dulac-Arnold 2021: Zhan et al. 2022, Dichl et al. 2023) and RL methods (Yu et al.) 2020; Kıdambi et al. 2020 Lu et al. (022) can use the learned Ma to obtain a policy. However, since D,, may not span the entire state-action space, Ma is unlikely to be globally accurate. Leaming/Planning without any safeguards against such
2

---

model inaceuracy can yield poor results. Yu et al. (2020); Kidambi et al. (2020): Lu al. (2022)
propose learning an ensemble of world models, using ensemble- based uncertainty estimations to construct a pessimistic MDP (P-MDP). and learning near-optimal policy atop it. Ideally, for any policy, the performance in the real environment is lower-bounded by the performance in the corre- sponding P-MDP (with high probability). thus avoiding being overly optimistic about an inaccurate model. Notably, none of these offline MBRL methods have modeled the problem as BAMDP, although Bayesian RL provides a principled framework for handling model uncertainty.
Bayes-adaptive planning: Approximate methods for solving BAMDPs, such as AAmmth et al. 2209 Sorg et al. (2010); Castro & Precup 20100;; & Littman (201T); Wang et al. (2012); Guez et al. (2013): Slade et al.J (2020) have been developed. As a representative work, BAMCP
(uuez et al. 2013) adopts Monte-Carlo Tree Search (MCTS) (Browne et al.] 2012) for Bayes-
adaptive planning and is shown to converge in probability to near Bayes-optimal policy at the root node of the search tree. However, all these methods cannot be applied to large-scale MDPs with continuous state and action spaces.
Also, these planning algorithms are not designed for offline MBRL. How to incorporate planning outcomes for policy improvement in RL and how to handle the inaccuracy of the learned world model during planning still require exploration.
METHODOLOGY
In this section, we present a novel offline MBRL algorithm based on Bayes Adaptive MCTS. In Section 4.2] we propose a Bayes Adaptive planning method (i.e., Continuous BAMCP) that can be applied to stochastic continuous control tasks. Then, in Section iteration framework, where the search results (from Continuous BAMCP) are distilled into policy and value networks for policy improvementdevaluutionn at each iteration. In this way, integrate offline MBRL with Bayes Adaptive MCTS. Both components require the use of an ensemble of world models for either practical implementation or uncertainty quantification, as detailed in Section
present search-based policy
4.1 THE KEY ROLE OF DEEP ENSEMBLES
Offline MBRL methods estimate world models Mo from static dataset Dr.- However, there could be various MDPs that behave identically on the limited set of states and actions in D,,, but their dynamics and reward functions may differ, especially on out-of-sample states and actions. Thus, are actually dealing with distribution of world models that follow prior distribution (0) A P(Me) (p) ). As introduced in Section 2, Bayesian RL handles such model uncertainty by explicitly including the belief over the models in its state representation. The belief updated with experience, providing measure of how the models' uncertainty has changed since the beginning of the episode.
The idea of deep ensembles (Lakshminarayanan et al., 2017) is to train multiple deep neural net- works as approximations of a function (for uncertainty quantification), each using different weight initialization and optimized with different mini-batch sequence. For offline MBRL, we can learn an ensemble of dynamics models (P).- PK } and reward models (R...... RK J from the dataset D, by minimizing the following supervised learning loss: (i
I,...KK)
L(P)) "(,,,,,') [log ()
((R)) = -E [log R2 (r(s,
{(P), Rb) } can be viewed as a set of independent and identically distributed (IID) samples from the prior P(Ma (D,)) and constitute 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 (P) R.).. In this case, reasonable prior distribution is (0) [1/K,... .... 1/ K). since these models are IID prior samples. After receiving a transition (s, a, r, s'), the belief can be updated as follows:
(4) This is a practical implementation of the Bayesian posterior (originally defined in Eq. D update based on deep ensembles, where b(0), (0). and (s (s. a)R3 (r|s, a) denote the prior, posterior distributions, and likelihood, respectively.
(0) & (1))
()), a), i i= 1....... K
3

---

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 static offline dataset. A typical solution for such uncer- tainty is to construct a pessimistic MDP, which discourages the agent from exploiting regions with high uncertainty. Another use of the deep ensemble is to construct Pessimistic Bayes-Adaptive MDP (P- (P-)..DP).
In particular, we apply reward penalty each transition (s),,) and get penalized reward f as follows:
X std (²t +
1Q (s*
(5)
Here, 8 ~ P a) and r² [- a)] denote the next state and expected reward from the i-th ensemble member. π and Q+ represent the policy and target Q-network. This penalty term the standard deviation of the one-step lookahead Q-value targets, computed across all ensemble members. The underlying intuition is that a high variation in these predictions indicates a high degree of uncertainty at this region. As demonstrated by Sun et al. | (023), employing this penalty to define pessimistic MDP yields an agent performance estimate that more closely matches real-world outcomes, compared to penalties based solely on ensemble disagreement in next-state prediction, as proposed in Yu et al. (2020); Kidambi et al. (2020).
4.2 CONTINUOUS BAMCP
BAMCP (uuzz et al. 2013) 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 a novel planning method to approximate the Bayes-optimal policy at a decision point (s, h) (h denotes the transition history that ends at s), which can be used to solve BAMDPs with continuous states/actions and stochastic transition kernels.
Double Progressive Widening (DPW): DPW (Couêtoux et al. 2011 Auger et al. 2013) is tech- nique to extend the use of MCTS to continuous state and action spaces. Instead of exploring all possible actions and next states, DPW maintains a finite list of options to search each decision node, incrementally adding new options to the list based on the visitation counts of that decision node. To be specific, for node (s, h)), new action a is sampled (with the current policy т) and added to its children set C(ss, f)), if [N((s, h)) 2 [C[(s, fi)], where a E (0, 1) is hyperparam- eter that controls the growth rate and N denotes the visitation counts of (s. h). Otherwise, an action is selected from existing options in C((s, h)) according to the UCT (Kocsis & Szepesvári 2006) rule. Similarly, to handle the infinitely many possible state transitions in stochastic environments, new next state s' is added to the children set C((ss, h).a) only if [N((s,)) a) 2 [C((s, a)l (3 E (0, 1)). Otherwise, the least visited child in C(ss, ft), a) will be selected as the next state. With DPW, the sets of possible actions or next states explore are finite, allowing deep tree search as in discrete scenarios.
Integration of DPW and BAMCP: Directly combining DPW and BAMCP (i.e., Algorithm 3) can- not solve BAMDPs with continuous state and action spaces. As introduced in Appendix A BAMCP relies on root sampling, which samples dynamics functions only at the root node and follows spe- cific dynamics function throughout a simulation rollout. However, the rationale of root sampling (i.e., Lemma A.1 does not hold when applying DPW As an alternative design, Polynomial Upper Confidence Tree (PUCT) (Auger et al. 2013). built upon DPW, is a provably consistent planning method for solving MDPs with infinite-scale state and action spaces and highly stochastic transition dynamics. Thus, propose recast P-BAMDPs as MDPs and solve them using PUCT, as novel Bayes planning method for continuous state and action spaces.
Proposed algorithm- Continuous BAMCP: the root node to an unvisited node, utilizing progressive widening when sampling actions or next states, as detailed in the ACTIONPW and STATEPW procedures. Compared PUCT, the main modifications include: (1) replacing (8, P, R) in MDPs with their extended definitions in BAMDPs,
In Algorithm 1 each simulation follows path from
'The last equality of Eq. 6 does not hold, since b(0)has') b(0) (s')s,)) * Po a) represents the distribution of next states when applying DPW, which differs from the true distribu- tion Po (s(' , a), dictated by the DPW rule.
(s)
4

---

Algorithm 1 Continuous BAMCP
Input: n, V, E, dmax. 7, a, B, Plik, Ri.K procedure SEARCH((s, h), b(0)
procedure ACTIONPW h))
if first visit then C((s.h) + 0
for e 1... do SIMULATE((S, (a,TT(((s,,))))) h), b(0), dmax)
if [[(((,))) C((s))) then
nt (s,,))))) C((s, h)) t C((s, (a) N((s,)), a), Q((s 0,0
end for (a, h)) & N((s,)).a). E h((s,))
else
= EeEC((HA)) o((a.h).) (((aah))
Q((s,
+ a-arg -aagmaX (((s,)),))
end if return end procedure procedure STATEPW ((s, (()),
return Tra, Pret end procedure
procedure SIMULATE((S. 1t), b(O), d)
if d = then return (((x, h))
if first visit then ((s,,)), <- if h), 2 (C((s))).a) then K
ACTIONPW((s. h))
r, s.s (0) + STATEPW ((,,)) b()))()).)) (((s h)) += 1, ((s, h),a) +=1
Ei ((0
if ((((s)) h)) > 1 then
s' ~ > b(0) ((PP(()))) Update b(0) V (0) using Eq.
SIMULATE(($* b(r))),
else
C(ssh), U
((ss')aarss))
(6r, s (0))
end if
[(ss' harrs))) + return s' (0)
Access or calculate i using Eq. 5
R - F + RR, cache F Q((s,)), a) += R- (a))()
end if
return the least visited node in C((s, h) end procedure
N((a.)),a)
return end procedure
ie., (S+ ,+, and (2) applying reward penalties to account for model uncertainty. To be spe- cific, in STATEPW, r and are sampled from the distribution predicted by all ensemble members, which is practical implementation of sampling from R+ and P as outlined in Eq. (1. After receiving the transition (s, a, r, s'), the belief vector b(0) is updated to (()0) following Eq. fin- ishing the transition in S+ from (s, b(0)) to (0)). The transition history h is then updated to h'' = hars'. Secondly, in SIMULATE, a penalized reward i (i.e., Eq. (5)) used, which effectively mitigates overfitting to inaccurate ensemble members by incorporating pessimism, discussed in Section 4.1
R*),
The following theorem establishes that our planner is reliable and consistent solver for P-BAMDPs defined in Section 4.1 The detailed proof is provided in Appendix B This consistency guarantee is crucial, it ensures that with sufficient computation, our planner's output converges the true optimal value within the pessimistic model, providing a reliable value target for our overall learning framework.
Theorem 4.1. Let (zo) be the value of the root node 20 estimated by Continuous BAMCP after E simulations. Under the regularity and value function conditions defined in Appendix B.] the error of this estimate with respect the P-BAMDP's optimal value Vi (20) is bounded as follows, exponentially surely in E:
(2 CoE-T0 + -.mmoo
+ EV,
where "oo Co > O are constants, is the discount factor, dmaa is the maximum search depth, and cy is the uniform bound on the value approximation error at the leaf nodes.
Algorithm I used to approximate the Bayes-optimal policy (s, h)), which is (s, h)) x [[ h), aj, E C(s, h) (Auger et al. 2013). However, we aim to solve the entire P-BAMDP
F
"The value consistency guaranteed by Theorem 4.1 implies that the returned policy was is near-optimal for the P-BAMDP. Since our planner derives Tua through pessimistic planning process, the pessimistic planning theory (Ii et 2020 Sun al. 2023) ensures that near-optimal policy in the P-BAMDP has bounded
5

---

Algorithm 2 BA-MCTS
Input: T. E1, Plikk R
procedure ACTOR
Initialize the actor π and critic Qu. D +- 0 procedure LEARNER
while true do
Sample s from P1. +- 8, +
c
Obtain the prior b(0) at
while true do {(s, h, Tee, F, s', h1), h));
fo H do
i= ~ D
Teet Vret <- SEARCH( (s. h). b(0))
Update π and Qu using SAC
a ~ Teer (-ss,-((s,11)))
++=1 1
Acquire s', (0) as in STATEPW
Update т and (defined on Q+) in
Calculate using Eq. (
ACTOR if
Append with ((s, h). Teets Vaa))
e%F= ==
8, 5,h,b(( b(0) <- 8', hars', (()) end for D Du end while end procedure
end while end procedure
offline, climinating the need for anything beyond simple inference using the policy network during deployment. This necessitates well-learned policy function each decision node, but we cannot execute Algorithm [ at every hs due to the scale of the state space. Therefore, we integrate the planning algorithm into a policy iteration framework as introduced in the next subsection. In this case, π and V in Algorithm denote denitt the policy and value functions from the previous learning iteration while Tet and Uret are the improved policy and value estimates for specific decision nodes. As additional details, multiple terms (labeled in blue) in Algorithm I Thave alternative designs across different literatures, which we elaborate on in Appendix C
4.3 OVERALL FRAMEWORK: BA-MCTS
Now, we present how integrate planning outcomes of continuous BAMCP into a policy iteration framework, to obtain a near Bayes-optimal policy, i.e., that can be directly referred to during execution. The pseudo code is shown as Algorithm [ For efficiency, a learner and a number of actors execute in parallel, reading from and sending data to the replay buffer D respectively. The actors update their copies of policy and value functions every Ej learner steps.
Each actor interacts with the learned world models to sample trajectory segments The starting states of these segments are sampled from the provided dataset D,.. Notably: (1) the segment length H is kept short minimize error accumulation when interacting with the learned world models; and (2) the prior b(0) at a starting state s is obtained by performing the Bayesian posterior update (i.e., Eq. (4) on the offline trajectory in DD, containing These belief updates are reliable, as the offline trajectories are from the real environment. At each decision node (s, h) of the segment, a SEARCH procedure (defined in Algorithm [ is executed. The search result Teet then indicates the action choice, i.c., ~ [ser(-((s, (s,)ss))), h)). Given a, the transition to the next state follows a BAMDP, where r R ()), ~ ~P*(->(s, (-((ss,, h), and the belief b(8) is adapted with the new transition.
With the collected segments we update the policy π and value function V. Our primary algorithm, BA-MCTS, employs an off-policy RL method: SAC (Haarnoja et al., (2018). The value function (critic) is updated using standard temporal difference learning based on the sam- pled transitions, while the policy (actor) is updated using the policy gradient derived from the critic. Additionally, we investigate an alternative approach, BA-MCTS-SL. where the policy is updated via supervised learning to mimic the search result गet by minimizing a cross-entropy loss: C(m, [(s, h, ) )BB
(((,))))ma)) THe log (-1(s,)))/ B, We compare these two approaches
sub- sub-optimality in the non- -on-eessiiiitic BAMDP. Finally, according to the foundational theory of HAMDPs/Gue al. 2013), a near-optimal policy for the RAMDP is, by definition, near-Bayes-optimal policy for the real environment.
at and V are functions of (s, h) because the states in BAMDPs include both and the corresponding helief b(0). b(0) is recursively updated using the history h and is function of h.

---

Table 1: Comparison of the prediction errors for next states and rewards in imaginary rollouts, with and without Bayesian belief adaptation. Ground truth for the imaginary rollouts is obtained by replaying the action sequences in the real simulators.
Prediction Lrvor on Next State
Prediction Error on Reward
Overall Prediction Error
Data Type
Environment
Adaptive
Uniform
Adaptive
Uniform
Adaptive
Uniform
random
HaliCheetah
0.339 (.021)
0.342 (.017)
0.016 (,001)
0,016 (.001)
0.177 (.011)
0.179 (.009)
random
Hopper
0,232 (.016)
0.189 (.008)
0.012 (.00)
0.022 (.008)
0.122 (.008)
0.106 (.008)
random
Walker2d
62.10 (11.0)
112.0 (19.0)
1.364 (.215)
4.908 (.708)
31.73 (5.61)
58.44 (9.83)
medium
HaltCheetah
1.387 (.040)
1,355 (.038)
0.057 (.001)
0.061 ((001)
0.722 (.021)
0.708 (.019)
medium
Hopper
0.429 (.033)
0.570 (.035)
19.03 (8.58)
68.24 (11.4)
9.727 (4.30)
34.40 (5.69)
medium
Walker2d
34.01 (.556)
34,26 (.142)
24.38 (4.31)
113.1 (3:47)
29.19 (1.89)
73.69 (1.78)
med-replay
HalfCheztah
0.677 (.027)
0.707 (.036)
0.115 (.005)
0.114 (.006)
0.396 (.016)
0.410 (.021)
med-replay
Hopper
0.170 (.022)
0.212 (,054)
1.177 (420)
3.320 (1.14)
0.674 (,221)
1.766 (.558)
med-replay
Walker2d
66.83 (7.39)
44.61 (1.34)
27.21(3.24)
60.52 (3.17)
47.02 (5.17)
52.57 (2.25)
med-cxpert
HaffCheztah
1,233 (.054)
LA67 ((046)
0.071 (002)
0.074 (002)
0.747 (.028)
0.771 (.024)
med-expert
Hopper
3.665 (2.56)
18.27 (2.09)
71.46 (79.9)
384.7 (37.6)
37.56 (41.2)
201.5 (19.4)
med-exxpet
Walker2d
55.69 (3.71)
51.89 (.481)
121.9 (12.6)
186.3 (3.86)
88.77 (8.18)
119.1 (2.08)
in our evaluation. In both methods, we seek to distill the superior, non-parametric search outcomes from Continuous BAMCP into the parameterized policy x.
5 EVALUATION
Section discusses the benefit of Bayesian model adaptation. Section presents prreesets comprehensive benchmarking results on D4RL MuJoCo tasks. Section 5.3 details our ablation studies. Finally, Section 5.4 validates our algorithm's applicability on a challenging Tokamak control task.
5.1 THE BENEFIT OF BELIEF ADAPTATION
The core idea of belief adaptation (using the BAMDP framework detailed in Section 4.1 is that by dynamically adjusting the belief over each ensemble member based on observed transitions, the agent can rely more on models that are accurate for the current trajectory, Icading to better planning. Our empirical results show that belief adaptation not only improves model prediction accuracy, as shown in Table I] but also translates to better final policy performance on offline RL tasks, as shown in Tables 2 and 4 Specifically, Table [ highlights that the adaptive ensemble achieves more accurate predictions for next states and rewards in 10 out of 12 benchmarking tasks, compared to treating each ensemble member uniformly as in previous offline MBRL methods. This suggests that Bayesian belief adaptation leads to more accurate world model for planning, which is a key factor underlying observed performance improvements. A more detailed study of this aspect, including visualization of the belief adaptation process within rollouts, is provided in Appendix F.
5.2 BENCHMARKING RESULTS ON D4RL MuJoCo
We evaluate our algorithm on a widely-used offline RL benchmark: D4RL MuJoCo (Fu et al. 2020), comparing its performance with state-of-the-art (SOTA) model-bassed/modll-fee methods, and search-based approaches.
Model-Based Methods. In Table 2 BA-MCTS achieves the highest average score across all twelve D4RL MuJoCo tasks when compared against a range of SOTA model-based basclines. This result confirms that the combination of principled Bayesian treatment of uncertainty and a powerful search-based planner leads to superior performance.
Model-Free Methods. We also compare our algorithms with series of model-free offline pol- icy learning methods, with detailed results presented in Appendix F.2 Our algorithms show significantly better performance, demonstrating the necessity of model-based learning in these en- vironments, particularly when data quality is low. Notably, APE-V (see Table 2 is model-free method built on BAMDPs, vet it shows inferior performance compared ours in 10 out of 12 tasks. As a summary, in Table EH we show comparisons against 19 strong model-based and model free baselines, where BA-MCTS achieves the highest average score. (See Appendix [3 for details of the selected baselines.) To quantify the significance of this improvement, we compute Cohen's 's d
'MOReL., IQI., and the model-free offline RL baselines in Appendix F2 did not report standard deviations. Additionally, IQL. and DT were evaluated D4RL tasks with random data.
7

---

Table 2: Comparisons of our algorithm against SOTA offline RL methods the D4RL benchmark. Each value represents the normalized score, as proposed in Fu et al. (2020), of the policy trained by the algorithm. For our algorithm, we report the average score of the final ten learning epochs and its standard deviation across three random seeds. To ensure fair comparison, the results for MOBILE and CBOP are from our own runs using their publicly codebases. The results for all other baselines
are taken from their respective original papers. Please check Appendix F
for the references.
Data Type
Environment
BA-MCTS (ours)
MOBILE
CBOP
RAMBO
APE-V
MAPLE
Optimized
random
HaliCheetah
(1.44)
39.27 (1.38)
32.95 (0.14)
40.0
29.9
41.5
31.7
random
Hopper
31.42 (0.10)
7.965 (0.43)
23.79 (11.2)
21.6
31.3
10.7
12.1
random
Walker2d
21.54 (0.06)
21.47 (0.03)
1,391 (0.06)
11.5
15.5
22.1
21.7
medium
HaliCheetah
77.31 (1.40)
76.42 (1.00)
69.98 (0.88)
77.6
69.1
48.5
45.7
medium
Hopper
103,9 (0.33)
102.5 (1.77)
98.68 (2.30)
92.8
44.1
69.3
medium
Walker2d
(0.78)
(2.79)
93.92 (1.92)
86.9
90.3
81.3
79.7
med-replay
HaliCheetah
71.24 (2.65)
64.67 (8.14)
63.36 (0.47)
68,9
64.6
69.5
58.0
med-replay
Hopper
106.4 (0.53)
105.7 (1.56)
101.3 (1.74)
96.6
98.5
85.0
90.8
med-replay
Walker2d
91.42 (1.54)
85.12 (10.1)
74.56 (1.29)
85.0
82.9
75.4
65.8
med-expert
HallCheetah
(2.27)
(4.00)
97.91 (5.26)
93.7
101.4
55.4
104.2
med-expert
Hopper
111.8 (0.82)
111.9 (0.34)
111.2 (0.77)
83.3
105.7
95.3
105.8
med-cxpert
Walker2d
116.0
115.0 (0.30)
117.6 (0.27)
68.3
110.0
107.0
97.1
Average Score
80,25
76.50
73.87
68.85
72.65
61.32
65.16
Cohen 1988) against all baselines for which standard deviations are available. The results show that the Cohen's is d value is greater than 1.8 in all cases. According to Cohen's Rule of Thumb, a value of d > 1 indicates a very large (statistically significant) difference. This comprehensive comparison demonstrates that our proposed method establishes new SOTA in offline RL.
Table 3: Comprehensive performance comparison on D4RL MuJoCo. Average scores for our algo- rithm and baselines are reported. For baselines that report standard deviations, compute Cohen's to quantify the statistical significance of the improvement achieved by our method. A value of d >>> indicates very large and statistically significant difference.
Algorithm
BA-MCTS
Optimized
COMBO
MOReL
MOPO
CQL
BEAR
Avg. Score
80.3 (1.I)
(5.4)
66.8 (3.4)
64.4
36.7 (11.0)
$4.9
39.8
Avg. (No Random)
96.5 (1.3)
Coben's d
3.5
4.6
4.7
Algorithm
BRAC-V
SAC
BC
MOBILE
CBOP
RAMBO
APE-V
Avg. Score
31.4
25.1
76.5 (2.7)
73.9 (2.2)
68.9 (2.8)
72.7 (L.8)
Avg. (No Random)
Cohen's d
4.8
Algorithm
MAPLE
TD3+BC
IQL.
SAC-N
EDAC
DT
Avg. Score
61.3 (1.8)
(1.8)
54.6 (2.4)
76.3 (0.6)
76.0
Avg. (No Random)
77.0
(1.2)
Coben's 's d
12.7
12.4
4.8
2.2
17.4
Search-Based Based Methods. MuZero also applies deep search to MBRL. To evaluate its performance on D4RL MuloCo, we use the open-source implementation and hyperparameter configurations of Sam- pled EfficientZero (Ye et al.) 2021) provided by LightZero (Niu et al.) 2023). Benchmarking results from LightZero indicate that Sampled EfficieniZero achieves the best performance on (online) Mu- JoCo tasks compared other MuZero variants. To adapt Sampled EfficientZero for offline learning, we employ the reanalyse technique proposed by (Shhrittwweeser et al. (2021). The evaluation results As shown, the results are significantly worse compared
are presented in Figure (Appendixi to the performance of offlin RL. methods listed in Table 2 despite Sampled EfficientZero's much higher computational cost. In Appendix F.4, we provide a detailed comparison of the compu- fational costs and algorithm designs between our method and Sampled EfficientZero. World model learning is the foundation of MBRL and can be particularly challenging in continuous con- trol and offline 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.
8

---

Table 4: Comparisons 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 (Ours)
BA-MCTS-SL (Ablation)
BA-MBRL (Ablation)
CQL
Optimized
Temperature
223.83
-21,16 (5.00)
-29.35 (4.72)
-59.62 (I.57)
-83.55 (10.56)
Rotation
-19.07 (5.85)
-14.14 (1.88)
-31.33 (11.54)
-85.48 (2.72)
-71.54 (9.88)
8n.
-18.93 (1.75)
-37.03 (17.98)
-23.40 (10.77)
-36.37 (1.17)
-57.84 (10.27)
Average
-20.61
-24.11
-28.03
-60.49
-70.98
5.3 ABLATION STUDY
For principled ablation study, we reimplement BA-MCTS and its ablated variants based on the codebase of Optimized (Lu et 2022). which thoroughly explores design choices in offline MBRL, making minimal changes to the code and hyperparameter settings (see Appendix D). Specifically, we first add a belief adaptation module to Optimized to obtain BA-MBRL, and then incorporate the Continuous BAMCP planner (Section 4.2) to create BA-MCTS. Finally, replace the policy distillation module of BA-MCTS from SAC to supervised learning (as detailed in Section 4.3 yielding BA-MCTS-SI. The detailed results are presented in Table [T] and Appendices BA-MBRL significantly outperforms Optimized, demonstrating the necessity of Bayesian model adaptation. BA-MCTS further improves upon the performance of BA-MBRL, underscoring the effectiveness of employing deep search as powerful policy improvement operator. Finally, BA- MCTS-SL achieves performance comparable to BA-MCTS. However, BA-MCTS-SL. can be less stable on certain complex tasks: it struggles on the Walker2d environments, where warm-up phase (using BA-MBRL) needed to establish a good initial policy for supervised learning to be effective. Finally, we conduct an ablation study to assess the necessity of the reward penalty. By removing this penalty term (i.c., setting À = 0 in Eq. (5), we observe a consistent drop in average performance. This demonstrates that the pessimistic reward penalty is crucial for stopping the agent exploiting inaccuracies in the learned world models. Results are in Table 12 in Appendix F.7
F.5
F.6
5.4 APPLICATION TO TOKAMAK CONTROL
We evaluate our algorithm on three target tracking tasks in tokamak control. The tokamak is one of the most promising confinement devices for achieving controllable nuclear fusion. We use a well- trained data-driven dynamics model provided by Char et al. (2024) as "ground truth" simulator for the nuclear fusion process during evaluation, and generate dataset (by this dynamics model) containing 725,270 transitions for offline RL. We select reference shot (i.e., an episode of a fu- sion process) from DIII-D (a tokamak device located in San Diego), and use its trajectories of Ion Rotation, Electron Temperature, and 3,, as targets for three tracking tasks. The tracking tasks have 28 28-dimensional state space and 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 sample from this model. For details the simulator, and the design of the state/action spaces and reward functions, please refer to Appendix E
We obtained access the DIII-D data for limited time window, which restricted our ability to provide evaluation results across a wide range of algorithms. For principled comparison with the baselines, we utilized the codebase from the ablation study, where our algorithm implementations are primarily based on Optimized, to ensure that performance improvements are attributable to al- gorithmic design rather than implementation differences. Additionally, we included evaluation of a widely-used model-free offline RL method: CQL. The average tracking eтrоrs over the final 10 training epochs are reported in Table where BA-MCTS outperforms all baselines. These results demonstrate the potential of our algorithm for application in complex, data-driven control problems.
6 CONCLUSION
We propose framing offline MBRL 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 MCTS. This planning process is integrated into policy iteration framework, enabling the derivation of a policy function for real-time execution. In our eval- uation, we test several variants of our algorithms to separately highlight the effectiveness of Bayesian

---

RL and deep search. Additionally, validate the superior performance of approach on both standard offline RL benchmarks and challenging continuous control tasks in Nuclear Fusion.
ACKNOWLEDGEMENT
This work was funded in part by Department of Energy Fusion Energy Sciences under grant DE- SC0024544.
REFERENCES
Joseph Abbate, Rory Conlin, and Egemen Kolemen. Data-driven profile prediction for diiid. Nu- clear Fusion, 61(4):046027, 2021.
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-bused offline re- inforcement learning with diversified g-ensemble, 2021. URL htttttss///rraai.rrr////bs
2110.01548
Ioannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas K. Hubert, and David Silver. Plan- ning stochastic enviromments with a learned model. In International Conference on Learning Representations. OpenReview.net, 2022.
Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. In International Con- ference on Learning Representations. OpenReview.net, 2021.
Marco Ariola, Alfredo Pironti, et al. Magnetic control of tokamak plasmas, volume 187. Springer, 2008.
John Asmuth and Michael L Littman. Approaching baycs-optimaliity using monte-carlo tree search. In International Conference on Automated Planming and Scheduling, 2011.
John Asmuth, Lihong Li, Michael L. Littman, Ali Nouri, and David Wingate. bayesian sampling approach exploration in reinforcement learning. In Uncertainty in Artificial Intelligence, pp- 19-26. AUAI Press, 2009.
David Auger, Adrien Couëtoux, and Olivier Teytaud. Continuous upper confidence trees with poly- nomial exploration consistency. In European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases, volume 8 8188, pp. 194 209. Springer, 2013.
Cameron B Browne, Edward Powley, Daniel Whitehouse, Simon M Lucas, Peter I Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. A survey of monte carlo tree search methods. IEEE Transactions Computational Intelligence and Al in Games, 4(1):1-43, 2012.
Pablo Samuel Castro and Doina Precup. Smarter sampling in model-based bayesian reinforce- ment learning. In Joint European Conference Machine Learning and Knowledge Discovery in Databases, volume 6321, pp. 200-214. Springer, 2010.
Ian Char, Joseph Abbate, Laszlo Bardoczi, Mark D. Boyer, Youngseog Chung, Rory Conlin, Keith Erickson, Viraj Mehta, Nathan Richner, Egemen Kolemen, and Jeff G. Schneider. Offline model- based reinforcement learning for tokamak control. In Learning for Dynamics and Control Con- ference, volume 211, pp. 1372, 2023.
Ian Char, Youngseog Chung, Joseph Abbate, Egemen Kolemen, and Jeff G. Schneider. Full shot predictions for the DIII-D tokamak via deep recurrent networks. CoRR, abs/2404.12416, 2024.
Jiayu Chen, Bhargav Ganguly, Yang Xu, Yongsheng Mei, Tian Lan, and Vaneet Aggarwal. Deep generative models for offline policy learning: Tutorial, survey, and perspectives on future dirce- tions. Transactions Machine Learning Research, 2024. ISSN 2835-8856.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, 2021a. URL http::s:///axxv..r//aab.ss.
01345.
10

---

Xiong- Kiong-Hui Chen, Yang Yu, Qingyang Li, Fan-Ming Luo, Zhiwei (Tony) Qin, Wenjie Shang, and Jieping Ye. Offline model- based adaptable policy learning. In Advances in Neural Information Processing Systems, pp. 8432-8848 -2244433, 2021b.
Jacob Cohen. Statistical Power Analysis for the Behavioral Sciences. Lawrence Erlbaum Associates, Hillsdale, NJ, 2nd edition, 1988.
Adrien Couêtoux, Rean-Baptiste Hoock, Nataliya Sokolovska, Olivier Teytaud, and Nicolas Bon- nard. Continuous upper confidence trees. In International Conference on Learning and Intelligent Optimization, volume 6683, pp. 433-445. Springer, 2011.
Christopher Diehl, Timo Sievernich, Martin Krüger, Frank Hoffmann, and Torsten Bertram. Uncertainty model-based offline reinforcement learning for automated driving. IEEE Robotics Automation Letters, 8(2:11667-1174, 2023.
Michael O'Gordon Duff. Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. University of Massachusetts Amherst, 2002.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu, minimalist approach to offline reinforcement learning, 2021. URLltttps: /
//axxvv http/rxivv.oo org. abs 0tlbb///2166666666
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, and Aviv Tamar. Bayesian reinforcement learning: A survey. Foundations and Trends@ in Machine Learning, 8(5-6)335-443) 2015.
Dibya Ghosh, Anurag Ajay, Pulkit Agrawal, and Sergey Levine. Offline RL policies should be trained to be adaptive. In International Conference on Machine Learning, volume 162, pp. 7513- 7530. PMLR, 2022.
Arthur Guez, David Silver, and Peter Dayan. Scalable and efficient learning based on monte-carlo tree search. Journal of Artificial Intelligence Research, 48:841- 883. 2013.
reinforcement
Arthur Guez, Nicolas Heess, David Silver, and Peter Dayan. Bayes-adaptive simulation-bail based search with value function approximation. In Advances in Neural Information Processing Sys- tems, pp. 451-459, 2014.
Kaiyang Guo, Yunfeng Shao, and Yanhui Geng. Model-based offline reinforcement learning with pessimisn-modulatee dynamics belief. In Advances in Neural Information Processing Systems, 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning, volume 80, pp. 1856-1865. PMLR, 2018.
Jessica B. Hamrick, Abram L. Friesen, Feryal M. P. Behbahani, Arthur Guez, Fabio Viola, Sims Witherspoon, Thomas Anthony, Lars Holger Buesing, Petar Velickovic, and Theophane Weber. On the role of planning in model-based deep reinforcement learning. In International Conference Learning Representations. OpenReview.net, 2021.
Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and planning in complex action spaces. In International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4476-4486. PMLR, 2021.
Jihwan Jeong, Xiaoyu Wang, Michael Gimelfarb. Hyunwoo Kim, Baher Abdulhai, and Scott Sanner. Conservative bayesian model- based value expansion for offline policy optimization. In Interna- tional Conference on Learning Representations. OpenReview.net, 2023.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? CoRR, abs/2012, 15085, 2020. URL https::::: /ttp///aiiiie wttp://wwwoor////////ooa 15085

---

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offiin reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In European Conference on Machine Learning, volume 4212, pp. 282-293, Springer, 2006.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. 2021. URL https: http:///awivv..or/////////////// http//b//2111e 06169
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Sys- tems, pp. 11761-11771, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, pp. -402-6413, 2017.
Jongmin Lee, Wonseok Jeon, Geon- Hyeong Kim, and Kee-Eung Kim. Monte-carlo tree search in continuous action spaces with value gradients. In The Thirty-Foure AAAI Conference on Artificial Intelligence. pp. 4561-4568. AAAI Press, 2020.
Geon-Hyeong
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives open problems. CoRR, abs/2005.01643, 2020.
Zichen Liu, Siyi Li, Wee Sun Lee, Shuicheng Yan, and Zhongwen Xu. Efficient offline policy optimization with a learned model. In International Conference on Learning Representations. OpenReview.net, 2023.
Cong Lu, Philip
J. Ball, Jack Parker-Holder, Michael A. Osborne, and Stephen Roberts. Revisiting design choices in offline model based reinforcement learning. In International Conference on Learning Representations. OpenReview.net, 2022.
Andrzej Maćkiewicz and Waldemar Ratajczak. Principal components analysis (pea). Computers & Geosciences, 19(3):303-344 1993.
Yazhe Niu, Yuan Pu, Zhenjie Yang, Xueyan Li, Tong Zhou, Jiyuan Ren, Shuai Hu, Hongsheng Li, and Yu Liu. Lightzero: A unified benchmark for monte carlo tree search in general sequential decision scenarios. Advances in Neural Information Processing Systems, 2023.
Yaniv Oren, Matthijs TJ Spaan, and Wendelin Böhmer. E-mets: Deep exploration in model-based reinforcement learning by planning with epistemic uncertainty. arXiv preprint arXiv:2210.13455, 2022.
A. Pironti and M. Walker. Fusion, tokamaks, and plasma control: an introduction and tutorial. IEEE Control Systems Magazine, 25(5):30-44) 43, 2005.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Marc Rigter, Bruno Lacerda, and Nick Hawes. RAMBO-RL: robust adversarial model- based offline reinforcement learning. In Advances in Neural Information Processing Systems, 2022.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-66 609, 2020.
Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barckatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with a learned model. In Advances in Neural Information Processing Systems, pp. 27580-27590 2021.
12

---

David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy P. Lillicrap, Karen Simonyan, and Demis Hassabis. Mastering chess and shogi by self-play with general reinforce- ment learning algorithm. CoRR, abs/1712.01815., 2017.
Patrick Slade, Zachary N. Sunberg, and Mykel J. Kochenderfer. Estimation and control using sampling crasee bayesian reinforcement learning. IET Cyber-Physical Systems: Theory Ap- plications, 5(1):127-135, 2020.
Jonathan Sorg. Satinder Singh, and Richard L. Lewis. Variance-based rewards for approximate bayesian reinforcement learning. In Uncertainty in Artificial Intelligence, pp. 564-571. AUAI Press, 2010.
Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. Model-bellman inconsistency for model-based offline reinforcement learning. In International Conference on Machine Learning, volume 202, pp. 33177-33194, PMLR, 2023.
Zachary Sunberg and Mykel J. Kochenderfer. Online algorithms for pomdps with continuous state, action, and observation spaces. In International Conference on Automated Planning and Scheduling, pp. 259-263. AAAI Press, 2018.
Yi Wang, Kok Sung Won, David Hsu, and Wee Sun Lee. Monte carlo bayesian reinforcement learning. In International Conference Machine Learning, 2012.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. 11361
CoRR, abs/1911.11361 2019. http://///axx
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In Advances in Neural Information Processing Systems, pp. 25476-25488, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. In Advances in Neural Information Processing Systems, 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. COMBO: conservative offline mosel-based policy optimization. In Advances in Neural Informa- tion Processing Systems, pp. 28954 28967, 2021.
Xianyuan Zhan, Xiangyu Zhu, and Haoran Xu. Model-based offline planning with trajectory prun- ing. In International Joint Conference on Artificial Intelligence, pp. 3716-3722. ijcai.org, 2022.
Chen Zhao, Yazhe Niu, Kaixin Huang, Yu Liu, and Chun Yuan. A perspective of improper dynamics on offline model-based planning, 2024. URL btteps://ppenteviee..ntttt
7kubdPr1RY
13

---

A BAMCP
Algorithm BAMCP
procedure SIMULATE((S, (1), if === 0 then
Input: T, E. dmax, R, 7/1 A, c
return 0
procedure SEARCH( 8, h), b(0)
end if
for e 1... do
if ((s,, h)) == 0 then
0 A~ b(.) SIMULATE((S, h)), 0, dmax)
for do
(((s,)),a). Q((s, a)) + 0,0 end for s' r R(s,
end for return arg max, (((s,)) a) end procedure
(()
0,0, R=++2RoLLOU(( has h, 0,,,0 1)
procedure ROLLOUT((s, if d = then return 0 end if
(((s,)) ((s,,, + 1, ((s, a)) + R
return It
end if
~ (s,(h)))
- arg max, (((s,)), + cV
lox N((s.)
((oo.))..)
s' Po(-(s.a)
s' ~ Po(-1s,a), r + R(s, a)
R(s,a)
+ has'),))s)))))))),)))) (,ass),,,,, [((s, h)) += 1, ((s, h), += (((s, h)) +=
R R+r+2RoLLOUT(e+ haa has'), 0,
return R end procedure
+= R R-Q((s,)).a)
N((ah).
return R end procedure
Bayes Adaptive Monte Carlo Planning (BAMCP) (Guez et al. 2013) a sample-based online plan- ning method, aiming to find the action a" that approximately maximizes the expected return at de- cision point (s, h) under the BAMDP. Its detailed pscudo 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 (Koesis) & Szepesvári 2006) 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
man- aging the tradcoff between exploration and exploitation. (As shown in Algorithm 3 N is the visit count, while Q denotes the accumulated reward.) Root sampling refers to sampling the dynamics model only at the root node (i.e., 0 ~ b(-)) and not adapting the belief b(.) according to the Bayes rule and experience during the search process, of which the rationality is justified in the following lemma.
used for selecting actions at non-leaf nodes, i.e., a + arg max: Q( h), cv
Jow N((a.)]] ([(,,))).r))
Lemma A.1. For all suffix histories h' of h, b(ol)) b(0)W"). Here, b(OW)) is the true posterior probability of at the decision point h', while b(0)h') the probability of experiencing 0 at ht when using root sampling.
Proof. This lemma be proved by induction.
Base case: When h' h, b(()h') b(0)h') b(0).
Step case:
b(0)has') P(has' (1hhas/(0)P(0)/((aaa)
(6)
=ZM(0)))P(($(ssaa)
-26(01))6688|1sc))
(o(00 a) b(0)has')
14

---

1/ l Po(s'|s, ( hab(() 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 at each node h is made independently of the sample 0. As for the last equality, to experience at has', the sample 0 needs to traverse ha (with probability b(0) ha)) and then the state s' needs to be sampled, which is with probability (s |s, a), so b(O has') x b(0)ha)Po (s
Here, Z f/ Po(s"|s, a)b(()))d) = Po(s'|s, hiih
B CONSISTENCY PROOF OF CONTINUOUS BAMCP
This appendix provides the detailed proof for Theorem 4.1 Our proof adapts the consistency proof methodology developed for Polynomial Upper Confidence Trees (PUCT) (Auger et al. 2013). To apply this framework, in the main text can be formally treated as a standard, albeit large-state-space MDP. Throughout this proof, we will refer to this mathematical object as the Pessimistic MDP (P-MDP). Our goal is to show that Continuous BAMCP is a consistent solver for this P-MDP.
first note that the Pessimistic Bayes-Adaptive MDP (P-BAMDP) defined
The remainder of this appendix is organized follows. Section B.1 formally defines the P-MDP and introduces the key notations, assumptions, and the induction property central to our proof. Sec- tion B.2 establishes the base case for our backward induction. Sections [33 and B.4 present the two core inductive steps, showing how the consistency property propagates upward through the random and decision nodes of the search tree, respectively. Finally, Section B.5 concludes the proof.
B. PREAMBLE AND DEFINITIONS
P-MDP Formalization The P-MDP that our planner solves is defined by the following compo- nents:
State Space S+: The augmented information state space, where each state is (s, h).
Action Space A:
The action space of the original problem.
Transition Funetion P+: The belief- weighted transition function over information states.
Reward Function F: The fixed, bounded, penalized reward function from Eq. 5
Relevant Notations We adopt the following notations:
A decision node, corresponding an information state (s, h).
w: A random node, corresponding to state-action pair a).
Vg (w): The optimal Bellman value of node and w under the P-MDP. ((=),n(w) The total number of visits node and during MCTS.
(=), (w): The empirical average return obtained at node and w.
V=)): The value of a node = estimated by the learned value function from the previous policy iteration.
d: The depth of a node in the search tree, integer for decision node while semi-integer for random node
Definition B.1 (Exponentially sure in.). We say that some property (P) depending on an integer n is exponentially sure in (denoted e.s.) there exists positive constant C, h such that the probability that (P) holds is at least
C Cexp(-hm).
Our proof relies on the following assumptions:
Assumption B.2 (Policy Regularity). For any decision node (s, h)) and any precision A 0, the policy x--lz) used to sample new actions has non-zero probability of sampling a -oppiial action within the current P-MDP. Formally, there exist fixed constants 0 0 and p > 1 such that,
VP (w (2,a)) (2(
A with probability at least min(1,027)
15

---

Assumption B.3 (Bounded Value Function Error). In algoritim, when Monte Carlo Tree Search simulation reaches a leaf node & of the current search tree (i.e., a node that has not yet been expanded), we use learned value function, V. estimate its value for the backup step. This function is trained as part of the overall policy iteration framework. We assume that V is an imperfect approximator of the P-MDP's true optimal value function, Vn, and that its approximation error is uniformly bounded by a constant EV 2 0 over the entire state space:
[V(2) Vi (2) cv,
Vz E
Furthermore, we assume the oulput of (=) is bounded.
The proof aims to establish the following property via backward induction: Definition B.4 (Induction Property Cons(r, f, d)). There exists constant C2 0 and function fa(:) such that for all decision node at depth d,
((=) - (z) S Can(z) + fa(Ev) e.s. in n(z)
and for all nodes w at semi-integer depth d +
(Vw)) (y(w)) & 70+4
fa+ (ev) e.s. in n(w)
where Ey the base approximation error of the value function (from Assumption B.3
B.2 BASE STEP
The backward induction begins at the maximum search depth, dmaa. We aim to establish the induc- tion property Cons(?,, f. d) for random nodes w at depth d dmaz 1 As our algorithm uses a 2- learned value function, V(2)), to evaluate leaf nodes 2, (children of w), this step relies the value function error bound established in Assumption B.3
Our goal is to bound the total error (V(w) - (w)l- Using the triangle inequality, we decompose this error:
(ww) - (y((w) < (w)) ([V(= + [=.) VØ (w:)
We bound the two terms on the right-hand side separately.
(a) Bounding the Sampling Error of the MCTS sampling process. Since the value function output V(2;, is bounded (per Assumption B.3), we can apply Hoeffding's s inequality, Following the original PUCT proof (Auger et al. 2013),
The first term, ([(s) [+:() is the concentration error
we set the deviation tolerance t = n(w) 1, which gives us the following bound:
V (w) - E[V(2))) < O(n(to)- $) e.s. in n(w)
(b) Bounding the Approximation Error The second term, [x[] Vg (w)), represents the inherent bias from using an approximate value function. We bound using Assumption B.3
[E]])(2)] V22 (w)) [V[ (z:) ([V((()))
V (2;) V (=))l
[[ev] = Ev
Conclusion for the Base Step By combining the bounds for the sampling error and the approxi- mation error, we obtain the bound on the total error:
cc (ww) - V (w)) s ((n(w)-b) k) + cv e.s. in n(w)
This result precisely matches the form required by our revised induction property, Cons(7, f, d). We have successfully shown that the property holds for the base case at depth dmax 2. with parameters:
Ydmas
SSmas (ev) EV
16

---

B.3 INDUCTIVE STEP: FROM DECISION TO RANDOM NODES
In this section, we perform the first inductive step, moving from decision nodes up to random nodes. We assume the induction property Cons(y, [, d) holds for all decision nodes at depth d. Our objective is to prove that, as consequence, the property also holds for their parent nodes, which are the random nodes w depth d 1.2
A key distinction in our algorithm is the special handling of the very first action sampled from decision node. This requires us to categorize all random nodes w (at depth d - ) based on their history:
Category 1:
First-Child Nodes (wa): random node that was the first child ever cre- ated from its parent decision node. Its initial value estimate was derived directly from our learned value function, V.
2. Category 2: Standard Nodes (wz. > O): All other random nodes. Their value estimates have always been computed through the standard MCTS backup process.
We will analyze these two categories separately show that the induction property holds universally.
B.3.1 ANALYSIS OF STANDARD RANDOM NODES (wy. j > 0)
For a standard random node w in this category, the proof structure closely follows Section 5 of Auger et al. 2013)). We aim to bound the total error (V(w) VØ (w))
Let in = [ be the newest child of w (a decision node at depth d), and let = iu be the number of "mature" children.(iee.. all children except for the newest). Let (1) ra + (1) denote the return through child i. We decompose the total error into four terms:
VV <
n(i) ((w) T
R. (3) +01++ Σ (Rn E( ())I
(a) Weight Deviation Fror
(b) Sampling Error of "Mature" Returns
n(io) (io)l
E (2) - V (w)) +
n(w)
(c) Bins/ Approximation Error
(d) "Nowest" Child Error
We now bound each term:
(a) Bounding the Weight Deviation Error This term is bounded by leveraging the fact that the children of random node are explored in a nearly uniform manner (a consequence of the DPW mechanism, as shown in Lemma I of Auger et al. (013).. Following their derivation, this term is bounded by O(n(w)-(1-0 + n(w)
(b) Bounding the Sampling Error of "Mature" Returns This term represents the concentra- tion error for the returns of the mature child nodes. The returns Ran (i) are i.i.d. bounded random variables. Therefore, we can apply Hoeffding's 's inequality. To align with the proof by Auger et al. (2013) and derive the resulting convergence rate for this layer, we select the deviation tolerance as:
:= n(w) 14374
With this choice, the application of Hoeffding's s inequality shows that the error for this term is bounded by (n((w) 1+3m%) and that this bound holds exponentially surely.
(c) Bounding the Bias/ Approximation Error This term is unique to our proof and requires our induction hypothesis.
E Ry ( = [Br, + +ッの Eir, + +V11 -)E V1
{) - VV
17

---

By the induction hypothesis for nodes at depth d, we have VG) - Vi Vo1 Can(j) fa(ev).
Substituting this into the inequality gives:
[m (i) Vi (w)) S :B(Cm((j)>>)) xSa(ev) ( e.s. in n(w)
The first part of this bound is decaying term, while the second part shows how the bias propagates.
(d) Bounding the "Newest" Child Error Following the same logic as Auger et al. (001)), this term is bounded by ((n(w) (a) because the newest child has been visited only small number of times.
Combining these bounds, the overall convergence rate is deter- mined by the coneest-on term. Following the derivation in the PUCT paper, setting the exploration parameter I43ha establishes the slowest convergence rate as 14334 The bias term holds.
Conclusion for Standard Nodes
374
is propagated from the layer below. Thus, for any standard random node, the Cons(y, f, d) property
B.3.2 ANALYSIS OF THE FIRST-CHILD RANDOM NODE (wa)
For the first child wo. the value estimate includes special term from its initial visit. On the first visit, the algorithm samples a next state, which we denote sy, and receives a corresponding reward ro- The next state's value is thus initially estimated using the learned function, V(s))- The error analysis for this random node, (wwg) (wo)), will contain the same four error terms as in the standard case (for the subsequent backups), plus an additional eπror term originating from the first visit:
Additional Error
n(w)) (ro + (V(s)) n(w))
Vi (wo)
Since (ss),, and (wo) are all bounded constants, this additional error term is bounded by O(n(wo)-1).
The overall convergence rate for we is determined by the minimum of the rates of all error com- ponents. From the PUCT derivation,
know that Ya $ which implies 7a-1/2 1+37a 1. Therefore, the convergence rate of the additional error term, ((n(wo)='), is faster than the bottle- neck rate of the standard MCTS eиror terms, O(n(w)) -7a-1/2). This special term does not change the overall convergence rate. Thus, the Cons(7; J, d) property also holds for the first-child random node wo with the exact same parameters.
B.3.3 CONCLUSION FOR THIS SECTION
We have shown that for any random node w at depth d - 2. if its children at depth d satisfy the Cons(r, f,d) d) property, then w itself also satisfies the property. The parameters of the property propagate upward according to the following recursive formulas:
7a- = I+3rg
fa-y (Ev) = rfalev)
B.4 INDUCTIVE STEP: FROM RANDOM TO DECISION NODES
In this section, we perform the second inductive step. We assume the induction property Cons(y, f, d) holds for all random nodes w at depth d+ 1 Our objective prove that the property also holds for their parent nodes, which are the decision nodes at depth d. The proof adapts the structure of Section Auger al. (2013), carefully accounting for the bias term fa+ (ev) from our induction hypothesis. First we establish an upper bound on ((2 (2), and then lower bound.
B.4.1 UPPER BOUND ON THE ERROR
To establish an upper bound on (22 V (z), we partition the children w, of a into two classes, with parameter to be determined later:
18

---

Class I: Children wy with few visits, n(w,) 1 ((2))_-0-
Class II: All other children, for which n(w;) > n(z)1-a-e,
For Class I, the number of children is at most r²z)". : so their total contribution is bounded by x(2)* - = m=2))= - For Class II, we apply the induction hypothesis for nodes at depth + 1: ((2)
n(w)) (V(w)) (x))
n(wi) n2())
Σ
Σ
(V(w;) - -V (w:))
.eClass II n(2)
iEClass
r(wws)
(C (n(ws)) -a+ + fatz (cv))
n²z)
ieClass II
-a+ 3 + sat (ev) Cafy (n(2)1-a--)
Combining the two parts, we get:
- - ( + Catyn(z) - ( + fat (cv)
74- (1-a)
To optimize this bound, balance the exponents by choosing final upper bound for the error:
This leads the
147
Yat (1-w)
((2) - (V)))V(())) (1+
1 + +5a+z(en
B.4.2 LOWER BOUND ON THE ERROR
The proof for the lower bound is more involved. It relies on the parameter definitions from the PUCT paper, well two key results regarding the UCT exploration mechanism, which we state here for completeness before proceeding.
Lemma B.5 (Adapted from Lemma 3, Auger et al. (2013)). Consider a bandit setting where the score of child node w, time n computed by the UCT rule:
g(n)
sca (wi) = (w:) +
n(w:))
where g(n) is non-decreasing exploration function. If denotes the tth constucted child, for all n > (a(1-a) we have
n(w,) min (g(n2-a),n) 2=²)
Corollary B.6 (Adapted from Corollary 3, Auger et (2013)). ), For the funncttoion g(n)
n' with O < < 1, we obtain
1 (11 ifi (1)-))
n(w;)) 2
(8)
With these results, we can define the parameters as functions of To+3 and p (from Assumption B.2 for this layer and proceed with the three-step proof structure.
Mat b
Progressive widening coefficient: a = 1+47
Exploration coefficient: e:= 2p 1+d7a+
Auxiliary proof parameters: = Ite (1-a) and A = (((())e((1-o)) - -741
19

---

Step 1: Existence of near-optimal child Based on Assumption B.2 (Policy Regularity), direct adaptation of the proof in Auger et al. (2013) shows that a oetrroltimal child node discovered with high probability. Formally: Exponentially surely in n(z), there exists time [T(2))(1 (()))(1)))) a child wo of such that
(wo)) (z) - A and in n(=){(1-o)a
Step 2: Score of children selected at later stages The objective of this step is to show that any child selected after time n(2)f must have high score. Let n' be such that n I n'' n. And, by Corollary B.6
1 (c(1-a) n' nwww))= 2 7"
1 En
a)
2 nn
4
First, according to Corollary B.6 the near-optimal child wy discovered in Step will be visited number of times n(wo) that grows polynomially with n²(2).
can therefore apply our induction hypothesis Cons(y, f.d) to wo, with the Eq. there exists
C' > 0 that:
(wo) (wo)) C (2 -hat
Sa+ (cv)
-=+ ((8))--fa-(a0 (cv)
Now, consider any child selected by the UCT rule at later time n'' 2 n2(25, Its score must be least high the score of wp- This gives the key inequality:
ne (wrr
n2
(w1)
V(wo))
n'(wo)
Combining these results, we can establish lower bound on the estimated value (wi) any child selected at late stage of the search.
(ww) + n'(w1) 2
As shown by Auger et al. (2013). this property holds exponentially surely in n(2).
Step 3: Bounding the value estimate (2) Consider child selected after n(2)s, By the previous step, exponentially n(2), this child must either satisfy
2 A
(10)
n(w))
or (ww) 2 V(=) (2 + C')A (ev)
(11)
We now bound the average value ((z) by partitioning the children of into three categories:
Children visited only before time n==)
2. Children visited after n==)5 satisfying 10
3. Children visited after n=)) satisfying
We bound the contribution of each category to the total error sum. The contribution from Category
1 is bounded by O(n(2)=-1) 1), since have
n(w) M cat.1 T²(2)
Σ
- (v(w)-
E in ca.1 nuw))
n(2) n22)
<
nl2)
20

---

For Category 2, the definition of the category implies n(w;) - n Since there are at most n²)) children, their total contribution to the error sum is bounded by ((n(2) (A²²)) For Category 3, the bound on (w:) from Step 2 introduces the crucial (f2+1/2(Ev) term.
Combining the bounds from all three categories gives the final lower bound on the error. Following the argument in Auger et al. (2013)
(ev) m(2)5-1_ n(2)ate-1
((z) - - 22) -
cat.3
cat.
cat.2
By comparing the magnitudes of the decaying terms with the chosen parameter definitions, is done in Auger et al. (2013), this complex expression simplifies to the desired form.
B.4.3 CONCLUSION FOR THIS SECTION
By combining the results from the upper and lower bound analyses, we conclude that the total error decision node is bounded as:
(v)) (2)1
(z)1 S Can(z) fev)(ev)
This demonstrates that the induction property Cons(7, f. d) holds for depth d. The parameters prop- agate upward according to the following recursive formulas, which are identical those derived by Auger al. (2013):
C. = (5 + 7444
na-3
la = 1477a1
falev) = f4+z (ev)
B.5 CONCLUSION OF THE PROOF
In the preceding sections, we have established all the necessary components for our backward in- duction argument.
In Section we established the base case, proving that the induction property Cons(7, J, d) holds for random nodes at depth d dmaz 2. with initial parameters
Idmar - 3/2 and famas ( EV.
In Section B.3
we proved the inductive step from decision nodes up to random nodes, 7d and fa (cv) 7fa(ev).
deriving the recursions 7a- 1+3ya
In Section
B.4
we proved the inductive step from random nodes up to decision nodes, 7 and falev) fat (ev).
deriving the recursions "d 1+77
With the base case proven and both inductive steps validated, conclude by backward induction that the property Cons(y, f, d) holds for all depths up to the root the search tree (d = 0).
By solving the recursion for the bias term, can find the accumulated error at the root. The base error at depth dmaz 2 is famae- (cv) = cy. The error term is multiplied by a discount factor each time propagates from a decision node layer up to a random node layer. This occurs dmaz times from the leaves to the root. Therefore, the bias term at the root node (d 0) is folev) dmaoee.
Specifically, for the root node 2o, its total error bounded by:
(=x)) - Vi (2o) & Con(zo) + adma
where 7o > and C. are constants derived from the recursions. This shows that as the number of simulations F. n=(2o) oc, the estimated value of the root node converges to ball of radius
21

---

Jamaryy cy around the true optimal value of the P-MDP. Notably, the size of this error ball, which stems from our function approximator, shrinks exponentially with the search depth dmaz.
This completes the proof of Theorem establishing that Continuous BAMCP is consistent planner for the P-MDP.
C ALTERNATIVE DESIGN CHOICES FOR CONTINUOUS BAMCP
The terms labeled in blue in Algorithm | have alternative design choices, which are introduced as below. Empirical comparisons among these alternatives are reserved for future work.
D((s,)),a) in the exploration strategy, i.e., a <- arg www.maxxxx..... Q((s, h),, could take var- ious forms. For instance, in Couêtoux et al.] (20T1); Guez. et al. (2013): Lee et al. (2020), (((s,)),z) = (((s,))x)) + cV
log (((ssh)), in PUCT (Auger et al. 2013). h((ss,),20 = [((oh)),a)
(((s,)),x) + (((a.)),))) where e(d) is schedule of coefficients related to the search depth d; in Sampled MuZero (a variant MuZero that can be applied in continuous action spaces (Hubert
et al. )021)), (((s,)),x) = Q((s, h), + (x](ss,)) h))
(N((a.)))
(c + +log
[(((sh))++ott
[+N((s)),,))
Here, e, C1, C2 are hyperparameters, # = 1-111/7 is a sample policy defined upon the real pol- icy In particular, at each decision point h), Sampled MuZero would sample M actions {a:} from the distribution 1/T and accordingly define (s) h)) Ex IM. where 0 is temperature hyperparameter. Thus, Sampled MuZero does not adopt progressive widening like ours. Following BAMCP, we adopt the first definition of (((s, ,), v). though it could potentially be improved with other choices. In addition, as an implementation trick (Hamrick al., (02T), the Q estimates are usually normalized into Q [0, 1] before being used to calculate Q above. The (((a.h).z)-amm. where Qmax and Qmin are normalized estimates can be computed as Q((s, h), x) = Qui Quis the maximum and minimum Q values observed the search tree so far.
As the planning/search result, Teet take multiple forms. In Guez et al. (201)) Sumberg & Kochenderfer (2018) Lee et al. (2020), Trei((s h)) = arg maxago((s.N)) Q[ h). aa; in (Sampled) N((e.h).a)*/+
MuZero, (a) h)) E.cort.a) 12) in ROSMO (a variant of MuZero with improved performance in offline scenarios (Liu et al. 2023)) (a) (s. h)) (a) h)) exp(Q((s,) (()(((s)))))))
((s, hi)). Here, 7 (o, 1] is a temperature parameter and decays with the training process, ensur- ing the action selection becomes greedier. We select the second form for Tee in Algorithm[] 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-ahear 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 a and | to be functions of the search depth d, while UCT-DPW (Couêtoux et al. 2011) utilizes a different set of conditions: [K,,N(s, A))a7 2 [C((s, h))l. K,N((s, h), a)d constant hyperparameters. Moreover, 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 random node can be sampled from C((s. h), a) according to distribution proportional the number of visits (following UCT-DPW). As shown Algorithm|] [ we follow the designs of PUCT, but keep and B as constants for simplicity in hyperparameter fine-tuning.
((s, where Ka, Ks, a, 3 are all
Finally, the condition for continuing the simulation procedure, i.e., NNss, h)) > 1, could potentially be replaced with N((s, a) > 1 or N((s', hars')) > O. These conditions indicate that the nodes (s, h). (s, h, a), and (8', hars') 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) h)) as the estimation.
The former approach is widely used in online planning algorithms (Guez et al.) 2013) Sunberg & Kochenderfer 2018 Lee et al. 2020). while the latter is used in policy iteration frameworks like ours.
22

---

D KEY HYPERPARAMETER SETUP
BA-MBRL
BA-MCTS
BA-MCTS-SL
Data Type
Environment
K
X
H
N
K
A
H
N
K
A
H
N
random
HalfCheetah
10
7
6
200
10
6
800
10
7
500
random
Hopper
50
47
700
50
47
800
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
5
800
12
6
500
medium
Hopper
40
42
200
40
42
800
12
40
200
medium
Walker2d
8
20
700
8
20
800
8
5
20
500
med-replay
HalfCheetah
11
40
10
300
11
40
10
800
11
40
10
500
med-replay
Hopper
5
700
5
800
5
500
med-replay
Walker2d
13
2.5
47
1000
13
2.5
47
800
13
2.5
47
500
med-expert
HalfCheetah
100
5
1000
7
100
5
800
100
1100
med-expert
Hopper
40
43
600
40
43
800
12
40
43
500
med-expert
Walker2d
20
37
400
6
20
37
800
6
20
37
500
Table 5: Key hyperparameters of the proposed algorithms for each evaluation task in the ablation
study. K; ensemble size, reward penalty coefficient, H: rollout horizon, N: number of training epochs.
Table lists the key hyperparameters of BA-MCTS and its ablated variants used in the ablation study. For each task, an ensemble of K dynamics and reward models 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 RL training iterations. In particular, 50000 states are randomly sampled from the offline dataset, with each state followed by rollout lasting H/ time steps. To mitigate overestimation, a reward penalty based on the discrepancy among the ensemble members is applied with a coefficient X, shown in Equation (5). The setups for K, x, and H almost the same across the three algorithms and primarily inherited from the baseline "Optimized" (Lu et al. 2022). to make sure the improvements are brought by the Bayesian RL and deep search components (rather than tricky hyperparameter setups).
sampled by inferacting with the simulator, followed by 1000
The policy is evaluated on the ground truth environment for 10 episodes at the end of each training epoch. For all tables involving benchmarking results on D4RL, we report the average scores across the final 10 training epochs of our algorithms. 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 [u et al.) (2022) 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 Table 9 which are from their original papers, could be further improved by fine-tuning the relevant hyperparameters.
BA-MCTS and BA-MCTS-SL utilize Continuous BAMCP (i.e., Algorithm T to collect samples for subsequent policy distillation. Instead of performing tree search at every state, domly select a proportion (i.e., p) of states from the available 50000 states at each rollout time step root nodes for tree search to decide on the optimal actions. For the remaining states, actions are sampled directly from the policy, i.e., a ~ T-1s)). The tree search proce- dure detailed in Algorithm []] with the number of MCTS iterations, E, set 50. Increasing and E' can potentially enhance performance, but it will also linearly increase the compu- tational cost. Table [ outlines the key hyperparameters related to the search process for each algorithm and task. (1) As described in Algorithm [ the parameters and / control the rate of double progressive widening. (3 is set as 0.5 across all tasks.) To encourage deeper search, we limit the number of actions sampled from state under n. and the number of next states sam- pled from an action under n,, respectively. Action selection follows the UCT rule, as discussed in Appendix A where > 0 balances the exploration and exploitation. Additionally, inspired by the success of MuZero in enhancing exploration, we introduce Dirichlet noise F. at the root nodes, where actions are sampled from mixture of distributions:
aad + (1 (-1(s
23

---

BA-MCTS
BA-MCTT-SL
Data Type
Environment
a
The
T1
71
Th
Nsr.
Np
random
HalfCheetah
0.1
0.5
2.5
0.3
20
0.5
2.5
20
random
Hopper
0.1
0.5
2.5
0.3
1
20
0.1
0.5
2.5
0.3
1
20
random
Walker2d
0.1
0.5
2.5
0.3
20
0.5
2.5
0.3
20
100
medium
HallCheetah
0.1
0.5
I.0
10
0.5
1.0
5
medium
Hopper
0.1
0.5
20
0.5
1.0
20
medium
Walker2d
0.1
0.5
2.5
20
2.5
20
100
med-replay
HalfCheetah
0.1
0.8
1.0
0.3
10
0.1
0.8
2.5
0.3
20
med-replay
Hopper
0.1
0.8
1.0
0.1
20
0.8
1.0
20
med-replay
Walker2d
0.1
0.8
2.5
0.3
20
0.1
0.8
2.5
0.3
20
200
med-expert
HallCheetah
0.1
0.8
1.0
0.3
10
0.1
0.8
2.5
0.3
20
med-expert
Hopper
0.1
0.8
1.0
0.3
1
20
0.1
0.8
1.0
0.3
1
20
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
Table 6: Important hyperparameters used in the search process.
and 17 controls the mixture rate. Notably, for (c, 7), n2 nz), we explore the set of possible com- binations: ((2.5, 0.3, 20, 1). 0.3, 20, 1). (1.0, 0.1, 20, 1), (1.0, 0.3, 10, 5), (1.0, 0.1, 5)) dur- ing 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, incorporat- ing buffer of samples from the past NsL epochs helps to stabilize the learning process. Also, in the Walker2d environment, BA-MCTS-33 SL requires a warm-up up training phase Np epochs using BA-MBRL. allowing the initial policy to generate effective signals for supervised learning.
E DETAILS OF THE TOKAMAK CONTROL TASKS
STATE SPACE
BR Internal Inductance, Line Averaged Density,
Scalar States
Loop Voltage, Stored Energy
Electron Density, Electron Temperature, Pressure, Safety Factor, Ion Temperature, Ion Rotation
Profile States
ACTION SPACE
Targets
Current Target, Density Target
Elongation, Top Triangularity, Bottom Triangularity, Minor Radius,
Shape Variables
Radius and Vertical Locations of the Plasma Center Power Injected, Torque Injected, Total Deuterium Gas Injection, Total ECH Power, Magnitude and Sign of the Toroidal Magnetic Field
Direct Actuators
Table 7: The state and action spaces of the tokamak control tasks.
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, releas- ing energy in the process. The primary challenge of fusion is confining plasma, i.c., an ionized gas of hydrogen isotopes, while heating it and increasing its pressure 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 force on the moving particles (Pironti & Walker 2005).
Char et al. (024) trained a deep recurrent network as a dynamics model for the DIII-D tokamak, 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 one-second ramp-up phase, a multi-second flat-top phase, and 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 demon- strate that the learned model predicts these measurements for entire shots with remarkable accuracy. Thus, this model "ground truth" simulator for tokamak control tasks.
24

---

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 outlined in Table 7 For detailed phys- ical explanations of their components, please refer to Abbatt et al. (:02T); Char et al. (2023); Ariola et al. (2008). The state space consists of five scalar values and six profiles which discretized measurements of physical quantities along the minor radius of the toroid. After applying principal component analysis (Maćkiewicz Ratajezak] 1993), the pressure profile reduced two dimen- sions, 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 con- sists 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 apply any domain-specifie 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 in as targets for three tracking tasks. Specifically, 8n 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, step, and the reward is normalized by the episode horizon (i.e., 251 time steps). Notably, for policy learning, the reward function provided rather than learned from the offline dataset as in D4RL tasks; and the dataset (for offline RL) does not include the reference shot or any nearby, similar shots, making pure imitation infeasible.
3r) each time
F ADDITIONAL EVALUATION RESULTS
This appendix provides detailed empirical results and analyses to supplement the findings presented in Section 5 Appendix offers both quantitative and qualitative analyses of the benefits of belief adaptation, presenting additional results for Section 5.1
The following appendices report additional experimental results for benchmarking on MuJoCo D4RL tasks (i.e., Section
Appendix
Evaluates our algorithm against a suite of model-free offline RL. methods.
Appendix Presents a comprehensive comparison against 19 offline RL. baselines, in-
cluding statistical significance analysis.
Appendix F44 Provides detailed comparison with MuZero and its variants, focusing on their algorithm design, empirical performance, and computational cost.
Finally, we provide complete results for the ablation study (i.e., Section 5.3
Appendix F.5 Demonstrates the significant performance enhancement achieved by the belief adaptation module and our MCTS-based planner.
Appendix F.6
Discusses the trade-off's between the two policy distillation mechanisms investigated (i.e., SAC and supervised learning).
Appendix E.7 Presents an ablation study on the necessity of using reward penalty in offline
MBRL.
F.1 QUANTITATIVE AND QUALITATIVE ANALYSIS OF THE EFFECTIVENESS OF BELIEF ADAPTATION
To understand why belief adaptation over an ensemble of models using a Bayes-Adaptive MDP (BAMDP) is effective, we analyze its impact the prediction accuracy of the learned world models.
25

---

1.0
1.0
0.8
0.8
Bellef
0.6
Belief
0.6
0.4
0.0
20
100
4
10
Time Step
Time Step
(a) Offline rollout
(b) Imaginary rollout (at step 0)
1.0
1.0
0.8
0.8
Belief
0.6
Belief
0.6
0.4
0.4
0.2
0.2
0.0
10
Time Step
Time Step
(c) Imaginary rollout (at step 22)
(d) Imaginary rollout (at step 495)
Figure 1: Belief adaptation during offline and imaginary rollouts. (a) shows the belief over twelve ensemble members. adapting an offline trajectory of Hopper-med-epper, (b), (c), and (d) illustrate the belief changes during imaginary rollouts which start from different points of the offline trajectory shown in (a).
The Bayesian adaptation, defined in Eq. uses observed state transition sequences to adjust the belief over each ensemble member, thereby enhancing the quality of predicted rollout trajectories.
For each benchmarking task in D4RL MuJoCo. compute the average transition likelihood (i.e., Pss'ss, a (rlss, a)) for the provided offline rollouts, comparing cases with and without belief adap- tation. The ratios of these likelihoods are presented in Table Results show that the offline rollouts, collected from real environments, are more likely under the adapted ensemble. This demonstrates that Bayesian belief adaptation serves as effective calibration method for the learned world mod- els.
he-med-expert
hc-med-replay
he-medium
he-random
0.7515
2.6315
2.2674
2.4561
hp-med-ere -e-peee-xppert
hp-med-replay
hp-medium
hp-random
14.064
4.2214
3,5913
1.7400
wk-med-expert
wk-med-replan
wk-medium
wk-random
2,0898
3,8169
34.264
1.0341
Table 8: Ratios of average transition likelihoods in offline data with and without belicf adaptation across ensemble members. A ratio greater than indicates that the adapted ensemble assigns higher likelihood to the true transitions.
Furthermore, we evaluate the quality imaginary rollouts generated by the learned models. Ta- in the main content presents the prediction errors for next states and rewards in these rollouts. The results reveal that the adaptive ensemble achieves more accurate predictions in out of 12 benchmarking tasks, indicating that belief adaptation leads to more accurate model for planning.
As a qualitative analysis, we plot the belief adaptation over an offline trajectory (from D₂)) in Hopper-med-experr in Figure [(a) Initially, all twelve ensemble members have the same belief. As the trajectory progresses, the beliefs of each member are updated based on the transition history, with the dominant model (the one with the highest belief) continuously changing. Figure T(b) i(d)
26

---

further show the belief adaption during imaginary rollouts that start from different time steps of that offline trajectory.
F.2
COMPARISON WITH MODEL-FREE METHODS ON D4RL MuJoCo
Data Type
Environment
BA-MCTS (ours)
CQL
BEAR
BRAC-v
SAC
BC
random
HalfCheetah
41,45 (1,40)
35.4
25.1
31.2
30.5
2.1
random
Hopper
31.42 (0.10)
10.8
11.4
12.2
11.3
1.6
random
Walker2d
21.54 (0.06)
7.0
7.3
1.9
4.1
9.8
medium
HalfCheetah
(7.31(1140)
44.4
41.7
46.3
-4.3
36.1
medium
Hopper
103.9 (0.33)
86.6
52.1
31.1
0.8
29.0
medium
Walker2d
88.23 (0.78)
74.5
59.1
81.1
0.9
6.6
med-replay
HalfCheetah
71.24 (2.65)
46.2
38.6
47.7
-2.4
38.4
med-replay
Hopper
106.4 (0.53)
48.6
33.7
0.6
3.5
11.8
med-replay
Walker2d
91.42 (1.54)
32.6
19.2
0.9
1.9
11.3
med-expert
HalfCheetah
102.3 (2.27)
62.4
53.4
41.9
1.8
35.8
med-expert
Hopper
111.8 (0.82)
111
96.3
0.8
1.6
111.9
med-expert
Walker2d
116.0 (1.49)
98.7
40.1
81.6
-0.1
Average Score
80.25
54.85
39.83
31.44
4.13
25.07
Table 9; Comparisons between the proposed algorithms and model-free offline policy learning meth- ods on the D4RL benchmark suite. Each value represents the normalized score, proposed in Fu er al. (0200), of the policy trained by the corresponding algorithm. These scores are undiscounted returns normalized to approximately range between and 100, where score of corresponds to random policy and score of 100 corresponds to expert-level policy. For our algorithms, report the average score of the final ten policy learning epochs and its standard deviation across three different random seeds.
Here, we compare our algorithms with series of model-free offline policy learning (Chen et al. 2024) methods in Table 9 We include SOTA model-free offline RL. methods: CQL (Kumar et al. 2020), BEAR Kumar et al. (2019), and BRAC-v (Wu et al.) 2019). Additionally, we show the performance of directly applying SAC Behavioral Cloning (BC) (Chen et al., (024) to the provided offline dataset in the last two columns. The mean performance of the baselines are taken from related works YYu al. 2020 Kidambi et al. 2020 Fu al. 2020). Our algorithm shows significantly better performance, demonstrating the necessity of model- mosed--baed learning in these envi- ronments.
F.3 COMPREHENSIVE BENCIMARKING ON D4RL MUJoCo
To firmly establish the state-of-the-arr (SOTA) performance of our algorithm, this section provides comprehensive comparison of BA-MCTS against 19 offline RL methods on the 12 D4RL MuJoCo tasks. To provide context, we briefly categorize the baselines:
The model-based (MBRL) methods include classic approaches like MOPO Yu et al, 2020), MOReL (Kidambi et al. 2020), and COMBO (Yu et al. 2021), as well as more recent/SOTA methods such as MOBILE (Sun et al. 2023). CBOP (Jeong et al. 2023). RAMBO (Rigter et al. 2022), and MAPLE (Chen al. 2001Tb)
The model-free (MFRL) methods cover wide range of technigues. We include promi- nent SOTA methods such as the value-based algorithm CQL (Kumar et al. 2020). the implicit Q-learning method IQL. (Kostrikov et al., [001)), ensemble-based methods like SAC-N and EDAC (An et al. 2021), the policy constraint method TD3+BC (Fujimoto & Gu 2021b, and the sequence modeling approach Decision Transformer (DT) (Chen et al. 2021a). We also compare against classic MFRL baselines, including BEAR (Kumar et al. 2019), BRAC-v (Wu et al.) 2019), and the direct application of SAC and Behavioral Cloning (BC) (Chen et al. 2024). Finally, APE-V V (Ghosh et al.) 2022) is included as a conceptually related model-free method that also leverages the BAMDP framework.
27

---

500
-200
-250
-400
Number of Samples Number of Samples
Tec
Number of Samples
166
Number of Samples
1e6
(a) ho-med-expert
(b) re-mee-replay
(c) he-medium
200
100
Return
155
-100
80
50
500
Number Samples
le6
Samples
Number Samples
(d) hc-random
(e) hp-med-expert
(1) hp op-med-replay
150
Return
Return
Return
100
0.0
3.5
1.0
15
2.0
2.5
30
0.0
0.2
0.4
0.6
0.8
1.0
Number
Samples
Number of Samples
Number Samples
(g) hp-medium
(h) hp-random
wk-med-exppert
300
103
Return
Return
Return
Am Mame
200
100
2
3 1e5
0.6
0.8
1.0
Number of Samples
1e5
Number of Samples
IeB
( we-meddreplay
(k) wk-medium
(l) wk-random
Figure 2: Performance of Sampled EfficientZero D4RL MuJoCo tasks. The results for HalfChee- tah, Hopper, and Walker2d are presented in the three rows, respectively. Each subfigure depicts the change undiscounted episodic return function of the number of training samples. Experiments are repeated three times with different random seeds, with the solid line representing the mean and the shaded indicating the 95% confidence interval. For reference, the expert-level episodic turns (corresponding to scores of 100) for HalfCheetah, Hopper, and Walker2d are 12135, 3234.3, and 4592.3, respectively.
The results, including statistical significance analysis using Cohen's d (Cohen 1988 are sum- marized in Table in the main content. Note that MOReL, IQL, and the model- froee--free offline RL baselines in Table did not report standard deviations. Additionally, IQL and Decision Transformer (DT) were not evaluated on D4RL tasks with random data.
shown in Table 3. BA-MCTS achieves the highest average score among 20 algorithms, both including and excluding the challenging 'random' datasets. To quantify the significance of this im- provement, we compute Cohen's 's against all baselines for which standard deviations are available. The results show that the Cohen's 's d value is greater than 1.8 in all cases, and for most baselines, it is significantly larger (e.g.. > 3.5). This indicates that the performance improvement of our algorithm not only substantial in magnitude but also statistically significant. This comprehensive compari- son provides strong evidence that our proposed method establishes a new state-of- of-t-e--rrtt in offline RL. on this widely-used benchmark.
28

---

BA-MCTS
BA-MCTS
BA-MBRL (ours)
Sampled
Data Type
Environment
-SL (ours)
(ours)
EfficientZero
random random
HalfCheetah
11.2 (2.6)
14.6 (2.2)
1.2 (0.4)
54.8 (2.6)
Hopper
48.7 (2.4)
65.2 (1.9)
50.)).8)
153.7 (19.4)
random
Walker2d
25.7 (2.0)
38.1 (1.1)
(12
175.7 (28.2)
medium
HalfCheetah
(3.0)
(1.3
(..)
54.7 (1.9)
medium
Hopper
(4.8)
67.5 (0.2)
(.2)
70.5 (1.0)
medium
Walker2d
24.0
33.3 (4.0)
(1.0
63.7 (1.5)
med-replay
HalfCheetah
24.7 (0.9)
21.2
2.1(0.4)
54.8 (1.8)
med-replay
Hopper
17.9 (0.0)
(1.2)
(0.2)
126.6 (12.0)
med-replay
Walker2d
40.3 (7.2)
77.0
8.3))
134.9 (4.2)
med-expert
HalfCheetah
32.6 (0.1)
12.6 (0.8)
(1.1
55.4 (1.3)
med-expert Hopper
61.5 (4.6)
76.9 (11.6)
5.3 (0.5)
66.1 (0.9)
med-expert
Walker2d
35.0 (6.5)
55.9 (1.8)
2.9 (0.9)
60.8 (1.8)
Table 10: Training time (in hours) of our algorithms and Sampled EfficientZero for each evaluation task. Results are presented as the mean and standard deviation from three repeated experiments.
F.4 COMPARISON WITH MUZERO AND ITS VARIANTS
Monte-Carle Tree Search (MCTS) Browne al. 2012) has been successfully integrated with RL, as exemplified by AlphaZero (Silver et al.J 2017) and MuZero (Schrittwieser et al.) (020). These methods have achieved superhuman performance in domains requiring highly sophisticated decision-making processes. AlphaZero relies on a given world model, 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.] 2021 Schrittwieser et al. 2021 Ye et al.) 2021 Antonoglou et al.| 22022] Oren al. 2022 Zhao et al. (2224,, most are designed for online MBRL. According to Niu et al. (023), applying MuZero to offline learning, especially for continuous control in highly stochastic environments, remains a significant open challenge.
Our algorithm design differs from MuZero (and its variants) in several key aspects: (1) MuZero integrates model learning and policy training into single stage, using a world model defined in latent state space. Our algorithm separately learns world model and then trains the policy top of it, aligning with the widely-adopted offline MBRL framework. (2) MuZero employs a single latent world model (rather than an ensemble) and does not explicitly account for uncertainty in dynamics or reward predictions. This makes it particularly vulnerable to modeling errors in offline MBRL. (3) We introduce double progressive widening (Auger al. (203)) and aayes-adaptive planning into MCTS, resulting in a fundamentally different planning algorithm compared to the one used in MuZero.
To evaluate the performance of MuZero-type methods on D4RL MuJoCo tasks, we use the open- source implementation and hyperparameter configurations of Sampled EfficientZero et 2021) provided by LightZero (Niu et al. 2023). Benchmarking results from LightZero indicate that Sam- pled EfficientZero achieves the best performance on (online) MuJoCo locomotion tasks compared to other MuZero variants. To adapt it for offline learning, we employ the reanalyse technique proposed by Schrittwieser al. (202T).
The evaluation results are presented in Figure As shown, the performance is significantly worse compared our methods listed in Figure 3 despite Sampled EfficientZero's higher computational cost. For reference, the expert-level episodic returns (corresponding to scores of 100) for HalfChee- tah, Hopper, and Walker2d are 12135, 3234.3, and 4592.3, respectively.
In Table TO, we report the training time of our proposed algorithms and Sampled EfficientZero on the D4RL MuJoCo tasks. The experiments were conducted on server with 40 Intel(R) Xeon(R) Gold 5215 CPUs and 4 Tesla V100-SXM2-32GB GPUs. Our search-based algorithm, BA-MCTS, requires considerably less training time than Sampled EfficientZero to achieve its superior perfor- mance. While BA-MCTS requires more computation than the non-search variant BA-MBRL. this extra cost is limited to the offline training phase; no MCTS is performed during deployment, ensur- ing real-time execution.
29

---

The performance gap may be attributed to the challenges of applying MuZero's design to the offline setting. World model learning is particularly difficult when the state-action space is vast but training data is limited. Sampled EfficientZero integrates model learning and policy training into single stage, which can significantly increase the learning difficulty compared to our decoupled approach.
F.5 ABLATION STUDY ON THE ENHANCEMENT FROM BELIEF ADAPTATION AND MCTS-BASED PLANNING
Data Type
Environment
Optimized
BA-MBRL
BA-MCTS
BA-MCTS- SL
random
HalfCheetah
31.7
22.76(1.16)
36.23 (1.04)
29.20 (2.00)
random
Hopper
12.1
31.47 (0.03)
31.56 (0.12)
33.83 (0.10)
random
Walker2d
21.7
21.45 (0.53)
21.59 (0.32)
21.89 (0.07)
medium
HalfCheetah
45.7
(5.20)
(3.811
70.47 (3.52)
medium
Hopper
69.3
98.25 (3.42)
96.70 (14.0)
97.75 (7.09)
medium
Walker2d
79.7
75.41 (4.17)
74.73 (3.25)
82.84 (1.85)
med-replay
HalfCheetah
58.0
62.50 (0.18)
65.45 (0.81)
61.16 (1.60)
med-replay
Hopper
90.8
93.91 (4.25)
101.8 (3.46)
106.3 (0.13)
med-replay
Walker2d
65.8
97.54 (1.93)
95.06 (2.11)
92.13 (5.13)
med-expert
HalfCheetah
104.2
90.52 (4.13)
76.16 (10.3)
80.53 (6.63)
med-expert
Hopper
105.8
107.8 (0.37)
108.3 (0.22)
112.2 (0.29)
med-expert
Walker2d
97.1
84.71 (0.87)
110.0 (L.74)
107.7 (0.82)
Average Score
65.2
71.06
74.45
74.62
Table 11: Performance comparison among our algorithm: BA-MCTS, ablated versions: BA- MBRL and BA-MCTS-SL, well the baseline: Optimized (Lu al. 2022). To ensure prin- cipled ablation study, we reimplement BA-MCTS and its ablated variants using the codebase of Optimized, which systematically investigates design choices in offline MBRL. We make only mini- mal modifications to the code and hyperparameter settings. This approach ensures that any observed performance improvements are attributable to algorithmic design, rather than implementation dif- ferences or tuning tricks.
For principled ablation study, we reimplement BA-MCTS and its ablated variants based on the codebase of Optimized et. 20.22, which thoroughly explores design choices offline MBRL, making minimal changes to the code and hyperparameter settings (see Appendix DD. The compari- son results are presented in Table and the training curves of BA-MCTS and its ablated variants shown in Figure
Firstly, by comparing the following two methods, any observed performance difference can be directly attributed to the benefits of belief adaptation in BAMDP formulation.
Optimized: A strong, non-adaptive offline modell-ased RL baseline.
BA-MBRL: Our ablation variant, which is minimal modification of Optimized, differing
only by the integration of our BAMDP framework that dynamically adapts the belief over the model ensemble via Bayesian updates (Eq.
As demonstrated, BA-MBRL outperforms the Optimized baseline on 8 out of 12 tasks and achieves higher average score, providing direct evidence that Bayesian belief adaptation leads to improved final policies.
Secondly, by comparing the following two methods, any observed performance difference can be directly attributed to the benefits of incorporating MCTS-based planning:
BAA-BBL:: Our ablation variant that incorporates the BAMDP belief adaptation frame- work but does not use MCTS. Its policy improvement relies solely on Soft Actor Critic.
5 'The performance of BA-Me MCTS decreases, relative to the results presented in Table 2 for two main reasons: (1) The original implementation in Table is based on a diffferntt codebase than the one used in Optimized; and (2) for the ablation results shown in Table [1 we adopt the reward penalty design from Optimized to enable ablation comparisons, rather than using the one described in Eq.
30

---

BA-MCTS: Our final algorithm, which builds upon BA-MBRL by integrating the Contin- uous BAMCP deep search planner.
This substantial gain highlights the effectiveness of using Continuous BAMCP as powerful policy improvement operator. While BA-MBRI relies on the implicit policy improvement provided by the temporal difference learning process, BA-MCTS leverages deep lookahead search to explicitly find stronger, non-parametric policy (Trt returned by Algorithm at each decision point. The actor- critic learner is then trained on the high-quality trajectories generated from this improved policy. This "RL+Search" paradigm, where the learner distills knowledge from stronger planner, is the primary reason for the observed performance leap. This result strongly supports our claim that the deep search component is crucial for achieving state-of-the afrthe-art performance in the offline model- based setting.
F.6 COMPARISON OF POLICY DISTILLATION MECHANISMS
This section addresses the research question: How can the search outcomes (ie.. Tha) be effec- tively used for policy improvement? We investigate this by comparing two different mechanisms for distilling the knowledge from the MCTS planner back into the parameterized policy network æ.
BA-MCTS: Our primary algorithm, which uses a policy gradient method (i.e., SAC) to update the policy network using trajectories generated by the planner.
BA--CTT-SSS An alternative variant where the policy network is updated via supervised learning (SL) to directly mimic the planner's 's output distribution Nea by minimizing cross- entropy loss.
Performance Comparison: As shown in Table 11, BA-MCTS-3 SL performs similarly to BA-MCTS in terms of the final average score, validating the effectiveness of both policy update mechanisms. This indicates that both policy gradient and supervised learning are viable methods for distilling the results of the deep search. However, note that BA-MCTS-SL can be less stable on certain complex tasks; for example, it struggles on the Walker2d environments, where a warm-up training phase (using BA-MBRL) is required establish good initial policy for the supervised learning to be effective.
Training Stability and Robustness: On the other hand, a key advantage of the SL.-based policy update is its training stability. As evident in the training plots (Figure B. BA-MCTS-SN often exhibits much smoother learning curves compared to the other two algorithms, indicating greater robustness in model selection for deployment. Furthermore, as shown in our ablation study on the reward penalty (Appendix F.7). we found that the SL-based policy update less sensitive to the degree of pessimism.
F.7
ABLATION STUDY ON THE REWARD PENALTY
BA-MCTS
BA-MCTS
BA-MCTS
BA-MBRL.
Data Type
Environment
BA-MCTS
BA-MBRL
-SL
SL
(A= 0)
0)
random
HalfCheetah
(2.000
(1.24
(1.16)
(1.39)
(1.65)
39.64
random
Hupper
33.83 (0.10)
31.56 (0.12)
31.47 (0.03)
9.16 (0.16)
(0.14)
6.97 (0.07)
random
Walker2d
21.89 (0.07)
21.59
21.45 (0.53)
(6.16)
21.53 (0.42)
21.1 (0.64)
medium
HallCheetah
70.47 (3.52)
(5.24) (3.81)
61.64 (4.58)
60.84 (2.00)
(2.29)
medium
Hopper
97.75 (7.09)
96.70 (14.0)
98.25 (3.42)
102.8 (2.29)
104.4 (1.88)
93.68 (11.4)
medium
Walker2d
82.24 (1.85)
74.23
(4.17)
82.61 (0.86)
57.01 (7.24)
57.97 (15.4)
med-replay
HalfCheetah
61.16 (1.60)
65.45 (0.81)
42.10 (0.18)
36.65 (2.39)
44,03 (7.35)
med-replay
Hopper
106.3 (0.13)
101.8 (3.46)
93.91 (4.25)
107.9 (0.07)
84.11 (2.97)
91.81 (11.5)
med-replat
Walker2d
92.13 (5.13)
95.06 (2.11)
(1.93)
88.61 (5.21)
97.33 (3.51)
98.19 (1.23)
med- med-expert
expert
HalfCheetah
80.33 (6.63)
76.16
(5.36) (4.13)
(1.46)
29.88 (2.28)
med-exped
Hopper
112.2 (0.29)
108.3 (0.22)
107.8 (0.37)
106.8 (6.34)
81.76 (6.45)
86.79 (18.7) 7)
mod-expe exp-exxert
Walker2d
107.7 (0.82)
110.0 (1.74)
84.71 (0.87)
110.8 (1.72)
110.2 (0.91)
53.35 (38.0)
Average Score
74.62
74.45
71.06
68.04
60.55
$5.43
Table 12: Comparison of the proposed algorithms with their corresponding versions without the reward penalty (i.c., À = 0). The definitions of the scores in this table are consistent with those in
Table g
31

---

12000
0010
10000
8000
6400
6030
aggrrttm
agorith m
4000
algoritiam
4000
BA-NBRL
SA-MAAI BA-HCTS
DA-MBRI GAwiii To DA GN-MCTS
2000
BA-MCT5 CL
SA MCTS SL
Epoches
Training Epeches
Epoches
(a) he-med-expert
(b) ne--mayreplay
(c) he-medium
2000
3000
3000
2000
E200
2000
algositbe
algnithh n
algoniti
84-MBRL
BA Hitt BA-MCIS
1000
Bi NBRL
1000
BA-METS
BA NCTS
MCTS
Ba MCTS SL
SA-NCIS-5L
200
400
603
800
200
430
600
810
Trinnin Epoches
Taining Epeches
Tairing Epeches
(d) hc-random
(e) hp-med-expert
(1) hp hp-med-replay
1200
5050
3000
1.000
4000
2000
3000
algortthm
aigortthm
B4- UJ
400
B4-MERL
200
BA-METS
2010
MCTS
430
600
BI0
200
400
603
800
430
600
610
training Epaches
Thiiiig Epoches
iraiting Epeches
hp-medium
(h) hp-random
wk-med-experr
4030
3000
4000
3000
800
3000
Baturi
600
2000
algerthmm BC MERL
algortthm
BA NBRL
400
1060
B4 MCTS
Bà MCTS
SA-MC13
B4-0CC1--1
200
DM-9ICT--5
800
Taining Epoches
Traniing Epoches
Taining Epochers
wk-med-replay
(k) wk-medium
(l) wk-random
Figure 3: Performance of our proposed algorithms on D4RL MuJoCo tasks in the ablation study. The results for HalfCheetah, Hopper, and Walker2d are presented in the three rows, respectively. Each subfigure depicts the change in the undiscounted episodic return as function of training epochs. Experiments resenting the mean and the shaded area indicating the 95% confidence interval. For reference, the expert-level episodic returns for HalfCheetah, Hopper, and Walker2d are 12135, 3234.3, and 4592.3, respectively. Note that the training epochs for each algorithm, as listed in Table have been linearly scaled 800 for better visualization.
repeated three times with different random seeds, with the solid line rep-
To demonstrate the necessity of incorporating the reward penalty in offline MBRL, we conduct an ablation study by setting
Eq. (5) to resulting in ablated versions of our proposed three al- gorithms. The results are presented in Table [2] First, the average performance of the algorithms with the reward penalty is consistently better, demonstrating the importance of using reward penal- ties in offline MBRL to prevent the overexploitation of the learned world models (which can be inaccurate). Second, the sppervised-leemiig--bbe affected by the absence of the reward penalty, compared to the seiicy-gradieet-baase methods. No- tably, BA-MCTS- SL (A = 0) and BA-MCTS-SL achieve comparable performance the Hopper and Walker2d tasks. This shows additional advantage of BA-MCTS-SL. its reduced sensitivity to model inaccuracies. Lastly, there are instances where superior performance achieved with set O. For example, BA-MCTS-SL (A 0) performs better than BA-MCTS-SI. in out of 12
algorithm (i.e., BA-MCTS-SL. (A 0)) less
32

---

tasks. This suggests that the performance of our algorithms in Tables 11 could be further improved by adjusting hyperparameters such as A.
G STATEMENT ON THE USE OF LARGE LANGUAGE MODELS
In the spirit transparency and accordance with the ICLR disclosure policy, we detail the use of Large Language Models (LLMs) in the preparation of this manuscript.
LLMs were utilized as an assistive tool in two main capacities: language refinement and technical support during the formalization of proofs. The specific roles were:
Language and Readability Enhancement: We employed LLM polish the manuscript's language. This included improving grammar, rephrasing sentences for bet- ter clarity, and ensuring a consistent and professional tone throughout the paper.
Assistance with Theoretical Proofs: An LLM provided support in the technical writing and verification of our theoretical proofs. This assistance was twofold:
Generating IЯTК source code for complex mathematical equations.
Acting as a verification aid by reviewing the logical steps our derivations to help identify potential inconsistencies or errors.
It is crucial that the LLM did contribute the original ideation, conceptual development, or the core derivations of the proofs themselves. The intellectual ownership and the scientific rea- soning behind all theoretical results remain entirely with the authors. The authors have meticulously reviewed, validated, and take full and final responsibility for all content, including the correctness and rigor of the proofs that were checked with L.M assistance.
33


Reviews:
[{"review_id": "qIPt5FfZPO", "reviewer": "Reviewer_jrky", "summary": "This paper introduces a novel framework for offline model-based reinforcement learning, which (1) trains an ensemble of models based on data from a behavior policy and builds a Bayes Adaptive MDP (BAMDP) of this ensemble (i.e., an MDP with uncertainty, where the belief about the transition and reward distributions is updated based on observed transitions); (2) introduces a shaped reward for the Bayes Adaptive MDP that penalizes actions with high uncertainty, called Pessimistic Bayes-Adaptive MDP; (3) extends an existing online-planning algorithm for BAMDP, called BAMCP, to handle continuous states and actions based on Double Progressive Widening; (4) proves that this extension, called Continuous BAMCP, approximates the optimal policy; and finally, (5) introduces an algorithm for efficiently learning a policy from the planning outcomes of Continuous BAMCP.", "questions": "* Can you please explain why the Bayesian update works in the offline setting?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The paper introduces a number of novel ideas for offline model-based RL: formulating an ensemble of learned models as a Bayes Adaptive MDP, extending BAMDP to handle continuous state and action spaces, and distilling the planning outcomes into a policy. According to the numerical results, most of these ideas are helpful in practice.\n* The proposed extension of BAMDP, called Continuous BAMDP, is proven to converge.\n* The proposed framework is compared to offline RL baselines on benchmark environments and data, and the numerical results are promising. There is also an ablation study.\n* The paper is generally well organized and easy to follow.\n* The source code is available, and it supported by some documentation.", "weaknesses": "* My main concern is that treating the ensemble of learned models as a Bayes Adaptive MDP for offline RL does not seem to make sense conceptually. Bayes Adaptive MDP is perfectly reasonable when the beliefs are updated based on transitions that are sampled from the real environment; each transition can provide new information, which we can use to update our beliefs (i.e., probabilities or probability densities associated to particular MDPs). However, in the proposed framework, the transitions are sampled from the learned models. What new information do these transitions provide? If we set our prior belief to be a uniform distribution, and then uniformly sample transitions according to these models, why is our expected updated belief different from uniform? For a particular sampled transition, the updated belief will of course be non-uniform; but why would it be non-uniform in expectation? And if it is, then are the belief updates just driven by randomness (e.g., probability of the model that was sampled uniformly at random would likely be increased)? How is this belief update useful?\n* That being said, Table 1 in the main text does show that the Bayesian update significantly improves predictions for imaginary rollouts (compared to ground-truth rollouts in the real environments for the same actions sequences). So, for some reason, the Bayesian update is useful; but it is not obvious why. What is the explanation? Does the improvement hinge on the particular ensemble of neural-network models? Does it hinge on the characteristics of these neural networks? Does it give an advantage to neural networks whose output distributions are more concentrated? If yes, could the Bayesian approach be replaced with simply favoring such models? Or does the Bayesian update help by simply concentrating the beliefs to a very small support after a few steps (as shown by the numerical results), which is somehow better for the simulation? If yes, could the Bayesian approach be replaced by simply using one particular randomly chosen model for each episode?\n* In the end, the numerical results do support the proposed approach. But it would be good to clarify why it works. (Appendix F.1 is interesting, but does not address the core question.)\n* The paper includes a number of ideas, addressing various challenges. While the paper is overall well organized and easy to read, it might benefit from a clearer focus. For instance, the proposed framework would be novel and complete even without Continuous BAMCP; it would simply be restricted to discrete environments. The policy learning also seems somewhat optional since the framework could be presented as an online planning approach (based on model learned from offline experiences); according to the numerical results (Table 11 in the appendix), BA-MCTS and BA-MCTS-SL are pretty close, so the policy learning does not seem that important compared to the other parts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "39RxPX9WLt", "reviewer": "Reviewer_Gn3V", "summary": "This paper reframes Offline Model-Based Reinforcement Learning (MBRL) as a Bayes-Adaptive Markov Decision Process (BAMDP) to formally handle uncertainty over learned dynamics models.  \nIt introduces Continuous BAMCP, a Monte-Carlo Tree Search planner extended to continuous, stochastic domains using Double Progressive Widening (DPW), and integrates it into a policy-iteration framework named BA-MCTS.  Maintain a belief distribution over ensemble models and update it via Bayesian posterior rules.  Provide a *consistency theorem* guaranteeing exponential convergence of the planner’s value estimate. Empirically achieve convincing results on D4RL MuJoCo tasks and 3 Tokamak control tasks.", "questions": "- What is the computational cost (simulations per decision, wall-time) relative to other baselines?\n- Could Continuous BAMCP be applied in learned latent spaces to improve scalability?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. Novel idea of combining BAMDP formulation, MCTS planning, and offline MBRL into a unified algorithmic pipeline.\n2. Provides a **consistency proof** for Bayes-adaptive planning in continuous spaces.\n3. Bayesian belief updates give interpretable, adaptive weighting over model ensembles\n4. Reward-variance penalty connects naturally to pessimistic (safe) RL literature\n5. Strong empirical performance   \n6. Generalization: Demonstrates transfer to real-world-style stochastic control (Tokamak tasks).", "weaknesses": "1. **Computation cost:** Continuous BAMCP remains expensive, many simulations per state, not scalable to large or long-horizon domains.\n2. **Limited data efficiency analysis:** No study of how ensemble size or dataset quality affect performance.\n3. **Offline-only limitation:** Once trained, policy does not update beliefs from online experience.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bxDDXeBJ6x", "reviewer": "Reviewer_u5So", "summary": "This paper views offline RL as a Bayes Adaptive Markov Decision Process (BAMDP) for effectively addressing uncertainty associated with a finite dataset, and an associated planning algorithm based on MCTS that is then folded back into learning a policy using a policy search procedure. Empirical simulations show that the proposed algorithm outperforms baselines in offline RL benchmark suites.", "questions": "- One thing that wasn't clear was is the output of the learner just the policy? \n- What do the tradeoffs look like if we tried to use the search procedure even at inference time, somewhat mirroring MPC style approaches? \n- For me, it wasnt totally clear what was the net positive from trying to use search in this setup, as it does appear as if one can effectively solve the problem using a mix of pessimism and dyna type sampling, and perhaps this isnt totally clear just from the current exposition.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Using search within offline RL based policy search is pretty interesting as a general approach, and this paper makes progress towards marrying these ideas in a clean manner. That said, my knowledge of this area isn't up to date particularly with the more recent efforts and so I do not know if there are similar efforts in offline RL that I am unaware of.", "weaknesses": "- The writing at times appears to come in thick and fast, and in a sense i understand why this happens given the different sub-routines but it does help to work a bit more to help clarity in exposition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a5eUU5dlAk", "reviewer": "Reviewer_JBcC", "summary": "This paper frames offline MBRL as a BAMDP to address model uncertainty inherent in learning from offline datasets. The authors propose a novel Continuous BAMCP planning algorithm based on Monte Carlo Tree Search with double progressive widening for continuous domains. Integrating this with a policy iteration scheme, they present a framework termed BA-MCTS that distills the search/planning output into policy/value networks for efficient real-time decision making. Empirical evaluations are provided on D4RL benchmark, showing strong performance relative to several baselines.", "questions": "1.In Table 1, should the mean prediction errors be calculated as weighted averages based on the dimensions of states and rewards? Why does the Walker2d environment exhibit such large prediction errors in Table 1, while showing very small performance variance in Table 2?\n2.Why are the results reported in Table 2 inconsistent with those in Table 12?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.Integrating BAMCP with offline MBRL and extending it to continuous spaces via the introduction of DPW is well-motivated and conceptually sound.\n2.The paper includes well-executed ablation studies that systematically evaluate the individual contributions of the algorithm’s components, highlighting the advantages of the BAMDP formulation, MCTS planning, and the search-guided policy learning approach.\n3.Theorem 4.1 establishes a consistency guarantee for the planning algorithm, proving its convergence to the optimal value in the pseudo-pessimistic BAMDP framework.", "weaknesses": "1.Unclear and inconsistent notation for augmented states: The notation used to represent augmented states—such as $(s, h)$ versus $(s, b)$—lacks clarity, and the distinction between history, belief, and memory storage in tree nodes remains vague. This ambiguity, evident across equations, algorithm steps (e.g., Algorithm 1, Page 5), and explanatory text, may cause confusion, particularly for readers less familiar with BAMCP/POMCP.\n2.Insufficient contextualization within Bayesian and offline MBRL literature: The paper would benefit from a more thorough discussion connecting the proposed method to existing frameworks that model offline RL as a Bayes Adaptive MDP or epistemic POMDP, such as those presented in Ghosh et al. (2022) and Dorfman et al. (2021).\nGhosh et al. (2022): \"Offline RL policies should be trained to be adaptive.\"\nDorfman et al. (2021): \"Offline meta reinforcement learning – identifiability challenges and effective data collection strategies.\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)