Learning Multimodal Behaviors from Scratch with Diffusion Policy Gradient

We thank the reviewer for their feedback. Due to the rebuttal limit and the number of questions, our responses are concise. We are happy to provide more detailed answers during the discussion period.

The paper is hard to follow ... for improvement.

Based on the reviewer's feedback, we will:

1. Moderate our claim about overcoming local minima and clarify DDiffPG's continuous exploration ability in Sec. 5.3.
2. Fix the typo in the intrinsic reward definition.
3. State clearly scenarios where an imperfect Q-function can lead to mode collapse.
4. Highlight the reclustering process in Sec. 4.2 and the exploration-exploitation balance in Sec. 4.3.
5. Clarify the information used in DTW computation, such as robot positions.

In section 3, ... in general.

In our method and concurrent works (Yang et al. 2023, Psenka et al. 2024), policy update for a diffusion process cannot follow the usual deterministic policy gradient derivation. The stochastic nature of diffusion models poses significant challenges in deriving formal theoretical guarantees of policy improvement. However, our approach updates target actions under the guidance of the Q-function. As the Q-function improves, these target actions are refined accordingly, which in turn trains the diffusion policy on these improved targets. While formal proof remains challenging, this intuitive process reflects an alignment between the policy and the evolving Q-function, suggesting effective policy improvement in practice.

I feel the paper ... moderate the claims a bit.

We apologize for any ambiguity in our claim. We mean that the continuous exploration capability of DDiffPG helps increase state coverage and overcome local minima. This capability comes from the exploratory mode, not from RND alone, as the same intrinsic reward was used for all baselines. DDiffPG maintains an exploratory mode throughout training, allowing it to discover versatile behaviors continuously. Unlike baselines that quickly converge to the first explored solution, DDiffPG keeps exploring, which is crucial for learning multimodal behaviors online. Fig. 6 and 12 support our claim, showing that DDiffPG achieves the highest state coverage, while overcoming the lower-reward goal, reaching the top-left goal with the higher reward, helping to escape the local minimum.

In line 157, is this a typo? ...

Yes, the intrinsic reward is defined as the novelty difference between $s'$ and $s$, which encourages to visit the boundary of explored region.

In section 4.2, ... distribution?

We agree that a perfect Q-function nicely capturing both modes could prevent mode collapse. However, in practice, achieving a perfect Q-function is unlikely since both the policy and Q-function are learned from scratch. This often leads to:

1. the policy being guided towards the first explored goal.
2. the policy converging to the solution with higher Q-values due to the discount factor, even if two modes are initially captured.

DDiffPG's design mitigates mode collapse through mode-specific Q-functions, each focused on a single mode, thereby isolating their Q-values. By sampling actions from the replay buffer rather than relying solely on policy-inferred actions, we ensure balanced improvement across all modes. However, policy improvement alone is not sufficient for multimodal behavior learning, as confirmed by the results of DIPO (Tab. 3 and 4).

Regarding the use of mode-specific ... during training?

Yes, we re-cluster trajectories every $F$ iterations, updating modes continuously. After re-clustering, new clusters inherit Q-functions and $a^\text{target}$ from the most overlapping cluster in the previous iteration. If the number of clusters increases, we initialize new Q-functions. This ensures a one-to-one correspondence between modes and Q-functions without needing to predefine a fixed number of Q-functions. This process is explained in Sec. 4.2 (lines 210-214), with further details in Alg. 2 and App. C. The implementation can also be found in our provided code to the Meta-reviewer under ddiffpg/utils/Q_scheduler.py.

It seems ... during RL?

We balance exploration-exploitation by adjusting the proportion of the exploratory mode during data collection. As explained in Sec. 4.3 (lines 230-233), we condition the diffusion policy on mode-specific embeddings, including an exploratory mode trained by $Q_\text{explore}$. We set the proportion for exploration to $p = \max(0.3, 1/\text{modes})$, ensuring at least 30% of actions are exploratory. Implementation details can be found in ddiffpg/algo/ddiffpg.py (lines 104-129).

In line 175, ... not applicable.

We apologise for any confusion. We only use the robot position for clustering, i.e., the Ant's position and the robot end-effector's position, which are accessible in real-world scenarios. We do not use any information about objects or goals in the environment. In the PDF, we provide an alternative clustering approach by training a VQ-VAE to learn trajectory representations and using the codevectors for clustering, which learns meaningful representations suitable for our approach.

The paper only ... by the paper.

We want to emphasize that the 8 robot locomotion and manipulation tasks are nontrivial. They are all high-dimensional and continuous control tasks with sparse rewards. In the AntMaze tasks, the agent must solve the Ant-locomotion task through joint control, a notoriously difficult continuous control problem, while also learning to steer locomotion gaits to find goals in complex mazes. All tasks are goal-agnostic. The agent must explore to find the goal, making it challenging to learn multimodal behaviors due to the vast exploration space. As the first work exploring online multimodal diffusion policies, we see our current tasks as a strong starting point and plan to extend our approach to long-horizon and vision-based tasks with greater multimodality in future work.

