PaperlayerView exact revision on arXiv

Source-linked AI summary

DeepMDP: Learning Continuous Latent Space Models for Representation Learning

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

arXiv:1906.02736v1cs.LGstat.ML

TL;DR

Reinforcement learning often uses high-dimensional observations containing redundant or irrelevant information, motivating compact continuous representations. The paper introduces DeepMDPs, trains them with reward and next-latent-state prediction losses, and reports theoretical guarantees alongside improved Atari performance.

  • Problem

    RL observations can contain substantial redundant information and task-irrelevant noise, while prior representation-learning methods either lack scalability and deep-learning compatibility or lack formal guarantees.

  • Method

    DeepMDPs are parameterized latent space models trained by minimizing tractable losses for reward prediction and next-latent-state distribution prediction, with guarantees also established for expectation-based deep-network losses.

  • Results

    DeepMDP learning yields theoretical guarantees for value functions and representation quality, while the best architecture produces nearly consistent performance improvements over C51 across 60 Atari 2600 games.

  • Takeaways & Limitations

    Using the Wasserstein metric gives DeepMDP representation learning a theoretically sound connection to bisimulation, while DeepMDP losses can serve as an auxiliary task for model-free RL.

  • Takeaways & Limitations

    In complex environments such as Atari, optimization can become trapped in low-information local minima, making useful representations difficult to discover.

Abstract

from arXiv · show

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.

1. Introduction

DeepMDP addresses the gap between theoretically grounded representation-learning methods that do not scale and deep-learning methods that lack formal guarantees. It learns latent space models through tractable reward and next-latent-state prediction losses, with guarantees tied especially to Wasserstein metrics.

  • Motivation: RL observations can contain redundant, task-irrelevant information, motivating mappings from high-dimensional states to simplified representations.For example, Atari frames contain 33,600 pixels, while PONG can be played from the positions and velocities of three objects.
  • Motivation: Representation learning seeks to improve RL learning efficiency by mapping states to simplified states.
  • Research gap: Existing theoretically grounded methods often do not scale or combine easily with deep learning, while deep-learning approaches often lack formal guarantees.
  • Approach: DeepMDP is a latent space model trained by predicting rewards and distributions over next latent states using two tractable losses.
  • Theory: With Wasserstein latent transition loss, DeepMDP analysis connects representation learning to bisimulation and yields scalable, modern-deep-network-compatible guarantees.

2. Background

The background defines MDPs and latent space models, then develops Wasserstein- and Lipschitz-based conditions for analyzing their value functions. The presentation emphasizes Wasserstein metrics while noting assumptions and limits on the generality of the theory.

  • 2.1. Markov Decision Processes: An MDP comprises states, actions, rewards, transitions, and a discount factor, with policies assigning action distributions conditioned on states.
  • 2.2. Latent Space Models: A latent space model pairs a continuous metric state space with an embedding from original states and shares the action space with the original MDP.
  • 2.2. Latent Space Models: The latent transition loss compares embedded next-state distributions with a probability-distribution metric, primarily the Wasserstein metric in this work.
  • 2.2. Latent Space Models: DeepMDP denotes a parameterized latent space model trained by minimizing reward and transition losses, including expectation-based losses for deep networks.
  • 2.3. Wasserstein Metric: Wasserstein distance measures the minimum transport cost between distributions under an underlying metric and also has a dual form over 1-Lipschitz functions.
  • 2.4. Lipschitz Norm of Value Functions: The analysis establishes Lipschitz-value properties for optimal, sufficiently smooth, and constant policies in Lipschitz DeepMDPs.
  • 2.4. Lipschitz Norm of Value Functions: The main Lipschitz analysis assumes a transition Lipschitz constant strictly below 1, although this assumption is unnecessary for finite-horizon terminating episodes.
  • 2.4. Lipschitz Norm of Value Functions: A general framework for understanding Lipschitz-valued policies remains lacking, leaving broader policy-class analysis as future work.

3. Global DeepMDP Bounds

The global DeepMDP losses provide theoretical guarantees for both model-value accuracy and representation quality across the whole state-action space. These results support using the learned embedding for value prediction.

  • Global losses: The global DeepMDP losses are defined as reward and transition losses based on suprema over the whole state-action space.The losses are introduced as the paper’s concrete global objectives.
  • Value difference bound: For Lipschitz-valued policies, the difference between original and latent-model value functions is bounded by the global DeepMDP losses.The bound compares values computed in the original MDP with those computed in the DeepMDP after embedding states.
  • Representation quality: When both global reward and transition losses are zero, the embedding cannot collapse states with different values for the considered policies.This establishes representation quality for value prediction.
  • Representation quality: The representation result justifies using the embedding function as a representation for predicting values.The paper relates this conclusion to earlier model-based representation results in the linear setting.
  • Policy performance: The paper also bounds the performance loss from running the optimal DeepMDP policy in the original MDP relative to the original optimal policy.This result is stated as Theorem 5 in Appendix A.

4. Local DeepMDP Bounds

Local DeepMDP losses evaluate reward and transition prediction under a policy’s state-action distribution, making the guarantees compatible with data-limited deep RL. The resulting bounds apply primarily to states and actions visited by that policy.

  • Local losses: Local DeepMDP losses measure reward and transition errors under an expected state-action distribution ξ rather than across the entire state space.This accommodates missing data from many state-space regions and limited model capacity.
  • Local losses: The local losses are compatible with stochastic-gradient methods used to train neural networks, connecting the theoretical analysis to practice.The paper presents this compatibility as a bridge between theory and deep-network optimization.
  • Value difference bound: For policies whose visited state-action pairs have small local reward and transition losses, the DeepMDP provides accurate value functions on states likely visited by those policies.The guarantee is policy-dependent and concerns the policy’s induced state-action distribution.
  • Representation quality: The local value-difference guarantee extends to a local representation-quality bound for predicting a policy’s value function.The paper states this as an analogue of its global representation result.
  • Representation quality: The local representation guarantee applies to states visited often by the policy.It does not extend automatically to states outside the policy’s frequently visited region.

5. Bisimulation

The paper connects Wasserstein-based DeepMDPs to bisimulation by relating latent embeddings and policies to behavioral equivalence. This connection explains why the analyzed policy class is theoretically meaningful.

  • Bisimulation: Bisimulation formalizes behavioral equivalence by requiring equivalent states to share rewards and transition distributions over equivalence classes.The maximal bisimulation groups states into the largest sets satisfying these conditions.
  • Bisimulation metrics: Bisimulation metrics soften bisimulation’s all-or-nothing relation by assigning pseudometric distances that quantify behavioral similarity.They address cases where states differ only slightly in rewards or transitions.
  • Bisimulation metrics: The bisimulation metric has a unique fixed point whose kernel is the maximal bisimulation relation, and it upper-bounds optimal value differences.The contraction and kernel properties connect metric distance to value behavior.
  • Bisimulation metrics: Bisimulation metrics remain computationally costly and poorly compatible with deep networks in large-scale settings.This limitation motivates the DeepMDP connection as a scalable alternative.
  • DeepMDP connection: With global Wasserstein DeepMDP losses, the bisimulation distance is bounded using the latent embedding distance and those losses.This parallels the value-difference guarantee while replacing value difference with bisimulation distance.
  • Policy classes: Lipschitz-deep policies approximately contain Lipschitz-bisimilar policies, so minimizing global losses yields a nearby deep policy for each Lipschitz-bisimilar policy.The result characterizes the policy class appearing in the earlier bounds.

6. Beyond the Wasserstein

The paper generalizes value-difference guarantees beyond Wasserstein to Norm-MMD metrics, whose associated function norms encode different value-function smoothness assumptions. Wasserstein remains special for representation guarantees and bisimulation connections.

  • Norm-MMD metrics: Norm-MMD metrics generalize the transition-loss analysis through function seminorms, and bound expectation differences for functions controlled by the associated seminorm.This family includes metrics with different behavioral and smoothness properties.
  • Metric choice: Choosing a metric whose smoothness matches the environment may tighten bounds, while simpler metrics may be preferable computationally.The Wasserstein metric has high computational cost and biased stochastic-gradient estimates; KL is offered as a potentially more convenient alternative.
  • Smoothness interpretations: Different Norm-MMD metrics correspond to different forms of value-function smoothness, including maximum magnitude, maximal change, and total change over the state space.The paper illustrates these interpretations using Total Variation, Wasserstein, and Energy distance.
  • Generalized guarantees: Global and local value-difference bounds, along with the suboptimality result, extend to any Norm-MMD metric.The full statements and proofs are provided in Appendix A.6.
  • Scope of guarantees: Representation-quality and bisimulation results do not generalize to Norm-MMD metrics, emphasizing Wasserstein’s special role for representation learning.The broader metric family therefore preserves some model-based guarantees but not all representation guarantees.

7. Related Work in Representation Learning

Representation learning spans state aggregation methods that reduce state-space dimensionality and feature-based approaches that seek task-relevant information. Deep-learning methods offer strong empirical results on complex domains, while earlier theoretically grounded methods face scalability or integration limitations.

  • State aggregation methods reduce dimensionality by joining states without sacrificing necessary information.
  • Other representation-learning approaches seek features that support the task’s relevant information.
  • Auxiliary-task methods combine model-free deep reinforcement learning with additional objectives and have shown benefits across diverse domains.
  • Distributional RL and alternative discounting techniques have also been studied as auxiliary-task mechanisms.

8. Empirical Evaluation

The empirical evaluation tests DeepMDP representations in a synthetic environment and as auxiliary objectives for Atari agents. DeepMDP recovers task-relevant latent structure, but optimization and architecture choices materially affect performance.

  • 8.1. DonutWorld Experiments: 32x32 greyscale observations in DonutWorld are generated from a simple 2D latent state, enabling tests of whether DeepMDP recovers the agent’s x-y coordinates.
  • 8.1. DonutWorld Experiments: DeepMDP representations recover the agent’s 2D position from high-dimensional pixels, whereas autoencoder representations are less meaningful despite near-perfect task performance.
  • 8.1. DonutWorld Experiments: In four identical tracks, DeepMDP merges states with indistinguishable value functions and becomes almost invariant to track identity.
  • 8.1. DonutWorld Experiments: Reward and transition losses compete during optimization: an uninformative uniform-zero representation makes transitions easy to predict but produces large reward loss.
  • 8.2. Atari 2600 Experiments: Convolutional transition models yield the best DeepMDP performance, and the architecture is highly consequential for the agent.
  • 8.4. DeepMDPs as an Auxiliary Task: DeepMDP produces nearly consistent performance improvements over C51 across 60 Atari 2600 games.
  • 8.5. Comparison to Alternative Objectives: Predicting next latent representations outperforms observation reconstruction, next-observation prediction, and next-logit prediction as an auxiliary objective.

9. Discussion on Model-Based RL

DeepMDPs also provide a principled basis for model-based reinforcement learning, including latent-space and standard state-space modeling. Their framework supports alternatives to KL-based transition modeling through Norm-MMD metrics.

  • DeepMDP guarantees extend to model-based RL when the embedding function is the identity, covering models that predict rewards and next-state distributions.The paper relates negative log-probability training to a one-sample KL estimate and interprets it through DeepMDP guarantees.
  • Norm-MMD metrics offer a research direction for model-based RL beyond KL-based objectives.

10. Conclusions

The paper introduces DeepMDPs as tractable latent-space models with theoretical guarantees for value functions and representations. Wasserstein-based training connects these guarantees to bisimulation and supports strong Atari results, while model-based planning remains future work.

  • DeepMDPs are parameterized latent-space models trained by minimizing tractable reward and transition-prediction losses.
  • Norm-MMD transition losses guarantee the quality of learned-model value functions, while Wasserstein losses additionally support bisimulation-based representation guarantees.
  • Wasserstein training guarantees that states differing in value for supported policies are not collapsed by the representation.
  • Strong Atari 2600 results show that minimizing DeepMDP losses can benefit model-free reinforcement learning as an auxiliary task.
  • Using DeepMDP transition and reward models for planning or exploration is identified as a future research direction.

Changes From ICML 2019 Proceedings

The updated document adds Norm-MMD generalization and associated mathematics, makes minor changes to Lemma 1, and substantially expands the bisimulation discussion. The supplied theorem passages state value, policy, and representation guarantees under metric and Lipschitz assumptions.

  • Changes From ICML 2019 Proceedings: The updated version adds the generalization to Norm-MMD metrics and associated mathematics in Section 6.
  • Changes From ICML 2019 Proceedings: The updated version makes minor changes to Lemma 1 and partially rewrites several sections, especially the expanded bisimulation discussion.
  • Theoretical guarantees: Under Lipschitz assumptions, theorems bound representation-induced value differences and the suboptimality of the learned model’s optimal policy.
  • Bisimulation connection: Wasserstein-based analysis connects embedding distance and DeepMDP losses to upper bounds on the MDP’s bisimulation distance.
  • Policy relationships: Theorems establish approximation relationships between Lipschitz-bisimilar policies and Lipschitz-DeepMDP policies.

B.1. Environment Specification

The synthetic DonutWorld environment uses a circular track with low-dimensional latent structure and, in one variant, high-dimensional pixel observations. Experiments learn a two-dimensional DeepMDP with neural embedding, reward, and transition models, while Atari configurations evaluate multiple latent and transition architectures.

  • Environment: DonutWorld places an agent on a circular track bounded between radii 3 and 6, with movement capped at 1 unit per timestep.
  • Environment: The four-track DonutWorld variant uses 64-by-64 pixel observations while preserving identical policy value functions across tracks.
  • Training setup: The synthetic environment is represented with a two-dimensional latent state, and training samples state-action pairs uniformly.
  • Model architectures: The embedding uses a convolutional neural network, while reward and transition functions use two-hidden-layer feed-forward networks.
  • Synthetic experiment: Loss-curve experiments compare empirical and theoretical value differences to assess whether neural optimization learns a good DeepMDP model.
  • Training setup: Training uses Adam for 30,000 steps, with discount factor γ = 0.9 and gradient-penalty coefficient λ = 0.01.
  • Atari architectures: Atari experiments evaluate single-layer fully connected, two-layer fully connected, and single convolutional transition models with FCLayer or ConvLayer latent states.
  • Atari experiments: Auxiliary-task experiments use a convolutional transition model and two-layer reward model on top of a C51 agent, varying transition-loss and reward-loss weights.

C.5. Representation Learning with DeepMDP Objectives

Minimizing DeepMDP objectives alone can collapse latent states, especially in complex environments where transition prediction dominates. Reconstruction supervision and larger batches help, but reconstruction is needed to reach baseline-level performance.

  • Latent transition-loss minimization collapses every latent state to zero because this perfectly predicts next latent states.The zero representation is optimal when only the latent transition loss is minimized.
  • Increasing batch size from 32 to 100 improves reward diversity, but reconstruction loss is needed to achieve baseline-level performance.The added reconstruction loss is the intervention that brings performance on par with the simple baseline.
  • In complex environments, transition loss appears to dominate, making it difficult to balance informative reward prediction against predictable latent dynamics.The authors conclude that additional work is required to balance the two losses.
  • DeepMDP optimization can become trapped in local minima where transition prediction favors uninformative representations and reward prediction remains poor.On DonutWorld, training eventually escaped this state, but the authors report that escape becomes harder as task difficulty increases.
  • Additional reconstruction supervision guides Atari representations toward explaining the environment's underlying dynamics.This provides the reported explanation for the observed benefit of reconstruction loss in the Arcade Learning Environment.
Loading 1906.02736v1…