PaperlayerView exact revision on arXiv

Source-linked AI summary

FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs

Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, Wen Sun

arXiv:2006.10814v2cs.LGstat.ML

TL;DR

Representation learning for RL asks how to discover features that make decision-making sample-efficient, but learning those features is intertwined with exploration. This paper characterizes low rank MDPs through latent-variable structure and develops FLAMBE for efficient exploration and system identification, with policy optimization as a corollary.

  • Problem

    Representation learning for RL lacks mature foundations, and learning useful features is challenging because it is intimately tied to exploration.

  • Method

    The paper relates low rank transition dynamics to latent-variable models and develops FLAMBE, a model-based algorithm that learns representations through exploration and system identification.

  • Results

    FLAMBE provably learns low rank MDP models in a statistically and computationally efficient manner under the stated realizability and reachability assumptions.

  • Takeaways & Limitations

    Low rank MDPs provide a more expressive representation-learning framework than prior models, while learned features support policy optimization for known rewards.

  • Takeaways & Limitations

    Some computationally tractable planning results require a reachability assumption that imposes an upper bound on latent-variable dimension.

Abstract

from arXiv · show

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.

1 Introduction

The paper frames representation learning in RL around defining useful low-dimensional features and learning them despite their entanglement with exploration. It studies these questions through low rank MDPs and positions its structural and algorithmic contributions against prior representation-learning settings.

  • A good RL representation should make value functions and policies simple, supporting sample-efficient learning with parametric function approximation.
  • Learning representations is difficult because comprehensive environmental experience is needed to learn them, while good representations may be needed for efficient exploration.
  • The comparison framework organizes representation-learning methods from block MDPs through low rank, Bellman-rank, and Witness-rank settings, with differing assumptions and generality.
  • Low rank MDPs provide a framework for representation learning by modeling transitions through low-dimensional state-action and next-state embeddings.
  • The paper contributes structural results connecting low rank MDPs to latent-variable models and algorithmic results for learning unknown representations.

2 Low Rank MDPs

Low rank MDPs factor each transition operator through low-dimensional embeddings, making learned features useful for Bellman backups and policy optimization. The formulation combines reward-free system identification with function-class realizability assumptions.

  • A low rank transition operator satisfies T(x′ | x, a) = ⟨φ⋆(x, a), µ⋆(x′)⟩ for state-action and next-state embeddings in R^d.
  • Realizability assumes the true stage-specific embeddings belong to the candidate function classes Φ and Υ used for model selection.
  • Low rank structure alone is insufficient for state-space-independent performance guarantees, so additional modeling assumptions encoded in Φ and Υ are required.
  • The learning goal is reward-free system identification: learn a model that uniformly approximates the environment’s one-step dynamics under policies’ induced state-action distributions.
  • System identification quantitatively guarantees that the learned representation can approximate Bellman backups for any value function and data distribution.
  • Linear function approximation with the learned features can fit policy Q functions for explicitly given rewards and supports dynamic programming for policy optimization.

3 Related work

The related-work discussion places low rank MDPs among linear and nonlinear RL models, emphasizing their use for unknown-feature representation learning. It also relates them to block MDPs and computationally efficient or inefficient alternatives.

  • Low rank models have been studied broadly in dynamics and controls, as well as for planning, estimation, and generative-model reinforcement learning.
  • Prior low rank MDP results typically assume known feature maps, motivating this paper’s representation-learning formulation for richer state spaces.
  • More general nonlinear methods trade off strong environment assumptions, strong function-class assumptions, linear class-size dependence, or computational tractability.
  • Block MDPs use finite latent states with decodable observations, and prior work learns those latent states jointly with exploration efficiently.

4 Expressiveness of low rank MDPs

