DeepMDP: Learning Continuous Latent Space Models for Representation Learning

By Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, Marc G. Bellemare. 2019

351 citations recorded by the source corpus.

Many reinforcement learning (RL) tasks provide the agent with high-dimensional observations that can be simplified into low-dimensional continuous states. To formalize this process, we introduce the concept of a DeepMDP, a parameterized latent space model that is trained via the minimization of two tractable losses: prediction of rewards and prediction of the distribution over next latent states. We show that the optimization of these objectives guarantees (1) the quality of the latent space as a representation of the state space and (2) the quality of the DeepMDP as a model of the environment. We connect these results to prior work in the bisimulation literature, and explore the use of a variety of metrics. Our theoretical findings are substantiated by the experimental result that a trained DeepMDP recovers the latent structure underlying high-dimensional observations on a synthetic environment. Finally, we show that learning a DeepMDP as an auxiliary task in the Atari 2600 domain leads to large performance improvements over model-free RL.

Subfields: Artificial Intelligence, Representation & Self-Supervised Learning, Probabilistic & Bayesian Methods

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

In Representation & Self-Supervised Learning

In Probabilistic & Bayesian Methods

Citations in the corpus

Mastering Atari, Go, chess and shogi by planning with a learned model (2019, methodology): “[11] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G.”
Provable Representation Learning for Imitation with Contrastive Fourier Features (2021, methodology): “We consider the standard MDP framework [32], in which the environment is given by a tuple M := (cid:104) S, A, R , P , µ, γ (cid:105) , where S is the state space, A is the action space, R : S × A → [ − R max , R max ] is the reward function, P : S × A → ∆( S ) is the transition function, 2 µ ∈ ∆( S”
Representation Matters: Offline Pretraining for Sequential Decision Making (2021, background): “…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., 2021).”
Dream to Control: Learning Behaviors by Latent Imagination (2019, background): “In principle, this could be achieved by simply learning to predict future rewards given actions and past observations (Oh et al., 2017; Gelada et al., 2019; Schrittwieser et al., 2019).”
Learning Domain Invariant Representations in Goal-conditioned Block MDPs (2021, methodology): “[6, 7, 50] implement the model-bisimulation framework [12] to acquire a minimal but sufficient representation for solving the MDP problem.”
Generalized Hidden Parameter MDPs Transferable Model-based RL in a Handful of Trials (2020, methodology): “In contrast, DeepMDP models a single task/MDP entirely in a latent space (Gelada et al. 2019), and can be merged to form a Deep-GHP-MDP.”
Local Information Agent Modelling in Partially-Observable Environments (2021, methodology): “…sample efﬁciency of RL algorithms using reconstruction-based [Corneil et al., 2018, Igl et al., 2018, Kurutach et al., 2018, Gregor et al., 2019, Gelada et al., 2019] or non-reconstruction-based methods [Laskin et al., 2020, Zhang et al., 2020], using contrastive learning or bi-simulation…”
Scalable methods for computing state similarity in deterministic Markov Decision Processes (2019, methodology): “Gelada et al. (2019) introduced DeepMDP latent models and established a close relationship with bisimulation metrics (speciﬁcally, Theorem 4 in their paper).”
Offline Reinforcement Learning with Pseudometric Learning (2021, background): “, 2021) or can be used as an auxiliary loss which leads to improvement performance on the Atari benchmark (Gelada et al., 2019).”
Contrastive Learning of Structured World Models (2019, background): “…the contrastive objective in Franc¸ois-Lavet et al. (2018), the mutual information objective in Anand et al. (2019), the distribution matching objective in Gelada et al. (2019) or using causality-based losses and physical priors in latent space (Jonschkowski & Brock, 2015; Ehrhardt et al., 2018).”
Neuroevolution of self-interpretable agents (2020, background): “[30, 84] suggest that not all details in the visual input are equally important, specifically they pointed out that rather than learning abstract representations that are capable of reconstructing the full observation, it is sufficient if the representation allows predicting”
Learning Invariant Representations for Reinforcement Learning without Reconstruction (2020, methodology): “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.”
Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning (2022, background): “Gelada et al. (2019) present results that their learned representation bounds the bisimulation metric and Zhang et al. (2021); Agarwal et al. (2021); Castro et al. (2021) use on-policy versions of the bisimulation metric (Castro, 2020) to learn an embedding space.”
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].”