Efficient Model-Free Exploration in Low-Rank MDPs
Authors: Zakaria Mhammedi
Adam Block
Dylan J Foster
Alexander Rakhlin
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement learning
Representation Learning
Low-rank MDPs
Model-Free Learning

Abstract: 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.

Efficient Model-Free Exploration in Low-Rank MDPs
Zakaria Mhammedi MIT mhammedi @mi t. edu
Adam Block
Dylan J. Foster Microsoft Research dyllanfosterrmnicroofft.u
Alexander Rakhlin MIT rakhlin@mi edu
MIT ablock@mit edu
Abstract
A major challenge in reinforcement learning is to develop practical, sample-efficient algorithms for exploration function approximation is required. Low-Rank Markov Decision Processes Processes-where transition probabilities admit a low-rank factorization based on an unknown feature embedding- offer simple, yet expressive framework for RI. with function approx- imation, yet existing algorithms either (1) are computationally intractable, or (2) require restrictive statistical assumptions such latent variable structure 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 compu- tationally 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 barycentric spanner for the feature embedding efficiently computable basis for exploration, performing efficient spanner computa- tion by interleaving representation learning and policy optimization subroutines. Our mallysss-hhee is appealingly simple and modular- carufar- combines sev- eral techniques, including computation, and a new analysis of certain minimax representation learning objective found in prior work.
ditheemmnniinnll domains where generalization and
approach to error-tolerant barycentric spanner
Introduction
In reinforcement learning and control, many of the most promising application domains require the agent navigate complex, high-dimensiooal state and action spaces, where generalization and function approximation necessary. The last decade has witnessed impressive empirical success in domains where data 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 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 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 nigh-dimensionall state spaces that are both computationally and statistically efficient raises three fundamental questions:
efficiently 1. What the right computational primitives for exploration? That is, how 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 discover such primitives in a computationally and statistically efficient fashion?
3. Given answers the first two questions, how can one efficiently interleave function approxima- tion and exploration to provide provably efficient algorithms?
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this paper. we investigate these questions through the Low-Rank MDP model [36, 45, 1]. In 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 E [H] satisfy
(1i41 Th: an.) pi-1 (I (Th..))))
(I)
where (%.()))) R and 1 (+) Rd are state-action and next-state embeddings. The low-rank structure in (1) facilitates tractable exploration: if the embedding o is known to the learner, one can efficiently learn near-optimal policy with sample complexity polynomial in the feature dimen- sion d. and independent of the size of the state space [23]; in this regard. oi 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 67 and 1h are unknown to the learner. This formulation generalizes well-known frameworks such the Block MDP (BMDP) model [12, 32], and necessitates the use of representation learning: the agent must learn an embedding that approxi- mates of it explores the environment, and must use this learned embedding drive subsequent exploration. This form of function approximation allows for great flexibility, as 0j can be an arbitrary, nonlinear function of the state; in practice, it common to model 0% 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 continuum of states, yet understanding of efficient algorithm design for this model is surprisingly limited. Existing algorithms suffer from at least of the following drawbacks:
1. Computational intractability [21, 24, 14, 43].
2. Strong modeling assumptions (e.g., ability to model ( (.), which facilitates application of model-based RL techniques) [1, 40, 10]; in this work, we aim for model-free methods that only require learning d1.
3. Restrictive structural assumptions (e.g.. non-negativity or latent variable structure for the embed- dings 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 dif- fieult acquire without having already explored the state space to begin with. Overcoming this challenge- patteenge-prtillllrryy where computational efficiency is foncerned-requiie 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) un- derstanding 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.
(1) representation
Contributions. We provide the first provably computationally efficient and model-free algorithm for general Low-Rank MDPs. Our algorithm, SpanRL, uses the notion of barycentric spanner [6] for the embedding 6% 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 layer-wise fashion, learning a new representation at each layer h using exploratory data gathered at previous layers, then uses this representation facilitate computation of a collection of exploratory policies (a policy cover), which act as an approximate barycentric spanner for the features at layer + 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 learn the representation d' (c.g.. neural nets or other flexible classes) whenever a certain minimax representation lcarning objective [35, 49] can be solved efficiently for the function class of interest. Compared efficient algorithms from prior work, SpanRL: (1) is model-free (i.e., only requires access to a function class & capable of modeling o", and does not need to model (i+1)) ). 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 secondary benefit, the algorithm 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 e-optimal policy. & denotes the feature class, and Y denotes an additional feature class capturing model-based function approximation. For approaches that require non-negative (latent variable) structure, dev [resp. 7 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, 7) denotes the reachability parameter.
Comp. efficient
Model-free
General low rank
Sample comp.
OLIVE [21] (see also [24, 14, 9, 43))
4 AAA [o] 22 a A 2222 (
x
FLAMBE [1]
/
Rep-UCB [40]
/
6² A² (
(sec also [10])
x Non-negativel latent variable
dev A y 1og|
MOFFLE [35]²
/
( (
BRIEE [49]
/
X Block MDP
15 A log
A A d (d sfayp
SpanRL (this paper)
and Kleinberg [6], and (2) new analysis of a minimax representation learning objective introduced in Modi et al. [35], Zhang 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 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 reinforce- ment 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 substan- tially improved of the algorithm that removes the reachability assumption the cost of larger (but still polynomial) sample complexity guarantee. We have included this algorithm and its analysis in the ArXiv version of this paper [30].
2
Problem Setting
2.1 Low-Rank MDP Model
We work in an episodic, Umite-horizee reinforcement learning framework, where H é N denotes the horizon. A Low-Rank MDP [36, 45, 1] is tuple M (X,A. (ej))e[r)) (M (eT)) consisting of state space X, action space A with [A| A. distribution over initial states p A(X), and mappings Wh+1 X Rdd and 67 X' x A + Rd Beginning with æ1 p, an episode proceeds in H steps, where for each step he [H]. the state 2t, evolves as a function of the agent's action ap. via
Th-1 Th Rn, as),
where Th is probability transition kernel, which assumed to factorize based on 6j and Hi- In detail, we assume that there exists o-fintte measure v on X' such that for all 1 < h s H 1. and for all é X and a A, the function z' - 1h. (z') o5 (z, a) is probability density with respect to (i.e. the function is everywhere non-negative and integrates to under v)) For any X' s S the probability that 1h+1 E X' under #h-1 (- 2h, ax) then assumed to follow the law
Tú (X' Zhi aw) fxx /i)+1()) 07 (Ih, an (z))).
(2)
For notational compactness, we assume (following, e.g., Jiang et al. (21]) that the MDP M is layered so that X X, u...u XH for X, n nX for all + j. where X, S &' is the subset of states in 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.
We compare the variant of MOFFLE that the same representation learning objective consider. Other variants have improved sample complexity, but make use stronger oracles.
We emphasize that neither μú nor Φi, known to the agent, in contrast to the linear MDP setting [44, 23].
3

---

that are reachable at layer h é [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 v, are a rigorous generalization of Low- Rank MDP formulations found in previous works [23, 11, 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 V as the counting measure yields (2" [In,an) Exex' (x)) oi, (Ih,, an), which is consistent with prior work.
Policies and occupancy measures. We define 1l. { + [ as the set of all randomized, Markovian policies. For a policy π Llx, we let 1 denote the law of (0),a)). under 0₁ ~ गт (xx)), ), and let E= denote the corresponding expectation. For any X' s Kn. we let PT PT 2 X"] denote the marginal law of æ under æ. For x é Nr we define the occupancy
(0H,
measure d' (x) dPP. (x) as the density of 1% with respect to v.
2.2 Online Reinforcement Learning and Reward-Free Exploration
We consider a standard online reinforcement learning framework where the Low-Rank MDP M is unknown, and the learning ag agent interacts with in episodes, where at each episode the agent executes policy of the form X + (AA) and observes the resulting trajectory (z1, a1),... ..... (æH ar). 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, aim to construct an policy cover, a collection of policies that can reach any state with near netr-optimal probability.
Definition 2.1 (Policy cover). For (0, 1]. subset V = Ih is an a-policy cover for layer h if
V Xn, max d' (x) 2 a max d'' (x). NeV
(3)
r'elly
We show (Appendix G) that given access to such a policy cover with constant a, it is possible to optimize any downstream reward function with polynomial sample complexity.
Assumptions. To facilitate learning policy cover, we make the following reachability assumption.
Assumption 2.1 (f-reachability). For any h E [H] and x E Nn, max d' 2 17 [z)]
Reachability is necessary if one aims to build 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 and was only very recently removed in the (more restrictive) BMDP setting [31, 491.
[32, 35, 5],
Remark 2.2 (Comparison to other reachability- assumptions). Assumption 2. generalizes and subsumes all previous reachability- reachability is implied by the notion of feature coverage [5] (used in the context of transfer learning in Low-Rant Rank MDPs), which asserts that suprell, (En [ (Th, an) an) oj (x TT) 2 1) for some mn O. It is also implied by explorability [46], which is similar to feature coverage, but involves the first moments of oj, 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].
conditions of which we are aware [33, 46, 1, 35]. Notably,
Beyond reachability, we assume (following [1, 351) for normalization that, for all h E [H] and (z, a) E Xn x A. 110], (z,a) S 1, and that for all g : Xi, + [0.
I1,,, ((()) s Va.
(4)
Function approximation and desiderata. We do not assume that the true features ( He]H or the mappings (Hi) ke[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 + c (6::2 X x A + Rd) that contains of,, for h E [H 1-1-1 1].
4

---

Assumption 2.2 (Realizabilitt), The feature class _ (o: x + B has oj e & for all h [H]. Moreover. for all & 2, é X, and E A. it holds that (0(2, a) 1.
The class & 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 O finite to simplify presentation, but extension to infinite classes is straightforwwar, as our results only invoke finiteness through standard uniform convergence arguments. Note that unlike model-base ap- proaches [1, 40, 10, 2], we do not assume access to a class capable of realizing the features /4; and our algorithm does not attempt to learn these features; this is why we distinguish our results as model-free.
For constant a, our goal is to learn an a-policy cover using poly (d, A, H, log/ n7 episodes of interaction. This guarantee scales with the dimension d of the feature map and the complexity log] Φ of the feature class but, critically, does not depend on the size of the state space X'; by [10], dependence on and A = |A] is necessary when ø unknown. Given such a guarantee, we show in Appendix G how to optimize any downstream reward function to error & with polynomial sample complexity.
Additional preliminaries. For any m, n N. we denote by [m..n] n2] the integer interval {m, ..... ...,...... .2.. We also let [n] = n.. For any sequence of objects Ø1, .2...., we define Owen == (o)))e(m...)) A partial policy is a policy defined over a contiguous subset of layers [f..r] S [H]. We denote by I = { Uh_e Xr, + A(A)) the set of all partial policies over layers & to r; note that Πl = II For a policy π E II and h E [[..r]. T(x)) denotes the action distribution for the policy at layer h when 24 E Nn is the current state. For 1 S t h H and any pair of partial policies v E IT n² E T we define π of x' E Th as the partial policy given by (# or x (1x) = (xx) for all & < and (x o n (xe) (xz) for all le ht.
We use the 2% I shorthand to indicate that 2; is drawn from the law PT, and likewise for (Ih, an) ~ π and so on. For a set of partial policies V = (-(0= : [N]]. we define u as the random partial policy obtained by sampling i ~ (n]]([]] and playing (1) We define Tunir E Ih as the random policy that selects actions in A uniformly at random at each layer. We use ||-I to denote the Euclidean norm, | se to denote the supremum norm on functions, and let B(r) Rd denote the Euclidean ball of radius r. We refer to a scalar e O as an absolute constant to indicate that it is independent of all problem parameters and use Ō(-) 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 (øi he() 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.
Let us first discuss the issue of computational efficiency. While there are number of algorithms all based on the principle of optimism in face of uncertainty that provide tight sample complexity guarantees for Low- MDPs in soward-bssed [21, 24, 14] and reward-free [9, 43] settings, these algorithms involve intractable optimization problems, and cannot be implemented efficiently even when the learner has access to an optimization oracle for the representation class [11]. This intractability arises because these algorithms implement optimism via a "global" approach, in which the algorithm explores at each round by choosing the most optimistic value function in a certain version space of candidate value functions; optimizing over this version space is challenging,
5

---

as it involves satisfying non-convex constraints with a complicated dependence on the learned representation, and because the constraints are coupled globally across layers h é [H]]
To avoid the intractability of global optimism, several works have restricted attention to a simpler model-based setting. Here, in addition to assuming that the feature maps (#)) be[H]] are realizable with respect to 4, one assumes access to second feature class T capable of modeling the mappings bet)): this facilitates direct estimation of the transition probability kernel (= x, a). For the model-based setting, it is possible to efficiently implement certain "local" forms of optimism [40, 10, 48], as well as certain non-optimistic exploration techniques based on policy covers [l]. Unfortunately, model- based realizability is restrictive assumption, and falls short of the model-free guarantees we aim for in this work; indeed, in general, one cannot hope to estimate the feature map 1 without sample complexity scaling with the number of states.*
When one moves from model-based learning to model-free learning, representation learning becomes substantially more challenging both for optimistic and non-optimistic approaches. Here, a key challenge is to develop representation learning procedures that are (1) efficient, yet (2) provide meaningful guarantees when the learned features are used downstream for exploration. To our knowledge, the only proposal for a representation learning procedure satisfying both desiderata comes from the work of Modi et al. [35], who introduced promising "minimax" representation learning objective (described in detail in the sequel; cf. Algorithm 5), which Zhang et al. [49] subsequently showed to have encouraging empirical performance. However, to provide guarantees for this objective, both works place substantial additional restrictions on the low-rank factorization. In particular, Modi et al. [35] make the so-called latent variable assumption [1], which asserts that 6% and μi are non-negative coordinate-we wise,, and Zhang et al. [49] further restrict to the Block MDP model [12, 32]. Non-negativity is a substantial restriction, as the best non non-negative factorization can have exponentially large dimension relative to the best unrestricted factorization [1], even when reachability is assumed (ef. Appendix H.I). The source of this restriction the problem of how to quantify how close a learned representation Φ is to the ground truth ø', which depends strongly on the downstream exploration strategy. In what follows, we show that with the right exploration strategy, this challenge can be ameliorated, but prior to our work it was unclear whether the minimax objective could lead to meaningful guarantees in the absence of non-negativity.
3.2 The SpanRL Algorithm
Our algorithm, SpanRL, is presented in Algorithm 1. The algorithm proceeds by building a policy cover layer-by-layer in an inductive fashion. For each layer h 2 2. SpanRL uses a policy cover (h) built at a previous iteration within a subroutine, RepLearn (Algorithm 5; deferred to Appendix B) to produce feature map ((h) that approximates o7. Using this feature map, the algorithm invokes second subroutine, RobustSpanner (Algorithm 2 in Appendix B) to produce a collection of policies T Wd that act as barycentric spanner for the feature map, ensuring maximal coverage in a certain sense; given these policies, a new policy cover for layer h + 2 is formed via ((h+2) { i 9 Tanif: : E [d]]. To invoke the RobustSpanner subroutine, SpanRL makes use of additional subroutines for policy optimization (PSDP; Algorithm 3 in Appendix B) and estimation of certain vector-valued functionals (EstVec; Algorithm 7 in Appendix B). We now describe each component of the algorithm in detail, highlighting how they allow us to overcome the challenges in the prequel.
Barycentric spanners. At the heart of SpanRL is the notion of barycentric spanner [6] as an efficient basis for exploration. We begin by defining a barycentric spanner for an abstract set.
Definition 3.1 (Awerbuch and Kleinberg [6]). Given a set W c T such that span(W) Re,, we say that a set (Wi,..., wa) E W is a (C,s) -apprrximmmtttt barycentric spanner for W if for every w e W. there exist $ Ba e [-C, C1 such that |w- Σi= S,wsl s E.²
The utility of barycentric spanners for reward-free exploration is highlighted in the following lemma.
Lemma 3.1. Suppose that Assumption 2.1 holds. If E Th is a collection of policies such that (E* [0]((z) an)) x E V Rd is a (C, -aprrxxiait barycentric spanner for Wh. == {E** [ an)] E I1,) with e S Z. then V is an a-policy cover for layer h+ 1 with (2dC)-1
*For example, in the special case of the Block MDP setting [12, 32], model-based realizability entails modeling a certain emission process, which is not required by model-free approaches. NNote that our definition is a slight generalization of [6. Definition 2. the latter is recovered with 0.

---

Lemma 3.1, proven in Appendix F.6, shows that to compute policy cover for layer h+ 1, it suffices to find a barycentric spanner for the set W,, == { [ (z ax)) π IIn] R.. Of course, even if oj is known, this observation is only useful if we can compute spanner without explicitly enumerating over the set TTy, since our goal is to develop an efficient algorithm. In what follows, we will show:
1. Using, Robus (pannerr,, a novel adaptation of the classical algorithm of Awerbuch and Kleinberg [6], holds that for any Φ Q, spanner computation for the set [E<[[(()) (2 ax)]] 1lx] can be performed efficiently whenever, for any 0 e B(1), one can (approximately) solve linear optimization problems of the form
arg max E" [0 o (nn a.)]
(5)
sell
2. Given access policy covers (1() for layers 1 h. one can efficiently solve the optimization
problem in (5) by appealing to the PSDP algorithm for policy optimization (Algorithm 3).
To handle the fact that 6 is unknown, Algorithm computes policies Ii: that induce barycentric spanner for the set ( [E** [E (ar,, a;)] x Πly), where (()) e is learned feature map. In what follows, we first give detailed explanation of the two above points, before showing how to complete argument by learning a feature map through representation learning.
Algorithm 1 SpanRL: Volumetric Exploration and Representation I Learning via Barycentric Spanner
Require: Feature class 4 and parameters c> and d (0,
1: Set ((()) Ø, (T2) {Munir)
2: Set TReeLearn - A log
3: Set Nesor = A² (og( (8))) and nestiee c - Plog 2log2
(d 1/)))
4: Define F= { : max ((1, 0 (())
5: Define G [g: (z, a) + (z, w OF 9 B(2vd))
6: for h = I, do
/* Learn feature representation for layer h.
7: Set (²)) = RepLearn(h, p(h), neeplearn), with p(h) unif( // Algorithm 6.
/* Computing an approximate spanner using learned features. */
8: For 9 E Rd and (z.a) e A, define
(1(h) (x, for h, otherwise.
rx(,,a;0)
O,
9:
For each te [h], set G. and p(t) unif(
10: For 0 e IRd define (i) = PSDP
(. (.:;)), 9), p(Lih), n1psor). For ll., define LinEst(m) = EstVec (N)), Mestvec).
// Algorithm 3.
11:
// Algorithm 7.
Set (II,..., Ra) Robust(panner((inopt)) LinEst(.), 2,
// Algorithm 2.
13: Set (b++2) [ PN+1 Tunir i [d]].
14:
Return: Policy cover (I:H))
Barycentric spanner computation via approximate linear optimization. To describe spanner computation in SpanRL, we take brief detour and consider an abstract approach barycentric spanner computation, which generalizes our problem. Suppose that we wish to compute a spanner for Z. The Z (which will be set to II when we return to RL) may be exponentially large and cannot be efficiently enumerated. In addition, given = Z, we cannot explicitly compute w and have to settle for noisy approximation.
an implicitly specified set W {w²) zez Rdd indexed by abstract
To allow for efficient spanner computation, we assume access to two oracles for the set W, a linear optimization oracle LinOpt (1) Z and an index-to-vecton oracle LinEst: Z Rdd We assume that for some E 0:
'While barycentric spanners have been used in a number of recent works sample-efficien RL [19, 20). the motivation for their use within algorithu and analysis are quite different; Appendix A.
7

---

1. For all 0 E Id with ||0]| = 1, the output 2/ := LinOpt(f) satisfies fw 2 supzez fw -E.
2. For all = Z, the output vi= Lin satisfies ||ıē= - w² s e.
The RobustSpannen algorithm (Algorithm 2) computes a ( -approiiio spanner for W using ((dlog(d/e) total calls to LinOpt and LinEst. RobustSpanner is an error-tolerant variant of the classical spanner computation algorithm of Awerbuch and Kleinberg [6], which was originally introduced and analyzed for spanner computation with an exact linear optimization oracle. Tolerance to approximation errors in the linear optimization oracle is critical for our application to RL, where additive errors will arise in sampling trajectories, as well as estimating the feature maps (di)ne(m)) Spanner achieves error tolerance by perturbing the vectors returned by LinOpt (0) in the direction of 0, which amounts to running the classical algorithm on an e-fattening of W, and is necessary in order to ensure that the approximation error of LinOpt does not swamp the signal in directions 0 in which W is too "skinny." This technique may be of independent interest; see Appendix C for additional details and formal guarantees.
Representation learning. Ideally, we would like to use Robus SSaannrr to construct a barycentric spanner for the set [I" (#n, an)] π [ly] with 2 [T.. Because do not have access to &h-
we instead apply Robust Spanner with W == ( (æ)) | π E Ily,, where ((th) is a learned representation. We now describe how the feature map (3(0) is learned, then show how to use these learned features to efficiently implement the oracles LinOpt(-) and LinEst(-).
To Iearn a representation for layer h, we use the RepLearn algorithm (Algorithm 5), which was originally introduced in Modi et al. [35], Zhang et al. [49]. The algorithm gathers a collection of triples (In, aj Ih+1) by rolling in to Iy with a policy sampled uniformly from the policy cover (A)) and selecting a uniformly at random. Using this dataset, the algorithm solves a sequence of adversarial training sub-prot problems (Line 9 of Algorithm 5) which involve the feature class t and an auxiliary discriminator class F: X' - IR. As we discuss in detail in the sequel, these sub-problems, described in (7), are amenable to standard gradient-basee based training methods. The sub-problems designed to approximate the following "idcalized" max-min-mie representation learning objective:
300) arg minsup inf (1(0))( dee faF
[( an) " ([[((( [ =n])"]
(6)
The intuition for this objective comes from the fact that in Low-Rank MDP, for any function f: + R. the quantity E[f(x)+) æb = Iz a, = a] is linear in oj (z, a). Thus, if F is sufficiently expressive, we may hope that ((h)) and 6² are close. We adopt the simple discriminator class F= (x maX 0 Φxx a) 0e B(1), & e Φ). We show that solving (6) with this choice for F, which is simpler than that in Modi et al. [35], Zhang et al. [49], yields an approximation guarantee for 3(h) that is suitable for downstream use in spanner computation for general Low-Rank MDPs.
Remark 3.1 (Improved analysis of RepLearn). To facilitate an analysis of SpanRL that does not require reachability assumptions, we use slightly different parameter values for Replearn than in Modi et al. [35], Zhang et al. [49], and provide a tighter sample complexity bound (Theorem 1) which may be of independent interest.
In more detail, prior work shows that the RepLearn algorithm solves variant of (6) with B(d poly (o)ly(= where e 0 is the desired bound on mean-squared error. Due the polynomial dependence on E such a guarantee would lead to vacuous guarantees when invoked within our analysis of SpanRL. Our improved analysis of RepLearn, which is based on determinantal potential argument, shows that w E ((poly(dd) d)) suffices. A secondary benefit of our improved bound is a faster rate with respect to the number of trajectories.
Putting everything together. Having learned ((h) using RepLearn, in SpanRL we apply RobustSpanner with W: == (E*() an)] π e [n), 2 [l, and C = 2; that is, we plug-in the learned representation ((()) for the true representation øj. Witt this choice, implementing LinOpt entails (approximately) solving arg max, RRTT [0T (OT (1hh for given 0 E B(1), and imple- menting the LinEst oracle entails estimating ( (In, an)] for a given π e ll.. We instantiate
Though the policies produced by the algorithm may not necessarily induce a spanner for Wi. [ET (æx, a.)] 11+) (this would require "point-wise" representation learning guarantees, which we do not have), analysis shows that they still suffice to build policy cover for layer h 2.
8

---

LinEst(T) as the Monte Carlo algorithm EstVec (Algorithm 7), which simply samples trajectories according to and returns the sample average of ((hh) (Ih, To implement LinOpt((). we appeal to PSDP (Algorithm 3). PSDP, given an arbitrary reward function risk : X x AA - and a function class G {g X' x A IR} capable of realizing all possible value functions induced by these re- wards,
the policy covers Ty( efficiently compute policy म (hD ri:h, 9, (T(ss)), .n) that approximately solves arg nel1, EE* [E Tr( a.)], and does so using polynomially many episodes; see Appendix D for details and formal guarantees. Thus, implementing LinOpt(O) is as simple as invoking PSDP with the rewards
{ 300) ((0 (z,a)'0, for = otherwise.
re(2,
0,
With this, we have all the ingredients needed for spanner computation, and the algorithm is complete.
3.3 Main Guarantee for SpanRL
The following result is the main sample complexity guarantee for SpanRL (Algorithm 1).
Theorem 3.2 (Main theorem for SpanRL), Let 8 e (0.1) be given, and suppose that realizability holds (Assumption 2.2) and that reachability (Assumption 2.1) is satisfied with parameter 1j O. n
If E = 36 and с = polylog(A, H. d, log(((//))) is sufficiently large, then the policies ((1HH) produced by SpanRL ( c. 6, are a ( (AAd -o-plic cover with probability at least 1 - 8. The total number of episodes used by SpanRL
most:
δ (A4d H (10+((//)))) 1/1/².
Theorem 3.2 is the first provable, model-free sample complexity guarantee for general Low-Rank MDPs that is attained by an efficient algorithm. Prior to work, all efficient model-free algorithms required non- non-negative features (latent variable structure) or stronger assumptions [35, 491, even the presence of similar reachability assumptions; see Table 1.
Remark 3.2 (On the reachability assumption). While the reachability assumption is shared by the best prior efficient algorithms [35], which require non-negativity in addition this assumption, is natural to ask to what extent reachability restricts the generality of the Low-Rank MDP model. In Appendix H, we show that when reachability holds, the embedding dimension in model can be exponentially smaller than the best embedding dimension for the best non-negative (latens variable) embedding [35]. Hence, our results are meaningfully more general than prior work.
While our guarantee is polynomial in all relevant problem parameters, improving the dependence further (e.g., to match that of the best known inefficient algorithms) is an interesting direction for future research, as is removing the reachability assumption.
Application to reward-based RL. By using the policy cover produced by SpanRL within PSDP (Al- gorithm 3), we can optimize any downstream reward function to error E using poly (d, A. H, bl) 1/e² episodes. See Appendix G for details.
Efficiency and practicality, We observe that SpanRL is simple and practical. Defining (D( f) E(n.a (((z,))) T (Ih, ah, Thi 0h-1) tuples, the algorithm is provably efficient whenever the adversarial objective
f(z'))² A//wl where D is dataset consisting of
p(t) E arg max max de
min (p( f. min (ō (o, n...
(7)
feF
in Line 9 of RepLearn (Algorithm 5), can be implemented efficiently (note that by the definition of LD, the "inner" minima over w and 7 in (7) can be solved in closed form). This objective was also assumed to be efficiently solvable in Modi al. [35], Zhang et al. [49] and was empirically shown to be practical in [49]; note that the objective is amenable to standard gradient-based optimization techniques, and that can be over wee-parameterized. While detailed experimental evaluation outside of the scope of this paper, we are optimistic about the empirical performance of the algorithm in light of the encouraging results based on the same objective Zhang et al. [49]
*This is the main place where the analysis the inductive hypothesis that (1ih) policy covers.

---

Outside of representation learning, the only overhead in SpanRL is the RobustSpanner subroutine, which has polynomial runtime. Indeed, Robust requires only polynomially many calls the linear optimization oracle, instantiated as PSDP, which is efficient whenever standard least-squares regression problems based on the class O can be solved efficiently, analogous to [33, 31].
Analysis and proof techniques. The proof of Theorem 3.2, which is given in Appendix F, is appealing in its simplicity and modularity. The crux the proof is to show that the representation learning guarantee in (6) is strong enough to ensure that the downstream spanner computation in RobustSpanner succeeds. It is straightforward to show that spanner computation would succeed
if we had access to an estimated representation that ((h) that approximates 0₁ point-wise (i.c., uniformly for all (z: a) pairs), but the key challenge is that the guarantee in (6) only holds on average under the roll-in distribution uni f (()())) ). Prior works that make use of the same representation learning objective (BRIEE [49] and MOFFLE [35]) do not make use of spanners; instead, they appeal to exploration strategies based on elliptic bonuses, addressing the issue of approximation errors through additional assumptions (non-negativity of the factorization for MOFFLE, and Block MDP structure for BRIEE). Perhaps the most important observation in our proof is that barycentric spanners are robust to the average-case approximation error guarantee in (6) as-is, without additional structural assumptions. Intuitively, this benefit seems to arise from the fact that the spanner property only concerns the first moment of the feature map o" while algorithms based on elliptic bonuses require approximation guarantees for the second moment; understanding this issue more deeply is an interesting question for future work. Another useful feature of our proof is to show that the notion of reachability in Assumption 2.1, which generalizes and extends all previous reachability conditions in the Low-Rand MDP and Block MDP literature [46, 13, 33, 1, 35, 31], is sufficient to build policy cover. We anticipate that this observation will find broader use.
10

---

Acknowledgments and Disclosure of Funding
We thank Noah Golowich, Dhruv Rohatgi, and Ayush Sekhari for several helpful discussions. Z.M and AR acknowledge support from the ONR through awards 100014-20-1-2333 and N00014-20-1-239 and ARO through award 1991NFF21-11035 AB acknowledges support from the National Science Foundation Graduate Research Fellowship under Grant No. 1122374.
References
[1] A. Agarwal, S. Kakade, A. Krishnamurthy, and W. Sun. FLAMBE: Structural complexity and representation learning of low rank MDPs. arXiv preprint arXiv:2006. 108/4, 2020.
[2] A. Agarwal, S. Kakade, and L. F. Yang. Model- reinforcement learning with generative model is minimax optimal. Conference on Learning Theory, pages 67-83. PMLR, 2020.
[3] A. Agarwal, M. Kakade, A. Krishnamurthy, and W. Sun. FLAMBE: structural complexity and representation learning of low rank mdps. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
[4] A. Agarwal, Y. Jin, and T. Zhang. Vo I: Towards optimal regret in model-free rl with nonlinear function approximation. arXiv preprint arXiv:2212.06069, 2022.
[5] A. Agarwal, Y. Song, W. Sun, K. Wang, M. Wang, and X. Zhang. Provable benefits of representational transfer in reinforcement learning. CoRR, 0et/2205..4577. 2022.
[6] B. Awerbuch and R. Kleinberg. Online linear optimization and adaptive routing. Journal of Computer and System Sciences, 74(1):97-114, 2008
[7] M. G. Azar, Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning. pages 263-272, 2017.
[8] J. Bagnell, S. M. Kakade, J. Schneider, and A. Ng. Policy search by dynamic programming. Advances in neural information processing systems, 16, 2003.
[9] Chen, Y. Bai, and S. Mei. Partially observable rl with b-stability: Unified structural condition and sharp sample-efficient algorithms. arXiv preprint arXiv:2209. 14990, 2022.
[10] Y. Cheng, R. Huang, Y. Liang, and J. Yang. Improved sample complexity for reward-free reinforcement learning under low-rank mdps. In The Eleventh International Conference on Learning Representations, 2023.
[11] C. Dann, N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. On oracle- efficient PAC RL with rich observations. In Advances in neural information processing systems, pages 1422-1432, 2018.
[12] S. S. Du, A. Krishnamurthy, N. Jiang, A. Agarwal, M. Dudik, and J. Langford. Provably efficient RL with rich observations via latent state decoding. arXiv preprint arXiv:1901.00 1901.09018... 2019.
[13] S. S. Du, Y. Luo, R. Wang, and Zhang. Provably efficient Q-learning with function ap- proximation via distribution shift error checking oracle. In Advances in Neural Information Processing Systems, pages 8060-8070, 2019.
[14] S. Du, M. Kakade, J. D. Lee, S. Lovett, G. Mahajan, W. Sun, and R. Wang. Bilinear classes: A structural framework for provable generalization in RL. arXiv preprint arXiv:2103. 10897, 2021.
[15] Y. Efroni, D. Misra, A. Krishnamurthy, A. Agarwal, and J. Langford. Provably filtering exogenous distractors using multistep inverse dynamics. In International Conference on Learning Representations, 2021.
[16] D. Ernst, P. Geurts, and L. Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6, 2005.
[17] D. J. Foster, S. M. Kakade, J. Qian, and A. Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487, 2021.
[18] D. J. Foster, N. Golowich, and Y. Han. Tight guarantees for interactive decision making with the decision-estimation coefficient. Conference on Learning Theory (COLT), 2023.
11

---

[19] N. Golowich, A. Moitra, and D. Rohatgi. Learning in observable pomdps, without computation- ally intractable oracles. In Advances Neural Information Processing Systems, 2022.
[20] A. Huang, J. Chen, and N. Jiang. Reinforcement learning in low-rank mdps with density features. International Conference on Machine Learning (ICML), 2023.
[21] N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. Contextual decision processes with low Bellman rank are PAC-learnable. In International Conference Machine Learning, pages 1704-1713, 2017.
[22] C. Jin, P. Netrapalli, R. Ge, S. M. Kakade, and M. I. Jordan. A short note on concentration inequalities for random vectors with subgaussian norm, arXiv preprint arXiv:1902.03776. 2019.
[23] C. Jin, Z. Yang, Z. Wang, and M. I. Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137-2143, 2020.
[24] C. Jin, Q. Liu, and S. Miryoosefi. Bellman eluder dimension: New rich classes of RL, problems, and sample-efficient algorithms. arXiv preprint arXiv:2102.00811 2021.
[25] S. M. Kakade. On the sample complexity of reinforcement learning. University of London, University College London (United Kingdom), 2003.
[26] J. Kober, J. A. Bagnell, and J. Peters. Reinforcement learning in roboties: A survey. The International Journal of Robotics Research, 32 12))11222444 2013.
[27] Li, W. Monroe, A. Ritter, D. Jurafsky, M. Galley, and Gao. Deep reinforcement learning for dialogue generation. In EMNLP, 2016.
[28] T. P. Lillicrap. J. J. Hunt, A. Pritzel, N. Heess, T. Frez, Y. Tassa, D. Silver, and D. Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: 150..025/, 2015.
[29] Z. Mhammedi, D. J. Foster, M. Simchowitz, D. Misra, W. Sun, A. Krishnamurthy, A. Rakhlin, and J. Langford. Learning the linear quadratic regulator from nonlinear observations. Advances in Neural Information Processing Systems, 1111422--45445 2020.
[30] 7. Mhammedi, A. Block, D. J. Foster, and A. Rakhlin. Efficient model-free exploration in low-rank mdps. arXiv preprint vrXvv22307.03997, 2023.
[31] Z. Mhammedi, D. J. Foster, and Rakhlin. Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observation rl. International Conference on Machine Learning (ICML., 2023.
[32] D. Misra, M. Henaff, A. Krishnamurthy, and J. Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. arXiv preprint 1rXiv:1911.05815, 2019.
[33] D. Misra, M. Henaff, A. Krishnamurthy, and J. Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In International conference on machine learning, pages 6961 -9711. PMLR, 2020.
[34] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep rein- forcement learning. Nature, 518(7540)5529, 2015.
[35] A. Modi, J. Chen, A. Krishnamurthy, N. Jiang, and A. Agarwal. Model-free representation learning and exploration in low-rank mdps. CORR, abs/2102.07035. 2021.
[36] S. Rendle, C. Freudenthaler, and L. Schmidt-Thieme. Factorizing personalized markov chains for naxk-basket recommendation. In Proceedings of the 19th International Conference on World Wide Web, WWW 2010, Raleigh, North Carolina, USA, April 26-30. 2010, pages 811-820. ACM, 2010.
[37] D. Russo and B. Van Roy. Eluder dimension and the sample complexity of optimistic exploration. In Advances in Neural Information Processing Systems, pages 2256-2264, 2013.
[38] D. Silver, A. Huang, J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):884, 2016.
[39] W. Sun, N. Jiang, A. Krishnamurthy, A. Agarwal, and J. Langford. Model-based RL in contextual decision processes: PAC bounds and exponential improvements over model-free approaches. In Conference learning theory, pages
PMLR, 2019.
12

---

[40] M. Uchara, X. Zhang, and W. Sun. Representation learning for online and offline RL in low-rank mdps. The Tenth International Conference on Learning Representations. ICLR 2022, Virtull Event, April 25-29, 2022, 2022.
[41] R. Wang. S. Du, L. Yang, and R. R. Salakhutdinov. On reward-free reinforcement learning with linear function approximation. Advances in neural information processing systems, 33: 17816-17826, 2020.
[42] R. Wang. R. Salakhutdinov, and L. F. Yang. Provably efficient reinforcement learning with general value function approximation. arXiv preprint arXhv:2005. 2020.
[43] T. Xie, D. J. Foster, Y. Bai, N. Jiang, and S. M. Kakade. The role of coverage in online reinforcement learning. arXiv preprint arXiv:22 10.0410 2022.
[44] I.. Yang and M. Wang. Sample-optimal parametric q-learning using linearly additive features. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 6995-7004, PMLR, 2019.
[45] H. Yao, C. Szepesvári, B. Á. Pires, and X. Zhang. Pseudo-mdps and factored linear action models. In 2014 IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning, ADPRL 2014, Orlando, FL, USA, December 9-12, 2014, pages 1-9. IEEE, 2014.
[46] A. Zanette, A. Lazaric, M. J. Kochenderfer, and E. Brunskill. Provably efficient reward- navigation with linear value iteration. Advances Neural Information Processing Systems, 33: 11756-11766, 2020.
[47] T. Zhang. Feel-good thompson sampling for contextual bandits and reinforcement learning. arXiv preprini arXiv:21 10.00870. 2021.
[48] T. Zhang, T. Ren, M. Yang, J. Gonzalez, D. Schuurmans, and B. Dai. Making linear mdps practical via contrastive representation learning. In International Conference on Machine Learning, pages 26447-26466. PMLR, 2022.
[49] X. Zhang, Y. Song, M. Uehara, M. Wang, A. Agarwal, and W. Sun. Efficient reinforcement learning block mdps: model-free representation learning approach. International Conference Machine Learning, pages 26517-26554 PMLR, 2022.
13

---

Contents of Appendix
Additional Related Work
14
Omitted Algorithms
15
C Generic Guarantee for RobustSpanner
18
D Generic Guarantee for PSDP
20
Guarantee for RepLearn
23
Analysis
26
F.1
Proof Strategy
26
F.2
Guarantee for PSDP as Subroutine for RobustSpanner
F33
Guarantee for RobustSpanner as Subroutine for SpanRL
28
F.4
Guarantee for RepLearn as Subroutine for SpanRL
F.5
Concluding the Proof of Theorem 3.2
F6
Proof of Lemma 3.1
G
Application to Haward-Based RL
32
H Properties of Reachability Assumption
33
H.I Comparison Latent Variable Model
H.2 Relation to Other Reachability Assumptions
35
A Additional Related Work
In this section, we discuss relevant related work not already covered.
Block MDPs. A particularly well-studied special case low-rank MDPs with the latent variable assumed in Modi et [35] (defined in Definition H. 1) is the Block MDP (BMDP) model Du al. [13], Misra et al. [32], Zhang et al. [49], Mhammedi et al. [31]. For this setting, Du et al. [13], Misra al. [32] provide algorithms that conduct exploration in provably oracle-efficient manner under reachability assumption. This reachability assumption was removed by subsequent work of Zhang al. [49] (with a suboptimal rate) and Mhammedi al. [31] (with optimal error dependence), but the analysis in these works is tailored to the BMDP model.
in which spanners. Huang et al. [20] consider a variant of the Low-Rank MDP framework given class T that realizes the next-state feature map yt" but do not have access to a class + for the feature map ø", which is unknown. Their algorithm, like SpanRL, is based on barycentric spanners, though the algorithm design considerations and analysis different. Notably, their algorithm is not computationally efficient, and their analysis takes advantage of the fact that realizability of μ² facilitates estimation of the occupancies (d' (e)) in é -errrr.. Barycentric spanners were also in the work of Golowich et al. [19] for reinforcement learning in Partially Observable MDPs (POMDPs). Their analysis substantially different from ours, and their algorithm appeals to the barycentric spanner computation approach in Awerbuch and Kleinberg [6] an off-the-shelf fashion.
Barycentric
significantly
14

---

B
Omitted Algorithms
Algorithm 2 RobustSpanner: Barycentric Spanner via Approximate Linear Optimization
Require:
Approximate linear optimization subroutine LinOpt: R + Z. Approximate index-to- subroutine LinEst: 2 + Rd
/* See Section 3.2 */
Parameters C,E>
1: Set W = (wi,..., wa) (eI,...,Ea).
2: for == 1, ,....,d do
Set 0, = (det(ej, W W-:)) E IR.. W.W., is defined to be W without the ith column 4: Set 2 (inOptt0:/|(00) and w LLiEstt="))
3:
5:
Set 2 Lin0pt( (00ppt(-0:/1.... (0,] and wi = LinEst(27))
6:
if 0. w 2 -0, w. then
7:
Set W, wT 2, 2t. and w, m. =0./10,1.
8:
else
9:
Set W: w/ 21 z and w; e0./10,l.
for
1,2,... do
11:
Set deti=1 1.
while d do
13:
Set 0g = (dj W-:)) e Rd.
14:
Set 2 LLinppt( [1/10: and w = LinEst(2))) ) (amoppt(0,//00))) and wj = LinEst(2,)) if 0, w/ + + ||01]| 2 C. |del ws, W-)) then
15:
Set 27
16:
Set W. w 2, af. and w, 0 + 0 0,|l-
18:
break
19:
else if 0. wj ||0:1| 2 C. det w,, W-)) then
20:
Set W. = wT. 24 = 27, and w; = T. - e 0,/10,1|
21:
break
Set i
23:
if d 1 then
24:
break
25:
Return: (z)....,24))
15

---

Algorithm 3 PSDP (h, Pi:h; G, (/(1sh), n): Policy Search by Dynamic Programming (variant of [8])
Require:
Target layer / [H]]
Reward functions ri:h-
Function class g.
Policy covers vr(1), (T(h)
Number of samples N.
1: for = 1,....,o do
2: D(t) Ø.
3: for n times do
4: Sample (It, a2 Ee_ re( az)) ~ unif(u((t)) or Tunir o441 ((++1),
5: Update dataset: D()) D(t) u (x1,, at,
at, Ei=
6: Solve regression:
(1) arg min E
(0(x;a) - R)².
ges
(s.o,R)eD(0)
7: Define (1) E & via
arg max an (0)) (I,
x E Nt,
(0) (x)
(8)
((1+1(
x Nt, le [t 1..h].
8: Return: Near-optimal policy #11) 1l..
Algorithm F,, n): Estimate E"[F ax)] for policy and function F: A Rd
Require:
Target layer h [H].
Vector-valued function F: 2' A - Rdd
Policy π Iln.
Number of samples N.
1: D Ø.
2: for times do
3: Sample (ær,, an))
4: Update dataset: D + Du {(0,, }.)).
5: Return: F = 1.2 E(r.))eD F(z,a).
16

---

Algorithm RepLearn(h, F, 2, n): Representation Learning for Low-Rank MDPs [35, 49].
Require:
Target layer [H].
Discriminator class F.
Feature class Φ
set of policies V.
Number of samples E N.
Set 100 dloog nog (1/6) 51 16
2: Initialize x= = Vn T= 2 and & = f1 + 2 (00) ④
3: D Ø
4: for n times do
5: Sample (æn, 0h1 Th+ unif(V) on Tunif.
6: Update dataset: D Du (E ar, xn+1)).
7: Define Lp(Ø, w. f) = Σ ((., w - f(x))² + A//wo/2
8: for t= , do
/* Discriminator selection */
9: Solve
((() arg max ā(J), where Σ(J) = max min (D((1)t) ,),,,,,,,- min Lz( 10 (0. feF
Φeb
10:
if {))()) S then
11:
Return ò- ((t),
/* Feature selection via least-squares minimization */
Solve
(((+1) e arg min aig min Σ mob
Lol ty, (().
1=0
17

---

C Generic Guarantee for RobustSpanner
In this section, we give generic guarantee for the RobustSpanne algorithm when invoked with oracles LinOpt and LinEst satisfying the following assumption.
Assumption C.1 (LinOpt and LinEst as approximate Linear Optimization Oracles). For some abstract set Z and collection of vectors {w² TR& i 2) indexed by elements 2, there exists e' such that for any TRd {0} and 2 Z, the outputs 20 LinOpt(8/(l0l) and i; = LinEst(())) satisfy
sup w² s f a + -10]], and |0= - w² - ss'. se2
Letting W: = {w= | } and assuming that W B(1), the next theorem bounds the number of under Assumption C.I, and shows that the
iterations of RobustSpanner output is approximate barycentric spanner for W (Definition 3.1). Our result extends those of Awerbuch and Kleinberg [6], in that only requires an approximate linear optimization oracle, which potentially of independent interest.
Proposition C.1. Fix C> 1 and (0, and suppose that {w> 2 E 2) = B(1). If RobustSpanner (Algorithm 2) run with parameters C,E> O and oracles LinOpt, LinEst satisfying Assumption C.I with e' e/2, then terminates after d + [ logc =f iterations, and requires most twice that many calls to each of LinOpt and LinEst. Furthermore, the output 2nd has the property that for all Z, there exist BI..... Bd [-C, CJ, such that
3Cd.E
w² 3,0 S
Proof of Proposition C.1. The proof will follows similar steps to those in Awerbuch and Kleinberg [6, Lemma 2.6], with modifications to account for the fact that linear optimization over the set W = (w 2 2) is only performed approximately.
Part I: Bounding the number of iterations. In Algorithm 2, there are two loops, both of which require two calls to LinOpt and LinEst per iteration. As the first loop has exactly d iterations, suffices to bound the number of iterations in the second loop.
Let M( (Wi...., Wi, Ci+ 1 ...................... ea) be the matrix whose columns the vectors at end of the ith iteration of the first loop (Line 2) of Algorithm 2; note that columns i + through are unchanged at this point in the algorithm. For i [d]. we define ((w) det(w: Net(w, (M ( and 0; ==
(det (ez, M je[d] E IRd, where we recall that for any matrix A, the matrix A_; defined as the result of removing the ith column from A. Note that & is linear in w, and in particular
& (w()) w TT..
Let (' (d) (Way.. ................................ wa), and let W denote the resulting matrix after iterations of the second loop (Line 10) of Algorithm We will show that for any 2 1,
det det(w())) s det(W(()) (100d)
((100)
(9)
2
By construction of the loop, have det(W 2 C- det(W))1)) for each j [J], and thus det det(W(()) 2 det det(W(")). C., Combining these two facts will establish the bound on the iteration complexity. We now prove (9).
Let u, = ct (A) (note that u, is a row vector) and let U denote the matrix whose ith row is u, We observe that for all e Rd.
(u(w)
N&W
E (ww))
18

---

where we note that (w())) by construction; indeed, the columns of M are basis for R because det( 0, and the equality holds on the columns, so the two linear functions must be equal. Now, since Assumption C.I holds with =/22, we have
0. T sup 0 w se2
1
and 0. W
inf 0. w² 222
11::1l
(10)
2
where w LinOpt(=f)) We will now show that
E |10.11- ((xe)) 2 2
(11)
There are two cases. First, suppose that 0J wej 2 - wT. corresponding to the conditional in Line 6 of Algorithm& being satisfied. Combining this with (10), we have
0. u
sup wl w .=e
-- 110,11
v (-0, w,))
2
-210.1)
sup 01 w 262
10:11
v sup 2E
w==
(by (10))
2
sup 0. w²
sup
0 u²
10:1.
xe2
x=
2
2
1011
(12)
Because the conditional is satisfied, w; i + To.1 _ and by plugging this into (12), we have
((w;) 0g ws 10.1-
The case that Of s -0,w is essentially identical, establishing (11). Now, recall that W = {wi e B} and let W @ B 3) {w + b we W and be B )} denote the Minkowski sum with 3e ). By Cauchy- Schwar,, it holds that for all == w + be W @ B (3
2
3s
(ww' 0. w' f. N 0 1
110.1l.
2
where we used that W s B(1) (by assumption). Thus, for any w/ W B (3r). have
|u:w1| (<ww') ((ww)
3s 2
$
We now observe that by construction and the fact that Assumption C.I holds with =' =/2, the kth column w/K of WJ) belongs to W B 35 for any [d]. Thus, the (i,k) entry u,w, of (W))) satisfies u,wrf [-11
2
2 1 + 3e ). and so the columns of UW(J) have Euelidean norm most 10va 2
Since the magnitude of the determinant of matrix is upper bounded by the product of the Euclidean norms of its columns, it holds that (det(UW(())] 10od
On the other hand, again by construction, we see that the columns Wi...., Wy of W(O) satisfy u,W;- for and u,w; 1. Thus, (W())) is upper-triangular matrix with Is the diagonal, and hence has determinant 1. Because determinants are multiplicative, this implies that det (U) # O. We now compute:
(e((( [det(UW(()) (d)(())(())) [det(U)]
100d
(de((UU((0))
2
Thus, the upper bound on (W((())))) holds and the claim is proven. Therefore, we have
100d $
and J s logo
(-
100d
19

---

Part I:: Spanner property for the output. Having shown that the algorithm terminates, we now show that the result is an approximate barycentric spanner for W. Let W (......, Wa) be the matrix at termination of the algorithm. By definition, if the second loop (Line 10) has terminated, then for all i [d],
max w -0]w, + e 10.11 s C (we W-:)l,
where 0, (det(ez, W W-)) e R.. On the other hand, by Assumption C.1. (10) holds, and so Vz 2, Vie [d], det(ww, WW:) 10, w=1 s max (0; w. -0, w + +001l
C- C-I W-+)l.
(13)
Now, fix a Z. Since det(W) # O, there exist BI.d R such that w Σde Sew, By plugging this into (13) and using the linearity of the determinant, we have
Vie [d], C C-Idet W-:) 2 det W-:) Σ 3r ded wy, W-.) | |we W-:)l- j==1
Therefore, [3:] s C, for all i E [d]. Now, by definition of widd and Wad. for all i [d], we have that |wo Wil $ Furthermore, by Assumption C.1, we also have that [7; w=l s e/2. Therefore, by the triangle inequality, we have
d
w² Σ
S
Σ
0
+ Σ i-1
Ball 80
Σ im1
1811
will s 3dCCs/2
This completes the proof.
D Generic Guarantee for PSDP
In this section, present generic guarantee for PSDP (Algorithm 3). We show that given any reward functions Tish : X' x A R and a function class g = {g : X' x A + RR} that "realizes" these reward functions (we formalize this the next definition), Tp (1=b) are a-policy covers for layers through h, then for sufficiently large n 2 and with high probability, the output á PSDP (h, Fis, G, (1=k), n) an approximate maximizer of the objective
h Σ [Er(or)
max T relly
To formalize this result, we define the notion of realizability we require for the function class G.
Definition D.1. We say that the function class & {g : X' A R} realizes the reward functions Tish X x A + IR for all e [h] and all n ITth
O G, where
E"
Σ re( It = = 2, 0.
(14)
[=++1
Note that Qf in (14) represents the state-action value function (Q-function) at layer [h] with respect to the rewards Tish and partial policy π.
In what follows, given function class {g: X' A A + R], use (=) to denote the e-covering number of 9 in la distance. With this, we now state a guarantee for PSDP.
Theorem D.2. Let 8 (0,1), B >>0 O, and h e [H]. Suppose reward functions Tisk X x + R, function class G, collection of policies (11:)), and a parameter 1 satisfy the following:
The function class realizes the reward functions Tish (in the sense of Definition D.I), functions in G are uniformly bounded by B, and lim n (o//)) 0.
For some 0 < a S for each S st h, it holds that (1) is an a-policy cover for layer t and moreover [T/(9]
The parameter n is chosen such that odHo- Estar (n, 6/H) S E, where Estar 8') =
/ BBnn (log.Ng(1/n) log(n/o)) and c 0 is a large enough absolute constant.
20

---

Then, with probability at least 1 1 6, the policy á PSDP(h, Fi4, 9, (0(11)), ,n) coming from Algo-
rithm 3, satisfies the following guarantee:
max
E"
Σ Er((z).a))
s E
E En(oz,ai)
relly
To prove the theorem, we need two intermediate results. The first shows that the Q function the Bayes-optimal predictor of the of rewards when rolling with policy и.
Lemma D.1. Let e [H], Tiah X x A R reward functions, π E Hly and let Foaye denote the Bayes-optimal predictor" for the sum of rewards under policy i.e.,
(
A Ere(Fz, as)
21
Poayes
arg min E² g:XixA-3
(15)
E=t
Then, Gaayes QT. where QF is the Q-function layer t e [h] with respect to the policy and rewards 7ith defined in (14).
Proof of Lemma D.1. The least-squares solution Paayes of the problem in (15) satisfies, for all a e A and x Xt,
Pbayes (z, Ed
Σ re( as) xt at aa
=B(ro(xt,, a.) 13 -r =a
Σ
re(
æi x, at =
1=++1
ar)( + EE
Σ ビ-t+1
re(wo, ar)
C; 21 ag
= a))
where the last step follows by the definition of the Q-function in (14).
We now show that the solution (*) the least-squares problem in (6) of Algorithm 3 is close to the Q-function in the appropriate sense.
Lemma D.2. Let e, 5 (0, B> O, and S S H. Further, let (Estat, Tick, G, (p(1:)), ,n)) be as in Theorem D.2. Then, the solution g(t) of the least-squares problem in (6) in Algorithm 3 when calling PSDP Tish, 9, (((i))), n) satisfies with probability at least 1-5,
(u(if(g(())
max (g (II, - -Qi."-)) (xi, a)) (z2.
Eat (n,
where # (t-1) 11 is in Algorithm 3, and is an absolute constant.
Proof of Lemma D.2. Fix t [h] and Dbayes = Playes be as in Lemma D.I with
(t)
uni (()(()())))))))))))))))))))))( it Tunif 9++1 ((+1)
(t) and reward functions 7ish in the lemma's statement. By Lemma D.I, 9 is the Bayes-optimal solution of the least-squares problem in (6) of Algorithm 3. Thus, since & realizes the reward functions Tis,, a standard uniform-convereee converenc guarantee for least-square regression (see e.g. Mhammedi et al. [29, Proposition B. with (independent of t, ht and any other problem parameters) such that with probability at least - 8,
almost surely) implies that there exists absolute constant 0
()()) max (e (x)) Ibayes (Ie, a)) s c² (ogN)((/n) + log(n/8)
(t) The desired result follows by the fact that Ihayes Q (1111)) see Lemma D. We also require the classical performance difference lemma from Kakade [25].
'Observe that because the loss is strongly convex, this predictor is unique up to sets of measure zero, justifying usage of "the" Bayes optimal reward.
21

---

Lemma D.3 (Performance Difference Lemma). Let TT, E Il be arbitrary, and Q7 be as defined in (14). Then, for any h 2
h Er(
h (a as)
E
E.
Σ
Σ tm1
R (17 (7 (1,, (x:))]
Using these results, we now prove Theorem D.2. Proof of Theorem D.2. We first show that for any t [h], there is an event &, of probability at least 1- 8/ H under which the learned partial policies $(1), ((++1) are such that
ET [t 1 (II, (x:)) 2₂" ( (I (x))) H
(16)
where arg mel, 1T Te( at) is the optimal policy and QF is the Q-function defined in (14). Once we establish (16) for all i e [h], we will apply the performance difference lemma (Lemma D.3) and the union bound to obtain the desired result.
For any LE [h.]. we have
TT Qr (111)
(1,, (x:)) - Q ((1))) (a,, * (x:)))
[o (a n, (xx)) - (x,, n. (ax)) + glth (x,, n, (x:)) Q+ Qf (a Q (2:))
ET.
E.
[o
(. - (0)
(I n. (a:))
(( (x,, # (xx)) Qi ((11))
(0,, x00 (ax:))
(xx)))
$ ET max aeA
Oi Qi (x), - (x),
(17)
where the penultimate inequality follows by the fact that ((0) (x) E arg max ((0 (x) a), for all
(x,, a) E X4 x A by its construction in (8). Now, using the assumption on v((1, we have for any & e [t].
E. [o 1()1)
(.:e)) - ai (1-1) (x,, t00 (x:))
{ (max
Q Q (x,a) 800) (z, (e)
2 Jx.
d= d"h (x)du(x),
aedd
21 Sr
(max aed
Ot (111 (z,a) -000g) (z.a)) max d' d<))d((x), TeWitt)
22 - f(L
max aeA
(I, -c,a)-g') (x,
(x, a)) ( (s(1)))))
2d. (unir(()())
max Q2 (2)1,0) ((0) (m),)))
(18)
asA
where the first inequality is by (17), the second inequality follows from the fact that (()) is an a-policy cover, the third inequality follows from the fact that [v(0)] d, and the equality is by definition. Now, by Lemma D.2, we have that for any & [h]. there is absolute constant (independent of and other problem parameters) and an event &, of probability at least 1 8/f under which the solution g()) of the least-squares regression problem on (6) of Algorithm 3 satisfies,
0
()))
max A
g)() (x, a) Qr (zt, a
s c. Estat (n, 4
QE
(19)
2dH
where the last inequality follows by the choice of in the theorem's 's statement. Combining (19) with (18) establishes (16) under the event &
Now, by the performance difference lemma (Lemma D.3), we have
R.
Σ re ( (Z..
- x Σ re (z., (2x,,a)) an)
Σ [-1
T.
[Q (z0, (2x)) Q ( (0),))
22

---

Algorithm 6 R, ,, P.n): Representation Learning for Low-Rank MDPs [35]
Require:
Target layer h [H].
Discriminator class F.
Feature class +
Policy distribution P ((I+).
Number of samples n N.
1: Set Estat Log for sufficiently absolute constant > 0 (see Appendix E).
2: Let (i(1) + be arbitrary.
3: Set D + Ø.
4: for times do
5: Sample π
6: Sample (I/, ans 21++1 ) π Tunif.
Update dataset: D Du (Eh ar #h+1)).
8: Define ( w. f) Σ ((..,='))
f(x'))2
9: for t= 1,2,... do
Discriminator selection Solve
10:
f(9) arg max A(f),
where î(f) = max deb
min
C w) GBB(2a))
min
co(0.0,11
2)
JeF
(3d(3))/2)
if (((()) s 16d1s stat then
12:
Return (11).
/* Feature selection via least-squares minimization */ Solve
cf(+1) e earg min
min
Σ
Lol
(Wis 1e (.)co.(
F-1
Thus, under the event & = U &,, we have that
Σ En an)
-E* # Σ ( am) SE.
T
The desired result follows by the fact that union bound implies P[E] 2 1 5.
E Guarantee for RepLearn
In this section, we give a generic guarantee for RepLearn (Algorithm 6). Compared to previous guarantees in Modi et al. [35], Zhang et al. [49], we prove a fast 1/n-type rate of convergence for RepLearn, and show that the algorithm succeeds even when the norm the weight in Line 10 does not grow with the number of iterations. We also use the slightly simpler discriminator class:
F f:x max (z, (ax 0 ((B()), 1
(20)
aed
The main guarantee for RepLearn is as follows.
23

---

Theorem E.1. Let E [ d' E (0,e"),, and n E N be given, and suppose that ph+1 satisfies the normalization assumption in Eq. (4). For any function e F, define
wy= = Sx... (1)dda.))
Let P с (TT,) be distribution over policies, FF as (20), and 4 be a feature class satisfying Assumption 2.2. With probability at least 1 - ö, RepLearn with input (h,F, ④ P.n) terminates after ts T= [ log3/2 [0ld iterations, and its output (()) satisfies
sup inf feF weB(33) BRoAPT (ww (wT an) w an))) keptearn (n, (21)
(x))
where
= ERpl carn (n, cTd log(( for some sufficiently large absolute constant 0.
To prove the theorem, we need a technical lemma, which follows from Modi et al. [35, Lemma 14].
Lemma E.I. Consider call to Replear earn(h,F,F p.n) (Algorithm 6) in the setting of Theorem E.
Further, let LD be as in Algorithm and define
(6 (1)
(--1)
(1)
arg min
Σ we, f(0).
(22)
de t (w ( BB2vaa
For any E (0, 1), there is an event g()) of probability at least 1 - 8 such that under ()), if
Algorithm 6 does not terminate at iteration t 2 1, then for w(() = w
( an Tn (2,,,, as )z) wn
RT F..pRR P..PBRAA Tuni
s $e te (n,0),
(23)
inf we (EBB(13/2)
E Brr.BBooee Tuesst
(0()) an ( an )w(())
8dte (n,
where Saa (n, 8) = cdLnn log( (6) and 2 1 is sufficiently large absolute constant.
With this, we prove Theorem E.I.
Proof of Theorem E.1. Let us abbreviate = (n, 6), with (nt 8) defined as in Lemma E. Further, let N == 1 + + (dloga/2(2dP22) 8 = 2N and define
Etat Estat (n.
(24)
Note that s Estat and N the latter inequality follows by the facts that the absolute constant in Lemma E.1 is at least 1 and (/8() 2 1. We define an event E = (8)) n ((A)) where (E)())), : are the success events in Lemma E.1. Note that P[E] condition on the event
where T the number of iterations in the theorem statement;
5/2 by the union bound. Throughout this proof,
To begin the proof, we define sequence of vectors (v (220 120 an inductive fashion, with of Rd for all i E [d] and 2 O. For l O, we let vf(0) ee./d, dd, for all i [d]. For & 1, we consider two cases:
. Case I: If
f(9 = { [d] |det(v (0)))) ( (1+C) (1+C)--ldet( (de((v((-)))] #2,
where ((6-1) == (o(t-1) uu..,,u
vi ((1)) E Rddd and (()) == wpn then we let j =
arg min argminye j' and define
5 =
if
v (6-1)
otherwise.
. Case II: (f) =, we let Di
(0)
(t-1)
for all i [d].

---

We first show that J(t) 2 at any iteration t [N] where RepLearn does not terminate. Let t [N] be an iteration where the algorithm does not terminate, and suppose that J(() Ø. This means that
Vj E [d],
/det(v ( s (1+
det (det(v(-))) (V()))]
(25)
Now, since det(((())) * 0 (note that [de)( is non-decreasing with t), we have that span(v = R Thus, there exist $1,..., Aa e IR be such that w(() Ei B:0g the linearity of the determinant and (25), we have
By
vjE E [d], (1+ [)) det(-d 2 det( (vee") ,w())),
det v((-1) d
Σ (,o(t-1) (t-1))
Σ BB det det( (((-1) ; (-1))) [c[a]
18y1 (ell-|eet (det(V((-1)))
This implics that |3,1 S (1+ + C)) for all j 6 [d]. Now, note that by the definition of (u, (((11)), ) we have that for any i [d] such that vs (1-1)
+ ses/d, there exists le [L- 1] such that w(() 0z (t-1) 1) Let
I = fie [d] ((--1) # ses/d},
and for any ie I(0) let la E [t 1] be such that w((t) = v ((111) Further, define
(
Σ
Bew(())
Σ B,u(t-1)
(26)
(eZ(s)
(eZ(e)
and note that by the triangle inequality and the fact that ((()) = Ed B,0, (1-1) we have
( - (0)
||
C)Estat.
(27)
Finally, with the notation in (22), define
T Σ 8,
(28)
iEI(O)
and note that
1 E (1 C))
since |8:] s (1+ C) for all e [d]. [I s d, and Tt( B(2Vd), for all e [t - 1]. Now, by Lemma E.1, in particular (23), have
Σ (Z
Eg.pPRoo Tunst
(()) as) (T(!
(1) 0% (mp, an) Tw w
stat
(29)
stat:
where Estat is as in (24). Using the expressions in Eqs. (26) and (28) with (29) and Jensen's 's inequality. we have that under E(O)
[0)00 (I), an) Tmm - (i an) T )()
Ex- Ex.pPExhhhuti Four
Σ |3,1
Σ En- (0(0 an) To((
o (2n, an) an) Teults))
()
(EZ(0
$ (1 C)dte
Now, using (27) and the facts that (a+b)2 2a² + 25 and 1 s 1, we have that
E p T Turir ((000(00)) an) To(" Φ7 an) T (0) 21
52 + 2(1++))=+++++++))))) C)dt stats
$ 2( 221 + 2(1 + C)dte?
stat.
25

---

Using that C= 1/2, we conclude that the right-hand side of this inequality is bounded by 8dt stat which is contradiction, since m. (1d)))(2d)2) B(3d
and by Lemma E.I, we must have
inf Ex-PP" To Tunir 3)
(6t an) NTu 0 (a an) ) StEstat
RepLearn does not terminate at round Therefore, we have that J(1) Ø, for any iteration t E [2..N] .J where RepLearn does not terminate.
We now bound the iteration count and prove that the guarantee in Eq. (21) holds at termination. Note that whenever J(t. Ø for /> we have by construction:
[de((V())) 3/2 (ve((v
Thus, if RepLearn runs for é [2 iterations. then
[det(V())) (1
(30)
On the other hand, since the determinant of matrix is bounded by the product of the norms of its columns and oid e B(2vd), we have
|dct(V())) s 24/2
Note also that (et(((())) = (=/d)². Plugging this into (30), we conclude that
(3/2)* /22d/2/))
Taking the logarithm on both sides and rearranging yields
1 d logg/22 (2d s
Thus, the algorithm must terminate after at most 1 iterations. Furthermore, by [35, Lemma 14], we have that with probability at least 2N the algorithm terminates at iteration t, then
max inf E B..PETo Manit (w' (wTT((0 (2n, ax) -ww -w000 (ar,, (a),,,n))) )n)
$ $31152 stat>
feF wEB(3d3/2)
32 1)
$ 32T stat
Applying a union bound completes the proof.
F Analysis
In this section, we prove the main guarantee for SpanRL (Theorem 3.2). First, outline proof strategy in Appendix F.1. Then, in Appendix F.2 and Appendix F.3, we present guarantees for the instances of PSDP (Algorithm 3) and RobustSpanner (Algorithm 2) used within SpanRL. We then combine these results Appendix F.5 to complete the proof of Theorem 3.2. self-contained guarantee for Robust RobustSpanner( 3.1) is given in Appendix F.6.
F.1 Proof Strategy
Like our algorithm, our analysis is inductive. For fixed h, we assume that the policy set (1sh+1) produced by SpanRL satisfies the property:
(1)
(h++) are ( -0-polyy covers for layers 1 through + 1, and max te(h+1]
|v(0 s d. (31)
Conditioned on this claim, we show that with high probability, the set (h+2) is Ind 0)-policy cover for layer h + 2. To prove this, we use the inductive assumption to show that PSDP acts as an approximate linear optimization oracle over W = { [<(*)] (x), Using this, we then instantiate the guarantee of Robus tSpanner from Lemma F.3 with LinOpt and LinEst instantiated with PSDP and EstVec. To conclude the proof of the inductive step, we the main guarantee for (Spannerr together with the main guarantee for RepLearn (Theorem E) along with a change of measure argument enabled by the assumption that (1=A) are policy covers (i.c. (31)).
(xh,, ax)] π Tl (Appendix F.2).
26

---

F.2 Guarantee for PSDP as Subroutine for RobustSpanner
We begin by showing that PSDP instantiates the approximate linear optimization oracle required by RobustSpanner particular, we fix a layer h and assume that (1:++1 satisfy (31) and apply the results of Appendix D.
More precisely, we need to show that, for any 0 E2 {0} and Φ e 4, PSDP approximately solves
max welly
O'R IET [o(an, (())n,as)] a-)]-
(32)
We can equivalently formulate (32) as, for fixed e IRd {0) and ee e Φ maximizing the sum of the reward functions TI:A (,) given by:
((z, a) E X' A, re(2,a;0,)) = o(x, a) Toj
for t h otherwise.
(33)
0,
Note that this matches the choice of reward functions in SpanRL (Algorithm 1) at iteration h with á = 35h), the feature map returned by RepLearn in Line 8. With these reward functions and the function class
g = =g (z, - a) w w we B(Vd)),
(34)
we show that the output मे PSDP r1:h (s:;0, Φ), ((Ii)), n) approximately solves (32) with high probability if n 2 1 is sufficiently large. We first verify that the class G realizes the reward functions specified in (33) the sense of Definition D.I.
Lemma F.1. Under Assumption 2.2, the function class 9 in (34) realizes the reward functions in (33), for any & & and 0 e IRd {0). Furthermore, we have that functions in uniformly bounded by & H, and log (o) s log [b] d log( H|)), where we recall that Ng(e) denotes the s-covering number G in loo distance.
Proof. Fix & e and 0 Rd {0), and let re(:.)) 0,(:,;0,6) 6), for & [h]. For = h, we clearly have that for any ι E IIhhh, Q. = re(.)) e g. For t h and I1 we have by the low-rank structure that
S Ke-i
Q7 (x),a)
E" ["(ra(ann an) x++1 y, a.+1 (v)) of (x. a (i (v)du(y),
= of (x, (Sx
E E"[ra(2hh an) xc+1 = 0++1 (y)1 pi (w)dv(y)).
(35)
Now, by the fact that ET Irn (rn,(an,,)) as) Ør+1 y, a.tt = (v)] [-1, 1]. for all y Xe.1 (since ((:+)) e B(1), for all & e ②), and the normalizing assumption made on (14) be[H) in Section 2.2 (i.e. that for all g: Ne.1 [0, 11. Mx.., (1) s vd), we have that
We == S
E. [rnn(<x an) x++1 = y, a41 {v)] (.. (v)du(y) e B(Vd).
X
This together with (35) and the fact that oj & (by Assumption 2.2), we have that Q7 e Combining Lemma E.I with Theorem D.2 results in the following bound on the quality of PSDP as an approximate linear optimization oracle over the space of policies.
O
Corollary F.I. Let 8,8 e (0,1) and h [H]. Further, let E Rdd X (0), & E o and = PSDP (h, (1:30.;;00)))) (1ih), n), where
The reward functions (1,60000,0 Φ) are as in (33).
The function class G is in (34).
The collection of policies (11))) satisfy (31).
The parameter is chosen such that cH ACd (n, 8/ S, where Estat (n, = ddn (dlog(nH)) log(()1/8)) and e> 0 is some large enough absolute constant. Then, under Assumption 2.2, with probability at least 1- we have that max T E ["((((,, an)] s 8 TT [' ax)] e/2. xell,
We emphasize that the inductive assumption that (1:h) is policy cover of bounded size enters only in the statement of Theorem D.2. We now give guarantee for RobustSpanner used in SpanRL.
27

---

Algorithm 7 EstVec(h,, n.): Estimate E= Require:
for given policy л and function F.
Target layer h [H].
Vector-valued function F: : A - Rd.
Policy π e Пl..
Number of samples N.
1: D Ø.
2: for times do
3: Sample (arn, 0%.)
4: Update dataset: D + D {(0 a.)).
5: Return: F= (x,,)eD F(2,a).
E3 Guarantee for Robust
Subroutine for SpanRL
In this section, we prove a guarantee for the instantiation of RobustSpanner in SpanRL, which we require in the proof of the main theorem (Theorem 3.2). We first show that the LinEst subroutine passed to BobustSpannes can be taken to be EstVec (Algorithm 7), which simply estimates the expected feature imbedding of (I), a,)) under policy by sampling sufficiently many trajectories and taking the empirical mean.
Lemma F.2 (Guarantee of EstVec). Let 8 E (0,1) and O. For h [H].
Φ, Tl.
and n e N such that n 2 & log(d/)) for some large enough absolute constant c O, the output t EstVec(h, 0, n) (Algorithm 7) satisfies, with probability at least 1-
1 - -on-E" E''''" an) s
Proof. By Hoeffding's inequality (see for example [22, Corollary 7)) and the fact that [|( a) s 1 for all X and A, there exists absolute constant such that with probability least 1 1-6 8,
log(d/5)
| J - -0 [ (x)(x))a))11 c.
Setting n as in the statement of the theorem concludes the proof.
In SpanRL, we instantiate Robust passing PSDP as LinOpt and EstVec as LinEst. Combin- ing Corollary F.1 and Lemma F.2 with the general guarantee of RobustSpanner in Proposition C.I, we have the following result.
Lemma F33. Consider iteration h [H] of SpanRL e, c) (Algorithm 1) with e 0, 8 E (0, 1), and feature class T satisfying Assumption 2.2. Further, let g(h) denote the feature map returned by RepLearn in Algorithm at iteration h. If (f-b)) satisfy (31) and polylog( A, d, H. log([ is large enough, then there is an event E,, with probability at least 2H such that
The number of iterations of Robusstpanner Line of Algorithm is at most N 1 loga
[100d] < co, and
c
. The output (Tı, va) of RobustSpanner has the property that for all n Hy, there exist $1..... 31 E [- such that
((0) (,b))) 3de, where ((h), EET { (2),
Proof. By Proposition C.I. the event that the instances of PSDP and EstVec used by RobustSpanner satisfy Assumption C.1 with e' 5. the two desiderata of the lemma hold." We
Here, instantiate the guarantee in Proposition C.I with C 2; this is what C' is set in Algorithu
28

---

claim that each call to PSDP and to EstVec satisfies Assumption C. with probability at least 1- SINH Because each of PSDP and EstVec get called at most 4dN times per iteration of RobustSpanner, union bound concludes the proof contingent on the above claim.
We prove the claim. First, note that the instance of PSDP that RobustSpanner uses within Algorithm is of the form:
PSD (.. (,,,,,,),,, v Resop)
with Pish and 9 as in Algorithm 1; this matches the form in Corollary F.1 (PSDP's guarantee) with Φ = ((h) Thus, by choosing
Tlesp c A ( + log(8dH12N(1//6))
g2
for = p:olllog((( (A, H. log({$//)) sufficiently large, the conditions of Corollary F.1 are satisfied, and its conclusion implies the claim for the PSDP instance used by RobustSpanmer. Similarly, the choice of Nesstvec in Algorithm ensures that the claim holds for the instance of EstVec that RobustSpanner uses by Lemma F.2. The result follows.
F.4 Guarantee for RepLearn as a Subroutine for SpanRL
In this section, we prove guarantee for the invocation of RepLearn within SpanRL
Recall that p(h) uni (()())) is the distribution over policies that SpanRL passes to RepLearn at iteration h [H 2] to compute feature map (A,), Thus, by invoking Theorem E. in Appendix E and using the choice of Naeeplear in Algorithm 1, we immediately obtain the following corollary.
Corollary F.2. Let 6,e e (0,1), 1), and F be as in Algorithm 1, and fix h e [H -2]. Suppose that the feature class & satisfies Assumption 2.2. Then, with probability at least 1 2 the instance of RepLearn in Line 9 of Algorithm 1 runs for S c- iterations for c= polylog( (A, d, H. log( (l))) sufficiently large, and returns output ((h) such that for all f F. there exists w satisfying
5
(h) B(3d112
Bunif(u((h))
Σ (o))) a)) 64A220
n/²
where Wy = SXL.. f(V) 1 (v)dv(y)
F.5 Concluding the Proof of Theorem 3.2
In this section, we conclude the proof of the main guarantee (Theorem 3.2). We derive the guarantee from the following inductive claim.
Theorem F.1. Consider iteration h E [H] of SpanRL($, c, 0) (Algorithm 1) with parameters e, > 0, 8 (0,1) and feature class & satisfying Assumption 2.2. Further, assume that:
The collection of policies ((1--=1) at the start of the hth iteration of SpanRL satisfy (31),
Assumption 2.1 (reachability) holds with 1
The input parameter e to SpanRL is set to E 364-/2-
The input parameter polylog (A, H, log({*1/))) is sufficiently large.
Then, with probability least 1 5 the policies (h produced by (o,nRL((D,,, the end of iteration h an (1 -)-llicy cover for layer +
With this, we can now prove Theorem 3.2.
Proof of Theorem 3.2. Note that it suffices prove that (31) holds for h H with probability at least - To do this, we proceed by induction over = 1,..., H 1. The base case of trivially holds because (11) Ø and (2) {munss) The induction step now follows by Theorem F.1 and the union bound (see e.g. [30, Lemma I.2]).
29

---

The number of trajectories used by SpanRL is dominated by calls to PSDP. Since PSDP is called O(d (d/=)) times at each iteration of SpanRL (Lemma F.3), and each call PSDP requires at most M trajectories, the total number of trajectories after H iterations of SpanRL is bounded by dmpspp) By plugging the choices for nosop and from the theorem statement, we obtain the claimed sample complexity.
Before proving Theorem F.I, make the following simple observation. Lemma F.4. For any л Ily h [H - 1], any x 2 we have
Ph-1 (z) " [ (zt, ==)]= d 0.
Proof of Lemma F.4. The equality follows by construction. The non-negativity of dl (x) follows by definition of a probability density.
We now prove Theorem F.I.
Proof of Theorem F.1. Let E, and E/ denote the success events in Lemma F.3 and Corollary F.2. respectively, and note that by the union bound, we have [& n 2 - 6/ H. For the rest of this proof, we will condition on & = & n.
Throughout, we denote
0/ E oo Th, an)), Vh e [H], VT e Il.
Because (I+h+++ satisfy (31) (i.e., are policy cover) it holds by Lemma F.4 that for all a e Xn.
max (x) 0 . TEDON
sup (z)) oi-11 for 4Ad
(36)
sell,
We will show that with probability at least layer h+ 2; that is, for all Xh+1,
H. the policy set v)(h=2) has the same property for
max (x) Phill Tew(h+2)
2 a sup (x) 0hi-1
"
(37)
xells
Again, by Lemma F.4 this is equivalent the statement that y is ( -p-lliy cover for layer h + 2.
For the remainder of the proof, we will fix e Xx42 and let arg mal, 2 (x) o Our goal is to show that the inequality Eq. (37) holds for x.
Preliminaries. Note that since Xn+2, we have [1))+(())| O. It will be convenient introduce a function f Xp+11 + defined by
(h-2(1)
f(y) (3, where 0g
Further, we define
Wg := SxL (w)d(())
(38)
By definition of T2, we have that for all y Kn+1,
0 (g. n (1)) max 0J oi
aeA z (h,
This together with the fact that |0 implies that
f x + max ó(x, a)) 0 o ((1), c aEA
(39)
the discriminator class in Line of SpanRL. Note also that since x E Xh+2, we have by reachability that
1
max T0 /+1 2 27> 0.
(40)
(h)2(x))
meth
30

---

Applying the guarantee for RepLearn. Moving forward, let g(a)) be the feature map returned by RepLearn at the hth iteration of Algorithm 1, and define g(h), == [ (T(()))) a,). for any x E 1ly.
Further, let wa (k) be the vector vo (h) in Corollary F.2 with f fz. and note that |0 s 3d
(41)
By Jensen's 's inequality, we compute
((w?))) 10 (h)) (wz) )
an) Tw( (h) 0f (2h, aaaa)="uu " ) (Jensen's inequality)
s T
(1, 10 (1) (y, (v)) wx)) (v) *r du()), (Low-Rank MDP)
S a
max #el((n)
(0.h) w (3 w. Mi o dv(), (by (36)
X
Σ
(1)
(1) o (by Lemma F.4)
tev(h))
Nr.
S Sa"
Σ Σ
(v. (v,a) (v. a)" ) (1) - dv(w),
(EV(R) aEA
N/.
(h) (ah; an) wxz) (16 (r, ) u Aa-'dd punir(v(n))
(42)
where the last step follows by the definition of (h) in Algorithm and that [T(()) d. Now, since wr _xx,, f1))i (y)dv(y) (see (38) and f F (see (39)); the guarantee for RepLearnin Corollary F.2 together with (42) implies that (conditioned on the event E)
Adn2 " $ 4"
- S
(43)
A4
Applying the guarantee for RobustSpanner. Letting WI:-.. RobustSpanner at iteration h of SpanRL, the guarantee of RobustSpanner in Lemma F.3 implies that there exist ....... 3r E [-2,2] such that
Fd be the policies returned by
(1))
Σ (
S 3de S
" 12/3/22
(44)
where the last inequality follows by the fact that triangle inequality, we get that
Combining (44) with (43) and using the
36/-/2
+ |100111
n
w 2
12d3/2
4'
1=1
$ Σ
3: 00
11 7 4
(by (41))
4'
1=1
$ 2d max w it[d]
T.
2
Combining this with (40) and rearranging implies
wIo WIth
Fs s 4d max w
(45)
sc[d]
On the other hand, by definition of Was have
max w. Jo max 0 0. m.
[d]
:c[d]
1
max E 181
[2)=z(x)]
Th+1 ( (2H4),
1ha2(0n (x))
[e|d]
max 1 1he Tunts (2)=2(x) thes T (2he1s ant)), (see below)
(x))
(e[d]
A
max 11+2(2)) 1 Tew(112)
(46)
31

---

where the inequality follows from the non-negativiyy of Mi+ (-) h- (2, for all (z,a) & X₁ x A (due Lemma F.4), and (46) follows from the definition of (1(22)) in Line 13 of Algorithm 1. Combining (45) and (46) then implies that
1 (h (hraat) 0u-1 0 = w mm
S 4d max w o
s
4Ad
(x))
(z)" Ph-1
(())
xew(h-2)
This, together with Lemma F.4, implies that (37) holds. Since this argument holds uniformly for all I E 2422, this completes the proof.
F.6 Proof of Lemma 3.1
By definition for x X1+1, have d= = maximizing d= (if no such maximizer exists, may pass maximizing sequence) and let (Ti,...,Tur Ta)- Then, we have for some $I,..., ECC,C,
ET [M/-1 (z) (1)'o), (z am)] Let Tue denote the policy
d (x) Min (T)
(
E Bi0h
+ +
-
d - 3,0 *
Cd max pha (x)T<a (E[d]
+ e (x)),
(Cauchy-Schwarz)
S Cd. max 1h (2) $5 | (x),
it[d]
where the inequality follows by the fact that Assumption 2.1 holds with s n/22. The result now follows by rearranging.
G Application to Reward-Based RL
In this section, we explain how the output (I:H)) of SpanRL (Algorithm 1) can be used to optimize downstream reward functions T1:1r: our treatment standard. Since the output of SpanRL a policy cover, one way to optimize the sum of rewards Sμ Σ rh is by first generating trajectories using policies in ((1#H), then applying an offline RL algorithm, e.g. Fitted Q-lteration (FQI) [16], to optimize Sh- It is also possible use PSDP with the policy cover (1:tt)) achieve the same goal. We will showcase the latter approach since we have already stated a guarantee for PSDP.
As in Appendix D, we assume access function class {g X A R} that realizes the rewards FI4H in the following sense: for all h [H] and all π II
Q% e G. where ( a) E
Σ r((2,, a:) 1-+++1
3hh = x, a₁ az =
Note that when the reward functions F1- are linear in the feature map Φi: that is, when for all h é [H] and (ar, a) M, A, ra (2,, a) f df (x. a) for some 9 B(1) (a common assumption in the context of RL in Low-Rank MDPs [32. 31, 49, 35]), then the function class
G = {g: (z, p(x a) Φ 2, we B(2HVd)).
realizes fi: Note that G is the same function class we used for the PSDP subroutine in Algorithm 3, albeit with larger ball for the w's. For the sake of generality, state the next result (which shows how to use a policy cover to optimize downstream reward function) for general f1:H and a function class that realizes ri-H.
Theorem G.1. Let E 0 and d E (0, 1). Suppose reward functions TH: X x A + R, a collection of policies (1:H), and a parameter n. 2 1 satisfy the following:
The function class G realizes the reward functions тı: (in the sense of Definition D.I), and lim n log NoggN/ O, where (1/1/)) to denote the --overrigg number of G in the supremum norm. Furthermore, we suppose that functions in G are uniformly bounded by H Hvd.
32

---

For some < S 1, for each 1 < h & H. it holds that (x(h) is an a-policy cover for layer and moreover [ s d.
The parameter is chosen such that cdHo Estat (n, E, where Estat (n, ==
VdH (logNV log(1/))) and is large enough absolute constant.
Then, with probability at least 1 the policy क PSDP H, F1.H, 9, p(..)) n) (where p(t)
unif (()(()))))))))))))))))))) for each t [hi]) coming from Algorithm 3, satisfies the following guarantee:
max ET Tell
Σ rh (2 an)
S E
Σ
rn an)
tE.
Moreover, the number of episodes used by PSDP in this case is
A 11 (log (log)o(s)) log(1/))) o
g²
Proof. This is simply a restatement of Theorem D.2 with h H. The number of trajectories follows by the fact that each call to PSDP requires Hn trajectories.
H Properties of Reachability Assumption
In this section, we compare 17-reachability (Assumption 2.1) ) to different reachability assumptions used in the literature in the context of RL in Low-Rank MDPs and show that ours is the weakest among those commonly assumed. In Appendix H.1, we demonstrate an exponential separation between our notion of reachability and that considered with respect the popular latent variable model [1, 35J. In Appendix H.2, we consider a number of other reachability assumptions made outside the latent variable model and show how they imply Assumption 2.1.
H.1 Comparison to Latent Variable Model
In this subsection, we show that our reachability assumption is implied reachability assumption used by [1, 35] in the latent variable/non-neetee negattive feature model, and show that our reachability assumption can hold even when the best possible latent variable embedding dimension is exponential in the dimension d. We begin by defining the latent variable model.
Definition H.1 (Latent variable representation). Givn a transition operator T: X' x - A(X), latent variable representation consists of a countable latent space Z and functions X AZ) and q: Z + A(X). such that T(- x. a) = Ez q( 2)v(21 æ,a). The latent variable dimension T, denoted dey is the cardinality of smallest latent space Z for which T admits a latent variable representation.
The interpretation for the latent variable model is follows:
Each (z, a) pair induces a distribution (xx, a) A(Z) over ZE Z.
2. The latent variable is sampled as z (2, a).
3. The next state is sampled x' x'~~0 9-(-1 2).
Note that in discrete state spaces, all transition operators admit trivial latent variable representation, as we may take (xx a) = T(- xx, a), but the dimension of such a representation is potentially infinite. latent variable representation certifies that there exists a factorization T(x' a) = (xxx, a) "q(aa) with embedding dimension [2]. and so dev, and hence gives an upper bound on the rank of the transition operator. On the other hand, compared with the general Law-Rank factorization, the latent variable factorization additionally requires that (xx a) and g(- 2) are probability distributions, and thus nen-negative, for all z Z and (x,, a) 6 X x A, implying that dlv is equivalent to the non-negative rank [1] of the transition operator.
Assuming that latent variable representation exists, [1, 35] consider the following notion of reacha- bility.
Definition H.2 (Reachability in latent variable model). There exists > such that sup
4h < [H 1], Vz e Zh-1.
|p [= =] 2 1).
(47)
xell.
33

---

We first show the latent variable reachability condition above implies our more general assumption.
Lemma H.1. Consider a Low-Rank MDP M with rank 2 1. Under the latent variable model in Definition H.1, if the latent variable reachability condition in (47) is satisfied for some 1) O, then, for all h [ the transition kernel Th M admits a factorization Th(:] x, a) 1 0f (v. a). where phe1 (-) TRdv and (--) TRdu, such that duy ddA/np and n² Eaa-reachability (in the sense
A/a
of Assumption 2.1) is satisfied.
Proof of Lemma H.1. Suppose that Assumption 2.1 (1-reachabilitt) holds. By Agarwal et al. [1, Proposition 4], the non-negative rank of M is bounded as dLy s d/²/11².
Letting q and vp be as in the definition of the latent variable representation in Definition H.1, we define ph+1 and 0f as: for all h [H- -1.
1 (+) = (g(- 2) 2z 18x, and (T((,)) = (v/(z (::)) e Rd
Now, fix h [H - 1] and 1x+1- For arg 2e2.1 ((= z), we have
sup d (x) = [*h+1 = x]
sup Σ q( 2) 1E [ on an)
mell,
nelly 24 2e
sup g(x 20) E (1) wh. an)];
melly
sup PP [zh-1 = s0].
xell
2 17 11h+1 (z)
(using reachability)
2 "
vdv | (x)ll.
We now complement the result above by showing that there exists low-rank MDPs for which our notion of reachability (Assumption 2.1) is satisfied with n polynomially small, yet the best possible latent variable embedding has dimension dly = 2n((), This contrasts the results in [1, Proposition 2], which show that latent variable reachability implies a polynomial bound on the latent variable dimension.
Theorem H.3. There exists a one-step Low-Rank-MDD of rank d 2 1, where n-reachability (Assump- 1 but where the non-negative rank satisfies dev =2º(4)_
tion 2.1) is satisfied with 1) 2/a'
Proof of Theorem H.3. Let N and d (*) + 1. As shown in the proof of Agarwal et al. [1, Proposition 2), there exists a horizon-two MDP M with the following properties:
The state spaces Xi and Xs layers and 2, respectively, finite.
The cardinality .A is d; i.. A= (a.....,aa). aa)."
The transition kernel Ti admits the factorization:
(- a) (i e A(X)), ((I,a)&X xA,
where for all X2, p (x') e I&sp and for all (2, e X x A, oi (z, a) E R
The non-negative rank of M is dey 20()
We augment this MDP by adding an extra state xg. and let X' == Xi u(o We define Φi : :x:xxxx A- R.g be the extension of di given by
e [d], (1) = and Vre X1, ((
"Techncallyy, the example in the proof of [1, Proposition 2] does not explicitly specify the number of actions. Instead, the example assigns a number of state-action pairs to vectors in RR, without specifying the number of actions. The number of actions in their example is a degree of freedom, which we set to d here without loss of generality.
34

---

where e; is the ith basis element in R.. We define the initial state distribution to have p( 2/2 and p(x) 2 for all X1. 12 We let M (A1 u A2 Φi (H) be p) denote the resulting MDP. Note that adding an extra state at layer in this fashion only adds d additional rows to the transition matrix T (viewed as (1 AI) x [2] matrix). Therefore, the non-negative rank of M is as least that of M.
We now show that reachability satisfied in M. Let n, the policy that always plays action a,. With this, we have that for any x' é X2,
sup d (x') 2 max d (x'): relly
se[d]
max u (x') E B[ (2i,
x[d]
max *[[]
[E[ Φ mo (=(2')
[[[{ + zo) p²(x') (z) a;]]}
2 max max))(=))_ कं (z).a;) a.).
(48)
x<[d]
where the last inequality follows by the fact that, for all (x) X, x A, 15
1z (xx)T (x,a) a) 2 (since μ (x')" o1 (2x a) is a conditional density). On the other hand, from the construction of ø and the fact that p2 (x') Rgo we have
max 12)("))))))) (vo, a:) 1x2(0') InG(s')IIVd.
<e[d]
Combining this with (48) and using that p(zo) = 1/2 implies that reachability 1/(2vd) is satisfied in M.
H.2
Relation Other Reachability Assumptions
In this subsection, show that Assumption 2.1 is implied by notion of feature coverage used in the context of transfer learning in Low-Rank MDPs [5], as well as a notion of explorability used in the context of reward-free RI in linear MDPs [46].
H.2.1 Feature Coverage
We first consider coverage condition used by Agarwal et al. [5], which involves the second moments of the feature map o
Definition H.4 (n--fetuure coverage). We say that the linear MDP with featurization o7 satisfies n-feature coverage for all h [H],
sup Amin rella
RE (R" (2n, dn (Ih, (1) 2
We show that ry-feature coverage implies (n/2) ( reachability, Thus, up to polynomial dependence, feature coverage is a special case of Assumption 2.1.
Lemma H.2. Suppose that an MDP satisfies n-feature coverage as in Definition H.4 for some 7i > O. 4J (,, a) e B(1) for all x, a, then the MDP satisfies Assumption 2.1.
in the sense of
Proof of Lemma H.2. Let h [H] and æ 2r+1 be given and define
Mh+1 (I)
0
|M (x)
To keep notation compact, define o; (#n, an). By f--feauure coverage, there exists Il such that
7 E [ (TT+))))))))))))))))))))))))))))
= 01,² <n/)) (0 Φ +1 ETTH(0 (1)² n/2) (0T 61) S n/2 + P7 (0 Φ; 2 n/22,
(49)
We note that [1] did specify the initial distribution, which not needed for the conclusion of their result.
35

---

where we have used that |0] 1 and (z, | s for all (z, 6 Xr x A. Rearranging (49) and using that 0J Φं 0 (it a scaled conditional density), have
[ Φ 2 n/2] P0 [(0'((00'' o;))² 2 n/2] n/2.
Now, by Markov's inequality, we have that
o = 1 Φ 2 P Φ 2 V n/2] 2 (n/2)
where we have once more used that A Φ; 2 0 almost surely.
H.2.2 Explorability
We now consider the explorability assumption of [46], which involves the first moment of the feature map o This notion is defined as follows.
Definition H.5 (1-explorabilitty. We say that a linear MDP satisfies n-explorability ffr any h e [H] and any A e Rd {0} it holds that
sup 1 [0 (Th, an)]l 2 |10]]
mellm
We now show that --explorability is special case of n-reachabiliity
Lemma H.3. Suppose that the explorability condition in Definition H.5 is satisfied with 17 > O. Then, qy-rachabiility satisfied.
Proof of Lemma H.3. Let x A4+1 and define 0= (= By explorability, we have that
sup d= (x) = sup E= rella
E" [ri-1 (x) oi (En, a+)),
rell,
sup tell,
|E" [1"( (x) of (Ih, an)] (5) a) is condition law)
sup (0"E"(00 (mn, an }}
rell,
2 n [0], (by explorability) = (1)+1
This shows that Assumption 2. is satisfied with parameter 7).
36


Reviews:
[{"review_id": "4IuHWO6Ujm", "reviewer": "Reviewer_VYmd", "summary": "The paper proposes a new algorithm for the low-rank MDPs setting. With a new reachability assumption, the paper shows that one can learn a policy cover in a statistically and computationally efficient manner. The proposed algorithm works in a modular manner, using barycentric spanner for exploration and the replearn algorithm for representation learning. The paper proves that the proposed algorithm can indeed learn a policy cover with in poly number of samples. ", "questions": "One argument that consistently appears across the paper is that the merit for a model-free algorithm is to avoid making assumption on an additional function class $\\Upsilon$, could you elaborate what is the practical significance of such condition? ", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The overall intuition of the proposed algorithm is reasonable, the design of the algorithm is well motivated, and the technical parts of the paper seems right. \n\n2. A computationally tractable model-free algorithm is a very good contribution for the low-rank MDPs setting. \n\n3. The reachability assumption is novel and general, but it might require a better explanation for the intuition of the such reachability assumption: while previous visitation density or feature coverage assumptions are more intuitive, using $\\mu$ seems less straightforward.\n\n4. The proposed algorithm is modular and utilizes previous algorithms well. ", "weaknesses": "1. Unlike previous papers such as BRIEE or MUSIK, the proposed computational efficient algorithm is not empirically tested, but this is a rather small shortcoming since the major contribution of the paper is theoretical.\n\n2. The paper requires more polishing. The ending of the paper seems rather sudden and abrupt, although plenty discussions have been given across the paper. There are other minor presentation issues, such as the FLAMBE paper is cited in two different versions. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments.\n\n**“Unlike previous papers such as BRIEE or MUSIK, the proposed computational efficient algorithm is not empirically tested…”** \nRegarding the lack of experiments, the contribution of the paper is primarily theoretical, and we believe that our theoretical contributions alone are significant enough to merit acceptance. Note however that because recent papers demonstrate the empirical success of BRIEE [Zhang et al. 2022] and MusIK [Mhammedi et al. 2023] and, because our computational oracles are identical or simpler than those found in these works, we believe there are strong reasons to expect that the assumptions underlying our work hold, and that the relevant oracles can be implemented in a straightforward and practically effective way.\n\n**“The paper requires more polishing. The ending of the paper seems rather sudden and abrupt, although plenty discussions have been given across the paper…”**\nThank you for the suggestion. In the revision, we will be sure to add a conclusion and improve any minor presentation issues.\n\n**“One argument that consistently appears across the paper is that the merit for a model-free algorithm…  could you elaborate what is the practical significance of such condition?”**\nThe fact that assuming access to a model class constitutes a non-trivial additional assumption is widely acknowledged in the literature (e.g., Zhang et al. 2022, Mhammedi et al. 2023). Consider the task of learning to control in robotics or atari with images (raw pixels) as observations. Here, model-based methods must learn a generative model mapping latent states to images, whereas model-free methods like SpanRL only need to learn a mapping from images to latent states; it is clear that the former task is significantly more challenging than the latter. Thus, it is of both theoretical and practical interest to develop an algorithm that does not rely on a model class in order to achieve a smaller sample complexity.\n\n**References:**\n\nXuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun. “Efficient reinforcement learning in block mdps: A model-free representation learning approach”. In International Conference on Machine Learning, pages 26517–26547. PMLR, 2022.\n\nZakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin. “Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observation RL”. International Conference on Machine Learning (ICML), 2023.\n"}, {"review_id": "zeh8G1QBaj", "reviewer": "Reviewer_4TUV", "summary": "This paper studies model-free RL under low-rank MDPs. This paper first generalizes previous fromulation of low-rank MDPs to a uncountable state case. Then this paper proposes a non-optimistic, model-free RL algorithm called SpanRL, which can provide a good policy cover with polynomial sample complexity. With policy cover produced by SpanRL,  this paper optimize any downstream reward function with polynomial sample complexity.", "questions": "As listed above.", "limitations": "There is no negative impact in this work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper provides a comprehensive summary of prior research on low-rank MDPs and presents a clear table that compares their sample complexity results with previous work.\n\n2. This paper introduces a novel model-free exploration algorithm for general low-rank MDPs, which improves previous work on model-free RL under low-rank MDPs in two aspects: requiring a less stringent reachability assumption and removing the non-negativity requirement for feature embeddings.\n", "weaknesses": "1. One advantage of the algorithm emphasized in this paper is its computational efficiency. However, there is no experiments to validate this claim.\n\n2. Due to the page limitations, several algorithms (subroutines) are deferred to the Appendix, which makes it hard to track the algorithmic steps. Additionally, I find there is no detailed descriptions and discussions of these algorithms. A high-level description of the algorithms will be helpful for reader to understand the novelty these algorithms. For instance, algorithm 2 is an error-tolerant variant of the classical spanner computation algorithm of [1], it would be beneficial if the authors can discuss in high level about the design modification and novelty of algorithm 2.\n\n3. The machinism why SpanRL can improve previous work is not explain well. For example, which design of algorithm helps to remove the non-nagetivity assumption in [2]? In other words, why does [2] need such assumption while SpanRL doesn't need such an assumption? \n\n[1] B. Awerbuch and R. Kleinberg. Online linear optimization and adaptive routing. Journal of Computer and System Sciences, 74(1):97–114, 2008.\n\n[2] A. Modi, J. Chen, A. Krishnamurthy, N. Jiang, and A. Agarwal. Model-free representation\nlearning and exploration in low-rank mdps. CoRR, abs/2102.07035, 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments and helpful suggestions!\n\n**“One advantage of the algorithm emphasized in this paper is its computational efficiency. However, there is no experiments to validate this claim.”**\nRegarding the lack of experiments, the contribution of the paper is primarily theoretical, and we believe that our theoretical contributions alone are significant enough to merit acceptance. Note however that because recent papers demonstrate the empirical success of BRIEE [Zhang et al. 2022] and MusIK [Mhammedi et al. 2023] and, because our computational oracles are identical or simpler than those found in these works, we believe there are strong reasons to expect that the assumptions underlying our work hold, and that the relevant oracles can be implemented in a straightforward and practically effective way.\n\n\n**“Due to the page limitations, several algorithms (subroutines) are deferred to the Appendix…”**\nThank you for the suggestion. In the revision, we will include a more detailed discussion of the motivation behind our algorithm design in the main body (especially Algorithm 2, which was deferred to the appendix due to space). Note that Algorithm 2—the robust variant of the barycentric spanner algorithm—is meaningfully different from the original algorithm of Awerbuch and Kleinberg. Most notably, the algorithm “fattens” the set in directions that would otherwise be too “skinny” for one to apply an approximate minimization oracle (as opposed to an exact one); note that without this fattening step ( i.e. naively applying the algorithm from Awerbuch and Kleinberg with an approximate minimization oracle), the algorithm would not work due to the presence of these “skinny” directions, necessitating our novel variant.\n\n\n**“The machinism why SpanRL can improve previous work is not explain well…?”**\nPrior works that make use of the same representation learning objective (BRIEE [Zhang et al. 2022] and MOFFLE [Modi et al. 2021]) do not make use of spanners; instead, they appeal to exploration strategies based on elliptic bonuses, addressing the issue of approximation errors through additional assumptions (non-negativity of the factorization for MOFFLE, and Block MDP structure for BRIEE). Perhaps the most important observation in our proof is that barycentric spanners are robust to the average-case approximation error guarantee in (5) as-is, without additional structural assumptions. Intuitively, this benefit seems to arise from the fact that the spanner property only concerns the first moment of the feature map, while algorithms based on elliptic bonuses require approximation guarantees for the second moment; understanding this issue more deeply is an interesting question for future work.\n\n\n\n**References:**\n\nXuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun. “Efficient reinforcement learning in block mdps: A model-free representation learning approach”. In International Conference on Machine Learning, pages 26517–26547. PMLR, 2022.\n\nAditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. “Model-free representation learning and exploration in low-rank mdps”. CoRR, abs/2102.07035, 2021.\n\nZakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin. “Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observationRL”l. International Conference on Machine Learning (ICML), 2023. \n"}, {"review_id": "MUDc7oYt5d", "reviewer": "Reviewer_y7X6", "summary": "The paper studies the reward-free exploration problem for low-rank MDPs. The approach is based on the combination of the barycentric spanner and a model-free representation module. The authors provide a theoretical analysis of the algorithm for both the reward-free exploration task and the downstream reward-based task. ", "questions": "1. Since the representation learning part is used as a plug-in module of the algorithm, one may wonder what the theoretical guarantee will be like when the true state-action feature is given. Furthermore, how does this result compare with the existing literature for known state-action features? \n\n2. The paper proposes a new reachability notion. While sufficient comparison has been given between it and existing notions, will this notion reduce to a more familiar concept when considering simple MDP settings, such as tabular MDP or block MDP? ", "limitations": "The paper addressed some of its limitations and proposed some open questions for future work. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The writing is clear and easy to follow. Both the representation learning problem and the reward-free exploration problem are important in RL literature. The result is significant, which applies to more general settings compared to existing works, and also has a comparable theoretical guarantee. The idea of combining barycentric spanner with representation learning is novel. ", "weaknesses": "1. While the paper studies the exploration problem of low-rank MDPs, its contribution mainly lies in the planning phase. The paper has limited contribution on the representation learning part, which simply reuses the algorithm and theoretical guarantee from existing RL papers. \n\n2. The main theorem provides a guarantee on the total number of episodes. It might make the result more clear if the number of samples / trajectories can also be stated, just as the theorems in previous papers. Since the algorithm repeatedly collects samples in each episode, some being hidden in the RobustSpanner module, it is not easy to figure out the sample complexity at first look. \n\n3. The reward-based downstream task is slightly different from existing reward-free framework. A typical reward-free exploration algorithm would only collect data in the exploration phase, and only use these collected data to construct a good policy for an arbitrary reward function. In the setting of this paper, the planning phase relies on the PSPD module, which seems to need additional interaction with the environment because the module itself (Algo 3, line 4-5) needs to sample data and construct a separate dataset. It's not clear whether this additional interaction is avoidable. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments and helpful suggestions!\n\n**“While the paper studies the exploration problem of low-rank MDPs, its contribution mainly lies in the planning phase…”**\nThank you for your review. We agree that our primary algorithmic contribution is not in the representation learning aspect (although our representation learning is slightly simpler than that found in earlier work due to the simpler discriminator class that we use), but rather in providing a rigorous algorithm for *exploration* in low-rank MDPs.  We emphasize that while the representation learning objective itself is not new, our analysis *is* new, and shows for the first time that the objective can lead to meaningful guarantees in general Low-Rank MDPs without latent variable structure or Block MDP structure. We view this as an important conceptual contribution.\n\n**“The main theorem provides a guarantee on the total number of episodes…”**\nWe use the term “episode” and “trajectory” interchangeably, and the sample complexity bound stated in the main theorem refers to the total number of trajectories. We will clarify this in the revision.\n\n**“The reward-based downstream task is slightly different from existing reward-free framework…”**\nOnline access to the MDP for optimizing downstream rewards is not actually required. In fact, it is possible to use Fitted-Q-Iteration (FQI) [Ersnt et al. 2005] instead of PSDP for a purely offline planning phase once trajectories are generated using the policy cover produced by our algorithm. The only reason we opted to present the reward-based results using PSDP was to keep the length of the paper down (as we already use PSDP within the analysis of SpanRL). We are happy to include a proof that FQI can be used to optimize downstream rewards for a purely offline planning phase in the final revision if this is of interest to the reviewer.\n\n**“Since the representation learning part is used as a plug-in module of the algorithm…”**\nNote that if the true feature map is given, then this reduces to the linear MDP setting and our algorithm can be run in the same way, but with the representation learning step removed.  The algorithm would produce a policy cover, which could then be used within PSDP to optimize any downstream rewards. However, compared to existing linear MDP algorithms such as LSVI-UCB [Jin et al. 2019], our algorithm would be non-optimistic and would require reachability as an additional assumption. We suspect that reachability can be removed using a more refined analysis and techniques from [Mhammedi et al. 2023] which specifically deals with this.\n\n\n**“The paper proposes a new reachability notion… will this notion reduce to a more familiar concept when considering simple MDP settings, such as tabular MDP or block MDP?”**\nEssentially all existing reachability assumptions for low-rank MDPs that we are aware of reduce to the standard reachability assumption for tabular/Block MDPs [Misra et al. 2020, Mhammedi et al. 2023]. This is perhaps most easily seen through our assumption (Assumption 2.1), which we now elucidate. Observe that in the Block MDP setting, we can take $\\mu_\\star(x)= q(x\\mid \\psi_\\star(x)) \\cdot e_{\\psi_\\star(x)} \\in \\mathbb{R}^S$ [Du et al. 2021], where $\\psi_\\star$ denotes the true decoder in the Block MDP setup and $q(x | \\psi_\\star(x))$ is the true emission distribution. Using this choice of $\\mu_\\star$, along with the fact that $d^{\\pi}(x)= q(x\\mid \\psi_\\star(x)) \\cdot d^{\\pi}(\\psi_\\star(x))$, one can see that our reachability assumption, which asserts that $\\max_{\\pi \\in \\Pi} d^{\\pi}(x) \\geq \\eta \\cdot \\lVert\\mu_{\\star}(x)\\rVert$, for all $x$ can equivalently be rewritten as $\\max_{\\pi \\in \\Pi} d^{\\pi}(\\psi_\\star(x)) \\geq \\eta$ for all $x$, after cancelling $q(x\\mid \\psi_\\star(x))$ and using that $\\lVert e_{\\psi_\\star}(x)\\rVert=1$. This is equivalent to the reachability assumption in the Block MDP setting [Misra et al. 2020, Mhammedi et al. 2023].\n\n\n**References:**\n\nDamien Ernst, Pierre Geurts, and Louis Wehenkel. \"Tree-based batch mode reinforcement learning.\" Journal of Machine Learning Research 6 (2005).\n\nDipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. \"Kinematic state abstraction and provably efficient rich-observation reinforcement learning.\" In International conference on machine learning, pp. 6961-6971. PMLR, 2020.\n\nChi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I. Jordan. \"Provably efficient reinforcement learning with linear function approximation.\" In Conference on Learning Theory, pp. 2137-2143. PMLR, 2020.\n\nZakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin. “Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observation RL”. International Conference on Machine Learning (ICML), 2023. \n"}, {"review_id": "6vF3Wqw0fL", "reviewer": "Reviewer_qxYM", "summary": "This paper studied model-free exploration in low-rank MDP. The authors proposed an algorithm based on the idea of barycentric spanner, which is provably efficient given a reachability assumption weaker than previous literatures.", "questions": "See weakness.", "limitations": "N.A.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper writing is clear. The comparsion with previous related work and the design idea of the algorithms are clear.\n\nThe paper provided results with a weaker reachability assumption than previous literatures.", "weaknesses": "I have doubt about whether the algorithm can be called \"efficient\". It still requires to solve a complex optimization problem in Eq.(6), especially the objective function can be non-convex/concave w.r.t. feature class $\\Phi$ and function class $\\mathcal{F}$, even though many previous literatures claimed that \"it can be efficiently solved\".\n\nThe sample complexity should be improved.\n\nAlthough this paper considered the model-free setting and realizability of $\\phi^*$ is much milder than the model realizability, it still requires reachability assumption in some sense, while the previous model-based algorithms (like [1]) does not.\n\n[1] Uehara et. al., Representation Learning for Online and Offline RL in Low-rank MDPs", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and useful suggestions!\n\n**“I have doubt about whether the algorithm can be called \"efficient\"...”**\nRegarding the computational efficiency of our algorithm, as noted, many other works use the same minimax oracle (in fact, our oracle is slightly weaker because the maximum is taken over a simpler function class). While considering alternatives to this oracle is an interesting direction for future research, several works have empirically demonstrated the viability of this minimax oracle. Note also, that even after assuming such an oracle, the problem of how to use this oracle to rigorously explore in the low-rank setting is far from trivial.\n\n**“The sample complexity should be improved.”**\nRegarding the sample complexity, with a slight modification to the current proof, we have figured out how to improve the sample complexity to $O(1/\\varepsilon^2)$. We will include this in the revision.\n\n**“Although this paper considered the model-free setting… it still requires reachability assumption in some sense, while the previous model-based algorithms (like [1]) does not.”**\nWhile it is true that model-based algorithms do not require reachability, note that assuming access to a model class is a nontrivial extra assumption. Indeed, it is often thought that learning the model is significantly more challenging than learning the feature map $\\phi_\\star$. Consider the task of learning to control in robotics or atari with images (raw pixels) as observations. Here, model-based methods must learn a generative model mapping latent states to images, whereas model-free methods like SpanRL only need to learn a mapping from images to latent states; it is clear that the former task is significantly more challenging than the latter. Thus, it is of both theoretical and practical interest to develop an algorithm that does not rely on a model class in order to achieve a smaller sample complexity. Removing reachability is a natural next step, but this is somewhat orthogonal to our main contribution. We suspect this can be achieved using a more refined analysis and techniques from [Mhammedi et al. 2023] which specifically looks at removing reachability in the Block MDP setting.\n\n**References:**\n\nZakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin. “Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observation RL”. International Conference on Machine Learning (ICML), 2023.\n"}]

Author rebuttal:

Meta-review: After discussion, all reviewers agree that the new model-free algorithm for low-rank MDP is interesting, and the paper should be accepted.
Decision: Accept (poster)