Pretraining Representations for Data-Efficient Reinforcement Learning

Max Schwarzer; Nitarshan Rajkumar; Michael Noukhovitch; Ankesh Anand; Laurent Charlin; Devon Hjelm; Philip Bachman; Aaron Courville

Source-linked AI summary

Pretraining Representations for Data-Efficient Reinforcement Learning

Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, Devon Hjelm, Philip Bachman, Aaron Courville

arXiv:2106.04799v1Jun 9, 2021cs.LG

TL;DR

Deep RL needs more data-efficient learning because tabula-rasa training can overfit and existing unsupervised exploration may require extensive interaction. SGI pretrains an encoder offline with latent dynamics, goal-conditioned RL, and inverse dynamics objectives before task-specific fine-tuning. On Atari 100k, it outperforms comparable representation-pretraining methods and approaches human-level sample efficiency, with stronger results from larger models and higher-quality data.

Problem
Deep RL often learns from scratch, leading to overfitting and poor sample efficiency, while existing exploration-based pretraining can require impractical amounts of interaction.
Method
SGI uses offline, reward-free pretraining with latent dynamics modeling, unsupervised goal-conditioned RL, and inverse dynamics modeling before downstream fine-tuning.
Results
On Atari 100k, SGI-M/L achieves a median HNS of 0.753, approaching human-level sample efficiency and outperforming all comparable methods except CPT.
Takeaways & Limitations
SGI provides a data-efficient offline alternative to exploration-based pretraining, with performance scaling robustly to higher-quality data and larger models.
Takeaways & Limitations
Naive fine-tuning can disrupt SGI representations, so careful reduction of fine-tuning learning rates is critical for optimal performance.

Abstract

from arXiv · show

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.

1 Introduction

Deep RL often learns from scratch, causing overfitting and poor sample efficiency, while existing unsupervised exploration requires impractical amounts of interaction. SGI addresses this with offline self-supervised representation pretraining that combines complementary RL-aligned objectives and improves Atari data efficiency.

Motivation and contribution: Tabula-rasa RL can suffer severe overfitting and poor sample efficiency because it lacks useful visual and dynamics priors.The introduction contrasts standard random-initialization training with humans’ use of prior knowledge about the world.
Motivation and contribution: Existing unsupervised exploration methods may require months or years of real-time experience, motivating practical and ideally offline pretraining.Real-world interaction can be limited by real-time execution, safety oversight, and parallelization costs.
Motivation and contribution: SGI combines offline self-supervised objectives to learn representations before downstream task-specific fine-tuning, using far less pretraining data than existing methods.The approach is intended to approach human-level data efficiency when fine-tuned on downstream tasks.
Method and findings: SGI pretrains with latent dynamics modeling, unsupervised goal-conditioned RL, and inverse dynamics modeling to capture diverse dynamical and temporal aspects of the environment.The objectives are combined because a collection of objectives may capture more information than any single objective.
Method and findings: SGI’s combined objectives outperform each objective in isolation and improve over representation-pretraining baselines such as ATC on Atari.Its results are competitive with exploration-based methods that require two to three orders of magnitude more pretraining data and online interaction.
Method and findings: SGI performance scales with data quality and quantity, and larger models can significantly exceed smaller networks after SGI pretraining.The reported scaling is strongest when data comes from better-performing or more exploratory policies.

2 Representation Learning Objectives

SGI builds a shared latent representation by combining forward dynamics, inverse dynamics, and goal-conditioned control objectives. These reward-free objectives can use arbitrary offline behavioral data while addressing complementary questions about the environment and agent interaction.

Objective design: The objectives are reward-free and environment-focused, enabling offline pretraining from arbitrary behavioral policies rather than task-specific rewards.Their questions concern environment structure rather than any particular policy.
Objective design: SGI combines forward dynamics, inverse dynamics, and goal-conditioned RL because the representation should answer complementary questions about transitions and control.The objectives respectively predict future states, infer actions from state transitions, and select actions to reach target states.
Objective design: All objectives operate in a shared latent space produced by one encoder, facilitating their joint use during pretraining and downstream learning.Figure 1 provides the corresponding shared-torso overview.
Self-Predictive Representations: SPR predicts future-state representations recursively in latent space from an encoded state and action sequence, comparing projected predictions with target representations.A target encoder and projection provide target representations, while a learned predictor maps predicted projections toward them.
Self-Predictive Representations: The SPR target-module parameters are updated as an exponential moving average of the online encoder and projection parameters.This target construction is part of the latent prediction objective.
Goal-Conditioned Reinforcement Learning: Goal-conditioned RL uses normalized encoder-sized goal vectors to create synthetic rewards and learn Q(s_t, a, g) for reaching goals under optimal future actions.Goals are sampled from other or future trajectory states and perturbed with noise to encourage semantic diversity.
Inverse Dynamics Modeling: Inverse dynamics predicts the action taken from consecutive states and is integrated into SPR rollouts to complement and stabilize the other objectives.In discrete control it is classification, while in continuous control it is regression and is trained with cross-entropy in the described rollout formulation.

3 Related Work

Prior RL pretraining methods use data efficiency, exploration, or visual self-supervision to learn representations, but many exploration approaches require online interaction. SGI instead supports offline, collection-policy-agnostic pretraining and reports stronger controlled data-efficiency results than related representation methods.

Data-Efficiency: RL data-efficiency research includes model-based and model-free methods, with later work showing that modified hyperparameters or image augmentation can substantially improve model-free performance.The Atari 100k benchmark limits agents to 100,000 environment-interaction steps.
Exploratory pretraining: Exploration-based pretraining commonly uses unsupervised environment interaction to maximize state diversity or learn skills before task adaptation.Examples include APT’s state-diversity objective and DIAYN or VISR-style skill learning.
Exploratory pretraining: SGI differs from these methods by supporting offline pretraining that is agnostic to how data was collected.If offline data is unavailable, exploration methods can instead generate a dataset for SGI.
Visual Representation Learning: RL representation learning has used latent-space dynamics, contrastive learning, predictive objectives, and inverse dynamics to exploit interaction structure.Some latent prediction approaches were reported as vulnerable to collapse or reliant on auxiliary reconstruction losses.
Visual Representation Learning: Controlled comparisons report that SGI’s representations are better for data efficiency than related reward-free temporal-contrastive methods.The related methods include ST-DIM-like approaches and SPR-based pretraining comparisons discussed in the paper’s context.

4 Experimental Details

Experiments evaluate SGI on Atari 100k under limited interaction and varied offline pretraining data, architectures, and training configurations. Performance is summarized using human-normalized scores across 26 games.

4.1 Evaluation: Agents receive only 100k environment-interaction steps on Atari, approximately matching the two hours of human experience used as a sample-efficiency reference.The benchmark contains 26 games and disables sticky actions.
4.1 Evaluation: Performance is measured with human-normalized score, averaged over 100 evaluation trajectories and 10 training seeds, reporting mean, median, IQM, and threshold counts across games.The metric normalizes agent scores between random-policy and human scores.
4.2 Pretraining Data: Performance increases with pretraining data quality and quantity, especially when data comes from better-performing or more exploratory policies.The datasets span random, exploratory, weak, and mixed-quality experience.
4.3 Training Details: Three representation-learning objectives are jointly optimized during pretraining, while finetuning retains reinforcement-learning and forward-dynamics losses with pretrained learning rates reduced by two orders of magnitude.Experiments compare multiple encoder architectures and dataset types, including SGI-R, SGI-E, SGI-W, and SGI-M.
4.3 Training Details: SGI scales robustly with model size, allowing larger pretrained models to significantly exceed the finetuning performance of smaller networks.The study varies standard convolutional, standard ResNet, and enlarged ResNet encoders.

5 Results and Discussion

SGI achieves strong Atari-100K data efficiency by combining multiple representation-learning objectives and benefits from higher-quality data, more pretraining, and larger models. Its best configuration approaches human-level sample efficiency, while ablations show that objective combination and careful finetuning are important.

