Latent Representation Prediction Networks

TL;DR

This paper introduces Latent Representation Prediction Networks (LARP), which learn representations optimized for predictability to improve planning and transferability in reinforcement learning tasks.

Contribution

The paper proposes a novel joint learning approach for representations and predictors, directly optimizing for predictability to enhance planning efficiency and transferability.

Findings

LARP outperforms pre-trained representations in planning tasks.

LARP is more sample-efficient than standard reinforcement learning methods.

The learned representations transfer successfully to dissimilar objects.

Abstract

Deeply-learned planning methods are often based on learning representations that are optimized for unrelated tasks. For example, they might be trained on reconstructing the environment. These representations are then combined with predictor functions for simulating rollouts to navigate the environment. We find this principle of learning representations unsatisfying and propose to learn them such that they are directly optimized for the task at hand: to be maximally predictable for the predictor function. This results in representations that are by design optimal for the downstream task of planning, where the learned predictor function is used as a forward model. To this end, we propose a new way of jointly learning this representation along with the prediction function, a system we dub Latent Representation Prediction Network (LARP). The prediction function is used as a forward model for search on a graph in a viewpoint-matching task and the representation learned to maximize predictability is found to outperform a pre-trained representation. Our approach is shown to be more sample-efficient than standard reinforcement learning methods and our learned representation transfers successfully to dissimilar objects.