Offline Reinforcement Learning with Closed-loop Policy Evaluation and Diffusion World-Model Adaptation
reinforcement learning;offline reinforcement learning;model-based reinforcement learning;diffusion model
Yes
No ethics review needed.
1. 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?
2. 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?
1. The proposed method is novel, well-motivated and clearly explained.
2. Theoretical results to support claims of bounding the return gap between world model and environment.
3. Strong results.
This work attempts to address the important issue of policy distribution shift in offline RL. The authors propose a novel method which uses a world model as a synthetic data generator for closed-loop policy evaluation, where the world model is adapted via importance sampling to the changing distribution of the learned policy. The proposed method is supported be theoretical bounds on the return gap and shows impressive performance on key offline tasks with suboptimal offline datasets.

\* Note that while this paper does fall squarely within my expertise, I was not able to give it the time it deserves and that is reflected in my confidence score.
1. While I like the method proposed by the paper, I don't see a clear dependency on diffusion. From my understanding, the method can be generalized to any world model with some form of uncertainty estimation. As such, I think the paper would be stronger if the method is generalized to any world model. However, I would also be satisfied by an ablation comparing diffusion to other world model types, and clearly showing why diffusion is necessary for the level of performance presented.
2. The paper can be made shorter and more concise to improve readability. I believe Section 3 is mostly unnecessary. You introduce the full notation and background of diffusion but barely sue it in the main text. I recommend shortening it significantly (possibly including it in Section 4) and leaving the full notation and explanation in the appendix as it is necessary for the proofs.
3. In Section 4.1, you state that "Introducing $s_{t+1}$ as an extra input significantly improves accuracy of reward prediction [...]". This is atypical in world model literature and I would recommend backing up this claim with an ablation. Terminating based on high uncertainty is also new to me and a great suggestion! I would also love to see an ablation of this as this changes the distribution of your trajectories significantly. 
4. The main results in Table 1 are poorly presented. I would recommend replacing the table with aggregated statistics as suggested by [(Agarwal et al, 2022)](https://arxiv.org/abs/2108.13264). Table 2 can also be bundled into this figure.


Minor remarks:
1. I found Figure 1 insufficient to grasp the proposed method. I was only able to grasp it after I read the full work at which point the figure has little value. I recommend using a simpler graph with only a few data points and simpler text annotations within the figure itself. (b3) is unnecessary, only showing the (b1) -> (b2) will improve readability.
2. Figure 2 can also be mode more self-explanatory and independent. The two replay buffers are confusing in the figure as they aren't sufficiently explained. The diffusion steps do not necessarily need to be visualized. You probably don't need to spell out each variable being sampled from the buffers. I would also suggest changing the blue box to 'Policy Evaluation within World Model'.
3. Line 228, (I assume) missing 'Section 3'.
4. Line 279, unclear what the loss $l$ is at this point of reading the text.
5. Figure 3, keep y axis the same between all subfigures. Move legend under figures. Possibly remove 'random' as it does not add to the paper.
Yes
No ethics review needed.
Q1: Can you elaborate on assumption 4.3 and why this is a reasonable assumption to make?  
Q2: Can you elaborate on what parts of section 4.3 are claimed to be novel in this work and which parts are taken from previous work?  
Q3: Can you elaborate on what part of the theory is specific to your algorithm and which parts you believe are generally true for all approximate models?  
Q4: Can you elaborate on why you chose standard deviation as a measure of dispersion?
I would like to preface this review by saying that I am not an expert in offline model-based RL. However, I am very familiar with the general online RL and theoretical RL landscape.

1. Clarity  
a) The language in the paper is largely clear and the text has a clear red line that the reader can follow.  
b) The visualizations in Figure 1 and 2 are helpful to understand the approach.

2. Related work  
a) From what I can tell, the work cites the most prominent approaches in offline (model-based) RL and provides a reasonable amount of related work to differentiate its contribution from prior art.

3. Novelty  
a) Based on my (incomplete) knowledge the idea to constrain a model based on the distribution shift of the policy based on offline data only seems novel enough to warrant publication. However, other reviewers may have more insight into this than I do.

