PaperlayerView exact revision on arXiv

Source-linked AI summary

Data-Efficient Reinforcement Learning with Self-Predictive Representations

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, Philip Bachman

arXiv:2007.05929v4cs.LGstat.ML

TL;DR

Limited-interaction deep RL remains difficult because high-performing agents often require enormous amounts of experience. SPR augments model-free RL with future latent-state prediction and augmentation consistency, using target representations from an EMA encoder. On Atari 100k, the full objective reaches a 0.415 median human-normalized score and exceeds human expert scores on 7 of 26 games.

  • Problem

    Deep RL's high sample complexity makes learning from limited interaction a key challenge, especially where collecting real-world data is costly.

  • Method

    SPR augments a model-free RL agent with a latent dynamics model that predicts future target representations from an EMA encoder and enforces consistency across augmented observations.

  • Results

    0.415 median human-normalized score is achieved on Atari 100k with augmentation, while SPR exceeds human expert scores on 7 of 26 games.

  • Takeaways & Limitations

    SPR improves data-efficient deep RL performance on the 100k-step Atari benchmark, with either temporal prediction or data augmentation alone still outperforming prior state of the art.

  • Takeaways & Limitations

    Median human-normalized scores can be unreliable because high human scores wash out differences and the median is influenced by only a few games.

Abstract

from arXiv · show

While deep reinforcement learning excels at solving tasks where large amounts of data can be collected through virtually unlimited interaction with the environment, learning from limited interaction remains a key challenge. We posit that an agent can learn more efficiently if we augment reward maximization with self-supervised objectives based on structure in its visual input and sequential interaction with the environment. Our method, Self-Predictive Representations(SPR), trains an agent to predict its own latent state representations multiple steps into the future. We compute target representations for future states using an encoder which is an exponential moving average of the agent's parameters and we make predictions using a learned transition model. On its own, this future prediction objective outperforms prior methods for sample-efficient deep RL from pixels. We further improve performance by adding data augmentation to the future prediction loss, which forces the agent's representations to be consistent across multiple views of an observation. Our full self-supervised objective, which combines future prediction and data augmentation, achieves a median human-normalized score of 0.415 on Atari in a setting limited to 100k steps of environment interaction, which represents a 55% relative improvement over the previous state-of-the-art. Notably, even in this limited data regime, SPR exceeds expert human scores on 7 out of 26 games. The code associated with this work is available at https://github.com/mila-iqia/spr

1 INTRODUCTION

Deep RL can achieve strong results but often requires impractically large amounts of interaction data. SPR addresses limited-data learning by combining temporally predictive, augmentation-consistent representations with RL, and is evaluated on Atari under a 100k-step budget.

  • Motivation: Deep RL agents can require 10–50 years of Atari experience or 45,000 years for OpenAI Five, making data efficiency important for costly real-world tasks.These interaction requirements are contrasted with the limited-data needs of many real-world settings.
  • Results: SPR remains ahead of prior methods on both mean and median human-normalized scores even without data augmentation.Figure 1 aggregates results across 26 Atari 100k games using the stated limited interaction budget.
  • Method: SPR trains representations to predict future latent states and remain consistent across augmented views of observations.A dynamics model predicts target representations generated by an exponential-moving-average encoder without reconstructing raw states.
  • Evaluation: 100k environment steps per Atari game produce 400k input frames, roughly corresponding to two hours of real-time experience.The benchmark gives agents and human experts a comparable interaction-time budget.
  • Results: 0.415 median human-normalized score is achieved by SPR with data augmentation on Atari 100k, outperforming prior methods by a significant margin.SPR also exceeds human expert scores on 7 of 26 games under roughly the same amount of in-game experience.

2 METHOD

The method combines Rainbow with Self-Predictive Representations, which predicts future latent states using an action-conditioned transition model and EMA target networks. Its auxiliary loss operates in latent space, optionally enforces augmentation consistency, and is optimized jointly with Q-learning.

  • Online and target networks: The online encoder represents observed states, while an EMA target encoder supplies future-state representation targets without gradient updates.The target encoder can use τ = 0, making it equal to the online encoder when augmentation already provides regularization.
  • Transition model: An action-conditioned transition model iteratively predicts K future latent representations from the current online representation and future actions.Observed future states are encoded by the target encoder, and prediction remains entirely in latent space rather than reconstructing pixels.
  • Prediction loss: Online and target projection heads map representations into a smaller latent space, while a prediction head maps online predictions toward target projections.The prediction and target representations are compared at each future timestep.
  • Optimization: The SPR loss sums cosine similarities between predicted and observed representations across future steps and is combined with the Rainbow loss.The implementation uses λ = 2 for the auxiliary-loss weighting and updates online parameters before applying the EMA target update.
  • Overall approach: SPR combines Rainbow’s Q-learning loss with an auxiliary self-supervised loss to improve sample efficiency.The combined objective updates the online encoder and task-specific components, while the target network is updated separately.
  • Augmentation and implementation: Data augmentation can be used with SPR to enforce representation consistency across views, and SPR does not require negative samples.The implementation uses random shifts and color jitter, while the transition model applies convolutional processing to spatial encoder outputs.

