# Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation

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

## Key facts

- Cited by 105 papers in the map (data through 14 Oct 2026).
- Builds on 69 of its 86 references that are in the map.
- First posted on arXiv 1 Jul 2021.

## Builds on (most cited first, 30 of 69)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Unsupervised Representation Learning by Predicting Image Rotations](https://mlsanity.com/paper/1803.07728/) (2018)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Context Encoders: Feature Learning by Inpainting](https://mlsanity.com/paper/1604.07379/) (2016)
- [Unsupervised Learning of Visual Representations by Solving Jigsaw Puzzles](https://mlsanity.com/paper/1603.09246/) (2016)
- [Contrastive Multiview Coding](https://mlsanity.com/paper/1906.05849/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Colorful Image Colorization](https://mlsanity.com/paper/1603.08511/) (2016)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Dueling Network Architectures for Deep Reinforcement Learning](https://mlsanity.com/paper/1511.06581/) (2015)

## Built on by (most cited first, 30 of 105)

- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [The Dormant Neuron Phenomenon in Deep Reinforcement Learning](https://mlsanity.com/paper/2302.12902/) (2023)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations](https://mlsanity.com/paper/2403.03954/) (2024)
- [Decomposing the Generalization Gap in Imitation Learning for Visual Robotic Manipulation](https://mlsanity.com/paper/2307.03659/) (2023)
- [MoDem: Accelerating Visual Model-Based Reinforcement Learning with Demonstrations](https://mlsanity.com/paper/2212.05698/) (2022)
- [Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning](https://mlsanity.com/paper/2212.08860/) (2022)
- [TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning](https://mlsanity.com/paper/2306.13229/) (2023)
- [On Pre-Training for Visuo-Motor Control: Revisiting a Learning-from-Scratch Baseline](https://mlsanity.com/paper/2212.05749/) (2022)
- [Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models](https://mlsanity.com/paper/2205.13817/) (2022)
- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)
- [Look Closer: Bridging Egocentric and Third-Person Views with Transformers for Robotic Manipulation](https://mlsanity.com/paper/2201.07779/) (2022)
- [Graph Inverse Reinforcement Learning from Diverse Videos](https://mlsanity.com/paper/2207.14299/) (2022)
- [Visual Reinforcement Learning with Self-Supervised 3D Representations](https://mlsanity.com/paper/2210.07241/) (2022)
- [Latent Action Learning Requires Supervision in the Presence of Distractors](https://mlsanity.com/paper/2502.00379/) (2025)
- [Finetuning Offline World Models in the Real World](https://mlsanity.com/paper/2310.16029/) (2023)
- [Learning to Manipulate Anywhere: A Visual Generalizable Framework For Reinforcement Learning](https://mlsanity.com/paper/2407.15815/) (2024)
- [Bootstrapped Transformer for Offline Reinforcement Learning](https://mlsanity.com/paper/2206.08569/) (2022)
- [Deep Transformer Q-Networks for Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2206.01078/) (2022)
- [Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning](https://mlsanity.com/paper/2207.09081/) (2022)
- [Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning](https://mlsanity.com/paper/2209.09203/) (2022)
- [On the Importance of Exploration for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2306.05483/) (2023)
- [RL-ViGen: A Reinforcement Learning Benchmark for Visual Generalization](https://mlsanity.com/paper/2307.10224/) (2023)
- [Conditional Mutual Information for Disentangled Representations in Reinforcement Learning](https://mlsanity.com/paper/2305.14133/) (2023)
- [GEVRM: Goal-Expressive Video Generation Model For Robust Visual Manipulation](https://mlsanity.com/paper/2502.09268/) (2025)
- [The Generalization Gap in Offline Reinforcement Learning](https://mlsanity.com/paper/2312.05742/) (2023)
- [A Comprehensive Survey of Data Augmentation in Visual Reinforcement Learning](https://mlsanity.com/paper/2210.04561/) (2022)
- [Seeing is not Believing: Robust Reinforcement Learning against Spurious Correlation](https://mlsanity.com/paper/2307.07907/) (2023)
- [Does Self-supervised Learning Really Improve Reinforcement Learning from Pixels?](https://mlsanity.com/paper/2206.05266/) (2022)
- [Vision-Guided Quadrupedal Locomotion in the Wild with Multi-Modal Delay Randomization](https://mlsanity.com/paper/2109.14549/) (2021)

## Abstract (arXiv)

While agents trained by Reinforcement Learning (RL) can solve increasingly challenging tasks directly from visual observations, generalizing learned skills to novel environments remains very challenging. Extensive use of data augmentation is a promising technique for improving generalization in RL, but it is often found to decrease sample efficiency and can even lead to divergence. In this paper, we investigate causes of instability when using data augmentation in common off-policy RL algorithms. We identify two problems, both rooted in high-variance Q-targets. Based on our findings, we propose a simple yet effective technique for stabilizing this class of algorithms under augmentation. We perform extensive empirical evaluation of image-based RL using both ConvNets and Vision Transformers (ViT) on a family of benchmarks based on DeepMind Control Suite, as well as in robotic manipulation tasks. Our method greatly improves stability and sample efficiency of ConvNets under augmentation, and achieves generalization results competitive with state-of-the-art methods for image-based RL in environments with unseen visuals. We further show that our method scales to RL with ViT-based architectures, and that data augmentation may be especially important in this setting.