Paper: ATraDiff: Accelerating Online Reinforcement Learning with Imaginary Trajectories (ICLR-2024)

Reviewer comment (Reproducibility):
**(Incomplete description of algorithm)** A huge amount of detail is missing when the authors claim they convert states to images and then back to states in Section 4.1. This is highly non-trivial and should be justified further. For example, can the authors present pictures of rendered states, what is the error from converting to image and back again? There is also no description of how Stable Diffusion was fine-tuned or how long online image generation takes. Furthermore, it is not explained how the authors would generate low-dimensional actions or obtain rewards from an image model.

Author response:
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.

Reviewer comment (Experiments):
**(Benefit of trajectories vs. transitions)** The paper claims “full trajectories offer a more comprehensive source of information, enabling RL agents to better learn from past experiences” without evidence in the introduction. No comparison is made to the transition-based baseline in [1] which also evaluates on the same MuJoCo locomotion environments in Figure 2. Furthermore, [3] also proposes a trajectory based method for generating synthetic data with diffusion models and is not discussed.

Author response:
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.

Reviewer comment (Experiments):
**(Why project states to images)** The paper claims that generating a collection of images outperforms single states in Figure 5. It is unclear why the simpler option of generating a sequence of states as in [2, 3] is not chosen. Given the image-based approach, it may also be asked why the authors do not focus on upsampling image datasets from the start, e.g. from [4].

Author response:
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.

Reviewer comment (Experiments):
**(Incomplete experimental description)** Unclear what experimental setup is used in Figure 2. Is each figure the average of 4 different D4RL datasets? No explanation for what environments or datasets are used in Figure 5.

Author response:
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.

Reviewer comment (Reproducibility):
**(Unclear description of sampling scheme)** There is an extremely vague description of what the reward or TD-error based pickup strategy means in Appendix A.1.1.

Author response:
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.

Reviewer comment (Writing):
In Sec. 4.2, is the synthetic trajectory actually $(\ldots,s_L,a_L)$ instead of $(\ldots, s_t,a_t)$?

Author response:
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.

Reviewer comment (Reproducibility):
How can the reward $R(s_i,a_i)$ of a synthetic state-action pair be calculated from the environment if the action is not executed in the environment, or is the reward function accessible to the agent?

Author response:
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 |

Reviewer comment (Experiments):
In the last part of Sec. 4.3, it states that "Hence, we will randomly drop some samples from the maintained subset regardless of their importance". Is there any ablation study to support the effectiveness of this random dropping strategy?

Author response:
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.

Reviewer comment (Experiments):
A recent work [1] about offline and online data augmentation using diffusion models has been published. Could you please include some comparison between your work and [1]? Experimental comparison will be better.

Author response:
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**.

Reviewer comment (Reproducibility):
The additional latency of using ATraDiff, and the continual training of the diffusion model on new experiences is not mentioned, but it's a rather important detail.

Author response:
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.

Reviewer comment (Evaluation):
Would have liked to see an analysis on the quality of the images generated by ATraDiff, particularly it would be interesting to see how temporally correlated are the trajectories.

Author response:
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.

Reviewer comment (Reproducibility):
Unsure how reproducible the paper is as many implementation details are missing.

Author response:
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.

Reviewer comment (Evaluation):
What is the added latency of using ATraDiff compared to only using one of the baseline RL algorithms? Regardless if it's significantly more computationally expensive, it would still be helpful to know quantitatively.

Author response:
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.

Reviewer comment (Experiments):
Not sure I understand why there is no comparison with its closest method SynthER?

Author response:
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**.

Reviewer comment (Experiments):
This approach assume the access to the ground truth reward function, which is a very impractical setting. Most of tasks evaluation of trajectories is expensive.

Author response:
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 |

Reviewer comment (Experiments):
This work lack of discussion on the assumptions of offline datasets that is already available to the approach. Is it a medium quality dataset? Is it a low-quality dataset? Does it contains trajectories from optimal policy? My understanding is that this will make substantial difference to the algorithm performance. The D4RL dataset contains different level of offline data. It would be good to clarify the setting of offline datasets.

Author response:
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).

Reviewer comment (Experiments):
In the online setting, the fair comparison would be using the offline data with the online algorithms. e.g., learning from demonstration approaches would be a better baseline. O.w. the proposed approach has an advantage of using more offline data. It's unclear to me whether the improvement comes from the access to the additional offline dataset or the proposed approach.

Author response:
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.

Reviewer comment (Experiments):
Did the author consider replacing the query. of reward function from env by training a reward function from offline trajectories? Please also refer to the weakness section.

Author response:
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 |

Reviewer comment (Evaluation):
There is no data about trade-offs associated with using ATraDiff.

Author response:
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.

Reviewer comment (Evaluation):
Diffusion model data generation can be pretty slow, raising the question of how it affects wall-clock training time.

Author response:
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.

Reviewer comment (Evaluation):
It would be good to see more discussions and data on the topic of wall-clock training time impact.

Author response:
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.