GTA: Generative Trajectory Augmentation with Guidance for Offline Reinforcement Learning

Thank you for your valuable feedback, which can enhance our manuscript.

(Weakness 1) Lack of theoretical guarantee to support the algorithm: The authors only provide empirical results, lacking theoretical analysis of the performance of the proposed method.

While existing work [1] proposes theoretical bounds between conditioning value and generated value, these proofs rely on stringent assumptions, which may not align with complex tasks in Offline RL benchmarks. Empirically, on GTA, we conducted an analysis of the correlation between the conditioning value and the rewards of generated trajectories. As illustrated in Figure 1, attached in pdf, we observed Pearson correlation coefficients of 0.55, 0.91, and 0.99 for halfcheetah-medium, medium-replay, and medium-expert datasets, suggesting that GTA can generate trajectories aligned with the conditioned return. Augmented datasets with high-rewarding trajectories can aid offline RL policies in reaching higher performance.

GTA has demonstrated consistent performance gains through extensive experiments across various environments. Thorough ablation studies have further proven the effectiveness of our proposed method. We would like to emphasize the empirical results presented in our paper, which underscore the effectiveness of our approach.

[1] Yuan, Hui, et al. "Reward-directed conditional diffusion: Provable distribution estimation and reward improvement." Advances in Neural Information Processing Systems 36 (2024).

(Question 1) How to select the hyperparameter $\alpha$ in the denoising with amplified return guidance in equation 8? (Question 3-1) Is the proposed method sensitive to the hyperparameter $\mu$?

As discussed in the general response, GTA outperforms other baselines in locomotion tasks even with a single hyperparameter configuration ($\mu=0.25, \alpha=1.1$), demonstrating the robustness of our method. We also provide a general recipe for selecting hyperparameters for new tasks: Increasing $\mu$ and $\alpha$ for datasets with low quality to promote further performance gain.

(Question 2) In the generation phase, do you generate whole trajectories or generate subsequences of trajectories?

Instead of generating the entire trajectory, we focus on generating subtrajectories. Please refer to Table 8 in Appendix E.1 for more details about subtrajectory generation.

(Question 3-2) Based on the description in Figure 4 and Appendix F.3, when \mu = 1, it only performs well on the dataset with low quality. Does GTA have trouble preserving and reconstructing the high-rewarding trajectories?

In addition to the dynamic MSE and oracle reward presented in Appendix F.4, we include data quality metrics for trajectories generated with a higher $\mu$ in the table below. The results indicate that while a $\mu$ of 0.75 results in a higher dynamic MSE than a $\mu$ of 0.25, it still maintains a lower value than SynthER. This demonstrates that GTA does not encounter significant issues in reconstructing high-rewarding trajectories.

Table 1. Dynamic MSE of generated data on halfcheetah-medium-expert-v2.

(Question 4) Can you provide comparison results with diffusion-based offline RL methods? For example, DecisionDiffuser and AdaptDiffuser.

We provide a comparison between diffusion planners [2, 3, 4] and GTA in the table 1, 2 attached in PDF. As shown in the table, GTA outperforms diffusion planners in locomotion tasks and demonstrates significant efficiency in test-time evaluation.

As mentioned in related work section, our method resembles diffusion planners in terms of generating trajectories with high return guidance. However, adopting diffusion as an augmentation method and utilize the model-free offline RL policies can transfer extensive computational burden of diffusion models from the evaluation stage to the data preparation stage. It can bypass the crucial challenge of diffusion planners, which requires extensive cost for sampling action.

[2] Janner, Michael, et al. "Planning with diffusion for flexible behavior synthesis." arXiv preprint arXiv:2205.09991 (2022).

[3] Ajay, Anurag, et al. "Is conditional generative modeling all you need for decision-making?." arXiv preprint arXiv:2211.15657 (2022).

[4] Liang, Zhixuan, et al. "Adaptdiffuser: Diffusion models as adaptive self-evolving planners." arXiv preprint arXiv:2302.01877 (2023).

(Question 5) What is the baseline algorithm in the ablation study of reweighted sampling?

We choose TD3BC for the ablation study of reweighted sampling because TD3BC requires minimal computational cost along the offline RL baselines (IQL, CQL, MCQ).

(Question 6) Can you also apply the reweighted sampling techniques for other baseline methods?

We conduct experiments on the effect of reweighted sampling for other baselines. Specifically, we sample transition from the buffer with probability proportional to its reward during policy training. Table 1 describes the performance of reweighted sampling on other baselines. We find that baselines with reweighted sampling underperform GTA.

Table 1. Experiment results of reweighted sampling on baselines.

Thank you for your critical reviews and valuable feedback.

(Weakness 1) The proposed approach is mainly a direct combination of current techniques.

Please note that we propose the hypothesis that augmenting the offline dataset with feasible and high-rewarding trajectories would boost the performance of offline RL algorithms. To achieve this, we made a novel combination of various techniques. GTA involves (1) generating trajectory-level data to capture sequential relationships between transitions and long-term transition dynamics, (2) introducing partial noising to control the exploration level of generated trajectories, and (3) guiding diffusion model with amplified return to find high-rewarding trajectories. We effectively balance exploration and exploitation with two strategies carefully tailored for diffusion models, as both strategies can seamlessly be integrated with the forward and reverse process of diffusion models. Our method consistently outperforms several baselines across different environments, highlighting the effectiveness of our proposed method. We want to emphasize that numerous machine learning studies continue to explore innovative combinations of established techniques.

(Question 1) In section 4.2, by denoising with amplified return guidance, the authors claim the trajectory can be guided towards high-rewarding regions, are there any theoretical guarantee for this statement?

There is existing work [1] proposing theoretical bounds between conditioning value and generated value. However, these proofs rely on stringent assumptions, making it challenging to directly apply them to GTA. Meanwhile, related works with return guidance [2, 3, 4] show that the conditionally sampled trajectories yield high rewards, demonstrating the effectiveness of return-guidance.

To verify that our conditional sampling indeed generates high-reward trajectories, we conducted analysis on the correlation between the conditioning value and the rewards of generated trajectories. As illustrated in Figure 1 attached in pdf, we observed Pearson correlation coefficients of 0.55, 0.91, and 0.99 for halfcheetah-medium, medium-replay, and medium-expert dataset, respectively. This indicates that diffusion model effectively samples trajectories aligned with the conditioned return.

[1] Reward-directed conditional diffusion: provable distribution estimation and reward improvement

[2] Janner, Michael, et al. "Planning with diffusion for flexible behavior synthesis." arXiv preprint arXiv:2205.09991 (2022).

[3] Ajay, Anurag, et al. "Is Conditional Generative Modeling all you need for Decision Making?." The Eleventh International Conference on Learning Representations.

[4] Liang, Zhixuan, et al. "AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners." International Conference on Machine Learning. PMLR, 2023.

(Question 2) How about the whole computation and time cost of the proposed approach when it is compared with other diffusion-free baselines?

As mentioned in general response, we compare the whole computation and time cost of the proposed approach with other diffusion-free baselines. For model-based methods, we need to train dynamic model and generate synthetic trajectories during policy training.

The table 2 attached in PDF summarizes the whole training and evaluation time cost of algorithms. While training dynamic models takes shorter time than training diffusion models, model-based RL takes much longer time for training policy as it rollouts synthetic trajectories during training.

(Question 3) As shown in the experiments, for different environments, different noise levels are required for the optimal performance, then how can we determine the noise level for a new task?

As mentioned in the general response, we find that even if we fix $\mu$ across different environments, we achieve higher performance than other baselines. We also observe that we can improve the performance further if we increase $\mu$ for exploration when we have low-quality dataset. Therefore, we recommend selecting $\mu$ based on dataset quality for a new task instead of extensive online tuning.

(Question 4) The proposed approach aims to deal with trajectories with reward signals, however, in reality, it may be difficult to obtain such signals, can the proposed approaches be extended to the cases where no reward signals can be provided?

Numerous offline RL methods [5, 6, 7, 8, 9] aim to learn decision-making policy that maximizes expected cumulative rewards. In this context, we want to emphasize that GTA focuses on enhancing the optimality of offline datasets while minimizing the degradation of dynamic plausibility to improve the performance of offline RL algorithms. Therefore, generating trajectories without any reward signal is beyond our scope, as they cannot be augmented toward high-rewarding trajectories and are not the objective of RL problems. Please acknowledge that we conduct extensive experiments on realistic and challenging tasks such as AntMaze, FrankaKitchen, Adroit, and pixel-based environments.

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

[6] Kumar, Aviral, et al. "Conservative q-learning for offline reinforcement learning." Advances in Neural Information Processing Systems 33 (2020): 1179-1191.

[7] Kostrikov, Ilya, Ashvin Nair, and Sergey Levine. "Offline Reinforcement Learning with Implicit Q-Learning." International Conference on Learning Representations.

[8] Lyu, Jiafei, et al. "Mildly conservative q-learning for offline reinforcement learning." Advances in Neural Information Processing Systems 35 (2022): 1711-1724.

[9] Yu, Tianhe, et al. "Combo: Conservative offline model-based policy optimization." Advances in neural information processing systems 34 (2021): 28954-28967.

Thank you for your positive review and valuable feedback!

(Weakness 1-1) Is there a principled argument on why diffusion models would learn the dynamics of the environment well? Especially where the goal is to create transitions outside the dataset. (Question 1-1) How do you ensure that the generated trajectories are plausible and adhere to the environment's dynamics?

As you correctly pointed out, diffusion model of GTA does not explicitly ensure dynamic plausibility. Instead, it implicitly learns the dynamics governing the trajectory through learning the data distribution.

GTA concentrates on enhancing the optimality of the dataset while reducing the degradation of the environment's dynamics of the augmented dataset to improve the performance of offline RL algorithms. Therefore, we will adjust our statement from "ensuring dynamic plausibility" to "minimizing the degradation of dynamic plausibility". We apologize for the overclaiming of our proposed method.

However, our method still has the potential to generate high-rewarding and dynamically plausible trajectories by virtue of two novel strategies: partial noising and denoising framework and amplified return guidance. As discussed in Appendix F.1, amplified return conditioning prevents the partially noised trajectories denoise conditioning on excessively high value, yielding lower dynamic MSE. Additionally, table 1 shows the dynamic mse of the generated dataset according to the differing level of $\mu$. The result demonstrates that low $\mu$ does help to lose dynamic information less.

Table 1. Dynamic MSE of generated data on walker-medium-v2.

To sum up, even though GTA does not explicitly enforce the model to generate dynamically plausible dataset, we introduce partial noising and denoising framework and amplified return guidance to minimize the degradation of dynamics of generated data.

(Weakness 1-2, Weakness 2) Figure 5 does not conclusively show that GTA is dynamically plausible, especially for tasks like Cheetah compared to methods like S4RL, yet GTA still achieves better rewards. It is unclear why this does not affect the final result.

As the reviewer noted, the trajectories generated by GTA exhibit larger dynamic MSE than S4RL in some cases like Halfcheetah. However, we observe that the dynamic plausibility of GTA is still better than other baselines such as SynthER. More importantly, GTA augments trajectories with higher rewards, which significantly affects the performance of offline RL policies.

(Weakness 3) The performance of GTA appears to depend heavily on finely-tuned hyperparameters. (Question 2) Does each task require a new &mu, &alpha?

To address the concern that the reviewer mentioned, we compared GTA with other baselines on single set parameter($\alpha=1.1, \mu=0.25$). The results demonstrate that fixed parameter configurations generally show performance gain across locomotion tasks, indicating that it may not require extensive online tuning. Furthermore, we propose a guideline for determining the appropriate $\mu$ and $\alpha$ for new tasks to boost performance. Please refer to the general response for a more detailed description.

(Question 1-2) What mechanisms to pick &alpha to handle cases where the return used for high-reward guidance is not reasonable? How do you ensure that the generated trajectories remain valid and useful?

As you mentioned, excessively guiding towards high rewards can result in invalid trajectories. This phenomenon occurs, as shown in Table 17 in Appendix F.1, when we generate trajectories conditioned on the maximum return that can be achieved from the environment. To prevent conditioning on unreasonable return, we introduced amplified return guidance, which involves multiplying the original trajectory by $\alpha$ for conditioning. It controls the exploitation level of GTA, and we observe that our method empirically generates valid and useful trajectories to improve the performance of RL algorithms.

