PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalization in Reinforcement Learning by Soft Data Augmentation

Nicklas Hansen, Xiaolong Wang

arXiv:2011.13389v2cs.LG

TL;DR

Vision-based RL struggles to generalize, and stronger augmentation can make optimization less efficient and stable. SODA decouples augmented representation learning from policy learning, using non-augmented data for RL and a soft consistency constraint for the encoder. Across DMControl-GB and robotic manipulation, it improves sample efficiency, generalization, and training stability over prior methods.

  • Problem

    Adding visual variation can make RL optimization challenging, reducing sample efficiency and stability while leaving the right augmentation strength task-dependent.

  • Method

    SODA uses augmented data only for auxiliary representation learning and strictly non-augmented data for policy learning, imposing a soft consistency constraint on the shared encoder.

  • Results

    SODA improves sample efficiency, generalization, and stability over prior methods across DMControl-GB and robotic manipulation evaluations.

  • Takeaways & Limitations

    Separating augmentation from policy learning preserves augmentation benefits while interfering minimally with the RL objective.

Abstract

from arXiv · show

Extensive efforts have been made to improve the generalization ability of Reinforcement Learning (RL) methods via domain randomization and data augmentation. However, as more factors of variation are introduced during training, optimization becomes increasingly challenging, and empirically may result in lower sample efficiency and unstable training. Instead of learning policies directly from augmented data, we propose SOft Data Augmentation (SODA), a method that decouples augmentation from policy learning. Specifically, SODA imposes a soft constraint on the encoder that aims to maximize the mutual information between latent representations of augmented and non-augmented data, while the RL optimization process uses strictly non-augmented data. Empirical evaluations are performed on diverse tasks from DeepMind Control suite as well as a robotic manipulation task, and we find SODA to significantly advance sample efficiency, generalization, and stability in training over state-of-the-art vision-based RL methods.

I. INTRODUCTION

Vision-based RL agents often fail to generalize across visually different environments, while stronger augmentation can make optimization less efficient and stable. SODA separates augmented-data representation learning from non-augmented policy learning to address this trade-off.

  • Vision-based RL agents frequently fail to generalize to new environments, especially with high-dimensional image observations.
  • Domain randomization can improve robustness, but adding more variation makes RL optimization harder and may reduce sample efficiency and stability.
  • Data augmentation can improve sample efficiency and generalization, yet effective augmentations and strengths are task-dependent and naive use may harm performance.
  • SODA decouples augmentation from policy learning: augmented data supports auxiliary representation learning, while policy optimization uses strictly non-augmented observations.
  • SODA learns representations that maximize shared information between augmented and non-augmented observations through a soft constraint on the shared encoder.
  • SODA improves sample efficiency and generalization over state-of-the-art vision-based RL methods in 9 of 10 DMControl-GB environments and all 3 robotic-manipulation test-environment types.

II. RELATED WORK

Prior work improves visual RL through self-supervision, contrastive learning, data augmentation, and domain randomization. These approaches demonstrate benefits but leave augmentation choices and policy-training effects task-dependent or unpredictable.

  • Self-supervised auxiliary learning has been combined with RL to improve sample efficiency and performance.
  • Contrastive learning with RL can nearly match state-based RL in sample efficiency, using auxiliary representation learning alongside policy optimization.
  • Domain randomization transfers visual representations across environments by varying scene factors such as textures.
  • Augmenting observations during policy training can promote invariance to visual perturbations, but the useful amount of randomization is unpredictable and task-dependent.

III. METHOD

SODA is a self-supervised auxiliary task that shares an encoder with an RL policy while separating representation learning on augmented data from RL on non-augmented data. Its architecture uses online and target projections with a prediction head and EMA updates.

  • SODA shares a common encoder with the RL policy and adds a self-supervised auxiliary task.
  • Representation learning uses augmented observations, whereas reinforcement learning uses non-augmented observations.
  • An augmented observation is encoded and projected into z′, while the original observation is encoded by an EMA target network into z⋆.
  • A prediction head maps the augmented representation toward the original representation through a consistency objective.

