A Minimalist Approach to Augmentation-based Self-supervised Representation Learning for On-policy Reinforcement Learning

Abstract

Data augmentation has been proven as an effective measure to improve generalization performance in reinforcement learning (RL). Generic approaches directly use the augmented data to learn the value estimate or regularize the estimation, often ignoring explicit learning representation. On the other hand, algorithms that employ self-supervised representation learning mechanisms through data augmentation come at the cost of additional complexity. Such RL algorithms introduce new hyperparameters and design choices, often requiring additional components such as negative samples, projection heads, etc. We identify incorporating such mechanisms with on-policy RL algorithms poses additional challenges. In this work, we aim to develop a deep RL algorithm that ensures the benefits of both data augmentation and representation learning while incorporating minimal changes to the pre-existing on-policy RL algorithm. We find that we can match the performance of the state-of-the-art data augmentation-based self-supervised RL algorithms just by adding a simple non-contrastive loss with least required components. Our proposed approach RAIR: Reinforcement learning with Augmentation Invariant Representation enables efficient representation learning and legitimate optimization sequence of objectives. We evaluate RAIR on all environments from the RL generalization benchmark Procgen. The experimental results indicate that RAIR outperforms PPO and achieves competitive or better performance with other data augmentation-based approaches.