(Question 3) What is the proportion of augmented trajectories to original ones? Is the entire offline dataset added augmented trajectories on top? Do all the offline datasets have the same size? Are 5M augmented transitions added regardless of the original dataset's size?

GTA augments 5M transitions for D4RL, and 1M transitions for VD4RL regardless of different sizes of original datasets. We follow the procedure of the prior method SynthER for a fair comparison. After the augmentation, we added augmented trajectories to the original dataset. There is no reason to exclude the original dataset, as it preserves the true dynamics of the environment. Please refer to Table 10 of Appendix G.4 for more discussion on the size of the augmented dataset.

(Question 4) Could you clarify if Figure 3 is based on a real example or if it is just an illustration of the method's intuition? If it is the latter, how can you ensure it accurately represents the method's practical application?

While Figure 3 of section 4.2 is presented for intuitive comprehension, our empirical results of table 2 clearly demonstrate the relationship between $\mu$, oracle reward, and the deviation from the original trajectory. Table 2 shows that both oracle reward and deviation from the original dataset tend to increase as $\mu$ becomes larger. These results align well with Figure 3, which illustrates that trajectories deviate from the original trajectories and towards high-rewarding regions as $\mu$ gets larger.

Table 2. Analysis of dataset generated by GTA with different $\mu$.

Thank you for your positive review and valuable feedback!

(Weakness 1) The authors could cite Ball et al 2021 "Augmented World Models" as another example of data augmentation in offline RL.

Thank you for pointing out crucial paper. The paper [1] augments learned model with simple transformations and approximate the augmentation for unseen environment in a self-supervised fasion for zero-shot generalization. While we focus on a single task offline RL problems, it is highly related in terms of data augmentation methods for offline RL problems. We will cite the paper in the related work section.

(Weakness 2) The authors could discuss how this relates to another paper building on SynthER, Policy-Guided Diffusion by Jackson et al. This is very recent work so makes sense it is not needed as a baseline, but the comparison may be worth including in text.

Thank you for highlighting an exciting concurrent work. Policy-Guided Diffusion (PGD) [2] generates synthetic trajectories for augmentation with classifier guidance from the target policy. On the other hand, GTA generates trajectories from the original trajectory, and introduce partial noising and denoising framework. To generate high-rewarding trajectories, GTA uses amplified return guidance. While the method is slightly different, both papers share ultimate objective, generating high-rewarding trajectories while retaining low dynamics error to improve offline RL algorithms. We will add comparison with PGD in the manuscript.

[1] Ball, Philip J., et al. "Augmented world models facilitate zero-shot dynamics generalization from a single offline environment." International Conference on Machine Learning. PMLR, 2021.

[2] Jackson, Matthew Thomas, et al. "Policy-guided diffusion." arXiv preprint arXiv:2404.06356 (2024).

Thank you for the insightful review!

(Weakness 1, Limitation 3) The performance of GTA is sensitive to the choice of hyperparameters, which might require extensive tuning for different tasks.

As we illustrated in tables in general response, we demonstrate that GTA outperforms other data augmentation baselines with a single set of hyperparameters across locomotion tasks. This result indicates that we may not need to extensively tune such hyperparameters. Furthermore, we also provide a general recipe for selecting hyperparameters for a new task: choosing high $\mu$ and $\alpha$ for low-quality datasets to promote exploration.

(Weakness 2-1) Dynamic plausibility is not explicitly enforced in GTA and might be a byproduct of generating high-rewarding trajectories.

As you mentioned, GTA does not explicitly enforce dynamic plausibility. Instead, it implicitly learns the dynamics through the trajectories which inherit true dynamics of environment, resulting in the generation of plausible trajectories.

We want to clarify that dynamic plausibility is not a byproduct of generating high-rewarding trajectories. Table 1 compares the dynamic MSE of total augmented trajectories to that of high-rewarding trajectories (top 10% in terms of reward). We found that the dynamic MSE of both cases are similar, validating our claim.

Table 1. Dynamic MSE of data generated with GTA using halfcheetah-medium-v2

(Weakness 2-2, 3, Limitation 2) GTA lacks validation in real-world applications where the dynamics and data distributions may differ significantly.

