Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning

Diffusion models have demonstrated highly-expressive generative capabilities in vision and NLP. Recent studies in reinforcement learning (RL) have shown that diffusion models are also powerful in modeling complex policies or trajectories in offline datasets. However, these works have been limited to single-task settings where a generalist agent capable of addressing multi-task predicaments is absent. In this paper, we aim to investigate the effectiveness of a single diffusion model in modeling large-scale multi-task offline data, which can be challenging due to diverse and multimodal data distribution. Specifically, we propose Multi-Task Diffusion Model (\textsc{MTDiff}), a diffusion-based method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis in multi-task offline settings. \textsc{MTDiff} leverages vast amounts of knowledge available in multi-task data and performs implicit knowledge sharing among tasks. For generative planning, we find \textsc{MTDiff} outperforms state-of-the-art algorithms across 50 tasks on Meta-World and 8 maps on Maze2D. For data synthesis, \textsc{MTDiff} generates high-quality data for testing tasks given a single demonstration as a prompt, which enhances the low-quality datasets for even unseen tasks.

Paper

Similar papers

Peer review

Reviewer gBBv7/10 · confidence 4/52023-06-29

Summary

This paper proposed Multi-Task Diffusion Model, a diffusion-based method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis in multitask offline settings. The performance of the proposed model on Meta-World and Maze2D benchmarks was shown.

Strengths

This paper is the first to achieve both effective planning and data synthesis for multi-task RL via diffusion model and GPT. This paper is well written and easy to follow.

Weaknesses

# The difference between the proposed model and PromptDT is not well explained. # Additional environment and ablation experiments may be more convincing.

Questions

# PromptDT showed performance on MuJoCo. How is the performance of the proposed model on MuJoCo? # This paper seemed to replace transformer in PromptDT with diffusion model. It’s better to explain the difference between the proposed model and PromptDT.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

The related work was introduced in Section 4. It would be clearer to introduce related works before Section 2 Preliminaries.

Area Chair qVd42023-08-17

Dear Reviewer, The author has posted their rebuttal, but you have not yet posted your response. Please post your thoughts after reading the rebuttal and other reviews as soon as possible. All reviewers are requested to post this after-rebuttal-response.

Reviewer nEaR6/10 · confidence 5/52023-07-01

Summary

The paper studies the use of diffusion models in offline multi-task reinforcement learning for planning and synthetic data generation. Both approaches use prompting to encode task-specific conditions for the generative model along with a transformer backbone. In the multitask setting, the approach outperforms prior SOTA baselines on MT50. In the synthetic data setting, the approach improves downstream TD3+BC performance.

Strengths

- Clear and well-written presentation of the method - A thorough set of baselines and strong evaluation on the challenging MT50 benchmark for MTDIFF-p - Elegant method of removing the requirement of one-hot task encoding by transition prompting - Strong results showing positive multi-task transfer in both planning and synthetic data generation, the algorithm neatly extends single-task equivalents in both areas

Weaknesses

- Only custom settings with large amounts of data are considered, it would be useful to understand what the minimum data needed to be effective is. Additionally, it would be useful to evaluate on pre-existing offline datasets for a more representative comparison. - In Figure 6, it would be useful to compare MTDIFF-s-single with the existing single task baseline, Synthetic Experience Replay. [1] - More thorough analysis of data quality would be valuable rather than just downstream performance, e.g. error of the transitions. Furthermore, downstream RL performance is only computed with one RL algorithm. - The two parts of the paper - multitask planning and synthetic data generation - are not necessarily connected. For example, the synthetic data (MTDIFF-s) is used for an entirely different RL algorithm, TD3+BC, and does not contribute to the performance of the planning algorithm (MTDIFF-p). Minor: - Line 242: Typo ‘fine-grind’ - Line 248: the author’s definition of ‘near-optimal’ is the same as D4RL ‘full-replay’ and ‘sub-optimal’ is the same as D4RL ‘medium-replay’. These descriptions may be more clear to offline RL practitioners - Figure 4: choosing one sampled trajectory for each approach is likely cherry-picking, I would suggest rendering N samples - Figure 5: missing standard deviation - Figure 6: missing the baseline without synthetic data - Line 15: Unclear what ‘high-quality’ and ‘low-quality’ mean as it seems MTDIFF-s models the original distribution [1] Synthetic Experience Replay. Cong Lu, Philip J. Ball, Yee Whye Teh, Jack Parker-Holder.

Questions

- How would the prompting perform if the tasks were not readily discernible from the initial transitions? E.g. multi-task settings with sparse task-specific reward - What is the minimum data required to make the algorithm work? - What is the speed of planning and sampling transitions with the transformer model? How does this compare to related methods? - How does MTDIFF-s-single compare to Synthetic Experience Replay? I.e. what is the benefit of modeling full trajectories instead of transitions?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Discussion of limitations hidden in Appendix, should be moved to the main paper.

Reviewer UScg6/10 · confidence 5/52023-07-06

Summary

