Pretraining the Vision Transformer using self-supervised methods for vision based Deep Reinforcement Learning

TL;DR

This paper investigates pretraining Vision Transformers with self-supervised methods for reinforcement learning, emphasizing the importance of temporal relations and demonstrating improved data efficiency and richer representations in Atari environments.

Contribution

It introduces a temporal order verification task to enhance self-supervised pretraining of Vision Transformers for RL, leading to better representations and performance.

Findings

Self-supervised pretraining improves RL data efficiency.

Temporal order verification enhances representation quality.

Pretrained encoder yields richer, more focused attention maps.

Abstract

The Vision Transformer architecture has shown to be competitive in the computer vision (CV) space where it has dethroned convolution-based networks in several benchmarks. Nevertheless, convolutional neural networks (CNN) remain the preferential architecture for the representation module in reinforcement learning. In this work, we study pretraining a Vision Transformer using several state-of-the-art self-supervised methods and assess the quality of the learned representations. To show the importance of the temporal dimension in this context we propose an extension of VICReg to better capture temporal relations between observations by adding a temporal order verification task. Our results show that all methods are effective in learning useful representations and avoiding representational collapse for observations from Atari Learning Environment (ALE) which leads to improvements in data efficiency when we evaluated in reinforcement learning (RL). Moreover, the encoder pretrained with the temporal order verification task shows the best results across all experiments, with richer representations, more focused attention maps and sparser representation vectors throughout the layers of the encoder, which shows the importance of exploring such similarity dimension. With this work, we hope to provide some insights into the representations learned by ViT during a self-supervised pretraining with observations from RL environments and which properties arise in the representations that lead to the best-performing agents. The source code will be available at: https://github.com/mgoulao/TOV-VICReg