Synthetic Data is Sufficient for Zero-Shot Visual Generalization from Offline Data

Augmented dataset has the same size as the Baseline dataset

We appreciate the reviewer’s positive feedback on the novelty of our approach, the clarity of our writing, and the comprehensive experimental analysis. Thank you for raising this important question about dataset size and the distribution of upsampled augmented data compared to the original augmented data. Our upsampled dataset extends the augmented data distribution to better align with the testing data, leveraging the broader diversity introduced by upsampling. While the diffusion model generates data based on the augmented data distribution, it also expands diversity beyond the original distribution, as shown in Figure 5 of the SynthER paper. Our JS divergence analysis supports this observation, demonstrating that the upsampled data moves closer to the testing data distribution while increasing overall diversity. At its core, our method integrates visual augmentation with upsampling into a unified approach to improve generalization when additional data is unavailable. Investigating the impact of dataset size by comparing pure augmentation with our combined method using equal-sized datasets would indeed be a valuable direction for future work. However, our current focus is on presenting a simple and effective solution for scenarios with limited data, where enhancing diversity through augmentation and upsampling is crucial.

the performance of this method might highly depend on the data augmentation used in the first phase

Thank you for pointing out this complementary aspect to the earlier discussion. To improve generalization, we began with all augmentations proposed by RAD [1] and systematically refined them. We eliminated augmentations that caused training instability and iteratively narrowed down to those most impactful on generalization performance. Through this process, we observed that both environments favored similar augmentations, and we fine-tuned their parameters to maximize stability and effectiveness (as detailed in Supplementary Section B.0.3). Our focus was on achieving stable training outcomes while balancing computational constraints, as this work was conducted on a single GPU academic setting. While we performed ablations during the selection process, we opted not to include an exhaustive set of results in the paper to maintain focus on the simplicity and effectiveness of combining visual augmentation with diffusion-based upsampling for generalization. We believe this strikes a balance between demonstrating the method’s impact and avoiding overemphasis on augmentation-specific studies in offline RL. We agree that a deeper exploration of the interplay between augmentations and diffusion-generated data is a valuable direction for future work.

[1] Laskin et al. Reinforcement Learning with Augmented Data. NeurIPS 2020

typos

Thank you for pointing out the minor typos. We corrected "invrease" (line 137) and "tecnic" (line 361) in the revised version of paper.

Are they augmented by the same image transformation or a randomly sampled image transformation?

Thank you for asking for the clarification of transformations applied to states. We realized that we didn’t add details about that and we updated section B.0.3 in supplementary material. To clarify, the augmentation function is applied to states s and s′, which uses independently sampled transformations from the same set of augmentations, each selected with equal probability. Within each state, the same transformation is consistently applied across all images in the stack to preserve temporal and spatial relationships, which is also used in RAD work. This ensures diversity across different states while maintaining structural integrity within each state.

Is it possible that, by incorporating a small amount of data from the Fixed Distraction Dataset along with the diffusion upsampling process?

We thank the reviewer for highlighting this intriguing aspect of our findings in section 5.1.1. Incorporating a small amount of data from the Fixed Distraction Dataset (FDD) alongside the diffusion upsampling process—and evaluating its performance without the initial data augmentation phase—is indeed an intriguing hypothesis. Our primary focus was to demonstrate the combination of augmentation and upsampling as a cohesive method, which is why we kept our approach consistent with the augmented and upsampled (ours) pipeline to ensure clarity and alignment with our key message. That said, we acknowledge the potential value of exploring this direction, which we have identified as an open problem in Section 5.1.1, encouraging further investigation by the research community. If the reviewer believes this analysis would provide significant additional insights and help improve the paper’s impact, we are happy to include an ablation study in the revised version to examine this hypothesis further.

We appreciate your suggestion and review of our responses. As mentioned in our initial response, our focus is on synthetic data generation rather than a thorough investigation of augmentation techniques. However, we emphasize that our method combines visual augmentation with upsampling to efficiently increase diversity and generalization in data-limited settings.

As discussed earlier, we methodically refined the augmentation decisions for stability and efficacy. Our JS divergence study demonstrates how the upsampled data better aligns with the testing distribution while improving diversity. We will consider your insightful recommendation to compare pure image augmentation with diffusion-augmentation under the same-size setting and to analyze the choice of image augmentation in future extension work

The discussion and analysis of chosen data augmentation techniques in Section 3.2 lacks sufficient depth

We thank the reviewer for recognizing the innovation in our approach, its computational efficiency, and the value of our JS divergence analysis. However, we respectfully disagree with the reviewer’s concern regarding the lack of sufficient depth in the data augmentation. While we greatly value the contributions of methods like DrAC and SVEA, these methods involve algorithmic changes specific to online RL, which differ from our objective of providing a simple, algorithm-agnostic solution for offline RL with visual observations. Our focus is on achieving non-algorithmic changes to offline RL using simple, effective visual augmentations, as demonstrated by RAD [1], with four specific techniques detailed in Supplementary Section B.0.3. These augmentations, combined with diffusion-based upsampling, are central to our goal of enhancing data diversity and generalization without altering the underlying RL algorithm. To address this further, we expanded the “Related Work” section to clarify these distinctions in more detail and included the missing papers from the three suggested by the reviewer.

