Representation Learning for Online and Offline RL in Low-rank MDPs
Authors: Masatoshi Uehara
Xuezhou Zhang
Wen Sun
Venue: ICLR 2022 Spotlight (2022)
Keywords: Provably sample efficient Reinforcement Learning
PAC bounds
Representation learning
Low-rank MDP

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), the state-of-art algorithm for learning representations in low-rank MDPs, we propose an algorithm REP-UCB Upper Confidence Bound driven Representation learning for RL), which significantly improves the sample complexity from $\widetilde{O}( A^9 d^7 / (\epsilon^{10} (1-\gamma)^{22}))$ for FLAMBE to $\widetilde{O}( A^4 d^4 / (\epsilon^2 (1-\gamma)^{2})  )$ with $d$ being the rank of the transition matrix (or dimension of the ground truth representation), $A$ being the number of actions, and $\gamma$ being the discounted factor. Notably, REP-UCB is simpler than FLAMBE, as it directly balances the interplay between representation learning, exploration, and exploitation, while FLAMBE is an explore-then-commit style approach and has to perform reward-free exploration step-by-step forward in time. For the offline RL setting, we develop an algorithm that leverages pessimism to learn under a partial coverage condition: our algorithm is able to compete against any policy as long as it is covered by the offline distribution.

Published as a conference paper at ICLR 2022
REPRESENTATION LEARNING FOR ONLINE AND OF- FLINE RL IN LOW-RANK MDPs
Masatoshi Uehara
Xuezhou Zhang
Department of Computer Science Cornell University, Ithaca, NY 14850, USA mu2230 cornel1
Department of Electrical and Computer Engineering
Princeton University, NJ 08544,USA
xz xz 8 w273922rrnnctton...
Wen Sun
Department of Computer Science, Cornell University, Ithaca, NY 14850, USA ws ws4550cornel edd
ABSTRACT
This work studies the question of Representation Learning in RL: how can we learn a compact www-dimensional representation such that top of the represen- tation we can perform RL procedures such as exploration and exploitation, in a sample efficient manner. We focus 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 FLAMMB((Agawaa et al. (020b) -the state-of-art algorithm for learning representations in low-rank MDPs, we propose an algo- rithm REP-UCB Upper Confidence Bound driven REPresentation learning for RL, which significantly improves the sample complexity from Ö(A7 (1c 2)) for FLAMBE Odd' A? (12(112111))) ( ))) with d being the rank of the transi- tion matrix (or dimension of the ground truth representation), A being the number actions, and y being the discount factor. Notably, REP-UCB simpler than FLAMBE, as it directly balances the interplay between representation learning, exploration, and exploitation, while FLAMBE is an sxplore-then-comee style ap- proach and has to perform reward-free exploration step-by-step forward in time. For the offline RI. setting, we develop an algorithm that leverages pessimism to learn under partial coverage condition: algorithm is able to compete against any policy as long as it is covered by the offline data distribution.
INTRODUCTION
When applying Reinforcement Learning (RL) to large-scale problems where data is complex and high-dimensional, learning effective transformations of the data, i.c., 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 for non-sequential and non- interactive 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 al. 2019b) or the optimal value functions (Wang et al. 2020 Weisz et al. 2021). RI. 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 Pro- cesses (MDPs). Despite the recent empirical success of representation learning in RL, its statistical guarantee and theoretical properties remain under-invesse 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 low-rank factoriza- tion, i.c., there exists two unknown mappings ju(s). ø(s, a), such that P(s js, a) = pus')" o(s, a) for all 8, a, s', where P(s'|ss a) is the probability of transiting to the state s' under the current

---

Published as a conference paper at ICLR 2022
state and action (s, a). The representation o in low-rank MDP not only linearizes the optimal state-action value function of the MDP (Jin et al. 2020a). but also linearizes the transition operator. A low-rankness assumption
large stochastic matrices is common and natural assumption and has enabled successful development of algorithms for real world applications such as movie recommendation systems (Koren et al. 2009). We note that low-rank MDP strietly generalizes the linear MDP model (Yang& & Wang 2020 Jin et al. (020a)) which assumes ø is known a priori. The unknown representation e 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 al. (02T) via leveraging rich function approximators. However, these algorithms are version space algorithms and are not computationally efficient. Recent work FLAMBE that learns in low-rank MDPs with a polynomial sample complexity, where the computation oracle is Maximum Likelihood Estimation (MLE) operating under the standard supervised learning style Empirical Risk Minimization (ERM) setting. In this work, we follow the same setup from FLAMBE (Agarwal et al. 2020b), and propose new algo- rithm - 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 Odd A /(e2 - -) for FLAMBE Odd (e2((221 - 7)"), where d is the rank of the transition matrix (or dimension of the true representation). A is the number of actions, c the suboptimality gap and y E [0. 1) is the discount factor in the MDP. Our algorithm is also arguably much simpler than FLAMBE: FLAMBE is an thpoor-then-coomiit algorithm, has to explore in a layer-by-layer forward way, and does not permit data sharing across different time steps. In contrast, UCP carefully trades exploration versus exploitation by combining the ward signal and exploration bonus (constructed using the latest learned representation), and enables data sharing across all time steps P 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 leamingg in Table
proposes an oracle-efficient algorithm {
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 specif- ically, we propose an algorithm REP-LCB-Lowe Confidence Bound driven Repreprexentation Learning for offline RL, that given an offline dataset, can learn to compete against any policy (including
policies) as long as it is covered by the offline data where the coverage measured using the relative condition number (Agarwal et al., 202T) associated with the ground truth representation. Thus, our offline RL result generalizes prior offline RL works on linear MDPs JJin et al.] 2020b; Zhang et al. (02lb) 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 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 representa- tion learning, exploration, and exploitation together, and significantly improves the sample complexity of the prior state-of-art algorithm FLAMBE;
'The oracle generally refers to supervised learning style empirical risk minimization oracle. We seek to design an algorithu that runs in polynomial time with each oracle call counting as O(1). The reduction to supervised learning has lead to many successful provable and practical algorithms in contextual bandit (Agarwal al. 2014 Dudik et al. 2017 Poster & Rakhlin 2020) and RI. (Du et al. 2019a] Misra et al. 2020).
Our algorithm and analysis can be easily extended to finite horizon non-stationary setting. We choose the discounted infinite horizon setting to contrast results to FLAMME: FLAMBE is not capahle of learning stationary policies under the discounted infinite horizon setting.
2

---

Published as a conference paper at ICLR 2022
Methods
Setting
Sample Complexity
Computation
OLIVE
Tiang et al.
2017
Low Bellman rank
d'
Inefficient
(-(1
Witness rank
Sun et al.
2019
Low Witness rank
dA
Inefficient
<²(1->)
BLin-UCB Du al.
2021
Bilinear Class
d'
Inefficient
Moffle (Modi et al.
2021
Low-nonnegaaive-rank MDP
A
Oracle efficient
FLAMBE Agarwal al.
20206
Low-rank MDP
d'
Oracle efficient
[20(1 7)
REP-UCB (Ours)
Low-rank MDP
d' A'
Oracle efficient
211-7>
Table 1: Comparison among different provable representation learning algorithms in online RL. Algorithms such as OLIVE, Witness rank, and BLin-UCB work for settings which are more gen- eral than low-rank MDPs and have tight sample complexity. However, these algorithms are version space algorithms and thus are not computationally efficient. Moffle is an oracle-efficien algorithm (with a much stronger oracle than the
in FLAMBE and ours), but the assumptions under which Moffle operates essentially imply that the MDP's transition has low non-negative matrix rank (nnr). Note that a nnr least large as and could be exponentially larger than the rank (Agarwal al., 2020b). Finally, FLAMBE operates under the same function approximation setting and the computa- tion oracle as ours. Our algorithm significantly improves the sample complexity from FLAMBE in all parameters. Note the horizon dependence is not exactly comparable as these prior works originally considered the finite horizon setting with nonstationary transition, and we convert their results the discounted setting by simply replacing the finite horizon H by ((1((1 ()--))
2. In the offline learning setting, we propose natural concentrability coefficient (i.c., 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 ones) under such a partial coverage condition.
2 RELATED WORK
Online Setting We list the comparison as follows, which is summarized in Table 1 Additional related works are discussed in Section
FL.AMBE (Agarwal et al. 2020b) was a state-of-the-art oracle-efficient algorithm for low-rank MDPs. In all parameters, the statistical complexity much worse than RCP Our algorithm and FLAMBE operate under the same computation oracle. FLAMBE does not balance exploration and exploitation, and uses explore-then-commee MDPs Brafman & Tennenholtz 2002 which results in its worse sample complexity.
style techniques (i.e., constructions of absorbing
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 there matrices whose nnr is exponentially larger than the rank (Agarwal et al. 2020b).
The sample complexity of Moffle can scale | /(2 (1-7)"1), where 77 is the reachability probability, and 1/1 could be as large as 1/2 (Proposition 4 in Agarwal et al. (020b)). which essentially means that Moffle has polynomial dependence on the nnr.
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 (c.g., O(d But these algorithms are computationally inefficient as they are version space algorithms. Dann et al. (2021) shows that with a policy class, solving low-rank MDP can take 12(2)) 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 not in contradiction to the result from Dann al. (202T)
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. (019a)) strong Bellman closure (Foster et al. 2020 and strong unsupervised learning oracles (Feng et al. 2020).
Offline Setting
We discuss related works in offline RL. Additional related works are discussed in
Section A
3

---

Published as a conference paper at ICLR 2022
Uehara & Sun (2021) obtained similar statistical results for offline RL low-rank MDPs. Though the sample complexity in their algorithm is slightly tighter, our algorithm is oracle-efficient, while the CPPO algorithm from [ehara & Sun (2021) is version space algorithm.
Xie et al. (00T)) propose a (general) pessimistic model-free algorithm in the offline setting. We can also apply their algorithm 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 num- ber only using the true representation, and whether they can compete with history-dependent poli-
cies. Thus, our result is still considered superior on low-rank MDPs. The detail is given in Section E
3 PRELIMINARIES
We consider an episodic discounted infinite horizon Markov Decision Process M (S, P.r. do) specified by a state space S. a discrete action space A. a transition model P P:S S A + A(S), reward function S A + R. discount factor 7 [0,1), 1), and an initial distribution da E A(S). To simplify the presentation, we assume r(s, a) and dg are known (e.g., when do is a probability mass only so, agent always starts from fixed initial state sof Following prior work (Jiang et al. [2117; Sun et al. 2019), we assume trajectory reward is normal- ized, i.e., for any trajectory {sh, an oc h= we have 2h= ("(( 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. For example, refer to Zhang et al. (2022).
We use the following notation. Given a policy x : S + A(A), we define the value function (s) E >" an) = s, P. to represent the expected total discounted reward of under P starting s. Similarly, we define the state-action Q function Qp(ss a) == r(s, a) + (E.--(()s.))) V (s'). The expected total discounted reward of policy π under transition and reward r is denoted as VF, == Eegndo V7 We define the discounted state-action occupancy distribution (s, = (1 - Li d'dp,,(s, a), where (s, a) is the probability of i visiting (s. at time step under π and P. We slightly abuse the notation, and denote (s) as the state visitation, which equal to Eaed dip(s,a) a). When P the true transition model Pt, we drop the subscript and simply use (-). Unless otherwise noted, 1I denotes the class of all polices {S + (A)). We denote total variation distance of P1 and P2 by || - Palli. Finally, given a vector define |a][B
(sn,,,
Va Ba. Co, C1.
universal constants.
We study low-rank MDPs defined as follows (Jiang et al. 2017: Agarwal et al. 2020b). The condi- tions on the upper bounds of the norm of o*, n=¹ are just for normalization.
Definition 1 (Low-rank MDP). A transition model P S x A + A(A) admits low rank decomposition with rank d E N if there exists two embedding functions 8² μ³ such that
Ys,8 s' ES P
= (s') (s,0))
any where ( (s, a)ll < ffor all (s, o.) and any function g S + [0. 4. J / (s))(s))(())) s Vd.
An MDP is a low rank MDP if p* admits such a low rank decomposition.
Low-rank MDPs capture the latent variable model (Agarwal et al. 2020b) where (s a) is a dis- tribution over discrete latent state space 2. 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 o² 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.
Episodic Online learning In online learning, our overall goal is to learn a stationary policy क so that it maximizes ViA. .r. where Pa 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 do (c.g., to emphasize the challenge from exploration, we can consider the special case where we can only reset to a fixed sp). In the episodic setting, given a policy म, sampling a state from the
Extension to the unknown case is straightforward. Recall the major challenging of RI. is due to the un- known transition model.
4

---

Published as a conference paper at ICLR 2022
state visitation di is done by the following roll-in procedure: starting at So dg, at every time step t, we terminate and return 8; with probability 1 y, and otherwise we execute at ~ т T((s) and move to + 1, i... $t+1 P( a.). Such sampling procedure is widely used in the policy gradient and policy optimization literature (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(0, at1), r()) s mihn I= ~ p(s, a)(r = r(s, ((*s s,a).
For simplicity, we assume d where To E + A(A)] is fixed behavior policy. We denote Ep[s(s, a, [)] = 1//n-11/..n/nnnn "(ea,""))) f(s a, s'). To succeed in offline RL, we in general need some coverage property of p. One common assumption is that p globally covers every possible policies' state-action distribution, i.e., max d'p. (s, (5,0a/ppe a) < xo (Antos et al. 2008). In this work, we relax such a global coverage assumption and work under the partial coverage condition where p 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 we precisely define the partial coverage condition using the concept of the relative condition number.
Function approximation setup and computational oracles use function classes to capture them. Our function approximation and computational oracles are ex- actly the same as the ones used in FLAMBE. For completeness, we state the function approximation and computational oracles below.
Since μ² and o* are unknown, we
Assumption 2. We have a model class M ( :
E V, & E }), where μ² E V, d' E 6.
Following the norm bounds on μe", 6* we similarly assume that the same norm bounds hold for our function approximator, i.e., for any μ E V E Q. (0(s, S 1, v(s,a) and If (s)))())d(=)Ilz S Vd, S - [0, 1]. and I (s"))(((s
v(s.a).
As for computational oracles, we use a supervised learning style MILE oracle.
Definition 3 (Maximum Likelihood Oracle (MLE)). Consider the model class M and a dataset D in the form of (s, a, the MLE oracle returns the maixmum likelihood estimator P == (β, 6) arg (.) ED In(u(s') á(s, a)).
We also invoke planning procedure for known linear MDPs with potentially nonlinear rewards, which can be done in polynomial time (we know that online learning linear MDPs can be done statistically and computationally efficient). Given reward and model P == (μ, Φ) with P(s'(s, a) = (us')" o(s, a) (i.e., a known linear transition with a known feature Φ). we can compute the optimal policy arg max, V by standard least square value iteration which uses linear regres- sion. A planning procedure for known linear MDP also used in FLAMBE, see Section 5.1 in Agarwal al. (2020b) how to implement this procedure with polynomial computation complexity.
REPRESENTATION LEARNING IN ONLINE SETTING
We consider the online episodic learning setting where the agent can only reset based on the initial state distribution dg To find near-optimal policy for low-rank MDP efficiently, need to carefully interleave representation learning, exploration, and exploitation.
4.1 ALGORITHM
We present our algorithm in the online setting in Algorithm 1 We first describe the data collection process. Every iteration, Algorithm 1 rollouts its current policy æ to collect a tuple (s, a, 8 a', where dz. a U(A, 8" ~ (**+ss, a). a' ~ (AA, uniform distribution over actions (note that we take two uniform actions here). Recall that to sample dp., we start at so ~ dg. at every time step t, terminate and return s, with probability 1-% and otherwise we execute 0r ~ π (8t) and move to + 1, i.e., 8t-1 ~ P*(-|se, at). Thus collecting one tuple requires exactly one roll-in, i.e., trajectory. We verify that with high probability the roll-in terminates with Õ((1 -) steps which is often called the effective horizon.
P*(. 8, a) where U (AA) is
After collecting new data and concatenating it with the existing data, we perform representation learning, i.e, learning factorization and a representation by MIE (line 6). set the bonus based on
5

---

Published as a conference paper at ICLR 2022
Algorithm UCB-driven representation learning, exploration, and exploitation (REP-UCB)
1: Input: Regularizer An. parameter Ctp, Models M {(u, : v E ) Iteration
2: Initialize To (-1 to be uniform; set Do 0, Di = 0
3: for episode n 1. 1...................... ,N do
4: Collect a tuple (s, a, s', a,', 3)) with
dp. a, U(A), s' ~ PM ('(,,,)),, ( P'(1s',a')
5: Update datasets by adding triples (s, 5') and ()):
Dn Dn-1+ ((,,,,))))
D2 D {(5' d', 8)}
6: Learn representation via ERM (i.e., MLE):
:= (Ans dn) arg (M ED,+D, [In (s))
7: Update empirical covariance matrix Er, Ex.ED, (1(8, + An
8: Set the exploration bonus:
(s == min (n)
(1)
9: Update policy nn = arg max, V Pa.rth.
10: end for Return T1. AN
the learned feature (Eq, D.. 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 respect to the learned feature o. and planning in a known linear MDP is known as computationally efficient (Jin et al. 2020a) (see the explanation after Definition as well).
Computation of the MLE oracle The MLE oracle in general could be non-convex optimization procedure when & and j are general nonlinear function approximators. However, this is a standard supervised learning ERM oracle and one can easily optimize via stochastic gradient descent style approaches if μu and á are differentiable. For special cases where the MDP is a tabular MDP, the MLE objective is convex and the optimal solution has closed-form For linear MDPs YYang & Wang 2020) where P* P*(s'(s.. |s. a) (* (s'))" M* (s, a) with known f and zi* but unknown M the MLE objective again is convex with respect to parameter M*. Thus when specializing to specific settings such as tabular MDPs and linear MDPs where computationally efficient approaches exist, Rep-UCB also provably computationally efficient. In contrast, more general approaches such as Olive (Jiang et al. 2017) are provably computationally inefficient even when specialized to tabular MDPs. We note that our setting does not directly capture the linear MDP setting from Jin et al. (2020a) since there μ² might not be captured by model class V that has bounded statistical complexity.
4.2 ANALYSIS
Theorem 4 (PAC Bound for REP-UCB), Fix 5 (0, (),1),c (0, 1). Let â be uniform mixture of T1. TN and πt" := arg max, Vp.,, as the optimal policy. By setting parameters as follows:
On O (V(IA) d²) In((M|n/8)) A, = (ddm((M))))))
with probability at least 1 - 8, have
Vp Vpr,, Se, where the number of collected samples is at most
(1 (1-7)662
where v only contains log terms and the dependence MM] is at most In(In((M))).
The theorem shows that REP-UCB learns in low-rank MDPs in a statistically efficient and oracle- efficient manner. To the best of our knowledge, this algorithm has the best sample complexity among

---

Published as a conference paper at ICLR 2022
all oracle efficient algorithms for low-rank MDPs. Extension to continuous function class v and + using statistical complexities such as covering dimension is possible since our analysis only uses standard uniform convergence analysis on y/ and Φ
Highlight of the analysis Below we highlight our key lemmas and proof techniques.
First, why is learning in a low-rank MDP harder than learning in models with linear structures?
Unlike standard linear models such linear MDPs (Yang & Wang 2020; [Jin et al. 2020a).
KNRs (Kakade et al. 2020, Abbasi-Yadkoni Yadkori & Szepesvári, 2011 Mania et al.
2020: Song &
Sun 2021). and GP kernel models (Chowdhury & Gopalan, 2019; Curi et al. 2020),
not 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: V's a (-|(,ss, a), (- a)) I a(s,, a) where (s,, a) is the uncertainty measure, and & is some distance metric (e.g.. & norm). With proper scaling, the uncertainty measure a a) is then used for the bonus. For example, in linear MDPs (i.e., low-rank MDP with known feature Φ*), given a, ,'), we can learn a non-parametric model ((s'(s,a) p(s') T (8, (8,a). and get pose uncertainty quantification:
dataset D = {s,a,**
M(s.a).1
(s)) S clØ (s' (s, - / (()) ( (s.
for some family of functions f S with Eg. - (s, as)" + (Lykouris al. 2021 Neu & Pike-Burke 2020). To the scaling properly, since o- is known a proo,, the lnnea regression analysis applies here, and one can apply Cauchy-Schwar inequality to the LHS of pull out 6² and get an upper bound the form of
(s, a I [(MMM(s))))))))))){))))) p" (s))s
(a)
(b)
where e set to the linear regression training error measured in the term (b) above.
However, when we jointly learn μ and o, since nonlinear function approximation is used, we can- not get point-wise uncertainty quantification via linear regression-based analysis. We stress that our bonus is not designed to capture the uncertainty quantification
the model error between P(-(s, a) = ja é(s, a) and (* a) = μ² (s a) in point- 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 8.
Vn [L, NI, II, VI
Pu.r+b. -p -C1
[A 18
We remark that the idea of optimism with respect to the initial state distribution has been used in prior works (Jiang et al.) 2017 (un 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 We start by using the simulation lemma (Lemma 20) inside the learned model which is important since our bonus bn uses on associated with the learned model P.::
& - VT... (1 -))-7)--1 2 Lag
((n - - P- ( (- a))l,,
from where we show that Bx,ands |P - P*((8,) a)ll as a whole nearly lower bounds the Rs
average bonus E. Esg [bm (s,,))) Thus the proof of optimism is fundamentally different from the proofs in tabular and linear MDPs which are done via induction in wisnt-wise manner. The detailed procedure is illustrated in Lemma in Appendix B

---

Published as a conference paper at ICLR 2022
Second, our bonus is using representation on that is being updated every episode, and our empirical covariance matrix E., is also updated whenever update on. 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 keep tracking potential function that is defined using the unknown ground truth representation o* i.e. the elliptical potential (s, a)ls
where
a) Pm = nalE (s.a)<pp ( (8 + A,
and a)) Σi=o dp. (s.a)/n. a)/n. Since this potential function uses the fixed representation o², 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 á,, to the potential defined with respect to the true feature 6²
function (1, a)ll -
Linking bonus under on to the elliptical potential function under o' With near optimism, using the simulation lemma (Lemma 20) inside the real model, we can upper bound the per-iteration regret as follows:
P-rr Vp*,r E(o
(p (s.
A
where Sn = (1/n),, and fn(s, a) == [ s. P. a)ll1. To connect the first term in the right-hand side of the above inequality to the elliptical potential under the fixed feature o², we show that for any function E S x A + [0, B) for B E R
((s,a))gg: (s(s a)) s (1 - 7) (L (d (s,(,,))lss a. /y/A/EE, (s, + +2 >An
V (1- (s,2 (s)
where (s) 1/n Ei=o d (s)u(a) and u(a) 1/| See Lemma 12 in Appendix B By
substituting g with bn + fa/(1 - 7). the first term of the RHS of the above inequality can be upper bounded as:
2(1 - 2(1-7)-² B Ess
a)) on.or
(2(
ig (s,() a))
And.
[1- -
(G1)
(9x)
In the term (G2), we expect n.,, (s, a)] to be O(1) as E,, (f, a)] is in order of 1/n due to the fact that it is the generalization bound of the MLE estimator P. which trained on the data drawn from A For n.E,p,, [6%, (s, a)), we expect to be the order of d as the (unnormalized) data covariance matrix Σn in the bonus 6,, uses training data from P/. i.e., we are measuring the expected bonus under the training distribution. In other words, the term (G2) scales in order of (oly(d). For the term (G1), since it contains the potential function based on d², the sum of the term (91) over all episodes can be controlled by the standard elliptical potential argument (see Lemma 18 and Lemma 19I. This concludes the proof sketch of our main theorem.
In summary, our analysis relies on the standard idea of optimism the face of uncertainty, but with novel techniques to achieve optimism under nonlinear function approximation with the MLE supervised learning style generalization bound, and to track regret under changing representations.
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 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 representation J. The second part is the construction of penalty term b. Using the learned representation g. we standard bonus in linear bandits the penalty term as if ò were the true feature. The third part is planning with the learned model P and reward r b.
8

---

Published as a conference paper at ICLR 2022
Algorithm 2 LCB-driven Representation Learning in offline RL (REP-LCB)
1: Input: Regularizer A, Parameter a, Model classes M = { 6: E V. E 2). Dataset D.
2: Learn model P by MLE: P = f ò arg maxpeM Epln s,a]].
3: Set the empirical covariance matrix Σ E(s.aED (s. AI.
((s,
4: Set the reward penalty:
b(s,a) min (
5: Solve स = arg max, V P
We present the PAC guarantee of REP-ICB. Before proceeding, define relative condition number as a mean to measure the deviation between a comparator policy π and the offline data:
Bap. (8' (8))) (s, as,a)]=
C# = sup ZERd x PT (s, a)ds T (s, a))a
In the special case where the MDP is just a tabular MDP (i.e., of is one-hot encoding vector), this reduced to a density ratio C maX a d'p. (s, p(s, a). The relative condition number C is always 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,
ready to state the main theorem for REP-LCB.
Theorem 6 (PAC Bound for REP-LCB), Let w maxa, (1/msla) s))- Denote 7 as the output of REP-L.B. There exists set of parameters such that with probability least 5, for any policy (including
non-Markovian policies),
MW (1-7)'n
Vi..- -V SOV
See Theorem 14 in Appendix for the detailed parameters. We explain several implications. First of all, this theorem shows that we can uniformly compete with any policy including distory-depender non-Markovian policies "a satisfying the partial coverage C < co. Particularly, if the optimal policy r is covered by the offfine data, i.c., C < cc, then our algorithm is able to compete against Note that assuming offline data covers क still a weaker assumption than the global coverage such as sup, (s,a) d'p. (s, a)/p(s, a) in prior offline RL works (Antos et al. 2008 Chen & Jiang, 2019). Second, our coverage condition is measured by relative condition number defined using the unknown ground truth representation di but not depending other features. Prior works that
relative condition numbers as measures of coverage are restricted to the settings where the ground truth representation o* known (Jin et al., 2020b Chang al. 2021 Zanette et al.) 2021b).
To sum up, our algorithm is the first oracle efficient algorithm which does not need to know 6", and requires partial coverage only in terms of ø" Note while [ehara & Sun (202T) has a similar guarantee on low-rank MDPs, their algorithm is not oracle-efficient as it a version space algorithm.
CONCLUSION
We study online/offline RL on low-rank MDPs, where the ground truth feature is not known a pri- ori. 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 lcarning works on low-rank MDPs forward step-by-step reward-free exploration framework, our algo-
and block MDPs that rely rithm interleaves representation learning, exploration, and exploitation together, and learns a single stationary policy. For offline RL, our 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.
*Given a = (T, 1 where n; depends on so, Gg 84, Vp-,, and di. (s, a) are still well-defined. 'We also require co, which is mild assumption since does not involve Pt. Indeed, it is much weaker than the global coverage type assumption 1/p(s, a) oc, V(s
9

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
The authors would like thank Alekh Agarwal, Praneeth Netrapalli and Ming Yin for valuable feedback.
Masatoshi Uchara partially supported by Masason foundation.
REFERENCES
Yasin Abbasi- Yadkori and Csaba Szepesvári. Regret bounds for the adaptive control of linear quadratic systems. In Proceedings of the 24th Annual Conference Learning Theory, pp. 1- 26. JMLR Workshop and Conference Proceedings, 2011.
Alekh Agarwal, Daniel Hsu, Satyen Kale, John Langford, Lihong Li, and Robert Schapire. Taming the monster: A fast and simple algorithm for contextual bandits. International Conference on Machine Learning, pp. 1638-1646. PMLR, 2014.
Alekh Agarwal, Mikael Henaff, Sham Kakade, and Wen Sun. Po-pg: Policy cover directed explo- ration for provable policy gradient learning. NeurlPS, 2020a.
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complex- ity and representation learning of low rank mdps. Advances in Neural Information Processing Systems, volume 33, pp. 20095-20107, 2020b.
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 12(98))1-76, 2021.
András Antos, Csaba Szepesvári, and Rémi Munos. Learning near-optimal policies with bellman- residual minimization based fitted policy iteration and single sample path. Machine Learning, 71:89- 129, 2008.
Ronen Brafman and Moshe Tennenholtz. R-max-a general polynomial time algorithm for near- optimal reinforcement learning. Journal of Machine Learning Research, 3(Oct):213-25) 2--11,, 2002.
Jonathan D Chang, Masatoshi Uchara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift in imitation learning via offline data without great coverage. 2021.
Jinglin Chen and Nan Jiang. Cooommation-thereeee considerations in batch reinforcement learning. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pp. 1042- 1051, 2019.
Sayak Ray Chowdhury and Aditya Gopalan. Online learning in kernelized markov decision pro- cesses. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 3197- 3205. PMLR, 2019.
Sebastian Curi, Felix Berkenkamp, and Andreas Krause. Efficient model learn- ing through optimistic policy search and planning. arXiv preprint urXiv:2006.08633 2020.
Christoph Dann, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert Schapire. On oracle-efficient pac rl with rich observations. arXiv preprint arXiv:1803.0006 2018.
Christoph Dann, Yishay Mansour, Mchryar Mohri, Ayush Sekhari, and Karthik Sridharan. Ag- nostic reinforcement learning with low-rank mdps and rich observations. arXiv preprint arXiv:2106.11579, 2021.
Simon Du Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik. and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning, pp. 1665-1674. PMLR, 2019a.
Simon Du, Sham M Kakade, Ruosong Wang, and Lin F Yang. Is a good representation sufficient for sample efficient reinforcement learning? arXiv preprint 1xxvv:1//0.03000 2019b.
10

---

Published as a conference paper at ICLR 2022
Simon S Du, Sham M Kakade, Jason D Lee, Shachar Lovett, Gaurav Mahajan. Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in rl. JCML, 2021.
Miroslav Dudik, Nika Haghtalab, Haipeng Luo, Robert E Schapire, Vasilis Syrgkanis, and Jen- nifer Wortman Vaughan. Oracle-efficient online learning and auction design. In 2017 ieee 58th amnual symposium on foundations of computer science (focs), pp. 528 539. IEEE, 2017.
Fei Feng, Ruosong Wang. Wotao Yin, Simon S Du, and Lin F Yang. Provably efficient exploration for reinforcement learning using unsupervised learning. arXiv preprint arXiv:2003. 2020.
Dylan Foster and Alexander Rakhlin. Beyond ucb: Optimal and efficient contextual bandits with regression oracles. In International Conference on Machine Learning, pp. 3199-3210. PMLR, 2020.
Dylan Foster, Alexander Rakhlin, David Simchi-Levi, and Yunzong Xu. Instance-dependent complexity of contextual bandits and reinforcement learning: A disagreement baase perspective. arXiv preprint urXiv:2010.03104, 2020.
Botao Hao, Yaqi Duan, Tor Lattimore, Csaba Szepesvári, and Mengdi Wang. Sparse feature selec- tion makes batch reinforcement learning more sample efficient. In International Conference on Machine Learning. pp. 4063-407, PMLR, 2021a.
Botao Hao, Tor Lattimore, Csaba Szepesvári, and Mengdi Wang. Online sparse reinforcement learn- ing. In International Conference on Artificial Intelligence and Statistics, pp. 316-324. PMLR, 2021b.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Con- textual decision processes with low bellman rank are pac-learnable. In International Conference on Machine Learning. 2017.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pp. 2137-2143. PMLR, 2020a.
Ying Jin, Zhuoran Yang. and Zhaoran Wang. Is pessimism provably efficient for offline rl? arXiv preprint ar Xiv:2201. 150855 2020b.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In In Proc. 19th International Conference on Machine Learning. Citeseer, 2002.
Sham Kakade, Akshay Krishnamurthy, Kendall Lowrey Motoya Ohnishi, and Wen Sun. Infor- mation theoretic regret bounds for online nonlinear control. In Advances in Neural Information Processing Systems, volume 33, pp. 15312-15322 2020.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In Advances in Neural Information Processing Systems, volume 33, pp. 21810-21823, Curran Associates, Inc., 2020.
Yehuda Koren, Robert Bell, and Chris Volinsky. Matrix factorization techniques for recommender systems. Computer, 42(8):30-37, 2009.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint arXiv:2006. 04779, 2020.
Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch off-policy reinforcement learning without great exploration. In Advances in Neural Information Processing Systems, volume 33, pp. 1264-1274, 2020.
Rui Lu, Gao Huang, and Simon S Du. On the power of multitask representation learning in linear mdp arXiv preprint arXiv:2106.08055 2021.
Thodoris Lykouris, Max Simchowitz, Alex Slivkins, and Wen Sun. Coruption-robust exploration in episodic reinforcement learning. In Conference on Learning Theory, pp. 3242-3245. PMLR, 2021.
11

---

Published as a conference paper at ICLR 2022
Horia Mania, Michael Jordan, and Benjamin Recht. Active learning for nonlinear system identifi- cation with guarantees. arXiv preprint arXiv:2006.1027) 10277, 2020.
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state ab- straction and provably efficient rich-observation reinforcement learning. In International confer- ence on machine learning, pp. 6961-6971. PMLR, 2020.
Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. Model-free representation learning and exploration in low-rank mdps. arXiv preprint arXiv:2102.07035, 2021.
Gergely Neu and Ciara BurkeB unifying view of optimism in episodic reinforcement learning. arXiv preprini arXiv:2007.0189), 2020.
Chengzhuo Ni, Anru Zhang, Yaqi Duan, and Mengdi Wang. Learning good state and action repre- sentations via tensor decomposition. arXiv preprint arXiv:2105.01 2021.
Matteo Papini, Andrea Tirinzoni, Marcello Restelli, Alessandro Lazaric, and Matteo Pirotta. Lever- aging good representations in linear contextual bandits. arXiv preprint arXiv:2104.0378) 2021.
Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline rein- forcement learning and imitation learning: tale of pessimism. arXiv preprint 2rXiv:2103 2021.
Devavrat Shah, Dogyoon Song, Zhi Xu, and Yuzhe Yang. Sample efficient reinforcement learning via low-rank matrix estimation. Advances in Neural Information Processing Systems, 33:12092- 12103, 2020.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanetot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play, Science, 362( 1144, 2018
Yuda Song and Wen Sun. Pc-mip: Model-based reinforcement learning with policy cover guided exploration. International Conference on Machine Learning, pp. 9801-9811. PMLR, 2021.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. arXiv preprint arXiv:2004.04136, 2020.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference Machine Learning, pp. 9870-9879. PMLR, 2021.
Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference on learning theory, pp. 2898 2933. PMLR, 2019.
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline rl: Pac bounds and posterior sampling under partial coverage. arXiv preprint arXiv:2107.06226, 2021.
Ruosong Wang, Dean P Foster, and Sham M Kakade. What are the statistical limits of offline rl with linear function approximation? arXiv preprint #:Xiv:2010.1189) 2020.
Gellért Weisz, Philip Amortila, and Csaba Szepesvári. Exponential lower bounds for planning in mdps with linearly-rallizable optimal action-value functions. In Algorithmic Learning Theory, pp. 1237-1264. PMLR, 2021.
Tengyang Xie, Ching-An An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistent pessimism for offline reinforcement learning, arXiv preprint arXiv:2106.069: 06926, 2021.
Lin Yang and Mengdi Wang. Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound. In Proceedings of the 37th International Conference on Machine Learning, pp- 10746-10756, 2020.
12

---

Published as a conference paper at ICLR 2022
Mengjiao Yang and Ofir Nachum. Representation matters: Offline pretraining for sequential deci- sion making, arXiv preprint arXiv:2102.05815, 2021.
Ming Yin, Yu Bai, and Yu-Xiang Wang. Near-optimal offline reinforcement learning via double variance reduction. arXiv preprint arXiv:2102.01748, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. In Advances in Neural Information Processing Systems, volume 33, pp. 14129 14142, 2020.
Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Learning near op- timal policies with low inherent bellman error. In International Conference Machine Learning, pp. 10978-10989. PMLR, 2020.
Andrea Zanette, Ching-An Cheng, and Alekh Agarwal. Cautiously optimistic policy optimization and exploration with linear function approximation. arXiv preprint 22tI::2003.12223, 202
Andrea Zanette, Martin Wainwright, and Emma Brunskill. Provable benefits of actor-critic meth- ods for offline reinforcement learning. arXiv preprint arXiv:2108.0888 2021b.
Weitong Zhang, Jiafan He, Dongruo Zhou, Amy Zhang, and Quanquan Gu. Provably efficient
representation learning in low-rank markov decision processes. arXiv preprint arXiv:2106.11935, 2021a.
Xuezhou Zhang, Yiding Chen, Jerry Zhu, and Wen Sun. Corruption-robust offline reinforcement learning. arXiv preprint arXiv:2106.00 066300 2021b.
Xuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Wen Sun, and Alekh Agarwal. Efficient reinforcement learning in block mdps: A model-free representation learning approach. arXiv preprint arXiv:2202.00066 2022.
13

---

Published as a conference paper at ICLR 2022
A
MORE RELATED WORK
Here, we mention several additional related works.
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 given exploratory distri- bution. Zhang et al. (2021a) consider how to choose the best representation among correct repre- sentations inspired by Papini et (202T) (i.e., the MDP is linear MDP under any representation in the function class $). Thus, it still falls into the linear function approximation setting, In contrast,
we only assume the MDP linear under some unknown o² E
Offiine setting In addition the work we mentioned, the pessimistic approach in offline RL has been extensively investigated. Empirically, it can work on simulation control tasks (Kidambi al., 2020 Yu et al., [2000] Kumar al., 2020, [Lu et al. 2020 Channg al., 2021). On the theo- retical side, pessimism allows us to obtain the PAC guarantee on various models when comparator policy is covered by offline data in some forms (Jin et al. 2000h) Rashidinejad et al. (02T]] Yin al.]
2021 Zanette et al., 2021b Zhang al. Chang et al. (022)). However, these algorithms and their analysis rely on a known representation and linear function approximation.
We mention works that tackle representation learning from different viewpoints and settings. Lu et al. (202T) consider multitask representation learning under generative model assumption. Hao al. aa2 study the feature selection problem in sparse linear MDPs and NN 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. [Shah al. (2020) studies learning the assumption that the optimal Q-function admits a low-rank structure on the generative model setting.
B PROOF OF THE THEORETICAL PROPERTY OF REP-UCB
We summarize the notations frrquently use. First of all, hereafter, assume Cp, C1, ... are some universal constants, and the notation
Notation
f(1((1 - 7), [AI, (m(1/)), In(IM)) d, n) S g(1(((1111_____1 (111(111)),) (n d)
means there exists some constant d₁ > 0, such that
f(1((1 - m), (A) In(10),, d, S cig(1((1 -2), -),, (1100, InIM)) for any 0 < < 1, |4, (1/5),
We define
(n(s) = ー (s).
With slight abuse of notation, we overload the above notation and use Pn for 1/n EnD1 (s) Next, define P/2 [S R] as marginal distribution of for a triple
(s.0,s)) ~ ((() ss)
We define three matrices as follows:
Enn xULA), = nB -U (s) AnI,
Σ Epn.b = nB nB (0(s, (s, + dal,
E nB nB (00 AnI.
Note that for fixed Φ, E is unbiased estimate of E (UA(()))
14

---

Published as a conference paper at ICLR 2022
Optimism
First, we prove the optimism at the initial distribution. This is proved by using sim- ulation 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 Br,and (s, a) is in the same order of the expected model enror E, aaddy [P,((I,, (* Note that the expectation is with respect to d%
Lemma (Almost Optimism the Initial Distribution). Consider an episode Ti S n2 N) and set
Ox = + Xn = 0 (dd =0 [n[[M/n/6)) (1m
With probability 1 8, we have
Vn E [1, (,,... NJ, VT IT, V Pairth, VB,,, (1A
Proof. In this proof, letting fn(s, a) = P: 18,0) -P a)]l we condition on the event
Vn, E a)) $ Sni Es-ph Eu s Sn. Vn, Vo, (l((o, a)ls & = O(]lo(s,
{(]lo(s,a)ll=]) aiلاs= x
From Lemma 10 and Lemma 17
this event happens with probability - 5. Then, for any policy ",
from simulation lemma 20
(1 (VF
rbn vr
-B(,,))dz (1, a) + IB P.. (VE.,('') -P [v']
B(saa)dd
min (anlónn as)
2) + + (5') 7.B [V (s])
Pa (u
(3)
where in the last step, we replaced the empirical covariance by the population covariance. Note the notation $ is to universal constants. Here, since [IV < (since we assume trajectory-wiss total reward is normalized between [0, 1). we have:
(saa)rd Pr.
{B
- (ewP))) [VP,, [VF,, < Blo.))o (s((s,a)}
The above is further bounded by Lemma 11
(aaaa)y (f)(s,a)} R(z.))ndz 11 a) ---1 x.
vTv (nlA B Be anU(A) [))(s 4ha 4nC,
+ (1- // (fR(s,)))
Then,
B(s.a)ayz (fn(s,a)3 S Va,,aas (5) (nxU(A))
(1 (4)
where
= + And + nsat S (1A) + dd) In([M]n/8).
Note we here use fr a) < E. o. & Sn and E ( Sn.
Combining all things together,
- Eg(xP-(a.) (s')
28 20 (fn(s,))}
I Van (5, dnxula).nn
+ - [1]A)AAAGGn
I anb (3, puxu(a)dnn
+ - (1-
where Or := Va (5)
15

---

Published as a conference paper at ICLR 2022
Going back to (3) have (1 (VP
Vin.r)
P.. mhrthhn
((,aa)d
min
(s,
Xu .2) + si -yBB
[V7
[6
min
,2 -min (allên
(11-2)
--((1-7)AA))
From the second line to the third line, we again use |VVp, O(1) and (4) This concludes the proof.
Next, we obtain the upper bound of Σ n= V p Recall π is the optimal policy. Though this form is the same as standard regret form, since we are not exactly deploying Tn episode (recall that we play a uniform action at the end of the episode), we cannot get the regret guarantec. However, suffices for the PAC guarantee.
Lemma 8 (Regret). With probability 1 - 5, we have
N
N
|Ald
Σ
V
Vp S Nln In
1+
In(N[M)/)) (1 (1-
A
Proof. Similar to Lemma event
letting (n(8, = PP s,a) - P*(- s. condition the
8,a)ll1, Sn. Vé, 10ls, "'="'''''."" (()o(ss.a)llss xux (6)
Vn,
E Eepp ULt)
From Lemma 10 and Lemma 17 this event happens with probability 5.
For any fixed episode n and any policy we have
VT., Vp
VE" Purtth Vpi,, + V|AIsn(11 - -1
Lemma 7
VTa
Puotthn -V + (1-1)
(T. = arg max P,,,+bb,
-1--)) " (s, a) use the 2nd form simulation Lemma 20 20 the last display.
TM 7IE (s))) Pw.rth, (s')] 7 (s'),,)) (V Pa. (8))l + (1 71
Then, noting 116,11oo I 2, we have I P.. 1l S 2/((1- - 7). Combining this fact with the above expansion, we have (VP",r Vprr.))
2 (I (1-
< (1-)))1 (laa)aaaa [n(ss
E +
(7)
First, we calculate the first term (a) in Inequality ((),, we have
Following Lemma 12 and noting the bonus i,
E B(s,a)adfa [in(s,a)]
KMCAAi 2)
S B(sa))aan min
(rrom (6)
$ 1( in ny
[1)(
+ drdn
enk (
+ BlanBBorph (1
2 (1-7).
16

---

Published as a conference paper at ICLR 2022
Note that we use the fact that B 2 when applying Lemma 12
In addition, we have
n., ((
[| a))
n Tr(E, p. xU (a) HnL xe + And3-1) S d.
Then,
(a (bn(sra) s 18 (3
on VidlAla + yddn + vdlAla - 7)/n.
Second, we calculate the term (b) in inequality Following Lemma 12 and noting fh a) is
upper-bounded by (i.e., B in Lemma 12). we have
(( (f (s)
Ba l (5,0))ls-1 on.0 V (mm (RR [P a)l} 4yAnn
B (f) -))
(( 10 (5, on.o.
vn|AlTGn 4yAnd + VIAlG((1-))
(( (10 (3,
x-1 o
+ (1
where in the second inequality, V7Vnl G, + A,d + nCC, S On-
E (
Sa. and in the last line, recall
Then, by combining the above calculation of the term (a) and term (b) in inequality have:
dlAl22 7)
Væ VpL,, S
(1 7)) (Eisa)=azz 11 (8," a)ll= Pn.or VdlAla + + +
+ (1 - 7)2 (F15,a)-d5a (s) 0>| pn.ar on + VIAIGn(1 (1)
1
Hereafter, we take the dominating term out. First, recall
On S ((NM110))
(IAMPM(NNIMM)) MIA
Second, we also use
E (3.a)mddp 11 (110))_) (3,
<
N
E (f (8, --- di (3,
pnise
(CS inequality)
N
N dete Σ
E (3 (3,
In det()
n=1
Lemma 18 and \ ... - Ax)
dN In 1+ dA1
(Potential function bound. Lemma 19 noting (s,"(s)))le S for any (s.
17

---

Published as a conference paper at ICLR 2022
Finally,
Σ Vp..r
N
Vp
(dNVin
1+ dAi
.
)
v VdlAlo + ddN + Σ
N
d|Ala (1 7)
(1 71
n=l
N
ddNIn
1+
ON
Σ VIAlG(11 1
(1-
di
N
N
2/
dN In
1+
V d.Ala
+
dN In
1+
aN
7)
dAi
(1 - >)
(Some algebra. We take the dominating term out.)
N
A 1/2
In (MIM11))
24
dN
1+
dλp
(1-7)²
This concludes the proof.
Using Lemma immediately obtain the PAC guarantee.
Theorem 9 (PAC guarantee of REP-UCB), By interacting with the environment for number of steps at most
Mlog(N10), = dd dd (1 7)2
( L In² (1+ ClA (11
(1 7)
with probability 1 - 6, we can ensure V - V Se.
Proof. From Lemma and
Lemma 22 when N is
o d [A) In( (M(MM In? (1+ d |A In((M118)
( (1 7)2
(P[APII((MM110)) ))
(1 7)²²2 ()
with probability 1 - 8, ensure
Σ n=1
V - Vpr,. Se.
N
With probability 1 1-6 6, we need (1-7)-1 ln interactions with the environment to get one tuple (s, a, s', a', 3) from one roll-in of π. Thus, the total sample complexity is O -) In(N/6))-
Next, we provide an important lemma to ensure the concentration of the bonus term. The version for fixed @ proved in Zanette et al. (2021a Lemma 39). Here, we take union bound over the whole feature & E Φb Recall
pn(:)
drp (-).
n
Lemma 10 (Concentration of the bonus term). Set X, (nlnnnn))// f for any n. Define
Epn.o n a)d (s, Ant, En Σ ((( aa0000 (al),at)) + A.
With probubility 1 5, we have
Vn E N+, Vo Φ (e,c)1l(s,,a)11 pa x $ [|b(,,a)]ss,, 5 scall((s,a))
18

---

Published as a conference paper at ICLR 2022
For any 9 E 8 A R. The next lemma shows that ((x.a)=d7 (o(s,a)} can be upper-bounded
as long as we have the convergence guarantee for using B(B,a)d7 (3, eniin
(s.aa Lemma 11 (One-step back inequality for the learned model). Take any 8 x A + R such that |1 I B. We condition on the event where the MLE guarantee (17)
Us) [0"(s. a)] and and (A) a--t---
Ex-pn -- (A() S Sn.
holds. Then, for any policy we have
1 (sa))ad {g(s,a)}
I B(z,))mdz, 11 R)b
V [n|A [s,( + B2 + nBa
(nXU(A)ion
51 () (0 (1 (s))
Recall E,, (.XU(A)) = a0
a)] + And.
Proof. First, we have an equality:
((a.))dd (o(s,a)} 7 - (. (g(s,a)} + (1 - -7) E) ana( (0(8,
(8)
The second term in (8) upper-bounded by
do(s) 8) (n(s)((a)
(1 v
max (s.a)
E. Eappo (n) (s)(),))) (os
Pn
do(s)T(a 5)
(1 (1 /
max
(s.a) (1 - ())(()uu( Eorpn an (s)( s v (1--))|A1B >) sn (n (s,(ssa)) a)}.
Next consider the first term in (8) By CS inequality, have
(3.)) p. (B. (g(s, a)} E(s (5,,a)
/Σ An(s)a((a s)g(s,a)d(s)
/ M/minoorte (1)00))))
EEBAAEER (b,(6,
Pm en
(U)))
Then,
Σ (n(s))(aa s)d(s) on xUCA).ón
E (n(s))(al ( s))() } (nterp.oweeslendd Aal {/Ennoot (
nB (nEeppaaaeli 1/Ennta ( (3,, s)d(s,
B² B²Annd
(Use the assumption
Ea m(a ())(s,a))lo and f (s) S v for any h: S 10, 1).)
nB -nBx-pn
E ( ())))))))))))))
(o(ss + B²X,
< nE nBs=pm
EE lo(s, a)] 2| + xnd + nB²cn
(MLE guarantee)
(s, { + Bdnn + B2nGn
(Jensen)
< <n|AA (B) (o²(s, a)}} + B² d + B²nG,
(Importance sampling)
< <n/ABB BIAR (2 (z) + B2 Ad + B2nGn
(Definition of (',)
19

---

Published as a conference paper at ICLR 2022
Then, the final statement is immediately concluded.
Below, we show a similar lemma as Lemma 11 ((s.))dz, {g(s,a)) })} instead of E (s,a))dy (g(s,a) For any S x A - R, this lemma shows that ((),a) a)) can be upper-bounded using ((g,))=dz. 11 (3, ->ا== on.6"
The difference is we aim for calculating
long as we
have the convergence guarantee for Es-p,. (UU((A) [g" a)]] Note comparing to Lemma 11 this not a probabilistic statement. Note that | (s, a)
1-- is the usual elliptical potential function mnige
under the fixed representation d²
Lemma 12 (One-step back inequality for the true model ) Take any E S x SXA- - R such that |0 S B. Then,
(( {)( & ((a,a)yyz, 10 (5, x)lج=- on.o. vrv / (g² a)] + dB,
( (El-))lllss-pn.oe (s,
Recall Epn.0' = (L, (3, And.
Proof. First, have
(ss,))az,, (g(s, a)} =yE (5.ā)=d7 St {g(s,a)) + (1 -1) oo0 (g(s,a))
The second term in (9) upper-hounded by
do(s)=(a 8)
(1 Dy max
Pa(())(() E Berp [s)(s,a)] S v (nU a)) (1 7).
(s:a)
By CS inequality, the first term (9) is further bounded as follows:
((g (2.0) (g(s, a)} = E(s,a)= o (s, / Ed((s) a))(s,
//uu(s)zle s)g(s,
I LB 2..
Here, we have
/ Σ (s)a(aa s)d(ss
2n13 // s))((s a))(s)
Σ (s) s))(s)))((s)
nB(ga)mpp (/ (s,, ())(()) s)((s) a)d(s) + An
< nn/ { (B(ss))pp (s) + And
(Jensen)
Therefore,
n (( P (5.a) (m))) (s, a)] } + AndB
< n|Al| (E (3d (0 [s) + dB,
(Importance sampling)
< m|| { Berpp 10 (s)
Xnd
20

---

Published as a conference paper at ICLR 2022
In the last line, use the following inequality:
E. (()))) (s,
Pn (2.0) (s) + -1 [+))
( P [g](s
Finally, we have
((ea))ad,, a)( s ((a.a)ddz. 1 (8, و]]|ا) Pri8. vTV {n|AABB an (o² d
+ v V|llBerpo (g) 71
This concludes the proof.
C PROOF THE THEORETICAL PROPERTY OF REP-LCB
This section provides the detailed proofs for our results in the offline setting.
Below we first prove that p is an almost pessimistic estimator of Vp. ..
Lemma (Almost Pessimism at the Initial Distribution). Let w = maxa,s 1/TO (a s). Set (In(LM//8)
(
a = (w (
A(d(dn((MM))))
With probability 1 for any policy x, we have
V Vp.,, s w(1 - 7) 11)-7))-((((M666))
n
Proof. We define
Ep. (BB(s,))~ J MI, = nED[de 11 AA.
where A = O(dIn(MM/6)). In this proof, letting f(s, 1|P(- s, - P*( (,, s,a)]lı, we
condition on the events:
(f S, VΦ E + (0((, =((l), a1ls=))))
(10)
where Ç O(Ln([MM//)/m)) From the offline version of Lemma 17 and Lemma this event happens with probability 1-6.
Then, from simulation lemma (Lemma 20p.
(1 (1-7)(/2,--- V...r)
B(e.o)=d3
b(, a) + +BB -BB (8'] PP [VP.,
mi < ((saa)=d3 min (allé(s, a)))=- , + (BBx)((e)) [V'), (Bsw))(s [s')) (From (10)
Here, we have
E(s.a)aazz ((or)
Bx-P-(ssa) (s'), s E(s,a)dzz {f(s,a)).
This is further bounded by Lemma 15 noting [|V <
B(ra)aa3 (f(s.a) a)} $ Va (1((5 a)ls=- 2. VwC(1- -7).
(11)
We remove In since is fixed in the offline setting.
21

---

Published as a conference paper at ICLR 2022
where
a' = nyw + + Ad + 72ng S (w d)) In(lM)/8).
Here, we use f(s, a) 2 2 in Lemma 13 and E B(s.a)=plf(ss.mmm s.
Thus,
Bleaanmd (BanPlOa) [V (N [ S Blaamay ({s,a)
(almma ||0 0)1-==1 + vw[(1- (1 77
=a 16 عووااا 29
+ Vws(l - >), a= Va'.
Going back to the simulation lemma 20 we have
(1 (V P.r
Vr..))
((.a))d
min (o11((8,))15= :2
E P( [ P
(a.o) [V [V')r( min (allā(s,a)lbz) Vwc(1 7):2 2) )
2id
(a.a))dd
min allô(s, ()lss,))lsss
(1lá(s,a))lx-1- 2.
< Vw((1- >)-
This concludes the proof.
With the above lemma, now we can proceed to prove the main theorem.
Theorem 14 (PAC guarantee REP-LCB). Set the parameters as in Lemma 13 With probability
1 - 6, for any comparator policy including history-dependent non non- non-Markovian policies, have
wd?
C= In(()
V VP (1 71-
n
where C* is the relative condition number under o":
x E' (s
C* == sup
ØER (s.)) (s.
Proof. In this proof, letting [(x, a) = -| 8. P*( 8, a)ll we condition on the events:
[s 1 (s, a)) S. VΦ E | a) بواا 0(1los,
(12)
From Lemma 10 and Lemma 17 this event happens with probability 1 1- 6.
For any policy n, we have
Vp.,, Vire,
< VV..r V Prr [w([1-7)"
Lemma 13)
< V V + W 71
1 1-1)
2 E(a + WC(1 ( -1
$ (1 7))- E(saa)= [b(s,a)]+
1-7)
(b)
Recall a() (1(
From the second line to the third line, note though is the argmax over Markoovian polices, A is also the argmax over all history-dependent polices. In the last line, we use simulation lemma 20
which is tailored to a time-inhomogeneous policy. We here use IV " S 2/((1 7)), noting
[lbl|>c O())
22

---

Published as a conference paper at ICLR 2022
We further caleulate the first term (a). Considering 16 and noting 116|10 2, we have
((aaa)=d7. [[ss, a)] S (E,ā)) 11 (5, a)) 1 7.b V nww {NIBLSnaoo [60(3.00] + +ndd
+ - (E,
From (12)
we have
n.E (> ()) (BB(o,a)mp
min (0?1) 4)
( a a)]
0.b
(13)
- Tr(nE (10 -
(14)
$ (r 100 Al) (nB 11 d.
(15)
Thus,
E(s.a)aep, [)) S Er..))dz..1l (B, a) (1, Vada2y + 7 + / - 7)
71
Second, we further calculate the second term (b). Considering the offline version Tf and noting (2 a) is upper-bounded by 4,
((aaa)d (s)
-.,,)) 10 (5, 3 nww (1II a)]] 47Ad / [F" a)l (1 1-71) 7)
S (( 10 (5, a) 3i V [175) + 7dd + v (w(117))
$ B 100 a) --ر a + /w((1= 7).
In the final line, recall Vw {nnc} + yAd + yn S
Finally, by combining the calculation of the first term (a) and the second term (b), we have
wa2d - 7)
Vp.rr V S (1 7) B(a,a)=d7, 11 (3,
3. Vdawwy + x +
Wc
- 7)2 E (ss,a)dd 11 (5, x)|ج== p.. V (1 - 733
71-
Now, we use the fact ((, 11 (3, ())a), upper-bounded as
E(z,))=zz (8, p.
$ V
Eis,ayma,, (10 mno
$ /
(-E a)l
(Refer to Lemma 21
VC-d/n.
(From (13)
Finally, we have
V v
C.d da + 7dd
Ctd +
wa wa?d(1 7)
W (1
7)
+
(1 (1-7) 7)
n
(1-7)
Crdd
C+d
$ (1 Viaawy+ + (1 7)
(Take out two dominating terms)
wift
C*I((MM/8))
71-
23

---

Published as a conference paper at ICLR 2022
The lemma below is key technical lemma for our proof. It shows that relate the expected value of any function (s, a) with respect dz (i.e., inside the learned model P) to the potential function with respect to d.p. i.., E(,āav=dz [ a)|ا== Pairing á and P is important since P is
2.
the low-rank transition model defined using 3. As have seen in the above analysis, when using the lemma below, instantiate f(s,a) == (P(-
(-Is,,)) Lemma 15 (One-step back inequality for the learned model in offline setting). Take any J C |l I B. We condition the event where the MLE guarantee holds:
A IR s.L.
B(s,a)~ s, - (- s,a))lf S.
Then, letting = maxx,a (1/ TT a)), for any policy n, have
(s,a))dy {f(s,a)}| s B (( مااا V (nwDE (f + > + my
WB (Ba (1
Proof. First, we have an equality:
B(aa)ady (f(s.a)} = (3. (f(s,a)3 + (1 - B (f(s,a))
(16)
The second term in (16) upper-bounded by
(x (S(s,a) a)} < (gndg,awa(so) (s WWBe - (
Next consider the first term in (16) By CS inequality, we have
E on P(B.) (f(s,a)) = Bis.ayay ((3,))))((,,)) / Ealomtaaa)) o))(ss
/ Σ ((s)T(a ()
$ (3, 1l00(, 0||=-
Then,
/ (())((a
} I - { /IMoo)maal (())((al>))))s))
Σ ((s)T(al ())(())))))))))))
(Baaa {1E (lss) ((s) δ(z,ā)T(a ( s)f(s,))d(s)
B2dd
(Use the assumption
Σ. S(s, a) ax S and II (()))()))(s)))a S Vd (0, 1])
(B.. [f(s, + B B2
(E. [f 2) |
Ba + nBB² (MLE guarantee and | (f) S B².)
<n (BGG))ops- (s, a) B² Ad + nB²C
(Jensen)
Finally, the first term (16) is upper-bounded by
+ XdB - nB²c n (3 [f(s, + AdB² + nB²
< nw {
(Importance sampling)
S nw { B [ (s,aa) + 2dB² nB²c
(Definition of p)
24

---

Published as a conference paper at ICLR 2022
In the last line, use the following equality:
((s.))= [s)) (. (B. (m,)))))m))) [f((s + (1 Bd, [f((s
Based on the above discussion, the final statement is immediately concluded.
We can prove the similar inequality for the true model. The proof is omitted since it is quite similar to the one of Lemmal 15
Lemma 16 (One-step back inequality for the true model in offline setting). Take any f CSXA R s.t. Il S B. Then, letting w = max,,o oax,,,(1mm) a s)), for any policy we have
(3, ((c.a)=dT. (s(8, s ((.a))a 13, a) sp... V (nwB (83(s, + 7
(a (1
D AUXILIARY LEMMAS
First, we present the MLE guarantee. Regarding the proof, refer to Agarwal et al. (20206 Theorem 21). Note Pn and nn are the quantities appearing in the proposed online algorithm. We can also immediately obtain the statement to the offline case.
Lemma (MLE guarantee). For a fixed episode n, with probubility 1 1, 6,
gg- (0.5,,... ).a) a) - (- 1. a.)) : In([M)/8)
n
As a straightforward corollary, with probability 1 6,
Vn E N+ E +0.50) (00 8,a) - (. ( (s,a) S 0.5G Sn == lm(lM/n/o)
(17)
The following is standard inequality to prove regret bounds for linear models. Refer to Agarwal ct al. [2000a] Lemma G.2.)
Lemma 18. Consider the following process. For L,... ,N, M, Mn 1 Gn, with Mo Aol and G, being a positive semidefiite matrix with eigenvalues
by 1. We have that:
2 In det(MN) 22 In
Tr
==1
Lemma 19 (Potential function lemma). Suppose Tr(G,) S B².
NB²
2 In det((M)) dee dln 1
dho
Proof. Let "., od be the set of singular values of MN recalling My is a positive semidefinite matrix. Then, by the AM-GM inequality,
d
d s Ind (a.//Do))
In (ddettMx)//dl(fA/e = In II (a,/\o) (T. :=1
Since we have E, Tr(MN) dig NB2 the statement is concluded.
Lemma 20 (Simulation lemma). Given two MDPs (P', +b) and (P,r). for any policy n, we have;
1 VPr,44 - V = B B(e.o)a;, (bs,a) + BB 17, (s', m) - PP (5,,, 1
and
VP
VPrr
E (s, (bs,,a) + "E 10 ()] -PTEp( 1 r4 x)]].
1-
25

---

Published as a conference paper at ICLR 2022
Proof. We use
1
VØ f50, T)
Ezz (r(s,)) + 1BMMaaaa m)] f(s,a)]
Then,
Prr++ Vp,
By (s.a) [m(s, a) + ( T]] -7,,
1
((o (b8,a) + ((s,, B1 10
Similarly,
Vpr Vp =
1 - E ((s + 11 7)] prr a)ll
7 E d3 (-b(s, + 7 (( 1P), x)] y1B 110 m)]l
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 o', we have:
C² = sup ((,,o) 10 (0 1
B(ea))dy, (e (s, a))
zERt x EE (s 2 (Ø²} 1
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/c In (1 + 1/' a1 In 12 a2) In 1/2(e2223 + az))
satisfies
a1 (1 (21+ agNN < cc.
where is constant independent of 02 01,02,00 ag.
Proof. We first have
1/N In (2 + 12 1/3 a)N) S max(In + 02) un) 22 + as). 1) I/Ni/Nln( + N).
Here, we use
In 122 + N) (ln(1+ 02) + In(1 N2 S (mlx( In(1 N). Then, we prove when = 1/e² x In ln2 1
11/NIn(1 + <)<
<6
This is proved by
WI Sex In(1 1/1+1/62 In (1 (1+1/)2)
In(1 + 1/22)
1n(1+1/63 + In(11 + In2 + 1/2))
Scx
In 11/2)
In(1+n
SetEx
In(1 + 1/(e)
Sctex 0.5[1+h2 + +1/2)))
In + 1/E2)
Sc.
From the third line to the fourth line, we use In(x) & 0/5(x - 1) for I > 0. Then, the final statement is concluded.
26

---

Published as a conference paper at ICLR 2022
E MORE COMPARISON TO XIE ET AL. (2021)
We briefly explain the guarantee when we use Algorithm first define new feature class $,-.
Xie et al. 2021).
For a given reward r,
Definition 23 (Augmented feature). Let ④ [01, .0a]
+ = :0 Φ). o, = 101. rd,
Then, we set
F aa o. |aal2 S L, E
where e is some suitable constant. Given the hypothesis class FF for the Q-function, we can run Algorithm in (Xie et al., 2021). We denote the output policy R.
We check two assumptions to ensure the algorithm works. The first assumption is realizability. This satisfied since for any policy II (II is the class of all Markovian polices), we have Qp.,. E F. The second assumption is completeness. This is also satisfied since Tp. C FF for any policy II where Tp.. .r
is Bellman-operator s.t.
T {S x A - - Rf (8) + P x)] E (8 ASSA--1 -)
where we denote [(s, TT) (ewx((s (8, a). Then, by invoking their Corollary 5, we have
Theorem 24 (PAC bound based on Xie et al. (221)). ). With probability 6,
V Ct.r
(d log(1/0) log b11) 1/5
VT E IT: V Vp.,r Sc
7)2
where
Eag. 1o {Ø laa
C
sup of - aE
sup
E E (ot }1aa
We compare the above result with our result in Theorem 6 First, since C includes and all pos- sible features in 4 this partial coverage condition is stronger than ours (recall our partial coverage condition only related to the representation ø²), and always have C Ct,', 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 hiseery-dependent policies.
27


Reviews:
[{"review_id": "yFL_9sAxaab", "reviewer": "Reviewer_MSqC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nThis paper improves the sample complexity bound of [FLAMBE](https://arxiv.org/abs/2006.10814) by doing the exploration, exploitation, and model selection at the same time. This interplay is much more practical and interesting than the previous explore and commit version. Also, in the offline algorithm (Rep-LCB), the proposed relative condition number provides some insight into the coverage of the offline dataset. It also makes sense by comparing the result of the offline algorithm with the policy covered by the offline dataset.\n\nWeakness:\n\nFirst, I'm concerned with the actual sample complexity of the algorithm. The authors calculate the sample complexity by counting $N$, which is the total episodes. However, it takes $poly((1 - \\gamma)^{-1})$ in expectation to sample state $s$ from distribution $d_{P^*}^{\\pi}$. Also, [FLAMBE](https://arxiv.org/abs/2006.10814) considered a time-inhomogeneous model for finite-time horizon setting, while here a time-homogeneous, infinite-time horizon setting is considered. I wonder if it is fair to directly convert the $H^{22}$ to $(1 - \\gamma)^{22}$ here.\n\nSecond, the realizability assumption (Assumption 2) sounds wired to me. Why we need **any** combination of $\\mu$ and $\\phi$ is a (signed) measure (i.e. $\\int_{s'} \\mu^\\top(s')\\phi(s, a)\\mathrm d(s') = 1, \\forall \\mu, \\phi, s, a$)? This assumption sounds too restrictive to me, and it has never appeared in the previous literature (like [FLAMBE](https://arxiv.org/abs/2006.10814)) to the best of my knowledge. It would be better if the authors can provide some examples satisfying that assumption.\n\nThird, though the previous literature made the same Maximum Likelihood Oracle as in Definition 3 here, I'm wondering if this optimization oracle can be solved efficiently for an arbitrarily large state space. It would be beneficial if the authors can provide some optimization solutions for this oracle.", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper studies the representation learning in linear MDP and significantly improves the previous sample complexity. However, as mentioned above, I'm concerned with the following aspects:\n- Sample complexity of the algorithm and fairness comparing two algorithms in different settings (time-inhomogeneous v.s. time-homogeneous)\n- Assumption 2 sounds too restrictive for me.\n- Definition 3, though appeared in the literature, I wonder if it can be solved efficiently\n\nBased on the concerns above, I lean towards rejecting this paper but I'm open to further discussions.\n\n-- post rebuttal\nI've updated my score since the authors clearly addressed my concerns.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "ninrxqaAx3B", "reviewer": "Reviewer_SjqZ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The almost optimism/pessimism seems not surprising to me. The key observation, as far as I understand, is that we can maintain the data distribution for the online setting via certain stop criteria in infinite horizon scenarios. That is to say, when we terminate the episode with probability $1-\\gamma$ at each step, the state distribution we sample is exactly the stationary state distribution. If we uniform sample the action at the state we terminates, by standard generalization arguments (e.g. Lemma 12), we only need to pay an additional $|\\mathcal{A}|$ factors, which will be fine. Also, such stationary distribution connects the value function with the sampling distribution of dataset, which is helpful for the theoretical derivation (e.g. in Lemma 9).\n\nAlthough the overall idea is interesting, I have several concerns on this paper, both from the high-level perspective and the technical details.\n\n1. If I understand correctly, the proposed algorithm can only be worked on the infinite horizon cases, as it heavily relies on the observation that when we terminate the episode with probability $1-\\gamma$ at each step, we can return the sample of state following the stationary distribution. As far as I know, this does not hold for any real ''episodic’’ case, where we always finish at a fixed number of step. I would like to ask, is it possible to construct a similar algorithm for the episodic setting that terminates at a fixed number of step, following the idea proposed in the paper, as the authors argue they follow the exactly same setup as Flambe in Table 1, but the basic setting already differs. And if no such generalization, I will suggest the authors revising such claim and making the application of this trick and the potential limitation of this trick more explicitly.\n\n2. I would also want to ask the question in the following technical way: how can we track the state visitation probability other than the policy cover methods in Flambe if we consider the fixed termination setting? I feel things can be even worse when the transitions can depend on the specific level $h\\in[H]$ where $H$ is the length of each episode (which is the exact setting of Flambe). Does this paper provide new insights on that? Or is this paper only provide a cheap way for the policy cover, utilizing the structure of infinite horizon setting?\n\n3. The proof part is not well organized and can be hard for the potential readers to understand. There are also some typos in the derivation, most of them are easy to fix. I want to name some typos that make me extremely confused at the first glance. At the end of Page 15, the authors claim to replace $\\pi_n$ with $\\bar{\\pi}_n$ via importance sampling. However, $d$ is the stationary distribution, whose density ratio cannot be directly bounded by $|\\mathcal{A}|$. Such transformation is needed for the application of ellipsoid potential lemma, and I guess here $d^{\\bar{\\pi}_n}$ is not the stationary distribution of $\\bar{\\pi}_n$, but the stationary distribution of $d^{\\pi_n}(s) \\times \\text{Uniform}(a)$. Overall, the $\\bar{\\pi}_n$ is never defined. I hope the authors can go through the proof again to eliminate all of these ambiguities.\n\n---------------\nUpdate after the rebuttal: I check the proof again and understand the generalization to the time-inhomogeneous setting now. The one-step back trick can be conducted at any time step $h$ conditioned on the visiting distribution at $h-1$, and with recursion we can go back to the initial distribution $d_0$, which I now feel is the most significant contribution to the community. In this way, we don't need to construct point-wise upper bound like UCBVI and LSVI, which can be beneficial for the learning of low-rank MDP (as if I understand correctly, Flambe, in some sense, aims at providing point-wise upper bound so that they need to make the policy cover with the ellipsoid planner). I have adjusted my score accordingly.\n", "summary_of_the_paper": "This paper consider the representation learning in low-rank MDPs under infinite-horizon settings. The authors provide the improved sample complexity and statistical error for the online and offline scenarios correspondingly, based on an observation called the almost optimism/pessimism at the initial state distribution.", "summary_of_the_review": "The authors don’t properly address the difference between the newly proposed method and Flambe, and I feel the proposed method, although delicate, has so many restrictions that the authors don’t properly address. Meanwhile, the proof is not well-polished. I hope the authors can refine the writing during the rebuttal and make the intuition much more clear, compared with the current highlight of analysis in Section 4.2, that implicitly use several convenient property of the infinite horizon setting.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "G0rtN1tFTvX", "reviewer": "Reviewer_cvhG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is well-motivated: representation learning is a very important question in RL, both empirically and theoretically. This paper is also well-written, and the technical part of this paper is easy-to-follow.\n\nOn the plus side:\n-\tThe REP-UCB algorithm is conceptually very simple: it alternates between learning a model (using MLE) and finding an optimistic policy. In terms of the computation complexity, the REP-UCB algorithm is oracle efficient, using the same computation oracle as FLAMBE.\n-\tThis paper significantly reduces the sota sample complexity in the online setting to a reasonable level. Although technically speaking, FLAMBE focuses on reward-free exploration, where this paper focus on a standard online setting.\n-\tFor the offline setting, this paper achieves polynomial sample complexity with partial coverage, where the coverage is measured by the ground-truth feature. In contrast, previous results in the offline setting assumes a known representation, or requires a coverage over all hypothesis. \n\nOn the minus side, I have the following questions/concerns:\n-\tThe sample complexity in this paper is measured by the number of episodes N. In each episode, the algorithm collects one (s,a,s’) tuple from distribution d^{\\pi}_{P^*} with uniformly random action. However, after taking a random action, the distribution of the state changes. Does it mean that, in order to collect one transition tuple, the algorithm needs to interact with the environment for multiple steps? How does this mechanism affect the sample complexity?\n-\t-\tIn the proof of Lemma 12, the last equation block in Page 18: Could you elaborate on the first inequality (the MLE guarantee)? If I understand correctly, the inequality is doing a change of measure (from \\hat{P} to P^\\star). But the MLE guarantee only upper bounds the expected one norm squared. So a naïve Jensen gives an $\\sqrt{\\zeta_n}$ upper bound, instead of \\zeta_n stated in the paper. Then $n\\sqrt{\\zeta_n}$ is on the order of \\sqrt{n}, and I expect that this bound is too loose for later steps.\n\nAdditional questions/remarks:\n-\tDid the authors try regret minimization in the online setting? It seems to me that the random action in data collection creates some technical difficulty.\n-\tIn Page 14, the second displayed equation in the proof of Lemma 8: it should be ${\\hat{\\Sigma}_{n,\\phi}^{-1}}$.\n\n======= after rebuttal =======\nThe authors' response addressed my concerns. I'll raise my score according.", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper is well motivated, well written, and the results are impressive. But I don’t fully understand some of the derivations (see main review). So at this point, I can only suggest a weak accept with low confidence. If my concerns were addressed, I’m happy to raise my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "4eyirdPhUbK", "reviewer": "Reviewer_99V8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written. \n\nThe paper provides a few tangible novel definitions such as partial coverage in offline RL.  \n\nHowever, the paper needs further work to be ready. \n\n1- The first major drawback of this paper is that the online algorithm just uses one sample per episode, no matter how long is the episode. I doubt anyone would implement this algorithm in practice and the insight is limited. I understand it makes the theory cleaner, but the authors are encouraged to provide a full study. \nAlso, it is not clear how this algorithm might be extended to undiscounted finite-horizon MDPs, a subset of episodic MDPs. \n\n2- It seems the work assumes M is finite. While extensin to infinite set might be straightforward, including it in the paper is necessary. In that case, is the MLE part efficient?\nNote that is it quite a limiting setting since state-space seems to be large.\n\n3-  The condition that for any mu and phi \\int_s' mu(s') phi(s,a)= 1 is very strong and very limiting. The authors are encouraged to relax it. \n\n\nI might have been wrong in the above statements.\n\nThe paper studies an important problem, but some details are still left to be developed and the authors are encouraged to address them. \n\n\n", "summary_of_the_paper": "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. \n\nThe algorithm provides PAC style bound on how many episodes are needed to learn epsilon optimal policy.\n\nThe 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.\n \n\n\n\n\n\n", "summary_of_the_review": "The paper is well written,\nAddress an important problem.\nHowever, there are pieces of this work that are still missing to complete the work. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper, the authors extend the FLAMBE to the infinite-horizon MDP and largely improved the sample complexity of the representation learning in FLAMBE. Meanwhile, the authors also consider the offline representation learning with the same framework. Although there is still some computational issue in MLE for the linear MDP, the paper completes a solid step towards making linear MDP for practice. The paper could be impactful for the RL community. 

As the reviewers suggested, there are still several minors to be addressed:

- The extension of the proposed algorithm for finite-horizon MDP should be added. 
- The directly comparison between the sample complexity of FLAMBE and the proposed algorithm in infinite-horizon MDP is not appropriate. The authors should clarify the difference here. 
- The organization of the proof is not clear. As reviewer suggested, the one-step back trick should be emphasized for better significance of the submission.
Decision: Accept (Spotlight)