PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Nicklas Hansen, Hao Su, Xiaolong Wang

arXiv:2107.00644v2cs.LGcs.CVcs.RO

TL;DR

Visual RL often overfits unseen environments, while strong augmentation can reduce sample efficiency or cause divergence. The paper analyzes augmentation-induced Q-target instability and introduces SVEA, which stabilizes off-policy RL through augmented and unaugmented data streams. Across ConvNet and ViT experiments, SVEA improves stability, efficiency, and generalization, though augmentation remains sensitive to the choice of transformation.

  • Problem

    Visual RL struggles to generalize beyond training environments, and extensive augmentation can reduce sample efficiency or cause divergence.

  • Method

    SVEA stabilizes off-policy RL by avoiding augmented bootstrapping targets, jointly estimating Q-values from augmented and unaugmented observations, and using unaugmented states for actor optimization.

  • Results

    SVEA improves ConvNet stability and achieves competitive or better efficiency and generalization across control and manipulation tasks, while scaling to ViT-based RL.

  • Takeaways & Limitations

    SVEA provides a computationally efficient augmentation framework, while augmentation appears particularly important for ViT architectures that overfit visual training environments.

  • Takeaways & Limitations

    Performance and stability remain dependent on the chosen augmentation, which can still cause divergence in Q-learning optimization.

Abstract

from arXiv · show

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.

1 Introduction

Visual RL can solve difficult tasks, but policies often overfit training environments and struggle to generalize to novel visuals. This paper identifies augmentation-induced instability and proposes SVEA to stabilize off-policy RL under strong augmentation.

  • Visual RL policies can overfit training environments, making generalization to novel environments especially difficult with image observations.
  • Data augmentation can improve sample efficiency through cropping and translation, yet stronger augmentation may reduce efficiency or cause divergence.
  • The paper attributes instability to high-variance Q-targets from indiscriminate augmentation and over-regularization from estimating Q-values only on augmented data.
  • SVEA applies augmentation only to current-state Q-value estimation, jointly trains on augmented and unaugmented observations, and optimizes actor policies on unaugmented data.
  • Across control and robotic manipulation benchmarks, SVEA improves ConvNet stability and achieves competitive or better sample efficiency, asymptotic performance, and generalization.

2 Related Work

Prior work uses representation learning, domain randomization, and augmentation to improve visual RL efficiency and generalization. Deep Q-learning research separately addresses instability from function approximation and off-policy temporal-difference learning.

  • Representation Learning: Self-supervised objectives and augmentation have been used to learn visual invariances that transfer to downstream tasks.
  • Visual Learning for RL: Visual RL methods pursue better sample efficiency and generalization through auxiliary tasks such as bisimulation, contrastive learning, and prediction objectives.
  • Data Augmentation and Randomization for RL: Domain randomization and data augmentation improve robustness by exposing models to varied visual transformations during training.
  • Improving Deep Q-Learning: Deep Q-learning methods address instability caused by function approximation and off-policy data with techniques including Double Q-learning, Dueling networks, and target-network updates.

3 Preliminaries

The paper formulates image-based RL as learning a policy over partially observed, frame-stacked states to maximize discounted return and generalize across MDPs. Deep Q-learning estimates action values from replay data, while actor-critic methods additionally learn policies for large or continuous action spaces.

  • Problem formulation: The MDP comprises states, actions, transitions, rewards, and a discount factor, with image-based states represented as sequences of consecutive frames.
  • Problem formulation: The policy should maximize discounted return and generalize to unseen MDPs, although this requires assumptions about the structure of the MDP space.
  • Deep Q-Learning: Deep Q-learning approximates the optimal state-action value function with a parameterized Q-function and minimizes a single-step Bellman residual.
  • Deep Q-Learning: Replay-buffer transitions support off-policy learning, and greedy actions are selected by maximizing the learned Q-function.
  • Deep Q-Learning: Target-network parameters can be updated using a slow-moving average, with momentum coefficient ζ ∈ (0, 1].
  • Actor-Critic Methods: Actor-critic algorithms learn a policy alongside the Q-function when maximizing over large or continuous action spaces is intractable.

4 Pitfalls of Data Augmentation in Deep Q-Learning