- The paper investigates the effectiveness of learning a diffusion model for modelling multi-task offline data. To do so, the paper introduces two variants of using a learned diffusion model: (a) by planning over a sequence of actions, (b) generating data, and using the generated data for offline policy optimisation/improvement. - The paper compares the proposed method to various different baselines, and shows decent improvements over all the baselines. - The paper shows the ability of the proposed model to generate useful data, by augmenting kow-quality datasets. - The paper can be seen as an extension of Decision Diffuser for multi-task scenario.

Strengths

- The paper is very well written. - The paper does a good job comparing to various strong baselines, as well as showing the usefulness of the generated synthetic data. - The proposed method is shown to be effective planner for solving multi-task problem on Maze2D and Meta-World benchmarks.

Weaknesses

No such weakness as such.

Questions

- For Table H in appendix, it will be useful to report results for other baselines too. - It will be useful to see how the performance changes as a result of varying the generated data. For example: MTDIFF-S synthesizes 2M transitions to expand the original 1M dataset. It will be helpful to report reports as to how the baselines perform by varying amount of data generated or augmented. - It will also be helpful to see how the performance changes as the MTDIFF-S is trained on less amount of tasks (i.e., instead of 45 tasks it's trained on 5/10/15/20/30 tasks, as the working hypothesis is MTDIFF-S can perform implicit data sharing, so will be useful to verify this by decreasing the number of tasks (figure 6 only shows the comparison for 2 tasks).

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

See Weaknesses and Questions.

Reviewer EBNs6/10 · confidence 3/52023-07-07

Summary

This paper extends diffusion-based planners to multi-task settings by combining prompt learning. Specifically, a few segments from expert demonstrations are used as task prompts to distinguish different tasks and guide the diffusion model to generate task-specific trajectories. A classifier-free guidance approach is used to sample trajectories that yield high expected returns. Unlike previous diffusion-based planners, the authors utilize a Transformer as the denoising network. The experiments conducted demonstrate a marked improvement in performance.

Strengths

This paper is easy to follow. I feel quite enjoyable while reading this paper.

Weaknesses

- The experiment section needs further revision: - Figure 5 x-label is not correct. - What are the quantitative results on Maze2D unseen map - Figure 6, what is the criterion for selecting these 2 tasks? This experiment compares the model training from scratch and the model with pre-training. The results are not that surprising since pre-training provides a better initialization in general. - line 334: The subtitle is confusing. This paragraph is mainly about figure 6, not comparing other augmentation methods. - In Table 2, can the authors add results with 2M random data expanded for S4RL and RAD? Just want to remove the effect of different data sizes. - In Table 2, what are the original results? Is it the maximal return of the dataset?

Questions

Please see weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Reviewer 8zHR7/10 · confidence 4/52023-07-08

Summary

The author propose a diffusion model MTDIFF for multi-task RL. The goal is to leverage diffusion process and transformer backbone to have a sota generative model for RL. The author demonstrates its effectiveness with generative planning on meta world and data augmentation.

Strengths

Combing transformer and diffusion together is a straightforward idea, but the paper excuted it very well. Also leveraging the power of prompting in generative planning is a cool idea. The analysis on how MTDIFF benefit from multi-tasking in data synthesis is interesting.

Weaknesses

The paper also seems unpolished. On table 1, row MTIQL is empty. Also the figure 5 should be "We evaluate these two methods on both seen maps and an unseen map. The average scores obtained in 8 training maps are referred to Figure 5", but I don't know which is PromptDT which is yours.

Questions

With generative planning, can the authors show the visualization of how model get from random noise to optimal trajectory? Like a GIF would be nice.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Area Chair qVd42023-08-17

Dear Reviewer, The author has posted their rebuttal, but you have not yet posted your response. Please post your thoughts after reading the rebuttal and other reviews as soon as possible. All reviewers are requested to post this after-rebuttal-response.

Reviewer nEaR2023-08-13

Thank you

Thank you to the authors for the clarifying response and additional experiments. I have raised my score accordingly.

Authorsrebuttal2023-08-14

Thank You!

We sincerely thank you for your recognition and for increasing your score! We really appreciate it! The constructive suggestions you gave during the rebuttal session are greatly helpful in improving the quality of our paper. Thanks to your time and hard work!

Authorsrebuttal2023-08-17

Thank you!

Thank you for your acknowledgment of the extra experiments we conducted! We really appreciate your effort to review our paper and your recognition of our work. We are glad to see that our paper's quality has greatly improved with your valuable suggestions. If you have any additional suggestions or questions, please feel free to let us know.

Reviewer EBNs2023-08-17

Response

Thank the authors for the clarification, especially in Figure 6. I'm happy to increase my score from 4 to 6.

Authorsrebuttal2023-08-18

Thank you!

We sincerely thank you for engaging with us and for increasing your score! We really appreciate your effort to review our paper and your recognition of our work! The constructive suggestions you gave during the rebuttal session are greatly helpful in improving the quality of our paper. Thanks again for your time and hard work!

Reviewer gBBv2023-08-18

Thanks the authors for the detailed response and additional performance experiments. I’ll keep my score “accept".

Reviewer 8zHR2023-08-18

Thanks for rebuttal

I ve read the rebuttal and my score maintains

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC