[Paper Review] FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs

Alekh Agarwal, Sham M. Kakade+2 more|arXiv (Cornell University)|Jun 18, 2020

Advanced Bandit Algorithms Research58 references36 citations

TL;DR

FLAMBE studies representation learning for low rank MDPs, proving that learned embeddings enable sample-efficient, model-based RL with provable guarantees, and introducing an exploration-driven algorithm for learning representations without known features.

ABSTRACT

In order to deal with the curse of dimensionality in reinforcement learning (RL), it is common practice to make parametric assumptions where values or policies are functions of some low dimensional feature space. This work focuses on the representation learning question: how can we learn such features? Under the assumption that the underlying (unknown) dynamics correspond to a low rank transition matrix, we show how the representation learning question is related to a particular non-linear matrix decomposition problem. Structurally, we make precise connections between these low rank MDPs and latent variable models, showing how they significantly generalize prior formulations for representation learning in RL. Algorithmically, we develop FLAMBE, which engages in exploration and representation learning for provably efficient RL in low rank transition models.

Motivation & Objective

• Motivate representation learning for RL under low rank dynamics to mitigate the curse of dimensionality.
• Formalize the relationship between low rank MDPs and latent variable models to justify representation learning.
• Develop Flambe, a model-based exploration algorithm that learns embeddings and achieves reward-maximizing policies.
• Provide statistical and computational guarantees for learning accurate environment models and representations under realizability.

Proposed method

• Formulate low rank MDPs with embedding functions φ and μ such that T(x'|x,a)=<φ(x,a), μ(x')> and assume bounded norms for φ and μ.
• Introduce two function classes Φ and Υ and a realizability assumption φ*_h ∈ Φ and μ*_h ∈ Υ to enable tractable learning.
• Define a reward-free exploration objective and a uniform accuracy criterion for learned model  M̂ = (φ̂, μ̂) in reproducing one-step dynamics across policies.
• Propose Flambe, an iterative algorithm that collects transitions, applies a maximum likelihood oracle (Mle) to learn embeddings, and uses a model-based planner to design exploratory policies.
• Develop a planning subroutine inspired by elliptical potentials to maximize coverage in learned feature space and build a mixture exploratory policy.
• Present two algorithmic variants: a general low rank MDP version and a simplex-factorization version with differing sample complexities.

Experimental results

Research questions

• RQ1 Can low rank MDPs be effectively learned through representation learning when the embeddings are unknown?
• RQ2 How can exploration be guided to achieve comprehensive coverage of the learned feature space in order to accurately identify the transition dynamics?
• RQ3 What are the statistical guarantees for learning both the environment model and the representation under realizability?
• RQ4 How do the expressiveness and limitations of low rank MDPs compare to block MDPs and latent variable models in terms of learnability and planning?
• RQ5 What are the computational and sample complexity implications of learning representations under different factorization assumptions (embedding dimension d vs latent dimension dLV)?

Key findings

• Flambe achieves statistically and computationally efficient learning of low rank MDP models under realizability.
• Low rank MDPs are more expressive than block MDPs, and Flambe leverages a latent-variable perspective to enable representation learning.
• Under Assumption 1 (realizability) and a model-based planner, Flambe attains a uniform-accuracy learned model with poly-time complexity and a trajectory count of roughly H^22 K^9 d^7 / ε^10 (up to logs).
• A variant of Flambe operating on simplex factorization yields improved dependence on horizon H, action space K, and ε, with complexity scaling as H^11 K^5 dLV^5 / ε^3 (up to logs).
• Reachability-based refinements show that, with a reachability assumption, the latent-variable dimension dLV is bounded and Flambe can operate without sampling or with reduced sampling requirements.