Task-Aware Exploration via a Predictive Bisimulation Metric

Accelerating exploration in visual reinforcement learning under sparse rewards remains challenging due to the substantial task-irrelevant variations. Despite advances in intrinsic exploration, many methods either assume access to low-dimensional states or lack task-aware exploration strategies, thereby rendering them fragile in visual domains. To bridge this gap, we present TEB, a Task-aware Exploration approach that tightly couples task-relevant representations with exploration through a predictive Bisimulation metric. Specifically, TEB leverages the metric not only to learn behaviorally grounded task representations but also to measure global intrinsic novelty over the learned latent space. To realize this, we first theoretically mitigate the representation collapse of degenerate bisimulation metrics under sparse rewards by internally introducing a simple but effective predicted reward differential term. Building on this robust metric, we further introduce potential-based global exploration bonuses over anchor states, which measure the relative novelty between observations over the latent space. Extensive experiments on MetaWorld and Maze2D show that TEB achieves superior exploration ability and outperforms recent baselines.

Learning from images is hard for reinforcement learning agents when useful reward feedback is rare while the scene contains many distracting details, such as lighting, textures, or backgrounds. We introduce TEB, a method that helps an agent focus on the parts of a scene that matter for completing a task, while exploring its environment more effectively. This prevents the agent from being misled by visual changes that are unimportant. In tests on robotic arm control tasks and maze navigation, our method learns faster and performs better than strong existing approaches. This can help make vision-based AI systems more reliable in difficult settings.