ContraDiff: Planning Towards High Return States via Contrastive Learning
Authors: Yixiang Shan
Zhengbang Zhu
Ting Long
Liang Qifan
Yi Chang
Weinan Zhang
Liang Yin
Venue: ICLR 2025 Poster (2025)
Keywords: Offline Reinforcement Learning
Decision Making
Diffusion Models
Machine Learning
Primary area: reinforcement learning
TL;DR: We propose CDiffuser, which devises a return contrast mechanism and enhances the trajectory generation-based diffusion RL by pulling the states in generated trajectories towards high-return states while pushing them away from low-return states.

Abstract: The performance of offline reinforcement learning (RL) is sensitive to the proportion of high-return trajectories in the offline dataset. However, in many simulation environments and real-world scenarios, there are large ratios of low-return trajectories rather than high-return trajectories, which makes learning an efficient policy challenging. In this paper, we propose a method called Contrastive Diffuser (ContraDiff) to make full use of low-return trajectories and improve the performance of offline RL algorithms. Specifically, ContraDiff groups the states of trajectories in the offline dataset into high-return states and low-return states and treats them as positive and negative samples correspondingly. Then, it designs a contrastive mechanism to pull the planned trajectory of an agent toward high-return states and push them away from low-return states. Through the contrast mechanism, trajectories with low returns can serve as negative examples for policy learning, guiding the agent to avoid areas associated with low returns and achieve better performance. Through the contrast mechanism, trajectories with low returns provide a ``counteracting force'' guides the agent to avoid areas associated with low returns and achieve better performance.
Experiments on 27 sub-optimal datasets demonstrate the effectiveness of our proposed method. Our code is publicly available at https://github.com/Looomo/contradiff.

Published as a conference paper at ICLR 2025
CONTRADIFF: PLANNING TOWARDS HIGH RETURN STATES VIA CONTRASTIVE LEARNING
Yixiang Shan, Zhengbang Zhu, Ting Long', Qifan Liang, Yi Chang; Weinan Zhang, Liang Yin 1 Jilin University, Shanghai Jiao Tong University (shanyx22, sianggf23) (zhengbangzhu, e jtu. eeu..du...uu..uu................ edu. cn (longting, yichang] edu. on yinla@apexs wtt.aapee.st......m
jlu. edu. cn
ABSTRACT
The performance of offline reinforeement learning (RL) sensitive to the propor- tion of high-return trajectories in the offline dataset. However, in many simulation environments and real-world scenarios, there are large ratios of low-return trajec- tories rather than high-return trajectories, which makes learning efficient policy challenging. In this paper, we propose a method called Contrastive Diffuser (Con- traDiff) to make full use of low-return trajectories and improve the performance of offline RL algorithms. Specifically, ContraDiff groups the states of trajectories in the offline dataset into high-return states and low-return states and treats them as positive and negative samples correspondingly. Then, it designs a contrastive mechanism to pull the planned trajectory of an agent toward high- states and push them away from low-return states. Through the contrast mechanism, trajec- tories with low returns can serve as negative examples for policy learning, guiding the agent to avoid areas associated with low returns and achieve better perfor- mance. Through the contrast mechanism, trajectories with low returns provide a "counteracting force" guides the agent to avoid areas associated with low returns and achieve better performance. Experiments on 27 sub-optimal datasets demon- strate the effectiveness of our proposed method. Our code is publicly available http:///ttuuuccco/oooooooo/o//// www./////oo////////oooo/oooooo/o cont
btttp::/githuu www. com
INTRODUCTION
Offline reinforcement learning (offline RL) (Levine et al. 2020; Prudencio et al.) 2023) is signifi- cant branch of reinforcement learning, where an agent is trained on pre-collected offline datasets and evaluated online. Since offline RI. avoids potential risks from interacting with the environment during policy learning, it has broad applications in numerous real-world scenarios, like commercial recommendation (Xiao & Wang 2021). health care (Fatemi al. 2022), dialog systems (Jaques) et al. 2020). and autonomous driving (Shi et al. 2021).
However, the performance of offline RL methods highly depends on the proportion of the high-return trajectories in the offline dataset. When the dataset contains large proportion of high-return trajec- tories, as is presented in Fb,, offline RL methods can easily learn the pattern of high-return trajectories such that they can achieve excellent performance when interacting with the environ- ment. In contrast, when the dataset has limited number of high-return trajectories, as is presented in Figure (c), offline RI methods struggle to learn a good pattern from the dataset to achieve high returns Hong et al., 2023b). Unfortunately, the issue limited high-return trajectories commonly exists in both simulation environments (e.g., Maze2d) and real-world scenarios (e.g.. robotics con- trol and medical diagnosis). As it is illustrated in Figure[[a) we visualize the probability density of trajectories' returns in Maze2d. We can observe that the number of high-return trajectories is much limited.
To address this issue, previous works (Hong et al. 2023b al typically reweight the importance of samples, assigning higher weights to high-return trajectories (as shown in Figure dd). However, these methods essentially focus on learning the patterns of high-return trajectories, encouraging the agent to replicate actions that lead to higher rewards during interactions with the environment. Nevertheless, if the agent starts low-return region unexpectedly falls into low-return region
"Corresponding Authors.

---

Published as a conference paper at ICLR 2025
leamedd trajectory low--retm trajectories dataset high-retum trajectories datuset
sturt
gozl
state
probubiitty dematy
0.15
Prolune
0.10
cularge
9.00
enlarge
30
(a)
(b)
(c)
(d)
(e)
(1)
Figure 1: (a) The probability density of trajectories' returns in Maze2d; (b) The learned trajectory when high-return trajectories are abundant; (c) The learned trajectory when the number of high- return trajectories is limited; (d) Previous address high-return trajectories limited issue by increasing the weight of high-return trajectories; (e) The example of our solution; (f) The contrastive learning applied by previous RL models.
during the interaction, as illustrated by state 8, in Figure [d), ignoring low-return areas during training makes hard for these methods to escape from 8t, since there's no existing path from 8, to goal. Therefore, in the case of sparse high-return trajectories, focusing solely on learning from high-retum trajectories is insufficient.
In fact, states in high-return trajectories indicate potential areas where high returns can be obtained, while states in low-return trajectories indicate the potential regions where agents might encounter low returns. Therefore, one promising way to fully utilize both low-return and high-return trajec- tories in an offline dataset, is using high-retur trajectories as examples to attract the agent to stay as close as possible, while using states from low-return trajectories to indicate potential low-return regions, guiding the agent to stay away from, as shown in Figure TT).
However, there are mature techniques to pull the states of a trajectory toward high-return states and push them away from low-return states, to the best of our knowledge. Fortunately, there is an analogous case: contrastive learning, which aims to bring a given sample close positive (i.e.. similar) samples and far from negative (i.e., dissimilar) samples (Xiao et al. 2020 Tian et al. 2020 Wang & Qil 2022 Khosla et al., (020). Inspired by that, better use the valuable insights imp implicated in low-return trajectories, we propose to treat states with high return in trajectories of offline dataset as positive samples and those with low return negative samples, and leverage contrastive learning to pull the states toward high-return states and push them away from low-return states, as Figure T() illustrates. It is worth noting that, unlike previous works (Qiu et al.) 2022] Laskin al., 2020 Yuan & Lu. 2022] Agarwal et al.] 2020). which apply contrastive learning to constrain the states of the same trajectory to similar representations and the states of different trajectories to dissimilar representations, as is illustrated in Figure [T) we aim to use contrastive learning to constrain policy toward high-return states and away from low-return states, as is illustrated in Figure dd). Furthermore, the criteria for distinguishing positive and negative samples here are based on the returns rather than the labels.
Through the contrast mechanism, trajectories with low returns serve counteracting force for policy learning, guiding the agent to avoid areas associated with low returns. Additionally, with the guidance of high-retum states, the agent ultimately achieves high returns. However, ordinary states are feedback from the environment rather than generated by the model, applying contrastive mech- anisms to these states produces gradient for policy optimization. Considering some diffusion- based RL methods generate subsequent trajectories for planning (Janner et al.) 2022 Ajay et al.) 2023 in which abundant states are generated by policy model, we build our constrastive mech- anism on those diffusion-based RL methods and propose method called Contrastive Diffuser (ContraDiff). Specifically, we first group the states of the trajectories in the offline dataset into high-return states and low-return states. Then, we learn a diffussion-baset trajectory generation model to generate the subsequent trajectories, and apply a contrastive mechanism to constrain the states of the generated trajectories by pulling them toward the high-return states and pushing them away from the low-return states in the offline dataset. With the contrastive mechanism constrained states for planning, the agent makes decisions towards the high-return states. Experiment results on 27 sub-optimal datasets demonstrate the superior performance of ContraDiff.
In summary, contributions are: (i) We propose a method called ContraDiff, which takes the advantage of low-return trajectories by pulling the states in trajectories toward to high-return states and pushing them away from low-return states. (ii) We perform contrastive learning to constrain the states in the agent's 's planned trajectory and enhance the policy learning. To the best of our knowl- edge, our work is the first which apply contrastive learning to directly improve the policy learning. (iii) Experimental comparison with SOTA methods on 27 sub-optimal datasets and thorough further investigations demonstrate the effectiveness of ContraDiff.
2

---

