Title: GTA: Generative Trajectory Augmentation with Guidance for Offline Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Offline Reinforcement Learning

Data Augmentation

Diffusion Models.
Abstract: Offline Reinforcement Learning (Offline RL) presents challenges of learning effective decision-making policies from static datasets without any online interactions. Data augmentation techniques, such as noise injection and data synthesizing, aim to improve Q-function approximation by smoothing the learned state-action region. However, these methods often fall short of directly improving the quality of offline datasets, leading to suboptimal results. In response, we introduce GTA, Generative Trajectory Augmentation, a novel generative data augmentation approach designed to enrich offline data by augmenting trajectories to be both high-rewarding and dynamically plausible. GTA applies a diffusion model within the data augmentation framework. GTA partially noises original trajectories and then denoises them with classifier-free guidance via conditioning on amplified return value. Our results show that GTA, as a general data augmentation strategy, enhances the performance of widely used offline RL algorithms across various tasks with unique challenges. Furthermore, we conduct a quality analysis of data augmented by GTA and demonstrate that GTA improves the quality of the data. Our code is available at https://github.com/Jaewoopudding/GTA

GTA: Generative Trajectory Augmentation with Guidance for Offline Reinforcement Learning

Abstract

Offline Reinforcement Learning (Offline RL) presents challenges of learning effective decision-making policies from static datasets without any online interactions. Data augmentation techniques, such as noise injection and data synthesizing, aim to improve Q-function approximation by smoothing the learned state-action region. However, these methods often fall short of directly improving the quality of offline datasets, leading to suboptimal results. In response, we introduce GTA , Generative Trajectory Augmentation, a novel generative data augmentation approach designed to enrich offline data by augmenting trajectories to be both high-rewarding and dynamically plausible. GTA applies a diffusion model within the data augmentation framework. GTA partially noises original trajectories and then denoises them with classifier-free guidance via conditioning on amplified return value. Our results show that GTA, as a general data augmentation strategy, enhances the performance of widely used offline RL algorithms across various tasks with unique challenges. Furthermore, we conduct a quality analysis of data augmented by GTA and demonstrate that GTA improves the quality of the data. Our code is available at https://github.com/Jaewoopudding/GTA

1 Introduction

Learning a decision-making policy through continual interaction with real environments is challenging when online interaction is costly or risky. Offline Reinforcement Learning [Offline RL, 1] emerges as a solution, focusing on training effective decision-making policies with static dataset gathered through an unknown policy [2]. However, offline data often does not provide enough coverage of state-action space, resulting in extrapolation error addressed as overestimation of Q-value [3, 4]. To mitigate extrapolation error, many previous works rely on adding explicit regularization terms [3-10] and have shown significant progress.

Moving beyond these mainstream methodologies, there exists an underexplored approach, data augmentation methods: traditional data augmentation and generative data augmentation. Traditional augmentation methods [11, 12] inject minimal noise to the state, preserving environment dynamics. Generative data augmentation [13, 14] builds a data synthesizer with a generative model to upsample offline data. Employing generative data augmentation approaches broadens the support of the data, thereby improving Q-function approximation [12].

While previously proposed data augmentation methods enhance the performance of offline RL, they still have limitations. As illustrated in Figure 1, traditional augmentation methods have difficulties in discovering novel states or actions beyond the existing offline data, limiting their role only to smoothing small local regions of the observed state. In the case of generative data augmentation, the reward distribution of generated data is constrained by the support of offline data, resulting in the

(a) RAD, S4RL

○

○

○

(b) SynthER

○

○

○

(c) GTA (Ours)

×

×

× ×

×

×

×

×

×

×

×

×

×

×

×

×

○

○

○

○

○

×

×

× ×

×

×

×

×

○

×

×

High-Rewarding Distribution

×

Augmented Datapoint

○

Original Datapoint

Original Distribution

generation of suboptimal data. These issues make the existing methods not align with the objective of offline RL, which aims to learn the most effective decision-making policy from the static dataset.

We introduce Generative Trajectory Augmentation (GTA), a novel approach that applies the conditional diffusion model to data augmentation, aiming to address the aforementioned limitations. GTA is designed to generate novel and high-rewarding trajectories while minimizing degradation of dynamic plausibility. Our approach consists of three main stages: (1) Training a conditional diffusion model that generates trajectory-level data , (2) Augmenting the offline data via partial noising with diffusion forward process and denoising with amplified return guidance , and (3) Training any offline RL algorithm with augmented data.

We train a conditional diffusion model that approximates the conditional distribution of trajectory given its return. Once the diffusion model is trained, we sample trajectories from the offline dataset and partially noise the trajectory with the diffusion forward process. Then, we denoise the noised trajectory with amplified return guidance , directing trajectories to the high-rewarding region. To this end, we can orthogonally integrate data augmented by GTA into any offline RL algorithm without any modification.

Our contributions . In this paper, we introduce GTA, a novel data augmentation framework that utilizes a conditional diffusion model to generate high-rewarding, novel, and dynamically plausible data. Through extensive experiments on commonly studied benchmarks [15, 16], we demonstrate that GTA significantly improves performance across various tasks with unique challenges, such as sparse reward tasks and high-dimensional robotics tasks. We thoroughly examine the impact of core design components of GTA, trajectory generation , partial noising with diffusion forward process and denoising with amplified return guidance . Furthermore, we assess the GTA augmented dataset with data quality metrics, proving its alignment with the objective of offline RL. These findings underscore the capability of GTA to efficiently augment trajectories, resulting in high-quality samples and improving the performance of offline RL algorithms.

2 Related Works

Data Augmentation for Reinforcement Learning . In Reinforcement Learning (RL), data augmentation enhances sample efficiency and improves Q-function approximation. In pixel-based RL, methods like CURL [17], RAD [11], and DrQ [18] have leveraged image augmentations such as cropping and translation to address sample-efficiency of RL. For the proprioceptive observations, S4RL [12] introduces variations into states by adding small-scale noise or adversarial gradients of Q-function under the similarity prior, that similar states yield similar rewards. AWM [19] improves robustness of augmented data by applying simple transformations to the learned dynamics model, enabling zero-shot generalization to unseen dynamics without requiring multiple test-time rollouts.

Recent strides in generative models have led to their adoption in generative data augmentation, as seen in SynthER [13], MTDiff-S [14], and PGD [20]. Among these, GTA, MTDiff-S, PGD, and SynthER all generate synthetic data using diffusion models but with distinct approaches. While GTA, MTDiff-S, and PGD generate at the trajectory level, MTDiff-S is designed for multi-task settings, focusing on generating trajectories for unseen tasks. PGD generates synthetic trajectory with classifier guidance of policy, similar to model-based offline RL. GTA, on the other hand, focuses on generating high-rewarding trajectories using return guidance. Focusing on single tasks, GTA and SynthER differ in what they generate. SynthER generates individual transitions, in contrast to the trajectory-level generation of GTA. This makes GTA compatible with any kind of offline RL model. Additionally, GTA initializes augmentation from the original trajectory, utilizing a conditional diffusion model, while SynthER starts augmentation from the Gaussian noise without conditional guidance.

Diffusion Models for Offline Reinforcement Learning . In focus on the diffusion planners, Diffuser [21] devises the RL problem as a trajectory-level generation with cumulative reward classifier guidance. Following this, the Decision Diffuser [22] replaces the classifier guidance of the Diffuser with classifier-free guidance [23]. Adaptdiffuser [24] alternates between generating trajectories with diverse reward functions and training the model using self-generated trajectories. However, diffusion planners require extensive time to sample actions, making their practical application challenging. GTA shifts the computational burden of the diffusion model from the decision-making step to the data-preparation step. This transfer allows GTA to leverage the advantages of diffusion models while avoiding extensive time costs during decision-making.

Model Based Offline Reinforcement Learning . Model-based Offline RL focuses on learning environment dynamics and reward functions from data, creating simulated trajectories to train both the critic and policy [25-28]. While both GTA and model-based offline RL generate synthetic trajectories, in the case of GTA, a key difference is that data generation and policy learning are separated and not done in an alternative cycle. This separation ensures that the quality of generated data is not affected by the training progress of critics or policies. It also mitigates the accumulating errors associated with single-step dynamics rollouts [22].

3 Preliminaries

3.1 Offline Reinforcement Learning

Reinforcement Learning (RL) is modeled with the Markov decision process (MDP) described by the tuple ( S , A , T , R , γ ) , consisting of state space S , action space A , transition function T : S×A → S , reward function R : S × A × S → R , and discount factor of future reward γ ∈ [0 , 1) [29]. At each timestep t , the agent selects an action a t according to the policy π given the state s t . Consequently, the agent receives a reward r t for the action a t taken in the state s t , leading to the next state s t +1 . The goal of RL is to learn policy π ∗ , which maximizes expected discounted return, J ( π ) = E π [ ∑ ∞ t =0 γ t r t ] .

In the offline RL setting, we can only have access to the fixed dataset D , which has been collected using unknown behavior policy π β . With insufficient and suboptimal offline data, offline RL aims to learn effective decision-making policy that surpasses the behavior policy.

3.2 Diffusion Models

Score-based diffusion models . Diffusion models are a family of generative models that approximate the data distribution p ( x ) with p θ ( x ) . When considering the data distribution as p ( x ) and the standard deviation of noise as σ , then the distribution of data with added noise is denoted as p ( x ; σ ) . The diffusion reverse process involves sequentially denoising from noise x K that randomly sampled from a Gaussian distribution N (0 , σ 2 max I ) , following a noise level sequence σ K = σ max > σ K -1 > · · · > σ 0 = 0 . Consequently, the endpoint x 0 of this process aligns with the original data distribution.

Considering the probability flow ordinary differential equations (probability flow ODE), noise is continuously added to data during the forward process and reduced in the reverse process. By scheduling the noise level at time k , represented as σ ( k ) , the reverse probability flow ODE is formulated as follows [30]:

where ∇ x log p ( x ; σ ( k )) denotes the score function, signifies the direction towards the data for a given noise level, and the dot represents the time derivative. The score function ∇ x log p ( x ; σ ( k )) is trained via denoising score matching. The denoising score matching loss for given denoiser D θ ( x ; σ ) is given by

When the denoiser D θ ( x ; σ ) is optimally trained, the score is calculated as

We sample data via solving Equation (1) with the learned denoising network.

StageA:TrainDiffusionModel

StageB:AugmentTrajectoryData

Stage C:Train OfflineRL

rnoised

Trajectory Generation

Eg(rnoised y)

Eg(rnoised, 0)

TD3BC

IQL

Noising k

Partial Noising

T

AmplifiedReturnConditioning

Denoising Network

rnoised~N(-｜°,(μ·K)I)

y=α

puosun9.M-puo3·(I+M)=θ2

CQL

MCQ

L(b) =Ile -Ee(t*k)I12

i=t

02

02

pasou

rdenoised or T0

ReweightedSampling

OfflineDataset

AugmentedDataset

D

DA

Conditional score-based diffusion model . In the domain of conditional diffusion models, two primary strategies are recognized: classifier guidance [31] and classifier-free guidance [23]. Classifier free guidance sets its guidance distribution ˜ p θ as ˜ p θ ( x | y ) ∝ p θ ( x | y ) · p θ ( y | x ) w . Subsequently, using the implicit classifier p θ ( y | x ) ∝ p θ ( x | y ) /p θ ( x ) and the equivalence relationship between score matching and denoising process, described as ∇ x log p θ ( x | y ) ∝ ϵ θ ( x , y ) [32, 23], the classifier free guidance score ˜ ϵ θ is formed as follows:

where w controls the strength of the guidance. The training objective of the classifier free guidance is to concurrently train the conditional score function and the unconditional score function as follows:

where y is a condition and β ∼ Bern( λ ) is a binary variable with dropout rate λ of condition y .

4 Method

In this section, we introduce GTA , G enerative T rajectory A ugmentation, which leverages the conditional diffusion model to generate high-rewarding, novel, and dynamically plausible trajectories for augmentation. As shown in Figure 2, our method is divided into three stages: (1) Train a conditional diffusion model at the trajectory level. (2) Augment trajectories using a partial noising and denoising framework with amplified return guidance . (3) Train any offline RL algorithm with the augmented dataset. We will discuss how each component works and their roles in effective data augmentation.

4.1 Stage A: Train Diffusion Model

Trajectory-level generation . The diffusion model for GTA is designed to generate the subtrajectory τ . This subtrajectory is a consecutive transition sequence of the state, action, and reward sampled from a trajectory ( s 1 , a 1 , r 1 , ..., s T , a T , r T ) . We represent subtrajectory τ with horizon H as follows:

We train a conditional diffusion model to approximate the conditional distribution p ( τ | y ( τ )) with offline dataset, where the condition y ( τ ) = ∑ T i = t γ i -t r i denotes the sum of the discounted return. The parameter θ is updated to maximize the expected log-likelihood of the data:

Using a diffusion model to generate trajectories offers multiple advantages. First, the diffusion model leverages sequential relationships between consecutive transitions while generating trajectories, allowing it to minimize the degradation of dynamic plausibility. Second, the diffusion model captures long-term transition dynamics, which is beneficial for environments with sparse rewards.

Diffusion Model Implementation . Integrating sequential dependencies within the subtrajectory into the model architecture is essential for trajectory-level generation. Therefore, we choose a denoising network that combines Temporal-Unet [21] with MLP-mixer [33] to exploit both local and global sequential information. We adopted the Elucidated Diffusion Model [30], known for its powerful performance in recent work [13]. The implementation details are elaborated in Appendix A.2.

4.2 Stage B: Augment Trajectory-level Data

We propose a novel approach to augment trajectories, partial noising and denoising framework with amplified return guidance. Partial noising modifies the original trajectory using the forward process of the diffusion model, thus providing exploration opportunities. The exploration level is adjusted by the nosing ratio µ . Following this, denoising with amplified return guidance refines the noised trajectory. During the denoising process, we guide the trajectory towards high-rewarding regions, promoting the exploitation of learned knowledge about the environment. We introduce the multiplier for conditioned return α to control the exploitation level. Figure 3 outlines the principle of data augmentation via partial noising and denoising framework.

Partial noising with forward process .

Let τ = τ 0 represent the original trajectory, and k denotes the diffusion timestep. We introduce a noising ratio µ (0 < µ = k K ≤ 1) to determine the extent of erasing information of the trajectory for exploration. We add noise to the original trajectory τ 0 , creating a noised trajectory denoted as τ µ · K ∼ N ( τ ; τ 0 , σ ( µ · K ) 2 I ) . The parameter µ controls the level of exploration. Small µ results in minimal exploration, thereby preserving much of the original information of the trajectory. Large µ facilitates broader exploration, potentially leading to generating novel trajectories while losing original information significantly.

Denoising with amplified return guidance .

After noising the trajectory, we reconstruct the trajectory with classifier-free guidance to push the trajectory towards the high-rewarding region to enhance optimality. We introduce amplified return guidance, which sets the conditioning value as the multiplied return of the original trajectory. It can prevent adverse effects that occur when significantly larger return values are conditioned during generation. Further analysis on this phenomenon is elaborated in Appendix F.1.

The amplified return is formally defined as follows:

where α is the control parameter for the exploitation level. We set α > 1 to make the conditioning value higher than the return of the original trajectory. α close to 1 induces mild conversion towards a high-rewarding trajectory region, while large α promotes significant drift, substantially exploiting the diffusion model. Our partial noising and denoising framework can be summarized as follows:

Noising Process

Guided Denoising Process

High Rewarding Region

Originality Preserving Region

!

!

!

!

!

!

⋅#

!

!

"

⋅#

!

!

#

⋅#

" #

!

"

" #

#

"

" #

$

"

Original Trajectory

4.3 Stage C: Offline RL Policy Training

The final stage of GTA is utilizing high-quality trajectories generated through previous stages for policy training. The trajectories generated by GTA can seamlessly integrate with existing offline RL methods. By introducing GTA, we can orthogonally integrate the expressiveness and controllability of diffusion models with existing offline RL algorithms.

4.4 Additional Technique: Reweighted Sampling

To concentrate on samples in high-rewarding regions, we additionally adopt a reweighting strategy [34, 35] during the sampling process. This approach prioritizes the sampling of subtrajectories with higher returns. Consequently, augmented trajectories are more densely distributed in high-rewarding regions. The complete specification of our reweighting strategy, including implementation details and parameter configurations, is provided in Appendix A.3.

5 Experiments

In this section, we present extensive experiments conducted across commonly studied benchmarks [15, 16] to evaluate the augmentation capability of GTA. Through experiments, we aim to address the following questions: 1) How much performance gain does GTA exhibit across various types of tasks? 2) What impacts do the design elements of GTA have on performance? 3) Are trajectories augmented via GTA high quality? 4) What is the preferable µ and α setting for the new tasks?

5.1 Experimental Setup

Datasets and environments . We demonstrate the versatility of GTA on the D4RL benchmark [15], from the standard continuous control tasks to the challenging tasks with unique difficulties. We validate that GTA enhances offline RL under the sparse rewards environments and overcoming high-dimensional, complex robotics tasks. Additionally, we demonstrate that GTA is effective with human-demonstrated, small-scale datasets, showing its adaptability to realistic settings. Finally, we extend GTA to pixel-based environments, specifically within VD4RL benchmark [16].

Data augmentation baselines . We compare GTA with existing augmentation methods. For traditional augmentation, we choose S4RL [12], which introduces Gaussian noise into states. In the domain of generative augmentation, we choose SynthER [13], which employs an unconditional diffusion model for transition-level generation.

Offline RL algorithms . To demonstrate the general efficacy of GTA on proprioceptive observations, we select four widely used offline RL algorithms: TD3BC [7], CQL [8], IQL [36], and MCQ [4]. For particularly challenging tasks, such as Maze2d, Antmaze, Adroit, and FrankaKitchen, we employ IQL, which provides all hyperparameter configurations across all tasks and consistently demonstrates stable performance. For pixel-based observations, we utilize DrQ+BC as the baseline algorithm, as established in the VD4RL benchmark by Lu et al. [16].

Data quality metrics . We propose data quality metrics to analyze whether GTA provides high-quality data. Following prior works [13, 37], we conduct a thorough analysis with three metrics: oracle reward, novelty, and dynamic MSE. Oracle reward, computed with the true reward of generated data, represents optimality. Novelty measures the ability of the augmentation method to discover novel states and actions not existing in the offline data. Finally, dynamic MSE evaluates how well the generated trajectories adhere to the dynamics of the environment. Formally, dynamic MSE and novelty are defined as follows:

where D A represent the augmented dataset and D is original offline dataset. ( s, a, r, s ′ ) denotes a single transition, and f ∗ denotes the true dynamic model of the environment.

5.2 Benchmark Evaluations

We conducted experiments to evaluate whether GTA can provide a performance boost when applied alongside existing offline RL algorithms in various tasks with its own unique challenges.

Gym locomotion . We experimentally demonstrate on Table 1 that offline RL algorithms with GTA outperform all other baselines across all algorithms on the average score with a statistically significant margin. This result highlights the versatility of GTA with various offline RL algorithms. We provide p -value for our experiments to validate the statistical significance of the results in Appendix F.2.

Maze tasks . To assess GTA in sparse reward tasks, we test GTA on Maze2d and AntMaze tasks. Table 1 reveals that GTA notably enhances the performance of IQL policy on the sparse reward tasks while SynthER often degrades the performance. This result suggests two insights. First, the trajectory-level generation of GTA helps capture long-term dynamics, effectively leveraging that information while augmenting sparsely rewarded trajectories. Second, augmenting the dataset with high-rewarding trajectories strengthens the goal-reaching ability by enriching the demonstrations with more successful trajectories.

Complex robotics tasks . We evaluate the effectiveness of the GTA on realistic, high-dimensional, challenging control tasks. Adroit-human datasets comprise 25 human-generated trajectories. The FrankaKitchen dataset, which involves multitasking behavior trajectories, requires a generalization ability to stitch trajectories. According to the results in Table 2, GTA effectively boosts the performance in high-dimensional, complex robotics tasks. The results on the Adroit-human dataset demonstrate that GTA effectively augments small-scale human demonstration data through the expressiveness of the diffusion model. Additionally, the performance improvements in kitchen tasks indicate that GTA aids offline RL in trajectory stitching [21].

Pixel-based Observations . We extend our approach to pixel-based observations by following the experimental setup detailed in [13]. Initially, we pretrain the policy and extract its visual encoder. Subsequently, we project offline pixel-based observations into the embedding space, followed by augmenting these embedded states with GTA. More detailed experiment setups are elaborated in Appendix C.4. The results presented in Table 3 show performance improvements on DrQ+BC across the environments and dataset qualities. It indicates that GTA can be further extended to pixel-based observations beyond proprioceptive observations.

μ=0.25

μ =0.5

μ =0.75

μ =1.0

65

110

Amplified conditioning

Normalizedscore

60

Normalized score

100

06

density

00

Fixed conditioning

Unconditioning

55

80

Probability

0.03

Offine data

50

70

0.02

60

45

0.01

50

40

40

0.00

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0

50

100

150

200

250

Train steps

×106

Train steps

×106

Oraclerewards sum of subtrajectory

(a)

(b)

(c)

5.3 Ablation Studies

We carry out thorough ablation studies to assess the effectiveness of each component within GTA. These studies are evaluated based on the performance scores of D4RL tasks and the aforementioned data quality metrics.

What are the benefits of trajectory-level generation? To investigate the impact of utilizing sequential relationships through trajectory-level generation, we examine performance differences as the trajectory length varied. We find that when H = 1 , the dynamic MSE becomes five times or more higher compared to trajectory-level generation with H longer than 16. Additionally, when training TD3BC with augmented data, the normalized score increased by 21% with trajectory generation compared to transition generation, where H = 1 . This result highlights sequential relationship between consecuive transitions in trajectory does help generation, minimizing degradation of dynamic plausibility. Detailed experiment results are in Appendix E.1

Is preserving original information via partial noising essential? We conduct experiments by varying the noising ratio on datasets of different optimality to explore the effect of the exploration on performance. In Figure 4a halfcheetah-medium, higher noise levels enhance performance, while in Figure 4b halfcheetah-medium-expert, optimal performance occurs at a noise level of µ = 0 . 25 , with declines at higher levels. We speculate that if the original data already contains high-rewarding trajectories, even minimal modifications can be highly effective, and excessive exploration might lead to unexpected outcomes. Therefore, preserving information of original trajectory by partial noising is crucial in achieving effective data augmentation while avoiding potential negative impacts from excessive exploration. Additionally, the terminal signal, which may be very sparse, could be lost significantly if the original information is not preserved. We further conduct experiments about the impact of partial noising on preserving terminal state information in Appendix F.3.

Does amplified return guidance improve the optimality? Figure 4c illustrates the superiority of the amplified return guidance on augmentation with respect to unconditioning [13], and fixed conditioning [22]. Augmented trajectories without any condition tend to replicate the original reward distribution, thereby failing to enhance data optimality. Fixed conditioning concentrates rewards near the maximum return offline dataset. Notably, amplified return conditioning demonstrates a superior ability to cover a broader range of rewards, especially on unseen high rewards, while leading on the average reward and D4RL score. Further analysis of the impact of the conditioning method on data quality can be found in Appendix F.1.

Oracle Reward

Novelty

Dynamic MSE

halfcheetah-m-e

halfcheetah-m-e

halfcheetah-m-e

hopper-Rr

hopper-Rr

halfcheetah-m-r

halfcheetah-m-r

halfcheetah-m-r

hopper-m-r

hopper-m-r

halfcheetah-m

halfcheetah-m

halfcheetah-m

hopper-m-e

hopper-m-e

hopper-m-e

walker-m-e

walker-m-e

walker-m-e

walker-m

walker-m

walker-m

walker-m-r

walker-m-r

walker-m-r

S4RL

SynthER

GTA

5.4 Data Quality Analysis

We introduce quality metrics specifically suitable for offline data analysis in Section 5.1. We compare GTA with other augmentation methods regarding not only how much it expands the coverage of offline data but also its optimality and dynamic MSE. Figure 5 presents the relative scale of each metric across baseline augmentation methods on gym locomotion tasks. Data augmented by GTA shows significantly higher oracle reward and novelty while maintaining a comparably low dynamic MSE. These outcomes indicate that the conditional diffusion model of GTA expands data coverage under dynamic constraints and discovers novel states and actions with higher rewards, effectively integrating environmental knowledge. We add further analysis in Appendix F.4.

5.5 Guideline for Selecting µ and α

In this section, we propose a guideline for selecting hyperparameters based on dataset quality and environment characteristics. First, when the known information about the dataset is limited and minimizing exploration risk is essential, a low exploration and exploitation approach can enhance performance without losing stability. As illustrated in Table 4, GTA with low exploration and exploitation ( µ = 0 . 25 , α = 1 . 1 ) surpasses baseline methods in gym locomotion tasks.

However, if we know that the offline dataset is suboptimal and there exists considerable room for improvement, a higher exploration and exploitation setting can yield superior performance by uncovering a broader range of high-quality behaviors. Table 5 shows that GTA with ( µ = 0 . 5 , α = 1 . 3 ) and ( µ = 0 . 75 , α = 1 . 3 ) significantly boosts performance on suboptimal datasets. These findings indicate that while conservative parameters with low µ and α offer stable gains, we can elevate values of µ and α by leveraging prior knowledge about dataset quality to achieve superior outcomes. Notably, GTA consistently outperforms baselines across all configurations without dataset-specific hyperparameter tuning.

5.6 Futher Experimental Results

We conduct extensive additional experiments to explore the further potential of the GTA.

GTA under realistic settings . We examined the augmentation ability of GTA for two realistic settings: one with a mixed dataset consisting of few expert datasets and the majority of low-performing trajectories in Appendix E.2, and the other with a small amount of offline dataset in the Appendix E.3. In both cases, We observed that GTA significantly enhances the performance, improving sample efficiency of the offline RL.

Larger batch size and training epochs . Since GTA offers more training data, it is worth exploring whether increasing batch size or training epochs can lead to additional performance gain. The results in Appendix G.1 demonstrate that we can enhance the performance by increasing batch size and training epochs while baselines do not benefit from the batch size and training epochs increment.

Extending to sequence modeling approach . We proved that GTA can be applied to sequence modeling approaches such as Decision Transformer [38] as GTA augments data at the trajectory level. We confirmed that GTA improves the performance of the sequence modeling approach, especially in sparse reward tasks. The detailed experiment setup and results are in Appendix E.4.

Sensitivity tests . Weconduct experiments on the sensitivity of parameters α in Appendix G.2 and µ in Appendix G.3. The sensitivity of GTA with the amount of augmented dataset is in Appendix G.4. Our results indicate that GTA maintains similar performance beyond one million augmented transitions, demonstrating robustness to the number of transition samples.

Additional ablations on design choices . We explore the impact of reweighted sampling introduced on Section 4.4 in Appendix E.5. For exploring alternative design choices, we adopt conditioning-byinpainting as an alternative conditioning approach instead of classifier-free guidance in Appendix E.6. We also test the model-based approach by replacing the reward of the generated trajectory with the prediction of the reward proxy model in Appendix E.7.

6 Conclusion

We propose GTA : Generative Trajectory Augmentation, a novel generative data augmentation approach for enriching offline data with high-rewarding and dynamically plausible trajectories. GTA combines data augmentation with the diffusion sampling process by partially adding noise to original trajectories and subsequently denoising them under amplified return guidance. Our extensive experiments on 31 datasets across 9 environments exhibit considerable performance improvements in four distinct offline RL algorithms, demonstrating the versatility and effectiveness of the GTA. We show that GTA successfully creates high-quality datasets from the sub-optimal offline dataset, which leads to learning effective decision-making policy.

Limitations and future works . In the GTA framework, we do not train additional transition model and reward model for simple augmentation framework. Therefore, in tasks where dynamic violations have a critical impact, we may not expect a significant performance boost. However, as shown in our experiments, dynamic violations are generally minimal. Additionally, while the main focus of GTA is on the offline setting, future work could explore off-to-online and online settings.

References

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the-art , pages 45-73. Springer, 2012.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing offpolicy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Jiafei Lyu, Xiaoteng Ma, Xiu Li, and Zongqing Lu. Mildly conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 35:1711-1724, 2022.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning , pages 2052-2062. PMLR, 2019.

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems , 34:20132-20145, 2021.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:11791191, 2020.

Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang, Ruslan Salakhutdinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. arXiv preprint arXiv:2105.08140 , 2021.

Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 , 2020.

Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems , 33:19884-19895, 2020.

Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4RL: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In Conference on Robot Learning , pages 907-917. PMLR, 2022.

Cong Lu, Philip J Ball, and Jack Parker-Holder. Synthetic experience replay. arXiv preprint arXiv:2303.06614 , 2023.

Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. arXiv preprint arXiv:2305.18459 , 2023.

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.

Cong Lu, Philip J Ball, Tim GJ Rudner, Jack Parker-Holder, Michael A Osborne, and Yee Whye Teh. Challenges and opportunities in offline reinforcement learning from visual observations. arXiv preprint arXiv:2206.04779 , 2022.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning , pages 5639-5650. PMLR, 2020.

Denis Yarats, Rob Fergus, and Ilya Kostrikov. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Representations, ICLR 2021 , 2021.

Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts. Augmented world models facilitate zero-shot dynamics generalization from a single offline environment. In International Conference on Machine Learning , pages 619-629. PMLR, 2021.

Matthew Thomas Jackson, Michael Tryfan Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson, and Jakob Foerster. Policy-guided diffusion. arXiv preprint arXiv:2404.06356 , 2024.

Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdiffuser: Diffusion models as adaptive self-evolving planners. arXiv preprint arXiv:2302.01877 , 2023.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems , 33:14129-14142, 2020.

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems , 33:21810-21823, 2020.

Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems , 34:28954-28967, 2021.

Marc Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. Advances in neural information processing systems , 35:1608216097, 2022.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems , 35:26565-26577, 2022.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems , 34:8780-8794, 2021.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlpmixer: An all-mlp architecture for vision. Advances in neural information processing systems , 34:24261-24272, 2021.

Aviral Kumar and Sergey Levine. Model inversion networks for model-based optimization. Advances in Neural Information Processing Systems , 33:5126-5137, 2020.

Siddarth Krishnamoorthy, Satvik Mehul Mashkaria, and Aditya Grover. Diffusion models for black-box optimization. arXiv preprint arXiv:2306.07180 , 2023.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

Minsu Kim, Federico Berto, Sungsoo Ahn, and Jinkyoo Park. Bootstrapped training of score-conditioned generator for offline design of biological sequences. arXiv preprint arXiv:2306.03111 , 2023.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov. Corl: Research-oriented deep offline reinforcement learning library. arXiv preprint arXiv:2210.07105 , 2022.

Yihao Sun. Offlinerl-kit: An elegant pytorch offline reinforcement learning library. https: //github.com/yihaosun1124/OfflineRL-Kit , 2023.

Cong Lu, Philip J Ball, Jack Parker-Holder, Michael A Osborne, and Stephen J Roberts. Revisiting design choices in offline model-based reinforcement learning. arXiv preprint arXiv:2110.04135 , 2021.

Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning , pages 9786-9801. PMLR, 2022.

Zhang-Wei Hong, Pulkit Agrawal, Rémi Tachet des Combes, and Romain Laroche. Harnessing mixed offline reinforcement learning datasets via trajectory weighting. arXiv preprint arXiv:2306.13085 , 2023.

Bernard L Welch. The generalization of 'student's'problem when several different population varlances are involved. Biometrika , 34(1-2):28-35, 1947.

Kerong Wang, Hanye Zhao, Xufang Luo, Kan Ren, Weinan Zhang, and Dongsheng Li. Bootstrapped transformer for offline reinforcement learning. Advances in Neural Information Processing Systems , 35:34748-34761, 2022.

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific antibody design and optimization with diffusion-based generative models for protein structures. Advances in Neural Information Processing Systems , 35:9754-9767, 2022.

Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. arXiv preprint arXiv:2108.01073 , 2021.

Brandon Trabucco, Kyle Doherty, Max Gurinas, and Ruslan Salakhutdinov. Effective data augmentation with diffusion models. arXiv preprint arXiv:2302.07944 , 2023.

Appendix

We provide further details of our paper in the appendix. Our code implementation can be found in https://github.com/Jaewoopudding/GTA

A Methodology Details

A.1 Subtrajectory Generation Details

In this section, we provide further details for our trajectory-level generation. As generating a whole episode that includes more than hundreds of timesteps at once is infeasible, we generate a subtrajectory with a horizon H , following prior works [21, 22]. Each subtrajectory consists of a sequence of state, action, and reward.

Note that we need full environment transitions, i.e., states, actions, rewards, next states, and terminals (if MDP is finite) for training conventional Q-learning-based Offline RL algorithms. To achieve this, we actually model the subtrajectory of length H +1 and use a generated state s t + H for the next state of H th transition.

In environments like Walker2d and Hopper, which have terminal transitions, each trajectory may have a different episode length. It is necessary to incorporate terminal information when setting the conditioning value for guiding generative augmentation. When it comes to directly generating terminal value, modeling the terminal variable is hard due to its binary nature and sparseness, as pointed out in [13]. The situation is worse in trajectory-level generation as it should be forbidden that the terminal variable is set to 1 in intermediate transitions. To address this, we follow the implementation in [22] by adjusting the return calculation and generation process: we subtract 100 from the reward of terminal transitions. This modification allows for the encoding of terminal information without the need for binary variable terminal generation.

A.2 Model Architecture Details

We represent the denoising network model with a temporal U-Net [22] with 6 repeated residual blocks. We choose MLP-Mixer [33] architecture for modeling each block instead of temporal convolution layers. In MLP-Mixer, we first concatenate each component (states, actions, and rewards) and project it into the MLP layer dimension-wise, then apply MLP across temporal locations and repeat the process. By doing this, we can efficiently capture both temporal and component-wise relationships.

A.3 Reweighted Sampling

We adopt reweighted sampling, which is suggested in [34, 35]. First, we segment the condition y , which is the return value of subtrajectory, into N B equal intervals. Each interval is weighted according to its value and the number of points in the interval, with the specific weight v j calculations applied as follows:

where ˆ y represents the maximum return value in the offline dataset D , | B j | denotes the count of data points in the j th bin, and y b j refers to the midpoint of the interval associated with bin B j . In the equation, u and q serve as smoothing parameters. The parameter q determines the weighting between high-score bins and low-score bins. Lowering q results in reduced weights for the low-score bins, and conversely, increasing q assigns higher weights to them. The detail setting of the u and q is provided in Appendix C.3.

B Hyperparameters

B.1 Diffusion Models

Size of the diffusion model and training details . We use the same size of network and hyperparameters for modeling conditional diffusion model across all experiments for ease of implementation. The hyperparameters we used for modeling and training are listed in Table 6. The number of trainable parameters of our diffusion model is about 29.6M, which is comparable with the size of the diffusion model used in SynthER ≈ 25.6M. When using the default network model, training takes approximately 30 hours for 1M training steps with batch size 256 on a 4 RTX3090 NVIDIA GPU.

Trajectory sampling from the diffusion model . For the sampling process, we use the Stochastic Differential Equation (SDE) sampler suggested in EDM [30] with the default hyperparameter setting used in SynthER [13]. As we employ classifier-free guidance, an additional hyperparameter guidance scale w has been introduced. We use the same guidance scale w = 2 . 0 for all experiments. Detailed hyperparameters are listed in Table 7.

To sample 5M transitions, it takes approximately 3.5 hours with 128 diffusion timesteps on a single RTX 3090 NVIDIA GPU. Note that our sampling time can be reduced as we adopt the partial noising and denoising techniques. For example, with µ = 0 . 5 , we only require 64 denoising steps and can accelerate the sampling process.

B.2 Detailed parameter settings in GTA

We provide hyperparameter search spaces for partial noising level µ , multiplier α , horizon length H , and reweighting parameter used for GTA in Table 8.

Note that in sparse reward settings where the reward is a binary variable, the reweighted strategy tends to predominantly sample subtrajectories near the goal location. It might not be helpful since it makes the distribution of generative data cover a narrow region. To address this, we sample episodes instead of subtrajectories with reweighting in sparse reward settings. We assign the weight c times the bigger value to the successful episodes. Once the episode is sampled via weighted sampling, subtrajectories for augmentation is sampled from the uniform distribution. For dense reward environments like Gym locomotion, we put details of reweighting parameters in Table 9

C Experiment Details

C.1 Data Augmentation Baselines

In this subsection, we describe the implementation details of the data augmentation baselines. Note that we try to follow the official implementation provided by authors, including hyperparameters, and re-implement only when public code is unavailable.

S4RL [12] : S4RL adds zero-mean Gaussian noise to input states as augmented data, i.e, ˜ s t ← s t + ϵ, ϵ ∼ N (0 , σI ) . We choose σ = 0 . 0003 , following the original paper.

SynthER [13] : SynthER trains an unconditional diffusion model that generates transitionlevel data (state, action, reward, next state, and terminals). We build SynthER based on code provided by authors 2 . For a fair comparison, we re-implement the setting for the offline RL with a larger diffusion network for all environments and sample the same number of transitions. The diffusion model of SynthER contains 25.6M trainable parameters. We only use generated data from SynthER to train offline RL methods for evaluation, following the original paper. Their code is opened at https://github.com/conglu1997/SynthER

C.2 Offline RL Algorithms

We employed Clean Offline Reinforcement Learning (CORL) library [39] and OfflineRL-Kit library [40] to adapt implementation code and hyperparameters of the Offline RL algorithms. We import CQL [8], IQL [36], TD3BC [7] and DT [38] from the CORL library, and MCQ [4] from the OfflineRLKit. The codebase of the CORL can be found at https://github.com/tinkoff-ai/CORL , and OfflineRL-Kit from https://github.com/yihaosun1124/OfflineRL-Kit .

C.3 Dataset Version and Evaluation Procedure

We employ v2 dataset for the Gym locomotion and AntMaze tasks, v1 for the Maze2d tasks, and v0 for the Adroit and FrankaKitchen tasks. We follow the standard evaluation procedure of D4RL benchmarks. For all the tasks we examined, we train TD3BC, IQL, CQL, MCQ for 1M training steps and report the score over the final evaluations. Especially, in case of the DT in Appendix E.4, we trained it for 100k steps. We reported the final evaluation scores of the experiments, using 10 evaluation steps for Gym locomotion tasks and 100 steps for Maze2d, AntMaze, Adroit, and Franka Kitchen tasks.

We normalize the score where 0 corresponds to a random policy, and 100 corresponds to an expert policy as proposed at D4RL benchmark [15]. All experiments of the main sections are conducted over 8 different random seeds. We set the batch size as 1024 across all Q-learning based offline RL algorithms. The justification for batch size configuration is covered at Appendix G.1

C.4 Pixel-based Environments

We conducted experiments on pixel-based environments using the VD4RL benchmark [16]. We follow the experiment protocol of the pixel-based experiment of Synther [13], firstly pretrain the CNN encoder with image observations for 1M steps and then fine-tune policy and Q-network with the augmented latent observations for 300k steps. We reported the mean and standard deviation of the last evaluation score. We calculate the standard deviation of the average score by averaging the standard deviation of each dataset, not identical with the standard deviation calculate protocol of Table 1 and Table 2, which utilize the standard deviation of mean performance of each seed. We used the DrQ+BC codebase and the offline dataset on the https://github.com/conglu1997/v-d4rl and re-implement the fine-tuning.

We used α = 2 . 0 and µ = 0 . 1 for the medium, medium-expert, and expert dataset and α = 1 . 3 and µ = 0 . 1 for the medium-replay dataset. For the reweighted sampling, we choose N B = 50 , u = 0 . 001 , and q = 5 . 0 across all datasets.

D Data Quality Metrics

In this section, we present details about the quality metrics of the augmented dataset. We sampled 1M of transitions among 5M generated transitions to effectively test these metrics.

D.1 Dynamic Plausibility

Dynamic plausibility, or Dynamic Mean Squared Error (Dynamic MSE), assesses the congruence of generated subtrajectories with the environment dynamics, following methodologies from prior works [41, 13]. Let D A represent the augmented dataset. ( s, a, r, s ′ ) denotes a single transition, and f ∗ denotes the true dynamic model of the environment. Then, dynamic MSE is computed by the error between the generated next state s ′ and the true next state f ∗ ( s, a ) . Prior to error computation, we normalize each state.

D.2 Novelty

Novelty metric, inspired by [42, 37], quantifies the uniqueness of the augmented trajectories compared to the original offline data. Novelty is calculated as the average l2 distance between each generated trajectory and its nearest offline data point. To gain a deeper understanding of the generated data, we categorized novelty into state novelty, action novelty, and state-action novelty. The calculation formula is as follows:

where D is offline dataset.

D.3 Optimality

Lastly, optimality reflects the ability of generated trajectories to achieve high rewards. We quantify optimality as an oracle reward, which is computed as the average of the reward values ˜ r obtained by querying the environment with the generated states and actions ( s, a ) . This metric serves as an indicator of the effectiveness of the generated trajectories in terms of reward maximization.

E Additional Experimental Results

In this section, we present additional experiment results to deepen our understanding on GTA.

E.1 Ablation on trajectory-level generation

Table 10 shows that reducing the trajectory generation to a single transition significantly raises the Dynamic MSE and degrades policy performance. Once the generation target becomes trajectory, horizon length marginally impacts both performance and Dynamic MSE. These results emphasize the importance of leveraging sequential relationships within transitions when augmenting trajectories.

E.2 Mixed-quality Dataset

It is widely known that the performance of offline RL significantly varies depending on the performance of the behavior policy used to collect the data [15]. However, offline RL algorithms often fail

to fully exploit the experience from high-performing trajectories [43] when the offline dataset is a mixed-quality dataset, which consists of few high-performing trajectories and many low-performing trajectories. Here, we demonstrate that GTA can effectively address mixed-quality dataset issues by augmenting high-performing trajectories using information from low-performing trajectories.

In this experiment, we mix the halfcheetah-medium-replay-v2 and halfcheetah-expert-v2 data with ratios of 1:20 and 1:10. We choose these two datasets as the performance gap between them is the largest among the Gym locomotion tasks. The halfcheetah-medium-replay-v2 dataset consists of 200k transitions, and the expert transitions in the mixed-quality dataset are 10k and 20k each.

We augment the mixed-quality dataset with various methods, including naive duplication, S4RL, SynthER, and GTA. Naive duplication is the setting where high-performing trajectories are duplicated to 5M transitions without any modification.

To prove that GTA can exploit information on low-performing trajectories for augmenting highperforming trajectories, we additionally trained TD3BC on the expert-only configuration where the dataset consists of only high-performing trajectories used for comprising mixed-quality datasets. We employ S4RL, SynthER, and GTA as an augmentation method for expert-only configuration. We hypothesize that If the result of the GTA on the mixed-quality outperforms GTA on expert-only, GTA exploits the information of low-performing trajectories while augmenting high-performing trajectories.

We train GTA on the mixed-quality or expert-only dataset and augment only the high-performing trajectories. Similarly, we train SynthER on both datasets and generate the new transitions from the diffusion model. By the partial noising and denoising technique , GTA can choose the specific data to augment, while SynthER cannot because SynthER starts generating the transition from the Gaussian noise without any guidance.

Table 11 demonstrates that GTA shows a significant performance boost on mixed-quality configurations. The result also exhibits that GTA on the mixed-quality configuration outperforms GTA on the expert-only configuration. From these results, we suggest that GTA can effectively augment the high-performing trajectories while exploiting the information on low-performing trajectories. Additionally, the exceptionally superior performance of GTA on expert-only configuration implies that GTA can enhance the sample efficiency considerably.

E.3 Effectiveness of GTA Under The Data Sparsity

To assess the effectiveness of GTA when the number of transitions in the offline data is small, we designed experiments that utilized only 5%, 10%, 15%, and 20% of the total offline data for the training diffusion model. We then evaluated GTA with TD3BC in the halfcheetah-medium-v2 environment. Table 12 demonstrate that Even with only 5% of the data, GTA outperforms the TD3BC baseline implemented with 100% of the data. This implies that GTA can efficiently improve the sample efficiency of offline RL algorithm when the volume of the offline dataset is limited.

E.4 Effectiveness of GTA on sequence modeling-based approaches

Unlike prior works that augment data at the transition level, GTA generates trajectories and therefore can be applied to sequence modeling approaches such as Decision Transformer [DT, 38]. For training DT, we should assign returns-to-go token for each timestep t , which is calculated as ˆ R t = ∑ T k = t r k . To model return-to-go for our generated subtrajectory, we set the return-to-go of the first timestep of the subtrajectory as ˆ R t , where ˆ R t can be computed from the original data. For subsequent timesteps, we deduct the reward that we generated.

To test the efficacy of GTA on sequence modeling, we train a DT on Gym locomotion tasks and Maze2d environments. The result in Table 13 demonstrates that DT with GTA outperforms baselines in terms of average performance at Gym locomotion tasks. And Table 14 shows that DT with GTA augmented data exhibits superior performance with respect to average score. These results highlight the effectiveness of trajectory-level data augmentation in preserving sequential relationships and leveraging the learned long-term dynamics.

E.5 Ablation on Reweighted Sampling

We conduct an ablation study on reweighted sampling, which is additionally introduced to make the diffusion model focus on high-reward regions. We conduct experiments on two environments (Hopper-medium and Walker2d-medium) to verify the effectiveness of the reweighted sampling. Table 15 presents evaluation scores and oracle rewards with and without reweighted sampling. As shown in the table, reweighted sampling achieves a higher D4RL score compared to its counterpart, showcasing its effectiveness. We also observe that the oracle reward of the generated dataset is also increased, indicating that our reweighted sampling strategy successfully positions more points in higher reward regions.

E.6 Ablation on Different Conditioning Strategies

We explored alternative design choices for GTA that exploit the inpainting ability of the diffusion model. GTA condition on amplified return value to improve the optimality of the offline dataset.

However, another design choice is to condition the diffusion model by inpainting each step of reward with amplified reward. We named this method as amplified reward inpainting. We implement reward inpainting by multiplying the reward of each trajectory by a multiplier α = 1 . 3 and continually replacing generated rewards with multiplied rewards of the original trajectory during the denoising process of the diffusion model. Note that we do not apply the return conditioning originally used in GTA when amplified reward inpainting is applied. Table 16 demonstrates that the trajectories augmented with reward inpainting yield lower returns than GTA, especially a significant performance drop on the Hopper-medium dataset.

Additionally, we conducted a comparative analysis of the augmented trajectories with different conditioning strategies, examining both the dynamic MSE and the oracle reward. As illustrated in Table 17, the dynamics MSE for data augmented through reward inpainting is significantly higher compared to the one augmented with GTA. Oracle reward of data augmented through reward inpainting is also lower than GTA. We hypothesize that reward inpainting introduces constraints that reduce the flexibility of states and actions, resulting in poor dynamic MSE and degrading the underlying offline RL policy.

E.7 Ablation on utilizing reward proxy model

Instead of using augmented reward, we additionally train the reward model to replace it. Table 18 compares performance using a reward function instead of augmented rewards.

Although the method that replaces augmented rewards with the predictions of the proxy reward model shows comparable performances, incorporating the reward model into the augmentation process would additionally involve a reward labeling step for augmented states and actions before training the offline RL algorithm. Therefore, we opted to generate all states, actions, and rewards at once to achieve a more straightforward yet more effective augmentation.

F Additional Analysis

In this section, we present additional analysis that are not included in the main section due to the page limit.

F.1 Comparing Conditioning Strategies

A common approach for providing classifier-free guidance is to condition on a fixed value, representing the desired return according to the environment [22, 14]. However, this strategy can lead to significant issues when jointly implemented with partial noising . In the scenario where a trajectory is denoised from the partially noised trajectory, conditioning on fixed return, which is excessively greater than the return of the original trajectory, collides with the pre-existing context.

When the partially noised trajectory is denoised with guidance towards the fixed excessively high return value, the guidance collides with the pre-existing context. Such forceful conditioning can lead to the generation of invalid trajectories. Furthermore, even without applying partial noising and denoising, using a fixed value tends to decrease the diversity of the samples because of the narrow extent of the conditioning value.

To this end, we analyze the fixed conditioning method in more detail. Table 19 shows the quality measurements of the dataset generated with a fixed conditioning. We observe that fixed-value conditioning is capable of generating novel and high-rewarding samples. However, the dynamic plausibility significantly increases due to the forceful conditioning.

F.2 Statistical Test Results for The Performance Boosting Effects of GTA

We investigate the statistical meaningfulness of the performance improvement compared to the ones of the SynthER. To compute the standard deviation for the average performance of each offline RL algorithm, we first organize the final performance by seed and calculate the mean score of each seed across the tasks of the table. Next, we calculate the standard deviation of these mean scores for each seed and report this value as the standard deviation of the algorithm's performance. This protocol for calculating standard deviation is applied at Table 1 and Table 2.

Based on the t-test results shown in Table 20, we suggest that GTA induces a statistically significant performance boost over SynthER, achieving p-values below 0.05 across all algorithms. The p-value result on Table 20 is calculated by the Welch's t-test for t-statistic [44] and the experiment results are sourced from the Table 1.

F.3 Impact of Partial Noising and Denoising Framework on Terminal States

We further explore the effects of partial noising and denoising framework , particularly focusing on how well they preserve original trajectory information, especially terminal states. Halfcheetah contrasts with hopper and walker2d environments in terms of the existence of episode termination.

While hopper and walker2d have terminals, the sparsity of terminal states in datasets, exemplified by their 0 . 0001% occurrence in walker2d transitions, underscores the challenge of preserving this critical information.

Our findings, illustrated in Figure 6, show that in Walker2d, a high noise level ( µ = 1 ) leads to unstable learning due to excessive deviation from the original trajectory. In contrast, a lower noise level ( µ = 0 . 5 ) maintains more of the original trajectory data, resulting in more stable learning outcomes. This highlights the importance of an approach to the balance of retaining original information in environments with terminal states.

hopper-medium

walker2d-medium

09

80

100

D4RL Normalized Score

58

70

08

56

60

09

54

50

40

20

52

40

0

50

30

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Training Steps (M)

Training Steps (M)

Training Steps (M)

μ=0.5

μ =1.0

F.4 Data Quality Table for the Figure 5

In this section, we provide the exact results in addition to the analysis of the data quality metric presented in Section 5.4. Table 21, Table 22, and Table 23 exhibit the optimality, novelty, and dynamic MSE of GTA and baseline augmentation methods. For a clear presentation of Figure 5, we take the logarithm of the Dynamic MSE. For novelty and oracle reward, we adjust the axis scale for each environment.

As shown in Table 21 and Table 22, GTA consistently achieves higher values than other augmentation methods at both the oracle reward and novelty across all locomotion environments. Simultaneously, the dynamic MSE remains comparably small. These results suggest that GTA successfully generates data that is novel, high-rewarding, and dynamically plausible.

Beyond comparing the novelty of state-action pairs, we measured the novelty by separating state and action to assess the ability of GTA to discover novel states as well as novel actions. As seen in Table 25 and Table 24, GTA consistently outperforms other methods in generating novel states and actions.

F.5 Comparisons of GTA with Offline Model-based RL and Diffusion Planners

In this section, we compare the performance and computational cost of GTA with offline model-based RL methods and diffusion planners. While GTA shares commonalities with both approaches as discussed in Section 2, it introduces key approaches that lead to improved performance and test-time computation efficiency.

Table 26 demonstrates superior performance of GTA compared to both model-based RL algorithms and diffusion planners in locomotion tasks. The results suggest that trajectory augmentation via diffusion models can be more effective than traditional single-transition dynamic models used in offline model-based RL. Additionally, GTA outperforms diffusion planners, which require a timeintensive sampling process during decision-making.

We also analyze the computational efficiency of GTA compared to baseline methods. As shown in Table 27, GTA exhibits significant computational advantages over diffusion planners during test-time evaluation. This advantage stems from shifting the computational burden from the sampling phase to the data preparation stage.

The model training process for both GTA and offline model-based RL algorithms [25, 27] consists of two stages: dynamic model training and policy training. GTA takes longer for dynamic model training due to its diffusion model with millions of parameters. However, offline model-based RL algorithms require significantly more time for policy training because model rollouts slow down this step. Overall, the total training time for COMBO and GTA is comparable, but GTA demonstrates superior performance, as shown in Table 26.

F.6 Correlation on Conditioning Value and Generated Reward

To evaluate the controllability of GTA in trajectory generation, we analyze the Pearson correlation between the sum of rewards for generated subtrajectories and the conditioning value. As shown in Figure 7, the correlation coefficients of Halfcheetah medium, Halfcheetah medium-replay and Halfcheetah medium-expert datasets are 0.55, 0.91, and 0.99, respectively. These results demonstrate that our diffusion model can effectively generate trajectories that align with the conditioned return values.

(a) halfcheetah-medium-replay

(b) halfcheetah-medium

(c) halfcheetah-medium-expert

250 -

Pearson Correlation:0.55

Pearson Correlation:0.99

200 -

==- y=x

y=x

400

200 -

150

Generated Values

Generated Values

150

Generated Values

300

100

100

200

50

50

100

50

Conditioning Values

100

150

200

50

Conditioning Values

100

150

200

250

100

Conditioning Values

200

300

400

F.7 Oracle Reward Distribution of Generated Trajectories

Wevisualize the reward of the original dataset, generated reward from GTA, and real reward computed from the environment using state and action generated from GTA to visualize that GTA can really shift the data distribution to high reward region while preserving environmental dynamics. We present the results conducted on Halfcheetah-medium-v2 in the main section. Other results are presented below. As shown in the figures, GTA mostly generates valid and high-reward trajectories across various environments and datasets.

0.0125

Probability density

0.020

0.010

0.0100

0.015 -

0.008

0.0075

0.006

0.010

0.0050

0.004

S000

0.0025

0.002

0000

0.0000

0.000

0

50

100

150

200

250

-50

0

50

100

150

200

0

100

200

300

400

hopper-medium

hopper-medium-replay

hopper-medium-expert

0.04

0.015

!!

!!

0.03

00

0.010

0.02

0.02

S000

100

0.00

20

40

60

80

100

120

140

160

0.000

0

50

001

150

0.00

25

50

75

100

125

150

175

walker2d-medium

walker2d-medium-replay

walker2d-medium-expert

0.04

0.025

0.010

Probability density

0.020

0.008

0.03

0.015

0.006

0.02

0.010

0.004

0.005

0.002

100

0.000

0.000

0.00

50

0

50

100

150

200

-50

0

50

100

150

200

50

0

50

100

150

200

Oracle rewards sum of subtrajectory

Oracle rewards sum of subtrajectory

Oracle rewards sum of subtrajectory

Offline data

Amplified

G GTA Sensitivity Test

In this section, we conduct a sensitivity test on hyperparameters introduced in GTA. To demonstrate the sensitivity of GTA on these hyperparameters, we choose Halfcheetah-medium-v2 dataset with TD3+BC as a baseline Offline RL algorithm. For all experiments, we train the algorithm with 1M gradient steps with 4 random seeds.

G.1 Batch Size and Training Epochs for Offline RL Training

We further verify that increasing batch size and training epochs affect the performance of our method. Most prior works fix batch size as 256 and training epochs as 1M, which is a reasonable choice where

the original dataset size is quite small (100K to 2M). However, as we upsample the dataset into 5M, it is natural to check whether increasing batch size or training epochs might lead to an improvement in the performance. To validate this, we perform experiments with the upsampled dataset by varying batch size and training epochs.

Figure 9 shows the performance of the TD3BC algorithm trained with an augmented dataset from GTA by varying batch size. We observe that increasing batch size significantly improves the performance. Table 28 illustrates the performance of offline RL algorithms trained with the original offline dataset. As illustrated in the table, there is no big difference in performance when we just naively scale up the batch size. We attribute this to the natural phenomenon of neural networks to generalize well with a larger batch size if there is more data available for training. To this end, we set the batch size as 1024 for all methods in the main experiments.

As shown in the Figure 10, we observe that as we increase the training epochs, GTA can achieve more performance gain. We also observe that using both large batch size and more training epochs continually improves the performance. After 10M training steps with a batch size of 1024, we achieve significant improvement compared to the TD3BC trained with the original dataset, trained 1M steps with a batch size of 1024.

100.0

D4RL Normalized Score

58

50

97.5

56

40

M/W

95.0

92.5

30

54

90.0

20

87.5

52

85.0

10

82.5

50

0.0

0.2

0.4

0.6

0.8

1.0

0

0.0

0.2

0.4

0.6

0.8

1.0

80.0

00

0.2

0.4

0.6

0.8

1.0

Training Steps (M)

Training Steps (M)

Training Steps (M)

batch size:256

batch size: 1024

60

D4RL Normalized Score

59

58

57

56

batch size:256

batch size:1024

55

0.0

2.0

4.0

6.0

8.0

10.0

Training Steps (M)

G.2 Scale of Multiplier α

The scale of multiplier α determines the strength of model exploitation towards the high reward region. If we set α high, GTA tries to generate high-reward trajectories compared to the original trajectories, which might lead to generating dynamically infeasible trajectories. If α leans towards 1, it would generate trajectories with similar return values to original trajectories and cannot push data distribution towards high-reward regions. Table 29 shows the performance of TD3+BC trained with generated data across different α values. We observe that there is a significant drop in performance where α is too high (e.g., 2.0). Also, we find that 1 . 1 ≤ α ≤ 1 . 5 generally exhibits good performance.

G.3 Scale of Partial Noising µ

The scale of partial noising ratio µ controls the balance between generating in-distribution and OOD trajectories. If we set µ ≈ 0 , GTA generates trajectories similar to the original dataset and stays close to the original data distribution. If we set µ ≈ 1 , it can generate novel trajectories via exploration while possessing the risk of violating environment dynamics. Intuitively, the high value of µ might lead to a significant performance drop due to the over-exploration of unknown regions. However, we observe that even when we choose relatively high µ (even set µ = 1 ), it occasionally generates not only novel but also dynamically plausible trajectories. Table 29 shows the D4RL score by varying µ . The table shows that data generated with relatively high µ exhibits higher performance. However, the aforementioned trend is not always true for all environments, as we have already discussed in Section 5.3.

G.4 Size of Augmented Dataset

In our main experiment section, we generate 5M transitions with the conditional diffusion model. We investigate the performance of our method in terms of the size of the upsampled dataset. To this end, we choose four different levels of size from the range [0.1M, 10M]. As shown in Figure 11, the performance gain achieved by the upsampled dataset increases as the size of the augmented dataset becomes larger and converges after 5M samples.

58.5

Score

58.0

57.5

D4RL NormalizedS

57.0

56.5

56.0

55.5

55.0

0.1

1.0

5.0

10.0

Size of Upsampled Dataset (M)

H Extended Related Works

H.1 Related Works in Offline RL

In this section, we discuss related works with our work but just briefly discuss them in the main paper to clearly define the position of our work in offline RL literature.

BooT : BooT [45] suggests a novel algorithm to train a Transformer-based model for offline RL by incorporating the idea of bootstrapping. BooT is a bootstrapping-based methodology that builds upon the Trajectory Transformer, reusing self-generated trajectories within the sequence model itself.

AdaptDiffuser : AdaptDiffuser [24] is a planning method with a diffusion model that can self-evolve by training with self-generated diverse synthetic data using guidance from reward gradients. The AdaptDiffuser generates trajectories guided by various reward functions and trains the diffusion planner itself, shifting the data distribution learned by the diffusion model towards high reward regions.

Main differences : Our focus is building a high-quality augmented dataset that can be used for training any offline RL algorithms in a plug-and-play manner. We develop an algorithm-agnostic method to boost the overall performance of offline RL, not confined to specific kinds of model. Unlike our work, BooT and AdaptDiffuser utilize generated trajectories only for training itself and thereby are categorized as planners, not data augmentation, which is the reason why they are not included as baselines of our approach. While our method can also utilize bootstrapping, i.e., further training diffusion model with generated trajectories, we leave it as a future work.

H.2 Related Works in Diffusion Models

We also notice that there are several approaches that resemble our approach, denoising from an intermediate state instead of pure random noise to preserve originality in various domains.

DiffAb : DiffAb [46] use diffusion model to optimize existing antibodies by first perturbing the CDR sequence and denoise it. They observe that optimizing antibodies with this procedure can improve binding energy while keeping the sequence similar to the original one.

SDEdit : SDEdit [47] highlights the key challenge of image editing as balancing the faithfulness and realism of synthetic images and proposes a methodology that leverages diffusion models to augment images. To enhance the realism of the generated images, it introduces partial noising and denoising, similar to the approach in GTA, where the editing is guided by the input image itself.

DA-Fusion : DA-Fusion [48] conducted data augmentation utilizing pretrained diffusion model. They generate synthetic images using the partially noised images and then denoise them with guidance toward the original image. DA-Fusion significantly improved few-shot classification performance on the Pascal and COCO datasets.

I Broader Impacts

