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

We thank the reviewer for their thoughtful feedback and for highlighting the strengths of our work, including the principled incorporation of Bayesian uncertainty estimation into offline model-based planning and its motivation for real-world scenarios. Below, we address your questions and concerns in detail.

”APE-V (Ghosh et al., 2022) seems like a valid baseline for adaptive offline algorithms, however the paper does not compare with it”

Thank you for suggesting this baseline. Unfortunately, APE-V does not have an official codebase, making direct comparisons challenging. However, MAPLE, included in our experiments, also claims to learn an adaptive policy in an offline manner using a model-based approach. As shown in Table 1, RefPlan significantly improves the test-time performance of MAPLE, demonstrating its ability to enhance prior policies.

”the hyperparameters need to be carefully tuned for each task, which might limit the usability of the proposed method.”

Thank you for raising this important point. As detailed in Appendix D.2, we tuned five key hyperparameters: horizon $H$, noise scale $\sigma$, inverse temperature $\kappa$, value penalty $p$, and the number of latent samples $\bar{n}$. Using grid search, we originally required 240 iterations for hyperparameter tuning.

To evaluate the practicality of this tuning process, we employed Bayesian optimization (BayesOpt) using Weights & Biases. Figure 11 in the appendix compares the number of BayesOpt iterations needed to achieve or exceed the best performance obtained via grid search. Using CQL as the prior policy on MR datasets from Hopper, HalfCheetah, and Walker2d, we observed that as few as 5 to 20 iterations were sufficient to match or surpass the performance of the grid search.

Notably, even the first BayesOpt iteration achieved performance levels comparable to or better than LOOP and the original prior policies. This demonstrates that the computational cost of hyperparameter tuning for RefPlan is manageable in practice.

”Section 5.1: How exactly do you test the prior policy in states sampled from the R dataset?”

Thank you for this question. To evaluate RQ1, we sampled a state randomly from the R dataset. During evaluation, this sampled state was used to override the internal state of the MuJoCo simulator when resetting the environment, allowing us to start the evaluation from this state. For fair comparison, we have used the same random seeds across compared methods to ensure they’re evaluated from the same initial states.

”how is the computational efficiency of the proposed RefPlan method, in training and in executing, respectively?”

We appreciate your question about computational efficiency. The training phase of RefPlan involves offline pre-training of a VAE dynamics model, which comprises two stages: encoder pre-training and decoder fine-tuning, as detailed in Appendix C.3. These stages rely entirely on supervised learning, making them computationally efficient. In our experiments, we performed 200 epochs of VAE pretraining and an additional 500 epochs of decoder fine-tuning.

The prior policies were trained using various model-free and model-based offline policy learning algorithms, a process orthogonal to RefPlan itself. At test time, compared to LOOP, RefPlan adds computational overhead due to the marginalization over latent variables in Equation 13. However, by maximizing GPU parallelization, the runtime efficiency of RefPlan is comparable to LOOP.

”Figure 4: I wonder what the performance will be like when you use the full dataset for LOOP and RefPlan. Maybe continuing the lines in the plots to 1M would help readers”

Thank you for this insightful suggestion. We have updated Figure 4 to include the performance of all compared methods when using the full dataset (1M transitions) for training.

Dear Reviewer F97j,

The discussion phase is about to conclude, and we kindly ask if you could review our responses and updated revisions. We’d greatly appreciate your feedback or reconsideration of your evaluation if our updates have addressed your concerns.

Thank you again for your time and thoughtful review.

We thank the reviewer for their thoughtful feedback and for highlighting the strengths of our work, including the principled and clear formulation of RefPlan, its soundness, and the clarity of our writing. Below, we address each of your comments and concerns in detail.

”It's questionable that RefPlan only fine-tunes a baseline policy. Either i) why the learnt model can’t be used to optimize a new policy”

Thank you for this suggestion. To investigate this, we conducted additional experiments using the learned VAE dynamics model (VariBAD’s model) for offline policy optimization. The results, included in Table 2 in Appendix B.3, compare the following setups:

1. Original prior policy learning and evaluation (Orig).
2. Using VariBAD’s model during offline policy training (NM (Train)).
3. Applying RefPlan to policies trained with VariBAD’s model (NM (Train) + RefPlan).
4. Applying RefPlan to original prior policies (RefPlan).