Published as a conference paper at ICLR 2025
2 PRELIMINARIES
2.1 DENOISING PROBABILISTIC MODELS
Denoising Diffusion Probabilistic Models (Diffusion Models) (Sohl-Dickstein et al. 2015 Song et al. Ho et al. 2020) are a group of generative models, which generate samples by denoising from Gaussian noise. A diffusion model composed of a forward process and a backward process. Given the original data æ q(z), the forward process transfers x into Gaussian noise by gradually adding N((x'; /I - F'''''''' , I), in which I is an identity matrix, 3 is the noise schedule measuring the proportion of noise added at each step, x." == æ is a sample from the offiine dataset, x xx? .... the latents of diffusion. The backward process recovers by gradually removing the noise at each step, which is formulated with a Gaussian distribution (Feller 1949)
noise, i.e., g(<0|< ()
parameterized by 0, i.e., pa pa(x"=1(<z) (x²) = (Neue(>,, i), (x i)), where μo (2'", i) = 1-a"
Vo-
Val TB' vo(a", i), a² -1 (1 - 3²) and (:. is a model to reconstruct z. The objective
1-4' function can be formulated as follows if we fix Cx 1) B,I (Ho et al. 2020
= R WI.N] [||x] tho 0)]])
2.2 CONTRASTIVE LEARNING
Contrastive learning (Schroff al. 2015 Sohn 2016 Khosla et al., 2020; Yeh et al. 2022 Oord et al., 2018) is class of sulf-supervised learning methods which aim pulling similar samples together and pushing dissimilar samples away from each other. Specifically, given a sample æ and similarity measure, the positive set S is defined the collection of samples similar to æ, while the negative set S is defined as the collection of samples dissimilar to x. Contrastive learning minimizes the distance of between and St, and maximizes the distance of and S That is, for each sample æ, select a positive sample x+ E S and negative samples x E ES, the learning loss is:
exp(sim(f(z)). ((xt)))
C C= log
[xxp(sim[f()), f()+)) + Ee es ()p(si((((x)) f(x))))
where f(-) is the function to map samples to latent space and sim(.,-) is the similarity measure.
2.3 OFFLINE RL PROBLEM DEFINITION
Considering system composed of three parts: policy, agent, and environment. The environ- ment in RL is usually formulated as a Markov Decision Process (MDP) (Sutton & Barto 2018) M = (S, A, P,,,,, where S is the state space, A is the action space, P(s' 8, a, is the transi- tion function, y represents the discount factor, is the instant reward of each step. At each step the agent responds to the state of environment 8; by action a according to policy me param- eterized by 0, and gets an instant return Tz. The interaction history is formulated as a trajectory = [(St, az, rt 2 0). In this paper, we define the cumulative discounted reward from step as 24- Eize rrr and call it the return of st.
We focus on the offline RL setting in this paper. Given an offline dataset D P {(8:, at, Tt, 8c+1) 0} consisting of transition tuples, and defining the return trajectory T as R(T) & Σiso E'son're our goal is learning Ty to maximize the expected return without directly interacting with the environ- ment, i.e..
Tg = arg max E E [R(T)]
3 METHODOLOGY
As we discussed previously, the performance of offline RL methods is suppressed when the number of high-retum trajectories is limited. Previous methods of increasing the weights of high-return samples fail to utilize the low-return samples and are ineffective in addressing situations where the agent gets stuck in regions with low returns. To address this challenge, we propose method called Constrastive Diffuser (ContraDiff), which introduces a contrastive mechanism to make full use of low-return trajectories and enhance the performance by constraining the states of the agent's trajectory towards high-retura states and away from low-return states. As is illustrated in Figure 2 Our ContraDiff is composed of two modules: (1) the Planning Module, which aims to generate subsequent trajectories; (2) the Contrastive Module, which is designed to constrain the states in generated trajectories within the high-return areas and away from low-return areas.
3

---

Published as a conference paper at ICLR 2025
Planning Module
origiinal states
MSt:
genented states
actionts)
hight-eetura stazes
pod
Jow-relurn stales
of Kd Loss 9sfv
trajectory
noised trajeciory
dillusion generated Injjectory
Contrastive Module
ContraDif
offline dataset
sampling
MLP
contrast
objectives
Figure The overall framework of ContraDiff. ContraDiff is composed of two modules: the Planning Module and the Contrastive Module. The Planning Module is designed to generate the subsequent trajectories, and the Contrastive Module is designed to pull the states in the generated trajectories toward the high-return states and push them away from the low-return states during the training phase.
3.1 PLANNING MODULE
N-step Denoising Diffusson Probabilistic Model (Sohl- Dickstein et al. 2013 Song et al. Ho et al. 2020). Given a state 8, at RL step f, the Planning Module first generates H-length subsequent trajectory 4. by alternately denoising generated trajectories and estimating trajectory returns, and then extract the action to be executed from fp, as is illustrated in Figure _ Specifically, first sample A.N from A( I), and replace s in +. with 8, as condition on the current observation:
We designed the Planning Module
N + (s1,, a,") (6 (141), 14.1)..... (& H a(tH))
4)
in which all the elements except 8; are pure Gaussian noise. We further feed +/" into the backward process of diffusion to generate the subsequent trajectory:
po Po(+1-11 ( -Nou (+). + pV J. ()
(5)
po no(Fis i) Va a²- -
Va--13i 21.0
Here 4, = 3/ ( i) represents the T constructed from Ωं at diffusion step i, (+ .) is network for trajectory generation, i [1, .N] is the diffusion step, P represents the guidance scale, Ja(:,)) .) is learned function predict the return given any noisy trajectory 4/ and the corresponding diffusion step i. We abbreviate + to ft for convenience, f; = {($,, â;), ($1+1, d.+1)..... ($1+H. a++H) }- t is considered the subsequent trajectory following 8,. We take out the in t. the action corresponding to the state 8t.
3.2 CONTRASTIVE MODULE
Although the Planning Module can independently generate the action responding to the environ- ment, its performance is suppressed when the number of high-return trajectories is limited, as the Planning Module alone fails in making full use of the valuable insights implicated low-return trajectories. As solution, constraining the states in a subsequent trajectory toward the high-return states and away from the low-return states. In the following parts, first introduce the construction of contrastive sample sets (i.e., sampling the positive and negative samples for contrasting), and then we explain how we perform the contrastive mechanism.
propose contrastive mechanism improve the performance by
3.2.
SAMPLE POSITIVE AND NEGATIVE STATES
The positive samples and negative samples are necessary before applying contrastive mechanism. For an arbitrary state 8; S in the offline dataset, we compute its return v; in advance, as is stated in Then, we propose two strategies to sample positive sets and negative sets an arbitrary
Section 2.3
state St:
Sampling according to return (SR). For an arbitrary state 8, in the trajectory Ω generated by the Planning Module, apply the theory Thoma al. 2020) to compute the possibility of an arbitrary state 8; E S in the offline dataset sampled as the positive sample and negative sample of state s;
De, (Uz) 1 (E())
4

---

Published as a conference paper at ICLR 2025
Ps. (v:) =
(8)
+ c
where v; denotes the return 8i. Pe, (1:) and Pa. (vz) denotes the probability of s; being grouped into positive sample and negative sample of 8t, correspondingly. o, E and Ç are the hyper-parameters extended from Thoma et al. (2020).
Sampling according to return and dynamic consistency (SRD). Though sampling according to return is easy to deploy, such a rough strategy neglects the transition probability between adjacent states in trajectories. Hence, for the sake of dynamic consistency, we additionally conduct MiniBatch K-Means clustering (Sculley) 2010) over states in the offline dataset. Given the cluster C, to which s; belongs, we take the next state all the states in C. as the positive candidate of st, U,,"" Note that use the entire state space S the negative candidate set of 8t, Ur- Formally, 4, {s'( a. 8') D & Ct), U, S. Then, the positive and negative samples of 8t can be sampled from Ut by Equation 7 and sampled from U by Equation (8), respectively.
have
3.2.2 CONSTRAIN THE TRAJECTORY WITH CONTRASTIVE LEARNING
To constrain the states in subsequent trajectories while avoiding the cost of running the whole back- ward denoising process, we leverage the noised trajectory in the diffusion backward process to
reconstruct neat trajectory, i.e., // {( 2.0 a a,:"), (8 a(+1). ..... ($1 az=H)) from T/ for
1.0
any arbitrary diffusion step i. Then, we extract states in 4. as S S4.." = {8 22.0 $4421..,, 21.0 st For
each state 8% E Syt.o, we sample x states as positive sample set St and k states as negative sample set Sy from the offline dataset.
Inspired by |cchroff et al. (2015) and Sohn 2016), to apply contrastive learning to the scenario of multiple positive samples and impose aggressive constraints, we removed the positive sample term from the denominator polynomial in Equation (2) and propose the following equation to pull the states in the generated subsequent trajectory toward the high-return states and away from the low- return states:
Cj - log Ex_ exp(sim(f($)) f(s
Exo exp(sim(f(5)) f(s, ))/T)
where st E S s, ST. s() represents the projection function, T represents the tempera- ture (Wang& & Liu 2021), and sim(,,) denotes the cosine similarity, which is computed as
sim a, b) = a
(10)
||l|| |bb||
It is worth noting that Equation 9 is differs from the standard InfoNCE loss (Oord et al. 2018). In practice, we found that removing positive samples from the denominator, as suggested by (Schroff) et al. 2015) and (Sohn 2016). allows for better utilization of negative samples (i.e., low-return trajectories) and leads to better performance.
3.3 MODEL LEARNING
Recall that the action responding to state 8t is one of the elements in the generated trajectory, which influenced by the return predictor Ja (+:-) and constrained by contrastive learning. Therefore, we optimize our method from the perspective of trajectory generation, return prediction, and contrastive leamnigg constrain.
Specifically, we optimize the trajectory generation by minimizing the Mean Square Error between the ground truth and neat trajectory predicted by ( given any intermediate noisy trajectories as input:
La ET.,D. [[]T: whe ( []
(11)
where denotes the step of diffusion, t/ is obtained in the i-th step of forward process.
We optimize the return predictor by minimizing the Mean Square Error between the predicted return To (T), i) and the ground -tuth return ey:
L.
-
(12)
5

---

are Published as a conference paper at ICLR 2025
We constrain the trajectory generation with weighted contrastive loss:
T+ H
C. E2 (.50,
Σ
C.
(13)
in which the coefficient T+I decreases h increases since the importance gradually diminishes as it approaches the end of the planning horizon.
Hence, the overall objective function of ContraDiff can be written as a weighted sum of the afore- mentioned loss terms:
L= Lu XcCe:
(14)
where >e is hyper- parameer which balances the importance the contrastive module. Please note that the return predictor J and (-) are independent, thus optimizing Ja(:.-) and (b)(;.) with identical to separately optimizing J=(:.) with C, and vo((,-) with La and Le Please refer the proof in Appendix A.8 for details.
The pscudo code of ContraDiff presented in Appendix A. the details of available in Appendix A.11
EXPERIMENTS
In this section, evaluate the performance of ContraDiff across wide variety of tasks. We first construct 27 sub-optimal datasets to evaluate the ability of ContraDiff and regular RL methods in cases where high-return samples traDiff by comparing ContraDiff with resampling- based methods. Further, we delve into more com- prehensive experiments to analyze the key designs of ContraDiff. Note that results tables of this section are reported over 50 random seeds. The best and the second-best results of each setting are marked as bold and underline, respectively.
insufficient. Next, we demonstrate the advantages of Con-
4.1 EXPERIMENT SETTINGS
Environments and datasets. We evaluate ContraDiff on the locomotion tasks. Specifically, evaluate the locomotion capability of ContraDiff on Halfcheetah, Hopper, Walker2d. For each en- vironment, train ContraDiff with various scales of offline datasets provided by D4RL. (Fu al., 2020), and test the performance ContraDiff on the corresponding environments.
Baselines. We compare ContraDiff with the resampling-basee methods, which focus on address- ing the limited high-return trajectory issue. These methods are advantage-weighting (AW) (Hong al., 2023a),, return-weijge weighting (RW) (Hong et al., (023a), density-ratio weighting with advantage (AW-DW) (Hong et al., 2023b) and density-ratio weighting with uniform (U-DW) (Hong et al.) 2023b). We also compare ContraDiff with other regular methods which focus on addressing the general offline reinforcement learning issues. These methods include newly proposed SOTA meth- ods Decision Transformer (DT) (Chen al. 2021). CDE (Cen al. 2024). Trajectory Transformer
(TT) (Janner et al. 2021), HD (Chen et al. 2024). Decision Stacks (Zhao & Grover 2024). ReDif-
fuser (He et al. 2024), Diffuser (Janner et al. 2022), Decision Diffuser (DD) (Ajay et al. 2023).
Implementation details. We adopt U-Net (Ronneberger al. 2015) as the denoise network v and the return predictor -)(:,),, and adopt a linear layer with Sigmoid as the activation function the projector f()). Our model is trained device with 4 NVIDIA A40 GPUs, Intel Gold 5220 CPU and 504G memory, optimized by Adam (Kingma & Ba 2014) optimizer.
4.2 MAIN RESULTS As we discussed previously, the proportion high-return trajectories has significant impact on the performance of the model. To further demonstrate the impact of the proportion of high-return sam- ples on model performance and evaluate the performance ContraDiff under varying proportions of high-return samples, especially in scenarios with sparse high-return samples, we combine data within the dataset and create three types of datasets:
M-Exp: mix the the trajectories in Medium and Expert.
MR-Exp: mix the the trajectories in Med-Replay and Expert.
Rand-Exp mix the trajectories in Expert and the trajectories sampled by random policy interact- ing with environment.

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of average normalized score between ContraDiff and regular offline RL meth- ods on sub-optimal datasets, with + denoting the standard deviation. The ratio denotes the ratio of trajectories from the Expert dataset.
42. 2 0.9
524 524--32
114-2:13
01U:333
515 21
46. L1
11.52 1.3
4.222:00
12
M-Fap
0.
01:-111 11
15.0-+1.5
257 177
637-05
065-000
016-4-13
17:-4-09
033
64.0 1.9
75.7 1.2
350 1.4
209 12
01441
$12:0-00
03.6 1.5
K2N 6 1.2
68:14-11
05
-11113
03
410-11 1.0
19.34-32
57. 4-18
55.94 155
449 4.0
1.
N7.10M
62% 1LS
7.
11:.11
13.8 036
16.1 1.1
4KB 1.99
Rax
02
63.
= 1.6
5.5 1.1
E9-002
50.1 0007
103:242
16.3 1.
T4A 1.5
6.:-03
12.3 0.7
66.0-111 13
103
64.9
4.2
5.119.1
21.5
13592.11
1.8
9.99
51.8:33 322
91
14
21.0 4.3
45
05
L1
M44ip
142
6.5
- 15
96 2.5
34M 2.0
45.7 100
1.5
SALL 1.0.1
1:0
143
118.1
RIM
104-3
21
96.12.00
29
196.0 144
Happer
MR-Exp
02
72.3=22 29
01
L.
71
1:331:33
28.2 11
46.:11 1.5
76.7111 1.1
31.5
60.5-32
604-2-43
ML3 2.6
130
513+008
29:1
43.6
322
02
11:000
21:00
1.
46
RaxdFap
B7
66.5 9.2
13
20.00
63.0--04
11::4
430.4 1.
T 1
0.3
83.5 1.1
21.0
21
2:1
70.9
1.5
21:
10.5
18:00
15
83.3+112
2.1
0.3
61.54-1..
99.11 03
11.7 31.74-31
10233 1.5
97.04-32
2i
MR-Exp
0.2
$19 1.9
0..40007
615-222 20
755-00 0.4
045:2 1.9
2.8
41.3
15.2.111111
4.29
0-
$4.6 1.4
MLL 2-22
759-+-07
00
4.
912.5 L1
910 1.1
32.22225
17.1
16:3::22
7066
2044- 1.1
24
02:000
10:2001
0.2
10:5001
511011
146:2008
203 13
147+119
Mazdktap
142
14
4
56.1.0
1:32
35.2.1:1
51.94100
03
23
4..2
153
1.
853 32
394.544
4:9
60.1 4:31
4..3
We select the ratio of trajectories from the Expert dataset from {0.1, 0.2, 0.3), resulting in 27 mixed sub-optimal datasets in total. Returns of sub-optimal Halfcheetah datasets are illustrated in Fig- ure [ From the dataset level, for Medium and Reed-Repaay, mixing with Expert data effectively increases the proportion of high-retum samples. Additionally, since Medium and Med-Replay are both collected based on trained policy, whereas Random is collected based on a randomly initialized policy (Wu et al. )0 after mixing with Expert data, Rand-Exp exhibits overall lower propor- tion of high-return samples compared to M-Exp and MR-Exp. Furthermore, within the datasets, a higher proportion of mixed Expert data corresponds to a higher proportion of high-return samples. Using these datasets, we compared ContraDiff with resampling- based methods and regular methods.
4.2.1 COMPARISON WITH REGULAR METHOSS
Results of ContraDiff and baselines on the datasets are illustrated in Table in which ContraDiff- SR denotes the states used for contrasting are sampled with SR strategy and ContraDiff-SRR de- notes the states used for contrasting are sampled with SRD strategy. From Table we can observe that our method ContraDiff achieves the optimal and sub-optimal results on 25 out of 27 datasets, showing the advantage of ContraDiff in situations with sub-optimal datasets. More specifically, (1) In most cases, the performance of offline RL methods declines as the ratio decreases, where smaller ratio indicates fewer high-return trajectories. Moreover, compared with the performance in M-Exp, most baselines have a performance decline when trained with Rand-Exp, which has fewer high-return trajectories than Med-Expert. These results validate that the performance of offline RL methods is influenced by the proportion of high-return trajectories in the dataset; (2) ContraDiff demonstrates more significant improvement in Rand-Exp than M-Exp and MR-Exp, for instance, it outperforms the best baseline by 2.2 on Hopper-M-Esy M-E2x-0.2, but it outperforms the best baseline by 8.8 on Hopper-Rame Raad--and------ee Exp-0.2. As a brief recap, ContraDiff makes better use of low-quality samples than the baseline method, and demonstrates significant superiority in the case with limited number of high-return trajectories.
4.2.2 COMPARISON WITH RESAMPIING METHODS
Some previous studies proposed to address the issue of limited high-return trajectories have proven effective. To further evaluate the effectiveness of ContraDiff, we compare it with these methods in more challenging scenarios where the agent is trapped in low-return states. Our main objective here to investigate the ability of ContraDiff to handle such situations. For fairness in comparison, we use Diffuser, which is also the backbone of ContraDiff, as the backbone for all methods. Specifically, we combined Diffuser with AW (Hong et al. 2023a). RW (Hong et al. 2023a), AW-DW (Hong et al.) 2023b) and U-DW (Hong al.. 2023b) our baselines.
Comparison results are summarized in Table 2 As can be observed, all resampling method aids in enhancing the effectiveness of Diffuser. Even though, the performance of ContraDiff exceeds all of the resampling methods in most cases, demonstrating its effectiveness under aub-optimal conditions. Moreover, similar to Table [. the proportion of high-return samples increases, ContraDiff shows a more significant performance improvement compared to the baseline methods, indicating that utilizing low-return samples can effectively help the model escape from low-return areas.
7

---

Published as a conference paper at ICLR 2025
Table Comparison of average normalized score between ContraDiff and resampling-bas -aaged base- line methods on various sub-optimal datasets, with + denoting the standard deviation. Ratio denotes the the ratio of trajectories from the Expert dataset.
Ensironment
Datmet
Mix Ratu
Diffuser
AifDeeee
RW-Diftser
AW-DWW-Difuser
U-DV-Diffuer
ContraDil--SR
ContraDiff-SRD
0.1
53.2 2.1
638 0.44
66.6 1.7
2.99129
35.8 2.1
63.2 2.9
69.6
M-Esp
0.2
66.3 04
77.5 2.2
1:.
39.7 1.9
32.9 0.
80.3 2: 1,3
67.5 0.9
Halfeheetah
MR Exp
0.1
0.2
29,9
A1.1 1.7
47.6 1.3
25.9 1.6
13441 23
40.3 1.2
48.4
0.3
44.6 1 288
52.7 1
59.9 1.7
46.3 2:2
20.1
519 1.6
50.2 1. 0.9
0.1
23.8 = 1.6
43.4 0.0
52.2
2.9
38 0.9
19.5 1.7
35.8 1.
Rand-Esp
11.3.
0.1
79.1 =2 2.7
N3.9 2.2.1
83.5 2.2
48.9 1.3
39.2.::11
77.1 4: 1.5
91.3
M-Eap
0.2
92.9 1 0.7
8O.1 1.7
94.3 44.1033 0.2
37.8 1.1100
38.1
1039 3.7
98.5 2.8
0.3
90.2
103.9
1008
47.9::002
45.7
104.9
1:26
1.0
23.
Hupper
MR-Eap Eap
2:1
0,3
62.5 L.
60.3 3.3.1
1.
35.9 1.7
7.9
614 1.3
65.8
0.1
23.4+
12.3++11
41.2++288
106++33
0.7+006
45.1 +-27
14.++-1
Rand-Eap
0.2
45.2 + 2.1
+.34
0.3
24
19
47.0 2.2
50.2 1.1
0.1
22
7
74.
M-Fsp
0.2
35.5 =2 U.7
90.3 3.2
21.3 1 0.6
14.7
92.8 21
69.2
0.5
$8.7 L. 1.9
95.2 2.
87.6 1.0.
69.
10,0 133
99.3 1.5
91.0 2.2
Walkerdd
0.1
68.3+
59.5
66.7 +.5.5
2.7 1.1
2.9 0.6
74.5 + 2.7
65.1 +
13.9
0.1
4.3
11.61:33 1:33
2:41
4.9 13
1.0100 0.2
12.2
30.1
Rand-Fap
0.2
10.3 2.0
$2.9 1.9
20.6 +.9
7.+ 08
2,3++ 3.1
55.4 + 1.7
104
0,3
36.1 =
$3.
S011 0.7
35.5
2.1 2.2
61.6 0.9
45.3
(a) Diffuser
Terminal
Ab
60
(c) ContraDiff
Figure 3: How different models escape from low-return regions on
Additionally, we visualized how Diffuser, AW and ContraDiff escape from low-return regions on Ballerrd-Raan--Expaae As shown in Figure 3 given similar initial state which is the verge of falling backward, only ContraDiff successfully recovers from this posture and achieves long- term, healthy interactions (lasting over 60 steps), indicating that the contrastive mechanism helps ContraDiff in escaping from low-return states. Even though AW can interact with the environment for a longer duration compared Diffuser, ultimately fails to achieve the more prolonged inter- actions similar to ContraDiff. Recalling that AW only focuses the high-return trajectories via reweighing, the comparison between AW and ContraDiff more intuitively highlights the importance of information in low-return samples.
4.3
FURTHER INVESTIGATION
To further investigate the performance of ContraDiff, conduct ablation study, and analyze the state-reward distribution as well as the reward distribution.
4.3.1
ABLATION STUDIES
We have the following variants to conduct ablation study:
ContraDifT-N: only apply the samples with high-return to train the model.
ContraDift-C: remove contrastive mechanism from ContraDiff, i.e., remove C. from Equa- tion (14).
The results summarized in Figure 4 from which can conclude the following key findings:
(1) Applying only the high-return samples in training diminishes benefits in some cases. Com- pared with ContraDiff. ContraDiff-N is trained solely with high-return samples. Surprisingly, it achieves a lower performance than ContraDiff in all 9 tasks, indicating that low low-return trajectories offer beneficial information for model training.
8

---

Published as a conference paper at ICLR 2025
ContraDalA
CoatraDiFiC
CoctraDif:
Performance 92
120
45
90
100
120
100
100
80-
100
80
70
90
40
allcchetahh
Hopper
Walker2d
Halccheeahh
Hopper
Walker2d
Halscheetah
Hopper
Walker24 2d
Expert
(b) Med
(c) Med Replay
Figure Results the ablation experiments on different variants.
Rewards in learned model
Rewards in dataset
2
ContraDitt
ContraDiffC
Datasel
Probublitty Dematy
0.8
0.4
0.6
0.3
0.2
0.2
0.1
0.0
0.0
5
Regards
Reggors
(a) Dataset.
(b) ContraDifi-C.
(c) ContraDiff.
Figure 5: The distribution of state and reward. is bet- Figure 6: Distribution of rewards (a) Walker2d-Med-Ree Replay and (b) Brricheetah-Mee-RReee
ter to view in color mode. ContraDiff achieves higher rewards out-of-distribution areas (circled with red).
(2) The contrastive mechanism benefits the performance. ContraDiff-C, which eliminates the contrastive mechanism from ContraDiff, exhibits poorer performance across all nine tasks than Con- traDiff. This suggests that the contrastive mechanism indeed provides benefits.
4.3.2 STATE REWARD DISTRIBUTION ANALYSIS To illustrate the advantages of ContraDiff more intuitively, we randomly collect the (state, reward) pairs from the offline dataset of Walker2d-Med-RRplee and the (state, reward) pairs collected when ContraDiff' and ContraDiff-CC interact with the environment. The results are shown in Figure 5 in which each scatter represents a state mapped by UMAP (Meinnes al. 2018), and its color denotes the reward gained in the corresponding state. From the results illustrated in Figure 5 we can observe that: (1) there are more red and yellow dots in Figure than (b). That indicates that the model achieves better rewards by utilizing low-return samples more thoroughly with contrasting mechanisms; (2) In distrrddittiitttion states (circled with red), ContraDiff gains higher rewards than ContraDiff-C, which indicates that contrastive mechanism has potential in tackling out-of- distribution issue.
4.3.3 REWARD DISTRIBUTION ANALYSIS
The contrast mechanism in ContraDiff plays significant role in more efficiently utilizing low- return samples. believe that a deeper reason lies the impact the contrast mechanism learning of the diffusion model. According to Section 3.3 three crucial components of ContraDiff are the trajectory generation of the diffusion model, return prediction, and the contrast mechanism. remove the return prediction of ContraDiff and ContraDiff-C (i.e., remove To from Equation (5). leverage them to generate subsequent trajectory. Then we apply the actions in the generated trajectory to interact with the environment HallCheetah-MMcdRRellaa and visualize the distribution of reward during the interaction, as shown in Figure 6 It can be observed that the ContraDiff has a higher probability density high rewards in both cases. That indicates the effect of the contrast mechanism essentially increasing the proportion of high-return trajectories generated by the diffusion model. As demonstrated in Equation (5). with the support of both generated high-return trajectories and the return predictor, ContraDiff achieves sound performance.
To validate our assumption,
Meateerrd-Me--Repllla and
We have also conducted experiments to investigate the computational requirements of our method (Appendix A.9), impact of the size of contrastive set (Appendix [0. hyper-parameters analy- sis (Appendix I.II). ablation study about the weighted contrastive loss (Appendix A. and the analysis of averaged completed tasks (Appendix A.13) please refer to the Appendix for details.
RELATED WORKS
5.1 DIFFUSION FOR DECISION- -AKINN
With diffusion models demonstrate their powerful generalization capabilities across different fields, many works have introduced diffusion models into reinforcement learning for decision-making Some methods use diffusion models for generating actions directly (Ada et al. 2023 Wang et al.)
2022 Chen et al. 2022a Chi et al. 2023). ) For instance. Diffusion Q-learning Wang et al. (022) de-
signs the policy as a diffusion model and improves it with double Q-learning architecture. Following
this, SRDPs (Ada et al. 2023) adopts diffusion policies for out-o--distributiinn problems. Some other

---

Published as a conference paper at ICLR 2025
methods predict the future trajectories with diffusion models for decision decision-making. Diffuser (Janner et al. 2022 models trajectories as sequences of state-action pairs. Based on this, Decision Dif- tuser (Ajay et al. 2023 proposes to predict state sequences with diffusion model conditioned on historical information, and adopts a reverse dynamic model to predict actions. Though these methods have gain significant achievements, they neglect the differences between high-return and low-return samples, and fails in making full use of low-return samples.
5.2 OFFLINE REINFORCEMENT LEARNING WITH IMBALANCED DATASETS
In many cases, offline RL datasets tend to have a limited number of high-return samples, while low-return samples are much more abundant, which is generally considered the data imbalance issue (Choi et al., 2024 Yan et al.) 2024). Many works have been proposed to tackle this problem, often by sample reweighing. Some works reweight samples at the trajectory level. For instance, Hong et al. (2023a) suggests that sampling can be reweighted to prioritize high-return trajectories, and proposes return-went weighted (RW) and sarrantgg--eiihhee (AW) as solutions. Based on this, Hong et al. (2023b) proposes Density-ratio Weighting (DW), which designs learnable weight for reinforcement learning objectives. Some other methods reweight samples at a more fine-grained level, i.c., at action or state level. LAPO (Chen al. 2022b) maps actions into latent space. which is further optimized through advantage-
advantage--eighhing weighting; and more recently, A2PR (Liu et al. 2024) ob- tains high-advantage actions from an augmented behavior policy. Differently, SAW (Lyu et al., 2022) reweights the future expected states and uses them as the goal for the policy. Although they focus on high-return samples through reweighing, they overlook the valuable information tained in low-return samples. In contrast, our method ContraDiff, leveraging the contrastive learning mechanism, fully exploits the information embedded in low-return samples.
5.3 CONTRASTIVE LEARNING IN REINFORCEMENT LEARNING
Most methods adopt contrastive learning (CL) to learn better representations, such as state represen-
tations (Laskin et al. 2020 Qiu et al., 2022) and task representations (Yuan Lu, 2022. Agarwal et al., (020) For instance, Laskin et al. (2020) proposes to learn image representations via CL., Qiu et al. (022) propose to learn the transition with CL, Yuan & Lu (2022) and Agarwal et al. (2000) applies CL to enhance the task representations to distinguish between different tasks. Some works apply CL in other ways. For instance, Laskin et al. (2022) utilize CL to learn behavior rep- resentations to encourage behavioral diversity, while QGPO (Lu al. (023) adopts CL for better score function of Diffusion-based RL. methods. In contrast, ContraDiff adopts CL for better usage of low-return samples.
6 CONCLUSION AND DISCUSSION
In this paper, we introduce ContraDiff for offline RI, which combines contrastive mechanism to make better uses low-return trajectories and addresses the challenge of limited high-return samples. Different from the previous works which apply contrastive learning to enhance the representation, we perform contrastive learning over the return of states. Specifically, we apply diffusion to gener- ate the subsequent trajectory for planning, and constrain the states in the generated trajectory toward the states with high returns and away from the states with low returns to improve the base distri- bution. In that way, the actions taken by the agent are always toward the high-return states, which makes the agent gain better performance in the online evaluation. We evaluate ContraDiff on 27 sub-optimal datasets, where the results demonstrate that our ContraDiff is able make better uses low-return samples and achieves outstanding performance. The ablation studies and investigations further substantiated the rationality of ContraDiff.
Although ContraDiff has achieved outstanding performance, there are still many areas for improve- ment and scalability. For instance, our current contrastive learning is applied to states, thus the ContraDiff framework can only be extended to trajectory generation methods for now, such as ac- tion or state-action pairs, or even at the latent of trajectories. We leave applying the contrastive mechanism to other other methods in the future work.
7 ACKNOWLEDGMENTS
This work supported by the National Natural Science Foundation of China (Grants No. 62307020 and 62322603), the National Key R&D Program of China (2022ZD01 14804), and the Shanghai Municipal Science and Technology Major Project (2011SHZD7XX0002),
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Suzan Ece Ada, Erhan Oztop, and Emre Ugur. Diffusion policies for out-of-distribution generaliza- tion offline reinforcement learning. arXiv preprint SeXii=2307.047266, 2023.
Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference on Learning Representations, 2020.
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference Learning Representations, 2023. https:/:::/ //opereeeeee.....e
htt./forumiidssse Fo2K9DFG
Zhepeng Cen, Zuxin Liu, Zitong Wang, Yihang Yao, Henry Lam, and Ding Zhao. Learning from sparse offine datasets via conservative density estimation. In The Twelfth International Conference on Learning Representations, 2024.
Chang Chen, Fei Deng, Kenji Kawaguchi, Caglar Gulcehre, and Sungjin Ahn. Simple hierarchical planning with diffusion. In The Twelfth International Conference on Learning Representations, 2024.
Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. arXiv preprint arXiv:2209. 145488 2022a.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Picter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances neural information processing systems, 34:1 5004--1979, 2021.
Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang. Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. Latent-variable advantage-weightee policy optimization for offline rl. arXiv preprint arXiv:2203.08949) 2022b.
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,, 20233
Dami Choi, Derrick Xin, Hamid Dadkhahi, Justin Gilmer, Ankush Garg, Orhan Firat, Chih-Kuan Yeh, Andrew M Dai, and Behrooz Ghorbani. Order matters in the presence of dataset imbalance for multilingual learning. Advances in Neural Information Processing Systems, 36, 2024.
Mehdi Fatemi, Mary Wu, Jeremy Petch, Walter Nelson, Stuart Conmolly, Alexander Benz, Anthony Carnicelli, and Marzyeh Ghassemi. Semi-markov offline reinforcement learning for healthcare. In Conference on Health, Inference, and Learning, pp. 119-137. PMLR, 2022.
William Feller. On the theory of stochastic processes, with particular reference applications. 1949.
1211077442
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Nantian He, Shaohui Li, Zhi Li, LIU Yu, and You He. Rediffuser: Reliable decision-making using diffuser with confidence estimation. In Forty-first International Conference on Machine Learning, 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6857 2020.
Zhang Wei Hong, Pulkit Agrawal, Remi Tachet des Combes, and Romain Laroche. Harness- ing mixed offline reinforcement learning datasets via trajectory weighting. In The Eleventh International Conference Learning Representations, 2023a.

---

Published as a conference paper at ICLR 2025
Zhang- Wei Hong, Aviral Kumar, Sathwik Karnik, Abhishek Bhandwaldar, Akash Srivastava, Joni Pajarinen, Romain Laroche, Abhishek Gupta, and Pulkit Agrawal. Beyond uniform sampling: Offline reinforeement learning with imbalanced datasets. Advances in Neural Information Processing Systems, 36:4985-5009, 2023b.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273-1286, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference Machine Learning, pp. 9902-9915. PMLR, 2022.
Natasha Jaques, Judy Hanwen Shen, Asma Ghandeharioun, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard. Human-centric dialog training via offline reinforcement learning. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 3985-4003, 2020.
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Advances in neural information processing systems, 33:1 (66616666187, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1iXiv:1412.6988 2014.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. International Conference on Machine Learning, pp. 5639- 5650. PMLR, 2020.
Michael Laskin, Hao Liu, Xue Bin Peng, Denis Yarats, Aravind Rajeswaran, and Pieter Abbeel. Unsupervised reinforcement learning with contrastive intrinsic control. Advances in Neural Information Processing Systems, 35:34478-3444 2022.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems, arXiv preprint arXiv:2005.01643, 2020.
Tenglong Liu, Yang Li, Yixing Lan, Hao Gao, Wei Pan, and Xin Xu. Adaptive advantage-guiee guided policy regularization for offline reinforcement learning, arXiv preprint wxXv:2245,199909 2024.
Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact gugggy diffusion sampling in offline reinforcement learning. arXiv preprint arXiv:2304.12824 2023.
Jiafei Lyu, Aicheng Gong, Le Wan, Zongging Lu, and Xiu Li. State advantage weighting for offline rl. arXiv preprint arXiv:2210.04251, 2022.
Leland McInnes, John Healy. Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold
approximation and projection. Journal of Open Source Software, 3(29), 2018.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arviv preprint arXiv: 0070007487... 2018.
Rafael Figueiredo Prudencio, Marcos ROA Maximo, and Esther Luna Colombini. A survey on offline reinforcement learning: Taxonomy, review, and open problems. IEBE Transactions on Neural Networks and Learning Systems, 2023.
Shuang Qiu, Lingxiao Wang. Chenjia Bai, Zhuoran Yang, and Zhaoran Wang. Contrastive ucb: Provably efficient contrastive self-supervised learning in online reinforcement learning. In International Conference on Machine Learning, pp. 18168-18210. PMLR. 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assiste Intervention-MIee 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234 241. Springer, 2015.
12

---

Published as a conference paper at ICLR 2025
Florian Schroff, Dmitry Kalenichenko, and James Philbin. Facenet: A unified embedding for face recognition and clustering, In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 815 823, 2015.
David Sculley. Web-scale k-means clustering. In Proceedings of the 19th international conference on World wide web, pp. 1177-1178, 2010.
Tianyu Shi, Dong Chen, Kaian Chen, and Zhaojjan Li. Offline reinforcement learning for au- tonomous driving with safety and exploration enhancement. arXiv preprint arXiv:21 10.07067, 2021.
Jascha Sohl- Dick-Diiittin Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2256-2265. PMLR, 2015.
Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. Advances in neural information processing systems, 29, 2016.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Janine Thoma, Danda Pani Paudel, and Lue V Gool. Soft contrastive learning for visual localization. Advances in Neural Information Processing Systems, 33:11119-11130, 2020.
Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning? Advances in neural information processing systems, 33:6827-6839) 2020.
Feng Wang and Huaping Liu. Understanding the behaviour of contrastive loss. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pp. 2495-2504, 2021.
Xiao Wang and Guo-Jun Qi. Contrastive learning with stronger augmentations. IEEE transactions on pattern analysis and machine intelligence, 45(5):5549-5555 2022.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou. Diffusion policies as expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2022.
Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua M Susskind, Jian Zhang, Ruslan Salakhutdi- nov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. In International Conference Machine Learning, pp. 11319-11333 1328. PMLR, 2021.
Teng Xiao and Donglin Wang. A general offline reinforcement learning framework for interactive recommendation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 4512-4520, 2021.
Tete Xiao, Xiaolong Wang, Alexei A Efros, and Trevor Darrell. What should not be contrastive in contrastive learning. In International Conference on Learning Representations, 2020.
Divin Yan, Gengehen Wei, Chen Yang, Shengzhong Zhang, et al. Rethinking semi-supervised im- balanced node classification from bias-variance decomposition. Advances in Neural Information Processing Systems, 36, 2024.
Chun-Hsiao Yeh, Cheng- Yao Hong, Yen-Chi Hsu, Tyng-Luh Liu, Yubei Chen, and Yann LeCun. Decoupled contrastive learning. In European Conference Computer Vision, pp. 668-684. Springer, 2022.
Cao
Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforceemnt learning via contrastive learning. In International Conference on Machine Learning, pp. 25747-25759. PMLR, 2022.
Siyan Zhao and Aditya Grover. Decision stacks: Flexible reinforcement learning via modular gen- erative models. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 PSEUDO CODE OF COFFRADIFF.
Algorithm T'aaining
1: while not converged do
2: Tv. Dj ~ D.
3: i~ [].
4: Generate T/.
5: Reconstruct Ti as 7/ va (T), i).
6: Calculate loss La with Equation
7: Calculate loss L, with Equation 12
8:
Extract states in 1.0 as S S4y0 t-1
0.0
8
9:
for P/- in Sqt.0 do
10:
Sample S and S according to Section 3.2.1 Calculate Lj, using Equation ().
11:
12:
end for
13:
Calculate L. using Equation 1
14:
Calculate & using Equation (14)
Update model by taking gradient decent with
16:
end while
Algorithm 2 Planning
Require: ContraDiff ())(,)), returntto-o predictor Tol:) :), guidance scale p. co co-variances Σ'
1: / 1.1.
2: while not done do
3: Observe state s;: sample T.
I)
4: for N, N - 1,..., do
5: Predict return-to-go with (+t, 1).
6: Sample - using Equation
7:
end for
8: Extract à form 70.
9:
Interact with environment using action å,-
10: I <- 1+1.
11: end while
A.2 ILLUSION OF EQUATION AND EQUATION
As can be observed in Figure our modified influence functions are designed to leave a blank the middle area deliberately, which different from (Thoma et al. 2020). The underlying reason is that not all states supposed be contrastive samples. Nevertheless, our modified influence functions collapse into modified influence functions in Thoma et al. 2020) if we set = Ç.
Example of Equation 7 and Equation (8) visualized in Figure 3
A.3 RESULTS ON REGULAR DATASETS
We have also compared ContraDiff with baseline methods regular datasets. The results of Con- traDiff and baseline methods are summarized in Table and Table
From Table F and Table 4 we can observe that: (1) Compared with all the baseline methods, Con- traDiff achieves the best or the second-best performance on out locomotion tasks (HalfChee- tah, Hopper, and Walker2d) and achieves the best or the second-best performance all the two high-dimensional manipulation tasks (Kitchen), demonstrating the outstanding performance of Con- traDiff under periodic settings. Moreover, ContraDiff achieves the best performance on 2 out of navigation tasks, demonstrating the excellent ability ContraDiff in long- term planning. (2) Com- pared to the methods with similar backbones, Diffuser outperforms Diffuser in all 14 tasks, and
14

---

Published as a conference paper at ICLR 2025
Accumulated Reward
10000
5000
0
Medium
M-Exp (Ratio=0.1) 0.1)
M-Exp (Ratio=0.2)
M-Exp (Ratio = 0.3)
Accumulaten lated Reward Accumulated Reward
1.0000
5000
Med-Replay
MR-Exp (Ratio=0.1)
MR-Exp (Ratio=0.2)
MR-Exp (Ratio=0.3)
1 10000
5000
0
50 Random
100
50
100
0
50
100
0
50
100
Rand-Exp (Ratio=0.1) Rand-Exp (Ratio=0.2) Rand-Esx (Ratio=0.3)
Figure 7: Returns distribution of Medium, Med-Replay and Random datasets of Halfcheetah mixed with different ratios of Expert data.
1.0
(+)),, o=25, 0= (=-0.5
0.8
Possibility
(+),,, o=25, 5=0.5
0.6
g'))), 0=12, {=-0.5
0.4
g + g+)), o=12, 5=0.5
0.2
g (+), o=25, (=-0.3
(+), 0=25, 5=3.3
0.0
-1.0
-0.5
0.0
0.5
1.0
Return
Figure 8: Illusion of Equation and Equation
g
15

---

Published as a conference paper at ICLR 2025
Table 3: Comparison of average normalized score Mujoco environments between ContraDiff and basclines, with + denoting the standard deviation.
1854
01.
171.
181 1113
137
1300
410
130
199
1004 4.5
1870033
1833
1111111
1124:13 1.3
003
043
406
011
16m
5
17
653
16
0131:11
763314
185:11
503:34
16::03
Sapar
614
983
123
16.
605
24
100.00
84
005-00
Table 4: Comparison of average normalized score on Maze2d and Kitchen between ContraDiff and baselines, with + denoting the standard deviation.
Datasel
Envirommett
COL
IQL
DT
TT
MOPO
Decision Stack
Difluser
DD
ContraDiT-SR
ContraDif-SSDD
U-Maze
MazeDd
5.7
47.4
9.2
25.4
13.6
111.3 + 12.2
113.9 3.1
0.0
142,9 2.2
122.1 #: 1.4
Medium
Mazedd
5.0
34,9
9.6
23.3
33.3
111.7 4
121.5 4: 2.7
0.0
140 1: 0.7
115.7 1: 0.7
Large
MaseDd
12.5
58.6
10.4
27.7
0.0
171.6 +: 134
123.0 6,4
0.0
131.5 + 3.2
129,9 1.2
Mixed
Kitchen
52.4
51.0
20.9
31.1
0.0
42.5
65.0 2.8
65.0 :: 1,3
32,5 :: 2.2
Partial
Kitchen
51.2
46.3
35.2
32.9
0.0
40.0 3.1
57.0 0 2.5
58.0 + 1.9
40.0 + 3.1
outperforms DD in 11 tasks, which demonstrates the benefits of introducing the constrastive mech- anism. (3) We can observe that ContraDiff exhibits more improvement in Medium and Med-Replay datasets than the Med-Expert datasets. Since Med-Expert datasets have more high-return samples, they offer abundant information for methods like Diffuser learn, thus they can achieve better re- sults. However, both Med and Med-Replay have more low-return samples than Med-Expeet, which increases the difficulty of learning a good policy. (4) Compared to the sub-optimal datasets in Ta- ble []] the improvement of ContraDiff-SRII in the regular scenarios is not as pronounced as that of ContraDiff-SR. Even though, ContraDiff-SRI outperforms Diffuser We believe this is because, in the regular dataset, it is easier for the model to learn the transitions between different states, making the guarantee of dynamic consistency (i.e., ContraDiff-SRD) rela- tively less important in this context. In other words, we recommend using ContraDiff-SR in regular scenarios, and use ContraDiff- -RD in challenging scenarios. For more detailed discussions, please refer to Appendix A.4
out of 14 situations.
A.4 WHICH STRATEGY TO SELECT FOR THE CONSTRUCTION OF CONTRASTIVE SAMPLE?
We provide two implementations to construct contrastiev samples, namely ContraDifT-SR and ContraDifT-SRD, As described in Section 3.2.1 ContraDifI-SR pulls the generated trajectories to- wards the globally high-return states, while Contraiff-SRD pulls the generated trajectories towards the transition-able high-return states. This may result in a slight decrease in comparative effective- ness, but it ensures dynamic consistency among contrastive samples.
Although the results indicate that each method has its advantages, there are also patterns to follow when making a choice. We first visualize the samples and then determine which implementation to choose based on the visualization results. For example, let's consider HalfcheetaahRRaddEEl 0.1 and Walker2d-Rand-Exp- crosses represent the expert data. From this visualization, we can obtain some prior information about dynamic consistency.
The states with their returns are visualized in Figure where
As shown in F most of the high-return states are far away from the low-return states, From this observation, we can infer that starting from any one of the original states from the Halfcheetah- Random dataset, it is difficult to transition to the expert data, i.e., the high-return states marked with a cross. Under this situation, pulling the generated trajectories with ContraDiff-SR will introduce uncertainty, the corresponding state transitions are unachievable. Therefore, we adopt ContraDiff- SRD for Halfcheetah- Raad-annd----ae
In contrast, as shown in Figure (b), the mixed expert data of Walker26-Rand-Expp shares sim- ilar distribution pattern as the original states from the Walker2d- Random dataset. this situation, adopt ContraDiff- SR for more direct constrain.
A.5 COMPATIBILITY STUDY
16

---

Published as a conference paper at ICLR 2025
1eb
500000
15-
1.0
400000
300000
0.6
Value
0.2
100000
-5
100
Figure 9: States distribution of (a) Haaficheetah-Rand-EEpp-- and (b) Walker2d-Rand-EXppee with
cross representing the expert data.
1.00
States Similarity
Trajectories
Trajectories
Trajectories
0.75
0.50
0.25
0.00
0.25
11 21
31
11
21
31
11
21
31
States in trajectory (a) Diffuser.
States trajectory
States in trajectory (c) ContraDiff.
(b) Decision Diffuser.
Figure 10: The similarities between the states in the generated trajectories and actual states. The generated states of ContraDiff are more similar with the actual states, demonstrating the better long- term dynamic consistency.
As we discussed Section 3.1 our ContraD- iff is build based Diffuser. To validate the compatibility of the contrastive mechanism of ContraDiff, transplant it Decision Dif- fuser (DD) (Ajay et al.) (203)), and evaluate and compare the improvement on three envi- ronments. The improvements are summarized in Table in which DD+ and Diffuser" denote the improvement of introducing contrast mech- anism in DD and Diffuser correspondingly. As we can observe, DD+ achieves noticeable im- provement in 2 out of tasks and Diffuser+ gains improvement in out of 3 tasks, which demonstrates the portability of the contrast mecha- nism of CoftraDiff. Interestingly, DD+ is unable to achieve any improvement in Halfcheetah-Mes Expert. This could be attributed to the separated training (sampling) of states and actions in DD, which results in a failure to effectively model their joint distribution.
Table 5: The improvements of the normalized score after transplanting the contrastive mecha- nism of ContraDiff to Decision Diffuser (DD+) and Diffuser (Diffuser*) with + denoting the variance.
Dataset
Environment DD
Diffusert
Med-Expert Med-Expert
Halfcheetah
3.1+-0.4
Hopper
5.44
9.1+1.2
Med-Expert Walker2d
6.5+10.9
1.340.4
A.6 CAN-EXXCUTION CONSISTENCY ANALYSIS.
We use the plan-execution consistency to denote the similarity between the states in the planned trajectory and the states encountered by the agent during its interaction with the environment. reflects the models' capability in modeling the environment. To investigate the plan-execution con- sistency of ContraDiff, we randomly take 24 trajectories generated by Diffuser, Decision Diffuser, and ContraDiff. For each generated trajectory, pute the similarity between each generated state and the actual state of the same step provided by the environment. Thus, there 24 32 similarity returns for each model, which corresponds to similarity matrix as the subgraphs Figure [0] illustrated. Each line in the subgraphs of Figure I0] represents generated trajectory, and the grids of each line represent the similarity of the states in
take the states of consecutive 32 steps and com-
17

---

observe Published as a conference paper at ICLR 2025
the generated trajectory and the states provided by the environment. From Figure 10.
that: (I) Most grids in Figure I0 (c) are blue, which denotes that most generated states are consistent with the actual states; (2) Figure I0 (c) contains more blue grids than Figure To (a) and (b), which denotes that ContraDiff has better plan-execution consistency than Diffuser and Decision Diffuser. Since the difference between ContraDiff and Diffuser is the contrastive module, combining Figure aland Figure I0] we can conclude that the contrative module benefits the plan-execution consistency of CoftraDiff and makes ContraDiff gain high rewards in both in-distribution and out-of-distribution situations.
A.7 VISUALIZATION OF POSITIVE AND NEGATIVE SAMPLES.
We randomly sample a subset of positive samples (states with high returns) and negative samples (states with low returns), is shown in Figure [1 can be observed that an agent in state corresponding to a high return tends to be in a position more conducive to walking or running, such as standing upright; correspondingly, an agent with state corresponding a position that is hard to walk, such as having already fallen down or about to fall down. This is reasonable, since poses such as standing upright are more conducive to walking or running, which causes the agent to continue moving and results in a higher return, while poses such as having fallen or about to fall cause the environment to give a stop signal, which results in a lower return.
low return will be in
A.8 OPTIMIZING J WITH EQUATION 14
Suppose we have the diffuison model (-)()) parameterized by 0, and the return predictor To param- eterized by Φ. Following Equation (14) we have
C La + L. Xche.
(15)
Further,
C. E.ED, [|[T vo(T)
(16)
C. = -.
21.
(17)
The training process can be viewed as procedure of calculating gradients of all the parameters and updating them, specifically,
VØ 0C
(18)
00
aLa
dC.
dC.
+ de
(19)
00
00
00
OLa dC.
(20)
00
80
OC V
(21)
dó UC.
0L,
QC.
+
No
(22)
Oo
do
Oc..
(23)
0
Thus, calculating the gradients of 0 with L is equal to caleulate 9 with La and Les calculating the gradients of o with is equal to calculate o with C., i.e., optimizing the return predictor Tx(:)) with Equation (14) is equal optimizing it with Equation (12 only.
18

---

Published as a conference paper at ICLR 2025
X
K
A
x
4
7
A
A
1
A
A
f
x
6
x
A
3
2
x
x
x
A
x
x
2
2
x
x
b
x
x
4
b
x
x
x
x
x
x
c
x
x
x
x
6
A
x
3
4
4
8
x
6
x
x
x
4
x
x
x
x
2
x
x
8
x
8
x
4
x
x
x
8
8
8
x
x
x
A
x
A
2
x
4
x
x
b.
4
8
x
x
x
8
x
x
4
x
x
4
x
x
x
x
x
x
x
x
x
&
x
x
x
f
x
x
x
x
x
x
x
x
h
x
b
A
A
x
x
x
x
b
h
A
A
x
x
A
b
x
(a) Agents with high-return states.
3
a
8
x
x
x
8
x
h
x
x
6
2
x
x
x
x
3
8
c
x
x
x
c
2
x
2
x
x
8
b
4
2
6
x
x
d
x
x
2
8
w
x
2
x
x
x
x
c
6
v
(b) Agents with low-return states.
Figure 11: Visualization of positive samples (states with high returns) and negative samples (states with low returns) in Maaeer2dd-Me-Repppaaa
Table Comparison of computational requirements.
Decision Diffuser
Diffuser
ContraDifT-SR
ContraDifi-SRD
CPU (9h)
107.6
106.4
107.7
112.4
Memory (GB)
6.9
7.0
13.9
Training time (Seconds every iterations)
98
112
122
A.9 COMPUTATIONAL REQUIREMENTS
We have quantified the impact of sample selection and the incorporation of contrastive learning on computational requirements in our method, shown in Table 6
can observed that even with the introduction of sampling and contrastive learning, ContraDiff- SR requires a comparable amount of resources and training time to other methods with similar backbone. For ContraDiff-SRR sampling based on dynamic consistency requires more resources. This is partly due to the dynamic consistency considerations, resulting in the requirement for more
19

---

Published as a conference paper at ICLR 2025
00.0
80
80
77.5
75.0
75
72.5
70.0
50
0.70
0.75
0.50
0.85
0.90
0.1
0.2
0.3
0.4
400
600
800
1000
1200
0.00
0.25
0.50
0.75
1.00
(a) Positive bound
Negative bound Ç.
Puzzy coefficient
Contrastive weight Ac.
Figure 12: The impact of whper-parameteer
CPU cores (5 more than other methods). Additionally, the introduction of cache matrices consumes more memory (7GB, but actually we can cache them CPU and conduct the sampling CPU if is necessary).
Despite that, the resources and training time are still manageable. Moreover, additional resources are only required during the training phase. During the inference phase, our resource requirements are consistent with other methods (e.g., Diffuser).
A.10 IMPACT OF THE SIZE OF CONTRASTIVE SET
Table 7: Impact of the size of contrastive set.
Mix Ratio
8
16
32
64
128
0.1
35.4 2.2
41.5 + 2.6
48.0 2.9
46.9 1.1
48.9 0.9
0.2
39.1
+ 2.9
69.6 2.1
77.6 0.9
73.9 + 0.9
78.5 0.6
0.3
73.9
1.7
84.8
88.7 0.9
90.4 1.7
88.3 1.3
Generally, we set the size of contrastive set to 32. We illustrated the impact the size of contrastive set on our method using Halfheetah-Rand-Ean in Table 7 We can observe, the performance of ContraDiff' gradually increases when increasing the size within 32. However, as the size of the contrastive samples increases further, the performance gradually stabilizes. Therefore, the size of the contrastive set is easy to determine.
A.11 HYPER-PARAMETERS
ContraDiff imports the following additional hyper-parameters positive bound (E), negative bound (9), fuzzy coefficient (σ), loss weight of contrastive learning (Ae). For locomotion tasks,
lect & E {0.55, 0.65, 0.85). select Ç E {0.05, 0.1, 0.2, 0.35, 0.4), select A E {0.1, 0.01, 0.001 select & E (0.1, 0.6, 5), select
is select between [4, 16 102. For Maze2d tasks, (0.01, 0.02, 0.2). We set as x 10ah and set A as 0.1 for Maze2d tasks. We use the other hyper-parameters as in the paper of Diffuser (Janner al. 2022).
Note that even there are 4 new hyper-parameters, searching for the optimal hyper- parameters is easy. On the one hand, these hyper-parameters are not interdependent, we can search for the optimal value of each hyper-parameter separately. On the other hand, analysis of the hyper- parammerrsestess indicates that their impact on model performance moderate, as discussed below.
We have visualized the impact of hyper-yppae parameters in Figure 12 As is summarized in Figure 12 (a)
The choice of hyper-parameters has a minimal impact on performance, with variations remaining within 7%. (b) The performance changes smoothly around the peak region as illustrated in Figure indicating tuning the
towards the increasing performance will obtain optimal per- formance. Therefore, tuning hyper-paramess of CDifffuser is relatively easy, costing little time.
(c) The performance of the model does indeed change with variations of pyper-parameters, as the selection of positive and negative samples and the weight of the contrastive loss influence the perfor- mance. However, as we discussed in (a), the hyper- hyper-parameters parameters are relatively easy to tune. Therefore, in most cases, hyper-parameters do not become bottleneck limiting model performance. Still, fur-
20

---

Published as a conference paper at ICLR 2025
kitchen-mixed
kitchen-partial
3.5
3.5
ContraDiff-C ContraDiff
ConkraDiff-C ContraDiff
Averaged Completed Tasks
3.0
3.0
2.5
2.5
2.0
1.5
1.5
1.0
1.0
0.5
0.5
0.0
0.0
50
100
150
200
250
50
100
150
200
250
Planning Steps
Planning Steps
Figure 13: The averaged completed tasks of ContraDiff-C and ContraDiff.
ther optimizing the design of the hyper- pareeeeaameterss is of the tasks will focus on in our future research.
A.12 ABLATION STUDY ABOUT EQUATION
13
The weighted contrastive loss Equation 13 is designed to emphasize the initial part of the gener- ation, the initial action in generated trajectories interact with the environment. Ad- ditionally, we have conducted ablation studies on the weighted contrastive loss, where ContraDiff- mean represents ContraDiff with werage--eiighted contrastive loss. As is shown in Table g the weighted contrastive loss brings better performance in all of the situations.
Table Performance comparison of ContraDiff-meet mean and ContraDiff. ContraDiff-nnd mean represents ContraDiff with weegage-weighted contrastive loss.
Mix Ratio
ContraDiff-mean
ContraDiff
0.1
34.2 + 1.2
48.0 2.9
0.2
71.1 1.9
72.3 0.7
0.3
82.9 + 0.3
88.7 0.9
A.13 ANALYSIS OF AVERAGED COMPLETED TASKS
We additionally compared the number of completed tasks of ContraDiff and ContraDiff-C C on Kitchen-Mixed and Mitchen-Partial to further demonstrate the advantages of ContraDiff. As is vi- sualized in Figure [3] as the planning steps increase, ContraDiff shows more completed tasks than ContraDiff-C. C, which intuitively demonstrates the advantages of the contrastive mechanism in sparse reward environments.
A. ANALYSIS OF RL LOSS AND CONTRASTIVE LEARNING LOSS
Recall the the diffusion loss (RL loss, Equation the contrastive loss Equation 13 in Section 3 Note that the diffusion loss (RL loss, Equation aims modeling the dataset distribution. How- ever, the closer the distance to high-return states and the farther the distance low-return states, the smaller the value of Equation [3
21

---

Published as a conference paper at ICLR 2025
ContraDiff
ContraDiff w/o CL
Diffusion Loss (RL Loss)
10-
5.5
RL Loss Eq.(13)
-60
Eq.(13)
-6.5
-6.0
0.6
02
Training Iteration1. 1eB (a)
Training Iteration le6
Training Ieeration tek
(b)
(c)
Figure 14: The (a) Diffusion loss (i.c., RL loss, Equation (b) value of Equation 13 and (c)
RL loss + value of Equation [3) on Walker2d-Ran E.p-
As is shown in Figure 14 a). regardless of the use of contrastive learning, the RL loss shows
downward trend and eventually stabilizes. However, for the value of Equation (13) shown in Fig- ure 14(b), when contrastive learning is not used, the contrastive loss remains large and fluctuates significantly. On the other hand, when contrastive learning is applied, the contrastive loss decreases. Moreover, ContraDiff shows lower value RL Loss Equation 13) than ContrasDiff w/o CL in Figure [4.c). This demonstrates that (1) the RL loss and contrastive loss are independent, with the contrastive loss serving meaningful regularization term; and (2) contrastive learning indeed maximizes the distance from low-return trajectories while minimizing the distance to high-return trajectories, further corroborated by the analyses in Section 4.3
22


Reviews:
[{"review_id": "fjzVQ7SjMb", "reviewer": "Reviewer_2zbW", "summary": "This paper introduces ContraDiff, a novel offline reinforcement learning method that leverages contrastive learning to make better use of low-return trajectories by pulling states towards high-return regions and pushing them away from low-return areas. While the approach shows promising results across various environments and demonstrates advantages in handling sub-optimal datasets, several limitations remain, including unclear theoretical justification for the relationship between contrastive and RL losses, heavy reliance on hyperparameters for positive/negative sample selection, and lack of trajectory-level temporal information in the contrastive learning design. Despite these limitations, the paper presents a well-structured contribution with comprehensive experiments, providing a new perspective on utilizing sub-optimal data in offline RL.", "questions": "See weakness", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-structured and easy to follow, with clear motivation and comprehensive explanations of each component in the proposed method.\n\n2. The empirical evaluation is thorough, covering 27 sub-optimal datasets and including extensive ablation studies, hyper-parameter analysis, and visualizations that help understand the method's behavior.\n\n3. The proposed method provides a novel perspective on utilizing low-return trajectories in offline RL, and the implementation is relatively simple with only a few additional hyperparameters compared to the base diffusion model.", "weaknesses": "1. The relationship between contrastive loss and RL loss needs further detailed justification. Intuitively, these two losses appear highly correlated, especially when considering low-return trajectories as negative samples. A numerical analysis comparing these losses across trajectories with different returns would be valuable to verify if they indeed show similar trends. If so, the contrastive loss might merely serve as an enhancement to the RL loss rather than a meaningful regularization term that needs to be balanced against the primary objective.\n2. The methodological differences illustrated in Figure 1(d)(e)(f) lack sufficient justification.This figure fails to demonstrate why pushing away from low-return states is fundamentally superior to upweighting high-return trajectories, or why return-based contrastive learning outperforms traditional trajectory-based approaches.\n3. The paper's method of determining positive and negative samples relies heavily on hyperparameters (ξ, ζ, σ) without theoretical justification for their value ranges. The paper lacks a principled approach to determine these thresholds across different environments. A more systematic study on how different return thresholds affect the distribution of positive/negative samples and their impact on learning dynamics would strengthen the method's foundation.\n4. The current state-level contrastive learning approach breaks the temporal correlation between states within trajectories, as it samples positive and negative states purely based on return values (SR) or clustering results (SRD). This design might lose important sequential information that could be better captured through trajectory-level contrastive learning. I encourage authors to explore trajectory-aligned state sampling strategies and conduct comparative experiments between state-level and trajectory-level contrastive learning to provide more solid empirical evidence for the design choices.\n5. The competitive performance of baselines like CDE and HD raises questions about the necessity of the proposed contrastive mechanism, especially considering these methods achieve comparable results without any specific designs for handling low-return trajectories. While ContraDiff shows improvements in certain scenarios, the performance overlap with these simpler methods suggests that the advantages of explicitly handling low-return trajectories might be less significant than claimed. This observation calls for a more thorough analysis of what unique benefits, if any, are brought by the contrastive mechanism for low-return trajectory utilization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "S6HtbgwfIZ", "reviewer": "Reviewer_DdY7", "summary": "The paper presents a method called ContraDiff, which leverages contrastive learning to guide an agent towards high-return states and pushing agents away from low-return states. The proposed method builds upon diffusion-based trajectory planning models and aims to solve offline reinforcement learning problems especially when high-return trajectories are limited in the offline dataset. Experiments are conducted on standard D4RL tasks and exhibit better performance as compared to baselines.", "questions": "Aside from mentioned in the Weakness section, my questions mostly come from experiments presented by the paper:\n1. When comparing Table 3 with Table 1, it seems like introducing the expert trajectory leads to a decrease in performance, unexpectedly. For instance, DT achieves a score of 36.6 at original HalfCheetah-MR dataset but 7.5, 6.7, 6.1 in three conditions where expert data is introduced. I was wondering if the same codebase is used in achieving these results, and if so, could the authors share insights on why DT’s performance drops so drastically with the added expert data?\n2. In most tasks ContraDiff-SR outperforms ContraDiff-SRD -- this seems unexpected given the author's intuition mentioned in section 3 that ContraDiff-SR may ignore the transition dynamics in its sampling process?  \n3. How would the proposed model perform on sparse-reward tasks like AntMaze? This benchmark is commonly used in offline RL evaluations but results for it are missing here.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The combination of contrastive learning and diffusion-based trajectory planning model is a novel idea, allowing the model to learn from both high- and low-return samples. The main figure is clear and well-organized, making it easy to understand the method. Experiments are extensively conducted across many test environments.", "weaknesses": "1. While the idea of using contrastive learning is well motivated, the section on how \"positive\" and \"negative\" examples are identified requires additional explanations. The paper did not specify how the clustering is performed -- if the clustering is to enforce \"dynamic consistency\", I would assume the information is already available in the offline dataset? For me it is unclear how well the clustering captures reachability or state transitions beyond just grouping next states together. This could be a limitation in environments with complex or non-linear dynamics.\n2. The paper lacks discussion on the complexity or runtime implications introduced by k-means clustering, especially for large offline RL datasets.\n3. The increases in performance are mostly marginal. For many tasks the increase falls within one standard deviation of the baseline methods. \n4. The authors mentioned  \"comparing ContraDiff with other regular methods\" but didn't include specific indication of where the results are presented. I later found the results in the appendix, but the authors may consider additional illustrations and explanations on the results. Furthermore methods like CQL are only tested on the standard D4RL datasets but not the mixture dataset proposed by the authors. More experiment results would be needed here to better demonstrate model performance. \n5. It is unclear how the baseline results were obtained -- whether they came from existing code, were reimplemented by the authors, or taken from other sources. \n6. The writing, especially in the experiments section, is not very clear to understand. Some sentences and typos make the paper harder to follow. For example: \n* L290: \"...high-return sample sparsity situations\"\n* L301: \"...which are focus on addressing...\"\n* L364: \"...declines with the ratio declines\"\n* Table3: \"Walekr2d\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "rIVYiUny7q", "reviewer": "Reviewer_Ku5d", "summary": "This paper introduces a novel approach, ContraDiff, designed to enhance offline reinforcement learning (RL) by leveraging low-return trajectories in a contrastive learning framework. Offline RL's performance often depends on the presence of high-return trajectories, which are frequently sparse in datasets. ContraDiff addresses this challenge by classifying trajectories into high- and low-return groups, treating them as positive and negative samples, respectively. The approach applies a contrastive mechanism that pulls the agent's trajectories toward high-return states while pushing them away from low-return ones. This strategy enables low-return trajectories to serve as a guiding force to avoid suboptimal regions, enhancing policy performance.", "questions": "* In the 3.3 model learning section, I noticed that the optimized trajectory generation by minimizing the Mean Square Error between the ground truth and neat trajectory predicted. Therefore, the diffusion should denoise the data from the noisy data completely. I think it's more expensive to train like this. Can you show some experiments comparing the cost of training?\n* ContraDiff planning towards high return states, leading policy improvements. It is very similar to some offline RL methods, such as SAW[1], A2PR[2], LAPO[3]. Can you add some discussion with these methods in the related works or more experiments comparison? \n* Did you experiment with setting different thresholds for what qualifies as “low” or “high” return? If so, how did varying these thresholds impact the learned policy or the model’s ability to generalize to novel tasks?\n\nReferences：\n[1] Lyu, Jiafei, et al. \"State advantage weighting for offline RL.\" arXiv preprint arXiv:2210.04251 (2022).\n\n[2] Liu, Tenglong, et al. \"Adaptive Advantage-Guided Policy Regularization for Offline Reinforcement Learning.\" In International Conference on Machine Learning (ICML). PMLR, 2024.\n\n[3] Chen, Xi, et al. \"Latent-variable advantage-weighted policy optimization for offline rl.\" arXiv preprint arXiv:2203.08949 (2022).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a compelling approach to enhancing offline reinforcement learning (RL) through a novel contrastive framework, which incorporates low-return trajectories in ways that previous approaches overlook. ContraDiff’s use of a contrastive learning framework to distinguish between high- and low-return trajectories is an innovative application of contrastive learning within the offline RL domain. Instead of merely reweighting high-return samples, the method leverages low-return trajectories as guiding forces, creating a “counteracting force” that helps the policy avoid suboptimal states. This approach represents a shift from traditional methods that focus primarily on high-return trajectories, expanding the value of underutilized low-return data.", "weaknesses": "* While the paper presents contrastive learning as a means to exploit low-return data, it could better clarify why this mechanism is theoretically optimal for avoiding low-return states compared to other potential approaches, such as weighting adjustments or imitation-based filtering. \n\n* The authors are not clear enough in describing the use of weighted contrast loss to constrain trajectory generation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)