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

2026-07-20Artificial Intelligence

Artificial IntelligenceMachine Learning
AI summary

The authors found that how you measure similarity between states in visual reinforcement learning (RL) impacts how well the AI learns. Using simple fixed measurements can limit the AI, while completely flexible ones can cause problems. To fix this, they created PAMD, a new way to measure similarities that adapts based on the pairs of states being compared. This method plugs into existing RL algorithms and improves their performance on visual control tasks.

reinforcement learningvisual representationlatent spacebehavioral distanceMahalanobis distancebisimulationMuJoCocontinuous controlmetric learning
Authors
Daegyeong Roh, Juho Bae, Han-Lim Choi
Abstract
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.