[Paper Review] Learning Invariant Representations for Reinforcement Learning without Reconstruction

Amy Zhang, Rowan McAllister+3 more|arXiv (Cornell University)|Jun 18, 2020

Reinforcement Learning in RoboticsComputer Science97 citations

TL;DR

This paper introduces Deep Bisimulation for Control (DBC), a non-reconstruction representation learning method that learns latent encodings aligned with bisimulation metrics to improve data-efficient RL from high-dimensional observations and ignore task-irrelevant details.

ABSTRACT

We study how representation learning can accelerate reinforcement learning from rich observations, such as images, without relying either on domain knowledge or pixel-reconstruction. Our goal is to learn representations that both provide for effective downstream control and invariance to task-irrelevant details. Bisimulation metrics quantify behavioral similarity between states in continuous MDPs, which we propose using to learn robust latent representations which encode only the task-relevant information from observations. Our method trains encoders such that distances in latent space equal bisimulation distances in state space. We demonstrate the effectiveness of our method at disregarding task-irrelevant information using modified visual MuJoCo tasks, where the background is replaced with moving distractors and natural videos, while achieving SOTA performance. We also test a first-person highway driving task where our method learns invariance to clouds, weather, and time of day. Finally, we provide generalization results drawn from properties of bisimulation metrics, and links to causal inference.

Motivation & Objective

• Motivate learning task-relevant representations that ignore distractors in high-dimensional observations.
• Propose a gradient-based method to train encoders so latent distances match bisimulation distances.
• Show that the learned representations enable robust downstream control across diverse visual distractors.
• Provide theoretical guarantees linking the learned representation to bounds on optimal value functions and causal feature sets.

Proposed method

• Define and optimize a latent space distance to match the on-policy bisimulation metric via a loss J(phi) that aligns |z_i - z_j|_1 with |r_i - r_j| + gamma W_2(P_hat(.|z_i,a_i), P_hat(.|z_j,a_j)) (Equation 4).
• Use a probabilistic dynamics model to compute the Wasserstein-2 distance between predicted next-state Gaussians for the latent transitions.
• Integrate the encoder phi with a reinforcement learning algorithm (SAC) so that the encoder, policy, and dynamics model are trained iteratively (Algorithm 1 and Algorithm 2).
• Argue and prove that the learned pi*-bisimulation metric converges to a fixed point, and show value-function Lipschitz bounds with respect to the bisimulation-based latent distance.
• Provide generalization insights linking bisimulation-based representations to causal ancestors of the reward and to task generalization.

Experimental results

Research questions

• RQ1 Can a non-reconstructive, bisimulation-based representation learn task-relevant latent encodings that ignore distraction in vision-based RL?
• RQ2 Does aligning latent-space distances to bisimulation distances improve data efficiency and robustness to visual distractors compared to reconstruction or contrastive methods?
• RQ3 What theoretical guarantees connect the learned representation to value function bounds and causal-inference concepts?
• RQ4 How well does the approach generalize across different reward functions and distraction types (e.g., clouds, weather, and time-of-day in driving tasks)?

Key findings

• DBC yields state-of-the-art robustness to background distractions in visual control benchmarks, outperforming reconstruction and contrastive baselines on distracted MuJoCo tasks.
• In driving-style tasks with high visual realism (CARLA), DBC significantly outperforms baselines and achieves higher task success and larger distance without collisions.
• The approach achieves a 46.8% improvement over the next best baseline in the driving experiment.
• Latent representations learned by DBC cluster similar task-relevant states together regardless of background changes, as shown by qualitative visualizations.
• Theoretical results connect the learned latent bisimulation metric to Lipschitz bounds on the optimal value function and to the sufficiency of causal ancestors for reward prediction.
• Empirical generalization: the learned representation generalizes to unseen distractors and to new reward functions sharing the same causal ancestors.