Behavior From the Void: Unsupervised Active Pre-Training

TL;DR

This work tackles reward-free pre-training for visual RL by introducing APT, which actively explores reward-free environments to learn diverse behaviors through maximizing entropy in a learned latent space using a nonparametric particle-based estimator. It combines this intrinsic objective with contrastive representation learning (SimCLR-style) and off-the-shelf RL algorithms (e.g., DrQ) to enable effective downstream fine-tuning when rewards are revealed. Empirically, APT achieves strong data efficiency and superior performance on Atari and DMControl benchmarks, including substantial gains in sparse-reward tasks, and outperforms prior unsupervised RL baselines. The results suggest that latent-space entropy, paired with robust representation learning, can enable scalable, task-agnostic pre-training with meaningful transfer to a broad class of downstream tasks; future work includes combining with model-based methods and improving fine-tuning stability.

Abstract

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.

Figures (3)

• Figure 1: Comparison of state-of-the-art pixel-based RL with unsupervised pre-training. APT (ours) and count-based bonus (both based on DrQ kostrikov2020image) are trained for a long unsupervised period (5M environment steps) without access to environment reward, and then gets exposure to the environment reward during testing. APT significantly outperform training DrQ from scratch, count-based bonus, and ImageNet pre-trained model.
• Figure 2: Diagram of the proposed method APT. On the left shows the objective of APT, which is to maximize the expected reward and minimize the contrastive loss. The contrastive loss learns an abstract representation from observations induced by the policy. We propose a particle-based entropy maximization based reward function such that we can deploy state-of-the-art RL methods to maximize entropy in an abstraction space of the induced by the policy. On the right shows the idea of our particle-based entropy, which measures the distance between each data point and its k nearest neighbors based on k nearest neighbors.
• Figure 3: Results of different methods in environments from DMControl. All curves are the average of three runs with different seeds, and the shaded areas are standard errors of the mean.

Theorems & Definitions (2)

• Proposition 1
• proof