[Paper Review] Evaluating Vision Transformer Methods for Deep Reinforcement Learning from Pixels

Tianxin Tao, Daniele Reda+1 more|arXiv (Cornell University)|Apr 11, 2022

Advanced Memory and Neural Computing4 citations

TL;DR

This paper evaluates Vision Transformers (ViT) with self-supervised pre-training methods—masking-based, contrastive, and prediction-based objectives—for deep reinforcement learning from pixels. Despite ViT's success in computer vision, the study finds that while auxiliary tasks improve ViT performance, the leading CNN-based RAD method still outperforms all ViT variants, with masking-based pre-training (e.g., MAE) significantly outperforming contrastive learning.

ABSTRACT

Vision Transformers (ViT) have recently demonstrated the significant potential of transformer architectures for computer vision. To what extent can image-based deep reinforcement learning also benefit from ViT architectures, as compared to standard convolutional neural network (CNN) architectures? To answer this question, we evaluate ViT training methods for image-based reinforcement learning (RL) control tasks and compare these results to a leading convolutional-network architecture method, RAD. For training the ViT encoder, we consider several recently-proposed self-supervised losses that are treated as auxiliary tasks, as well as a baseline with no additional loss terms. We find that the CNN architectures trained using RAD still generally provide superior performance. For the ViT methods, all three types of auxiliary tasks that we consider provide a benefit over plain ViT training. Furthermore, ViT reconstruction-based tasks are found to significantly outperform ViT contrastive-learning.

Motivation & Objective

• To investigate whether Vision Transformers (ViT) can effectively learn visual representations for deep reinforcement learning from raw pixels.
• To evaluate the impact of three self-supervised pre-training objectives—masking, contrastive learning, and future prediction—on ViT-based RL policies.
• To compare the performance of ViT-based methods against a strong CNN-based baseline, RAD, in standard control environments.
• To determine whether data augmentation and auxiliary tasks improve sample efficiency and final performance in ViT-based RL.
• To assess the transferability of self-supervised ViT pre-training to downstream RL control tasks with high-dimensional visual inputs.

Proposed method

• Adapts three self-supervised ViT pre-training methods: Masked Autoencoder (MAE), Data2Vec, and momentum contrastive learning (MoCo) as auxiliary tasks.
• Integrates the ViT encoder with a Soft Actor-Critic (SAC) agent, jointly training the policy and auxiliary objectives via shared representation learning.
• Applies random cropping as the primary data augmentation strategy, resizing 100×100 observations to 84×84 for training.
• Uses a lightweight ViT architecture with 4 layers, 8 attention heads, 12×12 patch size, and 128-dimensional latent space.
• Employs a two-branch decoder head for contrastive learning and a learnable mask vector for MAE, with separate optimization for contrastive objectives.
• Implements gradient blocking from actor to encoder (as in SAC+AE) to stabilize representation learning while allowing critic updates to the shared encoder.

Experimental results

Research questions

• RQ1 Can Vision Transformers (ViT) with self-supervised pre-training effectively improve sample efficiency and final performance in image-based deep reinforcement learning?
• RQ2 How do different self-supervised pre-training objectives—masking, contrastive learning, and future prediction—compare in their ability to guide ViT policy learning?
• RQ3 Does the performance of ViT-based RL agents surpass or fall short of state-of-the-art CNN-based methods like RAD in standard control benchmarks?
• RQ4 What is the impact of data augmentation (specifically random cropping) on ViT training stability and performance in RL settings?
• RQ5 Can auxiliary tasks mitigate the representation learning bottleneck in vision-based RL when using ViT architectures?

Key findings

• The CNN-based RAD method consistently outperforms all ViT-based approaches across all evaluated environments, indicating that ViT does not yet surpass established CNNs in this setting.
• All three auxiliary pre-training methods (MAE, Data2Vec, MoCo) improve ViT performance over plain pre-training without auxiliary objectives.
• Masking-based pre-training (MAE) significantly outperforms contrastive learning (MoCo) in ViT-based RL, suggesting that reconstruction tasks are more effective for policy learning.
• Data2Vec-based pre-training shows moderate improvement but underperforms both MAE and the RAD baseline.
• The choice of masking ratio (40% and 75%) affects performance, with 75% masking yielding better results in MAE, indicating that higher in-painting complexity improves representation quality.
• Joint training with auxiliary objectives and gradient blocking from actor to encoder stabilizes learning and improves final returns, confirming the value of architectural and training design choices.