Offline Reinforcement Learning with Closed-loop Policy Evaluation and Diffusion World-Model Adaptation

It is unclear whether the error bars report standard devision or standard derror. In table 2, the caption reads 'we show the standard deviation of the performance... Note that the standard error is usually large...'

It is clear that we report the standard deviation in our results like other relevant works in this area. We believe that there’s no issue in this caption since the high standard error is exactly the reason that causes the high standard deviation of our results.

I feel it is important to raise the significant issue of referencing in this paper as a weakness as well as in my ethics review. Buried in the appendix (line 903) there are 12 cited papers, many with common authorship and none with relevance to this paper or explanation. Either these papers are relevant to this work, and should be raised as related work with explanation, or they are not and thus should not be included. I assume these papers should not be included in this paper, but if they should be describing why is important.

This is a mistake. As the reviewer noted, some of these papers are on completely different topics like bio-data processing. These were legacies from the previous template/paper we used. We tried to remove them – successfully from the main text – but seem to have missed one place in the appendix. It's been removed in the revision and there's no need to read too much into it.

Minors:

Thank you so much for reading our work thoroughly and catching these. Most of the minors have been corrected in the revision. However, we disagree with the following ones:

In line 44, the world model does not interact with the policy! The policy interacts with teh world model.

This claim doesn't make any sense, since the word “interaction” naturally implies a two-way relationship.

The acronym of the algorithm (ADEPT) does not fit its name at all really.

That's not an issue, since the name of an algorithm is not necessarily its acronym.

The first sentence of the methodology does not make sense.

We didn't see any problem in that sentence. Could you make it more specific?

Line 346: I don't really know what this means - what is the discrepancy?

Discrepancy is already defined in section 4.1 text without an equation index.

Line 411: I don't know what $H$ is and it is not defined I don't think.

$H$ is already defined in Section 3 as the horizon, Paragraph named “Offline RL using World Models”.

In table 2, all bolded values are with a standard error of each so bolding, which implies significant improvement, is misleading.

This sentence seems somehow confusing. Nevertheless, we revised our paper to bold only the significant results.

Questions:

It seems a bit counterintuitive/non-obvious to compute your done flag when uncertainty reaches a limit. Does this effectively just lead to truncated rollouts? What kind of bias does this have on bias?

When the estimated uncertainty is larger than a threshold, a done flag will be set and truncate the rollout. A potential issue of this mechanism is that the truncated rollouts not only prevent synthetic trajectories from moving to uncertain states and action pairs on which the world model might have high prediction errors but also add a penalty in reward for those pairs, leading the target policy to be more conservative. Previous methods that sample initial states from any time step in the offline dataset could have another problem that the synthetic trajectories can’t move too far away from the existing trajectories. In our method, we sample initial states from both the offline dataset and the replay buffer consisting of previous synthetic trajectories to solve this problem.

I don't quite understand the assumption of line 340, that you can omit one of the inputs in the proof. To my understanding, the reward model comes after the denoising model, and as such would not intrinsically be able to ignore $\hat{s} _{t+1}$. Is this not correct?

For our method, it could be more accurate if we use $r_{\eta, \theta}(s_t, a_t) = \sum_{s’} r_{\eta}(s_t, a_t | s’)P_\theta(s’|s_t, a_t)$ to represent the reward prediction model in the theory, by treating the diffusion process and MLP as a whole. However, we simplified that as $r_{\eta}(s_t, a_t)$ for more conciseness and generalization, while such a simplification will not have any impact on the proof.

How were hyperparameters for the underlying learning method tuned? It says they were consistent for different methods - were they tuned for ADEPT or for one of the base algorithms, or are they taking default values.

We use the default hyperparameters for policy optimization (SAC) and most of the diffusion model training. We tuned the hyperparameters for uncertainty penalty($\lambda_r$, $\lambda_0$), the denoising steps ($K$), and the width of the diffusion network. The results can be found in the appendix.

How does this method compare to other approaches which compensate for overestimation bias? For example, how does it compare against policy guided diffusion ([1] above)?

We attached the comparison results in the response to Reviewer C8UH. Please check that.

