This is the support material for paper : "Don’t Touch What Matters: Task-Aware Lipschiz Data Augmentation for Visual Reinforcement Learning".

Abstract:

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 taskagnostic 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 for stability. We verify the effectiveness of our approach on DeepMind Control suite, CARLA and DeepMind Manipulation tasks. The extensive empirical results show that TLDA improves both sample efficiency and generalization; it outperforms previous state-of-the-art methods across 3 different visual control benchmarks.

Method：

Some examples of TLDA:

As shown in the above examples, blindly augmenting the image might touch the pixels in this area and cause catastrophic action/value changes. SVEA only uses strong augmentation but retains no raw pixel, while TLDA preserves the critical parts of the original observations.

Generalization Performance:

DMC-GB

CARLA

Manipulation Tasks

Code：https://github.com/gemcollector/TLDA