• W1: Table 1 performance

We would like to note that our method surpasses or parallels existing methods for 7 out of 10 tasks from the MBBL benchmark. We keep the same architecture and hyper-parameters when benchmarking our approach, so we expect better performances for the rest of the tasks if the architecture can be tuned specifically for them.

• W2: Diff-Rep lacks the flexibility of conditioning the policy at inference time?

We acknowledge that the current version of Diff-Rep lacks the flexibility of goal-conditional generation at inference time since we are modeling the transition model to handle common RL problems. However, our approach preserves the capability of conditional generation by using the goal information to condition the scoring network $\zeta$. We believe this presents an intriguing direction for future research.

• W3: Some minor derivations or explanations are needed for the approximation in Equation 12

In Equation 12, we parameterize the perturbed transition operator $\mathbb{P}(\tilde{s}'|s, a; \beta)$ as an EBM model, whose energy function $\psi(s, a)^\top \nu(\tilde{s}', \beta)$ is factorized by two representation functions $\psi$ and $\nu$. As noted in line 213 as well as in [1], such parameterization possesses the universal approximation capability as long as $\psi$ and $\nu$ are capable of universal approximation. Therefore, the parameterization in Equation 12 is sound since we are using neural networks for both $\psi$ and $\nu$.

[1] Hiroaki Sasaki and Aapo Hyvärinen. Neural-kernelized conditional density estimation. 2018.

• Q1: Runtime comparison between Diff-Rep and LV-Rep?

The runtime comparison is illustrated in the attached Figure of the general rebuttal. Note that Diff-Rep is more computationally efficient than both LV-Rep and PolyGRAD in most of the tasks from MBBL.

• Q2: showing the accuracy of the estimated Q function when using Diff-Rep?

Illustrating the accuracy of the estimated $Q$ functions in control tasks is generally challenging, as obtaining ground truth values necessitates Monte Carlo estimation, which suffers from great variance as the task horizon increases. Therefore, we conducted a toy experiment in a grid world environment, where the ground truth $Q$-values can be solved accurately via policy evaluation. Please refer to the general response for details.

• Q3: Suggest some methods to help interpret the latent representation?

Diff-Rep extracts $\psi(s, a)$ which captures the transition functions and constructs sufficient representations that can represent the $Q$-functions of any policy $\pi$ based on our theory. Thus, one way to interpret the quality of the representation is by inspecting the generation results to see whether $\psi$ faithfully captures the dynamics information. We included the generation results in Section F.2.

We apologize for any ambiguity regarding certain claims in the paper. Below, we will offer additional justifications for them and revise the corresponding sections. We also appreciate the grammar corrections and will ensure that the final version is free of such errors.

• W1: SVD decomposition of transition operator assumption

The SVD of the transition operator $T(s' | s, a)$ holds in general MDP settings, but potentially with infinite-dimensional spectrum. When the spaces $\mathcal{S}$ and $\mathcal{A}$ are finite, $T$ can be represented as a matrix $T\in \mathbb{R}^{(|\mathcal{S}|\cdot |\mathcal{A}|)\times |\mathcal{S}|}$, which admits a finite SVD. For continuous $\mathcal{S}$ and $\mathcal{A}$, if $T(s'|s, a)$ can be represented by a discrete latent variable model, i.e.,

$T(s'|s, a) = \sum_{i=1}^k p(s'|i)p(i|s, a),$

we obtain a low-rank spectral decomposition as $\phi^*(s, a)=[p(i|s, a)]_{i=1}^k$ and $\mu^*(s')=[p(s'|i)]\_{i=1}^k$. Generally, in cases where these spaces are infinite (for example, in continuous settings), a countably infinite SVD of the transition operator $T$ exists, provided that $T$ is compact. For additional information, please refer to [1].

[1] Jordan Bell. The singular value decomposition of compact operators on Hilbert spaces. 2014.

• W2: Why Diff-Rep is sufficiently expressive to represent the value function of any policy.

From line 82 to line 91, we show that with the SVD of the transition operator $\mathbb{P}(s'|s, a)=\langle\phi^*(s, a), \mu^*(s, a)\rangle$, we can represent $Q^\pi(s, a)$ as a linear function w.r.t. $r(s, a)$ and $\phi^*(s, a)$. Our paper derives such representations via diffusion (Section 3.2), thus enjoying the same property. We also provide a toy experiment in the general response to examine this property.

• W3: Claims about $\phi^{*}$

Un-normalized conditional density estimation:

As $\phi^*(s, a)$ comes from the SVD of the conditional density $\mathbb{P}(s'|s, a)$ (Equation 1), learning such representation translates into estimating the conditional density $\mathbb{P}(s'|s, a)$ using available samples. However, vanilla parameterization of $\phi^*$ and $\mu^*$ cannot guarantee the normalizing condition of probability distributions, i.e. $\int \mathbb{P}(s'|s, a)\mathrm{d}s'=1$, thus we refer to it as un-normalized density estimation.

Learning requires full-coverage data:

Learning an accurate $\phi^*$ requires data with full coverage to capture the underlying conditional density $\mathbb{P}(s' |s, a)$.

Exploration requires accurate $\phi^{*}$:

In the online setting, the data is progressively collected by the agent. To encourage the exploration of unvisited regions, existing approaches design exploration strategies that rely on $\phi^*$ (e.g., [1-2]). This dilemma has sparked a series of provable and practical RL algorithms [3-4], including our method. We will include relevant citations in the revision.

[1] Chi Jin, et al. Provably efficient reinforcement learning with linear function approximation. 2020.

[2] Zhuoran Yang, et al. Provably efficient reinforcement learning with kernel and neural function approximations. 2020.

[3] Alekh Agarwal, et al. Flambe: Structural complexity and representation learning of low rank mdps. 2020.

[4] Tongzheng Ren, et al. Latent variable representation for reinforcement learning. 2022.

• W4: Random Fourier features and explanations.

Thanks for pointing this out. Random Fourier feature [1] is an important technique to recover the SVD from EBM parameterization. We will provide a more detailed exposition of this in Section 2.

Specifically, $\langle \cdot, \cdot\rangle_{\mathcal{N}(\omega)}$ denotes the dot product under the distribution $\mathcal{N}(\omega)$. Equation 9 is obtained by a direct application of the random Fourier feature to Gaussian kernels:

$

\exp(-\frac{\\|x-y\\|^2}{2}) &=\int p(\omega)\exp(-\mathbf{i}\omega^\top(x-y))\mathrm{d}\omega \\\\
   &=\int p(\omega)\exp(-\mathbf{i}\omega^\top x)\exp(\mathbf{i}\omega^\top y)\mathrm{d}\omega \\\\
   &=\mathbb{E}_{\omega\sim p(\omega)} \left[\exp(-\mathbf{i}\omega^\top x)\exp(\mathbf{i}\omega^\top y)\right] \\\\

$

where $p(\omega)\sim\mathcal{N}(0, I_d)$.

[1] Ali Rahimi, Benjamin Recht, et al. Random features for large-scale kernel machines. NIPS. 2007.

• W5: Citations for EBMs

Thanks for pointing this out. There is a substantial amount of literature related to EBMs, and we have already cited some seminal works. We will include additional references in Section 3.2.

• W6: An explicit definition of $\nu$

In Equation 12, we parameterize the perturbed transition operator $\mathbb{P}(\tilde{s}'|s, a; \beta)$ as an EBM, whose energy function $\psi(s, a)^\top \nu(\tilde{s}', \beta)$ is composed by two representation functions $\psi$ and $\nu$. The representation function $\nu(\tilde{s}', \beta)$ receives the corrupted state $\tilde{s}'$ and the noise level $\beta$ as input, and outputs a possibly infinite-dimensional vector. We will update the paper to clarify this notation.

• Q1: The presence of $\sigma^2$ is unnecessary.

We apologize for the confusing expressions. The variance $\sigma^2$ should be removed and the corruption is defined by $\beta$ entirely. We will revise to remove $\sigma^2$ throughout the paper.

• Q2: Why does the second term equal 0 in line 591.

This is a direct application of the tower property. Specifically,

$

&\mathbb{E}_\beta \mathbb{E}_{(s, a, \tilde{s}', s')} [(\tilde{s}'+\beta\psi(s, a)^\top\zeta(\tilde{s}',\beta)-b)^\top (b-\sqrt{1-\beta}s')] \\ &=\mathbb{E}_\beta \mathbb{E}_{s, a, \tilde{s}'}[(\tilde{s}'+\beta\psi(s, a)^\top\zeta(\tilde{s}',\beta)-b)^\top (b-\sqrt{1-\beta}\mathbb{E}_{\mathbb{P}(s'|\tilde{s}', s, a, \beta)}[s'])] \\ &=0.

$

where the first equation is due to the independence between $b$ and $s'$; the second equation comes from the definition of $b$.

• W1 & Q1: Limited theoretical analysis.

The major contribution of this paper is developing an efficient algorithm for sufficient representation through diffusion, with intensive empirical evaluation. We would like to emphasize that our paper aligns with the existing algorithms [1][2] in leveraging the spectral representations for RL. The difference lies in that our approach efficiently extracts and constructs such representations via diffusion. Therefore, our method enjoys the same theoretical guarantee concerning sample complexity (e.g., Theorem 9 in [2]). In the next revision, we will include a remark that elaborates on the relationship between Diff-Rep and prior works, and refer interested readers to the theoretical analysis.

[1] Tongzheng Ren, et al. Latent variable representation for reinforcement learning. 2022.

[2] Tongzheng Ren, et al. Spectral decomposition representation for reinforcement learning. 2022.

• W2 & Q2: Lack of realistic or real-world tasks.

The primary objective of this research is to introduce a diffusion-based method for learning spectral representations. We followed the standard comparison protocol in [1,2] for both MDP and POMDP settings for fair comparison. The improvements in simulations demonstrate the effectiveness of our approach as our first step.

We are developing our method for applications in more realistic tasks as our future research. Given the capability of our method to extract sufficient representations for subsequent tasks (e.g. representing value functions), we envision its broader applications such as multi-task robotics control and preference alignment for large language models.

[1] Tongzheng Ren, et al. Latent variable representation for reinforcement learning. 2022.

[2] Hongming Zhang, et al. Provable Representation with Efficient Planning for Partially Observable Reinforcement Learning. 2023.

• Q3: Optimizing objective (18) along with model-free RL seems burdensome and other training strategies, such as pre-training the representations, may be exploited.

In our implementation, we built $Q$-functions on top of the representation network $\psi$, and used objective (18) to train the representation layers. Therefore, Diff-Rep incurs minimal additional costs compared to model-free algorithms. On the other hand, conventional model-based methods such as MWM involve independent representation learning, dynamics modeling and actor-critic learning processes, while our method leverages the benefits of dynamics from a representation learning perspective, thereby significantly reducing training costs.

In this paper, we mainly focus on the online RL settings to follow the evaluation protocols of existing representation-based RL methods. As the data is progressively collected by the agent, pre-training the representations beforehand is not feasible due to insufficient data. However in offline scenarios where extensive datasets are accessible, pre-training the representations is likely to bring more stable and improved outcomes.

• Q4: Connections to papers about RL for Diffusion.

Some studies within these two fields are connected in terms of methodology. Suppose the diffusion model is pre-trained on the data with distribution $p_{\text{data}}(x)$, both fields generally aim to sample from an enhanced distribution w.r.t. some evaluation metric $Q(x)$: $p_{\text{target}}(x)\propto p_{\text{data}}(x)\exp(Q(x)).$ In Diffusion for RL, typically $x$ is the trajectory $\tau$ while $Q(x)$ is the overall return of the trajectory $R(\tau)$. In the application of RL to Diffusion, $Q(x)$ can manifest as some specific metric, such as the aesthetic quality in image generation or bioactivity in biological sequence generation, as ELEGANT did.

Despite this, our method leverages the capabilities of diffusion models from the perspective of representation learning, which fundamentally differs from the aforementioned methods.

• W1: Some discussions on the representation quality are expected, e.g. to check whether Diff-Rep can recover latent state representations.

We would like to emphasize that Diff-Rep focuses on extracting representations that are sufficient to represent the $Q$-functions, rather than the latent representations of $(s, a)$. Therefore, the representations from Diff-Rep are generally different from the latent state of the MDP. To illustrate the representation quality, we included the generation results and also conducted experiments in a toy MDP where the ground truth $Q$-function can be solved for analysis. Please refer to Appendix F.2 and the general response for details.

• Q1: Why use gradient descent rather than closed-form solution?

While it is possible to optimize the $Q$-value functions through least squares regression, this approach would incur significant computational costs such as inverting the covariance matrix. Instead, we opt for gradient descent to optimize the $Q$-value functions for its simplicity and efficiency in our online setting. Note that when the representations are fixed, the objective is convex, and thus gradient descent gives the same solution as the closed-form solution.