Thank you for suggesting problem setting that enhaces our manuscript . To verify the effectivenes of GTA in real-world data, we conduct further experiments on the NeoRL benchmark [1]. NeoRL is composed of datasets with narrower data distributions, enforced stochasticity, and aleatoric uncertainty to evaluate how offline RL algorithms behave in realistic environments.

We train GTA on the halfcheetah-v3-medium-noise-1000 and halfcheetah-v3-low-noise-1000 of NeoRL benchmark to evaluate the performance. As shown in table 2, GTA enhances the performance of the base algorithm, indicating that GTA can also generate trajectories with real-world datasets.

Table 2. Experiment results on NeoRL benchmark. Experiments are conducted with 4 random seeds. Base algorithm is CQL.

[1] Qin, Rong-Jun, et al. "NeoRL: A near real-world benchmark for offline reinforcement learning." Advances in Neural Information Processing Systems 35 (2022): 24753-24765.

(Weakness 4) The paper does not include any experiments comparing GTA with model-based RL baselines.

Thank you for your insightful comment. As illustrated in the general response, we compare offline model-based RL methods, such as MOPO, MOReL, and COMBO, with GTA. We take the results of baselines from their original papers. For GTA, we report the score of TD3BC and CQL as base offline RL algorithms.

As shown in Table 1 of the attached PDF, GTA outperforms offline model-based RL algorithms in locomotion tasks. It suggests that using a diffusion model to augment trajectories is more beneficial than learning and exploiting a single-transition dynamic model [2, 3].

[2] Janner, Michael, et al. "Planning with Diffusion for Flexible Behavior Synthesis." International Conference on Machine Learning. PMLR, 2022.

[3] Jackson, Matthew Thomas, et al. "Policy-guided diffusion." arXiv preprint arXiv:2404.06356 (2024).

(Question 1) Why is IQL the only chosen baseline for the Adroit and FrankaKitchen?

For Adroit and FrankaKitchen tasks, IQL is the most reliable baseline among offline RL algorithms we adopted (TD3BC, MCQ, CQL). TD3BC exhibits near-zero performance on these tasks. MCQ does not provide any hyperparameter setting for these tasks, while the main hyperparameter $\tau$ in MCQ significantly influences the performance of the algorithm. Although CQL offers hyperparameter settings for these tasks, its training is notoriously unstable, leading to significant performance fluctuations.

(Question 2) Why is DrQ+BC the only chosen baseline and Cheetah-run the only chosen pixel-based observation task in Table 3? Why were results on the other baselines (IQL, TD3-BC, CQL, and MCQ) not presented for the pixel-based task(s)?

To evaluate GTA on more diverse algorithms and environments, we additionally conduct experiments on walker-walk environment and add BC algorithm for baseline as done in [4]. As shown in table 3, we achieve better performance than other baselines across both environments and algorithms, indicating the generalizability of our method.

Table 3. Experiment results on pixel-based observation tasks. The mean scores of table are calculated with 3 seeds.

[4] Lu, Cong, et al. "Synthetic experience replay." Advances in Neural Information Processing Systems 36 (2024).

(Question 3) How compute-intensive (GPU or hours) is running GTA on the pixel-based observation tasks?

Please note that the diffusion model for pixel-based observation tasks generates a latent vector rather than raw pixel observation. In a single RTX 3090, training GTA in a pixel-based environment for 500k gradient steps requires 50 hours.

We promise to fix minor issues raised by the reviewer of our manuscript.

Thank you for your detailed response and for providing additional supporting and clarifying empirical results. In line with reviewer WvGz, I am still skeptical about the paper's claims regarding GTA ensuring "dynamic plausibility". I would be more comfortable if such a claim is toned down; for instance, augmented trajectories generated from benchmark datasets using GTA exhibit dynamic plausibility.

I am satisfied with the author's responses to the rest of my questions and concerns and have raised my rating from 5 to 6.

We recently noticed that Table 3 was mistakenly uploaded as a duplicate of Table 1 in our general response. We sincerely apologize for any confusion this may have caused the reviewers and the area chair. We have now uploaded the correct Table 3, which shows a performance boost on medium-replay quality datasets, and we kindly ask you to review the updated table. Once again, we apologize for the inconvenience.

Table 3. D4RL normalized score on medium-replay quality locomotion environments with fixed $\alpha$ and $\mu$. The experiments are conducted with TD3BC.