We thank the reviewer for their positive feedback and insightful comments. We want to address the reviewer's concerns and questions as follows.

Several claims require additional support. ... enhance their reliability.

We thank the reviewer for the constructive comment. First, given the RL objective of maximizing the expected return, policies in goal-oriented tasks will converge to a solution once it is discovered. Tab. 3 and 4 experimentally show that traditional RL baselines fail to learn multimodal behaviors even when multiple solutions exist.

Second, DDiffPG maintains an exploratory mode throughout training, allowing it to continuously discover versatile behaviors. Fig. 6 demonstrates that DDiffPG achieves the highest state coverage, supporting our claim. Additionally, Fig. 12 illustrates a local minima scenario where the ant faces two goals. Baseline models often get trapped pursuing the easier, lower-reward goal, whereas DDiffPG continues to explore and eventually learns to reach the higher-reward goal, helping to escape the local minimum.

Third, DDiffPG achieves explicit mode discovery via hierarchical clustering with the DTW metric, preserves the multimodal action distribution by constructing multimodal training batches for policy updates, and improves the modes collectively with mode-specific Q-functions.

In Lines 263-267, ... significance of multiple trajectories.

We agree with the reviewer that presenting the best performance is beneficial. DDiffPG achieves the same success rate as TD3 in Fig. 4 but usually requires more steps to reach the goal in Tab. 3 and 4. In response, we provide the episode lengths that execute only the mode with the shortest path below. We will include these results in Tab. 3 and 4 in the revised version. Note that further training of DDiffPG would lead to further optimizing all paths. We report here the numbers obtained by the policies reported in the paper.

Why does directly ... better illustration.

We thank the reviewer for the insightful suggestions. Directly backpropagating the action gradient through the diffusion model for policy optimization is challenging and can cause instability due to the Markov chain in the diffusion process and its stochastic nature [1, 2, 3]. This issue is evident in the high variance of the baseline Diffusion-QL in Fig. 4. Additionally, we have included plots in the attached PDF to measure the gradient in Diffusion-QL, confirming the vanishing gradient issue.

[1] Psenka, Michael, et al. "Learning a diffusion model policy from rewards via q-score matching." ICML, 2024.

[2] Wallace, Bram, et al. "Diffusion model alignment using direct preference optimization." CVPR, 2024.

[3] Pascanu, Razvan, et al. "On the difficulty of training recurrent neural networks." ICML, 2013.

Unlike the offline setting ... (Lines 126-129)

In the proposed diffusion policy gradient, we can tune the step size $\eta$ of the action gradient for stability. In Section 5.4, we provided a hyper-parameter analysis on the step size and found that a step size of 0.03 can generally work well. However, an aggressive value, e.g., 0.05, could speedup the learning at the early stage but increases the variance. Consequently, we used a step size of 0.03 for all tasks without any further tuning.

Why does DIPO, ... (Lines 134-137).

According to Algorithm 2 in [4], DIPO stores transition $(s_t, a_t, s_{t+1}, r(s_{t+1}|s_t, a_t))$ into the replay buffer, performs gradient ascent on $a_t$, and replaces $a_t$ in the original buffer. Therefore the transition $(s_t, a_t, s_{t+1}, r(s_{t+1}|s_t, a_t))$ no longer aligns with the current MDP dynamics and reward function, due to the replacement of $a_t$. Given that DIPO is an off-policy algorithm, the reuse of these replaced transitions for training the Q-function could be problematic, as the agent is training values of actions that have not been actually played out in the environment, and their true outcome (reward and next state) are unknown.

[4] Yang, Long, et al. "Policy representation via diffusion probability model for reinforcement learning." arXiv, 2023.

Are the mode-specific ... standard policies?

We agree with the reviewer that the mode-specific Q-functions may be applicable to other candidate model parameterizations. The most straightforward example is to learn a separate unimodal actor for each Q-function. However, we advocate for the benefits of learning a unified model:

1. Learning a single model allows to share information (e.g., representations) across modes, which is particularly beneficial for future research, e.g., on image-based-observation tasks [5]. Our multimodal policy learning also draws an analogy to multitask RL, in which the objective is to solve multiple tasks with a single policy rather than separate policies per task, benefiting from knowledge sharing [6].
2. Learning separate unimodal policies would significantly increase the computational time and agent interactions with the environment. With separate policies, we need to iteratively update them; however, only one backpropagation is needed for a single multimodal policy.
3. When a new mode is discovered, our diffusion model continues learning, adding this mode in its landscape, without forgetting previous knowledge. However, if we have separate policies, we have to initialize a new policy and learn from scrach, otherwise additional techniques have to be introduced to transfer knowledge from previous training.

[5] Kalashnikov, Dmitry et al. “QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation.” CoRL, 2018.

