# Task-Aware Exploration via a Predictive Bisimulation Metric

arXiv: https://arxiv.org/abs/2602.18724 · HTML: https://mlsanity.com/paper/2602.18724/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 30 of its 44 references that are in the map.
- First posted on arXiv 21 Feb 2026.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 30)

- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Diversity is All You Need: Learning Skills without a Reward Function](https://mlsanity.com/paper/1802.06070/) (2018)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning](https://mlsanity.com/paper/2107.09645/) (2021)
- [Self-Supervised Exploration via Disagreement](https://mlsanity.com/paper/1906.04161/) (2019)
- [Efficient Exploration via State Marginal Matching](https://mlsanity.com/paper/1906.05274/) (2019)
- [Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2102.11271/) (2021)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [Explore, Discover and Learn: Unsupervised Discovery of State-Covering Skills](https://mlsanity.com/paper/2002.03647/) (2020)
- [CIC: Contrastive Intrinsic Control for Unsupervised Skill Discovery](https://mlsanity.com/paper/2202.00161/) (2022)
- [Lipschitz-constrained Unsupervised Skill Discovery](https://mlsanity.com/paper/2202.00914/) (2022)
- [TACO: Temporal Latent Action-Driven Contrastive Loss for Visual Reinforcement Learning](https://mlsanity.com/paper/2306.13229/) (2023)
- [QueST: Self-Supervised Skill Abstractions for Learning Continuous Control](https://mlsanity.com/paper/2407.15840/) (2024)
- [Closed-Loop Visuomotor Control with Generative Expectation for Robotic Manipulation](https://mlsanity.com/paper/2409.09016/) (2024)
- [Towards Robust Bisimulation Metric Learning](https://mlsanity.com/paper/2110.14096/) (2021)
- [DrM: Mastering Visual Reinforcement Learning through Dormant Ratio Minimization](https://mlsanity.com/paper/2310.19668/) (2023)
- [Behavior Contrastive Learning for Unsupervised Skill Discovery](https://mlsanity.com/paper/2305.04477/) (2023)
- [Curiosity-Driven Exploration via Latent Bayesian Surprise](https://mlsanity.com/paper/2104.07495/) (2021)
- [Robust Representation Learning by Clustering with Bisimulation Metrics for Visual Reinforcement Learning with Distractions](https://mlsanity.com/paper/2302.12003/) (2023)
- [Generalizing Consistency Policy to Visual RL with Prioritized Proximal Experience Regularization](https://mlsanity.com/paper/2410.00051/) (2024)
- [Understanding and Addressing the Pitfalls of Bisimulation-based Representations in Offline Reinforcement Learning](https://mlsanity.com/paper/2310.17139/) (2023)
- [SkiLD: Unsupervised Skill Discovery Guided by Factor Interactions](https://mlsanity.com/paper/2410.18416/) (2024)
- [Constrained Ensemble Exploration for Unsupervised Skill Discovery](https://mlsanity.com/paper/2405.16030/) (2024)
- [When Should We Prefer State-to-Visual DAgger Over Visual Reinforcement Learning?](https://mlsanity.com/paper/2412.13662/) (2024)
- [Constrained Intrinsic Motivation for Reinforcement Learning](https://mlsanity.com/paper/2407.09247/) (2024)
- [State Chrono Representation for Enhancing Generalization in Reinforcement Learning](https://mlsanity.com/paper/2411.06174/) (2024)
- [Spectral Representation-based Reinforcement Learning](https://mlsanity.com/paper/2512.15036/) (2025)

## Built on by (most cited first, 1 of 1)

- [SAVGO: Learning State-Action Value Geometry with Cosine Similarity for Continuous Control](https://mlsanity.com/paper/2605.00787/) (2026)

## Abstract (arXiv)

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 behaviorally 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. Building on this robust metric, we design potential-based exploration bonuses, which measure the relative novelty of adjacent observations over the latent space. Extensive experiments on MetaWorld and Maze2D show that TEB achieves superior exploration ability and outperforms recent baselines.