Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels

Ilya Kostrikov; Denis Yarats; Rob Fergus

Source-linked AI summary

Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels

Ilya Kostrikov, Denis Yarats, Rob Fergus

arXiv:2004.13649v4Apr 28, 2020cs.LGcs.CVeess.IVstat.ML

TL;DR

Sample-efficient pixel-based RL struggles with visual representation learning under limited interaction, motivating a method that regularizes value estimates through image transformations. The resulting DrQ approach improves model-free agents across continuous and discrete-control benchmarks, while remaining simple and computationally light.

Problem
Limited interaction, correlated samples, and sparse rewards make it difficult to train visual encoders and policies directly from pixels in sample-efficient RL.
Method
DrQ applies image transformations and regularizes Q-targets and Q-values using optimality-invariant transformations, without requiring auxiliary losses or a world model.
Results
DrQ achieves state-of-the-art performance on the DeepMind control suite and outperforms other methods on the Atari 100k median metric.
Takeaways & Limitations
The method provides a simple, broadly applicable way to train model-free agents directly from pixels without unsupervised auxiliary losses or a world model.
Takeaways & Limitations
Apart from action-repeat, the algorithm's hyperparameters are fixed across the six DeepMind control tasks.

Abstract

from arXiv · show

We propose a simple data augmentation technique that can be applied to standard model-free reinforcement learning algorithms, enabling robust learning directly from pixels without the need for auxiliary losses or pre-training. The approach leverages input perturbations commonly used in computer vision tasks to regularize the value function. Existing model-free approaches, such as Soft Actor-Critic (SAC), are not able to train deep networks effectively from image pixels. However, the addition of our augmentation method dramatically improves SAC's performance, enabling it to reach state-of-the-art performance on the DeepMind control suite, surpassing model-based (Dreamer, PlaNet, and SLAC) methods and recently proposed contrastive learning (CURL). Our approach can be combined with any model-free reinforcement learning algorithm, requiring only minor modifications. An implementation can be found at https://sites.google.com/view/data-regularized-q.

1 Introduction

Pixel-based, sample-efficient RL struggles to train useful visual representations, motivating augmentation-based regularization instead of auxiliary objectives or pre-training. The paper applies image transformations and value-function regularization to improve model-free learning from pixels.

Motivation: Large image encoders over-fit, whereas small encoders learn impoverished representations that limit task performance.Figure 1 illustrates worsening performance as encoder capacity increases for unmodified SAC.
Motivation: Limited environment interaction, correlated samples, and sparse rewards make jointly training convolutional encoders and policies difficult.The sample-efficient regime may provide only 10^4–10^5 transitions from a few hundred trajectories.
Approach: The paper focuses on data augmentation rather than self-supervised pre-training or auxiliary losses unrelated to the control task.Standard image transformations perturb observations while regularizing the critic so transformed images receive similar Q-function values.
Contributions: Straightforward pixel augmentation reduces over-fitting without changing the underlying reinforcement-learning algorithm.The approach is presented as requiring no additional auto-encoder, dynamics-model, or contrastive loss.
Contributions: Combined with vanilla SAC, the approach achieves state-of-the-art performance on the DeepMind control suite and also does so with a DQN-like agent on Atari 100k.The paper also provides an implementation combined with SAC and DQN.

2 Background

The paper formulates image-based control as a partially observable decision process and converts it to an MDP by stacking consecutive observations. It then describes SAC and DQN as the model-free foundations used in the work.

Image-based control: Image-based control is modeled as an infinite-horizon partially observable Markov decision process with image observations, actions, dynamics, rewards, and discounting.The observation space consists of high-dimensional image pixels.
Image-based control: The transition distribution conditions on observation history and the current action, while rewards map the history and action to a scalar.The discount factor satisfies γ ∈ [0, 1).
Image-based control: The POMDP is converted into an MDP by stacking several consecutive image observations into a state.The resulting transition and reward functions are redefined over the stacked state and current action.
Model-free algorithms: SAC learns a state-action value function, stochastic policy, and temperature by optimizing a discounted maximum-entropy objective.The formulation uses an MDP with state, action, transition, reward, and discount components.
Model-free algorithms: DQN approximates the Q-function with a convolutional network for discrete actions, where the policy can be inferred directly from Q-values.Practical DQN systems commonly combine refinements known as Rainbow.

3 Sample Efficient Reinforcement Learning from Pixels

The method starts with random-shift augmentation and extends it into DrQ, which regularizes Q-targets and Q-values using optimality-invariant image transformations. These mechanisms are integrated into generic off-policy actor-critic training with minimal algorithmic changes.

3.1 Image Augmentation: SAC trained with different encoder capacities exhibits over-fitting, motivating augmentation for sample-efficient pixel-based control.The experiment varies encoder architectures drawn from recent reinforcement-learning methods.
3.1 Image Augmentation: Random shifts are selected because they balance simplicity and performance while preserving the task better than many alternative transformations.Augmentation is applied to replay-buffer images, padding each side by 4 pixels and taking a random 84 × 84 crop.
3.1 Image Augmentation: Random-shift augmentation greatly reduces over-fitting and closes performance gaps between encoder architectures.The augmentation is applied during training rather than during environment sample collection.
3.2 Optimality Invariant Image Transformations: An optimality-invariant transformation preserves Q-values across transformed states, allowing multiple surrogate states with the same action values.For transformations f(s, ν), the paper defines Q(s, a) = Q(f(s, ν), a).
3.2 Optimality Invariant Image Transformations: The framework reduces Q-estimation variance by generating K transformed samples instead of relying on a single state sample.The transformation parameters are sampled from a set of possible parameters.
3.2 Optimality Invariant Image Transformations: DrQ regularizes Q-learning by augmenting target-value computation and by applying the same target to multiple augmentations of the current state.These are the two value-function regularization mechanisms introduced beyond basic image augmentation.
3.3 Our approach: Data-regularized Q (DrQ): DrQ combines input transformations, averaging Q-targets over K transformations, and averaging Q-values over M transformations.With K = 1 and M = 1, DrQ reduces to image transformations alone and can be applied to generic model-free algorithms.
3.3 Our approach: Data-regularized Q (DrQ): Experiments pair DrQ with SAC and DQN using image shifts with ν ± 4 and K = 2, M = 2 for target-Q and Q augmentation.Ablations show gains over unaugmented SAC, while Figure 2 compares combinations of the regularization techniques.

4 Experiments

DrQ is evaluated on DeepMind Control Suite benchmarks and Atari 100k, using fixed hyperparameters across tasks where stated. It achieves state-of-the-art results while improving efficiency and simplicity relative to competing approaches.

DeepMind Control Suite: DrQ outperforms other approaches at both 100k and 500k environment steps on the PlaNet benchmark.The table distinguishes data-efficient and asymptotic performance regimes; exploration-step accounting differs for SLAC and DrQ.
Evaluation setup: The experiments compare DrQ with leading model-free and model-based methods using results from corresponding authors and ten random seeds.Performance is averaged over ten episodes every 10,000 environment steps, with ±1 standard deviation shading.
DeepMind Control Suite: DrQ achieves state-of-the-art performance on all PlaNet benchmark tasks and compares favorably in wall-clock time.The comparison includes PlaNet, SAC-AE, CURL, SLAC, and SAC trained from internal states.
DeepMind Control Suite: DrQ outperforms Dreamer on 12 of 15 selected Dreamer benchmark tasks and approaches state-based SAC performance on many tasks.Dreamer remains superior on three tasks.
Atari 100k: On Atari 100k, DrQ combined with Efficient DQN achieves state-of-the-art performance and surpasses OTRainbow, Data Efficient Rainbow, and CURL.The setup evaluates median human-normalized episode returns after 125k environment steps.

5 Related Work

The paper situates DrQ among augmentation, regularization, generalization, and pixel-based reinforcement-learning methods. Unlike prior augmentation and contrastive approaches, it applies image transformations to control tasks and compares against model-based and model-free benchmarks.

Data augmentation: Computer-vision augmentation uses label-preserving transformations such as translations, scales, and color shifts to improve generalization.These transformations expand the training set without changing object-class semantics.
Data augmentation: Contrastive-learning methods impose invariance to image transformations for downstream recognition, whereas DrQ targets control tasks with different invariance requirements.The distinction is between recognition-oriented representation learning and control-oriented value learning.
RL regularization: Earlier reinforcement-learning regularization includes ℓ2 Q-function regularization, entropy regularization, dropout, and cutout.Entropy regularization adds causal entropy to rewards, making the Q-function smoother and facilitating optimization.
Generalization: RL generalization benchmarks and domain-randomization methods deliberately vary colors, textures, backgrounds, or viewpoints, while DrQ does not assume control of data generation.The cited benchmarks include Robot Learning in Homes, Meta-World, and ProcGen.
Pixel-based RL: Pixel-based continuous-control methods include model-based approaches that learn system dynamics and model-free approaches that learn directly from observations.The paper compares DrQ with methods from both groups in its benchmark experiments.

6 Conclusion

The method improves SAC trained directly from image pixels while adding negligible computational burden. It outperforms competing approaches across DeepMind control tasks and Atari 100k metrics, with robustness to hyper-parameter choices.

The regularization technique significantly improves SAC performance on standard continuous-control tasks from image pixels.The method adds a negligible computational burden.
On the DeepMind control suite, the method outperforms state-of-the-art approaches on the majority of tasks.
On Atari 100k, the method outperforms other methods on the median metric.
The method is robust to the choice of hyper-parameters.

B.4 PlaNet and Dreamer Benchmarks

The PlaNet and Dreamer benchmarks evaluate image-based control across DeepMind control tasks under differing task and hyper-parameter settings. The setups use stacked pixel observations and report comprehensive experimental configurations.

PlaNet Benchmark: The PlaNet benchmark contains six DeepMind control tasks with task-specific action-repeat hyper-parameters.
Dreamer Benchmark: The Dreamer benchmark uses a larger, more difficult task set and fixes action repeat to 2 across tasks.
Input Processing: Each observation is a 3-stack of 84 × 84 RGB frames from the 0th camera.
Experimental Configuration: The experiments report hyper-parameter settings for the DeepMind control suite and Atari 100k benchmarks.

D Full Atari 100k Results

The Atari 100k evaluation reports performance across 26 Atari games using episode returns averaged over random seeds. The method demonstrates better overall performance in the reported median metric.

26 Atari games are evaluated using median human-normalized episode returns and mean episode returns for each individual game.
The results are recorded at the end of training and averaged across 5 random seeds, except CURL, which uses 3 seeds.
The method demonstrates better overall performance on the Atari 100k benchmark.

E Image Augmentations Ablation

The ablation compares common image augmentations for pixel-based reinforcement learning. Random shifts are identified as the most effective technique for DeepMind control tasks and are implemented through padded random crops.

Compared Augmentations: The evaluated augmentations include random shifts, cutouts, flips, rotations, and intensity jittering.
Random Shift: Random shifts pad each 84 × 84 image by 4 pixels per side and randomly crop back to 84 × 84.
Cutout: Cutouts mask square image regions rather than individual pixels.
Flips: Horizontal and vertical flips are applied with probability 0.1.
Intensity: Intensity augmentation multiplies each image tensor by a scalar using µ = 1.0, σ = 0.1, and clipped Gaussian noise.
Ablation Result: Random shifts are the most effective augmentation for tasks from the DeepMind control suite.

F K and M Hyper-parameters Ablation

Increasing K and M generally improves agent performance, particularly on harder tasks such as Cheetah Run. The selected K=2, M=2 setting balances performance and computational demands.

Increasing K and M generally improves agent performance, especially on harder tasks such as Cheetah Run.
K=2, M=2 is chosen as a balance between performance and computational demands.

G Robustness Investigation

DrQ is generally robust to the studied hyper-parameter settings across three tasks, with low variation in performance. Its augmentation strategy also supports additional training transformations, although more updates per environment step can worsen over-fitting.

Low variance in the performance curves and heat maps indicates robustness to exact hyper-parameter settings.
DrQ is largely invariant to the studied mini-batch sizes, learning rates, and initial temperatures across three tasks.The study evaluates mini-batch sizes 128, 256, and 512; learning rates from 0.0001 to 0.005; and initial temperatures from 0.005 to 0.1, using three seeds.
The main experiments use Adam’s default learning rate of 0.001, although it is not always optimal.
Performing more updates per environment step can lead to worse over-fitting despite the available augmented observations.
DrQ generates varied transformations from each training observation, enabling investigation of more training updates per environment step for sample efficiency.