Paper: Learning a Diffusion Model Policy from Rewards via Q-Score Matching (ICLR-2024)

Reviewer comment (Experiments):
Compared to the rigorous theoretical aspect, the experimental section of the paper appears relatively weak. It only compares the proposed method with two relatively outdated baselines, SAC and TD3, on six simple tasks from the deepmind control suite.

Author response:
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.

Reviewer comment (Evaluation):
In the current experiments, the performance of QSM does not seem to be as impressive. Except for the "Hopper Hop" environment, the baselines seem to be able to outperform QSM in terms of convergence on other environments.

Author response:
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.

Reviewer comment (Novelty):
QSM is not the first algorithm that utilizes the Q function to guide policy updates. Diffusion-QL proposed by Wang et al. and QGPO proposed by Lu et al. both employ different approaches to leverage the Q function for policy learning. Although QSM is an off-policy online RL algorithm, it would be better to compare it with other methods in an offline setting to evaluate the strengths and weaknesses of different ways to utilize the Q function.

Author response:
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.

Reviewer comment (Theory):
**Theory-Practice Disparity** The theoretical framework suggests that the score of optimal policies $\Psi^*$ aligns with the action gradient of the optimal Q-function ($\nabla_a Q^{\Psi^*}$.). While conceptually sound, the challenge arises in practical scenarios where the optimal Q-function is unknown. In your case, you goal is to move $\Psi$ to $\Psi^*$, and you use the $\nabla _a Q$ but not $\nabla_a Q^*$ as the direction to update $\Psi$. The theorem, while insightful, does not appear to significantly contribute to the algorithm's practical efficacy.

Author response:
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.

Reviewer comment (Novelty):
**Comparative Methodology Concerns** The implementation closely resembles the one in Diffusion-QL [1], where actions are also sampled by denoising from noise through a trained score function, and this function is refined under the guidance of the Q-value. While Diffusion-QL optimizes the score function by directly maximizing the Q-value, your method minimizes the L2 norm between the score and the action gradient of Q. These methods seem to conceptually converge, raising questions about the distinctiveness of your approach.

Author response:
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.

Reviewer comment (Experiments):
**Insufficient Experimentation** The range of experiments conducted lacks breadth, especially in the domain of online continuous control tasks where environments like Ant and Humanoid are benchmarks. Incorporating these could enhance the empirical validation of the proposed method.

Author response:
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.

Reviewer comment (Writing):
Confused in the Algorithm box. What is the $\pi_\phi$ in the algorithm? The score network $\Psi$ depends on $(s, a)$. Why the sampling of $a_t$ doesn’t depend on $a_{t-1}$?

Author response:
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.