B. Representation Learning by SODA

SODA learns representations from augmented and non-augmented views while optimizing the RL policy only on non-augmented observations. Its auxiliary consistency task and alternating updates share the encoder with the policy.

  • Representation learning: SODA predicts the target projection of a non-augmented observation from an augmented observation’s latent feature.The augmented view passes through the shared encoder, projection, and predictor; the original view uses momentum-updated target components.
  • Representation learning: The method jointly optimizes the projection, predictor, and shared encoder with RL using a consistency loss.The representation-learning components are trained alongside the RL task.
  • Policy learning: RL uses the shared encoder on non-augmented observations without changing the policy architecture or inputs.Training alternates between the RL objective and the SODA objective, while test-time inference uses only the encoder and policy.
  • Optimization procedure: Algorithm 1 alternates RL updates on transition batches with SODA updates on augmented observation batches and momentum updates of target parameters.Each iteration samples transitions for L_RL, then observations for L_SODA, followed by ψ ← (1 − τ)ψ + τθ.

C. Data Augmentation as a Mutual Information Problem

SODA frames generalization as consistency between representations of original observations and augmented views. Its consistency objective approximates mutual-information maximization while encouraging the encoder to discard augmentation noise and irrelevant visual variation.

  • Mutual-information formulation: SODA seeks representations that preserve shared information between an observation and its augmented view.The encoder maps different views of the same underlying state to similar compact feature vectors.
  • Augmentations: Random convolution and random overlay are the two augmentations considered in the paper.Random convolution follows earlier work, while random overlay is introduced as novel here.
  • Mutual-information formulation: The mutual information between latent views is bounded by the mutual information between the original and augmented observations.The paper states I(z⋆; z′) ≤ I(o; o′).
  • Consistency objective: Because direct mutual-information optimization is impractical, SODA uses a consistency loss as an approximation.The approximation assumes the target encoder and projector maximally preserve information in the original observation.
  • Representation effect: Strong and varied augmentation can make the encoder ignore task-irrelevant variation and reduce observational overfitting.The paper qualifies the approximation by noting that maximal information preservation may not hold in practice.

IV. EXPERIMENTS

The experiments evaluate SODA on DMControl generalization tasks and simulated robotic manipulation under unseen visual conditions. The setup uses fixed training environments, randomized or video-based test distributions, and a SAC implementation with strong augmentations.

  • Evaluation setting: SODA is evaluated on five DeepMind Control tasks and a robotic manipulation task.The evaluation introduces DMControl-GB for vision-based RL generalization.
  • Evaluation setting: Agents train in fixed environments and are tested on distinct distributions with randomized colors or natural-video backgrounds.These correspond to the color hard and video easy DMControl-GB benchmarks.
  • Evaluation setting: The robotic manipulation evaluation uses fixed-environment training and tests randomized colors, video backgrounds, and additional random perturbations.The task is performed on a simulated robotic arm to better emulate deployment scenarios.
  • Implementation: SODA is implemented on top of Soft Actor-Critic with 100-dimensional projections and momentum coefficient τ = 0.005.The implementation uses 100 × 100 frame observations, with three frames for DMControl-GB and one for robotic manipulation.
  • Data augmentation: The study distinguishes weak augmentations that may improve sample efficiency from strong augmentations that improve generalization at a sample-efficiency cost.Temporally consistent random cropping is applied by default and treated as part of the non-augmented observation pipeline.
  • Data augmentation: Random convolution and random overlay are used as strong augmentation choices for the experiments.Random overlay interpolates observations with images sampled from the Places dataset using α = 0.5.

C. Baselines

