DyDiff: Long-Horizon Rollout via Dynamics Diffusion for Offline Reinforcement Learning

Authors: Hanye Zhao, Xiaoshen Han, Zhengbang Zhu, Minghuan Liu, Yong Yu, Weinan Zhang
Primary area: reinforcement learning
Keywords: reinforcement learning, diffusion model, dynamics model

Abstract:
With the great success of diffusion models (DMs) in generating realistic synthetic vision data, many researchers have investigated their potential in decision-making and control. Most of these works utilized DMs to sample directly from the trajectory space, where DMs can be viewed as a combination of dynamics models and policies. In this work, we explore how to decouple DMs’ ability as dynamics models in fully offline settings, allowing the learning policy to roll out trajectories. As DMs learn the data distribution from the dataset, their intrinsic policy is actually the behavior policy induced from the dataset, which results in a mismatch between the behavior policy and the learning policy. We propose Dynamics Diffusion, short as DyDiff, which can inject information from the learning policy to DMs iteratively. DyDiff ensures long-horizon rollout accuracy while maintaining policy consistency and can be easily deployed on model-free algorithms. We provide theoretical analysis to show the advantage of DMs on long-horizon rollout over models and demonstrate the effectiveness of DyDiff in the context of offline reinforcement learning, where the rollout dataset is provided but no online environment for interaction. Our code is at https://anonymous.4open.science/r/DyDiff.

Decision: reject
Reviewer scores: 6, 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1729285903855

thread_type:
official_review

<thread_object>
creator: reviewer_kq5g

<content>
summary:
The paper proposed Dynamics Diffusion (DyDiff) to improve offline RL by augmenting the dataset with a single-step dynamics model followed by an iterative refinement process.
DyDiff demonstrates better results on MuJoCo locomotion tasks with ablation studies.

strengths:
- The refinement process combining policy and trajectory diffusion models is novel.
- Theoretical analysis is provided to support the method and results.
- The paper is well-written and easy to follow.

weaknesses:
- The performance improvement is minor (within standard deviation) in most of the datasets.
- The method trains a single-step dynamics model to seed the trajectory diffusion model. The application is limited when the transition distribution coverage in the dataset isn't large enough to train the dynamics model.

questions:
- How does the accuracy of the single-step dynamics model affect the final performance? How would sub-optimal initial policies and dynamics models impact the result?
Are there some visualizations of how the rollouts evolve across different iterations? I recommend showing some visualization results to show how the iterative process improves the rollouts. It is hard to tell if the improvement comes from the iteration or other components with current results.
- In Figure 3a, the method seems sensitive to the iteration time $M$. How could one determine a good value by observing the dataset and the environment?
- In section 5.3, why does DyDiff improve much more than locomotion tasks? Does the improvement come from sparse rewards or easier dynamics in the environments?
- What is the training resource of DyDiff compared with baselines in terms of time and model sizes?

Overall, I have concerns about DyDiff's ability to capture the dynamics when the environments and datasets varied and what is the respective contribution of the seeding trajectory and the iterative process. I'm open to raising my score if the concerns are addressed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730110955629

thread_type:
official_review

<thread_object>
creator: reviewer_3lt8

<content>
summary:
In this work, the authors propose a novel model-based extension applicable to many existing offline Reinforcement Learning algorithms. Specifically, the authors propose utilizing both a single-step dynamics model and a diffusion model for trajectory modelling to extend the offline dataset based on on-policy rollouts. In their method, the single-step dynamics model together with the learned policy is used to generate a trajectory rollout from a state contained in the dataset. The actions selected by the policy during this rollout are then used to condition the diffusion model which predicts corresponding trajectory states with a lower error than the single-step model. Using an iterative scheme between predicting state trajectories, using the diffusion model, and corresponding actions, using the policy, the method generates new state-action-trajectories to be added to the training dataset. A learned reward model is employed to augment the trajectory with rewards and filter out low reward trajectories. The paper contains both experimental as well as theoretical analysis of the proposed method.

strengths:
•	The motivational example in Figure 1 is well done.

•	Nicely written Related Work, giving a good overview over both diffusion in RL and different model-based approaches in offline RL.

•	I like the more informal style of writing, even though clarity could be improved at points.

•	Nice overview in Figure 2. However, I’d suggest adding numbers 1), 2), 3) also to the figure to ease navigating the sketch. Also, variable k is not clear from the figure alone.

weaknesses:
•	Preliminaries are sufficient for a reader familiar with the field but might be brief for readers unfamiliar with either Diffusion Models or RL.

•	The notation in Section 4, especially the trajectory superscript, confused me. It appears the authors use the superscript both for the diffusion step i and “trajectory improvement iteration” k. If I conclude correctly, the superscript is in parenthesis if it regards the trajectory improvement iteration. However, this is not consistent, e.g. from line 274 on in section 4.2. Perhaps a brief introduction to the notation, e.g. in section 4.1, and more consistency could help.

•	I would have liked to see a bit more on what this paper contributes in comparison to concurrent work of Jackson et al.

•	The results in Section 5.2 indicate that DyDiff improves the performance in some settings, decreases performance in other settings and seems almost identical in the rest of settings. Without any significance tests the claim that “DyDiff significantly improves the performance of these algorithms without any additional hyperparameter tuning.” seems over the top. Please provide significance tests and/or rephrase the claim to be more precise.

•	Similar to the results in 5.2, the ablation studies in 5.4 seem to have little reliability due to high variance.

questions:
•	In line 232 you state that “Though we can now use DMs to generate state trajectories, the initial action trajectory for the condition is still left blank.” while Figure 2 already told the reader that the initial action sequence is initialized with the actions from the policy during one-step rollouts. For me the lines 182 to 186 thus caused a lot of confusion, which was then resolved later when reading section 4.2. Is there a particular reason for this structure or could it be adapted?

•	Rather a suggestion: The caption of Figure 2, and the corresponding description in lines 181-186 could be more clear.

•	In Section 5.2 you state that “From the perspective of different base policies, DyDiff exhibits relative incompatibility with CQL”, which seems intuitive, but nonetheless, CQL seems to have similar benefit from DyDiff than the other baselines in Section 5.3. Do you have an intuition on this?

--

Post-rebuttal: While I believe the current revision resolves most of the clarity issues, the reported probabilities of improvement, ranging from 0.52 to 0.6, are not convincing enough to warrant a more favorable evaluation as a "good paper." Therefore, I will maintain my initial evaluation of 6 (weak accept).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730448457209

thread_type:
official_review

<thread_object>
creator: reviewer_bgmc

<content>
summary:
To address the mismatch between the behavior policy and the learning policy, and to alleviate compounding errors induced by single-step dynamics models, in this paper, a method, namely DyDiff is proposed that can synthesize long-horizon on-policy rollouts for offline policy training, which claims to combine the advantages of both the rollout consistency of single-step dynamics models with arbitrary policies and the long-horizon generation of Diffusion models with less compounding error.

strengths:
The paper is well-motivated and aims to address some important concerns in the literature, the ideas presented in this paper are novel, and both experiments conducted on D4RL and theoretical analysis are provided to show the effectiveness and advantages of the proposed approach.

weaknesses:
While the whole framework is based on the claim that DMs have superior modeling capabilities compared with single-step dynamics models, no theoretical guarantees are provided for such claims, e.g., based on this claim, the conclusion of Theorem 1 shows that \epsilon_m \approx \epsilon_d, however, in the one hand, no theoretical are provided, in the other hand, in some cases, this may be not true. Also, it is unclear why just replacing the generated actions with the action of the learning policy, the mismatch issue can be addressed as the learning policy is also obtained from the real training dataset.

questions:
1. In theorem 1, the authors claimed that over long horizons, the single-step error bound \epsilon_m is even less than the accumulative multi-step error bound \epsilon_d as DMs has superior modeling capabilities, and then concluded that DMs are better for long-horizon rollout than single-step dynamics model. However, no theoretical guarantees are provided. 
2. In the proposed approach, states and actions are first generated simultaneously, then the generated actions are replaced with the real ones. Compared with DecisonDiffuser where only state sequences are generated, what are the advantages of the proposed method? 
3. In which conditions, can Assumptions 1 and 2 be satisfied? 
4. In the experimental results, only the training time of DyDiff is reported. The training time of other baselines should also be stated and compared with the training time of DyDiff.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730718603322

thread_type:
official_review

<thread_object>
creator: reviewer_u2eq

<content>
summary:
This paper introduces Dynamics Diffusion (DyDiff), a novel method that leverages diffusion models for synthesizing long-horizon, policy-aligned trajectories in offline reinforcement learning. DyDiff addresses the challenge of policy mismatch by iteratively refining generated trajectories to match the learning policy, thereby enhancing long-term rollout accuracy without relying on single-step dynamics models prone to compounding errors. Through theoretical analysis and experiments on D4RL benchmarks, DyDiff demonstrates its effectiveness in improving policy training across various tasks.

strengths:
1. The idea of this paper is very novel and interesting.
2. The writing is clear enough and easy to follow.
3. The theoretical proof of this paper looks precise and interesting.
4. The experimental results of this paper is solid.

weaknesses:
1. The ablation study result in Fig 3.a looks strange to me. If the core idea is to apply learning policy with the learned dynamics by diffusion models, more iteration times should lead to better performance intuitively. The author's explanation is that more iterations will lead to a higher probability of falling out of the data distribution. This is not so convincing to me. The initial trajectory is generated by the learning policy and the inaccurate dynamics (compared with diffusion models). Thus, the later iteration is more like to correct the dynamics with the diffusion model while remaining the policy to be the learning policy. In this case, having more iterations will not increase the risk of falling out of distribution, while it only improves the accuracy of dynamics. 

Could the authors provide further analyses and experiments to discuss this problem, including providing a more detailed analysis on how the state distributions and the dynamics consistency change during the iteration?

questions:
1. How do you get the learning policy $\pi$ in the generation process? Does it update iteratively with the learning process? If I understand correctly, does it mean that you first initialize the learning policy then generate trajectories with DyDiff and learn a new policy based on the generated trajectories? Also, what is the frequency of policy updates in relation to the trajectory generation process, and how this will affect the overall performance?

2. Why DyDiff can significantly improve the performance on Maze2d compared with Mojuco? Can the authors provide a more detailed analysis between the results on these two environments?

3. I'm wondering if DyDiff can still improve the performance with some other baselines like Decision Transformer.

4. It's better to include the comparison with other data augmentation methods like MTDiff.

5. How about directly comparing with some diffusion-based planning algorithms using the learning policy as guidance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732204163207

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We greatly appreciate all reviewers for your careful reviewing and valuable advice. Based on your suggestions, we have revised our paper, including:
- Revising the possibly confusing introduction to the framework of DyDiff in the caption of Figure 2, Line 182, and Line 232-239.
- Clearing the notations from Line 274, Section 4.2.
- Explanation of why $\epsilon_d$ and $\epsilon_m$ cannot be further decomposed in Line 336, Section 4.3.
- The online RL settings in Section 6 now include a time-consuming limitation.
- More details for preliminaries in Appendix B.
- Explanation of assumptions that explain why we make such assumptions and when they are likely to be satisfied, in Appendix D.4.
- Comparison to MTDiff-s in Appendix E.4.
- Explanation and empirical example of how the error of synthetic trajectories changes over the iteration times $M$, in Appendix E.5.
- Visualization of synthetic trajectories on Maze2D tasks in Appendix E.6.
- Training time and model sizes of DyDiff and SynthER in Appendix E.7.

We mark the modified text or the captions of added sections in red.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732442750504

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
The discussion will end in 4 days

comment:
Dear Reviewers,

We sincerely thank you for taking the time to review our manuscript and provide valuable feedback. Your insights have been instrumental in improving our work.

As the discussion deadline is now only 4 days away, please let us know if we missed anything or you still have confusion. We greatly appreciate your efforts and look forward to receiving your feedback. Any questions to our paper and rebuttal are always welcomed.

Thank you again for your time and support.

Best regards,

Authors of Submission5681
</content>
</thread_object>

</thread>
<thread>
creation date:
1733313742216

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
We would like to sincerely thank the reviewers for their thorough and insightful feedback on our paper. We are pleased to know that our responses have helped clarify some of the concerns and confusions regarding our work. In this paper, we introduce **Dynamics Diffusion (DyDiff)**, a novel approach for synthesizing long-horizon rollouts that are both accurate in terms of dynamics and consistent with the learning policy. We demonstrate the effectiveness of DyDiff both theoretically and empirically, with the theoretical analysis highlighting the potential of general non-autoregressive models in long-horizon generation.

Based on the reviewers’ feedback, our paper's strengths include:
- All reviewers agree that the idea is novel, and Reviewers U2eQ, 3Lt8, and kq5G also note that it is well-motivated.
- Reviewers U2eQ, 3Lt8, and kq5G find the paper well-written and easy to follow.
- Reviewers U2eQ and kq5G acknowledge that the theoretical analysis is precise and supports the method and experimental results.

For the weaknesses raised, we have revised the paper as follows:
- Clarified confusing notations, captions, and explanations of the theory.
- Provided additional background knowledge to support the context of our work.
- Expanded the empirical analysis, including a comparison of DyDiff with a new baseline, MTDiff-s, as well as the visualization of synthetic trajectories.

We would like to thank the reviewers once again for their invaluable constructive comments and suggestions. Their feedback has greatly helped us improve both the quality and clarity of our work. We believe that most of the concerns have been addressed through our responses and the revisions made to the paper. We truly appreciate the time and effort the reviewers have dedicated to evaluating our submission.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper attempts to address the problem in the use of diffusion models for offline RL where the generated trajectories don't align with the learned policy. The approach works by first training a single-step transition model and generating roll-outs using this model and the policy. Next, these trajectories are refined by sampling iteratively from the diffusion model and policy. Finally, generated trajectories are filtered and high-reward ones are stored in a synthetic data which can be used for offline training.

Strengths:
- Reviewers found the approach novel 
- Generating trajectory data that is consistent with the policy is desirable and the proposed approach makes sense
- Theoretical guarantees showed that the non-autoregressive approach works better when $\frac{\gamma}{1-\gamma} \frac{\epsilon_d}{\epsilon_m} > 1$, where $\epsilon_m$ is the TV error of the 1-step transition model against the true model, and $\epsilon_d$ is the same for the non-autoregressive model.

Weakness:
- The chief concern is gains are minor
- Given approach is computationally expensive

The main concern here is whether the gains justify the approach. This is raised both by reviewer 3Lt8 and reviewer kq5G. Further, the proposed approach is much more complex and computationally expensive. Training diffusion models is challenging. For these reasons, while I do see a promise here and an important problem, I am recommending weak rejection. I suggest authors improve their approach to get bigger gains, or simplify it significantly where those gains make sense.