Don't Touch What Matters: Task-Aware Lipschitz Data Augmentation for Visual Reinforcement Learning

by   Zhecheng Yuan, et al.
Tsinghua University
Stanford University

One of the key challenges in visual Reinforcement Learning (RL) is to learn policies that can generalize to unseen environments. Recently, data augmentation techniques aiming at enhancing data diversity have demonstrated proven performance in improving the generalization ability of learned policies. However, due to the sensitivity of RL training, naively applying data augmentation, which transforms each pixel in a task-agnostic manner, may suffer from instability and damage the sample efficiency, thus further exacerbating the generalization performance. At the heart of this phenomenon is the diverged action distribution and high-variance value estimation in the face of augmented images. To alleviate this issue, we propose Task-aware Lipschitz Data Augmentation (TLDA) for visual RL, which explicitly identifies the task-correlated pixels with large Lipschitz constants, and only augments the task-irrelevant pixels. To verify the effectiveness of TLDA, we conduct extensive experiments on DeepMind Control suite, CARLA and DeepMind Manipulation tasks, showing that TLDA improves both sample efficiency in training time and generalization in test time. It outperforms previous state-of-the-art methods across the 3 different visual control benchmarks.


page 4

page 5

page 6

page 7

page 12

page 13

page 15


Reinforcement Learning with Augmented Data

Learning from visual observations is a fundamental yet challenging probl...

Generalization of Reinforcement Learning with Policy-Aware Adversarial Data Augmentation

The generalization gap in reinforcement learning (RL) has been a signifi...

Enhancing Generalization and Plasticity for Sample Efficient Reinforcement Learning

In Reinforcement Learning (RL), enhancing sample efficiency is crucial, ...

Efficient Embedding of Semantic Similarity in Control Policies via Entangled Bisimulation

Learning generalizeable policies from visual input in the presence of vi...

Bootstrap Advantage Estimation for Policy Optimization in Reinforcement Learning

This paper proposes an advantage estimation approach based on data augme...

CROP: Towards Distributional-Shift Robust Reinforcement Learning using Compact Reshaped Observation Processing

The safe application of reinforcement learning (RL) requires generalizat...

Learning Better with Less: Effective Augmentation for Sample-Efficient Visual Reinforcement Learning

Data augmentation (DA) is a crucial technique for enhancing the sample e...

Please sign up or login with your details

Forgot password? Click here to reset