The results are also summarized below:

These experiments demonstrate that RefPlan consistently outperforms using the VAE model during offline training (NM (Train)). RefPlan dynamically adapts to epistemic uncertainty during test time using real-time history, which is not captured as effectively in offline training alone. Furthermore, applying RefPlan to NM (Train) policies significantly improves their performance, highlighting its ability to recover from suboptimal offline training.

”it's a bit unfair in terms of extra computation needed in comparisons to the baselines, including both model-free and model-base… the model learnt using model-based methods in the baseline will be discarded or unnecessarily unused in the fine-tuning stage of RefPlan.”

We appreciate this concern. To ensure fairness, we used LOOP as a primary baseline, as it also performs model-based planning during test time, introducing extra computation. The “Orig” column in Table 1 reflects the performance of prior policies without additional test-time planning, while LOOP and RefPlan include the cost of planning. While RefPlan incurs more computational overhead due to sampling latent variables $m_t$ and marginalizing via Monte Carlo, we leverage parallelized computation on a GPU to keep the runtime overhead sublinear and comparable to LOOP.

Thanks for the new results. However from the Table, RefPlan does not look like outperform it bconsistently the policy train offline with the VAE model.

We thank the reviewer for their feedback. Below, we address your comments and concerns in detail.

"lacks strong evidence that the agent has learned a near Bayes-optimal policy. It is recommended to add theoretical support or to include a navigation task, along with visualizations of the agent's behavior."

Thank you for raising this concern. RefPlan builds on the theoretical foundations of [1] (Proposition 5.1), which shows that the Bayesian offline RL objective— $J_{\mathrm{Bayes}}(\pi)=E_{\mathcal{M}\sim P(\mathcal{M}|\mathcal{D})}\left[J_{\mathcal{M}}(\pi) \right]$ ---is maximized by a policy conditioned on the agent’s belief over MDPs inferred from its history. While RefPlan employs approximate belief updates through variational inference, which may lead to suboptimal behaviors, it remains grounded in this theoretical principle.

We appreciate the suggestion to include a didactic navigation task and behavioral visualizations. These additions could help illustrate the agent’s performance and are planned for inclusion in the final version of the paper.

"the approach of offline model-based planning as probabilistic inference is common (see [1])"

Thank you for raising this point. The control-as-inference framework has indeed been used in various reinforcement learning contexts, including offline planning, as demonstrated in [1]. However, [1] does not account for leveraging prior action sampling distributions derived from offline-learned policies. In contrast, RefPlan seamlessly integrates model-free and model-based methods within a unified probabilistic framework, enabling prior policies learned through offline model-free algorithms to guide action sampling during test-time planning.

Moreover, RefPlan advances the field by adopting an epistemic POMDP perspective on offline RL, explicitly modeling and addressing the agent’s epistemic uncertainty during planning. This treatment of uncertainty allows RefPlan to handle out-of-distribution states more effectively, leveraging real-time history to adapt during deployment. To the best of our knowledge, this approach has not been considered in prior work, including [1].

By introducing the concept of prior policies into the control-as-inference framework and explicitly addressing epistemic uncertainty, RefPlan not only bridges model-free and model-based methods but also provides a novel and practical solution for robust offline RL.

"the proposed algorithm is merely a minor modification of VariBAD, lacking novelty"

We respectfully disagree with this assessment. While RefPlan builds on VariBAD’s VAE structure, its contributions extend beyond simple modification. Unlike VariBAD, which tackles meta-RL in multi-task settings, RefPlan focuses on single-task offline RL and introduces a unified probabilistic framework that explicitly considers epistemic uncertainty during test-time planning. This framework allows RefPlan to combine the strengths of offline model-free and model-based approaches: utilizing prior policies from offline RL and enhancing them through test-time planning with real-time uncertainty handling.

We also believe that connecting ideas across domains to tackle new problems constitutes meaningful and impactful research. RefPlan’s application of Bayesian modeling and offline planning to single-task offline RL offers a novel and practical solution to a challenging problem in reinforcement learning.

Dear Reviewer NMjJ,

The discussion phase is nearing its conclusion, and we would greatly appreciate it if you could review our responses and let us know if they have adequately addressed your concerns. And we hope our clarifications and additional results will assist in reconsidering your evaluation.

