Source-linked AI summary
Behavior From the Void: Unsupervised Active Pre-Training
Hao Liu, Pieter Abbeel
TL;DR
APT addresses the challenge that RL agents must otherwise train from scratch for each reward-defined task, while existing unsupervised methods struggle with diverse exploration and high-dimensional observations. It actively searches for novel states by maximizing nonparametric entropy in an abstract representation space during reward-free pre-training. Across Atari and DMControl, APT improves data efficiency and performance, including on tasks difficult to train from scratch, but its long pre-training is computationally intensive and fine-tuning can cause catastrophic forgetting.
Problem
RL agents typically train from scratch for each task, while existing unsupervised methods struggle to provide diverse data and explore sufficiently in high-dimensional visual environments.
Method
APT actively explores task-agnostic environments by maximizing particle-based nonparametric entropy in a learned abstract representation space during reward-free pre-training.
Results
APT improves data efficiency and performance across Atari and DMControl, outperforming DrQ from scratch and count-based exploration on every DMControl task and improving difficult sparse-reward tasks.
Takeaways & Limitations
A single task-agnostic pretrained model can tackle a broad set of visual RL tasks and achieve canonical supervised RL results using a small fraction of total samples.
Takeaways & Limitations
The long pre-training phase is computationally intensive, while fine-tuning the whole pretrained model can cause catastrophic forgetting.
Abstract
from arXiv · showhide
We introduce a new unsupervised pre-training method for reinforcement learning called APT, which stands for Active Pre-Training. APT learns behaviors and representations by actively searching for novel states in reward-free environments. The key novel idea is to explore the environment by maximizing a non-parametric entropy computed in an abstract representation space, which avoids challenging density modeling and consequently allows our approach to scale much better in environments that have high-dimensional observations (e.g., image observations). We empirically evaluate APT by exposing task-specific reward after a long unsupervised pre-training phase. In Atari games, APT achieves human-level performance on 12 games and obtains highly competitive performance compared to canonical fully supervised RL algorithms. On DMControl suite, APT beats all baselines in terms of asymptotic performance and data efficiency and dramatically improves performance on tasks that are extremely difficult to train from scratch.
1 Introduction
Unsupervised pre-training aims to let RL agents adapt efficiently to new reward-defined tasks, rather than relearning every task from scratch. APT addresses limited exploration and high-dimensional density estimation by actively seeking novel states in an abstract space.
- RL agents traditionally maximize extrinsic reward from scratch for each task, unlike creatures that reuse previously acquired behaviors.
- Unsupervised pre-training removes rewards during a long training phase, then exposes the agent to task-specific rewards for efficient downstream adaptation.
- Existing approaches struggle to provide diverse data because count-based exploration requires high-dimensional density estimation and ImageNet data is out-of-distribution for DMControl.
- APT actively explores unknown areas and maximizes nonparametric entropy over learned abstract representations to collect novel, task-agnostic data.
- 54% relative improvement in data efficiency is achieved on Atari 26 compared with DrQ, while the full Atari 57 suite reaches a 3× higher median human-normalized score than prior unsupervised RL methods and DQN.
2 Problem Setting
Unsupervised pre-training trains an agent without rewards before testing it on reward-defined tasks, seeking rapid zero-shot or few-shot adaptation. Prior mutual-information methods learn meaningful behaviors but lack explicit exploration and have limited visual-domain applicability.
- Pretrained RL uses a reward-free MDP for a long period, followed by testing with environment rewards to learn rapid task adaptation.
- The downstream objective is to maximize expected future rewards in a zero-shot or few-shot manner.
- Mutual-information methods maximize dependence between a policy-conditioning variable and state visitation, often through a variational posterior.
- Despite learning meaningful behaviors, these methods suffer from insufficient exploration because they contain no explicit exploration mechanism.
- Prior methods have only been shown to work from explicit state representations, motivating a nonparametric method for challenging visual-domain benchmarks.
3 Unsupervised Active Pre-Training for RL
APT combines contrastive representation learning with particle-based entropy maximization to explore reward-free environments through novel states. It converts nearest-neighbor distances in abstract representation space into intrinsic rewards optimized with standard RL algorithms.
- 3.1 Particle-Based Entropy Maximization: APT maps observations into an abstract representation space before maximizing entropy, avoiding direct density estimation in high-dimensional state spaces.The representation mapping is learned first, then a particle-based nonparametric method supplies the entropy objective.
- 3.1 Particle-Based Entropy Maximization: Particle-based entropy estimates distribution sparsity from distances between sampled particles and their k nearest neighbors.Each particle’s neighborhood volume reflects local sparsity, and the estimator is proportional to the average of these volumes.
- 3.1 Particle-Based Entropy Maximization: Averaging distances over all k nearest neighbors produces a more robust and stable entropy estimate than using only the k-th-neighbor distance.The resulting reward is based on the logarithm of neighborhood distances, with a constant used for numerical stability.
- 3.1 Particle-Based Entropy Maximization: APT applies the entropy objective to replay-buffer transitions by treating each next-state representation as a particle and using its entropy contribution as intrinsic reward.Intrinsic rewards are normalized by a running estimate of their mean before policy optimization.
- 3.2 Learning Contrastive Representations: Contrastive learning uses augmented views of replay-buffer images to learn representations that keep corresponding views close in latent space.APT uses SimCLR-style contrastive learning with DrQ’s random shifts and color jitter, while the encoder architecture follows DrQ with a reduced latent dimension.
- 3.3 Training APT: APT trains the encoder, computes particle-based intrinsic rewards, and updates the policy and critic through a standard actor-critic loop.The method is modular with respect to representation learning and is designed to work with existing RL algorithms.
4 Related Work
Related work studies entropy-based exploration, data-efficient RL, and unsupervised pre-training. These approaches include density estimation, model and augmentation techniques, and skill-learning interaction phases.
- State Space Entropy Maximization: State-space entropy maximization has been used for exploration through density-of-states estimation, including methods with provable efficiency under certain conditions.VAE-based entropy estimation has also been applied in lower-dimensional observation spaces.
- Data Efficient RL: Data-efficient RL combines model-based learning, algorithmic improvements, and data augmentation to improve performance when interaction data are limited.Prior work includes SimPLe, Rainbow-based combinations, and augmentation methods for vision-based RL.
- Unsupervised Pre-Training RL: Unsupervised pre-training RL adds a reward-free interaction phase in which agents learn representations or skills before learning a target task.Skill-learning methods train behaviors associated with different paths through the environment.
5 Results
APT improves reward-free pre-training and downstream adaptation across DMControl and Atari, with especially strong gains on sparse-reward tasks and challenging exploration settings. Its performance depends on learned representations, nearest-neighbor entropy, and the relationship between intrinsic and downstream rewards.
- DMControl results: APT significantly outperforms DrQ from scratch and count-based exploration on every evaluated DMControl task.It adapts quickly with few environment interactions and achieves higher return faster than prior methods.
- DMControl results: APT provides significantly higher data efficiency and asymptotic performance on sparse-reward tasks that are extremely difficult to train from scratch.The evaluation uses a long 5M-step unsupervised pre-training phase followed by downstream fine-tuning.
- Atari results: APT achieves super-human performance on 15 Atari games in the full suite, compared with a maximum of 12 for previous methods.The 26-game data-limited setting reports super-human performance on eight games, while Table 2 reports seven games under its evaluation protocol.
- Atari results: 54% increase in median score and 73% increase in mean score are reported for unsupervised pre-training on top of DrQ, with five additional human-level games.These gains surpass DQN despite DQN training on hundreds of millions of sampling steps.
- Comparisons and analysis: APT achieves comparable mean and median scores to SPR, while outperforming prior unsupervised methods and showing stronger performance in hard-exploration games than VISR.VISR achieves higher scores in dense-reward games, motivating possible combinations of the two approaches.
- Ablations: K-nearest neighbors are crucial for high scores, and a randomly initialized convolutional encoder significantly reduces performance while remaining above DrQ.APT also visits nearly 27 times more unique RAM states than the positive-reward variant, with downstream rewards of 687.1 versus 363.7.
6 Discussion
APT’s fine-tuning and pre-training introduce scale, optimization, computational, and forgetting challenges that motivate future improvements.
- Limitations: Fine-tuning works best when intrinsic and extrinsic rewards have similar scales.The authors associate scale discrepancies with suboptimal performance in dense-reward games.
- Limitations: The intrinsic reward is non-stationary and non-additive, creating an optimization challenge for reinforcement learning methods.The authors suggest that better optimization methods could further improve intrinsic-reward maximization.
- Conclusion: APT enables one task-agnostic pretrained model to address a broad set of visual reinforcement-learning tasks.The method uses an intrinsic reward derived from particle-based entropy maximization in abstract representation space.
- Future Work: APT’s long pre-training phase is computationally intensive because exhaustive exploration has high sample complexity.Combining APT with model-based reinforcement learning and search approaches is proposed as one way to reduce sample complexity.
- Future Work: Fine-tuning the whole pretrained model can cause catastrophic forgetting, motivating alternatives that preserve the pretrained model.One proposed direction combines a fixed pretrained model with a randomly initialized model.
A General Implementation Details
The experiments largely follow DrQ while specifying visual-network architectures, augmentation, replay, optimization, training durations, evaluation, and implementation resources.
- Architecture: APT experiments largely follow DrQ, using a three-layer convolutional policy network and an Impala neural encoder without its LSTM module.The encoder uses ELU nonlinearities and spectral normalization with five power iterations.
- Architecture: The convolutional network is followed by a LayerNorm-normalized fully connected layer with tanh applied to its output.
- Data Processing: Images are padded by 4 pixels on each side and randomly cropped back to 84 × 84 whenever sampled from the replay buffer.The replay buffer contains 100K entries.
- Training: Pre-training uses 5M environment steps on DMControl and 250M on Atari, followed by evaluation for 125K steps after 100K training steps.Contrastive learning uses learning rate 0.001, temperature 0.1, and a 64-dimensional projection output.
- Resources: The APT implementation is publicly available in the referenced GitHub repository.
C DeepMind Control Suite Details
The DMControl experiments specify environment-specific action repeats and direct readers to the corresponding hyperparameter tables.
- Environment Settings: Action-repeat hyperparameters are specified separately for each environment.
- Hyperparameters: DMControl experiment hyperparameters are provided in Tables 9 and 8.
D Asymptotic Behavior of Intrinsic Reward
Under finite-state and deterministic-encoder assumptions, the intrinsic reward decreases toward zero as observed state coverage grows, although the finite-state assumption may not hold in large environments.
- Asymptotic Behavior: The intrinsic reward decreases to 0 as more of the state space is visited, which is favorable for pre-training.
- Proposition 1: Under an episodic finite-state MDP with a deterministic encoder, an optimal intrinsic-reward policy reaches intrinsic reward 0 as sample size T approaches infinity.The proposition assumes the intrinsic reward uses k nearest neighbors with k ∈ N.
- Limitation: The finite-state assumption may not hold for large, complex environments such as Atari games.The proposition nevertheless provides insight into this intrinsic reward’s use for pre-training.
- Proof: The proof relies on visitation counts exceeding k as T approaches infinity because all states communicate in an episodic MDP.
E DeepMind Control Suite Sparse Environments
The evaluation adds three customized sparse-reward task families in DMControl: Jump Sparse, Reach Sparse, and Walker Turnover Sparse.
- Jump Sparse rewards HalfCheetah, Hopper, and Walker agents with 1 for jumping above a given height and 0 otherwise.
- Reach Sparse rewards HalfCheetah, Hopper, and Walker agents with 1 for reaching a target location and 0 otherwise.
- Walker Turnover Sparse starts Walker upside down and is included among the customized sparse-reward tasks.
F Scores on the full 57 Atari games
The full Atari evaluation compares APT with baseline methods across all 57 games using raw scores, with the highest score marked for each subset.
- APT achieves super-human performance on 12 Atari games, versus a maximum of 11 games for any previous method.
- APT achieves scores significantly higher than any previous methods across the individual Atari-game comparisons.
- Table 10 reports raw scores for each method on Atari games and marks the highest score in each subset in bold.
- Because VISR source code was unavailable, the authors made a best-effort attempt to reproduce its algorithm.