Goal-Conditioned Data Augmentation for Offline Reinforcement Learning
Authors: Xingshuai Huang
Di Wu
Benoit Boulet
Venue: Accepted by TMLR ()

Abstract: Offline reinforcement learning (RL) enables policy learning from pre-collected offline datasets, relaxing the need to interact directly with the environment. However, limited by the quality of offline datasets, it generally fails to learn well-qualified policies in suboptimal datasets. To address datasets with insufficient optimal demonstrations, we introduce Goal-cOnditioned Data Augmentation (GODA), a novel goal-conditioned diffusion-based method for augmenting samples with higher quality.  Leveraging recent advancements in generative modelling, GODA incorporates a novel return-oriented goal condition with various selection mechanisms. Specifically, we introduce a controllable scaling technique to provide enhanced return-based guidance during data sampling. GODA learns a comprehensive distribution representation of the original offline datasets while generating new data with selectively higher-return goals, thereby maximizing the utility of limited optimal demonstrations. Furthermore, we propose a novel adaptive gated conditioning method for processing noisy inputs and conditions, enhancing the capture of goal-oriented guidance. We conduct experiments on the D4RL benchmark and real-world challenges, specifically traffic signal control (TSC) tasks, to demonstrate GODA's effectiveness in enhancing data quality and superior performance compared to state-of-the-art data augmentation methods across various offline RL algorithms.

Published in Transactions on Machine Learning Research (08/2025)
Goal-Conditioned Data Augmentation for Offline Reinforcement Learning
Xingshuai Huang
ringshuai honpshuui.hungg@mmllmmooooo
Deportment of Electrical and Computer Engineering
McGill Universily
Di Wu
di. uru @maggll.ccc
Department Electrical and Computer Engineering McGill University
Benoit Boulet
benoit. Muwoil.boullt@@mccg..
Department of Electrical and Computer Engineering
McGill University
Reviewed on OpenReview: /ttps://openreeeee opw.review....eee.emmee net/ forum? id- 8K16dplph
Abstract
Offline reinforcement learning (RL) enables policy learning from pre-collected offline datasets, relaxing the need to interact directly with the environment. However, limited by the quality of offline datasets, it generally fails to learn well-qualified policies subop- timal datasots. To address datasets with insufficient optimal demonstrations, we introduco Goul-cOnditioned Dala A Augmentation (GODA), novel goal-conditioned diffusion-based method for augmenting samples with higher quality. Leveraging recent advancements in gen- erative modelling, GODA incorporates novel return-oriented goal condition with various selection mechanisms. Specifically, we introduce controllable scaling technique to pro- vide enhanced return-based guidance during data sampling. GODA learns comprehensive distribution representation of the original offline datasets while generating new data with selectively higher-return goals, thereby maximizing the utility of limited optimal demonstra- lions. Furthermore, we propose a novel adaptive gated conditioning method for processing noisy inputs and conditions, enhancing the capture of goal-oriented guidance. We conduct experiments on the D4RL benchmark and real-world challenges, specifically traffic signal control (TSC) tasks, to demonstrate GODA's effectiveness in enhancing data quality and superior performance compared to state-of-the-art data augmentation methods across vari- ous offline RL algorithms.
1 Introduction
Reinforcement learning (Sutton & Barto 2018 aims to learn a control policy from trial and error through interacting with the environment. While demonstrating remarkable performance in various domains, this approach typically requires vast. amounts of training data collected from these interactions. Such data- intensive requirements become impractical in applications where environmental interactions are costly, risky, time-consuming, such robotics (Bhateja et al. 2024 Tang al. 2025). autonomous driving (Li et all Taghavifar et al.] 2025), and TSC (Zhai et al.) 2025 Du et al. 2024). Offline RL offers feasible solution to these challenges by enabling policy learning directly from pre-collected historical datasets, thus significantly reducing the need to interact directly with the environment.
2024
Although offline RL makes policy learning less expensive, performance highly dependent on the quality of the pre-collected datasets and may suffer from lack of diversity, behavior policy bias, distributional shift, and suboptimal demonstrations (Prudencio et. al. 2023 The performance of offline RI. tends to decline

---

Published in Transactions on Machine Learning Research (08/2025)
Graph
Training data
GODA data
-6
Start point
2
Target point
4
6
6
Distance
Return-to go (RTG)
-8
2
-6
Figure 1: Illustrative examples of how GODA augments higher-return data with goal guidance. GODA utilizes scalable RTG-based goal conditions to generate samples with higher returns (shorter overall distance).
drastically when trained with suboptimal offline datasets. Previous studies have attempted to address these issues by constraining the learned policy to align closely with the behavior policy (Lyu al.] 2022) or by limiting out-of-dittribution action values (Kostrikov et al.) 2021). Although these approaches have shown performance improvements, they retain the inherent defects of offline datasets, remaining highly dependent on data quality.
Several studies have addressed the limitations of oflline RL using data augmentation methods to generate more diverse samples. One approach involves learning world models to mimic environmental dynamics and iteratively generate synthetic rollouts from a start state (Zhang et al.] 2023 Treven et al., 2024). While this method significantly improves sample efficiency and data diversity, it suffers from compounding errors and fails to control the quality of generated trajectories. Other research leverages generative models to capture the distributions of collected datasets and randomly sample new transition data (Lu et al.) 2024).. Although these methods demonstrate some performance improvements, they remain inefficient when dealing with datasets containing limited optimal demonstrations. This inefficiency stems from their inability to effectively control the quality of generated data.
We attempt to address this challenge by taking advantage of generative modelling to augment higher-quality data with directional goals. Unlike previous studies (Lu et al., 2024) that sample data unconditionally and randomly, we introduce GODA to incorporate representative goals, guiding the samples toward higher returns. Given the exceptional performance of diffusson models (Ho et al. 2020 Karras et al. 2022 in the field of generative artificial intelligence, GODA utilizes a diffusion model as its generative framework. GODA is trained to capture a comprehensive representation of the data distribution from the original dataset while sampling now data conditioned on selective high-return goals. This approach maximizes the utility of the limited well-performed trajectories in the original datasets.
Inspired by Decision Transformer (Chen et al., 2021), we define the 'goal' the return-to-go (RTG), which represents the cumulative rewards from the current step until the end, coupled with its specific timestep in trajectories. RTG explicitly indicates the expected future rewards for given behavior at the current timestep for a specific trajectory. We assume that at the same timestep across different trajectories, higher RTG signifies a higher goal. To generate samples that exceed the quality of the original dataset, we introduce three goal selection mechanisms and a scoling technique to control our expected goals during sampling.
We present an illustrative example of how GODA operates in Figure 1 path from the starting point to the target. By setting higher RTG goals during sampling, GODA can potentially discover a more efficient route that yields higher return (the RTG at the first timestep is equal to the return).
The task to identify the shortest
To better incorporate goal conditions, we further propose a novel adaplive gated condilioning approach. This method utilizes a condition-adaptive gated residual connection and an adaptive gated long skip connection to selectively capture multi-gramularity information effectively with the guidance of goals. GODA is an off the-shelf solution that can seamlessly integrate with other offline RL optimization approaches on various tasks achieve superior results. We summarize our contributions:
2

---

Published in Transactions on Machine Learning Research (08/2025)
1) We propose a goal-conditioned data augmentation method, namely GODA, for offiine RL. It achieved enhanced data diversity and quality for offline datasets with limited optimal demonstrations.
2) We introduce novel directional goals with selection mechanisms and controllable scaling to provide higher- return guidance for the data sampling process in our employed generative models. Additionally, we propose a novel adaptive gated conditioning opproach to better capture input information based on goal guidance.
3) We show GODA's competence through comprehensive experiments on the D4RL benchmark compared with state-of-the-art data augmentation methods across multiple offiine RL algorithms. We further evaluate the effectiveness of GODA on real-world application, i.e., traffic signal comtrol, with small-size datasets obtained from widely used controllers in real-world deployments. These evaluations verify GODA's effec- liveness in addressing various challenges, significantly enhancing the applicability of RI-based hods for real-world scenarios.
The remaining soctions are organized as follows: Some preliminarics about offline RL and diffusion models are introduced in Section 2 followed by details about our methodology in Section 3 Section 4 describes the experimental settings about D4RL and TSC tasks, baselines and evaluation algorithms. .Nxtt, we show the data quality measurement for our augmented datasets in Section and provide more detailed experimental results in Section 6 Finally, we conclude our research and discuss future directions in Section
7
The Related
Work section is in Appendix
2 Preliminaries
2.1 Offline Reinforcement Learning
In RL, the task environment generally formulated as a Markov decision process (MDP) (S, (Sutton & Barto 2018). E S, s' E S, a E A, r= R),,a), P(s'ls, a), and 7 [0,1) represent state, next state, action, reward function, state transition, and discount factor, respectively. RL aims to train agent to interact with the environment and learn a policy from experience. The objective of RL. is to maximize the expected discounted cumulative rewards over time:
J= Ex
Σ y R (st, ai)
(1)
1==0
where t denotes the timestep in a trajectory. For offline RL, the policy is learned directly from of- fline datasets pre-collected by other behavior policies, instead of environmental interactions. The offline dataset typically consists of historical experience described as tuples (s, and other environmental signals. After learning a policy xD) from dataset D, the performance is evaluated in online environment as Ex(D) 2 (st, a.)]
While offline RL eliminates reliance on interacting with the environment, it is highly restricted by the quality of offline datasets due to the lack of feedback from the environment. Our GODA aims to enhance the diversity and quality of the dalaset by upsampling the pre-collected data to augmented dataset D". The objective is to learn a policy T(D*) that outperform x(D) learned from original dataset D, such that
Σn 7' (st, an) t=
Σ T (st, an))
E ETD))
> ET(D)
(2)
1=0
2.2 Diffusion Models
Diffusion models (Sohl-Dickstein et al. 2015 Ho et al. 2020 Karras et al.] 2022). ) class of well-known
generative modeling methods, aim to learn comprehensive representation of the data distribution Paata (xx) with a standard deviation ddata from a given dataset. Diffusion models generally have two primary processes, the forward process, also known as the diffusion process, and the reverse/sane sampling process.
The forward process is characterized by a Markov chain in which the original data distribution x E Pdata (xN) is progressively perturbed with predefined i.i.d. Gaussian noise schedule gN < N-1
3

---

Published in Transactions on Machine Learning Research (08/2025)
Goal condition selection
Forwerd/Diffusion Process
Original Dataset
X
x
x
Denolser Network
Origbal Disjenory tnlers)) Tockecc
(5t,, de St-1. Fu Sr. O)
x
x
x
Maiin
Reverse/Sampling Process
Be
Goal Scaling Factor 1/A
Augmented Dataset
Sample
Eeer la
(6
x
x
Reverse/Sampling Process
Figure 2: Illustration of diffusion model training and data augmentation. The diffusson model first trained on the original dataset to learn the underlying data distribution, and is then used to generate higher-quality samples conditioned on selected controllable goals.
<..<0== = maxx Therefore, we can obtain a sequence of noised distributions x(x for each nose level of. where the last noised distribution p(x"; mmax) can be seen as pure Guassion noise when max >
Odata-
For reverse process, Denoising Diffusion Probabilistic Model (DDPM) (Ho et al. 2020 models it as a Markov chain that involves denoising an initial noise p (x") N (x":0.1 to the original data distribution with learned Gaussian transitions. Elucidated Diffusion Model (EDM) (Karras et al. 2022) formulates the forward and reverse processes as probability- ordinary differential equation (ODE), where the noise level can be increased or decreased by moving the ODE forward backward in time:
dx = - pox
ort) (t:))dtt
(3)
where at++) denotes derivative over denoise time and Vx log (t;)) is referred to as the score function (Song et al. 2020), which points towards regions of higher data density. It worth noting that we use 4 to denote the noise time to distinguish it from the trajectory timestep L. The ODE pushes the samples away from the data or closer to the data through infinitesimal forward or backward steps. The corresponding step soquence is (tu, f1, ...t,; , where ty =0 and N denotes the number of ODE solver itcrations.
EDM proposes to estie the score function using denoising score matching (Karras et al. 2022). Specifi- cally, denoiser neural network D,(x; o) is trained to approximate data x~ sampled from Pdaca by minimizing the Le denoising loss independently for each a:
min ExN Npp..-nm.nn.. |xxx + +n;o)= -x
Subsequently, the score function can be calculated as Vx logg(x;;; = (x;(x; Heun's 2nd order ODE solver (Ascher& Petzold 1998) to solve the reverse-time ODE by iteratively refining x toward tN 0, i.e., sampling data with the reverse process.
x)/o². EDM employs
3 Methodology
In this section, we introduce GODA, goal-conditioned data augmentation method utilizing generative modeling for augmenting higher-quality synthetic transition data. Our adopted diffusion model first learns comprchensive data distribution from the initial offline dataset, subsoquently sampling new data towards higher return with controllable, selective goal conditions, as shown in Figure 2 In this part, we define a representative goal for GODA in Section 3.1] and introduce different selective mechanisms for goal conditions in Section 3.2 To proactively control the sampling direction, a controllable goal scaling factor is introducod in Section 3.3 For better integrating goal conditions as guidance within the diffusion model, we further propose novel adaptive gated conditioning approach (Section 3.4 that rnduces condition-adaative gate
4

---

Published in Transactions on Machine Learning Research (08/2025)
mechanism into the long skip connection and the residual connection. Section 3.5 presents the pseudocode (Algorithm 1
and outlines the implementation details of our method.
3.1 Return-oriented goal
Prior diffusion-basse data augmentation work (Lu et al. 2024) lacks the ability to guide the sampling process toward desired outcome. In our approach, we introducc return-oriented goal as a conditioning signal for the diffusion model, analogons to how prompts are used to guide image generation in some diffusion models Peebles & Xie 2023: Li et al. 2025). As shown in Figure 3 the denoiser neural network not only takes as input the noised sample and noise level, as described in Equation 4 but also an extra goal condition. Inspired by Decision Transformer, we adopt return-to-go (RTG) (Chen et al.) [021). which represents the cunulative rewards from the current step till the end, as explicit goal condition
9: rv-
(5)
For each transition sample represented as tuple (s,,,s',)) within trajectory, the RIG serves unbiased measure of the future rewards corresponding to the current state-action pair. In the context of Decision Transformer, a higher RTG at a given timestep implies a more ambitious goal for the policy to achieve, thereby guiding the behavior of the agent loward more rewarding outcomes. Since the same behavior at different timesteps often yields varying RTGs across different trajectories, we combine the RTG with its corresponding timestep in the trajectory as the condition for each specific transition sample. The timestep signal acts as timestamp for each goal.
3.2 Selective goal conditions
To get the training goal conditions from the original dataset, we first organize offline samples into trajectories, compute the RTG 9. for each, and append timestep every sample. These goals along with the transition samples {(ee, f, 81, ae, r()) are then used to train the diffusion model to capture the underlying data distribution. For data augmentation (sampling procedure), however, aim to generate new samples condi- tioned on selectively chosen batch of goals. To fully leverage well-performing samples and augment samples with higher returns, we propose three distinct condition selection mechanisms: return-prior, RTG-prior, and random goal conditions.
(1) Return-prior goal condition. In this approach, we rank all trajectories based on their return values and select the top n trajectories. During the sampling procedure of the diffusion model, the RTG and timestep pairs t)),t) from these top trajectorics are selected the sampling goal conditions. This method filters high-return trajectories from the initial offline datasets and reuses them to sample more well-optimized transitions. The following example illustrates how we select goal conditions from the top two trajectories. of denotes the RTG value at timestep & from trajectory T.
Original Trajectory
t=1 t
t= 3
t t=
...
Return 100
1 (selected)
81 01
0
0
2 (sclected)
03
05
97
95
(excluded)
87
83
92
(excluded)
03
01
70
(2) RTG-prior goal condition. We group RTGs by their associated timesteps and then sort them to select the top RIGs along with their corresponding timesteps goal conditions. As illustrated in the example below, for each column, the RTG values are sorted in descending order from top to bottom. This approach selectively reuses high-RTG transitions for data augmentation, focusing on transitions that are most. likely to yield higher returns.
(3) Random goal condition. We randomly select m RTG and timestep pairs (gc. f) as sampling goal conditions for each batch of samples. This increases the diversity of the augmented data while paying less attention to the optimal trajectories for improving performance.
5

---

Published in Transactions on Machine Learning Research (08/2025)
Re-ordered Trajectory
(selected)
9]
92
9g
(selected)
91
02
0g
94
3 (excluded)
02
(excluded)
92
94
3.3 Controllable goal scaling
Selective goal conditions offer high-return guidance during the sampling process, but are limited in generating data with returns or quality beyond the initial offine datasets. To overcome this limitation, introduce controllable goal scaling factor, A, which can be multiplied with the goal values to represent higher relurn expectation. This approach enables flexible adjustment of goal values to drive the sampling process toward higher-quality data. As illustrated in Figure [ a higher RTG goal at each timestep directs the sampling process toward a trajectory with a greater overall return. Since RTG values can be either positive or negative in certain tasks, we propose multiplying positive goals by the scaling factor and dividing negative goals by it.
(A (Ge/A,t),
Ji >= 0 Ôt O.
goal
(6)
It worth noting that we use a goal scaling factor rather than an increment factor, i.e., goal = (g. X, t).
The reason is that incrementing each RTG J. = Er=rr with fixed value would be the same as simply incrementing the terminal reward er by À while leaving all other target rewards for t * T unchanged. However, assigning a higher RTG goal at cach timestep is intended to encourage the agent to generate trajectories that yield higher rewards throughout the trajectory, rather than only improving the terminal reward. Therefore, a goal scaling factor can be better choice than an increment factor.
3.4 Adaptive Gated Conditioning
To better capture goal guidance and scamlessly integrate conditions into the diffusion model, we propose novel adaptive gated conditioning approach, as shown in Figure This structure significantly enhances the ability to guide the diffusion and sampling processes using goal conditions. The conditional inputs include both the noise level condition and goal condition, which are embedded separately, then element- added, and fed into the neural network. The noise input is processed with several gated residual multi-layer perception (MLP) blocks with novel adaptive gated skip connections between shallow and deep layers.
3.4.1 Condition Embedding
The noise level for diffusson is encoded using Random Fourier Feature (Rahimi & Recht 2007 embedding. The RTG is processed with a linear transformation to get a hidden representation. The timestep of each RTG is embedded with Sinusoidal positional embedding (Vaswani et al. 2017). We concatenate the RTG and timestep embeddings form the goal condition, which is then element-wise added to the noise level embedding and used as the conditional input.
3.4.2 Adaptive Gated Long Skip Connection
As shown the left part of Figure 3, we adopt long skip connection similar to U-Net (Ronneberger et al., 2015) to connect MLP blocks at different levels. To capture different information with varying importance weights, we propose novel adaptive gated long skip connection structure by adding the previous information with an adaptive gate mechanism based on the given conditions.
Fout (1 - w) * Tskip + w *
where Zakip and r are outputs of a shallower block and the previous block, and w denotes learnable weight caleulated by regressing the conditional input with an MLP and sigmoid layer.

---

Published in Transactions on Machine Learning Research (08/2025)
Prediction
Adaptive Gated Residual
Linear and Reshape
Layer Norm
Dropout
Gated Residual MLP Block
2
Linear
Adaptive Gated
Dropout
Gated Residual MLP Block
Adaptive Gated Skip
Gated Residual MLP Block
Linear
Y.B
Gated Residual MLP Block
MLP
MLP
Scale, Shift
MLP
Layer Norm
Embedding
Embedding
Embedding
Embedding
Embedding
Noised Sample
Noise Level
Goal
Input
Noise Level
Goal
Figure 3: The denoiser neural network with adaptive gated conditioning architecture. The iindtioon-ddaptinn gated long skip connection, shown in the left part, captures both low- and high-level features, assigning varying importance weights each. The orange arrows to the left of each Gated Residual MLP Block represent the noise level and goal condition. The adaplive galed residual, depicted in the right part, further enhances the model by selectively preserving imput information based the given conditions.
3.4.3 Gated Residual MLP Block
The right part of Figure 3 depicts the structure of each gated residual MLP block. We mainly use simple MLP structure to process the inputs, given the non-sequence input format. We have also tried more complicated structures, e.g., self-attention and cross-atlention mechanisms, while getting degraded performance. We adopt the widely used adaptive layer normalization (adaLN) method (Peebles Xie 2023 to learn dimension-wise scale y and shift 3 basod on the conditional information. Besides, we explore modification of the residual connection (He et al., 2016) and introduce novel condition-adaptive gated residual connection to further enhance the model by selectively preserving input information. It also regresses the conditional input and gets learnable weight for adaptively preserving input information.
Tout = (1 - v) F(x) UUr)) U
where F. is the learned transformation.
3.5 Method Implementation
Given the strong ability of diffusion models capture the complex data distribution and generate high- dimensional data, adopt EDM (Karras et al. 2022) our generative model for augmenting offline data. The neural network equipped with adaptive gated conditioning, denoiser function. It takes input the noised sample x + noise level a, and goal condition c, and make predictions for the original sample x. Therefore, with goal conditions с and transition data x from the original datasets, the diffusion model Gg with a learnable denoiser neural network Do is trained by
illustrated in Figure 3 is used the
min Ex S0000.0.... (D n; -
(9)
Algorithm I outlines the overall process of our GODA method. We begin by training a conditional diffusion model to approximate the data distribution of the oflline dataset (Lines 2-5), using transition tuples paired with their corresponding goal conditions as training samples. Once the diffusion model is well trained, we use to generate new transition samples by conditioning on selected and scaled goals (Lines 6-11). The resulting augmented samples are then stored in the augmentation dataset D for subsequent policy training (Line 12). The learned policy finally evaluated target tasks (Line 13).

---

Published in Transactions on Machine Learning Research (08/2025)
Algorithm GODA: Goal-Conditioned Data Augmentatiio
Initialize generative model Gg and D 0
# Diffusion Model Training
2: Split initial offline dataset D into trajectories according to episode terminal information
3: Calculate RTGs for each transition sample in trajectories
4: Add RTGs and timesteps as goals into D
Train Gg on DD using Eq. by conditioning on goals
# Data Augmentation (Sampling)
repeat
Extract goal conditions for sampling according to the goal selection mechanism. Re-assign sampling goal conditions with goal scaling factor À using Eq.
9: Sampling a batch of new transition samples B* by conditioning on scaled goals
10:
D + + UB*
11: until [D* number of target samples
Policy Training and Evaluation
12: Train policy the final dataset D U
13: Evaluate policy target tasks
4 Experimental Settings
In this section, we provide comprehensive overview of the experiments conducted to evaluate the perfor- mance of our proposed GODA method.
4.1 Tasks and Datasets
4.1.1
D4RL Benchmark
We adopt three popular Mujoco locomotion tasks from Gym i.e., HalfCheetah, Hopper, and Walker2D, and navigation task, i.e., Maze2D (Fu al. 2020), as well more complex tasks, specifically the Pen and Door tasks from the Adroit benchmark (Rajeswaran et al. 2017 Fu et al. 2020 For locomotion tasks, we adopt four data quality levels: Random, Medium-Replay, Medium, and Expert- For Maze2D, three datasets collected from different maze layouts are adopted, i.e., Umaze, Medium, and Large. For the Adroit benchmark, we use two different datasets: Human and Cloned.
4.1.2 V-D4RL Benchmark
We further extend approach to the Cheetah-Run task from V-D4RT. La et al. 2023), high-dimensional pixel-based offline RL benchmark, with four data quality levels: Random, Repluy-Replay, Medium, and Medium-Experl. We implement our GODA by following the same settings adopted in SynthER Lu et al. (2024h. Specifically, the original image-basee based datasets are first used to pre-train a Behavior Cloning (BC) policy network composed of a CNN encoder, a trunk network, and a final fully connected layer. The CNN encoder and trunk network are then frozen and used to transform each observation and next observation, represented as raw image data size 84x84x3, into 50-dimensional latent. representation. These low- dimensional representation data, combined with rewards and actions, are then used train the diffusion model and augment transition samples in the latent space. These augmented samples subsequently used to fine-Lune only the final fully connected layer (output head), thereby avoiding the need to directly generate high-dimensions image data.

---

Published in Transactions on Machine Learning Research (08/2025)
True Reward Distribution
GODA Reward Distribution
0.8
SynthER
0.8
Aug
40
SynthER GODA
GODA
True
Dynamics MSE
0.6
0.6
Density
Density
20
0.4
10
0.2
0.2
0
4
0.0
0.0
Min L2 Distance
-2
2
-2
0
2
Figure 4: Data quality evaluation for SynthER and GODA Walker2D-Randoms V2. Left: Dynamics MSE and Min L2 Distance comparison. Smaller Dynamics MSE indicates better validity, and larger Min L2 Distance indicates higher diversity. Middle: Ground-truth reward distributions from the simulator for augmented datasets. Right: Ground-truth and augmented reward distributions for the GODA-angmented dataset.
4.1.3 Traffic Signal Control
To evaluate GODA's applicability to real-world challenges, we further test it on TSC tasks with much fewer training samples using the CityFlow simulator (Zhang et al.) 2019). TSC aims to optimize traffic flow by efficiently managing traffic signals to maximize overall traffic efficiency.
To evaluate our GODA, we select three real-world scenarios featuring 12 intterseeccton grid from Jinan (JN) city and two scenarios with 16-intersection grid from Hangzhou (HZ) city (Zhang & Dong 2023). These scenarios represent a variely of traflic patterns and intersection structures, allowing us to cover a wide range of traffic situations. To bridge the gap between simulation and real-world conditions, we use the widely adopted Fixed-Time (FT) controller as one of our behavior policies for generating the initial offline datasets. Additionally, we employ Advanced Max Pressure (AMP) (Zhang al. 2022) and Advanced CoLight (ACL) (Zhang al. 2022 to create higher-quality datasets for further evaluation. We present more details in Appendix B.1.4
4.2 Baseline Data Augmentation Methods
To verify the effectiveness of our proposed GODA, compare it with three state-of-the-arl data augmen- tation methods:
TATU Zhanng ct al.
)023), which learns world models to generate synthetic rollouts and truncates trajec-
tories with high accumulated uncertainty.
SynthER (Lu al.
2024). which employs diffusion models to unconditionally augment large amounts of new data based on the learned distribution of original datasets.
DiffStitch [Li et al. 2024), which augments data with a diffusion model and three MLPs, and actively connects low to high-reward trajectorics with stitch techniques.
4.3 Evaluation Algorithms
To verify the quality of datasets augmented by GODA, we follow the evaluation settings adopted in previous data augmentation studics. We train two widely-used offline RL algorithms, i.c., IQL (Kostrikov et al., 2021 and TD3+BC (Fujimoto Gu 2021), on datasets and evaluate the learned policy on D4RL tasks. For TSC tasks, we utilize BCQ (Fujimoto al. 2019h. CQL (Kumar et al. 2020), and DataLight (Zhang & Deng. 2023 as the evaluation algorithms. For V-D4RL tasks, the BC algorithm is adopted for evaluation.
4 https///ghhubb.ommm//gggell////h
9

---

goal conditions incorporated in the sampling Published in Transactions on Machine Learning Research (08/2025)
Table 1: Data quality evaluation metrics for SyuthER and GODA on Walker2D tasks. Smaller Dynamics MSE, larger Min L2 Distance, and larger Average Reward indicate better quality.
Dynamics MSE
Min 1,2 Distance
Average Reward
Task
SynthER
GODA
SynthER
GODA
SynthER
GODA
RD-Rooon--222
2.725.7
1.922.9
1920.6
2.420.66
0.120.6
0.620.5
Wa - Meplun-Replay-22
0,54177
0.411.1
1.620.9
1.72:0.7
2,5113
3.5:0.99
Balker2D-Medimma
0.3:1..
0.3:
0.8:0.55
0.81:0.6
3.41122
3.7.0..9
It is important to note that for GODA, we train the evaluation algorithms using a mix of the original datascts and the augmented datascts, whereas for the other basclinc methods, only the augmented datasets are used. This is because GODA focuses on augmenting samples from the high-reward zones, which may lead to reduced data, diversity. In contrast, the baseline methods, as reported in their respective papers (Lai et al. 2024) and our experiments, exhibit degraded similar performance when using mix of the original and augmented datasets, as illustrated in Section C.5
5 Augmented Data Quality Measurement
Since our GODA is built upon SynthER, we compare the quality of the datasets augmented by both SynthER, and GODA to assess whether the goal conditions incorporated by GODA enhance data quality. We adopt two metrics, L.e, Dynamics MSE and Min L2 Distance, from SynthER (Lu et al. 2024
M
Dynamics MSE
M ((1)1))1)) (8 (1)+)))))
1
Σ
(10)
M
1
Min L2 Distance
[[ at) at)l2
(11)
M
and introduce an Average Reward for evaluating reward distributions of augmented datasets
1 M
Average Reward
(12)
where M is the selected number of samples, si aj 1 ri denote the samples from angmented datasets, $ and ft denote the next state and reward generated by the simulator given states and actions from augmented datasots, and 8 and at are the state and action from original datasets.
Dynamics MSE measures how well the augmentation models capture the dynamics of the environment by learning patterns from the original datasets and generating data that aligns with those dynamics. Min L2 Distance assesses the models' exploration capabilities and data diversity by calculating the average L2 distance between each augmented sample and its nearest neighbor in the original dataset, reflecting how diverse the generated data is. Average Reward compares the ground-truth reward distributions produced by the simulator given states and actions in datasets augmented by SynthER and GODA.
The left part of Figure presents scatter plot 10K points sampled from the augmented datasets. Results show that datasets generated by GODA exhibit much lower Dynamics MSE and a wider range of Min L2 Distance values, indicating both better alignment with envirommental dynamics and greater diversity, The superior validity in performance may stem from the goal conditions (RTG-timeste pairs), which provide critical information for generating samples that better match the environment's 's dynamics. Mcanwhile, the increased diversity is likely due to the process.
The middle part demonstrates that GODA not only generates samples within a high-reward data zone but also extends the boundary of high rewards beyond the best demonstrations, compared SynthER. The right part shows that the rewards generated by GODA align closely with the ground-truth values.
10

---

Published in Transactions on Machine Learning Research (08/2025)
Table 2: Normalized scores of D4RL Gym and Maze2D tasks for GODA and baseline data augmentation methods. The results are calculated across random seeds. Values in bold represent the best performance (largest score). 'HC': HalfCheetah; 'WK: Walker2D; 'HP": Hopper; 'MZ' Maze2D.
IQL
TD3+BC
Task Data
Original
TATU
SynthER
DStitch
GODA
Original
SynthER
DStitch
Rand
15.2-1.2
17.74229
17.213.4
15.812.0
19.5
11.310.8
12.132.3
12.211.1
11.8114
12.511.
Mod-R
43.5-0.4
44.2+0.1
46.6+0.2
44.7+0.1
47.840.4
44.8+0.7
44.540.3
45.940.9
44.74-0.3
44,9+0.2
HC
Med
48.3=0.1
48.24-011
49.640.3
49.4+0.1
50.4+0.1
48.1
48.120.2
49.94112
50.44:0.5
48.5+0.1
Med-E
94.6##0.2
91.420.6
93.3222
914214
91011.1
90.827.0
89.323.9
87.2111.1
96.0+0.5
94.327.7
Rand
4.120.8
6.320.5
1.20033
46211
14.3:7.1
0.620.3
6.52:4.3
2.3:19
24210
4.211.8
Med-R
82.6=8.0
75.01:12.1
83.345.9
86.6122.8
96.114.9
85.6:4.66
62.11:10.4
90.54:.33
89.7114.2
93.015.6
Med
84.0-5.4
76.6++10.7
84.745.5
83.24-2.2
79.1
82.745.5
75.8+3.5
84.8+1.4
83.4+-177
86.2+10.7
Med-E
111.7+0.6
110.7:0.6
111.420.7
111.6#0.1
110.9+0.7
110.040.4
110.7=0.7
110.240.5
110.220.3
110.7=0.5
Rand
7.2+0.2
8.142.9
7.7+0..1
6.54-0.9
8.7:22.1
8.6+0.3
18.1
14.6+99.4
8.842..3
8.2+00.1
Med-R
84.6=13.5
79.647.6
408.240.41021111
102. 102.5+0.6
64.4224.8
61.1
53. 63.4415.5
79.6+13.5
63.0=12.8
HP
Med
62.8-6.0
60.3+-3.6
72.0+-4.5
71.0+-422
74.342.9
60.4++4.0
58.344.88
63.44-4,2
60.34-4.9
74.8+3.6
Med-E
106.2-6.1
93.4
90.8+11.99
110.9#2.9
96.9++12.3
101.1+10.5
99.0+
105.449.7
107.1#7.00
107.5+15..
Average
62.1
59.541-.0
63.74-3.5
66.24:2.9
60.04...1
62.04:3.1
Umnaze
37.1=2.0
33.02488
11.020.7
38.516.2
59.512.6
29.4214.2
37.7=10.9
37.611444
$8.417.5
46.418.3
MZ Med
35.5 1.0
35.11:1.
35.11:2.6
35.5
35.8112.6
59.5.11.99
73.81336.9
65.21:36.1
66.8:130.9
86.5+28.4
Large
49.6-22.0
843.33
12..6
109=16.5
93.1225.3
4436.2
104.3420.1
Average
8.39-8.3
46.728.2
46.6
47.6+6.8
68.1#6.6
62.0#28.2
68.2=10.6
16.17
65, 66,9221.
79.1=7.5
Table 3: Normalized scores of GODA and baseline data augmentation methods on Adroit tasks evaluated
using IQL.
Thsk
Dataset
Original
TATU
SynthER
DStitch
GODA
Human
79.1128.5
88.9122.6
96.818.6
87.418.6
75.6131.4
Pen
Cloned
52.5++27.9
45.3+-23.4
64.04-29.6
Average
62.4+29.2
70.7+25.2
71.0=1600
19.
70.24-26.0
Human
1612.1
7.0:1.66
8.31222
10.0::2.5
14.815.0
Door
-0.1+0.5
-0.14+0.3
4.44-0.4
Average
0.811.3
3.511.0
7.112.0
7.211.4
15.8155.5
The evaluation results for the three metrics in Table further demonstrate that GODA outperforms Syn- thER in terms of all data quality evaluation metrics across nearly all Walker2D tasks. Notably, GODA is particularly effective on random-level datasets with high randomness and sparse optimal demonstrations, where traditional augmentation methods often fall short.
Table 4: Evaluation of GODA on the Cheetah- Run task from the pixel-based V-D4RL benchmark using the BC algorithm, with sample augmentation performed in the latent space.
Eval Algorithm
Aug Methods
Rand
Med-R
Med
Med-E
Average
Original
0.0:0.0
25.013.6
31.621.4
67.516.3
33.5
BC
SynthER
0.140.0
27.9+33.4
52.2+1.2
99.
37.5+3.5
GODA
1.4:0.8
35.5116.1
54.922.0
67.6110.4
39.914.8
6 Experimental Results
After directly evaluating the data quality of the augmented datasets using various metrics, we further assess GODA by training policies these datasets and verifying their performance corresponding tasks. This section presents the overall performance comparison between our GODA and other baseline data augmen- tation methods.
6.1 Offline Reinforcement Learning Evaluation
6.1.1
Performance on D4RL
Gym tasks. Table P presents the normalized score comparison between GODA and other state-of-the-art data augmentation methods trained on the D4RL Gym and Maze2D. We adopt the results of Original and SyntheER from the SynthER paper (Luu et al.] 2024). and those of TATU and DStitch from the DStitch
11

---

Published in Transactions on Machine Learning Research (08/2025)
Table 5: Average travel time comparison on real-world TSC tasks. Smaller travel time indicates better traffic efficiency.
BCQ
CQL
DataLight SynthER
Traffic
Dataset
Original
SynthER
GODA
Original
SynthER
GODA
Original
GODA
FT AMP
269. 71111111
267.922.7
261.1
272.042.1
273.1425
271.74:.5
279 879.
2711116
1.
JN 1
271.51:399
264.0:6.1
259.7:0.9
261.82033
261.714.
260.6
298.1+3.1
299.242.0
301.7114
ACL
271.142.44
2.1
270.610.3
273.341.6
275.214.4
273.2:4.3
256.4133.
258.412.55
255.3:10.3
FT
267.643.6
265,5+1.2
266.625.1
269.3:0.3
275.346.9
272.51199
293.911.7
288.222.2
281.020.3
JN 2
AMP
250.7=10.7
254.71:266
252.113..4
251.914.0
249.415..
245.511.00
244.412..3
240.0:12.2
240.9:14.0
ACL
253.4+0.3
265.7415
262.0=1.
248.140.3
248.2422.
248.012.0
241 0.33
236.420.3
235.1=0.5
FT
206 966.
203.6+499
257.3+3.4
268.0+0.7
273.742.3
267.1
112
299.9+5.0
299.8+1.3
JN 3
AMP
263 8+3.1
259.3::0
263.2+4.4
251.5+2.9
247.7+41.9
342.5=3.5
239. 1.8
241.811.
232.5=1177
ACL
242.2+t:.00
245.720.6
243.2=1.6
242.1=1.4
244.5444.0
245.547.1
240.343.6
23..741.3
230.1=22.7
Average
367.923.2
267 1:2.3
263.942.7
261.219
265.72:12
202/4:3.2
007.522.
265 2..33
262.6#1.4
FT 324.527.3
313.2:2.0
310.520.8
317.4=5.7
315.52:10
307.1::2.7
290.120.6
290.820.3
2
HZ
AMP
295.8::4.6
302.711.9
301.7:4.1
300.010.3
295.11:0.3
285.4111.0
287.314..3
284.913.8
286.113.2
ACL
281.741.11
281.3+6.55
28.
288.643.3
286.342.8
278.9:12.1
284.9+55.6
282.1+1.77
278.0=
FT
340.0::..9
340.714.6
3.13
341.712.7
334.811.4
1.1.661..9
308.010.3
308.413.1
309. 20..233.0
HZ 2
AMP
332.5++0.3
324.8+3..0
316.9+1.9
318.0:10.66
321.9+3.3
321
312.3+33.6
310.4+2.2
308.4112.2
ACL
336.7+1.9
329.3+0.55
327.1+2.8
347.443.8
343.942.6
336.8+0.3
317.6+4.1
314.8+4.1
315.4+3.0
Average
317.3266
315.8233
313.612.3
319.142.1
316/112.1
310.821.8
302.0:36
300.143.0
299.1#3.0
Li et al. 2024 paper. We further conduct experiments for tasks not covered in the literature. As shown in the results, GODA consistently outperoorms other methods across most Gym locomotion tasks when evaluated with both IQL and TD3+BC, resulting in higher average scores. Notably, even for tasks using Random datasets, GODA successfully leverages limited high-quality samples to enhance data quality, leading to improved final performance. GODA boosts performance on Med and Med-R datasets; however, a notable gap still remains between the performance of algorithms trained on GODA-augmented Med/Med-R datasets and the performance of algorithms trained on the original Med-E datasets. These limitations suggest that while GODA is effective at maximizing the utility of suboptimal data, it cannot fully synthesize expert-level performance from non-expert data alone. Its strength lies in extracting and amplifying high-quality signals within existing data, not inventing expert behavior from scratch.
Maze2D tasks. For Maze2D tasks where rewards are sparse, GODA demonstrates significant improvements across all datasets, achieving average gains of 43.4% and 16.0% over the best bascline methods when evaluated with IQL and TD3+BC, respectively. The maximum improvement reaches 57.7% when applying GODA to the Maze2D Large dataset. This might be because Maze2D) tasks involve navigation from start to a goal location, similar to the illustrative example in Figure ! making the goal naturally explicit and spatially defined. GODA can directly condition the learning process on the desired goal position, which aligns perfectly with the task's structure. Besides, GODA can possibly reuse sub-trajectories and recombine transitions from different paths that end near the same goal, making data augmentation more effective and less likely to violate task constraints. These results highlights GODA's 's ability to effectively capture data distributions of various types of datasets and consistently augment high-quality samples.
Adroit tasks. Table 3 presents the normalized scores Adroit Pena and Door tasks, evaluated using the IQL algorithm, BC fails to perform on these tasks. As shown, GODA outperforms all baselines the Pen-Cloned dataset, although it underperforms on Pen-Human. This might be because Pen task involves high-dimensional, dexterous manipulation, where small perturbations in action can lead to disproportionately large changes in reward (c.g., minor deviations in gripper position can affect pen stability). GODA's 's reliance on RTG as a conditioning signal can be limiting such scenarios. However, for both datasets in the Door task, GODA demonstrates significant improvements over the best baseline methods. These results further demonstrate that GODA is capable of handling more complex tasks.
6.1.2 Performance on V-D4RL
Table 4 shows the results of the BC algorithm evaluated on the pixel-based Cheetah- task, with data augmentation performed the latent space. The results indicate that GODA achieves the best performance on three out of four dataset quality levels and obtains the highest average score. It outperforms the BC algorithm without data augmentation by 19.1%, and the BC with SynthFR-auggentted samples in the latent
12

