Source-linked AI summary
APS: Active Pretraining with Successor Features
Hao Liu, Pieter Abbeel
TL;DR
RL pretraining seeks data-efficient adaptation without access to environment rewards, but prior approaches separately trade off exploration and task-conditioned behavior learning. APS combines state-entropy maximization with variational successor features through a mutual-information lower bound, achieving state-of-the-art Atari performance and outperforming prior unsupervised-pretraining methods.
Problem
Prior unsupervised RL methods are limited because VISR explores inefficiently, while APT lacks task conditioning for quickly identifying downstream tasks.
Method
APS decomposes mutual information into H(s) −H(s|z), using APT for state-entropy exploration and variational successor features for task-conditioned behavior learning.
Results
APS achieves state-of-the-art Atari performance, including a 106% relative improvement on 26 games and a 3× higher human median score than DQN on Atari 57.
Takeaways & Limitations
APS demonstrates that combining state-entropy data with task-conditioned skill discovery improves unsupervised RL performance on Atari.
Takeaways & Limitations
VISR remains limited by inefficient exploration, while APT remains limited by its lack of task conditioning for downstream task identification.
Abstract
from arXiv · showhide
We introduce a new unsupervised pretraining objective for reinforcement learning. During the unsupervised reward-free pretraining phase, the agent maximizes mutual information between tasks and states induced by the policy. Our key contribution is a novel lower bound of this intractable quantity. We show that by reinterpreting and combining variational successor features~\citep{Hansen2020Fast} with nonparametric entropy maximization~\citep{liu2021behavior}, the intractable mutual information can be efficiently optimized. The proposed method Active Pretraining with Successor Feature (APS) explores the environment via nonparametric entropy maximization, and the explored data can be efficiently leveraged to learn behavior by variational successor features. APS addresses the limitations of existing mutual information maximization based and entropy maximization based unsupervised RL, and combines the best of both worlds. When evaluated on the Atari 100k data-efficiency benchmark, our approach significantly outperforms previous methods combining unsupervised pretraining with task-specific finetuning.
1. Introduction
Reinforcement learning typically trains from scratch for each reward-defined task, motivating unsupervised pretraining for data-efficient downstream adaptation. APS combines exploration through state entropy maximization with task-conditioned behavior learning through mutual-information-based successor features.
- RL algorithms commonly train from scratch to maximize extrinsic return for every encountered task.
- Unsupervised pretraining removes environment rewards during pretraining and targets data-efficient adaptation to downstream reward-defined tasks.
- Mutual-information methods learn distinguishable task-conditioned behaviors and adapt quickly through successor features, but VISR suffers from inefficient exploration.
- State-entropy methods encourage exploration, but APT is task-agnostic and lacks task conditioning for downstream adaptation.
- APS decomposes mutual information as H(s) −H(s|z), using APT for state entropy and successor features to optimize the conditional-entropy term.
- 106% relative improvement on 26 Atari games; on Atari 57, APS achieved a human median score 3× higher than DQN trained with 10M supervised environment steps.
2. Related Work
Related work explores entropy maximization and mutual-information objectives for unsupervised behavior discovery, with methods differing in exploration, task conditioning, and visual applicability.
- Entropy maximization has been used to encourage exploration in state-based and visual reinforcement learning.
- Table 1 compares unsupervised pretraining methods by exploration, off-policy learning, visual applicability, and task conditioning.
- Mutual-information objectives connect latent task variables with state visitation to discover meaningful behaviors.
3. Preliminaries
Reinforcement learning models an agent interacting with an uncertain environment to maximize cumulative reward. Successor features represent task values through feature expectations and task vectors, while APS combines exploration and task-conditioned behavior learning.
- Reinforcement learning seeks an optimal policy for an agent interacting with an uncertain environment and collecting rewards.
- An MDP specifies states, actions, transitions, initial-state distribution, reward function, and discount factor.
- The discounted return is the discounted sum of future rewards, which value-based methods estimate with a state-action value function.
- 3.1. Successor Features: Successor features assume rewards are linear in feature vectors, with task vector w specifying the desirability of each feature component.
- 3.1. Successor Features: The state-action value function decomposes into a task vector and successor features, generalizing Q-values to multidimensional feature rewards.
- APS combines state-entropy maximization for exploration with task-conditioned behavior learning through the conditional-entropy objective.
4. Method
APS combines entropy maximization for exploration with variational successor features for learning task-conditioned behavior. It uses a variational lower bound and shared representations to optimize mutual information between tasks and induced states.
- Motivation: APS combines APT-style state-entropy maximization with VISR-style variational successor features to address inefficient exploration and missing task information.APT explores through entropy maximization, while VISR provides task inference but suffers from inefficient exploration.
- Exploration: APT approximates state entropy in a lower-dimensional representation space with particle-based k-nearest-neighbor estimation.The representation encoder maps observations to lower-dimensional vectors, and entropy is estimated from distances to neighboring particles.
- Empirical setup: The gridworld experiments test whether methods can both explore multi-step environments and infer downstream tasks after unsupervised pretraining.The harder environment contains two key-passageway pairs, while Figure 4 separates zero-shot and fine-tuned testing performance.
- Objective: The method maximizes I(s; z) = H(s) − H(s|z), using state entropy to encourage exploration and conditional entropy to learn task-conditioned behaviors.The task variable z is sampled from a fixed distribution, and states are sampled from trajectories induced by the policy.
- Task-conditioned behavior: APS derives a variational lower bound and restricts task variables to successor-feature task vectors, yielding an intrinsic reward based on log q(s|w) = φ(s)^T w.The Von Mises-Fisher parameterization makes the variational approximation symmetric to its task-vector parameterization.
- Training: APS shares weights between the representation encoder f and successor-feature encoder φ, while combining the task-information reward with an exploration reward.The encoder is trained by minimizing the negative log likelihood of q(s|w), and the combined intrinsic reward is specified in the training procedure.
5. Results
APS is evaluated after reward-free pretraining followed by short reward-based adaptation on both the full Atari suite and the 26-game data-limited setting. It achieves strong performance in both settings, including super-human scores on multiple games.
- Evaluation: APS is evaluated on all 57 Atari games and the 26-game sample-efficient subset.The 26-game setting contains the easiest games in the Atari suite according to prior evaluation criteria.
- Evaluation: The evaluation uses 250M reward-free training steps followed by 100K environment steps with rewards.The 100K-step test phase equals 400K frames and is substantially shorter than the typical 500M-environment-step setting.
- Evaluation: During successor-feature evaluation, APS infers the task vector from 10 episodes or 40K steps, then fine-tunes for 60K steps.The task vector is inferred by linear regression from episodes conditioned on uniformly sampled points on a 5-dimensional sphere.
- Results: APS achieves super-human performance on eight games in the data-limited setting and exceeds previous state-of-the-art scores.Aggregate results for the 26-game subset are reported in Table 2.
- Results: APS achieves super-human performance on 15 games in the full suite, compared with a maximum of 12 for previous methods.It also achieves scores significantly higher than previous methods in this setting.
6. Analysis
The analyses examine how APS combines exploration and task-solving components, and whether fine-tuning and encoder sharing affect performance. Results indicate that both entropy maximization and variational successor features contribute, while fine-tuning and shared encoders can improve performance.
- Exploration and exploitation: Figure 5 compares methods and variants across 26 Atari games, including joint training of one method with data collected by another.The notation X →Y denotes training Y using data collected by X simultaneously.
- Aggregate results: APS outperforms prior methods on all aggregate metrics in the 26-game evaluation and exceeds expert human performance on 8 of 26 games.Scores are averaged over five random seeds for APS after 100K environment steps.
- Exploration and exploitation: Both entropy maximization and variational successor features are important to APS’s performance gain.The ablation is designed to separate contributions from exploration and fast task solving.
- Exploration and exploitation: APS →APT outperforms APT →APS on dense-reward games and significantly on exploration games, while APS →VISR also performs better in the reverse direction.These comparisons indicate that entropy maximization and variational successor features improve each other nontrivially.
- Fine-tuning: APS without fine-tuning outperforms all controlled baselines, including VISR with fine-tuning.The analysis also reports that applying the fine-tuning scheme can boost performance.
- Representation learning: Sharing the encoder can boost data efficiency because the shared representation better captures information useful for computing intrinsic reward.The paper leaves investigation of other representation-learning methods for future work.
7. Conclusion
APS addresses limitations of prior mutual-information and entropy-maximization methods by combining them in a new unsupervised RL pretraining method. The approach leverages state entropy maximization for exploration and variational successor features for task-conditioned skill discovery, achieving state-of-the-art Atari performance.
- APS is a new unsupervised pretraining method for RL that addresses limitations of mutual-information and entropy-maximization approaches.
- APS achieves state-of-the-art performance on the Atari benchmark, with significant improvements over prior work.
- APS leverages state entropy maximization data for task-conditioned skill discovery.
- APS decomposes mutual information as H(s) −H(s|z), optimizing the terms with particle-based entropy and variational successor features.
A. Experiment Details
The Atari experiments use image augmentation, generalized policy improvement, multiple random seeds, and a defined training and evaluation schedule. The ablation includes a shared-encoder variant aligned with APT's setup.
- Atari images are padded by 4 pixels on each side, randomly cropped to 84 × 84, and augmented whenever sampled from replay.
- The experiments use generalized policy improvement with 10 policies, average results over 5 random seeds, and evaluate for 125K steps after 100K training steps.
- The shared-encoder ablation reduces encoder output dimension to 5 to match default APS and uses a two-layer projection MLP with hidden size 128 and output size 64.
B. Scores Breakdown on 57 Atari games
APS is compared with baselines across all 57 Atari games using raw scores averaged over five random seeds. It achieves super-human performance on more games than any previous method and significantly higher scores overall.
- APS achieves super-human performance on 15 of the 57 Atari games, compared with a maximum of 12 for any previous method.
- APS achieves scores significantly higher than any previous methods across the individual-game comparison.
- Table 6 reports raw scores for each Atari method, averaged over five random seeds, with @N denoting fine-tuning RL interactions.