State Representation Learning for Goal-Conditioned Reinforcement Learning

TL;DR

This paper introduces a self-supervised state embedding method for reward-free MDPs that captures transition distances, aiding goal-conditioned policy learning without domain knowledge, validated on control and multi-goal tasks.

Contribution

It proposes a novel domain-agnostic state representation learning approach that encodes transition distances for goal-conditioned reinforcement learning.

Findings

Effective in classic control domains

Works well in multi-goal environments

Learns representations in large and continuous domains

Abstract

This paper presents a novel state representation for reward-free Markov decision processes. The idea is to learn, in a self-supervised manner, an embedding space where distances between pairs of embedded states correspond to the minimum number of actions needed to transition between them. Compared to previous methods, our approach does not require any domain knowledge, learning from offline and unlabeled data. We show how this representation can be leveraged to learn goal-conditioned policies, providing a notion of similarity between states and goals and a useful heuristic distance to guide planning and reinforcement learning algorithms. Finally, we empirically validate our method in classic control domains and multi-goal environments, demonstrating that our method can successfully learn representations in large and/or continuous domains.