Reflect-then-Plan: Offline Model-Based Planning through a Doubly Bayesian Lens
Authors: Jihwan Jeong
Xiaoyu Wang
Jingmin Wang
Scott Sanner
Pascal Poupart
Venue: ICML 2025 poster (2025)
Keywords: Offline reinforcement learning
Model-based planning
Bayesian inference
Bayesian reinforcement learning
Primary area: reinforcement_learning->batchoffline

Abstract: Offline reinforcement learning (RL) is crucial when online exploration is costly or unsafe but often struggles with high epistemic uncertainty due to limited data. Existing methods rely on fixed conservative policies, restricting adaptivity and generalization. To address this, we propose Reflect-then-Plan (RefPlan), a novel _doubly Bayesian_ offline model-based (MB) planning approach. RefPlan unifies uncertainty modeling and MB planning by recasting planning as Bayesian posterior estimation. At deployment, it updates a belief over environment dynamics using real-time observations, incorporating uncertainty into MB planning via marginalization. Empirical results on standard benchmarks show that RefPlan significantly improves the performance of conservative offline RL policies. In particular, RefPlan maintains robust performance under high epistemic uncertainty and limited data, while demonstrating resilience to changing environment dynamics, improving the flexibility, generalizability, and robustness of offline-learned policies.

Reflect-then-Plan: Offline Model-Based Planning through a Doubly Bayesian Lens
Jihwan Jeong' Xiaoyu Wang'
Jingmin Wang' Scott Sanner' Pascal Poupart"
Abstract
Most offline RL methods aim to learn conservative policy that stays close to the data distribution, thus steering away from high epistemic uncertainty.
Offline reinforcement learning (RL) is crucial when online exploration is costly or unsafe but often struggles with high epistemic uncertainty due to limited data. Most of existing methods rely on fixed conservative policies, restricting adaptiv- ity and generalization. To address this, we pro- pose Reflect-then then-Plan (RefPlan), novel doubly Bayesian offline model-based (MB) planning ap- proach. RefPlan unifies uncertainty modeling and MB planning by recasting planning as Bayesian posterior estimation. At deployment, it updates a belief over environment dynamics using real-time observations, incorporating uncertainty into MB planning via marginalization. Empirical results on standard benchmarks show that RefPlan signif- icantly improves the performance of conservative offiin RI. policies. In particular, RefPlan main- tains robust performance under high epistemic uncertainty and limited data, while demonstrat- ing resilience to changing environment dynamics, improving the flexibility, generalizability, and ro- bustness of offline-leamed policies.
While incorporating conservatism into offline learning has proven effective (Jin al.] 2021; Kumar et al.] [020), can result in overly restrictive policies that lack generaliz- ability. Most methods learn Markovian policy that relies solely on the current state, leading the agent to potentially take poor actions in unexpected states during evaluation. Model-based (MB) planning can enhance the agent's 's respon- siveness during evaluation (Sikchi et al. 2021: Argenson & Dulac-Aroold 2021] Zhan et al.
2022). but still primarily
addresses epistemic uncertainty through conservatism.
Noting this challenge, Chen et al. (2221) and Ghosh et al. (2022) propose to learn an adaptive policy that can reason about the environment and accordingly react at evaluation. Essentially, they formulate the offline RL problem as a partially observable MDP (POMDP)-whee) observability relates to the agent's epistemic uncertainty, aka Epistemic POMDP (Ghosh et al. (021). Thus, learning an adaptive policy involves approximately inferring the belief state from the history of transitions experienced by the agent and allowing the policy to condition on this belief state.
the partial
While learning an adaptive policy can help make the agent more flexible and generalizable, it still heavily depends on the training phase. Our empirical evaluation demonstrates that learned phethy-whether it be adaptive fixed can be significantly strengthened by incorporating MB planning. However, existing MB planning methods fall short in ade- quately addressing the agent's 's epistemic uncertainty, and it remains elusive how one can effectively incorporate the uncertainty into planning.
1. Introduction
Recent advances in offline reinforcement learning (RL) enable learning performant policies from static datasets (Levine et al. 2020 Kumar et al.] 2020), making it ap- pealing when online exploration is costly or unsafe.
The agent's inability to gather more experiences have severe implications. In particular, it becomes practically impossi- ble to precisely identify the true Markov decision process (MDP) with limited dataset, as it only covers a portion of the entire certainty for states and actions outside the data distribution.
We propose Reflect-then- (RefPlan), a novel doubly Bayesian approach for offline MB planning. RefPlan com- bines epistemic uncertainty modeling with MB planning in unified probabilistic framework, inspired by the control-as- inference paradigm (Levine) 2018). ReiPlan adapts Bayes- adaptive deep RL techniques (Zintgraf et al. 2020) Dorfman et al., 202T) to infer a posterior belief distribution from past experiences during test time (Reflect). To harness this uncer- tainty for planning, we recast planning as Bayesian posterior
space, leading to high epistemic un-
CCurrently at Google Research, Mountain View, CA. Unniversity of Tororto, Toronto University of Waterloo, Waterloo. Correspondence to: Jihwan Jeong <jihwanjeonn@e @ google.com>
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Reflect-then-Plang Offline Model-Based Planning through a Doubly Bayesian Lens
Reflect
Plan
Encoder
8t
{m/
jm1
Updating posterior belief the possible MDPs
p(T (0)
p(r|O, ma mm)p(m))ddie
Figure 1: Schematic illustration RefPlan. (Reflect) At time t, real-time experiences T:t (Sps ag .o.....5.) St) are used infer the posterior belief My over environments using variational autoencoder. Unlike prior methods, RefPlan learns diverse dynamics models conditioned on m,, capturing different transition and reward functions. (Plan) Offline planning is framed as probabilistic inference, where the posterior over optimal plans P(r[O) (O denotes optimality variables in the control-as- inference framework) is inferred. Importantly, prior P(T) is incorporated by learning via offline policy learning. By marginalizing out m,, RefPlan addresses epistemic uncertainty, enhancing my for better adaptivity and generalizability.
estimation (Plan). By marginalizing the agent's epis- temic uncertainty, RefPlan effectively considers range of possible scenarios beyond the agent's 's immediate knowl- edge, resulting in a posterior distribution over optimized plans under the learned model (Figure D.
[(b) Epl- (ls,,,,,,sssss a)). The learned model po then generates imaginary data Dmodel to aid offline policy learning.
Here, we focus on using learned models for rest-time plan- ning. MB planning typically employs model predictive control (MPC), where at each step, a trajectory optimiza- tion (TrajOpt) method re-plans to maximize the expected H H-step return under po., incorporating a value function Vi for long-term rewards (Lowrey et al. 2018):
In our experiments, we demonstrate that RefPlan can be in- tegrated with various offline RL policy learning algorithms to consistently boost their test-time performance in standard offline RL benchmark domains (Fu et al. 2020). RefPlan maintains robust performance under high epistemic uncer- tainty, demonstrating superior resilience when faced with diittiiiittiitttionn states, shifts in environment dynamics, or limited data availability, consistently outperforming com- peting methods in these challenging scenarios.
a argmax Ep (St am4+)),
(1)
En=o f acth) 7 (o($ is the return of a candidate sequence a H under P with Ag denoting the learned reward model.
where Rn (St, a14+H)
(Seths
2. Preliminaries
MPPI is a TrajOpt method that optimizes trajectories by sampling N action sequences, weighting them via a soft- max function with inverse temperature a based on re- turns (Nagabandi et al. 2019). The optimized action is afth ΣN exp(xRR;) ) where RD is the return of the (2)()/12))
We study RL in the framework of Markov decision processes (MDPs), characterized by a tuple MM- = AS, do, 7). The state (S) and action (A) spaces are continuous, T (s'[s, a) is the transition probability, r(s, is the reward function, do is the initial state distribution, and 7 [0, 1] is the discount factor. The model of the environment refers to the transition and reward functions. The goal of RL is to find an optimal policy n² which maximizes the expected discounted return, Ex (Eio (rr(st,
nith trajectory. MBOP (Argenson & Dulac-Arnodd 2021) adapts MPPI for offline settings by sampling actions from behavior-cloning (BC) policy with smoothing.
The control-as-inferencc framework The control-as- inference framework reformulates RL as a probabilistic inference problem (Levine 2018), introducing auxiliary binary variables O,, where O = denotes that (s,, a.) is optimal. The likelihood of optimality for a trajectory Tet+ (St, at...., 88+H) is:
Offline MB planning Given an offline dataset D = {(s,, a, rú. 1 collected by a behavior policy 3, model- based (MB) methods train predictive model ((sssrrrss,,) a) via maximum likelihood estimation (MLE), minimizing
2

---

Reflect-then-Plant Offline Model-Based Planning through a Doubly Bayesian Lens
Definition 1 (The optimality likelihood), For Te let O 0=11 1f 1fth =1 1h. Then,
unseen states or slightly altered dynamics during deploy- ment, leading arbitrarily poor performance.
To address these challenges, we provide a Bayesian take on the offline RL problem. enabling reasoning over the agent's s uncertainty through a prior belief by = p(M), updated to a posterior b. = p(M (Tt)) as new experiences Tür are gathered during deployment. However, computing the exact posterior belief is generally intractable. Therefore, in Section [ we tackle this challenge by approximating the belief distribution through variational inference techniques adapted from meta- RL approaches (Zintgraf et al. [0200 Dorfman et al. 2021).
p = (|T) x P = 1|scth: acth)-
(2)
h
Assuming p0 1, arth) x exp(x - reth), we obtain p(O]r) exp(k EATA Thus, expected return maximization reduces to posterior inference over trajectories given that all time steps are optimal under the probabilistic graphical model (PGM) (see Figure 2 left):
p(T(O) x p(T, O)
H II P(Ottt/sct auth (u4))p(s44+11see acth). (3)
3. RefPlan: a Probabilistic Framework for Offline Planning
p(si)
Typically, the prior over actions is assumed to be uniform or implicitly defined by rewards [Levine 2018 Piché et al. 2019). However, in Section 3.1 we explicitly model the action prior to formalize an offline MB planning framework, enabling policy refinement via MB planning-
In this section, we seek to address the following question:
How can we leverage a learned model at test time for enhancing an offline-trained agent?
To tackle this, we introduce RefPlan, a novel probabilistic framework for offline MB planning that allows agent to reason with its uncertainty during deployment. Our ap- proach is developed in three parts.
Epistemic POMDP, BAMDP, and Offline RL A par- tially observable MDP (POMDP) extends MDPs scenar- ios with incomplete state information. POMDPs can be reformulated as belief-state MDPs, where belief-a prob- ability distribution over states-repree represents the uncertainty over states given the agent's 's prior observations and actions KKelbling et al. 1998).
First, in Section 3.1 MB planning algorithm from the as-in--re-eee per- spective. This view crucial as provides principled foundation for our framework and formally justifies using an offline--tainee policy as a prior to guide planning. Next, in Section 3.2 we detail our approach for modeling the agent's 's epistemic uncertainty over the environment dynam- ies by adapting recent variational inference techniques from meta-RL. (Zintgraf et al., 2020: Dorfman et al.) 202T).
we derive a sampling- -bsed offline
Unlike an ordinary POMDP, epistemic POMDPs address generalization unseen test conditions in RL (Ghosh et al.) 2021). In this scenario, the agent experiences partial observ- ability entirely due to its epistemic uncertainty about the identity of the true environment M at test time. Specifically, at test time, the agent's goal is to maximize the expected return EM 7"r,] under the posterior p( (M) given the train data D. Thus, an epistemic POMDP is an instance of Bdyesvad MDP (BAMDP) (Dufi) 2002 Kaelbling et al. 1998) or Bayesian RL (Ghavanzadeh et al. 2015). Compared to standard BAMDP, epistemic POMDP puts focus on the agent's 's test time evaluation performance. For a thorough definition of BAMDPs, please check Ap- pendix A.I
While these individual components are built upon concepts from prior work, our primary conceptual novelty lies in their synthesis. In Section 3.3 we unify the probabilistic planning formulation with the learned epistemic uncertainty. This integration allows the agent to plan under the learned models and adapt in real-time while accounting for its uncertainty.
3.1. Offline Model-Based Planning as Probabilistic Inference
In this work, we view offline RL as epistemic POMDP, following Ghooh et al. (2022) drawing connections to Bayesian approaches. That is, limited coverage of the state- action space in the offline dataset induces epistemic uncer- tainty about dynamics beyond the data distribution. Failure to manage this uncertainty can result in catastrophic out- comes, particularly when an offline-trained agent encounters
We recast offline MB planning within the control-as- inference framework, treating planning as a posterior infer- ence problem. This approach enables the agent to optimize its actions by reasoning over the learned dynamics model and prior knowledge obtained from offline training. Central to this formulation is the use of a prior policy, which guides the agent's 's plans based on knowledge learned offline.
While we use exponential form for p(O) r). other mono- tonie functions possible (Okada & Taniguchi (020). We sim- plify notation by writing O, = as O2 and Tec+ f a8 T.
We start by formalizing the concept of a prior policy, which lay the basis for our Bayesian formulation of the offline MB
3

---

Reflect-then-P%ae Offline Model-Basee Planning through a Doubly Bayesian Lens
PGMs for the control-ast as inferreeeec framework, offline MB planning, and RefPlan. (Left) States evolve within the learned model, with actions and states influencing optimality. Optimality variables act like observations in hidden Markow model, framing planning as inferring the posterior over actions given optimality. (Middle) In offline MB planning, actions follow the prior policy mp- a ~ (-ps) 9). (Right) RefPlan uses past experiences ry to infer m. the agent's belief about the environment, and computes the expected optimal action sequence by marginalizing over ma.
Figure 2:
planning process.
Then, we can estimate the expected value of an arbitrary function f(au++ under p(r(O). That is,
Definition (Prior policy). A prior policy Tp S + P(A) a policy learned from an offline RL algorithm E using the dataset DD
J, (aaatH) p(r|O) (()++)))(((()))))))))))))))))))) dr
E Eptrio)
The prior policy, parameterized by 0, is provided by an offline learning algorithm Σ, such as CQL (Kumar et al. 2020) or BC, and must be considered by the offline MB planner when optimizing the planning objective in (T).
- ((aret)))
a- exexp
K6 En 7(th)
p(T) dt
p(O)
Ep(t) [((eect))
[x: En xnETieth)
exp
In the offline setting, aim to enhance the prior policy Tp via MB planning at test time by inferring the posterior over
=
Ep(r) exp K6 En rethh] En
(5)
a conditioned on the optimality observations Otth predicted by the learned model p... At time f, we seek to compute p(a (O), as shown in Figure (middle).
In the last step, we used p(O) S, (( a Ep(r) |exp (k En re+n)] and the proportionality coeffi- cient a > 0 cancels out.
The key distinction in this setup from the original control- as-inference framework is the inclusion of the prior policy, which serves source for action sampling during plan- ning. Given Tp and Pu. we can define the prior distribution over the trajectory as follows:
Thus, the posterior expectation over a 11 can be obtained with f(ar++ a as below.
Ep(r) a exp (K En Te+b)]
E Ep(rlo) [ax+*]]
(6)
Ep(r) Texp (1: En Toth))
H-1
p(T) = II
Tp (arhh (o (e+hh acpa)- (4)
exp (1: En reth)
a (7)
Li=1 exp (N En rt r(th))
Sampling trajectories from this prior, p(7), is straightfor- ward through forward sampling, where actions are drawn from Tp and state transitions are generated using Pv-
That is, we estimate the posterior mean by sampling N trajectories from p(T) with up and ps. then computing the weighted sum of the actions. Each weight w"
exp(x in (ttk))
Computing the exact posterior p(r|O) is intractable due to the difficulty of calculating the marginal pOO). However, importance sampling offers a practical method to estimate the posterior expectation over a To demonstrate, we first expand the posterior using Bayes' rule:
Σ exp(1 r244) is proportional to the exponentiated MB return of the nth trajectory, assigning higher weights to plans with better returns. This helps the agent select actions likely to improve on those from the prior policy.
We note that 77) can also be derived from an optimization perspective. Specifically, LOOP (iikchi et al. 2021) con- strains the distribution over plans by minimizing the KL. divergence from the prior policy. In LOOP, the variance of values generated by the model ensemble is penalized to mitigate uncertainty; however, the agent's 's epistemic uncer- tainty not explicitly modeled and fully addressed. By
H-1
O(T(O) & p(O)T)p(T) x exp
Σ h=0
reth)
II A=
Po (S Seeth Benr)mp (ac+ (4tth)

---

Reflect-then-Plann Offline Model-Based Planning through a Doubly Bayesian Lens
contrast, by viewing offline RL an epistemic POMDP and formulating it as a probabílistic inference problem, we can directly incorporate the agent's epistemic uncertainty into MB planning by approximately learning the belief dis- tribution, which we delve into in the next part.
Here, Py learns to predict future states and rewards condi- tioned the latent variable my- The encoder captures the agent's epistemic uncertainty, while the decoder provides predictions about the environment under different latent variables my- To sum up, we train a variational autoencoder (VAE) via maxø.v ED Ei=o ELBO, b) using tra- jectories sampled from the offline dataset D.
3.2. Learning Epistemic Uncertainty via Variational Inference
Unlike VariBAD, where the decoder is only used to train the encoder, we also use p for MB planning. To improve P 's accuracy, we employ a two-stage training procedure: first, the VAE is trained with ELBO; then, the encoder is frozen and pe is further finetuned using the MLE objective:
Although offline RL can be framed as a BAMDP, obtaining an exact posterior belief update is impractical. Inspired by Zintgraf et al. (2020) and Dorfmmnn et al. (002T), }. we intro- duce a latent variable m to approximate the underlying MDP. We assume that knowing the posterior distribution p(m) T.+) is sufficient for planning under epistemic uncertainty. As a result, transitions and rewards are assumed to depend on m, i.c., [s /5, a,, m. and r(s:, ar m). When p(m] Tia) is accurate and Ty is in-distribution, T and r will closely match the transitions in D. For out-of-distribution (OOD) Taa. the posterior over m allows modeling diverse scenarios for T and r.
T--1 Σ
L(t)= = E,
Emp |- P (5h+ Th (Sh: ap mw)l
Th
Trajectory segments of length H are sampled from D. At each step h E [0, H), the encoder q (-17%) samples mys enabling computation of the inner expectation and refining Py for improved predictions.
Given a trajectory Tat, consider the task of maximizing its likelihood, conditioned on the actions. Conditioning on the actions is essential because they are generated by a policy- 3 during training and Tp at evaluation and are not modeled by the environment. Although directly optimizing the like- lihood p(so, ro, $1, Ti....y St+1 |a. ....................... ay) intractable, can maximize the ELBO as in VariBAD by introducing an encoder 9,p and decoder P8:
3.3. Integrating Epistemic Uncertainty into Model-Based Planning
Building on the probabilistic inference formulation of offline MB planning and the representation of epistemic uncertainty via variational inference in the BAMDP framework, we introduce RefPlan. This offline MB planning algorithm integrates epistemic uncertainty into the planning process, improving decision decision-making and enhancing the performance of any offline-leamed prior policy during test time.
log p(So, ro....,, S141 ao....................... ax)
p(so, S141, me aa....,ay ar) dm, / m+
= log
Assume we have a sample m. ~ q (ml)..)) representing the agent's 's belief about the environment at time t. Our goal is to use this posterior to enhance test-time planning. In Section 3.1 we have computed p(T|O) using the learned models P and the prior policy Tp. By introducing the latent variable m to capture epistemic uncertainty, we extend the transition and reward functions to depend on m, giving the dynamics P (S++ 1 (8,, ar, m,) and rewards r(se, a,, m.). resulting in the following conditional trajectory distribution:
p(so, ro...., S141, me
a0: ,.....,, at)
= -log Emi
4o ( TTa)
2 E (oag (o S241 Im,, ap.....ar
- -K (mel
= EIBBO,(p,, v))
(8)
The encoder do is parameterized as an RNN followed by a fully connected layer that outputs Gaussian param- eters p(Ta)) and log (Tx). Thus, my (TT) = N (u(Ta), o² (T.))). The KL term regularizes the posterior with the prior p(m))). which is a standard normal at l 0 and the previous posterior qu (-((-T1--1 for subsequent time steps. The decoder Pa learns the transition dynamics and reward function of the true MDP. This becomes clear when we observe that the first term in F.I.O,, corresponds to the reconstruction loss, which can be decomposed as follows:
p(T]O, x (OFT,m.))(TT
H-1
kx)
r(s a1th,
x exp
h=0
H-1
II h=0
Pv (St+ (seth, acth, mp)mp(actt (sctte)
x
log (So, ro...., S141 |m,, ag...
ax)
(9)
Thus, we can apply the sampling- based posterior estimation in 7 to compute Ep(r O.m,) [atr++
= log p())mm))))) +)
log (Sh+ (sh, an, mc)
h=
A practical approach to handle epistemic uncertainty is to marginalize over the latent variable m,, effectively averag-
+ log (1h+1 1sh, an, mz)]
5

---

Reflect-then-Plan Offline Model-Based Planning through a Doubly Bayesian Lens
Table I: Normalized score performance when trained on ME and tested starting from OOD states sampled from R. For each prior policy, we report the performance of the prior policy (Prior), performance with LOOP, and performance with RefPlan. The "Orig' column shows the performance of the original prior policy on the ME dataset for reference. Best results under OOD conditions are bold.
Prior Policy
Environment
Orig
Prior
L.OOP
RefPlan
Hopper
111.4
84.21
90.01
89.39
CQI.
HalfCheetah
98.3
71.81
85.48
85.92
Walker2d
108.9
65.01
85.95
92.61
Hopper
46.9
39.60
45.78
42.50
MAPLE
HalfCheetah
64.0
62.31
91.41
91.84
Walker2d
111.8
66.82
74.05
87.82
Hopper
105.6
81.75
80.67
85,51
COMBO
HalfCheetah
97.6
57.89
79.57
84.49
Walker2d 108.3
62.54
66.88
72.82
ing over possible scenarios. This results in the marginal pos- terior P(T]O). Although directly computing this marginal posterior is challenging, we can estimate the expectation of optimal plans using the law of total expectation:
(RQ3) How does RefPlan perform when trained on limited offline datasets that increase epistemic uncertainty by re- stricting the datasets' coverage of the state-action space? (RQ4) How robust is RefPlan when faced with shifts in environment dynamics at test time?
Bp(r [ac++n] n] = -m (laa) (pp(r(0,m))
|a me]]
We evaluate these RQs using the D4RL benchmark (Fu ct al. 2020)) and its variations, focusing on locomotion tasks in HalfCheetah, Hopper, and Walker2d environments, each with five configurations: random (R), medium (M), medium- replay (MR), medium-exper expert (ME), and full-replay (FR).
The inner expectation follows 77. with states and rewards sampled from pv. conditional on my, The outer expectation over m, is computed using Monte Carlo sampling with ñ samples, giving us:
(p(((o) [u:++] R3
Metries: For RQ1-RQ3, we compare normalized scores averaged over 3 seeds, with 100 for online SAC and 0 for random policy, scaled linearly in between, using the D4RL library (Fu et al. 2020). For RQ4, we report average returns.
exp (K: Σn R.d
ab (10)
cxp (1: En r2th)
where T r(s a m²) and S44h+1 (- aPt,, m²). Figure (right) illustrates how Ref- Plan leverages the agent's past experiences Ty to shape epistemic uncertainty through the latent variable m, and enhances the prior policy Mp through posterior inference. Algorithm
Tn
Baselines: RefPlan is designed to improve any offline learned policy through planning. We have obtained prior policies using model-free methods (CQL, EDAC) and MB methods (MOPO, COMBO, MAPLE). Among offline MB planning methods, we use LOOP, which is designed to en- hance prior policies and outperforms methods like MBOP. Therefore, for each prior policy, we compare its original per- formance to its performance when augmented with LOOP or RefPlan for test-time planning.
2 in the appendix summarizes RefPlan ? Addi- tionally, following Sikchi al. (2021), we apply an uncer- tainty penalty based on the variance of the returns predicted by the learned model ensemble.
4. Experiments
4.1. Epistemic Uncertainty from OOD States
To address RQI we assessed RefPlan's robustness under high epistemic uncertainty caused by OOD initialization. Prior policies were trained on the ME dataset and evaluated the states from the R dataset. We tested three prior policies: CQL, MAPLE, and COMBO (Table
In this part, we answer the following research questions: (RQ1) How does RefPlan perform when the agent is initial- ized in way that induces high epistemic uncertainty due to OOD states? (RQ2) Can RefPlan effectively improve poli- cies learned from diverse offline policy learning algorithms?
Across all environments, RefPlan consistently mitigated per- formance degradation due to OOD initialization, with partic- ularly notable improvements in HalfCheetah and Walker2d. For instance, when MAPLE was used as the prior policy in
Direct planning with sampling methods like SIR (Skre et al. 2003) may be better for multi-modal problems. However, our approach using (10) yields strong empirical results, leave direct sampling for future work.
6

---

Reflect-then-Plan Plan: Offline Model-Based Planning through a Doubly Bayesian Lens
Table 2: Normalized scores of offline RL algorithms on D4RL MuloCo Gym environments (3 seeds). For each prior policy, we show its original performance and its performance augmented with LOOP or RefPlan (Ours) for MB planning during testing. Bold indicates the best performance, while underline denotes where confidence intervals significantly overlap between two methods.
CQL
EDAC
MOPO
COMBO
MAPLE
Orig
LOOP
Ours
Orig
LOOP
Ours
Orig
L.OPP
Ours
Orig
LOOPP
Ours
Orig
1.00P
Ours
R
1.0
1.1
1.2
23.6
23.5
23.5
32.2
32.4
32.4
6.3
6.2
6.0
31.5
31.8
31.6
Hepper
101.5
MR
94.6
97.5
98.1
100.4
101.0
101_1
90.3
93.6
94.5
101.1
101.4
101.8
61.0
77.7
82.6
ME
111.4
111.6
112.1
106.7
104.7
109.9
91.3
82.7
96.3
105.6
78.4
107.8
46.9
53.4
57.8
104.2
106.2
107.6
106.6
107.0
107.2
73.2
55.6
77.2
89.9
54,9
84.1
79.1
77.0
91.7
Halllhhellu
25.8
M
47.4
57.1
56.5
63.8
73.0
71.4
42.8
58.4
59.8
67.2
73.2
77.4
68.8
72.9
74,6
MR
47.0
52.1
54.1
61.8
66.9
66.5
70.6
71.8
73.8
73.0
71.2
75.0
71.5
74.7
76.3
ME
98.3
104.0
108.5
100.8
107.1
108.8
73.5
94.5
96.6
97.6
130.3
110.3
64.0
91.9
92.8
77.5
81.8
86.7
81.7
87.5
88.5
81.7
88.2
90.8
71.8
82.6
86.3
66.8
87.8
90.2
R
0.1
0.1
0.3
17.5
13.4
21.7
13.3
12.4
13.1
4.1
3.0
4.3
21.8
21.8
21.9
Walker2a
77.1
77.6
91.7
93.2
82.0
79.1
71.2
81.1
87.4
88.3
MR
63.5
81.9
93.6
85.0
86.0
86.4
81.7
85,2
88.3
88.0
89.5
93.3
85.0
89.5
91.2
MI
108.9
111.4
111.8
98.5
97.4
116.0
51.9
49.0
68.1
108.3
111.1
112.7
111.8
112.9
114.0
96.6
99.4
101.3
98.0
98.3
99.7
90.5
92.6
78.1
83.0
99.5
94.2
96.7
98.4
Median
IQM
Mean
Optimality Gap
RefPlan LOOP
0.875 0.900 0.925
0.88
0.90
0.92
0.94
0.850 0.875 0.900 0.12 0.14
0.16 0.18
Human Normalized Score
Figure 3: RLiable (Agarwal et al. 2022) comparison of RefPlan and LOOP. Across all four metrics (the higher the better for Median, IQM, and Mcan, while the lower the better for Optimality Gap), RefPlan demonstrates superior performance with non-overlapping confidence intervals, highlighting statistically significant improvements over LOOP.
HalfCheetah, RefPlan outperformed the original policy. Walker2d, RefPlan boosted performance by 16.4%, 31.4%, and 42.5% for COMBO. MAPLE, and CQL, respectively. Although the gains were modest in Hopper, RefPlan still reduced performance drops. Overall, RefPlan showed strong resilience under high epistemic uncertainty caused by OOD initialization.
using bootstrapping with stratified sampling, RLiable also estimates confidence intervals, ensuring that comparisons are not skewed by outliers noise.
We applied RLiable compare RefPlan and I.OOP across the tested environments and prior policy setups (Figure 3 Across all metrics, RefPlan consistently outperformed LOOP, with non-overlapping confidence intervals, indicat- ing statistically significant improvements. These results demonstrate RefPlan's superior ability enhance various offline policy learning algorithms by explicitly accounting for epistemic uncertainty during planning.
4.2. RefPlan Enhances Any Offline-learned Policies
To address RQ2 we evaluated the normalized score metric across the five offline policy learning algorithms. Table 2 shows that RefPlan outperformed baselines in 10 (CQL), (EDAC), 12 (MOPO). (COMBO), and 12 (MAPLE) of tasks, matching performance in the others. Both MB plan- ning methods, LOOP and RefPlan, improved performance, with RefPlan showing more substantial gain. On aver- age, RefPlan enhanced prior policy performance by 11.6%, compared to LOOP's 5.3%.
Another key question is whether the performance gains from RefPlan stem from its principled Bayesian framework or simply from a larger inference budget. To investigate this, we conducted an additional experiment using CQL as the prior policy the MR and FR datasets. We compared RefPlan to LOOP under equivalent computational load by allocating LOOP 16-fold increase in its sampling budget, matching the maximum budget used by RefPlan (i.c., when ñ = 16). The results, summarized in Table S demonstrate that while increasing the computational budget improves LOOP's performance, RefPlan consistently maintains its advantage across the tested configurations. This finding suggests that the superior performance of RefPlan is not merely a product of increased computation but is attributable
In order to make a more statistically rigorous comparison between RefPlan and LOOP, we leverage RLiable (Agarwal al. 2022), a framework designed for robust evaluation of RL algorithms. RLiable focuses statistically sound aggregate metrics, such as the median, interquartile mean (IQM). mean, and optimality gap, which provide compre- hensive view of algorithm performance across tasks. By

---

Reflect-then-Plan Offline Model-Based Planning through a Doubly Bayesian Lens
Hopper
HalfCheetah
Walker2d
110
90
105
85
100
105
95
80
90
100
75
85
70
80
95
(Orig)
CQL (Orig)
CQL (Orig)
COL LOOP
65
COL LOOP
75
COL LOOP
COL Ours
COL Ours
COL Ours
90
70
50K 100K 250K
500K Full (1M)
50K
100K 250K 500K Full (1M)
50K
100K
250K
500K Full (1M)
Figure Performance comparison of RefPlan and LOOP across different dataset sizes Hopper, HalfCheetah, and Walker2d environments using the FR dataset, which contains 1M samples. We use CQL as the prior policy learning algorithm, and the results represent the average and standard error calculated from three random seeds.
Table 3: RefPlan vs. LOOP with an equivalent inference budget.
Table 4: Average returns on HalfCheetah with dynamics changes.
Environment
Config
LOOP (16x)
RefPlan
Task
Orig
LOOP
Ours
Hopper
MR
97.8 1.1
98.1 0.5
joint hill
5295
6088
6190
FR
107.5 + 0.6
107.6 + 0.5
327.1
949.7
1224
Walker2d
MR
83.2 + 8.8
93.6
gentle
1087
2363
2435
FR
99.9 + 1.5
101.3 + 0.3
steep
2123
3245
6238
MR
53.2 + 0.1
54.1 0.6
field
1205
2774
3345
HalfCheetah
FR
83.1 + 0.8
86.7 1 0.7
shows that while MAPLE struggled with changed dynamics, MB planning methods improved performance. RefPlan achieved the best results across all variations but still faced notable drops, especially in the hill and gentle environments. Data augmentation for single-task offline RL could enhance adaptability, topic for future work.
to its explicit modeling and marginalization of epistemic uncertainty during planning.
4.3. Performance with Data with Different Sizes
With limited data, the agent faces increased epistemic un- certainty. A key question is whether RefPlan can better handle these scenarios with constrained data (RQ3). To explore this, we randomly subsample 50K, 100K, 250K, and 500K transition samples from the FR dataset for each environment. We then train the prior policy using CQL and compare its performance with that achieved when enhanced by either LOOP or RefPlan. As shown Figure 4 RefPlan consistently demonstrates greater resilience limited data, outperforming the baselines across all three environments.
5. Related Work
Offline RL Policy distribution shift in offline RL leads to instabilities like extrapolation errors and value overes- timation (Kumar et al.) 2019 Fujimoto al. 2019). To mitigate this, policy constraint methods limit deviation from the behavior policy (Wu al., 2019 Kumar et al. 2019
Fujimoto & Gul 2021), while value-based approaches pe- nalize OOD actions (Kumar et al. 2020 An et al.) [021]) Another strategy avoids querying OOD actions by learning values only from in-dataset samples and distilling a policy (Kostrikov al. 2022).
4.4. RefPlan More Robust Changing Dynamics?
To address RQ4 we evaluated RefPlan in the HalfCheetah environment under varying dynamics, including disabled joint, hill, slopes (gentle and steep), and field, following the approach of Clavera et al. (2019) (Appendix D). High epis- temic uncertainty arises when dynamics differ from those seen during prior policy training. We trained the prior policy using the FR dataset, which contains the most diverse tra- jectories, and used MAPLE for its adaptive policy learning.
MB offline policy learning trains a dynamics model from batch data to generate imaginary rollouts for dataset aug- mentation. To prevent exploiting model errors, ensemble- estimated uncertainty can be penalized in rewards (Yu et al. 2020; Kidambi et al. 2021 Lu et al.] (022). Alternatively, values of model--eeerraae samples can be minimized Yu al. 202T), or adversarial dynamics models can discourage
8

---

Reflect-then-Plan: Offline Model-Based Planning through a Doubly Bayesian Lens
selecting OOD actions (Rigter al. 2022
possible environments.
Offline policies are typically fixed after training, but Ghosh et al.] (2021) 2022) show they can fail under high epistemic uncertainty. underscoring the need for adaptivity. APE-V V (Ghosh et al. ) addresses this by using a value ensem- ble to approximate the distribution over environments, en- abling policy adaptation during evaluation. MAPLE (Chen ct al. 202T) employs an RNN to encode the agent's history into a dense vector for adaptive conditioning while leverag- ing an ensemble dynamics model to expose the policy to diverse simulated environments, improving robustness to uncertainty.
Relatedly, Ghosh al. 202I) introduced the epistemic POMDP, where an agent's 's epistemic uncertainty- from incomplete exploration or task umbiguity-induces partial observability. An epistemic POMDP is an instance of BAMDP, with a special focus on performance during a single evaluation episode rather than online learning and asymptotic regret minimization, Ghosh et al. (2022) further noted that offline RL in single-task setting can be viewed as an epistemic POMDP, as static offline datasets typically cover only a subset of the state-action space, leading to partial observability in environment dynamics beyond the dataset.
Model-based planning for offline RL MB planning en- hances responsiveness at test time. MBOP (Argenson & Arnold 202I) applies MPC with MPPI (Williams et al. 2015), adapting it for offline RL by using a BC policy for trajectory generation. Uncertain rollouts can be filtered based on ensemble disagreement (Zhan et al. 2022
We also adopt the epistemic POMDP perspective for single- task offline RL but focus on MB planning. Unlike prior approaches, our goal is to enhance offline-leareed policies by addressing epistemic uncertainty, enabling more effective generalization during deployment.
LOOP (Sikchi et al. 2022) improves offline-lcarne policies with planning, outperforming MBOP. It uses KL-regularized optimization for offiin planning but only addresses epis- temic uncertainty by penalizing ensemble variance in re- wards during TrajOpt. In contrast, RefPlan takes a Bayesian approach, explicitly modeling epistemic uncertainty for bet- ter generalization and performance.
6. Conclusion
In this paper, we introduced RefPlan (Reflect-then-Plane novel doubly Bayesian approach to offline model-based planning that integrates epistemic uncertainty modeling with boseel-basdd planning in a unified probabilistic framework. Our method enhances offline RI. by explicitly accounting for epistemic uncertainty, a common challenge in offline settings where data coverage is often incomplete. Through extensive experiments on standard offline RL benchmarks, we demonstrated that RefPlan consistently outperforms ex- isting methods, particularly under challenging conditions of OOD initialization, limited data availability, and changing environment dynamics, making it a valuable tool for more reliable and adaptive offline RL. Future work could extend RefPlan to more complex models and environments.
Probabilistie interpretation of MB planning The control-as-inferrence framework (Levine, 2018 Abdol
maleki et al.) 2018 provides probabilistic view of con- trol and RI. problems, naturally leading to sampling-based solutions in MB planning (Piché et al. 2019). Okada &
Taniguchi (220) showed that various sampling-based Tra- jOpt algorithms can be derived from this perspective. Jan- ner al. (2022) introduced a diffusion- planner using control-as-infeernnce to derive perturbation distribution, embedding reward signals into the diffusion sampling pro- cess. However, our knowledge, an offline MB planning algorithm that integrates an offline- learned policy as a prior in a Bayesian framework while explicitly modeling epistemic uncertainty within a unified probabilistic formulation.
are the first to propose
Bayesian RL and epistemic POMDP Bayesian RL (Ghavamzadeh et al.) 2015)) and BAMDPs (Duff, 2002) address learning optimal policies in unknown MDPs. A belief-state MDP, where sufficient statistic of the agent's his- tory (Guez et al. 2012), framing BAMDPs as a special case of partially observable MDPs (POMDPs) (Kaclbling let al. 1998). Zintgraf et al. (2020) extended this perspective to meta-RL. RL, introducing VariBAD, a variational inference- based method for approximating the belief distribution over
BAMDP can be reformulated
the belief serves
9

