Fast Task Inference with Variational Intrinsic Successor Features

Steven Hansen; Will Dabney; Andre Barreto; Tom Van de Wiele; David Warde-Farley; Volodymyr Mnih

Source-linked AI summary

Fast Task Inference with Variational Intrinsic Successor Features

Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, Volodymyr Mnih

arXiv:1906.05030v2Jun 12, 2019cs.LGcs.AIstat.ML

TL;DR

VISR addresses the limited generalization of distinguishability-based behavior learning and the feature-definition requirement of successor features. It combines BMI-derived controllable features with successor features for fast task inference, achieving notable performance across the full Atari suite after unsupervised pre-training.

Problem
Distinguishability-based behavior learning struggles to generalize beyond explicitly learned behaviors, while successor features require rewards linear in automatically defined features.
Method
VISR combines features learned by maximizing behavioral mutual information with successor features to support generalized behavior and fast task inference.
Results
VISR achieved notable performance on the full Atari task suite with unsupervised pre-training, outperforming all baselines and matching hundreds of millions of DQN interaction steps on some games.
Takeaways & Limitations
Bridging behavioral mutual information and successor features provides a method for learning controllable features and inferring tasks from limited feedback.
Takeaways & Limitations
The best practical approach for inferring task vectors remains unclear because extrinsic rewards are not guaranteed to be linear in the learned features.

Abstract

from arXiv · show

It has been established that diverse behaviors spanning the controllable subspace of an Markov decision process can be trained by rewarding a policy for being distinguishable from other policies \citep{gregor2016variational, eysenbach2018diversity, warde2018unsupervised}. However, one limitation of this formulation is generalizing behaviors beyond the finite set being explicitly learned, as is needed for use on subsequent tasks. Successor features \citep{dayan93improving, barreto2017successor} provide an appealing solution to this generalization problem, but require defining the reward function as linear in some grounded feature space. In this paper, we show that these two techniques can be combined, and that each method solves the other's primary limitation. To do so we introduce Variational Intrinsic Successor FeatuRes (VISR), a novel algorithm which learns controllable features that can be leveraged to provide enhanced generalization and fast task inference through the successor feature framework. We empirically validate VISR on the full Atari suite, in a novel setup wherein the rewards are only exposed briefly after a long unsupervised phase. Achieving human-level performance on 14 games and beating all baselines, we believe VISR represents a step towards agents that rapidly learn from limited feedback.

1 Introduction

Unsupervised RL pre-training can collect useful behavior before rewards are available, but existing behavioral mutual-information methods struggle to generalize and infer tasks quickly. VISR combines behavioral mutual information with successor features to address these limitations in a reward-sparse adaptation setting.

Motivation: Unsupervised RL pre-training removes rewards during initial interaction, then reinstates them to test data-efficient adaptation.The setup models environments where unstructured interaction is inexpensive but reward-defined tasks must later be learned quickly.
Existing methods: Behavioral mutual-information methods learn diverse policies by maximizing dependence between latent policies and their state-visitation behavior.These methods operate independently of extrinsic rewards and provide a natural basis for downstream adaptation.
Existing methods: Existing methods suffer from poor generalization and slow inference because the appropriate task may lie between learned latent behaviors.Conditional policies must generalize effectively to latent codes not encountered during unsupervised training.
Contribution: VISR combines behavioral mutual information with successor features to learn controllable features for fast task inference and improved generalization.The paper evaluates this combination in a 57-game Atari setting with a long unsupervised phase followed by reward-based adaptation.

2 reinforcement learning with unsupervised pre-training

The paper studies RL with an initial reward-free interaction phase followed by reward-observable learning. Because interaction costs are treated as significant only after rewards appear, unsupervised interaction is used to collect information that can accelerate later learning.

MDP setting: The agent-environment interaction is modeled as a Markov decision process with states, actions, transition dynamics, rewards, and discounting.The reward specifies the signal received on each transition, while the discount factor weights future rewards.
Two-stage training: The training process has an unsupervised phase without observed rewards and a reinforcement-learning phase in which rewards become observable.This two-stage setup defines the paper’s unsupervised pre-training scenario.
RL objective: During reinforcement learning, the agent seeks a policy maximizing expected discounted return.The return sums future rewards with progressively smaller weights determined by the discount factor.
Role of pre-training: Because interaction costs are assumed significant only during reinforcement learning, the unsupervised phase can collect information to speed subsequent learning.The agent is treated as essentially free to interact with the environment during pre-training.

3 Universal successor features and fast task inference

Successor features support rapid policy evaluation when rewards are linear in learned features, but implementing this framework requires suitable features and a set of policies. The paper proposes answering both questions through behavioral mutual information and generalized policy improvement.

Successor features: Successor features represent multidimensional value functions in which feature components play the role of rewards.They can be computed with standard reinforcement-learning algorithms.
Fast task inference: If successor features are available from unsupervised learning, early reward observations can estimate task weights and enable rapid policy evaluation.The resulting action values can support a policy-improvement step during the supervised phase.
Universal successor features: Universal successor features condition successor representations on an encoding of the policy, allowing values to be computed for multiple policies.This extends the single-policy successor-feature strategy to a policy family.
Generalized policy improvement: Generalized policy improvement selects actions using values from multiple policies and is a strict generalization of single-policy improvement.The resulting value is at least as large as the value of every candidate policy for each state-action pair.
Open design questions: Implementing the framework requires defining both the reward features and the policies used by generalized policy improvement.The paper treats these as complementary design questions addressed by its subsequent behavioral mutual-information construction.