4. Experiments  
a) The experiments are conducted with a sizable number of baselines to demonstrate the capabilities. I do think the experiments demonstrate that there might be benefits of the method in lower quality data regimes. However, I have several things that need to be addressed before I can make a certain statement about this. I will come to them later.
This paper operates in the setting of offline reinforcement learning. The paper proposes a new approach that uses an uncertainty-penalized diffusion model as a world model. This world model is used to update the offline policy by constraining a standard SAC's actions via uncertainty estimates. The world model is updated using importance sampling to address the distribution shift of the trained policy from the behavioral policy over time. The paper provides a theoretical analysis of error bounds as well as an experimental section highlighting the approaches performance in comparison with recent offline methods.
1. Mathematical rigor and theory  
a) It is a well-known fact in the learning theory literature that approximate models are sufficient to obtain high return. Analysis on this goes as far back as [1] which can easily be adjusted to fixed distributions. Given that the paper states that this analysis is based on prior work, it is unclear to me what is being claimed to be novel here. There might be subtleties I am missing due to unclarity of notation which I will outline next.     
b) In equation 5, and all following similar expectations, it is not clear what $s_t$ is sampled from. This is quite important given that we are talking about distribution shifts and without this notation being precise it is difficult to determine the correctness.  It is also not clear to me what an expectation over $\mathcal{M}$ means which seems to be a set.  
c) In equation 6, the TV distance is ill defined since $R(s_t, a_t)$ is not a distribution and there seem to be no assumptions on this function anywhere else.   
d) It is unclear to me, why assumption 4.3 is reasonable. I will ask for clarification.  
e) Theorem statement should generally be concise, but they should also be self-contained. In order to understand Theorem 4.5, one would have to read large parts of the paper just to understand the notation. I recommend adjusting this as needed for readability.  The provided proof is also not a proof, but it looks more like a sketch. I recommend stating it as a sketch and referring to the full proof.  

2. Experiments   
a) Experiments over 5 seeds in reinforcement learning can often be misleading given the high variance.    
b) Tables 1 and 2 have the maximum average bolded. This can be misleading as the reader might think these methods are superior as it is not uncommon to bold statistically significant results rather than max averages. I recommend the manuscript is switched to the latter to avoid confusion.    
c) To address the previous point, it is necessary to report variance measures for all baselines and not the presented algorithm. That should in general always be the case. In Table 1, all favorable results on the med-exp dataset are within variance of another approach, at least one of the favorable results of med-rep is within variance of another approach and it is unclear how many of the other results are significant. In Table 2, at least 5/6 results seem to be within variance.  Thus, the claim that the provided algorithm outperforms existing SOTA algorithms is not well supported.  
d) The paper does not provide any additional analysis besides best returns on D4RL and as a result, it is not clear when I should use this method as the results on lower quality datasets are not completely consistent. This makes things tricky because many of the other results may not be significant. One way to remedy the fact that the results are not necessarily much stronger in many cases would be to provide analysis as to *when* this method helps. This could include an experiment that validates the claims about lower distribution shift error or an ablation on the properties of the datasets on which the approach works well.  

Overall, I think this paper offers a neat new idea that can provide insights into how to build purely offline RL algorithms. However, I believe the theoretical section is the weakest part of the submission, and the paper might from this section being shortened. Further, precise notation is required should the authors intend to keep this section. The experiment section could be strengthened by additional analysis that helps understand when this method is useful. I do not think the claim that this method outperforms sota-algorithms is sufficiently supported. I do think that the paper provides an interesting new idea but in the current state I am recommending rejection.

[1] Michael J Kearns and Satinder P Singh. Finite-sample convergence rates for q-learning and indirect algorithms. NeurIPS 1999.
Yes
No ethics review needed.
**Important 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)?
- Are the importance-sample updates more important to final performance in environments other than halfcheetah?

**Less important questions**

- Where are the uncertainty intervals for the baselines in Tables 1 and 2?
- You set $H$ to 5, how does performance change with different values of $H$? 

