ATraDiff: Accelerating Online Reinforcement Learning with Imaginary Trajectories

4. Why projecting states into images

   • We would like to emphasize that our work develops a general plug-in data generation method to facilitate online RL algorithms, agnostic to the specific form of generation, which could either be pixel/image-level based or state-level based. For example, the important deployment method of our generator and online adaptation technique are agnostic to the form of generation. In our paper, we have investigated the two variants of our data generation method: at image or state levels, as shown in Fig. 6. And we found that 1) both image and state generation variants significantly outperform the baselines; 2) image generation works better under our paradigm. Our hypothesis for the second finding is that image generation can benefit from a pre-trained generation model (e.g., Stable Diffusion) while state generation is trained from scratch, though image generation needs additional modules to convert between images and states.

   • In principle, our method could still be applicable to those vision-based reinforcement learning tasks as suggested by the reviewer, not only restricted to the benchmarks we investigated in the paper. Following previous work [Ref2], this paper mainly focuses on traditional offline RL datasets like D4RL. Evaluating these vision-based tasks is an interesting direction for future work.

5. Incomplete experimental descriptions

   • For Fig. 3 (formerly Fig. 2 in the reviewer’s original comment – we included another figure in the updated manuscript), in the original manuscript, we followed prior studies [Ref4] and presented the average performance across different datasets from the D4RL benchmark. To provide a more comprehensive understanding, we have now included separate results for datasets in Section C.3 of the updated appendix.

   • In Fig. 6 (formerly Fig. 5 in the reviewer’s original comment – we included another figure in the updated manuscript), we conducted experiments on the D4RL Locomotion environments. We have included the missing description to the caption in the updated manuscript.

6. Unclear description of sampling scheme

   • We apologize for any confusion. We have thoroughly revised Section A.1 of the updated appendix, providing detailed explanation on the pick-up strategy. We also elaborate on the various indicators used and how they influence sample selection in our model. Below we briefly summarize the key points:

   • We investigated different indicators to represent the importance of the real samples collected from interacting with the environment when running online RL methods. These indicators are designed to discern different properties, influencing which samples are considered more crucial for effective learning and adaptation.

   • To ensure that the selected samples used to adapt the generator are suitable, we developed 'pick-up strategies' according to different indicators. These strategies involve a systematic approach to maintaining a data structure that stores real samples and selecting the most appropriate ones for online adaptation. The intention is to ensure that our model utilizes the most valuable samples, as indicated by our chosen metrics, to optimize learning outcomes.

   • The TD-error indicator uses the TD-error of a transition to represent its importance. It has the property of flexible importance during the training process. Therefore, we use a large buffer to store more samples, and select samples with the most importance each time when we need to update. In contrast, the reward indicator uses the total reward of a trajectory to represent its importance, which has the property of fixed importance. To avoid using fixed samples to update the generator too many times, we have designed a small buffer with a random dropping strategy.

7. Description in Appendix A.2

   • We apologize for this and appreciate the reviewer for pointing it out. We have rewritten this section as well as improved its clarity.

   • We hope that our clarifications have addressed the reviewer’s concern. We will also release our code upon acceptance. If there are any remaining details or concerns that require further clarification, please do not hesitate to let us know. We are committed to ensuring that our research is transparent and accessible.

References:

[Ref1] Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control. CoRL 2022.

[Ref2] Cong Lu, Philip J. Ball, Yee Whye Teh, and Jack Parker-Holder. Synthetic Experience Replay. NeurIPS, 2023.

[Ref3] Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, Xuelong Li. Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning. NeurIPS, 2023.

[Ref4] Philip J. Ball, Laura M. Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. ICML 2023.