[1] Laskin et al. Reinforcement Learning with Augmented Data. NeurIPS 2020

The proposed Generalization Performance metric $G_{perf} = \frac{T_{test} - B_{test}}{B_{train} - B_{test}}$ needs better justification.

Regarding the generalization metric, its design is grounded in established principles commonly used in reinforcement learning (RL) generalization studies. See the original Procgen paper [1] Section 2.2. They use the following $R_{norm} = \frac{R - R_{min}}{R_{max} - R_{min}}$. The idea being to normalize w.r.t. what is possible given the setup. Here $Rmin$ is the lowest possible score, i.e., the test performance of the baseline (with poor generalization), and $Rmax$ is the highest possible score, the train performance on the baseline training data, which is the least diverse and easiest to overfit. We added this discussion in the paper (section 4.3.1) as we agree it appears plucked from thin air in the current draft, so thank you for flagging.

[1] Cobbe et al. Leveraging Procedural Generation to Benchmark Reinforcement Learning. ICML 2020

The experimental results reveal a critical misalignment with the paper's claimed contribution to "Zero-Shot Visual Generalization.

Thank you for highlighting this critical aspect of our work on zero-shot generalization. We respectfully disagree with the reviewer’s assertion that the paper does not demonstrate improved zero-shot generalization, as we show this in Procgen (see aggregate performance added to Table 3). Additionally, we present the FDD approach (Table 2), where we observe improvement in the generalization gap for the DMC environments. That said, we understand that the improved performance in the original environment in Table 1 (not necessarily a bad thing!) could lead to confusion. We are happy to rephrase the title if you have a recommendation. One proposal could be “Synthetic Data Enables Training Robust Agents from Offline Data,” as our agents perform well across a wide range of settings. We also updated Tables 2 and 3 to include $Test/Train$ and $Train-Test$ results for both environments, aligning with the metrics suggested by the reviewer. Please let us know if this makes more sense now.

insufficient exploration of diffusion model design choices

Thank you for highlighting the need for deeper exploration of diffusion model design choices. Our diffusion model builds directly on SynthER’s design, which allows us to make consistent comparisons with the V-D4RL benchmark and ensures the reproducibility of our results. SynthER’s supplementary material (Section B) already provides extensive ablations on model architecture and denoiser parameters, and our findings closely align with these results. As such, duplicating those ablations in our work would have added redundancy without offering additional insights. Instead, we focused on demonstrating the effectiveness of our combined augmentation and diffusion-based upsampling method in the context of offline RL. To provide clarity, we have added references to SynthER’s ablations in the updated Supplementary Section C.3,1. Additionally, we conducted and included a table on latent space size ablations in the updated supplementary material, addressing a gap not explored in SynthER’s ablation work (Section C and Section F). While this analysis was initially omitted for brevity, we now provide it to offer further insights into the relationship between latent space dimensionality and performance, complementing SynthER's findings.

Dear Reviewer cYiQ,

We have carefully addressed all the points you raised in my rebuttal, and we believe the additions and explanations provided will help in evaluating the paper further. There is not much time left for us to address your additional comments; hence, we gladly ask you to evaluate our rebuttal at your earliest time, as your insights are important in ensuring a complete assessment of the work. We value the time you spent in this process and would be pleased to offer any more explanations should they be necessary.

Thank you for your time and consideration.

Authors

We appreciate your careful review and for noting our improvements as well as for raising your score. Your comments have been quite helpful in directing important clarifications and additions, including those seen in Tables 2 and 3.

For impact-to-complexity ratio, to the best of our knowledge, this is the first work to effectively implement this practical method in two different kinds of offline RL environments—one with continuous action spaces and one with discrete action spaces. By demonstrating its effectiveness across these settings, we believe our method provides a strong starting point for the research community, aiming to reduce the complexity of selecting augmentation strategies and tuning hyperparameters and making it easier for others to adopt and build upon this approach.

Thank you for your review, we are pleased to see you appreciate that our method is a simple approach to achieve improved generalization in two distinct domains. It appears your primary concerns relate to scaling beyond the domains shown, which we believe will be challenging for us to answer concretely in this rebuttal. We note that the domains used have been popular for existing works on data augmentation [1, 2] and are still being used by industry labs in recent publications [3].

See specific responses below:

Cost of diffusion modeling: in this paper we are focused on the offline RL setting, where there is bottleneck on the amount of available data but not necessarily the amount of time or compute to maximize performance with it. Further, it may be possible when scaling this method to use an open source foundation model which already has world knowledge. Finally, for what it is worth, this work was done with extremely constrained computation resources of a single GPU in an academic lab - yet we were able to get state of the art performance for visual generalization. We think that is a good sign for scalability!

