[Paper Review] Pretraining Representations for Data-Efficient Reinforcement Learning

Max Schwarzer, Nitarshan Rajkumar+6 more|arXiv (Cornell University)|Jun 9, 2021

Reinforcement Learning in Robotics60 references33 citations

TL;DR

SGI (self-supervised pretraining with SPR, goal-conditioned RL, and inverse dynamics) enables data-efficient RL by learning rich representations offline from unlabeled data and fine-tuning with small task data, achieving strong Atari-100k performance especially with larger models and higher-quality pretraining data.

ABSTRACT

Data efficiency is a key challenge for deep reinforcement learning. We address this problem by using unlabeled data to pretrain an encoder which is then finetuned on a small amount of task-specific data. To encourage learning representations which capture diverse aspects of the underlying MDP, we employ a combination of latent dynamics modelling and unsupervised goal-conditioned RL. When limited to 100k steps of interaction on Atari games (equivalent to two hours of human experience), our approach significantly surpasses prior work combining offline representation pretraining with task-specific finetuning, and compares favourably with other pretraining methods that require orders of magnitude more data. Our approach shows particular promise when combined with larger models as well as more diverse, task-aligned observational data -- approaching human-level performance and data-efficiency on Atari in our best setting. We provide code associated with this work at https://github.com/mila-iqia/SGI.

Motivation & Objective

• Motivate improving data efficiency in deep RL by leveraging unlabeled/offline data to learn representations.
• Propose a multi-objective pretraining scheme that captures diverse environment aspects.
• Evaluate how offline SGI pretraining scales with data quality and model size on Atari-100k.

Proposed method

• Pretrain a shared encoder with three self-supervised objectives: latent dynamics (forward predictions in latent space), unsupervised goal-conditioned RL, and inverse dynamics modeling.
• Use a two-stage pipeline: offline pretraining with a shared encoder, then finetune with task-specific RL losses while freezing or adapting the encoder with reduced learning rates.
• Combine objectives so they operate in the latent space and share parameters, with an exponential moving average target network for stability in SPR.
• Employ Rainbow-style Q-learning for the downstream task and augmented SPR loss to encourage robust representations.
• Evaluate under Atari-100k (100k environment steps) with varied pretraining data qualities (random, exploratory, weak, mixed, and offline DQN data).

Experimental results

Research questions

• RQ1 Does offline SGI pretraining improve data efficiency on Atari-100k compared to prior pretraining baselines?
• RQ2 How do data quality and model size affect SGI’s finetuning performance?
• RQ3 Is combining multiple self-supervised objectives more beneficial than any single objective?
• RQ4 How should finetuning be performed to preserve pretrained representations (e.g., learning rate scheduling, freezing)?

Key findings

• SGI achieves strong data efficiency on Atari-100k, with SGI-M/L achieving a median HNS of 0.753, approaching human-level performance.
• SGI-M (3M data) outperforms ATC-M on the same data; SGI-E matches or exceeds other exploratory pretraining baselines with far less pretraining data.
• Performance scales with data quality and model size: larger encoders benefit more when pretrained, and higher-quality offline data yields higher finetuning performance.
• Combining all three SSL objectives yields the best performance, and inverse dynamics modeling helps avoid representational collapse.
• Naive finetuning of pretrained representations degrades performance; reducing finetuning learning rates and occasionally freezing the encoder yields superior results.