[6] Hendawy, Ahmed et al. “Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts.” ICLR, 2024.

We thank the reviewer for the positive feedback and insightful comments. We address the reviewer's concerns and questions as follows.

The proposed diffusion training objective ... further application.

We proposed diffusion policy gradient, a method that combines RL and behavioral cloning (BC) for stable online diffusion policy updates. Inspired by methods like the Deterministic Policy Gradient (DPG) [1], the policy is optimized to follow the action-gradient $\nabla_{a} Q(s, a)$.

Directly backpropagating the action gradient through the diffusion model for policy optimization is known to be difficult and may cause the gradient vanishing problem and instability [2, 3]. This issue was evident in the high variance of the baseline Diffusion-QL in Fig. 4. We have also included additional plots to visualize this issue in the attached PDF. To mitigate this issue, we imitate the $a^{target}$ obtained via action gradient ascent. This way, we obtain an action that follows the updated Q and serves as the target for our BC objective. The solution we provided is general and requires no tuning. This can also be seen in the codebase we provided to the Meta-reviewer.

We validated our method on 8 challenging robotic tasks with high-dimensional and continuous action spaces, requiring no handcrafting of the policy update objectives. The proposed objective only introduces an additional hyper-parameter, which is the step size $\eta$ of the action gradient. In Sec. 5.4, we provided a hyper-parameter analysis on the step size and found that a step size of 0.03 can generally work well. As a result, we used this value for all tasks without any further tuning, proving the generality of the proposed method.

Besides the diffusion ... limited theoretical novelty.

While we acknowledge that some concepts used in our paper have been explored in prior works, the combination of these techniques in an online setting to learn multimodal behaviors using diffusion policy has not been studied before. For instance, DIPO [4] and QSM [2] focused only on the diffusion policy update objective. They do not study how to take advantage of the diffusion model to represent multimodal action distributions, nor do they consider the additional exploration challenges for discovering and learning multiple behavioral modes online. Our experiments show that naively learning a diffusion policy does not yield multimodal policies and might even be an "overkill" if multimodality is not a concern. Additionally, we highlighted the issue of the greedy RL objective in learning multimodal behaviors, which was not examined in prior works.

Code is not provided. ... the code is not convincing.

We agree with the reviewer on the importance of open-source code for the community. We have sent the code link of an anonymous repo to the AC, and kindly suggest reaching out to them for access to review it.

The ablation studies ... should be given.

First, we would like to emphasize that our main contribution is learning diffusion policies that capture multimodal behavioral modes. Tab. 3 and 4 show that our approach uniquely masters diverse behaviors. Despite exploring a larger space, Fig. 4 demonstrates that DDiffPG achieves comparable sample efficiency to baselines across all 8 tasks.

Second, we included DIPO as a baseline, which also serves as an ablation of our framework. As discussed in Sec. 3, Lines 130-137, we modified DIPO for consistency in MDP dynamics and reward function. The DIPO lacks clustering, mode-specific Q-functions, and multimodal batch. However, DIPO fails to explore and learn multimodal behaviors, proving the necessity of our design choices.

Third, we appreciate the reviewer's constructive suggestions and provide additional ablations on mode-specific Qs and the multimodal batch. The attached PDF shows that a single Q-function cannot address the greedy RL objective, guiding the policy toward the first explored solution. Moreover, in the ablation of multimodal batch, we find that while the diffusion policy learns multiple solutions, the modes are imbalanced due to the lack of enforced diversity in the batch, causing minority modes to appear less frequently. Thus, both design choices are crucial for learning multimodal behaviors.

The formatting ... in the appendix.

We thank the reviewer for pointing out these inconsistencies. We will update the caption of Figure 4 to "Performance of DDiffPG and baseline methods in the four AntMaze and robotic manipulation environments." Due to space limitations, we have placed Tables 3 and 4 in the appendix. Since the camera-ready version allows for one additional page, we will include these tables in the main body upon acceptance.

The diffusion ... fine and meaningful.

In our methodology and concurrent works [2, 4], policy update for a diffusion process cannot follow the usual deterministic policy gradient derivation. The stochastic nature of diffusion models poses significant challenges in deriving formal theoretical guarantees of policy improvement. However, intuitively, our approach updates target actions under the guidance of the Q-function. As the Q-function improves, these target actions are refined accordingly, which in turn trains the diffusion policy on these improved targets. While a formal proof remains challenging, this intuitive process reflects an alignment between the policy and the evolving Q-function, suggesting effective policy improvement in practice.

[1] Silver, David, et al. "Deterministic policy gradient algorithms." ICML, 2014.

[2] Psenka, Michael, et al. "Learning a diffusion model policy from rewards via q-score matching." ICML, 2024.

[3] Wallace, Bram, et al. "Diffusion model alignment using direct preference optimization." CVPR, 2024.

[4] Yang, Long, et al. "Policy representation via diffusion probability model for reinforcement learning." arXiv, 2023.