Learning a Diffusion Model Policy from Rewards via Q-Score Matching

Thank you for your review and for acknowledging the novelty and potential to inspire future research. We address the points you made a little out of order, as our responses to 1. and 2. build off of our response to point 3.

3: Both referenced papers rely on a behavior cloning term, while the main purpose of this paper is to make a theoretical framework connecting intrinsic structure of a diffusion model policy (here its direct representation of a distribution’s score) to Q-learning. While offline-RL experiments do not fall in this paper’s targeted domain, we do understand the point you make with respect to alternative methods of using the Q-function for diffusion policies. We also discuss sensible alternatives in our responses to Reviewers uSr1 and VG7o. To re-summarize here: a Diffusion-QL-type loss of backpropagating through the computation of Q through sample actions $a$ requires differentiating through a noise model $\phi$ applied on itself $K$ times (for a diffusion model with $K$ denoising steps), while Q-score matching methods only have to differentiate through $\phi$ itself, lending itself more robust and scalable to increases in diffusion model depth, even in the middle of training and evaluation. Policy gradient-type methods avoid the gradient of $\phi^{\circ k}$ problem, but become much less sample efficient as the new expectation (given in eq. (9) of our paper) requires many different samples from the diffusion model of that sampled action to converge to the true policy gradient, a luxury not afforded in online RL.

1: The experiments are not meant to provide a new state-of-the-art model, but to demonstrate validity of this update rule for diffusion model policies in more interpretable online RL environments, which to our knowledge is first demonstrated in literature here. There are plenty of interesting directions to improve on Q-score matching algorithms into a state-of-the-art continuous RL algorithm, but we feel the theory is dense and rich enough to deserve its own paper, and position it to lay the foundation for how to use diffusion model structure to design more efficient diffusion policy training algorithms and spark future research.

2: As discussed with Reviewer VG7o, the main demonstration for these results is not the maximum saturated reward return, but the number of samples it takes QSM to reach a high reward state. The emphasis of this paper is to provide a new theoretical framework that can more fully utilize the structure of diffusion models to write more data-efficient update rules for diffusion model policies.

If you feel any of this discussion should be included in the main body, please let us know. We hope this clarifies your questions and points. If so, please consider modifying your scores and rating to reflect this. And if not, we are of course happy to continue the discussion.

Thank you for your review and for acknowledging the novelty and potential for profound future impact. We are quite excited about this work and hope it sparks new directions in both the theory and practice of diffusion model agents.

To address your three questions, we start with a more general comment about the role of this paper and why we wrote it. To the best of our knowledge, all previous work using diffusion model agents was performed in the offline setting, utilizing a pre-existing dataset to train on. In such settings, it's reasonable to use a number of pre-existing methods. However, in online/off-policy RL, data is a much more valuable resource, and we need to make sure we’re exploiting both its structure and the structure of the agent, here a diffusion model, to the best of our ability. To this end, we develop a theoretical framework connecting the intrinsic structure of a diffusion model policy (here its direct representation of a distribution’s score) to Q-learning.

The experiments’ role is to then verify the theoretical framework’s validity in making diffusion model policies data-efficient; here this translates to the quickness of QSM to reach a high-reward state, which it shows more promising results with respect to. There are plenty of interesting directions to improve on QSM into a state-of-the-art continuous RL algorithm, but we feel the theory is dense and rich enough to deserve its own paper, and position it to lay the foundation for how to use diffusion model structure to design more efficient diffusion policy training algorithms and spark future research.

For the questions you raised:

1. Figure 4 was presented to show that QSM meets some baselines to be an interesting method to explore, rather than claiming novelty in that aspect of our algorithm or improvement over prior methods; here demonstrating that even converged policies still explore alternative paths without any explicit regularization terms. This also addresses the sub-question in question 2 about discrete state spaces: Section 4.3 is solely for pedagogy for the theoretical framework, and practice is solely focused on continuous state/actions spaces where diffusion models are more well-posed.

2. As discussed above, the main parts of these results that we find interesting is the sample-efficiency of QSM in an online RL setting (how fast the model gets to a high reward state). The online RL setting also proposed some unique challenges for comparable methods. Taking policy gradients for example, seeing eq. (9) in our paper, there is now an explicit dependency on inner-sampled actions in the expectation–if we expect finite sample estimations to converge to the true policy gradient, we need many different “diffusion paths” to the same final sampled action, of which we only get 1 in RL. Standard policy gradients are thus quite sample inefficient for diffusion model policy off-policy/online-RL, which we have empirically verified in the paper’s environments. As the purpose of our experiments is to demonstrate validity in Q-score matching as a learning framework, we did not include these in the paper, but if the reviewer finds these experiments would help strengthen the main claims of this paper, we can include these.

