Learning Invariant Representations for Reinforcement Learning without Reconstruction

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 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 . 1 Introduction Learning control from images is important for many real world applications . While deep reinforcement learning ( RL ) has enjoyed many successes in simulated tasks , learning control from real vision is more complex , especially outdoors , where images reveal detailed scenes of a complex and unstructured world . Furthermore , while many RL algorithms can eventually learn control from real images given unlimited data , data-efficiency is often a necessity in real trials which are expensive and constrained to real-time . Prior methods for data-efficient learning of simulated visual tasks typically use representation learning . Representation learning summarizes images by encoding them into smaller vectored representations better suited for RL . For example , sequential autoencoders aim to learn lossless representations of streaming observations—sufficient to reconstruct current observations and predict future observations—from which various RL algorithms can be trained ( Hafner et al. , 2018 ; Lee et al. , 2019 ; Yarats et al. , 2019 ) . However , such methods are taskagnostic : the models represent all dynamic elements they observe in the world , whether they are relevant to the task or not . We argue such representations can easily “ distract ” RL algorithms with irrelevant information in the case of real images . The issues of distraction is less evident in popular simulation MuJoCo and Atari tasks , since any change in observation space is likely task-relevant , and thus , worth representing . By contrast , visual images that autonomous cars observe contain predominately task-irrelevant information , like cloud shapes and architectural details , illustrated in Figure 1 . ∗Equal contribution . Corresponding author : amy.x.zhang @ mail.mcgill.ca Rather than learning control-agnostic representations that focus on accurate reconstruction of clouds and buildings , we would rather achieve a more compressed representation from a lossy encoder , which only retains state information relevant to our task . If we would like to learn representations that capture only task-relevant elements of the state and are invariant to task-irrelevant information , intuitively we can utilize the reward signal to help determine task-relevance , as shown by Jonschkowski & Brock ( 2015 ) . As cumulative rewards are our objective , state elements are relevant not only if they influence the current reward , but also if they influence state elements in the future that in turn influence future rewards . This recursive relationship can be distilled into a recursive task-aware notion of state abstraction : an ideal representation is one that is predictive of reward , and also predictive of itself in the future . We propose learning such an invariant representation using the bisimulation metric , where the distance between two observation encodings correspond to how “ behaviourally different ” ( Ferns & Precup , 2014 ) both observations are . Our main contribution is a practical representation learning method based on the bisimulation metric suitable for downstream control , which we call deep bisimulation for control ( DBC ) . We additionally provide theoretical analysis that proves value bounds between the optimal value function of the true MDP and the optimal value function of the MDP constructed by the learned representation . Empirical evaluations demonstrate our nonreconstructive approach using bisimulation is substantially more robust to task-irrelevant distractors when compared to prior approaches that use reconstruction losses or contrastive losses . Our initial experiments insert natural videos into the background of MoJoCo control task as complex distraction . Our second setup is a high-fidelity highway driving task using CARLA ( Dosovitskiy et al. , 2017 ) , showing that our representations can be trained effectively even on highly realistic images with many distractions , such as trees , clouds , buildings , and shadows . For example videos see https : //sites.google.com/view/deepbisim4control . 2 Related Work . Our work builds on the extensive prior research on bisimulation in MDP state aggregation . Reconstruction-based Representations . Early works on deep reinforcement learning from images ( Lange & Riedmiller , 2010 ; Lange et al. , 2012 ) used a two-step learning process where first an auto-encoder was trained using reconstruction loss to learn a low-dimensional representation , and subsequently a controller was learned using this representation . This allows effective leveraging of large , unlabeled datasets for learning representations for control . In practice , there is no guarantee that the learned representation will capture useful information for the control task , and significant expert knowledge and tricks are often necessary for these approaches to work . In model-based RL , one solution to this problem has been to jointly train the encoder and the dynamics model end-to-end ( Watter et al. , 2015 ; Wahlström et al. , 2015 ) – this proved effective in learning useful task-oriented representations . Hafner et al . ( 2018 ) and Lee et al . ( 2019 ) learn latent state models using a reconstruction loss , but these approaches suffer from the difficulty of learning accurate long-term predictions and often still require significant manual tuning . Gelada et al . ( 2019 ) also propose a latent dynamics model-based method and connect their approach to bisimulation metrics , using a reconstruction loss in Atari . They show that ` 2 distance in the DeepMDP representation upper bounds the bisimulation distance , whereas our objective directly learns a representation where distance in latent space is the bisimulation metric . Further , their results rely on the assumption that the learned representation is Lipschitz , whereas we show that , by directly learning a bisimilarity-based representation , we guarantee a representation that generates a Lipschitz MDP . We show experimentally that our non-reconstructive DBC method is substantially more robust to complex distractors . Contrastive-based Representations . Contrastive losses are a self-supervised approach to learn useful representations by enforcing similarity constraints between data ( van den Oord et al. , 2018 ; Chen et al. , 2020 ) . Similarity functions can be provided as domain knowledge in the form of heuristic data augmentation , where we maximize similarity between augmentations of the same data point ( Laskin et al. , 2020 ) or nearby image patches ( Hénaff et al. , 2019 ) , and minimize similarity between different data points . In the absence of this domain knowledge , contrastive representations can be trained by predicting the future ( van den Oord et al. , 2018 ) . We compare to such an approach in our experiments , and show that DBC is substantially more robust . While contrastive losses do not require reconstruction , they do not inherently have a mechanism to determine downstream task relevance without manual engineering , and when trained only for prediction , they aim to capture all predictable features in the observation , which performs poorly on real images for the same reasons world models do . A better method would be to incorporate knowledge of the downstream task into the similarity function in a data-driven way , so that images that are very different pixel-wise ( e.g . lighting or texture changes ) , can also be grouped as similar w.r.t . downstream objectives . Bisimulation . Various forms of state abstractions have been defined in Markov decision processes ( MDPs ) to group states into clusters whilst preserving some property ( e.g . the optimal value , or all values , or all action values from each state ) ( Li et al. , 2006 ) . The strictest form , which generally preserves the most properties , is bisimulation ( Larsen & Skou , 1989 ) . Bisimulation only groups states that are indistinguishable w.r.t . reward sequences output given any action sequence tested . A related concept is bisimulation metrics ( Ferns & Precup , 2014 ) , which measure how “ behaviorally similar ” states are . Ferns et al . ( 2011 ) defines the bisimulation metric with respect to continuous MDPs , and propose a Monte Carlo algorithm for learning it using an exact computation of the Wasserstein distance between empirically measured transition distributions . However , this method does not scale well to large state spaces . Taylor et al . ( 2009 ) relate MDP homomorphisms to lax probabilistic bisimulation , and define a lax bisimulation metric . They then compute a value bound based on this metric for MDP homomorphisms , where approximately equivalent state-action pairs are aggregated . Most recently , Castro ( 2020 ) propose an algorithm for computing on-policy bisimulation metrics , but does so directly , without learning a representation . They focus on deterministic settings and the policy evaluation problem . We believe our work is the first to propose a gradient-based method for directly learning a representation space with the properties of bisimulation metrics and show that it works in the policy optimization setting . 3 Preliminaries . We start by introducing notation and outlining realistic assumptions about underlying structure in the environment . Then , we review state abstractions and metrics for state similarity . We assume the underlying environment is a Markov decision process ( MDP ) , described by the tuple M = ( S , A , P , R , γ ) , where S is the state space , A the action space , P ( s′|s , a ) the probability of transitioning from state s ∈ S to state s′ ∈ S , and γ ∈ [ 0 , 1 ) a discount factor . An “ agent ” chooses actions a ∈ A according to a policy function a ∼ π ( s ) , which updates the system state s′ ∼ P ( s , a ) , yielding a reward r = R ( s ) ∈ R. The agent ’ s goal is to maximize the expected cumulative discounted rewards by learning a good policy : maxπ EP [ ∑∞ t=0 [ γ tR ( st ) ] . While our primary concern is learning from images , we do not address the partial-observability problem explicitly : we instead approximate stacked pixel observations as the fully-observed system state s ( explained further in Appendix B ) . Bisimulation is a form of state abstraction that groups states si and sj that are “ behaviorally equivalent ” ( Li et al. , 2006 ) . For any action sequence a0 : ∞ , the probabilistic sequence of rewards from si and sj are identical . A more compact definition has a recursive form : two states are bisimilar if they share both the same immediate reward and equivalent distributions over the next bisimilar states ( Larsen & Skou , 1989 ; Givan et al. , 2003 ) . Definition 1 ( Bisimulation Relations ( Givan et al. , 2003 ) ) . Given an MDP M , an equivalence relation B between states is a bisimulation relation if , for all states si , sj ∈ S that are equivalent under B ( denoted si ≡B sj ) the following conditions hold : R ( si , a ) = R ( sj , a ) ∀a ∈ A , ( 1 ) P ( G|si , a ) = P ( G|sj , a ) ∀a ∈ A , ∀G ∈ SB , ( 2 ) where SB is the partition of S under the relation B ( the set of all groups G of equivalent states ) , and P ( G|s , a ) = ∑ s′∈G P ( s′|s , a ) . Exact partitioning with bisimulation relations is generally impractical in continuous state spaces , as the relation is highly sensitive to infinitesimal changes in the reward function or dynamics . For this reason , Bisimulation Metrics ( Ferns et al. , 2011 ; Ferns & Precup , 2014 ; Castro , 2020 ) softens the concept of state partitions , and instead defines a pseudometric space ( S , d ) , where a distance function d : S × S 7→ R≥0 measures the “ behavioral similarity ” between two states1 . Defining a distance d between states requires defining both a distance between rewards ( to soften Equation ( 1 ) ) , and distance between state distributions ( to soften Equation ( 2 ) ) . Prior works use the Wasserstein metric for the latter , originally used in the context of bisimulation metrics by van Breugel 1Note that d is a pseudometric , meaning the distance between two different states can be zero , corresponding to behavioral equivalence . & Worrell ( 2001 ) . The pth Wasserstein metric is defined between two probability distributions Pi and Pj as Wp ( Pi , Pj ; d ) = ( infγ′∈Γ ( Pi , Pj ) ∫ S×S d ( si , sj ) p dγ′ ( si , sj ) ) 1/p , where Γ ( Pi , Pj ) is the set of all couplings of Pi and Pj . This is known as the “ earth mover ” distance , denoting the cost of transporting mass from one distribution to another ( Villani , 2003 ) . Finally , the bisimulation metric is the reward difference added to the Wasserstein distance between transition distributions : Definition 2 ( Bisimulation Metric ) . From Theorem 2.6 in Ferns et al . ( 2011 ) with c ∈ [ 0 , 1 ) : d ( si , sj ) = max a∈A ( 1− c ) · |Rasi −R a sj |+ c ·W1 ( P a si , P a sj ; d ) . ( 3 )

Summary: The paper focuses on how learning state-representations that encode information relevant to the task can improve reinforcement learning from pixels. Often, observations in an MDP can contain information that are irrelevant (“distractors”) to the task at hand and can likely “distract” the downstream RL algorithm used. Unlike existing reconstruction based approaches (which don’t explicitly incentivize ignoring task-irrelevant information), the authors propose Deep Bisimulation Control (DBC) that relies on bi-simulation metrics (as the task-aware criterion) that encode behavioral similarity b/w states with respect to the reward structure. Instead of explicitly learning a bi-similarity distance function, authors enforce the representations which under L1 distances correspond to bi-simulation metrics. DBC demonstrates learning these representations in conjunction with the control policy, reward and a dynamics model. Furthermore, the authors highlight connections to causal inference which can hopefully further provide insights into which “new” reward structures can the learned representations generalize to (since bi-similarity metrics themselves are heavily dependent on the reward structure). Results obtained by the authors demonstrate that DBC can learn task specific representations and a control policy in a robust manner in the presence of distractors on the Deepmind Control Suite and the CARLA simulator. Additionally, the authors also demonstrate how DBC can generalize to new reward functions on Mujoco.