---

Published in Transactions on Machine Learning Research (08/2025)
Walker2D Medium-Replay
Walker2D Medium--Repae
Walker2D Medium Medium-Replay
120
120
Normalized Score
110
RTG
Return Random
Normalized Score
120
IQL
TD3+BC
Normalized Score
110
B80 adaLN
GODA
In-Context
100
100
100
90
90
80
70
70
60
0.8
1.0
1.2
1.4
1.6
1.8
2.0
60
TD3+BC
Scaling Factor
TD3+BC
Figure 5: Ablation studies on condition selection mechanism, goal scaling factor, and adaptive gated condi- tioning from left to right. Yellow and purple horizon lines represent results for SynthER.
Walker2d Random
Walker2d Medium-Replay
Walker2d Medium
Normalized Score
30
IQL
TD3+BC
Normalized Score
110
IQL
TD3+BC
Normalized Score
100
IQL
TD3+ BBC
20
100
90
10
90
80
80
70
10
30 50
70
90 110130
10
30
50 70
90 110130
10
30
70 90 110130 Top Conditions
50
Top Conditions
Top Conditions
Figure 6: Ablation study on top conditions.
space by 6.4%. Notably, GODA delivers substantial improvements on the Random and Medium Meeium--eeplyy datasets, highlighting its suitability for offline scenarios with limited expert demonstrations.
6.1.3 Performance on Traffic Signal Control
Table presents a comparison of average travel times across different methods for TSC tasks. As shown, while SynthER achieves modest. improvements
the performance of models trained on the original datasets,
it fails to surpass the original datasets on JN tasks when using the CQL algorithm. In contrast, GODA consistently outperforms both the original datasets and SynthER across most tasks and all average evaluations. These extended experiments on TSC tasks further validate that GODA is not only applicable to diverse tasks but also capable of improving performance in real-world challenges.
6.2 Ablation Study
To validate the effectiveness of GODA's 's components, we conduct experiments using different configurations.
6.2.1 Ablation on Condition Selection
We test three condition selection mechanisms described in Section random goal conditions. As shown in the left part of Figure 5 the return-prior method demonstrates superior performance compared to the other two approaches. Morcover, GODA with the return- and RTG- prior conditions outperforms SynthER when lested on two offline RL algorithms. In contrast, the random- prior method shows results comparable to SynthER. This suggests that high-goal conditions identified by the return- and RTG-prior methods enable GODA generate samples boyond the original data distribution. Randomly selected goal conditions, however, fail target high-reward regions, producing similar results SynthER.
3.2 return-prior, RTG-prior, and
13