[1] Policy Guided Diffusion, Jackson et al. 2024 [2] Lu, Cong, et al. "Synthetic experience replay." Advances in Neural Information Processing Systems 36 (2024). [3] Janner, Michael, et al. "When to trust your model: Model-based policy optimization." Advances in neural information processing systems 32 (2019).

We thank the reviewer for their thoughtful comments. We address your concerns below.

Weakness:

I found figure 1 quite confusing, and personally think that if a figure needs such a long caption the figure is probably quite unclear. For instance, I don't follow what (b3) demonstrates. It seems like finding a way to show this without showing all the data points etc., for visual purposes, might make things clearer.

Long captions are often used to make a figure self-contained. It is a common practice and, in our opinion, shouldn’t be considered as a weakness. We decided to show the data points in Figure 1 because the diffusion model is trained on the samples reweighted by importance sampling. Further explanations can be found in our main text. Nevertheless, we have adjusted Figure 1 in the revision, also based on the opinions of Reviewer PaFJ. We hope that the new figure will be easier to understand.

In related work, stating that [1] doesn't 'provide explicit mechanisms to solve the distributional shift issue' is fundamentally false - that is the entire and explicit basis of the method. Besides this, I found the related work relatively thorough; one other relevant work would be [2].

The distribution shift problem considered in our paper is quite specific. As training goes on, not only is the target policy being updated and continuing to deviate from the behavior policy for data collection, the diffusion model – which is trained on offline data collected by behavior policy – also continues to deviate and leads to increased error in evaluating the target policy. Paper [1] did not consider this problem. In fact, taking a closer look at Paper [1], you can see that it didn’t show significant improvement compared with the unguided diffusion generated datasets or original datasets in nearly 2/3 of the environments, nor did it show higher performance than the previous work SyntheR [2] on the same dataset. The second paper you mentioned considers online RL, whereas we consider offline RL in this paper. We are happy to clarify this in the updated version to avoid such misunderstanding.

I found the description of the architecture hard to interpret. I would clarify that the MLP predicts the reward when the architecture is first interpreted. Similarly, the way the inputs are introduced ('we replace $x_0$ with ...') was a bit confusing and could be worded better.

We feel the fact that the MLP generates the reward prediction is quite straightforward (so is the the inputs of the NNs), but can add a sentence on it as requested by the reviewer.

Despite spending a long time with it, and attempting to verify the appendix proofs, I found I had a tough time with this maths and didn't find it intuitive to follow. It is also not made clear to me how the derivations in the paper lead to the reward gap guarantee at the top. Note I am not coming from a theoretical background, but imagine that others might also find this difficult.

Sorry to hear that the reviewer was not able to go through the entire proof in the appendix. As the reviewer mentioned, it may not be easy for someone who is not used to this type of analysis in this area. We’d suggest the reviewer start with the proofs in MOPO[1], MOReL[2], and [3], which have considered some partial special cases of our problem. Due to space limitations, we couldn’t add too much explanation to the paper. However, if the reviewer has specific questions about any steps of our proof, we will be happy to provide further information.

It feels this should be compared to other existing methods for compensating for overestimation bias. The key baselines here should be comparing the same policy optimisation algorithm with different approaches for reducing overcompensation bias. This is not what is shown in this paper.

We strongly disagree with this comment. First of all, we are not sure where the terms “overestimation bias” or “overcompensation bias” are coming from. These are not related to our paper. As we introduced in the paper, the distributional shift problem is the fundamental reason for overestimation of returns, which is well known for offline reinforcement learning. We propose a new offline RL algorithm to address this problem. All baselines we considered in the evaluation, including model-free and model-based offline RL algorithms, are designed with different approaches to address this problem. Therefore, the claim that “This is not what is shown in this paper” is totally false.

There are no error bars for any of the results besides ADEPT's, meaning it hard to see overlapping of confidence intervals.

We explain the reason for no error bars in the original paper in the response to Reviewer TBhS and Reviewer C8UH. As it's required by other reviewers, we have included variances of the baselines in the revision.

Dear authors,

Thank you for your very thorough response. Below, I go through step-by-step:

we have adjusted Figure 1 in the revision

Given it is at the beginning of the paper, it is not helpful for understanding the frequency of things changing - for instance, why $\hat{\mathcal{M}}$ changes only once. I personally do not think it is as helpful for conveying the idea of the paper as a simpler figure would be. I do however look forward to seeing the new version.

