Diffusion Spectral Representation for Reinforcement Learning

Summary

This paper proposes an efficient and novel method for integrating diffusion models into reinforcement learning (RL). One major drawback of existing diffusion models is the high computational cost at inference time due to iterative sampling. This paper utilizes a diffusion model to learn representations of latent structures for the transition function. These representations can then be used for the value function in the RL framework. The authors evaluated their method on both MDP and POMDP tasks, showing significant benefits on the POMDP tasks compared to the baselines.

Strengths

The writing is generally good to follow though some editing on the notations could make the reading easier especially in section 3.1 and 3.2. And the author did a good job on motivating from LV-Rep. The proposed method, which integrates the diffusion model into the LV-Rep framework, is novel and offers a different perspective on existing diffusion-based RL methods. This method outperforms current diffusion-based RL techniques on POMDP tasks while consuming less runtime.

Weaknesses

In Table 1, for the MDP setting, Diff-Rep only outperforms other relevant baselines in the HalfCheetah task while showing similar or worse performance compared to LV-Rep and PolyGRAD. It might be beneficial to emphasize the advantages of the proposed algorithm in the POMDP setting. In addition, an explanation for this lower performance would be necessary instead of claiming "significantly better of comparable performance to all baseline methods in most tasks except for Humanoid-ET". While the runtime of the proposed method is lower than other diffusion-based methods, Diff-Rep loses the flexibility of conditioning the policy at inference time. Minor Some derivations/explanations are needed for the approximation in Eq 12.

Questions

Could the author provide a runtime comparison between Diff-Rep and LV-Rep? Could the author provide some plots showing the accuracy of the estimated Q function when using the proposed representation? Could the author suggest any methods that might help interpret the latent representation?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

## Main summary The authors propose Diff-Rep, an algorithmic framework leveraging diffusion models for learning spectral representations of MDPs, from which relevant quantities such as the Q-function can be linearly decoded. The main methodological contribution of the paper is showing that spectral representations can be learned using a diffusion-like loss function, obtained by exploring the connection between energy-based models (EBMs) and diffusion models. The resulting method, however, does not require sampling from the diffusion model, and so requires less wall time. The authors also show how to perform online policy optimization using Diff-Rep. They evaluate their method on Gym-MuJoCo locomotion tasks, which are MDP-based, and on image-based partially observable MDP tasks. Diff-Rep achieves similar or superior performance to baseline methods on most locomotion tasks, and consistently surpasses the performance and wall-time of PolyGRAD, a recent diffusion-based trajectory prediction model. ## More detailed summary of methods Their framework assumes the transition function $\mathbb{P}(s’ | s, a)$ can be written as an inner product of state-action features $\phi^*(s, a)$ and next-state features $\mu(s’)$: $\mathbb{P}(s’ | s, a) = \langle \phi^*(s, a), \mu(s’) \rangle$. The authors propose modeling this transition function using an Energy-Based Model (EBM), from which they obtain explicit forms for $\phi^*$ and $\mu^*$ as random Fourier features. Next, by considering a noised next state $\tilde{s}’$, they obtain a relation that must be satisfied by the energy term of this EBM. They then derive from this relation a diffusion-like loss which can be used to train the parameters of the EBM. To learn the state-action features $\phi^*$, they further include a regularization term incentivizing orthonormality of the entries of $\phi^*$. $\phi^*$ can then be optimized jointly with the EBM, sharing many similarities with the training of Denoising Diffusion Probabilistic Models. Finally, the authors show how to jointly learn representations $\phi^*$, a Q-function, and a policy $\pi$ online, leveraging $\phi^*$ to parameterize the Q-function. They further propose including an Upper Confidence Bound (UCB) bonus to rewards during optimization to incentivize exploration.

Strengths

1. Significance of research question: leveraging the flexibility and representation power of diffusion models has been actively researched recently, and their expensive sampling loop has proved a challenge (e.g. Janner et al. 2022, section 5.4). The authors consider the problem of reaping the benefits of diffusion in decision-making, but without incurring in this inefficiency. 1. Originality: while learning spectral representations for MDPs was already present in prior work, the authors’ approach of parameterizing the transition function and spectral representations using an EBM and training it via a diffusion-like loss function is original. 1. Quality of exposition and overall quality: I believe these can be significantly improved. See Weaknesses. 1. Quality of evaluation: the authors compare the performance of Diff-Rep on Gym-MuJoCo locomotion tasks (e.g. Hopper, HalfCheetah) and on a partially-observable image-based task from the Meta-World benchmark. They consider several model-based and model-free baselines, as well as other techniques based on learning representations. They demonstrate performance gains on most MDP and POMDP tasks, and also highlight a gain in wall time compared to PolyGRAD (Rigter et al. 2024). 1. In addition, they give qualitative results showing the diffusion model trained via Diff-Rep is able to reconstruct scenes with good fidelity in the setting of their POMDP experiment.

Weaknesses

