Improved Sample Complexity for Reward-free Reinforcement Learning under Low-rank MDPs
Authors: Yuan Cheng
Ruiquan Huang
Yingbin Liang
Jing Yang
Venue: ICLR 2023 poster (2023)
Keywords: Reward Free Exploration
Representation Learning
Sample Complexity
Model-Based RL

Abstract: In reward-free reinforcement learning (RL), an agent explores the environment first without any reward information, in order to achieve certain learning goals afterwards for any given reward. In this paper we focus on reward-free RL under low-rank MDP models, in which both the representation and linear weight vectors are unknown. Although various algorithms have been proposed for reward-free low-rank MDPs, the corresponding sample complexity is still far from being satisfactory. In this work, we first provide the first known sample complexity lower bound that holds for any algorithm under low-rank MDPs. This lower bound implies it is strictly harder to find a near-optimal policy under low-rank MDPs than under linear MDPs. We then propose a novel model-based algorithm, coined RAFFLE, and show it can both find an $\epsilon$-optimal policy and achieve an $\epsilon$-accurate system identification via reward-free exploration, with a sample complexity significantly improving the previous results. Such a sample complexity matches our lower bound in the dependence on $\epsilon$, as well as on $K$ {in the large $d$ regime}, where $d$ and $K$ respectively denote the representation dimension and action space cardinality. Finally, we provide a planning algorithm (without further interaction with true environment) for RAFFLE to learn a near-accurate representation, which is the first known representation learning guarantee under the same setting.
 

