Task-Aware Lipschitz Confidence Data Augmentation in Visual Reinforcement Learning From Images

Biao Wu, Haitao Wang, Hejun Wu

Abstract: Visual reinforcement learning is a technique that learns effective policies from image pixels. Data augmentation is widely adopted in visual reinforcement learning to improve the generalization of the learned policies as data augmentation increases data diversity. However, applying data augmentation to all pixels simultaneously results in a divergence in action distribution as well and degrades the training stability in turn. Additionally, existing methods compute task weights for each pixel and apply augmentation methods separately based on tasks. As a result, they require a significant amount of computational resources. To enhance both the training stability and computational efficiency in the computation of task weights, we propose a Task-Aware Lipschitz Confidence (TALC) data augmentation method for visual reinforcement tasks. TALC calculates the task-aware confidence of all pixels on the image at once, only enhancing low confidence pixels to increase data diversity. We have conducted experiments on DeepMind Control suite tasks and the results demonstrate that TALC not only improves the training efficiency, but also enhances the generalization ability during testing. Overall, TALC out-performs existing methods in most different visual control benchmarks.