RTDiff: Reverse Trajectory Synthesis via Diffusion for Offline Reinforcement Learning
Authors: Qianlan Yang
Yu-Xiong Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement Learning
Diffusion Model
Reverse Synthesize
Primary area: reinforcement learning

Abstract: In offline reinforcement learning (RL), managing the distribution shift between the learned policy and the static offline dataset is a persistent challenge that can result in overestimated values and suboptimal policies. Traditional offline RL methods address this by introducing conservative biases that limit exploration to well-understood regions, but they often overly restrict the agent's generalization capabilities. Recent work has sought to generate trajectories using generative models to augment the offline dataset, yet these methods still struggle with overestimating synthesized data, especially when out-of-distribution samples are produced. To overcome this issue, we propose RTDiff, a novel diffusion-based data augmentation technique that synthesizes trajectories *in reverse*, moving from unknown to known states. Such reverse generation naturally mitigates the risk of overestimation by ensuring that the agent avoids planning through unknown states. Additionally, reverse trajectory synthesis allows us to generate longer, more informative trajectories that take full advantage of diffusion models' generative strengths while ensuring reliability. We further enhance RTDiff by introducing flexible trajectory length control and improving the efficiency of the generation process through noise management. Our empirical results show that RTDiff significantly improves the performance of several state-of-the-art offline RL algorithms across diverse environments, achieving consistent and superior results by effectively overcoming distribution shift.

Published as a conference paper at ICLR 2025
RTDIFF: REVERSE TRAJECTORY SYNTHESIS VIA DIF- FUSION FOR OFFLINE REINFORCEMENT LEARNING
Qianlan Yang Yu-Xiong Wang University of Illinois Urbana Champaign (qianlan2 y loninniii nois edu
ABSTRACT
In offline reinforcement learning (RL), managing the distribution shift between the learned policy and the static offline dataset a persistent challenge that can result in overestimated values and suboptimal policies. Traditional offline RL meth- ods address this by introducing conservative biases that limit exploration to well- understood regions, but they often overly restrict the agent's generalization capa- bilities. Recent work has sought to generate trajectories using generative models augment the offline dataset, yet these methods still struggle with overestimat- ing synthesized data, especially when -i-ooiiiiiiittiionn samples are produced. To overcome this issue, we propose RTDiff, novel diffusion-basee data augmen- tation technique that synthesizes trajectories in reverse, moving from unknown to known states. Such reverse generation naturally mitigates the risk of overestima- tion by ensuring that the agent avoids planning through unknown states. Addition- ally, reverse trajectory synthesis allows us to generate longer. more informative trajectories that take full advantage of diffusion models' generative strengths while ensuring reliability. We further enhance RTDiff by introducing flexible trajec- tory length control and improving the efficiency of the generation process through noise management. Our empirical results show that RTDiff significantly improves the performance of several ofate-of-the-art offline RL algorithms across diverse environments, achieving consistent and superior results by effectively overcom- ing distribution shift. Our code can be found at https:/yangaal.iiubbi/ROR github.io/RTDif..
INTRODUCTION
Deep reinforcement learning (RL) has become a powerful tool for tackling complex challenges across a wide array of fields, including mastering board games (Silver al., 2016), achieving su- perhuman performance in video games (Mnih et al., 2015), and improving continuous control in robotics (Lillicrap et al., 2016). The success of deep RL algorithms primarily due to their abil- ity to interact with and learn from extensive datasets generated through environmental interactions. However, in real-world applications, accumulating such a vast amount of exploratory data is often impractical and costly. In critical domains like healthcare (Gottesman et al., 2019) and autonomous driving (Yu et al., 2018), 8), every interaction carries significant costs or risks, making online data col- lection infeasible and unsafe. Offline RL (Lange et al., 2012: Levine et al., 2020) offers solution by training agents on pre-existing datasets, thereby avoiding the risks and costs of online data gener- ation. Nevertheless, the transition from online offline RL is challenging. Directly applying online RL techniques to offline RL tasks usually results in poor performance (Fujimoto et al., 2019; Wu al., 2019). This is mainly due to the distribution shift between the policy derived from the offline data and the policy that originally generated the data. Such a shift can lead to overestimation of the unseen data the offline dataset, producing inaccurate value estimates and suboptimal policies.
A typical solution is to develop advanced offline RL algorithms that incorporate a conservative bias into the learning process. These algorithms limit the policy search to regions within the offline dataset where there is high confidence. Model-free offline RL approaches, such those proposed by Fujimoto et al. (2019) and Wu et al. (2019), embed this bias directly into their policy or value functions through conservative regularizations or specialized network architectures. Although these strategies effectively mitigate the problems associated with distribution shift, they may overly restrict the policy search, limiting the agent's ability to generalize beyond the specific confines of the offline

---

Published as a conference paper at ICLR 2025
Initall States
Ageat
States offiine dataset
Offline Dataset D
Goal!
RTDiff
Knowwn area
Senthetic Trglestociss for treining
Real Trajectrries for training
Dungerous area
Tajortories ofliine dataret
Agent
Reverse syothesls
synthesis
(a) The main paradigm of RTDifF.
(b) Illustrative example.
Figure 1: Illustration of RTDifT. (a) RTDiff is used to augment the training dataset for reinforce- ment learning agents. It receives initial states from the offline dataset D and synthesizes reverse trajectories. Both real and synthesized trajectories are utilized to train the RL agent. (b) An illustra- tive example is shown. In this figure, blue arrows represent trajectories inside the offline dataset, red arrows represent trajectories synthesized the forward direction, and green arrows represent tra- jectories synthesized in the reverse direction. Forward trajectories include paths from known areas to dangerous areas, which might hinder performance. contrast, reverse trajectories only include paths from dangerous areas to known areas, which do not adversely affect the agent.
dataset. Conversely, model-based offline RL methods (Kidambi & Rajeswaran, 2020; Yu et al., 2020; Lee et al., 2021; Yu et al., 2021) adopt a different strategy. These methods begin by learning forward dynamics model that integrates conservative estimates from the offline dataset. They then use this model to generate imaginary trajectories, thereby expanding the dataset with high- confidence synthetic data.
Inspired by the promise of data synthesis, we turn our focus to state-of-the-art generative models. Diffusion models, known for their generative capabilities in computer vision and natural language processing, have been noticed by the RL researchers. Janner et al. (2022) and Ajay et al. (2023) introduce a foundational approach for decision-making in RL by generating full trajectories through single denoising process. Building on the use of diffusion models in planning, Lu et al. (2023) pro- pose synthesizing transitions using diffusion models trained on offline datasets to augment available data. More recent work (Yang & Wang, 2024; He et al., 2023) has expanded these ideas, applying diffusion models to generate trajectories that accelerate RL training. However, these approaches face the persistent challenge of distribution shift in data synthesis, which can lead to overestimation of values when -iddtr--iittiittionn data are synthesized relative to the offline dataset.
To address this issue, in this paper, we propose simple yet effective strategy overlooked by the community: generating reverse trajectories instead of forward ones. Our key insight lies in intuition that overestimating trajectory that begins from unknown state and moves state does not affect performance, as the agent will not pass through the unknown state during plan- ning. Here, known states refer states that close the distribution of states in the offline dataset, while unknown states are relatively far from this data distribution. During the planning phase, the agent always starts from known state, where overestimated trajectory originating from such states may risk leading the agent to unknown states, ultimately harming performance. However, by replacing such trajectories with those starting from wnknown states and moving to known states, these risks are naturally reduced. Therefore, reverse synthesis directly addresses the issue of distribution shift. Additionally, by eliminating concerns about overestimating -iittriiittiiitio trajectories, we can generate longer trajectories, providing the agent with more information and leveraging the ability of diffusion models to produce long, reliable trajectories.
known
Based on this insight, our work introduces a novel paradigm in offline RL by incorporating a dif- fusion model with reverse synthesis, termed Reverse Trajectory Diffuser (RTDiff). As illustrated in Fig. I, our approach generates trajectories that converge towards target states already present in the offline dataset. By incorporating flexible generation length control, we can extend the trajectory length much possible to maximize the benefits of the diffusion model while ensuring the reli- ability of the generated data. Furthermore, by controlling the initial noise received by the diffusion model, we enhance the generation efficiency of our model.
2

