SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zakaria Mhammedi MIT
mhammedi@mit.edu
Adam Block MIT
ablock@mit.edu
Dylan J. Foster
Microsoft Research anfoster@microsoft.com
Alexander Rakhlin  
MIT
A major challenge in reinforcement learning is to develop practical, sample-efficient algorithms for exploration in high-dimensional domains where generalization and function approximation is required. Low-Rank Markov Decision Processes—where transition probabilities admit a low-rank factorization based on an unknown feature embedding—offer a simple, yet expressive framework for RL with function approximation, yet existing algorithms either (1) are computationally intractable, or (2) require restrictive statistical assumptions such as latent variable structure or access to model-based function approximation. In this work, we propose the first provably sample-efficient algorithm for exploration in Low-Rank MDPs that is both computationally efficient and model-free, allowing for general function approximation while requiring no structural assumptions beyond a reachability condition that we show is substantially weaker than that assumed in prior work. Our algorithm, SpanRL, uses the notion of a barycentric spanner for the feature embedding as an efficiently computable basis for exploration, performing efficient spanner computation by interleaving representation learning and policy optimization subroutines. Our analysis—which is appealingly simple and modular—carefully combines several techniques, including a new approach to error-tolerant barycentric spanner computation, and a new analysis of a certain minimax representation learning objective found in prior work.
In reinforcement learning and control, many of the most promising application domains require the agent to navigate complex, high-dimensional state and action spaces, where generalization and function approximation is necessary. The last decade has witnessed impressive empirical success in domains where data are abundant [34, 38, 26, 28, 27], but when data are limited, ensuring efficient exploration in large domains is a major research question. For statistical efficiency, the foundations have recently begun to take shape, with a line of research providing structural conditions that facilitate sample-efficient exploration, as well as fundamental limits [37, 21, 39, 42, 14, 24, 17, 18]. Computational efficiency, however, remains a major challenge: outside of simple settings [7, 23], existing algorithms with provable sample complexity guarantees are computationally inefficient, and typically require solving intractable non-convex optimization problems [21, 11, 24, 10]. The prospect of developing practical algorithms for exploration in high-dimensional state spaces that are both computationally and statistically efficient raises three fundamental questions:
1. What are the right computational primitives for exploration? That is, how can one efficiently represent and compute exploratory policies that allow the learner to explore the state space and gather useful data?
2. How should one leverage function approximation—for example, via representation learning—to discover such primitives in a computationally and statistically efficient fashion?
3. Given answers to the first two questions, how can one efficiently interleave function approximation and exploration to provide provably efficient algorithms?
In this paper, we investigate these questions through the Low-Rank MDP model [36, 45, 1]. In a Low-Rank MDP, the state space is large and potentially continuous, but the transition probabilities admit an (unknown) low-rank factorization. Concretely, for a finite-horizon Low-Rank MDP with horizon  $H$ , the transition densities for layer  $h \in [H]$  satisfy
$$
T _ {h} \left(x _ {h + 1} \mid x _ {h}, a _ {h}\right) = \mu_ {h + 1} ^ {\star} \left(x _ {h + 1}\right) ^ {\top} \phi_ {h} ^ {\star} \left(x _ {h}, a _ {h}\right), \tag {1}
$$
where  $\phi_h^{\star}(\cdot ,\cdot)\in \mathbb{R}^d$  and  $\mu_{h + 1}^{\star}(\cdot)\in \mathbb{R}^{d}$  are state-action and next-state embeddings. The low-rank structure in (1) facilitates tractable exploration: if the embedding  $\phi_h^\star$  is known to the learner, one can efficiently learn a near-optimal policy with sample complexity polynomial in the feature dimension  $d$ , and independent of the size of the state space [23]; in this regard,  $\phi_h^\star$  can be thought of as a low-dimensional representation that enables sample-efficient RL. Following Agarwal et al. [1], we consider the challenging setting in which both  $\phi_h^\star$  and  $\mu_{h + 1}^{\star}$  are unknown to the learner. This formulation generalizes well-known frameworks such as the Block MDP (BMDP) model [12, 32], and necessitates the use of representation learning: the agent must learn an embedding that approximates  $\phi_h^\star$  as it explores the environment, and must use this learned embedding to drive subsequent exploration. This form of function approximation allows for great flexibility, as  $\phi_h^\star$  can be an arbitrary, nonlinear function of the state; in practice, it is common to model  $\phi_h^\star$  as a neural net [49].
The Low-Rank MDP is perhaps the simplest MDP structure that demands systematic exploration and nonlinear function approximation while allowing for a continuum of states, yet understanding of efficient algorithm design for this model is surprisingly limited. Existing algorithms suffer from at least one of the following drawbacks:
1. Computational intractability [21, 24, 14, 9, 43].
2. Strong modeling assumptions (e.g., ability to model  $\mu_{h+1}^{\star}(\cdot)$ , which facilitates application of model-based RL techniques) [1, 40, 10]; in this work, we aim for model-free methods that only require learning  $\phi_h^*$ .
3. Restrictive structural assumptions (e.g., non-negativity or latent variable structure for the embeddings in (1)) [35, 49].
At the root of these limitations is the complex interplay between exploration and representation learning: the agent must learn a high-quality representation to guide in exploring the state space, but learning such a representation requires gathering diverse and informative data, which is difficult to acquire without having already explored the state space to begin with. Overcoming this challenge—particularly where computational efficiency is concerned—requires (1) representation learning procedures that lead to sufficiently expressive representations for downstream applications, (2) efficient exploration procedures that are robust to errors in learned representations, and 3) understanding the interaction between these procedures, which must be interleaved. In this work, we propose an algorithm that addresses each of these challenges, as detailed below.
Contributions. We provide the first provably computationally efficient and model-free algorithm for general Low-Rank MDPs. Our algorithm, SpanRL, uses the notion of a barycentric spanner [6] for the embedding  $\phi_h^*$  as an efficiently computable basis for exploration, and combines this with a minimax representation learning approach [35, 49]. SpanRL interleaves exploration with representation learning in a layer-wise fashion, learning a new representation at each layer  $h$  using exploratory data gathered at previous layers, then uses this representation to facilitate computation of a collection of exploratory policies (a policy cover), which act as an approximate barycentric spanner for the features at layer  $h + 1$ , ensuring good coverage for subsequent iterations. SpanRL is simple and modular, and its analysis is surprisingly compact given the greater generality compared to prior work [49, 35, 31].
SpanRL can accommodate general-purpose function approximation to learn the representation  $\phi^{\star}$  (e.g., neural nets or other flexible classes) whenever a certain minimax representation learning objective [35, 49] can be solved efficiently for the function class of interest. Compared to efficient algorithms from prior work, SpanRL: (1) is model-free (i.e., only requires access to a function class  $\Phi$  capable of modeling  $\phi^{\star}$ , and does not need to model  $\mu_{h+1}^{\star}$ ), and (2) applies to general Low-Rank MDPs, replacing strong additional assumptions such as non-negativity of the feature embeddings (so-called latent variable structure) or block structure (see Table 1) with a reachability assumption that we show is substantially weaker than that assumed in prior work (see Appendix H). As a secondary benefit, the algorithm is reward-free. Our analysis carefully combines several new techniques, including (1) an error-tolerant variant of the classical barycentric spanner computation algorithm of Awerbuch
Table 1: Comparison of sample complexity required learn an  $\varepsilon$ -optimal policy.  $\Phi$  denotes the feature class, and  $\Upsilon$  denotes an additional feature class capturing model-based function approximation. For approaches that require non-negative (latent variable) structure,  $d_{\mathrm{LV}}$  [resp.  $\gamma$ ] denotes the latent variable dimension [resp. the reachability parameter in the latent representation], and for BMDPs,  $|S|$  denotes the size of the latent state space. For SpanRL,  $\eta$  denotes the reachability parameter.
Comp. efficient	Model-free	General low rank	Sample comp.
OLIVE [21] (see also [24, 14, 9, 43])	X	✓	✓	d3AH5log|Φ| /ε2
FLAMBE [1]	✓	X	✓1	d7A9H22log(|Φ||Y|) /ε10
Rep-UCB [40] (see also [10])	✓	X	✓	d4A2H5log(|Φ||Y|) /ε2
MOFFLE [35]2	✓	✓	X Non-negative/ latent variable	d19V A32H19log|Φ| / (ε6γ3 ∧ γ11)
BRIEE [49]	✓	✓	X Block MDP	|S|8A14H9log|Φ| / ε4
SpanRL (this paper)	✓	✓	✓	A4d9H4(d+log|Φ|) / ε2 ∧ η2
and Kleinberg [6], and (2) a new analysis of a minimax representation learning objective introduced in Modi et al. [35], Zhang et al. [49], which shows for the first time that this objective can lead to meaningful guarantees in general Low-Rank MDPs without latent variable structure; this increased generality is meaningful, as we show in Appendix H that there is an exponential separation between our guarantees and those that require such a structure.
Organization. Section 2 formally introduces the Low-Rank MDP model and the online reinforcement learning framework we consider. In Section 3, we highlight challenges faced by previous approaches, introduce our main algorithm, SpanRL, and show how it overcomes these challenges, and then present its main sample complexity guarantee.
Comparison to ArXiv Version. After the initial submission of this work, we developed a substantially improved of the algorithm that removes the reachability assumption at the cost of a larger (but still polynomial) sample complexity guarantee. We have included this algorithm and its analysis in the ArXiv version of this paper [30].
2.1 Low-Rank MDP Model
We work in an episodic, finite-horizon reinforcement learning framework, where  $H \in \mathbb{N}$  denotes the horizon. A Low-Rank MDP [36, 45, 1] is a tuple  $\mathcal{M} = (\mathcal{X},\mathcal{A},(\phi_h^\star)_{h\in [H]},(\mu_h^\star)_{h\in [H]},\rho)$  consisting of a state space  $\mathcal{X}$ , action space  $\mathcal{A}$  with  $|\mathcal{A}| = A$ , distribution over initial states  $\rho \in \Delta(\mathcal{X})$ , and mappings  $\mu_{h+1}^{\star}:\mathcal{X}\to \mathbb{R}^{d}$  and  $\phi_h^\star:\mathcal{X}\times \mathcal{A}\to \mathbb{R}^d$ . Beginning with  $x_1 \sim \rho$ , an episode proceeds in  $H$  steps, where for each step  $h \in [H]$ , the state  $x_h$  evolves as a function of the agent's action  $a_h$  via
$$
\boldsymbol {x} _ {h + 1} \sim T _ {h} (\cdot | \boldsymbol {x} _ {h}, \boldsymbol {a} _ {h}),
$$
where  $T_{h}$  is a probability transition kernel, which is assumed to factorize based on  $\phi_h^\star$  and  $\mu_h^\star$ . In detail, we assume that there exists a  $\sigma$ -finite measure  $\nu$  on  $\mathcal{X}$  such that for all  $1 \leq h \leq H - 1$ , and for all  $x \in \mathcal{X}$  and  $a \in \mathcal{A}$ , the function  $x' \mapsto \mu_{h + 1}^\star(x')^\top \phi_h^\star(x, a)$  is a probability density with respect to  $\nu$  (i.e. the function is everywhere non-negative and integrates to 1 under  $\nu$ ). For any  $\mathcal{X}' \subseteq \mathcal{X}$ , the probability that  $\boldsymbol{x}_{h + 1} \in \mathcal{X}'$  under  $\boldsymbol{x}_{h + 1} \sim T_h(\cdot | x_h, a_h)$  is then assumed to follow the law
$$
T _ {h} \left(\mathcal {X} ^ {\prime} \mid x _ {h}, a _ {h}\right) = \int_ {\mathcal {X} ^ {\prime}} \mu_ {h + 1} ^ {\star} (x) ^ {\top} \phi_ {h} ^ {\star} \left(x _ {h}, a _ {h}\right) \mathrm {d} \nu (x). \tag {2}
$$
For notational compactness, we assume (following, e.g., Jiang et al. [21]) that the MDP  $\mathcal{M}$  is layered so that  $\mathcal{X} = \mathcal{X}_1 \cup \dots \cup \mathcal{X}_H$  for  $\mathcal{X}_i \cap \mathcal{X}_j = \emptyset$  for all  $i \neq j$ , where  $\mathcal{X}_h \subseteq \mathcal{X}$  is the subset of states in  $\mathcal{X}$ .
For the stated sample complexity, FLAMBE requires access to a sampling oracle for the learner model. Without this oracle, the results require additional latent variable structure and a reachability assumption.
2We compare to the variant of MOFFLE that uses the same representation learning objective we consider. Other variants have improved sample complexity, but make use of stronger oracles.
3We emphasize that neither  $\mu_h^\star$  nor  $\phi_h^\star$  is known to the agent, in contrast to the linear MDP setting [44, 23].
that are reachable at layer  $h \in [H]$ . This can be seen to hold without loss of generality (modulo dependence on  $H$ ), by augmenting the state space to include the layer index.
Remark 2.1 (Comparison to previous formulations). Our formulation, in which the transition dynamics (2) are stated with respect to a base measure  $\nu$ , are a rigorous generalization of Low-Rank MDP formulations found in previous works [23, 1], which tend to implicitly assume the state space is countable and avoid rigorously defining integrals. We adopt this more general formulation to emphasize the applicability our results to continuous domains. However, in the special case where state space is countable, choosing  $\nu$  as the counting measure yields  $T_{h}(\mathcal{X}^{\prime}\mid x_{h},a_{h}) = \sum_{x\in \mathcal{X}^{\prime}}\mu_{h + 1}^{\star}(x)^{\top}\phi_{h}^{\star}(x_{h},a_{h})$ , which is consistent with prior work.
Policies and occupancy measures. We define  $\Pi_{\mathbb{M}} = \{\pi : \mathcal{X} \to \Delta(\mathcal{A})\}$  as the set of all randomized, Markovian policies. For a policy  $\pi \in \Pi_{\mathbb{M}}$ , we let  $\mathbb{P}^{\pi}$  denote the law of  $(x_1, a_1), \ldots, (x_H, a_H)$  under  $a_h \sim \pi(x_h)$ , and let  $\mathbb{E}^{\pi}$  denote the corresponding expectation. For any  $\mathcal{X}' \subseteq \mathcal{X}_h$ , we let  $\mathbb{P}_h^\pi[\mathcal{X}'] := \mathbb{P}^\pi[x_h \in \mathcal{X}']$  denote the marginal law of  $x_h$  under  $\pi$ . For  $x \in \mathcal{X}_h$ , we define the occupancy measure  $d^\pi(x) := \frac{\mathrm{d}\mathbb{P}_h^\pi}{\mathrm{d}\nu}(x)$  as the density of  $\mathbb{P}_h^\pi$  with respect to  $\nu$ .
2.2 Online Reinforcement Learning and Reward-Free Exploration
We consider a standard online reinforcement learning framework where the Low-Rank MDP  $\mathcal{M}$  is unknown, and the learning agent interacts with it in episodes, where at each episode the agent executes a policy of the form  $\pi : \mathcal{X} \to \Delta(\mathcal{A})$  and observes the resulting trajectory  $(x_1, a_1), \ldots, (x_H, a_H)$ . While the ultimate goal of reinforcement learning is to optimize a policy with respect to a possibly unknown reward function, here we focus on the problem of reward-free exploration, which entails learning a collection of policies that almost optimally "covers" the state space, and can be used to efficiently optimize any downstream reward function [12, 33, 15, 31]. To wit, we aim to construct an policy cover, a collection of policies that can reach any state with near-optimal probability.
Definition 2.1 (Policy cover). For  $\alpha \in (0,1]$ , a subset  $\Psi \subseteq \Pi_{\mathbb{M}}$  is an  $\alpha$ -policy cover for layer  $h$  if
$$
\forall x \in \mathcal {X} _ {h}, \quad \max  _ {\pi \in \Psi} d ^ {\pi} (x) \geq \alpha \cdot \max  _ {\pi^ {\prime} \in \Pi_ {\mathfrak {M}}} d ^ {\pi^ {\prime}} (x). \tag {3}
$$
We show (Appendix G) that given access to such a policy cover with constant  $\alpha$ , it is possible to optimize any downstream reward function with polynomial sample complexity.
Assumptions. To facilitate learning a policy cover, we make the following reachability assumption.
Assumption 2.1 ( $\eta$ -reachability). For any  $h \in [H]$  and  $x \in \mathcal{X}_h$ ,  $\max_{\pi \in \Pi_{\mathfrak{M}}} d^{\pi}(x) \geq \eta \cdot \| \mu_h^*(x)$ .
Reachability is necessary if one aims to build a policy cover that satisfies (3) uniformly for all states; without such a condition, one gives up on covering hard-to-reach states. Some notion of reachability is required in essentially all prior work on efficient model-free algorithms for Low-Rank MDPs [32, 35, 5], and was only very recently removed in the (more restrictive) BMDP setting [31, 49].
Remark 2.2 (Comparison to other reachability-like assumptions). Assumption 2.1 generalizes and subsumes all previous reachability-like conditions of which we are aware [33, 46, 1, 35]. Notably, reachability is implied by the notion of feature coverage [5] (used in the context of transfer learning in Low-Rank MDPs), which asserts that  $\sup_{\pi \in \Pi_{\mathfrak{M}}}\lambda_{\min}(\mathbb{E}^{\pi}[\phi_h^\star (\pmb {x}_h,\pmb {a}_h)\phi_h^\star (\pmb {x}_h,\pmb {a}_h)^\top ])\geq \eta$  for some  $\eta >0$ . It is also implied by explorability [46], which is similar to feature coverage, but involves the first moments of  $\phi_h^\star$ . Our reachability assumption is also weaker than that used in [1, 35] under the latent variable model, and generalizes that made for BMDPs [33]. See Appendix H for details, as well as an exponential separation between our assumptions and analogous assumptions in [1, 35].
Beyond reachability, we assume (following [1, 35]) for normalization that, for all  $h \in [H]$  and  $(x, a) \in \mathcal{X}_h \times \mathcal{A}$ ,  $\| \phi_h^*(x, a) \| \leq 1$ , and that for all  $g: \mathcal{X}_h \to [0,1]$ ,
$$
\left\| \int_ {\mathcal {X} _ {h}} \mu_ {h} ^ {*} (x) g (x) \mathrm {d} \nu (x) \right\| \leq \sqrt {d}. \tag {4}
$$
Function approximation and desiderata. We do not assume that the true features  $(\phi_h^\star)_{h\in [H]}$  or the mappings  $(\mu_h^\star)_{h\in [H]}$  are known to the learner. To provide sample-efficient learning guarantees we make use of function approximation as in prior work [3, 35], and assume access to a feature class  $\Phi \subseteq \{\phi : \mathcal{X} \times \mathcal{A} \to \mathbb{R}^d\}$  that contains  $\phi_h^\star$ , for  $h \in [H - 1]$ .
Assumption 2.2 (Realizability). The feature class  $\Phi \subseteq \{\phi : \mathcal{X} \times \mathcal{A} \to \mathbb{R}^d\}$  has  $\phi_h^\star \in \Phi$  for all  $h \in [H]$ . Moreover, for all  $\phi \in \Phi$ ,  $x \in \mathcal{X}$ , and  $a \in \mathcal{A}$ , it holds that  $\| \phi(x, a) \| \leq 1$ .
The class  $\Phi$  may consist of linear functions, neural networks, or other standard models depending on the application, and reflects the learner's prior knowledge of the underlying MDP. We assume that  $\Phi$  is finite to simplify presentation, but extension to infinite classes is straightforward, as our results only invoke finiteness through standard uniform convergence arguments. Note that unlike model-based approaches [1, 40, 10, 2], we do not assume access to a class capable of realizing the features  $\mu_h^{\star}$ , and our algorithm does not attempt to learn these features; this is why we distinguish our results as model-free.
For constant  $\alpha$ , our goal is to learn an  $\alpha$ -policy cover using  $\mathrm{poly}(d, A, H, \log |\Phi|, \eta^{-1})$  episodes of interaction. This guarantee scales with the dimension  $d$  of the feature map and the complexity  $\log |\Phi|$  of the feature class but, critically, does not depend on the size of the state space  $\mathcal{X}$ ; by [10], dependence on  $H$  and  $A = |\mathcal{A}|$  is necessary when  $\phi^{\star}$  is unknown. Given such a guarantee, we show in Appendix G how to optimize any downstream reward function to error  $\varepsilon$  with polynomial sample complexity.
Additional preliminaries. For any  $m, n \in \mathbb{N}$ , we denote by  $[m..n]$  the integer interval  $\{m, \ldots, n\}$ . We also let  $[n] := [1..n]$ . For any sequence of objects  $o_1, o_2, \ldots$ , we define  $o_{m:n} := (o_i)_{i \in [m..n]}$ . A partial policy is a policy defined over a contiguous subset of layers  $[\ell..r] \subseteq [H]$ . We denote by  $\Pi_{\mathbb{M}}^{\ell:r} := \{\pi : \bigcup_{h=\ell}^{r} \mathcal{X}_h \to \Delta(\mathcal{A})\}$  the set of all partial policies over layers  $\ell$  to  $r$ ; note that  $\Pi_{\mathbb{M}} \equiv \Pi_{\mathbb{M}}^{1:H}$ . For a policy  $\pi \in \Pi_{\mathbb{M}}^{\ell:r}$  and  $h \in [\ell..r]$ ,  $\pi(x_h)$  denotes the action distribution for the policy at layer  $h$  when  $x_h \in \mathcal{X}_h$  is the current state. For  $1 \leq t \leq h \leq H$  and any pair of partial policies  $\pi \in \Pi_{\mathbb{M}}^{1:t-1}, \pi' \in \Pi_{\mathbb{M}}^{t:h}$ , we define  $\pi \circ_t \pi' \in \Pi_{\mathbb{M}}^{1:h}$  as the partial policy given by  $(\pi \circ_t \pi')(x_\ell) = \pi(x_\ell)$  for all  $\ell < t$  and  $(\pi \circ_t \pi')(x_\ell) = \pi'(x_\ell)$  for all  $\ell \in [t..h]$ .
We use the  $x_{h} \sim \pi$  as shorthand to indicate that  $x_{h}$  is drawn from the law  $\mathbb{P}^{\pi}$ , and likewise for  $(\pmb{x}_h, \pmb{a}_h) \sim \pi$  and so on. For a set of partial policies  $\Psi \coloneqq \{\pi^{(i)} : i \in [N]\}$ , we define  $\mathrm{unif}(\Psi)$  as the random partial policy obtained by sampling  $i \sim \mathrm{unif}([N])$  and playing  $\pi^{(i)}$ . We define  $\pi_{\mathrm{unif}} \in \Pi_{\mathsf{M}}$  as the random policy that selects actions in  $\mathcal{A}$  uniformly at random at each layer. We use  $\| \cdot \|$  to denote the Euclidean norm,  $\| \cdot \|_{\infty}$  to denote the supremum norm on functions, and let  $\mathcal{B}(r) \subseteq \mathbb{R}^{d}$  denote the Euclidean ball of radius  $r$ . We refer to a scalar  $c > 0$  as an absolute constant to indicate that it is independent of all problem parameters and use  $\widetilde{O}(\cdot)$  to denote a bound up to factors polylogarithmic in parameters appearing in the expression.
3 SpanRL: Algorithm and Main Results
In this section, we present the SpanRL algorithm. We begin by describing challenges in deriving efficient, model-free algorithms using existing approaches (Section 3.1). We then formally describe SpanRL (Section 3.2) and build intuition as to how it is able to overcome these challenges, and finally state our main sample complexity guarantee (Section 3.3).
3.1 Challenges and Related Work
Designing algorithms with provable guarantees in the Low-Rank MDP setting is challenging because of the complicated interplay between representation learning and exploration. Indeed, while there are many efficient algorithms for the so-called linear MDP setting where the feature maps  $(\phi_h^\star)_{h\in [H]}$  are known (removing the need for representation learning) [23, 47, 4, 41], these approaches do not readily generalize to accommodate unknown features. For Low-Rank MDPs, previous algorithms suffer from at least one of the following three drawbacks: (1) the algorithms are computationally inefficient; (2) the algorithms are model-based; or (3) the algorithms place strong assumptions on the MDP that are unlikely to hold in practice. To motivate the SpanRL algorithm, we briefly survey these results, highlighting several key challenges in avoiding these pitfalls.