Multi-scale Predictive Representations for Goal-conditioned Reinforcement Learning

Goal-conditioned reinforcement learning (GCRL) requires agents to learn effective state and goal representations, which represents a challenging problem, especially in high-dimensional vision-based environments, as differences in the observations can be uncorrelated with dynamical distances. Classical deep reinforcement learning techniques often fail to capture the alignment between state and goal spaces, requiring additional representation learning techniques. To address this, we propose $\textit{Ms.PR}$, a representation learning framework that augments model-free GCRL methods with a multi-scale predictive architecture. Leveraging predictive dynamics learning, the latent embedding space captures both physical causality and temporal distances between states. Furthermore, by learning information at multiple timescales, the agent acquires a better understanding of how close and distant goals relate to a given state. We demonstrate that Ms.PR leads to improved representation quality and strong performance on the OGBench benchmark, both on vision and state-based tasks.