[论文解读] SMART: Self-supervised Multi-task pretrAining with contRol Transformers

Yanchao Sun, Shuang Ma+4 more|arXiv (Cornell University)|Jan 24, 2023

Domain Adaptation and Few-Shot Learning被引用 4

一句话总结

SMART 提出了一种基于控制变压器（Control Transformer, CT）的自监督多任务预训练框架，用于序列决策任务。该框架从离线轨迹中学习与奖励无关的表征。通过结合前向动力学、逆向动力学和掩码事后控制目标，SMART 在多种模仿学习（IL）与强化学习（RL）任务中实现了最先进（SOTA）的迁移性能，即使在低质量或随机的预训练数据下也表现优异。

ABSTRACT

Self-supervised pretraining has been extensively studied in language and vision domains, where a unified model can be easily adapted to various downstream tasks by pretraining representations without explicit labels. When it comes to sequential decision-making tasks, however, it is difficult to properly design such a pretraining approach that can cope with both high-dimensional perceptual information and the complexity of sequential control over long interaction horizons. The challenge becomes combinatorially more complex if we want to pretrain representations amenable to a large variety of tasks. To tackle this problem, in this work, we formulate a general pretraining-finetuning pipeline for sequential decision making, under which we propose a generic pretraining framework extit{Self-supervised Multi-task pretrAining with contRol Transformer (SMART)}. By systematically investigating pretraining regimes, we carefully design a Control Transformer (CT) coupled with a novel control-centric pretraining objective in a self-supervised manner. SMART encourages the representation to capture the common essential information relevant to short-term control and long-term control, which is transferrable across tasks. We show by extensive experiments in DeepMind Control Suite that SMART significantly improves the learning efficiency among seen and unseen downstream tasks and domains under different learning scenarios including Imitation Learning (IL) and Reinforcement Learning (RL). Benefiting from the proposed control-centric objective, SMART is resilient to distribution shift between pretraining and finetuning, and even works well with low-quality pretraining datasets that are randomly collected.

研究动机与目标

• 解决为具有不同动力学、奖励和动作空间的多样化序列决策任务预训练通用可迁移表征的挑战。
• 克服预训练与微调阶段之间的数据分布偏移问题，尤其在预训练数据由次优或随机策略收集时。
• 开发一种统一的表征学习框架，支持下游模仿学习（IL）与强化学习（RL）任务，而无需针对特定任务进行调整。
• 设计一种以控制为中心的预训练目标，从高维观测中捕捉短期与长期动态，且不依赖奖励信号。
• 确保对未见任务与领域的鲁棒性与泛化能力，即使预训练数据质量较低或非专家数据亦可。

提出的方法

• 引入控制变压器（CT），一种基于因果注意力机制的架构，用于从高维观测中建模状态-动作交互。
• 设计一种新颖的以控制为中心的预训练目标，包含三个组成部分：前向动力学预测、逆向动力学预测和掩码事后控制（Mask-Ctl）。
• 利用掩码事后控制预测未来状态下的动作，以鼓励长期时间建模和与策略无关的动力学理解。
• 在无奖励信号或专家演示的情况下，对 CT 进行自监督方式的离线轨迹数据训练。
• 通过聚焦于转换动态而非奖励建模，实现表征学习与奖励建模的解耦，从而兼容 IL 与 RL 的下游任务。
• 在部分变体中使用动量编码器进行对比监督，但主实验结果优先采用核心的以控制为中心的目标。

实验结果

研究问题

• RQ1 自监督预训练框架是否能有效学习多样化序列决策任务的可迁移表征，且无需奖励监督？
• RQ2 所提出的以控制为中心的预训练目标与基于视觉或语言启发的目标相比，在下游迁移性能上表现如何？
• RQ3 当预训练数据由随机或非专家策略收集时，SMART 在未见任务与领域上的泛化能力有多强？
• RQ4 短期（前向/逆向动力学）与长期（掩码事后控制）动力学建模对整体性能的相对贡献如何？
• RQ5 添加辅助损失（如掩码状态预测或对比损失）是否能提升泛化能力？在何种条件下有效？

主要发现

• SMART 显著提升了在 DeepMind Control Suite 中模仿学习（BC）与强化学习（RL）下游任务的学习效率，优于从零开始训练和先前的 SOTA 预训练方法。
• 消融实验表明，预训练目标的三个组成部分——前向动力学、逆向动力学和掩码事后控制——均不可或缺，任一组件的移除均导致性能下降。
• SMART 展现出强大的分布偏移鲁棒性：即使在随机策略收集的预训练数据下，仍能有效泛化，优于基线方法。
• 当预训练数据具有探索性时，不添加额外损失的原始 SMART 目标实现最佳整体性能，表明核心以控制为中心的设计已足够且稳定。
• 在随机预训练数据下，添加掩码状态预测或对比损失等辅助损失可提升性能，表明这些损失在分布偏移较高时有助于学习更优的视觉表征。
• 该框架在 128 个实验设置（4 种变体 × 2 种学习场景 × 2 种数据集选择 × 8 个下游任务）中保持一致的性能，证实其多功能性与可靠性。