2. Benefits of trajectories vs. transitions: Comparison with SynthER

   • First, we would like to clarify that the task settings differ between our work and SynthER. Our work primarily focuses on leveraging offline data to enhance online reinforcement learning performance. In contrast, SynthER [Ref2] aims to upsample the offline dataset using their data synthesizer. A crucial distinction lies in their methodology for online experiments, where they do not integrate the offline dataset, which is a key component in our approach. This fundamental difference in task settings initially led us to exclude a direct comparison with SynthER.

   • Additionally, at the time of our submission to ICLR, SynthER was not yet published, appearing later in NeurIPS 2023. Given these circumstances, a direct comparison was not feasible in our initial submission.

   • Following the reviewer’s suggestion, we now include comparisons between our method and SynthER. Consistent with the discussion above, we consider two types of comparisons: (I) Employ SynthER in our task setting, which we care most about. To this end, we have conducted a comparative analysis in an online setting within the Meta-World Environments. This analysis is crucial to demonstrate the adaptability and effectiveness of our approach in dynamic and complex settings. Our results indicate a significant performance advantage of our method over SynthER in these online environments, as summarized in the table below. For more detailed results, including full performance curves, please refer to Section C.5 of the updated appendix.

   Task Name	SPiRL+SynthER	SPiRL+ours	RLPD+SynthER	RLPD+ours
   multitask	0.45	0.61	0.52	0.75
   sequencetask	0.49	0.57	0.62	0.80

   • (II) Employ our method in SynthER’s task setting, which is not the focus of our paper. To this end, we have included an additional comparison in Section C.1 of the updated appendix. This comparison, conducted in the D4RL Locomotion environment under offline settings, demonstrates that our method can achieve comparable or superior results to SynthER.

   Task Name	TD3+BC	TD3+BC+SynthER	TD3+BC+ours	IQL	IQL+SynthER	IQL+ours
   halfcheetah-random	11.3	12.2	12.5	15.2	17.2	17.1
   halfcheetah-medium	48.1	49.9	52.3	48.3	49.6	53.1
   halfcheetah-replay	44.8	45.9	46.5	43.5	46.7	49.2
   halfcheetah-expert	90.8	87.2	93.6	94.6	93.3	95.2
   hopper-random	8.6	14.6	15.2	7.2	7.7	8.1
   hopper-medium	60.4	63.4	65.7	62.8	72.0	72.4
   hopper-replay	64.4	53.4	64.7	84.6	103.2	103.6
   hopper-expert	101.1	105.4	111.2	106.2	90.8	113.6
   walker-random	0.6	2.3	2.1	4.1	4.2	4.3
   walker-medium	82.7	84.8	87.5	84.0	84.7	89.1
   walker-replay	85.6	90.5	86.3	82.6	83.3	85.4
   walker-expert	110.0	110.2	111.2	111.7	111.4	111.7

   • In summary, our aforementioned experimental results validate that, in both our task setting and SynthER’s task setting, our trajectory-level generation outperforms the transition-level generation of SynthER.

3. Benefits of trajectories vs. transitions: Difference from [Ref3]

   • Our work is different from [Ref3] in two important ways: 1) [Ref3] actually addresses a task different from our work. They focus on the planning phase, while we focus on the online training aspect. 2) Different from the one-time generation in their work, we leverage multiple generations from different start points with different tasks which deal with the unique challenges associated with flexible generation and distribution shift. We have discussed these differences in the Related Works section.

Thanks a lot for your insightful and inspiring comments! We provide the following clarifications in response to your concerns:

1. Incomplete description of algorithm

   • Thank you for pointing out the missing details in our methodology. To address this, we have revised our manuscript thoroughly and included the previously omitted methodology and implementation details, including those pointed out by the reviewer (e.g., the conversion method between images and states) and other details (e.g., the architecture of the generation models and the pick-up strategy). These details are discussed in the updated main paper, and a detailed explanation is available in Section A of the updated appendix. In addition, we will release our code upon acceptance. We hope that the revision improves the clarity and provides a thorough understanding of our method and procedure. For ease of reference, below we briefly summarize the missing details raised by the reviewer, with more in-depth information available in the revision.

   • (converting between images and states) For the renderer converting states to images, we follow the previous work [Ref1]. This involves running the agent in the environment and capturing observations as rendering results. The recognizer, which converts images back to states, comprises a general convolutional component and separate MLPs. These MLPs are tasked with estimating the 3D positions of key points, tailored to different environments.

   • (Pictures of rendered states) In Appendix D of the updated appendix, we have included representative visualizations of our generated image trajectories. We present a variety of images generated under different conditions, including varying initial states and a range of tasks. The results demonstrate the robustness and effectiveness of our image generation method, consistently performing well across diverse scenarios and producing high-quality and temporally coherent image trajectories.

   • (Conversion error between images and states) 1) The error of converting states to images and then back to states is relatively small. For example, during the training of our recognizer in Meta-World, we found that the loss, which represents the error, is less than $0.054$. 2) Important, we would like to point out that in our case, we are using the synthesized data to train RL methods, not for planning. This means that we do not require the conversion to be very accurate. Instated, we can tolerate the errors within this process and even benefit increased robustness from this.

   • (Description of the generation model and fine-tuning Stable Diffusion) We have set the generation lengths for multiple generators at 5, 10, 15, 20, and 25, respectively. For fine-tuning each generator corresponding to length $k$, we use the offline dataset to generate training data, creating sub-trajectories from points in full trajectories. Sub-trajectories shorter than length $k$ are padded to length $k$ using the last frame. We use LoRA to fine-tune the Stable Diffusion based generators. Additionally, we have conducted an experiment on the wall-clock running time of our generator, detailed in Section C.4 of the updated appendix. We would like to note that in the context of online reinforcement learning, the most significant cost often lies in interacting with the real environment and data collection. Consequently, our primary focus with ATraDiff has been on optimizing sample efficiency to minimize the time required for environmental interaction.

   • (Action and reward generation) We would like to clarify that the action and reward are not generated directly from images. Instead, they are derived from states obtained via the recognizer. More concretely, actions are calculated based on adjacent states, specific to the task, while rewards are generated through interaction with the ground-truth environment. Note that in the context of online training, RL methods inherently involve accessing the environment. Additionally, in the updated manuscript, we have also included a variant of our method that trains a reward prediction function from an offline dataset. The effectiveness of this reward predictor has been demonstrated in our offline experiments, the results of which are included in Section C.1 of the updated appendix.

4. Lack of comparison with SynthER

   • First, we would like to clarify that the task settings differ between our work and SynthER. Our work primarily focuses on leveraging offline data to enhance online reinforcement learning performance. In contrast, SynthER aims to upsample the offline dataset using their data synthesizer. A crucial distinction lies in their methodology for online experiments, where they do not integrate the offline dataset, which is a key component in our approach. This fundamental difference in task settings initially led us to exclude a direct comparison with SynthER.

   • Additionally, at the time of our submission to ICLR, SynthER was not yet published, appearing later in NeurIPS 2023. Given these circumstances, a direct comparison was not feasible in our initial submission.

   • Following the reviewer’s suggestion, we now include comparisons between our method and SynthER. Consistent with the discussion above, we consider two types of comparisons: (I) Employ SynthER in our task setting, which we care most about. To this end, we have conducted a comparative analysis in an online setting within the Meta-World Environments. This analysis is crucial to demonstrate the adaptability and effectiveness of our approach in dynamic and complex settings. Our results indicate a significant performance advantage of our method over SynthER in these online environments, as summarized in the table below. For more detailed results, including full performance curves, please refer to Section C.5 of the updated appendix.

   Task Name	SPiRL+SynthER	SPiRL+ours	RLPD+SynthER	RLPD+ours
   multitask	0.45	0.61	0.52	0.75
   sequencetask	0.49	0.57	0.62	0.80

   • (II) Employ our method in SynthER’s task setting, which is not the focus of our paper. To this end, we have included an additional comparison in Section C.1 of the updated appendix. This comparison, conducted in the D4RL Locomotion environment under offline settings, demonstrates that our method can achieve comparable or superior results to SynthER.

   Task Name	TD3+BC	TD3+BC+SynthER	TD3+BC+ours	IQL	IQL+SynthER	IQL+ours
   halfcheetah-random	11.3	12.2	12.5	15.2	17.2	17.1
   halfcheetah-medium	48.1	49.9	52.3	48.3	49.6	53.1
   halfcheetah-replay	44.8	45.9	46.5	43.5	46.7	49.2
   halfcheetah-expert	90.8	87.2	93.6	94.6	93.3	95.2
   hopper-random	8.6	14.6	15.2	7.2	7.7	8.1
   hopper-medium	60.4	63.4	65.7	62.8	72.0	72.4
   hopper-replay	64.4	53.4	64.7	84.6	103.2	103.6
   hopper-expert	101.1	105.4	111.2	106.2	90.8	113.6
   walker-random	0.6	2.3	2.1	4.1	4.2	4.3
   walker-medium	82.7	84.8	87.5	84.0	84.7	89.1
   walker-replay	85.6	90.5	86.3	82.6	83.3	85.4
   walker-expert	110.0	110.2	111.2	111.7	111.4	111.7

   • In summary, our aforementioned experimental results validate that, in both our task setting and SynthER’s task setting, our trajectory-level generation outperforms the transition-level generation of SynthER.

