Revisiting Bisimulation Metric for Robust Representations in Reinforcement Learning

TL;DR

This paper improves the bisimulation metric for reinforcement learning by addressing its limitations with a new, adaptive measure that enhances representation quality, supported by theoretical guarantees and empirical validation on benchmark tasks.

Contribution

It introduces a revised bisimulation metric with a more precise reward gap and adaptive updates, providing convergence guarantees and better representations in RL.

Findings

Enhanced representation distinctiveness demonstrated on benchmarks.

Theoretical convergence guarantees established.

Improved performance over conventional bisimulation metrics.

Abstract

Bisimulation metric has long been regarded as an effective control-related representation learning technique in various reinforcement learning tasks. However, in this paper, we identify two main issues with the conventional bisimulation metric: 1) an inability to represent certain distinctive scenarios, and 2) a reliance on predefined weights for differences in rewards and subsequent states during recursive updates. We find that the first issue arises from an imprecise definition of the reward gap, whereas the second issue stems from overlooking the varying importance of reward difference and next-state distinctions across different training stages and task settings. To address these issues, by introducing a measure for state-action pairs, we propose a revised bisimulation metric that features a more precise definition of reward gap and novel update operators with adaptive coefficient. We also offer theoretical guarantees of convergence for our proposed metric and its improved representation distinctiveness. In addition to our rigorous theoretical analysis, we conduct extensive experiments on two representative benchmarks, DeepMind Control and Meta-World, demonstrating the effectiveness of our approach.