Task-agnostic Pre-training and Task-guided Fine-tuning for Versatile Diffusion Planner

Dear Reviewer V32c,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

W1: ... the summation of the log probabilities in the reverse diffusion process in Equation (11) lacks clarity, as it differs from the standard $p(a|s)$ formulation commonly used in reinforcement learning.

A1: According to the formulation in Eq.(9), $\pi\_{\theta}(a|s)=p\_{\theta}(\symbfit{a}\_{t}\^{k-1}|\symbfit{a}\_{t}\^k,\symbfit{s}\_{t})$ which is consistent with the standard policy formulation. In other words, for the denoising MDP, we model the policy network $\pi$ as our pre-trained diffusion denoising network $p_{\theta}$. Similar to DDPO[1], when using DDIM sampler in our experiments, the denoising policy $\pi$ simplifies to well-defined isotropic Gaussian distributions through this formulation.

W2: In Equation (12), the trust region loss is applied to the diffusion process rather than directly addressing the reinforcement learning process.

A2: Since our diffusion functions as an action planner, the actions used for interacting with RL environments are sampled from the pre-trained diffusion model $p\_{\theta}$. Specifically, the action distribution is represented as $p(\symbfit{a}\_t|\symbfit{s}\_t)=p\_{\theta}(\symbfit{a}\_t^{0:K}|\symbfit{s}\_t)$. After applying the importance sampling trick, the old policy distribution becomes $p\_{\text{old}}(\symbfit{a}\_t|\symbfit{s}\_t)=p\_{\theta\_{\text{old}}}(\symbfit{a}\_t^{0:K}|\symbfit{s}\_t)$. Consequently, the trust-region loss for the RL MDP is equivalent to the loss for the diffusion MDP as defined in Eq. (12).

W3: Over-reliance on BC Regularization ... consider an experiment with only the BC regularization term and no fine-tuning term ...

A3: Thank you for your insightful suggestions. We conduct experiments using only Eq.(14) for fine-tuning, with the results presented below. Directly BC leads to poorer performance, as BC lacks reward labels to effectively guide exploration. Although BC during the fine-tuning phase can access dynamic actions, it is limited to 'imitation' rather than 'evolution,' since the model lacks the ability to differentiate between good and bad actions.

Q1: I assume $\sum_{k=1}^{K} \nabla\log p_{\theta}(a_{t}^{k-1}|a_{t}^k,s_{t})$ is modeling $\nabla p(a_t|s_t)$ in policy gradient for the RL MDP. The probability of a specific sample is known to be intractable in the diffusion process.

A4: As discussed in A1, $p_{\theta}$ in Eq.(11) is our pre-trained diffusion denoising network. Therefore, $\sum_{k=1}^{K} \nabla\log p_{\theta}(a_{t}^{k-1}|a_{t}^k,s_{t})$ is modeling $\nabla \pi(a^k|s^k)$ in policy gradient for the diffusion MDP. Therefore, the calculation for the log-likelihoods is feasible since $\pi=p_{\theta}$ is a well-defined Gaussian distribution.

Q2: In equation (12), is "generating new samples" referring to samples in the RL MDP or the diffusion MDP?

A5: This refers to the samples in the diffusion MDP. Consequently, training with the loss function in Eq.(12) is equivalent to optimizing the diffusion model to focus on higher-reward actions.

References

[1] Black et al. "Training diffusion models with reinforcement learning." In ICLR 2024.

W4.4: ... unfairness ... baselines are trained with sub-optimal data ... offline-RL and BC baselines were trained with expert data or even allowed to interact with the environment ...

A4.4:

• To ensure a fair comparison, we conducted fine-tuning for 100k steps per task in our method, collecting online interaction samples during the fine-tuning process. These samples were then incorporated as a supplementary dataset alongside the original data. Consequently, the resulting dataset expanded from 50M to 50M+100k×50. We subsequently trained the baseline methods on this augmented dataset so that the data used for our method and baseline methods are same. The experimental results, presented below, demonstrate that our method still outperform the baseline methods under this configuration. For MTDIFF, we utilize the default parameters provided by the authors. However, we observe a performance decline on these new datasets, which may be attributed to the increased inclusion of sub-optimal data during our online interaction phase.