Published as a conference paper at ICLR 2023
IMPROVED SAMPLE COMPLEXITY FOR REWARD- FREE REINFORCEMENT LEARNING UNDER LOW-RANK MDPs
Yuan Cheng
Ruiquan Huang"
University of Science and Technology of China cy168mahe wsw.cscccoom edu.
The Pennsylvania State University #zh5514@psuu edu
Jing Yang
Yingbin Liang
The Pennsylvania State University yangjjing@psu. edu edu
The Ohio State University 1 liang. 8898osu. edu
ABSTRACT
reward-free reinforcement learning (RL), an agent explores the environment first without any reward information, in order to achieve certain learning goals afterwards for any given reward. In this paper we focus on reward-free RL under low-rank MDP models, in which both the representation and linear weight vectors are unknown. Although various algorithms have been proposed for reward-free low-rank MDPs, the corresponding sample complexity is still far from being isfactory. In this work, we first provide the first known sample complexity lower bound that holds for any algorithm under low-rank MDPs. This lower bound implies it is strictly harder to find a near-optimal policy under low-rank MDPs than under linear MDPs. We then propose novel model-based algorithm, coined RAFFLE, and show it can both find an e-optimal policy and achieve e-accurate system identification via reward-free exploration, with a sample complexity sig- nificantly improving the previous results. Such sample complexity matches our lower bound in the dependence on e, as well as on in the large d regime, where d and K respectively denote the representation dimension and action space cardi- nality. Finally, we provide a planning algorithm (without further interaction with true environment) for RAFFLE to learn near-accurate representation, which is the first known representation learning guarantee under the same setting.
INTRODUCTION
Reward-free reinforcement learning, recently formalized by Jin al. (2020b). arises as a powerful framework to accommodate diverse demands in sequential learning applications. Under the reward- free RL framework, an agent first explores the environment without reward information during the exploration phase, with the objective to achieve certain learning goals later on for any given re- ward function during the planning phase. Such a learning goal can be to find an e-optimal policy, to achieve an e-accurate system identification, etc. The reward-free RL paradigm may find broad application in many real-world engineering problems. For instance, reward- free exploration can be efficient when various reward functions taken into consideration over single environment, such as safe RL (Mryyoosefi & Jin [02 Huang et al. 2022), multi-objective RL (Wu et al.) (02T), multi-task RI. (Agarwal et al.) |022) Cheng et al. 2022)) ) etc. Studies of reward-free RI. the theo- retical side have been largely focused on characterizing the sample complexity to achieve a learning goal under various MDP models. Specifically, reward-free tabular RL has been studied in J al. (2220)): Ménard et al. (202 Kaufmann et al. (2021); I): Zhang et al. (2020). For reward-free RL with function approximation, Wang al. (2020) studied linear MDPs introduced by Jin et al. (202bb, where both the transition and the reward are linear functions of a given feature extractor, Zhang et al.
"Equal contribution

---

Published as a conference paper at ICLR 2023
2021b) studied linear mixture MDPs introduced by Ayoub et al. (2020), and Zanette et al. (0020b
considered a classes of MDPs with low inherent Bellman emorr introduced by Zanette et al. (020a)
In this paper, focus a decomposition into two embedding functions that map to low dimensional spaces. Compared with linear MDPs, the feature functions (i.e., the representation) under low-rank MDPs are unknown, hence the design further requires representation learning and becomes more challenging. Reward- free RL under low-rank MDPs was first studied by [Agarwal al. (2020), and the authors introduced a provably efficient algorithm FLAMBE, which achieves the learning goal of system identification
reward-free RL under low-rank MDPs, where the transition kernel admits
with sample complexity of Ō =² Here d, H and K respectively denote the representation dimension, episode horizon, and action space cardinality. Later on, Modi et al.) (2021) proposed model-free algorithm MOFFLE for reward-free RL under low-nonnegative-re
MDPs (where
feature functions are non-negative), for which the sample complexity for finding an e-optimal policy scales as Ō Hº K*d2v ) (which is rescaled under the condition of Σ Th S for fair comparison). c2n
Here, dry denotes the non-negative rank of the transition kernel, which may exponentially larger than d as shown in AAarwal et al. (2020), and 1) denotes the positive reachability probability to all states, where 1/11 can be as large as vdLV as shown in Uchara et al. (2022b). Recently, a reward- free algorithm called RFOLIVE has been proposed under non- monerrnnearaee MDPs with low Bellman Eluder dimension (Chen et al. 2022b), which can be specialized to low-rank MDPs. However, RFOLIVE is computationally more costly and considers a special reward function class, making their complexity result not directly comparable to other studies on reward-free low-rank MDPs.
This paper investigates reward-free RL under low-rank MDPs address the following important open questions:
For low-rank MDPs, none of previous studies establishes a lower bound on the sample com- plexity showing a necessary sample complexity requirement for near-optimal policy finding.
The sample complexity of previous algorithms in Agarwal et al. (2020); Modi et al. (202T) on reward-free low-rank MDP is polynomial in the involved parameters, but still much higher than desirable. It is vital to improve the algorithm to further reduce the sample complexity.
Previous studies on low-rank MDPs did not provide estimation accuracy guarantee on the learned representation (only on the transition kernels). However, such a represcntation learning guarantee can be very beneficial to reuse the learned representation in other RL environment.
1.1 MAIN CONTRIBUTIONS
We summarize our main contributions in this work below.
Lower bound: We provide the first-known lower bound 52 Hadk) on the sample complexity that holds for any algorithm under the same low-rank MDP setting. Our proof lies in novel construction of hard MDP instances that capture the necessity of the cardinality of the action space on the sample complexity. Interestingly, comparing this lower bound for low-rank MDPs with the upper bound for linear MDPs in Wang et al. (2020) further implies that it is strictly more challenging to find opar-optimal policy under low-rank MDPs than linear MDPs.
Algorithm: We propose new model-based reward-free RL algorithm under low-rank MDPs. The central idea of RAFFLE lies in the construction of a novel exploration-driver reward, whose corresponding value function serves as an upper bound on the model estimation error. Hence, such a pseudo-reward encourages the exploration to collect samples over those state-action space where the model estimation error is large so that later stage of the algorithm can further reduce such an error based on those samples. Such reward construction is new for low-rank MDPs, and serve as the key reason for our improved sample complexity.
Sample complexity: We show that our algorithm can both find an e-optimal policy and achieve an e-accurate system identification via reward-free free exploration, with a sample complexity of which matches our lower bound in terms of the dependence on as well as on K in the large regime. Our result significantly improves that of Ō Ō ""*"""" in Agar- wal al. (2020) achieve the same goal. Our result also improves the sample complexity of
Ō H k(d
(""K"d.v in Modi et al. (2021) in three aspects: order on K is reduced; d can be exponen- Ō <2n
tially smaller than dry as shown in Agarwal al. (2020); and no introduction of 1, where 1 1/n
2

---

Published as a conference paper at ICLR 2023
can be large vdgv. Further, result reward-free RL naturally achieves the goal of reward-known RL, which improves that of Ô H in Uehara et al. (2022b) by (
Near-accurate representation learning: We design planning algorithm that exploits the exploration phase of RAFFLE to further learn provably near-accurate representation of the transition kernel without requiring further interaction with the environment. To the best of our knowledge, this is the first theoretical guarantee representation learning for low-rank MDPs.
2 PRELIMINARIES AND PROBLEM FORMULATION
Notation. For any H E N, we denote [H] == (1...., HJ. For any vector I and symmetric matrix A. we denote |,0|| as the Ez norm of x and ||=|| Vx Ar. For any matrix A, we denote || as its Frobenius norm and let (A) be its i-th largest singular value. For two probability measures P. Q St, use | Qlirv denote their total variation distance.
2.1 EPISODIC MDPS
We consider an episodic Markov decision process (MDP) M (S,A.H,P,r), where S can be an arbitrarily large state space: A is finite action space with cardinality K; H is the number of steps in each episode;
x x + [0, 1] the time-dependent transition kernel, where Pa (#k+1181, (%) denotes the transition probability from the state-action pair (sh; (%,) at step h to state 8h+1 in the next step; Гд S A - [0, [] denotes the deterministic reward function at step h; We further normalize the summation of reward function as Σx yy < 1. A policy ı is set of mappings [Th S + ((A)) where (AA) the all probability distributions over the action space A. Further, (AA indicates the uniform selection of an action a from A.
In each episode of the MDP, we assume that fixed initial state is drawn. Then, at each step [H]. the agent observes state s S, takes an action aş AA under a policy тhi and receives reward (s,, (,)) ) (in the reward-known setting), and then the system transits to the next state 94+1 with probability (s Sh, an). The episode ends after H steps.
As standard in the literature, we use under the transition kernel for h steps. If the previous state-action pair (#)-1,4)-1) given, we use 8h P to denote that 84. follows the distribution PA. (-18-18-1 (n-1). We use the notation (.)e.))))((..)))
(P, to denote state sampled by executing the policy π
[]] denote the expectation over states 8h (P, T) and actions on ~ π.
For given policy æ and MDP M (S,,.H,P,r), we denote the value function starting from state 8n step h (s.) = ((s,)sn))) En= Σ#=n (sw, ag) We VØ,, to denote VI, (s) for simplicity. Similarly, we denote the action-value function starting from state V (9+1))) an
action pair (sh, an) at step h as (s,,,, an) (5n, an) EESTIP
We use Pa to denote the transition kernel of the true environment and for simplicity, denote (.. as E. Given reward function r, there always exists an optimal policy " that yields the optimal value Vp.,- = sup, V
2.2 LOW-RANK MDPS
This paper focuses on the low-rank MDPs (Agarwal al. 2020) defined follows.
Definition 1. (Low-rank MDPs). A transition probability Pž S x A - A(A) admits low-rank decomposition with dimension d E N if there exists two embedding functions oi S x A Rd and pj S + ad such that (), ( a) = (0) (s, a), (s')) Vs, E S,a A.
For normalization, we assume [s, [o, 1]. PT. (s))s))) < Vd. An MDP M a low-rank MDP with dimension for each E H. Ph. admits low-rank decomposition with dimension d. We use o" = {dr he(n and f= = (nEA) to denote the embeddings for P*.
I for all (s, a). and for any function g
We remark that when di is revealed to the agent, low-rank MDPs specialize to linear MDPs (Wang al. 2020 [ii et al. 2020b). Essentially, low-rank MDPs do not assume that the features (
3

---

Published as a conference paper at ICLR 2023
known priori. The lack of knowledge features in fact invokes nonlinear structure, which makes the model strietly harder than linear MDPs or tabular models. Since it is impossible to learn model in polynomial time if there is no assumption features On and Ma. adopt the following conventional assumption from the recent studies on low-rank MDPs.
Assumption 1. (Realizability). A learning agent can access to a model class {(0, V) that contains the true model, i.e., the embeddings o²* 4 u E V.
While we assume the cardinality of the function classes to be finite for simplicity, extensions to infi- nite classes with bounded statistical complexity (such as bounded covering number) are not difficult (Sun al. 2019
Agarwal et al. 2020).
2.3
REWARD-FREE RL AND LEARNING OBJECTIVES
Reward-free RL typically has two phases: exploration and planning. In the exploration phase, an agent explores the state space via interaction with the true environment and can collect samples over multiple episodes, but without access to the reward information. In the planning phase, the agent is no longer allowed to interact with the environment, and for any given reward function, required to achieve certain learning goals (elaborated below) based on the outcome of the exploration phase.
The planning phase may require the agent to achieve different learning goals. In this paper, we focus on three of such goals. The most popular goal in reward-free ffrer RL is to find near-optimal policy that achieves the best value function under the true environment with e-accuracy, as defined below.
Definition 2. (e-optimal policy). Fix é > O. For any given reward function r, learned policy is c-optimal if it satisfies Vp... V S.
For model-based learning, Agarwal et al. (2020) proposed system identification another useful learning goal, defined as follows.
Definition (e-accurate system identification). Fix e > O. Given a model class (L, V), a learned model (ő, ()) is said to achieve c-accurate system identification if it uniformly approximates the true model P", i.e., VT, E [H]. E; (on (.) - P7 (-1s%, an) TV
S
Besides those two common learning goals, we also propose an additional goal on near-accurate rep- resentation learning. Towards that, we introduce the following divergence-based metric to quantify distance between two representations, which has been used in supervised learning (Du et al. 2021b).
Definition 4. (Divergence between two representations). Given a distribution over S x and two representations o, ø' E 4, define the covariance between 6 and ø' w.r.t & as [(e,o) E [o(s, (s, a) [] Then, the divergence between à and ø' with respect is defined
d')
Dy( Φ0) E(s,))= (6'.6))) (0) E(sn))
(E(s ((). "()))))))) (0) 0')
It can be verified that Dy(0, Φ') ) 2 0 (i.e., positive semidefinite) and Dy Φ) = 0 for any 0 E T.
3 LOWER BOUND ON SAMPLE COMPLEXITY
In this section, we provide a lower bound on the sample complexity that all reward- RL algo- rithms must satisfy under low-rank MDPs. The detailed proof can found Appendix C
Theorem (Lower bound). For any algorithum that can output an e-optimal policy (as Definition) 2 if H > max(24e, S 2 6, K 2 3 and 8 < 1/16, then there exists low-rank MDP model M such that the number of trajectories sampled by the algorithm is at least Ω (HaK)
To the best of our knowledge, Theorem establishes the first lower bound for learning low-rank MDPs in the reward-free setting. More importantly, Theorem I shows that it strictly more costly in terms of sample complexity to find near-optimal policies under low-rank MDPs (which have unknown representations) than linear MDPs (which have known representations) by at least factor of Q(K). . This can be seen as the lower bound in Theorem [ for low-rank MDPs has an additional term K compared to the upper bound ò H
provided in Wang et al. (000)) for linear MDPs.
4

---

Published as a conference paper at ICLR 2023
This can be explained intuitively as follows. In linear models, all the representations Φ SxAA - R are known. Then, requires at most O(d) actions with linearly independent features to realize all transitions (d, p)). However, learning low-rank MDPs requires the agent to further select ((K actions to access the unknown features, leading to a dependence on K.
Our proof of the struction of hard MDP instances. a) We divide the actions into two types. The first type of actions is mainly used to form large state space through a tree structure. The second type of actions is mainly used to distinguish different MDPs. Such a construction allows us to separately treat the state space and the action space, so that both state space and action space can be arbitrarily large. b) We explic- itly define the feature vectors for all state-action pairs, and more importantly, the dimension is less than or equal the number of states. These two ingredients together guarantee that the number of actions K can be arbitrarily large and independent with other parameters d and S, which indicates that the dependence on the number of actions K unavoidable.
lower bound mainly features the following two novel ingredients in the con-
4 THE RAFFLE ALGORITHM
In this section, we propose RAFFLE (see Algorithm I for reward-free RL under low-rank MDPs.
Summary of design novelty: The central idea of RAFFLE lies in the construction of novel exploration-driven reward, which is desirable because its corresponding value function serves as an upper bound on the model estimation error during the exploration phase. Hence, such pseudo- reward encourages the exploration to collect samples over those state-action space where the model estimation error large so that later stage of the algorithm can further reduce such an error based on those samples. Such reward construction are new for low-rank MDPs, and serve as the key enabler for our improved sample complexity. They also necessitate various new ingredients in other steps of algorithms, as elaborated below.
Exploration and MILE model estimation. each iteration n during the exploration phase, for each h E [H), the agent executes the exploration policy πT-1 (defined in the previous iteration) up to step Different from FLAMBE (Agarwal et al. 2020) that collects large number of samples for each episode, our algorithm uses each exploration policy to collect only one sample trajectory at each episode, indexed by h). Hence, the sample complexity of RAFFLE is much smaller than that of FLAMBE. In fact, such an efficient sampling together with our new termination idea introduced later benefit sample complexity.
1, after which it takes two uniformly selected actions, and stops after step + 1.
Then, the agent estimates the low-rank components 3(m) and (((") via the MLE oracle with given model class ( V) and a dataset Dj as follows
MLF(DE) == arg max ((.,,,)) log (0n a), un (s))).
Design of exploration reward. The agent updates the empirical covariance matrix (x()) as
ún()) = And + Er=1 ((m) (s) af ( (s) a
(T,
(T.h+1) (()) (r,++)) (1.h+1)))
(1)
where (s/ .h
(r,h+1) (r,+1) (rh+1)) is collected at iteration t, episode (h + 1), and step h.
8441
Next, the agent uses both Thr²) and (() to construct an exploration-driven reward function as (s, a) = min (1 (s. ())0))))))))-
(2)
where an is pre-determined parameter. We note that although individual ((x) (s, a) for each step may not represent point-wise uncertainty as indicated in [Uehara et al. (2022b). we find its total cumulative version Vi V7to, can serve as a trajectory-wise uncertainty mea- sure to select exploration policy. To see this, it can be shown that for any т and h, B; [|K (hh (sh,
p(o,,jo jicns + V where is a constant and On = Ollog)) n). As itcration number n grows, the second term diminishes to zero, which indi- cates that Th ((m) (under the reward of ((m) serves as good upper bound on the estimation error
(sh, an), (/""))))) P- an) rv
|
5

---

Published as a conference paper at ICLR 2023
for the true transition kernel. Hence, exploration guided by maximizing V (cn) (oo) will collect more trajectories over which the learned transition kernels are not estimated well. This will help to reduce the model estimation error in the future.
Algorithm RAFFLE (RewArd-Free Feature LEarning)
1: Input: 6ns Sn > O, E (0, 1), regularizer A,, model classes {(1, E 10 $).
2: Initialize To(-)s (s) be uniform; set Dj 0.
3: Phase I: Exploration Phase
4: for 1, 1.......... do
5: for h = 1, 1..... ......H H do
6: Use mn- roll into 8h-1, uniformly choose ap-1, Oh, enter into Shs 8h-1-
7: Collect data of ((",A), ( ......h (on,))
(o.h)
(n,h) (n.k)
"
8: Add the triple (n,A) (m.h) (n.hh)
(s/m
ah -8/+1 to the dataset Dj =D - (n.h) (n.h) (onk))
Learn (0h) (1)))) M))((D))
9:
Update transition dynamics p(n) as (s'(s,)) (0th) (s, a) (n) (s'))
10:
end for
12: Update empirical covariance matrix (7(00) as in Equation
13: Define exploration-driven reward function (/7)) in Equation 2
14:
Define an estimated value function V
Pin) joo) based on p(n) and ((m) as in Equation 3
Find exploration policy nn arg maxx p(n), B(m)"
VT
16:
if 2
((6),6(e) + +2 then
17:
Terminate Phase I: Exploration Phase and set p= = -p(n), -(0))
Tai Th
end if
19:
end for
20: Phase II: Planning Phase
21: Option (leam near-optimal policy): Receive reward function
(rA | 1=11 and compute
policy π = arg max, Vi..,
22: Option 2 (system identification): let f p= pe.
23: Option 3 (learn near-accurate representation): call Algorithm 2 of Repleerr and obtain Ø.
24: Output: policy म, learned transition dynamics pe, learned representation ø.
Design of exploration policy. The agent defines truncated value function iteratively using the estimated transition kernel and the exploration-die driven reward as follows:
@T prns 860 (Sh; an) = min { B(m)) (sh, an) ) + p(n) h01, goo (sh, an)
a N.,peo),Bon (SL) = E ( pon) Boo (sh, an)
The truncation technique here is important for the improvement of the sample complexity on the dependence of H. The agent finally finds an optimal policy maximizing Vig Poo (ion) and uses this policy the exploration policy for the next iteration.
Novel termination criterion. RAFFLE does not require a pre-determined maximum number of iterations as its input. Instead, it will terminate and output the current estimated model if the optimal value function VTT p(o) ((m) plus a minor term is below a threshold. Such a termination criterion essentially guarantees that the value functions under the estimated and true models are close to each other under any reward and policy, hence the exploration can be terminated in finite steps. Such termination criterion enables algorithm to identify accurate model and later find near-optimal policy with fewer sample collections than FLAMBE in Agarwal et al. (2020) as we discuss in the exploration phase. Additionally, our termination criterion provides strong performance guarantees on the output policy and estimator from the last iteration. On the contrary, |ehaaa et al.) 2022h) can only provide guarantees random mixture of the policies obtained over all iterations.

---

Published as a conference paper at ICLR 2023
Planning phase. Given any reward function r, the agent finds near-optimal policy by planning with the learned transition dynamics p and the given reward Note that such planning with known low-rank MDP is computationally efficient by assumption.
UPPER BOUNDS ON SAMPLE COMPLEXITY
In this section, we first show that the policy returned by RAFFLE an e-optimal policy with respect to any given reward in the planning phase. The detailed proof can be found in Appendix A Theorem 2 (c-optimal policy). Assume M is low-rank MDP with dimension d, and Assumption holds. Given any 6, ô E (0, 1}, and any reward function let it and Pr be the oulput of RAFFLE and x¹ == arg maX V be the optimal policy under the true model p Set ān =Ö(vK + d') and An = (d). Then, with probability at least 1 1-8, 8, we have V Vp S E, and the total number of trajectories collected by RAFFLE is upper bounded by Ō KKP 6
We note that the upper bound in Theorem 2 matches the lower bound in Theorem in terms of the dependence on as well as on K in the large d regime.
Compared with MOFFLE (Modi et al., 202I), which also finds an e-optimal policy in reward-free RL, our result improves their sample complexity of Ô 11 K x>d",x der on K is reduced. Second, the dimension dev of the underlying function class in MOFFLE can be exponentially larger than reachability assumption, leading to a factor 1/n in the sample complexity, which can be as large as vdLv. Further, Theorem 2 naturally achieves the goal of Soward-knownr RL with the same sample complexity, which improves that of Ô 11 [H*
in three aspects. First, the or-
shown (Agarwal et al. (2020). Finally, MOFFLE requires
in Uehara al. (022b)) by a factor of O(H²).
Proceeding to the learning objective of system identification, the learned transition kernel output by Algorithm achieves the goal of e-accurate system identification with the same sample complexity as follows. The detailed proof can be found in Appendix B
Theorem (e-accurate system identification). Under the same condition Theorem ( p be the output of RAFFLE. Then, with probability at least 1 - 5, io achieves e-accurate system identification, i.e. for any x and h: Ex IlK 0s, (sh, an), RF, ()) - P an) |lv S
2 Lnd les Pe
and the number of trajectories collected by RAFFLE is upper bounded by Ō H +K)).
Theorem 3 significantly improves the sample complexity of Ö [1 xd in Agarwal et al. 2020 100
on the dependence of all involved parameters for achieving e-accurate system identification.
6 NEAR-ACCURATE REPRESENTATION LEARNING
In low-rank MDPs, it is of great interest to learn the representation & accurately, because other similar RL environments can very likely share the same representation (Rusu et al.) 2016 /hhu cet al.) 2020 Dayan 1993) and hence such learned representation can be directly reused in those environments. Thus, the third objective of RAFFLE in the planning phase is to provide an accurate estimation of 6. We note that although RAFFLE provides estimation of ó during its execution, such an estimation does not come with an accuracy guarantee. Besides, Theorem B on system identification does not provide the guarantee on the representation ő. but only on the entire transition kernel P. Further, none of previous studies of reward-free RL under low-rank MDPs (Agarwal et al. 2020
Modi et al. 2021
Uchara et al. 2022b) established the guarantee on ô.
6.1
THE REPLEARN ALGORITIIM
In this section, we present the following algorithm of RepLearn which exploits the learned transition kernel from RAFFLE and leams near-accurate representation without additional interaction with
the environment. The formal version of RepLeam, Algorithm 2 is delayed in Appendix [D] We explain the main idea of Algorithm 2 as follows. First, for each h E [H], E [T], where 7 the
7

---

Published as a conference paper at ICLR 2023
number of rewards, Ny pairs of state-action ($h; (b)) are generated based distribution Qh- Note that the agent does not interact with the true environment during such data generation. Then, for any h, if we set the reward at step h to be zero, Qp., A.r (Sh: an) can have linear structure in terms of the true representation 0% (Sh, ap). Namely, there exists wh, decided by r, P and a such that TF.,,, (n)) = (0) (sh; (n), wa)) Then, with the estimated transition kernel P that RAFFLE provides and the well-designed rewards phot 2 P. (Sh, on) can be computed efficiently and serve as a target to learn the representation via the following regression problem:
arg min w ER Eterry Ln,,,nn)D (Q P. (5h, ax) - (#x an), w)
(4)
The main difference between our algorithm and that in Lu et al. 2021)) for representation learning is that, our data generation is based on P from RAFFLE, which carries natural estimation error but requires no interaction with the environment, whereas their algorithm assumes generative model to collect data from the ground-truth transition kernel.
6.2 GUARANTEE ON ACCURACY
In order to guarantee the learned representation by Algorithm 2 is sufficiently close to the ground truth, we need to employ the following two somewhat necessary assumptions.
First, for the distributions of the state-action pairs { in Algorithm 2 it is desirable to have P(-[s,a) a) approximates true p* (-Is, a) well over those distributions, so that (,p,,,(sh. ax) can ap- proximate the ground truth well. Intuitively, if some state-action pairs can hardly be visited under any policy, the output P(-|s, a) of Algorithm [ cannot approximate true p* (* a) well over these state-action pairs. Hence, we assume reachability type assumptions for MDPs as following so that all state-action pairs is likely be visited by certain policy. Such reachability type assumption is common in relevant literature (Modi et al. [2021] Agarwal et al., (020)).
As discussed in Section 6 intuitively, if some state action pairs can be hardly visited by any policy, the output of Algorithm P(-|s, a) can not approximate true (-(-|s, a) well over these state action pairs, so a standard reachability type assumption is necessary so that all states can be visited.
Assumption 2 (Reachability). For the true transition kernel 1", there exists a policy Tº such that min P7," (s) 2 Mmin. where IP7" (-) S + R is the density function over S using policy x to roll into state s at timestep h.
We further assume the input distributions {qA H
h=1 are bounded with constant CB. Then, together for any (s, a) E S A, we have (s, a) S Cmin (s, as, where Cmin Tuain
with Assumption 2
Cc..
Next, we assume that the rewards chosen for generating target Q-functions are sufficiently diverse so that the target Qp.,,, (St,, aλ) spans over the entire representation space to guarantee accurate representation learning in Equation (4) Such an assumption is commonly adopted in multi-task
representation learning literature (Du al. 202 Yang et al. 2021 Lu al., (02II. To formally
state the assumption, for any h E [H]] let (r
JeElT be a set of rewards (where T d) satisfying
Th h,t = As result, for each t,, given n², there exists w E IRd such that @ (8h, an) (sh, ox), wf," )) Let Wi;- = [[w], ...... "[""]] TRdxT.
Assumption 3 (Diverse rewards). The smallest singular value (a(W))) of W defined above satis- fies o3 2 22( ()), i.e., there exists constant Cp > 0 such that (2()) Cp.
We next characterize the accuracy of the output representation of Algorithm 2 in terms of the di- vergence Definition 4 between the learned and the ground truth representations in the following theorem and defer the proof to Appendix D
Theorem 4 (Guarantee for representation learning). Under Assumption 7 and Assumption 3 for any e,6 E (0, 1), any h E [H], and sufficiently large N,, let P be the oulput transition kernel of RAFFLE that satisfies Theorem 3 With probability at least 1 - 8, the output of RAFFLE satisfies
| (0h, 1/2 = =0( smmCmmi + Co
Tog
(5)
Cp
TN,
F
8

---

Published as a conference paper at ICLR 2023
We explain the bound in Equation (5) follows. The first term arises from the system identification error e by using the output of RAFFLE, and can be made as small as possible by choosing appropriate The second term is related to the randomness caused by sampling state-action pairs from imput distribution (ga he(H and vanishes as N becomes large. Note that Nj is the number of simulated samples in Algorithm| 2 which does not require interaction with the true environment. Hence, it can be made sufficiently large to guarantee a small error.
Theorem 4 shows RAFFLE can learn a near-accurate representation, based on which learned repre- sentations can be further used in other RL environments sharing common representations, similarly in spirit to how representation learning has been exploited in supervised learning (Du et al.) 2021b).
RELATED WORK
Reward-free RL. While various studies (Oudeyer al. 2007 Bellemare et al. 2016 Burda al.,
2018 Colas et al., 2018 Nair et al.| 2018 Eysenbach et al. 2018 Co-Reyes et al. 2018: Hazan et al. 2019 Du et al. 2019] Pong et al. 2019 Misra et al.) (020) proposed exploration algorithms the state space without using explicit reward signals, theoretically speaking, the paradigm of reward-free RL was first formalized by Jin et al. (2020a), where they provided both upper and lower bounds the sample complexity. For tabular case, several follow-up stud- ies (Kaufmann et al. 2021 Ménard et al.) (021) further improved the sample complexity. and Zhang let al.] (2020) established the minimax optimality guarantee. Reward-free RL was also studied with function approximation. Wang et al.] (2020) studied linear MDPs, and ZZang et al. (2021b) studied linear mixture MDPs. Further, /Zanette et al. (2020b) considered class of MDPs with low inherent Bellman error introduced by Zanette et al. (2020a). Chen et al. (2022)) proposed reward-free al- gorithm called RFOLIVE under non-linear MDPs with low Bellman Eluder dimension. In addition, Miryooseñi & Jin (021) proposed a reward-free free approach to solving constrained RL problems with any given reward- free RL oracle under both the tabular and linear MDPs.
for good coverage
Reward-free RL under low-rank MDPs. As discussed in Section I reward-free RL under low- rank MDPs have been studied recently (Agarwal et al., 2020, Modi et al. 202T). and our result significant improves the sample complexity therein. When finite latent state space is assumed, low- rank MDPs specialize block MDPs, under which algorithms termed as PCID (Du et al.) (2019) and HOMER Misra et al. 2020) achieved sample complexities of Ō x*H2x4
min(r""2 and O dHH min(ny".**)
respectively. Our result the general low-rank MDPs can be utilized to further improve those results for block MDPs.
Reward-known RL under low-rank MDPs. For reward-known RL. Uehara et al. (2022b) pro- posed a computationally efficient algorithm REP-UCB under low-rank MDPs Our design of the reward-free algorithm for low-rank MDPs is inspired by their algorithm with several new ingre- dients as discussed in Section 4 and improves their sample complexity on the dependence of H. Meanwhile, algorithms have been proposed for MDP models with low Bellman rank (Jiang et al.) 2017]. low witness rank (Sun et al. 2019). bilinear classes (Du et al. 2021a) and low Bellman cluder dimension (Jin et al. 2021), which can specialize low-rank MDPs. However, those algorithms are computationally more costly as remarked in [ehaaa et al. (202b)). although their sample complex- ity may have sharper dependence on d, K or H. Specializing to block MDPs, Zhang et al. (2022) proposed an algorithm called BRIEE, which empirically achieves the state-of-art sample complexity for block MDP models. Besides, Zhang et al. (2021a) proposed an algorithm coined ReLEX for slightly different low-rank model, and obtained a problem dependent regret upper bound.
CONCLUSIONS
In this paper, we investigate the reward-free reinforcement learning, in which the underlying model admits low-rank structure. Without further assumptions, we propose an algorithm called RAFFLE which significantly improves the state-of-the-art sample complexity for accurate model estimation and near-optimal policy identification. We further design algorithm to exploit the model learned by RAFFLE for accurate representation learning without further interaction with the environment. Although e-accurate system identification can easily induce H Ho-optimal policy, the relationship be- tween these two learning goals under general reward-free exploration setting remain under-explored, and is interesting topic for future investigation.
9

---

Published as a conference paper at ICLR 2023
REFERENCES
Alckh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complex- ity and representation learning of low rank mdps. In Advances in Neural Information Processing Systems, volume 33, 2020.
Alckh Agarwal, Yuda Song. Wen Sun, Kaiwen Wang, Mengdi Wang, and Xuezhou Zhang. Provable benefits of representational transfer in reinforcement learning, arXiv preprint arXiv:2205.14571, 2022.
Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin Yang. Model-based reinforcement learning with value-targeted regression. In International Conference on Machine Learning, pp- 463-474. PMLR, 2020.
Mare Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos.
Unifying count-based exploration and intrinsic motivation. Advances in neural information pro- cessing systems, 29:1471-1479, 2016.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv: 12894, 2018.
Fan Chen, Yu Bai, and Song Mei. Partially observable with b-stability: Unified structural condition and sharp sample-efficient algorithms. arXiv preprint arXiv:2209. 14990, 2022a.
Jinglin Chen, Aditya Modi, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. On the statisti- cal efficiency of reward-free exploration in non-linear RL. CoRR, abs/2206. 10770020611770, 2022b.
Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang. Provable benefit of multi- task representation learning in reinforcement learning. CoRR, 0abs/2006005000 2022.
John Co-Reyes, YuXuan Liu, Abhishek Gupta, Benjamin Eysenbach, Picter Abbeel, and Sergey Levine. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajec- tory embeddings. In International Conference on Machine Learning, pp. 1009-1018. PMLR, 2018.
Cédric Colas, Pierre Fournier, Olivier Sigaud, and Pierre-Yves Oudeyer. CURIOUS: intrinsically motivated multi-task, multi-goal reinforcement learning. CoRR, 11../810 2018.
Christoph Dann, Tor Lattimore, and Emma Brunskill. Unifying pac and regret: Uniform pac bounds for episodic reinforcement learning. arXiv preprint arXiv:1703.077100 2017.
Peter Dayan. Improving generalization for temporal difference learning: The successor representa- tion. Neural Comput., 5(4):613 (1.444616624 1993.
Omar Darwiche Domingues, Pierre Ménard, Emilie Kaufmann, and Michal Valko. Episodic rein- forcement learning in finite mdps: Minimax lower bounds revisited, 2020.
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, 2019.
Simon S Du, Sham M Kakade, Jason Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in rl. arXiv preprint arXiv:2103. 103.99.700,0 2021a.
Simon Shaolei Du, Wei Hu, Sham M. Kakade, Jason D. Lee, and Qi Lei. Few-shot learning via learning the representation, provably. In 9th International Conference on Learning Representa- tions, ICLR 2021. OpenReview.net, 2021b.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint 102,000000 2018.
Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. In International Conference Machine Learning, pp. 2681-2691, PMLR, 2019.
10

---

Published as a conference paper at ICLR 2023
Ruiquan Huang. Jing Yang, and Yingbin Liang. Safe exploration incurs nearly additional sample complexity for reward-fiee -ree arXiv preprint WwXiv:2206.140577 2022.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Con- textual decision processes with low bellman rank are pac-leamable. International Conference on Machine Learning. pp. 1704-1713, PMLR, 2017.
Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward-free exploration for reinforcement learning. In Proceedings of the 37th International Conference Machine Learning, pp. 4870-4879, 2020a.
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, 2020b.
Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of problems, and sample-efficient algorithms, 2021.
Emilie Kaufmann, Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Edouard Leurent, and Michal Valko, Adaptive reward-free exploration. In Algorithmic Learning Theory, pp. 865- 891. PMLR, 2021.
Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvari, and Chi Jin. Optimistic mle-a generic model-based algorithm for partially observable sequential decision making. arXiv preprint arXiv:2209. 24t.:22799.4997 2022.
Rui Lu, Gao Huang, and Simon S Du. On the power of multitask representation learning in linear mdp. arXiv preprint arXiv:2106.000 100.0055,, 2021.
Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Emilie Kaufmann, Edouard Leurent, and Michal Valko. Fast active learning for pure exploration reinforcement learning. In Inter- national Conference on Machine Learning, pp. 7599-7608. PMLR, 2021.
Sobhan Miryoosefi and Chi Jin. simple reward-fres approach to constrained reinforcement learn- ing, 2021.
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-6977 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.
Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine, Visual reinforcement learning with imagined goals. arXiv preprint arXiv:1807.04742. 2018.
Pierre Yves Oudeyer, Frdric Kaplan, and Verena V Hafner. Intrinsic motivation systems for au- tonomous mental development. IEEE transactions evolutionary computation, 11(2):26--866. 2007.
Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew- State-covering self-supervised reinforcement learning. arXiv preprint urXiv:1903.0344 2019.
Tongzheng Ren, Tianjun Zhang, Lisa Lee, Joseph Gonzalez, Dale Schuurmans, and Bo Dai. Spec- tral decomposition representation for reinforcement learning, arXiv preprint arXiv:2208.089 09515, 2022.
Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. CoRR, 2016.
11

---

Published as a conference paper at ICLR 2023
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, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. arXiv preprint arXiv:2110.0460 2021.
Masatoshi Uehara, Ayush Sekhari, Jason D Lee, Nathan Kallus, and Wen Sun. Provably efficient reinforcement learning in partially observable dynamical systems. arXiv preprint arXiv:2206. 12020, 2022a.
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline RL in low-rank mdps. In The Tenth International Conference on Learning Representations, ICLR. OpenReview.net, 2022b.
Lingxiao Wang. Qi Cai, Zhuoran Yang, and Zhaoran Wang. Embed control partially observed systems: Representation learning with provable sample efficiency. CoRR, abs/2205. 13422220,,347,,, 2022.
Ruosong Wang, Simon S Du, Lin F Yang, and Ruslan Salakhutdinov. On reward-free reinforcement learning with linear function approximation. arXiv preprint arXiv:2006. 1274, 2020.
Jingfeng Wu, Vladimir Braverman, and Lin Yang. Accommodating picky customers: Regret bound and exploration complexity for multi-objective reinforcement learning. Advances in Neural In- formation Processing Systems, 34:131 12-1132244, 2021.
Jiaqi Yang, Wei Hu, Jason D. Lee, and Simon Shaolei Du. Impact of representation learning in linear bandits. In 9th International Conference on Learning Representations, ICLR 2021. Open- Review.net, 2021.
Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Learning near op- timal policies with low inherent bellman error. In International Conference on Muchine Learning, pp. 10978-10989) PMLR, 2020a.
Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Provably efficient reward-agnostic navigation with linear value iteration. arXiv preprint arXiv:2008.07737, 2020b.
Andrea Zanette, Ching-An Cheng, and Alekh Agarwal. Cautiously optimistic policy optimization and exploration with linear function approximation. Conference on Learning Theory, pp. 4473- 4525. PMLR, 2021.
Wenhao Zhan, Masatoshi Uchara, Wen Sun, and Jason D Lec. Pac reinforcement learning for predictive state representations. arXiv preprint arXiw:2207.05738, 2022.
Weitong Zhang, Jiafan He, Dongruo Zhou, Amy Zhang, and Quanquan Gu. Provably efficient representation learning in low-rank markov decision processes. CoRR, abs/2106.1 11955, 2021a.
Weitong Zhang, Dongruo Zhou, and Quanquan Gu. Reward-free model-based reinforcement learn- ing with linear function approximation. Advances in Neural Information Processing Systems, 34, 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.0006. 2022.
Zihan Zhang, Simon S Du, and Xiangyang Ji. Nearly minimax optimal reward-free reinforcement learning, arXiv preprint arXiv:2010.05901, 2020.
Zhuangdi Zhu, Kaixiang Lin, and Jiayu Zhou. Transfer learning deep reinforcement learning: A 2020.
survey. CoRR,
12

---

Published as a conference paper at ICLR 2023
A PROOF OF THEOREM 2
We first provide a proof outline to highlight our key ideas in the analysis of Theorem and then provide the detailed proof. To simplify the notation, denote the total variation distance (sn, ah) - (sk. an) TV by fr. (sh,
Proof Outline for Theorem 2 Step provides an upper bound on the difference value functions under estimated model p(n) and the true model P for any given policy and reward r, as given in the following proposition (see proposition 3 in appendix A.2
Proposition 1 (Informal). There exist constants On = Ollog n) and d' = O(1) such that for any policy π and reward r, with high probability, we have
V - V
p S
Vi 860
+
Cn
This proposition is inspired by Uehara al. (2022b), while generalizes to non-stationary setting and arbitrary reward scenario from infinite stationary MDP and fixed reward. The main proof idea first notice that for any reward, we have
Vp.r
Vp
V pom f(n)
Then, we show that V p(n), ,(m) S V Bro),B.o) V
Cu
VT
Bin)
+
due the optimality of Eon
the exploration policy T..
Step 2 shows sublinearity of the summation of V T pin,, jom3" which is an upper bound of value function difference under the true model with the estimated model (see proposition [ in appendix (A.3).
Proposition 2 (Informal). Under the same setting Proposition 7 with high probability, the
mation of value functions pon, jim) under exploration policies Tn Ininn(N) with exploration-driver reward functions B(m) is sublinear, given by
Σ ME [M
The pio
860) 8600 SŌ 5Ō (HddK((K+
Step combines Step 2 and Step 1. We are able to conclude that RAFFLE will terminate with polynomial sample complexity such that, the value function difference of P and the returned envi- ronment p at most e.
Proposition 3. With high probability, RAFFLE terminates after at most Ô (H-PK(( K)/E)
iterations, and the oulput model ie satisfies
VP
S Ving pie (ine) + On. S <//2,
ne
where n. is the iteration number where RAFFLE lerminates, i.e. pe = p(m))
Finally, with some algebraic operations, Proposition concludes the proof of Theorem
A.1 SUPPORTING LEMMAS
We first present the high probability event.
Lemma 1. We define HI to be uniform mixture of previous exploration policies:
II, Ld(TT Ti...., Tn-1).
Denote the total variation of p(m) and P*. and the expected matrix of (7(1) follows.
13

---

Published as a conference paper at ICLR 2023
f/t) (sh. an) Ph P (-1sh, an)
TV
U/e) E MameP.IA an)ve(())
[0(s (1(s):a)((((n))) (0)sn
2x1
(8)
E Whog (sn ( II..)
[Ø(8n, (n) (n(sn an)) Ant.
where A, (21 and A3 (1) is constant coefficient. Suppose Algorithm
runs N iterations and let events Eo and & be defined as follows.
[1 (sh. (sh. an)) S Sn)
& Yn E [M], [H) ES, EA,
E
aha(Pr, "a
&
Vn [ E SS E,,
1 53 0/ (s, a) (00))) )-1 s 1100 16.02 (0
(s, (3) (
where Sn log (2) 1 v / /n. We further denote Ç = NÇN log (2
H/NH// for sim-
plicity. Denote & Eg n & as the intersection of the two events. Then, P[E) 2
Proof. By Corollary 2 in Appendix we have P[E]] 2 5/2. Further, by Lemma 39 Zanette et al. 2021) for the version of fixed @ and Lemma 11 in Uehara al. (2022b). we have P|
0/2. Therefore, Piε] 2 1-
Based on Lemma] we bound the exlporation-driven reward in RAFFLE as follows. Corollary 1. Given that the event occurs, the following inequality holds for any n E [N], h E [H]. 8h S, an E
ân
(sn, an) I
} < (1m) (sh: an) 30h 1 (sh, and
min-
where ān = 253/23n (K cd).
Proof. Recall (/()) (sh, an) = min ân (3 (s.)) Applying Lemma we can im-
mediately obtain the result.
The following lemma extends the Lemmas 12 and 13 under infinite discount MDPs in Uchara et al. 2022b) to episodic MDPs. We provide the proof for completeness.
Lemma 2. Let Pa-1- (Øn-1, (h-1) be generic MDP model, and II be arbitrary and possibly mixture policy. Define an expected Gram matrix as follows
M4-1,6 A +
(1-1($%-1,,n-1) 0n-1) 10 on-))))
Further, let 44-1) be the total variation between ph- and Ph-1 at time step h 1. Suppose E x A + I is bounded by B E c0), i.e., llglloo I B. Then, Wh. 2 2, V policy The [0(8,, ar. 18h-)
E
< |1 (8h an 1
nK
E
(sh, an)] + dn, + nB²
E
[fn- (0--188-11-1 an-1)?l.
"y (P*
"h "h11<(PM1 -h-10l1
14

---

Published as a conference paper at ICLR 2023
Proof. We first derive the following bound:
E
[g(s,, an)]sn-1, an-1]
( -h- (sh))
bh an
$ 1000--1(0-11,,0--)00000------- /
Σ an
(n)(aann)) ( (sh)dsh
Mr. 1.0
where the inequality follows from Cauchy-Schward inequality. We further expand the second term in the RHS of the above inequality as follows.
|
alt.. M4-1.8
27
E
E (on
+ And
27
E Aho-Ph Bn 7T
g(sh, an) (n
+ And
(P m
"g
(18)
2
2n
E P ug
g(sn sh-1, ah-1
And
11-10(p-ap m)
+ BnB2
E
ifn-1 (h-1(99-1,.0-11)
*h (p0 0p-1 wil
(181)
E
an) 84-11 ( an-1
And
(P m
+ 2n B2
E
[Jn 182-1,0h-117
#h-1
11-1
(40)
2nK
E
[g(sn ? + An, + 2nB
E
(8-1,,11,,
where (i) follows from the assumption that ||g]loc < B, (fi) due to that fn- (1h-1,, total variation between P and Ph-1 at time step h 1 and the fact that (a + b)² < 2a + 26 (iit) follows from Jensen's 's inequality, and (iv) is due to importance sampling. This finishes the proof.
the
we summarize three useful inequalities which bridges the total variation f(()) and the exploration-driie driven reward ((m))
Based on Lemma 2
Lemma 3. Define
W(n h.d
E
[d(Sh: (6)(),)))))))) A
(10)
15

---

Published as a conference paper at ICLR 2023
where An Byd (2n ( 1/6)) Given that the event & occurs, the following inequalities hold.
For any n, when 2
( (sh, an)
I an | (sn ag l1ogo ann
E
84-1, an-1
)-1
(11)
Am1
E *A A Pno
f( (s an)
84-1, 1,
I an - 18 ag 1l i
(12)
E
((m) ( (shr an)
8h-1, an-1
S Im 1/ (sh an an-1) - )-1
(13)
AN Pk ok mi
where
"n = 4583nG,KK(K de.). On = 2233nnn + dd,,
Specially, when h = 1,
E (81, S VKGn
E
b(s1,a)) S 15cm
dK
(14)
Proof. We start by developing Equation as follows. Given that the event & occurs, for h 2 have
E A-1
(h (sh: an) 10 an-1
an-
(m) hh-1 (sh-1, an-1))
)- x
1 (sh, And +n
(1 [ (s (3 ah-1)
nk
E *h- 1PP mln)
*h-1 mnn)
ril
Ph- (sh-1, an
(n)
)- x
n
E
(1" (sh, [n) + And + nK
E
(n) h-1 (8h-1, ah-1)
*h (PP mln)
(156)
(()) (sh-1,
hー
an
(200) ( )
2/15nK + Bon
< on [ (sn-1,
an an-1)
(Urin) )-1
where (1) follows from Lemma 2 and the fact that f(") an) < (ti) follows from importance sampling at time step h 2, and (Mi) follows from Lemma
Equation follows from the arguments similar to the above.
To obtain Equation
13) we first apply Lemma 2 and obtain
((m) (sn, an) 8h-1, ah
E "i-Pk_1 21 1
1 (8h-1. (1-11) (1- 1.r))
nK
[1. (sh, an)) 1] + And,
"A PP ma) on
16

---

Published as a conference paper at ICLR 2023
where we the fact that 8 (sh. an)
1.
We
further
bound
the
term
E Ea mn) [(8]] (sh, an))?] follows:
(1))) (sh, an))
E
*A p "
Uh (sh+ an 1100o
((m)
n1
E
d
"N mln
wid
11
9 on (sp,
un) 100000
96
n
(sh, an)a
((10) (sn;an)
E >n-(PP mla) un mile
Jan (s 3000
(sn, az) (sn,an) + dx!
2h >k-(P*. mln "g-ie
96 (r(T) 9
where (i) follows from Lemma | and we use tr(A) to denote the trace of any matrix A. Hence,
Blmh (sn, (sn, an) 8h an-1
1 an an-1)ll(w 1- V /9Ka + X,
"A -PP Oi 1017
S mn
an-1) (-))
where the last inequality follows from that an = 5a, and the definition of in In addition, for = 1, have
E a P
{(m
(s1.
10
K E aimll
(s1,
s VKG.
and
(0) b(s)
E aiaaa
S āny
K
|1 a1)ll (o
arall
S 38n v
K
|121 (1)),,
airill
25Ka3d
dk
< 3/
15
A.2 PROOF OF PROPOSITION 1
Equipped with Lemma 3 the following proposition provides an upper bound on the difference of value functions under the estimated model p(n) and the true model p* for any given policy ı and reward r.
Proposition 4 (Restatement of Proposition D. For all E [N), policy and reward r, given that the event & occurs, we have
VB..r VPcm)r
Ve pin), + / KSn.
Proof. Step 1. We first show that Vp.,- - V p Vi plo)
17

---

Published as a conference paper at ICLR 2023
Recall the definition of estimated value functions V/ pow), (5h) and pin (sn, an):
pio (sh, an) = min { I, rn and + ((m) ir hi,, (5h, (sh, V h,pew,, (5h) = E Pho (ohs (shsan) an)
We develop the proof by induction. For the base case V
H + 1, we have
"H4,, (SH+1) V , (sH+1)
V (to, (10 (SH+1).
Assume that
Va 14n, (56+1) Vtt,p (8h+1) S Va
((n), f(n) (sh+1) holds for any 8h+1.
Then, from Bellman equation, we have,
Pon (sh, an) Qh,
p(n) VM,VV,,,,, h,P()), (sh, an) - Ph V Vht1, (sh.
p(")) (V - V Vi+ti.p..)) (sh+ and + (P/m)) P/ Vr, and
(0)
S min (1 (1,,))(sn))))) P((n) V h(1, Vh WIt,p (1th)aa))
(a)
min { {1, (sn,
1h (sn; an) + P/ vx
1,
P(n), yon for (sh, as)
= - pinn, stoo (8h, an),
(15)
where (i) follows from the (action) value function is at most 1, and (it) follows from the induction hypothesis.
Then, by the definition of क. h.n (sn), (sn), we have
V , (8h) (8b) -V (s()
EB @h (shs an) - E (sk. an)
E
O Ti, (on, an)
QT. (8n, (n))
2g Bens, find (9h: an)
V. Mep (sh),
where (i) follows from Equation 15) Therefore, by induction, we have
Vp..r
Vi
VT
Step 2. Then, we show that V ((m)
VKG.
By Equation (I]) and the fact that the total variation distance is upper bounded by 1, with probability at least 1 - 8/22, we have
dn-1) | -رر .1) 1)
((m) (sh, (sh, an)ss-)
9h-1
<
E 0h-1
min On (h-1) (sn
1sh-
wh 2.
p(n)
(16)
18

---

Published as a conference paper at ICLR 2023
Similarly, when h 1,
E A (s1, S V
K E anil
(s)))))
2
S VKS
(17)
a a
Based on Corollary
1 Equation
(66) and On 5 we have
(m) h
(sp; un)
84
E
min
((m) (on, an) lo0g an
IE p(n),
(18)
For the base case h = H, we have
Vx
H,ppen (tn)
fl (sH (ss SH-1
p(
p()) x
(n) DH
(8H-1, -1111 1)18
1, E @0 (5 1, sH-1 }
S min
# 110, (81-1).
oe Assume that Eptri 1+1, Awm, yor (sh+1) sn Jensen's inequality, we obtain
S V h.Pin jen) (8n) holds for step h + 1. Then, by
E
pin, (on) (sn) 8h-1
p(n),
min
(sn, ((m) yrn
1+1, p(w), (sh; an) sn-1
M
(8h-1: 0n an-1) ( + E
(n)
Vr New peaz (8n+z) 8h
min
1, n
E Pin),n
sh-1
pin),
min
1, E
(1-1
(sh
+
E p(n),
vin A.Pin), (sn) sn-1
min
1,
-h pow Ben (8--1,
P0n joo (Sh-1),
where (i) follows from Equation 18). and (ii) is due to the induction hypothesis. By induction, we conclude that
Vr 2,pen pim) (82) 81
V
pin) = E (( (51 (s1, + E
p
S KGr + Viren pen)
Combining Step and Step 2, we conclude that
VI VÆ. Vitm)rl S KVK +
p(o) Blm)
19

---

Published as a conference paper at ICLR 2023
A.3 PROOF OF PROPOSITION 2
The following lemma is key to ensure that RAFFLE terminates in finite episodes.
Proposition (Restatement Proposition 2 Given that the event occurs, Ç log ( the summation of the truncated value functions Vi, policies [Tn]ne[NN] is sublinear, i.e., the following bound holds:
pin) jow) under exploration
Σ MEIN]
TTh p(a) + +KKG S 3322Hd 22CHd/Baa(BB K)N.
Ba
Proof. Note that V
< 1 holds for any policy æ and h E [H]. We first have
pin, ((n)
The VTn
tr
2, 2,n, ((() (81, a1) P.V... 2, Blo) (s1
(in),
Pa, Bin)
m/.
(P/m)) pi)
D
2,A(m) B(o) (01, + PT (Trn 2,P100,80m
V
2, 8(m) (01,
Wn
E
(( (s (81.01) PP
2,P(m),b(m))
Vaa 2,P (i0)
nm
....
H (m) (sn, an)
< ((s sta)
= V
Te which implies Pons,Bto
I VPa
R p(sj(s)
Applying the Equation Tie
and Equation
14
we obtain the following bound on the value function
p=p.o 800
Bn (n (hh) an)
Vin Pa pan)(n)
E 2y (an) By met
h=1
H
dK
Yo 100 ,-) + 15a
" pepp xmn) any 1 mi
7
h=
dK
E P mer) metter
In 11o (n)laan (ooo)))
+ 15an V n
Similarly, we obtain
TTa
( (shn an)
Pr. pr,
"A (ppp)a)) and ayoup
h=1
an (8h- (1 1 KKn
" ( vn)
h=2
an (l0h an (a)):oo)a))))
KKn.
ayonn
20

---

Published as a conference paper at ICLR 2023
Then, taking the summation of V (r,88n)+ over
[M) we have
Van pe, ((n))))) + VKGn
NEIN]
dk
Σ
15an
+2 E VKG +
Σ
In (sk mel h.g.)
72
(p
nE[N]
DE[N!
DEIN] hm
ox me
On (8p an an)l
(
N h=1
H
17aN VdKN IN
N
Σ
[1ot(sh and an)
h=1
NEIN
H
+ ON
KN
Σ nE[N]
E mil
an)
h=1
(8))
232dK(K + +)N 1 V 455 + +FDV
H1 HV + dWdKKNG
$ 32(Hd) (dK KKN,
where (1) follows from Cauchy-Sehhe -Scwarrz inequality and importance sampling, and (it) follows from Lemma 10 Hence, the statement of Proposition 5] verified.
A.4 PROOF OF PROPOSITION 3
Based on Proposition we argue that with enough number of iterations, RAFFLE can find p satisfying the condition in line 15 of Algorithm I
Proposition 6 (Restatement of Proposition 3 Fix any E (0, 0. Suppose the algorithm runs forsfor N 2 BaH &' K(d +K) log" (2)3) d K'd K)/(8e)) iterations, with probability at least
1- 8, RAFFLE can find an n. S. N in the exploration phase such that 2V pinn) gines + +2 Kre S In other words, Algorithm can output P p(mm) satisfying the condition in line 15. In addition
VT-.r
V pinc),r c/2
Proof. We show that the algorithm terminates by contradiction. If it does not stop, applying Propo- sition 5 we have
eN/2 < Σ
(V p(n), ((n) + /KGr)
nE[N]
32 22CHd
32cHd/BBKK(d B
Therefore,
2 3g H2 HK KYÇ²
N
2
Recall <
NCN log (2 V NH/8). Using the fact that S clog2 = s
4clog"(an,)) (an,e) Vc 2 e² 1, R+ + be concluded that
N < 2 2 + K) log
which is contradiction.
21

---

Published as a conference paper at ICLR 2023
Therefore, there exists an or O p(n.) satisfies
such that pe
2V
The pines blore
+ BVKKn S.
Combining Proposition 4 we finish the proof.
Proof of Theorem Recall that is is the output of RAFFLE in the nı-iteration. Then, by Proposi-
tion 6
Vp., Vp..
S pa, Vi + t
(
V per - -V + +e/2
$ c/2 + e/2
=,
where (i) follows from the definition of न. The number of trajectories n.H is at at most
HAK + log" (100010 | ( ())())2))
B PROOF OF THEOREM 3
In this section, we adopt the same notations as in Appendix A The following lemma provides an upper bound for the estimation error of any learned model from the true model.
Lemma 4. Fix E (0, 1). for any [H), N+, any policy n, with probability least 8/2,
(sn, an) $ 22/ $2/KGn + 2v
E *A (P*.)
Proof. Recall that f/rh (s, a)
P) (-|s,a) -
TV'
Fix any policy n, for any h 2,
have
E
@h,pien,amm (s),
E
[p(m) vr
h. B80 (sh- an-1)
(N))(P()),,) ()) dn
" (n(pt)) .yy "p
min (sh-1,,, an-1)
p(o) Dit
pin, Blo)),io) (8h-1, (92-1,0h-1) 0n-1)
*h-1 ((((n) m))
[ (Mm,1,Mm,,B00) (3h-1, an-1)
E *h-1m(P(m) m)) me
< (on (1),
aI
V
(19)
P(n) B(a)
22

---

Published as a conference paper at ICLR 2023
Hence, for h 2 2, we have
(2)
E an)
((() h (8h- (sh- an (
(pm(P(R),m) ,m) aya
((((()) my
1 100 Bin) (9h-1,
E
(-(b()) ob- 24-1
VT p(n), Bin)
(20)
where (7) follows from Equation 18
(ii) follows from the definition of 2 1, (lw) (8n 1, ah)
and (tii follows from Equation 19
E Bn(ps,a) (m) S
f/m (sh, (sh, an) +
E
(sh.
E
(shn
ah >) "p
(p) an
"A m m)) o
(V P(n) poo),BiM) Bim) + VKGn)= + (VKGn + (pans,bcm)
(18 S 2 KKKGn -2
B(M)
where the first term in (1) is due Equation 20 and Equation (14). the second term in (t) is due to Proposition 4 and (wi) follows from the definition of kn-
Proof of Theorem By Proposition 6. let n. = 11 +K) Kylog'
2²
with than H. H trajectories, RAFFLE can learn model p., bonus òe and policy n, the n. -h iteration satisfying 2V + 2 KKn,,
222KKnn S e. Then, following Lemma we have
p.,8-
[1 an) (-1ss,, s <2/KGn + +2V p.. Se.
C PROOF OF THEOREM (LOWER BOUND ON SAMPLE COMPLEXITY)
C.I STEP I: CONSTRUCTION OF HARD MDP INSTANCES
Our hard MDP instances is inspired by Domingues et al. 2020 for tabular MDPs. However, the lower bound in Domingues et al. 2020) requires S 2 K, where S, denote the cardinality of state and action space respectively. Our hard MDP instances remove the assumption that S > K by constructing the action set with two types of actions. The first type of actions is mainly used form a large state space through a tree structure. The second type of actions is mainly used to distinguish different MDPs. Such a construction allows to separately treat the state space and the action space, so that both state and action spaces can arbitrarily large. We then explicitly define the feature vectors for all state-action pairs and show our hard MDP instances have a low- rank structure with dimension d S. In a nutshell, we construct family of HdK MDPs that are hard to distinguish in KL divergence, while the corresponding optimal policies are very different as shown in Figure !
First, we define reference MDP Mo as follows. We start with the construction of the state space and the action space A.
Let S = {sw, So: Sg. $p} Uve()).eR--nf(+ where Swi So, Sg and Sp denote 'wait- ing state', 'outlier state', 'good state', and 'bad state', respectively. The states in UiED)) form binary tree, where sfu denotes the j-th branch node of the layer i.
23

---

Published as a conference paper at ICLR 2023
4
S
#
s
a²
8
2.1
s
D.
D.2
So
65
S
S
Figure 1: Hard MDP instances.
Let A (aw, a a²) Uo. where aw denotes 'waiting action', a a² two unique actions that form the binary tree, and Aol = K 3 3.
Then, the transition probabilities of Mu specified through the following rules.
The initial state is the waiting state Swe
If the agent takes the waiting action aw before time step H, waiting state sw stays on itself. Otherwise, transits to the root state s2.1 of the binary tree.
Mathematically, (s and] 1 and [s |8m: a] = au or h> H-
When i < D, for states ghd in the binary tree, we have the following transition rules:
If the agent takes actions or ghd deterministically transits to its children or 112), respectively.
1++1,2j-1 1
Mathematically. P - J a = 1, and 1[s 1,55, a²] = 1.
If the agent takes any action other than a² or a², the agent will reach the outlier state
8g-
Mathematically. Palsalsti, a] = Va * a..aa a²
Leaf state s P.j uniformly transits to good state sy and bad state no matter what action the agent takes.
Mathematically, Ps J. a) = P3 a] _ Va A.
Good state 8g. bad state and outlier state are absorbing states.
Now, define the features follows, which are S-dimensional vectors.
24

---

Published as a conference paper at ICLR 2023
(w; Gw) 0,000,0000000000)
[4 (. = (1 (IgH. Slg.a.0,00ss,004p M4
Φn (S a) (0.1,0g=5.0.0),00
(1 1,,0m,1,,00,,,,000.
D
(nt )²) = (0.0....12//-222000000..... 1,2 @n a) = (0.0000-=,1100,,,,
Nh (sh..) =
D
Mi (s.) = (0.0,03--,1,0,00)
sD
@n a) (0,0,0s-s:0,
(%(s,) (0.0,0s-s,0,100)
8g. 8b
(n a) = (sg), (n. = (56)
(%+)) = (0,0,0 0,0,0sss.0,011)
(5. 0) (so):
where 0s-5 E R denotes the S - 5 dimension vector with all zeros and eij E R denotes the one-hot vector that is zero everywhere except the (2 + - 2)-th coordinate. Here 2 S i D,1 S 2
#+D) x [20-1] x A. we define an MDP Mik- ,e", ,a*) through Mo- Specifically, the only difference of M(s ,0,,0") from MM is that the transition probability from the leaf state S. and action a" to the good state 8g increases cg, i.e. PA- [sg lgl .a'] 1 + ED> and Ph [s aa]
For each (h",e", "","""a" E (1+D.
co. where co will be specified later. We note that the features of Mah ,et aa)) are the same those of Mo except that (s aaa) = (0, Os 5, 0, t0. -n)-
We remark here that the cardinality of A can be arbitrarily large, so that the resulting lower bound will hold for both d < K and d > K regimes. In addition, although in our hard instances, d S, it is straightforward to generalize it to the regime with S > d if we set the outlier state S. to be a set of outlier states S..
Definition of reward: the reward can only be attained in two special states: the good state 8g and the outlier state so at the last stage H.
1
VS A, a,) (-==s,,h=1) m) + 2 (a=a,
and (s, a) still belongs to [o, 1].
C.2 STEP 2: ANALYSIS OF HARD MDP INSTANCES
Proof of Theorem 7
Let co = 2c.
For any MDP MM*,a*, the optimal policy is to take action aw to stay at state sw until stage h - D, and then take the corresponding action to the only state 5 D,E" at stage h². At state sD the agent takes the only optimal action a*. The optimal value function V Maaa.r function of the output policy of Alg is given by
+ EO. and the value
V Man
+ < P
Man.rra-)) [8k. = = ah- ...-= a"],
(21)
pal- where MMa- at)) is the probability distribution over the states and actions (sh, an) following the Markov policy म. in the MDP M(k..r.a)) We remark that the reward of outlier state are specially designed to be 1/2 at stage H to make Equation (1) hold for policy म, falling into the outlier state.
Hence,
V Minaeaa
V Missera c0(1 MMa (ay) [sh- = ah , =0 a²l)
p
pt. 2c(1 Min (a) [sh. sh", ,a,- a")).
and
V Mar.e..r
P SCC
Mar. [sh = 5 an. = a N 2
(22)
The transitions of all MDPs are the same when the leaves states are reached. We define the event
Puitrrerr [sk- =t" a. , a') 2 1 25
p
25

---

Published as a conference paper at ICLR 2023
From Equation 22 the event is equal to the event {V Mar (a
V Men
S As result,
(e Ma
P e ())) fiv (r,))
PP ( (( .0t) (Vuureea
V
S 2 6.
Recall that (M ( (r,)) = En= (==)===)== .at} such that Exx (' aa) S T. This
inequality holds because the agent is likely to fall into the outlier state Sn- We denote Po and Eg to be with respect to Mo-
Now, invoke intermediate result in the proof of Theorem 7 in Domingues al. 2020 to conclude that
Eo Nix ra,-) ((a))))))))
1
(1 (1-Po) Po (((7)))010))))) log
[
log
16E² Summing over all (h", e" ,a)), have
[o[T] 2 (h",t",a")
Σ
Ro (TT
1
HLK-
E (b (a"))";aa))
Po e (e log(()
HLL MoLKlog? 2
(23)
1662
Notice that
Σ (h',e",a")
Po (a Eo
Σ (h",t' (a))
1 ( te an ax==a-]23))
51.
(24)
Substituting Equation
24
into Equation
23 yields
Eo[T]
( 1) log(5) HILK log2
16e2
2
HLK HLKlog( ()
32c²
where the fact that 8 < 1/16. With the assumption of K 2 23,5 we have S. Taking
H # and with the assumption of D < H/3, we have
HdK
E [7]
log (
Then following the analysis similar that for Corollary 8 in Domingues et (2020) with proba- bility at least 1 8, the number of iterarion is at least
HdK
log 1
Ω
D ALGORITHM 2 REPLEARN AND PROOF OF THEOREM
We first present the full algorithm in Section 6 below as Algorithm
D.I SUPPORTING LEMMAS
We first show that Q P. can approximate Qp- hrr well over distribution (g, hee by following
two lemmas.
Lemma 5. Given any 8 E (0,1). Let P be the oulput of Algorithm for any policy and rewards
r. with probability at least 1 8. we have
QF (Sh, and - Q 2. an)
E (s..0))~(P.))
Sc
E ((P
, (Sh, an) - Q7 -,
prr ( ar
Se.
26

---

Published as a conference paper at ICLR 2023
Algorithm 2 RepLearn: Representation Learning in Planning Phase of RAFFLE
1: Input: Sample size Ny. state-action pair distributions {qn fg special designed reward func-
tion {ra*() he and policy {x* hejm) and estimated transition kernel P from the output of Algorithm model class: Φ.
2: Initialize 0/
3: for 1, ................................ , do
4: for h = ,...., H do
5: for / 1..... T do
6:
Choose (s ax a...)
Uh and add (sh 2 a. dataset D Di,
1,t U (8h ,0 ).
7: end for
8: end for
9: end for
for h
H do Learn (Øh, 2h ................................ 1n7)) in Equation (4
12: end for
Output: / (dn hemm--
Proof. Define Jn (sn,an)
= (r an) P*(,|on, an) an>lyy and f is a collection of all Su. i.e.
{ hE[H]
B (sn p (Sh, - Qp (5h anl
E (snaamm(et)) (0 an)
(1)
S Vp.
(m))
e/2,
(25)
where (i) follows from Equation from the proof of Theorem
115 (di) follows from the definition of and Q and (iit) follows
Then for the second inequality,
[ (om, - O7 (sh, -, am))
E (shaa))(PP
Opr,er an) 27, pr,, (sh,
E (sh,o))((P,n)
(th.an)) (or -Bren 1 (on (mn.0))]]
(8)
$ c/2++V/ V
(1)
S.
where the first term in (i) is due to Equation 255 and Equation (44 the second term in (f) is due to Step in Proposition 4 and (fi) follows from the proof of Theorem
Lemma Given any d,e (0, 1) and the output of Algorithm 7 under Assumption 2 for
any policy π and rewards let the input distributions VH bounded with constant CB. Denote Cuin
Cp
Then with probability at least 1 5, for each h E [H]
Quin
(een | (sh, an) - QT h.r (Sh, an) < eCCinn
27

---

Published as a conference paper at ICLR 2023
Proof. First, together with Assumption for any (s, E S A, have (s, 1 Cmiu Ph (s, a), then
Bosneanman (o an) 27 -. (sh, an)
/ on (sh, and Mo and O7 Vp (sh an) dsndan
( I Cmiu Ph (sh) @7 pp an) 2% P,h,r (8h, ax) dondan
(sh.
Cmin
E (PP
eCmin
where (i) follows Assumption 2 and (it) follows Lemma
The lemma above shows that for any reward r. 27 h. (Sh, an) can be the target Qp- Jur (Shr an)) when (sh, a%) are chosen from given distribution Qh-
We then show that for any h, when reward r is set to zero at step h, QP.,,.. Ahr has linear structure w.r.t di.
Lemma 7. For any & E [H], policy and given (sh)ah) E E x A, given any such that r is set to be zero at step h, i.e. Th = O, then Q7- hrr (Sh, an) is linear with respect to o" (Sh, a²), i.e. there exist a wf, such that Q7+ fhr,,,, (on, an) = (0% (Sh, wf,)
Proof.
(pr,,rr(h, an) Tn(sh, an) + Ba (V,,, (ot))) anl
= / (*s+21lsh VF. s+1, (8h41)dsh+1
(snn an). / ( be, (Sh41
(St4 )dsutz
(0) un), wh))
where W = 14 1ds
D.2 PROOF OF THEOREM 4
Proof of Theorem 4 We allow & and T to apply to all the samples in a dataset matrix si- multaneously, i.e. (DN (0n(5) ....., (...)........... (5 "")) E NNxx and Q7',,,,. (D)) (.t) ( )..., (s ( E RN,, After we estimating on and then taking as known, from the property of linear regression Equation (4, for any reward function ,hit any policy oe we got
TH = (A (D) (n TT Op (D, tn (D) () = si Pp (D) Nyit
28

---

Published as a conference paper at ICLR 2023
"TT, represents the projection
where P. (D operator to the column spaces of (D), Then
p1 2 P (D) /
T (D N )00h ( 191) ||?
Σ
AETT
(0)
(E[T]
@p , Q p nhe .0 aitt)
(T(
(it)
[ ,ay 27' ( a ( )
2
TN₂ log
2
E Esnnan Jmag 1E|T]
(11))
TN log 2
N5
Been,on)~an
(sh
not. (ahh)) ap :ah +1
2
TE(T]
(iv)
TN, log 2 eC NyT
(26)
2
where (i) follows from minimality of (0% heepyy and (wi heepy, (ii) follows Hoeffding's incquality, (iit) follows that Qp- Op- (8) 0%
Q pn (6% a n.t) S 1 and (iv) follows Lemma 6
As result:
Σ
pl
(D (0 (D No Juh"""
(E(T]
Σ rE[T]
pL (D) )il) Mppe (Ditsit))
|ppt
11 Fi (D) (D). /|?
(D por (D.) 27 Pip (DN)..)) {?]
$ Σ
+ pd
(E[T]
Σ ECmin
TN, log
ofp 11 1 (D) Q3 O.A.n (D)1,
2
TE(T)
tE(T]
(12 < 2cCmin NyT+ +
2 2TN log
8*
where (i) follows from | - (1(4) ||l|| and (ii) follows from that 0 (P-
Nnye)) < 1 and
the process to derive Equation (26).
29

---

CDT Published as a conference paper at ICLR 2023
Table 1: Comparison among provably efficient RI algorithms under low-rank MDPs
METHODS
SETTING
SAMPLE COMPLEXITY
FLAMBE
LAGARWAL ET AL.
2020
LOW-RANK MDP
Ō R
MOFFLE
MODI AL.
2021
LON-KOOWEGAATVE--ANN
HOMER
MISRA ET AL.
2020
BLOCK MDP
OL )
Ō H
RAFFLE (OURS)
LOW-RANK MDP
We do include reward-know RL under low-rank MDPs in this table and only focus reward-free RL. The detailed discussion of neward-knowe RL under low-rank MDPs can be found in Section
We finally use the technique in Du et al. (202bb) to derive the super population guarantee.
2
2x 2cCmin NT + 2 log
P. "D (2 (D No.t
"
6E[2]
(T- (D ) (Ōn (on(D) TTn "n( (.
E )T 0% (DN (1 (1- (On
(D) (1 ""W
[E[
Ny(whh) TT Diyy ōn)uh,"
EE
(
2 0.9
Nr(wh, Dan (0h, w
(E[T]
-0.9,
Σ (E(T]
Dos (0%, (1) wh
2 0.9N,
Dan (oh, 2.)
0.9,,
| (01, 5n)
((w))
2 0.9, Dmn. (0h. )
d
where (1) follows from lemma B. in Du al. 2021b) and N, can be large enough, and (ii) follows from that ||AB|] 2 || min (B), where miiu (B) denote the smallest singular value of matrix B. Then
Pys (0% 1023 <
2edCmin d 0.9Cp 0.9Cp V TN,
2 log 2
+
E MORE DISCUSSION ON RELATED WORK
In this section, we first summarize the directly comparable work in Appendix E
E. LOW-RANK MDPS IN EXTENDED RL SETTINGS
Many studies have been developed on various extended low-rank models. Wang et al. (002); Uchara
al. 2022a) studied partially observable Markov decision process (POMDP) with latent low-rank
30

---

Published as a conference paper at ICLR 2023
structure. Zhan et al. (2022) studied predictive state representations model, and applied their results to POMDP with latent low-rank structure. [heng et al. (2022) Agarwal et al. (2022) studied benefits of multitask representation learning under low-rank MDPs. Huang al. (2022) proposed general safe RI. framework and instantiated to low-rank MDPs. Ren et al. (2022) studied reward-known RL under low-rank MDPs and proposed spectral method replace the computation oracle. We note that even given those further developments, our results on lower bound and representation learning are still completely new, and our algorithm design and result on sample complexity are so far still the best-known result for standard reward-fiee low-rank MDPs.
E.2 DISCUSSION ON OPTIMISTIC MLE-BASED APPROACH FOR DIFFERENT SETTINGS
There is some concurrent work also using an optimistic MLE-based approach for different settings (POMDP) (Liu et al. 2022 Chen ct al. 2022a). We elaborate the key differences between our paper and Liu et al. 2022); Chen et al. 2022a) as follows.
The two POMDP papers mentioned above study reward-known setting, whereas focus here is on the reward-free setting.
Although optimistic MLE-based approach is used in both settings, the design of exploration policy in the two settings different. Our algorithm identifies which estimated model is used for exploration policy design and then calculate the exploration policy based on bonus terms designed for the value function. However, the POMDP papers construct a confidence set about the true model and solve an optimization problem within this generic model set. Such an oracle may not be easy realize computationally.
Due the hardness of POMDP, MLE approach only guarantees that the estimated distri- bution the trajectories is close to the true one. In contrast, in low-rank MDP we study here, we show that the estimation error for the transition probabilities controlled each time step.
F AUXILIARY LEMMAS
Recall ((()) (s, a) = (.1s, (-ss,, lzv represents the estimation error in the n-th itera- tion at step h, given state and action in terms of the total variation distance. By using Theorem 21 in Aggawaa et al. (2020), we are able to guarantee that under all exploration policies, the estimation be bounded with high probability.
Lemma 8. (MLE guarantee). Given 6 E 10, have the following inequality holds for any
n, 2 2 with probability at least 8/2:
m-1
(sh, an) S nSn; where Sa = log (2 V7 1//H/6
Σ
(x))
7
(ap (ndrtA)
4N- --192-1 1
In addition, for h 1,
Σ
E [fT(s,,))) s nCn. T ((A)
Divide both sides of the result of lemma Lemma by n, and define 11, = have the following corollary which will be intensively used in the analysis.
we Corollary 2. Given 8 e (0, 1), the following inequality holds for any n. h 2 2 with probability at least 8/2:
(sh, < Ca.
an 1P m.d) an- 1. on ( (1--119--))
In addition, for h 1,
E a (((A)
(1) Sn.
31

---

Published as a conference paper at ICLR 2023
The following lemma (Dann al. 2017) will be useful to measure the difference between
two value functions under two MDPs and reward functions. We define Ps. Pr,V++1 as) = (.s,,) [V(s)] for shorthand.
Lemma 9. (Simulation lemma). Suppose Pi and P2 are two MDPs and vi, 12 are the corresponding reward functions. Given a policy n, we have,
Pi (5h) Par (sh)
H
Σ
(ri(sar, - (2(ss)aa)) + (PI.N -22 P2,n)VVH,,P. VI, an))
M/=h
11
Σ
(ri(snr T2(shr aw) + (PIN P2.w (M.n) an) [1s))
M=h
The following lemma a standard inequality in regret analysis for linear models in reinforcement learning (see Lemma G.2 in Agarwal et al. (2020) and Lemma 10 in Uehara et al. (202T)).
Lemma 10. (Elliptical potential lemma). Consider sequence of dx positive semidefinite matrices
X1, X1,..., X with Lr(X,) for all n E [N). Define Mo Aol and M, = Mn-1 + Xn Then
N
(1+
N
Σ tr (X) S det - 2 log det (Mo) S 2d 2d
dAo
If we choose any subset of the set 1X, )N 1= 1,-M)11
still get a sublinear summation as follows.
32


Reviews:
[{"review_id": "pzqXOqy71V", "reviewer": "Reviewer_kHZp", "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": "", "summary_of_the_paper": "This paper studies the sample complexity of reward-free learning in low-rank MDP and achieves sharp upper and lower bounds. The learned representations can also be reused in a different MDP.", "summary_of_the_review": "The paper achieves the SOTA upper and lower bounds for reward-free learning in low rank MDP. This enhances the understanding of low rank MDP of the RL theory community. Certain clarification of technical contribution will be super useful to further improve the quality of the paper.", "strength_and_weaknesses": "Pros:\n-The paper achieves the SOTA upper and lower bounds for reward-free learning in low rank MDP. This enhances the understanding of low rank MDP of the RL theory community.\n\nCons:\n-I think the authors may want to emphasize more on the difference between different setup in function approximation. For example, citing the Bellman rank paper [1] for the setting of low-rank MDP may not be very accurate. As I understand low-rank MDP is a special case of low Bellman-rank MDP? As a result, comparing results in different setting may involve further complications. The authors may want to elaborate that a little bit.\n- The table given is very clear. It would be even better if the sample complexities of algorithms in the normal learning setting can also be compared with the proposed algorithm. This will be helpful as it showcase the difference between the normal setting and reward-free setting.\n- It's not clear technique-wise what's the new ingredient that makes the proposed algorithm performs better then the existing ones. The novel termination rule is explained but I believe it is only part of the story. More concretely, it makes me think of the algorithm is a combination of the technique from zero-reward exploration in tabular MDP ([2][3]) and the MLE techniques in FLAMBE. Am I missing anything important here?\n- There is some concurrent work also using an optimistic MLE-based approach for different settings (POMDP), such as [4][5]. I know this is not necessary, but it would be a plus if you could also compare it with the algorithmic design and proof technique in these papers.\n\n[1] Contextual Decision Processes with Low Bellman Rank are PAC-Learnable\n[2] Task-agnostic Exploration in Reinforcement Learning\n[3]  A Sharp Analysis of Model-based Reinforcement Learning with Self-Play.\n[4] Optimistic MLE—A Generic Model-based Algorithm for Partially Observable Sequential Decision Making\n[5] Partially Observable RL with B-Stability: Unified Structural Condition and Sharp Sample-Efficient Algorithms", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "aq2pMcJOaH", "reviewer": "Reviewer_rgZQ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "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": "", "summary_of_the_paper": "This work focused on reward-free RL with low-rank MDP. The author proposed a new model-based reward-free (RAFFLE) algorithm and provided a theoretical guarantee with polynomial sample complexity. In addition, the author offers a lower sample complexity bound for any algorithm within the low-rank MDP. Finally, the author shows it is possible further to learn a provably near-accurate representation of the transition kernel.\n\n", "summary_of_the_review": "This work proposes a novel algorithm (RAFFLE) for learning reward-free low-rank MDPs and improving the sample complexity. However, it seems that the lower bound ignores an essential requirement, and the author makes a wrong claim, \"reward-free RL under low-rank MDPs is strictly harder than linear MDPs is incorrect.\" as a contribution. In addition, I am concerned about why it is possible to distinguish two representations in Theorem 4, when they correspond to the same transition kernel $P$.", "strength_and_weaknesses": "Strength:\n\nThe author proposed a novel model-based reward-free (RAFFLE) algorithm and improved the sample complexity in low-rank MDP. In addition, the theoretical lower bound also supports the efficiency of the RAFFLE algorithm.\n\nWeakness:\n\n1. The lower bound on sample complexity (Section 5.2) ignores an essential requirement that $d\\ge K$. More specifically, according to figure 1 in Appendix C, the first layer have $L/K$ different states, and the second layer has $L$ states. In this case, the number of different states is at least $K$ since the first layer should have at least one state. Thus, the dimension $d=L\\ge K$ and the hard-to-learn instance only hold when $d\\ge K$.\n\n2. Based on weakness 1, the claim that reward-free RL under low-rank MDPs is strictly harder than linear MDPs is incorrect. To make this claim, it is necessary to create an instance that shows the lower bound for low-rank MDPs is strictly larger than the upper bound for linear MDPs. However, according to weakness 1, the lower bound only holds for $d\\ge K$ (large $d$ regime), and the lower bound does not contradict the $O(d^3H^3/\\epsilon^2)$. With a similar argument, we can easily extend the lower bound of linear MDP $O(d^2/\\epsilon^2)$ to $O(d|S|/\\epsilon^2)$ with the regime $d\\ge |S|$ and mentions \"linear MDPs is strictly harder than low-rank MDPs since the lower bound to have an additional dependency on $|S|$.\"\n\n3. The comparison with Uehara et al. (2022b) is incorrect. Uehara et al. (2022b) calculate the sample complexity with the number of steps, while this work measures the sample complexity with the number of trajectories. There always exists a gap of H between trajectories and steps. Thus, the RAFFLE algorithm only improves a factor of $H$ rather than $H^2$.\n\n4. It would be better if the author could provide more intuition on why Lemma 2 holds. The author mentions it is an extension of Lemma 12 in Uehara et al. (2022b) to episodic MDPs. However, in Uehara et al. (2022b), both the triple at stage $h-1$ and state $h$ are used to estimate the model (MLE). It would be better if the author could provide more intuition as to why only using the data at stage $h$ can maintain a similar performance or the uniform choose action $a_{h-1}$ is unnecessary.\n\n5. It is strange why the RAFFLE can always find the correct representation $\\phi^*$ in Theorem 4. More specifically, even for a known transition kernel $P$, there may still exist multi-representation for $P$ in the lower-rank MDP. For instance, if we switch the first and second dimensions of $\\phi_1,\\mu_1$ and obtain $\\phi_2,\\mu_2$, then two representations will lead to the same transition kernel. Under this situation, it seems impossible to learn the actually representation $\\phi^*$ from $\\phi_1$ and $\\phi_2$.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "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": "W64xDbIUQJ", "reviewer": "Reviewer_2hcZ", "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": "", "summary_of_the_paper": "This paper proposes a new model-based algorithm, coined RAFFLE, and show that it can both find an $\\epsilon$-optimal policy and achieve an $\\epsilon$-accurate system identification via reward-free exploration, with\na sample complexity of $\\tilde{O}(H^3d^2K(d^2+K))$, where $d,H$ and $K$ respectively denote dimension, episode horizon, and action space cardinality. This significantly improves the sample complexity in Agarwal et al. (2020) for the same learning goals.\n\nWe further provide a sample complexity lower bound that holds for any reward-free algorithm under low-rank MDPs, which matches our upper bound in the dependence on $\\epsilon$, as well as on $K$ in the large $d$ regime.", "summary_of_the_review": "Please answer my major concern.", "strength_and_weaknesses": "Strength: propose a new model-based reward-free RL algorithm under low-rank MDPs. The algorithm can both find an ε-optimal policy and achieve an ε-accurate system identification via reward-free exploration, with a sample complexity of $\\tilde{O}(H^3d^2K(d^2+K))$. A lower bound on the sample complexity that holds for any reward-free algorithm under low-rank MDPs.\n\nWeakness: \n\nThe main weakness of this paper is the algorithm seems close related to [1]. Note reward-free multi-task is a more general setting than the current study and the exploration of RAFFLE seems to be very similar to REFUEL in [1] (with $T=1$). In addition, the main contribution of this paper seems to have overlap with [1]. For instance, (6) of Theorem 4.2. is a system identification result and (7) is the reward-free learning result. At this moment, I suspect the result of the current paper might be a corollary of [1]. It is likely I might be wrong, so please correct me if I misunderstood something.\n\n[1] Provable benefit of multitask representation learning in reinforcement learning, NeurIPS 2022\n\n\n", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ATK1Q2M729", "reviewer": "Reviewer_exQF", "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": "", "summary_of_the_paper": "This paper combines the REP-UCB and RF-UCRL to obtain an improved sample complexity for reward-free reinforcement learning under the structural assumption on Low-rank MDPs.\n", "summary_of_the_review": "For me the combination of REP-UCB and RF-UCRL is not so striking. The most interesting part for me is the lower bound, where the authors argue that we cannot avoid the dependency on the number of actions if we need to learn the representation. However, as the proof lacks some rigor, I'm not sure if the proof is correct. For me, intuitively the lower bound is correct, as for the instance the authors consider, if we know the exact feature, we already know which action is different. I believe the authors need to make the proof of the lower bound much more formal. The analysis in Section 5.3 should be much more well-motivated.\n", "strength_and_weaknesses": "### Strength:\n* The idea is simple and easy to follow.\n* The interpretation of the lower bound is interesting.\n\n### Weakness:\n* The proof of the lower bound lacks some rigor.\n* Section 5.3 is not well-motivated under the context of RL. In other words, I don’t quite understand how this divergence will help to analyze the performance of reinforcement learning.\n", "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": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster