Diffusion Spectral Representation for Reinforcement Learning
Authors: Dmitry Shribak
Chen-Xiao Gao
Yitong Li
Chenjun Xiao
Bo Dai
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion Models
Reinforcement Learning
Representation Learning
Primary area: reinforcement_learning
TL;DR: We propose an algorithm that harnesses the flexibility of diffusion models for representation learning to achieve efficient policy optimization, while avoiding the time-consuming sampling process typically associated with diffusion models.

Abstract: Diffusion-based models have achieved notable empirical successes in reinforcement learning (RL) due to their expressiveness in modeling complex distributions. Despite existing methods being promising, the key challenge of extending existing methods for broader real-world applications lies in the computational cost at inference time, i.e., sampling from a diffusion model is considerably slow as it often requires tens to hundreds of iterations to generate even one sample. To circumvent this issue, we propose to leverage the flexibility of diffusion models for RL from a representation learning perspective. In particular, by exploiting the connection between diffusion models and energy-based models, we develop Diffusion Spectral Representation (Diff-SR), a coherent algorithm framework that enables extracting sufficient representations for value functions in Markov decision processes (MDP) and partially observable Markov decision processes (POMDP). We further demonstrate how Diff-SR facilitates efficient policy optimization and practical algorithms while explicitly bypassing the difficulty and inference cost of sampling from the diffusion model. Finally, we provide comprehensive empirical studies to verify the benefits of Diff-SR in delivering robust and advantageous performance across various benchmarks with both fully and partially observable settings.

Diffusion Spectral Representation for Reinforcement
Learning
Dmitry Shribak" Georgia Tech shribak@gatech.. edu
Chen-Xiao Gao" Nanjing University gaocxelamda nju. edu. en
Yitong Li
Georgia Tech
yli3277@gatech.co edu
Chenjun Xiao
Bo Dai Georgia Tech
CUHK(SZ)
chenjunx@cuhk, edu.
bodai@cc.gatechh edu
Abstract
Diffusion-basee models have achieved notable empirical successes in reinforce- ment learning (RL) due to their expressiveness in modeling complex distributions. Despite existing methods being promising, the key challenge of extending exist- ing methods for broader real-world applications lies in the computational cost at inference time, i.e., sampling from a diffusion model is considerably slow as it often requires tens to hundreds of iterations to generate even one sample. To circumvent this issue, we propose leverage the flexibility of diffusion models for RI. from representation learning perspective. In particular, by exploiting the connection between diffusion models and energy-based models, we develop Diffusion Spectral Representation (Dfff-SR), coherent algorithm framework that enables extracting sufficient representations for value functions Markov decision processes (MDP) and partially observable Markov decision processes (POMDP). We further demonstrate how Diff-SR facilitates efficient policy optimization and practical algorithms while explicitly bypassing the difficulty and inference cost of sampling from the diffusion model. Finally, we provide comprehensive empirical studies to verify the benefits of Diff-SR in delivering robust and advantageous performance across various benchmarks with both fully and partially observable settings.
Introduction
Diffusion models have demonstrated remarkable generative modeling capabilities, achieving signifi- cant success producing high-qualit, samples across various domains such as images and videos [Ramesh et al., 2021, Saharia et al., 2022, Brooks et al., 2024]. In comparison to other generative approaches, diffusion models stand out for their ability to represent complex, multimodal data distri- butions, a strength that can be attributed to two primary factors. First, diffusion models progressively denoise data by reversing a diffusion process. This iterative refinement process empowers them to capture complicated patterns and structures within the data distribution, thus enabling the generation of samples with unprecedented accuracy [Ho et al., 2020]. Second, diffusion models exhibit impres- sive mode coverage, effecttivly addressing common issue of mode collapse encountered in other generative approaches [Song al., 2020].
The potential of diffusion models is increasingly being investigated for sequential beiision-making tasks. The inherent flexibility of diffusion models accurately capture complex data distributions
Equual Contribution. Correspondence Bo Dai <bodai @ea.@eecc.ccccecee gatech.edu> 38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

