Towards Robust Bisimulation Metric Learning

Abstract

Learned representations in deep reinforcement learning (DRL) have to extract task-relevant information from complex observations, balancing between robustness to distraction and informativeness to the policy. Such stable and rich representations, often learned via modern function approximation techniques, can enable practical application of the policy improvement theorem, even in high-dimensional continuous state-action spaces. Bisimulation metrics offer one solution to this representation learning problem, by collapsing functionally similar states together in representation space, which promotes invariance to noise and distractors. In this work, we generalize value function approximation bounds for on-policy bisimulation metrics to non-optimal policies and approximate environment dynamics. Our theoretical results help us identify embedding pathologies that may occur in practical use. In particular, we find that these issues stem from an underconstrained dynamics model and an unstable dependence of the embedding norm on the reward signal in environments with sparse rewards. Further, we propose a set of practical remedies: (i) a norm constraint on the representation space, and (ii) an extension of prior approaches with intrinsic rewards and latent space regularization. Finally, we provide evidence that the resulting method is not only more robust to sparse reward functions, but also able to solve challenging continuous control tasks with observational distractions, where prior methods fail.

arXiv commentAccepted to NeurIPS 2021

References in corpus (16)

• In Defense of the Triplet Loss for Person Re-Identification
• Soft Actor-Critic Algorithms and Applications
• DeepMind Control Suite
• Large-Scale Study of Curiosity-Driven Learning
• Learning Latent Dynamics for Planning from Pixels
• Visual Reinforcement Learning with Imagined Goals
• Incentivizing Exploration In Reinforcement Learning With Deep Predictive Models
• Unifying Count-Based Exploration and Intrinsic Motivation
• Variational Intrinsic Control
• SOLAR: Deep Structured Representations for Model-Based Reinforcement Learning
• Learning to Poke by Poking: Experiential Learning of Intuitive Physics
• Surprise-Based Intrinsic Motivation for Deep Reinforcement Learning
• A survey on intrinsic motivation in reinforcement learning
• DeepMDP: Learning Continuous Latent Space Models for Representation Learning
• Decoupling Representation Learning from Reinforcement Learning
• Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction