Reinforcement Learning with Augmented Data

Introduction

To solve sequential decision-making problems Reinforcement Learning (RL) has achieved remarkable success, still it suffers from a fundamental limitation, as large amount of interaction data is needed to learn effective policies. Here we (along with my team mate @coderishiraj )are writing this blog to go through and build an understanding of the paper — Reinforcement Learning with Augmented Data by authors — Michael Laskin and Kimin Lee and Adam Stooke and Lerrel Pinto and Pieter Abbeel and Aravind Srinivas.

This paper aims to prove that data augmentation works for Reinforcement learning formally phrased as — The paper Reinforcement Learning with Augmented Data (RAD) proposes a simple and powerful idea: improve both data efficiency and generalisation by incorporating data augmentation into Reinforcement Learning. The central idea is that the optimal action or value should not be changed by different transformations of the same observations

The paper puts forward the following Hypothesis -

1. Data Augmentation is Sufficient: The authors want to prove that simple data augmentations alone can enable RL algorithms to outperform complex, state-of-the-art methods without changing the underlying RL algorithm or adding auxiliary losses.

2. Enhanced Generalization: They aim to prove that RAD significantly improves an agent’s ability to perform in unseen environments (zero-shot generalization), specifically within the OpenAI ProcGen benchmark suite.

3. Universal Applicability: The paper intends to show that RAD is a general framework that works for both on-policy (e.g., PPO) and off-policy (e.g., SAC) reinforcement learning.

Adding to these the authors also introduces 2 new kinds of data augmentations namely — Random Translation and Random Amplitude Scaling.

PROBLEM FORMULATION

Reinforcement Learning is typically modeled as a Markov Decision Process (MDP) (S, A, P, γ). The objective of RL is to maximize the expected discounted return:

the agent receives observations oₜ ∈ Rᵏ instead of the true state sₜ, in pixel based RL. Hence, a representation learning problem is introduced, where the agent must learn both a policy π(a|oₜ) and a feature mapping f(oₜ). The model may over-fit to irrelevant pixel-level details which leads to poor generalization, without additional constraints.

ROLE OF CNN AND IMAGE DATA

In RAD images serve as the dimensional sensory data and CNN acts as a critical feature extractor for the Policy Network.

1. High-Dimensional Observations: The agent receives image-based observations (oₜ), which are indirect, high-dimensional representations of the environment’s true state (sₜ​).
2. Temporal Information: To help the agent understand motion and velocity, multiple consecutive image frames are “stacked” together to form a single observation.
3. Target for Augmentation: Images are the subject of ten different stochastic transformations — such as random crop, translate, and color jitter — designed to inject “priors” like translation in variance into the learning process.

The CNN functions as an encoder that processes the augmented pixel mini batches into a lower-dimensional latent space for the RL algorithm.

This image has been created using Google Gemini

STANDARD RL OBJECTIVE

The critic is trained by minimising the Bellman error in the Soft Actor Critical (SAC) framework.

and valued function is defined as:

The policy is optimised using:

All these objectives depend directly on the observation oₜ. However, the objective does not enforce consistency between them, if two observations represent the same state but differ visually, which leads to poor robustness.

RAD: Augmented Data in RL

A stochastic augmentation function T applied to observations is introduced by RAD. Instead of training on (oₜ, aₜ, rₜ, oₜ₊₁), RAD trains on:

Where Oₜ is the Observation, aₜ is the action, rₜ is the reward at time t and Oₜ₊₁ is the observation at time t + 1. T(.) denotes the transformation applied.

INTUITION BEHIND RAD : Optimal value function depends only on underlying state s, not on its observation o. Both o and T (o) correspond to the same state only when T preserves the semantics of the state. So, the optimal function satisfies:

To improve generalisation and reducing over-fitting, RAD enforces this property implicitly by training on augmented data.

EXPERIMENTAL RESULTS

The researchers used DeepMind Control Suite which is the standard set of virtual environments that are used to train agents to control robots, on various parameters such as Locomotion, Balance and Goal Reaching. Where RAD was compared with various benchmark models, here is the reported metrics for performance of RAD on various tasks.

Authors summarised the main findings from this table as the following —

1. RAD proved to be the state of the art model which beats all the other models in both the 500K steps and 100K steps score, thus establishing its performance benchmark.
2. RAD matches the scores of state based SAC that is the ultimate benchmark at DMCS as it receives not just the image feed but also sensor signals directly to make decisions.
3. Random Translation and Random Crop as transformations performed the best for adaptability, among these 2 Random Crop proved to be most effective.

DMCS is a good benchmark for performance related tasks but to prove the generalisation it is important to test the model performance on some other kind of data where there are some changes in the training and test data, here the authors have used OpenAI ProcGen, as one of the key characteristics is that the training and test environments differ in both the visual appearance and structure. Authors used zero-shot learning, and PPO for policy optimisation. Here are the results reported by Authors-

Comparison Result table taken from the paper

Here is how the authors have summarised the finding:

1. Data augmentation methods like Random crop and coutout are proved to be better at generalisation.
2. RAD trained on 100 levels outperforms PPO on 200 levels signify that data augmentation can be effective in generalisation.
3. In the navigation task the gain from data augmentation was not as significant.

To prove that data augmentation works not only for images but also for raw data (numbers) there must be a study of augmented RL algorithm and state of the art models, this is where the authors have used OpenAI Gym which gives data input such as velocity, direction to train the models. Here RAD was compared with state of the art models and here are the comparison table for the study.

Comparison table taken from paper

Here in this part of experiment the table provides evidences that RAD is comparable to the best models even surpassing many of them in terms of performance, here as well RAD performs better than most of the models except POPLIN-A and POPLIN-P in specific parameters.

EXPERIMENTAL ANALYSIS Across multiple benchmarks, RAD achieves state-of-the-art performance. With up to 4× gains over baseline methods, one key result is a significant improvement in data efficiency,. This can be interpreted as an effective increase in dataset size:

Additionally, RAD improves generalization to unseen environments which indicates that it learns robust and invariant representations.

Figure Taken from the paper

CRITICAL ANALYSIS

Despite effectiveness, RAD has several limitations. First, it does not address structural generalisation, such as changes in environment dynamics or layout. Second, inappropriate transformations can harm performance and the choice of augmentation is crucial. Third, the method is primarily supported by empirical results and lacks theoretical guarantees. Finally, other challenges such as exploration or reward sparsity has not been addressed by RAD.

FUTURE DIRECTIONS

There are few improvements which can be explored. First one is to learn the augmentation function T_θ adaptively. Second one is to combine RAD with explicit consistency regularization:

Extending RAD to other domains such as language models and incorporating causal reasoning and are also promising directions.

CONCLUSION

RAD shows that data augmentation can significantly improve reinforcement learning by enforcing invariance in learned representations.It achieves better data efficiency and generalization without increasing model complexity, by modifying only the input distribution, which highlights the importance of inductive bias in designing effective learning systems.