3 RELATED WORK

Related work improves reinforcement-learning sample efficiency through pixel-level models, latent reconstruction, augmentation, and contrastive objectives. SPR differs from these approaches through its latent predictive objective, EMA target encoder, and focus on low-data single-task performance.

  • Sample-efficient Atari RL: SiMPLe generates simulated Atari training data with a pixel-level transition model but requires several weeks of training.DER and OTRainbow instead adapt Rainbow for sample efficiency in the 100k-frame setting.
  • Data augmentation: DrQ and RAD show that modest image augmentation can substantially improve reinforcement-learning sample efficiency.These augmentation-based methods reportedly outperform prior model-based approaches in continuous-control settings.
  • Contrastive representation learning: CURL combines image augmentation with a contrastive loss, while follow-up results suggest augmentation provides most of its benefit.This motivates separating augmentation effects from contrastive objectives when evaluating representation learning methods.
  • Temporal objectives: CPC, CPC|Action, ST-DIM, and DRIML optimize temporal contrastive losses, whereas SPR is evaluated against these alternatives through ablation.Kipf et al. instead learn object-oriented contrastive representations with a graph-based transition model.
  • Predictive latent models: DeepMDP predicts future representations with an unnormalized L2 loss and reward prediction, but uses online-encoder targets and adds reconstruction to mitigate collapse.SPR uses a target encoder and remains self-supervised in latent space.
  • Predictive representation learning: PBL uses two gradient-trained target networks for future representation prediction, while SPR uses one target encoder updated without backpropagation.PBL studies asymptotic multitask generalization, whereas SPR targets single-task performance in low-data regimes.

4 RESULTS

SPR is evaluated on Atari under a 100k-step interaction budget, where it achieves state-of-the-art aggregate performance but exposes limitations in standard evaluation metrics and protocols.

  • 4 RESULTS: 100,000 environment steps provide 400,000 frames, or just under two hours, versus the typical 50,000,000-step training standard.The benchmark is designed to assess performance under a strongly data-limited setting.
  • 4 RESULTS: 0.415 median human-normalized score is a new state-of-the-art result on the Atari 100k task.The score is obtained when SPR is combined with data augmentation.
  • 4 RESULTS: 7 out of 26 games reach super-human performance, compared with a maximum of two for previous methods.SPR also scores higher than DrQ on 23 out of 26 games.
  • 4 RESULTS: Table 1 reports results averaged over 10 random seeds for SPR, 20 for CURL, and 5 for other methods.The table compares methods after 100k environment steps across the 26 Atari games.
  • 4 RESULTS: Human-normalized scores subtract random performance and divide by the human-minus-random range, then aggregate across games by mean or median.This normalization accounts for different score scales between games.
  • 4.1 EVALUATION: Human-normalized medians can be unreliable because high human scores wash out differences and the median is influenced by only a handful of games.The authors therefore also report DQN-normalized mean and median scores and plot score distributions.
  • 4.1 EVALUATION: 500,000-frame evaluations may include as few as four complete episodes, so the authors recommend evaluating 100 episodes and averaging over ten seeds.This addresses concerns about episode-return estimation and unreliable comparisons from few seeds.

5 ANALYSIS

The analysis identifies target encoders, multi-step dynamics modeling, normalized prediction losses, and projection networks as important components of SPR. Ablations show that removing or altering these components can substantially reduce performance.

  • Target encoder: A separate target encoder is vital: allowing gradients through online-encoder targets reduces the augmented median human-normalized score from 0.415 to 0.278.The target encoder is updated through an exponential moving average rather than backpropagation.
  • Dynamics modeling: Extended dynamics modeling consistently improves performance up to roughly prediction depth K = 5.Larger depths can help some games but increase computation.
  • Comparison with contrastive losses: SPR consistently outperforms temporal and non-temporal contrastive-loss variants, including CURL.This comparison is reported in the appendix experiments.
  • Using a quadratic loss causes collapse: Using an un-normalized quadratic loss performs only slightly better than random, unlike SPR’s cosine-similarity objective.The ablation tests whether the loss normalization is important for learning future latent dynamics.
  • Projections are critical: Removing projection and prediction networks produces inferior performance.The authors suggest this may limit how much of the network the SPR objective affects and may impose invariance too early in the representation hierarchy.

6 FUTURE WORK