4 Behavioral Mutual Information

Behavioral mutual information provides an intrinsic objective for learning controllable behavior without task rewards. VISR connects this objective to successor-feature task representations, using a discriminator-derived signal and policy-conditioned features for downstream inference.

Feature design: The feature representation should make downstream task rewards likely to be simple functions of the learned features.Because rewards are unavailable during unsupervised learning, this relies on an inductive bias toward features relevant to many reasonable tasks.
BMI objective: Behavioral mutual information is maximized between a policy-conditioning variable and a function of the induced trajectory.The objective optimizes policy parameters so that behavior contains information about the conditioning variable.
BMI objective: The conditioning variable is commonly sampled from a fixed or state-independent distribution, simplifying the objective to conditional-entropy minimization.This choice is described as useful for stability.
Behavior representation: The behavior function may sample a state from a trajectory, while sufficiently long trajectories approximate the policy’s steady-state distribution.The paper uses a single uniformly sampled state for simplicity.
VISR implementation: The VISR diagram includes the task vector as an input to the successor-feature model, enabling generalized policy improvement.For the random-feature baseline, the discriminator is frozen after initialization while the same objective trains the successor-feature model.
Optimization: The discriminator can be used to train the policy through an intrinsic reward derived with the REINFORCE trick.Its variational parameters learn to predict the correct conditioning variable from behavior.

5 Variational Intrinsic Successor Features

VISR combines behavioral mutual information with successor features by identifying BMI conditioning vectors with successor-feature task vectors. This enables diverse unsupervised behaviors while supporting regression-based task inference and generalized policy improvement.

VISR combines successor features’ rapid task inference with BMI methods’ ability to learn diverse behaviors without rewards.
VISR identifies BMI conditioning vectors z with successor-feature task vectors w, making the SF and BMI reward formulations coincide.
Unit-length task vectors and features with a unit-scale Von Mises-Fisher discriminator satisfy the required BMI–SF correspondence.
The conditional policy is factorized into policy-conditional successor features and task vectors, using a USFA representation.
VISR selects a regression-derived base task vector and samples additional vectors from a Von Mises-Fisher distribution centered on it for generalized policy improvement.

6 Experiments

Experiments evaluate VISR on all 57 Atari games under unsupervised pre-training and brief reward-enabled adaptation, alongside unsupervised and low-data baselines. VISR’s reward-regression adaptation substantially improves over random search and outperforms the reported unsupervised alternatives.

6.1 Reinforcement Learning With Unsupervised Pre-training: VISR is evaluated on 57 Atari games after 250M reward-free steps followed by 100k reward-enabled steps.
6.1 Reinforcement Learning With Unsupervised Pre-training: The adapted DIAYN baseline remains significantly weaker than even ablated VISR versions when using five-dimensional conditioning vectors.
6.2 Unsupervised approaches: VISR can vastly outperform existing fully unsupervised methods across all reported criteria when given 100k steps of reinforcement learning.
6.4 Fast inference: Linear reward regression is not guaranteed to be appropriate because extrinsic rewards may fail to be linear in the learned features.
6.4 Fast inference: 109.16 mean performance and 8.99 median performance were achieved with reward-regression task inference, versus 63.57 and 3.45 for random search.
6.4 Fast inference: VISR with reward regression outperformed random search on 41 of 57 games, with one tie, using identical task-inference data.

7 Conclusions

VISR achieves notable performance across the full Atari suite after unsupervised pre-training and few-step reinforcement learning, while combining behavioral mutual information features with successor features. The results also identify limitations in the fully unsupervised version and motivate further improvement.

VISR achieved notable performance on all 57 Atari games after unsupervised pre-training and few-step reinforcement learning, outperforming all baselines.The evaluation used a long unsupervised phase followed by a short reward-exposed phase.
The fully unsupervised version of VISR produced underwhelming results, leaving substantial room for improvement.The authors suggest incorporating curiosity-based intrinsic rewards to encourage exploration of controllable policies.
VISR’s successor-feature formulation bridges behavioral mutual information and successor features to address feature construction and task inference.

8 Appendix

The appendix derives VISR’s intrinsic-reward training procedure and evaluates its learned features, sampled rewards, value functions, and implementation on gridworld and Atari settings.

Derivation of Intrinsic Reward: The mutual-information lower bound is optimized with a variational approximation q, whose parameters are trained by back-propagation.The policy parameters require a REINFORCE-style score-function estimator because the environment is non-differentiable.
Derivation of Intrinsic Reward: The intrinsic reward is log q(z | s), allowing standard reinforcement-learning algorithms to optimize the policy.The paper explicitly uses log q(z|s) as the reward signal, with n-step Q-learning used throughout.
Qualitative Grid-World Results: VISR learns five-dimensional features in a 10-by-10 gridworld, and the resulting structure is attributed to the mutual-information objective rather than input correlations.The gridworld states were represented as one-hot vectors.
Qualitative Grid-World Results: VISR’s feature space produces 49 varied reward-function partitionings, generated by taking inner products between φ and uniformly sampled vectors on the 5-sphere.The authors use this diversity as evidence that many externally defined rewards may lie near the learned feature space.
Qualitative Grid-World Results: Generalized policy improvement over 10 uniformly sampled policies yields value functions that correspond clearly to the 49 sampled reward functions.The authors report that VISR does not significantly suffer from underfitting despite representing an infinite space of value functions.