I MPROVED S AMPLE C OMPLEXITY FOR R EWARDFREE R EINFORCEMENT L EARNING UNDER L OW - RANK MDP S
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
While reward-free RL has been well studied under the tabular setting with minimax optimal sample complexity being achieved, theoretical study of reward-free RL with complicated function approximation is still limited.
In this paper we focus on reward-free RL under low-rank MDP models, which capture the representation learning in RL.
We propose a new model-based algorithm, coined RAFFLE, and show that it can both find an ϵ -optimal policy and achieve an ϵ -accurate system identification via reward-free exploration, with a sample complexity of ˜ O
( H 3 d 2 K ( d 2 + K ) ϵ) , where d , H and K respectively denote the representation dimension, episode horizon, and action space cardinality.
This significantly improves the sample complexity of ˜ O
( H 22 K 9 d 7 ϵ) in Agarwal et al.
(2020) for the same learning goals.
We further provide a sample complexity lower bound of ˜Ω( HdKϵ 2 ) that holds for any reward-free algorithm under low-rank MDPs, which matches our upper bound in the dependence on ϵ , as well as on K in the large d regime.
Comparing this lower bound for low-rank MDPs with the upper bound for linear MDPs in Wang et al.
(2020), it implies that reward-free RL under low-rank MDPs is strictly harder than linear MDPs.
Finally, we complete our study by reusing RAFFLE to learn representation.
We estimate the representation individually with only access to the learned transition kernels from RAFFLE and without interacting with true environment, and then theoretically characterize the closeness between the learned and the ground truth representation.
The learned representation can be further used for few shot RL as in supervised learning (Du et al., 2021b).
I NTRODUCTION
[Table]
reinforcement learning, recently formalized by Jin et al.
(2020b), arises as a powerful framework to accommodate diverse demands in sequential learning applications.
Under the rewardfree 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 reward during the planning phase.
The learning goal can be to find an ϵ -optimal policy, to achieve an ϵ -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 are taken into consideration over a single environment, such as safe
RL (Miryoosefi & Jin, 2021; Huang et al., 2022), multi-objective RL (Wu et al., 2021), multi-task
RL (Agarwal et al., 2022; Cheng et al., 2022), etc.
Studies of reward-free RL on the theoretical 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 Jin et al.
(2020a); M´enard et al.
; Kaufmann et al.
; Zhang et al.
For reward-free RL with function approximation, Wang et al.
(2020) studied linear MDPs introduced by Jin et al.
(2020b), where both the transition and the reward are linear functions of a given feature extractor, Zhang et al.
(2021b) studied linear mixture MDPs introduced by Ayoub et al.
(2020), and Zanette et al.
(2020b) considered a classes of MDPs with low inherent Bellman error introduced by Zanette et al.
(2020a).
In
[Table]
[Table]
[Table]
[Table]
this paper
[Table]
we focus on reward-free RL under low-rank MDPs, where the transition kernel admits a decomposition into two embedding functions that map to low dimensional spaces.
Compared with linear MDPs, the feature functions (i
[Table]
e., the representation) under low-rank MDPs are unknown, hence the design further requires representation learning and becomes more challenging.
Rewardfree RL under low-rank MDPs was first studied by Agarwal et al.
(2020), and
the authors introduced a provably efficient algorithm FLAMBE, which achieves the learning goal of system identification with a sample complexity of ˜ O
( H 22 K 9 d 7 ϵ) , where d , H and K respectively denote the representation dimension, episode horizon, and action space cardinality.
Later on, Modi et al.
(2021) proposed a model-free algorithm MOFFLE for reward-free RL under low-nonnegative-rank MDPs (where feature functions are non-negative), for which the sample complexity for finding an ϵ -optimal policy scales as ˜ O
( H 5 K 5 d 3 LV ϵ 2 η ) (which is rescaled under the condition of (cid:80) Hh =1 r h ≤ 1 for fair comparison).
Here, d LV denotes the non-negative rank of the transition kernel, which may be exponentially larger than d as shown in Agarwal et al.
(2020), and η denotes the positive reachability probability to all states, where 1 /η can be as large as √ d LV as shown in Uehara et al.
(2022b).
Recently, a reward-free algorithm called RFOLIVE has been proposed under non-linear MDPs with low Bellman Eluder dimension by 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 further investigates reward-free RL under low-rank MDPs to address the following three important open questions:
• Can we design a reward-free algorithm for low-rank MDPs that improves the sample complexity of FLAMBE in Agarwal et al.
(2020) and MOFFLE in Modi et al.
(2021), and achieves the goal of system identification as in Agarwal et al.
and
finds near-optimal policy as in Modi et al.
• What is the fundamental limit on the required sample complexity to learn a low-rank MDP?
While this question can be answered by a lower bound on the sample complexity for low-rank MDPs, none of the previous studies on low-rank MDPs establish such a lower bound.
• Is it possible to reuse the learned representation (i e., features) of the transition kernels in other RL environment?
Intuitively, the features can be utilized if learned accurately.
However, previous studies on low-rank MDPs did not provide any estimation accuracy guarantee on the learned representation, but only on the transition kernels.
1. M AIN C ONTRIBUTIONS
We summarize our main contributions in this work below.
• We propose a new model-based reward-free RL algorithm under low-rank MDPs.
Our algorithm features carefully designed exploration-driven reward functions, exploration policies, and exploration termination criterion, and hence enjoys provably much better sample efficiency than the existing reward-free RL algorithms under low-rank MDPs.
• We show that our algorithm can both find an ϵ -optimal policy and achieve an ϵ -accurate system identification via reward-free exploration, with a sample complexity of ˜ O
( H 3 d 2 K ( d 2 + K ) ϵ) , which significantly improves that of ˜ O
( H 22 K 9 d 7 ϵ) in Agarwal et al.
(2020) to achieve the same goal.
Our result also improves the sample complexity of ˜ O
( H 5 K 5 d 3 LV ϵ 2 η ) in Modi et al.
(2021) in three aspects: order on K is reduced; d can be exponentially smaller than d LV as shown in Agarwal et al.
(2020); and no introduction of η , where 1 /η can be as large as √ d LV .
Further, our result on reward-free RL naturally achieves the goal of reward-known RL, which improves that of ˜ O (cid:16) H 5 d 4 K 2 ϵ(cid:17) in (Uehara et al., 2022b) by a factor of Θ( H 2 ) .
• We provide a lower bound ˜Ω( HdKϵ 2 ) on the sample complexity that holds for any reward-free algorithm under low-rank MDPs, which matches our upper bound in terms of the dependence on ϵ as well as on K in the large d regime.
Comparing this lower bound for low-rank MDPs with the upper bound for linear MDPs in Wang et al.
(2020) implies that reward-free RL under low-rank MDPs is strictly harder than linear MDPs.
• We further design an algorithm to exploit the learned transition kernel from our reward-free algorithm to further learn a 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 work that provides such guarantee on representation learning in reward-free RL.
R ELATED W ORK
Reward-free RL.
While various studies (Oudeyer et al., 2007; Bellemare et al., 2016; Burda et 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., 2020) proposed exploration algorithms for good coverage on 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 on the sample complexity.
For tabular case, several follow-up studies (Kaufmann et al., 2021; M ´ enard et al., 2021) further improved the sample complexity, and Zhang et al.
(2020) established the minimax optimality guarantee.
Reward-free RL was also studied with function approximation.
Wang et al.
(2020) studied linear MDPs , and Zhang et al.
(2021b) studied linear mixture MDPs .
Further, Zanette et al.
(2020b) considered a class of MDPs with low inherent
Bellman error introduced by Zanette et al.
(2020a).
Chen et al.
(2022b) proposed a reward-free algorithm called RFOLIVE under non-linear MDPs with low Bellman Eluder dimension.
In addition,
Miryoosefi & Jin (2021) proposed a reward-free approach to solving constrained RL problems with any given reward-free RL oracle under both the tabular and linear MDPs.
Reward-free RL under low-rank MDPs.
As discussed in Section 1, reward-free RL under lowrank MDPs have been studied recently (Agarwal et al., 2020; Modi et al., 2021), and our result significant improves the sample complexity therein.
When finite latent state space is assumed, lowrank MDPs specialize to block MDPs, under which algorithms termed as PCID (Du et al., 2019) and
HOMER (Misra et al., 2020) achieve sample complexities of ˜ O
( d 4 H 2 K 4 min( η 4 γ 2 n ,ϵ 2 ) ) and ˜ O ( d 8 H 4 K 4 min( η 3 ,ϵ 2 ) ) , respectively.
Our result on 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) proposed 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 ingredients 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 eluder dimension (Jin et al., 2021), which can specialize to low-rank MDPs.
However, those algorithms are computationally more costly as remarked in Uehara et al.
(2022b), although their sample complexity 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 a slightly different low-rank model, and obtained a problem dependent regret upper bound.
Low-rank MDPs in extended RL settings.
After this work was initially submitted for potential publication early this year, many concurrent studies have been developed on various extended RL settings with low-rank MDP models.
Wang et al.
; Uehara et al.
(2022a) studied partially observable Markov decision process (POMDP) with latent low-rank structure.
Zhan et al.
(2022) studied predictive state representations model, and applied their results to POMDP with latent lowrank structure.
Cheng et al.
; Agarwal et al. (2022) studied benefits of multitask representation learning under low-rank MDPs.
Huang et al.
(2022) proposed a general safe RL framework and instantiated it to low-rank MDPs.
Ren et al.
(2022) studied reward-known RL under low-rank MDPs and proposed a spectral method to replace the computation oracle.
This work focuses on the most fundamental setting in the study of low-rank MDPs, and our novel algorithm design and theoretical analysis can be adopted to further improve the sample complexity in many of these studies.
P RELIMINARIES AND P ROBLEM F ORMULATION
Notation.
For any H ∈ N , we denote [ H ] := { 1 , . . .
, H } .
For any vector x and symmetric matrix
A , we denote ∥ x ∥ 2 as the ℓ 2 norm of x and ∥ x
∥
A := √ x ⊤ Ax .
For any matrix A , we denote ∥
A ∥ F as its Frobenius norm and let σ i
( A ) be its i -th largest singular value.
For two probability measures
P, Q ∈ Ω , we use ∥ P − Q ∥ TV to denote their total variation distance.
3. E PISODIC MDP S
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 a finite action space with cardinality K ; H is the number of steps in each episode; P : S ×
A × S → [0 , 1] is the time-dependent transition kernel, where
P h ( s h
+
| s h , a h ) denotes the transition probability from the state-action pair ( s h , a h ) at step h to state s h +1 in the next step; r h : S ×
A → [0 , 1] denotes the deterministic reward function at step h ; We further normalize the summation of reward function as (cid:80)
Hh r h ≤ 1 .
A policy π is a set of mappings { π h : S → ∆( A ) } h ∈ [ H ] , where ∆( A ) is the set of all probability distributions over the action space A .
Further, a ∼ U ( A ) indicates the uniform selection of an action a from A .
In each episode of the MDP, we assume that a fixed initial state s 1 is drawn.
Then, at each step h ∈ [ H ] .
the agent observes state s h ∈ S , takes an action
a h ∈
A under a policy π h , and receives a reward r h ( s h , a h ) (in the reward-known setting), and then the system transits to the next state s h + with probability P h
( s h
+
| s h , a h ) .
The episode ends after H steps.
As standard in the literature, we use s h ∼ ( P, π ) to denote a state sampled by executing the policy π under the transition kernel P for h − 1 steps.
If the previous state-action pair ( s h − 1 , a h − 1 ) is given, we use s h ∼ P to denote that s h follows the distribution P h ( ·| s h − 1 , a h − 1 ) .
We use the notation E ( s h ,a h ) ∼ ( P,π ) [ · ] to denote the expectation over states s h ∼ ( P, π ) and actions a h ∼ π .
For a given policy π and an MDP M = ( S , A , H, P, r ) , we denote the value function starting from state s h at step h as V πh,P,r
( s h ) :
=
E ( s h ′ ,a h ′ )
∼ ( P,π ) (cid:104)(cid:80) Hh ′ =
h r h ′
( s h ′ , a h ′ ) | s h (cid:105) .
We use V πP,r to denote V π 1 ,P,r ( s 1 ) for simplicity.
Similarly, we denote the action-value function starting from state action pair ( s h , a h ) at step h as Q πh,P,r
( s h , a h ) :
= r h ( s h , a h ) + E s h +1 ∼ P (cid:104) V πh +1 ,P,r ( s h +1 ) | s h , a h (cid:105) .
We use P ⋆ to denote the transition kernel of the true environment and for simplicity, denote E ( s h ,a h ) ∼ ( P ⋆ ,π ) [ · ] as E ⋆π
Given a reward function r , there always exists an optimal policy π ⋆ that yields the optimal value V ⋆P
⋆ ,r = sup π V πP ⋆ ,r .
3. L OW - RANK MDP S
This paper focuses on the low-rank MDP (Jiang et al., 2017; Agarwal et al., 2020) defined as follows.
Definition 1.
(Low-rank MDP).
A transition probability P ⋆h : S ×
A → ∆( A ) admits a low-rank decomposition with dimension d ∈ N if there exists two embedding functions ϕ
⋆h : S ×
A → R d and µ ⋆h : S → R d such that
[Equation]
For normalization, we assume ∥ ϕ
⋆h ( s, a ) ∥ 2 ≤ 1 for all ( s, a ) , and for any function g : S → [0 , 1] , ∥ (cid:82) µ ⋆h ( s ) g ( s ) ds ∥ 2 ≤ √ d .
An MDP M is a low-rank MDP with dimension d if for each h ∈ [ H ] , P h admits a low rank decomposition with dimension d .
We use ϕ ⋆ = { ϕ ⋆h } h ∈
[ H ] and µ ⋆ = { µ ⋆h } h ∈
[ H ] to denote the embeddings for P ⋆ .
We remark that when ϕ h is revealed to the agent, low-rank MDPs specialize to linear MDPs (Wang et al., 2020; Jin et al., 2020b).
Essentially, low-rank MDPs do not assume that the features { ϕ h } h are known a priori.
The lack of knowledge on features in fact invokes a nonlinear structure, which makes the model strictly harder than linear MDPs or tabular models.
Since it is impossible to learn a model in polynomial time if there is no assumption on features ϕ h and µ h , we 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 { (Φ , Ψ) } that contains the true model, i.e., the embeddings ϕ
⋆ ∈
Φ , µ ⋆ ∈ Ψ .
While we assume the cardinality of the function classes to be finite for simplicity, extensions to infinite classes with bounded statistical complexity (such as bounded covering number) are not difficult (Sun et al., 2019; Agarwal et al., 2020).
We further assume that for a given low-rank MDP with any reward function, there exists an algorithm that returns the optimal policy in polynomial time.
Such a computational oracle is commonly adopted in the literature (Agarwal et al., 2020; Uehara et al., 2022b).
3. R EWARD - FREE RL AND L EARNING O BJECTIVES
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, is required to achieve certain learning goals (to be elaborated below) based on the outcome of exploration phase.
The planning phase may require the agent to achieve different learning goals.
In this paper, we focus on two of such goals.
The most popular goal in reward-free RL is to find a near-optimal policy that achieves the best value function under the true environment with ϵ -accuracy, as defined below.
Definition 2.
( ϵ -optimal policy).
Fix ϵ > 0 .
For any given reward function r , a learned policy π is ϵ -optimal if it satisfies
[Equation]
For model-based learning, Agarwal et al.
(2020) proposed system identification as another useful learning goal, defined as follows.
Definition 3 ( ϵ -accurate system identification) .
Fix ϵ > 0 .
Given a model class (Φ , Ψ) , a learned model (ˆ ϕ, ˆ µ ) is said to achieve ϵ -accurate system identification if it uniformly approximates the true model P ⋆ , i.e., ∀ π, h ∈
[ H ] , E ⋆π (cid:104) (cid:13)(cid:13)(cid:13)(cid:68)
ˆ ϕ h ( s h , a h ) , ˆ µ h ( · ) (cid:69) − P ⋆h ( ·| s h , a h ) (cid:13) (cid:13)(cid:13) TV (cid:105) ≤ ϵ .
T HE RAFFLE A LGORITHM
In this section, we describe our proposed algorithm RAFFLE for reward-free RL under low-rank
MDPs.
The algorithm is presented in Algorithm 1.
Exploration and MLE model estimation.
In each iteration n during the exploration phase, for each h ∈ [ H ] , the agent executes the exploration policy π n − 1 (defined in the previous iteration) up to step h − 1 , after which it takes two uniformly selected actions, and stop after step h + 1 .
Differently from FLAMBE that collects a large number of samples for each episode, our algorithm uses each exploration policy to collect only one sample trajectory at each episode, indexed by ( n, 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.
Then, the agent estimates the low-rank components ˆ ϕ ( n ) h and ˆ µ ( n ) h via the MLE oracle with given model class (Φ , Ψ) and a dataset D nh as follows
[Equation]
Update transition dynamics
Update empirical covariance matrix ˆ U n ) h as
Define exploration-driven reward function ˆ b ( n ) h as in Equation (3).
Define an estimated value function ˆ V π ˆ P ( n ) , ˆ b ( n ) based on ˆ P ( n ) and ˆ b ( n ) as Equation (4).
Find exploration policy π n = arg max
π ˆ V π ˆ P ( n ) , ˆ b ( n ) .
16: if 2 ˆ V π n ˆ P ( n ) , ˆ b ( n ) + 2 √ Kζ n ≤ ϵ then 17:
Terminate Phase I: Exploration Phase and set ˆ P ϵ
=
ˆ P ( n ) , ˆ b ϵ =
ˆ b ( n ) , π ϵ = π n , n ϵ = n .
18: end if 19: end for 20: Phase II: Planning Phase 21:
Receive reward function r = { r h } Hh =1 , and compute policy ¯ π = arg max π V π ˆ P ϵ ,r .
22: Output: policy ¯ π , learned transition dynamics ˆ P ϵ .
The agent then updates the estimated transition kernel with ˆ ϕ ( n ) h and ˆ µ ( n ) h .
Design of exploration reward.
The agent updates the empirical covariance matrix ˆ U ( n ) h as
[Equation]
where { s ( τ,h +1) h , a ( τ,h +1) h , s ( τ,h +1) h + } is collected at iteration τ , episode ( h + 1) , and step h .
Next, the agent uses both ˆ ϕ ( n ) h and ˆ U ( n ) h to construct an exploration-driven reward function as
[Equation]
where ˆ α n is a pre-determined parameter.
We note that although individual ˆ b ( n ) h ( s, a ) for each step may not represent point-wise uncertainty as indicated in Uehara et al.
(2022b)
, we find its total cumulative version
ˆ V π ˆ P ( n ) , ˆ b ( n ) can serve as a trajectory-wise uncertainty measure to select exploration policy.
To see this, it can be shown that for any π and h ,
E ⋆π (cid:104)(cid:13)(cid:13)(cid:13)(cid:68)
[Equation]
ˆ V
[Equation]
n ˆ P ( n ) , ˆ b ( n ) + (cid:112) c n n , where c ′ is a constant and c n = O (log n ) .
As iteration number n grows, the second term diminishes to zero, which indicates that ˆ V π n ˆ P ( n ) , ˆ b ( n ) (under the reward of ˆ b ( n ) ) serves as a good upper bound on the estimation error for the true transition kernel.
Hence, exploration guided by maximizing ˆ V π ˆ P ( n ) , ˆ b ( n ) 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.
Design of exploration policy.
The agent defines an truncated value function iteratively using the estimated transition kernel and the exploration-driven reward as follows
[Equation]
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 ˆ V π ˆ P ( n ) , ˆ b ( n ) , and uses this policy as 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 ˆ V π ˆ P ( n ) , ˆ b ( n ) 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 step.
Such a termination criterion enables our algorithm to identify an accurate model and later find a 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, Uehara et al.
(2022b) can only provide guarantees on a random mixture of the policies obtained over all iterations.
Planning phase.
Given any reward function r , the agent finds a near-optimal policy by planning with the specially selected learned transition dynamics ˆ P ϵ and the given reward r .
Note that such planning with a known low-rank MDP is computationally efficient by assumption.
M AIN R ESULTS 5. U PPER B OUNDS ON S AMPLE C OMPLEXITY
In this section, we first show that the policy returned by RAFFLE is an ϵ -optimal policy with respect to any given reward r in the planning phase.
The detailed proof can be found in Appendix A.
Theorem 1 ( ϵ -optimal policy) .
Assume M is a low-rank MDP with dimension d , and Assumptionholds.
Given any δ, ϵ ∈ (0 , 1) , and any reward function r , let ¯ π
and ˆ P ϵ be the output of RAFFLE and π ⋆ :=
arg max π V πP ⋆ ,r be the optimal policy under the true model P ⋆ .
Set ˆ α n = ˜
O
( (cid:112) ( K + d 2 )) and λ n = ˜
O
( d ) .
Then, with probability at least 1 − δ , we have V π ⋆ P ⋆ ,r − V ¯
πP ⋆ ,r ≤ ϵ , and the total number of trajectories collected by RAFFLE is upper bounded by
[Equation]
Compared with MOFFLE (Modi et al., 2021), which also finds an ϵ -optimal policy in reward-free
RL, our result improves their sample complexity of ˜ O (cid:16) H 5 K 5 d 3 LV ϵ 2 η (cid:17) in three aspects.
First, the order on K is reduced.
Second, the dimension d LV of the underlying function class in MOFFLE can be exponentially larger than d as shown in Agarwal et al.
Finally, MOFFLE requires reachability assumption, leading to a factor 1 /η in the sample complexity, which can be as large as √ d LV .
Further, Theorem 1 naturally achieves the goal of reward-known RL with the same sample complexity, which improves that of ˜ O (cid:16) H 5 d 4 K 2 ϵ(cid:17) in (Uehara et al., 2022b) by a factor of O ( H 2 ) .
Proceeding to the learning objective of system identification, the learned transition kernel output by
Algorithm 1 achieves the goal of ϵ -accurate system identification with the same sample complexity as follows.
The detailed proof can be found in Appendix B.
Theorem 2 ( ϵ -accurate system identification) .
Under the same condition of Theorem 1 and let ˆ P ϵ
= { ˆ ϕ ϵh , ˆ µ ϵh } be the output of RAFFLE.
Then, with probability at least 1 − δ , ˆ P ϵ achieves ϵ -accurate system identification, i.e. for any π and h : E ⋆π (cid:104)(cid:13)(cid:13)(cid:13)(cid:68)
ˆ ϕ ϵh ( s h , a h ) , ˆ µ ϵh ( · ) (cid:69) − P ⋆h ( ·| s h , a h ) (cid:13)(cid:13)(cid:13)
[Equation]
and the number of trajectories collected by RAFFLE is upper bounded by
[Equation]
Theorem 2 significantly improves the sample complexity of ˜ O
( H 22 K 9 d
[Equation]
) in Agarwal et al.
(2020) on the dependence of all involved parameters for achieving ϵ -accurate system identification.
5. L OWER B OUND ON S AMPLE C OMPLEXITY
In this subsection, we provide a lower bound on the sample complexity that all reward-free RL algorithms must satisfy under low-rank MDPs.
The detailed proof can be found in Appendix C.
Theorem 3 (Lower bound) .
For any algorithm Alg , which can output an ϵ -optimal policy satisfying Equation (1) , if H > max(24 ϵ, 4) , S ≥ 6 , K ≥ 3 and δ < 1 / 16 , then there exists a low-rank MDP model M such that the number of trajectories sampled by Alg is at least Ω (cid:0) HdKϵ 2 (cid:1) .
Theorem 3 has two important implications.
First, the lower bound in Theorem 3 matches the upper bound in Theorem 1 in terms of the dependence on ϵ as well as on K in the large d regime.
Second, Theorem 3 also shows that reward-free RL under low-rank MDPs is strictly harder than linear MDPs by at least a factor of Ω( K ) , because the lower bound in Theorem 3 for low-rank
MDPs has an additional term K compared to the upper bound ˜ O (cid:16) d 3 H 4 ϵ 2 (cid:17) provided in Wang et al.
(2020) for linear MDPs.
This can be explained intuitively as follows.
In linear models, all the representations ϕ : S ×
A → R d are known.
Then it requires at most O ( d ) actions with linearly independent features to realize all transitions ⟨ ϕ, µ ⟩ .
However, learning low-rank MDPs requires the agent to further select O ( K ) actions to access to unknown features, leading to a dependence on K .
5. R EPRESENTATION L EARNING VIA R EWARD -F REE RL
In low-rank MDPs, the representation ϕ , if learned accurately, can be reused in other RL environments.
Although RAFFLE provides an estimation of ϕ during its execution, such an estimation does not come with an accuracy guarantee.
Note that Theorem 2 on the system identification does not provide the guarantee on the representation ˆ ϕ , but only on the entire transition kernel ˆ P .
None of previous studies of reward-free RL under low-rank MDPs (Agarwal et al., 2020; Modi et al., 2021;
Uehara et al., 2022b) established the guarantee on ˆ ϕ .
In this section, we design an algorithm, which exploits the learned transition kernel from RAFFLE and learns an accurate representation without additional interaction with the environment.
We first introduce the following divergence-based metric to quantify the distance between two representations, which has been used in supervised learning (Du et al., 2021b).
Definition 4.
(Divergence between two representations).
Given a distribution q over S × A and two representations ϕ, ϕ ′ ∈
Φ , define the covariance between ϕ and ϕ ′ w.r.t q as Σ ( s,a ) ∼ q ( ϕ, ϕ ′ ) =
E (cid:2) ϕ ( s, a ) ϕ ′ ( s, a ) ⊤ (cid:3) .
Then, the divergence between ϕ and ϕ ′ with respect to q is defined as
[Equation]
It can be verified
that D q ( ϕ, ϕ ′ ) ⪰ 0 (i.e., positive semidefinite) and D q ( ϕ, ϕ ) =
for any ϕ, ϕ ′ ∈ Φ .
Algorithm 2 for representation learning via RAFFLE.
The formal version of Algorithm 2 is presented in Appendix D.
We explain the main idea of Algorithm 2 as follows.
First, for each h ∈ [ H ] , t ∈ [ T ] , where T is the number of rewards, N f pairs of state-action ( s h , a h ) are generated based on a distribution { q h } .
Note that the agent does not interact with the true environment during such data generation.
Then, for any h , if we set the reward r at step h to be zero, Q πP ⋆ ,h,r ( s h , a h ) can have a linear structure in terms of the true representation ϕ
⋆h ( s h , a h ) .
Namely, there exists a w h decided by r , P ⋆ and π such that Q πP ⋆ ,h,r
( s h , a h ) = ⟨ ϕ ⋆h ( s h , a h ) , w h ⟩ .
Then, with the estimated transition kernel ˆ P from RAFFLE and the well-designed rewards r h,t , Q π t ˆ P,h,r h,t ( s h , a h ) can be computed efficiently and serve as a target to learn the representation via the following regression problem:
[Equation]
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 a natural estimation error but requires no interaction with true environment, whereas their algorithm assumes a generative model to collect data from the ground truth transition kernel.
In order to guarantee the learned representation by Algorithm 2 is close enough to the ground truth, we need to employ the following two somewhat necessary assumptions.
First, for the distributions { q h } Hh =1 in Algorithm 2 of the state-action pair, it is desirable to have ˆ P ( ·| s, a ) approximates true P ⋆ ( ·| s, a ) well over distribution { q h } Hh =1 , so that Q π ˆ P,h,r ( s h , a h ) can approximate the ground truth well.
Intuitively, if some state-action pairs can hardly be visited under any policy, the output of Algorithm 1 ˆ P ( ·| s, a ) cannot approximate true P ⋆ ( ·| s, a ) well over these state-action pairs.
Hence, ˆ P ( ·| s, a ) naturally only performs well on the input distributions with small density on these state-action pairs.
We formalize such an assumption as follows.
Assumption 2.
For the distributions { q h } Hh =1 in Algorithm 2, there exists a policy π 0 and a constant
C min such that for any ( s, a ) ∈ S × A , q h ( s, a ) ≤ C min P π 0 h ( s, a ) , where P π 0 h ( · , · ) : S ×
A → R is the density over S ×
A induced by using policy π 0 to roll into state-action pair ( s, a ) at time h .
Assumption 2 is easy to satisfy if each state-action pair under the true MDP can be explored well by some policies.
In fact, under standard reachability type assumptions such as in Modi et al. (2021),
Assumption 2 holds for any distributions with finite bound (see Appendix D.1 for details).
Next, we assume that the rewards chosen for generating target Q -functions are sufficiently diverse so that the target Q πP ⋆ ,h,r ( s h , a h ) spans over the entire representation space to guarantee accurate representation learning in Equation (5).
Such an assumption is commonly adopted in multi-task representation learning literature (Du et al., 2021b;
Yang et al., 2021; Lu et al., 2021).
To formally state the assumption, for any h ∈ [ H ] , let { r h,t } t ∈
[ T ] be a set of T rewards (where T ≥ d ) satisfying r h,th = 0 .
As a result, for each t , given π t , there exists w th⋆ ∈ R d such that Q πP ⋆ ,h,r h,t ( s h , a h ) =
⟨ ϕ ⋆h ( s h , a h ) , w th⋆ ⟩ .
Let W ⋆h =
[ w 1 h⋆ , . .
, w Th ⋆ ] ∈ R d × T .
Assumption 3 (Diverse rewards) .
The smallest singular value σ d ( W ⋆h ) of W ⋆h defined above satisfies σ 2 d ( W ⋆h ) ≥ Ω( Td ) , i.e., there exists a constant C D > 0 such that σ 2 d ( W ⋆h ) ≥ C D Td .
We next characterize the accuracy of the output representation of Algorithm 2 in terms of its closeness to the ground truth representation in the following theorem and defer the proof to Appendix D.
Theorem 4 (Guarantee for representation learning) .
Under Assumptions 1 - 3, for any δ, ϵ ∈ (0 , 1) , any h ∈ [ H ] , and sufficiently large N f , let ˆ P be the output transition kernel of RAFFLE that satisfies
Theorem 2.
Then, with probability at least 1 − δ , the output of Algorithm 2 ˜ ϕ
satisfies
[Equation]
We explain the bound in Equation (6) as follows.
The first term arises from the system identification error ϵ 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 input distribution { q h } h ∈ [ H ] , and vanishes as N f becomes large.
Note that N f 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 that Algorithm 2 can learn a near-accurate representation, based on which the learned representation can be further used in other RL environments that share a common representation, similar to how representation learning has been exploited in supervised learning (Du et al., 2021b).
C ONCLUSIONS
In this paper, we investigate the reward-free reinforcement learning, in which the underlying model admits a 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 an algorithm to exploit the model learned by RAFFLE for accurate representation learning without further interaction with the environment.
Although ϵ -accurate system identification can easily induce Hϵ -optimal policy, the relationship between these two learning goals under general reward-free exploration setting remain under-explored, and is an interesting topic for future investigation.
Appendices A P ROOF OF T HEOREM
We first provide a proof outline to highlight our key ideas in the analysis of Theorem 1, and then provide the detailed proof.
To simplify the notation, we denote the total variation distance (cid:13)(cid:13)(cid:13) ˆ P ( n ) h ( s h , a h ) − P ⋆h
( s h , a h ) (cid:13)(cid:13)(cid:13) TV by f ( n ) h ( s h , a h ) .
Proof Outline for Theorem 1.
Step 1 provides an upper bound on the difference of 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 4 in appendix A.2).
Proposition 1 (Informal) .
There exist constants c n = O (log n ) and c ′
=
O (1) such that for any policy π and reward r , with high probability, we have
[Equation]
This proposition is inspired by Uehara et al.
(2022b), while it generalizes to non-stationary setting and arbitrary reward scenario from infinite stationary MDP and fixed reward.
The main proof idea is to first notice that for any reward, we have
[Equation]
Then, we show that ˆ V π
[Equation]
n due to the optimality of the exploration policy π n .
Step 2 shows sublinearity of the summation of ˆ V π n ˆ P ( n ) , ˆ b ( n ) , which is an upper bound of value function difference under the true model with the estimated model (see proposition 5 in appendix A.3).
Proposition 2 (Informal) .
Under the same setting of Proposition 1, with high probability, the summation of value functions V π n ˆ P ( n ) , ˆ b ( n ) under exploration policies { π n } n ∈ [ N ] with exploration-driven reward functions ˆ b ( n ) is sublinear, as given by
[Equation]
Step 3 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 environment ˆ P ϵ is at most ϵ .
Proposition 3.
With high probability, RAFFLE terminates after at most ˜ O (cid:0) H 2 d 2 K
[Equation]
d 2 + K ) /ϵ
[Equation]
iterations, and the output model ˆ P ϵ satisfies
[Equation]
where n ϵ is the iteration number where RAFFLE terminates, i.e. ˆ P ϵ
= ˆ P ( n ϵ ) .
Finally, with some algebraic operations, Proposition 3 concludes the proof of Theorem 1.
A. S UPPORTING L EMMAS
We first present the high probability event.
Lemma 1.
We define Π n to be a uniform mixture of previous n − 1 exploration policies:
[Equation]
Denote the total variation of ˆ P ( n ) and P ⋆ , and the expected matrix of ˆ U ( n ) h as follows.
[Equation]
where λ n = β 3 d log(2 nH | Φ | /δ )) and β 3 =
O (1) is a constant coefficient.
Suppose Algorithmruns N iterations and let events E
and E 1 be defined as follows.
[Equation]
[Table]
where ζ n = log (2 | Φ || Ψ | nH/δ )
/n .
We further denote ζ = Nζ N = log (2 | Φ || Ψ |
NH/δ ) for simplicity.
Denote E = E 0 ∩ E 1 as the intersection of the two events.
Then, P [ E ] ≥ 1 − δ .
Proof.
By Corollary 2 in Appendix F, we have P [ E 0 ] ≥ 1 − δ/ 2 .
Further, by Lemma 39 in Zanette et al.
(2021) for the version of fixed ϕ
and Lemma 11 in Uehara et al.
(2022b)
, we have P [ E 1 ] ≥ 1 − δ/ 2 .
Therefore, P [ E ] ≥ 1 − δ .
Based on Lemma 1, we can bound the exlporation-driven reward in RAFFLE as follows.
Corollary 1.
Given that the event E occurs, the following inequality holds for any n ∈ [ N ] , h ∈ [ H ] , s h ∈ S , a h ∈ A : min [Equation]
[Equation]
h
[Equation]
( U ( n ) h, ˆ
[Equation]
where ˆ α n = 5 (cid:112)
β 3 nζ n ( K + d 2 ) .
Proof.
Recall ˆ b ( n ) h
( s h , a h ) = min (cid:26)
ˆ α n (cid:13)(cid:13)(cid:13)
ˆ ϕ ( n ) h ( s, a ) (cid:13)(cid:13)(cid:13) (ˆ U ( n ) h ) − 1 ,(cid:27) .
Applying Lemma 1, we can immediately obtain the result.
The following lemma extends the Lemmas 12 and 13 under infinite discount MDPs in Uehara et al.
(2022b) to episodic MDPs.
We provide the proof for completeness.
Lemma 2.
Let P
h − 1 = ⟨ ϕ h − 1 , µ h − 1 ⟩ be a generic MDP model, and Π be an arbitrary and possibly mixture policy.
Define an expected Gram matrix as follows
[Equation]
Further, let f h − 1 ( s h − 1 , a h − 1 ) be the total variation between P ⋆h − 1 and P h − 1 at time step
h − 1 .
Suppose g ∈ S ×
A → R is bounded by B ∈ (0 , ∞ ) , i.e., ∥ g ∥ ∞ ≤ B .
Then, ∀ h ≥ 2 , ∀ policy π h ,
[Equation]
Proof.
We first derive the following bound:
[Equation]
where the inequality follows from Cauchy-Schwarz inequality.
We further expand the second term in the RHS of the above inequality as follows.
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where ( i ) follows from the assumption that ∥ g ∥ ∞ ≤ B , ( ii ) follows from Jensen’s inequality, ( iii ) is due to that f h − 1 ( s h − 1 , a h − 1 ) is the total variation between P ⋆h − 1 and P h − 1 at time step
h − 1 , and ( iv ) is due to importance sampling.
This finishes the proof.
Based on Lemma 2, we summarize three useful inequalities which bridges the total variation f ( n ) h and the exploration-driven reward ˆ b ( n ) h .
Lemma 3.
Define
[Equation]
where λ n = β 3 d log(2 nH | Φ | /δ ) .
Given that the event E occurs, the following inequalities hold.
For any n , when h ≥ 2 ,
[Equation]
where
[Equation]
Specially, when h = 1 ,
[Equation]
Proof.
We start by developing Equation (12) as follows.
Given that the event E occurs, for h ≥we have
E
[Equation]
where ( i ) follows from Lemma 2 and the fact that f ( n ) h ( s h , a h ) ≤ 1 , ( ii ) follows from importance sampling at time step h − 2 , and ( iii ) follows from Lemma 1.
Equation (13) follows from the arguments similar to the above
[Equation]
To obtain Equation (14)
[Equation]
we first apply Lemma 2 and obtain
E
[Equation]
where we use the fact that ˆ b ( n ) h ( s h , a h ) ≤ 1 .
We further bound the term n E sh ∼ ( P⋆,
[Equation]
n ) ah ∼U [(ˆ b ( n ) h ( s h , a h )) 2 ] as follows:
[Equation]
where ( i ) follows from Lemma 1, and we use tr( A ) to denote the trace of any matrix A .
Hence,
E
[Equation]
where the last inequality follows from that ˆ α n = 5 α n and the definition of γ n In addition, for h = 1 , we have
[Equation]
and
E a
[Equation]
A. P ROOF OF P ROPOSITION
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 1) .
For all n ∈ [ N ] , policy π and reward r , given that the event E occurs, we have
[Equation]
Proof.
Step 1.
We first show that (cid:12)(cid:12)(cid:12) V πP ⋆ ,r − V π ˆ P ( n ) ,r (cid:12)(cid:12)(cid:12) ≤
ˆ V π ˆ P ( n ) ,f ( n ) .
Recall the definition of estimated value functions ˆ V h, ˆ P ( n ) ,r
[Equation]
s h
[Equation]
and ˆ Q h, ˆ P
[Equation]
n )
[Equation]
r
[Equation]
h
[Equation] [Equation]
We develop the proof by induction.
For the base case h =
H + 1 , we have
[Equation]
Assume that (cid:12)(cid:12)(cid:12) V πh +1 , ˆ P ( n ) ,r ( s
[Equation]
V πh
[Equation]
⋆
[Equation]
( s
[Equation] [Equation]
(cid:12)(cid:12)(cid:12) ≤
ˆ V πh +1 , ˆ P ( n )
[Equation]
f ( n ) ( s h +1 ) holds for any s h
+1 .
Then, from Bellman equation, we have,
[Equation]
where ( i ) follows from the (action) value function is at most
[Equation]
and ( ii ) follows from the induction hypothesis
[Equation]
Then, by the definition of ˆ V πh, ˆ P ( n ) ,r ( s h )
[Equation]
we have
[Equation]
where ( i ) follows from Equation (16)
[Equation]
Therefore, by induction, we have
[Equation]
Step 2.
Then, we show that ˆ V π ˆ P
[Equation]
By Equation (12) and the fact that the total variation distance is upper bounded by 1, with probability at least 1 − δ/ 2 , we have
[Equation]
h
[Equation]
( s h
[Equation] [Equation]
h
[Equation]
1 ) (cid:13)(cid:13)(cid:13) ( U ( n ) h − 1 , ˆ ϕ ) −,(cid:33)(cid:35) , ∀ h ≥ 2 .
Similarly, when h = 1 ,
[Equation]
Based on Corollary
[Equation]
Equation (17) and α n = 5ˆ α n , we have
[Equation]
For the base case h
= H
[Equation]
we have
[Equation]
Assume that E ˆ P ( n ) ,π
[Equation]
n )
[Equation]
s h ) holds for step h
+ 1 .
Then, by
Jensen’s inequality
[Equation]
we obtain
E ˆ P ( n ) ,π
[Equation]
where ( i ) follows from Equation (19)
[Equation]
and ( ii ) is due to the induction hypothesis.
By induction
[Equation]
we conclude that
[Equation]
Combining Step 1 and Step 2, we conclude that
[Equation]
A. P ROOF OF P ROPOSITION
The following lemma is key to ensure that RAFFLE terminates in finite episodes.
Proposition 5 (Restatement of Proposition 2) .
Given that the event E occurs, ζ = log (2 | Φ || Ψ |
NH/δ )
the summation of the truncated value functions ˆ V π n ˆ P ( n ) , ˆ b ( n ) under exploration policies { π n } n ∈
[ N ] is sublinear, i.e., the following bound holds:
[Equation]
Proof.
Note that ˆ V πh, ˆ P ( n ) , ˆ b ( n ) ≤ 1 holds for any policy π and h ∈
[ H ] .
We first have ˆ
[Equation] [Equation]
which implies ˆ V π n ˆ P ( n
[Equation]
, ˆ b
[Equation]
n P
[Equation]
Applying the Equation (14) and Equation (15)
[Equation]
we obtain the following bound on the value function
V π n P ⋆ , ˆ b ( n ) :
[Equation]
n
[Equation]
Similarly
[Equation]
we obtain
[Equation]
Then, taking the summation of V π n P ⋆ , ˆ b ( n ) + f ( n ) over n ∈ [ N ] , we have
[Equation]
where ( i ) follows from Cauchy-Schwarz inequality and importance sampling, and ( ii ) follows from
Lemma 10.
Hence, the statement of Proposition 5 is verified.
A. P ROOF OF P ROPOSITION
Based on Proposition 5, we argue that with enough number of iterations, RAFFLE can find ˆ P ϵ satisfying the condition in line 15 of Algorithm 1.
Proposition 6 (Restatement of Proposition 3) .
Fix any δ ∈ (0 , 1) , ϵ > 0 .
Suppose the algorithm runs for N =
[Equation]
iterations, with probability at least 1 − δ , RAFFLE can find an n ϵ ≤ N in the exploration phase such that 2 ˆ V π nϵ ˆ P ( nϵ ) , ˆ b ( nϵ
[Equation]
+2 (cid:112) Kζ n ϵ ≤ ϵ .
In other words, Algorithm 1 can output ˆ P ϵ
= ˆ P ( n ϵ ) satisfying the condition in line 15.
In addition
[Equation]
Proof.
We show that the algorithm terminates by contradiction.
If it does not stop, applying Proposition 5, we have
[Equation]
Therefore,
[Equation]
ϵ
[Equation]
Recall ζ = Nζ N = log (2 | Φ || Ψ |
NH/δ ) .
Using the fact that n ≤ c log 2 ( α n n ) ⇒ n ≤ 4 c log 2 ( α n c ) , ∀ c ≥ e 2 , n ≥ 1 , α n ∈ R + , it can be concluded that
[Equation]
which is a contradiction.
Therefore, there exists an n ϵ
=
O ( H 2 d 2 K ( d 2 + K )log 2 ( |
[Equation]
||
[Equation]
|
[Equation] [Equation]
) such that ˆ P ϵ
= ˆ P ( n ϵ ) satisfies
[Equation]
Combining Proposition 4, we finish the proof.
Proof of Theorem 1.
Recall that ˆ P ϵ is the output of RAFFLE in the n ϵ -iteration.
Then, by Proposition
[Equation]
where ( i ) follows from the definition of ¯ π .
The number of trajectories n ϵ H is at at most
[Equation]
B P ROOF OF T HEOREM
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 δ ∈ (0 , 1) , for any h ∈
[ H ] , n ∈ N + , any policy π , with probability at least 1 − δ/ 2 ,
[Equation]
Proof.
Recall that f ( n ) h
[Equation]
s, a
[Equation]
=
[Equation]
[Table]
(cid:13)(cid:13)(cid:13)
E
[Equation]
Hence, for h ≥ 2 , we have
[Equation]
[Equation]
where ( i ) follows from Equation (19), ( ii ) follows from the definition of ˆ Q πh − 1 , ˆ P ( n ) , ˆ b ( n ) ( s h − 1 , a h − 1 ) and ( iii ) follows from Equation (20)
[Equation]
where the first term in ( i ) is due to Equation (21) and Equation (15), the second term in ( i ) is due to
Proposition 4 and ( ii ) follows from the definition of π n .
Proof of Theorem 2.
By Proposition 6, let n ϵ
= O (cid:16) H 2 d 2 K
( d 2 + K )log 2 ( | Φ || Ψ |
NH 3 d 2 K ( K + d 2 ) / ( δϵ 2 )) ϵ(cid:17) , with no more than n ϵ H trajectories, RAFFLE can learn a model ˆ P ϵ , bonus ˆ b ϵ and policy π ϵ at the n ϵ -th iteration satisfying 2 V
π ϵ ˆ P ϵ , ˆ b ϵ + 2 (cid:112)
Kζ n ϵ ≤ ϵ .
Then, following Lemma 4, we have
[Equation]
C P ROOF OF T HEOREM 3 (L OWER B OUND ON S AMPLE C OMPLEXITY ) C. S TEP 1: C ONSTRUCTION OF H ARD 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 ≥ K , where S, K 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 to 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 us to separately treat the state space and the action space, so that both state and action spaces can be arbitrarily large. We then explicitly define the feature vectors for all state-action pairs and show our hard MDP instances have a lowrank structure with dimension d = S .
[-+
 | s h ,-]In a nutshell, we construct a family of HdK MDPs that are hard to distinguish in KL divergence, while the corresponding optimal policies are very different as shown in Figure 1.
First, we define a [-) -th coordinate of µ ( s h +1 ) is the transition probability P ( s h +
-]reference MDP M 0 as follows.
We start with the construction of the state space S and the action space A . • Let S =
{ s w , s o , s g , s b } (cid:83) i ∈
[ D ] ,j ∈ [2 i − 1 ] { s i,j } , where s w , s o , s g and s b denote ‘waiting state’, ‘outlier state’, ‘good state’, and ‘bad state’, respectively.
The states in (cid:83) i ∈ [ D ] ,j ∈ [2 i − 1 ] { s i,j } form a binary tree, where s i,j denotes the j -th branch node of the layer i .
• Let A = { a w , a 1 , a 2 } (cid:83) A 0 , where a w denotes ‘waiting action’, a 1 , a 2 are two unique actions that form the binary tree, and |A 0 | =
K − 3 . Then the transition probabilities of M 0 are specified through the following rules • The initial state is the waiting state s w .
• If the agent takes the waiting action a w before time step ¯ H , waiting state s w stays on itself. Otherwise, s w transits to the root state s 1 , 1 of the binary tree. Mathematically, P h [ s w | s w , a w ] = 1 h ≤ ¯ H , and P h [ s 1 , 1 | s w , a ] = 1 a ̸ = a w or h> ¯ H .
• When i < D , for states s i,j in the binary tree, we have the following transition rules: – If the agent takes actions a 1 or a 2 , s i,j deterministically transits to its children s i +1 , 2 j − 1 or s i +1 , 2 j , respectively. Mathematically, P h [ s i +1 , 2 j − 1 | s i,j , a 1 ] = 1 , and P [ s i +1 , 2 j | s i,j , a 2 ] = 1 .
– If the agent takes any action other than a 1 or a 2 , the agent will reach the outlier state s o .
Mathematically, P h
[ s o | s i,j , a ] = 1 , ∀ a ̸ = a 1 , a 2 . • Leaf state s D,j uniformly transits to good state s g and bad state s b no matter what action the agent takes. Mathematically, P h [ s g | s D,j , a ] =
P
h [ s b | s D,j , a ] = 12 , ∀ a ∈ A . • Good state s g , bad state s b , and outlier state s o are absorbing states. Now we define the features as follows, which are S -dimensional vectors.
[Equation]
where 0 S − 5 ∈ R S − 5 denotes the S − 5 dimension vector with all zeros and e i,j ∈ R S − 5 denotes the one-hot vector that is zero everywhere except the (2 i − 1 + j − 2) -th coordinate.
Here 2 ≤ i ≤
D, 1 ≤ j ≤ 2 i − 1 .
For each ( h ∗ , ℓ ∗ , a ∗ ) ∈ { 1+ D, . .
, ¯ H + D }× [2 D − 1 ] ×A , we define an MDP M ( h ∗ ,ℓ ∗ ,a ∗ ) through
M 0 .
Specifically, the only difference of M ( h ∗ ,ℓ ∗ ,a ∗ ) from M 0 is that the transition probability from the leaf state s D,ℓ ∗ and action a ∗ to the good state s g increases ϵ 0 , i.e. P h ∗ [ s g | s D,ℓ ∗ , a ∗ ] = 12 + ϵ 0 , and P h ∗ [ s b | s D,ℓ ∗ , a ∗ ] = 12 − ϵ 0 , where ϵ 0 will be specified later.
We note that the features of
M ( h ∗ ,ℓ ∗ ,a ∗ ) are the same as those of M 0 except that ϕ h ∗ ( s D,ℓ ∗ , a ∗ )
= (0 , 0 , 0 S − 5 , 0 , 12 + ϵ 0 , 12 − ϵ 0 ) .
We remark here that the cardinality K 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 o to be a set of outlier states S o .
Definition of reward : the reward can only be attained in two special states: the good state s g and the outlier state s o at the last stage H .
[Equation]
and r h ( s, a ) still belongs to [0 , 1] .
C. S TEP 2: A NALYSIS OF H ARD MDP I NSTANCES
Proof of Theorem 3.
{+Let ϵ 
+}
= 2 ϵ .
For any MDP M h ∗ ,a ∗ , the optimal policy is to take action a w to stay at state s w until stage h ∗ − D , and then take the corresponding action to the only state s D,ℓ ∗ at stage h ∗ .
At state s D,ℓ ∗ , the agent takes the only optimal action a ∗ .
The optimal value function V ∗M ( h ∗ ,ℓ ∗ ,a ∗ ) = 12 + ϵ 0 , and the value function of the output policy of Alg is given by
[Equation]
22) where P ˆ π τ M ( h ∗ ,ℓ ∗ ,a ∗ ) is the probability distribution over the states and actions ( s h , a h ) following the
Markov policy ˆ π τ in the MDP M ( h ∗ ,ℓ ∗ ,a ∗ ) .
We remark that the reward of outlier state are specially designed to be 1 / 2 at stage H to make Equation (22) hold for policy ˆ π τ falling into the outlier state. Hence, [Equation]
and [Equation] The transitions of all MDPs are the same when the leaves states are reached. We define the event
[Equation]
From Equation (23), the event is equal to the event { V ∗M ( h ∗ ,ℓ ∗ ,a ∗ ) − V
ˆ π τ M ( h ∗ ,ℓ ∗ ,a ∗ ) ≤ ϵ } .
As a result,
[Equation]
Recall that N
[Equation]
[Equation]
)=
[Equation]
∗ ,a ∗ ) } such that (cid:80) ( h ∗ ,ℓ ∗ ,a ∗ ) N τ
( h ∗ ,ℓ ∗ ,a ∗ ) ≤ τ .
This inequality holds because the agent is likely to fall into the outlier state s o .
We denote P 0 and E 0 to be with respect to M 0 .
Now, we invoke an intermediate result in the proof of Theorem 7 in Domingues et al.
(2020) to conclude that [Equation] Summing over all ( h ∗ , ℓ ∗ , a ∗ ) , we have [Equation]
Notice that
[Equation]
Substituting Equation (25) into Equation (24) yields
[Equation]
where we use the fact that δ < 1 / 16 .
With the assumption of K ≥
[Equation]
S ≥
[Equation] we have d = S . Taking ¯ H =
H 3 and with the assumption of D ≤ H/ 3 , we have [Equation] Then following the analysis similar to that for Corollary 8 in Domingues et al. (2020), with probability at least 1 − δ , the number of iterarion is at least
[Equation]
D A LGORITHM FOR R EPRESENTATION L EARNING IN S ECTION 5.3 AND P ROOF OF T HEOREM
We first present the full algorithm in Section 5.3 below as Algorithm 2.
D. D ISCUSSION ON R EACHABILITY A SSUMPTION : A SSUMPTION
As discussed in Section 5.3, intuitively, if some state action pairs can be hardly visited by any policy, the output of Algorithm 1 ˆ P ( ·| s, a ) can not approximate true P ⋆ ( ·| s, a ) well over these state action pairs, so the input distribution naturally requires restrictions.
On the other hand, it is easy for input distributions { q h } Hh =1 to satisfy Assumption 2 if each state action pair under true MDPs can be explored well by some policies.
In fact if we assume following reachability assumption, Assumption 2 holds with for any { q h } Hh =1 with finite bound C B .
Such reachability type assumption is common in relevant literature (Modi et al., 2021; Agarwal et al., 2020).
[Equation]
Assumption 4 (Reachability) .
For the true transition kernel P ⋆ , there exists a policy π 0 such that min s ∈S P π 0 h ( s ) ≥ η min , where P π 0 h ( · ) :
S → R is the density function over S using policy π 0 to roll into state s at timestep h .
Under Assumption 4, for any bounded input distributions { q h } Hh =1 with bound C B , Assumptionholds with C min =
C B η min .
D. S UPPORTING L EMMAS
We first show that Q
π ˆ P,h,r can approximate Q
πP ⋆ ,h,r well over distribution { q h } h ∈
[ H ] by following two lemmas.
Lemma 5.
Given any δ ∈ (0 , 1) .
Let ˆ P be the output of Algorithm 1, for any policy π and rewards r , with probability at least 1 − δ , we have
[Equation]
Proof.
Define ˆ f h
( s h , a h ) =
(cid:13)(cid:13)(cid:13) ˆ P h ( ·| s h , a h ) − P ⋆ ( ·| s h , a h ) (cid:13)(cid:13)(cid:13) TV and ˆ f is a collection of all ˆ f h , i.e.
(cid:110)
ˆ f h (cid:111)
h ∈
[ H ] .
[Equation]
[Equation]
where ( i ) follows from Equation (16), ( ii ) follows from the definition of ˆ V and ˆ Q , and ( iii ) follows from the proof of Theorem 1.
Then for the second inequality,
[Equation]
where the first term in ( i ) is due to Equation (26) and Equation (15), the second term in ( i ) is due to
Step 1 in Proposition 4 and ( ii ) follows from the proof of Theorem
[Equation]
Lemma 6.
Given any δ, ϵ ∈ (0 , 1) and the output of Algorithm 1, under Assumption 2, for any policy π and rewards r , let { q h } Hh =1 be any distribution satisfying Assumption 2.
Then with probability at least 1 − δ , for each h ∈ [ H ] , E ( s h ,a h ) ∼ q h (cid:104)(cid:12)(cid:12)(cid:12) Q πP ⋆
[Equation]
h
[Equation]
r
[Equation]
s h
[Equation]
a h )
[Equation]
Q π ˆ P
[Equation]
h
[Equation]
r ( s h
[Equation]
a h ) (cid:12)(cid:12)(cid:12)(cid:105)
[Equation]
min .
Proof.
[Equation]
where ( i ) follows Assumption 4 and ( ii ) follows Lemma 5.
The lemma above shows that for any reward r , Q π ˆ P,h,r ( s h , a h ) can be the target of Q πP ⋆ ,h,r
( s h , a h ) when ( s h , a h ) are chosen from given distribution q h .
We then show that for any h , when reward r is set to be zero at step h , Q πP ⋆ ,h,r has a linear structure w.r.t ϕ
⋆ h .
Lemma 7.
For any h ∈ [ H ] , policy π and given ( s h , a h ) ∈ S × A , given any r such that r is set to be zero at step h , i.e. r h = 0 , then Q πP ⋆ ,h,r ( s h , a h ) is linear with respect to ϕ
⋆
( s h , a h ) , i.e. there exist a w ⋆h such that Q πP ⋆ ,h,r
( s h , a h ) = ⟨ ϕ ⋆h ( s h , a h ) , w ⋆h ⟩ .
Proof.
[Equation]
where w ⋆h = (cid:82) µ ⋆h ( s h
+
V πP ⋆ ,h +1 ,r [Equation] s h +
[Equation]
) ds h +1 .
D. P ROOF OF T HEOREM
Proof of Theorem 4.
We allow ϕ and Q π t ˆ P,h,r h,t to apply to all the samples in a dataset matrix simultaneously, i.e. ϕ h ( D N f ,t h )
=
( ϕ h ( s 1 ,th ) , . . . , ϕ h ( s N f ,t h ))
⊤ ∈ R N f × d and Q π t ˆ P,h,r h,t ( D N f ,t h ) =
( Q π t ˆ P,h,r h,t ( s 1 ,th ) , . . .
, Q π t ˆ P,h,r h,t ( s N f ,t h ))
⊤ ∈ R N f .
After we estimating ˜ ϕ h and then taking it as known, from the property of linear regression in Equation (5), for any reward function r h,t , any policy π t we got
[Equation]
) , where P ˜ ϕ
h ( D Nf,th ) =
˜ ϕ
h ( D N f ,t h )(˜ ϕ h ( D N f ,t h )
⊤ ˜
ϕ
h ( D N f ,t h )) † ˜ ϕ
h ( D N f ,t h ) ⊤ , represents the projection operator to the column spaces of ˜ ϕ
h ( D N f ,t h ) .
Then (cid:88) t ∈ [ T ] (cid:13)(cid:13)(cid:13) (cid:13)
P ⊥ ˜ ϕ
h ( D Nf,t h ) Q
π t ˆ P,h
[Table]
[Equation] [Equation]
where ( i ) follows from minimality of { ˜ ϕ th } t ∈ [ T ] and { (cid:101) w th } t ∈
[ T ] , ( ii ) follows Hoeffding’s inequality, ( iii ) follows that (cid:12) (cid:12)(cid:12) Q π t P ⋆ ,h,r h,t ( s n,t h , a n,t h ) − Q π t ˆ P,h,r h,t ( s n,t h , a n,t h ) (cid:12) (cid:12)(cid:12) ≤ 1 and ( iv ) follows Lemma 6.
As a result: (cid:88) t ∈
[Equation] [Equation]
=
[Equation]
t
[Equation]
[ T ]
[Equation]
≤ (cid:88) t ∈ [
[Equation] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where ( i ) follows from ∥ Av ∥ 2 ≤ σ 1 ( A ) ∥ v ∥ 2 and ( ii ) follows from that σ 1 ( P
⊥ ˜
ϕ
h ( D Nf,th ) ) ≤ 1 and the process to derive Equation (27).
We finally use the technique in Du et al.
(2021b) to derive the super population guarantee.
[Equation]
[Table]
[Equation]
where ( i ) follows from lemma B [Equation] 1 in Du et al.
(2021b)
and N f can be large enough
[Equation]
and ( ii ) follows from that ∥ AB ∥ 2 F ≥ ∥
A ∥ 2 F σ 2min ( B ) [Equation] where σ min ( B ) denote the smallest singular value of matrix
B .
Then
[Equation]
E D ISCUSSION 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 et al., 2022a).
We elaborate the key differences between our paper and Liu et al.
; Chen et al.
(2022a) as follows.
• The two POMDP papers mentioned above study reward-known setting, whereas our 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 are 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 to realize computationally.
• Due to the hardness of POMDP, MLE approach only guarantees that the estimated distribution of 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 is controlled at each time step.
F A UXILIARY L EMMAS
Recall f ( n ) h ( s, a ) =
∥
ˆ P ( n ) h ( ·| s, a ) − P ⋆h
( ·| s, a ) ∥
TV represents the estimation error in the n -th iteration at step h , given state s and action a , in terms of the total variation distance.
By using Theoremin Agarwal et al.
(2020), we are able to guarantee that under all exploration policies, the estimation error can be bounded with high probability.
Lemma 8.
(MLE guarantee).
Given δ ∈ (0 , 1) , we have the following inequality holds for any n, h ≥ 2 with probability at least 1 − δ/ 2 :
[Equation]
where ζ n := log (2 | Φ || Ψ | nH/δ ) n .
In addition, for h = 1 ,
[Equation]
Divide both sides of the result of lemma Lemma 8 by n , and define Π n =
U ( π 1 , . . .
, π n − 1 ) , we have the following corollary which will be intensively used in the analysis.
Corollary 2.
Given δ ∈ (0 , 1) , the following inequality holds for any n, h ≥ 2 with probability at least 1 − δ/ 2 :
[Equation]
In addition, for h = 1 ,
[Equation]
The following lemma (Dann et al., 2017) will be useful to measure the difference between two value functions under two MDPs and reward functions.
We define P h V
h +1 ( s h , a h ) =
E s ∼ P h ( ·| s h ,a h )
[ V ( s )] for shorthand.
Lemma 9.
(Simulation lemma).
Suppose P 1 and P 2 are two MDPs and r 1 , r 2 are the corresponding reward functions.
Given a policy π , we have,
[Equation]
The following lemma is 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.
Lemma 10.
(Elliptical potential lemma).
Consider a sequence of d × d positive semidefinite matrices
X 1 , . . .
, X N with tr( X n ) ≤ 1 for all n ∈ [ N ] .
Define M 0 = λ
I and M n = M n − 1 + X n .
Then
N
[Equation]
If we choose any subset of the set { X n M − 1 n − 1 } N n =1 , we can still get a sublinear summation as follows.