3. More empirical investigation needs to happen in order to push Q-score learning based methods for competitive performance. The purpose of this paper is to establish a theoretical framework for using the underlying structure of diffusion models (here the score-based structure) in order to design more efficient algorithms for training such policies. Much of this rich theory is developed and presented in the proofs in the appendix; while important for reading, they are quite long and dense, and would unfortunately greatly hinder readability if crammed in the main body.

We hope we have clarified your questions and points. If you feel any of these points should be added or emphasized in the main paper, please let us know. And if any points remain, we are of course happy to continue discussion.

Thank you for your review and for acknowledging the innovation and thought-provoking aspect. Hopefully we can address your concerns, as there seem to be some misconceptions about the paper.

1. There seems to be a misconception of the main theorem, as we make no dependency on knowing $Q^*$. This is in fact a primary motivation for the theory as described in the 3rd paragraph of the introduction: we want to “score match” to something unknown, so we need to develop new theory that doesn’t rely on knowing anything about $\pi^*$. The main theorem is this, as seen through its contrapositive (and more clearly in its proof): if the score $\Psi$ and action gradient of the Q function with respect to $\Psi$, $\nabla_a Q^\Psi$, are misaligned, then any amount of aligning $\Psi$ to $\nabla_a Q^\psi$ will strictly increase the Q-function towards $Q^*$, and is thus a valid update rule. A large part of this paper's novelty comes from the new math developed to prove this theorem in the appendix, which makes the theoretical role of the QSM algorithm much more clear. We understand the appendix is not required reading, so we have updated the main text to increase clarity.

2. Note in eq. (6) that the state and action are time-separated systems: for every state step, there are many action steps. For any Langevin dynamics-based method (including diffusion models) to work, there need to be enough steps so that the initialization at each state time step makes little difference. While the state-time initialization $a_{t-1}$ is less important, the internal diffusion steps using $\Psi$ are important, and our theory is applied to develop the QSM algorithm to update $\Psi$ based on these internal steps. In summary: initializing between $a_{t-1}$ and Gaussian noise makes little difference, but the latter has some practical benefits, such as the computational graph not connecting through time, and allowing for bang-bang control, etc.

3. The main practical difference: all previous works apply diffusion model policies to offline RL problems, while our work focuses on online/off-policy RL. While both try to maximize Q in some way, these are fundamentally different approaches under the hood. For a diffusion model of depth $K$ and denoising model $\phi$, Diffusion-QL needs to backpropagate through $\phi$ composed on itself $K$ times, resembling gradient structures and issues with RNN’s (vanishing/exploding gradients), while QSM only needs to differentiate with respect to $\phi$ itself (for arbitrary depth $K$) – using our theory to utilize the “inner optimization” of diffusion models allows methods like QSM by design to, unlike Diffusion-QL, train well with deeper diffusion models and improve as $K$ is increased, even during training or evaluation. QSM has implicit entropy regularization as described through the pedagogical example in 4.3, while Diffusion-QL needs explicit regularizers. These benefits come from our presented methodology, and the example algorithm QSM, being derived from and respecting the underlying score-based structure of diffusion models, which the direct Q optimization objective does not capture.

4. The emphasis of this paper is a new theoretical framework to update diffusion model policies, as well as a standard initial implementation in an off-policy RL setting to demonstrate that the theory works as intended and shows its promise, namely when there are less training samples. Much of this rich theory is developed and presented in the proofs in the appendix; while important for reading, they are quite long and dense, and would greatly hinder readability if crammed in the main body. There are plenty of interesting directions to improve on QSM into a state-of-the-art continuous RL algorithm, and many other interesting ways to implement a Q-score matching algorithm (e.g. using $\nabla_a Q$ directly to sample via Langevin dynamics), but we feel the theory is dense and rich enough to deserve its own paper, and position it to lay the foundation for how to use diffusion model structure to design more efficient diffusion policy training algorithms and spark future research. If there are particular insights you hope to gain from Ant and Humanoid, please let us know.

Questions: Note $\pi_\phi$ is not modeled explicitly in diffusion models, but we do give how actions are sampled in the first line within the loop. The other questions should be addressed from the above response.

Hopefully this has clarified some points in the paper and addressed your concerns. If so, please consider modifying your scores and rating to reflect this. If not, we are of course happy to continue the discussion.