Year: 2021

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

heading: 1 INTRODUCTION
text: When applying Reinforcement Learning (RL) to large-scale problems where data is complex and high-dimensional, learning effective transformations of the data, i.e., representation learning, can often significantly improve the sample and computation efficiency of the RL procedure. Indeed, several empirical works have shown that leveraging representation learning techniques developed in supervised or unsupervised learning settings can accelerate the search for good decision-making strategies (Silver et al., 2018; Stooke et al., 2021; Srinivas et al., 2020; Yang & Nachum, 2021). However, representation learning in RL is far more subtle than it is for non-sequential and noninteractive learning tasks (e.g., supervised learning). Prior works have shown that even if one is given the magic representation that exactly linearizes the optimal policy (Du et al., 2019b) or the optimal value functions (Wang et al., 2020; Weisz et al., 2021), RL is still challenging (i.e., one may still need exponentially many samples to learn). This indicates that an effective representation that permits efficient RL needs to encode more information about the underlying Markov Decision Processes (MDPs). Despite the recent empirical success of representation learning in RL , its statistical guarantee and theoretical properties remain under-investigated.
In this work, we study the representation learning question under the low-rank MDP assumption. Concretely, a low-rank MDP assumes that the MDP transition matrix admits a low-rank factorization, i.e., there exists two unknown mappings µ(s′), φ(s, a), such that P (s′|s, a) = µ(s′)>φ(s, a) for all s, a, s′, where P (s′|s, a) is the probability of transiting to the next state s′ under the current state and action (s, a). The representation φ in a low-rank MDP not only linearizes the optimal stateaction value function of the MDP (Jin et al., 2020a), but also linearizes the transition operator. A lowrankness assumption on large stochastic matrices is a common assumption and has enabled successful development of algorithms for real world applications such as movie recommendation systems (Koren et al., 2009). We note that a low-rank MDP strictly generalizes the linear MDP model (Yang
& Wang, 2020; Jin et al., 2020a) which assumes φ is known a priori. The unknown representation φ makes learning in low-rank MDPs much more challenging than that in linear MDPs since one can no longer directly use linear function approximations. On the other hand, the fact that linear MDPs can be solved statistical and computational efficiently if φ is known a priori implies that if one could learn the representation of the low-rank MDP, one could then efficiently learn the optimal policy.
Indeed, prior works have shown that learning in low-rank MDPs is statistically feasible (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021) via leveraging rich function approximators. However, these algorithms are version space algorithms and are not computationally efficient. Recent work FLAMBE proposes an oracle-efficient algorithm1 that learns in low-rank MDPs with a polynomial sample complexity, where the computation oracle is Maximum Likelihood Estimation (MLE). In this work, we follow the same setup from FLAMBE (Agarwal et al., 2020b), and propose a new algorithm — Upper Confidence Bound driven Representation Learning, Exploration and Exploitation (REP-UCB), which can learn a near optimal policy for a low-rank MDP with a polynomial sample complexity and is oracle-efficient. Comparing to FLAMBE, our algorithm significantly improves the sample complexity from O(d7A9/( 10(1− γ)22) for FLAMBE to O(d4A4/( 2(1− γ)3), where d is the rank of the transition matrix (or dimension of the true representation),A is the number of actions, is the suboptimality gap and γ ∈ [0, 1) is the discount factor in the MDP. Our algorithm is also arguably much simpler than FLAMBE: FLAMBE is an explore-then-commit algorithm, has to explore in a layer-by-layer forward way, and does not permit data sharing across different time steps. In contrast, REP-UCB carefully trades exploration versus exploitation by combining the reward signal and exploration bonus (constructed using the latest learned representation), and enables data sharing across all time steps. Our sample complexity nearly matches the ones from those computationally inefficient algorithms (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021). We summarize the comparison with the prior works that study representation learning in Table 1.
In addition to the online exploration setting, we also show that our new techniques can be directly used for designing offline RL algorithms for low-rank MDPs under partial coverage. More specifically, we propose an algorithm REP-LCB—Lower Confidence Bound driven Reprepresentation Learning for offline RL, that given an offline dataset, can learn to compete against any policy (including history-dependent policies) as long as it is covered by the offline data where the coverage is measured using the relative condition number (Agarwal et al., 2021) associated with the ground
1The oracle generally refers to empirical risk minimization oracle employing non-bilevel optimization. We seek to design an algorithm that runs in polynomial time, with each oracle call counting as O(1). This algorithmic framework has lead to many successful algorithms in the literature of contextual bandit (Agarwal et al., 2014; Dudı́k et al., 2017; Foster & Rakhlin, 2020).
truth representation. Thus, our offline RL result generalizes prior offline RL works on linear MDPs (Jin et al., 2020b; Zhang et al., 2021b) which assume representation is known a priori and use linear function approximation. Computation-wise, our approach uses one call to the MLE computation oracle, and hence is oracle-efficient. REP-LCB is the first oracle efficient offline algorithm for low-rank MDP enjoying the aforementioned statistical guarantee. See Section 2 for a more detailed comparison with the existing literature on representation learning in offline RL.
Our contributions. We develop new representation learning RL algorithms that enable sample efficient learning in low-rank MDPs under both online and offline settings:
1. In the online episodic learning setting, our new algorithm REP-UCB integrates representation learning, exploration, and exploitation together, and significantly improves the sample complexity of the prior state-of-art algorithm FLAMBE;
2. In the offline learning setting, we propose a natural concentrability coefficient (i.e., relative condition number under the true representation) that captures the partial coverage condition in low-rank MDP, and our algorithm REP-LCB learns to compete against any policy (including history-dependent ones) under such a partial coverage condition.

heading: 2 RELATED WORK
text: Online Setting We list the comparison as follows, which is summarized in Table 1.
FLAMBE (Agarwal et al., 2020b) was a state-of-the-art oracle-efficient algorithm for low-rank MDPs. In all parameters, the statistical complexity is much worse than REP-UCB . Our algorithm and FLAMBE operate under the same computation oracle. FLAMBE does not balance exploration and exploitation, and uses explore-then-committee style techniques (i.e., constructions of absorbing MDPs (Brafman & Tennenholtz, 2002)) which results in its worse sample complexity.
With a more complex oracle, Moffle (Modi et al., 2021) is a model-free algorithm for low-rank MDPs, with two additional assumptions: (1) the transition has low non-negative rank (nnr), and (2) reachability in latent states. The first assumption significantly restricts the scope of low-rank MDPs as there are matrices whose nnr is exponentially larger than the rank (Agarwal et al., 2020b). The sample complexity of Moffle can scale O(d6|A|13/( 2η5(1− γ)5)), where η is the reachability probability, and 1/η could be as large as nnr1/2 (Proposition 4 in Agarwal et al. (2020b)), which essentially means that Moffle has a polynomial dependence on the nnr. Thus, Moffle needs the nnr of the transition matrix to be small.
OLIVE (Jiang et al., 2017), Witness rank (Sun et al., 2019) and Bilinear-UCB (Du et al., 2021), when specialized to low-rank MDPs, have slightly tighter dependence on d (e.g., O(d2/ 2)). But these algorithms are computationally inefficient as they are version space algorithms. Dann et al. (2021) shows that with a policy class, solving a low-rank MDP can take Ω(2d) samples. In this work, similar to Witness rank (Sun et al., 2019) and FLAMBE, we use function approximators to model the transition. Thus our positive result is not in contradiction to the result from Dann et al. (2021).
VALOR (Dann et al., 2018), PCID (Du et al., 2019a), HOMER (Misra et al., 2020), RegRL (Foster et al., 2020), and the approach from Feng et al. (2020) are algorithms for block MDPs which is a more restricted setting than low-rank MDPs. These works require additional assumptions such as deterministic transitions (Dann et al., 2018), reachability (Misra et al., 2020; Du et al., 2019a), Bellman completion (Foster et al., 2020), and strong unsupervised learning oracles (Feng et al., 2020).
Offline Setting We discuss related works in offline RL.
Uehara & Sun (2021) obtained similar results for offline RL on low-rank MDPs. Though the sample complexity in their algorithm is slightly tighter, our algorithm is oracle-efficient, while the CPPO algorithm from Uehara & Sun (2021) is a version space algorithm.
Xie et al. (2021) propose a (general) pessimistic model-free algorithm in the offline setting. We can also apply their algorithm to low-rank MDPs and show some finite-sample guarantee. However, it is unclear whether the final bounds in their results can be characterized by the relative condition number only using the true representation, and whether they can compete with history-dependent policies. Thus, our result is still considered superior on low-rank MDPs. The detail is given in Section E.
In addition to the above two works, the pessimistic approach in offline RL has been extensively investigated. Empirically, it can work on simulation control tasks (Kidambi et al., 2020; Yu et al., 2020; Kumar et al., 2020; Liu et al., 2020; Chang et al., 2021). On the theoretical side, pessimism allows us to obtain the PAC guarantee on various models when a comparator policy is covered by offline data in some forms (Jin et al., 2020b; Rashidinejad et al., 2021; Yin et al., 2021; Zanette et al., 2021b; Zhang et al., 2021b; Chang et al., 2021). However, these algorithms and their analysis rely on a known representation. We emphasize the differences in Section 5.

heading: 3 PRELIMINARIES
text: We consider a Markov Decision Process M = 〈S,A, P, r, γ, d0〉 specified by a state space S, a discrete action space A, a transition model P : S ×A → ∆(S), a reward function r : S ×A → R, a discount factor γ ∈ [0, 1), and an initial distribution d0 ∈ ∆(S). To simplify the presentation, we assume r(s, a) and d0 are known (e.g., when d0 is a probability mass only on s0, agent always starts from a fixed initial state s0)2. Following prior work (Jiang et al., 2017; Sun et al., 2019), we assume trajectory reward is normalized, i.e., for any trajectory {sh, ah}∞h=0, we have ∑∞ h=0 γ
hr(sh, ah) ∈ [0, 1]. Since the ground truth P ? is unknown, we need to learn it by interacting with environments in an online manner or utilizing offline data at hand. We remark that the extension of our all results to the finite horizon nonstationary case is straightforward.
We use the following notation. Given a policy π : S → ∆(A), which maps from state to distribution over actions and some model P , we define the value function V πP (s) = E [∑∞
h=0 γ hr(sh, ah)|s0 = s, P, π
] to represent the expected
total discounted reward of π under P starting at s. Similarly, we define the state-action Q function QπP (s, a) := r(s, a) + γEs′∼P (·|s,a)V πP (s′). The expected total discounted reward of a policy π is denoted as V πP,r := Es0∼d0V πP (s0). We define the discounted state-action occupancy distribution dπP (s, a) = (1−γ) ∑∞ t=0 γ tdπP,t(s, a), where d π P,t(s, a) is the probability of π visiting (s, a) at time step t under π and P . We slightly abuse the notation, and denote dπP (s) as the state visitation, which is equal to ∑ a∈A d π P (s, a). When P is the ground truth transition model P ?, we drop the subscript and simply use dπ to denote its state-action distribution. Unless otherwise noted, Π denotes the
class of all polices {S → ∆(A)}. We denote total variation distance of P1 and P2 by ‖P1 − P2‖1. Finally, given a vector a, we define ‖a‖2 = √ a>a, ‖a‖B = √ a>Ba. c0, c1, · · · are universal constants.
We study low-rank MDPs defined as follows (Jiang et al., 2017; Agarwal et al., 2020b). The conditions on the upper bounds of the norm of φ?, µ? are just for normalization.
Definition 1 (Low-rank MDP). Given function classes Φ ⊂ S × A → Rd and Ψ ⊂ S → Rd, a transition model P ? : S × A → ∆(A) admits a low rank decomposition with dimension d ∈ N if there exists two embedding functions φ? ∈ Φ, µ? ∈ Ψ such that
∀s, s′ ∈ S, a ∈ A : P ?(s′ | s, a) = µ?(s′)>φ?(s, a)
where ‖φ∗(s, a)‖2 ≤ 1 for all (s, a) and for any function g : S → [0, 1], ‖ ∫ µ?(s)g(s)d(s)‖2 ≤ √ d.

heading: An MDP is a low rank MDP if P ? admits such a low rank decomposition.
text: Low-rank MDPs capture the latent variable model (Agarwal et al., 2020b) shown in Figure 1 where φ?(s, a) is a distribution over a discrete latent state space Z . The block-MDP model (Du et al., 2019a) is a special instance of the latent variable model with φ?(s, a) being a one-hot encoding vector. Note the linear MDPs (Yang & Wang, 2020; Jin et al., 2020a) assume φ? is known.
Next, we explain two settings: the online learning setting and the offline learning setting. Then, we present our function approximation setup and computational oracles.
2Extension to the unknown case is straightforward. Recall the major challenging of RL is due to the unknown transition model.
Episodic Online learning In online learning, our overall goal is to learn a stationary policy π̂ so that it maximizes V π̂P?,r, where P
? is the ground truth transition. We assume that we operate under the episodic learning setting where we can only reset to states sampled from the initial distribution d0 (e.g., to emphasize the challenge from exploration, we can consider the special case where we can only reset to a fixed s0). In the episodic setting, given a policy π, sampling a state s from the state visitation dπP is done by the following roll-in procedure: starting at s0 ∼ d0, at every time step t, we terminate and return st with probability 1−γ, and otherwise we execute at ∼ π(st) and move to t+1, i.e., st+1 ∼ P (·|st, at). Such a sampling procedure is widely used in the episodic learning setting for discounted infinite horizon MDPs (e.g., (Kakade & Langford, 2002; Agarwal et al., 2021; 2020a)).
Offline learning In the offline RL, we are given a static dataset in the form of quadruples:
D = {s(i), a(i), r(i), s′(i)}ni=1 ∼ ρ(s, a)δ(r = r(s, a))P ?(s′ | s, a).
For simplicity, we assume ρ = dπbP? , where πb ∈ S → ∆(A) is a fixed behavior policy. We denote ED[f(s, a, s′)] = 1/n ∑ (s,a,s′)∈D f(s, a, s
′). To succeed in offline RL, we in general need some coverage property of ρ. One common assumption is that ρ globally covers every possible policies’ state-action distribution, i.e., maxπ,s,a dπP?(s, a)/ρ(s, a) < ∞ (Antos et al., 2008). In this work, we relax such a global coverage assumption and work under the partial coverage condition where ρ may not cover distributions of all possible policies. Instead of competing against the optimal policy under the global coverage, we aim to compete against any policies covered by the offline data. In section 5, we will precisely define the partial coverage condition using the concept of the relative condition number (Agarwal et al., 2021).
Function approximation setup and computational oracles Since µ? and φ? are unknown, we use function classes to capture them. Our function approximation and computational oracles are exactly the same as the ones used in FLAMBE. For completeness, we state the function approximation and computational oracles below. Assumption 2. We have a model classM = {(µ, φ) : µ ∈ Φ, φ ∈ Φ}, where µ? ∈ Φ, φ? ∈ Φ.
Following the norm bounds on µ?, φ? we similarly assume that the same norm bounds hold for our function approximator, i.e., for any µ ∈ Φ, φ ∈ Φ, ‖φ(s, a)‖2 ≤ 1, ∀(s, a) and ‖ ∫ µ(s)g(s)d(s)‖2 ≤ √ d,∀g : S → [0, 1], and ∫ µ>(s′)φ(s, a)d(s′) = 1, ∀(s, a).
As for computational oracles, we use a supervised learning style MLE oracle. Definition 3 (Maximum Likelihood Oracle (MLE)). Consider the model classM and a dataset D in the form of (s, a, s′), the MLE oracle returns the maixmum likelihood estimator P̂ := (µ̂, φ̂) = arg max(µ,φ)∈M ED ln(µ(s′)>φ(s, a)).
We also invoke a planning procedure for known linear MDPs with potentially nonlinear rewards, which can be done in polynomial time (after all, we know that online learning in linear MDPs can be done statistically and computationally efficient). Given a reward r and a model P := (µ, φ) with P (s′|s, a) = µ(s′)>φ(s, a) (i.e., a known linear transition with a known feature φ), we can compute the optimal policy arg maxπ V πP by standard least square value iteration which uses linear regression. A planning procedure for a known linear MDP is also used in FLAMBE, see Section 5.1 in Agarwal et al. (2020b) for details of how to implement this procedure for a known linear MDP with polynomial computation complexity.

heading: 4 REPRESENTATION LEARNING IN ONLINE SETTING
text: We consider the online episodic learning setting where the agent can only reset based on the initial state distribution d0. Thus, to find a near-optimal policy for a low-rank MDP efficiently, we need to carefully interleave representation learning, exploration, and exploitation.

heading: 4.1 ALGORITHM
text: We present our proposed algorithm in the online setting. This is described in Algorithm 1. Every iteration, Algorithm 1 rollouts its current policy π to collect a triple (s, a, s′) where s ∼ dπP? , a ∼ U(A), s′ ∼ P ?(·|s, a) where U(A) is a uniform distribution over actions. Note that to sample s ∼ dπP? , we start at s0 ∼ d0, at every time step t, we terminate and return st with probability 1 − γ, and otherwise we execute at ∼ π(st) and move to t + 1, i.e., st+1 ∼ P ?(·|st, at). After
Algorithm 1 UCB-driven representation learning, exploration, and exploitation (REP-UCB) 1: Input: Regularizer λn, parameter αn, ModelsM = {(µ, φ) : µ ∈ Ψ, φ ∈ Φ}, Iteration N 2: Initialize π0(· | s) to be uniform; set D0 = ∅ 3: for episode n = 1, · · · , N do 4: Collect a transition (s, a, s′) where s ∼ dπn−1P? , a ∼ U(A), s′ ∼ P ?(·|s, a) 5: Add (s, a, s′) to the dataset: Dn = Dn−1 + {s, a, s′} 6: Learn representation via MLE: P̂n := (µ̂n, φ̂n) = arg max(µ,φ)∈M EDn [ lnµ>(s′)φ(s, a)
] 7: Update empirical covariance matrix Σ̂n = ∑ s,a∈Dn φ̂n(s, a)φ̂n(s, a)
> + λnI 8: Set the exploration bonus:
b̂n(s, a) := min ( αn √ φ̂n(s, a)>Σ̂ −1 n φ̂n(s, a), 2 ) (1)
9: Update policy πn = arg maxπ V π P̂n,r+b̂n
10: end for 11: Return π1, · · · , πN
collecting new data and concatenating it with the existing data, we perform representation learning, i,e, learning a factorization and a representation by MLE (line 6), set the bonus based on the learned feature (Eq. 1), and update the policy via planning inside the learned model with the bonus-enhanced reward (Line 9). Note the learned transition P̂ from MLE is linear with the learned feature φ̂, and planning in a known linear MDP is known as computationally efficient (Jin et al., 2020a) (see the explanation after Definition 3 as well). Our computation MLE oracle is the same as the one used in FLAMBE (Agarwal et al., 2020b).

heading: 4.2 ANALYSIS
text: We show that our algorithm has the following PAC bound. Theorem 4 (PAC Bound for REP-UCB). Fix δ ∈ (0, 1), ∈ (0, 1). Let π̂ be a random output of π1, · · · , πN and π? = arg maxπ V πP?,r. By setting parameters as follows:
αn = O (√ (|A|+ d2) γ ln(|M|n/δ) ) , λn = O (d ln(|M|n/δ)) ,
with probability at least 1− δ, we have
V π ?
P?,r − V π̂P?,r ≤ , where the number of collected samples is at most
N log(N/δ), N := O
( d4|A|4 ln(|M|/δ)
(1− γ)3 2 · ln2
( 1 +
d4|A|4 ln(|M|/δ) (1− γ)2 2
)) .
The theorem shows that REP-UCB learns in low-rank MDPs in a statistically efficient and oracleefficient manner. To the best of our knowledge, this algorithm has the best sample complexity among all oracle efficient algorithms for low-rank MDPs.
Highlight of the analysis Below we highlight our key lemmas and proof techniques.
First, why learning in a low-rank MDP is harder than learning in models with linear structures? Unlike standard linear models such as linear MDPs (Yang & Wang, 2020; Jin et al., 2020a), KNRs (Kakade et al., 2020; Abbasi-Yadkori & Szepesvári, 2011; Mania et al., 2020; Song & Sun, 2021), and GP / kernel models (Chowdhury & Gopalan, 2019; Curi et al., 2020), we cannot get uncertainty quantification on the model in a point-wise manner. When models are linear, one can get the following style of point-wise uncertainty quantification for the learned model P̂ : ∀s, a : `(P̂ (·|s, a), P ?(·|s, a)) ≤ σ(s, a) where σ(s, a) is the uncertainty measure, and ` is some distance metric (e.g., `1 norm). With proper scaling, the uncertainty measure σ(s, a) is then used for the bonus. For example, in linear MDPs, given a dataset D = {s, a, s′}, we can learn a model P̂ (s′|s, a) := µ̂(s′)>φ?(s, a) , and get point-wise uncertainty quantification:
∀(s, a), | ∫ f(s′)µ̂>(s′)φ?(s, a)d(s′)− ∫ f(s′)µ?>(s′)φ?(s, a)d(s′)| ≤ c‖φ?(s, a)‖Σ−1
φ? (2)
for some family of functions f : S → R with Σφ? = ∑ s,a∈D φ
?(s, a)φ?(s, a)+λI (Lykouris et al., 2021; Neu & Pike-Burke, 2020). To set the scaling c properly, since φ? is known a priori, the linear regression analysis applies here, and one can apply Cauchy-Schwarz inequality to the LHS of (2) to pull out φ? and get an upper bound in the form of
‖φ?(s, a)‖Σ−1 φ?︸ ︷︷ ︸
(a)
‖ ∫ f(s){µ̂(s)− µ?(s)}d(s)‖Σφ?︸ ︷︷ ︸
(b)
where c is set to be the linear regression training error measured in the term (b) above.
However, when we jointly learn µ and φ, since nonlinear function approximation is used3, we cannot get point-wise uncertainty quantification via linear regression-based analysis. We stress that our bonus is not designed to capture the uncertainty quantification on the model error between P̂ (·|s, a) = µ̂>φ̂(s, a) and P ?(·|s, a) = µ?>φ?(s, a) in a point-wise way, which is not tractable as P̂ and P ? does not even share the same representation. Instead, the bonus is carefully designed so that it only provides near-optimism at the initial state distribution. This is formalized as follows.
Lemma 5 (Almost Optimism at the Initial State Distribution). Set the parameters as in Theorem 4. With probability 1− δ,
∀n ∈ [1, · · · , N ],∀π ∈ Π, V π P̂n,r+b̂n − V πP?,r ≥ −c1
√ |A| ln(|M|n/δ)(1− γ)
n .
We remark that the idea of optimism with respect to the initial state distribution has been used in prior works (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021; Zanette et al., 2020). However, these algorithms are not computationally efficient (i.e., they use version space instead of reward bonus), and their version-space based analysis is different from ours.
Proof sketch for Lemma 5 We start by using the simulation lemma (Lemma 20) inside the learned model which is important since our bonus b̂n uses φ̂n associated with the learned model P̂n:
V π P̂n,r+b̂n − V πP?,r = Es,a∼dπ P̂n [ b̂n(s, a) + Es′∼P̂n(·|s,a)V π P?(s ′)− Es′∼P?(·|s,a)V πP?(s′) ]
≥ Es,a∼dπ P̂n
[ b̂n(s, a)− ‖P̂n(·|s, a)− P ?(·|s, a)‖1 ] ,
from where we show that Es,a∼dπ P̂n ‖P̂ (·|s, a) − P ?(·|s, a)‖1 as a whole nearly lower bounds Es,a∼dπ P̂n [b̂n(s, a)]. The detailed procedure is illustrated in Lemma 8 in Appendix B.
Second, our bonus is using representation φ̂n that is being updated every episode, and our empirical covariance matrix Σ̂n is also updated whenever we update φ̂n, which means that standard elliptical potential based analysis (i.e., analysis used in linear bandits/MDPs with known features) cannot work here as our feature changes every episode. Instead, in our analysis, we have to keep tracking a potential function that is defined using the unknown ground truth representation φ?, i.e., the elliptical potential ‖φ?(s, a)‖2
Σ−1 ρn,φ?
, where
Σρn,φ? = nE(s,a)∼ρnφ ?(s, a)φ?(s, a)> + λnI,
and ρn is the marginal state-action distribution of Dn at an iteration n, i.e., ρn(s, a) =∑n−1 i=0 d πi P?(s)U(a)/n. Since this potential function uses the fixed representation φ
?, we can apply the standard elliptical potential argument to track the progress that our algorithm makes during learning. Below we illustrate the procedure of linking the bonus under φ̂n to the potential function ‖φ?(s, a)‖2
Σ−1 ρn,φ?
defined with respect to the true feature φ?.
3There are nonlinear models using Eluder dimension (Russo & Van Roy, 2014) as the complexity measure. However, to the best of our knowledge, the only known examples that admit low Eluder dimension are (generalized) linear models.
Algorithm 2 LCB-driven Representation Learning in offline RL (REP-LCB) 1: Input: Regularizer λ, Parameter α, Model classesM = {µ>φ : µ ∈ Ψ, φ ∈ Φ}, Dataset D. 2: Learn a model P̂ by MLE: P̂ = µ̂>φ̂ = arg maxP∈M ED[lnP (s′ | s, a)]. 3: Set the empirical covariance matrix Σ̂ = ∑ (s,a)∈D φ̂(s, a)φ̂
>(s, a) + λI . 4: Set the reward penalty:
b̂(s, a) = min ( α √ φ̂(s, a)>Σ̂−1φ̂(s, a), 2 ) .
5: Solve π̂ = arg maxπ V π P̂ ,r−b̂.
Linking bonus under φ̂n to the elliptical potential function under φ? With near optimism, using the simulation lemma (Lemma 20), we can upper bound the per-iteration regret as follows: c where fn(s, a) := ‖P̂n(· | s, a) − P ?(· | s, a)‖1. To connect the right-hand side of the above inequality to the elliptical potential under the fixed feature φ?, we show that for any function g ∈ S ×A → [0, B] for B ∈ R+,
E(s,a)∼dπn P? [g(s, a)] ≤ E(s,a)∼dπn P? [ ‖φ?(s, a)‖Σ−1
ρn,φ?
]√ nγ|A|Eρn [g2(s, a)] + γλndB2
+ √
(1− γ)|A|Eρng2(s, a).
See Lemma 13 in Appendix B. By substituting g with b̂n + fn/(1− γ), the first term of the RHS of the above inequality can be upper bounded as:
2E(s,a)∼dπn P? [ ‖φ?(s, a)‖Σ−1
ρn,φ? ] ︸ ︷︷ ︸
(G1)
√ n|A|Eρn [ f2n(s, a)
(1− γ)2 + b̂2n(s, a) ] + λnd︸ ︷︷ ︸
(G2)
.
In the term (G2), we expect nEρn [f2n(s, a)] to be O(1) as Es,a∼ρn [f2n(s, a)] is in order of 1/n due to the fact that it is the generalization bound of the MLE estimator P̂n under the training distribution ρn. For nEρn [b̂2n(s, a)], we expect it to be in the order of d as the (unnormalized) data covariance matrix Σ̂n in the bonus b̂n uses training data from ρn, i.e., we are measuring the expected bonus under the training distribution. In other words, the term (G2) scales in order of poly(d). For the term (G1), since it contains the potential function based on φ?, the sum of the term (G1) over all episodes can be controlled by the standard elliptical potential argument (see Lemma 18 and Lemma 19). 5 REPRESENTATION LEARNING IN OFFLINE SETTING
In this section, we study representation learning in the offline setting. We consider the setting where the offline data does not have a full global coverage. We present our algorithm Lower Confidence Bound driven Representation Learning in offline RL (REP-LCB) in Algorithm 2. Our proposed algorithm consists of three parts. The first part is MLE which learns a model P̂ and a representation φ̂. The second part is the construction of a penalty term b̂. Using the learned representation φ̂, we use a standard bonus in linear bandits as the penalty term as if φ̂ were the true feature. The third part is planning with the learned model P̂ and reward r − b̂.

heading: 5.1 ANALYSIS
text: We present the PAC guarantee of REP-LCB. Before proceeding, we define a relative condition number as a mean to measure the deviation between a comparator policy π and the offline data:
C?π = sup x∈Rd
x>Edπ P? [φ?(s, a)φ?>(s, a)]x
x>Eρ[φ?(s, a)φ?>(s, a)]x .
In the special case where the MDP is just a tabular MDP (i.e., φ? is a one-hot encoding vector), this is reduced to a density ratio C?∞ = maxs,a d π P?(s, a)/ρ(s, a). The relative condition number C ? π is
always no larger than the density ratio and could be much smaller for MDPs with large state spaces. Note that we quantify the relative condition number using the unknown true representation φ?. With the above setup, now we are ready to state the main theorem for REP-LCB.
Theorem 6 (PAC Bound for REP-LCB). Let ω = maxa,s(1/πb(a | s)). Denote π̂ as the output of REP-LCB. There exists a set of parameters such that with probability at least 1− δ, for any policy π (including history-dependent non-Markovian policies),
V πP?,r − V π̂P?,r ≤ c √ d4ω2C?π log(|M|/δ) (1−γ)2n .
See Theorem 15 in Appendix B for the detailed parameters. We explain several implications. First of all, this theorem shows that we can uniformly compete with any policy including history-dependent non-Markovian policies 4 satisfying the partial coverage C?π <∞. Particularly, if the optimal policy π? is covered by the offline data, i.e., C?π? < ∞, then our algorithm is able to compete against it 5. Note that assuming offline data covers π? is still a weaker assumption than the global coverage such as supπ sup(s,a) d π P?(s, a)/ρ(s, a) in prior offline RL works (Antos et al., 2008; Chen & Jiang, 2019). Second, our coverage condition is measured by a relative condition number defined using the unknown ground truth representation φ? but not depending on other features. Prior works that use relative condition numbers as measures of coverage are restricted to the settings where the ground truth representation φ? is known (Jin et al., 2020b; Chang et al., 2021; Zanette et al., 2021b).
To sum up, our algorithm is the first oracle efficient algorithm which does not need to know φ?, and requires partial coverage only in terms of φ?. Note while Uehara & Sun (2021) has a similar guarantee on low-rank MDPs, their algorithm is not oracle-efficient as it is a version space algorithm.
Highlight of the analysis Prior offline RL works that use reward penalties (Rashidinejad et al., 2021; Zhang et al., 2021b; Chang et al., 2021) all assume the representation φ? is known a priori, which allows them to use linear regression analysis to derive point-wise model uncertainty quantification which in turn serves as a penalty. In contrast, our function class is nonlinear. As in the online setting, what we can show is near pessimism in the initial state distribution d0.
Lemma 7 (Almost Pessimism at the Initial State Distribution). There exists a set of parameters that with probability 1− δ, we have
∀π ∈ Π : V π P̂ ,r−b − V π P?,r ≤ c1
√ ω log(|M|/δ)(1−γ)
n .
We leave the detailed proof to Lemma 14 in the Appendix C, which is similar to the proof of Lemma 5. Though the pessimism in the initial state distribution was recently also used in Xie et al. (2021); Zanette et al. (2021b), the derivation is totally different.
Similar to the online setting, to obtain a result only depending on the relative condition number using the ground truth representation φ? but not the learned feature φ̂, we again need to translate the penalty defined with φ̂ to the potential function ‖φ?(s, a)‖Σ−1
ρ,φ? . The same techniques that we used
for the online setting can be leveraged here to achieve the above goal.

heading: 6 CONCLUSION
text: We study online/offline RL on low-rank MDPs, where the ground truth feature is not known a priori. For online RL, our new algorithm REP-UCB significantly improves the sample complexity of the piror state-of-the-art algorithm FLAMBE in all parameters while using the same computational oracles. REP-UCB has the best sample complexity among existing oracle efficient algorithms for low-rank MDPs by a margin. Comparing to prior representation learning works on low-rank MDPs and block MDPs that rely on a forward step-by-step reward-free exploration framework, our algorithm interleaves representation learning, exploration, and exploitation together, and learns a single stationary policy. For offline RL, our new algorithm REP-LCB is the first oracle efficient algorithm for low-rank MDPs that has a PAC guarantee under a partial coverage condition measured by the relative condition number defined with the true feature representation.
4Given π = {πi}∞i=0 where πi depends on s0, a0, . . . si, V πP?,r and dπP?(s, a) are still well-defined. 5We also require ω < ∞, which is a mild assumption since it does not involve P ?. Indeed, it is much
weaker than the global coverage type assumption 1/ρ(s, a) <∞,∀(s, a).

heading: A MORE RELATED WORK
text: Here, we mention several related works in addition to Section 2
Online setting We mention works that tackle representation learning in quite different settings.
Hao et al. (2021b) consider feature selection in sparse linear MDPs with a given exploratory distribution. Zhang et al. (2021a) consider how to choose the best representation among correct representations inspired by Papini et al. (2021) (i.e., the MDP is a linear MDP under any representation in the function class Φ). In contrast, we only assume the MDP is linear under some φ? ∈ Φ.
Offline setting we mention works that tackle representation learning from different viewpoints and settings. Lu et al. (2021) consider multitask representation learning under a generative model assumption. Hao et al. (2021a) study the feature selection problem in sparse linear MDPs and Ni et al. (2021) study dimensionality reduction in a given kernel space, under the assumption that the offline data admits some form of full coverage condition.

heading: B PROOF OF THE THEORETICAL PROPERTY OF REP-UCB
text: Hereafter, we assume c0, c1, · · · , are some universal constants, and the notation
f(1/(1− γ), |A|, ln(1/δ), ln(|M|), d, n) . g(1/(1− γ), |A|, ln(1/δ), ln(|M|), d, n)
means there exists some constant c1 > 0, such that
f(1/(1− γ), |A|, ln(1/δ), ln(|M|), d, n) ≤ c1g(1/(1− γ), |A|, ln(1/δ), ln(|M|), d, n)
for any 0 ≤ γ < 1, |A|, ln(1/δ), ln(|M|), d, n. We first prove a lemma to show the optimism. Hereafter, we use the following notation to denote the data sampling distribution
dπ̄(s, a) := dπP?(s) · u(a).
where u(a) denotes the uniform distribution over actions.
First, we prove the optimism at the initial distribution. This is proved by using a simulation lemma inside the learned model which is important since both the bonus and the learned model use φ̂. In high level, we will show that the expected bonus Es,a∼dπ
P̂n b̂n(s, a) is in the same order of the
expected model error Es,a∼dπ P̂n ‖P̂n(·|s, a)− P ?(·|s, a)‖1. Note that the expectation is with respect to dπ P̂n .
Lemma 8 (Almost Optimism at the Initial Distribution). Consider an episode n(1 ≤ n ≤ N) and set
αn = O( √ (|A|+ d2) γ ln(|M|n/δ)), λn = O (d ln(|M|n/δ)) , ζn = O (
ln(|M|n/δ) n
) .
With probability 1− δ, we have
∀n ∈ [1, · · · , N ],∀π ∈ Π, V π P̂n,r+b̂n
− V πP?,r ≥ − √ (1− γ)|A|ζn.
Proof. We define
ρn = 1
n n∑ j=0 dπ̄j (s, a), Σρn,φ = nE(s,a)∼ρn [φφ >] + λnI.
In this proof, letting fn(s, a) = ‖P̂n(· | s, a)− P ?(· | s, a)‖1, we condition on the event
∀n, E(s,a)∼ρn [f 2 n(s, a)] ≤ ζn, ∀n, ∀φ, ‖φ(s, a)‖Σ̂−1n ,φ = Θ(‖φ(s, a)‖Σ−1ρn,φ),
where
Σ̂n,φ = nE(s,a)∼Dn [φφ >] + λnI.
From Lemma 11 and Lemma 17, this event happens with probability 1− δ. Then, for any policy π, from simulation lemma 20,
V π P̂n,r+b̂n − V πP?,r
= E(s,a)∼dπ P̂n [ b̂n(s, a) + γEs′∼P̂n(s,a) [ V πP?,r(s ′) ] − γEs′∼P?(s,a) [ V πP?,r(s ′) ]]
& E(s,a)∼dπ P̂n
[ min ( αn‖φ̂n(s, a)‖Σ−1
ρn,φ̂n
, 2 ) + γEs′∼P̂n(s,a) [ V πP?,r(s ′) ] − γEs′∼P?(s,a) [ V πP?,r(s ′) ]]
(3)
where in the last step, we replaced the empirical covariance by the population covariance. Note the notation . is up to universal constants. Here, since ‖V πP?,r‖∞ ≤ 1 (since we assume trajectory-wise total reward is normalized between [0, 1]), we have:∣∣∣E(s,a)∼dπ
P̂n
{ Es′∼P̂n(s,a) [ V πP?,r(s ′) ] − Es′∼P?(s,a) [ V πP?,r(s ′) ]}∣∣∣ ≤ E(s,a)∼dπ
P̂n
{fn(s, a)} .
This is further bounded by Lemma 12:
E(s,a)∼dπ P̂n
{fn(s, a)} . √ α′nE(s̃,ã)∼dπ
P̂n ‖φ̂n(s̃, ã)‖Σ−1 ρn,φ̂n
+ √ |A|ζn(1− γ). (4)
where
α′n = γ{n|A|ζn + λnd+ nζn} . γ ( |A|+ d2 ) ln(|M|n/δ).
Note we here use fn(s, a) ≤ 2 and Eρn [fn(s, a)2] ≤ ζn. Combining all things together,∣∣∣E(s,a)∼dπ
P̂n
{ Es′∼P̂n(s,a) [ V πP?,r(s ′) ] − Es′∼P?(s,a) [ V πP?,r(s ′) ]}∣∣∣ ≤ 2E(s,a)∼dπ
P̂n
{fn(s, a)} . √ α′nE(s̃,ã)∼dπ
P̂n ‖φ̂n(s̃, ã)‖Σ−1 ρn,φ̂n
+ √ (1− γ)|A|ζn
≤ αnE(s̃,ã)∼dπ P̂n ‖φ̂n(s̃, ã)‖Σ−1 ρn,φ̂n
+ √ (1− γ)|A|ζn, where αn := √ α′n. (5)
Going back to (3), we have
V π P̂n,r+b̂n − V πP?,r
& E(s,a)∼dπ P̂n
[ min ( αn‖φ̂n(s, a)‖Σ−1
ρn,φ̂n
, 2 ) + γEs′∼P̂n(s,a) [ V πP?,r(s ′) ] − γEs′∼P?(s,a) [ V πP?,r(s ′) ]]
≥ E(s,a)∼dπ P̂n
[ min ( αn‖φ̂n(s, a)‖Σ−1
ρn,φ̂n
, 2 ) −min ( αn‖φ̂n(s, a)‖Σ−1
ρn,φ̂n
+ √ (1− γ)|A|ζn, 2 )]
≥ − √ (1− γ)|A|ζn.
From the second line to the third line, we again use ‖V πP?,r‖∞ ≤ O(1) and (4). This concludes the proof. Next, we obtain the upper bound of ∑N n=0 V π? P?,r − V πn P?,r. Recall π
? is the optimal policy. Though this form is the same as a standard regret form, since we are not exactly deploying πn in episode n (recall that we play a uniform action at the end of the episode), we cannot get the regret guarantee. However, it suffices for the PAC guarantee. Lemma 9 (Regret). With probability 1− δ, we have
N∑ n=1 V π ? P?,r − V πn P?,r .
√ N ln ( 1 +
N
d2 ln(|M|/δ)
) ln(N |M|/δ) |A| 2d2
(1− γ) .
Proof. Similar to Lemma 8, letting fn(s, a) = ‖P̂n(· | s, a) − P ?(· | s, a)‖1, we condition on the event
∀n, E(s,a)∼ρn [f 2 n(s, a)] ≤ ζn, ∀n, ∀φ, ‖φ(s, a)‖Σ̂−1n ,φ = Θ(‖φ(s, a)‖Σ−1ρn,φ). (6)
From Lemma 11 and Lemma 17, this event happens with probability 1− δ. For any fixed episode n and any policy π, we have
V π ? P?,r − V πn P?,r
≤ V π ? P̂n,r+b̂n − V πnP?,r +
√ |A|ζn(1− γ) (Lemma 8)
≤ V πn P̂n,r+b̂n − V πnP?,r + √ |A|ζn(1− γ) (πn = arg maxπ V πP̂n,r+b̂n )
= E(s,a)∼dπn P? [b̂n(s, a) + γEP̂n(s′|s,a)[V πn P̂n,r+b̂n (s′)]− γEP?(s′|s,a)[V πnP̂n,r+b̂n(s ′)]] +
√ |A|ζn(1− γ).
We use the 2nd form of simulation Lemma 20 in the last display.
Then, noting ‖b̂n‖∞ ≤ 2, we have ‖V πnP̂n,r+b̂n‖∞ ≤ 2/(1− γ). Combining this fact with the above expansion, we have
V πP?,r − V πn P?,r
≤ E(s,a)∼dπn P? [b̂n(s, a)] +
( 2
1− γ
) E(s,a)∼dπn
P? [fn(s, a)] +
√ |A|ζn(1− γ)
≤ |A|E(s,a)∼dπ̄n P?
[ b̂n(s, a) ] ︸ ︷︷ ︸
(a)
+
( 2
1− γ
) |A|E(s,a)∼dπ̄n
P? [fn(s, a)]︸ ︷︷ ︸
(b)
+ √ |A|ζn(1− γ) (7)
where in the last inequality we use the notation dπ̄P?(s, a) := d π P?(s) ·U(a) and an important weighting step.
First, we calculate the first term (a) in Inequality 7. Following Lemma 13 and noting the bonus b̂n is O(1), we have
E(s,a)∼dπ̄n P?
[ b̂n(s, a) ] . E(s,a)∼dπ̄n
P?
[ min ( αn‖φ̂n(s, a)‖Σ−1
ρn,φ̂n
, 2 )] ( From (6))
. E(s̃,ã)∼dπ̄n P? ‖φ?(s̃, ã)‖Σ−1
ρn,φ?
√ nγ|A|α2nE(s,a)∼ρn [ ‖φ̂n(s, a)‖2Σ−1
ρn,φ̂n
] + dγλn
+ √ |A|α2nE(s,a)∼ρn [ ‖φ̂n(s, a)‖2Σ−1
ρn,φ̂n
] (1− γ).
Note that we use the fact that B = 2 when applying Lemma 13. In addition, we have
nE(s,a)∼ρn [ ‖φ̂n(s, a)‖2Σ−1
ρn,φ̂n
] = nTr(Eρn [φ̂nφ̂>n ]{nEρn [φ̂nφ̂>n ] + λnI}−1) ≤ d.
Then,
E(s,a)∼dπ̄n P?
[ b̂n(s, a) ] ≤ E(s̃,ã)∼dπ̄n
P? ‖φ?(s̃, ã)‖Σ−1
ρn,φ?
√ γd|A|α2n + γdλn + √ d|A|α2n(1− γ)/n.
Second, we calculate the term (b) in inequality 7. Following Lemma 13 and noting f2n(s, a) is upper-bounded by 4 (i.e., B = 4 in Lemma 13), we have E(s,a)∼dπ̄n
P? [fn(s, a)]
≤ E(s̃,ã)∼dπ̄n P? ‖φ?(s̃, ã)‖Σ−1
ρn,φ?
√{ n|A|γE(s,a)∼ρn [f2n(s, a)] } + 4γλnd+ √ |A|E(s,a)∼ρn [f2n(s, a)(1− γ)]
≤ E(s̃,ã)∼dπ̄n P? ‖φ?(s̃, ã)‖Σ−1
ρn,φ?
√ n|A|γζn + 4γλnd+ √ |A|ζn(1− γ)
≤ E(s̃,ã)∼dπ̄n P? ‖φ?(s̃, ã)‖Σ−1 ρn,φ? αn +
√ |A|ζn(1− γ),
where in the second inequality, we use Es,a∼ρf2n(s, a) ≤ ζn, and in the last line, recall√ γ √ n|A|ζn + λnd+ nζn . αn.
Then, by combining the above calculation of the term (a) and term (b) in inequality 7, we have:
V πP?,r − V πn P?,r . |A| ( E(s̃,ã)∼dπ̄n
P? ‖φ?(s̃, ã)‖Σ−1
ρn,φ?
√ d|A|α2n + dλn +
√ d|A|α2n(1− γ)
n
)
+ |A|
(1− γ)
( E(s̃,ã)∼dπ̄n
P? ‖φ?(s̃, ã)‖Σ−1 ρn,φ? αn +
√ |A|ζn(1− γ) ) .
Hereafter, we take the dominating term out. First, recall αn . √ {|A|+ d2} ln(N |M|/δ)) . √ |A|d2 ln(N |M|/δ).
Second, we also use
N∑ n=1 E(s̃,ã)∼dπ̄n P? ‖φ?(s̃, ã)‖Σ−1 ρn,φ? ≤ √√√√N N∑ n=1 E(s̃,ã)∼dπ̄n P? [φ?(s̃, ã)>Σ−1ρn,φ?φ ?(s̃, ã)]
(CS inequality)
. √√√√N (ln det( N∑ n=1 E(s̃,ã)∼dπ̄n P? [φ?(s̃, ã)φ?(s̃, ã)>])− ln det(λ1I) ) (Lemma 18 and λ1 ≤ · · · ≤ λN )
≤ √ dN ln ( 1 + N
dλ1
) .
(Potential function bound, Lemma 19 noting ‖φ?(s, a)‖2 ≤ 1 for any (s, a).) Finally,
N∑ n=1 V πP?,r − V πn P?,r . |A|
(√ dN ln ( 1 + N
dλ1
)√ d|A|α2N + dλN + N∑ n=1 √ d|A|α2n(1− γ) n )
+ |A|
(1− γ)
(√ dN ln ( 1 + N
dλ1
) αN + N∑ n=1 √ |A|ζn(1− γ) )
. |A| √ dN ln ( 1 + N
dλ1
)√ d|A|α2N + |A| (1− γ) √ dN ln ( 1 + N
dλ1
) αN
(Some algebra. We take the dominating term out.)
. √ dN ln ( 1 + N
dλ1
) |A|2d3/2 ln1/2(N |M|/δ)
(1− γ) .
This concludes the proof.
Using Lemma 9, we can immediately obtain the PAC guarantee. Theorem 10 (PAC guarantee of REP-UCB). By interacting with the environment for a number of steps at most
N log(N/δ), N := O
( d4|A|4 ln(|M|/δ) (1− γ)3 2 ln2 ( 1 + d4|A|4 ln(|M|/δ) (1− γ)2 2 )) .
with probability 1− δ, we can ensure V π?P?,r − V π̂P?,r ≤ .
Proof. From Lemma 9 and Lemma 22, when N is
O
( d4|A|4 ln(|M|/δ) (1− γ)2 2 ln2 ( 1 + d4|A|4 ln(|M|/δ) (1− γ)2 2 )) ,
with probability 1− δ, we can ensure
1
N N∑ n=1 V π ? P?,r − V πn P?,r ≤ .
With probability 1− δ, we need ln(1/δ) interactions with the environment to get one triple (s, a, s′) from dπP?(s, a). Thus, the total sample complexity is O(N(1− γ)−1 ln(N/δ)).
Next, we provide an important lemma to ensure the concentration of the bonus term. The version for fixed φ is proved in Zanette et al. (2021a, Lemma 39). Here, we take a union bound over the whole feature φ ∈ Φ. Recall
ρn = 1
n n−1∑ i=0 dπ̄iP?(s, a), (s (i), a(i)) ∼ dπ̄iP?(s, a).
The extension to the offline case is again straightforward. Recall the definition of dπ̄(s, a) in the beginning of Appendix. Lemma 11 (Concentration of the bonus term). Set λn = Θ(d ln(n|Φ|/δ)) for any n. Define
Σρn,φ = nE(s,a)∼ρn [φ(s, a)φ >(s, a)] + λnI, Σ̂n,φ = n−1∑ i=0 φ(s(i), a(i))φ>(s(i), a(i)) + λnI.
With probability 1− δ, we have ∀n ∈ N+,∀φ ∈ Φ, c1‖φ(s, a)‖Σ−1ρn,φ ≤ ‖φ(s, a)‖Σ̂−1n,φ ≤ c2‖φ(s, a)‖Σ−1ρn,φ .
Next, we show our key lemma leveraging low-rank MDP structures. For any g ∈ S ×A → R, This lemma shows that E(s,a)∼dπ
P̂n {g(s, a)} can be upper-bounded using E(s̃,ã)∼dπ P̂n ‖φ̂n(s̃, ã)‖Σ−1 ρn,φ̂n as long as we have the convergence guarantee for E(s,a)∼ρn [ g2(s, a) ] . Lemma 12 (One-step back inequality for the learned model). Take any g ∈ S × A → R s.t. ‖g‖∞ ≤ B. We condition on the event where the MLE guarantee (15):
E(s,a)∼ρn [‖P̂n(· | s, a)− P ?(· | s, a)‖21] . ζn,
holds. Then, for any policy π, we have
|E(s,a)∼dπ P̂n {g(s, a)} | ≤ E(s̃,ã)∼dπ P̂n ‖φ̂n(s̃, ã)‖Σ−1 ρn,φ̂n
√ γ √{ n|A|E(s,a)∼ρn [g2(s, a)] } + λndB2 + nB2ζn
+ √
(1− γ)|A|E(s,a)∼ρn [g2(s, a)].
Recall Σρn,φ̂n = nE(s,a)∼ρn [φ̂n(s, a)φ̂ > n (s, a)] + λnI .
Proof. First, we have an equality:
E(s,a)∼dπ P̂n {g(s, a)} = γE(s̃,ã)∼dπ P̂n ,s∼P̂n(s̃,ã),a∼π(s) {g(s, a)}+ (1− γ)Es∼d0,a∼π(s0) {g(s, a)} ,
(8)
The second term in (8) is upper-bounded by
(1− γ) √ max (s,a) d0(s)π(a | s) ρn(s, a) E(s,a)∼ρn [g2(s, a)]
≤ (1− γ) √ max (s,a) d0(s)π(a | s) (1− γ)d0(s)u(a) E(s,a)∼ρn [g2(s, a)] ≤ √ (1− γ)|A|E(s,a)∼ρn [g2(s, a)].
Next we consider the first term in (8). By CS inequality, we have
E(s̃,ã)∼dπ P̂n ,s∼P̂n(s̃,ã),a∼π(s) {g(s, a)} = E(s̃,ã)∼dπP̂n φ̂n(s̃, ã)
> ∫ ∑
a
µ̂n(s)π(a | s)g(s, a)d(s)
≤ E(s̃,ã)∼dπ P̂n ‖φ̂n(s̃, ã)‖Σ−1 ρn,φ̂n ∥∥∥∥∥ ∫ ∑
a
µ̂n(s)π(a | s)g(s, a)d(s) ∥∥∥∥∥ Σρn,φ̂n .
Then, ‖ ∫ ∑
a
µ̂n(s)π(a | s)g(s, a)d(s)‖2Σρn,φ̂n
≤ {∫ ∑ a µ̂n(s)π(a | s)g(s, a)d(s) }> { nE(s,a)∼ρn [φ̂nφ̂ > n ] + λnI }{∫ ∑ a µ̂n(s)π(a | s)g(s, a)d(s) }
≤ nE(s̃,ã)∼ρn  [∫ ∑
a
µ̂n(s) >φ̂n(s̃, ã)π(a | s)g(s, a)d(s) ]2+B2λnd (Use the assumption ‖ ∑ a π(a | s)g(s, a)‖∞ ≤ B and ∫ ‖µ̂n(s)h(s)d(s)‖2 ≤ √ d for any h : S → [0, 1].)
= nE(s̃,ã)∼ρn [{ Es∼P̂n(s̃,ã),a∼π(s) [g(s, a)] }2] +B2λnd
= nE(s̃,ã)∼ρn [{ Es∼P?(s̃,ã),a∼π(s) [g(s, a)] }2] +B2λnd+ nζnB 2 (MLE guarantee)
≤ nE(s̃,ã)∼ρn,s∼P?(s̃,ã),a∼π(s) [ g2(s, a) ] +B2λnd+ nζnB 2. (Jensen) ≤ n min{u(a)} { E(s̃,ã)∼ρn,s∼P?(s̃,ã),a∼U(A) [ g2(s, a) ]} + λndB 2 + nζnB 2
(Importance sampling) ≤ n|A| { 1
(γ) E(s,a)∼ρn
[ g2(s, a) ]} + λndB 2 + nζnB 2,
Then, the final statement is immediately concluded.
Here, we show a similar lemma as Lemma 12. The difference is we aim for calculating E(s,a)∼dπ
P? {g(s, a)} instead of E(s,a)∼dπ
P̂n {g(s, a)} . For any g ∈ S × A → R, this lemma shows that E(s,a)∼dπ
P? {g(s, a)} can be upper-bounded using E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1 ρn,φ̂? as long as we have the convergence guarantee for E(s,a)∼ρn [ g2(s, a) ] . Note comparing to Lemma 12, this is not a probabilistic statement. Note that ‖φ?(s, a)‖Σ−1 ρn,φ?
is the usual elliptical potential function under the fixed representation φ?. Lemma 13 (One-step back inequality for the true model ). Take any g ∈ S × A → R such that ‖g‖∞ ≤ B. Then,
E(s,a)∼dπ P? {g(s, a)} ≤ E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖
Σ−1 ρn,φ?
√ γ √ n|A|E(s,a)∼ρn [g2(s, a)] + λndB2
+ √
(1− γ)|A|E(s,a)∼ρn [g2(s, a)].
Recall Σρn,φ? = nE(s,a)∼ρn [φ?{φ?}>] + λnI .
Proof. First, we have E(s,a)∼dπ
P? {g(s, a)} = γE(s̃,ã)∼dπ P? ,s∼P?(s̃,ã),a∼π(s) {g(s, a)}+ (1− γ)Es∼d0,a∼π(s0) {g(s, a)} .
(9)
The second term in (9) is upper-bounded by
(1− γ) √ max (s,a) d0(s)π(a | s) ρn(s, a) E(s,a)∼ρn [g2(s, a)] ≤ √ |A|E(s,a)∼ρn [g2(s, a)] (1− γ).
By CS inequality, the first term in (9) is further bounded as follows:
E(s̃,ã)∼dπ P? ,s∼P?(s̃,ã),a∼π(s) {g(s, a)} = E(s̃,ã)∼dπ P? φ?(s̃, ã)> ∫ ∑ a µ?(s)π(a | s)g(s, a)d(s)
≤ E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1
ρn,φ? ∥∥∥∥∥ ∫ ∑
a
µ?(s)π(a | s)g(s, a)d(s) ∥∥∥∥∥ Σρn,φ? .
Here, we have ‖ ∫ ∑
a
µ?(s)π(a | s)g(s, a)d(s)‖2Σρn,φ?
≤ {∫ ∑ a µ?(s)π(a | s)g(s, a)d(s) }> { nE(s,a)∼ρn [φ ?(s, a){φ?(s, a)}>] + λnI }{∫ ∑ a µ?(s)π(a | s)g(s, a)d(s) }
≤ nE(s̃,ã)∼ρn  [∫ ∑
a
µ?(s)>φ?(s̃, ã)π(a | s)g(s, a)d(s) ]2+ λndB2 ≤ n { E(s̃,ã)∼ρn,s∼P?(s̃,ã),a∼π(s) [ g2(s, a) ]} + λndB 2. (Jensen)
Therefore,
n { E(s̃,ã)∼ρn,s∼P?(s̃,ã),a∼π(s) [ g2(s, a) ]} + λndB
≤ n min{u(a)}
{ E(s̃,ã)∼ρn,s∼P?(s̃,ã),a∼U(A) [ g2(s, a) ]} + λndB 2
≤ n|A| { 1
γ E(s,a)∼ρn
[ g2(s, a) ]} + λndB 2.
Finally, we have
E(s,a)∼dπ P? {g(s, a)} ≤ E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1ρn,φ? √{ n|A| γ E(s,a)∼ρn [g2(s, a)] } + λndB2
+ √ |A|E(s,a)∼ρn [g2(s, a)] (1− γ).
This concludes the proof.

heading: C PROOF OF THE THEORETICAL PROPERTY OF REP-LCB
text: This section provides the detailed proofs for our results in the offline setting.
Below we first prove that V π P̂ ,r−b̂ is an almost pessimistic estimator of V π P?,r.
Lemma 14 (Almost Pessimism at the Initial Distribution). Let ω = maxa,s 1/πb(a | s). Set
α = c1 √ (ω + d2) γ ln(|M|/δ), λ = O(d ln(|M|/δ)), ζ = O (
ln(|M|/δ) n
) .
With probability 1− δ, for any policy π, we have
V π P̂ ,r−b̂ − V π P?,r ≤
√ ω ln(|M|/δ)(1− γ)
n .
Proof. We define
Σρ,φ = nE(s,a)∼ρ[φφ>] + λI, Σ̂φ = nED[φφ>] + λI.
where λ = O(d ln(|M|/δ)). In this proof, letting f(s, a) = ‖P̂ (· | s, a) − P ?(· | s, a)‖1, we condition on the events:
E(s,a)∼ρ[f2(s, a)] ≤ ζ, ∀φ ∈ Φ : ‖φ(s, a)‖Σ̂−1φ = Θ(‖φ(s, a)‖Σ−1ρ,φ). (10)
where ζ = O(ln(|M|/δ)/n). From the offline version of Lemma 17 and Lemma 11 6, this event happens with probability 1− δ. Then, from simulation lemma (Lemma 20),
V π P̂ ,r−b̂ − V π P?,r
= E(s,a)∼dπ P̂ [ −b̂(s, a) + γEs′∼P̂ (s,a) [ V πP?,r(s ′) ] − γEs′∼P?(s,a) [ V πP?,r(s ′) ]]
. E(s,a)∼dπ P̂
[ −min ( α‖φ̂(s, a)‖Σ−1
ρ,φ̂
, 2 ) + γEs′∼P̂ (s,a) [ V πP?,r(s ′) ] − γEs′∼P?(s,a) [ V πP?,r(s ′) ]] .
(From (10))
Here, we have∣∣∣E(s,a)∼dπ P̂ { Es′∼P̂ (s,a) [ V πP?,r(s ′) ] − Es′∼P?(s,a) [ V πP?,r(s ′) ]}∣∣∣ ≤ E(s,a)∼dπ P̂ {f(s, a)} ,
noting ‖V πP?,r‖∞ ≤ 1. This is further bounded by Lemma 16:
E(s,a)∼dπ P̂ {f(s, a)} .
√ α′E(s̃,ã)∼dπ
P̂ ‖φ̂(s̃, ã)‖Σ−1
ρ,φ̂
+ √ ωζ(1− γ). (11)
where
α′ = nγωζ + γ2λd+ γ2nζ . ( ω + d2 ) γ ln(|M|/δ).
Here, we use f(s, a) ≤ 2 in Lemma 16 and E(s,a)∼ρ[f2(s, a)] ≤ ζ. Thus,∣∣∣E(s,a)∼dπ
P̂
{ Es′∼P̂ (s,a) [ V πP?,r(s ′) ] − Es′∼P?(s,a) [ V πP?,r(s ′) ]}∣∣∣ ≤ E(s,a)∼dπ P̂ {f(s, a)}
≤ √ α′E(s̃,ã)∼dπ
P̂ ‖φ̂(s̃, ã)‖Σ−1
ρ,φ̂
+ √ ωζ(1− γ)
= αE(s̃,ã)∼dπ P̂ ‖φ̂(s̃, ã)‖Σ−1
ρ,φ̂
+ √ ωζ(1− γ), α = √ α′.
Going back to the simulation lemma 20, we have
V π P̂ ,r−b̂ − V π P?,r
. E(s,a)∼dπ P̂
[ −min ( α‖φ̂(s, a)‖Σ−1
ρ,φ̂
, 2 ) + Es′∼P̂ (s,a) [ V πP?,r(s ′) ] − Es′∼P?(s,a) [ V πP?,r(s ′) ]]
≤ E(s,a)∼dπ P̂
[ −min ( α‖φ̂(s, a)‖Σ−1
ρ,φ̂
, 2 ) + min ( α‖φ̂(s, a)‖Σ−1
ρ,φ̂
+ √ ωζ(1− γ), 2 )] ≤ √ ωζ(1− γ).
This concludes the proof.
With the above lemma, now we can proceed to prove the main theorem.
Theorem 15 (PAC guarantee of REP-LCB). Set the parameters as in Lemma 14. With probability 1− δ, for any comparator policy π including history-dependent non-Markovian policies, we have
V πP?,r − V π̂P?,r . ωd2
(1− γ)
√ C?π ln(|M|/δ)
n ,
where C?π is the relative condition number under φ ?:
C?π := sup x∈R
x>E(s,a)∼dπ P? [φ?(s, a){φ?(s, a)}>]x x>E(s,a)∼ρ[φ?(s, a){φ?(s, a)}>]x .
6We can remove lnn since n is fixed in the offline setting.
Proof. In this proof, letting f(s, a) = ‖P̂ (· | s, a)− P ?(· | s, a)‖1 we condition on the events:
E(s,a)∼ρ[f2(s, a)] ≤ ζ, ∀φ ∈ Φ : ‖φ(s, a)‖Σ̂−1φ = Θ(‖φ(s, a)‖Σ−1ρ,φ). (12)
From Lemma 11 and Lemma 17, this event happens with probability 1− δ. For any policy π, we have
V πP?,r − V π̂P?,r ≤ V πP?,r − V π̂P̂ ,r−b̂ + √ ωζ(1− γ) (Lemma 14)
≤ V πP?,r − V πP̂ ,r−b̂ + √ ωζ(1− γ)
. E(s,a)∼dπ P? [b̂(s, a)]︸ ︷︷ ︸ (a) +
( 1
1− γ
) E(s,a)∼dπ
P? [f(s, a)]︸ ︷︷ ︸
(b)
+ √ ωζ(1− γ).
Recall f(s, a) = ‖P̂ (· | s, a)− P ?(· | s, a)‖1. From the second line to the third line, note though π̂ is the argmax over Markoovian polices, π̂ is also the argmax over all history-dependent polices. In the last line, we use a simulation lemma 20, which is tailored to a time-inhomogeneous policy. We here use ‖V π
P̂ ,r−b̂‖∞ ≤ 2/((1 − γ)). noting ‖b̂‖∞ = O(1). We further calculate the first term (a). Considering the offline version of Lemma 13 and noting ‖b̂‖∞ ≤ 2, we have
E(s,a)∼dπ P? [b̂(s, a)] . E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1
ρ,φ?
√ nω { γE(s,a)∼ρ [ b̂2(s, a) ]} + γλd
+ √ ω(1− γ){Eρ[b̂2(s, a)]}1/2.
From (12), we have nE(s,a)∼ρ [ b̂2(s, a) ] ≤ nE(s,a)∼ρ [ min ( α2‖φ̂(s, a)‖2
Σ−1 ρ,φ̂
, 4 )] ≤ nE(s,a)∼ρ [ α2‖φ̂(s, a)‖2
Σ−1 ρ,φ̂ ] (13)
≤ Tr[nE[φ̂φ̂>]{nE[φ̂φ̂>] + λI}−1] ≤ Tr[n(E[φ̂φ̂>] + λI){nE[φ̂φ̂>] + λI}−1] ≤ d. Thus,
E(s,a)∼dπ P? [b̂(s, a)] ≤ E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1
ρ,φ?
√ ωdα2γ + γλd+
√ ωdα2(1− γ)
n .
Second, we further calculate the second term (b). Considering the offline version of Lemma 13 and noting f2(s, a) is upper-bounded by 4,
E(s,a)∼dπ P? [f(s, a)]
= E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1
ρ,φ?
√ nω { γE(s,a)∼ρ [f2(s, a)] } + 4γλd+ √ ωE(s,a)∼ρ [f2(s, a)] (1− γ)
. E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1
ρ,φ?
√ ω {nγζ}+ γλd+ √ ωζ(1− γ)
. E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1 ρ,φ? α+
√ ωζ(1− γ).
In the final line, recall √ ω {nγζ}+ γλd+ γnζ ≤ α.
Finally, by combining the calculation of the first term (a) and the second term (b), we have
V πP?,r − V π̂P?,r . E(s̃,ã)∼dπP?‖φ ?(s̃, ã)‖Σ−1
ρ,φ?
√ dα2ωγ + γλd+
√ ωα2d(1− γ)
n
+ α
(1− γ) E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1 ρ,φ? +
√ ωζ
(1− γ)
Now, we use the fact E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1 ρ,φ? is upper-bounded as
E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖Σ−1 ρ,φ? ≤ √ E(s̃,ã)∼dπ P? ‖φ?(s̃, ã)‖2 Σ−1 ρ,φ? ≤ √ C?E(s̃,ã)∼ρ‖φ?(s̃, ã)‖2Σ−1 ρ,φ?
(Refer to Lemma 21) ≤ √ C?d/n. (From (13))
Finally, we have
V πP?,r − V π̂P?,r
.
√ C?d
n
√ dα2ωγ + γλd+
α
(1− γ)
√ C?d
n +
√ ωα2d(1− γ)
n +
√ ωζ
(1− γ)
.
√ C?d
n
√ dα2ωγ +
α
(1− γ)
√ C?d
n (Take out two dominating terms)
. ωd2
(1− γ)
√ C? ln(|M|/δ)
n .
The lemma below is a key technical lemma for our proof. It shows that one can relate the expected value of any function f(s, a) with respect to dπ
P̂ (i.e., inside the learned model P̂ ) to the potential
function with respect to dπ P̂ , i.e., E(s̃,ã)∼dπ P̂ ‖φ̂(s̃, ã)‖Σ−1
ρ,φ̂
. Pairing φ̂ and P̂ is important since P̂ is
the low-rank transition model defined using φ̂. As we have seen in the above analysis, when using the lemma below, we instantiate f(s, a) := ‖P̂ (·|s, a)− P ?(·|s, a)‖1. Lemma 16 (One-step back inequality for the learned model in offline setting). Take any f ⊂ S × A → R s.t. ‖f‖∞ ≤ B. We condition on the event where the MLE guarantee holds:
E(s,a)∼ρ‖P̂ (· | s, a)− P ?(· | s, a)‖21 . ζ.
Then, letting ω = maxs,a(1/πb(a | s)), for any policy π, we have
|E(s,a)∼dπ P̂ {f(s, a)} | ≤ E(s̃,ã)∼dπ P̂ ‖φ̂(s̃, ã)‖Σ−1
ρ,φ̂
√{ nωE(s,a)∼ρ [f2(s, a)] } + γ2λdB2 + nγ2ζB2
+ √ ωE(s,a)∼ρ [f2(s, a)] (1− γ).
Proof. First, we have an equality:
E(s,a)∼dπ P̂ {f(s, a)} = γE(s̃,ã)∼dπ P̂ ,s∼P̂ (s̃,ã),a∼π(s) {f(s, a)}+ (1− γ)Es∼d0,a∼π(s0) {f(s, a)} .
(14)
The second term in (14) is upper-bounded by
Es∼d0,a∼π(s0) {f(s, a)} ≤ Es∼d0,a∼π(s0) { f2(s, a) } }1/2 = √ ωE(s,a)∼ρ [f2(s, a)] /(1− γ).
Next we consider the first term in (14). By CS inequality, we have∣∣∣E(s̃,ã)∼dπ P̂ ,s∼P̂ (s̃,ã),a∼π(s) {f(s, a)} ∣∣∣ = ∣∣∣∣∣E(s̃,ã)∼dπP̂ φ̂(s̃, ã)> ∫ ∑
a
µ̂(s)π(a | s)f(s, a)d(s) ∣∣∣∣∣ ≤ E(s̃,ã)∼dπ
P̂ ‖φ̂(s̃, ã)‖Σ−1
ρ,φ̂
‖ ∫ ∑
a
µ̂(s)π(a | s)f(s, a)d(s)‖Σρ,φ̂ .
Then, ‖ ∫ µ̂(s)π(a | s)f(s, a)d(s, a)‖2Σρ,φ̂
≤ {∫ ∑ a µ̂(s)π(a | s)f(s, a)d(s) }> { nE(s,a)∼ρ[φ̂φ̂>] + λI }{∫ ∑ a µ̂(s)π(a | s)f(s, a)d(s) }
≤ nE(s̃,ã)∼ρ  [∫ ∑
a
µ̂(s)>φ̂(s̃, ã)π(a | s)f(s, a)d(s) ]2+B2λd (Use the assumption ‖ ∑ a f(s, a)‖∞ ≤ B and ‖ ∫ µ̂(s)h(s)d(s)‖2 ≤ √ d for h : S → [0, 1].)
= nE(s̃,ã)∼ρ{Es∼P̂ (s̃,ã),a∼π(s) [f(s, a)] 2}+B2λd
= nE(s̃,ã)∼ρ{Es∼P?(s̃,ã),a∼π(s) [f(s, a)] 2}+B2λd+ nB2ζ (MLE guarantee and ‖Ea∼π(·)[f2(·, a)]‖∞ ≤ B2.) ≤ n { E(s̃,ã)∼ρ,s∼P?(s̃,ã),a∼π(s) [ f2(s, a) ]} +B2λd+ nB2ζ. (Jensen)
Finally, the first term in (14) is upper-bounded by n { E(s̃,ã)∼ρ,s∼P?(s̃,ã),a∼π(s) [ f2(s, a) ]} + λdB2 + nB2ζ
≤ nω { E(s̃,ã)∼ρ,s∼P?(s̃,ã),a∼πb(s) [ f2(s, a) ]} + λdB2 + nB2ζ (Importance sampling)
≤ nω { 1
γ E(s,a)∼ρ
[ f2(s, a) ]} + λdB2 + nB2ζ. (Definition of ρ)
The final statement is immediately concluded.

heading: D AUXILIARY LEMMAS
text: First, we present the MLE guarantee. Regarding the proof, refer to Agarwal et al. (2020b, Theorem 21). Note P̂n and π̄n are the quantities appearing in the proposed online algorithm. We can also immediately obtain the statement to the offline case.
Lemma 17 (MLE guarantee). For a fixed episode n, with probability 1− δ, n−1∑ j=0 E(s,a)∼dπ̄j [‖P̂n(· | s, a)− P ?(· | s, a)‖21] . nζ, ζ := ln(|M|/δ) n .
As a straightforward corollary, with probability 1− δ,
∀n ∈ N+, n−1∑ j=0 E(s,a)∼dπ̄j [‖P̂n(· | s, a)− P ?(· | s, a)‖21] . nζn, ζn := ln(|M|n/δ) n . (15)
The following is a standard inequality to prove regret bounds for linear models. Refer to Agarwal et al. (2020a, Lemma G.2.)
Lemma 18. Consider the following process. For n = 1, · · · , N ,Mn = Mn−1+Gn withM0 = λ0I and Gn being a positive semidefinite matrix with eigenvalues upper-bounded by 1. We have that:
2 ln det(MN )− 2 ln det(λ0I) ≥ N∑ n=1 Tr(GnM −1 n−1).
Lemma 19 (Potential function lemma). Suppose Tr(Gn) ≤ B2. 2 ln det(MN )− 2 ln det(λ0I) ≤ d ln ( 1 + NB2
dλ0
) .
Proof. Let σ1, · · · , σd be the set of singular values of MN recalling MN is a positive semidefinite matrix. Then, by the AM-GM inequality,
ln det(MN )/det(λ0I) = ln d∏ i=1 (σi/λ0) ≤ ln d
( 1
d d∑ i=1 (σi/λ0))
)
Since we have ∑ i σi = Tr(MN ) ≤ dλ0 +NB2, the statement is concluded.
Lemma 20 (Simulation lemma). Given two MDPs (P ′, r+b) and (P, r), for any policy π, we have:
V πP ′,r+b − V πP,r = 1
1− γ E(s,a)∼dπ P ′ [b(s, a) + γEP ′(s′|s,a)[QπP,r(s′, π)]− γEP (s′|s,a)[QπP,r(s′, π)]]
and
V πP ′,r+b − V πP,r = 1
1− γ E(s,a)∼dπP [b(s, a) + γEP ′(s′|s,a)[Q π P,r+b(s ′, π)]− γEP (s′|s,a)[QπP ′,r+b(s′, π)]].
Proof. We use
V πP − f(s0, π) = 1
1− γ EdπP [r(s, a) + γEP (s′|s,a)[f(s
′, π)]− f(s, a)]]
Then,
V πP ′,r+b − V πP,r = 1
1− γ E(s,a)∼dπ P ′ [r(s, a) + b(s, a) + γEP ′(s′|s,a)[QπP,r(s′, π)]−QπP,r(s, a)]]
= 1
1− γ E(s,a)∼dπ P ′ [b(s, a) + γEP ′(s′|s,a)[QπP,r(s′, π)]− γEP (s′|s,a)[QπP,r(s′, π)]].
Similarly,
V πP,r − V πP ′,r+b = 1
1− γ E(s,a)∼dπP [r(s, a) + γEP (s′|s,a)[Q π P ′,r+b(s ′, π)]−QπP ′,r+b(s, a)]]
= 1
1− γ E(s,a)∼dπP [−b(s, a) + γEP (s′|s,a)[Q π P ′,r+b(s ′, π)]− γEP ′(s′|s,a)[QπP,r(s′, π)]].
The following lemma is used to deal with the distribution shift in the offline setting. For the proof, refer to Chang et al. (2021).
Lemma 21 (Distribution shift lemma). Consider any policy π and state-action distribution ρ, and any representation φ?, we have:
E(s,a)∼dπ P? [φ?(s, a){φ?(s, a)}>] ≤ C?Eρ[φ?(s, a){φ?(s, a)}>], C? := sup x∈R
x>E(s,a)∼dπ P? [φ?{φ?}>]x x>E(s,a)∼ρ[φ?{φ?}>]x .
This is some auxiliary lemma to convert the finite sample error bound into the sample complexity.
Lemma 22 (Conversion of finite sample error bounds into sample complexities). By taking
N = 1/ ′2 × ln2(1 + 1/ ′2), ′ = a1 ln 1/2(e+ a2) ln 1/2(e+ a3) .
It satisfies
a1 √ 1/N ln1/2(1 + a2N) ln 1/2(1 + a3N) < c .
where c is a constant independent of a1, a2, a3.
Proof. We first have a1 √ 1/N ln1/2(1 + a2N) ln 1/2(1 + a3N) ≤ a1 max(ln1/2(1 + a2) ln1/2(1 + a3), 1) √ 1/N ln(1 +N).
Here, we use ln1/2(1 + a2N) ≤ {ln(1 + a2) + ln(1 +N)}1/2 ≤ √ max(1, ln(1 + a2)) ln(1 +N). Then, we prove when N = 1/ 2 × ln2(1 + 1/ 2).√ 1/N ln(1 +N) < .
This is proved by√ 1/N ln(1 +N) ≤ × ln(1 + 1/
2 × ln2(1 + 1/ 2)) ln(1 + 1/ 2)
≤ × ln(1 + 1/ 2) + ln(1 + ln2(1 + 1/ 2))
ln(1 + 1/ 2)
≤ + × ln(1 + ln 2(1 + 1/ 2))
ln(1 + 1/ 2)
≤ + × 0.5{1 + ln 2(1 + 1/ 2))}1/2 − 1
ln(1 + 1/ 2)
. .
From the third line to the fourth line, we use ln(x) ≤ 0.5(x1/2 − 1) for x > 0. Then, the final statement is concluded.

heading: E MORE COMPARISON TO XIE ET AL. (2021)
text: We briefly explain the guarantee when we use Algorithm 1 (Xie et al., 2021). For a given reward r, we first define a new feature class Φ+r . Definition 23 (Augmented feature). Let φ = [φ1, · · · , φd].
Φ+r = {φ+r ;φ ∈ Φ}, φ+r = [φ1, · · · , φd, r].
Then, we set F = {a>φ+r | ‖a‖2 ≤ c √ d+ 1, φ+r ∈ Φ+r }. where c is some suitable constant. Given the hypothesis class F for the Q-function, we can run Algorithm 1 in (Xie et al., 2021). We denote the output policy as π̂.
We check two assumptions to ensure the algorithm works. The first assumption is realizability. This is satisfied since for any policy π ∈ Π (Π is the class of all Markovian polices), we haveQπP?,r ∈ F . The second assumption is completeness. This is also satisfied since T πP?,rF ⊂ F for any policy π ∈ Π where T πP?,r is a Bellman-operator s.t. T πP?,r : {S × A → R} 3 f 7→ r(s, a) + γEs′∼P?(s,a)[f(s′, π)] ∈ {S ×A → R}, where we denote f(s, π) = Ea∼π(s)f(s, a). Then, by invoking their Corollary 5, we have Theorem 24 (PAC bound based on Xie et al. (2021)). With probability 1− δ,
∀π ∈ Π : V πP?,r − V π̂P?,r ≤ c
√ C†π,r
(1− γ)2
( (d+ 1) log(1/δ) log |A|
n
)1/5 .
where
C†π,r = sup φ+r ∈Φ+r sup a∈Rd+1
a>Edπ P? [φ+r {φ+r }>]a a>Eρ[φ+r {φ+r }>]a .
We compare the above result with our result in Theorem 6. First, since C†r includes r and all possible features in Φ, this partial coverage condition is stronger than ours (recall our partial coverage condition is only related to the true representation φ?), and we always have C?π ≤ C†π,r. Secondly, the dependence on n is much worse. Third, it is unclear whether the learned policy can compete against any history-dependent policy. Recall in Theorem 6, we show that our algorithm can compete with any history-dependent policies.

Review summaries:
This paper studies low-rank episodic MDPs when the reward is deterministic and the reward function is known. The paper proposes a method that collects data, uses the data to estimate the low-rank MDP, uses this estimate, confidence bound around it, to come up with the policy to be used for the next episode.   The algorithm provides PAC style bound on how many episodes are needed to learn epsilon optimal policy.  The authors extend the scope of this work to offline low-rank MDP where a logged data set is given where an RL algorithm is needed to learn and optimize to come up with a new policy.        

This paper studies the representation learning for linear MDP. The authors provide an online algorithm and its offline counterpart. Theoretical analysis justifies that the proposed algorithms are sample efficient.

This paper studies representation learning in low rank MDPs in both online and offline setting. This paper proposes REP-UCB algorithm, which improves the sample complexity of FLAMBE significantly. In addition, REP-UCB is much simpler than FLAMBE.  For the offline setting, the REP-LCB algorithm also achieves polynomial sample complexity with partial coverage. In particular, the coverage is measured by the ground-truth feature.