Paper [1] did not consider this problem

To my understanding, this is the exact problem considered by [1]. While I appreciate that your results may be better (in fact, they are - and I appreciate the additional comparison), that is not an excuse for falsehoods in related work comparison.

We feel the fact that the MLP generates the reward prediction is quite straightforward

Perhaps this could be clearer by stating 'and reward function $R(s_t,a_t)$ respectively'?

However, if the reviewer has specific questions about any steps of our proof, we will be happy to provide further information.

Note that I have not counted anything in the proof within my review. I will spend some time in the next 24 hours to spend longer checking the maths to ensure my understanding is correct.

First of all, we are not sure where the terms “overestimation bias”

I apologise for my missing reference, the term overestimation bias comes from: "The Edge-of-Reach Problem in Offline Model-Based Reinforcement Learning", Sims et al.

Therefore, the claim that “This is not what is shown in this paper” is totally false.

I apologise for any lack of clarity; by this, I mean that since SAC is the policy optimisation algorithm in this work it feels like there should be additional comparison of SAC with other methods for preventing overestimation of values, such as adding pessimism to the value (as in the paper I referenced above).

As it's required by other reviewers, we have included variances of the baselines in the revision.

Thank you for this. I will wait until the new version of the pdf is uploaded before making any edits.

We believe that there’s no issue in this caption since the high standard error is exactly the reason that causes the high standard deviation of our results.

The confusion arose for table 2 as the caption discusses both standard deviation and error. For clarity's sake, I would refer consistently to standard deviation but will not hold this against the paper.

It's been removed in the revision and there's no need to read too much into it.

I am glad to know that this was a mistake and has been corrected.

since the word “interaction” naturally implies a two-way relationship.

Sure, but conventionally the 'agent interacts with the environment' makes more sense due to the suggestion of an active subject and passive object. While it is true that the environment in turn interacts wiht the agent, it feels clunky to write this was as the agent is the active one in the relationship.

That's not an issue, since the name of an algorithm is not necessarily its acronym.

I only flagged this for stylistic reasons and do not consider it in my score (and agree that it would be wrong to as it is not necessary for the anem and acronym to match).

We didn't see any problem in that sentence.

Apologies. I think my confusion arises due to the : introducing a list, but then the sentence continuing after; due to the colon, it reads as three things in a list, and doesn't amke grammatical sense to continue the sentence after the second item in the list. I think: ' As shown in Figure 2, the two key components in ADEPT are policy evaluation on a guided diffusion world model and importance-sampled world model updates. These work in a closed-loop operation through the whole training process' would read better personally, for instance.

Discrepancy is already defined in section 4.1

Thank you for clarifying.

already defined in Section 3 as the horizon, Paragraph named “Offline RL using World Models”.

I apologise for missing this, $H$ is correctly defined and I missed this on my read through.

This sentence seems somehow confusing. Nevertheless, we revised our paper to bold only the significant results.

I apologise for any confusion caused, but appreciate your new bolding of signifcance and look forward to seeing the new manuscript.

(Response to all questions)

These are enlightening and fit what I was expecting. THank you for your clarification.

We attached the comparison results in the response to Reviewer C8UH. Please check that.

Thank you. It certainly seems like ADEPT has significant improvement over PGD.

We thank you for evaluating our work. We now answer the following concerns raised in the review.

Minor feedback: Thanks for catching these! We have addressed these in the new revision.

Weakness:

Though the results are impressive, my primary concern is that the method appears to be very similar to that proposed in [1,2]. It seems the key difference from prior works is the proposed mechanism for realigning the diffusion model's predictions with the changing policy during training. However, when this mechanism is ablated in Figure 3, it appears to only significantly improve performance in one of the four datasets on halfcheetah (medium-replay). The reader would be able to better understand the performance gains expected from this method w.r.t. the methods from [1,2] if they were implemented as baselines, but unfortunately they are not. A more thorough comparison of the authors proposed method with those from [1,2] would improve the paper, and leave the reader more confident that their proposals are a concrete step forward from these works.