The paper identifies future use of large unlabeled corpora, planning with SPR's transition model, and generalization to unseen tasks or environment changes as directions for future work.

  • Large corpora of unlabeled data from multiple MDPs or raw videos could be used to improve RL performance in low-data regimes.
  • It remains unclear whether SPR's learned transition model can support planning effectively in low-data regimes.
  • Future work could test whether self-supervised methods like SPR generalize to unseen tasks or changes in the environment.

7 CONCLUSION

The paper introduces SPR to improve deep RL data efficiency by learning representations predictive of future states and consistent across observation views. SPR achieves state-of-the-art Atari performance and remains stronger than prior work when either major component is disabled.

  • SPR predicts future-state representations from a target encoder while enforcing consistency across different views of environment observations.
  • SPR achieves state-of-the-art performance on the Atari 100k-step benchmark and significant improvements over prior work.
  • SPR outperforms the previous state of the art when either data augmentation or temporal prediction is disabled.

A.1 FULL RESULTS

The full-results section documents hyperparameters, controlled baselines, and Atari results after 100k environment steps. SPR exceeds prior methods on aggregate metrics and human performance on seven games, while contrastive alternatives provide little improvement over controlled Rainbow.

  • A.1 FULL RESULTS: Table 4 reports mean episodic returns after 100k environment steps, averaged over 10 random seeds.
  • A.1 FULL RESULTS: SPR outperforms prior methods on all aggregate metrics and exceeds expert human performance on 7 out of 26 games.
  • A.2 CONTROLLED BASELINES: Table 5 compares controlled Rainbow variants with and without augmentation against previous methods, including a high DQN-normalized score influenced by Private Eye.
  • A.2 CONTROLLED BASELINES: Controlled experiments use identical hyperparameters and random seeds for baselines to test whether minor baseline changes explain SPR's gains.
  • B COMPARISON WITH A CONTRASTIVE LOSS: Table 6 contains scores for contrastive alternatives to SPR, all implemented with data augmentation.

B COMPARISON WITH A CONTRASTIVE LOSS

The paper compares SPR with several contrastive objectives based on InfoNCE and related formulations. Although these variants beat CURL, none substantially improves the controlled Rainbow baseline.

  • The comparison includes a view-based contrastive loss similar to CURL.
  • The comparison includes single-step temporal contrastive prediction with augmentation, equivalent to single-step CPC.
  • A dynamics-model contrastive variant predicts up to five future steps and uses other state encodings as negative samples.
  • A soft contrastive variant combines an SPR-like augmentation-invariance objective with a uniformity objective on the unit hypersphere.
  • All contrastive variants use the same augmentation and Rainbow hyperparameters as SPR with augmentation for a fair comparison.
  • Although the variants outperform CURL, none substantially improves performance over the controlled Rainbow baseline.

C THE ROLE OF THE TARGET ENCODER IN SPR

SPR remains trainable across target-encoder settings, while augmentation changes the preferred EMA behavior. The authors hypothesize that augmentation can make EMA stabilization redundant, but leave schedules for future work.

  • Target encoder variants: Allowing gradients through online-encoder target representations causes drastic performance reductions, attributed to representational collapse.This variant performs poorly both with and without augmentation.
  • EMA coefficient: With augmentation, performance peaks at τ = 0, corresponding to a target encoder without EMA-based smoothing.The analysis evaluates nine logarithmically interpolated τ values on ten Atari games using ten seeds per game.
  • EMA coefficient: Without augmentation, SPR appears less sensitive to τ, and the authors use τ = 0.99 because of its reasonable performance and consistency with prior work.The y-axis scales differ between the augmentation and no-augmentation analyses.
  • EMA coefficient: Overall, SPR does not appear overly sensitive to τ, and no tested setting fails to train.The authors contrast this robustness with purely unsupervised methods such as BYOL.
  • Discussion: The authors hypothesize that augmentation makes EMA stabilization redundant and note that an increasing-τ schedule could potentially combine both benefits.They leave schedule tuning for future work because it would introduce an additional hyperparameter.

D WALL CLOCK TIMES

SPR has favorable wall-clock runtimes on a single Atari game, with augmentation adding overhead. Its runtime is substantially shorter than the cited SimPLe training time.

  • Runtime: A complete 100K-step SPR run with augmentation takes around 4.5 hours on a single game.The runtime is reported for training and evaluation on a P100 GPU.
  • Runtime: Without augmentation, SPR can complete a run in roughly 3 hours.The authors identify data augmentation as the source of the added overhead.
  • Comparison: SPR’s runtime compares favorably with SimPLe, which requires roughly three weeks to train on a comparable GPU.The comparison concerns complete training and evaluation runs for a single Atari game.
Loading 2007.05929v4…