PAMD: Structured Adaptive Distances for Bisimulation Representations in Visual Reinforcement Learning

In visual reinforcement learning (RL), many algorithms learn state representations by aligning a latent distance with a behavioral distance—a metric derived from reward and transition similarities. Existing methods typically use fixed, hand-designed metrics like L1 or L2 norms. This paper investigates whether these rigid choices limit the expressivity of the learned representation and whether a more flexible, adaptive distance metric can improve performance without falling into degenerate solutions.

The authors introduce Pairwise Adaptive Mahalanobis Distance (PAMD). Unlike fixed norms, PAMD parameterizes a positive-definite quadratic form that is conditioned on the specific pair of states being compared. By using a lightweight neural network (MetricNet) to output the Cholesky decomposition of a matrix, the method ensures the distance remains symmetric and positive-definite. To prevent the distance module from simply absorbing the supervision signal (a common failure mode in unconstrained metric learning), the authors apply trace normalization and a small ridge penalty, ensuring the distance remains a structured, meaningful surrogate for behavioral similarity.

PAMD acts as a drop-in replacement for existing bisimulation-based RL algorithms. Experiments on pixel-based DeepMind Control Suite tasks demonstrate that equipping standard algorithms like DBC, MICo, and SimSR with PAMD leads to substantial improvements in final task returns. The authors show that PAMD effectively reduces representation under-training, where the encoder fails to capture relevant features because the distance module is too flexible or too rigid. By providing a middle ground between fixed norms and unconstrained learned metrics, PAMD allows the representation to better adapt to the specific geometry of the environment's behavioral distance.