If the authors answer the important questions I'm more than happy to update my score.
- The experimental results are strong. Most of the required baselines (see weaknesses) are implemented and the proposed method outperforms them in aggregate.
- The theoretical analysis is additive, and as far as I can tell, sound.
- The paper is generally well-written and well-motivated.
- The appendix provides some interesting additional results that support the main body.
This paper proposes a model-based offline RL algorithm leveraging diffusion models. Unlike past works that use pre-trained, frozen diffusion models for generating synthetic rollouts, this work proposes to iteratively align the model's predictions with the evolving policy during training. A theoretical analysis is provided that provides an upper bound on the return gap between an optimal policy trained inside the diffusion model versus the real environment, and experiments are performed on D4RL that show an improvement over canonical offline RL methods.
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. 

**Minor feedback**

- Lines 120-121: Ha & Schmidhuber's paper introduced world models for online RL, not offline RL
- Section 3 title should be "Preliminaries"
- Line 151: "Agent [that] acts..."
- Line 153: transits -> transitions
- Line 161: "real dynamics $P$"?
- Line 272: collects -> collected
- Section 5 title should be "Experiments"
- Line 490 there should be text following _i.e._
- Missing closed brackets in Equations 7 and 8
- Missing brackets in Equation 12
- Line 398: "practically"

**References**

[1] Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng. Diffusion world model. arXiv preprint
arXiv:2402.03570, 2024.

[2] Matthew Thomas Jackson, Michael Tryfan Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson,
and Jakob Foerster. Policy-guided diffusion. arXiv preprint arXiv:2404.06356, 2024
Yes
I am concerned about the referencing used by the authors in this paper. In the appendix (B.1 Baselines) without any detail, there is a long list of cited works (12 papers) which seem to hold no relevance to the paper and all have some degree of common authorship. These are not mentioned elsewhere in the paper and are cited without explanation. This is done in line 903 of the paper:

> In addition, we cite to the following works. (Zhang et al., 2024b; Zou et al., 2024; Gao et al., 2024; Zhang et al., 2024a; Wáng et al., 2023; Fang et al., 2022; 2023; Zhou et al., 2022; 2023; Mei et al., 2023; Chen et al., 2023; 2021)

Observing some of these papers, it is clear that they have no relevance to this paper. For example, one of the cited papers is (line 679):

> Yì Xiáng J Wáng, Zhi-Hui Lu, Jason CS Leung, Ze-Yu Fang, and Timothy CY Kwok. Osteoporotic-
like vertebral fracture with less than 20% height loss is associated with increased further vertebral
fracture risk in older women: the mros and msos (hong kong) year-18 follow-up radiograph results.
Quantitative Imaging in Medicine and Surgery, 13(2):1115, 2023.

Besides the above, a large amount of the papers are about reinforcement learning but in completely different areas to this work (such as MARL). It would be one thing if these papers had been motivated with explanation, and as such had some link to the research at hand. However, their means of introduction and the way they are buried in a subsection of the appendix makes me believe that this is a case of academic dishonesty; in particular, this is an example of self-citation to boost the author's citation count rather than being relevant to the paper, and also risks affecting the paper's anonymity.
Yes, Research integrity issues (e.g., plagiarism, dual submission)
- It seems a bit counterintuitive/non-obvious to compute your done flag as when uncertainty reaches a limit. Does this effectively just lead to truncated rollouts? What kind of bias does this have on bias?
- 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?
- 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.
- 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)?
In general, I find the work reports good results and, while in places is hard to parse, has writing of a reasonable standard. I raise a number of strengths for this paper below.

- I find the related work has generally good coverage over crucial areas, besides what I would class as a significant error in description for [1] (below) and an omission.
- I like the approach of penalising overestimation by relying on the uncertainty of the world model itself. This is quite elegant.
- Using importance sampling to account for distribution shift is also intuitive (though based on the ablation has a relatively minimal impact on performance). I guess it just makes things theoretically sounder.
- I am grateful to see an ablation study, which I think is very inforamtive. However, I think making claims about the significance of important sampling is hard, given that reported values are often in confidence. The ablation study does clearly suggest that the uncertainty penalty in ADEPT contributes to performance improvement.
This paper deals with synthetic data generation in offline RL using diffusion models. Their principal contribution revolves around two methodological changes; introducing a penalty in the reward calculation for uncertain states, and employing importance sampling to account for policy distribution shift. They motivate their design decisions theoretically, and demonstrate strong performance in D4RL, a standard offline RL benchmark.
I have a number of strong concerns about this paper.

- 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.
- 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].
- 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 $\mathbf{x}_0$ with ...') was a bit confusing and could be worded better.
- 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.
- 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.
- There are no error bars for any of the results besides ADEPT's, meaning it hard to see overlapping of confidence intervals.
- 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...'
- 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.


There are also a small number of minor points and typos to highlight:
- In line 37, stating that offline datasets typically consist of limited transitions is tautological; the offline dataset can't be infinite by nature.
- In line 44, the world model does not interact with the policy! The policy interacts with teh world model.
- The acronym of the algorithm (ADEPT) does not fit its name at all really.
- Defining in line 181 that $\hat{P}$ is the transition distribution defined by the world model would be worthwhile.
- Line 190 'is a certain various schedule' doesn't make sense and I am not sure what it is meant to say.
- Line 190 'the diffusion model define another' - firstly, this should be 'defines another'. Secondly, the diffusion model is composed of a forward and backward process, rather than this being defined by the diffusion model itself.
- Line 199 does not make sense.
- The first sentence of the methodology does not make sense.
- Line 346: I don't really know what this means - what is the discrepancy?
- Line 357: 'of between $\pi$' should not have 'of'
- Line 381: 'there existing $\delta$' should read 'there exists $\delta$
- Line 398: 'piratically' is not the correct word I asume.
- Line 411: I don't know what $H$ is and it is not defined I don't think.
- In table 2, all bolded values are with a standard error of each so bolding, which implies significant improvement, is misleading.


[1] Policy Guided Diffusion, Jackson et al. 2024
[2] World Models via Policy-Guided Trajectory Diffusion, Rigter et al. 2024
This paper proposes a new model-based offline RL algorithm ADEPT adopting uncertainty-penalized diffusion world model and importance-sampled world model adaptation, with theoritical analysis and experimental results demonstrated.
@inproceedings{
anonymous2024offline,
title={Offline Reinforcement Learning with Closed-loop Policy Evaluation and Diffusion World-Model Adaptation},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=1zuJZ1jGvT},
note={under review}
}
Generative models, particularly diffusion models, have been utilized as world models in offline reinforcement learning (RL) to generate synthetic data, enhancing policy learning efficiency. Current approaches either train diffusion models once before policy learning begins or rely on online interactions for alignment. In this paper, we propose a novel offline RL algorithm, Adaptive Diffusion World Model for Policy Evaluation (ADEPT), which integrates closed-loop policy evaluation with world model adaptation. It employs an uncertainty-penalized diffusion model to iteratively interact with the target policy for evaluation. The uncertainty of the world model is estimated by comparing the output generated with different noises, which is then used to constrain out-of-distribution actions. During policy training, the diffusion model performs importance-sampled updates to progressively align with the evolving policy. We analyze the performance of the proposed method and provide an upper bound on the return gap between our method and the real environment under an optimal policy. The results shed light on various key factors affecting learning performance. Evaluations on the D4RL benchmark demonstrate significant improvement over state-of-the-art baselines, especially when only sub-optimal demonstrations are available -- thus requiring improved alignment between the world model and offline policy evaluation.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
reinforcement learning offline reinforcement learning model-based reinforcement learning diffusion model
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/febf2c0d26abc654b1de3055d063946dd9d7bcfa.pdf
reinforcement learning
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/fc7d733920ca70158e6c6ea182a0c65cbac95ddb.zip
Offline Reinforcement Learning with Closed-loop Policy Evaluation and Diffusion World-Model Adaptation
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission