Saliency-Guided Representation with Consistency Policy Learning for Visual Unsupervised Reinforcement Learning

TL;DR

This paper introduces SRCP, a framework that enhances zero-shot visual unsupervised reinforcement learning by improving successor representations through saliency guidance and consistency policy learning.

Contribution

It proposes a novel decoupled representation learning approach with saliency guidance and a consistency policy for better generalization in visual URL.

Findings

SRCP achieves state-of-the-art zero-shot generalization on 16 tasks.

The framework improves successor measure accuracy and skill controllability.

Compatible with various successor representation methods.

Abstract

Zero-shot unsupervised reinforcement learning (URL) offers a promising direction for building generalist agents capable of generalizing to unseen tasks without additional supervision. Among existing approaches, successor representations (SR) have emerged as a prominent paradigm due to their effectiveness in structured, low-dimensional settings. However, SR methods struggle to scale to high-dimensional visual environments. Through empirical analysis, we identify two key limitations of SR in visual URL: (1) SR objectives often lead to suboptimal representations that attend to dynamics-irrelevant regions, resulting in inaccurate successor measures and degraded task generalization; and (2) these flawed representations hinder SR policies from modeling multi-modal skill-conditioned action distributions and ensuring skill controllability. To address these limitations, we propose Saliency-Guided Representation with Consistency Policy Learning (SRCP), a novel framework that improves zero-shot generalization of SR methods in visual URL. SRCP decouples representation learning from successor training by introducing a saliency-guided dynamics task to capture dynamics-relevant representations, thereby improving successor measure and task generalization. Moreover, it integrates a fast-sampling consistency policy with URL-specific classifier-free guidance and tailored training objectives to improve skill-conditioned policy modeling and controllability. Extensive experiments on 16 tasks across 4 datasets from the ExORL benchmark demonstrate that SRCP achieves state-of-the-art zero-shot generalization in visual URL and is compatible with various SR methods.