Data-Efficient Reinforcement Learning with Self-Predictive Representations

Abstract

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

arXiv commentThe first two authors contributed equally to this work. v4 includes new ablations and reformatting for ICLR camera ready

References in corpus (30)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• A Simple Framework for Contrastive Learning of Visual Representations
• Representation Learning with Contrastive Predictive Coding
• Dueling Network Architectures for Deep Reinforcement Learning
• Learning deep representations by mutual information estimation and maximization
• Unsupervised Data Augmentation for Consistency Training
• Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results
• Temporal Ensembling for Semi-Supervised Learning
• Dota 2 with Large Scale Deep Reinforcement Learning
• Data-Efficient Image Recognition with Contrastive Predictive Coding
• Noisy Networks for Exploration
• Contrastive Multiview Coding
• DeepMind Control Suite
• Big Self-Supervised Models are Strong Semi-Supervised Learners
• Rainbow: Combining Improvements in Deep Reinforcement Learning
• Model-Based Reinforcement Learning for Atari
• Learning with Pseudo-Ensembles
• Deep Reinforcement Learning that Matters
• Learning Latent Dynamics for Planning from Pixels
• State Representation Learning for Control: An Overview
• Challenges of Real-World Reinforcement Learning
• Agent57: Outperforming the Atari Human Benchmark
• Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels
• When to use parametric models in reinforcement learning?
• Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model
• rlpyt: A Research Code Base for Deep Reinforcement Learning in PyTorch
• Bootstrap Latent-Predictive Representations for Multitask Reinforcement Learning
• Deep Reinforcement and InfoMax Learning
• Learning Invariances for Policy Generalization

Cited by in corpus (19)

• Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning
• Decoupling Representation Learning from Reinforcement Learning
• Reinforcement Learning with Prototypical Representations
• Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation
• Learning Vision-Guided Quadrupedal Locomotion End-to-End with Cross-Modal Transformers
• Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning
• BYOL works even without batch statistics
• Generalization in Reinforcement Learning by Soft Data Augmentation
• Mine Your Own vieW: Self-Supervised Learning Through Across-Sample Prediction
• S4RL: Surprisingly Simple Self-Supervision for Offline Reinforcement Learning
• HyAR: Addressing Discrete-Continuous Action Reinforcement Learning via Hybrid Action Representation
• Improving Few-Shot Learning with Auxiliary Self-Supervised Pretext Tasks
• Towards Deeper Deep Reinforcement Learning with Spectral Normalization
• Expert Q-learning: Deep Reinforcement Learning with Coarse State Values from Offline Expert Examples
• Reinforcement Learning with Latent Flow
• Measuring Progress in Deep Reinforcement Learning Sample Efficiency
• Learning Representations for Pixel-based Control: What Matters and Why?
• Learning State Representations from Random Deep Action-conditional Predictions
• Is High Variance Unavoidable in RL? A Case Study in Continuous Control