5. Results and Discussion: 0.753 median HNS: SGI-M/L approaches human-level sample-efficiency and outperforms all comparable methods except CPT on Atari-100K.SGI-M achieves 0.679 median HNS with less data and a smaller model, significantly outperforming ATC-M; SGI-E reaches 0.456, matching or exceeding APT, VISR, and ATC-E.
5. Results and Discussion: SGI achieves strong performance with 6M pretraining transitions, compared with 250M for APT and 4B for CPT.The 6M transitions correspond to roughly 4.5 days of experience, versus approximately 193 days for APT and 8.45 years for CPT.
5.1 Data quality matters: Finetuning performance improves with offline-data quality: SGI-R helps on 16 of 26 games, SGI-E raises median HNS by 33%, and SGI-M reaches a 98% improvement over SGI-None.Task reward engagement correlates positively with finetuning performance, with diminishing returns at higher engagement.
5.2 Pretraining unlocks the value of larger networks: Larger networks benefit disproportionately from SGI pretraining, whereas the standard three-layer encoder provides only minor gains over training from scratch.The results suggest increasing model size alongside the amount of pretraining data.
5.3 Combining SGI’s objectives improves performance: Performance monotonically increases as more SSL objectives are combined, supporting multiple objectives for broader representational coverage of the MDP.Inverse dynamics modeling combined with either other objective performs respectably well; configurations including it avoid the representational collapse observed for S, G, and S+G.
5.4 Naively finetuning ruins pretrained representations: Reduced finetuning learning rates for pretrained parameters outperform both unrestricted finetuning and freezing, consistent with early gradient disruption of pretrained representations.The authors relate this disruption to a phenomenon analogous to catastrophic forgetting.

6 Conclusion

SGI combines self-supervised objectives to improve data-efficient reinforcement learning, with reported gains on Atari 100k and benefits from larger models, better pretraining data, and careful fine-tuning.

SGI combines latent dynamics modeling, unsupervised goal-conditioned reinforcement learning, and inverse dynamics modeling to capture multiple aspects of environment dynamics.
SGI improves Atari 100k data efficiency, particularly compared with unsupervised exploration approaches requiring orders of magnitude more pretraining data.
All three self-supervised objectives contribute to SGI’s success, while performance scales with higher-quality pretraining data and larger models.
Careful reduction of fine-tuning learning rates is critical for optimal SGI performance.
Deep Q-learning approximates action values with a parameterized Qθ, and Rainbow combines several DQN improvements for control.

B Uncertainty-aware comparisons

The paper uses IQM scores and bootstrap uncertainty to make aggregate comparisons more robust, finding that nearly all reported comparisons are statistically significant except one fine-tuning comparison.

The analysis uses raw per-run data for the authors’ experiments because prior work did not report the data needed for these comparisons.
IQM human-normalized scores and bootstrap confidence intervals summarize variants and controls across the reported comparisons.
All Figure 5 comparisons are statistically significant at p < 0.05 except SGI-M versus SGI-M with SGI fine-tuning, where p ≈0.4.

C Implementation Details

The implementation materials document the codebase, comparison and ablation layouts, uncertainty visualization, and aggregate evaluation over Atari 100k games.

The implementation is based on released SPR code and uses rlpyt, NumPy, and PyTorch.
Figure 5 includes comparisons with behavioral cloning and ATC, dataset ablations, and ablations over SGI’s self-supervised objectives.
The reported uncertainty analysis uses bootstrapping distributions for IQM measurements.
Table 6 reports IQM, median, and mean human-normalized scores for SGI variants and controls across 26 Atari 100k games.

C.1 Training

SGI training combines several self-supervised components with synthetic goal generation and FiLM-based goal conditioning, while acknowledging a potential representation-collapse failure mode.

