[Paper Review] Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation

Nicklas Hansen, Hao Su+1 more|arXiv (Cornell University)|Jul 1, 2021

Reinforcement Learning in RoboticsComputer Science63 references37 citations

TL;DR

SVEA stabilizes Q-value estimation under data augmentation for off-policy RL, improving both stability and generalization for ConvNets and ViT architectures on DMControl and robotic tasks.

ABSTRACT

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.

Motivation & Objective

• Motivate learning robust policies from visual observations that generalize to unseen environments.
• Identify causes of instability when applying data augmentation to off-policy RL.
• Propose a simple, effective stabilization framework that minimizes augmentation-induced variance and over-regularization.
• Demonstrate scalability to ConvNets and Vision Transformers across diverse tasks and benchmarks.

Proposed method

• Analyze sources of instability when applying data augmentation to off-policy Q-learning.
• Introduce SVEA: a data-augmentation framework with two data streams (augmented and unaugmented) and a mixed Q-learning objective.
• Apply augmentation only to current-state Q-value estimation, not to bootstrapped target Q-values.
• Use a data-mixing objective that jointly optimizes Q-values on augmented and unaugmented data.
• Maintain the actor (if present) to learn from unaugmented data, with encoder sharing to promote generalization.
• Provide an implementation that requires no extra forward passes or learnable parameters.

Experimental results

Research questions

• RQ1 What are the primary causes of instability when using data augmentation in off-policy Q-learning with neural networks?
• RQ2 Can a simple augmentation-aware training objective stabilize Q-value estimation without increasing computational cost?
• RQ3 How well does the proposed stabilization method scale to ConvNets and Vision Transformers on image-based RL benchmarks?
• RQ4 Does the method improve generalization to unseen visuals in DMControl-related benchmarks and robotic manipulation tasks?

Key findings

• SVEA substantially improves stability and sample efficiency of Q-learning with strong augmentations on DMControl tasks.
• SVEA achieves competitive or superior asymptotic performance and generalization compared to state-of-the-art image-based RL methods.
• The two key components— restricting augmentation to current-state Q-value estimation and a dual-stream, mixed objective—reduce variance and over-regularization.
• SVEA scales to ViT-based architectures, where augmentation is particularly beneficial for avoiding overfitting.
• DrQ-style baselines with more augmentation (via higher K,M) incur higher computational cost, while SVEA matches or exceeds performance with lower cost.