Source-linked AI summary
Loss is its own Reward: Self-Supervision for Reinforcement Learning
Evan Shelhamer, Parsa Mahmoudieh, Max Argus, Trevor Darrell
TL;DR
Reward in RL can be delayed and sparse, limiting supervision for representation learning and end-to-end optimization. The paper adds self-supervised losses over transitions, using pre-training and joint optimization to improve policy returns and data efficiency. The reported results include a 1.4× average speed-up to 95% of the best Atari return, with joint optimization improving further.
Problem
Delayed and sparse reward provides a difficult, impoverished signal for representation learning during end-to-end RL.
Method
The paper augments RL with self-supervised losses for state, dynamics, inverse dynamics, and reward, using representation transfer through pre-training and joint optimization.
Results
Self-supervised pre-training reaches 95% of the best Atari return with a 1.4× average speed-up, and joint optimization improves further.
Takeaways & Limitations
Self-supervision supplies instantaneous environmental supervision from transitions, including without extrinsic reward, and can scaffold policy learning.
Takeaways & Limitations
Observation-only self-supervised losses had no effect in preliminary experiments and were not pursued because temporal dependence may require supervision across transitions.
Abstract
from arXiv · showhide
Reinforcement learning optimizes policies for expected cumulative reward. Need the supervision be so narrow? Reward is delayed and sparse for many tasks, making it a difficult and impoverished signal for end-to-end optimization. To augment reward, we consider a range of self-supervised tasks that incorporate states, actions, and successors to provide auxiliary losses. These losses offer ubiquitous and instantaneous supervision for representation learning even in the absence of reward. While current results show that learning from reward alone is feasible, pure reinforcement learning methods are constrained by computational and data efficiency issues that can be remedied by auxiliary losses. Self-supervised pre-training and joint optimization improve the data efficiency and policy returns of end-to-end reinforcement learning.
1. Introduction
End-to-end RL must learn both representations and policies, but reward-focused optimization can make representation learning a bottleneck. The paper augments reward with self-supervised losses to provide broader supervision and improve RL efficiency.
- Motivation: Representation learning is a bottleneck because end-to-end RL delegates it to backpropagation while focusing on policy optimization.The paper separates representation learning from reinforcement learning and argues that reward-bound supervision is insufficiently broad.
- Motivation: Destroying policy and value outputs leaves most representation parameters intact, allowing the policy to recover faster than during initial training.This recovery gap illustrates that learned representations can support rapid relearning even when policy distribution and value estimates are lost.
- Self-Supervision: Self-supervised losses use states, actions, and successors to generate instantaneous, ubiquitous gradients from transitions, including those without reward.Reward captures the task, whereas self-supervision captures environmental signals and supplies surrogate annotations from unlabeled inputs.
- Self-Supervision: The paper studies discriminative auxiliary losses for state, dynamics, inverse dynamics, and reward, transferring representations through pre-training and joint optimization.These losses broaden the supervision available to RL beyond delayed and sparse reward.
- Results: A 1.4× speed-up on average to 95% of the best return is achieved with Atari pre-training, while joint optimization improves further.The reported gains include higher returns and better data efficiency.
2. Preliminaries
RL optimizes policies through sampled experience, commonly using actor-critic policy gradients. The paper extends this framework with auxiliary self-supervised gradients and shared-encoder architectures.
- Reinforcement Learning: RL optimizes policies on Markov decision processes defined by states, actions, transitions, rewards, and discounting.The transition distribution specifies successor states, while the reward function and discount define the task objective.
- Reinforcement Learning: Policy-gradient methods estimate expected-return gradients from trajectories generated by executing the policy in the environment.Actor-critic methods can scale the policy gradient using an estimated advantage rather than the return itself.
- Self-Supervision: This work augments the policy gradient with auxiliary gradients from self-supervised tasks.The approach adds learning signals beyond the reward-based policy objective.
- Self-Supervision: Self-supervised learning creates surrogate losses from unannotated data, unlike historically generative or reconstructive unsupervised objectives.In RL, these losses can mine additional supervision from data already collected by the policy.
- Architectures: The self-supervised architectures share the actor-critic encoder and add task-specific decoders and losses.Shared encoders support transferability between self-supervised tasks and RL.
3. Self-Supervision of Policies
The paper defines self-supervised policy losses over transition components to enrich representations and improve reinforcement learning beyond delayed reward. It explores auxiliary tasks, transfer settings, and intrinsic-reward extensions.
- Auxiliary tasks: Self-supervision mines surrogate annotations from transition components such as states, actions, rewards, and successors.These losses use the same transition data available to reinforcement learning methods.
- Representation transfer: The approach attaches task-specific decoders to a common agent-compatible encoder, then transfers the shared representation into reinforcement learning.Pre-training discards the decoder after optimizing auxiliary losses.
- Reward: Reward self-supervision predicts instantaneous reward bins, reducing noise from policy stochasticity and bootstrapping error.The task bins r_t into positive, zero, and negative classes and corresponds to one-step or zero-discount value estimation.
- Dynamics and inverse dynamics: Dynamics self-supervision verifies state-successor relationships, while inverse dynamics infers intervening actions from observation histories.Both tasks can be defined from single transitions, and temporal difficulty can vary with time span and stride.
- Observation cues: Observation-only auxiliary losses had no effect in preliminary experiments and were not pursued because policy actions span temporally dependent transitions.The paper therefore emphasizes self-supervision involving transition structure.
- Intrinsic rewards: Self-supervisory losses may act as intrinsic rewards by directing policies toward novel, unlearned states and rewarding learning progress.This unifies representation tuning and policy guidance, beyond static pre-training on a fixed random-policy distribution.
4. Results
Self-supervised pre-training improves reinforcement-learning optimization and data efficiency while generally preserving or improving converged returns. Joint optimization further improves over pre-training, while probing shows that pure RL representations are narrowly tuned to the RL task.
- Policy Pre-training: Self-supervised pre-training generally reaches returns nearly equal to or better than baseline while improving optimization efficiency.The sole exception is divergent pre-training, which is described as simple to diagnose.
- Policy Pre-training: Multi-task self-supervision ranks second by data efficiency across environments, while the baseline has median rank 4.5 out of 8.The paper presents multi-task self-supervision as a practical default, although individual task ordering varies across environments.
- Policy Pre-training: 2.7× improvement in early data efficiency is obtained from multi-task self-supervision during the first 10M iterations, versus 1.3× over 60M iterations.Data efficiency is measured by the area under the score/iteration curve relative to baseline.
- Probing Policy Representations: Most proxy-task accuracies drop by more than 30% when decoded from fixed RL features, indicating that pure RL representations are narrowly tuned to the RL task.The VAE diverges in several environments.
- Joint Policy and Auxiliary Optimization: Joint optimization improves over pre-training for every task, including inverse dynamics, which fails when pre-trained but improves over baseline when trained jointly.Joint optimization sums reinforcement-learning and auxiliary gradients while optimizing auxiliary losses on the policy distribution.
5. Related Work
Related work develops state representations and auxiliary objectives for reinforcement learning, including auto-encoding, dynamics modeling, robotic priors, and task-specific supervision. This paper distinguishes its discriminative transition losses and end-to-end transfer from prior approaches.
- State representation learning: State representation learning maps histories of observations, actions, and rewards to sufficient states for efficient policy learning.The mapping is formalized as φ, producing s_t from o_1:t, a_1:t, and r_1:t.
- State representation learning: Prior representation-learning methods use auto-encoding with dynamics modeling or spatial smoothness to obtain task-relevant latent representations from images.Examples include embed to control and the deep spatial auto-encoder.
- Transfer and auxiliary objectives: Earlier deep-RL transfer approaches learned observation mappings or successor models, whereas this work transfers all layers and jointly optimizes the actor-critic network.The passage contrasts generative successor self-supervision and fixed transferred representations with the paper’s discriminative, end-to-end approach.
- Transfer and auxiliary objectives: Robotic priors define auxiliary losses for temporal coherence, repeatability, proportionality, and causality using distances conditioned on action and reward.The paper instead defines discriminative losses directly on transition tuples (s, a, r, s′).
- Concurrent auxiliary-loss methods: Concurrent work augments reinforcement learning with instantaneous reward prediction, pseudo-reward control tasks, depth regression, and path-recognition objectives.These methods target value estimation, navigation, spatial representations, and path representations.
6. Discussion
Self-supervision improves reinforcement learning with or without reward: pre-training helps especially early, while joint training further improves data efficiency by following the evolving policy distribution.
- Self-supervision with and without reward can improve reinforcement learning.
- Self-supervised pre-training alone improves optimization, especially during early iterations.The pre-training distribution is closest to the policy distribution near the beginning of optimization.
- Self-supervised representation learning is task-agnostic and provides a policy scaffold regardless of the reward.The scaffold can be developed while the policy interacts with the environment without extrinsic reward.
- Transitions without reward can still contribute to representation learning when reinforcement learning is augmented with self-supervision.