References:

[Ref1] Ksenia Konyushkova, Konrad Zolna, Yusuf Aytar, Alexander Novikov, Scott Reed, Serkan Cabi, and Nando de Freitas. Semi-supervised reward learning for offline reinforcement learning. NeurIPS Workshop, 2020.

[Ref2] Cong Lu, Philip J. Ball, Yee Whye Teh, and Jack Parker-Holder. Synthetic Experience Replay. NeurIPS, 2023.

Thanks a lot for your insightful and inspiring comments! We provide the following clarifications in response to your concerns:

1. Technique detail in Section 4.2

   • In Sec. 4.2, the synthesized result is actually $(\dots, s_t, a_t)$, not $(\dots, s_L, a_L)$. Here $L$ represents the expected length of the generated trajectories, which is not the real length. In practice, the generated length can be flexible, due to the multiple generators and trajectory pruner.

2. Reward generation

   • The reward is calculated by accessing the ground truth environment. It is important to note that in the context of online training, RL methods inherently involve accessing the environment. Our work focuses on the paradigm that utilizes the offline data to improve the performance of online reinforcement learning, aiming to develop a general tool that can be applied to any online reinforcement methods. Therefore, our pipeline involves accessing the online environment in the online training phase, which might be different from traditional offline learning methods.

   • We would like to also highlight that our ATraDiff's interaction with the environment is deliberately constrained in our method. As detailed in Section 4.2, the environment access probability is set as $\frac{\rho}{(1-\rho)L}$. Here $L$ represents the expectation length of generation, $T$ represents the real sample time, and $\rho$ is the controlling variable. This ensures that environment interactions with ATraDiff occur only $\frac{\rho} * T$ times, a constant and limited number, thereby controlling the potential increase in time costs.

   • In the updated manuscript, we have also included a variant of our method that trains a reward prediction function from an offline dataset. To achieve this, we have followed the established methodology from prior work [Ref1], and implemented supervised training with the offline dataset. The effectiveness of this reward predictor has been demonstrated in our offline experiments, the results of which are included in Section C.1 of the updated appendix as well as shown in the table below. In our future work, we aim to adapt this reward predictor for online scenarios and further improve the performance.

   Task Name	TD3+BC	TD3+BC+ours_with_learned_reward	IQL	IQL+ours_with_learned_reward
   halfcheetah-random	11.3	11.7	15.2	16.8
   halfcheetah-medium	48.1	49.3	48.3	50.2
   halfcheetah-medium-replay	44.8	45.3	43.5	45.6
   halfcheetah-medium-expert	90.8	91.9	94.6	94.5
   hopper-random	8.6	14.8	7.2	8.3
   hopper-medium	60.4	64.5	62.8	72.2
   hopper-medium-replay	64.4	65.2	84.6	97.0
   hopper-medium-expert	101.1	108.7	106.2	108.3
   walker-random	0.6	3.4	4.1	4.0
   walker-medium	82.7	85.0	84.0	85.3
   walker-medium-replay	85.6	87.5	82.6	84.9
   walker-medium-expert	110.0	111.1	111.7	111.6

