Title: Efficient Model-Free Exploration in Low-Rank MDPs
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

Amajor 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.

1 Introduction

In reinforcement learning and control, many of the most promising application domains require the agent to navigate complex, high-dimensional state and action spaces, where generalization and function approximation is necessary. The last decade has witnessed impressive empirical success in domains where data are abundant [34, 38, 26, 28, 27], but when data are limited, ensuring efficient exploration in large domains is a major research question. For statistical efficiency , the foundations have recently begun to take shape, with a line of research providing structural conditions that facilitate sample-efficient exploration, as well as fundamental limits [37, 21, 39, 42, 14, 24, 17, 18]. Computational efficiency , however, remains a major challenge: outside of simple settings [7, 23], existing algorithms with provable sample complexity guarantees are computationally inefficient, and typically require solving intractable non-convex optimization problems [21, 11, 24, 10]. The prospect of developing practical algorithms for exploration in high-dimensional state spaces that are both computationally and statistically efficient raises three fundamental questions:

What are the right computational primitives for exploration? That is, how can one efficiently represent and compute exploratory policies that allow the learner to explore the state space and gather useful data?

How should one leverage function approximation-for example, via representation learning-to discover such primitives in a computationally and statistically efficient fashion?

Given answers to the first two questions, how can one efficiently interleave function approximation and exploration to provide provably efficient algorithms?

In this paper, we investigate these questions through the Low-Rank MDP model [36, 45, 1]. In a Low-Rank MDP, the state space is large and potentially continuous, but the transition probabilities admit an (unknown) low-rank factorization. Concretely, for a finite-horizon Low-Rank MDP with horizon H , the transition densities for layer h ∈ [ H ] satisfy

where ϕ ⋆ h (⋅ , ⋅) ∈ R d and µ ⋆ h + 1 (⋅) ∈ R d are state-action and next-state embeddings. The low-rank structure in (1) facilitates tractable exploration: if the embedding ϕ ⋆ h is known to the learner, one can efficiently learn a near-optimal policy with sample complexity polynomial in the feature dimension d , and independent of the size of the state space [23]; in this regard, ϕ ⋆ h 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 ϕ ⋆ h and µ ⋆ h + 1 are unknown to the learner. This formulation generalizes well-known frameworks such as the Block MDP (BMDP) model [12, 32], and necessitates the use of representation learning : the agent must learn an embedding that approximates ϕ ⋆ h as it explores the environment, and must use this learned embedding to drive subsequent exploration. This form of function approximation allows for great flexibility, as ϕ ⋆ h can be an arbitrary, nonlinear function of the state; in practice, it is common to model ϕ ⋆ h as a neural net [49].

The Low-Rank MDP is perhaps the simplest MDP structure that demands systematic exploration and nonlinear function approximation while allowing for a continuum of states, yet understanding of efficient algorithm design for this model is surprisingly limited. Existing algorithms suffer from at least one of the following drawbacks:

Computational intractability [21, 24, 14, 9, 43].

Strong modeling assumptions (e.g., ability to model µ ⋆ h + 1 (⋅) , which facilitates application of model-based RL techniques) [1, 40, 10]; in this work, we aim for model-free methods that only require learning ϕ ⋆ h .

Restrictive structural assumptions (e.g., non-negativity or latent variable structure for the embeddings in (1)) [35, 49].

At the root of these limitations is the complex interplay between exploration and representation learning: the agent must learn a high-quality representation to guide in exploring the state space, but learning such a representation requires gathering diverse and informative data, which is difficult to acquire without having already explored the state space to begin with. Overcoming this challenge-particularly where computational efficiency is concerned-requires (1) representation learning procedures that lead to sufficiently expressive representations for downstream applications, (2) efficient exploration procedures that are robust to errors in learned representations, and 3) understanding the interaction between these procedures, which must be interleaved. In this work, we propose an algorithm that addresses each of these challenges, as detailed below.

Contributions. Weprovide the first provably computationally efficient and model-free algorithm for general Low-Rank MDPs. Our algorithm, SpanRL , uses the notion of a barycentric spanner [6] for the embedding ϕ ⋆ h as an efficiently computable basis for exploration, and combines this with a minimax representation learning approach [35, 49]. SpanRL interleaves exploration with representation learning in a layer-wise fashion, learning a new representation at each layer h using exploratory data gathered at previous layers, then uses this representation to facilitate computation of a collection of exploratory policies (a policy cover ), which act as an approximate barycentric spanner for the features at layer h + 1 , ensuring good coverage for subsequent iterations. SpanRL is simple and modular, and its analysis is surprisingly compact given the greater generality compared to prior work [49, 35, 31].