makes them exceptionally suitable for both model-free and model- -ased methods in reinforcement learning (RL). There are three main approaches that attempts to apply diffusion models, including diffission policy [Wang et al., 2022, Chi et al., 2023], diffusion-bused planning [Janner et al., 2022, Jackson et al., 2024, Du et al., 2024], and diffiasion world model [Ding et al., 2024, Rigter et al., 2023]. Empirical results indicate that diffusion-
approaches can stabilize the training process and enhance empirical performance compared with their conventional counterparts, especially in environments with high-dimenniooal inputs.
The flexibility of diffusion models, however, also comes with a substantial inference cost: generating even a single sample from diffusion model is notably slow, typically requiring tens to thousands of iterations [Ho al., 2020, Lu et al.. 2022, Zhang and Chen, 2022, Song al., 2023]. Furthermore, prior works also consider generating multiple samples to enhance quality. which further exacerbates this issue [Rigter et al., 2023]. The computational demands are particularly problematic for RI,, since whether employing a diffusion-based policy or diffusion world model, the learning agent must frequently query the model for interactions with the environment during the learning phase or when deployed in the environment. This becomes the key challenge when extending diffusion-biis -bsed methods for broader applications with more complex state spaces. Meanwhile, the planning with exploration issue has not been explicitly considered in the existing diffusion-based RL algorithms. The flexibility of diffusion models in fact induces extra difficulty in the implementation of the principle of optimism in the face of uncertainty in the planning step to balance the inherent trade-off between exploration vs. exploitation, which is indispensable in the online setting to avoid suboptimal policies [Lattimore and Szepesvári, 2020, Amin et al., 2021].
In conclusion, there has been insufficient work considering both efficiency and computation tractabil- ity for planning and exploration in a unified and coherent perspective, when applying diffusion model for sequential decision-making This raises a very natural question, i.e.,
Can we exploit the flexibility of diffusion models with efficient planning and exploration for RL?
In this paper, we provide an affirmative answer to this question, based on our key observation that diffusion models, beyond their conventional role as generative tools, can play a crucial role in learning sufficient representations for RL. Specifically,
By exploiting the energy-based model view of the diffusion model, we develop a coherent algorith- mic framework Diffusioo Spectral Representation (Diff-SR), designed to learn representations that capture the latent structure of the transition function in Section 3.2:
We then show that such diffusion-based representations are sufficiently expressive to represent the value function of any policy, which paves the way for efficient planning and exploration, circumventing the need for sample generation from the diffusion model, and thus, avoiding the inference costs associated with prior diffusion- methods in Section 3.3;
We conduct comprehensive empirical studies to validate the benefits of Diff-SR, in both fully and partially observable RL settings, demonstrating its robust, superior performance and efficiency across various benchmarks in Section 5.
2 Preliminaries
In this section, we briefly introduce the Markov Decision Processes, as the standard mathematical abstraction for RL, and diffusion models, as the building block for our algorithm design.
Markoy Decision Processes (MDPs). We consider Markov Decision Processes [Puterman, 2014] specified by the tuple M (S, I, Y, Ho). where S is the state space, A is the action space, P: P:SX x A A(S) is the transition functioon r S x A R is the reward function, 7 E [0. 1) is the discount factor, 14 A(S) is the initial state distribution?. The value function specifies the discounted cumulative rewards obtained by following a policy т S - A(A), V=(s) = E" EE 7'r(s), ss) = s), where E" denotes the expectation under the distribution induced by the interconnection of π and the environment. The atate-action value function is defined by
Q=(ss, a) = r(s, a) + 7BB-eP())ea)) [V*(s']
The goal of RL is to find an optimal policy that maximizes the policy value, i.e., =² argmax, Egwpo (s)).
We use the standard notation ((X) to denote the set of probability distributions finite set X

---

For any MDP, one can always factorize the transition operator through the singular value decomposi- tion (SVD), i.e.,
P(s'ss,) (0" (),
(1)
with ( defined the inner product. Yao et al. [2014], Jin et al. [2020], Agarwal et al. [2020a] considered a subset of MDPs, in which di (s,a) E R with finite d, which is known as Lineur/Low- rank MDP. The subclass then generalized for infinite spectrum with fast decay eigenvalues [Ren et al., 2022a]. Leveraging this specific structure of the function class, this spectral view serves as an instrumental framework for examining the statistical and computational attributes of RL algorithms in the context of function approximation. In fact, the most significant advantage of exploiting said spectral structure is that we can represent its state-value function Q" (s a) as a linear function with respect to [r(s, a), 6" ("(("" a)] for any policy n
(s, r0:==((s00) r(s,a) + PPE [VT = r(s, ((050))((((0000) (01 10 (''6')
(2)
It's important to highlight that in many practical scenarios, the feature mapping 6" is often unknown. In the meantime, the learning of the of is essentially equivalent to un-normalized conditional density estimation, which is notoriously difficult [LeCun et al., 2006, Song and Kingma, 2021, Dai et al., 2019]. Besides the optimization intractability in learning, the coupling of exploration to learning also compounds the difficulty: learning 6 requires full-coverage data capture IP(s'|ss a). while the design of exploration strategy often relies on an accurate o² [Jin et al., 2020, Yang et al., 2020]. Recently, range of spectral representation learning algorithms has emerged to address these challenges and provide an estimate of ø" in both online and offline settings [Uehara et al., 2021]. However, existing methods either require designs of negative samples [Ren et al., 2022b, Qiu et al., 2022, Zhang et al., 2022], or rely on additional assumptions on o* [Ren et al., 2022c,a].
Diffusion Models. Diffusion models (Sohl-Dickstein et al., 2015, Ho et al., 2020, Song et al., 2020] are composed of a forward Markov process gradually perturbing the observations a target distribution 2T qr (z) with corruption kernel 4c+1/ts and a backward Markov process recovering the original observations distribution from the noisy aT. After T steps, the forward process forms a joint distribution,
Pu (z) to
T-1
qu:r (IO:T) = po (2o) II t=0
2o41 (Ze4alz2).
The reverse process can be derived by Bayes' rule from the joint distribution, i.e.,
4tj++1 (2+(01+1) 1fe (Ie+1 (22)) qi (2r)
4++11 (I++1)
with qe (xz) as the marginal distribution at t-step. Although we can obtain the expression of reverse kernel (-I-+ from Bayes's 's rule, it is usually intractable. Therefore, the reverse kernel is usually parameterized with neural network, denoted p" (are(æe+1). Recognizing that the diffusion models are a special class of latent variable models [Sohl- Dickssic et al., 2015, Ho et al., 2020], maximizing the ELBO emerges as a natural choice for learning,
letbo (0) = Epo(z) DKL (orjoller) - -E= Enyo DKL (()++)))))))+))] Busio [Log Pon (zolæ1)]
For continuous domain, the forward process of corruption usually employs Gaussian noise, N( /1 1 (c411), where I (O...., 1}. The kernel for backward process is also Gaussian and can be parametrized as p" (xz(#141) =
i.e., (c=1)e(Fe441xx,)
N Tei -iーيi (Iet1 + B24100 (1441, 1)) B.421) where 8 (Ir+1, 1) denotes the score network with parameters 0. The EL.BO can be specified as
T
Σ (1) E Begio [|.* (Ite, VV, Vz, log Zoo (mra)))))
Eum (0)
where 04 -I",, B.). With the learned s (ar, t), the samples can be generated by sampling JT (N0, I). and then following the estimated reverse Markov chain with dd (xz)ate+) iteratively. To ensure the quality of samples from diffusion models, the reverse Markov chain requires tens to thousands of iterations, which induces high computational costs for diffusion model applications.
3

---

Random Fourier Features. Random Fourier features [Rahimi and Recht, 2007, Dai et al.. 2014] al- low us approximate adinite-dimensional kernels using unite-dimensional feature vectors. Bochner's theorem states that a continuous function of the form k=x,, y) k(x y) can be represented by Fourier transform of a bounded positive measure [Bochner, 1932]. For Gaussian kernel k(x - y). consider the feature z (x)) exp(-iw kw x), with w I);
T k(x- y) / p(w exp( in Ts - y))d
exp(ww y))l
(4)
Ew /zw (z))z (z)( 2u (v))x(()))
where (ww)) is shorthand for Eu.-No. By sampling w1, W2, WN NO I). can approximate k(x y) with the inner product of inite-dimentions vectors (x)
VN (Zm, (2), Aws (I)...., (x)) and ww = VN (Zw, (y), Rws (V). .,.., Aux (1)).
3 Diffusion Spectral Representation for Efficient Reinforcement Learning
is well known that the generation procedure of diffusion models becomes the major barrier for real-world application, especially in RL. Moreover, the complicated generation procedure makes the uncertainty estimation for exploration intractable. The spectral representation 6² (s, a) provides an efficient way for planning and exploration, Spectral Representation (Diff-SR) for RL, as our answer to the motivational question. As we will demonstrate below, the representation view of diffusion model enjoys the flexibility and also enables efficient planning and exploration, while directly bypasses the cost of sampling. For simplicity, we introduce Diff-SR in MDP setting in the main text. However, the proposed Diff-SR is also applicable for POMDPs as shown in Appendix B. We first illustrate the inherent challenges of applying diffusion models for representation learning in RL.
illustrated in Section 2, which inspires our Diffusion
3.1 An Impossible Reduction to Latent Variable Representation [Ren et al., 2022a]
In the latent variable representation (LV-Rep) [Ren et al., 2022a], the latent variable model is exploited for spectral representation 6² (s, a) in (I), by which arbitrary state-value function can be linearly represented, and thus, efficient planning and exploration is possible. Specifically, in the LV-Rep, one considers the factorization of dynamics as
P(' a) / a)p(ss)) = (p(zls, ().s.a). p(s')))))
(5) By recognizing the connection between (5) and SVD of transition operator (1), the learned latent variable model p(z|s, a) can be used as (8 a) for linearly representing the Q"-function for an arbitrary policy m.
Since the diffusion model can be recast as a special type of latent variable model [Ho et al., 2020], the first straightforward idea extend LV-Rep with diffusion models for a2) We consider the following forward process that perturbs each 24-1 with Gaussian noises:
p(zz/2i-1, 8, a), Vi = 1,....k. .k
with 2q = 8'. Then, following the definition of the diffusion model, the backward process can set
\<2,-1/5;, 8, a), === O.
which are Gaussian distributions that denoise the perturbed latent variables. Thus, the dynamics model be formulated
P(s' f (((2 s, a)q(s'(21, a)d 1=2
Indeed, Equation (7) converts the diffuusion model to a latent variable model for dynamics modeling. However, the dependency of (s, a) in q(s' [2 s, a) correspondingly induces the undesirable depen- dency of (s, a) into uuss') in (1), and therefore the factorization provided by the diffusion model cannot linearly represent the Q" Q"-funntion the vanilla LV-Rep reduction from diffusion models is impossible.
3.2 Diffusion Spectral Representation from Energy-based View
Instead of reducing to RV-Rep, in this work we extract the spectral representations by exploiting the relationship between diffusion models and energy-bascd models (EBMs).

---

Spectral Representation from EBMs. We parameterize the transition operator P(s'|s, a) using an EBM, i.e.,
P(s||s, exp
- log 2 (s, , (s,a) = -x (s(s, (s))) ds' (8)
By simple algebra manipulation,
(s, (ss,a)"v(s'' (s') = (IW - 11 (s, - || [6)
we obtain the quadratic potential function, leading to exp (-11 (s, (s')) exp (Iw( (s'))
P(s (8))) x exp (IIv (e)aa)))// ((+))) is the Gaussian kernel, for which we apply the random Fourier
(9)
The term exp
feature [Rahimi and Recht, 2007, Dai et al., 2014] and obtain the spectral decomposition of (8), IP(8'18, a) = (ow (#, a) , (8 N(w)))
(10)
where w NO(,, I). and Du (s, u) = exp (-iw T db (s, a)) exp (IV (s, ())) log log (s,a),
(11)
Ho (s') exp (-iw in T y (s') exp || (s))
(12)
This bridges the factorized EBMs (8) to SVD, offering a spectral representation for efficient planning and exploration, as will be shown subsequently.
Exploiting the random feature to connect EBMs to spectral representation for RL was first proposed by Nachum and Yang [2021] and Ren et al. [2022c] but only Gaussian dynamics p(s' (ssls, has been considered for its closed-form 7 (s. a) and tractable MI.E. Equation (8) is also discussed in [Zhang al., 2023b, Ouhamma al., 2023, Zheng et al., 2022] for exploration with UCB-style bonuses. Due to the notorious difficulty in MLE of EBMs caused by the intractability of Z (s, a) [Zhang et al., 2022], only special cases of (8) have been practically implemented. How efficiently exploit the flexibility of general EBMs in practice still remains an open problem.
Representation Learning via Diffusion. We revisit the EBM understanding of diffusion models, which not only justifies the flexibility of diffusion models, but more importantly, paves the way for efficient learning of spectral representations (11) through Tweedie's 's identity for diffusion models.
Given (s, a), we consider perturbing the samples from dynamics s' F(s'|s,, a) with Gaussian noise, i.e., (3 :) N (VI Bs', BT). Then, we parametrize the corrupted dynamics as F(3' 1P |s a; B) F33 () Is, a)ds' x exp (vs.
(s. a)T .(31
(13)
where (s, a) is shared across all noise levels B, and P(3' a.) P(s'(s, with 3' + s', as 3 - 0, there is no noise corruption on s'.
Proposition (Tweedie's Identity [Efron, 2011]). For arbitrary corruption IP (3'1s'; B) and B in P (3 (s, a; 3), we have
Vg log P( (s a; B) = Ep(s 3,88888) V log P(''P (8)
(14)
This can be easily verified by simple calculation. i.e.,
V (P(''s' F)s(s' a)ds'
V y log ((3''1s,; a;
TP ((3'ss,a;3)
1(55|18, a; ()
/ V P(s (33) do' = Excor
V log P()))1)8))
(15)
P(''|s,;;7
For Gaussian perturbation with (13), Tweedie's identity (14) is applied
VI Bs'
5
1 (s. a) T V V (8', 3) Ep(s' s.a58)
3
38' + Bu (8,
V v (8', v1 pp a. [s]].
(16)
Let ( (3:: = V ()) we can learn vb (s, a) and ( (3': () by matching both sides of (16),
||3 Br (s. < (8', (8,,)) 9) - v VI- BB a.:8)) [s|]
27
min EgL $/5
(17)
5

---

Algorithm Diffusion Spectral Representation (Diff-SR) Training
Input: representation networks v, <, noise levels (8 1 replay buffer D
2: for update step = 1, .... Nep do
3: Sample a batch of transitions {(5), atz s() ~ D
4: Sample noise schedules for each transition 18:12- Uniform (8), 8 ,..., .T.
5: Corrupt the next states 8 f- v1 - B, + vB,e,, where e, (0,
6: Optimize v,, S via gradient descent by minimizing Eq (18)
7: end for
8:
Return J.5
which shares the same optimum of
| + By (s,a) ((3 V1 -Bs'
min lair (t), {) := Eg
(18)
thx
The equivalence of (17) and (18) provided in Appendix A.
Diffusion Spectral Representation for Q-function. The loss described by (18) estimates the score function (s, a) < (3',, (3) for diffusion models. the context of generating samples, the score function suffices drive the reverse Markov chain process. However, when deriving the random feature ew defined in (11), (I), the partition function (s, a) is indispensable. Furthermore, the random feature ow (s, a)) is only conceptual with infinite dimensions where - (0, I). Next, will proceed to analyze the structure of (s, a) and finally construct the spectral representation with the learned vi (s.
We first illustrate (s, a) also linearly representable by random features of /a (s, a).
Proposition 2. Denote Pw (s, a) == exp in T Th (s, a) + 1| (s, (ssaa)l///) /2). the partition function is linearly representable by Pn a), ie., Z (s.a) = (pw (s. u) N(W))
Proof. We have &w a)
Puls.a) 2(s,a)
and P (s'[s,a) a) = p Z[s.a] Mu (8)) N(ww)
which implies
[P(('ls,a) =1
palaa) Z(s.a)
Ne (s') ds'
(pw (8, wiwww) Z (s)
Plug this into (11) and (2), we can represent the Q" function for arbitrary π as (exp (-iw"y vi (s, log (exp (-ww ot (s, u) w) .5))
Q (s,a)=
Aelam)) S² N(w
x (19)
which eliminates the explicit partition function calculation.
We thereby construct Diffuusion Spectral Representation (Diff-SR), tractable finite-dimensional representation, by approximating Pa,u (8, a) with some neural network upon a (s, a). Specifically, in the definition of Pu,u (s, a) in (19), it contains Fourier basis exp (-iw" T a (s, a)), suggesting the use of trigonometry functions [Rahimi and Recht, 2007] upon vi (s, a). i.e., sin (W, & (s) Meanwhile, it also contains product with (exp(- in ")) suggesting the additional non-
linearity over sin [W] (s, o)). Therefore, we consider the snite-dimensional neural network de (s, a) = elu (W₂ sin (WT y (s, a))) E IRdd with learnable parameters (W, W2), to approxi- mate the unfnite-dimmenional yp (s) with w ~ (O,TT).
Remark (Connection to sufficient dimension reduction [Sasaki and Hyvärinen, 2018)): The theoretical properties of the factorized potential function in EBMs (8) have been investigated in [Sasaki and Hyvärinen, 2018]. Specifically, the factorization is actually capable of universal approximation. Moreover, z6 (s, a) also constructs an implementation of sufficient dimension reduc- tion [Fukumizu et al., 2009], i.e., given v (s, a), we have P(s'|s, a) = PP (s'| (s. a)), or equivalently s' I. (s, d)lb a). These properties justify the expressiveness and sufficiency of the learned & (x, a). However, d (8, a) in [Sasaki and Hyvärinen, 2018] is only estimated up to the partition function Z a), which makes it not directly applicable for planning and exploration in RL discussed.

---

Algorithm 2 Online RL with Dift-SR
1: Initialize networks т. (f1. B1), ($2. 02), and D =
2: for timestep 1= do
3: ar ( (118))
4: Ft r(; at); 8 (( 0t)
5:
Compute bonus b(st; ac) using (23) (Optional)
6:
D (5t, ae, Te, s6)
7:
Update & with D by Algorithm 1
8: Update the critic (E1, 01), (E2. B2) by (21)
9: Update the policy by max, EgwD a (1,2) (1.2) Qs,o(s,a)
10: end for
11:
Return
3.3 Policy Optimization with Diffusion Spectral Representation
With the approximated snite-dimensionas representation Φ9, the Q-function can be represented as linear function Φ9 and a weight vector { E RRd
Qg,0(8,
do (s T
(20)
This approximated value function can be integrated into any model-free algorithm for policy opti- mization. In particular, we update (5. 0) with the standard TD learning objective
Ex,,,,'' (r + [Qg, (3((''' a')] Q5,0(8, a))² Qc,)(s,a))
(aritic(
(21)
where ı is the learning algorithm's 's current policy, (E. () is the target network, D the replay buffer. We apply the double Q-network trick to stabilize training [Fujimoto et al., 2018). In particular, two weights (E1, 01): (E2,0), are initialized and updated independently according to (21). Then the policy is updated by considering max, Ey-D.aw> min, (1.2) Qe..o, (s, a.). Algorithm 2 presents the pseudocode of online RL with Diff-SR. This learning framework is largely consistent with standard online reinforcement learning (RL) approaches, with the primary distinction being the incorporation of diffusion representations. As more data collected, Line 7 specifies the update of the diffusion representation de using Algorithm and the latest buffer D.
Remark (Exploration with Diffusion Spectral Representation); Following [Guo al., 2023],
even exp (-iw" (s)a)) is not enough for linearly representing Q", we still can exploit exp (-iw & (s) for bonus calculation to implement the principle of optimism in the face of
uncertainty for exploration in RL, i.e.,
b(s,a) 1 - K (s, (K +1) K (s,a),
(22)
where D is dataset of state action pairs, & ((s, a), a')) exp (-1 ((s.) )1?
((s, a) == [b<((s, a), (s),,) and K ())(s,))))(s,)),) The bonus (22) derivation is straightforward by applying the connection between random feature and kernel, similar to [Zheng et al., 2022]. In practice, we can also approximate UCB bonus with Diff-SR as -1
b(s, a) = a)) (E., (EstarED 05 a)og(s, a)T + +1 os(s,a).
(23)
These bonuses are subsequently added to the reward in (21) to facilitate exploration.
Remark (Comparison to Spectral Representation): Both the proposed Diff-SR and the existing spectral representation for RL [Zhang et al., 2022, Ren al., 2022b,a, Zhang et al., 2023b] are extracting the representation for Q Q-function. However, we emphasize that the major difference lies in that existing spectral representations seek low-rank linear representations, while our representation is sufficient for representing Q-function, but in a nonlinear form, as we revealed in (19). This nonlinearity fact comes from the partition function (s, a), which constrained to be in [Zhang et al., 2022, 2023b), and thus, less flexible for representation. Meanwhile, even with nonlinear representation, has been shown that the corresponding bonus still enough for exploration [Guo et al., 2023], without additional computational cost.
7

---

4 Related Work
Representation learning in RI. Learning good abstractions for the raw states and actions based on the structure of the environment dynamics is thought to facilitate policy optimization. To effectively capture the information in said dynamics, existing representation learning methods employ various techniques, such as reconstruction [Watter et al., 2015, Hafner et al., 2019b, Fujimoto et al., 2023], successor features [Dayan, 1993, Kulkarni et al., 2016, Barreto et al., 2017], bisimulation [Ferns et al., 2004, Gelada et al., 2019, Zhang et al., 2021], contrastive learning [Oord et al., 2018, Nachum and Yang, 2021], and spectral decomposition [Mahadevan and Maggioni, 2007, Wu et al., 2018, Duan al., 2019, Ren et al., 2022b]. Previous works also leverage the assumption that the transition kernel possesses a low-rank spectral structure, which permits linear representations for the state-action value function and provably sample-efficient reinforcement learning [Jin et al., 2020, Yang and Wang. 2020, Agarwal et al., 2020b, Uehara et al., 2022]. Based on this, there have been several attempts towards both practical and theoretically grounded reinforcement learning algorithms by extracting the spectral representations from the transition kernel [Ren et al., 2022c, Zhang et al., 2022, Ren et al., 2022a, Zhang et al., 2023a]. Our approach aligns with this paradigm but distinguishes itself by learning these representations via diffusion and enjoying the flexibility of mnergy-based modeling.
Diffusion model for RL. By virtue of their ability to model complex and multimodal distribu- tions, diffusion models present themselves as well-suited candidates for specific components in reinforcement learning. For example, diffusion models can be utilized synthesize complex behav- iors [Janner et al., 2022, Ajay et al., 2022, Chi et al., 2023, Du et al., 2024], represent multimodal policies [Wang al., 2022, Estruch al., 2023, Chen et al., 2023b], or provide behavior regularizations [Chen et al., 2023a]. Another line of research utilizes the diffusion model as the world model. Among them, DWM [Ding al., 2024] and SynthER (Lu al., 2023] train diffusion model with off-policy dataset and augment the training dataset with synthesized data, while Poly- GRAD [Rigter et al., 2023] and PGD [Jackson al., 2024] sample from the diffusion model with policy guidance to generate near on- on-policy trajectory. On larger scale. UniSim [Yang et al., 2023] employs a video diffusion model to learn a real-world simulator that accommodates instructions in various modalities. All of these methods incur great computational costs because they all involve iteratively sampling from the diffusion model to generate actions or trajectories. contrast, our method leverages the capabilities of diffusion models from the perspective of representation learning, thus cireumventing the generation costs.
5 Experiments
We evaluate our method with based MDP tasks (Gym-MuJoC) locomotion [Tooorov et al., 201 21) and mmaagebased POMDP tasks (Meta-World Benchmark [Yu et al., 2020]) in this section. Besides, we also provide experiments with state-based POMDP tasks in Appendix E. Our code is publicly released at the project website.
8000
TVS
2000
SSC
SAC
SPC
4000
5000
2000
100
5.0 022
350
035
1.00
1.25
1.50
1.75
2.00
6:0
0225
0.50
1.50
1.2
2000
8.00
0,5
850
8.75
1.00
1.55
1.50
1.75
2000
000
0.5
1.58
025
100
1,2
150
Figure I: The performance curves of the Diff-SR and baseline methods on MBBL tasks, We report the mean (solid line) and one standard deviation (shaded area) across 4 random seeds.
5.1 Results of Gym-MuJoCo Tasks
In this first set of experiments, we compare Diff-SR against both model-based and model-free baseline algorithms, using the implementations provided by MBBL [Wang et al., 2019]. For representation- based RL algorithms, we include LV-Rep [Ren et al., 2022a], SPEDE [Ren et al., 2022c] and Deep
8

---

Table Performances of Diff-SR and baseline RL algorithms after 200K environment steps. Results are averaged across 4 random seeds and window size of 10K steps. Results marked with are taken from MBBL [Wang et al., 2019] and i are taken from LV-Rep [Ren et al., 2022a].
HaltCheetah
Reacher
Humanoid-ET
Pendulum
I-Pendulum
ME-TRPO*
2283.7 + 900.4
13.4
5.2
72.9 + 8.9
177.3 + 1.9
- -126.2 86.6
PETS-RS
966.9 171.6
-9.1
109.6 102.6
167.9 +
25.1
Model-Based
PETS-CEM*
2793.3 # 879.9
-12.3 5.2
110.8 = 90.1
167.4 53.0
-20.5 28.9
Best MBBL
3639.0 1135.8
0.1
1377.0 4: 150.4
177.3 1:
0.0
PolyGRAD
2563.5 + 2012
-20.7
1026.6 58.7
166.3 1: 6.3
33.5 4.8
PPО*
84.4
0.7
+ 39.1
163.4
40.8 21.0
Model-Free RL
TRPO*
-12.0+
0.6
280.8 + 5.2
166.7 + 7.3
27.6 15.8
SAC* (3-layer)
4000.7 202.
-6.
1794.4 + 458.3
168.2 + 9.5
0.2 +.1.1
DeepSHf
4180.4 113.8
3.6
168.6 5.1
168.6 5:1
-0.2
SPEDE
4210.3 1: 92.6
1.1
$86.9 1: 93.2
169.5 1: 0.6
1:0::00 U.O
Representation RL
LV-Rept
5657.6 + 439.5
-5.8
1086 278.2
167.1 3.1
0.0 0.0
Diff-SR
7223.53 437.1
-6.5 2.0
821.3 + 118.9
162.3 + 3.0
0.0
Ant-ET
Hopper-ET
S-Humanoid-ET
CartPole
Walker-ET
ME TRPO*
42.6 21.1
1272.5 4: 500.9
154.9 +: 534.3
160.1 +: 69.1
1600.3 4: 657.5
PETS-RS*
130.0 + 148.1
205.8 36.5
320.7 + 182.2
195.0 + 28.0
312.5 + 490.4
Model-Based RL
PETS-CEM
81.6 145.8
129.3 36.0
355.1 157.1
195.5 3.0
260.2 + 536.9
Best MBBL*
275.1 309.1
1272.5 + 500.9
1084.3 77.0
200.00 : 0.0
312.5 + 493.4
PolyGRAD
433.6 + 158.6
1151.6 + 182.3
1110.1 = 181.8
193.1 # 11.5
268. 4 + 77.3
PPQ*
80.1 17.3
758.0 62.0
454.3 36.7
86.5 + 7.8
300.1 + 17.2
Model-Free RL
TRPO*
116.8+4472 47.3
33.5
281.3 10.9
15.7
229.5
SAC* (3-layer)
2012.7 + 571.3
1815.5 + 655.1
834.6 + 313.1
199.4 + 0.4
2216.4 + 678.7
DeepSFt
768.1 1 11.1
548.9 253.3
533.8 = 154.9
1915 5.8
165.6 :: 127.9
Representation RL.
SPEDE
806.2
T32.2
154.7
4:
5016
LV-Repi
2511.8 =: 460.0
2201. + 196.0
963.1 = 45.1
200.7 0.2
2523.5 = 333.9
Diff-SR
4788.6 1: 623.1
2800.5 95.4
1160.3 + 150.1
199.9 1.1
3722.2 1 406
Successor Feature (DeepSF) [Barreto al., 2017] baselines. Note that the SPEDE is special case of Gaussian EBM. We also include PolyGRAD [Rigter et al., 2023], a recent method that utilizes diffusion models for RL as additional baseline. All algorithms executed for 200K environment steps and we report the mean and standard deviation of performances across 4 random seeds. More implementation details including the hyper-parameters can be found in Appendix C.
Table 1 presents the results, demonstrating that Diff-SR achieves significantly better or comparable performance to all baseline methods in most tasks except for Humanoid-ET. Specifically, in Ant and Walker, Diff-SR outperforms the second highest baseline LV-Rep by 90% and 48%. Moreover, Diff-SR consistently surpasses PolyGRAD in nearly all environments. We also provide the learning curves of Diff-SR and baseline methods in Figure for an illustrative interpretation of the sample efficiency Diff-SR brings.
Runtime of All Gym Envs
Computational Efficiency and Runtime Comparison Compared to other diffusion-baate based RL algorithms, Diff- SR harnesses diffusion's 's flexibility while circumventing the time-consuming sampling process. To showcase this, we record the runtime of Diff-SR and PolyGRAD on MBBL tasks using workstations equipped with Quadro RTX 6000 cards. Results in Figure illustrate that Diff-SR about 4 x faster than PoDyGRAD, and such advantage is consistent across all environments. We provide a per-task breakdown of the runtime results in Appendix 4 due to space constraints.
Hours
DiF-SR
LV-Rap
PolyGRAD
Figure Runtime comparison between Diff-SR vs. IV-Rep vs. diffusion-based RL (PolyGRAD).
5.2 Results of Meta- World Tasks
As the most difficult setting, we evaluate Diff-SR with 8 visual-input tasks selected from the Meta- World Benchmark. Rather than directly diffuse over the space of raw pixels, we resort to techniques similar to the Latent Diffusion Model (LDM) [Rombach et al., 2022] which first encodes the raw

---

bin- bin-pickiin
box- close
button-press
docr-open
1.00
1.00
100
100
SurrwAN Batts (%0
Succass Rate (%0
Succusss Ratu (%6)
75
75
50
50
50
25
1.0
0.2
0.8
0.2
0.6
Envtronment Steps
Environment Stops
1eG
Environment
Steps
Environment Stops
drawer-close
hammer
sti ck
100
100
Success hate [84
Success hatt [K)
Success hata (N3)
75
MuM
50
Drsameev2
DrQ-+2
25
25
0.0
0.2
0.4
DE
0.8
1.0 lee
0.0
0.2
0.4
0.6
0.8
1.0 0.0
0.2
0.4
0.6
0.9
1.0
0.0
0.2 04
0.6 0.8
Enstronmant Stops
Environment Stops
led
Environment Stops
Environment Stops
166
Figure 3: Performance curves for image-based POMDP tasks from Meta-World. World. We report the mean
(solid line) and the standard deviation (shaded area) of performances across random seeds.
observation image to 1-D compact latent vector and afterward performs the diffusion process the latent space. To deal with the partial observability, we truncate the history frames with length L, encode these L. frames into their latent embeddings, concatenate them together, and treat them the state of the agent. The learning of diffusion representation thus translates into predicting the next frame's latent embedding with the action and L-step concatenated embedding. Following existing practices [Zhang et al., 2023a], we set = 3 for all Meta-World tasks. We use DrQ- V2 [Yarats et al., 2021] as the base RL algorithm, and more details of the implementations are deferred to Appendix
The performance curves of Diff-SR and the baseline algorithms are presented in Figure 3. We see that Diff-SR achieves greater than 90% success rate for seven of the tasks, 4 more tasks than the second best baseline yLL--Rp. Overall, Diff-SR exhibits superior performance, faster convergence speed, and stable optimization in most of the tasks compared to the baseline methods. Finally, although Diff-SR does not require sample generation, we present the reconstruction results to validate the efficacy of the score function vb as (3', in Figure 6.
6 Conclusions and Discussions
We introduce Diff-SR, a novel algorithmic framework designed to leverage diffusion models for reinforcement learning from representation learning perspective. The primary contribution of our work lies in exploiting the connection between diffusion models and energy-based models, thereby enabling the extraction of spectral representations of the transition function. We demonstrate that such diffusion-based representations can sufficiently express the value function of any policy, facilitating efficient planning and exploration while mitigating the high inference costs typically associated with iiffusion-based methods. Empirically, we conduct comprehensive studies to validate the effectiveness of Diff-SR in both fully and partially observable sequential decision-makinn
problems. Our results underscore the robustness and advantages Diff-SR across various benchmarks. However, the main limitation of this study is that Diff-SR has not yet been evaluated with real- wolldorld and multi-task data. Future work will focus on testing Diff-SR on real-world applications, such as robotic control.
10

---

References
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in neural information processing systems, 33:20095-20115 2020a.
Alekh Agarwal, Sham M. Kakade, Akshay Krishnamurthy, and Wen Sun. FLAMBE: structural complexity and representation learning of low rank mdps. In Advances in Neural Information Processing Systems 33: Anmual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020b.
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. con- ditional generative modeling you need for decision decision-making? arXiv preprint arXiv:2211.15657, 2022.
Susan Amin, Maziar Gomrokchi, Harsh Satija, Herke van Hoof, and Doina Precup. A survey of exploration methods in reinforcement learning. arXiv preprint arXiv:2109.00 00157, 2021.
André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, David Silver, and Hado van Hasselt. Successor features for transfer in reinforcement learning. In Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 4055-4065, 2017.
Salomon Bochner. Vorlesungen uber fouriersche integrale. In Akademische Verlagsgesellschaf, 1932.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models world simulators. 2024. URL http:://ppeaal..oo////earrccc//e aesels-sswworrd-siiin
Huayu Chen, Cheng Lu, Zhengyi Wang, Hang Su, and Jun Zhu. Score regularized policy optimization through diffusion behavior. In The Twelfth International Conference on Learning Representations, 2023a.
Huayu Chen, Cheng Lu, Chengyang Ying. Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net. 2023b.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- ran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine Deep reinforcement learning handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- DisacBianchi, and R. Garnett, editors, Advances in Neural Infor- mation Processing Systems, volume 31. Curran Associates, Inc., 2018
Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and Pieter Abbeel. Model-based reinforcement learning via meta-policy optimization. In Conference on Robot Learning, pages 617-636 629. PMLR, 2018.
Bo Dai, Bo Xie, Niao He, Yingyu Liang, Anant Raj, kernel methods via doubly stochastic gradients, Advances in neural information processing systems, 27, 2014.
F Balcan, and Le Song. Scalable
Bo Dai, Zhen Liu, Hanjun Dai, Niao He, Arthur Gretton, Le Song, and Dale Schuurmans. Exponential family estimation via adversarial dynamics embedding, Advances Neural Information Processing Systems, 32, 2019.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Comput., 5(4):613-624, 1993.

---

Mare Peter Deisenroth and Carl Edward Rasmussen. Pilco: A model-based and data-efficien approach to policy search. Lise Getoor and Tobias Scheffer, editors, Proceedings of the 28th International Conference on Machine Learning, ICML. 2011, Bellevue, Washington, USA, June 28 July 2, 2011, I, pages 465-472, Omnipress, 2011.
Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng. Diffusion world model. arXiv preprint arXiv:2402.03577, 2024.
Yilun Du, Sherry Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Josh Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. Advances in Neural Information Processing Systems, 36, 2024.
Yaqi Duan, Zheng Tracy Ke, and Mengdi Wang. State aggregation learning from markov transition data. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 4488-4497, 2019.
Bradlcy Efron. Tweedic's s formula and selection bias. Journal of the American Statistical Association, 106(496);1602-1614, 2011.
Yonathan Efroni, Chi Jin, Akshay Krishnamurthy, and Sobhan Miryoosefi. Provable reinforcement learning with a short-term memory. In International Conference on Machine Learning, pages 5832-5850. PMLR, 2022.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In UAI '04, Proceedings of the 20th Conference in Uncertainty in Artificial Intelligence, Banff, Canada, July 7-11, 2004, pages 162-169. AUAI Press, 2004.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference on machine learning, pages 1587-1596. PMLR, 2018.
Scott Fujimoto, Wei-Di Chang, Edward J. Smith, Shixiang Gu, Doina Precup, and David Meger. For SALE: state-action representation learning for deep reinforcement learning. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrlPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023.
Kenji Fukumizu, Francis R Bach, and Michael Jordan, Kernel dimension reduction in regression. 2009.
Tanmay Gangwani, Joel Lehman, Qiang Liu, and Jian Peng. Learning belief representations for imitation learning in pomdps. In uncertainty in artificial intelligence, pages 1061-1071. PMLR, 2020.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California. USA, volume 97 of Proceedings of Machine Learning Research, pages 2170-2179. PMLR, 2019.
Jiacheng Guo, Zihao Li, Huazheng Wang, Mengdi Wang, Zhuoran Yang, and Xuezhou Zhang. Prov- ably efficient representation learning with tractable planning in low-rank pomdp. In International Conference on Machine Leurning, pages 1967-11997. PMLR, 2023.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A Pires, and Rémi Munos. Neural predictive belief representations. 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 1861-1870, PMLR, 10-15 Jul 2018
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. 2019a
12

---

Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 2555 2565. PMLR, 2019b.
Danijar Hafner, Timothy P. Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Philippe Hansen-Estruch, Iya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 6851, 2020.
Matthew Thomas Jackson, Michael Tryfan Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson, and Jakob Foerster. Policy-guided diffusion. arXiv preprint arXiv:2404.0633 2024.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.099) 0.999,, 2022.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on learning theory, pages 2137-2143. PMLR, 2020.
Tejas D Kulkarni, Ardavan Saeedi, Simanta Gautam, and Samuel J Gershman. Deep successor reinforcement learning. arXiv preprint
2016.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization, 2018
Tor Lattimore and Csaba Szepesvári. Bandit algorithms. Cambridge University Press, 2020.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic:
Deep reinforcement learning with a latent variable model. volume 33, pages 741-752, 2020.
Sergey Levine and Pieter Abbeel. Learning neural network policies with guided policy search under unknown dynamics. Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu, Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-577 5787, 2022.
Cong Lu, Philip Ball, Yee Whye Teh, and Jack Parker-Holdes Synthetic experience replay. In Advances in Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10- -0-18 16, 2023, 2023.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for -osel deep reinforcement learning with theoretical guarantees, 2021.
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A laplacian framework for learning representation and control in markov decision processes. J. Mach. Learn. Res., 8:2169-223, 2007.
13

---

Ofir Nachum and Mengjiao Yang. Provable representation learning for imitation with contrastive fourier features. Advances in Neural Information Processing Systems, 34:30100-301 12, 2021.
Anusha Nagabandi, Gregory Kahn, Ronald S. Fearing, and Sergey Levine. Neural network dy-
namics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 7559-7566, 2018. doi: 10.1 10. 8866388....
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free can be a strong baseline for many pomdps. 2021.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03888 2018.
Reda Ouhamma, Debabrota Basu, and Odalric Maillard. Bilinear exponential family of mdps: frequentist regret bound with tractable exploration & planning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pages 9336 9344, 2023.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Shuang Qiu, Lingxiao Wang, Chenjia Bai, Zhuoran Yang, and Zhaoran Wang. Contrastive ucb: Prov- ably efficient contrastive self-supervised learning in online reinforcement learning. In International Conference on Machine Learning, pages 18168-18210. PMLR, 2022.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International conference on machine Pmir, 2021.
learning, pages 8821-8831.
Tongzheng Ren, Chenjun Xiao, Tianjun Zhang, Na Li, Zhaoran Wang, Sujay Sanghavi, Dale Schuurmans, and Bo Dai. Latent variable representation for reinforcement learning. arXiv preprint arXiv:2212.08765, 2022a.
Tongzheng Ren, Tianjun Zhang, Lisa Lee, Joseph E Gonzalez, Dale Schuurmans, and Bo Dai. Spectral decomposition representation for reinforcement learning. arXiv preprint arXiv:2208.099 09515,, 2022b.
Tongzheng Ren, Tianjun Zhang, Csaba Szepesvári, and Bo Dai. A free lunch from the noise: Provable and practical exploration for representation learning. In Uncertainty in Artificial Intelligence, pages 1686-1696. PMLR, 2022c.
Mare Rigter, Jun Yamada, and Ingmar Posner. World models via policy- siiiee trajectory diffusion. arXiv preprint / 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 10674-10685. IEEE, 2022.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems, 35:36479-36444 2022.
Hiroaki Sasaki and Aapo Hyvärinen. kerneie-ereliized conditional density estimation. arXiv preprint arXiv: (0::::8060000 0754, 2018.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1889-1897, Lille, France, 07-09 Jul 2015. PMLR.
14

---

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-2265. PMLR, 2015.
Yang Song and Diederik P Kingma. How to train your energy-based models. arXiv preprint arXiv:2101.03288, 2021.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXhv:2011.13356 2020.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01465 2023.
Yuval Tassa, Tom Erez, and Emanuel Todorov. Synthesis and stabilization of complex behaviors through online trajectory optimization. 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 4906-4913, 2012. doi: 10. 10.1109/ROS.AA (/00//ROS.2012..63666666 2012.6386025.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for modell-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012. doi: R..100/RRS..2012.63888
Masatoshi Uchara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. arXiv preprint arXiv:2110.0460 2021.
Masatoshi Uchara, Ayush Sekhari. Jason D. Lee, Nathan Kallus, and Wen Sun. Provably efficient reinforcement learning in partially observable dynamical systems. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022,
NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benehmarking model-based reinforcement learning, 2019.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou Difffusion policies as expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06799 061993,, 2022.
Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin A. Riedmiller. Embed to control: locally linear latent dynamics model for control from raw images. In Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-1 2015, Montreal, Quebec, Canada, pages 2746-273 2754, 2015.
Stephan Weigand, Pascal Klink, Jan Peters, and Joni Pajarinen. Reinforcement learning using guided observability, 2021.
Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in rl: Learning representations with 2018.
efficient approximations. arXiv preprint
Lin Yang and Mengdi Wang. Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound. In Proceedings of the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 10746 10756. PMLR, 2020.
Mengjiao Yang, Yilun Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint arXiv:23 (0.06/ 14, 2023.
Zhuoran Yang, Chi Jin, Zhaoran Wang, Mengdi Wang, and Michael Jordan. Provably efficient reinforcement learning with kernel and neural function approximations. Advances in Neural Information Processing Systems, 13113933-3916, 2020.
15

---

Hengshuai Yao, Csaba Szepesvári, Bernardo Avila Pires, and Xinhua Zhang. Pseudo-mdpp and factored linear action models. In 2014 IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL), pages 1-9. IEEE, 2014.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control:
Improved data-augmented reinforcement learning, 2021.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1 100. PMLR, 2020.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Hongming Zhang, Tongzheng Ren, Chenjun Xiao, Dale Schuurmans, and Dai. Provable repre- sentation with efficient planning for partially observable reinforcement learning, arXiv preprint 1wX1::2311.22244, 2023a.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204. 10.000..33900 /3902, 2022.
Tianjun Zhang, Tongzheng Ren, Mengjiao Yang. Joseph Gonzalez, Dale Schuurmans, and Bo Dai. Making linear mdps practical via contrastive representation learning. In International Conference on Machine Learning. pages 26447-26466. PMLR, 2022.
Tianjun Zhang. Tongzheng Ren, Chenjun Xiao, Wenli Xiao, Joseph Gonzalez, Dale Schuurmans, and Bo Dai. Energy-based predictive representations for partially observed reinforcement learning. In Uncertainty Artificial Intelligence, pages 2477-2487. PMLR, 2023b.
Sirui Zheng, Lingxiao Wang, Shuang Qiu, Zuyue Fu, Zhuoran Yang, Csaba Szepesvari, and Zhaoran Wang. Optimistic exploration with learned features provably solves markov decision processes with neural dynamics. The Eleventh International Conference on Learning Representations, 2022.
16

---

A Detailed Derivations
Proposition 3. Equation (17) shares the same optimal solutions with Equation (18).
Proof. Denote b = v1 (s'1 [s'] for simplicity, and we have
8' + as, (3' VI- Bs'
33 (s) (('' b b b- v1 - 356 3501
3 (s, (33 - bl2
2 BB (BB,,,,,,3
(3' (s'+ (s, a ($' b) T
(b-vi- Bs')
(1 Eg ||s Ep 200 [31[1]]
constant The second term equals to comes from the definition of Moreover, since is independent w.r.t. 8', we have
3 + Bu (5, (s) ((E - -bl = E 8' 3'+ (s, aaaaaaa ((3 - 61
We conclude that
|| 3 + B (s, (35 v1- Bs 1
Eg EpL(na,
| 3 Bub (s, ((3, - [ + constant.
= (s.a.
Therefore, we obtain the claim.
B Generalization for Partially Observable RL
In this section we discuss how to generalize Diff-SR to Partially Observable MDP (POMDP), We follow the definition of a POMDP given [Efroni al., 2022], which is formally denoted tuple P = (S, A, O, H, IP, (), where is the state space, A is the action space, and O is the observation space. The positive integer H denotes the horizon length, Po is the initial state distribution. r O x A + [0, 1] is the reward function, P--ls, a) : S x A + (SS) is the transition kernel capturing dynamics over latent states. and O(-(s) S + ((O) is the emission kernel, which induces an observation from given state.
The agent starts state so drawn from (o(s). At each step h, the agent selects action from A. This leads to the generation of a new state 8441 following the distribution P(-[s), 4z), from which the agent observes 0 according to O(-[s441). The agent also receives a reward r(os+ ah++ 1)- Observing o instead of the true state s leads to non-Markovian transition between observations, which means need consider policies Fą : O (A x 0)h hh + A(A) that depend the entire H}.
history, denoted by тh = {oo, ag. ..... Let [ ==
The value functions defined by
V (bb) E
Σ r(o.. ab)lb
(b, an) r(on, an) + Ep, (b++1()))))
(24)
Definition (L-decodability [Efroni et al., 2022]). Yh E [H], define
IA E (0 A) O, 2n = =0 Cz 1. 00.). ,.)))
(25)
A POMDP is L-decodable there exists a decoder p" X - A(S) such that p" p=(xx) b(Tn).
17

---

That is, under L-decodability assumption. it is sufficient to recover the belief state by an L-step memory 24, rather than the entire history тhe This implies that can parameterize the Q-value as function of the observation history Q% ( rather that the unknown belief state Q (bs, (2%) as). By exploiting this T-decodabilltty Zhang al. [2023a] propose probable efficient linear function approximation of Q% (Zn, ah)) by considering a L-step prediction IP" (*x+L (,,
Inspired by this, we apply EBM for Pa (# an),
Pa [xb+L The an) = exp (1b22,, an) (#N+L) log Z (Ih, an)).
(26)
Z (Ih, an) / exp (1((T,, an) (#n+1.) d d
(27)
We then apply the techniques presented in Section 3.2 to learn diffusion representation ((T,, aq) E Rd as an approximation the random features
Sw (æn, an) = exp (-iw & (In, an)) exp ( (In, an) an) /1//2 log (an, (n)
(28)
The learned representation is subsequently utilized to parameterize the value function Q7, (In, an) for policy optimization demonstrated by [Zhang et al., 2023a]. In particular, consider value function approximation Qg,0 an.) = &0 (Jh, an) f and update it by
Periii (5) = ,,ya, (r
(Qg,o aa') Qg, ")
(29)
The policy, now conditioned on the L-step history I. is updated by max, Bgw a (1.)) Qe (r, a)]. We refer interested readers [Zhang et al., 2023a] for more details on representation learning in POMDPs.
C Details and Analysis for Fully Observable MDP Experiments
Table 2: Hyperparameters used for Diff-SR in state-based MDP environments.
Hyperparameter
Value
Actor Learning Rate
0.003
Critic Learning Rate
0.0003
Learning Rate for v/, 1,,5,0
0.0001
Actor Hidden Layer Dimensions
(256, 256)
Diff-SR Representation Dimension
256
Discount factor
0.99
Critic Soft Update Factor
0,005
Batch Size
1024
Number of Noise Levels
1000
Network Width
256
& Network Hidden Depth
1
C Network Width
512
< Network Hidden Depth
1
C.1 Baseline Methods
For baseline methods, include ME-TRPO [Kurutach et al., 2018], PETS [Chua et al., 2018), and the best model-based results among Luo et al. [2021], Deisenroth and Rasmussen [2011]. Heess et al. [2015]. Clavera et al. [2018], Nagabandi et al. [2018]. Tassa et [2012] and Levine and Abbeel [2014] from MBBL [Wang et al., 2019]. For model-free algorithms, we include PPO [Schulman al., 2017], TRPO [Schulman al., 2015] and SAC [Haarnoja al., 2018].
C.2 Experiment Setups
We implemented our algorithm based on Soft Actor-Critic (SAC). use feature update ratio to denote the frequency of updating the diffusion representations compared to critic updates. We
18

---

Table Performance on various continuous control problems with partial observation. We average results across 4 random seeds and window size of 10K. Diff-SR achieves a similar or better performance compared the baselines. Here, Best-FO denotes the performance Diff-SR using full observations inputs, providing reference on how well an algorithm can achieve most in our tests.
HalfCheetah
Humanoid
Walker
Ant
Hopper
Pendulum
Diff-SR
3864.2 = 482.3
650.1 = 57.4
1860.5 912.1
3189.7 = 720.7
1357.6 = 506.4
167.4 = 4.4
Vl.V-Rpp
3596.2 874.5
806.7 120.7
1298.1 276.3
1621.4 472.3
1096.4 130.4
168.2 + 5.3
Dreamer-v2
2863.8 + 386
672.5 + 36.6
1305.8 234.2
1252.1 284.2
758.3
115.8
172.3 + 8.0
SAC-MLP
1612.0 + 223
242.1 +: 43.6
736.5 + 65.6
1612.0 = 223
614.15 + 67.6
163.6 1 9.3
SLAC
3012.4 724.6
387.4 1 69.2
536.5 123.2
1134.8 326.2
739.3 98.2
1 11.2
PSR
2679.75 =386
534.4 = 36,6
862.4 = 355.3
1128.3 166.6
818.8 87.2
159.4 = 9.2
PolyGRAD
987.1 # 374.6
764.5 = 91.2
211.4 = 100.4
493.0 = 95.4
$27.8 97.7
174.0 = 6.1
Best-FO
7223.5:1437.15
823.1=118.9
3722.211:06.1
6738.6:1623.1
2800.5495.4
162.343.0
sweep the value of this parameter within [1, 3, 5, 10, 20] for all MuJoCo experiments and report the configuration that achieved the best results. Other evaluation, test all methods every 5,000 environment steps by simulating 10 episodes and recording the cumulative return. The reported results are the average return over the last four evaluations and four random seeds.
are listed in Table For
D Computational Efficiency and Runtime Comparison
Full results are presented in Figure 4. We use exactly the same experimental setups as used in Table 1, which has been described in 5.1. We observe that in all cases, Diff-SR is about 4x faster than PolyGRAD. Such efficiency can be attributed to the fact that while we utilize diffusion to train the representations, we do not have to sample from the diffusion model iteratively, which is the main computational bottleneck for SOTA diffusion-based RL algorithms.
AmH-T
CartPole
HarhCrewthh
PeppweeE
Hamamat-ET
Den
Rup
Prt-RRAD
Do SR
LVP
Po
Us-Rp
Pay/RA
Ui-tpp
Ps
Figure Per-task runtime of Diff-SR, LV-Rep and PolyGRAD on tasks from MBBL.
E Details and Analysis for State-based Partially Observable MDP Experiments
E.1 Implementation Details
We also experiment Diff-SR with state-based Partially Observable MDP (POMDP) tasks. We contruct partially observable variant based on OpenAl gym Muloco [Todorov et al., 2012]. Adhereing to standard methodology, we mask the velocity components within the observations presented to the agent, effectively rendering the tasks partially observable [Ni et al., 2021, Weigand et al., 2021, Gangwani al., 2020]. Under this masking scheme, a single observation insufficient for decision- making. Therefore, the agent must aggregate past observations to infer the missing information and select appropriate actions. The reconstruction of missing information can done by aggregating a past history of length L, demonstrated in Definition 1.
19

---

The POMDP experiment follows similar setup to the fully observable one described in Appendix C. In the partially observable setting, velocity information masked, and we concatenate the past observations follow [Zhang et al., 2023a]. The universal hyperparameters used for POMDP are listed in Table 4. For each individual environment, we explored an array of parameters and chose the highest- pereerstpproooingggg configuration. Specifically, the critic and actor learning rates were varied across {0.0015, 0.00015), the model learning rate was varied across {0.0001, 0.0003, 0.0008} and the feature update ratio was varied across {1, 3,5, 10, 20).
We evaluate six baselines in our experiments: a diffusion approach, PolyGRAD [Rigter et al., 2023], two model-based approaches, Dreamer [Hafner et al., 2019a, 2021] and Stochastic Latent Actor-Critic (SLAC) [Lee et al., 2020]. model-free baseline, SAC-MLP, that concatenates history sequences (past four observations) as input to an MLP layer for both the critic and policy, and the neural PSR [Guo al., 2018]. We also compared to representation- based baseline, Jul. |Zhang et al., 2023a]. All methods are evaluated using the same procedure as in the fully observable setting. As reference, we also provide the best performance achieved in the fully observable setting (without velocity masking), denoted as Best-FO,, which serves benchmark for the optimal result algorithm can achieve in our tests.
E.2 Results and Analysis
Table presents all experiment results, showing effectiveness of Diff-SR in partially observable continuous control tasks. The proposed method delivers superior results in 4 out of 6 tasks (HalfChee- tah, Walker, Ant, Hopper). It significantly outperforms other algorithm in Walker and Ant, and achieved comparable result with the lowest standard deviation on Pendulum, indicating consistent performance.
The wall time comparison between Diff-SR and PolyGRAD is shown in Figure 5. In the Humanoid task, Diff-SR is approximately 3 times faster than PolyGRAD, whereas the other tasks, about 4 times faster.
HalfCheetah
Humanaid
Valker
18
Hour
HOurS
HOUrS
30
0-
SR
PelyRRDD
PolyGRAD
DrfisR
PolyGRAD
AnL
Hapwer
Per dulum
10
HOUTS
Hours
2-
Drf-st
PayOMAD
DIFFSR
RoyG/DD
DIf-SR
PalyoR/D
Figure 5: Per-task running time of DifT-SR and PolyGRAD on tasks from MBBL with partial observation.
E.3 Ablations And Modifications
Masking Observations Observations consist of both velocities and positions. Masking velocities, rather than positions, more accurately reflects real-world scenarios where positions are typically directly observed, whereas velocities must be inferred. In the humanoid environment, we specifically mask only the gevelocity component.
20

---

Instead of directly diffusing over the raw pixel space, used the VAE structure from Zhang et al. [2023a] to first encode the pixel observation o into a 1-D latent embedding e. To deal with the partial observability, we also follow Zhang et al. [2023a] observations (denoted as o²) together as the state of the agent, denoted as e². Let the next frame be of and its embedding be e', the representation learning objective thus translates to fitting the score Table 4: Hyperparameters used for Diff-SR state-based POMDP experiments.
Hyperparameter
Value
Actor Hidden Layer Dimensions
(512, 512)
Diff-SR Representation Dimension
512/ 128 (cheetah)
Discount Factor
0.99
Critic Soft Update Factor
0.005
Batch Size
1024
Number of Noise Levels
1000
Network Width
512
Network Hidden Depth
Network Width
256
C Network Hidden Depth
At the beginning of training and evaluation, we randomly sample actions from the action space. This is necessary because we use concatenated observations to sample actions from the policy. This approach might explain the larger standard deviation observed in POMDP compared to the fully observable case.
Random Action
F Image-based Partially Observable MDP Experiment
F1
Implementation Details
concatenate the embeddings of the past three
function ré(es, a) and ((ē', (3), i.e. performing the diffusion in the latent space. In the following paragraphs, we will detail the architectures for Diff-SR.
Diffusion Representation Learning. Unlike previous state-based experiments, we formalize the network &a and ( using the L.N_ResNee architecture proposed by IDQL [Hansen-Estruch al., 2023]. Compared to standard MLP networks, LN Le_ReeNeeett is equipped with layer normalization and skip connections, making expressive enough for diffusion modeling. For representation learning, is worthwhile to note that, apart from the loss objectives defined in Eq (18), we also preserve the gradients of the representation networks and train them with the critic's loss defined in (21). This design choice aligns with previous works [Zhang et al., 2023a] and encourages the representations to contain task-relevant information.
RL optimization. We develop our code based on DrQ-V2 [Yarats et al., 2021], model-free RL algorithm designed for tasks with visual inputs. We preserved most of the design choices from DrQ-V2, except for the architectures of the actor and the critic. For the actor network, it receives the concatenated embeddings ex and outputs an action. The critic networks receive the diffusion representation dy as defined in the main text and predict the Q-values.
The hyper- parameters for the score functions vb S, the actor and the critic networks are listed in Table
F.2 Generation Results
In addition to the learning curves, we also present the qualitative generation results using the learned score functions. Figure 6 illustrates the progression of the diffusion process over time. the figure, we sample batch of data (o², a, o') from the replay buffer, and the first row depicts the ground-trufe target image o'. Starting from the second row, we sample a random Gaussian noise o100 ~ N /) and iteratively apply the learned reverse diffusion process using the guidance of (e3 a) to generate the latent embeddings at various stages of the diffusion. For every 200 steps, we pass the latent embedding to the decoder to obtain reconstructed images o; and visualize them in the figure.
21

---

Table 5: Hyperparameters used for Diff-SR image-based POMDP experiments.
Hyper-parameters
Value
Actor Learning Rate
0.0001
Critic Learning Rate
0.0001
Learning Rate for Φ, Ç.0
0.0003
Actor Hidden Layer Dimensions
(1024, 1024)
Dift-SR Representation Dimension
1024
Discount Factor y
0.99
Critic Soft Update Factor
0.01
Batch Size
1024
Number of Noise Levels
1000
LN_ResNet Layer Width for +
512
LN Re_eeNeettt Layer Width for <
512
# of LN T_RReNet Layers for 4b
of LN IN_ReeNNet Layers for
original obs
01000
6go0
Diffusion
%,00
Process
0400
0200
Figure 6: The generation results of Diff-SR.
The second row corresponds to the initial noisy embeddings and exhibits significant distortion and noise. As the denoising steps progress, we observe gradual reduction in noise and an increasing clarity in the generated images. By the time we reach 0000 the overall structure of the scene becomes more discernible, although some artifacts remain. Further along the denoising process, the images exhibit substantial improvements in terms of detail and realism. The final output, og closely resembles the original observations, indicating the effectiveness of our score functions in capturing the underlying information about the data.
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The main claims are that (1) by exploiting the energy- model view of the diffusion model, we develop an algorithm to learn latent variable representation of the transition function (Section 3.2) and (2) we demonstrate that said representation sufficiently expressive to linearly represent the value function, paving the way for efficient planning and exploration, cireumventing the need for sample generation from the diffusion model, and thus, avoiding the inference costs associated with prior diffusion-based methods (Section 3.3).
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discuss the assumptions in the derivations and experiments. We also address limitations in the conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
23

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide full sets of assumptions in the main text body and if an additional proof is required, we supplement the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide representation training algorithm implementation description in Algorithm 2 and full online pipeline implementation in Algorithm 1. We provide experi- mental setup and procedures 5, We also provide additional implementation details and hyperparameter set in C.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
24

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We have included all experimental details; code is released on our website.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines for more details.
(https://nips.c httpp:://iis.cc//ppbiicg///uu// cccpuuullguuddssoooeellmsissssoo
While we encourage the release of code and data, we understand that this might be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https://nips.cr (tttps:/nips.cc/ppubccquuu poublicgguidesoooeSSubiieseee
for more details.
The authors should provide instructions on data access and preparation, including h ow to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We provide full algorithm description in Algorithm and 2, and experiment details in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: For all experimental results, we run over multiple random seeds and provide error bars for each result.
Guidelines:
25

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errоrs is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We experiments Quadro RTX 6000 and provide time results. Our experiments can be done on standard graphics cards and do not require excessive computational capabilities.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: We acknowledge and affirm the importance of the NeurIPS Code of Ethics. We have carefully considered the ethical implications of our work throughout the research and development process.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
26

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We develop reinforcement learning algorithms with broad applicability. These algorithms can be used for diverse applications, some with positive outcomes and others with potential risks. Our contribution lies in the general method, not its specific use case.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploymennt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA].
Justification:
This project does not use data or models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have credited creators and original owners of assets (e.g.. code, data, models) used in the paper. The license and terms of use are explicitly mentioned and properly respected.
27

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com pnperssithcode.comddlaaasst has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The assets are not released at this time.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA].
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA].
28

---

Justification:
The paper does not involve IRB approval, crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
29


Reviews:
[{"review_id": "fHtIYH0rnQ", "reviewer": "Reviewer_TNcM", "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.", "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?", "limitations": "/", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "* **W1: Some discussions on the representation quality are expected, e.g. to check whether Diff-Rep can recover latent state representations.** \n\nWe 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.\n\n* **Q1: Why use gradient descent rather than closed-form solution?**\n\nWhile 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xaesMmDpGV", "reviewer": "Reviewer_2EpY", "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.", "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?\n\n2. 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?\n\n3. 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.\n\n4. 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.\n\n\n\n[1] https://arxiv.org/abs/2305.16381\n\n[2] https://arxiv.org/abs/2402.15194", "limitations": "The authors adequately addressed limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. 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. \n\n3. 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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "* **W1 \\& Q1: Limited theoretical analysis.** \n\nThe major contribution of this paper is developing an efficient algorithm for sufficient representation through diffusion, with intensive empirical evaluation.\nWe 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. \n\n[1] Tongzheng Ren, et al. Latent variable representation for reinforcement learning. 2022.\n\n[2] Tongzheng Ren, et al. Spectral decomposition representation for reinforcement learning. 2022.\n\n* **W2 \\& Q2: Lack of realistic or real-world tasks.**\n\nThe 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.\n\nWe 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. \n\n[1] Tongzheng Ren, et al. Latent variable representation for reinforcement learning. 2022.\n\n[2] Hongming Zhang, et al. Provable Representation with Efficient Planning for Partially Observable Reinforcement Learning. 2023.\n\n\n* **Q3: Optimizing objective (18) along with model-free RL seems burdensome and other training strategies, such as pre-training the representations, may be exploited.**\n\nIn 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.\n\nIn 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.\n\n\n* **Q4: Connections to papers about RL for Diffusion.**\n\nSome 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)$:\n$$p_{\\text{target}}(x)\\propto p_{\\text{data}}(x)\\exp(Q(x)).$$\nIn 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. \n\nDespite this, our method leverages the capabilities of diffusion models from the perspective of representation learning, which fundamentally differs from the aforementioned methods.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hPrbMQAFqz", "reviewer": "Reviewer_sdAM", "summary": "## Main summary\n\nThe 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. \n\nThe 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.\n\nThey 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.\n\n## More detailed summary of methods\n\nTheir 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$. \n\nThe 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.\n\nNext, 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.\n\nTo 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.\n\nFinally, 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.", "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$?\n\n1. 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.", "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.\n\nMy 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.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "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.\n\n\n1. 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.\n\n1. Quality of exposition and overall quality: I believe these can be significantly improved. See Weaknesses.\n\n1. 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). \n\n1. 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.\n\n1. 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.\n\n1. 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.\n\n1. 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.\n\n1. Further, appropriate citations for EBMs seem to be missing in Section 3.2.\n\n1. An explicit definition of $\\nu(\\tilde{s}’, \\beta)$ is missing in Equation 12.\n\nClarity and writing:\n\nThe 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:\n\n1. 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.”\n\n1. Line 50: “to avoid suboptimal policy” -> “to avoid a suboptimal policy”\n\n1. Line 99: “the existing methods either require [...]  or relies on” -> “the existing methods either require [...]  or rely on” \n\n1. Line 105: “After T-step” -> “After T steps”\n\n1. 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:” \n\n1. Line 117: “as a neural networks” -> “as a neural network”\n\n1. Line 118: “the samples can be generated from [...], and then, following [...]” -> “the samples can be generated by sampling [...], and then following [...]”\n\n1. Line 155: “We apply the EBM for transition operator” -> “We parameterize the transition operator using an EBM”\n\n1. 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”\n\n1. 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)$”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n* **W1: SVD decomposition of transition operator assumption**\n\nThe SVD of the transition operator $T(s' | s, a)$ holds in general MDP settings, but potentially with infinite-dimensional spectrum. \nWhen 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. \nFor continuous $\\mathcal{S}$ and $\\mathcal{A}$, if $T(s'|s, a)$ can be represented by a discrete latent variable model, i.e., \n\n$$T(s'|s, a) = \\sum_{i=1}^k p(s'|i)p(i|s, a),$$\n\nwe 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].\n\n[1] Jordan Bell. The singular value decomposition of compact operators on Hilbert spaces. 2014. \n\n* **W2: Why Diff-Rep is sufficiently expressive to represent the value function of any policy.**\n\nFrom 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. \n\n* **W3: Claims about $\\phi^{*}$**\n\n**Un-normalized conditional density estimation**: \n\nAs $\\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. \n \n**Learning requires full-coverage data**:  \n\nLearning an accurate $\\phi^*$ requires data with full coverage to capture the underlying conditional density $\\mathbb{P}(s' |s, a)$. \n\n**Exploration requires accurate $\\phi^{*}$**:\n\nIn 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. \nWe will include relevant citations in the revision. \n\n[1] Chi Jin, et al. Provably efficient reinforcement learning with linear function approximation. 2020.\n\n[2] Zhuoran Yang, et al. Provably efficient reinforcement learning with kernel and neural function approximations. 2020.\n\n[3] Alekh Agarwal, et al. Flambe: Structural complexity and representation learning of low rank mdps. 2020.\n\n[4] Tongzheng Ren, et al. Latent variable representation for reinforcement learning. 2022.\n\n* **W4: Random Fourier features and explanations.**\n\nThanks 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. \n\nSpecifically, $\\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:\n\\begin{align}\n    \\exp(-\\frac{\\\\|x-y\\\\|^2}{2}) &=\\int p(\\omega)\\exp(-\\mathbf{i}\\omega^\\top(x-y))\\mathrm{d}\\omega \\\\\\\\\n       &=\\int p(\\omega)\\exp(-\\mathbf{i}\\omega^\\top x)\\exp(\\mathbf{i}\\omega^\\top y)\\mathrm{d}\\omega \\\\\\\\\n       &=\\mathbb{E}_{\\omega\\sim p(\\omega)} \\left[\\exp(-\\mathbf{i}\\omega^\\top x)\\exp(\\mathbf{i}\\omega^\\top y)\\right] \\\\\\\\\n\\end{align}\n\nwhere $p(\\omega)\\sim\\mathcal{N}(0, I_d)$. \n\n[1] Ali Rahimi, Benjamin Recht, et al. Random features for large-scale kernel machines. NIPS. 2007. \n\n* **W5: Citations for EBMs**\n\nThanks 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.\n\n* **W6: An explicit definition of $\\nu$**\n\nIn 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.\n\n* **Q1: The presence of $\\sigma^2$ is unnecessary.**\n\nWe 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.\n\n* **Q2: Why does the second term equal 0 in line 591.**\n\nThis is a direct application of the tower property. Specifically,\n\n\\begin{align}\n&\\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')] \\\\\\\\\n&=\\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'])] \\\\\\\\\n&=0.\n\\end{align}\n\nwhere the first equation is due to the independence between $b$ and $s'$; the second equation comes from the definition of $b$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z1xmbPlrYH", "reviewer": "Reviewer_iUqH", "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.", "questions": "Could the author provide a runtime comparison between Diff-Rep and LV-Rep?\n\nCould the author provide some plots showing the accuracy of the estimated Q function when using the proposed representation?\n\nCould the author suggest any methods that might help interpret the latent representation?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\nThe 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\".\n\nWhile 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.\n\nMinor\nSome derivations/explanations are needed for the approximation in Eq 12.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "* **W1: Table 1 performance**\n\nWe 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.  \n\n\n* **W2: Diff-Rep lacks the flexibility of conditioning the policy at inference time?**\n\nWe 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.\n\n\n*  **W3: Some minor derivations or explanations are needed for the approximation in Equation 12**\n\nIn 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$. \n\n[1] Hiroaki Sasaki and Aapo Hyvärinen. Neural-kernelized conditional density estimation. 2018.\n\n\n * **Q1: Runtime comparison between Diff-Rep and LV-Rep?**\n\nThe 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. \n\n\n* **Q2: showing the accuracy of the estimated Q function when using Diff-Rep?**\n\nIllustrating 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. \n\n* **Q3: Suggest some methods to help interpret the latent representation?**\n\nDiff-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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all our reviewers for their effort and time in providing constructive suggestions for our paper. We are delighted that our reviewers recognized the relevance of our problem and the novelty of our method. 

We note that both the reviewer adAM and iUqH raised concerns regarding the expressiveness of Diff-Rep in representing the $Q$-functions of any policy $\pi$. Therefore, we conduct experiments with a toy environment, FrozenLake-v1 (\url{https://www.gymlibrary.dev/environments/toy_text/frozen_lake/}). FrozenLake is a grid world where an agent navigates a frozen lake from the Start (S) to the Goal (G) while avoiding Holes (H). The environment features a 4x4 discrete state space and a discrete action space with 4 possible movements: left, down, right, and up. The agent's movement can be affected by the slippery surface, adding stochasticity to the transition. Rewards are given for reaching the goal (+5) or the holes (-1). 

We first train an optimal policy in this environment, and then at each step, we let the policy take random actions with different probabilities $\epsilon_i\in\\{0,0.2,0.4,0.6,0.8,1.0\\}$ to simulate different policies $\pi_i$. Then we conduct policy evaluation to obtain ground truth values of $Q^{\pi_i}$. 

Our objective is to find out whether the diffusion representation in our paper can sufficiently represent each $Q^{\pi_i}$ with sufficiently low error, as our theory suggests. The diffusion representation proposed in our paper is given by $\phi(s, a) = \texttt{elu}(W\sin(\Omega^\top \psi(s, a)))$, where $\mathbb{P}(\tilde{s}'|s, a; \beta)\propto \exp(\psi(s, a)^\top \nu(\tilde{s}, \beta))$, $\Omega$ is composed by $\omega\sim\mathcal{N}(0, I)$, and $W$ is a learnable matrix. To obtain such representations, we treat the coordinates of the grid as continuous values and perform diffusion training based on the coordinates. Once we obtain $\psi$, we can construct the diffusion representations $\phi$, fix them, and use them to regress the ground truth $Q^{\pi_i}$.

The residual errors of the regression are listed in the following table. As an approximation to the spectral representations, Diff-Rep can accurately represent $Q$-values of any policy $\pi_i$, which validates its expressiveness of representing $Q$-values.

|$Q$-function|$Q^{\pi_0}$|$Q^{\pi_1}$|$Q^{\pi_2}$|$Q^{\pi_3}$|$Q^{\pi_4}$|$Q^{\pi_5}$|
|---|---|---|---|---|---|---|
|Residual Error|$1.2\times 10^{-8}$|$3.6\times 10^{-7}$|$1.2\times 10^{-7}$|$4.9\times 10^{-9}$|$2.0\times 10^{-8}$|$9.2\times 10^{-8}$|

Meta-review: The paper presents an interesting new idea by observing that:

1. it is well known that better representations can lead to better RL
   performance.
2. diffusion models are well recognized for their ability to represent
   high-dimensional multi-modal distributions and that this can translate
   into SOTA RL performance in some cases.
3. however, standard use of diffusion models in RL requires computationally
   expensive sampling which limits their practical use.

It then proposes a new representation learning method that uses insights
and loss functions adapted from diffusion models along with random fourier
features.  The representation learning is integrated into a relatively
standard RL algorithm that learns value functions and policies as linear
functions of the learned representation.  The method is compared against a
variety of diffusion-based and more traditional deep RL methods on 10
standard RL benchmarks with the proposed method performing at or above all
other methods on most of them.  There are POMDP examples in the appendix
which are presented more like examples rather than comparisons.

Note that the MDP comparison's training iterations are capped before
training has converged in any of the methods.  This shows the sample
efficiency of the early learning phase but gives no insight into the
converged performance of any of the methods.  Although this practice is not
uncommon in the RL literature, it does limit the evaluation.  Ideally, 
comparisons at convergence would also be included.

Overall, the paper proposes an interesting new algorithm with good insights
from recent literature and strong empirical results against a broad set of
comparison algorithms.
Decision: Accept (poster)