This paper is for the advancement of offline Reinforcement Learning. While it opens up various potential applications, any specific societal impacts are not immediately apparent and require further exploration.

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We clearly state the main claims in the abstract and introduction, accurately reflecting the paper's contributions and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss limitations in Section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: We do not include theoretical result.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide detailed experiment settings in the Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide open access to the data and code to reproduce the paper.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All of these can be found in the Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We conduct all experiments with multiple random seeds and report error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide information on the computer resources in the Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We use D4RL, and VD4RL which does not contain harmful or offensive content.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the broader impact of the paper in Appendix I.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We do not use controversial dataset.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We mention the license in the README.md of our code.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We offer codebase of GTA via Anonymous github.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: GTA does not involve any crowdsourcing and human-subject elements.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: GTA is about data augmentation without any crowdsourcing or human subjectrelated experiments.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the challenges of Offline RL, which involves learning effective decision-making policies from static datasets without online interactions. The authors introduce Generative Trajectory Augmentation (GTA), a novel approach that uses a diffusion model to enhance offline data by augmenting trajectories to be both high-rewarding and dynamically plausible. GTA partially noises original trajectories and denoises them with classifier-free guidance based on amplified return values. The experiment results indicate that GTA, as a general data augmentation strategy, enhances the performance of widely used offline RL algorithms in both dense and sparse reward settings.\n\nSTRENGTHS:\n(1) Interesting and important topic: The guided data augmentation for offline RL is interesting. The success in this domain will further benefit the real-world applications of offline RL algorithms.\n\n(2) Extensive experiments: the authors provide extensive experiments to validate their proposed method.\n\nWEAKNESSES:\n(1) Lack of theoretical guarantee to support the algorithm: The authors only provide empirical results, lacking theoretical analysis of the performance of the proposed method.\n\nFor other potential weaknesses, please see my questions.\n\nQUESTIONS:\nI will adjust my rating score based on the rebuttal results.\n\n(1) How to select the hyperparameter \\alpha in the denoising with amplified return guidance in equation (8). Can you also provide ablation study results for different selections of \\alpha?\n\n(2) In the generation phase, do you generate whole trajectories or generate subsequences of trajectories?\n\n(3) Is the proposed method sensitive to the hyperparameter \\mu? Based on the description in Figure 4 and Appendix F.3, when \\mu = 1, it only performs well on the dataset with low quality. For high-quality datasets with abundant expert demonstration, does your method have trouble preserving and reconstructing the high-rewarding trajectories?\n\n(4) Can you provide comparison results with diffusion-based offline RL methods? For example, Decision Diffuser [1] and AdapDiffuser [2].\n\n(5) In the ablation study on reweighted sampling, what base offline RL do you use for D4RL score calculation?\n\n(6) Since the reweighted sampling can implicitly change the dataset distribution, can you also apply the reweighted sampling technique as well for other baseline methods?\n\nReference:\n\n[1] Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T., & Agrawal, P. (2022). Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657.\n\n[2] Liang, Z., Mu, Y., Ding, M., Ni, F., Tomizuka, M., & Luo, P. (2023). Adaptdiffuser: Diffusion models as adaptive self-evolving planners. arXiv preprint arXiv:2302.01877.\n\nLIMITATIONS:\nYes, the authors discussed the limitations in their paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo improve the quality of offline datasets, in this paper, a generative data augmentation approach is proposed by leveraging the diffusion models. Moreover, with the adoption of partial noising and denoising framework with amplified return guidance, the trajectory can be guided towards high-rewarding region. Finally, with the generated trajectories, existing offline RL methods can be utilized to learn the optimal policy.\n\nSTRENGTHS:\nThe paper is clearly organized and the experimental results show the effectiveness of the proposed approach. Also, the separation of the data generation with diffusion models and the policy learning mitigates the possible time and computation cost in the policy learning. Moreover, the trajectory-level data generation can capture the transition dynamics, which is beneficial for environments with sparse rewards.\n\nWEAKNESSES:\nThe proposed approach is mainly a direct combination of current techniques, e.g., the diffusion model for data generation, the adding noise for the exploration and the existing offline RL techniques for policy learning, the theoretical contribution of the paper is trivial. Also, experimental results are still not sufficient, e.g., the lack of the comparations of time and computation costs between different methods, and only two data augmentation baselines are compared. Moreover, the proposed approach is limited to trajectories with reward signals, which is not the case in many real-world applications. Furthermore, it is unclear to me whether there are any theoretical guarantees to push the generated trajectory toward the high-rewarding region.\n\nQUESTIONS:\n1. In section 4.2, by denoising with amplified return guidance, the authors claim the trajectory can be guided towards high-rewarding regions, are there any theoretical guarantee for this statement?\n2. How about the whole computation and time cost of the proposed approach when it is compared with other diffusion-free baselines? \n3. As shown in the experiments, for different environments, different noise levels are required for the optimal performance, then how can we determine the noise level for a new task? \n4. The proposed approach aims to deal with trajectories with reward signals, however, in reality, it may be difficult to obtain such signals, can the proposed approaches be extended to the cases where no reward signals can be provided?\n\nLIMITATIONS:\nThe authors have addressed some limitations, but for some other limitations, e.g., the computation and time cost of the proposed approach, the application of the proposed approach in some real-world applications with no reward signal provided should be further stated.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper builds on ideas from SynthER but adds classifier free guidance to boost the returns of the generated trajectories. This makes sense as the resulting data has higher quality and the paper is a totally sensible next step in the series of works building upon SynthER, which is an exciting direction for research. The paper is well executed and a solid contribution.\n\nSTRENGTHS:\n* A highly relevant area building on synthetic data generation for offline RL. Out of all the current active areas in RL research this one benefits the most from the current foundation model/large data regime, and scaling offline RL has been shown to be highly impactful in areas such as Robotics (E.g. RTX). \n* The method is not overly complicated and builds upon the recent SynthER paper, presented at NeurIPS 2023. This is an example of a simple idea that makes a great deal of sense, and it is well executed. \n* The paper reads well.\n* The experiments go beyond what was done in SynthER, including some new environments. It is great to see this - as RL needs to keep pushing for more complex benchmarks and not just sticking to D4RL and Atari \"because thats what the previous paper did\".\n\nWEAKNESSES:\nThere are no major weaknesses here, it is a solid paper taking a nice step in an exciting general direction of research. The below points are fairly minor:\n* The authors could cite Ball et al 2021 \"Augmented World Models\" as another example of data augmentation in offline RL.\n* The authors could discuss how this relates to another paper building on SynthER, Policy-Guided Diffusion by Jackson et al. This is very recent work so makes sense it is not needed as a baseline, but the comparison may be worth including in text.\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nNice to see this in the main body rather than the Appendix. It would be good to add a discussion on scalability if possible.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces Generative Trajectory Augmentation (GTA), a data augmentation approach for Offline Reinforcement Learning (RL) that enhances the quality of static datasets by generating high-rewarding and dynamically plausible trajectories using a conditional diffusion model. GTA partially noises original trajectories and then denoises them with classifier-free guidance via conditioning on amplified return value. The authors demonstrate that GTA improves the performance of various offline RL algorithms across several benchmark tasks.\n\nSTRENGTHS:\n- GTA integrates seamlessly with existing offline RL methods. It builds on previous advancements, and new offline RL algorithms could also benefit from it, given its agnostic nature towards the specific RL method used.\n- The ability to generate high-return trajectories that do not exist in the logged data is a significant advantage. This feature potentially improves the performance of offline RL algorithms by enriching the dataset with valuable transitions, as supported by the experimental results.\n- The experimental results presented in the paper show significant improvements across various offline RL algorithms and environments. Additionally, thorough ablations are provided, highlighting the impact of different components and hyperparameters on the overall performance. This comprehensive evaluation demonstrates the practical effectiveness of GTA.\n\nWEAKNESSES:\n- The claim that GTA-generated data adheres to the dynamics of the environment (lines 145-148) seems unfounded. Is there a principled argument on why diffusion models would learn the dynamics of the environment well? Especially where the goal is to create transitions outside the dataset. Figure 5 does not conclusively show that GTA is dynamically plausible, especially for tasks like Cheetah compared to methods like S4RL.\n- Compared to S4RL, the augmented trajectories appear less plausible. This issue is evident in tasks like HalfCheetah, yet GTA still achieves better rewards. It is unclear why this does not affect the final results.\n- Based on Appendix B.2, the performance of GTA appears to depend heavily on finely-tuned hyperparameters, as suggested by the different values used for different environments in Table 6. This raises concerns about the generalizability of the method, indicating that it might rely on online manual tuning.\n\nQUESTIONS:\n- How do you ensure that the generated trajectories are plausible and adhere to the environment's dynamics?  What mechanisms to pick &alpha to handle cases where the return used for high-reward guidance is not reasonable? How do you ensure that the generated trajectories remain valid and useful? \n- Does each task require a new &mu, &alpha hyperparameter? How does this impact the generalizability and practical application of GTA?\n- What is the proportion of augmented trajectories to original ones in the final dataset? Is the entire offline dataset kept before training, with augmented trajectories added on top? Do all the offline datasets have the same size? Are 5 million augmented tranitions added regardless of the original dataset's size?\n- Could you clarify if Figure 3 is based on a real example or if it is just an illustration of the method's intuition? If it is the latter, how can you ensure it accurately represents the method's practical application?\n\nLIMITATIONS:\nThe authors discuss the limitations and potential impacts in their paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents Generative Trajectory Augmentation (GTA) that is aimed at improving offline reinforcement learning (RL). GTA uses a diffusion model conditioned on high returns to generate high-rewarding trajectories. These trajectories are used to augment the static datasets used to train offline RL algorithms. The augmentation process involves partially adding noise to existing trajectories and then refining them using a denoising mechanism guided by an amplified return value. The authors demonstrate that GTA improves the performance of popular offline RL algorithms across tasks in the D4RL benchmark. They also analyze the quality of the augmented trajectories, showing improvements in both data optimality and novelty over two baselines (S4RL and SynthER).\n\nSTRENGTHS:\n**Originality:** The use of a conditional diffusion model for trajectory augmentation is a novel approach that adds computation overhead on the data creation stage instead of the policy learning stage.\n\n**Quality:** The paper includes experiments across tasks from the D4RL benchmark, demonstrating the effectiveness of GTA in various settings, including dense and sparse reward tasks, high-dimensional robotics tasks, and a pixel-based observation task.\n\n**Clarity:** The paper is well-organized and presents the methodology, experiments, and results clearly and logically. The authors provide an anonymized link to their code enhancing the reproducibility and transparency of the proposed method.\n\n**Significance:** GTA is shown to be compatible with offline RL algorithms, making it a flexible solution for data augmentation in offline RL. By adding novel high-rewarding trajectories to the datasets in the D4RL benchmark, GTA shows improvement in the performance of offline RL algorithms.\n\nWEAKNESSES:\n1. As shown in Figures 4 (a)(b) and Table 25, the performance of GTA is sensitive to the choice of hyperparameters, such as the noising ratio ($\\mu$) and the multiplier for conditioned return ($\\alpha$) which might require extensive tuning for different tasks.\n2. While empirical results show that the generated trajectories using GTA are dynamically plausible (Figure 5), it is not explicitly enforced in GTA and might be a byproduct of generating high-rewarding trajectories. The reviewer believes that further evaluation of GTA on real-world data is needed to make claims regarding GTA's dynamic plausibility such as in lines 6-9 (“In response, we introduce … and dynamically plausible”), lines 39-40 (“GTA is designed … dynamic plausibility”).\n3. Related to weakness #2, while the paper shows GTA's effectiveness on standard benchmarks, it lacks validation in real-world applications where the dynamics and data distributions may differ significantly from simulated environments. This poses a question on the real-world applicability of GTA.\n4. The paper does not include any experiments comparing GTA with model-based RL baselines. Although the authors state that both approaches perform data augmentation at different stages (data generation vs policy learning), the reviewer believes that a comparative study should be included. This is because both approaches learn a separate model to generate augmented data.\n\nMinor:\n- Line 48: “... any offline RL algorithms …” should be “... any offline RL algorithm …”.\n- A pink curve is present in Figure 4 (b), but its legend is missing. The blue curve denoting $\\mu$ = 1.0) is missing in this figure,\n\nQUESTIONS:\n1. Why is IQL the only chosen baseline (and not TD3-BC, CQL, MCQ) for the tasks listed in Table 2 (Adroit and FrankaKitchen)?\n2. Why is DrQ+BC the only chosen baseline and Cheetah-run the only chosen pixel-based observation task in Table 3? Why were results on the other baselines (IQL, TD3-BC, CQL, and MCQ) not presented for the pixel-based task(s)?\n3. How compute-intensive (GPU or hours) is running GTA on the pixel-based observation tasks?\n\nLIMITATIONS:\n1. As the authors have mentioned, while empirical results on the D4RL benchmark tasks show that GTA has low dynamic MSE, in tasks where dynamic violations have a critical impact, the performance boost may not be significant.\n2. The current evaluation is limited to simulated tasks, and the effectiveness of GTA in real-world offline RL tasks remains to be validated.\n3. As shown in Figures 4 (a)(b) and Table 25, the performance of GTA is sensitive to the hyperparameters ($\\mu$ and $\\alpha$) which might limit its applicability.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-27"}