Reflect-then-Plan: Offline Model-Based Planning through a Doubly Bayesian Lens

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.

1. 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.

2. 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.

3. 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.

4. 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.

5. 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.

6. 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.

[1] Lee et al. (2022). “Offline-to-Online Reinforcement Learning via Balanced Replay and Pessimistic Q-Ensemble”

Thank you for the detailed, constructive review. We appreciate the opportunity to address your feedback, particularly on statistical significance and budget comparisons.

1. 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.

2. Reporting raw scores: Normalized scores are from the D4RL codebase; we will clarify this on revision.

3. 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.

4. 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.

5. Answers to specific questions:

   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.
   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.
   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$.
   4. Comparison with same inference budget: This is a crucial point. We conducted new experiments to directly address this:

• 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).

Table 4: RefPlan vs. LOOP-16x budget, RQ2

• 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.

Table 5: RefPlan vs. LOOP-16x budget, RQ3

Hopper FR:

• 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).

Table 6:

We 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.

[1] Agarwal et al. (2021). “Deep Reinforcement Learning at the Edge of the Statistical Precipice”

We sincerely thank Reviewer CFgz for taking the time to thoroughly analyze our paper and provide constructive and insightful feedback. We address the main points raised below:

1. 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).
2. 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.

[Table 1]

The table shows strong performance is maintained without the penalty, indicating gains primarily stem from Bayes-adaptive planning utilizing epistemic uncertainty, not conservatism.

3. 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.
4. 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$.

[Table 2]

In 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.

[Table 3]

5. 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.

6. 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.

We 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.

[1] Ghosh et al. (2021) “Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability”