PaperlayerView exact revision on arXiv

Source-linked AI summary

Representation Learning for Online and Offline RL in Low-rank MDPs

Masatoshi Uehara, Xuezhou Zhang, Wen Sun

arXiv:2110.04652v3cs.LGcs.AIstat.ML

TL;DR

RL representation learning asks how to obtain a compact representation that supports sample-efficient decision-making, but low-rank MDP representations are not known in advance. This paper proposes REP-UCB for online RL and REP-LCB for offline RL. REP-UCB improves sample complexity over FLAMBE, while REP-LCB provides a guarantee under partial coverage.

  • Problem

    In low-rank MDPs, the transition representation is unknown, and learning a representation that supports efficient RL remains a theoretically under-investigated problem.

    Source: §1 · p. 1, §6 · p. 11

  • Method

    The paper develops REP-UCB, which interleaves representation learning, exploration, and exploitation, and REP-LCB, which uses pessimism for offline RL under partial coverage.

    Source: §1 · pp. 2–3, §4.1 · p. 7

  • Results

    REP-UCB improves sample complexity over FLAMBE in all parameters using the same computational oracle, and REP-LCB has a PAC guarantee under partial coverage measured by the true representation’s relative condition number.

    Source: §1 · p. 2, §6 · pp. 11–12

  • Takeaways & Limitations

    The algorithms provide sample-efficient online and offline learning guarantees for low-rank MDPs without assuming the ground-truth feature representation is known in advance.

    Source: §6 · pp. 11–12

    The offline guarantee assumes a finite relative condition number ω, an assumption described as weaker than requiring global coverage 1/ρ(s,a) < ∞ for every state-action pair.

    Source: §6 · p. 11

Abstract

from arXiv · show

This work studies the question of Representation Learning in RL: how can we learn a compact low-dimensional representation such that on top of the representation we can perform RL procedures such as exploration and exploitation, in a sample efficient manner. We focus on the low-rank Markov Decision Processes (MDPs) where the transition dynamics correspond to a low-rank transition matrix. Unlike prior works that assume the representation is known (e.g., linear MDPs), here we need to learn the representation for the low-rank MDP. We study both the online RL and offline RL settings. For the online setting, operating with the same computational oracles used in FLAMBE (Agarwal et.al), the state-of-art algorithm for learning representations in low-rank MDPs, we propose an algorithm REP-UCB Upper Confidence Bound driven Representation learning for RL), which significantly improves the sample complexity from $\widetilde{O}( A^9 d^7 / (\epsilon^{10} (1-\gamma)^{22}))$ for FLAMBE to $\widetilde{O}( A^2 d^4 / (\epsilon^2 (1-\gamma)^{5}) )$ with $d$ being the rank of the transition matrix (or dimension of the ground truth representation), $A$ being the number of actions, and $\gamma$ being the discounted factor. Notably, REP-UCB is simpler than FLAMBE, as it directly balances the interplay between representation learning, exploration, and exploitation, while FLAMBE is an explore-then-commit style approach and has to perform reward-free exploration step-by-step forward in time. For the offline RL setting, we develop an algorithm that leverages pessimism to learn under a partial coverage condition: our algorithm is able to compete against any policy as long as it is covered by the offline distribution.

1 Introduction

The paper studies representation learning for online and offline RL in low-rank MDPs, where the transition representation is unknown. It introduces REP-UCB for online learning and REP-LCB for offline learning under partial coverage.

Motivation

  • Representation learning in RL remains theoretically subtle because linearizing the optimal policy or value function alone may still leave learning exponentially sample-intensive.The paper argues that efficient RL representations need to encode more information about the underlying MDP.Source: §1 · p. 1

Problem setting

  • A low-rank MDP has an unknown transition factorization P(s′|s,a) = μ(s′)ᵀφ(s,a), so its representation must be learned rather than assumed known.The representation linearizes both the transition operator and the optimal state-action value function.Source: §1 · p. 2

Online setting

  • REP-UCB significantly improves FLAMBE’s sample complexity from O(d^7 A^9/(ϵ^10(1 − γ)^22)) to O(d^4 A^2/(ϵ^2(1 − γ)^5)), where d is transition-matrix rank, A actions, ϵ the suboptimality gap, and γ∈[0,1) the MDP discount factor.REP-UCB integrates representation learning, exploration, and exploitation by combining the reward signal with an exploration bonus constructed using the latest learned representation.Source: §1 · pp. 2–3
  • REP-UCB combines reward and an exploration bonus based on the latest learned representation, sharing data across time steps.Unlike FLAMBE’s layer-by-layer explore-then-commit strategy, REP-UCB balances exploration and exploitation as it learns.Source: §1 · p. 2

Offline setting

  • REP-LCB uses pessimism and a relative condition number under the true representation to compete with any offline-covered policy, including history-dependent policies.It uses one MLE computation-oracle call and is presented as the first oracle-efficient offline low-rank-MDP algorithm with this guarantee.Source: §1 · pp. 2–3

2 Related Work

Prior work on low-rank MDPs spans computationally inefficient version-space methods, oracle-efficient approaches with additional assumptions, and offline methods that often assume a known representation. The paper positions its algorithms against these differences in computation, assumptions, and coverage guarantees.

Online setting

  • FLAMBE uses the same oracle and function-approximation setting as REP-UCB, but its explore-then-commit strategy does not balance exploration and exploitation and has worse sample complexity.The paper’s comparison states that REP-UCB improves sample complexity in all parameters.Source: §2 · p. 3
  • Moffle is oracle-efficient but assumes low non-negative rank and latent-state reachability, restricting its scope and potentially making its complexity depend polynomially on non-negative rank.The cited discussion notes that non-negative rank can be exponentially larger than ordinary rank.Source: §2 · pp. 3–4
  • OLIVE, Witness rank, and Bilinear-UCB can have tighter dependence on d in low-rank MDPs, but their version-space algorithms are computationally inefficient.The cited work also notes that solving low-rank MDPs with a policy class can require Ω(2^d) samples.Source: §2 · p. 4

Offline setting

  • Earlier offline pessimistic guarantees rely on a known representation, while Uehara and Sun report similar low-rank-MDP statistical results with a slightly tighter bound but a version-space algorithm.The paper says guarantees for Xie et al.’s method are unclear on true-representation relative condition numbers and history-dependent policies.Source: §2 · p. 4

3 Preliminaries

The preliminaries define low-rank MDPs through a transition factorization with unknown embeddings, then specify the online and offline learning settings and modeling tools.

Modeling and oracles

  • The setup assumes known rewards and initial distribution, represents unknown embeddings with function classes, and uses MLE and known-linear-MDP planning procedures.Source: §3 · pp. 4, 6

Low-rank MDPs

  • For every state s, next state s′, and action a, the transition probability P ⋆(s′ | s, a) factors as µ⋆(s′)⊤φ⋆(s, a).Source: §3 · p. 5
  • Low-rank MDPs capture latent-state models; block MDPs are the special case where φ⋆ is one-hot, while linear MDPs assume φ⋆ is known.Source: §3 · p. 5

Learning settings

  • Online learning uses episodic interaction from d0, whereas offline learning uses a static dataset and targets policies covered by its distribution.The partial coverage condition need not cover every policy’s state-action distribution.Source: §3 · pp. 5–6

4 Representation Learning in Online Setting

REP-UCB interleaves representation learning, exploration, and exploitation in online low-rank MDPs, with an oracle-efficient analysis that handles changing learned representations.

4 Representation Learning in Online Setting

  • REP-UCB carefully interleaves representation learning, exploration, and exploitation to find a near-optimal policy from episodic resets.Source: §4 · p. 6

4.1 Algorithm

  • Each iteration collects a tuple, learns a transition factorization by MLE, sets a bonus from the learned feature, then plans with the bonus-enhanced reward.The data collection roll-in terminates in about 1/(1−γ) steps with high probability.Source: §4.1 · pp. 6–7

4.2 Analysis

  • REP-UCB achieves the best sample complexity among oracle-efficient algorithms for low-rank MDPs, while remaining statistically and oracle efficient.The supplied theorem statements do not provide the bound’s full expression.Source: §4.2 · pp. 7–8
  • The analysis obtains optimism under nonlinear MLE approximation and tracks regret despite representations that change during learning.It uses a potential function based on the fixed ground-truth representation for tracking progress.Source: §4.2 · pp. 9–10

5 Representation Learning in Offline Setting

REP-LCB combines a learned model, a learned-feature penalty, and planning to compete with policies covered by offline data, under a relative condition number defined using the true representation.

5 Representation Learning in Offline Setting

  • REP-LCB learns a model by MLE, constructs a linear-bandit-style penalty from its learned representation, and plans with the learned model and penalized reward.Source: §5 · p. 10

5.1 Analysis

  • REP-LCB can compete uniformly with any policy, including history-dependent policies, when the policy satisfies the partial coverage condition.If the optimal policy is covered, this guarantee applies without requiring global coverage of all policies.Source: §5.1 · p. 11
  • The coverage measure is a relative condition number defined using the unknown ground-truth representation, rather than the learned features.The measure can be smaller than a density ratio in MDPs with large state spaces.Source: §5.1 · pp. 10–11
  • REP-LCB is described as oracle-efficient without knowing the true representation, while requiring partial coverage only in terms of that representation.Source: §5.1 · p. 11

6 Conclusion

The paper studies online and offline RL in low-rank MDPs when the ground-truth representation is unknown. It presents REP-UCB for online learning and REP-LCB for offline learning under partial coverage.

6 Conclusion

  • REP-UCB improves sample complexity over FLAMBE in all parameters while using the same computational oracles.The paper describes it as having the best sample complexity among existing oracle-efficient algorithms for low-rank MDPs.Source: §6 · pp. 11–12
  • REP-UCB interleaves representation learning, exploration, and exploitation rather than using forward, step-by-step reward-free exploration.It learns a single stationary policy.Source: §6 · p. 12
  • REP-LCB gives an oracle-efficient PAC guarantee under partial coverage measured by the relative condition number defined with the true feature representation.The paper states that its condition requires finite ω, which is weaker than global coverage requiring 1/ρ(s,a) < ∞ for every state-action pair.Source: §6 · pp. 11–12

A Proof of the theoretical property of REP-UCB

The REP-UCB proof relates its bonus to learned-model error and uses this analysis to establish a PAC guarantee. The resulting bound is sufficient for PAC, although the proof does not yield a regret guarantee.

A Proof of the theoretical property of REP-UCB

  • At the initial distribution, the proof compares the expected bonus with the expected learned-model error.The passage states that these quantities are of the same order under the proof’s analysis.Source: A Proof of the theoretical property of… · p. 16
  • The derived bound has a regret-like form but does not guarantee regret because the algorithm plays a uniform action at the end of each episode.The passage says this bound nevertheless suffices for the PAC guarantee.Source: A Proof of the theoretical property of… · p. 18
  • The proof obtains a PAC guarantee for REP-UCB after bounding the required number of environment interactions.The interaction count is O(N(1 − γ)^−1 ln(N/δ)) when N is the number of episodes used by the guarantee.Source: A Proof of the theoretical property of… · p. 20

B Proof of the theoretical property of REP-LCB

The REP-LCB proof establishes a PAC guarantee for any comparator policy, including history-dependent policies. Its analysis uses a one-step-back inequality to relate learned-model expectations to a potential function.

B Proof of the theoretical property of REP-LCB

  • The offline proof uses a concentration event that holds with probability 1 − δ and sets ζ = O(ln(|M|/δ)/n).The passage notes that ln n can be removed because n is fixed in the offline setting.Source: B Proof of the theoretical property of… · p. 24
  • With probability 1 − δ, REP-LCB’s PAC guarantee applies to any comparator policy, including history-dependent non-Markovian policies.The theorem is stated with parameters set as in Lemma 14.Source: B Proof of the theoretical property of… · p. 25
  • The proof’s one-step-back inequality relates expectations under the learned model to a potential function paired with the learned representation.It instantiates the function as the learned transition model’s error relative to the true model.Source: B Proof of the theoretical property of… · p. 26

C Auxiliary lemmas

The auxiliary lemmas support the paper’s analysis by supplying a maximum-likelihood estimation guarantee, tools for distribution shift, and a conversion from finite-sample error bounds to sample complexities.

C Auxiliary lemmas

  • The MLE guarantee is presented as a component of the online algorithm and is stated to extend immediately to the offline case.The proof is referred to Agarwal et al. (2020b, Theorem 21).Source: C Auxiliary lemmas · p. 28
  • The simulation lemma compares two MDPs, one with reward r + b and another with reward r, for any policy.Source: C Auxiliary lemmas · p. 29
  • The distribution shift lemma is introduced to handle distribution shift in the offline setting.It is stated for any policy, state-action distribution, and true representation.Source: C Auxiliary lemmas · p. 29
  • A separate auxiliary lemma converts finite-sample error bounds into sample complexities.Source: C Auxiliary lemmas · p. 29

D More comparison to Xie et al. (2021)

The appendix compares a PAC guarantee obtained via Xie et al. (2021)'s Algorithm 1 with Theorem 6, noting stronger partial coverage and worse n dependence for the comparison. Unlike that comparison, Theorem 6 guarantees competitiveness against any history-dependent policy.

D More comparison to Xie et al. (2021)

Loading 2110.04652v3…