Accounting for the Sequential Nature of States to Learn Features for Reinforcement Learning

TL;DR

This paper addresses the challenge of learning useful state representations in reinforcement learning environments with non-overlapping states by leveraging the sequential nature of states to improve variational autoencoder performance.

Contribution

The authors introduce a method that uses the sequential order of states to enable VAEs with triplet loss to learn effective features without extra supervision.

Findings

VAEs fail in environments with non-overlapping states.

Sequential state information can be used to approximate a distance metric.

Modified VAEs with triplet loss learn useful features in challenging environments.

Abstract

In this work, we investigate the properties of data that cause popular representation learning approaches to fail. In particular, we find that in environments where states do not significantly overlap, variational autoencoders (VAEs) fail to learn useful features. We demonstrate this failure in a simple gridworld domain, and then provide a solution in the form of metric learning. However, metric learning requires supervision in the form of a distance function, which is absent in reinforcement learning. To overcome this, we leverage the sequential nature of states in a replay buffer to approximate a distance metric and provide a weak supervision signal, under the assumption that temporally close states are also semantically similar. We modify a VAE with triplet loss and demonstrate that this approach is able to learn useful features for downstream tasks, without additional supervision, in environments where standard VAEs fail.