SpanRL can accommodate general-purpose function approximation to learn the representation ϕ ⋆ (e.g., neural nets or other flexible classes) whenever a certain minimax representation learning objective [35, 49] can be solved efficiently for the function class of interest. Compared to efficient algorithms from prior work, SpanRL : (1) is model-free (i.e., only requires access to a function class Φ capable of modeling ϕ ⋆ , and does not need to model µ ⋆ h + 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 a secondary benefit, the algorithm is reward-free. Our analysis carefully combines several new techniques, including (1) an error-tolerant variant of the classical barycentric spanner computation algorithm of Awerbuch

∧

and Kleinberg [6], and (2) a new analysis of a minimax representation learning objective introduced in Modi et al. [35], Zhang et al. [49], which shows for the first time that this objective can lead to meaningful guarantees in general Low-Rank MDPs without latent variable structure; this increased generality is meaningful, as we show in Appendix H that there is an exponential separation between our guarantees and those that require such a structure.

Organization. Section 2 formally introduces the Low-Rank MDP model and the online reinforcement learning framework we consider. In Section 3, we highlight challenges faced by previous approaches, introduce our main algorithm, SpanRL , and show how it overcomes these challenges, and then present its main sample complexity guarantee.

Comparison to ArXiv Version. After the initial submission of this work, we developed a substantially improved of the algorithm that removes the reachability assumption at the cost of a larger (but still polynomial) sample complexity guarantee. We have included this algorithm and its analysis in the ArXiv version of this paper [30].

2 Problem Setting

2.1 Low-Rank MDP Model

We work in an episodic, finite-horizon reinforcement learning framework, where H ∈ N denotes the horizon. A Low-Rank MDP [36, 45, 1] is a tuple M=(X , A , ( ϕ ⋆ h ) h ∈[ H ] , ( µ ⋆ h ) h ∈[ H ] , ρ ) consisting of a state space X , action space A with ∣A∣ = A , distribution over initial states ρ ∈ ∆ (X) , and mappings µ ⋆ h + 1 ∶ X → R d and ϕ ⋆ h ∶ X × A → R d . 3 Beginning with x 1 ∼ ρ , an episode proceeds in H steps, where for each step h ∈ [ H ] , the state x h evolves as a function of the agent's action a h via

where T h is a probability transition kernel, which is assumed to factorize based on ϕ ⋆ h and µ ⋆ h . In detail, we assume that there exists a σ -finite measure ν on X such that for all 1 ≤ h ≤ H -1 , and for all x ∈ X and a ∈ A , the function x ′ ↦ µ ⋆ h + 1 ( x ′ ) ⊺ ϕ ⋆ h ( x, a ) is a probability density with respect to ν (i.e. the function is everywhere non-negative and integrates to 1 under ν ). For any X ′ ⊆ X , the probability that x h + 1 ∈ X ′ under x h + 1 ∼ T h (⋅ ∣ x h , a h ) is then assumed to follow the law

For notational compactness, we assume (following, e.g., Jiang et al. [21]) that the MDP M is layered so that X = X 1 ∪ ⋅ ⋅ ⋅ ∪ X H for X i ∩ X j = ∅ for all i ≠ j , where X h ⊆ X is the subset of states in X

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 ν , are a rigorous generalization of LowRank MDP formulations found in previous works [23, 1], which tend to implicitly assume the state space is countable and avoid rigorously defining integrals. We adopt this more general formulation to emphasize the applicability our results to continuous domains. However, in the special case where state space is countable, choosing ν as the counting measure yields T h (X ′ ∣ x h , a h ) = ∑ x ∈X ′ µ ⋆ h + 1 ( x ) ⊺ ϕ ⋆ h ( x h , a h ) , which is consistent with prior work.

Policies and occupancy measures. We define Π M = { π ∶ X → ∆ (A)} as the set of all randomized, Markovian policies. For a policy π ∈ Π M , we let P π denote the law of ( x 1 , a 1 ) , . . . , ( x H , a H ) under a h ∼ π ( x h ) , and let E π denote the corresponding expectation. For any X ′ ⊆ X h , we let P π h [X ′ ] ∶= P π [ x h ∈ X ′ ] denote the marginal law of x h under π . For x ∈ X h , we define the occupancy measure d π ( x ) ∶= d P π h d ν ( x ) as the density of P π h with respect to ν .

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 agent interacts with it in episodes , where at each episode the agent executes a policy of the form π ∶ X → ∆ (A) and observes the resulting trajectory ( x 1 , a 1 ) , . . . , ( x H , a H ) . While the ultimate goal of reinforcement learning is to optimize a policy with respect to a possibly unknown reward function, here we focus on the problem of reward-free exploration , which entails learning a collection of policies that almost optimally 'covers' the state space, and can be used to efficiently optimize any downstream reward function [12, 33, 15, 31]. To wit, we aim to construct an policy cover , a collection of policies that can reach any state with near-optimal probability.

Definition 2.1 (Policy cover) . For α ∈ ( 0 , 1 ] , a subset Ψ ⊆ Π M is an α -policy cover for layer h if

We show (Appendix G) that given access to such a policy cover with constant α , it is possible to optimize any downstream reward function with polynomial sample complexity.

Assumptions. To facilitate learning a policy cover, we make the following reachability assumption.

Assumption 2.1 ( η -reachability) . For any h ∈ [ H ] and x ∈ X h , max π ∈ Π M d π ( x ) ≥ η ⋅ ∥ µ ⋆ h ( x )∥ .

Reachability is necessary if one aims to build a policy cover that satisfies (3) uniformly for all states; without such a condition, one gives up on covering hard-to-reach states. Some notion of reachability is required in essentially all prior work on efficient model-free algorithms for Low-Rank MDPs [32, 35, 5], and was only very recently removed in the (more restrictive) BMDP setting [31, 49].

Remark 2.2 (Comparison to other reachability-like assumptions) . Assumption 2.1 generalizes and subsumes all previous reachability-like conditions of which we are aware [33, 46, 1, 35]. Notably, reachability is implied by the notion of feature coverage [5] (used in the context of transfer learning in Low-Rank MDPs), which asserts that sup π ∈ Π M λ min ( E π [ ϕ ⋆ h ( x h , a h ) ϕ ⋆ h ( x h , a h ) ⊺ ]) ≥ η , for some η > 0 . It is also implied by explorability [46], which is similar to feature coverage, but involves the first moments of ϕ ⋆ h . Our reachability assumption is also weaker than that used in [1, 35] under the latent variable model , and generalizes that made for BMDPs [33]. See Appendix H for details, as well as an exponential separation between our assumptions and analogous assumptions in [1, 35].

Beyond reachability, we assume (following [1, 35]) for normalization that, for all h ∈ [ H ] and ( x, a ) ∈ X h ×A , ∥ ϕ ⋆ h ( x, a )∥ ≤ 1 , and that for all g ∶ X h →[ 0 , 1 ] ,

Function approximation and desiderata. We do not assume that the true features ( ϕ ⋆ h ) h ∈[ H ] or the mappings ( µ ⋆ h ) h ∈[ 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 Φ ⊆ { ϕ ∶ X × A → R d } that contains ϕ ⋆ h , for h ∈ [ H -1 ] .

Assumption 2.2 (Realizability) . The feature class Φ ⊆ { ϕ ∶ X ×A → R d } has ϕ ⋆ h ∈ Φ for all h ∈ [ H ] . Moreover, for all ϕ ∈ Φ , x ∈ X , and a ∈ A , it holds that ∥ ϕ ( x, 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 Φ is finite to simplify presentation, but extension to infinite classes is straightforward, as our results only invoke finiteness through standard uniform convergence arguments. Note that unlike model-based approaches [1, 40, 10, 2], we do not assume access to a class capable of realizing the features µ ⋆ h , and our algorithm does not attempt to learn these features; this is why we distinguish our results as model-free .

For constant α , our goal is to learn an α -policy cover using poly ( d, A, H, log ∣ Φ ∣ , η -1 ) 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 H and A = ∣A∣ is necessary when ϕ ⋆ is 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 ] the integer interval { m,...,n } . We also let [ n ] ∶= [ 1 .. n ] . For any sequence of objects o 1 , o 2 , . . . , we define o m ∶ n ∶= ( o i ) i ∈[ m..n ] . A partial policy is a policy defined over a contiguous subset of layers [ ℓ .. r ] ⊆ [ H ] . We denote by Π ℓ ∶ r M ∶= { π ∶ ⋃ r h = ℓ X h → ∆ (A)} the set of all partial policies over layers ℓ to r ; note that Π M ≡ Π 1 ∶ H M . For a policy π ∈ Π ℓ ∶ r M and h ∈ [ ℓ .. r ] , π ( x h ) denotes the action distribution for the policy at layer h when x h ∈ X h is the current state. For 1 ≤ t ≤ h ≤ H and any pair of partial policies π ∈ Π 1 ∶ t -1 M , π ′ ∈ Π t ∶ h M , we define π ○ t π ′ ∈ Π 1 ∶ h M as the partial policy given by ( π ○ t π ′ )( x ℓ ) = π ( x ℓ ) for all ℓ < t and ( π ○ t π ′ )( x ℓ ) = π ′ ( x ℓ ) for all ℓ ∈ [ t .. h ] .

We use the x h ∼ π as shorthand to indicate that x h is drawn from the law P π , and likewise for ( x h , a h ) ∼ π and so on. For a set of partial policies Ψ ∶= { π ( i ) ∶ i ∈ [ N ]} , we define unif ( Ψ ) as the random partial policy obtained by sampling i ∼ unif ([ N ]) and playing π ( i ) . We define π unif ∈ Π M as the random policy that selects actions in A uniformly at random at each layer. We use ∥⋅∥ to denote the Euclidean norm, ∥⋅∥ ∞ to denote the supremum norm on functions, and let B( r ) ⊆ R d denote the Euclidean ball of radius r . We refer to a scalar c > 0 as an absolute constant to indicate that it is independent of all problem parameters and use ̃ O (⋅) 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 ( ϕ ⋆ h ) h ∈[ H ] are known (removing the need for representation learning) [23, 47, 4, 41], these approaches do not readily generalize to accommodate unknown features. For Low-Rank MDPs, previous algorithms suffer from at least one of the following three drawbacks: (1) the algorithms are computationally inefficient; (2) the algorithms are model-based; or (3) the algorithms place strong assumptions on the MDP that are unlikely to hold in practice. To motivate the SpanRL algorithm, we briefly survey these results, highlighting several key challenges in avoiding these pitfalls.

Let us first discuss the issue of computational efficiency. While there are a number of algorithms-all based on the principle of optimism in the face of uncertainty -that provide tight sample complexity guarantees for Low-Rank MDPs in reward-based [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,

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 ( ϕ ⋆ h ) h ∈[ H ] are realizable with respect to Φ , one assumes access to a second feature class Υ capable of modeling the mappings ( µ ⋆ h ) h ∈[ H ] ; this facilitates direct estimation of the transition probability kernel T h (⋅ ∣ 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 [1]. Unfortunately, model-based realizability is a 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 µ ⋆ h + 1 without sample complexity scaling with the number of states. 4

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 a 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 ϕ ⋆ h and µ ⋆ h are non-negative coordinate-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-negative factorization can have exponentially large dimension relative to the best unrestricted factorization [1], even when reachability is assumed (cf. Appendix H.1). The source of this restriction is 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 , SpanRL uses a policy cover Ψ ( h ) built at a previous iteration within a subroutine, RepLearn (Algorithm 5; deferred to Appendix B) to produce a feature map ˆ ϕ ( h ) that approximates ϕ ⋆ h . Using this feature map, the algorithm invokes a second subroutine, RobustSpanner (Algorithm 2 in Appendix B) to produce a collection of policies π 1 , . . . , π d that act as a 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 ○ h + 1 π unif ∶ i ∈ [ 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 a 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 ⊂ R d such that span (W) = R d , we say that a set { w 1 , . . . , w d } ⊆ W is a ( C,ε ) -approximate barycentric spanner for W if for every w ∈ W , there exist β 1 , . . . , β d ∈ [-C,C ] such that ∥ w - ∑ d i = 1 β i w i ∥ ≤ ε . 5

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 Ψ ⊆ Π M is a collection of policies such that { E π [ ϕ ⋆ h ( x h , a h )] ∣ π ∈ Ψ } ⊆ R d is a ( C,ε ) -approximate barycentric spanner for W h ∶= { E π [ ϕ ⋆ h ( x h , a h )] ∣ π ∈ Π M } with ε ≤ η 2 , then Ψ is an α -policy cover for layer h + 1 with α = ( 2 dC ) -1 .

Lemma 3.1, proven in Appendix F.6, shows that to compute a policy cover for layer h + 1 , it suffices to find a barycentric spanner for the set W h ∶= { E π [ ϕ ⋆ h ( x h , a h )] ∣ π ∈ Π M } ⊆ R d . Of course, even if ϕ ⋆ h is known, this observation is only useful if we can compute a spanner without explicitly enumerating over the set Π M , since our goal is to develop an efficient algorithm. In what follows, we will show: 6

Using, RobustSpanner , a novel adaptation of the classical algorithm of Awerbuch and Kleinberg [6], it holds that for any ϕ ∈ Φ , spanner computation for the set { E π [ ϕ ( x h , a h )] ∣ π ∈ Π M } can be performed efficiently whenever, for any θ ∈ B( 1 ) , one can (approximately) solve linear optimization problems of the form

Given access to policy covers Ψ ( 1 ∶ h ) for layers 1 to 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 ϕ ⋆ h is unknown, Algorithm 1 computes policies π 1 ∶ d that induce a barycentric spanner for the set { E π [ ˆ ϕ ( h ) ( x h , a h )] ∣ π ∈ Π M } , where ˆ ϕ ( h ) ∈ Φ is a learned feature map. In what follows, we first give a detailed explanation of the two above points, before showing how to complete the argument by learning a feature map through representation learning.

Algorithm 1 SpanRL : Volumetric Exploration and Representation Learning via Barycentric Spanner Require: Feature class Φ and parameters ε, c 0 and δ 0 , 1 .

Barycentric spanner computation via approximate linear optimization. To describe spanner computation in SpanRL , we take a brief detour and consider an abstract approach to barycentric spanner computation, which generalizes our problem. Suppose that we wish to compute a spanner for an implicitly specified set W = { w z } z ∈Z ⊆ R d indexed by an abstract set Z . The set Z (which will be set to Π M when we return to RL) may be exponentially large and cannot be efficiently enumerated. In addition, given z ∈ Z , we cannot explicitly compute w z , and have to settle for a noisy approximation.

To allow for efficient spanner computation, we assume access to two oracles for the set W , a linear optimization oracle LinOpt ∶ B( 1 ) → Z and an index-to-vector oracle LinEst ∶ Z → R d . We assume that for some ε > 0 :

For all θ ∈ R d with ∥ θ ∥ = 1 , the output ˆ z θ ∶= LinOpt ( θ ) satisfies θ ⊺ w ˆ z θ ≥ sup z ∈Z θ ⊺ w z -ε .

For all z ∈ Z , the output ˆ w z ∶= LinEst ( z ) satisfies ∥ ˆ w z -w z ∥ ≤ ε .

The RobustSpanner algorithm (Algorithm 2) computes a ( C,ε ) -approximate spanner for W using O ( d log ( d / ε )) 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 ( ϕ ⋆ h ) h ∈[ H ] . RobustSpanner achieves error tolerance by perturbing the vectors returned by LinOpt ( θ ) in the direction of θ , which amounts to running the classical algorithm on an ε -fattening of W , and is necessary in order to ensure that the approximation error of LinOpt does not swamp the signal in directions θ 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 RobustSpanner to construct a barycentric spanner for the set { E π [ ϕ ⋆ h ( x h , a h )] ∣ π ∈ Π M } with Z = Π M . Because we do not have access to ϕ ⋆ h , we instead apply RobustSpanner with W ∶= { E π [ ˆ ϕ ( h ) ( x h , a h )] ∣ π ∈ Π M } , where ˆ ϕ ( h ) is a learned representation. We now describe how the feature map ˆ ϕ ( h ) is learned, then show how to use these learned features to efficiently implement the oracles LinOpt (⋅) and LinEst (⋅) .

To learn 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 ( x h , a h , x h + 1 ) by rolling in to x h with a policy sampled uniformly from the policy cover Ψ ( h ) and selecting a h uniformly at random. Using this dataset, the algorithm solves a sequence of adversarial training sub-problems (Line 9 of Algorithm 5) which involve the feature class Φ and an auxiliary discriminator class F ∶ X → R . As we discuss in detail in the sequel, these sub-problems, described in (7), are amenable to standard gradient-based training methods. The sub-problems are designed to approximate the following 'idealized' max-min-max representation learning objective:

The intuition for this objective comes from the fact that in a Low-Rank MDP, for any function f ∶ X → R , the quantity E [ f ( x h + 1 ) ∣ x h = x, a h = a ] is linear in ϕ ⋆ h ( x, a ) . Thus, if F is sufficiently expressive, we may hope that ˆ ϕ ( h ) and ϕ ⋆ are close. We adopt the simple discriminator class F = { x ↦ max a ∈A θ ⊺ ϕ ( x, a ) ∣ θ ∈ B( 1 ) , ϕ ∈ Φ } . 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 ˆ ϕ ( 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 E.1) which may be of independent interest.

In more detail, prior work shows that the RepLearn algorithm solves a variant of (6) with w ∈ B( d 1 / 2 ⋅ poly ( ε -1 )) , where ε > 0 is the desired bound on mean-squared error. Due to the polynomial dependence on ε -1 , such a guarantee would lead to vacuous guarantees when invoked within our analysis of SpanRL . Our improved analysis of RepLearn , which is based on a determinantal potential argument, shows that w ∈ B( poly ( 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 π [ ˆ ϕ ( h ) ( x h , a h )] ∣ π ∈ Π M } , Z = Π M , and C = 2 ; that is, we plug-in the learned representation ˆ ϕ ( h ) for the true representation ϕ ⋆ h . 7 With this choice, implementing LinOpt entails (approximately) solving arg max π ∈ Π M E π [ θ ⊺ ˆ ϕ ( h ) ( x h , a h )] for a given θ ∈ B( 1 ) , and implementing the LinEst oracle entails estimating E π [ ˆ ϕ ( h ) ( x h , a h )] for a given π ∈ Π M . We instantiate

LinEst ( π ) as the Monte Carlo algorithm EstVec (Algorithm 7), which simply samples trajectories according to π and returns the sample average of ˆ ϕ ( h ) ( x h , a h ) . To implement LinOpt ( θ ) , we appeal to PSDP (Algorithm 3). PSDP , given an arbitrary reward function r 1 ∶ h ∶ X × A → R and a function class G ⊆ { g ∶ X × A → R } capable of realizing all possible value functions induced by these rewards, can use the policy covers Ψ ( 1 ∶ h ) to efficiently compute a policy ˆ π = PSDP ( h, r 1 ∶ h , G , Ψ ( 1 ∶ h ) , n ) that approximately solves arg max π ∈ Π M E π [∑ h t = 1 r t ( x t , a t )] , and does so using polynomially many episodes; see Appendix D for details and formal guarantees. 8 Thus, implementing LinOpt ( θ ) is as simple as invoking PSDP with the rewards

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 δ ∈ ( 0 , 1 ) be given, and suppose that realizability holds (Assumption 2.2) and that reachability (Assumption 2.1) is satisfied with parameter η > 0 . If ε = η 36 d 5 / 2 and c = polylog ( A,H,d, log (∣ Φ ∣/ δ )) is sufficiently large, then the policies Ψ ( 1 ∶ H ) produced by SpanRL ( Φ , ε, c , δ ) are a ( 1 4 Ad , 0 ) -policy cover with probability at least 1 -δ . The total number of episodes used by SpanRL is at most:

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 our work, all efficient model-free algorithms required non-negative features (latent variable structure) or stronger assumptions [35, 49], even in 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 to this assumption, it is natural to ask to what extent reachability restricts the generality of the Low-Rank MDP model. In Appendix H, we show that even when reachability holds, the embedding dimension in our model can be exponentially smaller than the best embedding dimension for the best non-negative (latent 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 (Algorithm 3), we can optimize any downstream reward function to error ε using poly ( d, A, H, log ∣ Φ ∣) ⋅ 1 / ε 2 episodes. See Appendix G for details.

Efficiency and practicality. We observe that SpanRL is simple and practical. Defining L D ( ϕ, w, f ) ∶= ∑ ( x,a,x ′ )∈D ( ϕ ( x, a ) ⊺ w -f ( x ′ )) 2 + λ ∥ w ∥ 2 , where D is a dataset consisting of ( x h , a h , r h , x h + 1 ) tuples, the algorithm is provably efficient whenever the adversarial objective

in Line 9 of RepLearn (Algorithm 5), can be implemented efficiently (note that by the definition of L D , the 'inner' minima over w and ˜ w in (7) can be solved in closed form). This objective was also assumed to be efficiently solvable in Modi et 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 F can be over-parameterized. While a detailed experimental evaluation is 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 in Zhang et al. [49]

Outside of representation learning, the only overhead in SpanRL is the RobustSpanner subroutine, which has polynomial runtime. Indeed, RobustSpanner requires only polynomially many calls to the linear optimization oracle, instantiated as PSDP , which is efficient whenever standard least-squares regression problems based on the class Φ 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 of 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 ϕ ⋆ h point-wise (i.e., uniformly for all ( x, a ) pairs), but the key challenge is that the guarantee in (6) only holds on average under the roll-in distribution unif ( Ψ ( h ) ) . 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 ϕ ⋆ , 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-Rank MDP and Block MDP literature [46, 5, 13, 33, 1, 35, 31], is sufficient to build a policy cover. We anticipate that this observation will find broader use.

Acknowledgments and Disclosure of Funding

We thank Noah Golowich, Dhruv Rohatgi, and Ayush Sekhari for several helpful discussions. ZM and AR acknowledge support from the ONR through awards N00014-20-1-2336 and N00014-20-1-2394, and ARO through award W911NF-21-1-0328. AB acknowledges support from the National Science Foundation Graduate Research Fellowship under Grant No. 1122374.

References

A. Agarwal, S. Kakade, A. Krishnamurthy, and W. Sun. FLAMBE: Structural complexity and representation learning of low rank MDPs. arXiv preprint arXiv:2006.10814 , 2020.

A. Agarwal, S. Kakade, and L. F. Yang. Model-based reinforcement learning with a generative model is minimax optimal. In Conference on Learning Theory , pages 67-83. PMLR, 2020.

A. Agarwal, S. 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.

A. Agarwal, Y. Jin, and T. Zhang. Vo q l: Towards optimal regret in model-free rl with nonlinear function approximation. arXiv preprint arXiv:2212.06069 , 2022.

A. Agarwal, Y. Song, W. Sun, K. Wang, M. Wang, and X. Zhang. Provable benefits of representational transfer in reinforcement learning. CoRR , abs/2205.14571, 2022.

B. Awerbuch and R. Kleinberg. Online linear optimization and adaptive routing. Journal of Computer and System Sciences , 74(1):97-114, 2008.

M. G. Azar, I. Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning , pages 263-272, 2017.

J. Bagnell, S. M. Kakade, J. Schneider, and A. Ng. Policy search by dynamic programming. Advances in neural information processing systems , 16, 2003.

F. 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.

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.

C. Dann, N. Jiang, A. Krishnamurthy, A. Agarwal, J. Langford, and R. E. Schapire. On oracleefficient PAC RL with rich observations. In Advances in neural information processing systems , pages 1422-1432, 2018.

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.09018 , 2019.

S. S. Du, Y. Luo, R. Wang, and H. Zhang. Provably efficient Q-learning with function approximation via distribution shift error checking oracle. In Advances in Neural Information Processing Systems , pages 8060-8070, 2019.

S. S. Du, S. 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.

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.

D. Ernst, P. Geurts, and L. Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research , 6, 2005.

D. J. Foster, S. M. Kakade, J. Qian, and A. Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487 , 2021.

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.

N. Golowich, A. Moitra, and D. Rohatgi. Learning in observable pomdps, without computationally intractable oracles. In Advances in Neural Information Processing Systems , 2022.

A. Huang, J. Chen, and N. Jiang. Reinforcement learning in low-rank mdps with density features. International Conference on Machine Learning (ICML) , 2023.

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 on Machine Learning , pages 1704-1713, 2017.

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.03736 , 2019.

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.

C. Jin, Q. Liu, and S. Miryoosefi. Bellman eluder dimension: New rich classes of RL problems, and sample-efficient algorithms. arXiv preprint arXiv:2102.00815 , 2021.

S. M. Kakade. On the sample complexity of reinforcement learning . University of London, University College London (United Kingdom), 2003.

J. Kober, J. A. Bagnell, and J. Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research , 32(11):1238-1274, 2013.

J. Li, W. Monroe, A. Ritter, D. Jurafsky, M. Galley, and J. Gao. Deep reinforcement learning for dialogue generation. In EMNLP , 2016.

T. P. Lillicrap, J. J. Hunt, A. Pritzel, N. Heess, T. Erez, Y . Tassa, D. Silver, and D. Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971 , 2015.

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 , 33:14532-14543, 2020.

Z. Mhammedi, A. Block, D. J. Foster, and A. Rakhlin. Efficient model-free exploration in low-rank mdps. arXiv preprint arXiv:2307.03997 , 2023.

Z. Mhammedi, D. J. Foster, and A. Rakhlin. Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observation rl. International Conference on Machine Learning (ICML) , 2023.

D. Misra, M. Henaff, A. Krishnamurthy, and J. Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. arXiv preprint arXiv:1911.05815 , 2019.

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-6971. PMLR, 2020.

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 reinforcement learning. Nature , 518(7540):529, 2015.

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.

S. Rendle, C. Freudenthaler, and L. Schmidt-Thieme. Factorizing personalized markov chains for next-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.

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.

D. Silver, A. Huang, C. 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):484, 2016.

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 on learning theory , pages 2898-2933. PMLR, 2019.

M. Uehara, X. Zhang, and W. Sun. Representation learning for online and offline RL in low-rank mdps. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 , 2022.

R. Wang, S. 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.

R. Wang, R. Salakhutdinov, and L. F. Yang. Provably efficient reinforcement learning with general value function approximation. arXiv preprint arXiv:2005.10804 , 2020.

T. Xie, D. J. Foster, Y. Bai, N. Jiang, and S. M. Kakade. The role of coverage in online reinforcement learning. arXiv preprint arXiv:2210.04157 , 2022.

L. 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.

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.

A. Zanette, A. Lazaric, M. J. Kochenderfer, and E. Brunskill. Provably efficient reward-agnostic navigation with linear value iteration. Advances in Neural Information Processing Systems , 33: 11756-11766, 2020.

T. Zhang. Feel-good thompson sampling for contextual bandits and reinforcement learning. arXiv preprint arXiv:2110.00871 , 2021.

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.

X. Zhang, Y. Song, M. Uehara, M. Wang, A. Agarwal, and W. Sun. Efficient reinforcement learning in block mdps: A model-free representation learning approach. In International Conference on Machine Learning , pages 26517-26547. PMLR, 2022.

Contents of Appendix

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 al. [35] (defined in Definition H.1) is the Block MDP (BMDP) model Du et al. [13], Misra et al. [32], Zhang et al. [49], Mhammedi et al. [31]. For this setting, Du et al. [13], Misra et al. [32] provide algorithms that conduct exploration in a provably oracle-efficient manner under a reachability assumption. This reachability assumption was removed by subsequent work of Zhang et al. [49] (with a suboptimal rate) and Mhammedi et al. [31] (with optimal error dependence), but the analysis in these works is tailored to the BMDP model.

Barycentric spanners. Huang et al. [20] consider a variant of the Low-Rank MDP framework in which we are given a class Υ that realizes the next-state feature map µ ⋆ , 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 are significantly 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 π (⋅)} π ∈ Π M in ℓ 1 -error. Barycentric spanners were also in the work of Golowich et al. [19] for reinforcement learning in Partially Observable MDPs (POMDPs). Their analysis is substantially different from ours, and their algorithm appeals to the barycentric spanner computation approach in Awerbuch and Kleinberg [6] in an off-the-shelf fashion.

B Omitted Algorithms

Algorithm 3 PSDP ( h, r 1 ∶ h , G , Ψ ( 1 ∶ h ) , n ) : Policy Search by Dynamic Programming (variant of [8])

Require:

Target layer h ∈ [ H ] .

Reward functions r 1 ∶ h .

Function class G .

Policy covers Ψ ( 1 ) , . . . , Ψ ( h ) .

Number of samples n ∈ N .

1: for t = h, . . . , 1 do

2: D ( t ) ←∅ .

3: for n times do

6: Solve regression:

7: Define ˆ π ( t ) ∈ Π t ∶ h M via

8: Return: Near-optimal policy ˆ π ( 1 ) ∈ Π M .

Algorithm 4 EstVec ( h, F, π, n ) : Estimate E π [ F ( x h , a h )] for policy π and function F ∶ X × A → R d .

Require:

Target layer h ∈ [ H ] .

Vector-valued function F ∶ X × A → R d .

Policy π ∈ Π M .

Number of samples n ∈ N .

1: D ← ∅ .

2: for n times do

3: Sample ( x h , a h ) ∼ π .

4: Update dataset: D ← D ∪ {( x h , a h )} .

5: Return: ¯ F = 1 n ∑ ( x,a )∈D F ( x, a ) .

Algorithm 5 RepLearn ( h, F , Φ , Ψ , n ) : Representation Learning for Low-Rank MDPs [35, 49].

Require:

Target layer h ∈ [ H ] .

Discriminator class F .

Feature class Φ .

A set of policies Ψ .

Number of samples n ∈ N .

3: D ← ∅ ,

4: for n times do

/* Discriminator selection */

9: Solve

/* Feature selection via least-squares minimization */

12: Solve

C Generic Guarantee for RobustSpanner

In this section, we give a generic guarantee for the RobustSpanner 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 a collection of vectors { w z ∈ R d ∣ z ∈ Z} indexed by elements in Z , there exists ε ′ > 0 such that for any θ ∈ R d ∖ { 0 } and z ∈ Z , the outputs ˆ z θ ∶= LinOpt ( θ /∥ θ ∥) and ˆ w z ∶= LinEst ( z ) satisfy

Letting W ∶= { w z ∣ z ∈ Z} and assuming that W ⊆ B( 1 ) , the next theorem bounds the number of iterations of RobustSpanner ( LinOpt (⋅) , LinEst (⋅) , ⋅ , ⋅) under Assumption C.1, and shows that the output is an approximate barycentric spanner for W (Definition 3.1). Our result extends those of Awerbuch and Kleinberg [6], in that it only requires an approximate linear optimization oracle, which is potentially of independent interest.

Proposition C.1. Fix C > 1 and ε ∈ ( 0 , 1 ) and suppose that { w z ∣ z ∈ Z} ⊆ B( 1 ) . If RobustSpanner (Algorithm 2) is run with parameters C,ε > 0 and oracles LinOpt , LinEst satisfying Assumption C.1 with ε ′ = ε / 2 , then it terminates after d + ⌈ d 2 log C 100 d ε 2 ⌉ iterations, and requires at most twice that many calls to each of LinOpt and LinEst . Furthermore, the output z 1 ∶ d has the property that for all z ∈ Z , there exist β 1 , . . . , β d ∈ [-C,C ] , such that

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 z ∣ z ∈ Z} 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, it suffices to bound the number of iterations in the second loop.

Let M ( i ) ∶= ( w 1 , . . . , w i , e i + 1 , . . . , e d ) be the matrix whose columns are the vectors at end of the i th iteration of the first loop (Line 2) of Algorithm 2; note that columns i + 1 through d are unchanged at this point in the algorithm. For i ∈ [ d ] , we define ℓ i ( w ) ∶= det ( w,M ( i ) -i ) and θ i ∶= ( det ( e j , M ( i ) -i )) j ∈[ d ] ∈ R d , where we recall that for any matrix A , the matrix A -i is defined as the result of removing the i th column from A . Note that ℓ i is linear in w , and in particular

Let W ( 0 ) ∶= M ( d ) = ( w 1 , . . . , w d ) , and let W ( j ) denote the resulting matrix after j iterations of the second loop (Line 10) of Algorithm 2. We will show that for any J ≥ 1 ,

By construction of the loop, we have det ( W ( j ) ) ≥ C ⋅ det ( W ( j -1 ) ) for each j ∈ [ J ] , and thus det ( W ( J ) ) ≥ det ( W ( 0 ) ) ⋅ C J . Combining these two facts will establish the bound on the iteration complexity. We now prove (9).

Let u i = e ⊺ i ( M ( i ) ) -1 (note that u i is a row vector ) and let U denote the matrix whose i th row is u i . We observe that for all w ∈ R d ,

where we note that ℓ i ( w i ) ≠ 0 by construction; indeed, the columns of M ( i ) are a basis for R d because det ( M ( i ) ) ≠ 0 , and the equality holds on the columns, so the two linear functions must be equal. Now, since Assumption C.1 holds with ε ′ = ε / 2 , we have

where w ± i = LinOpt ( z ± i ) . We will now show that

There are two cases. First, suppose that θ ⊺ i w + i ≥ -θ ⊺ i w -i , corresponding to the conditional in Line 6 of Algorithm 2 being satisfied. Combining this with (10), we have

Because the conditional is satisfied, w i = w + i + ε ⋅ θ i ∥ θ i ∥ , and so by plugging this into (12), we have

The case that θ ⊺ i w + i ≤ -θ ⊺ i w -i is essentially identical, establishing (11). Now, recall that W ∶= { w z ∣ z ∈ Z} and let W⊕B( 3 ε 2 ) ∶= { w + b ∣ w ∈ W and b ∈ B ( 3 ε 2 )} denote the Minkowski sum with B( 3 ε 2 ) . By Cauchy-Schwarz, it holds that for all w ′ ∶= w + b ∈ W ⊕ B ( 3 ε 2 ) ,

where we used that W ⊆ B( 1 ) (by assumption). Thus, for any w ′ ∈ W ⊕ B ( 3 ε 2 ) , we have

We now observe that by construction and the fact that Assumption C.1 holds with ε ′ = ε / 2 , the k th column w ′ k of W ( J ) belongs to W⊕B( 3 ε 2 ) , for any k ∈ [ d ] . Thus, the ( i, k ) entry u i w ′ k of UW ( J ) satisfies u i w ′ k ∈ [-1 -3 ε 2 , 1 + 3 ε 2 ] , and so the columns of UW ( J ) have Euclidean norm at most 10 √ d ε . Since the magnitude of the determinant of a matrix is upper bounded by the product of the Euclidean norms of its columns, it holds that ∣ det ( UW ( J ) )∣ ≤ ( 100 d ε 2 ) d 2 .

On the other hand, again by construction, we see that the columns w 1 , . . . , w d of W ( 0 ) satisfy u i w j = 0 , for j < i , and u i w i = 1 . Thus, UW ( 0 ) is an upper-triangular matrix with 1 s on the diagonal, and hence has determinant 1 . Because determinants are multiplicative, this implies that det ( U ) ≠ 0 . We now compute:

Thus, the upper bound on ∣ det ( W ( J ) )∣ holds and the claim is proven. Therefore, we have

and so J ≤ ⌈ d 2 log C ( 100 d ε 2 )⌉ .

Part II: 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 ∶= ( w 1 , . . . , w d ) be the matrix at termination of the algorithm. By definition, if the second loop (Line 10) has terminated, then for all i ∈ [ d ] ,

where θ i = ( det ( e j , W -i )) j ∈[ d ] ∈ R d . On the other hand, by Assumption C.1, (10) holds, and so

Now, fix z ∈ Z . Since det ( W ) ≠ 0 , there exist β 1 ∶ d ∈ R such that w z = ∑ d i = 1 β i w i . By plugging this into (13) and using the linearity of the determinant, we have

Therefore, ∣ β i ∣ ≤ C , for all i ∈ [ d ] . Now, by definition of w 1 ∶ d and ̃ w 1 ∶ d , for all i ∈ [ d ] , we have that ∥ w i - ̃ w i ∥ ≤ ε . Furthermore, by Assumption C.1, we also have that ∥ ̃ w i -w z i ∥ ≤ ε / 2 . Therefore, by the triangle inequality, we have

This completes the proof.

D Generic Guarantee for PSDP

In this section, we present a generic guarantee for PSDP (Algorithm 3). We show that given any reward functions r 1 ∶ h ∶ X × A → R and a function class G ⊆ { g ∶ X × A → R } that 'realizes' these reward functions (we formalize this in the next definition), if Ψ ( 1 ∶ h ) are α -policy covers for layers 1 through h , then for sufficiently large n ≥ 1 and with high probability, the output ˆ π = PSDP ( h, r 1 ∶ h , G , Ψ ( 1 ∶ h ) , n ) is an approximate maximizer of the objective

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 ⊆ { g ∶ X × A → R } realizes the reward functions r 1 ∶ h ∶ X × A → R if for all t ∈ [ h ] and all π ∈ Π t ∶ h M ,

Note that Q π t in (14) represents the state-action value function ( Q -function) at layer t ∈ [ h ] with respect to the rewards r 1 ∶ h and partial policy π .

In what follows, given a function class G ⊆ { g ∶ X ×A → R } , we use N G ( ε ) to denote the ε -covering number of G in ℓ ∞ distance. With this, we now state a guarantee for PSDP .

Theorem D.2. Let ε, δ ∈ ( 0 , 1 ) , B > 0 , and h ∈ [ H ] . Suppose reward functions r 1 ∶ h ∶ X × A → R , function class G , a collection of policies Ψ ( 1 ∶ h ) , and a parameter n ≥ 1 satisfy the following:

The function class G realizes the reward functions r 1 ∶ h (in the sense of Definition D.1), functions in G are uniformly bounded by B , and lim n →∞ n -1 ⋅ log N G ( 1 / n ) = 0 .

For some 0 < α ≤ 1 , for each 1 ≤ t ≤ h , it holds that Ψ ( t ) is an α -policy cover for layer t and moreover ∣ Ψ ( t ) ∣ ≤ d .

The parameter n is chosen such that cdHα -1 ⋅ ε stat ( n, δ / H ) ≤ ε , where ε stat ( n, δ ′ ) ∶= √ B 2 n -1 ⋅ ( log N G ( 1 / n ) + log ( n / δ )) and c > 0 is a large enough absolute constant.

Then, with probability at least 1 -δ , the policy ˆ π = PSDP ( h, r 1 ∶ h , G , Ψ ( 1 ∶ h ) , n ) coming from Algorithm 3, satisfies the following guarantee:

To prove the theorem, we need two intermediate results. The first shows that the Q function is the Bayes-optimal predictor of the sum of rewards when rolling out with policy π .

Lemma D.1. Let t ∈ [ H ] , r 1 ∶ h ∶ X × A → R be reward functions, π ∈ Π M , and let g π bayes denote the Bayes-optimal predictor 9 for the sum of rewards under policy π , i.e.,

Then, g π bayes = Q π t , where Q π t is the Q -function at layer t ∈ [ h ] with respect to the policy π and rewards r 1 ∶ h defined in (14).

Proof of Lemma D.1. The least-squares solution g π bayes of the problem in (15) satisfies, for all a ∈ A and x ∈ X t ,

where the last step follows by the definition of the Q

-function in (14).

We now show that the solution ˆ g ( t ) of the least-squares problem in (6) of Algorithm 3 is close to the Q -function in the appropriate sense.

Lemma D.2. Let ε, δ ∈ ( 0 , 1 ) , B > 0 , and 1 ≤ t ≤ h ≤ H . Further, let ( ε stat , r 1 ∶ h , G , Ψ ( 1 ∶ h ) , 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 ( h, r 1 ∶ h , G , Ψ ( 1 ∶ h ) , n ) satisfies with probability at least 1 -δ ,

where ˆ π ( t + 1 ) ∈ Π t + 1 ∶ h M is as in Algorithm 3, and c > 0 is an absolute constant.

Proof of Lemma D.2. Fix t ∈ [ h ] and let g ( t ) bayes = g π bayes be as in Lemma D.1 with

and reward functions r 1 ∶ h as in the lemma's statement. By Lemma D.1, g ( t ) bayes is the Bayes-optimal solution of the least-squares problem in (6) of Algorithm 3. Thus, since G realizes the reward functions r 1 ∶ h , a standard uniform-convergence guarantee for least-square regression (see e.g. Mhammedi et al. [29, Proposition B.1] with e = 0 almost surely) implies that there exists an absolute constant c > 0 (independent of t, h , and any other problem parameters) such that with probability at least 1 -δ ,

t ( t + 1 )

The desired result follows by the fact that g ( ) bayes ≡ Q ˆ π t ; see Lemma D.1. We also require the classical performance difference lemma from Kakade [25].

Lemma D.3 (Performance Difference Lemma) . Let π ⋆ , π ∈ Π M be arbitrary, and Q π t be as defined in (14). Then, for any h ≥ 1

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 E t of probability at least 1 -δ / H under which the learned partial policies ˆ π ( t ) , ˆ π ( t + 1 ) are such that

where π ⋆ ∈ arg max π ∈ Π M E π [∑ h t = 1 r t ( x t , a t )] is the optimal policy and Q π t is the Q -function defined in (14). Once we establish (16) for all t ∈ [ h ] , we will apply the performance difference lemma (Lemma D.3) and the union bound to obtain the desired result.

For any t ∈ [ h ] , we have

where the penultimate inequality follows by the fact that ˆ π ( t ) ( x ) ∈ arg max a ∈A ˆ g ( t ) ( x, a ) , for all ( x, a ) ∈ X t ×A by its construction in (8). Now, using the assumption on Ψ ( t ) , we have for any t ∈ [ t ] ,

where the first inequality is by (17), the second inequality follows from the fact that Ψ ( t ) is an α -policy cover, the third inequality follows from the fact that ∣ Ψ ( t ) ∣ ≤ d , and the equality is by definition. Now, by Lemma D.2, we have that for any t ∈ [ h ] , there is an absolute constant c > 0 (independent of t and other problem parameters) and an event E t of probability at least 1 -δ / H under which the solution ˆ g ( t ) of the least-squares regression problem on (6) of Algorithm 3 satisfies,

where the last inequality follows by the choice of n in the theorem's statement. Combining (19) with (18) establishes (16) under the event E t .

Now, by the performance difference lemma (Lemma D.3), we have

Algorithm 6 RepLearn ( h, F , Φ , P, n ) : Representation Learning for Low-Rank MDPs [35]

Require:

Target layer h ∈ [ H ] .

Discriminator class F .

Feature class Φ .

Policy distribution P ∈ ∆ ( Π M ) .

Number of samples n ∈ N .

1: Set ε stat = O ( √ cd 2 n -1 log (∣ Φ ∣/ δ )) for sufficiently absolute constant c > 0 (see Appendix E).

2: Let ϕ ( 1 ) ∈ Φ be arbitrary.

3: Set D ← ∅ .

4: for n times do

5: Sample π ∼ P .

6: Sample ( x h , a h , x h + 1 ) ∼ π ○ h π unif .

7: Update dataset: D ← D ∪ {( x h , a h , x h + 1 )} .

8: Define L D ( ϕ, w, f ) = ∑ ( x,a,x ′ )∈D ( ϕ ( x, a ) ⊺ w -f ( x ′ )) 2 .

9: for t = 1 , 2 , . . . do

/* Discriminator selection */

10: Solve

11: if ̂ ∆ ( f ( t ) ) ≤ 16 dtε 2 stat then

12: Return ϕ ( t ) .

/* Feature selection via least-squares minimization */

13: Solve

Thus, under the event E ∶= ⋃ h t = 1 E t , we have that

The desired result follows by the fact that a union bound implies P [E] ≥ 1 -δ .

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 of the weight w in Line 10 does not grow with the number of iterations. We also use the slightly simpler discriminator class:

The main guarantee for RepLearn is as follows.

Theorem E.1. Let h ∈ [ H ] , δ ∈ ( 0 , e -1 ) , and n ∈ N be given, and suppose that µ ⋆ h + 1 satisfies the normalization assumption in Eq. (4). For any function f ∈ F , define

Let P ∈ ∆ ( Π M ) be a distribution over policies, F be as (20), and Φ be a feature class satisfying Assumption 2.2. With probability at least 1 -δ , RepLearn with input ( h, F , Φ , P, n ) terminates after t ≤ T ∶= ⌈ d log 3 / 2 ( 2 nd -1 / 2 )⌉ iterations, and its output ϕ ( t ) satisfies

where ε 2 RepLearn ( n, δ ) ∶= cTd 3 n -1 log (∣ Φ ∣/ δ ) , for some sufficiently large absolute constant c > 0 .

To prove the theorem, we need a technical lemma, which follows from Modi et al. [35, Lemma 14].

Lemma E.1. Consider a call to RepLearn ( h, F , Φ , P, n ) (Algorithm 6) in the setting of Theorem E.1. Further, let L D be as in Algorithm 6 and define

For any δ ∈ ( 0 , 1 ) , there is an event E ( t ) ( δ ) of probability at least 1 -δ such that under E ( t ) ( δ ) , if Algorithm 6 does not terminate at iteration t ≥ 1 , then for w ( ℓ ) ∶= w f ( ℓ ) :

where ε 2 stat ( n, δ ) ∶= cd 2 n -1 log (∣ Φ ∣/ δ ) and c ≥ 1 is a sufficiently large absolute constant. With this, we prove Theorem E.1.

Proof of Theorem E.1. Let us abbreviate ε ∶= ε stat ( n, δ ) , with ε stat ( n, δ ) defined as in Lemma E.1. Further, let N ∶= 1 +⌈ d log 3 / 2 ( 2 d 3 / 2 / ε )⌉ , δ ′ ∶= δ 2 N , and define

Note that ε ≤ ˜ ε stat and N -1 ≤ T , where T is the number of iterations in the theorem statement; the latter inequality follows by the facts that the absolute constant c in Lemma E.1 is at least 1 and log (∣ Φ ∣/ δ ) ≥ 1 . We define an event E ∶= E ( 1 ) ( δ ′ ) ∩ ⋅ ⋅ ⋅ ∩ E ( N ) ( δ ′ ) , where (E t (⋅)) t are the success events in Lemma E.1. Note that P [E] ≥ 1 -δ / 2 by the union bound. Throughout this proof, we condition on the event E .

To begin the proof, we define a sequence of vectors ( v ( ℓ ) 1 ∶ d ) ℓ ≥ 0 in an inductive fashion, with v ( ℓ ) i ∈ R d for all i ∈ [ d ] and ℓ ≥ 0 . For ℓ = 0 , we let v ( 0 ) i = εe i / d , for all i ∈ [ d ] . For ℓ ≥ 1 , we consider two cases:

· Case I: If

where V ( ℓ -1 ) ∶= ( v ( ℓ -1 ) 1 , . . . , v ( ℓ -1 ) d ) ∈ R d × d and w ( ℓ ) ∶= w f ( ℓ ) , then we let j ∶= arg min j ′ ∈J ( ℓ ) j ′ and define

Case II : If J ( ℓ ) = ∅ , we let v ( ℓ ) i = v ( ℓ -1 ) i , for all i ∈ [ d ] .

We first show that J ( t ) ≠ ∅ 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 ( t ) = ∅ . This means that

Now, since det ( V ( t -1 ) ) ≠ 0 (note that ∣ det ( V ( t ) )∣ is non-decreasing with t ), we have that span ( V ( t -1 ) ) = R d . Thus, there exist β 1 , . . . , β d ∈ R be such that w ( t ) = ∑ d i = 1 β i v ( t -1 ) i . By the linearity of the determinant and (25), we have

This implies that ∣ β j ∣ ≤ ( 1 + C ) for all j ∈ [ d ] . Now, note that by the definition of ( v ( t -1 ) i ) , we have that for any i ∈ [ d ] such that v ( t -1 ) i ≠ εe i / d , there exists ℓ ∈ [ t -1 ] such that w ( ℓ ) = v ( t -1 ) i . Let

and for any i ∈ I ( t ) , let ℓ i ∈ [ t -1 ] be such that w ( ℓ i ) = v ( t -1 ) i . Further, define

and note that by the triangle inequality and the fact that w ( t ) = ∑ d i = 1 β i v ( t -1 ) i , we have

Finally, with the notation in (22), define

since ∣ β i ∣ ≤ ( 1 + C ) for all i ∈ [ d ] , ∣I ( t ) ∣ ≤ d , and ̂ w ( t ) ℓ ∈ B( 2 √ d ) , for all ℓ ∈ [ t -1 ] . Now, by Lemma E.1, in particular (23), we have

where ˜ ε stat is as in (24). Using the expressions in Eqs. (26) and (28) with (29) and Jensen's inequality, we have that under E ( t ) ,

Now, using (27) and the facts that ( a + b ) 2 ≤ 2 a 2 + 2 b 2 and ∥ ϕ ⋆ h ∥ 2 ≤ 1 , we have that

and note that

Using that C = 1 / 2 , we conclude that the right-hand side of this inequality is bounded by 8 dt ˜ ε 2 stat which is a contradiction, since ̂ w ( t ) t ∈ ( 1 + C )B( 2 d 3 / 2 ) = B( 3 d 3 / 2 ) and by Lemma E.1, we must have

if RepLearn does not terminate at round t . Therefore, we have that J ( t ) ≠ ∅ , for any iteration t ∈ [ 2 .. N ] 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 ( ℓ ) ≠ ∅ for ℓ > 1 , we have by construction:

Thus, if RepLearn runs for t 2 .. N iterations, then

On the other hand, since the determinant of a matrix is bounded by the product of the norms of its columns and v ( t ) 1 ∶ d ∈ B( 2 √ d ) , we have

Note also that ∣ det ( V ( 0 ) )∣ = ( ε / d ) d . Plugging this into (30), we conclude that

Taking the logarithm on both sides and rearranging yields

Thus, the algorithm must terminate after at most N -1 iterations. Furthermore, by [35, Lemma 14], we have that with probability at least 1 -δ 2 N , if the algorithm terminates at iteration t , then

Applying a union bound completes the proof.

F Analysis

In this section, we prove the main guarantee for SpanRL (Theorem 3.2). First, we outline our 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 in Appendix F.5 to complete the proof of Theorem 3.2. A self-contained guarantee for RobustSpanner (Lemma 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 Ψ ( 1 ∶ h + 1 ) produced by SpanRL satisfies the property:

Conditioned on this claim, we show that with high probability, the set Ψ ( h + 2 ) is a ( 1 4 Ad , 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 = { E π [ ϕ ( h ) ( x h , a h )] ∣ π ∈ Π M } (Appendix F.2). Using this, we then instantiate the guarantee of RobustSpanner 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 RobustSpanner together with the main guarantee for RepLearn (Theorem E.1), along with a change of measure argument enabled by the assumption that Ψ ( 1 ∶ h ) are policy covers (i.e. (31)).

F.2 Guarantee for PSDP as a Subroutine for RobustSpanner

We begin by showing that PSDP instantiates the approximate linear optimization oracle required by RobustSpanner . In particular, we fix a layer h and assume that Ψ ( 1 ∶ h + 1 ) satisfy (31) and apply the results of Appendix D.

More precisely, we need to show that, for any θ ∈ R d ∖{ 0 } and ϕ ∈ Φ , PSDP approximately solves

We can equivalently formulate (32) as, for fixed θ ∈ R d ∖{ 0 } and ϕ ∈ Φ , maximizing the sum of the reward functions r 1 ∶ h (⋅ , ⋅ ; θ, ϕ ) given by:

Note that this matches the choice of reward functions in SpanRL (Algorithm 1) at iteration h with ϕ = ˆ ϕ ( h ) , the feature map returned by RepLearn in Line 8. With these reward functions and the function class

we show that the output ˆ π = PSDP ( h, r 1 ∶ h (⋅ , ⋅ ; θ, ϕ ) , G , Ψ ( 1 ∶ h ) , n ) approximately solves (32) with high probability if n ≥ 1 is sufficiently large. We first verify that the class G realizes the reward functions specified in (33) in the sense of Definition D.1.

Lemma F.1. Under Assumption 2.2, the function class G in (34) realizes the reward functions in (33), for any ϕ ∈ Φ and θ ∈ R d ∖{ 0 } . Furthermore, we have that functions in G are uniformly bounded by h ≤ H , and log N G ( ε ) ≤ log ∣ Φ ∣ + d log ( H / ε ) , where we recall that N G ( ε ) denotes the ε -covering number of G in ℓ ∞ distance.

Proof. Fix ϕ ∈ Φ and θ ∈ R d ∖ { 0 } , and let r ℓ (⋅ , ⋅) ≡ r ℓ (⋅ , ⋅ ; θ, ϕ ) , for ℓ ∈ [ h ] . For t = h , we clearly have that for any π ∈ Π h ∶ h M , Q π t (⋅ , ⋅) = r t (⋅ , ⋅) ∈ G . For t < h and π ∈ Π t ∶ h M , we have by the low-rank structure that

Now, by the fact that E π [ r h ( x h , a h ) ∣ x t + 1 = y, a t + 1 = π ( y )] ∈ [-1 , 1 ] , for all y ∈ X t + 1 (since ϕ (⋅ , ⋅) ∈ B( 1 ) , for all ϕ ∈ Φ ), and the normalizing assumption made on ( µ ⋆ h ) h ∈[ H ] in Section 2.2 (i.e. that for all g ∶ X t + 1 →[ 0 , 1 ] , ∥ ∫ X t + 1 µ ⋆ t + 1 ( y ) g ( y ) d ν ( y )∥ ≤ √ d ), we have that

This together with (35) and the fact that ϕ ⋆ t ∈ Φ (by Assumption 2.2), we have that Q π t ∈ G . Combining Lemma F.1 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.

Corollary F.1. Let ε, δ ∈ ( 0 , 1 ) and h ∈ [ H ] . Further, let θ ∈ R d ∖ { 0 } , ϕ ∈ Φ , and ˆ π = PSDP ( h, r 1 ∶ h (⋅ , ⋅ ; θ, ϕ ) , G , Ψ ( 1 ∶ h ) , n ) , where

The reward functions r 1 ∶ h (⋅ , ⋅ ; θ, ϕ ) are as in (33).

The function class G is as in (34).

The collection of policies Ψ ( 1 ∶ h ) satisfy (31).

The parameter n is chosen such that cHACd 2 ⋅ ε stat ( n, δ / H ) ≤ ε , where ε stat ( n, δ ′ ) ∶= √ dn -1 ⋅ ( d log ( nH ) + log (∣ Φ ∣/ δ ′ )) and c > 0 is some large enough absolute constant.

Then, under Assumption 2.2, with probability at least 1 -δ , we have that

We emphasize that the inductive assumption that Ψ ( 1 ∶ h ) is a policy cover of bounded size enters only in the statement of Theorem D.2. We now give a guarantee for RobustSpanner as used in SpanRL .

Algorithm 7 EstVec ( h, F, π, n ) : Estimate E π [ F ( x h , a h )] for given policy π and function F .

Require:

Target layer h ∈ [ H ] .

Vector-valued function F ∶ X × A → R d .

Policy π ∈ Π M .

Number of samples n ∈ N .

1: D ← ∅ .

2: for n times do

3: Sample ( x h , a h ) ∼ π .

4: Update dataset: D ← D ∪ {( x h , a h )} .

5: Return: ¯ F = 1 n ∑ ( x,a )∈D F ( x, a ) .

F.3 Guarantee for RobustSpanner as a 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 RobustSpanner can be taken to be EstVec (Algorithm 7), which simply estimates the expected feature imbedding of ( x h , a h ) under policy π by sampling sufficiently many trajectories and taking the empirical mean.

Lemma F.2 (Guarantee of EstVec ) . Let δ ∈ ( 0 , 1 ) and ε > 0 . For h ∈ [ H ] , ϕ ∈ Φ , π ∈ Π M , and n ∈ N such that n ≥ c ε 2 log ( d / δ ) for some large enough absolute constant c > 0 , the output ϕ h = EstVec ( h, ϕ, π, n ) (Algorithm 7) satisfies, with probability at least 1 -δ ,

Proof. By Hoeffding's inequality (see for example [22, Corollary 7]) and the fact that ∥ ϕ ( x, a )∥ ≤ 1 for all x ∈ X and a ∈ A , there exists an absolute constant c > 0 such that with probability at least 1 -δ ,

Setting n as in the statement of the theorem concludes the proof.

In SpanRL , we instantiate RobustSpanner passing PSDP as LinOpt and EstVec as LinEst . Combining Corollary F.1 and Lemma F.2 with the general guarantee of RobustSpanner in Proposition C.1, we have the following result.

Lemma F.3. Consider iteration h ∈ [ H ] of SpanRL ( Φ , ε, c , δ ) (Algorithm 1) with ε, c > 0 , δ ∈ ( 0 , 1 ) , and feature class Φ satisfying Assumption 2.2. Further, let ˆ ϕ ( h ) denote the feature map returned by RepLearn in Algorithm 1 at iteration h . If Ψ ( 1 ∶ h ) satisfy (31) and c = polylog ( A,d,H, log (∣ Φ ∣/ δ )) is large enough, then there is an event E h with probability at least 1 -δ 2 H such that

The number of iterations of RobustSpanner in Line 12 of Algorithm 1 is at most N = ⌈ d 2 log 2 ( 100 d ε )⌉ < ∞ , and

The output ( π 1 , . . . , π d ) of RobustSpanner has the property that for all π ∈ Π M , there exist β 1 , . . . , β d ∈ [-2 , 2 ] such that

Proof. By Proposition C.1, on the event that the instances of PSDP and EstVec used by RobustSpanner satisfy Assumption C.1 with ε ′ = ε 2 , the two desiderata of the lemma hold. 10 We

claim that each call to PSDP and to EstVec satisfies Assumption C.1 with probability at least 1 -δ 8 dNH . Because each of PSDP and EstVec get called at most 4 dN times per iteration of RobustSpanner , a union bound concludes the proof contingent on the above claim.

We now prove the claim. First, note that the instance of PSDP that RobustSpanner uses within Algorithm 1 is of the form:

with r 1 ∶ h and G as in Algorithm 1; this matches the form in Corollary F.1 ( PSDP 's guarantee) with ϕ = ˆ ϕ ( h ) . Thus, by choosing

for c = polylog ( A,d,H, log (∣ Φ ∣/ δ )) sufficiently large, the conditions of Corollary F.1 are satisfied, and its conclusion implies the claim for the PSDP instance used by RobustSpanner . Similarly, the choice of n EstVec in Algorithm 1 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 a guarantee for the invocation of RepLearn within SpanRL

Recall that P ( h ) = unif ( Ψ ( h ) ) is the distribution over policies that SpanRL passes to RepLearn at iteration h ∈ [ H -2 ] to compute feature map ϕ ( h ) . Thus, by invoking Theorem E.1 in Appendix E and using the choice of n RepLearn in Algorithm 1, we immediately obtain the following corollary.

Corollary F.2. Let δ, ε ∈ ( 0 , 1 ) , and F be as in Algorithm 1, and fix h ∈ [ H -2 ] . Suppose that the feature class Φ satisfies Assumption 2.2. Then, with probability at least 1 -δ 2 H , the instance of RepLearn in Line 9 of Algorithm 1 runs for t ≤ c ⋅ d iterations for c = polylog ( A,d,H, log (∣ Φ ∣/ δ )) sufficiently large, and returns output ϕ ( h ) such that for all f ∈ F , there exists w ( h ) f ∈ B( 3 d 3 / 2 ) satisfying

where w f ∶= ∫ X h + 1 f ( y ) µ ⋆ h + 1 ( y ) d ν ( 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 ∈ [ H ] of SpanRL ( Φ , ε, c , δ ) (Algorithm 1) with parameters ε, c > 0 , δ ∈ ( 0 , 1 ) and a feature class Φ satisfying Assumption 2.2. Further, assume that:

The collection of policies Ψ ( 1 ∶ h + 1 ) at the start of the h th iteration of SpanRL satisfy (31).

Assumption 2.1 (reachability) holds with η > 0 .

The input parameter ε to SpanRL is set to ε = η 36 d 5 / 2 .

The input parameter c = polylog ( A,d,H, log (∣ Φ ∣/ δ )) is sufficiently large.

Then, with probability at least 1 -δ H , the set of policies Ψ ( h + 2 ) produced by SpanRL ( Φ , ε, c , δ ) at the end of iteration h is an ( 1 4 Ad , 0 ) -policy cover for layer h + 2 .

With this, we can now prove Theorem 3.2.

Proof of Theorem 3.2. Note that it suffices to prove that (31) holds for h = H -1 with probability at least 1 -δ . To do this, we proceed by induction over h = 1 , . . . , H -1 . The base case of h = 1 trivially holds because Ψ ( 1 ) = ∅ and Ψ ( 2 ) = { π unif } . The induction step now follows by Theorem F.1 and the union bound (see e.g. [30, Lemma I.2]).

The number of trajectories used by SpanRL is dominated by calls to PSDP . Since PSDP is called O ( d log ( d / ε )) times at each iteration of SpanRL (Lemma F.3), and each call to PSDP requires at most Hn PSDP trajectories, the total number of trajectories after H iterations of SpanRL is bounded by ̃ O ( H 2 dn PSDP ) . By plugging the choices for n PSDP and ε from the theorem statement, we obtain the claimed sample complexity.

Before proving Theorem F.1, we make the following simple observation.

Lemma F.4. For any π Π , h H 1 , any x , we have

Proof of Lemma F.4. The equality follows by construction. The non-negativity of d π ( x ) follows by definition of a probability density.

We now prove Theorem F.1.

Proof of Theorem F.1. Let E h and E ′ h denote the success events in Lemma F.3 and Corollary F.2, respectively, and note that by the union bound, we have P [E h ∩ E ′ h ] ≥ 1 -δ / H . For the rest of this proof, we will condition on E ∶= E h ∩E ′ h .

Throughout, we denote

Because Ψ ( 1 ∶ h + 1 ) satisfy (31) (i.e., are a policy cover) it holds by Lemma F.4 that for all x ∈ X h ,

We will show that with probability at least 1 -δ H , the policy set Ψ ( h + 2 ) has the same property for layer h + 2 ; that is, for all x ∈ X h + 1 ,

Again, by Lemma F.4 this is equivalent to the statement that Ψ ( h + 2 ) is an ( 1 4 Ad , 0 ) -policy cover for layer h + 2 .

For the remainder of the proof, we will fix x ∈ X h + 2 and let π x ∈ arg max π ∈ Π M µ ⋆ h + 2 ( x ) ⊺ ϕ ⋆ ,π h + 1 . Our goal is to show that the inequality Eq. (37) holds for x .

Preliminaries. Note that since x ∈ X h + 2 , we have ∥ µ ⋆ h + 2 ( x )∥ > 0 . It will be convenient to introduce a function f ∶ X h + 1 → R defined by

Further, we define

By definition of π x , we have that for all y ∈ X h + 1 ,

This together with the fact that ∥ θ x ∥ = 1 implies that

the discriminator class in Line 4 of SpanRL . Note also that since x ∈ X h + 2 , we have by reachability that

Applying the guarantee for RepLearn . Moving forward, let ϕ ( h ) be the feature map returned by RepLearn at the h th iteration of Algorithm 1, and define ϕ ( h ) ,π ∶= E π [ ϕ ( h ) ( x h , a h )] , for any π ∈ Π M . Further, let w ( h ) x be the vector w ( h ) f in Corollary F.2 with f = f x , and note that

By Jensen's inequality, we compute

where the last step follows by the definition of Ψ ( h ) in Algorithm 1 and that ∣ Ψ ( h ) ∣ = d . Now, since w x = ∫ X h + 1 f ( y ) µ ⋆ h + 1 ( y ) d ν ( y ) (see (38)) and f ∈ F (see (39)); the guarantee for RepLearn in Corollary F.2 together with (42) implies that (conditioned on the event E )

Applying the guarantee for RobustSpanner . Letting π 1 , . . . , π d be the policies returned by RobustSpanner at iteration h of SpanRL , the guarantee of RobustSpanner in Lemma F.3 implies that there exist β 1 , . . . , β d ∈ [-2 , 2 ] such that

where the last inequality follows by the fact that ε = η 36 d 5 / 2 . Combining (44) with (43) and using the triangle inequality, we get that

Combining this with (40) and rearranging implies

On the other hand, by definition of w x , we have

where the inequality follows from the non-negativity of µ ⋆ h + 1 (⋅) ⊺ ϕ ⋆ h + 1 ( x, a ) , for all ( x, a ) ∈ X h ×A (due to Lemma F.4), and (46) follows from the definition of Ψ ( h + 2 ) in Line 13 of Algorithm 1. Combining (45) and (46) then implies that

This, together with Lemma F.4, implies that (37) holds. Since this argument holds uniformly for all x ∈ X h + 2 , this completes the proof.

F.6 Proof of Lemma 3.1

By definition for x ∈ X h + 1 , we have d π ( x ) = E π [ µ ⋆ h + 1 ( x ) ⊺ ϕ ⋆ h ( x h , a h )] . Let π x denote the policy maximizing d π ( x ) (if no such maximizer exists, we may pass to a maximizing sequence) and let Ψ = { π 1 , . . . , π d } . Then, we have for some β 1 , . . . , β d ∈ [-C,C ] ,

where the inequality follows by the fact that Assumption 2.1 holds with ε ≤ η / 2 . The result now follows by rearranging.

G Application to Reward-Based RL

In this section, we explain how the output Ψ ( 1 ∶ H ) of SpanRL (Algorithm 1) can be used to optimize downstream reward functions r 1 ∶ H ; our treatment is standard. Since the output of SpanRL is a policy cover, one way to optimize the sum of rewards S H ∶= ∑ H h = 1 r h is by first generating trajectories using policies in Ψ ( 1 ∶ H ) , then applying an offline RL algorithm, e.g. Fitted Q-Iteration ( FQI ) [16], to optimize S H . It is also possible to use PSDP with the policy cover Ψ ( 1 ∶ H ) to 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 to a function class G ⊆ { g ∶ X × A → R } that realizes the rewards r 1 ∶ H in the following sense: for all h ∈ [ H ] and all π ∈ Π h ∶ H M ,

Note that when the reward functions r 1 ∶ H are linear in the feature map ϕ ⋆ h ; that is, when for all h ∈ [ H ] and ( x, a ) ∈ X h ×A , r h ( x, a ) = θ ⊺ h ϕ ⋆ h ( x, a ) for some θ h ∈ B( 1 ) (a common assumption in the context of RL in Low-Rank MDPs [32, 31, 49, 35]), then the function class

realizes r 1 ∶ H . Note that G is the same function class we used for the PSDP subroutine in Algorithm 3, albeit with a larger ball for the w 's. For the sake of generality, we state the next result (which shows how to use a policy cover to optimize a downstream reward function) for general r 1 ∶ H and a function class G that realizes r 1 ∶ H .

Theorem G.1. Let ε > 0 and δ ∈ ( 0 , 1 ) . Suppose reward functions r 1 ∶ H ∶ X × A → R , a collection of policies Ψ ( 1 ∶ H ) , and a parameter n ≥ 1 satisfy the following:

The function class G realizes the reward functions r 1 ∶ H (in the sense of Definition D.1), and lim n →∞ n -1 ⋅ log N G ( 1 / n ) = 0 , where N G ( 1 / n ) to denote the 1 n -covering number of G in the supremum norm. Furthermore, we suppose that functions in G are uniformly bounded by H √ d .

For some 0 < α ≤ 1 , for each 1 ≤ h ≤ H , it holds that Ψ ( h ) is an α -policy cover for layer h and moreover ∣ Ψ ( h ) ∣ ≤ d .

The parameter n is chosen such that cdHα -1 ⋅ ε stat ( n, δ / H ) ≤ ε , where ε stat ( n, δ ′ ) ∶= √ dH 2 n -1 ⋅ ( log N G ( 1 / n ) + log ( 1 / δ )) and c > 0 is a large enough absolute constant.

Then, with probability at least 1 -δ , the policy ˆ π = PSDP ( H,r 1 ∶ H , G , P ( 1 ∶ H ) , n ) (where P ( t ) ∶= unif ( Ψ ( t ) ) , for each t ∈ [ h ] ) coming from Algorithm 3, satisfies the following guarantee:

Moreover, the number of episodes used by PSDP in this case is

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 η -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 to the popular latent variable model [1, 35]. 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 a reachability assumption used by [1, 35] in the latent variable/non-negative 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 × A → ∆ (X) , a latent variable representation consists of a countable latent space Z and functions ψ ∶ X ×A → ∆ (Z) and q ∶ Z → ∆ (X) , such that T (⋅ ∣ x, a ) = ∑ z ∈Z q (⋅ ∣ z ) ψ ( z ∣ x, a ) . The latent variable dimension of T , denoted d LV is the cardinality of smallest latent space Z for which T admits a latent variable representation.

The interpretation for the latent variable model is as follows:

Each ( x, a ) pair induces a distribution ψ ( x, a ) ∈ ∆ (Z) over z ∈ Z .

The latent variable is sampled as z ∼ ψ ( x, a ) .

The next state is sampled as x ′ ∼ q (⋅ ∣ z ) .

Note that in discrete state spaces, all transition operators admit a trivial latent variable representation, as we may take ψ ( x, a ) = T (⋅ ∣ x, a ) , but the dimension of such a representation is potentially infinite. A latent variable representation certifies that there exists a factorization T ( x ′ ∣ x, a ) = ψ ( x, a ) ⊺ q ( x ′ ) with embedding dimension ∣Z∣ , and so d LV , and hence gives an upper bound on the rank of the transition operator. On the other hand, compared with the general Low-Rank factorization, the latent variable factorization additionally requires that ψ ( x, a ) and q (⋅ ∣ z ) are probability distributions, and thus non-negative, for all z ∈ Z and ( x, a ) ∈ X ×A , implying that d LV is equivalent to the non-negative rank [1] of the transition operator.

Assuming that a latent variable representation exists, [1, 35] consider the following notion of reachability.

Definition H.2 (Reachability in latent variable model) . There exists η > 0 such that

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 d ≥ 1 . Under the latent variable model in Definition H.1, if the latent variable reachability condition in (47) is satisfied for some η > 0 , then, for all h ∈ [ H ] , the transition kernel T h in M admits a factorization T h (⋅ ∣ x, a ) = µ ⋆ h + 1 (⋅) ⊺ ϕ ⋆ h ( x, a ) , where µ ⋆ h + 1 (⋅) ∈ R d LV and ϕ ⋆ h (⋅ , ⋅) ∈ R d LV , such that d LV ≤ dA 2 / η 2 and η 2 A √ d -reachability (in the sense of Assumption 2.1) is satisfied.

Proof of Lemma H.1. Suppose that Assumption 2.1 ( η -reachability) holds. By Agarwal et al. [1, Proposition 4], the non-negative rank of M is bounded as d LV ≤ dA 2 / η 2 .

Letting q and ψ be as in the definition of the latent variable representation in Definition H.1, we define µ ⋆ h + 1 and ϕ ⋆ h as: for all h ∈ [ H -1 ] ,

Now, fix h ∈ [ H -1 ] and x ∈ X h + 1 . For z 0 ∈ arg max z ∈Z h + 1 q ( x ∣ z ) , we have

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 η polynomially small, yet the best possible latent variable embedding has dimension d LV = 2 Ω ( d ) . 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-MDP of rank d ≥ 1 , where η -reachability (Assumption 2.1) is satisfied with η = 1 2 √ d , but where the non-negative rank satisfies d LV = 2 Ω ( d ) .

Proof of Theorem H.3. Let n ∈ N and d ∶= ( n 2 ) + 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 X 1 and X 2 at layers 1 and 2, respectively, are finite.

The cardinality of A is d ; i.e. A={ a 1 , . . . , a d } . 11

The transition kernel T 1 admits the factorization:

where for all x ′ ∈ X 2 , µ ⋆ 2 ( x ′ ) ∈ R d ≥ 0 , and for all ( x, a ) ∈ X 1 ×A , ϕ ⋆ 1 ( x, a ) ∈ R d ≥ 0 .

The non-negative rank of M is d LV = 2 Ω ( d ) .

We augment this MDP by adding an extra state x 0 , and let X 1 ∶= X 1 ∪{ x 0 } . We define ϕ ⋆ 1 ∶ X 1 ×A→ R d ≥ 0 be the extension of ϕ ⋆ 1 given by

where e i is the i th basis element in R d . We define the initial state distribution to have ρ ( x 0 ) = 1 2 and ρ ( x ) = 1 2 ∣X 1 ∣ , for all x ∈ X 1 . 12 We let M=(X 1 ∪X 2 , A , ϕ ⋆ 1 , ( µ ⋆ h ) h ∈[ 2 ] , ρ ) denote the resulting MDP. Note that adding an extra state at layer 1 in this fashion only adds d additional rows to the transition matrix T (viewed as a (∣X 1 ×A∣) × ∣X 2 ∣ matrix). Therefore, the non-negative rank of M is as least that of M .

We now show that reachability is satisfied in M . Let π i the policy that always plays action a i . With this, we have that for any x ′ ∈ X 2 ,

where the last inequality follows by the fact that, for all ( x, a ) ∈ X 1 × A , µ ⋆ 2 (⋅) ⊺ ϕ ⋆ 1 ( x, a ) = µ ⋆ 2 ( x ′ ) ⊺ ϕ ⋆ 1 ( x, a ) ≥ 0 (since µ ⋆ 2 ( x ′ ) ⊺ ϕ ⋆ 1 ( x, a ) is a conditional density). On the other hand, from the construction of ϕ ⋆ 1 and the fact that µ ⋆ 2 ( x ′ ) ∈ R d ≥ 0 , we have

Combining this with (48) and using that ρ ( x 0 ) = 1 / 2 implies that reachability 1 /( 2 √ d ) is satisfied in M .

H.2 Relation to Other Reachability Assumptions

In this subsection, we show that Assumption 2.1 is implied by a 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 RL 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 ϕ ⋆ h .

Definition H.4 ( η -feature coverage) . We say that the linear MDP with featurization ϕ ⋆ h satisfies η -feature coverage if for all h ∈ [ H ] ,

We show that η -feature coverage implies ( η / 2 ) 3 / 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 η -feature coverage as in Definition H.4 for some η > 0 . If ϕ ⋆ h ( x, a ) ∈ B( 1 ) for all x, a , then the MDP satisfies ( η / 2 ) 3 / 2 -reachability in the sense of Assumption 2.1.

Proof of Lemma H.2. Let h ∈ [ H ] and x ∈ X h + 1 be given and define

To keep notation compact, we define ϕ ⋆ h ∶= ϕ ⋆ h ( x h , a h ) . By η -feature coverage, there exists π ∈ Π M such that

where we have used that ∥ θ ∥ = 1 and ∥ ϕ ⋆ h ( x, a )∥ ≤ 1 for all ( x, a ) ∈ X h × A . Rearranging (49) and using that θ ⊺ ϕ ⋆ h ≥ 0 (it is a scaled conditional density), have

Now, by Markov's inequality, we have that

where we have once more used that θ ⊺ ϕ ⋆ h ≥ 0 almost surely.

H.2.2 Explorability

We now consider the explorability assumption of [46], which involves the first moment of the feature map ϕ ⋆ h . This notion is defined as follows.

Definition H.5 ( η -explorability) . We say that a linear MDP satisfies η -explorability if for any h ∈ [ H ] and any θ ∈ R d ∖{ 0 } it holds that

We now show that η -explorability is a special case of η -reachability:

Lemma H.3. Suppose that the explorability condition in Definition H.5 is satisfied with η > 0 . Then, η -reachability is satisfied.

Proof of Lemma H.3. Let x ∈ X h + 1 and define θ ∶= µ ⋆ h + 1 ( x ) . By explorability, we have that

This shows that Assumption 2.1 is satisfied with parameter η .

Peer reviews:
comments: 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.

STRENGTHS:
The paper writing is clear. The comparsion with previous related work and the design idea of the algorithms are clear.

The 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".

The sample complexity should be improved.

Although 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.

[1] Uehara et. al., Representation Learning for Online and Offline RL in Low-rank MDPs

QUESTIONS:
See weakness.

LIMITATIONS:
N.A.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-15

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe paper addressed some of its limitations and proposed some open questions for future work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nAs listed above.\n\nLIMITATIONS:\nThere is no negative impact in this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nOne 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-17"}