[Paper Review] 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 Learning4 citations

TL;DR

SMART proposes a self-supervised multi-task pretraining framework for sequential decision-making using a Control Transformer (CT) that learns reward-agnostic representations from offline trajectories. By combining forward dynamics, inverse dynamics, and masked hindsight control objectives, SMART achieves state-of-the-art transfer performance across diverse imitation and reinforcement learning tasks, even with low-quality or random pretraining data.

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.

Motivation & Objective

• Address the challenge of pretraining generic, transferable representations for diverse sequential decision-making tasks with varying dynamics, rewards, and action spaces.
• Overcome data distribution shift between pretraining and fine-tuning, especially when pretraining data is collected from suboptimal or random policies.
• Develop a unified representation learning framework that supports both imitation learning (IL) and reinforcement learning (RL) downstream without task-specific adaptation.
• Design a control-centric pretraining objective that captures both short-term and long-term dynamics from high-dimensional observations without relying on reward signals.
• Ensure robustness and generalization to unseen tasks and domains, even when pretraining data is low-quality or non-expert.

Proposed method

• Introduce the Control Transformer (CT), a causal attention-based architecture that models state-action interactions from high-dimensional observations.
• Design a novel control-centric pretraining objective with three components: forward dynamics prediction, inverse dynamics prediction, and masked hindsight control (Mask-Ctl).
• Use masked hindsight control to predict actions given future states, encouraging long-term temporal reasoning and policy-invariant dynamics understanding.
• Train the CT in a self-supervised manner on offline trajectory data without access to reward signals or expert demonstrations.
• Decouple representation learning from reward modeling by focusing on transition dynamics, enabling compatibility with both IL and RL downstream.
• Apply momentum encoders for contrastive supervision in some variants, but prioritize the core control-centric objective for main results.

Experimental results

Research questions

• RQ1 Can a self-supervised pretraining framework effectively learn transferable representations for diverse sequential decision-making tasks without reward supervision?
• RQ2 How does the proposed control-centric pretraining objective compare to vision- or language-inspired objectives in terms of downstream transfer performance?
• RQ3 To what extent can SMART generalize to unseen tasks and domains when pretraining data is collected from random or non-expert policies?
• RQ4 What is the relative contribution of short-term (forward/inverse dynamics) versus long-term (masked hindsight control) dynamics modeling to overall performance?
• RQ5 Does adding auxiliary losses (e.g., masked state prediction or contrastive loss) improve generalization, and under what conditions?

Key findings

• SMART significantly improves learning efficiency across both imitation learning (BC) and reinforcement learning (RL) downstream tasks in the DeepMind Control Suite, outperforming training from scratch and prior SOTA pretraining methods.
• The ablation study confirms that all three components of the pretraining objective—forward dynamics, inverse dynamics, and masked hindsight control—are essential, as removing any leads to performance degradation.
• SMART demonstrates strong robustness to distribution shift: it generalizes effectively even when pretraining data is collected via random policies, outperforming baselines in such low-quality data regimes.
• When pretraining data is exploratory, the original SMART objective without additional losses achieves the best overall performance, indicating that the core control-centric design is sufficient and stable.
• Adding auxiliary losses like masked state prediction or contrastive loss improves performance only under random pretraining data, suggesting they help learn better visual representations when distribution shift is high.
• The framework maintains consistent performance across 128 experimental settings (4 variants × 2 learning scenarios × 2 dataset selections × 8 downstream tasks), confirming its versatility and reliability.