Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner

Dear Reviewer 4vpc,

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

W1: would offline+online RL transformers work better?

R1: We extend the multi-task Decision Transformer (DT) following [1] by pre-training it on the same sub-optimal data and fine-tuning each task online for 1M steps on the Meta-World 10 tasks. We consider two types of DT: (1) the vanilla DT, which uses a deterministic policy, and (2) the variant proposed in [1], which employs a stochastic policy. The results are presented below.

W2: ...gradients are back-propagated across the whole denoising rollout... analysis on training and testing efficiency...

R2: Similar to previous work [2], we do not require gradients to be back-propagated through the entire denoising chain at once. From Eq. (11) and (12), each step only requires the final reward $r(a^0)$ to compute the gradient. This allows us to back-propagate the gradient for each step separately, which improves memory efficiency. Furthermore, we employ DDIM sampling during the fine-tuning process, requiring only 10 denoising steps, which reduces both training and inference time. We present the computational efficiency results below. All results are obtained using a single 4090 GPU.

Previous works on fine-tuning text-to-image diffusion models [3] propose a more efficient approach that directly back-propagates the reward gradient through the denoising chain. This method bypasses the need to model the denoising process as a MDP and reduces memory consumption by truncating backpropagation. These reward-backpropagation methods rely on differentiable rewards. However, such differentiable rewards are impractical in RL because reward functions vary across different environments and cannot be easily predicted by a differentiable reward model. Therefore, we use the raw, non-differentiable environment reward and employ PPO for fine-tuning.

Recent studies [4,5] have improved RL-based diffusion model fine-tuning by simplifying the original PPO or introducing dense reward signals at each denoising step, leading to greater effectiveness and better performance. We leave the exploration of optimizing SODP for future work.

Q1: ... the performance of pre-training a Decision transformer with the same sub-optimal trajectories and fine-tuning it with online RL

R3: See R1.

Q2: What is the training and testing efficiency?

R4: See R2.

Q3: I'm wondering if the behavior cloning regularization is needed, if you consider a preference-based RL objectives?

R5: Preference serves as a weak label when the reward function is hard to obtain and must be converted into an explicit or implicit reward signal [6,7] for model fine-tuning. In our setting, the model interacts with online environments and can readily obtain rewards to evaluate the performance of given actions. Therefore, we can directly fine-tune the model using these ground-truth reward labels without the additional steps of constructing and translating preference labels. The recent success of DeepSeek-R1 [8] further demonstrates the effectiveness of using rewards to improve LLM performance. Whether through preferences or rewards, the objective of fine-tuning is a KL-constrained reward maximization problem [6,7], defined as: $L=\mathbb{E} _ {s,a}[R(s,a)]-\beta D _ {\rm KL}[\pi_\theta(a|s)\|\pi_{\rm ref}(a|s)].$ Preference-based methods such as DPO offer an alternative approach to injecting update gradients into the model but still require constraints to regulate the update step size. However, directly computing the KL-divergence term for diffusion models is not feasible due to their complexity. Therefore, we employ BC-regularization as an alternative to achieve similar effects, ensuring that the learned policy preserves the desirable properties of the reference policy while adapting to the specific task.

[1] Online decision transformer.

[2] DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models.

[3] Directly fine-tuning diffusion models on differentiable rewards.

[4] A Simple and Effective Reinforcement Learning Method for Text-to-Image Diffusion Fine-tuning.

[5] Aligning Few-Step Diffusion Models with Dense Reward Difference Learning.

[6] Training language models to follow instructions with human feedback.

[7] Direct Preference Optimization: Your Language Model is Secretly a Reward Model.

[8] DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning.

Dear Reviewer aUqS,

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

C1: The computational costs of pre-training (500k steps) are omitted

R1: Thank you for your question. The data used for pre-training are sub-optimal trajectories. Compared to scarce expert demonstrations, collecting sub-optimal data is more cost-effective. Thus, the pre-training stage is data-efficient, allowing us to use relatively inexpensive training data to obtain a reasonable initial policy. The pre-training stage is a supervised learning process in which the diffusion model is trained offline. This setup enables the use of data parallelism to efficiently fit the action distribution. In contrast, online interaction is considerably more expensive because data must be collected and used iteratively to optimize the model. Consequently, our primary focus is on the RL fine-tuning process, aiming to improve the efficiency of data utilization and reduce the costs associated with online data collection. Overall, large-scale offline data can be leveraged for pre-training, making total computational cost less of a concern. The primary challenge in RL is not the overall computational expense but rather the efficiency of data utilization, as the cost of interaction in the fine-tuning stage is significantly higher.

C2: No experiments on entirely novel tasks to evaluate true generalization

R2: We fine-tune a model on the hammer task and evaluate it on tasks that were not included in the pre-training dataset. The results are presented below.

C3: I think Appendix C is very important to support the claims. The authors had better merge them with the results in main paper.

R3: Thank you for your suggestion. Appendix C indeed provides further evidence supporting the advantages of our methods. We conduct additional ablation studies on various fine-tuning strategies, such as fine-tuning with high-quality offline data and directly applying behavior cloning on online transitions. Additionally, we fine-tune models pre-trained on different datasets to validate the scalability of our fine-tuning approach. We will merge and discuss these ablation studies in the revised version to enhance clarity.

W1: Insufficient analysis of scaling with task complexity.

R4: As shown in Figure 8, all methods achieve a high success rate in the simple movement task handle-pull-side, which requires the arm to pull a handle upward. However, in more complex two-stage tasks such as basketball and hammer, which demand precise arm control to interact with small objects, both MTDIFF and HarmoDT exhibit a success rate below 10%.

Despite this challenge, our method consistently outperforms the baselines, achieving higher performance (e.g., 80% success rate in the hammer task). This improvement can be attributed to the broad action prior learned during pre-training, which enables effective knowledge transfer across tasks. This is particularly advantageous for challenging tasks where task-specific feedback is limited, requiring the model to leverage skills acquired from other tasks. As shown in Figure 7, fine-tuning from scratch results in a substantial decline in success rate for the hammer task, dropping from 80% to 20%.

Q1: How does SODP's computational efficiency compare to baselines in terms of training time and inference overhead?

R5: Please refer to R2 for Reviewer 4vpc.

Q2: the performance of MTDIFF declines markedly. How could it happen? Is it because we do not use a good way to distill the new information into the model?

R6: Thank you for your question. We hypothesize that this is due to the 'mode shift' problem identified in previous theoretical work [1,2], which suggests that diffusion models struggle to fit the target distribution when distant modes exist. The original dataset was collected using an SAC agent with a Gaussian policy, while the newly introduced data was collected using our pre-trained model with a diffusion policy. As a result, the data distributions obtained by these policies differ significantly. When combined into a single dataset, the overall distribution contains two distant modes, which may severely hinder the model's ability to accurately capture such complex distributions.

Q3: Are there specific task categories where SODP struggles that would help identify the advantage of this paper?

R7: As shown in Fig. 8, other baselines often struggle with tasks that require the arm to pick and move a small object (e.g., basketball, sweep), whereas our method achieves higher performance.

[1] On the Generalization Properties of Diffusion Models. NeurIPS 2023.

[2] Statistical Efficiency of Score Matching: The View from Isoperimetry. arXiv:2210.00726.

Dear Reviewer w6L1,

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

W1: Different sets of MetaWorld tasks are selected for different ablation studies ... provide some context information or rationale behind the choice of particular tasks ...

R1: Thank you for your question. Our main experiments were conducted on Meta-World 50 tasks, demonstrating the effectiveness of our method compared to baseline approaches. For ablation studies, such as investigating the impact of different pre-training datasets and evaluating performance on image-based Meta-World, we conducted experiments on Meta-World 10 tasks. This subset is a well-established benchmark that includes tasks of varying difficulty levels and is widely used in multi-task learning research [1,2]. We find it sufficient to assess the effect of the pre-training dataset and the model’s generalization ability to high-dimensional image observations. For other ablation studies, we primarily selected tasks from those used to visualize the learning curves. These tasks were chosen to cover a range of difficulty levels, ensuring a comprehensive comparison of our method's performance.

W2: Table 1 should include more baselines that (1) are diffusion-based and (2) support offline re-training to online RL fine-tuning procedures for better comparison

R2: We extend two diffusion-based methods, DQL [3] and IDQL [4], to multi-task offline-to-online training on Meta-World 10 tasks, following [5]. Specifically, we first pretrain a diffusion-based actor using behavior cloning on an offline dataset and then fine-tune both the actor and critic in an online environment. However, as the results below show, the diffusion actor can still be misled by inaccurate value network estimations, causing rapid forgetting of actions learned from offline pretraining. Furthermore, since goals in Meta-World are randomized, traditional methods struggle to learn how to complete them effectively.

W3: ... provide ablations on additional environments, such as image-based Adroit ...

R3: Thank you for your suggestion. We conduct ablation studies on three Adroit environments, fine-tuning the same pre-trained model with and without BC regularization separately. As shown below, the performance improvement of the model fine-tuned without BC is marginal, whereas the model fine-tuned with BC achieves significantly better results. This indicates that BC can also facilitate learning in high-dimensional single-task settings. This improvement can be attributed to the update stride constraint imposed by BC, which ensures that updates are not too aggressive, thereby preventing the model from becoming trapped in a suboptimal region [6,7].

S1: For offline-to-online settings ... show the performance before and after online finetuning

R4: Thank you for your suggestion. We evaluate the performance of the pre-trained model over 50 episodes across three seeds and report the results below. The pre-training performance of SODP surpasses that of other baselines, demonstrating the advantage of diffusion models in capturing multi-modal action distributions. Additionally, with RL-based reward fine-tuning, the performance improvement exceeds that of traditional offline-online methods, highlighting the effectiveness of our fine-tuning strategy in generalizing to downstream tasks more efficiently.

Q1: In MetaWorld, are dense rewards or sparse success signals used for RL finetuning?

R5: We use dense rewards for RL finetuning because it is necessary to provide the reward signal $r(a _ 0)$ for each samling process to guide the update of $p _ \theta$, as shown in Eq.(12). However, as shown in Table 2, we achieve good performance with only 100k online steps, which is 10% of the unlabeled data used for pre-training.

[1] Hard Tasks First: Multi-Task Reinforcement Learning Through Task Scheduling. ICML 2024.

[2] Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization. NeurIPS 2024.

[3] Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv:2304.10573

[4] Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning. ICLR 2023.

[5] Diffusion Policy Policy Optimization. arXiv:2409.00588.

[6] Trust Region Policy Optimization. ICML 2015.

[7] Proximal Policy Optimization Algorithms. arXiv:1707.06347.

Dear Reviewer aKBX,

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

W1: ...there are areas where it could improve...:

W1.1: ... analyzing ... performance degradation over long adaptation periods

R1.1: As shown in Tables 1 and 2, increasing the fine-tuning steps from 100k to 1M does not degrade performance. After 1M steps, the learning curve converges and remains stable.

W1.2: BC regularization prevents catastrophic forgetting is not explicitly tested ...

R1.2: We fine-tune two models on the same button-press-topdown task using the same pre-trained model, with and without BC regularization, respectively. We then evaluate them on other tasks from the pre-training dataset. As shown below, the performance of the model fine-tuned with BC regularization does not decline significantly, whereas the model without BC exhibits a noticeable performance drop, indicating that it forgets the acquired skills.

W1.3: ...transfer efficiency between tasks

R1.3: As shown in R1.2, the model fine-tuned on the button-press-topdown task also improves its performance on button-press-topdown-wall and button-press-wall, as these tasks share similar goals, allowing the learned skills to be transferred.

W2：...theoretical motivation for pre-training on sub-optimal data ...

R2：Thank you for your question. The theoretical motivation of our method is grounded in the provable efficient reward-free exploration in RL [1,2]. This approach employs a two-stage learning process for developing an RL policy: a reward-free exploration phase and a reward-based optimization phase.

Specifically, in the reward-free stage, the agent optimizes a purely exploration reward, namely, the upper-confidence bound (UCB) term of the value function. Leveraging a UCB-based bonus, the agent is able to collect diverse state-action pairs with broad coverage, thereby providing sufficient information for the subsequent planning phase. In the second stage, the agent performs least-square value iteration (LSVI) by interacting with the environment. The goal of this stage is to conduct value and policy updates for a specific task with a specific reward function. This two-stage learning paradigm, which combines reward-free exploration and reward-based optimization, can achieve polynomial sample complexity according to theoretical results [1]. We believe this approach provides a solid theoretical foundation for our method's empirical success.

Our algorithm essentially follows this theoretical motivation but incorporates more practical implementations. The reward-free exploration stage in [1] is similar to count-based exploration, which aims to collect diverse datasets in the first stage [3]. In SODP, we adopt an alternative exploration method (entropy-based exploration) via a non-expert SAC policy to provide sufficient exploration in the data collection stage. We also pre-train a diffusion planner on this exploratory dataset rather than directly performing reward-based policy optimization. We empirically demonstrate that SODP requires few online samples to achieve strong performance in the second stage. This finding verifies the theoretical results of [1,2], which suggest that such a two-stage learning paradigm is sample-efficient.

Concerning the BC-regularization, we would like to provide further clarification. We follow the principled RL-tuning formulation used in LLMs fine-tuning [4], which is defined as:

$L=\mathbb{E} _ {s,a}[R(s,a)]-\beta D _ {\rm KL}[\pi_\theta(a|s)\|\pi _ {\rm ref}(a|s)].$

Here, the first term aims to maximize the reward during fine-tuning, while the second term ensures that the learned policy $\pi _ \theta$ remains close to the reference policy $\pi _ {\rm ref}$. In the context of SODP, the $\pi _ {\rm ref}$ is the diffusion planner trained on a large-scale non-expert dataset in the first stage, and $\pi_\theta$ is the policy learned for a specific task in the second stage. However, directly calculating the KL-divergence term for diffusion models is not feasible due to their complex nature. Therefore, we adopt BC-regularization as an alternative to achieve similar effects, ensuring that the learned policy retains the desirable properties of the reference policy while adapting to the specific task.

W3：... lacks computational efficiency analysis.

R3: Please refer to R2 for Reviewer 4vpc.

[1] On reward-free reinforcement learning with linear function approximation.

[2] Reward-free exploration for reinforcement learning.

[3] Count-based exploration with neural density models.

[4] Training language models to follow instructions with human feedback.