1. Is the singular value decomposition of the transition operator taken as an assumption, or do the authors claim it holds for any MDP? Other works (e.g. Ren et al. 2022a) list this decomposition as an assumption, but the phrasing in line 82 (“one can always factorize [...]”) seems to claim that it holds in general, provided one does not require the representations to take values in finite dimensions. If the authors do claim this, I believe either a citation or a proof would be warranted. 1. On a similar note, the authors claim on line 63 that they show “diffusion-based representations are sufficiently expressive to represent the value function of any policy”. It is not clear to me that either the exposition of the methods or the experiments prove this in such generality. 1. The claim on line 93 that “the learning of the ϕ* is essentially equivalent to un-normalized conditional density estimation” seems unsupported by a citation or an explanation. Similarly for the claim that “learning requires full-coverage data, while the exploration strategy requires an accurate ϕ*” in line 96. 1. Considering the central role of EBMs and random Fourier features in the derivation of the method, I believe a brief exposition about them in Section 2 would be warranted, in particular to clarify the $\langle \cdot, \cdot, \rangle_{\mathcal{N}(\omega)}$ notation in Equation 9. 1. Further, appropriate citations for EBMs seem to be missing in Section 3.2. 1. An explicit definition of $\nu(\tilde{s}’, \beta)$ is missing in Equation 12. Clarity and writing: The paper contains a relatively large number of what seem to me like grammar mistakes. Here is a non-exhaustive list of examples and suggestions for edits: 1. Line 29: “Exploring the potential of diffusion models for sequential decision-making is increasingly being investigated.” -> “The potential of diffusion models for sequential decision-making is increasingly being investigated.” 1. Line 50: “to avoid suboptimal policy” -> “to avoid a suboptimal policy” 1. Line 99: “the existing methods either require [...] or relies on” -> “the existing methods either require [...] or rely on” 1. Line 105: “After T-step” -> “After T steps” 1. Line 110: “Recognize the diffusion models are [...], the ELBO naturally as a choice for learning,” -> “Recognizing that diffusion models are [...], maximizing the ELBO arises as a natural choice:” 1. Line 117: “as a neural networks” -> “as a neural network” 1. Line 118: “the samples can be generated from [...], and then, following [...]” -> “the samples can be generated by sampling [...], and then following [...]” 1. Line 155: “We apply the EBM for transition operator” -> “We parameterize the transition operator using an EBM” 1. Line 176: “we consider the samples from dynamics s ′ ∼ P(s ′|s, a) is perturbed with Gaussian noise” -> “we consider perturbing the samples from dynamics s ′ ∼ P(s ′|s, a) with Gaussian noise” 1. Line 178: “$\mathbb{P}(\tilde{s}’|s, a; \alpha) \to \mathbb{P}(s’|s, a)$” -> “$\mathbb{P}(\tilde{s}’|s, a; \beta) \to \mathbb{P}(s’|s, a)$”

Questions

1. In line 177, why is noising chosen to be $\mathbb{P}(s’|s, a) = \mathcal{N}(\sqrt{1-\beta}s’, \beta \sigma^2 I)$, as opposed to $\mathbb{P}(s’|s, a) = \mathcal{N}(\sqrt{1-\beta}s’, \beta I)$, as is normally done, e.g. in Section 2? The presence of the $\sigma^2$ factor means the variance need not be preserved. Is this intentional? If so, what drove this design decision, and how do you choose $\sigma^2$? 1. In line 591, it seems to me that the second term equals 0 by applying the tower property conditioning on $\tilde{s}’$, $s$ and $a$. Is this correct? It might be good to make this explicit.

Rating

6

Confidence

3

Soundness

2

Presentation

2

Contribution

3

Limitations

The authors include one sentence addressing limitations in the conclusion, referencing the fact that their method has not been evaluated on real-world data. My view is that the main limitation of this work is exposition and clarity, which would warrant significant revision before publication. Particularly critical are unjustified claims, lack of clarity regarding the assumptions of the method, and several grammar/writing issues, as outlined in the Weaknesses section.

Summary

This paper proposes Diffusion Representation (Diff-Rep). This approach leverages diffusion models to learn representations for value functions in reinforcement learning while avoiding the high inference cost of sampling from diffusion models.

Strengths

1. Diff-Rep provides a novel and principled approach to leveraging the flexibility of diffusion models for reinforcement learning without incurring their high inference costs. The key technique used here is exploiting the energy-based model view of diffusion models. 2. Empirically, this approach demonstrates solid performances across various continuous control tasks. Diff-Rep consistently outperforms competitive baselines, including prior diffusion-based RL methods and state-of-the-art model-based and model-free algorithms. 3. The writing of this paper is smooth and easy to follow.

Weaknesses

1. While the paper provides some results on the expressivity of the learned diffusion representations, the theoretical analysis is still somewhat limited. More formal characterizations or guarantees on the representation power, such as error bounds on the Q-function approximation, would further strengthen the technical contributions. 2. The experimental evaluation, while extensive, focuses primarily on simulated benchmarks. The lack of more realistic or practical tasks is a disadvantage. Such results could help illustrate the approach's scalability and robustness under more complex and noisy conditions, which cannot be evaluated in a relatively easy environment.