3. Random dropping strategy

   • In the online adaptation phase of our pipeline, we employed the total reward importance indicator to assess the significance of each trajectory. An inherent characteristic of this method is that the importance assigned to any given trajectory remains constant throughout the training process. Consequently, trajectories with high importance are perpetually prioritized for generator updates.

   • This consistent prioritization led us to implement a random dropping strategy. By introducing this strategy, we aimed to inject variability into the training process, mitigating potential biases that might arise from the continuous use of high-importance trajectories. We believe that this approach not only diversifies the training data but also enhances the robustness of our generator.

   • To evaluate the effectiveness of the random dropping strategy, we have conducted an ablation study. The result and discussion are detailed in Section C.2 of the updated appendix, showing that the random dropping strategy significantly improves the performance.

3. Lack of comparison with SynthER

   • First, we would like to clarify that the task settings differ between our work and SynthER. Our work primarily focuses on leveraging offline data to enhance online reinforcement learning performance. In contrast, SynthER aims to upsample the offline dataset using their data synthesizer. A crucial distinction lies in their methodology for online experiments, where they do not integrate the offline dataset, which is a key component in our approach. This fundamental difference in task settings initially led us to exclude a direct comparison with SynthER.

   • Additionally, at the time of our submission to ICLR, SynthER was not yet published, appearing later in NeurIPS 2023. Given these circumstances, a direct comparison was not feasible in our initial submission.

   • Following the reviewer’s suggestion, we now include comparisons between our method and SynthER. Consistent with the discussion above, we consider two types of comparisons: (I) Employ SynthER in our task setting, which we care most about. To this end, we have conducted a comparative analysis in an online setting within the Meta-World Environments. This analysis is crucial to demonstrate the adaptability and effectiveness of our approach in dynamic and complex settings. Our results indicate a significant performance advantage of our method over SynthER in these online environments, as summarized in the table below. For more detailed results, including full performance curves, please refer to Section C.5 of the updated appendix.

   Task Name	SPiRL+SynthER	SPiRL+ours	RLPD+SynthER	RLPD+ours
   multitask	0.45	0.61	0.52	0.75
   sequencetask	0.49	0.57	0.62	0.80

   • (II) Employ our method in SynthER’s task setting, which is not the focus of our paper. To this end, we have included an additional comparison in Section C.1 of the updated appendix. This comparison, conducted in the D4RL Locomotion environment under offline settings, demonstrates that our method can achieve comparable or superior results to SynthER.

   Task Name	TD3+BC	TD3+BC+SynthER	TD3+BC+ours	IQL	IQL+SynthER	IQL+ours
   halfcheetah-random	11.3	12.2	12.5	15.2	17.2	17.1
   halfcheetah-medium	48.1	49.9	52.3	48.3	49.6	53.1
   halfcheetah-replay	44.8	45.9	46.5	43.5	46.7	49.2
   halfcheetah-expert	90.8	87.2	93.6	94.6	93.3	95.2
   hopper-random	8.6	14.6	15.2	7.2	7.7	8.1
   hopper-medium	60.4	63.4	65.7	62.8	72.0	72.4
   hopper-replay	64.4	53.4	64.7	84.6	103.2	103.6
   hopper-expert	101.1	105.4	111.2	106.2	90.8	113.6
   walker-random	0.6	2.3	2.1	4.1	4.2	4.3
   walker-medium	82.7	84.8	87.5	84.0	84.7	89.1
   walker-replay	85.6	90.5	86.3	82.6	83.3	85.4
   walker-expert	110.0	110.2	111.2	111.7	111.4	111.7

   • In summary, our aforementioned experimental results validate that, in both our task setting and SynthER’s task setting, our trajectory-level generation outperforms the transition-level generation of SynthER.

4. Missing analysis on the quality of generated images

   • Thanks for your valuable suggestion. In Section D of the updated appendix, we have included representative visualizations of our generated image trajectories. We present a variety of images generated under different conditions, including varying initial states and a range of tasks. The results demonstrate the robustness and effectiveness of our image generation method, consistently performing well across diverse scenarios and producing high-quality and temporally coherent image trajectories.

5. Missing implementation detail and reproducibility

   • Thank you for pointing out the missing implementation details. To address this, we have revised our manuscript and included the previously omitted methodology and implementation details, such as the architecture of the generation models, the conversion between images and states, and the pick-up strategy in Section A of the updated appendix. We hope that the revision improves the clarity and provides a thorough understanding of our method and procedure.

   • Importantly, we will release our code upon acceptance to ensure reproducibility.

   • If there are any remaining implementation details that require further clarification, please do not hesitate to let us know.