Low rank MDPs generalize latent-variable and block-MDP representations, sometimes achieving small embedding dimension despite requiring exponentially larger latent spaces. These connections also relate low rank structure to Bellman and Witness rank, though existing rank-based algorithms may be computationally intractable.

  • Latent variable representations: Low rank MDP transitions admit a latent-variable interpretation in which each (x, a) induces a distribution over z, followed by an emission distribution over x′.The resulting factorization uses simplex features ψ and ν, with latent-variable dimension dLV equal to the smallest feasible latent-space cardinality.
  • Block MDPs: Block MDPs impose disjoint emission supports, making them a restricted subclass of low rank MDPs.In block MDPs, latent variables correspond to decodable latent states and partition the next-state space.
  • Block MDPs: For any d ≥2 and M, some M-state environments have low rank embedding dimension d but require M latent states in every block-MDP representation.Thus, block-MDP latent-state complexity can scale with the full state-space size even when low rank remains fixed.
  • Simplex features: For any even n, some MDPs have embedding dimension O(n^2) but latent-variable dimension dLV ≥2^Ω(n), showing latent-variable representations can be exponentially weaker.This separation follows because dLV corresponds to non-negative rank, which can greatly exceed ordinary rank; negative-valued embedding functions add expressiveness.
  • Bellman and Witness rank: Low rank MDPs have Bellman rank at most d and can yield low Witness rank, but OLIVE and the analogous Witness-rank algorithm are computationally intractable because they enumerate function classes.These results motivate the subsequent development of computationally tractable algorithms.

5 Main results

The paper develops FLAMBE for exploration and representation learning in low rank MDPs, using computational oracles to identify a model efficiently. Its guarantees cover richer models than block MDPs, while highlighting trade-offs involving latent dimension, actions, reachability, and sampling assumptions.

  • Algorithm design: FLAMBE iteratively fits transition embeddings with maximum likelihood and plans exploratory policies to collect increasingly informative transition data.The algorithm updates a learned model from accumulated trajectories, then uses model-based planning to construct the next exploratory policy.
  • Algorithm design: The model-based planner requires no environment interaction and can be implemented efficiently with the sampling oracle.The planner optimizes over policies because Bellman backups in low rank MDPs are linear in the learned features.
  • Guarantees: FLAMBE learns low rank MDP models with statistical and computational efficiency under Assumptions 1 and 2, with policy optimization as a corollary.The paper states that the result has comparable dependencies on d, H, K, and ε to prior work.
  • Conceptual advances: FLAMBE applies to a significantly richer class of models than block MDP algorithms and does not require reachability assumptions in its main setting.The paper identifies these as key advances over prior block MDP results.
  • Comparisons and limitations: Relative to prior low rank results, FLAMBE addresses unknown embeddings, but its representation-learning bound scales with the number of actions K.The authors believe additional structural assumptions on Φ are needed to remove this dependence.
  • Comparisons and limitations: Compared with some more general theories, FLAMBE gains computational efficiency but covers a less general model class and has worse polynomial sample-complexity factors.This comparison is stated relative to Jiang et al. (2017) and Sun et al. (2019).
  • Guarantees: The simplex-factorization guarantee trades dependence on embedding dimension d for polynomial dependence on latent dimension dLV while scaling more favorably with H, K, and ε.This version requires simplex constraints in the function class Φ.
  • Real-world planning: An environment-interaction variant avoids SAMP but requires reachability and has sample complexity polynomial in dLV, H, K, 1/ηmin, 1/ε, and log(|Φ||Υ|/δ).The result uses H calls to MLE and is presented for simplex representations.

6 Discussion

The paper interprets low rank MDPs through latent-variable representations, establishes their greater expressiveness, and develops FLAMBE for efficient system identification. It identifies model-free representation learning and empirical validation as future directions.

  • Discussion: The paper gives low rank MDPs an intuitive latent-variable interpretation and proves structural results showing greater expressiveness than prior representation-learning models.These results frame low rank MDPs as a broader setting for representation learning in RL.
  • Discussion: FLAMBE provides a computationally and statistically efficient model-based algorithm for system identification in low rank MDPs.Policy optimization follows as a corollary of the learned model guarantee.
  • Future work: Future work includes provably efficient model-free representation learning and practical algorithms for complex reinforcement-learning tasks.The paper presents both theoretical and empirical extensions as open directions.

A Proofs for the structural results

This appendix proves the paper’s structural results, beginning with Lemma 1 and then addressing realizability, reachability, separation, and rank connections.

  • The appendix provides proofs for the paper’s structural results.It begins with Lemma 1 and organizes subsequent proofs by topic.
  • Lemma 1: The proof of Lemma 1 fixes a stage and a bounded value function while using the low-rank MDP property.The notation suppresses stage dependence during the proof.
  • Lemma 1: The proof then combines its intermediate result with equation (1) for any policy.

A.2 On realizability and reachability

The appendix establishes realizability and reachability properties, and proves that low-rank structure alone cannot guarantee state-independent sample complexity.

  • Lower bound: Rank-2, horizon-2 low-rank MDPs can require Ω(M) samples to find a near-optimal policy with moderate probability.For n ≤ M/2 trajectories, some instance leaves the output policy at least 1/8-suboptimal with probability at least 1/8.
  • Lower bound: The lower bound embeds binary classification into a one-step MDP whose two actions lead deterministically to good or bad terminal states.The instance is indexed by a binary vector specifying the good action at each state.
  • Lower bound: The collected trajectories are information-theoretically equivalent to uniformly sampled state indices paired with their binary labels.
  • Lower bound: The classification lower bound transfers because predictor error equals policy suboptimality under the construction.
  • Reachability: For latent-state reachability, the analysis rewrites policy-induced latent distributions through action-selection matrices and the low-rank dynamics.It iteratively selects a sign vector while peeling off reachable latent variables, gaining (ηmin/K)^2 per reachable variable.
  • Reachability: When all dLV latent variables are reachable, the iterative construction yields a lower bound accumulated across all latent variables.

A.3 Separation results

The separation results show that low matrix rank does not ensure compact block-MDP representations and can coexist with exponentially large simplex-feature dimension.

  • Block-MDP separation: A rank-2 transition operator can require N latent states in any block-MDP representation.The construction makes all N next-state transition columns pairwise linearly independent, while block membership would require pairwise dependence.
  • Block-MDP separation: The construction uses a horizon-2 MDP with one initial state, two actions, and N possible next states.Its transition operator has rank 2.
  • Simplex-feature separation: The perfect matching polytope has exponentially many vertices and exponentially many constraints.Its vertices correspond to perfect matchings, while the constraints arise from odd-sized subsets.
  • Simplex-feature separation: The associated slack matrix has polynomial ordinary rank but non-negative rank at least 2^Ω(n).The non-negative-rank lower bound follows from the polytope’s extension complexity and its equivalence to slack-matrix non-negative rank.
  • Simplex-feature separation: Positive diagonal rescaling preserves non-negative rank, so the transition matrix inherits the exponential lower bound.For row-stochastic matrices, this lower bound equals the smallest internal dimension of a row-stochastic factorization.

A.4 On Bellman and Witness rank

The appendix connects low-rank MDPs to Bellman and Witness rank, showing low rank and realizability for suitable function classes under explicit modeling assumptions.

  • Bellman rank: Under φ⋆0:H−1 ∈ Φ, OLIVE has a stated sample-complexity guarantee for known rewards.
  • Connections: The analysis treats low Bellman and Witness rank as concrete connections between low-rank MDPs and prior algebraic RL frameworks.
  • Bellman rank: In a low-rank MDP, the average Bellman error admits a factorization through d-dimensional representations.The associated normalization bounds include ∥νh(π)∥2 ≤ 1 and ∥ξh((g,π))∥2 ≤ (2B+1).
  • Bellman rank: The Bellman rank is at most d with normalization parameter O(Bd) for bounded function classes.
  • Bellman rank: The optimal Q function is linear in the true embedding functions, so realizability holds for the induced value and policy classes.For bounded value functions, the linear coefficient vector has norm bounded by H.
  • Witness rank: The Witness rank is at most d, and the model-based algorithm has the stated sample complexity under the corresponding realizability conditions.The witness misfit measures one-step total-variation error on a policy-induced data distribution.

B Analysis of FLAMBE

FLAMBE’s analysis combines maximum-likelihood model estimation with exploratory policies and a covariance-based potential argument. Simulation and escaping-probability lemmas then yield accuracy and sample-complexity guarantees for sparse rewards.

  • Model estimation: The learned model’s accuracy is controlled by second-moment matrices induced by previous exploratory policies.These matrices also define the covariance quantities used in the potential argument.
  • Model estimation: With probability 1 −δ, maximum-likelihood estimation produces feature and transition estimates satisfying the stated error guarantee under data collected by exploratory roll-ins and uniform actions.Each iteration contributes n transition triples for the relevant time step.
  • Simulation argument: The absorbing MDP terminates trajectories that leave the known state-action set, allowing estimated-model error to be related to escaping probability.The learned model approximates the absorbing MDP on the known set, while value differences depend on the probability of escape.
  • Exploration progress: If an exploratory policy has large escaping probability, FLAMBE identifies a time step where the corresponding unknown-set indicator has large expected value.This connects exploration failure directly to the potential-based progress argument.
  • Final guarantee: After the prescribed number of iterations, the final model satisfies a value guarantee for any sparse reward function and policy.The guarantee is established under the per-round estimation-error assumptions used throughout the analysis.
  • Final guarantee: The analysis derives a final sample-complexity expression after balancing estimation and exploration-error terms.The bound is stated while ignoring logarithmic factors.

B.1 Refined analysis for simplex representations.

For simplex representations, FLAMBE tracks coverage over latent variables rather than general feature covariance. Its planner activates latent coordinates, adds sufficiently covered variables to the known set, and obtains refined guarantees.

  • Latent-variable coverage: The simplex analysis defines coverage through cumulative probabilities of reaching latent variables and marks variables with probability at least ∆ as known.The quantity used for coverage is a sum of probability distributions, not itself a distribution.
  • Estimation: Maximum-likelihood estimation provides high-probability error control from datasets collected by exploratory roll-ins followed by uniformly random actions.The guarantee applies for each time step and iteration under the defined latent-variable coverage measure.
  • Latent-variable coverage: A large escaping probability implies that some previously unknown latent variable is reached with probability at least ∆.This pigeonhole step is the simplex analogue of the general known-set exploration argument.
  • Guarantees: After Jmax = HdLV + 1 iterations, the final model satisfies the stated sparse-reward value guarantee.The iteration bound follows because each sufficiently exploratory round adds a latent variable to the known set.
  • Guarantees: With a sampling oracle, the simplex-feature guarantee is preserved up to constants using polynomially many oracle calls.The additional optimization error is absorbed into the existing additive error term.
  • Exploratory planning: The latent-coordinate planner forms a mixture of policies, with each component targeting one coordinate of the learned simplex feature map.Each component can be computed using dynamic programming such as LSVI.

C.3 Elliptical planner

The elliptical planner iteratively searches for policies that produce large quadratic forms under an inverse covariance matrix, updating a policy mixture until exploration is sufficient. Its sample-based version preserves the qualitative guarantees.

  • Elliptical planning: Each iteration searches for a policy maximizing quadratic forms involving the inverse current covariance matrix.A policy that cannot achieve a sufficiently large form triggers termination; otherwise it is added to the mixture.
  • Elliptical planning: The exact planner terminates after at most T + 1 iterations, with T ≤4d log(1 + 4/β)/β.Its output mixture satisfies the performance guarantee stated in Lemma 17.
  • Sampling version: The sample-based planner terminates after T + 1 iterations with T ≤O(d log(1 + 1/β)/β).It outputs an estimated covariance matrix and a policy with probability at least 1 −δ.
  • Sampling version: Setting εopt = O(β) leaves both the iteration complexity and performance guarantee unchanged.The empirical covariance approximation remains accurate because it averages accurate policy-specific covariance estimates.
  • Forward variant: FLAMBE.F learns dynamics one time step at a time, collecting exploratory data, applying MLE, and passing learned features and policies to a subsequent planner.This forward procedure uses environment interaction rather than a model sampling oracle.
  • Forward variant: The forward analysis maintains a coverage condition comparing the probability of each latent state under an arbitrary policy with that under the exploratory policy.The stated condition bounds the former by κ times the latter.

D.1 Optimizing a fixed reward function

To optimize a fixed reward, the environment planner uses learned features and exploratory coverage within LINEAR-FQI. Bellman-backup approximation and least-squares generalization then support a value guarantee.

  • LINEAR-FQI: LINEAR-FQI uses learned feature maps and exploratory policies to perform least-squares dynamic programming for the fixed reward.The algorithm proceeds backward from the final time step and returns a greedy policy.
  • Bellman approximation: The learned-model Bellman backup is linear in the estimated feature map, with bounded value functions enabling statistical error control.The bounded range is ˆVh+1(xh+1) ∈[0, H].
  • Statistical control: With probability at least 1 −δ, the squared-error bound holds simultaneously for all time steps after a union bound over h ∈[H].The analysis uses ∆n := 24H2d 2(d log n + log(2H/δ))/n.
  • Statistical control: The coverage condition transfers squared prediction error from the exploratory distribution to distributions induced by other policies.This transfer is used in the induction underlying the planning guarantee.
  • Value guarantee: The learned policy is compared with the optimal policy through greediness with respect to the learned Q function and a telescoping value-difference argument.The resulting bound combines the learned-model error terms across the horizon.
  • Value guarantee: The sample size is chosen so the ∆n contribution is at most H2εTV, balancing statistical and model-estimation errors.A sufficient condition yields the stated result.

D.2 Instantiating the reward functions.

The section establishes FLAMBE's inductive guarantees by learning simplex features and measures, constructing exploratory policies with LINEAR-FQI, and transferring distributional guarantees through importance weighting. It then derives sample complexity contributions from MLE and LINEAR-FQI and invokes a realizable sequential MLE theorem, with an alternative generative-adversarial objective also discussed.

  • Inductive estimation: FLAMBE inductively learns feature and measure estimates using Theorem 21, producing an estimated transition model ˆT = ⟨ˆφ_h, ˆµ_h⟩ with probability at least 1 − δ.The argument identifies this as the step using the optimization oracle and MLE.
  • Exploration: At each iteration, LINEAR-FQI produces dLV policies that approximately maximize reaching individual coordinates of the learned feature representation.The exploratory policy is formed by using these policies and induces a distribution over the next state.
  • Distribution transfer: Importance weighting transfers guarantees from arbitrary policies to the exploratory policy, yielding π P[z_h+1 = z] ≤ dLVK · P_ρh+1[z_h+1 = z] + dLVεstat + (1 + dLV)√εTV.The transfer relies on the nonnegative error and measure terms together with randomized actions in the training policy.
  • Sample complexity: The calls to MLE require nH samples, while the calls to LINEAR-FQI require nHdLV samples in total.The section combines these components to obtain the total sample complexity.
  • MLE guarantee: Theorem 21 provides a high-probability sequential MLE guarantee under a finite realizable function class, and its proof establishes squared Hellinger convergence that implies the stated total-variation guarantee.The analysis uses a tangent-sequence decoupling argument and relates total variation to Hellinger distance.
  • Alternative estimation: A discriminator-based generative-adversarial objective is presented as an alternative to MLE when the discriminator class is sufficiently expressive.The stated sufficient condition is inclusion of indicators for the Scheffe sets of all model pairs.
Loading 2006.10814v2…