Questions

1. As mentioned, regarding why the diffusion representations can expressively approximate the Q-function, are there any theoretical results or analyses that could formalize this representation power more precisely? What are the potential difficulties in providing such results? 2. What are some example real-world applications where you think Diff-Rep could provide the most benefit over prior approaches? Do you consider such extensions for the next/future version? 3. Optimizing the representation objective (18) alongside the model-free RL updates seems burdensome. Could you comment on the computational/sample complexities of training this? Moreover, did you explore any other training strategies, such as pre-training the representation? I guess it necessarily means some trade-offs. Adding discussing and comparing the tradeoffs here might be interesting. 4. While this paper mainly studies Diffusion for RL, the energy-based formulation and the KL-based training objectives presented in this paper are somewhat similar to those of papers studying (KL-regularized) RL for Diffusion (e.g., [1-2]). It might be interesting to discuss the methodological relations/differences between the two interdisciplinary directions. [1] https://arxiv.org/abs/2305.16381 [2] https://arxiv.org/abs/2402.15194

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed limitations.

Summary

The paper proposed a representation learning method based on diffusion model. The paper developed the method using the EBM setting of the transition probability, and proposed a finite dimension approximation of the state-action representation by minimizing an orthormal regularization term. The performance of the approach is supported with various experiments.

Strengths

The paper has sufficient comparison with the existing works. The writing is clear and easy to follow. The derivation of the diffusion setting from EBM is novel and the experiment results are convincing.

Weaknesses

To make the paper more complete, the authors can make some discussions on the representation quality, such as running the algorithm on a toy latent-state-MDP and show if the algorithm is able to find the latent state-representation.

Questions

As the representation learning enables the linear representation of the Q-value, one would have closed-form solution to Equation (20). On the other hand, the algorithm still uses gradient descent with the double-Q-network trick. Is there any specific reason of doing so?

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

/

Response to rebuttal

I thank the reviewers for their careful response to my reported concerns. A few follow-up questions: # W1 - What do the authors mean by spectrum of T in the fully general case with no assumptions on S and A? - What do the authors mean by an infinite-dimensional spectrum? Do the authors mean instead an infinite spectrum? - My understanding is that the transition operator maps a pair (s, a) to a probability measure over s', which can potentially be represented as a density (i.e. a function of s, a and s') if it is absolutely continuous with respect to some reference measure. Since S and A need not be vector spaces in general, and since the space of probability measures over a set is not a vector space (due to the need that such measures integrate to 1), I don't see how T could be regarded as an operator between e.g. normed vector spaces in the fully general setting. Does the SVD you reference hold when S and A are allowed to be a completely arbitrary set, and T is allowed to be a completely arbitrary kernel mapping S x A to a probability measure over S? - One can consider a "continuous" setting such as the control of a point mass inside a bounded region. In this case, the state space is infinite, but is not a vector space. In what way do you apply results of spectral theory in Hilbert spaces in this situation? - What does it mean for T to be compact in this case (S is "continuous" but not a vector space)? Does it involve something like "mapping bounded sets in S x A to relatively compact sets of probability measures on S"? In general, I believe it would be most efficient to find a sufficiently restricted setting in which you wish to claim the existence of the SVD in a way that remains relevant for this work. The argument for full generality is unfortunately not yet clear to me. # W2 It seems to me that whether this argument follows depends on the generality of the existence of the SVD, of which I am not yet convinced. It would be great if the authors could provide further clarification (ideally including specific references to literature on this result). # W3 - Un-normalized conditional density estimation: I am satisfied with this explanation. - Learning requires full-coverage data: it seems intuitive to me that data with full coverage would be a sufficient condition for learning the underlying conditional density, but not so much that it is impossible to learn $\mathbb{P}(s' |s, a)$ with acceptable accuracy without full coverage. In general, the notions of full coverage and necessary conditions for "learning $\phi^*$" don't seem to be defined precisely in the context of this statement. Because of this, wording emphasizing the qualitative nature of the statement seems more appropriate (e.g. "learning $\phi^*$ can be challenging without full-coverage data"). - Exploration requires accurate $\phi^*$: do the referenced works show that exploration without accurate $\phi^*$ is impossible, or that exporation with $\phi^*$ is possible? # W4, W5, W6, Q1 and Q2 Thank you for agreeing to include these in the paper. I believe it will help clarity and accessibility, and make it more self-contained. With the above, I believe most formatting and clarity concerns will be addressed. My remaining queries mostly concern the mathematical and scientific precision of certain statements, but I think these can be easily remedied by simply making the claims in less generality or in a more qualitative way (e.g. "learning might be a challenge in practice" instead of "learning would be impossible in principle"). Because of this, I will raise my score.