Thank you once again for your time and thoughtful review.

”Q2: Given concerns of W3 (offline RL vs meta-learning), , rather little information is learn until "much data" is gathered. As a result, it feels as if any additional performance from online planning could have been done offline: refine policy by doing control-as-inference offline on Varibad's model. Do you have any idea how well that could or would perform?”

Thank you for this insightful question. The control-as-inference framework can indeed be applied in various ways for policy learning, as explored in [5, 6]. Extending this framework to offline RL with VariBAD’s model could potentially involve novel algorithmic design choices. While this is an interesting direction, developing such an offline RL algorithm would constitute a separate study beyond the scope of this work.

To address your concern more directly, we investigated a related baseline where VariBAD’s model is used during offline training for policy optimization. Specifically, we compared the following setups:

1. Original prior policy learning and evaluation (Orig).
2. Using VariBAD’s model during offline policy training (NM (Train)).
3. Applying RefPlan to policies trained with VariBAD’s model (NM (Train) + RefPlan).
4. Applying RefPlan to original prior policies (RefPlan).

These experiments, detailed in Appendix B.3 (Table 2), indicate that planning with VariBAD’s model at test time (as in RefPlan) consistently outperforms using it during offline training. RefPlan leverages real-time history to dynamically adapt to epistemic uncertainty during deployment, which is not captured as effectively when the model is used exclusively in offline training. Moreover, combining RefPlan with NM (Train) policies leads to significant performance improvements, underscoring its ability to recover from limitations in offline training.

We sincerely thank the reviewer for their thoughtful feedback and for highlighting the relevance of our work, its non-trivial contributions, and its strengths. Below, we address each weakness and question in detail.

”W1: One room of improvement, especially for someone with lack of background, is the accessibility of the background.”

Thank you for pointing this out. We have revised the main text to include clear explanations of epistemic POMDPs, BAMDPs, and control-as-inference in Section 3. Additionally, we have expanded on BAMDP in Appendix A.1. These additions aim to make the background more accessible and help readers distinguish the novel contributions of our work from prior methods.

”W2: A major concern, in my opinion, is the lack of experiments for online learning.”

We appreciate this concern and would like to clarify our problem setup. Our work focuses on offline RL, where a policy is trained using an offline dataset and subsequently deployed in a testing environment. The core goal is to address the agent’s epistemic uncertainty during deployment to achieve robust performance, rather than improving online learning efficiency.

As introduced by [1], epistemic POMDPs explicitly model the train-test split in RL. While BAMDPs emphasize online learning, epistemic POMDPs focus on single-episode evaluation at test time, prioritizing the agent’s adaptivity during deployment. Our experiments are consistent with this framework and intentionally focus on the deployment phase to isolate and evaluate the benefits of Bayesian modeling for epistemic uncertainty.

While our work focuses on the deployment phase, we acknowledge that the Bayesian framework introduced in RefPlan has the potential to inform offline-to-online RL methods. By explicitly modeling the agent’s epistemic uncertainty, it could help improve efficiency and adaptivity in online learning. This aligns with directions explored in the offline-to-online RL literature, such as [2, 3], and represents an important avenue for future research beyond the scope of this work.

”W3: VariBad tackles meta-learning: it is assumed (and exploited) that the training data set is generated from different tasks. … not clear whether the "Bayesian" argument holds here: Varibad's encoder might just collapse - as all trajectories come from the same environment”

Thank you for this insightful comment. Unlike VariBAD, which addresses meta-RL across multiple tasks, our work targets offline RL with a single task. Here, the source of epistemic uncertainty arises from incomplete state-action coverage in the offline dataset, as discussed in [1] and [4]. This aligns offline RL with the epistemic POMDP framework, where uncertainty is due to unexplored parts of the environment rather than task variability.

To investigate whether the latent distribution collapses to a deterministic one, we examined the impact of the $\bar{n}$ hyperparameter (the number of latent samples) on evaluation performance. Specifically, for CQL as a prior policy and with the FR dataset, we fixed all other hyperparameters and varied $\bar{n} \in {1, 4, 8, 16}$, measuring the performance across three random seeds for each configuration. The resulting correlations between $\bar{n}$ and performance in each environment were as follows:

These non-zero correlations suggest that the latent distribution does not collapse; instead, it retains uncertainty information that improves performance as $\bar{n}$ increases.

