Reinforcement Learning with Augmented Data

Overview

The paper proposes Reinforcement Learning with Augmented Data (RAD), a simple framework that applies stochastic data augmentations to observations before using them in standard RL updates. The main idea is that different augmented views of the same observation should preserve the task-relevant information and therefore lead to consistent policy and value predictions. Training on these views discourages the agent from relying on irrelevant details, improving both data efficiency and generalization.

RAD does not introduce an auxiliary representation-learning loss or modify the underlying RL objective. Instead, it can be added directly to both off-policy algorithms, such as SAC, and on-policy algorithms, such as PPO. The paper studies augmentations for both image and state observations.

Method

During training, RAD samples a mini-batch from a replay buffer or a recent on-policy trajectory, randomly augments its observations, and then performs the original RL update normally. Thus, the action, reward, and learning objective remain unchanged; only the observations supplied to the policy and value networks are transformed.

Image-Based Observations

The paper investigates the following augmentations for pixel-based observations:

• Random crop: extracts a randomly positioned patch from the original image, encouraging translation invariance by changing where objects appear in the resulting observation.
• Random translation: places the complete image at a random position inside a slightly larger empty frame. Unlike cropping, this changes object position without removing information from the original image.
• Random window: masks a randomly selected boundary region, reproducing the information-removal component of random cropping without translating the remaining content.
• Grayscale: converts the RGB observation to grayscale with a specified probability, encouraging the agent to rely less on color.
• Cutout: replaces a randomly selected rectangular region with black pixels, encouraging robustness to partial occlusion.
• Colored cutout: replaces a randomly selected rectangular region with a randomly generated color rather than black pixels.
• Horizontal flip: reflects the observation across its vertical axis with a specified probability.
• Rotation: randomly rotates the observation by $0^\circ$, $90^\circ$, $180^\circ$, or $270^\circ$.
• Random convolution: passes the observation through a randomly initialized convolutional layer, changing its colors and visual appearance while retaining its spatial structure.
• Color jitter: converts the image to HSV and perturbs its hue, saturation, and value channels to produce color variation.

When an observation contains a stack of consecutive frames, the same randomly sampled augmentation is applied to every frame in the stack. This preserves temporal information such as object motion and velocity; applying a different transformation to each frame would create artificial motion.

State-Based Observations

The paper considers two transformations for proprioceptive state vectors containing features such as positions and velocities:

• Random amplitude scaling:
  $s' = s \odot z, \qquad z \sim \operatorname{Uniform}(\alpha,\beta)$
  where the state is multiplied by either a scalar or an element-wise random vector. This changes the magnitude of the inputs while preserving intrinsic information such as their signs.
• Gaussian noise:
  $s' = s + z, \qquad z \sim \mathcal{N}(0,I)$
  which encourages robustness to small offsets in the state features.

As with stacked image observations, the randomization is temporally consistent: the same sampled transformation is applied to the current and next states of a transition.