Learning Invariant Representations for Reinforcement Learning without Reconstruction

By Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, Sergey Levine. 2020

593 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

Citations in the corpus

Provable Representation Learning for Imitation with Contrastive Fourier Features (2021, methodology): “…learning objectives in practice in Section 5, but for now we note that the expectation over d oﬀ already leads to a connection between theory and practice much closer than exists in other works, which often resort to supremums over state-action errors [30, 41, 1] or zero errors globally [17].”
Representation Matters: Offline Pretraining for Sequential Decision Making (2021, methodology): “While these methods originally required explicit knowledge of the reward and dynamics functions of the environment, a number of recent works have translated these ideas to stochastic representation learning objectives using deep neural networks (Gelada et al., 2019; Zhang et al., 2020; Agarwal et al”
Adapt to Environment Sudden Changes by Learning a Context Sensitive Policy (2022, methodology): “Self-Supervised Representation Learning has been used in several RL domains in recent work (Laskin, Srinivas, and Abbeel 2020; Zhang et al. 2021; Stooke et al. 2021).”
Local Information Agent Modelling in Partially-Observable Environments (2021, methodology): “This baseline was included because recent works [Laskin et al., 2020, Zhang et al., 2020] in state representation found that non-reconstruction methods tend to perform better than reconstruction methods, especially in domains that have pixel-based observations.”
Learning Markov State Abstractions for Deep Reinforcement Learning (2021, methodology): “In Figure 4, we compare against unmodiﬁed RAD, as well as contrastive methods CURL (Laskin et al., 2020b) and CPC (Van den Oord et al., 2018), bisimulation methods DeepMDP (Gelada et al., 2019) and DBC (Zhang et al., 2021), and pixel-reconstruction method SAC-AE (Yarats et al., 2019).”
Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning (2021, background): “There are other ways to approximate the Wasserstein distance in bisimulation metrics ( e.g., Ferns et al., 2006; 2011; Castro, 2020; Zhang et al., 2021).”
Contrastive Learning as Goal-Conditioned Reinforcement Learning (2022, background): “However, prior work has found it challenging to uphold this promise [7, 79, 118, 125], prompting many prior methods to employ separate objectives for representation learning and RL [31, 43, 70, 74, 89, 91, 100, 101, 117, 140, 143].”
Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction (2020, methodology): “Deep bisimulation for control (DBC) [26] and discriminative particle ﬁlter reinforcement learning (DPFRL) [27] are other types of cutting edge MFRL methods, which utilize different concepts of representation learning without reconstruction.”
Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning (2022, methodology): “(4) Adding a reward decoder model R ( φ ( s, g ) , φ ( s (cid:48) , g )) and (5) learned dynamics model as in DBC (Zhang et al., 2021).”
Offline Reinforcement Learning with Pseudometric Learning (2021, background): “, 2012) has been shown to be beneficial to learn controllable representations that eliminates unnecessary details of the state (Zhang et al., 2021) or can be used as an auxiliary loss which leads to improvement performance on the Atari benchmark (Gelada et al.”
Inventing Relational State and Action Abstractions for Effective and Efficient Bilevel Planning (2022, background): “Our work continues a long line of research on learning state abstractions [2, 3, 4, 5, 6, 46, 47, 48, 49, 50] and action abstractions [27, 51, 52] for planning.”
R3M: A Universal Visual Representation for Robot Manipulation (2022, methodology): “Such methods include contrastive learning with data augmentation [25, 26, 27], dynamics prediction [28, 29], bi-simulation [30], temporal or goal distance [31, 32], or domain specific information [33].”
Monte Carlo Tree Search With Iteratively Refining State Abstractions (2021, background): “However, while there has been recent progress in learning bisimulation metrics [24, 7], doing so in large MDPs with arbitrary transition distributions remains a challenging problem.”
A Metaverse: Taxonomy, Components, Applications, and Open Challenges (2022, unknown): “Zhang et al. [331] studied how expressive learning can accelerate reinforcement learning from rich observations (e.g., images) without relying on domain knowledge.”
Structural Similarity for Improved Transfer in Reinforcement Learning (2022, unknown): “In [28], Zhang et al. also utilized a bisimulation metric in order to disassociate task-irrelevant information between observations within MuJoCo [22] tasks and high-ﬁdelity highway driving tasks using CARLA, an open urban driving simulator [8].”
Diffusion policy: Visuomotor policy learning via action diffusion (2023, unknown): “But interestingly, it need only learn the task-relevant dynamics (Subramanian and Mahajan, 2019; Zhang et al., 2020); in the linear dynamics example here, if K has a column that is all zeros, then leftmost multiplication by K means that the associated state dynamics are not relevant to the task and…”
Learning Domain Invariant Representations in Goal-conditioned Block MDPs (2021, methodology): “We include MISA [6] and DBC [7] in our comparison as they are the most successful implementations for high-dimensional tasks.”