Additionally, as shown in Figure 9 in the Appendix, the variance of the optimized actions averaged across an episode decreases as $\bar{n}$ increases, and the performance generally improves. This provides further evidence that the latent distribution is capturing meaningful uncertainty, enabling RefPlan to leverage this information effectively during planning.

”Q1: Is 4.and 4.2 background (control as inference & varibad), or are there particular extensions / modifications hidden in there?”

Section 4.1 adapts control-as-inference to offline MB planning, specifically incorporating a prior policy and emphasizing deployment. This distinction is highlighted just above Equation 3. Section 4.2 leverages the ideas from VariBAD but modifies its use to focus on single-task offline planning. We discuss these differences at the end of Section 4.2. To further clarify, we added additional preliminary content in the main text and appendix to help readers distinguish between prior work and our contributions.

Dear Reviewer ezET,

The discussion phase is nearing its conclusion, and we would greatly appreciate it if you could review our responses and let us know if they have adequately addressed your concerns. And we hope our clarifications and additional results will assist in reconsidering your evaluation.

Thank you once again for your time and thoughtful review.

We sincerely thank the reviewer for taking the time to provide valuable feedback and constructive suggestions. We appreciate your recognition of the importance of our problem, the novelty of our contribution, and the strengths of our work. Below, we address each of your concerns in detail:

”Background / Improving Clarity: The paper could be improved with more clarity on a lot of the background…”

Thank you for highlighting this point. We have added a concise discussion of the control-as-inference framework, epistemic PODMP, and BAMDP in Section 2 to enhance the reader’s understanding. Additionally, we expanded on BAMDP in Appendix A.1.

”RQ3 and RQ4 seem to be a superset of RQ1.”

Thank you for pointing out the need for clarification. Each RQ focuses on high epistemic uncertainty arising from different causes:

• RQ1 evaluates RefPlan’s ability to handle uncertainty caused by OOD initialization. Specifically, the agent is offline-trained using a medium-expert dataset with limited state-action coverage and is then initialized in a state sampled from the random dataset, creating significant epistemic uncertainty due to minimal overlap between the datasets.
• RQ3 addresses the epistemic uncertainty resulting from limited data availability during offline training. By subsampling the full-replay dataset, we vary dataset sizes, with smaller datasets leading to greater epistemic uncertainty at test time about the environment’s dynamics.
• RQ4 evaluates the scenario with high epistemic uncertainty due to changing environment dynamics at test time.

Although RQ1, RQ3, and RQ4 all assess performance under high epistemic uncertainty, they do so in different settings. To clarify this, we have revised the text to explicitly define the sources of epistemic uncertainty in each RQ and ensure that their scopes are distinct.

”the last comment in alg2 refers to line 5 in alg1, but there are no line numbers, line 5 specifically is the beginning of a loop”

Thank you for catching this oversight. We have added line numbers to Algorithm 1 and corrected the comment in Algorithm 2 to accurately describe the referenced portion of Algorithm 1.

”what are the error bars used in the experiments?”

The error bars in Figure 4 represent the standard error computed over three random seeds. All experiments presented in Section 5 use three random seeds. We have updated the figure caption to clarify this.

”bold vs underline meaning in Table 1?”

We appreciate the opportunity to clarify. Bold numbers indicate the best performance for each prior policy learning algorithm. Underlined numbers indicate the top two results when their confidence intervals overlap significantly. Since RefPlan is designed to enhance the performance of offline-learned prior policies during test time, the comparison is made per prior policy. For instance, in the Hopper environment with a medium dataset, RefPlan boosts the performance of the CQL prior policy from 66.9 to 85.1. We have revised the text to make this clearer.

”H-step is mentioned without being defined.”

Thank you for noting this potential source of confusion. We define the H-step return in Equation 1 as the discounted sum of model-predicted rewards over H steps, and we consistently use H to denote the prediction horizon. If there are additional areas where this definition is unclear, we would appreciate further clarification and will make the necessary revisions.

Dear Reviewer sNjS,

The discussion phase is nearing its conclusion, and we would greatly appreciate it if you could review our responses and let us know if they have adequately addressed your concerns. And we hope our clarifications and additional results will assist in reconsidering your evaluation.

Thank you once again for your time and thoughtful review.