# Don't Touch What Matters: Task-Aware Lipschitz Data Augmentation for Visual Reinforcement Learning

arXiv: https://arxiv.org/abs/2202.09982 · HTML: https://mlsanity.com/paper/2202.09982/

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 23 references that are in the map.
- First posted on arXiv 21 Feb 2022.

## Builds on (most cited first, 18 of 18)

- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [CARLA: An Open Urban Driving Simulator](https://mlsanity.com/paper/1711.03938/) (2017)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/1806.10293/) (2018)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Target-driven Visual Navigation in Indoor Scenes using Deep Reinforcement Learning](https://mlsanity.com/paper/1609.05143/) (2016)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [A Study on Overfitting in Deep Reinforcement Learning](https://mlsanity.com/paper/1804.06893/) (2018)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Assessing Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/1810.12282/) (2018)
- [Self-Supervised Policy Adaptation during Deployment](https://mlsanity.com/paper/2007.04309/) (2020)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [Training robust neural networks using Lipschitz bounds](https://mlsanity.com/paper/2005.02929/) (2020)
- [SECANT: Self-Expert Cloning for Zero-Shot Generalization of Visual Policies](https://mlsanity.com/paper/2106.09678/) (2021)

## Built on by (most cited first, 21 of 21)

- [Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning](https://mlsanity.com/paper/2212.08860/) (2022)
- [Learning to Manipulate Anywhere: A Visual Generalizable Framework For Reinforcement Learning](https://mlsanity.com/paper/2407.15815/) (2024)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [RL-ViGen: A Reinforcement Learning Benchmark for Visual Generalization](https://mlsanity.com/paper/2307.10224/) (2023)
- [A Comprehensive Survey of Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2210.04561/) (2022)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)
- [Improving Generalization in Visual Reinforcement Learning via Conflict-aware Gradient Agreement Augmentation](https://mlsanity.com/paper/2308.01194/) (2023)
- [A Recipe for Unbounded Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2405.17416/) (2024)
- [DAG-Plan: Generating Directed Acyclic Dependency Graphs for Dual-Arm Cooperative Planning](https://mlsanity.com/paper/2406.09953/) (2024)
- [Normalization Enhances Generalization in Visual Reinforcement Learning](https://mlsanity.com/paper/2306.00656/) (2023)
- [Generalizable Visual Reinforcement Learning with Segment Anything Model](https://mlsanity.com/paper/2312.17116/) (2023)
- [Focus On What Matters: Separated Models For Visual-Based RL Generalization](https://mlsanity.com/paper/2410.10834/) (2024)
- [CtrlFormer: Learning Transferable State Representation for Visual Control via Transformer](https://mlsanity.com/paper/2206.08883/) (2022)
- [Learning Generalizable Agents via Saliency-Guided Features Decorrelation](https://mlsanity.com/paper/2310.05086/) (2023)
- [Train a Real-world Local Path Planner in One Hour via Partially Decoupled Reinforcement Learning and Vectorized Diversity](https://mlsanity.com/paper/2305.04180/) (2023)
- [Sequential Action-Induced Invariant Representation for Reinforcement Learning](https://mlsanity.com/paper/2309.12628/) (2023)
- [Merging and Disentangling Views in Visual Reinforcement Learning for Robotic Manipulation](https://mlsanity.com/paper/2505.04619/) (2025)
- [Efficient Training of Generalizable Visuomotor Policies via Control-Aware Augmentation](https://mlsanity.com/paper/2401.09258/) (2024)
- [SeMOPO: Learning High-quality Model and Policy from Low-quality Offline Visual Datasets](https://mlsanity.com/paper/2406.09486/) (2024)
- [Salience-Invariant Consistent Policy Learning for Generalization in Visual Reinforcement Learning](https://mlsanity.com/paper/2502.08336/) (2025)
- [Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization](https://mlsanity.com/paper/2607.00796/) (2026)

## Abstract (arXiv)

One of the key challenges in visual Reinforcement Learning (RL) is to learn policies that can generalize to unseen environments. Recently, data augmentation techniques aiming at enhancing data diversity have demonstrated proven performance in improving the generalization ability of learned policies. However, due to the sensitivity of RL training, naively applying data augmentation, which transforms each pixel in a task-agnostic manner, may suffer from instability and damage the sample efficiency, thus further exacerbating the generalization performance. At the heart of this phenomenon is the diverged action distribution and high-variance value estimation in the face of augmented images. To alleviate this issue, we propose Task-aware Lipschitz Data Augmentation (TLDA) for visual RL, which explicitly identifies the task-correlated pixels with large Lipschitz constants, and only augments the task-irrelevant pixels. To verify the effectiveness of TLDA, we conduct extensive experiments on DeepMind Control suite, CARLA and DeepMind Manipulation tasks, showing that TLDA improves both sample efficiency in training time and generalization in test time. It outperforms previous state-of-the-art methods across the 3 different visual control benchmarks.