• Following [1], we apply MTBC to the entire replay buffer, which can be considered a near-optimal dataset due to the significant increase in the proportion of expert data. We adopt the MTBC results reported in [1] as follows. Our method achieves competitive performance compared to MTBC, which is trained on a larger amount of expert data.

W5.1: Why wasn’t image observation utilized in Meta-World to create an image-based environment there as well? Including more image-based tasks ...

A5.1: Thank you for your insightful suggestions. Since no existing image-based sub-optimal dataset for Meta-World is available, we collect data for 10 tasks by training separate SAC agents for each task and rendering the environments to obtain image data. We then follow the same procedure as in Adroit to convert the images into point clouds and use the DP3 encoder to extract visual features. We use MTDIFF_3D, as described in A2, as our baseline, along with DP3 trained on our sub-optimal dataset. The experimental results are presented below.

W5.2: “training on inferior data is more difficult, and the baselines trained on expert data perform better in this case”—is not entirely convincing ...

A5.2: Thank you for your insightful questions. We observed that incorporating additional expert data during pretraining improved performance on the Hammer task. This suggests that the poor performance may be attributed to an insufficient prior for the Hammer task, which could be an isolated case. We have revised our paper to clarify this and prevent any potential misunderstandings.

W5.3: provide task-specific performance results from Table 1 instead of just the average success rate?

A5.3: Thank you for your insightful questions. The varying difficulty levels of tasks impact their corresponding success rates. We provide several example tasks to compare the performance of our method with current baselines. While the success rates of all baselines decline when faced with more challenging tasks, our method consistently outperforms them. Moreover, our method can complete some tasks that baselines fail (e.g. basketball). The detailed performance for each task of our method is presented in Appendix G of our revised paper.

Dear Reviewer 5i69,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

W1.1: There is no experiment using only single-task sub-optimal data for pre-training.

A1.1: Thank you for your insightful suggestion. There are two main reasons that we adpot multi-task data for per-training instaed of single-task:

• We report the success rates achieved after fine-tuning models pre-trained on single-task sub-optimal data (SODP_single) as follows. For some challenging tasks, directly pre-training with single-task sub-optimal data can result in poorer performance and may fail to facilitate online fine-tuning (e.g. basketball).

• We conduct experiments to evaluate the performance of the pre-trained model on fine-tuning for unseen tasks. Specifically, the model is pre-trained on MT-10 data (SODP_mt10) as well as basketball data (SODP_bas). We then fine-tune the pre-trained model on 3 unseen tasks. The success rates achieved after fine-tuning is shown below:

In conclusion, pre-training on multi-task data not only facilitates downstream task learning but also improves generalizability to unseen tasks, as multi-task data provide a broader range of action distribution priors compared to single-task data.

W1.2: The choice of using 50% and 30% of experiences ... results from varying the ratios of experience while keeping the total number of pre-training transitions ...

A1.2: Thank you for your insightful suggestions.

• We utilize 30% experiences as it is sufficient to achieve strong performance on Adroit tasks. We also use 50% experiences for pre-training and the success rates of the fine-tuned models are shown below. We make a trade-off between fine-tuning performance and the quality of pre-training data

• To maintain a consistent total number of pre-training transitions, we select the final 50% of the converged SAC data, instead of the initial 50% (used as the sub-optimal dataset), as a near-optimal pre-training dataset for the Meta-World 10 tasks. Following the procedure outlined in our paper, we report the success rates across three seeds as follows. Incorporating a higher proportion of optimal transitions can enhance performance.

Thank you for the responses! They addressed some of my questions, and I have adjusted my score accordingly. However, a few answers remain unclear, and I have additional questions. If these points can be resolved, I would happily reconsider my score further.

Q1:

Could you provide more details on how multi-task pre-training is conducted on Adroit? You explicitly mentioned that "the action spaces of different tasks vary," but further clarification on how this is managed during pre-training would be helpful. Additionally, given the variation in action spaces across tasks in Adroit, how does multi-task pre-training transfer knowledge when the meaning of actions differs across tasks?

Q2:

I am unclear on how fine-tuning is performed across different environments. Please correct me if I misunderstand your evaluation methods. Specifically, are SODP fine-tuning experiments on state—and image-based Meta-World conducted in a single-task or multi-task setting? Suppose fine-tuning is done in a single-task setting. In that case, comparable baselines trained under single-task settings should also be included, as multi-task BC or offline RL methods do not always outperform single-task methods. This would make the comparisons with SODP fairer, especially when fine-tuned with single-task rewards.

Q3:

Why is DP [4], one of the state-of-the-art baselines for imitation learning, not included as a baseline in the Meta-World state-based setting? My understanding is that DP can be applied in both state- and image-based scenarios.

Q4:

For some challenging tasks, directly pre-training with single-task sub-optimal data can result in poorer performance and may fail to facilitate online fine-tuning (e.g., basketball).

Can you clarify whether most tasks perform better in multi-task pre-training than in single-task pre-training? Are the reported results specific to the state-based setting of Meta-World?

Q5:

We apply MTBC to the entire replay buffer, which can be considered a near-optimal dataset due to the significant increase in the proportion of expert data

For Meta-World, all baselines and our pre-training stage are trained on the same dataset.

Based on these statements, it is unsurprising that BC baselines underperform, as "50% of experiences" consist of both near-optimal and random actions, which is not how BC is typically used. BC baselines should be trained exclusively on expert data. SODP benefits significantly from its ability to collect data with rewards during the online RL fine-tuning stage, giving it an inherent advantage over BC. Even when BC is trained on pure expert data, SODP is expected to outperform it due to this advantage. However, given that MT-BC, trained on the entire replay buffer, performs comparably to SODP—despite the presence of suboptimal data—it raises an important question: If BC were trained solely on pure expert data, could it outperform SODP?

Q6:

Could you provide the return values of other baseline methods for a comprehensive comparison in Table 9? Including the success rates for SODP compared to all baselines would also be appreciated.

Q8: Why does the "drawer-close" task perform worse after fine-tuning with more optimal data? Could you provide an explanation?

A8: According to the results shown in A1.2, the success rate of "drawer-close" are the same (100.0). Did you mean the "dial-turn"? If so, we suspect that this may result from "overfitting" to the reward function. Specifically, at the beginning of the fine-tuning stage, the return and success rate of the sub-optimal pre-trained model are 0.68 and 1056, respectively, while those of the near-optimal pre-trained model are 0.38 and 2477. Since the replay buffer is collected by a SAC agent whose goal is to maximize return, models pre-trained on the near-optimal dataset achieve higher returns. However, higher reward does not necessarily equate to a higher success rate, even though there is a positive correlation between the two. Furthermore, we increased the number of fine-tuning steps and found that fine-tuning on the near-optimal dataset for the "dial-turn" task requires more steps to achieve the same performance as sub-optimal pre-trained model.

Q9: Have you explored pre-training using only the first 10% or 20% of the data from the replay buffer in Meta-World? If so, could this approach still enhance fine-tuning performance or potentially outperform the baselines?

A9: Thank you for your insightful questions. We conduct experiments using the first 10% and 20% of the data from the replay buffer across 5 Meta-World tasks. The success rates after pre-training and fine-tuning are presented below. Pre-training on a smaller percentage of data can still enhance downstream fine-tuning. As the amount of pre-training data increases, the fine-tuning success rate also improves.

Dear Reviewer 5i69,

Thank you so much for raising the score and for providing further feedback. Below, we would like to address your concerns step by step.

Q1: Could you provide more details on how multi-task pre-training is conducted on Adroit?

A1: As discussed in W4.2 and A4.2, we currently perform single-task pre-training for each task in Adroit. The purpose of this implementation is to verify the scalability of our method in complex, high-dimensional observation and action spaces. Adroit serves as a typical benchmark, controlling a more complex dexterous hand, as opposed to the simpler arm models in Meta-World.