Thanks a lot for your insightful and inspiring comments! We provide the following clarifications in response to your concerns:

1. Additional latency of using ATraDiff

   • Our ATraDiff is designed with a trade-off between sample efficiency and computational resource demands. We would like to emphasize that in the context of online reinforcement learning, the most significant cost often lies in interacting with the real environment and data collection. Consequently, our primary focus with ATraDiff has been on optimizing sample efficiency to minimize the time required for environment interaction.
   • Following the reviewer’s suggestion, we quantify the time impact of ATraDiff, by recording the total training time during experiments on the D4RL Locomotion dataset, running on a single NVIDIA RTX 4090TI GPU. The results are shown in the table below, which are also included in Table 4 of the updated Appendix. The comparison between SAC/REDQ and SAC/REDQ with ATraDiff shows that integrating ATraDiff with RL methods leads to an extension of the training duration. However, this is counterbalanced by the improved efficiency in data utilization.

   	walker2d	hopper	halfcheetah
   SAC	6.5h	7h	7h
   SAC_with_ATraDiff	15h	17h	16.5h
   REDQ	6h	6.5h	6.5h
   REDQ_with_ATraDiff	14.5h	16.5h	16h

   • Moving forward, we are committed to enhancing the time efficiency of ATraDiff. Our goal is to reduce both training and inference times, while maintaining its current performance levels. We appreciate the reviewer’s valuable feedback and view it as a guiding direction for our future work.

2. Line names ordering and color in figures

   • Thanks for your valuable feedback. We have adjusted the line name ordering and color in Figures 3, 4, and 5 of the updated manuscript.

Thanks for addressing my concerns and answering the questions. I am satisfied by the answers and the updated paper, which I believe to be now a stronger submission, I have thus increased my rating.

One more suggestion that I have for appendix D of the generated images, would be to show them side by side with real trajectories from the simulator, so that people that are not very familiar with the task can compare them more easily.

3. Unfair comparison with pure online methods

   • We would like to clarify that our comparison is fair, because both our method and the baselines use the same offline datasets. Specifically, as mentioned in Section 5.2 of the original manuscript, our comparison primarily focuses on offline-to-online methods like SPiRL, RLPD, and PARROT, which, like ours, leverage offline datasets in conjunction with online learning. The distinction between our method and these baselines lies in our unique utilization of offline data. Therefore, our improvement indeed comes from our methodology.

   • In addition, the reviewer suggested using learning from demonstration methods as baselines. However, such methods usually require high-quality and task-specific trajectories as a training dataset, which may not achieve good performance with low-quality datasets and would be worse than the offline-to-online baselines we compared with.

   • Regarding the specific ways of utilizing offline data, the Related Work section details how our method significantly differs from these existing offline-to-online methods.

References:

[Ref1] Ksenia Konyushkova, Konrad Zolna, Yusuf Aytar, Alexander Novikov, Scott Reed, Serkan Cabi, and Nando de Freitas. Semi-supervised reward learning for offline reinforcement learning. NeurIPS Workshop, 2020.

[Ref2] Philip J. Ball, Laura M. Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. ICML 2023.

Thanks a lot for your insightful and inspiring comments! We provide the following clarifications in response to your concerns:

1. Having access to the ground truth reward

   • We understand the reviewer’s concern on the potential costs associated with environment interaction and access to ground truth rewards. However, we would like to emphasize that in the context of online training, RL methods inherently involve accessing the environment, which aligns with the operational framework of our method. The core contribution of our work is the development of a versatile plug-in method, designed for integration with various online RL techniques.

   • We would like to also highlight that our ATraDiff's interaction with the environment is deliberately constrained in our method. As detailed in Section 4.2, the environment access probability is set as $\frac{\rho}{(1-\rho)L}$. Here $L$ represents the expectation length of generation, $T$ represents the real sample time, and $\rho$ is the controlling variable. This ensures that environment interactions with ATraDiff occur only $\rho * T$ times, a constant and limited number, thereby controlling the potential increase in time costs.

   • We appreciate the reviewer’s suggestion to train a reward prediction function from an offline dataset. To achieve this, we have followed the established methodology from prior work [Ref1], and implemented supervised training with the offline dataset. The effectiveness of this reward predictor has been demonstrated in our offline experiments, the results of which are included in Section C.1 of the updated appendix as well as shown in the table below. In our future work, we aim to adapt this reward predictor for online scenarios and further improve the performance.

   Task Name	TD3+BC	TD3+BC+ours_with_learned_reward	IQL	IQL+ours_with_learned_reward
   halfcheetah-random	11.3	11.7	15.2	16.8
   halfcheetah-medium	48.1	49.3	48.3	50.2
   halfcheetah-medium-replay	44.8	45.3	43.5	45.6
   halfcheetah-medium-expert	90.8	91.9	94.6	94.5
   hopper-random	8.6	14.8	7.2	8.3
   hopper-medium	60.4	64.5	62.8	72.2
   hopper-medium-replay	64.4	65.2	84.6	97.0
   hopper-medium-expert	101.1	108.7	106.2	108.3
   walker-random	0.6	3.4	4.1	4.0
   walker-medium	82.7	85.0	84.0	85.3
   walker-medium-replay	85.6	87.5	82.6	84.9
   walker-medium-expert	110.0	111.1	111.7	111.6

2. Quality of offline datasets

   • We acknowledge that the quality of the offline dataset impacts overall performance in reinforcement learning. However, we would like to emphasize that our work introduces a versatile tool for online reinforcement learning that operates effectively regardless of data quality.

   • In our experiments, we observe that 1) high-quality data significantly increase the performance gap between our method and the baseline; 2) importantly, our method is designed to be robust and it does not require exceptionally high-quality data to function effectively, as shown in the experiment with medium-quality data in Section C.3; 3) our method is comparable to or better than the baseline with low-quality data; and 4) even datasets with non-optimal trajectories are sufficient for our method to achieve reliable performance, as shown in the multi-task experiment in Fig. 5.

   • Regarding the experiments on D4RL, in the original manuscript, we followed prior studies [Ref2] and presented the average performance across different datasets from the D4RL benchmark. Following the reviewer’s suggestion and to provide a more comprehensive understanding, we have now included separate results for datasets of varying quality in Section C.3 of the updated appendix. The results in C.3 show how data quality impacts the effectiveness of our method and baselines across a spectrum of scenarios – while our method is comparable to or better than the baseline with low-quality data (like random data), the performance gap between our method and the baseline becomes much more pronounced with medium/high-quality data (like medium-expert data).

Thanks a lot for your insightful and inspiring comments! We provide the following clarifications in response to your concerns:

1. Lack of wall-clock training time in experiments

   • We acknowledge the concern regarding the increased wall-clock training time associated with the use of diffusion model data generation in ATraDiff. ATraDiff is designed with a trade-off between sample efficiency and computational resource demands. We would like to emphasize that in the context of online reinforcement learning, the most significant cost often lies in interacting with the real environment and data collection. Consequently, our primary focus with ATraDiff has been on optimizing sample efficiency to minimize the time required for environment interaction.

   • Following the reviewer’s suggestion, we quantify the time impact of ATraDiff, by recording the total training time during experiments on the D4RL Locomotion dataset, running on a single NVIDIA RTX 4090TI GPU. The results are shown in the table below, which are also included in Table 4 of the updated Appendix. The comparison between SAC/REDQ and SAC/REDQ with ATraDiff shows that integrating ATraDiff with RL methods leads to an extension of the training duration. However, this is counterbalanced by the improved efficiency in data utilization.

   	walker2d	hopper	halfcheetah
   SAC	6.5h	7h	7h
   SAC_with_ATraDiff	15h	17h	16.5h
   REDQ	6h	6.5h	6.5h
   REDQ_with_ATraDiff	14.5h	16.5h	16h

   • Moving forward, we are committed to enhancing the time efficiency of ATraDiff. Our goal is to reduce both training and inference times, while maintaining its current performance levels. We appreciate the reviewer’s valuable feedback and view it as a guiding direction for our future work.