The GAN and VAE comparison is a great point - however - this comparison was already made in the SynthER paper and we do not have any reason to believe it would not hold in a more challenging setting. Please check out Table 1: https://arxiv.org/pdf/2303.06614, the performance is drastically better for Diffusion which makes sense, it is now the dominant paradigm in generative modeling. What we are showing here is that it may also make it possible to achieve additional visual generalization benefits, if set up correctly with augmenting the data, which was not obvious to us initially and thus we believe is a valuable contribution to the community.

We agree that it is definitely important to make sure the data does not deviate too far from the ground truth distribution. This was shown in the SynthER paper in Figure 5 and note there has been additional work in this space, such as Policy Guided Diffusion [4], which likely improves our synthetic data generation pipeline. The main contribution in this paper is showing these methods can aid visual generalization, which had not been shown before.

[1] Yarats et al. Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels. ICLR 2021

[2] Laskin et al. Reinforcement Learning with Augmented Data. NeurIPS 2020

[3] Ortiz et al. DMC-VB: A Benchmark for Representation Learning for Control with Visual Distractors. NeurIPS 2024 Datasets and Benchmarks Track

[4] Jackson et al. Policy-Guided Diffusion. NeurIPS 2023 Workshop on Robot Learning

Dear Reviewer uuZk,

We have carefully addressed all the points you raised in my rebuttal, and we believe the additions and explanations provided will help in evaluating the paper further. There is not much time left for us to address your additional comments; hence, we gladly ask you to evaluate our rebuttal at your earliest time, as your insights are important in ensuring a complete assessment of the work. We value the time you spent in this process and would be pleased to offer any more explanations should they be necessary.

Thank you for your time and consideration.

Authors

the method requires tuning of the data augmentation parameters, which limits it's applicability.

We sincerely thank you for your valuable feedback and are glad you found our method "simple and easy to understand." We acknowledge that tuning data augmentation techniques and diffusion model parameters can be challenging, as noted in our limitations discussion in Section 7. To address this, we used all augmentations from RAD [1] to systematically narrow down the augmentation options to reduce computational cost for effectiveness and training stability. For diffusion model parameters, we started with hyperparameters proposed in the SynthER paper that led to a proper balance between generalization and overfitting. To mitigate the time-consuming nature of tuning, we employed JS divergence analysis and distribution visualizations to assess alignment between the upsampled and baseline datasets. Specifically, this was achieved by systematically varying the size of the denoiser network and the number of training steps, allowing us to identify configurations that produced the best alignment. This approach allowed us to reduce the need for repeated RL training while efficiently improving data diversity and generalization. We have updated Supplementary Section B.0.3 to include a detailed explanation.

[1] Laskin et al. Reinforcement learning with augmented data. 2020

The experiments only include DrQ and CQL. Since this paper deals with extending the data and can be applied to many different methods, it would make the method more compelling if there were more methods like in SynthER, e.g. IQL, TD3+BC, EDAC.

We selected DrQ+BC and CQL to align with the benchmark datasets from V-D4RL and Offline Procgen (Lu et al., 2023a; Mediratta et al., 2024), respectively, to ensure fair comparisons with existing results. DrQ+BC was chosen because its authors highlighted generalization challenges for model-free algorithms on distracting datasets, providing an opportunity to demonstrate how our method effectively addresses these issues. CQL was selected due to its underperformance in offline generalization tasks compared to other algorithms, making it an ideal case to showcase the potential of our method. Note that DrQ+BC is largely the same as TD3+BC but for visual observations. While additional algorithms like IQL and EDAC could be explored in future work, our focus was on demonstrating that our method is algorithm-agnostic and applicable across diverse environments, including both continuous and discrete action spaces, rather than comparing the relative performance of various algorithms.

Do I understand correctly that your 'Upsampled' method is akin to SynthER? If that's so, I would put that in parenthesis. If not, could you provide that comparison?

Thank you for the suggestion. In the Method Section (Subsection 3.1), we have updated the "Upsampling with Diffusion Models" item to explicitly reference SynthER as the method we employed for upsampling.

Response to "Writing"

Thank you for your valuable feedback on the writing and figure presentation. Regarding Figures 4, 5, and 6, we aimed to balance readability and organization by combining multiple charts into single figures to effectively summarize the data. While we recognize the figures are somewhat large, this approach minimizes disruption to the paper's structure. For the heatmaps, we chose this format for concise and quick visual interpretation. Although adding exact values on the squares could provide additional information, it reduced visual clarity due to interference with the color scheme. We opted for heatmaps instead of bar plots to maintain readability but appreciate the suggestion and will explore alternative formats, such as annotated heatmaps, in future work. Finally, we corrected the spelling of "tecnic" to "technique" in the revised version and are glad you liked the original phrasing.

Dear Reviewer zkst,

We have carefully addressed all the points you raised in my rebuttal, and we believe the additions and explanations provided will help in evaluating the paper further. There is not much time left for us to address your additional comments; hence, we gladly ask you to evaluate our rebuttal at your earliest time, as your insights are important in ensuring a complete assessment of the work. We value the time you spent in this process and would be pleased to offer any more explanations should they be necessary.

Thank you for your time and consideration.

Authors