[论文解读] 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 Computing被引用 4

一句话总结

本文评估了基于自监督预训练方法（掩码式、对比式和预测式目标）的视觉变换器（ViT）在从像素进行深度强化学习中的表现。尽管ViT在计算机视觉领域表现优异，但本研究发现，尽管辅助任务能提升ViT性能，领先的基于CNN的RAD方法仍优于所有ViT变体，其中基于掩码的预训练方法（如MAE）显著优于对比学习方法。

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.

研究动机与目标

• 探究视觉变换器（ViT）是否能有效学习原始像素输入下的视觉表征以用于深度强化学习。
• 评估三种自监督预训练目标——掩码、对比学习和未来预测——对基于ViT的强化学习策略的影响。
• 在标准控制环境中，将基于ViT的方法与强大的基于CNN的基线方法RAD进行性能比较。
• 确定数据增强和辅助任务是否能提升基于ViT的强化学习中的样本效率和最终性能。
• 评估自监督ViT预训练在高维视觉输入的下游强化学习控制任务中的可迁移性。

提出的方法

• 采用三种自监督ViT预训练方法：掩码自编码器（MAE）、Data2Vec和动量对比学习（MoCo）作为辅助任务。
• 将ViT编码器与软 actor-critic（SAC）智能体结合，通过共享表征学习联合训练策略和辅助目标。
• 以随机裁剪作为主要数据增强策略，将100×100的观测图像缩放为84×84用于训练。
• 采用轻量级ViT架构，包含4层、8个注意力头、12×12的图像块大小和128维潜在空间。
• 为对比学习使用双分支解码头，为MAE使用可学习的掩码向量，并对对比学习目标进行独立优化。
• 实施从智能体到编码器的梯度阻断（如SAC+AE中所用），以在允许评论网络更新共享编码器的同时稳定表征学习。

实验结果

研究问题

• RQ1 基于自监督预训练的视觉变换器（ViT）能否有效提升图像强化学习中的样本效率和最终性能？
• RQ2 在引导ViT策略学习方面，不同自监督预训练目标（掩码、对比学习和未来预测）的性能如何比较？
• RQ3 基于ViT的强化学习智能体在标准控制基准测试中的表现是否超越或落后于最先进的基于CNN的方法（如RAD）？
• RQ4 数据增强（特别是随机裁剪）对ViT在强化学习设置中的训练稳定性和性能有何影响？
• RQ5 在使用ViT架构时，辅助任务能否缓解视觉强化学习中的表征学习瓶颈？

主要发现

• 基于CNN的RAD方法在所有评估环境中的表现均持续优于所有基于ViT的方法，表明在此设置下ViT尚未超越成熟的CNN模型。
• 所有三种辅助预训练方法（MAE、Data2Vec、MoCo）均在无辅助目标的预训练基础上提升了ViT性能。
• 基于掩码的预训练（MAE）在基于ViT的强化学习中显著优于对比学习（MoCo），表明重建任务对策略学习更有效。
• 基于Data2Vec的预训练表现出中等程度的改进，但性能仍低于MAE和RAD基线。
• 掩码比率的选择（40%和75%）影响性能，其中75%的掩码在MAE中表现更优，表明更高的图像修复复杂度可提升表征质量。
• 与辅助目标联合训练并结合从智能体到编码器的梯度阻断，可稳定学习过程并提升最终回报，证实了架构与训练设计选择的价值。