---

Reflect-then-Plane Offline Model-Based Planning through a Doubly Bayesian Lens
Acknowledgements
pdf
This work was supported by the Institute Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korean Government (MSIT) (No. RS- 2024-00457882, National AI Research Lab Project).
Clavera, I,, Nagabandi. A., Liu, S., Fearing, R. S., Abbeel, P., Levine, S., and Finn, C. Learning to adapt in dynamic, real-world environments through meta-reinforcement learning. In International Conference on Learning Rep- resentations, 2019. URL https:://pennevieet hwl.rooummid--yytxeea
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning, particularly Reinforcement Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Dorfman, R., Shenfeld, I., and Tamar, A. Offline meta reinforcement learning identifiability challenges and effective data collection strategies. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL ptttt::://ppeeeeveeww net/ forum?
References
Abdolmaleki, A., Springenberg. T., Tassa, Y., Munos, R., Heess, N., and Riedmiller, M. Maximum pos- teriori policy optimisation. In International Confer- ence on Learning Representations, 2018. https //openreview net. wwt.:/et/toooeee forum? idddd=sllm ANX0WO0
id- IIBEffLvve
Duff, M. O. Optimal Learning: Computational Procedures for Bayes-Adaptis Markov Decision Processes. PhD thesis, University of Massachusetts Amherst, Amherst, USA, 2002.
Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A., and Bellemare, M. G. Deep reinforcement learning at the edge of the statistical precipice, 2022. URL https: 7/arxiv. org/ abs/ / 08 3264
Fu, J.. Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
Fujimoto, S. and Gu, S. A minimalist approach to offline reinforcement learning. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https
An, G., Moon, S., Kim, J.-H., and Song, H. O. Uncertainty- based offline reinforcement learning with diversified q- ensemble. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural In- formation Processing Systems, 2021. URL https: 12h3
//openreview net/ umw (tt.://o
Fujimoto, S., Meger, D., and Precup, D. Off Off-policy deep reinforcement learning without exploration. In Proceed- ings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learn- ing Research, pp. 2062-2062, PMLR, 09-15 Jun 2019. URR.htttps://proceeeii wtt.::s....p...r....mmmmm pr9ss/v97 Fujimotol9a.httm
Argenson, and Dulac-Arnold, G. Model-based offline planning. In International Conference on Learning Rep- resentations, 2021. URL /tttps://epenreview net fert/oorm?2i http:/oorumn1a-oO0BBBGxxxxd
m
Biewald, L. Experiment tracking with weights and biases, 2020.
wandb. com/ Software available from wandb.com.
Ghavamzadeh, M., Mannor, S., Pineau, J., and Tamar, A. Bayesian reinforcement learning: A survey. Found. Trends Mach. Learn., 8(5-6):359 483, November 2015. ISSN 1935-8237.
Chen, X.-H., Yu, Y., Li, Q., Luo, F.MM., Qin, Z. T., Wen- jie, S., and Ye, J. Offline model-based adaptable pol- icy learning. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, W. (eds.), Advances in Neural In- formation Processing Systems, 2021. URL https: //openreview.ngg net. wtwp///wwwwe.//ouuimiiaa foormmiid=].. rdxc11 jm6
Ghosh, D., Rahme, J., Kumar, A., Zhang, A., Adams, R. P., and Levine, S. Why generalization in is difficult: Epistemic pomdps and implicit partial observability. In Ranzato, M.. Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 25502-255155 Curran Associates, Inc.,
Chua, K., Calandra, R., McAllister, R., and Levine, S. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, vol- ume 31, 2018 URL neurips. ct/ http://c..ccppppp//////////1111
2021. URL httt:s::/proceeeings.nneet
CC [cc//ppeee_ffeee/pppprr//////111 05FF18377881117733299877822--
pdf
10

---

Reflect-then-Plang Offline Model-Based Planning through a Doubly Bayesian Lens
Ghosh, D., Ajay, A., Agrawal, P., and Levine, S. Of- fline RL policies should be trained to be adaptive. In Chaudhuri, K., Jegelka, S., Song. L., Szepes- vari, C., Niu, G., and Sabato. S. (eds.), Proceed- ings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 7513-7530. PMLR, 17-23 Jul 2022. htlllttpss:/proceeediiee vi 2162/ghosh22..hm h
via bootstrapping error reduction. In Advances in Neural Information Processing Systems, vol- ume 32, 2019. URL https::/ppoceeetee
neur neurips.cc/ Peu.:/ps.cc//ppeer//01/////
pdf
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Con- servative q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
mwr......m.....o press/
Guez, A., Silver, D., and Dayan, P. Efficient bayes- adaptive reinforcement learning using sample-based search. In Pereira, F., Burges, C., Bottou, L., and Weinherger, K. (eds.), Advances in Neural Information Processing Systems, volume 25, Curran Associates, Inc., 2012. URL http::://proeeeddnggg.. neurips CC cc/paper files/ppppr/2012/EE
Levine, S. Reinforcement learning and control as proba- bilistic inference: Tutorial and review, 2018.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. CoRR, 0bs/2005.01643, 2020.
pdf
Lowrey, K., Rajeswaran, A., Kakade, S., Todorov, E., and Mordatch, I. Plan online, learn offline: Efficient learning and exploration via model- control, 2018. URL https://axxiv.org/e/ee www.aa////1111111111
Janner, M., Fu, J.. Zhang, M., and Levine, S. When to trust your model: Model-baser policy optimization. In Ad- vances in Neural Information Processing Systems, 2019.
Lu, C., Ball, P., Parker-Holder, J., Osborne, M., and Roberts, S. Revisiting design choices in offline model based reinforcement learning, 2021. URL https
Janner, M., Du, Y., Tenenbaum, J., and Levine, S. Plan- ning with diffusion for flexible behavior synthesis. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th In- ternational Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 9902-9915. PMLR, 17-23 Jul 2022.
//openreview. www.:neee///ouuuoo forum? http:e//fformmmmoooooueevew d=Ufx6vkvht.
Nagabandi, A., Konoglie, K., Levine, S., and Kumar, V. Deep dynamics models for learning dexterous manipula- tion, 2019. httt:ts:::///tt wwtpii.orr//bbb/100/0
Jin, Y., Yang. Z. and Wang, Z. Is pessimism provably ef- ficient for offline rl? In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Ma- chine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 5084-5094 PMLR, 18-24 Jul
1 11652
Okada, M. and Taniguchi, T. Variational inference mpe for bayesian model-based reinforcement learn- ing. In Kaelbling, L. P., Kragic, D., and Sugiura, K. (eds.), Proceedings of the Conference Robot Learn- ing, volume 100 of Proceedings of Machine Learn- ing Research, pp. 258-272 PMLR, 30 Oct-01 Nov 2020. URI wRt.httss: http:://///occeeddee/e iig. ml r.rsss// v v100/ html
2021. URL BtLhhtps::/ httptttppppppppp...poooooooooeee
mlr ptt.s/gss...pp.p............
v1 v139/ v139/jin21e, html
Kaelbling, L. P., Littman, L., and Cassandra, R. Plan- ning and acting in partially observable stochastic domains. Artif. Intell., 101(1-2):99-134 may 1998. ISSN 0004- 3702.
Piché, A., Thomas, V., Ibrahim, C., Bengio, Y., and Pal, C. Probabilistic planning with sequential monte carlo meth- ods. In International Conference on Learning Represen- tations, 2019. URL
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. Morel: Model-based offline reinforcement learning. In Advances in Neural Information Processing Systems, 2021.
net/
Rigter, M., Lacerda, B., and Hawes, N. Rambo-rl: Robust adversarial model-based offline reinforcement learning. arXiv preprint 1eXXi::2244115555 2022.
Kostrikov, I., Nair, A., and Levine, S. Offline reinforcement learning with implicit q-learning. International Confer- ence on Learning Representations, 2022. URL //openreview.ro wtw.:neeee.oommmdddddmdmdm..
Sikchi, H., Zhou, W., and Held, D. Learning off-cy with online planning. In 5th Annual Conference on Robot Learning, 2021. URL ht lttpp::// http://///p/eeee/////mmr net/ forum?id= 3 95eJ
Kumar, A., Fu, J., Soh, M., Tucker, G., and Stabilizing off-policy q-learning
Levine, S.

---

Reflect-then-Plann Offline Model-Based Planning through a Doubly Bayesian Lens
Skare, Ø., Bølviken, E., and Holden, L. Improved sampling- importance resampling and reduced bias importance sampling, Scandinavian Journal of Statistics, 30(4):
719-737, 2003. doi: http:///oo...rgg/..1111111444110
00360. URL Utt. ley
Com /ttp:///bb/...1111/1611111/11/1/
00360
Snoek, J. Larochelle, H., and Adams, R. P. Practical bayesian optimization of machine learning algorithms. Advances in neural information processing systems, 25,
2012.
Williams, G., Aldrich, A., and Theodorou, E. Model pre- dictive path integral control using covariance variable importance sampling, 2015.
Wu, Y., Tucker, G., and Nachum, O. Behavior regularized offline reinforcement learning. CoRR, URL http://arxiv.org/
abs/1911.1 1361, 2019.
abs 191 1911.11361
Yu, T. Thomas, G., Yu, L., Ermon, S., Zou, J., Levine, S., Finn, C., and Ma, T. Mopo: Model-basee offline policy optimization. In Advances in Neural Information Processing Systems, 2020.
Yu. T., Kumar, A., Rafailov, R., Rajeswaran, A., Levine, S., and Finn, C. COMBO: conservative offline model-based policy optimization. CoRR, abs/2102.08363, 2021.
Zhan, X., Zhu, X., and Xu, H. Model-based offline plan- ning with trajectory pruning. In Raedt, L. (ed.), Proceedings of the Thirty-FFrst International Joint Con- ference on Artificial Intelligence, IJCAI-22, pp. 3716- 3722. International Joint Conferences on Artificial Intel- ligence Organization, 7 2022. doi: 10.24966////ccca2222222
516. URL https:t//////....... (too.oorg/11..296/1117
2022/516 6
Main Track.
Zintgraf, L.. Shiarlis, K., Igl, M., Schulze, S., Gal, Y., Hof- mann, K., and Whiteson, S. Varibad: A very good method for bayes-adaptive deep rl via meta-learning, 2020. URL wttp:/iiii.rgg////////b/// http:://aaxx..orrg////////11/
12

---

Reflect-then-Plang Offline Model-Based Planning through a Doubly Bayesian Lens
A. Additional Background
A.I. Bayes-Adaptive Markoy Decision Processes
layes-Adaptive Markov Decision Processes (BAMDPs) Duuff 2002) extend the standard MDP framework by explicitly incorporating uncertainty over the transition and reward functions. In a BAMDP, instead of assuming that the transition dynamies T (s' js, a) and reward function r(s, a) are known and fixed, we assume that they drawn from an unknown distribution. The agent maintains posterior belief over these functions and updates it new data are collected through interaction with the environment.
To illustrate, consider simple case where we have finite and discrete state and action spaces with |S] = n, and |A| = Rai hence, a state can be represented with an integer, i.e., 8 = i for i == 1, 1,..., f2y and similarly for the actions. While the reward function r(s, a) assumed to be known, we are uncertain about the transition probabilities s(s'|s, s, a). We can model this uncertainty by placing a prior distribution over the transition probabilities, typically using Dirichlet prior, which is conjugate to the multinomial likelihood of observing transitions between states.
For each state-action pair (8, a) S A. the transition probabilities T a) are parameterized by multinomial distribution:
T('|s, a) Multinomial(Osa..))
(11)
where 8.. (0 s.o.1:.... s. represents the probabilities of transitioning from state to any state s' E S under action a. These parameters follow a Dirichlet distribution:
2..a Dirichlet(asaa),
(12)
where 0's.a = (O (0s,,,,1....... Os,a,n.) > O are the Dirichlet hyperparameters.
Initially, the agent holds a prior belief about the transition probabilities, represented by the Dirichlet hyperparameters as,a for all state-action pairs. As the agent interacts with the environment and observes transitions of the form (s, a, s'), it updates its posterior belief by simply updating the corresponding Dirichlet hyperparameters. Specifically, when the agent observes a transition from state s to state s' under action a, the corresponding Dirichlet hyperparameter is updated as:
O'g,0,s'
(13)
while all other Dirichlet hyperparameters remain unchanged. This process of updating the Dirichlet hy hyprraarameters fully captures the agent's experiences; hence, these hyperparameters act as sufficient statistics for the agent's belief about the environment.
By transforming the BAMDP into belief-state MDP, where the belief state b, p(0 Ti)) is distribution over transition probabilities conditioned on the observed trajectory 7. = (8p, 4g, . the agent can solve the problem using standard MDP solution methods. The augmented state space, or hyper-state space, includes both the physical state s E S and the belief state & B. In this simple finite state-action example, the belief state corresponds to the Dirichlet hyperparameters a.
The transition dynamies of the resulting belief-state MDP are fully known and can be written as:
T(s' T(s,, a s, T(5'|s,a,,,) a)laaa,s,)p(a
(14)
ag.o,s' (a as + +1), "'vrs as,a.""
(15)
where I(() the indicator function. This transformation turns the BAMDP into a fully observable MDP the hyper-state space, which allows the use of standard, e.g.. DP methods to compute an optimal policy.
However, the computational complexity of solving the BAMDP grows quickly with the number of states and actions. If the states are fully connected (i.e., p(s' a) > O, Vs, s'), the number of reachable belief states increases exponentially over time, making exact solutions intractable for even moderately sized problems.
For a comprehensive overview of solution methods for BAMDPs, we refer readers to the survey by Ghavamzadeh et al. (2013). In this work, we have utilized variational inference techniques from Zintgraf et al. (2020) and Dorfman et al. 22221) to approximate the agent's 's posterior belief over the environment dynamics, p(b) T:e), based on past experiences.
13

---

Sellect-then-Plan Offline Model-Based Planning through a Doubly Bayesian Lens
Table 5: Performance comparison of RefPlan against baseline methods on Hopper, HalfCheetah, and Walker2d tasks using MOPO and COMBO for offline policy optimization. The table evaluates original policies (Orig), policies trained with Non-Markovian (NM) dynamics models (NM (Train)), NM-trained policies combined with RefPlan for planning (NM (Train) + RefPlan), and RefPlan using original policies as priors. Results demonstrate RefPlan's ability to improve test-time performance across different dynamics models and environments.
NM (Tainn)
NM (Train)
NM (Train)
NM (Train) +ReiPlan
Orig
RetPlan
Orig
RefPlan
+RetPlan
Hopper
M
66.9
67.7
Hopper
M
60.9
52.2
62.30
77.2
90.3
93.2
101.1
101.8
ME
91.3
96.5
ME
105.6
27.3
39.23
107.8
HalfCheetah
Walker2d HallCheetah
42.8
40.6
66.45
59.8
M
67.2
30.3
41.61
77.4
MR
70.6
53.2
72,46
73.8
MR
73.0
47.6
59.54
75.0
ME
73.5
71.6
100.34
96.6
MR
97.6
93.5
109.25
110.3
Walker2d
82.0
60.6
72.73
M
71.2
87.4
MR
81.7
53.3
79.75
8883
MR
88.0
80.4
91.01
93,3
108.3
B. Additional Results
B.I. Performance comparison:
non-Markovian dynamics model for training vs. planning
The experiments presented in Table [ aims to evaluate the effectiveness of RefPlan in leveraging the VAE dynamics- consisting of the variational encoder 4e and the probabilistic ensemble decoder Pv (Figure for planning at test time. Specifically, these experiments compare the following approaches:
"Orig": the original prior policy trained using MOPO or COMBO.
"NM (Train)": the policy trained using non-Markovian (NM) VAE dynamics model during offline policy optimization via MOPO or COMBO.
"NM (Train) RefPlan the RefPlan agent that uses the policies trained using NM dynamics models as priors.
"RefPlan" ": the RefPlan agent that uses the original prior policies priors.
The results demonstrate several key findings. First, RefPlan consistently outperforms NM (Train) across environments and datasets, confirming that the VAE dynamics models are significantly more effective when used for planning at test time rather than during offline policy training. This highlights RefPlan's 's ability to explicitly handle epistemic uncertainty, leveraging the agent's s real-time history to infer the underlying MDP dynamics.
Second, in MOPO results, NM (Train) diverged or underperformed in several cases. This suggests that the heuristically estimated model uncertainty used in MOPO is not well-suited for integrating with the VAE dynamics models during offline training. Even with large penalty parameters, the value function diverged in the Hopper tasks, indicating a fundamental limitation in using NM models with MOPO for policy optimization. By contrast, COMBÓ results did not exhibit these issues, suggesting that COMBO's framework better equipped incorporate such dynamics models during training.
Finally, applying RefPlan policies trained with NM dynamics models (NM (Train) RefPlan) further boosted test-time performance, often by substantial margins. This demonstrates that even when NM dynamics models introduce suboptimality during offline training, RefPlan can recover and enhance the policy's 's performance through effective planning at test time. Across all environment-datase combinations, RefPlan provides robust improvements over both the original and NM (Train)-optimized policies, further validating its capability to address epistemic uncertainty and improve the generalization of offline- policies.
B.2. Evaluating the impact of the number of latent samples on variances and performance
This experiment evaluates how the sample variance of the marginal action posterior mean from (00p changes with the number of latent samples (7) used the outer expectation. each time step, we compute the posterior mean K times, calculate its
14

---

Reflect-then-Peet Offline Model-Basse Planning through a Doubly Bayesian Lens
X10
x10
Sample Variance
4.05
98.1
Sample Variance
1.20
108
Normalized lzdd Score
97.8
104
4.00
Norma 12
Pormalizee
1.14
97.5
100
1.08
3.95
1.02
96
3.90
96.9
0.96
92
2
21
2²
2
2
2
2
2
Number Latent Samples
Number of Latent Samples ñ
(a) Hopper (MR)
(b) Hopper (FR)
x10-2
x10
Sample Variance
2.10
54.1
Norma aizealizee Score
Sample Variance.
1.44
Normalizer Score
86.50
1.95
54.0
1.36-
1.80
86.25
53.9
1.28
1.65
1.20
86.00
53.8
1.50
1.12
85.75
20
21
2
2
24
20
21
22
2
24
Number of Latent Samples
Number of Latent Samples ń
(c) HalfCheetah (MR)
(d) HalfCheetah (FR)
x10-2
x10 -2
100.60
Sample Variance
92.5
Normalized Score
Sample Variance
Normalized Score
1.8-
2.1
100.45
90.0
1.6
1.8
100.30
87.5
1.4
1.5
100.15
1.2
85.0
1.2
100.00
2 Number of Latent Samples ń
2 Number of Latent Samples
2
(c) Walker2d (MR)
(f) Walker2d (FR)
Figure 5: The sample variance and the performance vs. the number of latent samples of RefPlan, evaluated from three environments with the MR and FR datasets using CQL as the prior policy.
variance averaged across action dimensions, and report the running average over 1,000-step episode. Results are averaged over three random seeds, with CQL as the prior policy, across three environments (Hopper, HalfCheetah, Walker2d) and two dataset configurations (MR, FR).
The figures show that increases, the average sample variance decreases, with ñ consistently yielding the highest variance. Performance, measured as normalized scores, generally improves with increasing ñ. suggesting positive correlation between reduced variance and higher performance. However, while reduced variance likely contributes to this improvement, further investigation is needed to confirm causality, as other factors may also play role.
C. Algorithm Details
C.1. Algorithm Summary
RefPlan designed to enhance any offline RL policy by incorporating MB planning that accounts for epistemic uncertainty. The algorithm operates in two primary stages: pretraining (Appendix and test-time planning.
15

---

Reflect-then-Plan Offline Model-Based Planning through a Doubly Bayesian Lens
Algorithm 1 Offline MB Planning
Input: Va. D, me, & Train Ps with D via MLE
Train VP and To with 2 and 'D
1+1
repeat
Observe St
a TTaajOOp Pv. mo, Vo)
Take a observe 1,7
t
until episode terminates
Algorithm RefPlan: Offline MB Planning as Probabilistic Inference
Input: T. (So: ag, ro:-...--// 8.)....,)).... Pei 4g- pp- Q, H, N,
Pes di (-Ixa) (Get the Gaussian parameters)
{m] (pe, of) (Sample latent vectors from the approximate posterior)
for n = 1....,,. do
for = 0,.... H do
a Tp (p(((s(4k) (Sample prior action sequence)
St+b+1 (+ (b) (Sample the next state from model using yt
end for for j 1...., n do
nJ f- St
for h = O.....
do
St+h+1 ~ (-1s ai m²) { (Sample next state from model using mi
rith +- rss a.+nmm) (Compute the reward using my
end for
end for
end for
Compute (o() [a with 100)
return Ep(((o) [ (Return the plan be used in line of Algorithm
!
Pretraining stage The first step is to train prior policy Vp using any offline RL algorithm. In parallel, a VAE is trained using the ELBO objective in (8). where the encoder captures the agent's 's epistemic uncertainty and the decoder learns the environment dynamics. See Appendix C.3 for more details.
During evaluation, the agent employs MPC (Algorithm D. where RefPlan serves as the trajectory optimization subroutine. At each time step the agent gathers its history T. and encodes it into latent variable m, using the pretrained encoder (line 2 of Algorithm] 2). This latent variable encapsulates the agent's 's current belief about the environment, reflecting epistemic uncertainty.
Test-time planning stage
Then, we first generate N prior plans with the prior policy and the learned model (lines 5-8). Each plan has the length of H, and we use μrr to condition Py at this stage. Optionally, we add a Gaussian noise to the actions sampled by xp. following Argenson & Dulac-Amoll (202T) [ikchi et al. (202T).
Once the prior plans are prepared, rollout the plans under the learned model to generate multiple trajectories. That is, for each sampled m. we obtain N trajectories (lines 9-14). These trajectories are then used to estimate the optimal plan, conditioned on m². We marginalize out the latent variable via Monte-Carlo expectation using the law of total expectation.
Finally, the first action from the optimized plan selected and executed in the environment. This process repeats at each subsequent time step, with the agent continuously updating its belief state and re-optimizing its plan based on new observations.
16

---

Reflect-then-Plan Offline Model-Based Planning through a Doubly Bayesian Lens
Table 6: Hyperparameters for Model Architecture and Training
Architecture Hyperparameters
Value
Task Embedding Dimension
State Embedding Dimension
Action Embedding Dimension
16
Reward Embedding Dimension
GRU Hidden Dimension
256
Decoder Network Architecture
Fully connected, [200, 200, 200, 200J with skip connection
Decoder ensemble size
20
Decoder number of elite models
Training Hyperparameters
Value
KL Weight Coefficient
Input Normalization
True
Learning Rate
0.001
Weight Decay
0.01
Optimizer
AdamW W
Batch Size
64
C.2. Architecture
Encoder
Decoder
2t
(90)
St
S141
8
p
F4
a/-1
GRU
172
P111
log oe
P1-1
Figure 6: A schematic illustration of the architecture of RefPlan. We use the same encoder architecture in VariBAD (Zintgraf et al. 2020), which consists of GRU model and a fully connected layer. Unlike VariBAD, which uses the decoder only for training the encoder, we employ a two-stage training procedure (Appendix(C.. (..3 to learn decoder that is directly used for planning at test time. The decoder network reconstructs the past trajectory and predicts the next state but does not attempt to predict the entire future trajectory as in the prior work (see also Eq.
Figure illustrates the architecture of RefPlan. For the encoder, we adopt the architecture from VariBAD (Zintgraf et al.) 2020), with few minor modifications the hyperparameters. The encoder utilizes GRU network encode the agent's history and outputs the parameters of Gaussian distribution representing the latent variable my.
At time t = 0, we initialize z-1 0 and a 0. The state st, the previous action a;-1, and the previous reward rt-1 are first embedded into their respective latent spaces using distinct linear layers, each followed by ReLUU activation. These embedded vectors, along with the hidden state from the previous time step 2t-1, are then processed by the GRU, which outputs the updated hidden state z,, This hidden state subsequently linearly projected onto the task embedding space to obtain the mean (Pt) and log variance (log of) of the Gaussian distribution for the latent variable at the current time step.
Since the decoder plays critical role in test-time planning, follow established practices from prior work and implement the decoder using probabilistic ensemble network (Chua et al.) 2018 Janner et al. 2019 Yu et al. 2020 2021 Chen et al.) 2021). Specifically, the ensemble consists of 20 models, from which we select the 14 elite models that achieve the lowest validation loss during training. The decoder network conditions on a latent sample m, ~ A(4,, of) along with s, and a,, to predict the next state S1+1 and reward rt+1- The hyperparameters associated with the architecture are summarized in Table
17

---

SellectthennPlang Offline Model-Based Planning through a Doubly Bayesian Lens
Table Reproducing the reported performances of offline policy learning algorithms on the D4RL MuJoCo tasks, "Numbers reported An et al.
2021
CQL
EDAC
МОРO
COMBO
MAPLE
Paper
Rep.
Paper
Paper
Rep.
Paper
Paper
Rep.
R
10.8
1.0
25.3
23.6
11.7
32.2
17.9
10.6
31.5
Hopper
M
86.6
66.9
101.6
101.5
28.0
66.9
97.2
60.9
21.1
29.4
ME
111.0
111.4
110.7
106.7
23.7
91.3
111.1
105.6
42.5
46.9
101.9"
104.2
73.2
79.1
HallCheetal
R
35.4
19.9
28.4
22.5
35.4
29.8
38.8
40.3
38.4
33.5
MR
46.2
47.0
61.3
61.8
53.1
70.6
55.1
73.0
59.0
71.5
62.4
98.3
106.3
73.5
63.5
FR
76.9*
77.5
84.6
81.7
81.7
71.8
66.8
R
7.0
0.1
16.6
17.5
13.6
13.3
7.0
21.7
21.8
WalkerZd
M
74.5
77.1
92.5
77.6
17.8
82.0
81.9
71.2
56.3
88.3
MR
39.0
ME
98.7
108.9
114.7
98.5
44.6
51.9
103.3
108.3
73.8
111.8
94.2"
96.6
99.8
98.0
90.5
78.1
94.2
C.3. Pretraining
RefPlan requires two stages of pretraining. First, an off-the-shelf offline RL algorithm to train prior policy Tp. In our experiments, we evaluated several algorithms, including CQL (Kumar et al. 2020), EDAC (An et al. 2021), MOPO Yuu
et al. 2020), COMBO Yu et al. 2021). and MAPLE (Chen et al. 2021). though any offline RL policy learning algorithm could be utilized.
Second, we train the encoder 46 and the decoder P The encoder do is trained using the ELBO loss as defined in (8. The decoder Pa is trained reconstruct the past and to predict the next state, conditioned the sample m, the current state s,, and the action a, This training constitutes the first phase of dynamics learning. During this step, the encoder learns latent representation that captures essential information for reconstructing the trajectory. Unlike VariBAD, where the decoder is trained to reconstruct the entire trajectory including future states, we found that focusing the past and the next state improves the decoder's performance.
After completing the first training phase, we freeze the encoder network parameters and proceed to the second phase. In this phase, we fine-tune the decoder network Py to accurately predict the next state given my St. and a,. This is achieved using the loss function that reiterate here for clarity:
I( ENND
Σ Emi rn) lo (Sh+1, Th Psls ag. m,)]
(16)
The second training phase ensures that the learned dynamics model, Pu. accurately predicts the next state. This two-stage approach allows RefPlan maintain effective dynamics moel for planning at test time, unlike VariBAD, where the decoder is discarded after training the VAE.
D. Experimental Details
D.I. Experimental settings
D4RL MuJoCo environments datasets
We use the v2 version for each dataset as provided by the D4RL library FFu
al. 2020).
Evaluation under high epistemic uncertainty due to OOD initialization (RQ1) To address RQ1 assessed policy trained on the ME dataset of each MuJoCo environment by initializing the agent from state randomly selected from the R dataset. The results, presented in Figure [ and [nd[an are averaged over 3 seeds. For fair comparison. the same initial state was used across all methods being compared the prior policy, LOOP, and Re/Plan-under the same random seed.
Benchmarking on D4RL tasks (RQ2) To generate the benchmark results shown in Table 2, we first trained the five baseline policies on each dataset across the three environments. The focus of our analysis on the performance
18

---

Deflect-then-Plan Offline Model-Based Planning through a Doubly Bayesian Lens
CQL RefPlan (BayesOpt)
CQL RefPlan (BayesOpt)
CQL + RefPlan (BayesOpt)
CQL Refplan (Grid)
CQL Refplan (Grid)
CQL + Refplan (Grid)
Best Performance
Best Performance
Best Performance
55
95
100
95
50
90
45
85
5
10
15 Iteration
20
25
30
5
10
15
20
25
30
5
10
15
20
25
30
Iteration
Iteration
(a) Hopper (MR)
(h) HalfChectah (MR)
(c) Walker2d (MR)
Figure 7: Best performance vs. the number of BayesOpt iterations, using CQL prior policy on the MR datasets across
three environments.
improvements of these prior policies when augmented with either LOOP RefPlan MB planning algorithm during the evaluation phase. Thus, approach designed be complementary any offline policy learning algorithms, making the relative performance gains more relevant than the absolute performance of each algorithm. Nevertheless, we aimed to closely replicate the original policy performance reported in prior studies. Table compares our reproduced results with those originally reported. Overall, our implementation closely matches the original performances, often exceeding them significantly across various datasets. However, some cases, our reproduced policy checkpoints underperformed compared to the originally reported results, such as CQL on the datasets, EDAC on Walker2d M and ME datasets, COMBO on the Hopper R and M datasets, and MAPLE on the Hopper MR dataset. We will make our code publicly available upon acceptance.
Varying dataset sizes (RQ3) Figure we present the normalized average return scores for CQL and its enhancements with either LOOP RefPlan vary the dataset size from 50K 500K. We conducted these experiments using the FR dataset across three environments, which originally contains 1M transition samples. To create the smaller datasets, randomly subsampled trajectories. If the subsampled data exceeded the desired dataset size, we trimmed the last trajectory accordingly. For CQL training, we applied the same hyperparameters as those used for the full FR dataset.
we adapted the HalfCheetah environment following the ap- proach of Clavera et al. (2019), introducing five vari- ations: disabled joint, hill, gentle slope, steep slope, and field. These variations were implemented using the code fromatttss///gitaat wwtlbb.cooiicccc wtw.iiaaeecom/miiaeee..eee Learning_to_sdapp
Changing dynamics (RQ4) To explore RQ4
Table 8: Environment configuration for the HalfCheetah varia- tions used in RQ4 experiments, showing the original and modified height parameter values for cach task.
Task
Original
Modified
hill 0.6
0.2
gentle
0.2
Unlike the original work, which focuses on meta-RL, fline RL problem within a single task framework. Hence, to make the tasks easier, we modified the height param- eter for most variations, excluding the disabled joint task. The specific adjustments to the height parameters are detailed in
steep
0.5
study addresses of-
field
Uniform(0.2.1 Uniiform(0.05 0.4)
Table g These changes were intended to create more manageable tasks while still providing a meaningful challenge for the offline RL algorithms.
D.2. Hyperparameters
Table 9 1 outline the hyperparameters used for RefPlan across the five prior policies discussed in Section We conducted a grid search over the following hyperparameters: the planning horizon H E (2,4) 4), the standard deviation of the Gaussian
19

---

Reflect-then-PPen PhennPlan: Offline Model-Based Planning through a Doubly Bayesian Lens
Table Hyperparameters used MAPLE + RefPlan used on D4RL MuJoCo Gym environments
Hopper
HalfCheetah
Walker2D
H
H
8
R
random
2
0.01
10.0
1
1.0
4
0.05
10.0
16
1.0
0.05
0.1
16
1,0
medium
0.01
0.1
8
0.1
4
0,01
5.0
0.5
0.05
10.0
1
med replay
0.01
5.0
0.1
4
10.0
0.1
0.05
16
med-expeat
0.01
0.1
1.0
4
0.01
10.0
0.1
0.01
10.0
8
0.5
full-an full-replay
0.05
Table 10: Hyperparameters used for COMBO + RefPlan used on D4RL MuloCo Gym environments
Hopper
HallChectah
Walker2d
H
A
H
R
random
0.05
0.1
16
0.1
2
0.05
0.1
0.1
0.01
16
1,0
medium
2
0.01
10.0
16
0,5
4
0.05
5.0
0.1
0.05
1.0
16
0.5
med-replay
med expert
0.01
0.5
8
1.0
2
0.05
5.0
16
0.1
0.01
10.0
16
0.1
full-replay
Table 11: Hyperparameters used for MOPO + RefPlan used on D4RL MuloCo Gym environments
Hopper
HalfCheetah
Walker2d
H
P
random
0.01
5.0
10.0
0.1
0.05
medium
2
0.05
5.0
1
0.1
4
0.05
10.0
0.1
0.05
5.0
4
1.0
mod-replay
16
0.1
2
10.0
0.1
0.01
med-expert
0.05
1.0
8
0.1
4
0.01
10.0
0.1
0.01
10.0
16
1,0
full replay
0.05
5.0
16
4
0.05
5.0
0.1
0.05
10.0
Table 12: Hyperparameters used for CQL + RefPlan used on D4RL MuJoCo Gym environments
Hopper
HallCheetah
Walker2d
H
P
5.0
medium
0.01
5.0
16
0,5
2
0.01
5.0
0.5
0.05
10.0
16
1,0
med-replay
0.05
0.1
8
0.1
4
0.01
5.0
0.1
0.05
1.0
4
med expert
0.01
1.0
0,5
2
5.0
0.1
0.01
full-replay
0.05
10.0
16
0.1
2
0.01
5.0
0.1
0.01
5.0
Table 13: Hyperparameters used for EDAC
RefPlan used on D4RL MuJoCo Gym enviromments
Hopper
HalfCheetah
Walker2d
P
H
P
P
random
10.0
16
0.5
4
5.0
0.1
0.05
10.0
4
medium
0.01
16
0.1
0.05
10.0
0.1
0.05
10.0
mod-replay
2
0.05
10.0
1
0.1
4
0.05
10.0
0.1
0.05
5.0
16
0.1
med-expert
0.05
full replay 0.05
10.0
4
0.1
2
0.05
10.0
0.1
0.05
10.0
1
0.1
noise a E (0.01,0.05]) the inverse temperature parameter и E 1.0, 10.0), the number of latent samples 41, 8, 16), and the value uncertainty penalty E 0.5, 1.0). Our findings indicate that k. and are the most influential hyperparameters, while the others have a comparatively minor effect on performance. For LOOP, we conducted similar grid search over the same hyperparameters, excluding which is specific RefPlan.
In addition, used Bayesian optimization (BayesOpt, Snoek et al. (2012)). implemented in W&B (Biewald 2020), to explore the challenge of identifying optimal hyperparameters for ReiPlan. Figure 7 compares the number of iterations required for BayesOpt to achieve or surpass the performance the best hyperparameter configuration found via grid search in each environment. Specifically, we used CQL as the prior policy and the MR dataset from three environments. Notably, BayesOpt required fewer than iterations to exceed the performance reported in Table
20

---

Rellect-ten-Plan: Offline Model-Basce Planning through Doubly Bayesian Lens
Table 14: Per-epoch runtimes for VAE pretraining on the ME dataset.
Hopper
HaltCheetah
Walker2d
4
Pe
94
Pe
4e
Ps
55.3s
39.8s
53.2s
40.7s
54.6s
40.7s
Table 15: Runtime per environment step for RefPlan during evaluation in the HalfCheetah environment.
f
2
4
H
7.9 10 10
8.7 x 10 's
9.3 x
10
1.0 X 10
4
1.5 10 10
1.6 10 3s
1 10 3
1.9
10
D.3. Computational Costs of RefPlan
In this section, we provide a detailed discussion of the computational costs associated with deploying RefPlan. outlined in Appendix C.3 RefPlan requires the following pretrained components: a prior policy Tp. encoder qá. and decoder p Since the prior policy is trained using standard offline policy learning algorithms (e.g.. CQL, EDAC, MOPO, COMBO, and MAPLE), which not contributions, focus reporting the computational costs associated with training the VAE model and executing the planning stage. All experiments were conducted on single machine equipped with an RTX 3090 GPU.
VAE Pretraining Table 14 presents the per-epoch runtimes for VAE pretraining in the three environments. The reported runtimes correspond datasets with 2M transition samples, the largest dataset size used in our experiments. Both the VAE pretraining and decoder fine-tuning phases were executed for up 200 epochs until the validation loss ceased to improve for 5 consecutive epochs, whichever occurred first.
Test- Test-Time Planning At test time, planning with RefPlan involves selecting hyperparameters as detailed in Appendix D.2 Among these, the planning horizon H and the number of latent samples ń influence runtime. Specifically, the computational cost scales linearly with H, which is an inherent property of planning algorithms. However, the cost increases sub-linearly with n, as shown in Table 15 For example, with H = 4 and ń = 4, the agent achieves approximately 53 environment steps per second. We hypothesize that further optimization of PyTorch tensor operations to fully exploit GPU parallelism could yield better computational performance, particularly with respect to
21


Reviews:
[{"review_id": "kjETpE4TKe", "reviewer": "Reviewer_CFgz", "summary": "This paper combines ideas from _adaptive_ and _online planning_ Offline RL to achieve a conceptually nice framework and reasonable performance improvements. They are able to improve upon (a) epistemically adaptive methods with no online computation, and (b) online computation methods with no explicit epistemic adaptability.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer CFgz for taking the time to thoroughly analyze our paper and provide constructive and insightful feedback.\nWe address the main points raised below:\n1. Epistemic POMDP vs. BAMDP: Thank you for this important point. We agree there's no fundamental conceptual difference and \"reformulation\" was imprecise. Our use of the Epistemic POMDP perspective aimed to highlight our specific offline pre-training -> single online evaluation setting, following [1]. As noted (lines 185-186), an Epistemic POMDP is an instance of a BAMDP; the distinction is mainly perspective. We apologize for the confusion and commit to revising the text for the camera-ready version to clarify this relationship and correct related statements (e.g., line 130).\n2. Ablation on uncertainty penalty: This is an excellent suggestion. To disentangle the benefits of Bayes-adaptivity from conservatism, we conducted new ablation experiments removing the uncertainty penalty, focusing on the CQL prior with medium-replay (MR) and full-replay (FR) datasets.\n\n[Table 1]\n| Env | Config | RefPlan| RefPlan w/o penalty |\n|----|----|----|----|\n| Hopper | MR | 98.1 ± 0.5 | 98.26 ± 0.5 |\n| Hopper | FR | 107.6 ± 0.5 | 107.71 ± 0.5 |\n| Walker2d | MR | 93.6 ± 0.3 | 93.71 ± 0.2 |\n| Walker2d | FR | 101.6 ± 1.1 | 100.35 ± 1.4 |\n| Halfcheetah | MR| 54.1 ± 0.6 | 54.34 ± 0.3 |\n| Halfcheetah | FR | 86.7 ± 0.7 | 87.42 ± 0.9 |\n\nThe table shows strong performance is maintained without the penalty, indicating gains primarily stem from Bayes-adaptive planning utilizing epistemic uncertainty, not conservatism.\n\n3. Ablation with random prior policy: We appreciate the suggestion. Early experiments showed random priors performed very poorly due to distribution shift. Given that even stronger priors like BC policies (e.g., in MBOP) were outperformed by LOOP, we decided a full sweep with a random prior would be less informative than other comparisons.\n4. Number of latent samples: Thank you for this keen observation regarding Figure 8 and the potential benefits of using more latent samples. We ran new RQ2 experiments with 32 and 64 latent samples (CQL prior, MR/FR datasets). Table 2 compares these with $\\bar{n}=16$.\n\n[Table 2]\n|Env|Config|16|32|64|\n|---|---|---|---|---|\n|Hopper|MR|98.3±0.5|96.5±0.3|98.0±0.3|\n| |FR|107.6±0.5|107.6±0.5|107.6±0.6|\n|Walker2d|MR|92.8±0.7|88.1±2.0|87.1±3.8|\n| |FR|100.1±0.8|100.3±0.9|100.3±1.3|\n|Halfcheetah|MR|54.3±0.3|54.1±0.3|54.4±0.1|\n| |FR|87.4±0.9|87.1±0.9|87.2±0.7|\n\nIn this standard setting, increasing $\\bar{n}$ beyond 16 yielded no significant gains, possibly due to lower epistemic uncertainty. We hypothesized the effect would be clearer with higher uncertainty. To test this, we ran additional RQ3 experiments using models trained on subsampled datasets.\n\n[Table 3]\n|Env|Data Size|16|32|64|\n|---|---|---|---|---|\n|Hopper|50k|99.5±12.9|102.7±4.0|**106.2±3.3**|\n|Hopper|100k|107.0±0.7|106.9±0.4|106.9±0.4|\n|Hopper|250k|106.8±0.6|103.2±3.3|106.7±0.2|\n|Hopper|500k|107.7±0.4|107.7±0.5|107.5±0.6|\n|Walker2d|50k|82.1±9.5|85.3±10.9|**93.9±1.0**|\n|Walker2d|100k|96.4±2.1|89.5±6.2|84.4±5.9|\n|Walker2d|250k|96.8±1.2|96.8±0.6|96.4±0.6|\n|Walker2d|500k|100.8±0.8|99.2±0.3|99.9±0.4|\n|Halfcheetah|50k|68.5±1.6|69.1±1.0|**69.3±1.0**|\n|Halfcheetah|100k|75.7±1.0|76.2±0.3|76.1±0.6|\n|Halfcheetah|250k|81.8±0.3|82.0±0.2|82.1±0.5|\n|Halfcheetah|500k|83.9±1.4|83.8±1.1|84.0±1.0|\nAs shown in the table, with higher uncertainty (limited data), increasing $\\bar{n}$ to 64 does indeed lead to more noticeable performance gains, particularly with the 50k dataset. This confirms the value of more samples in these settings.\n\n5. Stochastic environments: We agree that evaluating on stochastic environments would be a valuable extension. We acknowledge this limitation and plan to explore this direction in future work.\n\n6. Minor points & typos: Thank you for pointing out the areas for clarification and the typos. We will address all these points in the camera-ready version.\n\nWe hope these clarifications, new experimental results, and our commitment to revise the paper address your concerns. We are grateful for the constructive feedback, which has immensely helped us strengthen the paper. We believe the results, particularly the new ablations and latent sample experiments, further validate the effectiveness of our approach.\n\n[1] Ghosh et al. (2021) “Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability”", "claims_and_evidence": "The claims made are broadly sensible and supported, other than:\n\n1. The claim that \"existing methods rely on fixed conservative policies\" (abstract lines 15-17) is an oversimplification and is contradicted by later discussion of several adaptive ORL methods. The authors should clarify this point so the claims are not overstated.\n1. I don't agree with the way Epistemic POMDPs are compared to BAMDPs (lines 126-155, 185-191).\n    - There is no real conceptual difference between the Epistemic POMDP and BAMDP formulation, or the methods used to solve them. The difference is purely perspective/assumption-based, where epistemic POMDPs were pitched with the offline->online deployment setting in mind. The description in the text (\"Epistemic POMDPs prioritize performance during a single evaluation episode\") makes it seem as if there is a fundamental difference in their representative capabilities, which is not the case. This is even said later in the paper -- \"an epistemic POMDP is an instance of a BAMDP\" (lines 185-186).\n    - Along the same lines: line 202-203 says \"we can leverage the BAMDP reformulation of epistemic POMDPs\". I don't think there's any reformulation needed.\n    - I don't think it’s accurate to say a “BAMDP can be reformulated as a belief MDP” (line 130); a BAMDP is itself a special case of a belief MDP [Zintgraf et al. VariBAD: A very good method for bayes-adaptive deep rl via meta-learning, ICLR 2020.]\n    - Ultimately I think this confusion arises from the original Epistemic POMDP paper [Ghosh et al. 2021], where I do wish they'd originally posed the problem as a BAMDP. But there's no need to perpetuate that confusion here.", "methods_and_evaluation_criteria": "The chosen combination of methods is sensible, and the evaluation domains and baselines are standard for offline RL.", "theoretical_claims": "I checked the derivations in the text, which are relatively simple (the result of minor changes to existing methods). These derivations seem correct to me.", "experimental_designs_or_analyses": "The basic experimental approach seems sound.\n\n### Ablations\n\nI would have liked to see ablation studies or some kind of qualitative comparison to justify the design choices and components of the method:\n\n1. \"Additionally, following Sikchi et al. (2021), we apply an uncertainty penalty based on the variance of the returns predicted by the learned model ensemble\" (316-320), would be nice to see an ablation of this. It would be more elegant if all behaviour was Bayes-adaptive rather than some conservatism.\n1. Ablating prior policy with random policy instead\n\n### Statistical significance\n\nAlthough 3 seeds is not very many for the experiments, experiments are repeated several times over different prior policy methods, so the results in Figure 5 are based on a reasonable number of samples. It would be nicer if this plot was in the main body of the paper.\n\n### Experiment improvements\n\n- I would like to see more latent sample values being tried (Figure 8), rather than 16 being the maximum number.\n    - The \"normalised score\" hasn't obviously plateaued in the range of latent sample values shown in Figure 8.\n    - MCTS-based Bayes-adaptive planning methods such as BAMCP [Guez et al., \"Scalable and efficient Bayes-adaptive reinforcement learning based on Monte-Carlo tree search\". JAIR 2013.] effectively sample a new MDP (equivalent to latent sample) for each MCTS trial, which would be >> 16 samples.\n    - Given everything seems to run relatively fast (Table 13) I don't see why this couldn't be done with 32 or 64 samples for example.\n- It would be nice if stochastic MDPs were evaluated, as well as the D4RL deterministic continuous control ones. Evaluating only on deterministic environments is common for offline RL papers however.", "relation_to_broader_scientific_literature": "The work is a combination of ideas from two main areas in offline RL: adaptive behaviour and online planning. The contribution of this paper is to combine these ideas in a way that is conceptually nice.", "essential_references_not_discussed": "I did not know of or find any related work that was essential to be cited in this paper.", "other_strengths_and_weaknesses": "1. The flow and clarity of the paper is good. Given that this paper combines ideas/algorithms from several fields, the authors achieved the task of succinctly covering the important concepts and how they are combined into their method.", "other_comments_or_suggestions": "## Minor / clarity\n\n1. Figures 3/6/7 seem like they could be combined somehow, with a more information-dense presentation than bar charts. It could also do with more discussion in the text, as the performance improvement over LOOP seems weakest in this experiment setting.\n1. I think clarity would be aided by always showing subscripts on $\\mathcal{O}$ in (2) and (3)\n1. The list of evaluation configurations (lines 351-355) could make it clearer that it is the behaviour policy that is changing between configurations. It would be helpful to have a brief description of how these configurations vary, or a pointer to somewhere where the variants are more fully described, especially for a reader without familiarity with D4RL.\n\n\n### Typos\n\n- Line 43: “have severe implications” should be “has severe implications”\n- Line 164: “model of the environment refers to the transition and reward functions”. The difference between the ground-truth T and r, and the estimated models of them, should be made clearer here.\n- Line 973 “moel” -> “model”", "questions_for_authors": "1. Did you carry out ablations or comparisons to justify your design choices discussed in my ablation comments? As a combination-of-methods paper I think it is important to discriminate between the contributions of the individual components, and ablation studies or more discussion would help with this.\n1. Do you agree with my comments on the Epistemic POMDP/BAMDP confusion? I could be convinced otherwise if you have a good argument.", "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": "xJGpjn8MC6", "reviewer": "Reviewer_GXYS", "summary": "The authors introduce RefPlan, a doubly Bayesian method for offline model-based RL.\nRefPlan combines two existing methods (1) the probabilistic control-as-inference formulation of MB planning (using a policy prior) with (2) the variational representation of epistemic uncertainty. \nAt inference time, RefPlan marginalizes over the latent variable representing the environment, to consider several possible MDPs and improve inference-time planning.\nThe authors show that on D4RL, RefPlan is robust to OOD states, to environment changes and performs well with limited data.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed, constructive review. We appreciate the opportunity to address your feedback, particularly on statistical significance and budget comparisons.\n1. Statistical significance: While runs use 3 seeds, Figure 5 uses RLiable [1] for robust aggregate analysis (Appx B.1). Figure 5 demonstrates that RefPlan consistently outperforms LOOP across aggregate metrics with non-overlapping confidence intervals, indicating statistically meaningful improvements. We will add error bars to Figure 3 on revision.\n2. Reporting raw scores: Normalized scores are from the D4RL codebase; we will clarify this on revision.\n3. Essential references: We will add discussion on Dyna/Dreamer. Re: SMC, we cited Piche et al. (lines 114, 166) but will expand discussion on its relation to RefPlan in revision.\n4. Paper structure & highlighting contributions: Thank you for the thoughtful suggestions on restructuring. We agree that improving the flow to present novel aspects earlier could strengthen the paper. While major restructuring is challenging, we commit to revising the presentation for the camera-ready version to improve clarity and better foreground RefPlan's core contributions.\n5. Answers to specific questions:\n\n     1) Novelty in Sec 4.1 & 4.2: While algorithmically related to MBOP/VariBAD, our primary contribution is the conceptual framing and synthesis. Deriving MB planning via control-as-inference justifies the prior and enables the Bayesian uncertainty treatment (Sec 4.3). The novelty lies in integrating these for offline RL in RefPlan. We will clarify this. We will refine the writing to make this relationship clearer.\n     2) $\\bar{n}$ and $\\bar{N}$ values: For the main results (RQ2), we used $\\bar{n}=16$ and the same $\\bar{N}$ as LOOP. Full hyperparameter details are in Appendix D.2, Table 7.\n    3) Usefulness of latent variable: Fig 8 shows benefits up to $\\bar{n}=16$. New results with $\\bar{n}=32,64$ (in our response to Reviewer CFgz, Tables 2-3), show further gains, particularly under high uncertainty due to limited-data (RQ3), suggesting the value of marginalizing over $m$.\n    4) Comparison with same inference budget: This is a crucial point. We conducted new experiments to directly address this:  \n* LOOP with increased budget: We compared RefPlan against LOOP using 16x its original sampling budget (matching RefPlan's max $\\bar{n}$). Table 4 shows RefPlan generally maintains an advantage, though LOOP benefits from compute (esp. w/o penalty in HalfCheetah-FR).\n\nTable 4: RefPlan vs. LOOP-16x budget, RQ2\n|Env|Config|LOOP|LOOP w/o penalty|RefPlan|\n|---|---|---|---|---|\n|Hopper|MR|97.8 ± 1.1|97.8 ± 0.8|98.1 ± 0.5|\n| |FR|107.5 ± 0.6|107.5 ± 0.6|107.6 ± 0.5|\n|Walker2d|MR|83.2 ± 8.8|78.4 ± 7.1| 93.6 ± 1.1 |\n| | FR | 99.9 ± 1.5 | 99.9 ± 1.5 | 101.3 ± 0.3 |\n| Halfcheetah | MR | 53.2 ± 0.1 | 55.1 ± 0.4 | 54.1 ± 0.6 |\n| | FR | 83.1 ± 0.8 | 89.4 ± 0.7 | 86.7 ± 0.7 |\n* LOOP with increased budget (limited data): We repeated this comparison in the RQ3 setting (Table 5). Here, RefPlan's advantage over the high-budget LOOP was often more pronounced, especially with smaller datasets, suggesting RefPlan better handles the higher epistemic uncertainty.\n\nTable 5: RefPlan vs. LOOP-16x budget, RQ3\n\nHopper FR:\n| Data Size | LOOP| RefPlan|\n|----|----|----|\n| 50k | 101.7 ± 4.9| 99.5 ± 12.9|\n| 100k| 106.9 ± 0.3| 107.0 ± 0.7|\n| 250k | 107.2 ± 0.3| 106.8 ± 0.6|\n| 500k | 104.5 ± 3.6| **107.7 ± 0.4**|\nWalker2d FR:\n| Data Size | LOOP| RefPlan|\n|----|----|----|\n| 50k| 70.6 ± 16.4 | **82.1 ± 9.5**|\n| 100k| 95.9 ± 1.4 | 96.4 ± 2.1|\n| 250k| 96.1 ± 0.8 | 96.8 ± 1.2|\n| 500k| 98.6 ± 0.9 | **100.8 ± 0.8**|\nHalfcheetah FR:\n| Data Size | LOOP | RefPlan |\n|----|----|----|\n| 50k| 63.8 ± 0.2| **68.5 ± 1.6** |\n| 100k| 71.4 ± 0.4| **75.7 ± 1.0**|\n| 250k| 77.9 ± 0.3 | **81.8 ± 0.3**|\n| 500k| 82.4 ± 0.7 | **83.9 ± 1.4**|\n* RefPlan with reduced budget: We matched RefPlan's budget to LOOP's by increasing its replanning interval. Table 6 shows RefPlan still performed competitively, often better than LOOP (except HalfCheetah-FR).\n\nTable 6:\n|Env|Config|LOOP|RefPlan-ReducedFreq|\n|----|----|----|----|\n| Hopper| MR| 97.5 ± 0.5| **98.7 ± 0.6**|\n|  | FR  | 106.2 ± 0.7| **107.8 ± 0.5**|\n| Walker2d| MR| 81.9 ± 3.0| **89.3 ± 2.3**|\n| | FR  | 99.4 ± 0.3| **100.1 ± 0.5** |\n| Halfcheetah | MR| 52.1 ± 0.0| 52.1 ± 0.2 |\n| | FR | **81.8 ± 1.0**| 79.4 ± 0.7|\nThese results suggest that RefPlan's performance gains are not merely due to increased computation but stem from the way it handles epistemic uncertainty via the doubly Bayesian approach, leading to more effective planning.\n\nWe hope these responses, new experiments directly addressing the budget comparison, and our commitments to revision clarify the contributions and robustness of RefPlan. We are grateful for your constructive feedback and hope we have provided the evidence needed to reconsider the initial evaluation.\n\n[1] Agarwal et al. (2021). “Deep Reinforcement Learning at the Edge of the Statistical Precipice”", "claims_and_evidence": "The experiments ask interesting questions and study whether Replan\n - mitigates performance drop due to OOD states (5.1)\n - enhances performance and outperforms other methods (conservative policies) (5.2)\n - performs well with subsampled datasets (5.3)\n - is robust to changing dynamics (5.4)\n\nHowever, I have two comments about the experiments:\n\n(1) I wish the authors had used more than 3 seeds, which is very little and insufficient to claim statistical significance. Also, their error bars are missing in Table 1, Table 2 and Figure 3.\n\n(2) In the current setup, RefPlan increases the inference budget by a factor $\\bar{n}$ (the number of times the latent variable m is sampled). Consequently, the authors should show that the performance gains reported do not simply come from this increase, but from combining the strengths of their two Bayesian frameworks.\nSee my questions (2), (3), (4) below.", "methods_and_evaluation_criteria": "The proposed framework which combines (1) a probabilistic inference for MP planning, using a learned policy as prior and (2) a latent variable for modeling the underlying MDP makes sense.", "theoretical_claims": "The paper does not contain theoretical claims.", "experimental_designs_or_analyses": "See \"Claims And Evidence\"\n\nAlso, the authors only report normalized scores (100 being online SAC). They should report the SAC scores, so that future work can compare to them.", "relation_to_broader_scientific_literature": "Increasing run-time compute for improving the performance of pre-trained agents is a critical problem for developing general agents. The paper proposes an interesting approach in this direction.", "essential_references_not_discussed": "Since the authors discuss MB offline policy learning (also known as background planning) they should cite the original Dyna work\n  - Sutton, Richard S. \"Integrated architectures for learning, planning, and reacting based on approximating dynamic programming.\" Machine learning proceedings 1990. Morgan Kaufmann, 1990. 216-224.\nand the Dreamer line of work:\n   - Hafner, Danijar, et al. \"Mastering diverse domains through world models.\" arXiv preprint arXiv:2301.04104 (2023).\n\nIn the context of planning as inference, the authors should cite the sequential Monte Carlo work, which estimates the probability of following an optimal trajectory from the current timestep and add a resampling step.\n  - Piché, Alexandre, et al. \"Probabilistic planning with sequential monte carlo methods.\" International Conference on Learning Representations. 2018.", "other_strengths_and_weaknesses": "I enjoyed reading the paper, which was well written and easy to follow.\n\nHowever, I think that the paper would benefit from highlighting more the authors' contribution. Here is one option:\n\n - As far as I understand, Sections 4.1 and 4.2 are summarizing the MBOP and VariBAD work, and the main novelty of the paper is 4.3. I find the current 4.1 and 4.2 to be misleading: I think the authors should consider moving these section to the preliminaries Section 3. They should also consider compressing Section 2 and 3, so that their contribution starts before the 6th page---which is currently the case.\n\n - In this new organization, Section 4 would be focused on RefPlan. The authors may also consider moving Algorithm 2 to Section 4.\n\n - The authors should keep more space for their experiments. For instance I am not sure that the results for applying RefPlan to MAPLE and COMBO should be in the appendices (there could be a table in the main text).", "other_comments_or_suggestions": "(1) Equation (1): $\\hat{r}_{\\psi}$ has not been defined. I assume it is a learned reward model. \n\n(2) The optimism bias that occurs with the RL as inference framework should probably be mentioned.", "questions_for_authors": "(1) Is there any technical novelty in Sections 4.1 and 4.2? As far as I understand, 4.1 is similar to MBOP which also uses a BP policy as action prior, and 4.2 is summarizing VariBAD?\n\n(2) What is the value of $\\bar{n}$ and $\\bar{N}$ used in the experiments? How does it compare to LOOP?\n\n(3) Could the authors show that introducing the latent variable $m$ is useful? One option would be to compare: RefPlan vs. planning as inference without epistemic uncertainty (Equation 7) using using $\\bar{n} \\times \\bar{N}$ samples\n\n(4) Similarly, could the authors compare RefPlan to others method for $\\textbf{the same}$ inference budget? Two options could do that could be:\n  - using $\\bar{n} \\times \\bar{N}$ samples for LOOP (same as RePlan)\n  - only replanning the sampled every $\\bar{n}$ step when using RefPlan \n\nI am willing to increase my score if the authors can provide a stronger evidence that their gains come from their proposed RefPlan and not from simply increasing the inference budget of existing methods.", "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": "r8XtWK7Mqz", "reviewer": "Reviewer_BPa5", "summary": "This paper proposes learning an approximate Bayesian model for offline RL and use planning guided by an offline RL learned policy prior for action selection. Standard practices of ensemble architecture and variance penalty are used for planning. Experiment shows improved performance over offline RL only policy and comparable or better performance over LOOP which is a similar method but does not use a VAE based model.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer BPa5 for dedicating time to review our paper and for the feedback. We are glad you found the experiments demonstrate improved performance and the theoretical motivations sensible, and we appreciate the opportunity to clarify aspects of our problem setting and methodology that seemed unclear.\n1. __Problem setting (offline vs. offline-to-online):__ Thank you for this clarifying question. Our setting is distinct from typical _\"offline-to-online\" RL (e.g., [1]) which often involves substantial online training iterations_. We focus on the scenario where an agent is trained entirely offline and then deployed online for evaluation, using planning to adapt its behavior based on the offline-learned model and uncertainty estimate. **No further training or model updates occur during this online deployment phase.** This \"offline pre-train, deploy with planning\" setting is important for applications where online interaction for training is limited, costly, or unsafe, yet adapting behavior based on learned uncertainties during deployment is crucial for robustness.\n\n2. __Training phases (phase 1 / phase 2):__ To clarify the phases mentioned: both phase 1 (training the policy prior $\\pi_{\\mathrm{p}}$ and the VAE encoder/decoder) and phase 2 (fine-tuning the decoder) are performed entirely offline using the static dataset before any online interaction. The subsequent online phase is purely for evaluation without any training updates.\n\n3. __Section 5.1 initialization:__ Your understanding is correct. For evaluation, the agent is initialized in a state sampled from the random (R) D4RL dataset and then interacts with the environment online.  \n\n4. __Planning objective (vs. Bayes-adaptive):__ RefPlan performs planning at each step using the learned model ensemble. And regarding the objective, as also clarified for Reviewer CFgz, we view the problem through the lens of an epistemic POMDP, which is an instance of a BAMDP. Therefore, indeed the underlying objective is Bayes-optimal behavior. RefPlan approximates this by sampling multiple latent model hypotheses (m) from the VAE's posterior approximation and planning within each sampled MDP using the prior-guided planner. The posterior over the MDP is continuously updated through the encoder during this evaluation time. By effectively marginalizing plans across these model samples, RefPlan seeks actions robust to the epistemic uncertainty captured by the VAE, thereby aiming to maximize performance in the evaluation episode, implicitly accounting for model uncertainty.\n\n5. __Decoder fine-tuning.__ To clarify further based on your comment, the decoder fine-tuning mentioned (phase 2) occurs offline during pre-training, not during the online test/evaluation episode.\n\n6. __Planner implementation (iterations):__ While Algorithm 2 presents a simplified view, our planner implementation does support multiple optimization iterations, similar to standard MPPI-style planners (e.g., resampling trajectories based on importance weights). We kept the pseudocode concise for clarity but will release our code upon acceptance to provide full implementation details. However, our earlier experiments showed no significant performance improvements from multiple sampling iterations, leading us to not pursue this approach.\n\n[1] Lee et al. (2022). “Offline-to-Online Reinforcement Learning via Balanced Replay and Pessimistic Q-Ensemble”", "claims_and_evidence": "Yes, the experiments demonstrate improved performance across different environments and offline RL algorithms.", "methods_and_evaluation_criteria": "It's not super clear to me the exactly problem the authors are trying to tackle, mostly the two phases the authors are referring to. It's not offline-to-online RL, or is it? \n\n**On data sampling and evaluation**\n\nAs far as I understand:\n* In phase 1: train policy prior use any offline RL algo, train encoder and decoder using the VAE objective\n* In phase 2: fine-tune decoder on new data, with encoder frozen.\n\nSection 5.1 mentions the new data are states from R. It is using R to set the initial state of the agent and the agent interacts with the environment from there on? \n\n**On planning objective**\n* The planning objective used by the authors is more of a posterior sampling approach than a Bayes adaptive approach which would require the agent to plan by rolling out $m_{t:t+h}$? The authors mentioned on line 143 that the evaluation target is single episode performance. Why not use the Bayes adaptive objective to actively reduce uncertainty and thus enhance performance in a single episode?\n* The single episode objective also seems to be a bit conflicting with fine-tuning the decoder, where fine-tuning is usually done between episodes for model based RL.", "theoretical_claims": "The paper is largely empirical. The theoretical motivations from the Bayesian perspective make sense.", "experimental_designs_or_analyses": "The experiment designs make sense. The chosen tasks and environments largely follow standard practices.", "relation_to_broader_scientific_literature": "The paper sits within the broad offline RL literature. It takes a Bayesian perspective on offline RL which has been considered before by e.g.,  [Chen et al.](https://proceedings.neurips.cc/paper/2021/hash/470e7a4f017a5476afb7eeb3f8b96f9b-Abstract.html). Introducing online planning in this context seems to be the main novelty.", "essential_references_not_discussed": "I am not aware of missed references.", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "* On planner implementation: most MPPI style planner can optimize action sequences for multiple iterations. Algorithm 2 shows that you only use a single iteration. Is that correct?", "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": {"ethical_review_concerns": "NA"}}]

Author rebuttal:

Meta-review: The RefPlan method is a useful approach for offline RL. It uses both online planning (guided by a policy prior) and a way to handle uncertainty about the environment (using latent MDPs). Tests showed it got better results and was reliable on hard problems. Reviewers initially had some questions about the experiments and explanations. However, the authors addressed some of the experiment problems in their response, which helped support their results. Since the approach is new and works well, the overall view is positive and the paper should be accepted.
Decision: Accept (poster)