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

Many visual reinforcement learning (RL) algorithms learn representations by matching latent distances to a behavioral distance induced by reward and transition similarity. In practice, the choice of the latent distance can strongly affect performance: using a fixed, pre-specified global norms (e.g., $\ell_p$ norms or other hand-designed metrics) may be overly restrictive to capture the behavioral distance. In contrast, unconstrained pairwise distances may admit degenerate solutions that drive the metric loss down without improving the representation. To address this gap, we introduce **PAMD: Pairwise Adaptive Mahalanobis Distance**, which parameterizes a positive-definite, pair-conditioned metric for measuring latent state similarity. PAMD is a simple plug-in for existing bisimulation-based methods, offering a more expressive yet structured alternative to fixed, pre-specified latent distances. We empirically validate our method on visual MuJoCo continuous-control tasks, where final performance of several recent bisimulation-based RL algorithms is substantially improved when equipped with the distance we propose.

Many AI agents learn to act from images, such as camera views in a robot or game-like environment. To make good decisions, the agent must learn which visual differences matter for the task and which differences are irrelevant. For example, two images may look different because of lighting or background changes, but they may still require the same action. This paper studies how to help reinforcement learning agents compare what they see in a more useful way. Existing methods often use a fixed rule for measuring similarity between internal representations, which can be too rigid. We propose PAMD, a method that lets the agent adapt how it measures similarity depending on the pair of situations being compared, while still keeping this comparison structured and well constrained. Across several image-based control tasks, adding PAMD to existing reinforcement learning methods improves final task performance. These results suggest that better ways of measuring similarity can help agents learn more useful representations from visual observations.