---

Published in Transactions on Machine Learning Research (08/2025)
Table 6: Ablation study on adaptive gated long skip connection and adaptive gated residual connection. AdaLN removes both the adaptive gated long skip connection and adaptive gated residual connection but only retains the AdaLN structure; w/o AG-LSC removes the adaptive gated long skip connection; w/o AG-RC removes the adaptive gated residual connection; w/o Gate-LSC retains the long skip connection, but without any gating mechanism; and w/o Gate-RC retains the residual connection, but without galing.
Task
Eval Alg.
AdaLN
w/o AG-LSC
w/o AG-RC
w/o Gate-LSC
w/o Gate-RC
GODA
WK2D
IQL
89.1
90.2482
91.2+0.1
92.428.6
96.0=E1.7
96.1
TD3+BC
91.7#5.9
92.24:99
92.527.9
91.926.5
92.9
93.0#5.6
Average
90.423.8
91.216.6
93.4144.
92.227.6
94.0110.4
94.627.5
Morcover, return-prior generally outperforms the RTG-prior. This may be because using a return-prior aligns better with the distribution of the original trajectories. However, directly applying the RTG prior can introduce inconsistencies when reordering trajectories, as illustrated in Section 3.2 Specifically, the resulting rewards (calculated as the difference between the RTGs of the current and next timesteps) in the re-ordered trajectories may sometimes lose important high-reward signals, since RTGs are extracted from different original Irajectories and may no longer stay in high-reward regions.
6.2.2 Ablation on Goal Scaling Factor
We further examine the effect of different scaling factors on D4RL lasks, testing values ranging from 0.8 to 2.0. As seen in the middle part Figure 5 the performance improves as the scaling factor increases, but slightly degrades when the scaling factor exceeds 1.1. Scaling factors below 1.0 shrink the selected goals, leading to suboptimal samples. Conversely, scaling factors above 1.1 push the selected goals too far beyond the training data distribution, resulting in diminished performance. Further ablation studies can be seen in Appendix C.4
6.2.3 Ablation on Adaptive Gated Conditioning
Moreover, we evaluate the impact of the adaptive gated conditioning method. We compare GODA with two variants: one using only adaLN conditioning (Peebles & Xiel 2023). and another using in-context conditioning, where condition embeddings part of Figure 5 it is clear that GODA with adaptive gated conditioning achieves the best results, and the adaLN and in-context conditioning show similar performance. Additionally, all three methods outperform SynthER, which lacks goal conditions. This demonstrates that the inclusion of goal conditions is crucial for guiding the sampling process toward high returns, and our adaptive gated conditioning method enhances the model's ability fully utilize these conditions.
directly appended to the input embeddings. From the right
We conduct further ablation experiments both the condition-adaative gated long skip connection (Section 3.4.2 and aandition-adapiie gated residual connection (Section 3.4.3). and each adaptive gating component well. The results in Table D show that removing either the adaptive gated long skip connection or the adaptive gated residual connection leads to performance degradation, with the adaptive gated long skip connection presenting a more pronounced impact. Replacing the adaptive gated long skip connection with a vanilla long skip connection, i.e., without any adaptive gating, also results a significant performance drop, though it still performs slightly better than having no long skip connection at all. This highlights the importance of incorporating goal information into the conditioning structure. A similar trend observed for the adaptive gated residual connection.
6.2.4 Ablation on Top Conditions Selection
Given that the original datasets contain varying numbers of trajectories, the number of top conditions selected for sampling may differ across datascts. We compare different selections of the top n conditions for each D4RL datasct. Based on the results shown in Figure D we empirically select 50 top conditions for the Random and Medium-Replay datasets, and 40 for the Medium datasets.
14

---

Published in Transactions on Machine Learning Research (08/2025)
While this hyperparameter may seem less critical than X, it still meaningfully influence performance. For example, as shown in Figure 7, the performance difference between the best and worst top-n settings on Walker2D Medium-RRppaay using IQL reaches 8.8, which is a notable gap. The top-n selection primarily acts as a filter, ensuring that data augmentation focuses on high-return regions. However, increasing n can introduce more diversity, allowing the diffusion model to learn from a broader range of transitions. This represents a classic trade-off between data quality and diversity: smaller n prioritizes high-return samples, while larger n can improve generalization through diversity.
7 Conclusion and Discussion
This paper proposes a novel goal-conditioned data augmentation method for offline RL, namely GODA, which integrates goal guidance inlo the data augmentation process. We define the easily obtainable return- to-go signal, along with its corresponding timestep in a trajectory, as the goal condition. To achieve high- return augmentation, we introduce several goal selection mechanisms and a scaling method. Additionally, we propose novel adaptive gated conditioning structure to better incorporate goal conditions into our diffusion model. We demonstrate that data augmented by GODA shows higher quality than SynthER without goal conditions on different evaluation metrics. Extensive experiments the D4RL benchmark confirm that GODA enhances the performance of classic offline RL methods when trained on GODA-augmented datasets. Furthermore, we evaluate GODA on real-world traffic signal control tasks. The results demonstrate that GODA is highly applicable to TSC problems, even with very small real-world training datasets, making RL-based methods more practical for real-world applications.
We acknowledge that while GODA is effective at maximizing the utility of suboptimal data, it cannot fully synthesize expert-level performance from non-expert data alonc. Its strength lies in extracting and amplifying high-quality signals within existing data, not inventing expert behavior from scratch. Besides, GODA can sometimes fail to perform well on some datasets with high reward sensitivity, such as Adroit Pen-Human. Moreover, our current method is not directly applicable to high-dimensional, pixel-based visual RL or real- world robotics datasets, as it requires transforming image data into Oww-dimensional latent representations before applying difflssion-bssed data augmentation. Since our approach operates low-dimensional feature space, this additional transformation step can complicate the implementation and introduce dependencies on the architecture of the policy network. In future work, we aim to develop more direct methods for augmenting pixel-based based datasets by improving the sampling efficiency of the diffusion model. We also plan to explore the application of GODA to more complex, real-world scenarios.
References
Cri M. Ascher and Linda R Petzold. Computer methods for ordinary differential equations and differential- algebraic equations. SIAM, 1998.
Chethan Bhateja, Derek Guo, Dibya Ghosh, Anikait Singh, Manan Tomar, Quan Vuong, Yevgen Chebotar, Sergey Levine, and Aviral Kumar. Robotic offline rl from internet videos via value-function learning. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 16977 16984. IEEE, 2024.
Chang Chen, Fei Deng, Kenji Kawaguchi, Caglar Gulcehre, and Sungjin Ahn. Simple hierarchical planning with diffusion. ICLR, 2024.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Ar- avind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in newral information processing systems, 34:15084 15097, 2021.
Nicholas E Corrado, Yuxiao Qu, John U Balis, Adam Labiosa, and Josiah P. Hanna. Guided data aug- mentation for offline reinforcement learning and imitation learning. arXio preprint orXiv:2310.182, 18247, 2024.
Xinqi Du, Ziyue Li, Cheng Long, Yongheng Xing, S Yu Philip, and Hechang Chen. Felight: Fairness-aware traffic signal control via samplo-efficient reinforcement learning. IEEE Transactions on Knowledge and Data Engineering, 2024.
15

---

Published in Transactions on Machine Learning Research (08/2025)
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. approximation in reinforcement learning. Neural networks, 107:3-11, 2018.
linear units for neural network function
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data- driven reinforcement learning. arXiv preprint arXiv:2004.072.1, 2020.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132 20145, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp. 2052-2062. PMLR, 2019.
Chen-Xiao Gao, Chenyang Wu, Mingiun Cao, Rui Kong, Zongzhang Zhang, and Yang Yu. Act: Empowering decision transformer with dynamic programming via advantage conditioning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 12127 12135, 2024.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870. Pmlr, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Aduances in neural information processing systems, 33:6840-6855 2020.
Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt-tuning decision transformer with preference ranking. arxiv preprint ar 2rv: 09648, 2023.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205. 09991, 2022.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Flucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35:26565 26577, 2022.
llya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofrr Nachum. Offline reinforcement learning with fisher divergence critic regularization. International Conference on Machine Learning, pp. 5774 5783. PMLR, 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for oflline reinforce- ment learning. Advances in Neural Information Processing Systems, 33:1179-111 1191, 2020.
Guanghe Li, Yixiang Shan, Zhengbang Zhu, Ting Long, and Weinan Zhang, Diffstitch: Boosting oflline reinforcement learning with diffusion- based trajectory stitching. ICML, 2024.
Zijie Li, Henry Li, Yichun Shi, Amir Barati Farimani, Yuval Kluger, Linjie Yang, and Peng Wang. Dual diffusion for unified image generation and understanding. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 2779-2790, 2025.
Cong Lu, Philip J Ball, Tim GJ Rudner, Jack Parker-Holder, Michael A Oshorne, and Yee whye Teh. Challenges and opportunities in offline reinforcement learning from visual observations. Transactions Machine Learning Research, 2023.
Cong Lu, Philip Ball, Yoe Whye Teh, and Jack Parker-Holder. Synthetic experience replay. Advances in Newral Information Processing Systems, 36, 2024.
Jiafei Lyu, Xiaoteng Ma, Xiu Li, and Zongqing Lu. Mildly conservative q-laarning for offline reinforcement learning. Advances in Neural Information Processing Systems, 35:1711-1724, 2022.
16

---

Published in Transactions on Machine Learning Research (08/2025)
Yixiu Mao, Qi Wang, Yun Qu, Yuhang Jiang, and Xiangyang Ji. Doubly mild generalization for offinee
reinforcement learning. Advances in Neural Information Processing Systems, 37:51436-511 51473, 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195-4205, 2023.
Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Newral Information Processing Systems, 33:3976-3990, 2020.
Rafael Figueiredo Prudencio, Marcos ROA Maximo, and Esther Luna Colombini. A survey on offline rein- forcement learning: Taxonomy, review, and open problems. IEEE Transactions on Neural Networks and Learning Systems, 2023.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines, Advances in neural information processing systems, 20, 2007.
Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. arXiv preprinl or Xiv: 110770990000.. 2017.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation, In Medical image computing and computer- Aitervention-MMCCCA 2015: 18th inter- national conference, Munich, Germang, October 5-9, 2015, proceedings, part III 18, pp. 234 241. Springer, 2015.
Jascha Sohl- Doch-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2256-22655 PMLR, 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXio preprinl arXiv:2011.13456, 2020.
Richard S Sutton and Andrew G Barto. Reinforcement leurning: An introduction. MIT press, 2018
Hamid Taghavifar, Chuan Hu, Chongfeng Wei, Ardashir Mohammadzadeh,, and Chunwei Zhang, Behaviorally-aware multi-agent rl with dynamic optimization for autonomous driving. IEEE Transac- tions on Automation Science and Engineering, 2025.
Chen Tang, Ben Abbatemattco, Jiaheng Hu, Rohan Chandra, Roberto Martin-Marth, and Peter Stone. Deep reinforcement learning for robotics: A survey of real-world successes. In Proceedings of the AAAI Conference Artificial Intelligence, volume 39, pp. 28694 28698, 2025.
Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov. Corl: Research-oriented deep offline reinforcement learning library. Advances in Neural Information Processing Systems, 36, 2024.
Lenart Treven, Jonas Hübotter, Florian Dorfler, and Andreas Krause. Efficient exploration in continuous- time model-based reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
Yuch-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances in Neural Information Processing Systems, 36, 2024.
Haoran Xu, Li Jiang, Li Jianxiong, and Xianyuan Zhan. A policy-guided imitation approach for offline reinforcement learning. Advances in Neural Information Processing Systems, 85:4085-40099 2022.
17

---

Published in Transactions on Machine Learning Research (08/2025)
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural informalion processing systems. 34:28954 28967, 2021.
Ziyang Zhai, Ruru Hao, Boyang Cui, and Siyi Wang- Hgat and multi-agent rl-based method for multi- intersoction traffic signal control. IEEE Transactions Intelligent Transportation Systems, 2025.
Huichu Zhang, Siyuan Feng, Chang Liu, Yaoyao Ding, Yichen Zhu, Zihan Zhou, Weinan Zhang, Yong Yu, Haiming Jin, and Zhenhui Li. Cityflow: A multi-agent reinforcement learning environment for large scale city traffic scenario. In The world wide web conference, pp- 3620 0024, 2019.
Junjie Zhang, Jiafei Lyu, Xiaoleng Ma, Jiangpeng Yan, Jun Yang, Le Wan, and Xiu Li. Uncertainly- driven trajectory truncation for data augmentation in offline reinforcement learning. In ECAI 2023, pp. 3018-3025. IOS Press, 2023.
Liang Zhang and Jianming Deng. Data might be enough: Bridge real-world traffic signal control using offline reinforcement learning, arXiv preprint ar Xiv:23303 10828, 2023.
Liang Zhang, Qiang Wu, Jun Shen, Linyuan Lü, Bo Du, and Jianqing Wu. Expression might be cnough: representing pressure and demand for reinforcement learning based traffic signal control. In International PMLR, 2022.
Conference on Machine Learning, pp.
18

---

Published in Transactions on Machine Learning Research (08/2025)
A Related Work
A.1 Offline Reinforcement Learning
Offline RL involves learning policies from pre-collected offline datasets comprising trajectory rollouts gener- ated by behavior policies. This approach is promising because it avoids the costs and risks associated with direct interactions with the environment. Conventional offline RL methods aim to alleviate the distributional shift problem, i.c., a significant drop in performance duc to deviations between the learned policy and the behavior policy used for gemeraling the ofliin data (Hu et al. have been employed, including explicit correction (Xu et al.
2023 To address this issue, various stralegies
2022 such as constraining the policy to a restricted action space (Kumar al.) 2019), and making conservative estimates of the value function (Yu)
et al. 2021 Kumar al. (020), with the aim of aligning the behavior policy with the learned policy.
Some recent studies exploit the strong sequence modeling ability of Transformer models to solve offline RL with trajoctory optimization. For instance, Decision Transformer (Chen et al. 2021 and its variants (Wul et al. 2024 Gao et al., 2024) utilize a GPT model to autoregressively predict actions given the recent subtrajectories composed of historical RTGs, states, and actions. These approaches integrate hindsight return information, i.e., RTG, with sequence modeling, eliminating the necessity for dynamic programming.
Diffusion models have also been adopted in offline RL, given their exceptional capability of multi-modal distribution modeling. Diffuser (Janner et al., 2022) employs diffusion models for long-horizon planning, effectively bypassing the compounding errors associated with classical model-based RL. Hierarchical Diffuser (Chen et al. 2024) enhances this approach by introducing hierarchical structure, specifically a jumpy planning method, to improve planning effectiveness further.
A.2 Data Augmentation in Offline Reinforcement Learning
Rather than passively reusing data and concentrating on enhancing learning algorithms, data augmentation proactively generates more diverse data to improve policy optimization. Some model-based RL methods employ learned world models from offline datasets to sinulate the environment and ileratively generate synthetic trajectories, facilitating policy optimization (Zhang al. 2023). For instance, TATU (Zhang et al. 2023) uses world models to produce synthetic trajectories and truncates those with high accumulated uncertainty. However, model-based RL often suffers from compounding errors in the learned world models. Some basic data augmentation functions (DAFs), i.e., translate, rotate, and reflect, are also applied to augment trajectory segments (Pitis et al.) (020). GuDA (Corrado et al. (024)) further introduces human guidance into these DAFs to improve data quality, while human intervention costly and lacks scalability. Diffusion models are also directly applied to data augmentation through the sampling process. SynthER (Lu et al.] (024) the first work that employs diffusion models learn the distribution of initial offline datasets and unconditionally augment large amounts of new random data. However, it fails control the sampling process to steer toward high-return directions actively. DiffStitch (Li et al. 2024) attempts to enhance the quality of generated data by actively connecting low-reward trajectories to high-reward ones using stitching technique.
We propose enhancing the quality of generated data from a different perspective by introducing a controllable directional goal into our generative modeling. This approach selectively reuses optimal trajectories to guide the sampling process toward achieving higher returns.
B Experiment Details
B.1 Hyperparameters
We show more details about the hyperparameter settings of the GODA model.
19

---

Published in Transactions on Machine Learning Research (08/2025)
Training Data
GODA Data
Training Distribution
GODA Data Distribution
High-Reward Zone
Trajectory
Sample
Figure 7: An illustrative example of how GODA utilizes higher goals to steer the sampling process toward higher-reward data distribution region.
Table 7: Hyperparameters for denoiser network.
Hyperparameter
Value
embedding dimension
128
MLP width
512
MLP activation
SiLU
gate activation
Sigmoid
learning rate
0.0003
batch size
256
learning rate schedule
cosine annealing
optimizer
Adain
gradient update steps
100K
B.1.1 Denoising Network
The denoising neural network utilizes the adaptive gated conditioning architecture, as shown in Figure ref fig:gac. Table 7 details the associated hyperparameters. We use Random Fourier Feature embedding (Rahimi & Recht 2007 and Sinusoidal positional embedding (Vaswani et al. 2017| to process the noise level and timestep of each RTG, respectively, with an embedding dimension of 128. The width of the linear layers in the MLP block is to 512, with SiLU (Elfwing et al.] (018)) as the activation function. The total number of trainable parameters for the denoiser neural network approximately 3.3 M. We train our GODA model with 100K steps of gradient updates, with batch size of 256 and learning rate of 0.0003.
B.1.2 Elucidated Diffusion Model
2022 as our diffusion model and follow the original settings from SynthER LL et al. 2024), with the default hyperparameters shown in Table EDM employs Heun's 2nd order ODE solver [Ascher & Petzold 1998) to solve the reverse-time ODE, enabling data sampling through the reverse process. The diffusion timestep is set to 128 for higher-quality results. All training and sampling conducted on AMD Ryzen 7 7700X 8-Core Processor and single NVIDIA GeForce RTX 4080 GPU. Training GODA
We adopt EDM Karras et al.
Table 8: Hyperparameters for the diffusion model.
Hyperparameter
Value
number of diffusion steps
128
Schrn
80
Stmin
0.05
S Smaa
50
nolae
1.003
T Omin
0.002
Tmax
80
20

---

Published in Transactions on Machine Learning Research (08/2025)
ft
1
(a) Intersection
(b) Phases
Figure 8: A standard signalized intersection with four 3-lane approaches. phase settings.
The right part shows eight different
for 100K steps takes approximately 14 minutes, while generating 5M samples with a sampling batch size of 250K requires about 300 seconds.
B.1.3 D4RL Tasks
In the case of Gym tasks where dense rewards are available, we employ three distinct data configurations from the D4RL datasets: Random, Kedium-RRepayy and Medium. To elaborate, Random datasets contain transition data generated by randomly initialized policy. Medium datasets consist of million data points gathered using& policy that achieves one-third of the performance of an expert approach. Medium-replay datasets contain the stored experience in replay buffer during the training of a policy until it reaches the score Modium datasets. Medium-Expert datasets consist of 50/50 mixture of expert demonstrations and suboptimal trajectories.
For Maze2D, 2D agent is trained reach goal position utilizing minimal feedback, i.e., single point for success, zcro otherwise. Three datasets collected from different maze layouts adopted, i.c., Umaze, Medium, and Large.
The Pen and Door tasks from the Adroit benchmark (Rajeswaran et al. 2017 Fu et al. 2020 involve manipulating pen and opening a door using a 24-DoF simulated hand robot. For each task, we use two different datasets: Human and Cloned. The Human dataset consists of trajectories from human demonstra- tions, while the Cloned datasot generated by applying an imitation policy trained from a mix of human and expert demonstrations and combining the resulting trajectories with the demonstrations in 50/50 split.
We augment 5M samples for each D4RL task. For the sampling process, condition selection method and set the goal scaling factor to 1.1 for all tasks. We use the implementation of IQL and TD3+BC from the Clean Offline Reinforcement Learning (CORL) codebase (Tarasov et al. 2024 for D4RL tasks.
the return-prio goal
B.1.4 TSC Tasks
As shown in Figure g a signalized intersection in TSC problems is composed approaches with several lanes in each approach. The controller manages the phase as shown in the top right part of Figure S which determines the activated traffic signals for different directions, to control the orderly movement. of vehicles.
We formulate the TSC problem MDP and define the state, action, and the reward function follows:
21

---

Published in Transactions on Machine Learning Research (08/2025)
Table 9: Performance comparison on Gym tasks using DMG as the evaluation algorithm. Values in bold represent the best performance (largest normalized score).
DMG
Task
Data
Original
SynthER
GODA
Random
27.3+2.5
27.8
31.5114.3
Medium- Replay
51.4#1.6
47.1#2.4
50.9=1.7
HalfChectah-v2
Medium
55.2+1.8
57.0+0.6
60.5
Medium-Expe Expert
41.3+2.9
44.2111.3
46.21:0..6
Random
5.3+1.2
3.6=1.1
13.24:.99
Medium-Replay
90.9+-4.9
91.2+5.1
94.0=4.6
2alker2D-v2
Medium
90.5+5.0
95.7+2.5
88,2=4.7
Expert
114.0:#0.8
115.3+0.5
112.5++0.8
Random
9.9-F4.44
10.8+3.55
11.2:42.9
Medium Medium-Replay
101.5:1..5
95.4+3.3
99.7=2.8
Hopper-v2
Medium
100.4+1.0
102.2:0.9
102.2
Mediun-Experl
37.61-33.3
39.1
36.5=30.1
Average
60.4
60.845.0
62.2.16.0
State. For behavior policies (AMP and ACL), the state representation includes the current phase, traffic movement efficiency pressure, and the number of effective running vehicles (Zhang ct al. 2022). For eval- uation algorithms, BCQ and CQL use the same state representation as AMP, while DataLight adopts the number of vehicles, along with the total velocity saturation and unsaturation degrees (Zhang & Deng 2023
Action. The action is generally defined as the phase selection for the next time period.
Reward. AMP, BCQ, CQL and DataLight use pressure (Zhang et al. 2022 as the reward while ACL uses queue length. It is worth noting that we use the opposite of these metrics as the final reward function.
We generate a total of 24K samples for each dataset using the behavior policies for each task. Addi- tionally, we augment 20K samples for each task using GODA model. The task horizon for each TSC scenario is set to 3600 seconds, with control step of 15 seconds. For the sampling process, employ the return-prior goal condition selection method and set the goal scaling factor to 1.2 in TSC tasks. For evaluation methods, we employ the implementation of BCQ, CQL, and DataLight from https://github.com/Tiie /iinnghhagg waaaangg996//// (996/DaaaLigh..
C Further Experimental Results
In this section, we show some more experimental results for our GODA model.
C.1 Further Evaluation on Gym tasks
To evaluate the generalizability of our proposed method, we further conduct experiments by evaluating GODA with a state-of-the-art offline RL method, i.c., Doubly Mild Generalization (DMG) (Mao et al. 2024), on Gym tasks. DMG introduces a novel generalization strategy by selecting actions from a const neighborhood of the training data, thereby maximizing Q-values while avoiding excessive extrapolation. DMG also mitigates generalization propagation and aims to balance generalization and conservatism in offline RL.
As can be seen in Table D GODA-cnhanced DMG outperforms both the original DMG and the SynthER- enhanced DMG in 7 out of 12 evaluation scenarios, and achieves the highest average normalized score overall. It's worth noting that we were unable to reproduce the results for BerCCheeah-MMediin Expert and Hopper- Medium-Exper for DMG. SynthER and our GODA also fail to improve the results of DMG on these two datasets to normal levels. We suspect this may be due hyperparameter sensitivity in DMG rather than issues with the datasets themselves. Overall, these results support the claim that GODA can effectively
22

---

Published in Transactions on Machine Learning Research (08/2025)
Algorithm 2 Online Reinforcement Learning with GODA
1: Input: Real-to-synthetic data ratio ç, diffusion model update frequency
2: Initialize policy Tæ: generative model Gø, real replay buffer Dreal < 0, synthetic buffer D + 0
3: for t= to 7 do
4: Deploy policy To to collect real transitions and store them in Dreal
if mod K 0 then
Update diffusion model Gg using samples from Dreat (see Line 2-5 in Algorithm
Generate synthetic samples using Gg and add them to D (soe Line 6-11 in Algorithm
8: end if
9:
Update policy Ts using samples from both Dreal and D with a real-to-synthetic ratio of &
10: end for
HalfCheetah-v2
Walker2d-v2
Hopper-v2
12500
5000
3000
Eval Return
4000
7500
3000
2000
5000
1000
1000
0
60000 Timesteps
80000
100000
20000
40000
60000
80000
100000
20000
40000
100000
Timesteps
Timestepa
SAC
SynthER
GUDA
Figure 9: Online evaluation with SAC Gym tasks. GODA significantly outperforms the baselines in the early stages of training, and continues to match outperform baselines training progresses.
generalize to more advanced offline RI. algorithms and remains promising data augmentation framework for enhancing state-of-the-ar methods.
C.2 Online Evaluation
While primary focus is data augmentation for offline RL, framework can be naturally extended to the online RL setting, similar to the approach laken in SynthER (Lu et al. 2024).
Algorithm shows the pseudocode, In this online extension, we periodically train the diffusion model using transitions collected from ongoing environmental interactions. The learned goal-conditinned diffusion model is then used to augment the dataset with synthetic transitions. The online policy is updated using a hybrid buffer containing both real and synthetic transitions in a 1:1 ratio.
To evaluate this setting, implemented our method using the Soft Actor-Critie (SAC) (Haarnoja et all 2018) am and conducted experiments on three Gym tasks: HalfCheetah, Hopper, and Walker2D. We train SAC with an Update-To- TotData (UTD) ratio of 1 while training SAC (SynthER) and SAC (GODA) with a UTD of 20. For every 10,000 new real data points gathered, we enhance the dataset by augmenting it with million additional synthetic samples.
Bocause the distribution and quality of the real buffer improve over time during online training, we adaptod our hyperparameters accordingly. Specifically, we applied decaying goal scaling factor, from 1.8 (effective in low-quality, random offline settings) to 1.1 (effective in expert-level offline settings), and an increasing Top-n condition selection, from 50 to 120, to reflect growing confidence in higher-quality goals as training progresses.
We compare SAC (GODA) with both vanilla SAC and SAC (SynthER), and present online learning curves on Figure showing evaluation returns over 100K interaction steps. As shown, GODA significantly outperforms the baselines in the early stages of training, where real transitions tend be suboptimal. the later stages, the quality and quantity of real data improve, GODA continues to match slightly outperform SAC
23

---

Published in Transactions on Machine Learning Research (08/2025)
(SynthER), with the performance gap gradually narrowing. These results demonstrate that GODA is not only effective in offline settings, but also promising approach for enhancing online RL through adaptive, goal-conditioned data augmentation.
Table 10:
GODA's sensitivity to dataset complexity. We assess the complexity of offline datasets with several Lipschitz constant metrics and use their 95th percentile values quantify reward sensitivity. The Pen- dataset exhibits the highest sensitivity score, making it particularly challenging for GODA to accurately capture the underlying data distribution and effectively extrapolate to higher-reward regions. Bold values indicate the highost scnsitivity.
Dataset
LocalLip95
GlobalLip95
GlobalMax
RewardStd
TrajLipa8
GODASensitivity
Door- Clor-Cloned
3.361
0.016
5.9
4.0313
7.447
10.16
Door-Human
3.016
1.798
4.8
5.6160
4.757
13.45
Pen-Cloned
29.728
15.896
796.5
28.394
325.811
106.03
Pen-Human
20.626
9.450
1474.1
29.0234
407.446
C.3 Failure Case Analysis of GODA
We observe that GODA fails to perform well on some scenarios, such as the Adroit Pen-Human datasot, as shown in Table Therefore, we develop a comprehensive reward sensitivity measurement to examine which offline RL dataset characteristies make GODA effective problematic. Our method quantifies reward landscape sensitivity using multiple Lipschitz constant metrics across four D4RL Adroit datasots. Local Lipschitz constants are computed using k-nearest neighbors and weighted least squares to estimate how rapidly rewards change with small state-action perturbations at each data point. Gloial Lipschitz analysis calculate pairwise ratios of reward differences to distances across all neighbor pairs to capture worst-case sensitivity spikes throughout the datasct. Trajectory-loe analysis examinc within-trajectory scnsitivity by computing maximum Lipschitz ratios between trajectory steps. These metries were combined into composite GODA Sensitivity weighted by local sensitivity (40%), global sensitivity (30%), reward variability (20%), and trajectory sensitivity (10%), with data normalization and efficient subsampling for computational tractability.
Ir(s, a) r(5',a') (s.a)<(35aa) [[(s, - - (8', a2)ll
T sup
(13)
As can be seen in Table 10 the analysis revealed that task complexity plays an important role on GODA's success. Door manipulation lasks achieved excellent GODA Sensitivity with low local Lipschitz constants, stable global behavior, and manageable trajectory sensitivity, making them highly preferable for GODA. Conversely, pen manipulation tasks exhibited catastrophic sensitivity with GODA Sensitivity exceeding 100, extreme local Lipschitz values, global sensitivity spikes reaching 800-1500, and volatile trajectory patterns, making them fundamentally challenging for GODA. The critical insight is that the difference between Human and Cloned datascts are minimal within the same task type, but inst-intrinsic reward sensitivity creates a significant gap. GODA's 's failure on Pen-Human may also stem from the fact that the original datasel contains only 25 human demonstration trajectories, which lack sufficient diversity for GODA to effectively learn the underlying data distribution and extrapolate out-of-distribution high-reward regions.
C.4 Ablation on Scaling Factor on Different Datasets
To evaluate the practical effectiveness of goal scaling, we conducted additional ablation studies on the Walker2D-Random, Medium, and Medium-Expert datasets.
Figure 10 shows that goal scaling does have tangible impact: On the Random dataset, performance improved significantly with scaling factors of 1.8 for IQL and 1.4 for TD3+BC. For instance, the score gap between À =1.8 and À =1.0 in IQL reached 14.6. On Medium and Medium-Replay, a scaling factor of 1.1 consistently yielded the best results for IQL, with gains up to 4.2 compared to À =1.0, In contrast, for the Medium-Expert. dataset, the difference between X =1.0 and X =1.1 was marginal. This is likely because datasets with more expert demonstrations have less room for improvement, and higher scaling can push goals outside the feasible return range.
24

---

Published in Transactions on Machine Learning Research (08/2025)
Walker2D Random
Walker2D Medium-Replay
40
120
IQL
TD3+BC
IQL
TD3+BC
Normalized Score
30
Normalized Score
110
100
20
90
10
70
0.8
1.0
1.2
1.4
1.6
1.8
2.0
0.8
1.0
1.2
1.4
1.6
1.8
2.0
Scaling Factor
Scaling Factor
Walker2D Medium
Walker2D Medium-Expert
100
IQL
TD3+BC
TD3+BC
Normalized Score
Normalized Score
115
90
80
110
70
105
60
0.8
1.0
1.2
1.4
1.6
2.0
0.8
1.2
1.6
2.0
Scaling Factor
Scaling Factor
Figure 10: Ablation on scaling factor on different datasets. Goal scaling is particularly beneficial in lower- quality datasets while still necessary for higher-quality datasets.
Halfcheetah
Walker2d
Hopper
MIX
100
MIX
100
MIX
AUG
AUG
AUG
40
75
75
50
50
20
25
25
0
Rand
Med-R
Med
Rand
Med-R
Med
Rand
Med-R
Med
Figure 11: Ablation on mixed datascts for IQL evaluation.
These results suggest that goal scaling is particularly beneficial in low-quality datascts, where enables exploration of high-return regions beyond what the original data distribution offers. While À =1.1 performs robustly across all quality levels and is used as default, further gains can be achieved through task-specific hyperparameter tuning.
C.5
Ablation on Mixed Datasets
Since GODA primarily augments samples from high-reward regions of the data distribution, which might result. in a lack of diversily, we use a mix of both the original and augmented datasets for training. In this section, we compare the performance of our default setting (mixed datasets) with the use of only augmented datascts. As shown Figures I1. removing the original datasets leads to slight performance degradation across most tasks. Therefore, combining our augmented datasets with the original datasets helps increase the diversity and extend the reward boundary.
25

---

Published in Transactions on Machine Learning Research (08/2025)
Table 11: Ablation study on mixed datasets for baseline methods.
IQL
TD3+BC
Task
Datasct
TATU
SynthER
DStitch
GODA
TATU
SyuthER DStitch
GODA
w/o Original
46.2+44.3
52.1422.4
51.5+155
52.5+2.99
43.3-4.2
46.3+4.7
47.9+3.3
47.0+4.6
Gym Avg.
w/ Original Data
46.0=3.5
51.2414
52.141.6
54.7=2.2
41.8=1,9
46,241.5
46,840
48.4=3.9
w/o Original Data
45.748.2
45.6422.
47.5+6.8
66.7:99.3
68.2-10.8
65.1429.7
65.0424.9
74.
Maze2D Avg.
w/ Original Data
47.126.8
44.849.2
46.048.4
68.1--6.6
68.0=10.9
65.5437.6
66.14:3377
79.1
BCQ
CQL
DataLight
Average Travel Time (5]
400
Original
Original
Original
GODA
GODA
GODA
350
300
250
5
10
15
20
10 15
20
5
10
15
20
Dataset Size (K)
Dataset Size (K)
Dataset Size (K)
Figure 12: Performance comparison on JN1 task with difference sizes of original real-world training datasets. The parts exceeding the maximum travel time are not displayed.
We further present the results for baseline data augmentation methods when combining original datasets with augmented datasets, shown in Table 11 For Gym tasks, TATU and SynthER cxhibit degraded performance when the original datasets are included during training, while their performance remains com- parable on Maze2D tasks. In contrast, DStitch demonstrates the opposite trend, performing comparably on Gym tasks but worse on Maze2D tasks when using the mixed datasets. Our GODA shows reduced scores when the original datasets are removed; however, its final results remain superior to other baselines. GODA only delivers slightly worse performance than DStitch on Gym tasks when evaluated with TD3+BC, but still outperforms TATU and SynthER across all tasks.
C.6 Ablation on Size of Original Dataset
Our experiments on TSC tasks utilize only 24K samples, compared to 2M samples for D4RL tasks. The superior performance across both task types demonstrates that GODA not only excels in tasks with large training datasets but also addresses real-world challenges elfectively with much smaller datasets. Considering the greater limitations of original datasets in real-world TSC scenarios, we further evaluate GODA TSC datascts with reduced sizes, ranging from 24K to 2.4K samples.
Figure 12 illustrates the performance of different evaluation algorithms trained separately on the original datasets and on the mixture of original and augmented datasets. Notably, we only reduce the size of the original datasets used for directly training the evaluation algorithus or for training GODA, while maintaining the same number of augmented samples (20K) across all cases. Specifically, for the 'Original" method, the evaluation algorithms are trained solely the reduced original datasets. For GODA, the GODA model is first trained on the reduced original datasets, and then 20K augmented samples are generated. The evaluation algorithms are subsequently trained on the mixed datasets.
The results indicate that when the original training datascts contain fewer than 9.6K samples, all three eval- nation algorithms trained directly on the reduced original datasets fail to learn effective policies, though their performance improves as the dataset size increases. In contrast, GODA effectively augments high-quality data, enabling the training of qualified evaluation policies for most sizes of original datasets. Remarkably, GODA achieves this even with only 2.4K samples in the original datasets, failing in just one case cach for BCQ and CQL, while consistently demonstrating strong performance otherwise.
26


Reviews:
[{"review_id": "ImxrTT64qf", "reviewer": "Reviewer_zNDP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The idea of goal-conditioned data generation is meaningful and well-motivated in the RL setting, where high quality transitions could in principle make a difference.\n- EDM is a solid choice for the base diffusion model. The proposed optimizations, such as goal-conditioning, adaptive gated conditioning, and selective goal conditions are interesting and have some degree of novelty.\n- The experiments contained both standard tasks as well as tasks inspired by the real world (TSC). The evaluation shows quite clearly and convincingly that the proposed methods usually outperform other data augmentation methods, under both IQL and TD3+BC. The figures also show how GODA is able to sample from the high-return regions, promoting higher quality but also enhanced diversity.\n- GOAD can perform well even under small datasets, as results on TSC demonstrate.\n- The ablation study is extensive and provides insights into various practical aspects of the proposed framework and algorithms.\n\nWeaknesses\n- There is no theoretical guarantee of any sort. This is to be expected to some extent, given the inherent difficulty and uncertainty of offline RL. That being said, there are settings where GODA performs worse than the state of the art (e.g., on Pen-Androit). It would be great, if the authors could perhaps provide possible insights into the settings that are more challenging to GODA. Perhaps settings where small perturbations in the transition can result in very different rewards, or something along these lines?\n- I was not so clear about the controllable goal scaling part. In practice, it seems values close to 1 are good enough. Does it even make a difference if we just always set it to 1? I mean, intuitively Equation (6) makes sense, but in practice I got the impression that a value of 1 will typically be close to optimal, so that scaling by some $\\lambda$ is not really necessary. Similarly, I was not sure whether the Top n conditions make a big difference, especially given the high variance that we observe in Figure 7.\n- I think it would have been nice to get more explanations/insights into some of the experimental results. For instance: (a) Why the return-prior generally outperforms the RTG-prior? The latter intuitively seems a strong baseline, so what could be the possible reasons for that? (b) Why does DStitch suffer from more diverse datasets? (c) Why does sometimes Med-R perform better than Med, and other times the opposite trend  is observed? For instance, Med-R is worse than Med for HC but better for WK2D (with IQL, Table 4). Do the authors have insights into which type of dataset is preferable and why?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors  propose a goal-conditioned data augmentation (GODA) framework, which augments offline reinforcement learning (RL) datasets with synthetically generated samples of higher quality. The generative model is based on the Elucidated diffusion model (EDM), which additionally conditions on a return-oriented goal. To account for noisy inputs and conditions GODA relies on an adaptive gated conditioning method, which significantly enhances the ability to guide the diffusion and sampling process. The idea is for GODA to learn the distribution of the original offline dataset, while generating new data of higher quality/return. This can improve the performance of offline RL, especially in the presence of small and suboptimal offline datasets. Experiments on the D4RL benchmark and traffic signal control tasks demonstrate the superior performance of the proposed framework compared to the state-of-the-art data augmentation methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I generally do not have major concerns with this work. However, I would appreciate it if the authors could respond to the previous questions, as well as to some new ones:\n- In Table 4, HP with Med-R has extremely different performance under IQL than under TD3+BC. Why is there such a dramatic gap between the two algorithms? Fora RAND setting, I would have understood such a fluctuation, but what about the Med-R setting? Does that perhaps show that Med-R is inherently unstable?\n- I was wondering whether this framework could be used even in the context of online RL to get higher-quality data. The main difference would be that Algorithm 1 would be run periodically, and we would periodically re-train the goal-conditioned neural net. Have the authors considered such a use case?\n- IQL and TD3+BC are not very recent methods. Have the authors considered more recent baselines? It would be interesting to compare the different settings under state-of-the-art offline RL methods.\n- It would have been interesting to see experiments with High-quality datasets, as opposed to Med and Med-R. This would allow one to see whether goal-conditioned data augmentation can raise Med or Med-R performance close to the optimum, or whether the gap remains. Ideally, we would like to reach a performance similar to a high-quality offline dataset after goal-conditioned data augmentation. Is that really what is happening? Insights into that could be very helpful.", "broader_impact_concerns": "I do not have any broader concerns that are not addressed in the manuscript.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "VpoGlovYy3", "reviewer": "Reviewer_ZSer", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper addresses an important problem in offline reinforcement learning which is the learned policy is limited to the quality of the data. The authors introduces GODA, an approach that generate new high-return data to improve the effectiveness of the learned policies. This is relevant for many applications and has a wide audience. I think the paper is generally easy to follow (except the method's section which needs a lot of work to improve its clarity). \n\nThe main issue in the paper is the absence of the limitations section.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The performance of offline reinforcement learning (RL) methods is dependent on the quality of the collected data (in contrast to online RL), which in most cases can lead to suboptimal policies. This paper introduces GODA, an approach that augment the data using the output of a conditional diffusion model that model the distribution of the transitions given some metric to generate high-return data. The authors did a good job in demonstrating the higher data quality generated by GODA and the better performance on various tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- The paper lacks a **discussion on the limitations** of their approach. I think it's a critical to the paper and I request that the authors to add it. For example, after reading the paper, it's not clear to me how general your approach can be. Can it be applied to visual RL datasets or robot datasets? If not, what can possibly be improved or modified to support those complex tasks? etc.\n\n- **Improve clarity of writing of the methods section** to make it easier to follow. This is a critical issue that needs to be addressed.\n\n- I think the paper is generally easy to follow although it's rather long without a good reason for the length. I suggest the authors try to make writing more concise and reduce the size of the manuscript. This is only a recommendation and not a critical issue.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "oytZqcjM3G", "reviewer": "Reviewer_YYPb", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed algorithm, GODA, is a simple and effective method for generating high-quality synthetic data for offline RL. The authors conducted experiments on multiple benchmarks, and GODA has outperformed existing data-augmentation baselines on most of them. They also performed extensive ablation studies, which helped the reader understand the contribution of each component. However, there is room for improvement in terms of clarity and formatting, which I will discuss in detail in the next section. Also, it would be interesting to see the contribution of adaptive gated long skip connection, as the paper only provides a comparison between GODA, adaLN, and in-context conditioning.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a diffusion-based data augmentation algorithm for offline RL called Goal-cOnditioned Data Augmentation (GODA). GODA trains a diffusion model conditioned on return-to-go (RTG) and employs a novel goal-selection mechanism and controllable scaling to guide the model towards producing near-optimal samples with high RTGs. The paper also introduces new neural network architectures, including the adaptive gated long skip connection and the gated residual MLP block. GODA outperforms existing data augmentation offline RL approaches on various benchmarks. Extensive ablation studies were conducted to demonstrate that every component of the algorithm contributes to the overall performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "All of the following are non-critical.\n\n1. The paper is rather long. Please consider moving some unnecessary parts, such as hyperparameter settings, implementation details, and environment details, to the appendix.\n\n2. The last sentence, \"It is worth noting that a ... rewards unchanged[,]\" on page 6 is difficult to understand.\n\n3. In computing the L2 distance metric (11), it is unclear how the two samples $(s_i, a_i)$ and $(\\bar{s}_i, \\bar{a}_i)$ are selected. Lu et al. (2024) appear to compute the minimum distance, but the paper does not explicitly state this.\n\n4. The authors have used textual citations for all of their citations, but most of them should be parenthetical instead. For example, in the first sentence of the Introduction, \"Reinforcement Learning Sutton & Barto (2018)\" should be changed to \"Reinforcement Learning (Sutton & Barto, 2018)\".\n\n5. In the second last sentence on page 7, \"... introduce a condition-adaptive gated residual connection to further enhances the ...\" should be changed to \"... introduce a condition-adaptive gated residual connection to further [enhance] the ...\"", "broader_impact_concerns": "I do not have any concerns about the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)