Summary
The authors introduce a new algorithm for continuous RL environments, Diffusion Q-Sampling (DQS).
DQS makes use of an existing method, iterated Denoising Energy matching (iDEM).
The key idea is to use iDEM to learn a score function which can be used in a reverse diffusion process to sample actions.
The score function is trained such that the reverse diffusion process approximately samples from a Boltzmann distribution with respect to the Q-function of the current policy.
The authors give two theoretical results, corresponding to policy improvement and policy iteration respectively, to justify their choice of training rule for the action-value function and the diffusion model.
The authors then give experimental results for their method, DQS.
In the first set of results, they compare DQS to SAC (soft actor-critic) and QSM (Q-score matching) in terms of the diversity of behaviors learned. They demonstrate that, in a goal reaching maze environment, DQS can successfully learn a diverse set of solutions, while SAC and QSM learn a more concentrated set of solutions.
In the second set of results, they compare DQS to SAC and QSM on 8 tasks from the DeepMind control suite. They demonstrate that on many of these tasks, QSM dominates the other methods.
Strengths
Originality - The application of iDEM is (to this reviewer's knowledge) novel; although other methods seek to use diffusion model policies, they typically use other methods for fitting the diffusion model. The application of iDEM is novel.
Quality - The empirical results given are strong. The first set of results demonstrates well that DQS can indeed learn a policy which has support on multiple different solution types for problems. The second set of results shows that DQS can learn well, and outperform baseline methods in terms of sample efficiency.
Clarity - In general, the authors writing is clear. The method is well-explained, and seems reproducible.
Significance - The authors propose an effective new algorithm for continuous control. This algorithm seems particularly useful for the setting where compute is not a bottleneck, and multimodal policies are explicitly desired.
Weaknesses
046 - The authors give methods of policy representations in the continuous setting. I would suggest that they mention SQL, which allows for the training of expressive policies which come from neither noise injection nor parametric family. These are trained via Stein-variational gradient descent.
071 - The claim is made that "[Diffusion models] have been extensively applied to solve sequential decision-making tasks, especially in offline settings where they can model multimodal datasets from suboptimal policies or diverse human demonstrations." No citations are given for these techniques - please include citations to the literature to which you are referring.
191 - I would encourage the authors to say more about the role of the reverse SDE (3) in generation. Specifically, please be clear about how (3) is used to generate samples, rather than assuming this knowledge on the part of the reader.
205 - Missing tildes over the x's in the expectations in Eq. (4).
210 - Subscript below the S in equation (5) should be a capital K.
260 - Lemma 1 is false, and its proof is invalid. Lemma 1 states that, for any action-value function, the policy which is Boltzmann with respect to that action-value function has a dominating action-value function. This statement is incorrect, and obviously so. Let $\pi^*$ be the optimal policy, with action-value function $Q^*$. Then we know that $Q^*$ satisfies the Bellman optimality operator, $T^* Q(s,a) = r + \gamma \mathbb{E}[ \max_{a'} Q(s',a') | s, a]$, where the expectation is taken over next-states $s'$ conditional on state-action pair $s,a$. If Lemma 1 were true, it would mean that the Boltzmann policy $\pi_B$ with respect to $Q^*$ has an action-value function which dominates $Q^*$. But note that $T^* Q^*(s,a) \geq T^{\pi_B} Q^*(s,a)$, which can be seen by expanding definitions and using the fact that the maximum over $a'$ dominates any expectation with respect to $a'$, except if that expectation only places mass on the argmax actions. From this it also follows that this inequality is strict somewhere provided $Q^*$ is non-uniform somewhere. But since $T^* Q^* = Q^*$, it follows that $Q^* \geq T^{\pi_B} Q^*$. But monotonicity of the Bellman operator, it follows that, for all $n$, $Q^* \geq [T^{\pi_B}]^n Q^*$. Taking limits as $n \to \infty$, we obtain that, $Q^* \geq Q^{\pi_B}$, with strict inequality somewhere provided $Q^*$ is not flat. This contradicts the stated result.
We now turn to the proof given in A.1, and examine the error of reasoning. In the first two lines of (10), the expectation of $\log(\pi_{new})$ is taken with respect to $\pi_{new}$, and the expectation of $\log(\pi_{old})$ is taken with respect to $\pi_{old}$. However, in the third line of (10), the expectation of both terms is taken with respect to $\pi_{new}$. This allows the authors to express this term as a KL-divergence, a step critical to their proof. However, the term should instead be a difference of entropies, which in general is not non-negative (as the KL-divergence is).
265 - The proof of Theorem 1 is invalid. The proof relies heavily on the same argument as in Lemma 1, which is faulty.
In general, it seems like the authors fail to appreciate that results from the entropy regularised setting and the classical setting cannot be freely interchanged. The optimal policy is Boltzmann only if an entropy regularisation term is included in the Bellman backup, (7). When there is no such entropy term in the backup, the optimal policy will simply be the classical optimal policy, which in general is deterministic (or has support only on argmax actions). Similarly, the Boltzmann improvement map only gives improvement with entropy regularisation. Otherwise it can result in a strictly worse policy, as explained above.
I would suggest that the authors either cut their theoretical results entirely, or think about replacing the Bellman backup in (7) with the entropy regularized backup - however this would result in a substantial change to the algorithm, which may be too late at this stage.
Questions
074, 122 - The claim is made twice that for Q-Score Matching (Psenka et al. 2023), "the exact form of the policy is unspecified and it is unknown what distribution the diffusion models sample from". But this is equally true for your method. Both DQS and QSM train score functions which are used to generate samples of the policy. And both DQS and QSM aim for these score functions to allow for sampling from the Boltzmann distribution with respect to the current action-value function. So it is unclear what this comment is meant to mean, or what advantage you are supposing DQS has over QSM. Can the authors please clarify this?
461 - You mention that diffusion based policies have an increased runtime compute requirement compared to parametric policy methods. Can you give an indication of the ratio of runtime for your method vs. SAC? Are there experiments you can run which demonstrate that DQS outperforms SAC when normalised for compute time?
Will a codebase be made available to accompany the paper?