Pretraining uses batch size 256, exponential-moving-average target networks with τ = 0.99, and a budget equivalent to 10 epochs over 6M samples.
Goals are sampled from near-future target representations, mixed with Gaussian noise, and sometimes exchanged across minibatch states.
The synthetic goal-conditioned reward uses target-encoder representations and potential-based reward shaping.
The reward changes as the target encoder changes, although the slower-moving target encoder may provide stability.
If all target representations collapse to one constant vector, every synthetic reward becomes zero and the goal-conditioned task is trivially solvable.
The goal-conditioned Q-function uses FiLM to generate per-channel scales and biases that modulate DQN features.

C.3 Model Architectures

SGI uses a residual-style encoder, shared augmentations, and a pretraining procedure combining latent dynamics, inverse modeling, and goal-conditioned RL losses. The procedure samples unlabeled sequences, augments observations, constructs noisy hindsight goals, and optimizes the corresponding objectives.

C.3 Model Architectures: The larger encoder uses inverted residual blocks, batch normalization, and three residual-block groups with 32, 64, and 64 channels.This architecture builds on designs from Espeholt et al., EfficientNets, and MobileNetv2.
C.3 Model Architectures: All experiments use random crops with 4-pixel padding and 84x84 crops together with image-intensity jittering during pretraining and fine-tuning.The augmentations follow those used in SPR and DrQ.
C.3 Model Architectures: For comparison, ATC is adapted to SGI’s encoder, optimizer, augmentation, pretraining data, epochs, and fine-tuning procedure while retaining its minibatch structure.The retained structure samples 32 subsequences of eight steps, producing a batch size of 512 and defining InfoNCE negatives.
C.3 Model Architectures: SGI pretraining combines SGI, SPR, inverse-modeling, and goal-conditioned RL losses over sampled unlabeled sequences.The pseudocode separates these objective computations and accumulates their contributions into the training loss.
C.3 Model Architectures: The training pipeline augments observations, samples hindsight goals, injects noise, and sometimes permutes goals to increase their difficulty.Goals are encoded from hindsight states, mixed with noise, and permuted with probability 0.2.

E Full Results on Atari100k

The paper reports SGI’s full Atari100k results across all 26 games and compares them with reproduced SPR scores and an ATC implementation. The reported evaluation averages returns over 10 seeds after 100k interaction steps, while one cited baseline contains an impossible Pong score.

E Full Results on Atari100k: The authors omit APT and VISR per-game scores because they believe the currently available paper version may contain errors.They specifically flag a claimed VISR Pong score below −21 as impossible under standard settings.
E Full Results on Atari100k: Table 7 reports mean episodic returns for SGI across all 26 Atari100k games after 100k steps, averaged over 10 random seeds.The table compares SGI with reproduced SPR results and ATC scores from the authors’ implementation.
E Full Results on Atari100k: Table 8 evaluates modified SGI fine-tuning variants on the same 26 Atari100k games after 100k steps using averages over 10 seeds.The caption identifies these as versions discussed in Section 5 and includes reproduced SPR scores.
E Full Results on Atari100k: Table 9 evaluates combinations of SGI’s pretraining objectives on the 26 Atari100k games after 100k steps, with scores averaged over 10 random seeds.This isolates objective-combination variants rather than reporting only the final SGI configuration.

F Transferring Representations between Games

The paper tests whether representations pretrained on semantically and visually similar Atari games transfer to a held-out game. Across the clique experiments, extra cross-game frames are mostly unhelpful and often reduce fine-tuning performance.

F Transferring Representations between Games: The transfer study groups games into five cliques based on similar semantics and visual elements, using 750k frames per game before fine-tuning on one game.A four-game clique therefore contributes 3M pretraining frames.
F Transferring Representations between Games: Pretraining on extra frames from semantically similar game cliques is mostly unhelpful for fine-tuning performance.Only Kangaroo shows a modest improvement; several games show no difference, while most decline.
F Transferring Representations between Games: The authors suggest Atari may be less suitable for transferable representations, consistent with earlier negative results using Atari for this purpose.This interpretation is offered as an explanation for the predominantly negative cross-game transfer results.