Naïve augmentation can improve generalization but destabilize Q-learning by increasing target variability and creating a trade-off between accurate value estimation and generalization. The paper motivates separating augmented and unaugmented data streams to address these pitfalls.

  • Augmentation Assumptions: Optimality-invariant augmentation preserves Q-values across transformed states and is intended to support generalization to unseen visual states.
  • Augmentation Assumptions: Naïve augmentation may harm Q-value estimation because it substantially enlarges the observation distribution seen during learning.
  • Pitfall 1: Non-deterministic Q-target: Augmenting both current and successor states makes the bootstrapped Q-target depend on random augmentation parameters rather than remaining deterministic.
  • Pitfall 1: Non-deterministic Q-target: Higher Q-target variance increases target volatility and can contribute to divergence, particularly in actor-critic algorithms whose targets also depend on policy-sampled actions.
  • Pitfall 1: Non-deterministic Q-target: The resulting trade-off is between accurate Q-value estimation and generalization, and optimization may still diverge depending on the augmentation choice.

5 Method

SVEA stabilizes off-policy RL under augmentation by separating augmented current-state estimation from unaugmented bootstrapping and mixing augmented and unaugmented data streams. It is designed as a general framework compatible with standard off-policy algorithms without additional parameters or forward passes.

  • Architectural Overview: SVEA operates with standard off-policy RL algorithms without changing the policy network, adding learnable parameters, auxiliary tasks, or additional forward passes.The framework uses two data streams and can be applied to image-based RL without assumptions about state structure.
  • Architectural Overview: SVEA applies augmentation only when estimating the current state's Q-value, leaving successor-state targets unaugmented for bootstrapping.This addresses erroneous bootstrapping caused by augmented data.
  • Learning Objective: The modified Q-objective jointly optimizes Q-estimation over augmented and unaugmented data streams to mitigate over-regularization.The objective uses a nonnegative linear combination whose coefficients balance the two streams.
  • Learning Objective: The mixed objective oversamples unaugmented data as an implicit variance-reduction strategy.The paper describes data mixing as working together with the modified bootstrapping scheme.
  • Learning Objective: For α = β, SVEA evaluates the objective in a single batched forward pass by concatenating augmented and unaugmented inputs along the batch dimension.The implementation commonly uses α = 0.5 and β = 0.5.
  • Algorithm: Algorithm 1 samples replay transitions, applies stochastic augmentation, optimizes the networks, and updates target parameters using an exponential moving average.When applicable, the policy is optimized from unaugmented states.

6 Experiments

Experiments evaluate SVEA against augmentation and generalization baselines across DMControl, ViT encoders, and robotic manipulation. SVEA improves stability and sample efficiency while maintaining competitive generalization across these settings.

  • Experimental setup: The evaluation covers ConvNets and Vision Transformers on DMControl benchmarks and robotic manipulation tasks under distribution shifts.Test environments vary colors, backgrounds, lighting, camera poses, and video appearance.
  • Stability: SVEA improves sample efficiency in 27 of 30 augmentation instances, while DrQ requires approximately 6x wall-time to match its efficiency.Both SVEA components contribute to its success, and DrQ degrades substantially under most augmentations.
  • Generalization: SVEA outperforms the considered methods in 12 of 15 DMControl-GB comparisons and improves DistractingCS generalization by 42% at low randomization intensity.Its generalization degrades more slowly than DrQ as test-time randomization intensity increases.
  • RL with Vision Transformers: With ViT encoders, SVEA improves generalization by 706% on Walker, walk and 233% on Cartpole, swingup over DrQ.SVEA reaches 877 mean episode return on Walker, walk color_hard, compared with 760 using CNNs.
  • RL with Vision Transformers: ViT experiments use a four-layer encoder over 144 space-time patches and evaluate training and color_hard test performance across three tasks.The ViT architecture is designed to roughly match the CNN encoder in learnable parameters.
  • Robotic manipulation: In robotic manipulation, SVEA with convolution augmentation has training efficiency similar to unaugmented DrQ, while DrQ with convolution augmentation fails on push.The comparison includes both CNN and ViT encoders and randomized test environments.

A Ablations

SVEA’s formulation outperforms alternative design variants, while ViT encoders improve SVEA but expose severe overfitting and instability in baselines.

  • Ablations: SVEA outperforms all other ablation variants in test performance by a large margin.The comparison includes encoder, objective, augmentation, and target-augmentation choices.
  • Ablations: ViT encoders further improve SVEA’s training and test performance compared with CNN encoders.
  • Ablations: Using a ViT encoder decreases DrQ’s test performance by a factor of 5.DrQ’s training performance improves slightly with ViT, but its test performance does not benefit.
  • Ablations: DrQ and RAD become unstable under strong augmentation.

B Stability under Data Augmentation

Across five DMControl tasks and six augmentations, SVEA improves stability whenever augmentation impairs DrQ, whose behavior depends strongly on task and augmentation choice.

  • Stability under Data Augmentation: SVEA improves stability in all 27 instances where DrQ is impaired by data augmentation.
  • Stability under Data Augmentation: DrQ’s stability under augmentation is highly sensitive to both the augmentation choice and the task.DrQ is relatively unaffected in Walker, stand but significantly unstable across augmentations in Cartpole, swingup.
  • Stability under Data Augmentation: SVEA can eliminate costly trial-and-error when applying data augmentation.