The study compares SODA with SAC-based baselines using domain randomization or direct data augmentation. Results indicate that decoupling augmentation from RL preserves training behavior while improving generalization and end performance.

  • Baselines: The baselines include SAC/RAD, SAC with domain randomization, SAC with random convolution, and SAC with random overlay.All baselines use random cropping, while SODA variants use convolution or overlay in the auxiliary formulation.
  • Random convolution: With random convolution, SODA improves generalization without destabilizing training and has sample efficiency similar to SAC.SAC converges on the training environment but generalizes poorly, while SAC (DR) and SAC (conv) fail or become sub-optimal.
  • Soft augmentation comparison: Applying augmentation to both representation learning and RL improves over augmenting only RL, but the proposed SODA formulation performs better on sample efficiency and end performance.The comparison is reported for walker walk and walker stand training returns.

E. Generalization to Unseen Environments

SODA is evaluated on unseen visual environments and outperforms prior methods across most benchmark instances. It also improves sample efficiency and generalization in random-overlay and robotic-manipulation comparisons.

  • Random-overlay SODA achieves better sample efficiency than SAC (overlay) and similar generalization to SODA (conv) on random colors.This occurs despite minimal visual similarity between random overlays and the random color environment.
  • SODA outperforms previous methods in 9 out of 10 instances, by as much as 81% on ball in cup catch with video backgrounds.It also outperforms the section IV-C baselines in all tasks except walker stand on video backgrounds.
  • In robotic manipulation, SODA outperforms all baselines in both sample efficiency and generalization while reducing variance.The comparison uses average returns from 5 runs, with shaded standard deviations.

F. Robotic manipulation

The robotic manipulation evaluation tests whether SODA generalizes beyond fixed training visuals to randomized colors, video backgrounds, and deployment-relevant visual changes. SODA outperforms baselines across all considered test environments while improving sample efficiency.

  • The task requires a robotic arm to push a yellow cube to a red disc using dense rewards over 50-step episodes.Each successful timestep gives reward 1, with an additional distance-proportional penalty.
  • SODA outperforms baselines in both sample efficiency and generalization across random colors, video backgrounds, and randomized camera, lighting, and texture.These test variations are intended to simulate real-world deployment conditions.
  • SODA’s reported improvements support its use for reducing observational overfitting in robotic-manipulation policy deployment.The paper presents this as a suggestion based on the observed test-environment results.

APPENDIX I DMCONTROL GENERALIZATION BENCHMARK

DMControl-GB trains agents in fixed environments and evaluates them on unseen, visually diverse test distributions. SODA performs strongly on easier benchmark settings but retains room for improvement on the most challenging video-background tasks.

  • DMControl-GB evaluates agents trained in fixed environments on unseen test distributions with randomized colors or natural-video backgrounds.The benchmark includes progressively broader background replacement, from skybox-only changes to full-background replacement.
  • SODA outperforms previous state-of-the-art methods in 9 out of 10 environments and achieves near-optimal generalization on color easy.The benchmark comparison uses standardized architectures and hyper-parameters across implementations.
  • On video hard, SODA shows substantial improvements over baselines but still has room for improvement on the majority of tasks.The paper associates the training-to-test performance gap with observational overfitting and notes reliance on shadows correlated with reward signals.
  • Table III reports average returns for SODA and state-of-the-art methods on color easy and video hard, with means and standard deviations over 5 runs.Results for color hard and video easy are reported in Table I.

APPENDIX II IMPLEMENTATION DETAILS

The implementation uses standardized SAC-based architectures and hyper-parameters across DMControl-GB experiments. It evaluates random convolution and random overlay as the principal augmentation choices, with specified update settings.

  • SODA and baselines use Soft Actor-Critic with the architecture and hyper-parameters adopted from prior work.The same algorithmic setup is used for robotic manipulation, with different frame-stack sizes across domains.
  • DMControl-GB uses 3-frame observation stacks, whereas robotic manipulation uses single-frame observations.The implementation keeps algorithms, architectures, and hyper-parameters otherwise identical across these settings.
  • The evaluated augmentations are random convolution and random overlay, with overlay interpolating observations and Places images using α = 0.5.Places provides 1.8M diverse scenes for the overlay augmentation.
  • The target components fψ and gψ use momentum coefficient τ = 0.005, while SODA updates occur after every second RL update.Both LRL and LSODA are optimized with Adam.
Loading 2011.13389v2…