---

Published as a conference paper at ICLR 2025
Our contributions are three-fold. (i) We propose RTDiff, novel diffusion- approach that synthesizes reverse trajectories to address distribution shifts in offline reinforcement learning. which is general and can be integrated into variety of offline reinforcement learning algorithms to improve their performance. (ii) We develop an out-of-distribation detector-basee method that automatically adjusts the length of the generated trajectories. (iii) We reduce the number of generated samples by improving the generation efficiency through noise control. Empirical evaluation shows that RTDiff consistently achieves state-of-the-ar performance across a variety of offline reinforcement learning environments.
2 RELATED WORK
Offiine RL. In the realm of offline reinforcement learning, researchers focus on generalizing to out- of-distribution (OOD) data and avoiding the overestimation bias. These methods can be categorized into two main types: model-free and model-based algorithms. Model-free offline RL methods typ- ically restrict their policy search to the offline dataset. This can be achieved through various ways, such as explicitly constraining the learning policy to remain close to the dataset (Fujimoto et al., 2019; Liu et al., 2020b), learning conservative value function (Kumar et al., 2020), and applying importance sampling-base based algorithms (Precup et al., 2001; Liu et al., 2020a). Additionally, they may estimate uncertainty quantification for the value function (Levine et al., 2020; Agarwal al., 2020).
On the other hand, model-basee offline RL methods have explored several different strategies. These include model-uncertainty quantification (Kidambi & Rajeswaran, 2020; Yu et al., 2020; Ovadia al., 2019), representation learning (Lee al., 2021), and constraining the policy to imitate the behavioral policy (Matsushima et al., 2020). They also use conservative estimation of the value function to enhance performance (Yu al., 2021). Recently, new strategies have emerged for solving offline RL, such as treating RL problems as sequence generation problem (Chen et al., 2021; Janner et al., 2021) or repurposing diffusion models for planning (Janner et al., 2022; Ajay et al., 2023).
Diffusion Models in RL. Diffusion models have demonstrated impressive capabilities in rein- forcement learning. particularly in enhancing long-term planning and policy expressiveness. Prior work (Janner et al., 2022; Ajay et al., 2023) introduces paradigm where diffusion models are used to construct full trajectories through conditioned sampling, guided by various criteria such as re- wards, goal-oriented navigation, and skill deployment. These methods exploit the unique ability of diffusion models to generate extensive trajectories, effectively addressing challenges like long hori- zons and sparse rewards in RL. Other work (Du et al., 2023; He et al., 2023) extends this paradigm to generate visual-based data using diffusion models. Lu et al. (2023) employs a diffusion model to generate transitions that supplement the replay buffer, offering distinctive strategy compared to earlier approaches. He et al. (2023) further positioned diffusion models as data synthesizers for gen- erating trajectories solve multi-task RL problems. More recently, ATraDiff (Yang & Wang, 2024) introduces a general framework that leverages offline data to generate full synthetic trajectories, improving the performance of online RL methods.
Data Augmentation in RL. Data augmentation has become a key technique for improving RL performance. Traditional methods (Yarats et al., 2021; Laskin et al., 2020; Sinha et al., 2021) in- corporate various data augmentations like adding noise or random translation on observations for visual-based RL. Such approaches aim to help the agent to learn on multiple views of the same ob- servation to improve the robustness. Some other works focus on generating synthetic data based on the models learned from the offline dataset (Yu al., 2020; Wang et al., 2021). Recently, different from disturbing existing data points, researchers have focused on generating new data points from models learned from the original datasct and upsampling the original dataset. Representative work includes using diffusion models generate data (Lu et al., 2023; Ajay et al., 2023; Du et al., 2023; He et al., 2023; Yang & Wang, 2024).
3 BACKGROUND
Markov Decision Process. In this paper, we explore sequential decision-making tasks that can be represented as a Markov Decision Process (MDP), denoted by .M = (5, A, R, 7). In this frame- work, S denotes the set of states, A denotes the set of actions, and 7 E [0, 1) is the discount factor.
3

---

Published as a conference paper at ICLR 2025
The functions [(s a) and R(s, a) describe the transition dynamics and reward structure, respec- tively. At each time step t. the agent selects an action a E A, resulting in a new state s' based on the transition function [ss' (s, a) and receives immediate reward R(s, a). A trajectory in this context a sequence of states and actions, expressed as (81,41 #2, 02, r2, 8t, 4z, rz,),, where 8t, ai and rı represent the state, action, and reward at time step respectively. Similarly,we define the verse trajectory starting from the state is the sequence (st, Qt, Fl: 91-1, 71-,,,,........... $1, 01, TI).
Diffusion Models.
Diffusion probabilistic models conceptualize the process of generating data as an iterative denoising sequence, represented by p=(r"-17"). This sequence reverses forward diffusion mechanism, q(r'|r'=), which incrementally adds noise to the data over N steps, thus degrading its structure. The resulting data distribution is described by:
pa (=(x p(a
p(ax> (T)) II po po (T))dT
where )) acts as a standard Gaussian prior, and p(r") corresponds to the original noiseless data. The model parameters, 0, are optimized by minimizing a variational bound on the negative log-likes log-likelihood of the denoising process: 0 = arg ming -E,o (log (o (= This reverse process is typically modeled as a Gaussian distribution with fixed, timestep-specific covariances:
Po (xx (T)=="() =
(T), E*).
Offiine RL. Offline RL is a type of reinforcement learning algorithm where the agent will be offered with offline dataset D {(s, a, a, a,); and try to learn policy TD from the offline dataset D. In offline RL, the agent will not be allowed to interact with the environment for online data collecting. The offline dataset is usually collected through multi-source policies.
4 METHOD
We now present our approach to augmenting the offline dataset by training our generative model RTDiff on the offline dataset to synthesize trajectories reversely. We begin by introducing how we design and train RTDiff (Sec. 4.1) and then decide the length of the generated trajectories to support flexible generation space (Sec. 4.2). Finally, we introduce how to further improve the generation efficiency by adjusting the input noise of the diffusion model (Sec. 4.3).
4.1 TRAJECTORY GENERATOR
To capture the trajectory data distribution of the offline dataset, we train a diffusion model to solve conditional generation problem:
max (x(")
where (T) is the final generated reverse trajectory and y(T) is the generation condition. In the experiments our paper focuses on, we use proprioceptive information as the state. So, the gener- ated trajectories should be relatively low-dimensional. Specifically, we define the generated reverse trajectory at I-th diffusion step with length I as the following 2D array:
st.
s"
x' <'(r)
=
ag
a
a²
In Eqn. 2, the i + 1-th column of (T) is concatenation of state 8-j, action a-; and reward r-is which index 0 means the current step, -i means the i-th last step. The generation condition is set to be the initial state y(T) 80-
Training. To obtain the training dataset of the diffusion model, we first sample trajectories with fixed length L from the offline dataset. For trajectory with length in the offline dataset, we divide it into n + trajectories with length L so that the whole dataset becomes a larger dataset with fixed length L. For the training of the diffusion model with N denoising steps, we use the following loss:
T(0)= MM [[]e (o(=)())) (7), 011
(3)
4

---

Published as a conference paper at ICLR 2025
Algorithm 1 Augment Offline Dataset D with RTDiff
Require:
Offline dataset D, augmentation size M
1: Initialize the synthesized dataset D. -
2: Train RTDiff Pa with the dataset D
3: while [D2
< M do
4: Sample state so D from the offline dataset
5: Generate the trajectory т = (So. 00,Y0: 8 8-1, a_1, --., --1,,-1,7-11
with RTDiff
6: Cut the trajectory with the OOD detector
(So; on, ro, 8-1,0-1, 1-,,, --1, -1,51)))))
7: Add all the transitions of the trajectory T to the synthesized dataset D, = D, UT
8: end while
9: Combine the synthesized dataset to the original dataset D = DUD,
where u(N) is the uniform distribution on {1, .1,
Architecture. Since our experimental tasks focus on proprioceptive environments, the generation content is relatively low-dimensional compared with pixel-based generation, and thus do not use a similar U-Net architecture as image works. We parameterize 69 with an MLP with skip connections from the previous layer followed as (Lu et al., 2023) but increase the network width to 4096. For the sampling process, we use EDM (Karras et al., 2022) as the sampling method and we set the diffusion steps 128.
Deployment. As stated in Algorithm 1, to augment the given offline dataset, we first train our RTDiff on the given offline dataset. Then we repeatedly sample a state so the initial state from the original dataset D. and synthesize a reverse trajectory from the initial state After that, we will cut the trajectory by an out-of-distribution (OOD) detector which will be introduced in the next section, and add all the transitions in the trajectory to the synthesized dataset. We repeat this process until the size of the synthesized dataset reaches our expected size M.
4.2 GENERATION LENGTH CONTROL
RTDiff is limited to generating fixed-length trajectories due to the inherent characteristics of diffu- sion models. Our investigation indicates that the quality of the augmented dataset is sensitive to the length of these trajectories. When the generated trajectories are too short, the agent gains limited benefit from the diffusion model's capacity produce extended, reliable trajectories, thereby stricting RTDiff's performance enhancement, Conversely, if the generated trajectories are too long, they will go into -isttiibiitiitti regions, causing hallucinations. An excess of these unrealistic generations can significantly decrease performance. For a deeper empirical analysis of the perfor- mance of different generation lengths, please refer to Sec. 5.2.
Therefore, we aim to design a flexible length control mechanism that can automatically adjust the generation length. This mechanism will ensure that each generated trajectory is long as possible without excessively entering out-of -istriiiitiontoonnn regions.
We train an out-of-distribution detector d(s) on the state space with the offline dataset, where measures the distance of state s to the distribution of the given offline dataset. Ideally, we con- sider d(s) > 1 to represent that s is an out-of-distribttion state. To cut the generated trajectory (so, 0o, To, 8-1. ". 8- a,,, rI), we find the smallest / + I that d(s-(1+1)) > dis then we drop the trajectory after 8-1- Here dis M is a hyperparameter of our method, the details about how choose this hyperparameter are included in Appendix C.3
Distri-Dittiiitiinn Detector. We train the OOD detector only by using the data points in the offline dataset. We design the OOD detector following the classic OOD work SSD (Schwag et al., 2021). The main training scheme is: we first train a feature representation method to get the features of the states. Then, partition the offline dataset into clusters based the trained features, denoting each cluster as Zm. Finally, the distance of state is defined as d'(s) = min Dm (8, Zm) where Dm the Mahalanobis distance. We normalize the distance by dividing the maximum d'(s) distance of states in the offline dataset D, d(s)) max,rep d' to make this threshold generalizable to all the environments.
5

---

Published as a conference paper at ICLR 2025
4.3 IMPROVING GENERATION EFFICIENCY
Although RTDiff can enhance the performance of offline RL by augmenting the dataset, the number of generated samples can still be quite large, resulting in increased time and cost. To address this, we propose a technique aimed at improving the generation quality of RTDiff, thereby reducing the number of samples required in the data augmentation process and increasing sample efficiency.
Intuitively, if we generate two similar trajectories from the same initial state, the information car- ried by the two trajectories will be less than two independent trajectories. Therefore, we need to maximize the dissimilarity between all generated trajectories, which means reducing the correlation between the generated trajectories.
To this end, we propose simple strategy that controls the noise .:" (r) we choose for the diffusion x. (r) for one single initial model. Instead of separately generating noises of (r),, (7).-.., , the state, we generate these noises together that they will be evenly located within the space. That is we generate the noises by minimizing:
x (r)
j=1
which repulses noise vectors from each other.
Conceptually, this strategy is effective. As the diffusion model itself has the ability to generate diversified trajectories at a sufficiently high probability (this means the diffusion model will not always generate the same trajectory for all noises), can assume that: for trajectory тΛ generated by noise N (T), with high probability, another trajectory тn generated from a randomly sampled noise a (7) ~ (0, 1) is sufficiently different from trajectory TA, say d(TA, TB) > do, where de constant. We regard the full sampling process as a funetion from noise to trajectory. As the diffusion model's network and noise samplers are continuous, such sampling process is a continuous function. Using such continuity, we reach the conclusion that if we sample noise N (T) only from the neighborhood of noise x²A (r), with high probability the generated trajectory tB from x (T) will be close to trajectory TA- Therefore, intuitively if choose N (T) which is located far from x (T), we can avoid the neighborhood region (that fails at high probability) and have higher probability to get a different trajectory.
5 EXPERIMENTS
We conduct a variety of experiments to verify the effectiveness of RTDiff. First, we validate that our approach can improve the performance of state-of-the-ee offline RL algorithms in various envi- ronments (Sec. 5.1). Moreover, we conduct ablation studies to validate the effectiveness of different components in our approach (Sec. 5.2). Finally, we conduct additional experiments on visual RL tasks to show that our method can be extended to risual-based methods (Sec. 5.3). For evaluation, all results in this section are presented with the median performance over 5 random seeds.
5.1 MAIN RESULTS
We first show the overall performance of RTDiff, validating that RTDiff can improve the perfor- mance of oftee-of-the-an art offline RL algorithms. We conduct experiments on 4 different environ- ments from D4RL (Fu et al., 2020): Maze2d, AntMaze, Locomotion, and Kitchen. For comparison, we selected IQL. (Kostrikov et al., 2022) and TD3+BC (Fujimoto & Gu, 2021) as benchmark of- fline RI. algorithms, both widely regarded as state-of-the-ard baselines. Additionally, we include experimental results for CQL (Kumar et al., 2020) and Decision Transformer (Chen et al., 202 though these are omitted the main paper due to space constraints; full results can be found in Appendix C. We select several data augmentation methods as baselines, including SynthER (Lu et al., 2023) and ATraDiff (Yang & Wang, 2024).
As illustrated by the results shown in Tables and 2, our RTDiff consistently improves the per- formance of the offline RL algorithms in a wide range of environments. Mcanwhile, our method outperforms the other data augmentation baselines including SynthER (Iu et al., 2023) and ATraD-

---

Published as a conference paper at ICLR 2025
Overall results of D4RL Maze, Antmaze, and Kitchen environments. The number de- notes the performance increase by the data augmentation method compared to the original result. RTDiff consistently enhances the performance of offline reinforcement learning algorithms in all these environments. In particular, RTDiff significantly improves the performance in maze and kitchen environments.
Table I:
IQL (Kostrikov al., 2022)
TD3+BC (Fujimoto Gu, 2021)
Environment
Data Type
RTDiff
SynthER
ATraDiff
RTDiff
SynthER
ATraDiff
8.3
4.3
5.645.0
10.2
8,324.3
9.3=0.
umaze
maze2d
medium
3.342.7
0.8
2.145.6
9.8=2.5
6,342.
7.4
large
14.3
11.442.8
12.4
7.7
4.843.8
4.3
fixed
5.2
4.943.7
4.44...0
5.743.5
5,443.8
2.5=4.6
antmaze-umaze
diverse
1.342.7
4.3
4.7
4.2
3.9
3.5
play
7.9:::.2
7.543.6
6.6
8.343.2
8.4
7.6
antmaze-medium
diverse
9.2#:3.8
8,5
8.8+4.6
8.9
8.343.6
8.723.1
play
6.5
5.4+2.8
5.744.0
5.4
4.8+2.0
5,0=1.5
antmaze-large
diverse
6.343.4
5.742.6
4.4
6.2
4.7462
4.9=7.1
complete 6.6
3.4+9.3
5,449.2
5,3+7.2
3.6+0.5
5.647.5
kitchen
partial
13.6.6.3
8.3+7.2
11.3
14.247.8
6.4
12.3
mixed
11.3#8.55
6,2
8,449.4
10.3
7,247.7
9.3
Table Overall results of D4RL Locomotion environment. The number denotes the performance increased by the data augmentation method compared to the original result. RTDiff improves the performance of various offline RL methods in different tasks, achieving the state of the art performance.
IQL (Kostrikov et al., 2022)
TD3+BC (Fujimoto & Gu, 2021)
Environment
Data Type
RTDiff
SynthER
ATraDiff
RTDiff
SynthER
ATraDiff
random
0.
0,1
0.341.3
0.0
0,1
0,0
mixed
1.342.3
0.744..3
0.445.1
4.342.5
4.2428
3.7
walker2d
medium
1.1447
0.743.7
0.8
3.142.6
2.6
1.942.3
medexp
0.2
0.0
0.4
0.2#0.2
0.3
0.1
random
1.2#0.2
0.5+0.3
0.6
5.340.4
3.640.5
1.5+0.8
mixed
17.2#1.7
18.4#2.4
13.5
7.4
5.4
3.625.3
hopper
medium
10.7#6.0
9.8
7.8
8,24339
4.6
2.5+8.3
medexp
3.6+444
2.4
3.8
7.2
5.8+1.3
2.9+8.7
random
3.5+112
2.1
2.7+2.3
1.6
1.8+2.0
1.2+2.8
mixed
4.2
3.3+0.4
3.8+1.7
4.2#1.2
3.7
2.4+1.9
halfcheetah
medium
2.4
1.5+0.3
0.9+1.
1.4
1,20011
0.6+0.3
medexp
0.6
0.240.1
0.540.3
0.5
0.1
0.0
iff (Yang & Wang. 2024), which shows that RTDiff could generate data with the higher quality to enhance the offline dataset.
In particular, RTDiff significantly enhances the performance of offline RL algorithms in the Maze2D and Kitchen environments. We hypothesize two main reasons for this improvement. First, for tasks with long horizons, RTDiff leverages diffusion models to generate extended yet reliable trajecto- ries, which are crucial for performance gains. Second, in environments where the state space large relative to the states covered by the offline dataset-resulting in sparsely populated offline dataset-RTDDin can more effectively utilize reverse synthesis to in the gaps, thus providing greater benefits.
5.2 ABLATION STUDIES
To verify the effect of different components in our RTDiff, conduct several ablation studies.

---

Published as a conference paper at ICLR 2025
Is reverse synthesis important to the performance? The key contribution of our work is to syn- thesize trajectories reversely. We argue that reverse synthesis is better than normal synthesis for data augmentation. To show the importance of reverse synthesis, paring our method with a normal synthesis algorithm. Here the method 'Normal' represents the algorithm that synthesizes trajectories in the forward order, with all other details the same with RT- Diff. And the method 'Short' represents the algorithm synthesizing trajectories in the forward order, but the generation length has been set to be 3, which we found to be the best choice in forward synthesis. The results shown in Table 3 illustrate that reverse synthesis significantly outperforms normal synthesis. Meanwhile, we found that we perform normal synthesis, we cannot benefit from long trajectory generation, the performance of synthesizing 'Short' trajectories outperforms synthesizing 'Normal' trajectories.
conduct an ablation study com-
Table 3: Ablation study on reverse synthesis and normal synthesis with different generation lengths. The results indicate that reverse synthesis significantly outperforms normal synthesis. Additionally, increasing the generation length in normal synthesis leads to a decrease in performance.
IQL (Kostrikov et al., 2022)
TD3+BC (Fujimoto & Gu, 2021)
Environment
Data Type
RTDifF
Normal
Short
RTDiff
Normal
Short
8.3
4.3+10
5.8+3.8
10.2
3.5+11
8.5
umaze
maze2d
medium
3.3+2.7
2.4+32
5,9+41
9.8+2.5
3.4+30
6.1 6.1
large
14.3
3.243.5
7.8
7.7
2.4
8.1
complete
6.647.4
1.5+7.0
3.1
5.3
0.8
2.6
kitchen
partial
13.6
3.0
8,246.8
14.2
3.547.1
5.947.4
mixed
11.3
1.9+8.0
5,4+8.3
10.3
3,2+7.7
7.7+7.9
Effect of the generation length. The flexible generation length control is also an important com- ponent of our method. We found that the quality of the generated data is very sensitive to the length of the trajectories. Here we conduct an ablation study show the performance of RTDiff with different fixed generation lengths. From the experimental results in Table 4, we can see that the performance of different generation lengths varies significantly and our generation length control strategy outperforms every fixed generation length setting.
The effect of noise control. We conduct an ablation study on the noise control component of RTDiff. We test RTDiff with and without the noise control part, using three different quantities of generated samples. The results, shown in Table 5, indicate that the noise control component enhances overall performance. Notably, the number of generated samples decreases, the perfor- mance improvement due to noise control becomes more noticeable.
Is RTDiff better than model-based methods? We include more results of baselines including model-based RL methods MOPO (Yu et al., 2020) and model-based reverse imagination method ROMI (Wang al., 2021). The results shown in Table 6 illustrate that RTDiff still outperforms those model-based baselines.
5.3 VISUAL REINFORCEMENT LEARNING
Finally, demonstrate the applicability of RTDiff to visual reinforcement learning tasks, using WerldWorld (Yu et al., 2019) as the benchmark. We selected tasks with varying difficulty levels, ranging from easy to hard. To adapt RTDiff for image-based inputs, we employed different ar- chitectures to generate visual trajectories. Building on the general pipeline from ATraDiff (Yang & Wang, 2024), which synthesizes pixel-based trajectories, we used Stable Diffusion to directly generate trajectory images. Actions and rewards were then predicted from these generated images using diffusion layers as the feature map. In addition, we modified the training images used in ATraDiff to synthesize reverse trajectories by training in reverse order. For evaluation, we applied offline RL algorithms, including CQL (Kumar et al., 2020), TD3+BC (Fujimoto & Gu, 2021), 1), and IQL (Kostrikov al., 2022). We directly compare RTDiff with ATraDiff for the Visual RL setting. The results, summarized in Table 7, demonstrate that RTDiff generalizes well to visual reinforce- ment learning, consistently enhancing the performance of offline RI. methods.
8

---

Published as a conference paper at ICLR 2025
Table 4: Ablation study on the effect of generation lengths. Our method significantly outperforms approaches using fixed generation lengths. We observed that as the generation length increases, performance improves as well. However, the generation length becomes too long, performance begins to decline.
CQL (Kumar et al., 2020)
Environment
Data Type
RTDiff
L=1
L=3
L=5
L=10
12.343.5
8,343.0
9,4+3.2
10.4#3.4
6,4428
umaze
maze2d
medium
8.3
6,2#25
6,642.6
7,2
5,3+24
large
11.343.3
8.3
9.2
10.4#3.2
7.2+2.9
complete
6,6474
3.840.5
4.540.8
5,847.1
4.346.3
kitchen
partial
13.6#6.3
9.4
10.345.8
11.4#6.0
10.5
mixed
11.3#8.5
7.4
8.5+78
9.3
8,4
Table 5: Ablation study the effect of noise control. The noise control component indeed improves the overall performance of RTDiff, especially when the number of generation samples is smaller.
IQL (Kostrikov al., 2022)
TD3+BC (Fujimoto Gu, 2021)
Environment
Noise Control
0.5M
1M
3M
0.5M
IM
3M
Yes
4.1
6.7420
8.343.5
7.2421
10.3
10.2
maze2d-umaze
No
2.3+10
3.1+18
6,3+1.1
4.8+15
8.6#2.8
8,3
Yes
4.3
5.8+1.9
3.3
3.5
8.5#27
9.8+2.
maze2d-medium
No
1.4+10
2,9+18
2.8+3.2
3,4+10
6.1
6.342.5
Yes
4.6+1.4
7.3#22
14.3
5.2#1.6
9.4
7.724.5
maze2d-large
No
3.4
6.9=2.1
11.4
4.3#13
8.6#2.6
4.8
Table 6: Overall results D4RL Maze2D environments. The number denotes the performance increase by the data augmentation method compared to the original result.
IQL (Kostrikov al., 2022)
TD3+BC (Fujimoto & Gu, 2021)
Environment
Data Type
RTDiff
MOPO
ROMI
RTDif
MOPO
ROMI
8.3
5.1
5.4
10.2
9.6
9.6
umaze
maze2d
medium
3.3
1.7+19
2.1+18
9.8
8,8+2.0
9.4+177
large
14.3:3.33
5.942.
8.1
7.7
2.613.8
3.5
Table 7: Results of the offline experiments on Meta- World (Yu et al., 2019). The number denotes
the success rate of completing the task. RTDiff consistently improves the performance of various offline reinforcement learning algorithms different tasks and outperforms ATraDiff.
TD3+BC
CQL
IQL
Task Name
Original
RTDiff
ATraDiff
Original
RTDiff
ATraDiff
Orignal
RTDiff
ATraDiff
Basketball
0.13
0.21
0.16
0.02
0.05
0.04
0.16
0.24
0.18
Box Close
0.03
0.07
0.04
0.00
0.01
0.00
0.08
0.12
0.09
Push Wall
0.06
0.12
0.08
0.02
0.03
0.02
0.09
0.14
0.12
Coffee Push
0.62
0.46
0.58
0.51
0.61
0.72
Sweep
0.19
0.24
0.22
0.16
0.21
0.17
0.24
0.26
0.24
VISUALIZATION AND ANALYSIS
To better understand why RTDiff improves the performance of offline reinforcement learning algo- rithms and why reverse synthesis avoids issues present in normal synthesis, we conduct an analysis using an illustrative environment, We design task similar to Maze2D but simpler and cleaner for better analysis. In this task, the agent starts from the bottom-left corner and aims to reach the target in the lef-left corner of the map. The agent receives reward of for approaching the goal. The left part of the map contains dangerous area; if the agent enters this area, the episode ends immediately, and the agent receives a reward of -00. The offline dataset does not contain full trajectory from the start to the end. Instead, there is small area in the middle of the map that is not covered by any trajectories the dataset.

---

Published as a conference paper at ICLR 2025
(a)
(b)
(c)
Figure 2: Left: The visualization of the illustrative environment. Middle: The trajectory of the offline RL agent with nomaal synthesis. Right: The trajectory of the offline RL agent with reverse synthesis. Normal synthesis will lead a trajectory going across the dangerous area, while Reverse synthesize can help the agent to remain in the known area and reach the goal.
We test both reverse synthesis and normal symthesis in this environment. As illustrated in Fig. 3, normal syn- thesis generates trajectories that start from the lower area and enter the middle dangerous area. If these hallucinated trajectories relatively higher rewards, they can mislead the agent into the dangerous area, ultimately resulting in failed ex- ecution. Conversely, reverse synthesis generates trajectories that move from the dangerous area to the upper or lower areas. Although these trajectories are also hallucinations, they do not impact the decision-making process because the agent will never make decisions from within the dangerous area. Consequently, the agent successfully navigate to the goal location.
In2Out
Our2In
In2In
Out2Out
Reverse
2.7%
21.6%
48.0%
27.3%
Normal (wla OOD)
11.2%
5.8%
71.8%
11.2%
Normal (w/ OOD)
18.2%
6.2%
52.1%
23.5%
Table 8: Ratios of different types of transitions generated by reverse synthesis and normal synthesis on the Maze2D- large environment. Normal synthesis generates much more In20ut transitions than reverse synthesis.
assigned
To further analyze this phenomenon in larger environments, we evaluated the ratio of four types of transitions: In2In, Out2Out, In2Out, and Out2In. In particular, In2Out transitions, which move from an in-distribution state to an out-of-distribution (OOD) state, are especially risky, as they may mislead the agent and degrade performance. In contrast, Out2In transitions, where the agent moves back from an OOD state to an in-distribution state, are less risky. We measure these ratios in the Maze2D-large environment for both reverse and normal synthesis methods, with and without an OOD detector. As shown in Table 8, normal synthesis produces significantly more In20ut transitions compared to reverse synthesis, even when an OOD detector applied. This finding underscores the robustness of reverse synthesis in avoiding risky transitions, ultimately improving the performance of RI. agents.
7 CONCLUSION
We introduce RTDiff, a diffusion-based offline reinforcement learning data augmentation method that synthesizes reverse trajectories. By incorporating the diffusion models to generate trajectories, we reversely generate long trajectories to augment the dataset. By using an out-of-distribution detec- tor, automatically adapt the length of the generated trajectories, improving the generation quality. By controlling the noise of the diffusion model, we remove redundant generations and improve the efficiency of generation. We test RTDiff in various environments and found that RTDiff consistently improves offline reinforcement learning algorithms, especially in some long-horizon, complicated environments.
Limitation and Future Work. Our work introduces novel reinforcement learning data augmen- tation method that mitigates the issue distributional shifts. However, our method still necessitates generating large number of samples achieve optimal performance, which results in extended in- ference running time. In future work, we aim to enhance the noise control component and improve sample efficiency, enabling to achieve comparable results with significantly fewer samples.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported in part by NSF Grant 2106825, NIFA Award 2020-67021-32799, the Toyota Research Institute, the IBM-IIlinois Discovery Accelerator Institute, the Amazon-Illinow Center on AI for Interactive Conversational Experiences, Snap Inc., and the Jump ARCHES en- dowment through the Health Care Engineering Systems Center at Illinois and the OSF Foundation. This work used computational resources, including the NCSA Delta and DeltaAl supercomputers through allocations CIS220014 and CIS230012 from the Advanced Cyberinfrastructure Coordina- tion Ecosystem: Services & Support (ACCESS) program, as well as the TACC Frontera super- computer and Amazon Web Services (AWS) through the National Artificial Intelligence Research Resource (NAIRR) Pilot.
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. ICML, 2020.
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision decision-making? ICLR, 2023.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. 2021.
Yilun Du, Mengjiao Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Joshua B. Tenenbaum, Dale Schu- urmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. In NeurIPS, 2023.
E. C. Fieller, H. O. Hartley, and E. S. Pearson. Tests for rank correlation coefficients. Biometrika, 1957.
Justin Fu, Aviral Kumar, Olir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv: 2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach offline reinforcement learning. In NewrlPS, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML, 2019.
Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale DDosii-elez, and Leo Anthony Celi. Guidelines for reinforcement learning in healthcare. Nature medicine, 25 2019.
Abhishck Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. In CoRL, 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICML, 2018.
Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xue- long Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. NeurIPS, 2023.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. NeurIPS, 2021.
Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In ICML, 2022.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In NeurIPS, 2022.

---

Published as a conference paper at ICLR 2025
Rahul Kidambi and Aravind Rajeswaran. MORel: Model-based offline reinforcement learning. In NeurIPS, 2020.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine, Offline reinforcement learning with implicit q- learning. ICLR, 2022.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In NeurIPS, 2020.
Sascha Lange. Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning, In Reinforce- ment learning, pp. 45-73. Springer, 2012.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. NeurIPS, 2020.
Byung-Jun Lee, Jongmin Lee, and Kee-Eung Kim. Representation balancing offline model-based reinforcement learning. In ICLR, 2021.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv:2005.01643 2020.
Timothy Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In ICLR, 2016.
Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with stationary distribution correction. In Uncertainty in Artificial Intelligence, pp. 1180-1190. PMLR, 2020a.
Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch reinforce- ment learning without great exploration. arXiv:2007.08202, 2020b.
Cong Lu, Philip Ball, Yee Whye Teh, and Jack Parker-Holder. Synthetic experience replay. In NeurlPS, 2023.
Frank J. Massey Jr. The kolmogorov-smimoov test for goodness of fit. Journal of the American Statistical Association, 1951.
Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment- efficient reinforcement learning via model-based offline optimization. arXiv:2006.0334 03647, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 2015.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, David Sculley, Sebastian Nowozin, Joshua V Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. arXiv:/906.025380 2019.
Doina Precup, Richard S Sutton, and Sanjoy Dasgupta. Off-policy lemporal-differene learning with function approximation. In ICML, 2001.
Vikash Sehwag, Mung Chiang, and Prateek Mittal. SSD: A unified framework for self-supervised outlier detection. In ICLR, 2021.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):444-88,, 2016.
Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4RI.: surprisingly simple self-supervision for offline reinforcement learning in robotics. In CoRL, 2021.
Jianhao Wang, Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-based imagination. In NewrIPS, 2021.
12

---

Published as a conference paper at ICLR 2025
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv:1911.11611, 2019.
Qianlan Yang and Yuxiong Wang. Atradiff:: Accelerating online reinforcement learning with imag- inary trajectories. In ICML, 2024.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. In ICLR, 2021.
Fisher Yu, Wenqi Xian, Yingying Chen, Fangehen Liu, Mike Liao, Vashisht Madhavan, and Trevor Darrell. Bdd100k: diverse driving video database with scalable annotation tooling. PrXivv1805..0 0068044877 2018.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In CoRL, 2019.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. arXiv:2005. 13239, 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-basee policy optimization. arXiv:2102.08363, 2021.
13

---

Published as a conference paper at ICLR 2025
Appendices
A ENVIRONMENTS TESTED
Following are the environments we evaluated in Sec. 5:
D4RL Maze2d (Fu et al., 2020). The maze2d task is navigation task that requires 2D agent to reach a fixed goal location in the maze. This task jusitifies the ability of offline RL algorithms to stitch previously collected subtrajectories to get the shortest path to the goal location. There three layouts this task, including umaze, medium and large. The dataset of this environment is generated by selecting waypoints randomly and using a planner which could generate subtrajectories among the waypoints.
D4RL AntMaze (Fu et al., 2020). The Antmaze task is navigation task that replaces the 2D ball from Maze2D with a 8-Dof Ant quadraped robot. This task combines the challenges of controlling the robot and navigting the robot the goal location. There are three different layouts in this environment, including umaze, medium, and large. The environment also contains3 three flavors of datasets, including fixed, diverse, and play, wich differs in the chosen of the start and goal locations.
D4RL Locomotion (Fu et al., 2020). The Locomotion environment contains three different types of tasks (walker2d, hopper, and halfcheetah), including 12 different offline data with varying levels of expertise (random, medium, medium-reppe ated by policy trained with early-stopping SAC (Haarnoja et al., 2018). The random datasets are generated by random initilized policy. The medium-replay datasets consist of samples in the re- play buffer during the training until the policy reaches the medium performance. The medium expert dataset contains part of the expert demonstrations and part of the suboptimal trajectories.
and cadium--eppert) The medium datasets are gener-
D4RL Kitchen (Fu et al., 2020). The Kitchen task involves a simulated environment where a 9-DoF robot manipulates various objects, such as sliding cabinet door, switching an overhead light, and opening a microwave. Initially introduced by (Gupta et al., 2019), this task requires the robot to complete sequence of multiple subtasks, each rewarded with sparse, binary reward upon successful completion. The offline dataset provided includes only portions of the complete sequence, necessitating that the agent learn to assemble these sub-trajectories effectively.
Meta- World (Yu et al., 2019). Meta-World is an extensive platform created to assess and enhance algorithms in both reinforcement learning and multi-task learning. With 50 unique robotic manip- ulation tasks, it provides a varied and demanding setting for evaluating how well algorithms can generalize and rapidly learn new skills.
B HYPERPARAMETERS
We list all the hyperparameters here, which are applied to all the environments. In addition, we will release our code upon acceptance.
Hyperparameter Batch Size
Value
16
Training Steps
10c
Optimizer
Adam
Learning Rate Trajectory Length
2 x 10 4
10
Distance Threshold
1.5
Diffusion Steps
128
Number of Generations
5 x 106
Table 9: Hyperparameter settings used in our experiments.
14

---

Published as a conference paper at ICLR 2025
Table 10: Results of CQL and Decision Transformer on the D4RL Maze, Antmaze, and Kitchen en- vironments. The numbers denote the performance increase by the data augmentation method compared the original result. RTDiff consistently improves the performance of offlie rein- forcement learning algorithms in all these environments.
CQL (Kumar al., 2020)
DT (Chen et al., 2021)
Environment
Data Type
RTDiff
SynthER
ATraDiff
RTDiff
SynthER
ATraDiff
12.3
6,3+4.1
7.124.0
17.2
8,324.3
9.0#4.1
umaze
maze2d
medium
8.3+2.7
5.8432
6,243.1
9.8
6,342.5
5.9#26
large
11.3
7,4428
7,842.9
12.7
7,843..
7.5
fixed
5.2
4.9
4.843.
5.7
5.4
5.5+3.7
antmaze-umaze
diverse
4.3
4.3+311
4.3
4.2
3,922.8
4.0*29
play
7.9
7.5+3.6
7,4
8,3
8.4
8,2#28
antmaze-medium
diverse
9.2
8.543.
8.8
8.9
8.3
8.543.4
play
6.5+t3.5
5,4228
5.6+330
5.4
4.8-22.0
4.6*2.1
antmaze-large
diverse
6.3
5,722.
5.9#2.7
5.8+5.5
4.726.2
5,0#6.0
complete
6.6
3,428.3
4.0
5.3
3,6:26.5
4.2
kitchen
partial 13.6:16.3
8,3+7.2
9.0+7.0
14,2#7.8
6,4468
7.0#6.9
mixed
11.348.0
6.249.1
6.0
10.3
7.2
7.5+7.6
Table 11: Results of CQL and DT on the D4RL Locomotion environment. The numbers denote the performance increase by the data augmentation method compared to the original result. RTDiff improves the performance of these reinforcement learning methods in different tasks.
CQL (Kumar et al., 2020)
DT (Chen al., 2021)
Environment
Data Type
RTDiff
SynthER
ATraDiff
RTDiff
SynthER
ATraDiff
mixed
5.2+2.3
4.9=4.
5.1
2.2
2.4
2.2
walker2d
medium
2.6
2.3-37
2.5
2.3
2.1
2,2223
medexp
0.1
0.0=0.4
0.1
0.6
0.4+0.7
0.5
mixed
16.4
18.4
17.6=2.1
11.2#5.3
13.6
13.2
hopper
medium
6.3
5.8=4.8
6.1
4.3
3.5+2.3
4.0
medexp
5.3
3.6
4.9*4.
1.6
1.3
1,541.9
mixed
2.4
1.9=0.5
2,3+0.6
2.4+0.8
1.9+0.5
2.340.6
halfcheetah
medium
0.9
0.6=0.1
0.8+0.1
0.9t0.3
0,6+0.4
0,840.4
medexp
1.3
0.0=0.6
1.0=0.7
1.3
0.0+0.6
1,0+0.7
C ADDITIONAL EXPERIMENTAL RESULTS
In this section, we show more experimental results support the conclusion of our paper.
C.1 RESULTS WITH DIFFERENT BASIC RL ALGORITHMS
To illustrate that our RTDiff indeed improves the performance general offline RL methods, here include more experiments involving Decision Transformer (Chen et al., 2021) and CQL (Kumar et al., 2020), which are representative sequence modeling baseline and model-free baseline. The results shown in Tables 10 and 11 illustrate that our method consistently improves the performance of different offline RI methods.
C.2 ORIGINAL PERFORMANCE REPORT
The performance increase reported in Section 5.1 is measured by the difference between the normal- ized score with data augmentation and the original normalized score without any data augmentation methods. The original results shown in Table 12.
15

---

Published as a conference paper at ICLR 2025
Table 12: Original normalized return of the methods we used in our paper on the D4RL Locomotion environment.
Environment
Data Type
CQL
TD3+BC
DT
IQL
mixed
73.1=
8..6
8.9
82.2 3.0
walker2d
medium
80.8 3.3
4.87
1..1
3.2
medexp
109.6
1.0.0
110.4
111.7
mixed
95.1 5.3
21.5
2.79
9.44
hopper
medium
3.8
6..4
67.6
3.8
medexp
95.1 - 5.3
101.2
107.1
1.7.4
mixed
45.0 0.3
4.6
3.5
44.5
halfchectah
medium
0.2
4..
4.32
48.3
medexp
9..6
90.
91.
9..7
C.3
MORE ABLATION STUDIES
Threshold of the OOD detector. We select the value of this threshold with the following method, using D4RL Locomotion environment as the representative environment: We use grid search to find the best choice of the hyperparameter, and then do cross-validation of the representative environment to ensure its robustness. After selecting the threshold, we directly apply this threshold all the environments threshold, we conduct a further ablation study on the environment maze2d. The results shown in Table 13 illustrate that this threshold dis 1.5 is reasonable across different environments.
used, without any further tuning. To demonstrate the robustness of
Table 13: Performance of maze2d environments under different thresholds. dism = 1.5 achieves the overall best performance compared with other threshold choices.
Threshold
1.0
1.3
1.5
2.0
maze2d-umaze umaze
7.2
12.5
12.3
4.1
maze2d-medium
5.7
7.9
8.3
2.8
maze2d-large
7.1
9.6
11.3
3.7
C.4 RESULTS OF OTHER BASELINES WITH OOD DETECTOR AND NOISE CONTROI.
To verify the effectiveness of different components in RTDifT, incorporate them with other base- lines to show the results.
Specifically, in our method, the OOD detector is used to control the length of the generated trajec- tories. This is crucial because we aim to generate trajectories that extend beyond the offline data distribution to provide new information to the agent, while ensuring they do not deviate too far, which could reduce their usefulness and increase risk. The effectiveness of our OOD detector is demonstrated in Tables 4 and 13. Our results show that the OOD detector outperforms any fixed- length generation strategy. Additionally, we observed that setting the threshold too high or too low negatively impacts performance.
To better understand why reverse synthesis is useful, we want to further show that this OOD detec- is useful for other synthesis methods. First of all, SynthER only performs transition-level synthesis, which inherently does not need to control the generation length. Also we want to em- phasize that with forward synthesis like ATraDiff, generating trajectories going out of the offline data distribution is more risky, as transitions going from inside to outside may lead to performance degradation. To validate this, we show the results of combining ATraDiff with an OOD detector of different thresholds. From the results shown in Table 14, we found that those forward synthesis methods derive limited benefits from the OOD detector, which validated the unique effectiveness of reverse synthesis in framework.
Our proposed noise control is largely independent of the specific generation method and can be applied to any data augmentation approach using inffusion-based framework. It demonstrates par- ticularly effective when the number of examples limited. Notably, this technique also applicable to SynthER and ATraDiff, demonstrated by the results shown in Table 15. We use SynthER and
16

---

Published as a conference paper at ICLR 2025
ATraDiff to generate both 1M data with and without the random generation. The results in Table 15 illustrate that the noise control can consistently improve the performance.
Table 14: Performance of ATraDiff with the OOD detector across different thresholds in the Maze2D environments. Forward synthesis cannot benefit from the OOD detector.
Method
RTDiff
ATraDiff
disM =
1.0
disM =
1.3
dis = 1.5
disM = 2.0
maze2d-umaze 12.3
7.1
7.4
7.0
6.3
6.0
maze2d-mediug
8.3
6.2
6.6
6.3
5.6
5,4
maze2d-large
11.3
7.8
7.7
7.8
7.4
7.1
Table 15: Performance of different baselines with and without noise control in the Maze2D envi- ronments. The number of generated samples is 1M. The noise control technique can consistently improve the performance of different data augmentation methods.
SynthER Noise control
SynthER w/o Noise control
ATraDiff w/ Noise control
ATraDiff w/o Noise control
maze2d-umaze
2.3
1.7
3.1
2.8
maze2d-medium
0.6
0.3
1.6
1.2
maze2d-large
7.3
6.3
8.1
5.2
C.5 QUANTITATIVE EVALUATION OF THE GENERATED SAMPLES
In this section, we conduct quantitative evaluation of the generated samples RTDiff. To measure the fidelity of the generated samples, we follow the previous works using two statisties: Marginal: Mean Kolmogorov-Smiing -mimmv (Massey Jr., 1951) and Correlation: Mean Correlation Sim- ilarity (Fieller al., 1957). To measure the model error of the generated samples, calculate the normalized error of the synthesized states and the real states after transition, which (T a) 8/2²
for transition (s, a, s²). The results are presented in Tables 16 and 17.
As expected, the results show that RTDiff does not aim to generate more realistic trajectories, but rather to produce more diverse samples that lie outside the distribution, thereby benefiting the RL performance. This is because RTDiff generates adaptive, longer trajectories compared with other baselines, attributed
proposed OOD detector and reverse synthesis model.
Therefore, while fidelity an important factor in assessing data generation in general, our focus in this paper is more on the "usefulness" of the generated data, specifically how it improves RL performance.
Table Performance comparison across different datasets for RTDiff, SynthER, and ATraDiff in terms of Marginal and Correlation metrics.
RTDiff
SynthER
ATraDiff
Dataset
Marginal t Correlation 1
Marginal t
Correlation t
Marginal f
Correlation t
hopper-medium
0.932
0.983
0.985
0.998
0.967
0.994
hopper-medexp
0.953
0.989
0.958
0.992
0.963
0.994
hopper-expert 0.941
0.985
0.934
0.982
0.953
0.991
C.6 COMPARISON WITH VANILLA DIFFUSION MODEL
To clarify that vanilla diffusion models alone cannot match the performance of our RTDiff, we conducted an additional ablation study. This study involved using diffusion models to generate verse trajectories without incorporating the other components of our method. The results, presented in Table 18, show that na've diffusion models significantly underperform compared to RTDiff.
17

---

Published as a conference paper at ICLR 2025
Table 17: Model errors for RTDiff, SynthER, and ATraDiff across different Maze2D environments.
RTDiff
SynthER
ATraDiff
maze2d-umaze
0.05
0.02
0.03
maze2d-medium
0.06
0,03
0.03
maze2d-large 0.11
0.07
0.08
Table 18: Overall results of D4RL Maze2D environments with vanilla diffusion models. The num- ber denotes the performance increase by the data augmentation method compared to the orig- inal result.
IQL (Kostrikov et al., 2022)
TD3+BC (Fujimoto Gu, 2021)
Environment
Data Type
RTDiff
DM(vanilla)
ROMI
RTDiff
DM(vanilla)
ROMI
umaze
8.3
4.342.6
5,4+2.4
10.2
9.3#44
9.64
maze2d
medium
3.3
2.3
2,1
9.8
8.9
9.4
large
14.3 3.3
9.0+3.1
8,1
7.7
4.3#26
3,5
Furthermore, when comparing the performance of vanilla diffusion models with ROMI, the improve- ments, if any, are inconsistent and relatively minor. Therefore, the results validate our contribution, demonstrating that it extends beyond the straightforward application of diffusion models.
D ADDITIONAL VISUALIZATIONS
In this section, show some additional visualization results in the D4RL (Fu al., 2020) Maze2D- umaze environment. We collect offline data consisting of trajectories going from green ball to red ball and use train the data synthesizer. The normal synthesis will generate trajectories entering the obstacle area, while reverse synthesis can avoid this problem.
18

---

Published as a conference paper at ICLR 2025
Figure 3: Top: The trajectory in the offline dataset. Middle: The trajectory of the offline RL agent with normal synthesis. Bottom: The trajectory of the offline RL agent with reverse synthesis.
19


Reviews:
[{"review_id": "9WxREqrg5W", "reviewer": "Reviewer_LuMS", "summary": "To address the distribution shift issue in offline reinforcement learning, this paper proposes a novel diffusion-based data augmentation technique, namely RTDiff, where rather than generating forward trajectories, the reverse one is synthesized, which is also the paper's main contribution. Furthermore, the performance of RTDiff is enhanced with the introduction of trajectory length control and noise management. Experimental results show the effectiveness of the proposed approach.", "questions": "1. Are the OOD detector and noise management incorporated with other data augmentation baselines, e.g., SynthER and ATraDiff?\n\n2.In Section 6, a very specific environment is adopted to show the advantages of reverse synthesis over normal synthesis, what is the reason/possible explanation that normal synthesis with OOD detector even produces significantly more In2Out than normal synthesis without OOD detector? (18.2%11.2)\n\n3. For other questions, please see questions raised in Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The ideas in this paper are interesting and novel, where to my knowledge, this is the first work that utilizes the concept of generating the reverse trajectories to address the distribution shift issue. The paper is clearly written and well-motivated.  The effectiveness of the proposed method is verified in various environments, and ablation studies are also conducted to validate the effectiveness of different components of the proposed method.", "weaknesses": "The main concern of the paper is whether the reverse synthesis can actually address the issue of distribution shift. As it is not state clearly whether the OOD detector is incorporated with other data augmentation baselines, e.g., SynthER and ATraDiff, it is not very certain whether the better performance achievement of RTDiff is due to the reverse synthesis or the using of the OOD detector. In Section 6, an analysis is conduct by using a simple illustrative environment to show why reverse synthesis avoids issues present in normal synthesis. However, on the one hand, it is better to directly use an environment used in the experiments for the analysis, on the other hand, it is confusing why the reverse synthesis generates trajectories that move from the dangerous area to the upper or lower areas while normal synthesis generates trajectories that start from the lower area and enter the middle dangerous area. Do these two approaches both start from a state in the offline data, and then generate the trajectories in different ways? If the OOD detector is used in the normal synthesis, can the dangerous areas also be avoided? Moreover, the theoretical contribution of the proposed is not very significant. \nIf the concerns can be addressed, I would like to raise my score.", "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": "v1tYbXlH12", "reviewer": "Reviewer_FUCB", "summary": "This paper proposes RTDiff, a novel diffusion-based data augmentation technique that synthesizes trajectories in a reverse direction. Such reverse generation naturally mitigates the risk of overestimation by ensuring that the agent avoids planning through unknown states. RTDiff also introduces some other tricks including flexible trajectory control and noise management to improve sythesis quality. Emprirical results show the advantage of reverse generation over forward generation.", "questions": "1. Can the authors discuss in detail the differences between this paper and previous similar works (like ROMI)?\n2. This method is built on IQL, TD3BC, and CQL. Have there been any adjustments to the hyperparameters of these methods after using data augmentation?\n3. Is it possible to conduct a quantitative evaluation of the quality of the generated trajectories? For example, assessing the model error of the generated trajectories, etc.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The article is well-written, with clear expression and logic, effectively reflecting the main argument.\n2. The experiments are very comprehensive. The authors validate the effectiveness of their method across a series of tasks, including both proprioceptive observations and visual observations.", "weaknesses": "1. I believe the authors miss a key related work (called ROMI [1]), which is the first to propose using reverse trajectory generation in the field of offline reinforcement learning. The motivation described for reverse trajectory generation in these two work is also very similar. Therefore, considering that this paper is neither the first to propose the use of reverse trajectory generation in offline RL nor the first to use diffusion for data augmentation, I would say the novelty of this paper is quite limited.\n2. I think this paper lacks comparisons with model-based offline RL methods in the experimental section. According to my understanding, using a diffusion model for data augmentation essentially falls under the same category as previous model-based offline RL methods. For instance, MOPO [2] essentially generates a batch of synthetic samples to supplement the original samples. Therefore, the authors should compare some model-based offline RL methods, such as MOPO, RAMBO [3], etc.\n\n[1] Wang et al. \"Offline Reinforcement Learning with Reverse Model-based Imagination\" (NeurIPS'21)\n[2] Yu et al. \"MOPO: Model-based Offline Policy Optimization\" (NeurIPS'20)\n[3] Rigter et al. \"RAMBO-RL: Robust Adversarial Model-Based Offline Reinforcement Learning\" (NeurIPS'22)", "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": "QQJAPT9YxU", "reviewer": "Reviewer_NMzW", "summary": "This paper introduces a novel diffusion-based data augmentation method, RTDiff, for offline reinforcement learning. First, RTDiff mitigates the data distribution shift issue present in previous data augmentation methods by generating reverse trajectories instead of forward ones. Second, RTDiff trains an out-of-distribution (OOD) detector to truncate the OOD segments of generated trajectories, further enhancing sample authenticity. Finally, the authors propose a new noisy control method to improve sample generation efficiency. Experimental results validate the effectiveness and efficiency of RTDiff and different components across both vector-based and pixel-based tasks.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is well-written and easy to follow.\n2. RTDiff introduces reverse trajectory generation, an OOD detector, and a noisy control method to achieve efficient and reliable sample generation. These approaches are intuitively sound.\n3. The authors conduct extensive benchmark and ablation experiments. Experimental results demonstrate that RTDiff outperforms previous baselines on both vector-based and pixel-based tasks, and each component—reverse trajectory generation, the OOD detector, and noisy control—exhibits its effectiveness.", "weaknesses": "1. Reverse trajectory generation is not new to offline RL, and the paper lacks a discussion and experimental comparison of prior works, such as [1].\n2. The paper lacks a clear and well-reasoned explanation of the issues with previous data augmentation methods. In lines 88-90, this paper claims that previous data augmentation methods suffer from the data distribution shift, potentially leading to value overestimation. However, since all of these works use offline RL algorithms, I think data distribution shift is not the key issue. On the contrary, data augmentation is only effective when the generative model can produce samples that differ from the training data.\n3. Generated data fidelity is a more critical factor. The paper lacks a quantitative evaluation of the fidelity of generated samples.\n\n[1] Offline Reinforcement Learning with Reverse Model-based Imagination. NeurIPS, 2021.", "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": "0YfwUQhJ7g", "reviewer": "Reviewer_exWn", "summary": "Traditional offline reinforcement learning methods often introduce conservative biases to limit exploration to familiar regions, but this can restrict an agent's ability to generalize. While recent approaches use generative models to expand offline datasets, they can overestimate synthesized data, particularly when it includes out-of-distribution samples. RTDiff is introduced to address this, which is a diffusion-based data augmentation technique that creates trajectories in reverse, moving from unknown to known states. This reverse approach reduces the risk of overestimation by ensuring the agent avoids planning through unfamiliar regions. It also supports generating longer trajectories, utilizing diffusion models effectively while maintaining reliability. RTDiff further optimizes the process with flexible trajectory length control and noise management to improve generation efficiency.", "questions": "Please see weaknesses.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The reverse generation method, which has been empirically verified, effectively reduces the augmentation of data in risky regions. This concept is intuitively illustrated in the diagram.", "weaknesses": "1.Previous research [1] introduced a reverse data generation approach using a transition model. In the current work, the vanilla model is replaced with a diffusion model, yet the fundamental concept remains unchanged, limiting the overall contribution.\n\n2.The explanation relies on an intuitive diagram, but it would be more effective to demonstrate several specific cases. Identifying which states are risky.  For example, some states are prone to be overestimated and easily generated by a forward model, but the reverse model effectively avoids generating them.\n\n3.Minor errors are present, such as in line 156, where \"dat\" should be corrected to \"data.\"\n\n4.Environments based on visual data should be included in the analysis.\n\n[1] Wang, Jianhao, et al. Offline reinforcement learning with reverse model-based imagination.", "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)