C Data Augmentation in RL

Data augmentation can improve generalization but may reduce sample efficiency or destabilize learning. SVEA stabilizes training across augmentations, while generalization remains augmentation- and task-dependent.

  • Data Augmentation in RL: Small random crops and shifts improve sample efficiency, whereas strong augmentations can improve generalization at the expense of sample efficiency.
  • Data Augmentation in RL: SVEA exhibits comparable stability and sample efficiency across all six evaluated augmentations.
  • Data Augmentation in RL: SVEA’s generalization on color_hard depends strongly on augmentation choice, with conv, overlay, and affine-jitter generally performing best.
  • Data Augmentation in RL: The augmentations with the best generalization also cause the most instability for the DrQ-plus-augmentation baseline.
  • Data Augmentation in RL: The study evaluates random shift, random convolution, random overlay, random cutout, Gaussian blur, affine-jitter, and random rotation.Random convolution applies a randomly initialized convolutional layer rather than a convolution operation in the usual sense.
  • Data Augmentation in RL: SVEA provides similar benefits when implemented with RAD, whose additional convolution augmentation is generally more unstable than DrQ’s.
  • Data Augmentation in RL: SVEA consistently outperforms DrQ across DistractingCS intensities and tasks, except Ball in cup, catch at the highest intensity.
  • Data Augmentation in RL: DMControl generalization benchmarks train agents in fixed environments and test them on novel environments with varying visual difficulty and variation factors.

F Additional Results on DMControl-GB

SVEA achieves competitive performance across all five tasks in the video_hard DMControl-GB benchmark.

  • Additional Results on DMControl-GB: SVEA achieves competitive performance across all 5 tasks in the video_hard benchmark.

G Robotic Manipulation Tasks

The robotic manipulation experiments evaluate visual generalization from a fixed training environment to unseen environments with randomized visual conditions. They use three goal-conditioned tasks implemented in MuJoCo and compare methods on challenging visual shifts.

  • Tasks: The study evaluates reach, reach moving target, and push using a simulated Kinova Gen3 robotic arm.The tasks cover positioning the gripper, tracking a moving target, and pushing a cube to a marked goal.
  • Training and evaluation: Each episode lasts 50 time steps, and observations are stacks of three RGB frames with no access to state information.The maximum episode return is therefore 50, while the lower return has no strict bound.
  • Baselines and benchmarks: The broader evaluation compares methods trained in fixed environments against unseen visual conditions, including video_hard backgrounds in DMControl-GB.The video_hard benchmark replaces the floor and background with natural videos and reports episode-return means and standard deviations across five runs.
  • Generalization setup: Agents train in a fixed environment but are tested with randomized camera position, colors, lighting, and backgrounds.The camera orientation continues to follow the gripper during testing, while additional visual factors are randomized.

H Implementation Details

The implementation combines standard off-policy RL components with augmentation-aware Q-learning, using CNN and ViT encoders across DMControl and robotic manipulation tasks. SVEA avoids augmenting bootstrap targets and jointly trains on augmented and unaugmented current-state observations.

  • Network architecture: The DMControl encoder is an 11-layer CNN that maps stacked 84×84×3 RGB frames to 32×21×21 feature maps.The architecture and hyperparameters follow Hansen and Wang without changes for a fair comparison.
  • Network architecture: The ViT encoder replaces the CNN with 144 image patches, 128-dimensional embeddings, four Transformer encoders, and eight attention heads.It is trained jointly with the Q-function using Adam, without pretraining or weight decay.
  • Augmentation baseline: Naive augmentation applies transformations to both the current and successor states, but strong augmentation can harm sample efficiency and stability.This approach is described as the standard augmentation operation in off-policy RL.
  • SVEA: SVEA applies strong augmentation to the current-state Q estimate but not to successor states used for bootstrapping or policy learning.It jointly optimizes the Q-function over augmented and unaugmented current-state data streams, using α = β = 0.5 by default for CNNs.
  • SVEA: The SVEA critic update concatenates each state with its augmented version while duplicating the action and Q-target across the resulting 2B batch.The critic is then updated from predictions on both state streams against the same target.
  • Experimental setup: Experiments span diverse DMControl skills and robotic manipulation tasks implemented with MuJoCo.The DMControl set includes Walker and Cartpole tasks, while the manipulation suite targets fundamental visuomotor skills.
Loading 2107.00644v2…