Although ADEPT, DWM and PGD both use diffusion models as the world model in offline reinforcement learning, there are still many fundamental differences. The uniqueness of ADEPT comes from 2 key components as we mentioned in the paper: the uncertainty estimation and penalty based on the diffusion model itself, and the world model alignment via importance sampling. To the best of our knowledge, we are the first to apply these two methods in model-based offline RL and provide theoretical analysis for that. Besides, DWM and PGD design their diffusion model as multi-step planners, while we adopt it as a single-step planner for policy evaluation.

Questions:

How does the performance of ADEPT compare to past works that utilise diffusion models as world models for policy training on D4RL (namely DWM (Ding et al., 2024) or PGD (Jackson et al., 2024)?

The performance comparison of ADEPT and two works mentioned are as follows:

The results are from their original paper cited by the reviewer. Note that PGD didn’t report their performance on medium-expert dataset, while DWM didn’t report their performance on random dataset and their results are rounded to the nearest whole number. Compared to these two results, our method shows significant advantages in every environment and dataset.

Are the importance-sample updates more important to final performance in environments other than halfcheetah?

We provided additional results in the revision presented with aggregated statistics as suggested by Reviewer PaFJ. We observed a similar phenomenon in the other two environments that importance sampling mainly improves the performance on medium-replay dataset, where the behavior policy is more stochastic and the policy shifting problem is more severe ($\hat{\varepsilon}_p (\pi)$ is high), while having limited effect on datasets collected by more determined and nearly optimal policies as medium-expert ($\hat{\varepsilon}_p (\pi)$ is low). These results further support our theory that the importance sampling method could largely solve the distributional shift problem by continuously aligning the world model with the new distribution under the target policy.

Where are the uncertainty intervals for the baselines in Tables 1 and 2?

Thanks for highlighting that. The intervals are not included at first because we follow the presentation style of existing works on diffusion-based offline RL [1-3] and try to make the table look more concise. In the revision, we add the uncertainty intervals for baselines in Table 1. However, the original papers on baselines didn’t provide the uncertainty intervals for environments in Table 2. Thus, we decide to move it to the Appendix as secondary results.

You set $H$ to 5, how does performance change with different values of $H$?

That’s a very good point. The value of $H$ could affect the distribution of synthetic policy significantly since it generally determines how far the trajectory could deviate from the offline dataset. Setting a large H has its pros and cons. On the one hand, the model-generated data could be closer to the true distribution under the target policy. However, compounding errors will grow rapidly even with uncertainty estimation, which could counteract the benefit and get exploited by the policy to gain plausibly high returns. We select different values of H and show the results as follows:

Based on these results we found that the best pick of $H$ varies on different datasets and environments. For this study, the value of horizon is not the emphasis in our work, so we simply choose $H=5$ as a reasonable value for all datasets.

[1] Janner, Michael, et al. "Planning with Diffusion for Flexible Behavior Synthesis." International Conference on Machine Learning. PMLR, 2022.

[2] Ajay, Anurag, et al. "Is Conditional Generative Modeling all you need for Decision Making?." The 11th International Conference on Learning Representations (ICLR), 2023.

[3] Wang, Zhendong, Jonathan J. Hunt, and Mingyuan Zhou. "Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning." The 11th International Conference on Learning Representations (ICLR) 2023.

Hi authors, thanks very much for your hard work addressing my (and the other reviewers') questions.

You have responded convincingly to my important questions which I appreciate. In particular, the extended ablation study (Section 5.2 / Figure 4) is a significant improvement on the original paper. I notice that you have not included the above comparison to PGD and DWM in the updated paper; I would recommend that you add this, even if just as an appendix, because these are important recent works that readers will be keen to see a comparison against.

Because you have addressed my important questions I will update my score from 5 to 6 accordingly.

Finally, I was concerned by Reviewer g96b's findings that you had cited a series of irrelevant papers in the appendix in what looked like an attempt to boost citations. Clearly this behaviour can't help but raise doubts about the scientific integrity of the rest of the paper. I appreciate you have stated this was a mistake, but it not clear how such citations could appear without the authors' knowledge. I'm not sure what the protocol is here, so I will leave it to the AC to decide how to proceed.

We thank you for your valuable feedback. We address your comments in the following.

Q1/weakness 1-d: Assumption 4.3 is based on the theory that the diffusion world model prediction error, or intuitively the uncertainty (LHS of equation 8, 9) could be estimated by the diffusion ensemble variance, defined as the discrepancy (RHS of equation 8, 9). Similar assumptions are already used in classic model-based RL algorithms, such as the “assumption 4.3” in MOPO[1] and the “equation (3)” in MOReL[2], while both papers use MLP as the world model and bootstrap ensembles as the uncertainty estimator, and provide detailed theoretical and empirical analysis. Furthermore, the usage of bootstrap ensembles for uncertainty estimators has been justified theoretically in [3] and empirically in [4], while diffusion models also show a high correlation between uncertainty and discrepancy, presented empirically in recent papers [5-6]. Based on these facts, we believe that assumption 4.3 is a reasonable assumption. Besides, we didn’t give upper bounds for \alpha_m and \alpha_r in the paper, meaning that they could be set large enough to satisfy Equation 8 and Equation 9 for all possible states and actions.

Q2/weakness 1-a: When using a world model to guide policy evaluation and learning, there are three sources of errors that may affect the resulting return gap, known as the model transition error, reward prediction error, and the policy distributional error. Our work is the first to consider all three sources of errors in diffusion world models and to quantify their joint impact on return. This result is stated in our main theorem. Existing work has only considered a subset of these errors. In particular, the work in [7] assumes that the reward function is known but in our paper, we take the reward prediction error and uncertainty penalty into account. Besides, MOPO[1] directly adds the uncertainty penalty to the reward, and MOReL[2] only uses it to determine the terminal signal. Our work provides a theoretical analysis with a new bound on the return gap.

Q3: As we responded in Q2, the part of the uncertainty penalty on rewards is specific to our algorithm. If assumption 4.3 is removed, then the proof could lead to another looser bound which is generally true for all other models.

Q4: Not sure we fully understand this question. The word “dispersion” is never used in our paper. If you mean “discrepancy”, it is defined in Section 4.1 as $d_\theta (s_t, a_t)$ and it’s not the standard deviation. Additionally, the same definition can be seen in MOReL, while MOPO did use the standard deviation as discrepancy.

Weakness 1-b: We stated that $M$ is a Markov Decision Process (MDP) in the first sentence of Section 3, which includes transition probability $P$, initial state distribution $\mu _ 0$, and other elements. Therefore, in Equation 5 $s_t$ is sampled from the marginal probability distribution of s in time step $t$ under MDP $M$ and policy $\pi$, with $s_t \sim P(\cdot |s_{t-1}, a_{t-1})$, $a_{t-1} \sim \pi( \cdot | s_{t-1})$ and $s_0 \sim \mu_0$. If it’s necessary, we can modify our paper and clarify this in the main text.

Weakness 1-c: We don’t agree that equation 6 is ill-defined. The TV-distance is based on the joint distribution of s_t and a_t instead of the reward function itself.

Weakness 1-e: Thanks for your advice. We will adjust the theoretical part to improve readability. Also in fact, we have stated that “we give an outline of our proof and display the details in the Appendix” in the first sentence of Section 4.3.

Weakness 2-a, b, c: Thanks for the question. We modified our paper by bolding the statistically significant results instead of the max averages. The variances of baselines are not included at first because we follow the presentation style of existing works on diffusion-based offline RL [8-10] and try to make the table look more concise. We have added the variances of the baselines in Table 1 as the main results in the revision. The original papers of baselines in Table 2 didn’t provide their variances. Thus we left Table 2 unchanged and moved it to Appendix. It’s not surprising that the improvements achieved by our solution fluctuate across different environments and datasets. We would like to note that almost all previous work in this area has demonstrated similar (if not lower) improvements over existing methods. However, we need to point out that the average score is still significant in Table 1, which supports the claim that our algorithm outperforms other SOTA algorithms. Table 2 shows the performance of other environments with sparse reward functions, in which high variance is almost unavoidable.

Weakness 2-d: We did analyze in the ablation study that the uncertainty penalty can improve the stability of training in all kinds of datasets, and importance sampling has more impact on suboptimal datasets such as “medium-replay” in which the distribution is more scattered while having limited effect on nearly optimal datasets. Besides, although we use the same hyper-parameters of uncertainty penalty for all datasets, the optimal values of these hyper-parameters vary among different datasets, since Figure 4 in the Appendix shows different scales of uncertainty and discrepancy in different datasets. However, an analysis of the dataset itself is not the emphasis of this paper.

[1]Yu, Tianhe, et al. "Mopo: Model-based offline policy optimization." Advances in Neural Information Processing Systems 33 (2020): 14129-14142.

[2]Kidambi, Rahul, et al. "Morel: Model-based offline reinforcement learning." Advances in neural information processing systems 33 (2020): 21810-21823.

[3]Peter J Bickel and David A Freedman. Some asymptotic theory for the bootstrap. The annals of statistics, pages 1196–1217, 1981.

[4]Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, pages 4754–4765, 2018.

[5]Berry, L., Brando, A. & Meger, D.. (2024). Shedding Light on Large Generative Networks: Estimating Epistemic Uncertainty in Diffusion Models. Proceedings of the Fortieth Conference on Uncertainty in Artificial Intelligence, PMLR 244:360-376, 2024.

[6] Shu, Dule, and Amir Barati Farimani. "Zero-Shot Uncertainty Quantification using Diffusion Probabilistic Models." arXiv preprint arXiv:2408.04718 (2024)

[7] Janner, Michael, et al. "When to trust your model: Model-based policy optimization." Advances in neural information processing systems 32 (2019).

[8] Janner, Michael, et al. "Planning with Diffusion for Flexible Behavior Synthesis." International Conference on Machine Learning. PMLR, 2022.

[9] Ajay, Anurag, et al. "Is Conditional Generative Modeling all you need for Decision Making?." The 11th International Conference on Learning Representations (ICLR), 2023.

[10] Wang, Zhendong, Jonathan J. Hunt, and Mingyuan Zhou. "Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning." The 11th International Conference on Learning Representations (ICLR) 2023.

Dear authors,

thank you for the clarifications and the thorough response. I have also read the other reviews. It seems that reviewer PaFJ had similar concerns about the clarity of section 4 and questions the dependence on diffusion models. Reviewer g96b seems to have similar concerns about statistical validity.

I appreciate you working on updating the statistical significance in the plots. Also, thank you for explaining the notation of the expectation. I read the commata as dividers. I believe it might be clearer if parenthesis are used around (s, a). After reading the responses to my questions, I will provide some more detail to my review.

Assumption 4.3. The reason why I asked specifically about this assumption is because it bounds the TV distance of 2 distributions by the expected distances over states. This seems to disregard the presence of aleatoric uncertainty and only works with epistemic uncertainty. This is the reason why I am also asking about what benchmarks the method is applicable to. I believe the current manuscript does not differentiate between these uncertainties. The model’s performance could be perfectly accurate but the distance between two next states could still be arbitrarily large since there are no assumptions on the MDPs that are being considered. Thus, it seems that this method can only work in deterministic MDPs or possibly MDPs with Lipshitz transitions which would require proof. Lipshitz transitions which have been studied before [1]. This would be a big limitation and should be addressed in future versions of the paper.

An example: Consider an MDP where the largest reward can be obtained in using a state-action pair whose true distribution over next states has high variance. The method would actively penalize the reward for going to these high reward states.

Other works attempt to get around this issue by obtaining a distribution $Q$ over transitions $P$ and measuring the variance of that distribution $Q$. That is the ensemble methods that are being referred to in the manuscript. I don’t think it is easily possible to achieve such estimates by sampling from the same distribution $P$.

Theoretical Bounds First, as I mentioned, the theory of sample and computation complexity is old and goes back at least 25 years. I do not agree that a method from 2020 is a classical method. After a closer look I realized that this paper only cites work from 2018 and newer. If the manuscript ought to make a strong theoretical contribution, I recommend it include relevant literature that covers model-based sample complexities and error bounds which would allow for comparisons. Lemmata such as 4.7 and 4.8 are standard tools in these works and the presented manuscript does not make it clear that these tools exist and are being reused.

On the novelty of the result: It is well known that if the TV distance on transitions and returns can be bounded, one can achieve low error. Again this literature goes back to at least the paper I cited from 1999. A fundamental challenge is to develop algorithms that effectively minimize this distance.

1. The presented work assumes that the error is bounded by some arbitrarily large quantity. Given a high dimensional state-space, the distance between two states generated by a random model ought to be larger than 1 which would make the assumed bound vacuus. I can always get a guarantee that the difference in returns is lower bounded by the maximum return. (i.e. if the TV distance is bounded by 1 this is basically the second term on the RHS of equation 11). Stating that alpha can be picked arbitrarily large does not resolve this issue; it might even make it worse.

2. In order to make a theoretical contribution, it would be relevant to demonstrate that the presented method can minimize the true TV distance which I am not convinced of. This is due to the issues with uncertainty that I elaborated on earlier.

3. The reward is often assumed to be known as estimating the reward is information theoretically easy if one can estimate the transitions. Usually, no novel technical extensions are needed to include it.

Equation 6 I’m not sure what joint distribution the rebuttal is referring to. In the manuscript R is defined as a function. If the notation is being overloaded here, that should probably be clarified.

Ablation study I appreciate the pointer to the ablation study but it is still a study on maximum rewards. To give an example for the point that I raised. Since the paper is trying to mitigate o.o.d. estimation, it would be interesting to see how for instance Q-values of o.o.d. state-action pairs behave. This would be an experiment related to reviewer g96b's point on overestimation.

Dispersion I was referring the the variance measure in the experiments. The paper reports mean and standard deviation. I was asking why standard deviation is used.

[1] Lipschitz Continuity in Model-based Reinforcement Learning. Kavosh Asadi et al.,ICML 2018

We thank you for your insightful and encouraging review, which inspires us for our future research. We address your concerns in the following.

Weakness 1: The adoption of diffusion models in this paper is mainly to utilize its denoising procedure to estimate the uncertainty and augment the generalization ability. The uncertainty estimation implies is highly correlated with the prediction error of the diffusion model. By applying penalties on uncertainty, the target policy is prevented from exploring unseen states and action pairs with overestimated returns. The method of uncertainty estimation in this paper is exclusive to diffusion models. Nevertheless, it’s true that other world model types with traditional uncertainty estimation methods such as ensembles and dropouts could also be combined with importance sampling. We would consider it as a future work and add some discussions in the paper. Thanks for your valuable suggestion.

Weakness 2: Thanks for your suggestion, we have shortened Section 3 in the revision to provide more readability.

Weakness 3: We are glad that you noticed these. The effectiveness of these two mechanisms is well supported by our experimental results. We test the average reward prediction error over 100k samples, random policy, and the same diffusion world model with a horizon of 5, and show the results as follows:

Based on these results, we found that introducing $\hat{ s } _ {t+1}$ into the reward model could largely reduce the prediction error in all datasets by 10% to 40%. A possible explanation for this is that in most environments, including Mujoco locomotion tasks, the true reward functions are defined directly with $s_t$ and $s_{t+1}$. One common example is defining the moving distance between two states as the reward. Therefore, by introducing $s_{t+1}$ in the input, the trained reward model becomes more similar to the true reward function, thus has more generalizing capability.

The terminating based on high uncertainty is first introduced in MOReL[1] for offline reinforcement learning, and it’s also a crucial component to ensure training stability in our algorithm. Figure 3 in the ablation study shows that without uncertainty penalty (including termination), the training fails in all medium, medium-replay, and medium-expert datasets in halfcheetah environment. Same phenomenon is observed in hopper and walker2d environments. We also notice that a threshold too low also leads to poor performance, since the synthetic trajectories are strictly constrained near the distribution of offline dataset.

Weakness 4: Thank you for this suggestion. We add the aggregated statistics for both Table 1 and the ablation study in the revision. However, since D4RL is a widely used benchmark in offline reinforcement learning, it would be convenient for the community to compare the performance with exact numbers present in the form of tables. Therefore, we decided to keep Table 1 and Table 2 like many previous works did and move Table 2 to Appendix.

Minor Remarks: We appreciate the reviewer for highlighting these, we adjusted Figure 1 and Figure 2 to make it easier to understand. Figure 3 is changed into aggregated statistics including the results in the other two environments. The mentioned typos are also addressed in the revision.

Questions:

I like the idea of using IS to adapt the world model but the reasoning is not exactly clear to me. Usually, we want to use IS to estimate some variable under an unknown distribution by reweighing it by some other known distribution. However, in this case, we can estimate both distributions very well. If my understanding is correct, can you motivate the use of IS more?

This is a good point. Our method is indeed estimating a form of expectation under an unknown distribution. We can consider two trajectory distributions $P_b$ and $P_t$ in the real environment: $P_b$ is obtained by the behavior policy to collect the offline dataset, and $P_t$ is obtained by the target policy we are optimizing. When training the diffusion model, we should minimize an expected loss by sampling trajectories according to $P_t$, so that the learned diffusion model can minimize the error of evaluating the target policy. However, since P_t is not known during data collection and training, the diffusion model is trained using data sampled according to $P_b$ instead. As the target policy further deviates from the behavior policy during the training, this problem becomes more serious. As pointed out by the reviewer, our algorithm leverages IS to address this problem and reweight the samples based on $P_t$. It enables continual alignment between the target policy and world model during training.

From my understanding, IS is a poor technique to use when the two (policy) distributions are very different and that is a completely plausible scenario in your problem setting. Can you explain how you avoid this? Furthermore, have you considered alternative techniques such as MCMC?

Main reasons: 1. IS has been adopted by many existing policy-based offline RL algorithms. It is a simple yet effective solution. We are using IS on the world model to support continual alignment with the target policy. 2. The use of the uncertainty penalty could constrain the policy from deviating too much away from the behavior policy since unseen state and action pairs will get punishment.

Tricks: We set a maximum for the IS weight.

An alternative technique using MCMC seems a very interesting idea, we will consider it in our future research. Thanks for this suggestion.

[1] Kidambi, Rahul, et al. "MOReL: Model-based offline reinforcement learning." Advances in neural information processing systems 33 (2020): 21810-21823.

Weakness 2: Thanks for your suggestion, we have shortened Section 3 in the revision to provide more readability.

I do not see substantial changes in the revised PDF.

Weakness 3: We are glad that you noticed these. The effectiveness of these two mechanisms is well supported by our experimental results. We test the average reward prediction error over 100k samples [...]

Thank you for sharing the results and your following hypothesis. Your statement makes the subtle assumption that more accurate reward prediction will result in better policy learning. While a fair assumption, it is still unclear if it is a valid one [1] [2]. As such, I am not convinced and would like to see an ablation.

Weakness 4: [..] However, since D4RL is a widely used benchmark in offline reinforcement learning, it would be convenient for the community to compare the performance with exact numbers present in the form of tables.

In my opinion, poor historical choices shouldn't be a valid reason for poor future decisions. The new figures used in the revised version (e.g. Fig 3) significantly increase the readability of your work, and I urge you to do the same for Table 1. You can still keep the table in the appendix to make it more easily comparable with prior work. This comment does not affect my score.

This is a good point. Our method is indeed estimating a form of expectation under an unknown distribution. We can consider two trajectory distributions $P_b$ and $P_t$ in the real environment [..]

This point is still unclear to me. I understand that you can treat $P_t$ as an unknown distribution and estimate it via IS. However, what is stopping you from estimating $P_t$ by simply rolling out the policy within your world model?

Main reasons: 1. IS has been adopted by many existing policy-based offline RL algorithms. It is a simple yet effective solution. We are using IS on the world model to support continual alignment with the target policy. 2. The use of the uncertainty penalty could constrain the policy from deviating too much away from the behavior policy since unseen state and action pairs will get punishment.

Thank you for the answer. Unfortunately, I do not find (1) to be valid reasoning, and (2) appears to be an unverified hypothesis.

I will keep my score as is for the time being.

[1] Lambert et al. (2020) Objective Mismatch in Model-based Reinforcement Learning

[2] Wei et al. (2023) A Unified View on Solving Objective Mismatch in Model-Based Reinforcement Learning

Dear all reviewers,

We have updated the revision of our paper. For ease of comparison, we highlight all major changes in blue.

Summary of revisions:

Shorten Section 3.

Rewritten the theoretical part and the proof in the Appendix.

Adjust Figure 1 and Figure 2 as suggested by reviewer PaFJ and g96b.

Add the variances of baselines in Table 1.

Use aggregated statistics to report the results, and extend the compared environments of ablation study to hopper and walker2d.

Add the results of all required additional experiments to the appendix.

Fix all the minors mentioned by the reviewers.

Thanks again to all of the reviewers for providing valuable feedback and helping us improve our work.