Q2: ... are SODP fine-tuning experiments on state—and image-based Meta-World conducted in a single-task or multi-task setting? ... comparable baselines trained under single-task settings should also be included ...

A2: Thank you for your insightful questions. We fine-tune the pre-trained model in a single-task setting. We train a SAC in a single-task setting on MT-10 tasks and the results are shown below.

Q3: Why is DP, one of the state-of-the-art baselines for imitation learning, not included as a baseline in the Meta-World state-based setting?

A3 Thank you for your insightful questions. Directly using DP corresponds to our pre-training stage, which performs significantly worse than our method after fine-tuning. However, as noted in Q5, we previously applied DP only on a sub-optimal dataset. To address this, we conducted experiments using DP on pure expert data from the MT-10 tasks. The results are presented below.

Q4: Can you clarify whether most tasks perform better in multi-task pre-training than in single-task pre-training? Are the reported results specific to the state-based setting of Meta-World?

A4: Yes, we observe that pre-training on multiple tasks facilitates fine-tuning more effectively than single-task pre-training, resulting in higher performance on average. We also find the same trend in image-based Meta-World environments.

Q5: ... SODP benefits significantly from its ability to collect data with rewards during the online RL fine-tuning stage, giving it an inherent advantage over BC ... If BC were trained solely on pure expert data, could it outperform SODP?

A5: Thank you for your insightful questions. We train an MTBC model using pure expert data from the MT-10 tasks, and the results are presented below.

Dear Reviewer kFg1,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

W1: ... no specialized design or methodological advancement in handling such data. The pre-training stage follows standard diffusion model training procedures ...

A1: Thank you for your insightful questions. In pre-training, our objective is to learn useful priors from extensive transition data to facilitate downstream task learning. Given the powerful multi-modal modeling capabilities of diffusion models, we found that a standard diffusion model is sufficient for effectively learning from our multi-task dataset. Thus, we avoided incorporating additional complexities, ensuring simplicity, scalability, and ease of use.

W2: ... the fairness of experimental ... (1) the learning curves in Figure 4 ... (2) whether baseline methods undergo pre-training or have access to online data collection during training.

A2:

• The pre-training phase is also depicted in Figure 4. The learning curve corresponding to pre-training is shown prior to the vertical black dashed line, while the segment between the black dashed line and the orange dashed line represents the fine-tuning phase.
• To ensure a fair comparison, we conducted fine-tuning for 100k steps per task in our method, collecting online interaction samples during the fine-tuning process. These samples were then incorporated as a supplementary dataset alongside the original data. Consequently, the resulting dataset expanded from 50M to 50M+100k×50. We subsequently trained the baseline methods on this augmented dataset so that the data used for our method and baseline methods are same. The experimental results, presented below, demonstrate that our method still outperform the baseline methods under this configuration. For MTDIFF, we utilize the default parameters provided by the authors. However, we observe a performance decline on these new datasets, which may be attributed to the increased inclusion of sub-optimal data during our online interaction phase.

W3: ... control experiments ...(1) fine-tuning with high-quality offline data and (2) direct training with high-quality data without pre-training.

A3: Thank you for your insightful suggestions. We fine-tuned the models for 100k steps across five tasks and collected 200 successful episodes (equivalent to 100k steps) for each task. These datasets were subsequently utilized to train five models independently in an offline setting, using both our pre-trained model (SODP_off) and a model trained from scratch without pre-training (SODP_off scratch). We report the success rates across 3 seeds below. Without pre-training, the model lacks the action priors required to efficiently discover high-reward action distributions. Moreover, directly fine-tuning with high-quality offline data is inadequate, as static reward labels may not offer sufficient guidance for the model in dynamic environments and fail to facilitate efficient exploration.

Dear Reviewer 4UTs,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

W1: Action and State Space Limitations: SODP requires tasks to share the same action space, limiting adaptability to tasks with different action spaces. It’s unclear how the method handles varying state space dimensions.

A1: Thank you for your insightful questions. Following prior work, our goal is to enable a single agent to perform multiple tasks[1][2], so the action space for one agent is consistent (e.g., a 7-dimension action space for a 7-DOF robotic arm). However, recent research has explored cross-embodied robotic control, where the action spaces of different robotic arms may vary. These approaches typically address the discrepancy by mapping the distinct action spaces of each robot into a unified action space, effectively standardizing the action representation across different robots[3]. Then the different action space can also be viewed as one same action space. For variations in state dimensions, visual inputs can be handled straightforwardly by resizing images to a consistent format and processing them with a visual encoder. Similarly, low-dimensional state representations can be transformed into a unified state space, akin to the method used for harmonizing action spaces.

W2: Typographical Errors: There are minor errors...

A2: Thank you for your comment. We have corrected it and checked the entire manuscript again.

W3: Dependence on Online Fine-tuning: SODP relies on an online environment for fine-tuning, which may limit practical applications. The paper also lacks details on the number of fine-tuning steps and associated costs.

A3: Thank you for your insightful questions. Online interaction can indeed be costly and challenging, particularly in high-risk real-world environments where collecting sufficient trajectories is difficult. However, by leveraging existing offline data, a world model can be trained to simulate the transitions and feedback of real environments. Since our method only requires the judgment of actions planned by the model, it is well-suited for integration with a world model. Consequently, the entire learning process can occur within the model's simulated imagination, eliminating the need for direct access to the actual environments.

In Meta-World, we fine-tune each task for $1e^6$ steps whcih requires about 15 hours on single A100 GPU. In Adroit, we fine-tune each task for $3e^3$ steps whcih requires about 3 hours on single A100 GPU.

Dear Reviewer NpAc,

We sincerely appreciate your precious time and constructive comments. In the following, we would like to answer your concerns separately.

W1: Was the fine-tuning stage also applied to the baseline methods? ... this might represent an unfair comparison...

A1: Since the baseline methods (MTDIFF and HarmoDT) are offline algorithms, we did not apply them to online interaction settings. To ensure a fair comparison, we conducted fine-tuning for 100k steps per task in our method, collecting online interaction samples during the fine-tuning process. These samples were then incorporated as a supplementary dataset alongside the original data. Consequently, the resulting dataset expanded from 50M to 50M+100k×50. We subsequently trained the baseline methods on this augmented dataset so that the data used for our method and baseline methods are same. The experimental results, presented below, demonstrate that our method still outperform the baseline methods under this configuration. For MTDIFF, we utilize the default parameters provided by the authors. However, we observe a performance decline on these new datasets, which may be attributed to the increased inclusion of sub-optimal data during our online interaction phase.

W2: Regarding the Adroit task, was the 3D visual feature process (from DP3) also employed for the Diffusion Policy in Table 2? Moreover, what is the performance of SODP when this process is not utilized?

A2:

• DP3 is the variant of Diffusion Policy that incorporates the 3D visual feature processing module, whereas the original Diffusion Policy employs a ResNet18 encoder for visual feature extraction.
• We replace the DP3 encoder with the same visual encoder used in original Diffusion Policy and conduct experiments on Adroit 3 tasks. The experimental results are shown below. Our methods still outperform Diffusion Policy using original image inputs.

Q1: Is the 3D visual feature process not compatible with MTDIFF? Have the authors tried combining the 3D visual feature process with MTDIFF on the Adroit task?

A3: The 3D visual feature processing module can be integrated into MTDIFF. However, due to the incompatibility of MTDIFF with the different action spaces of different tasks in Adroit, we applied this integration to the image-based Meta-World 10 task instead. The experimental results, shown below, further validate the effectiveness of our method in complex high-dimentional scenarios.

Q2: What does the vertical orange dashed line indicating in Figure 4?

A4：The vertical orange dashed line indicates the completion of our fine-tuning process. Our method requires approximately 0.65 million gradient steps in total, encompassing both the pre-training and fine-tuning stages, whereas the baseline methods necessitate 1.5 million gradient steps for training.