R EPRESENTATION L EARNING FOR O NLINE AND O FFLINE RL IN L OW - RANK MDP S
Anonymous authors
Paper under double-blind review
A BSTRACT
This work studies the question of Representation Learning in RL: how can we learn a compact low-dimensional representation such that on top of the representation we can perform RL procedures such as exploration and exploitation, in a sample efﬁcient manner.
We focus on the low-rank Markov Decision Processes (MDPs) where the transition dynamics correspond to a low-rank transition matrix.
Unlike prior works that assume the representation is known (e.g., linear MDPs), here we need to learn the representation for the low-rank MDP.
We study both the online RL and ofﬂine RL settings.
For the online setting, operating with the same computational oracles used in F LAMBE (Agarwal et al., 2020b)—-the state-of-art algorithm for learning representations in low-rank MDPs, we propose an algorithm R EP -UCB—Upper Conﬁdence Bound driven R EP resentation learning for
RL, which signiﬁcantly improves the sample complexity from (cid:101) O
( A 9 d 7 / ( (cid:15) 10 (1 − γ ) 22 )) for F LAMBE to (cid:101) O
( A 4 d 4 / ( (cid:15) 2 (1 − γ ) 3 )) with d being the rank of the transition matrix (or dimension of the ground truth representation),
A being the number of actions, and γ being the discount factor.
Notably, R EP -UCB is simpler than
F LAMBE , as it directly balances the interplay between representation learning, exploration, and exploitation, while F LAMBE is an explore-then-commit style approach and has to perform reward-free exploration step-by-step forward in time.
For the ofﬂine RL setting, we develop an algorithm that leverages pessimism to learn under a partial coverage condition: our algorithm is able to compete against any policy as long as it is covered by the ofﬂine data distribution.
I NTRODUCTION
When applying Reinforcement Learning (RL) to large-scale problems where data is complex and high-dimensional, learning effective transformations of the data, i.e., representation learning, can often signiﬁcantly improve the sample and computation efﬁciency of the RL procedure.
Indeed, several empirical works have shown that leveraging representation learning techniques developed in supervised or unsupervised learning settings can accelerate the search for good decision-making strategies (Silver et al., 2018; Stooke et al., 2021; Srinivas et al., 2020; Yang & Nachum, 2021).
However, representation learning in RL is far more subtle than it is for non-sequential and noninteractive learning tasks (e.g., supervised learning).
Prior works have shown that even if one is given the magic representation that exactly linearizes the optimal policy (Du et al., 2019b) or the optimal value functions (Wang et al., 2020; Weisz et al., 2021), RL is still challenging (i.e., one may still need exponentially many samples to learn).
This indicates that an effective representation that permits efﬁcient RL needs to encode more information about the underlying Markov Decision Processes (MDPs).
Despite the recent empirical success of representation learning in RL , its statistical guarantee and theoretical properties remain under-investigated.
In this work, we study the representation learning question under the low-rank MDP assumption.
Concretely, a low-rank MDP assumes that the MDP transition matrix admits a low-rank factorization, i.e., there exists two unknown mappings µ ( s (cid:48) ) , φ ( s, a ) , such that P ( s (cid:48) | s, a ) = µ ( s (cid:48) )
(cid:62) φ ( s, a ) for all s, a, s (cid:48) , where P ( s (cid:48) | s, a ) is the probability of transiting to the next state s (cid:48) under the current state and action ( s, a ) .
The representation φ in a low-rank MDP not only linearizes the optimal stateaction value function of the MDP (Jin et al., 2020a), but also linearizes the transition operator.
A lowrankness assumption on large stochastic matrices is a common assumption and has enabled successful development of algorithms for real world applications such as movie recommendation systems (Koren et al., 2009). We note that a low-rank MDP strictly generalizes the linear MDP model (Yang
[Table]
Algorithms such as OLIVE, Witness rank, and BLin-UCB work for settings which are more general than low-rank MDPs and have tight sample complexity.
However, these algorithms are version space algorithms and thus are not computationally efﬁcient.
Mofﬂe is an oracle-efﬁcient algorithm (with a much stronger oracle than the one in F LAMBE and ours), but the assumptions under which Mofﬂe operates essentially imply that the MDP’s transition has low non-negative matrix rank (nnr) (see the detailed discussion in the related work section).
Note that a nnr is at least as large as and could be exponentially larger than the rank (Agarwal et al., 2020b).
Finally, F LAMBE operates exactly under the same setting of ours, i.e., the same function approximation and computation oracle.
Our algorithm signiﬁcantly improves the sample complexity from F LAMBE in all parameters .
Note the horizon dependence is not exactly comparable as these prior works originally considered the ﬁnite horizon setting with nonstationary transition, and we convert their results to the discounted setting by simply replacing the ﬁnite horizon H by Θ(1 / (1 − γ )) .
& Wang, 2020; Jin et al., 2020a) which assumes φ is known a priori.
The unknown representation φ makes learning in low-rank MDPs much more challenging than that in linear MDPs since one can no longer directly use linear function approximations.
On the other hand, the fact that linear MDPs can be solved statistical and computational efﬁciently if φ is known a priori implies that if one could learn the representation of the low-rank MDP, one could then efﬁciently learn the optimal policy.
Indeed, prior works have shown that learning in low-rank MDPs is statistically feasible (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021) via leveraging rich function approximators.
However, these algorithms are version space algorithms and are not computationally efﬁcient.
Recent work
F LAMBE proposes an oracle-efﬁcient algorithm 1 that learns in low-rank MDPs with a polynomial sample complexity, where the computation oracle is Maximum Likelihood Estimation (MLE).
In this work, we follow the same setup from F LAMBE (Agarwal et al., 2020b), and propose a new algorithm — Upper Conﬁdence Bound driven Representation Learning, Exploration and Exploitation ( R EP -UCB ) , which can learn a near optimal policy for a low-rank MDP with a polynomial sample complexity and is oracle-efﬁcient.
Comparing to F LAMBE , our algorithm signiﬁcantly improves the sample complexity from O ( d 7 A 9 / ( (cid:15) 10 (1 − γ ) 22 ) for F LAMBE to O
( d 4 A 4 / ( (cid:15) 2 (1 − γ ) 3 ) , where d is the rank of the transition matrix (or dimension of the true representation)
, A is the number of actions, (cid:15) is the suboptimality gap and γ ∈ [0 , 1) is the discount factor in the MDP.
Our algorithm is also arguably much simpler than F LAMBE :
F LAMBE is an explore-then-commit algorithm, has to explore in a layer-by-layer
forward way,
and does not permit data sharing across different time steps.
In contrast, R EP -UCB carefully trades exploration versus exploitation by combining the reward signal and exploration bonus (constructed using the latest learned representation), and enables data sharing across all time steps.
Our sample complexity nearly matches the ones from those computationally inefﬁcient algorithms (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021).
We summarize the comparison with the prior works that study representation learning in Table 1.
In addition to the online exploration setting, we also show that our new techniques can be directly used for designing ofﬂine RL algorithms for low-rank MDPs under partial coverage.
More specifically, we propose an algorithm R EP -LCB— Lower Conﬁdence Bound driven Reprepresentation
Learning for ofﬂine RL , that given an ofﬂine dataset, can learn to compete against any policy (including history-dependent policies) as long as it is covered by the ofﬂine data where the coverage is measured using the relative condition number (Agarwal et al., 2021) associated with the ground truth representation.
Thus, our ofﬂine RL result generalizes prior ofﬂine RL works on linear MDPs (Jin et al., 2020b; Zhang et al., 2021b) which assume representation is known a priori and use linear function approximation.
Computation-wise, our approach uses one call to the MLE computation oracle, and hence is oracle-efﬁcient.
R EP -LCB is the ﬁrst oracle efﬁcient ofﬂine algorithm for low-rank MDP enjoying the aforementioned statistical guarantee.
See Section 2 for a more detailed comparison with the existing literature on representation learning in ofﬂine RL.
Our contributions.
We develop new representation learning RL algorithms that enable sample efﬁcient learning in low-rank MDPs under both online and ofﬂine settings:
In the online episodic learning setting, our new algorithm R EP -UCB integrates representation learning, exploration, and exploitation together, and signiﬁcantly improves the sample complexity of the prior state-of-art algorithm F LAMBE ; 2.
In the ofﬂine learning setting, we propose a natural concentrability coefﬁcient (i
e., relative condition number under the true representation) that captures the partial coverage condition in low-rank MDP, and our algorithm R EP -LCB learns to compete against any policy (including history-dependent ones) under such a partial coverage condition.
R ELATED W ORK
Online Setting We list the comparison as follows, which is summarized in Table 1.
F LAMBE (Agarwal et al., 2020b) was a state-of-the-art oracle-efﬁcient algorithm for low-rank
MDPs.
In all parameters, the statistical complexity is much worse than R EP -UCB .
Our algorithm and F LAMBE operate under the same computation oracle.
F LAMBE does not balance exploration and exploitation, and uses explore-then-committee style techniques (i.e., constructions of absorbing
MDPs (Brafman & Tennenholtz, 2002)) which results in its worse sample complexity.
With a more complex oracle, Mofﬂe (Modi et al., 2021) is a model-free algorithm for low-rank
MDPs, with two additional assumptions: (1) the transition has low non-negative rank (nnr) , and (2) reachability in latent states.
The ﬁrst assumption signiﬁcantly restricts the scope of low-rank MDPs as there are matrices whose nnr is exponentially larger than the rank (Agarwal et al., 2020b).
The sample complexity of Mofﬂe can scale O ( d 6 |A| 13 / ( (cid:15) 2 η 5 (1 − γ ) 5 )) , where η is the reachability probability, and 1 /η could be as large as nnr 1 / 2 (Proposition 4 in Agarwal et al.
(2020b)), which essentially means that Mofﬂe has a polynomial dependence on the nnr.
Thus, Mofﬂe needs the nnr of the transition matrix to be small.
OLIVE (Jiang et al., 2017), Witness rank (Sun et al., 2019) and Bilinear-UCB (Du et al., 2021), when specialized to low-rank MDPs, have slightly tighter dependence on d (e.g., O ( d 2 /(cid:15) 2 ) ).
But these algorithms are computationally inefﬁcient as they are version space algorithms.
Dann et al.
(2021) shows that with a policy class, solving a low-rank MDP can take Ω(2 d ) samples.
In this work, similar to Witness rank (Sun et al., 2019) and F LAMBE , we use function approximators to model the transition.
Thus our positive result is not in contradiction to the result from Dann et al.
VALOR (Dann et al., 2018), PCID (Du et al., 2019a), HOMER (Misra et al., 2020), RegRL (Foster et al., 2020), and the approach from Feng et al. (2020) are algorithms for block MDPs which is a more restricted setting than low-rank MDPs.
These works require additional assumptions such as deterministic transitions (Dann et al., 2018), reachability (Misra et al., 2020; Du et al., 2019a), Bellman completion (Foster et al., 2020), and strong unsupervised learning oracles (Feng et al., 2020).
Ofﬂine Setting We discuss related works in ofﬂine RL.
Uehara & Sun (2021) obtained similar results for ofﬂine RL on low-rank MDPs.
Though the sample complexity in their algorithm is slightly tighter, our algorithm is oracle-efﬁcient, while the CPPO algorithm from Uehara & Sun (2021) is a version space algorithm.
Xie et al.
(2021) propose a (general) pessimistic model-free algorithm in the ofﬂine setting.
We can also apply their algorithm to low-rank MDPs and show some ﬁnite-sample guarantee.
However, it is unclear whether the ﬁnal bounds in their results can be characterized by the relative condition number only using the true representation, and whether they can compete with history-dependent policies.
Thus, our result is still considered superior on low-rank MDPs.
The detail is given in Section E.
In addition to the above two works, the pessimistic approach in ofﬂine RL has been extensively investigated.
Empirically, it can work on simulation control tasks (Kidambi et al., 2020; Yu et al., 2020; Kumar et al., 2020; Liu et al., 2020; Chang et al., 2021).
On the theoretical side, pessimism allows us to obtain the PAC guarantee on various models when a comparator policy is covered by ofﬂine data in some forms (Jin et al., 2020b; Rashidinejad et al., 2021;
Yin et al., 2021; Zanette et al., 2021b; Zhang et al., 2021b; Chang et al., 2021).
However, these algorithms and their analysis rely on a known representation .
We emphasize the differences in Section 5.
P RELIMINARIES
We consider a Markov Decision Process M
=
(cid:104)S , A , P, r, γ, d 0 (cid:105) speciﬁed by a state space S , a discrete action space A , a transition model P : S
×
A → ∆( S ) , a reward function r : S ×
A → R , a discount factor γ ∈ [0 , 1) , and an initial distribution d 0 ∈ ∆( S ) .
To simplify the presentation, we assume r ( s, a ) and d 0 are known (e.g., when d 0 is a probability mass only on s 0 , agent always starts from a ﬁxed initial state s 0 ) 2 .
Following prior work (Jiang et al., 2017; Sun et al., 2019), we assume trajectory reward is normalized , i.e., for any trajectory { s h , a h } ∞
h =0 , we have (cid:80) ∞
h =0 γ h r ( s h , a h ) ∈ [0 , 1] .
Since the ground truth P (cid:63) is unknown, we need to learn it by interacting with environments in an online manner or utilizing ofﬂine data at hand.
We remark that the extension of our all results to the ﬁnite horizon nonstationary case is straightforward.
Here φ (cid:63) ( s, a ) is a distribution over a discrete latent state space Z .
Note that this is still a Markovian model as there is no transition among latent states.
We use the following notation.
Given a policy π : S → ∆( A ) , which maps from state to distribution over actions and some model P , we deﬁne the value function V πP ( s ) =
E (cid:2)(cid:80) ∞
h =0 γ h r ( s h , a h ) | s
= s, P, π (cid:3) to represent the expected total discounted reward of π under P starting at s .
Similarly, we deﬁne the state-action Q function Q πP ( s, a ) :
= r ( s, a ) + γ E s (cid:48) ∼ P ( ·| s,a ) V πP ( s (cid:48) ) .
The expected total discounted reward of a policy π is denoted as V πP,r :=
E s 0 ∼ d 0 V πP ( s 0 ) .
We deﬁne the discounted state-action occupancy distribution
d
πP ( s, a ) =
(1 − γ ) (cid:80) ∞
t =
γ t d πP,t ( s, a ) , where d πP,t ( s, a ) is the probability of π visiting ( s, a ) at time step t under π and P .
We slightly abuse the notation, and denote d πP ( s ) as the state visitation, which is equal to (cid:80) a ∈A d πP ( s, a ) .
When P is the ground truth transition model P (cid:63) , we drop the subscript and simply use d π to denote its state-action distribution.
Unless otherwise noted, Π denotes the class of all polices { S → ∆( A ) } .
We denote total variation distance of P 1 and P 2 by (cid:107) P 1 − P 2 (cid:107) 1 .
Finally, given a vector a , we deﬁne (cid:107) a (cid:107) 2 = √ a (cid:62) a, (cid:107) a (cid:107) B = √ a (cid:62) Ba .
c 0 , c 1 , · · · are universal constants.
We study low-rank MDPs deﬁned as follows (Jiang et al., 2017; Agarwal et al., 2020b).
The conditions on the upper bounds of the norm of φ (cid:63) , µ (cid:63) are just for normalization.
Deﬁnition 1 (Low-rank MDP) .
Given function classes Φ ⊂ S ×
A → R d and Ψ ⊂ S
→ R d , a transition model P (cid:63) : S × A → ∆( A ) admits a low rank decomposition with dimension d ∈ N if there exists two embedding functions φ (cid:63) ∈ Φ , µ
(cid:63) ∈ Ψ such that
[Equation]
where (cid:107) φ ∗ ( s, a ) (cid:107) 2 ≤ 1 for all ( s, a ) and for any function g : S → [0 , 1] , (cid:107) (cid:82) µ (cid:63) ( s ) g ( s ) d
( s ) (cid:107) 2 ≤ √ d .
An MDP is a low rank MDP if P (cid:63) admits such a low rank decomposition.
Low-rank MDPs capture the latent variable model (Agarwal et al., 2020b) shown in Figure 1 where φ (cid:63) ( s, a ) is a distribution over a discrete latent state space Z .
The block-MDP model (Du et al., 2019a) is a special instance of the latent variable model with φ (cid:63) ( s, a ) being a one-hot encoding vector.
Note the linear MDPs (Yang & Wang, 2020; Jin et al., 2020a) assume φ (cid:63) is known.
Next, we explain two settings: the online learning setting and the ofﬂine learning setting.
Then, we present our function approximation setup and computational oracles.
Episodic Online learning In online learning, our overall goal is to learn a stationary policy ˆ π so that it maximizes V ˆ πP (cid:63) ,r , where P (cid:63) is the ground truth transition.
We assume that we operate under the episodic learning setting where we can only reset to states sampled from the initial distribution d 0 (e.g., to emphasize the challenge from exploration, we can consider the special case where we can only reset to a ﬁxed s 0 ).
In the episodic setting, given a policy π , sampling a state s from the state visitation d πP is done by the following roll-in procedure: starting at s 0 ∼ d 0 , at every time step t , we terminate and return s t with probability 1 − γ , and otherwise we execute a t ∼ π ( s t ) and move to t +1 , i.e., s t +1 ∼ P ( ·| s t , a t ) .
Such a sampling procedure is widely used in the episodic learning setting for discounted inﬁnite horizon MDPs (e.g., (Kakade & Langford, 2002; Agarwal et al., 2021; 2020a)).
Ofﬂine learning In the ofﬂine RL, we are given a static dataset in the form of quadruples:
[Equation]
For simplicity, we assume ρ = d π b P (cid:63) , where π b ∈ S → ∆( A ) is a ﬁxed behavior policy.
We denote E
D [ f ( s, a, s (cid:48) )]
= 1 /n (cid:80) ( s,a,s (cid:48) ) ∈D f
( s, a, s (cid:48) ) .
To succeed in ofﬂine RL, we in general need some coverage property of ρ .
One common assumption is that ρ globally covers every possible policies’ state-action distribution, i.e., max π,s,a d πP (cid:63) ( s, a ) /ρ
( s, a ) < ∞
(Antos et al., 2008).
In this work, we relax such a global coverage assumption and work under the partial coverage condition where ρ may not cover distributions of all possible policies.
Instead of competing against the optimal policy under the global coverage, we aim to compete against any policies covered by the ofﬂine data.
In section 5, we will precisely deﬁne the partial coverage condition using the concept of the relative condition number (Agarwal et al., 2021).
Function approximation setup and computational oracles Since µ (cid:63) and φ (cid:63) are unknown, we use function classes to capture them.
Our function approximation and computational oracles are exactly the same as the ones used in F LAMBE .
For completeness, we state the function approximation and computational oracles below.
Assumption 2.
We have a model class M = { ( µ, φ ) : µ ∈ Φ , φ ∈ Φ } , where µ (cid:63) ∈ Φ , φ (cid:63) ∈ Φ .
Following the norm bounds on µ (cid:63) , φ (cid:63) we similarly assume that the same norm bounds hold for our function approximator, i.e., for any µ ∈ Φ , φ ∈ Φ , (cid:107) φ ( s, a ) (cid:107) 2 ≤ 1 , ∀ ( s, a ) and (cid:107) (cid:82) µ ( s ) g
( s ) d ( s ) (cid:107) 2 ≤ √ d, ∀ g : S → [0 , 1] , and (cid:82) µ (cid:62) ( s (cid:48) )
φ ( s, a ) d ( s (cid:48) )
= 1 , ∀ ( s, a ) .
As for computational oracles, we use a supervised learning style MLE oracle.
Deﬁnition 3 (Maximum Likelihood Oracle (MLE)) .
Consider the model class M and a dataset D in the form of ( s, a, s (cid:48) ) , the MLE oracle returns the maixmum likelihood estimator ˆ P :=
(ˆ µ, ˆ φ ) = arg max ( µ,φ ) ∈M E D ln( µ ( s (cid:48) )
(cid:62) φ ( s, a )) .
We also invoke a planning procedure for known linear MDPs with potentially nonlinear rewards , which can be done in polynomial time (after all, we know that online learning in linear MDPs can be done statistically and computationally efﬁcient).
Given a reward r and a model P :=
( µ, φ ) with P ( s (cid:48) | s, a )
=
µ ( s (cid:48) ) (cid:62) φ ( s, a ) (i.e., a known linear transition with a known feature φ ), we can compute the optimal policy arg max π V πP by standard least square value iteration which uses linear regression.
A planning procedure for a known linear MDP is also used in F LAMBE , see Section 5.
in Agarwal et al.
(2020b)
for details of how to implement this procedure for a known linear MDP with polynomial computation complexity.
R EPRESENTATION L EARNING IN O NLINE S ETTING
We consider the online episodic learning setting where the agent can only reset based on the initial state distribution d 0 .
Thus, to ﬁnd a near-optimal policy for a low-rank MDP efﬁciently, we need to carefully interleave representation learning, exploration, and exploitation.
We present our proposed algorithm in the online setting.
This is described in Algorithm 1.
Every iteration, Algorithm 1 rollouts its current policy π to collect a triple ( s, a, s (cid:48) ) where s ∼ d πP (cid:63) , a ∼ U ( A ) , s (cid:48) ∼ P (cid:63) ( ·| s, a ) where U ( A ) is a uniform distribution over actions. Note that to sample s ∼ d πP (cid:63) , we start at s 0 ∼ d 0 , at every time step t , we terminate and return s t with probability 1 − γ , and otherwise we execute a t ∼ π ( s t ) and move to t + 1 , i.e., s t +1 ∼ P (cid:63) ( ·| s t , a t ) .
After n
(ˆ µ n ˆ φ n ) = arg max ( µ,φ ) ∈M E D n (cid:2) ln µ (cid:62) ( s (cid:48) ) φ ( s a ) (cid:3) 7:
Update empirical matrix ˆΣ n = s,a ∈D n ˆ φ n ( s, a )ˆ φ n ( s, a ) (cid:62) + λ n I
ˆ b n ( s, a ) :
= min (cid:18) α n (cid:113) ˆ φ n ( s, a ) (cid:62) ˆΣ − 1 n ˆ φ n ( s, a ) , 2 (cid:19) (1) 9: π n arg max
π V π ˆ P n ,r +ˆ b n 10:
11: Return collecting new data and concatenating it with the existing data, we perform representation learning, i,e, learning a factorization and a representation by MLE (line 6), set the bonus based on the learned feature (Eq. 1), and update the policy via planning inside the learned model with the bonus-enhanced reward (Line 9).
Note the learned transition ˆ P from MLE is linear with the learned feature ˆ φ , and planning in a known linear MDP is known as computationally efﬁcient (Jin et al., 2020a)
(see the explanation after Deﬁnition 3 as well).
Our computation MLE oracle is the same as the one used in F LAMBE (Agarwal et al., 2020b).
4. A NALYSIS
We show that our algorithm has the following PAC bound.
Theorem 4 (PAC Bound for R EP -UCB) .
Fix δ ∈ (0 , 1) , (cid:15) ∈ (0 , 1) .
Let ˆ π be a random output of π 1 , · · · , π N and π (cid:63) = arg max π V πP (cid:63) ,r .
By setting parameters as follows:
[Equation]
with probability at least 1 − δ , we have
[Equation]
where the number of collected samples is at most
N log( N/ [Equation]
[Equation]
The theorem shows that R EP -UCB learns in low-rank MDPs in a statistically efﬁcient and oracleefﬁcient manner.
To the best of our knowledge, this algorithm has the best sample complexity among all oracle efﬁcient algorithms for low-rank MDPs.
Highlight of the analysis Below we highlight our key lemmas and proof techniques.
First, why learning in a low-rank MDP is harder than learning in models with linear structures?
Unlike standard linear models such as linear MDPs (Yang & Wang, 2020; Jin et al., 2020a),
KNRs (Kakade et al., 2020; Abbasi-Yadkori & Szepesv´ari, 2011; Mania et al., 2020; Song &
Sun, 2021), and GP / kernel models (Chowdhury & Gopalan, 2019; Curi et al., 2020), we cannot get uncertainty quantiﬁcation on the model in a point-wise manner.
When models are linear, one can get the following style of point-wise uncertainty quantiﬁcation for the learned model ˆ P : ∀ s, a : (cid:96) ( ˆ P ( ·| s, a ) , P (cid:63) ( ·| s, a )) ≤ σ ( s, a ) where σ ( s, a ) is the uncertainty measure, and (cid:96) is some distance metric (e.g., (cid:96) 1 norm).
With proper scaling, the uncertainty measure σ ( s, a ) is then used for the bonus.
For example, in linear MDPs, given a dataset D = { s, a, s (cid:48) } , we can learn a model ˆ P ( s (cid:48) | s, a ) :
= ˆ µ ( s (cid:48) ) (cid:62) φ (cid:63) ( s, a ) , and get point-wise uncertainty quantiﬁcation: ∀
[Equation]
a
[Equation]
for some family of functions f : S → R with Σ φ (cid:63) =
(cid:80) s,a ∈D φ (cid:63) ( s, a ) φ (cid:63) ( s, a )+ λI (Lykouris et al., 2021;
Neu & Pike-Burke, 2020).
To set the scaling c properly, since φ (cid:63) is known a priori, the linear regression analysis applies here, and one can apply Cauchy-Schwarz inequality to the LHS of (2) to pull out φ (cid:63) and get an upper bound in the form of
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
where c is set to be the linear regression training error measured in the term ( b ) above.
However, when we jointly learn µ and φ , since nonlinear function approximation is used 3 , we cannot get point-wise uncertainty quantiﬁcation via linear regression-based analysis.
We stress that our bonus is not designed to capture the uncertainty quantiﬁcation on the model error between ˆ
P ( ·| s, a ) =
ˆ µ (cid:62) ˆ φ ( s, a ) and P (cid:63) ( ·| s, a ) = µ (cid:63) (cid:62) φ (cid:63) ( s, a ) in a point-wise way, which is not tractable as ˆ P and P (cid:63) does not even share the same representation.
Instead, the bonus is carefully designed so that it only provides near-optimism at the initial state distribution .
This is formalized as follows.
Lemma 5 (Almost Optimism at the Initial State Distribution) .
Set the parameters as in Theorem 4.
With probability 1 − δ ,
[Equation]
We remark that the idea of optimism with respect to the initial state distribution has been used in prior works (Jiang et al., 2017; Sun et al., 2019; Du et al., 2021; Zanette et al., 2020).
However, these algorithms are not computationally efﬁcient (i.e., they use version space instead of reward bonus), and their version-space based analysis is different from ours.
Proof sketch for Lemma 5 We start by using the simulation lemma (Lemma 20) inside the learned model which is important since our bonus ˆ b n uses ˆ φ n associated with the learned model
ˆ P n :
[Equation]
from where we show that E s,a ∼ d π ˆ Pn
[Equation]
P
[Equation]
[Equation] s [Equation] ·| s, a ) (cid:107) 1 as a whole nearly lower bounds
E s,a ∼ d π ˆ Pn [ˆ b n ( s, a )] .
The detailed procedure is illustrated in Lemma 8 in Appendix B.
Second, our bonus is using representation ˆ φ n that is being updated every episode, and our empirical covariance matrix ˆΣ n is also updated whenever we update ˆ φ n , which means that standard elliptical potential based analysis (i.e., analysis used in linear bandits/MDPs with known features) cannot work here as our feature changes every episode.
Instead, in our analysis, we have to keep tracking a potential function that is deﬁned using the unknown ground truth representation φ (cid:63) , i.e., the elliptical potential (cid:107) φ (cid:63) ( s, a ) (cid:107) 2Σ − 1 ρn,φ(cid:63) , where
[Equation]
and ρ n is the marginal state-action distribution of D n at an iteration n , i.e., ρ n ( s, a ) = {+(cid:80) n −i =
+}
d π i P (cid:63) ( s ) U ( a ) /n .
Since this potential function uses the ﬁxed representation φ (cid:63) , we can apply the standard elliptical potential argument to track the progress that our algorithm makes during learning.
Below we illustrate the procedure of linking the bonus under ˆ φ n to the potential function [Equation] deﬁned with respect to the true feature φ (cid:63) .
Linking bonus under ˆ φ n to the elliptical potential function under φ (cid:63)
With near optimism, using the simulation lemma (Lemma 20), we can upper bound the per-iteration regret as follows: c where f n ( s, a ) :
= (cid:107) ˆ P n ( · | s, a ) − P (cid:63) ( · | s, a ) (cid:107) 1 .
To connect the right-hand side of the above inequality to the elliptical potential under the ﬁxed feature φ (cid:63) , we show that for any function g ∈
S × A → [0 , B ] for B ∈ R + ,
[Equation]
See Lemma 13 in Appendix B. By substituting g with ˆ b n + f n / (1 − γ ) , the ﬁrst term of the RHS of the above inequality can be upper bounded as:
[Equation]
In the term ( G 2 ), we expect n E ρ n [ f 2 n ( s, a )] to be O (1) as E s,a ∼ ρ n [ f 2 n ( s, a )] is in order of 1 /n due to the fact that it is the generalization bound of the MLE estimator ˆ P n under the training distribution ρ n .
For n E ρ n [ˆ b 2 n ( s, a )] , we expect it to be in the order of d as the (unnormalized) data covariance matrix ˆΣ n in the bonus ˆ b n uses training data from ρ n , i.e., we are measuring the expected bonus under the training distribution.
In other words, the term ( G 2 ) scales in order of poly ( d ) .
For the term ( G 1 ) , since it contains the potential function based on φ (cid:63) , the sum of the term ( G 1 ) over all episodes can be controlled by the standard elliptical potential argument (see Lemma 18 and Lemma 19).
R EPRESENTATION L EARNING IN O FFLINE S ETTING
In this section, we study representation learning in the ofﬂine setting.
We consider the setting where the ofﬂine data does not have a full global coverage.
We present our algorithm Lower Conﬁdence Bound driven Representation Learning in ofﬂine RL (R EP -LCB) in Algorithm 2.
Our proposed algorithm consists of three parts.
The ﬁrst part is MLE which learns a model ˆ P and a representation ˆ φ .
The second part is the construction of a penalty term ˆ b .
Using the learned representation ˆ φ , we use a standard bonus in linear bandits as the penalty term as if ˆ φ were the true feature.
The third part is planning with the learned model ˆ P and reward r − ˆ b .
5. A NALYSIS
We present the PAC guarantee of R EP -LCB.
Before proceeding, we deﬁne a relative condition number as a mean to measure the deviation between a comparator policy π and the ofﬂine data:
[Equation]
In the special case where the MDP is just a tabular MDP (i.e., φ (cid:63) is a one-hot encoding vector), this is reduced to a density ratio C (cid:63) ∞ =
max s,a d π P (cid:63) ( s, a ) /ρ
( s, a ) .
The relative condition number C (cid:63)π is always no larger than the density ratio and could be much smaller for MDPs with large state spaces.
Note that we quantify the relative condition number using the unknown true representation φ (cid:63) .
With the above setup, now we are ready to state the main theorem for R EP -LCB.
Theorem 6 (PAC Bound for R EP -LCB) .
Let ω = max a,s (1 /π b ( a | s )) .
Denote ˆ π as the output of
R EP -LCB .
There exists a set of parameters such that with probability at least 1 − δ , for any policy π (including history-dependent non-Markovian policies),
[Equation]
See Theorem 15 in Appendix B for the detailed parameters.
We explain several implications.
First of all, this theorem shows that we can uniformly compete with any policy including history-dependent non-Markovian policies 4 satisfying the partial coverage C (cid:63)π < ∞ .
Particularly, if the optimal policy π (cid:63) is covered by the ofﬂine data, i.e., C (cid:63)π (cid:63) < ∞ , then our algorithm is able to compete against it.
Note that assuming ofﬂine data covers π (cid:63) is still a weaker assumption than the global coverage such as sup π sup ( s,a )
d πP (cid:63) ( s, a ) /ρ
( s, a ) in prior ofﬂine RL works (Antos et al., 2008; Chen & Jiang, 2019).
Second, our coverage condition is measured by a relative condition number deﬁned using the unknown ground truth representation φ (cid:63) but not depending on other features.
Prior works that use relative condition numbers as measures of coverage are restricted to the settings where the ground truth representation φ (cid:63) is known (Jin et al., 2020b; Chang et al., 2021; Zanette et al., 2021b).
To sum up, our algorithm is the ﬁrst oracle efﬁcient algorithm which does not need to know φ (cid:63) , and requires partial coverage only in terms of φ (cid:63) .
Note while Uehara & Sun (2021) has a similar guarantee on low-rank MDPs, their algorithm is not oracle-efﬁcient as it is a version space algorithm.
Highlight of the analysis Prior ofﬂine RL works that use reward penalties (Rashidinejad et al., 2021;
Zhang et al., 2021b;
Chang et al., 2021) all assume the representation φ (cid:63) is known a priori, which allows them to use linear regression analysis to derive point-wise model uncertainty quantiﬁcation which in turn serves as a penalty.
In contrast, our function class is nonlinear.
As in the online setting, what we can show is near pessimism in the initial state distribution d 0 .
Lemma 7 (Almost Pessimism at the Initial State Distribution) .
There exists a set of parameters that with probability 1 − δ , we have
[Equation]
We leave the detailed proof to Lemma 14 in the Appendix C, which is similar to the proof of
Lemma 5.
Though the pessimism in the initial state distribution was recently also used in Xie et al.
; Zanette et al.
(2021b), the derivation is totally different.
Similar to the online setting, to obtain a result only depending on the relative condition number using the ground truth representation φ (cid:63) but not the learned feature ˆ φ , we again need to translate the penalty deﬁned with ˆ φ to the potential function (cid:107) φ (cid:63) ( s, a ) (cid:107) Σ − 1 ρ,φ(cid:63) .
The same techniques that we used for the online setting can be leveraged here to achieve the above goal.
C ONCLUSION
We study online/ofﬂine RL on low-rank MDPs, where the ground truth feature is not known a priori.
For online RL, our new algorithm R EP -UCB signiﬁcantly improves the sample complexity of the piror state-of-the-art algorithm F LAMBE in all parameters while using the same computational oracles.
R EP -UCB has the best sample complexity among existing oracle efﬁcient algorithms for low-rank MDPs by a margin.
Comparing to prior representation learning works on low-rank MDPs and block MDPs that rely on a forward step-by-step reward-free exploration framework, our algorithm interleaves representation learning, exploration, and exploitation together, and learns a single stationary policy.
For ofﬂine RL, our new algorithm R EP -LCB is the ﬁrst oracle efﬁcient algorithm for low-rank MDPs that has a PAC guarantee under a partial coverage condition measured by the relative condition number deﬁned with the true feature representation. Given π = { π i } ∞ i =0 where π i depends on s 0 , a 0 , . . . s i , V πP (cid:63) ,r and d πP (cid:63) ( s, a ) are still well-deﬁned. We also require ω < ∞ , which is a mild assumption since it does not involve P (cid:63) .
Indeed, it is much weaker than the global coverage type assumption 1 /ρ ( s, a ) < ∞ , ∀ ( s, a ) . A M ORE R ELATED W ORK Here, we mention several related works in addition to Section
Online setting We mention works that tackle representation learning in quite different settings.
Hao et al.
(2021b) consider feature selection in sparse linear MDPs with a given exploratory distribution.
Zhang et al.
(2021a) consider how to choose the best representation among correct representations inspired by Papini et al.
(2021) (i.e., the MDP is a linear MDP under any representation in the function class Φ ).
In contrast, we only assume the MDP is linear under some φ (cid:63) ∈ Φ .
Ofﬂine setting we mention works that tackle representation learning from different viewpoints and settings.
Lu et al.
(2021) consider multitask representation learning under a generative model assumption.
Hao et al.
(2021a) study the feature selection problem in sparse linear MDPs and Ni et al.
(2021) study dimensionality reduction in a given kernel space, under the assumption that the ofﬂine data admits some form of full coverage condition.
B P ROOF OF THE THEORETICAL PROPERTY OF R EP -UCB
Hereafter, we assume c 0 , c 1 , · · · , are some universal constants, and the notation f (1 / (1 −
[Equation]
) , |A
[Equation]
, ln
[Equation] [Equation]
) , ln
[Equation]
|M
[Equation]
[Equation]
, n ) (cid:46) g (1 / (1 − γ ) , |A| , ln(1 /δ ) , ln( |M| ) , d, n ) means there exists some constant c 1 > 0 , such that f (
[Equation]
[Equation]
[Equation]
| , ln
[Equation]
[Equation]
[Equation]
M
[Equation] [Equation] [Equation] [Equation]
[Equation]
[Equation] [Equation] [Equation] |M|
[Equation]
, d, n ) for any 0 ≤ γ < 1 , |A| , ln(1 /δ ) , ln( |M| ) , d, n .
We ﬁrst prove a lemma to show the optimism.
Hereafter, we use the following notation to denote the data sampling distribution
[Equation]
where u ( a ) denotes the uniform distribution over actions.
First, we prove the optimism at the initial distribution.
This is proved by using a simulation lemma inside the learned model which is important since both the bonus and the learned model use ˆ φ .
In high level, we will show that the expected bonus E s,a ∼ d π ˆ Pn ˆ b n ( s, a ) is in the same order of the expected model error E s,a ∼ d π ˆ Pn (cid:107) ˆ P n ( ·| s, a ) − P (cid:63) ( ·| s, a ) (cid:107) 1 .
Note that the expectation is with respect to d π ˆ P n .
Lemma 8 (Almost Optimism at the Initial Distribution) .
Consider an episode n (1 ≤ n ≤ N ) and set
[Equation]
With probability 1 − δ , we have
[Equation]
Proof.
We deﬁne
[Equation]
In this proof, letting f n ( s, a ) = (cid:107) ˆ P n ( · | s, a ) − P (cid:63) ( · | s, a ) (cid:107) 1 , we condition on the event ∀ n,
[Equation]
n
[Equation]
, a )] ≤
[Equation]
n
[Equation]
where
[Equation]
From Lemma 11 and Lemma [Equation] this event happens with probability 1 − δ .
Then, for any policy π , from simulation lemma 20,
[Equation]
where in the last step
[Equation]
we replaced the empirical covariance by the population covariance. Note the notation (cid:46) is up to universal constants.
Here, since (cid:107) V πP (cid:63) ,r (cid:107) ∞ ≤ 1 (since we assume trajectory-wise total reward is normalized between [0 , 1] ), we have:
[Equation]
This is further bounded by Lemma 12:
[Equation]
where α
[Equation]
Note we here use f n ( s, a ) ≤ 2 and E ρ n [ f n ( s, a ) 2 ] ≤ ζ n .
Combining all things together,
[Equation]
where α n := (cid:112) α (cid:48) [Equation]
Going back to (3)
[Equation]
we have
[Equation]
≥ − (cid:112) (1 − γ ) |A| ζ n .
From the second line to the third line, we again use (cid:107) V πP (cid:63) ,r (cid:107) ∞ ≤ O (1) and (4).
This concludes the proof.
Next, we obtain the upper bound of (cid:80) Nn =
V π (cid:63) P (cid:63) ,r − V π n P (cid:63) ,r .
Recall π (cid:63) is the optimal policy.
Though this form is the same as a standard regret form, since we are not exactly deploying π n in episode n (recall that we play a uniform action at the end of the episode), we cannot get the regret guarantee.
However, it sufﬁces for the PAC guarantee.
Lemma 9 (Regret) .
With probability 1 − δ , we have N
[Equation]
Proof.
Similar to Lemma 8, letting f n ( s, a ) = (cid:107) ˆ P n [Equation]
[Equation]
, a )
[Equation]
P
[Equation]
| [Equation] , a
[Equation]
1 , we condition on the event
[Equation]
From Lemma 11 and Lemma 17, this event happens with probability 1 − δ .
For any ﬁxed episode n and any policy π , we have
[Equation] [Equation]
We use the 2nd form of simulation Lemma 20 in the last display.
Then, noting (cid:107) ˆ b n (cid:107) ∞ ≤ 2 , we have (cid:107) V π n ˆ P n ,r +ˆ b n
[Equation]
≤ [Equation] [Equation] − γ ) .
Combining this fact with the above expansion, we have
[Equation]
where in the last inequality we use the notation d ¯ πP (cid:63) ( s, a ) :
= d πP (cid:63) ( s ) · U ( a ) and an important weighting step.
First, we calculate the ﬁrst term (a) in Inequality 7.
Following Lemma 13 and noting the bonus ˆ b n is O (1) , we have
[Equation]
Note that we use the fact that B = 2 when applying Lemma 13.
In addition, we have
[Equation]
Then,
[Equation]
Second, we calculate the term (b) in inequality 7.
Following Lemma 13 and noting f 2 n ( s, a ) is upper-bounded by 4 (i
[Equation]
e., B = 4 in Lemma 13), we have [Equation]
[Equation]
where in the second inequality, we use E s,a ∼ ρ f 2 n ( s
[Equation]
a ) ≤ ζ n , and in the last line, recall √
[Equation]
Then, by combining the above calculation of the term (a) and term (b) in inequality 7, we have:
[Equation]
Hereafter, we take the dominating term out.
First, recall
[Equation]
Second, we also use
[Equation]
Finally,
[Table]
n
[Equation]
(Some algebra.
We take the dominating term out.)
[Equation]
This concludes the proof.
Using Lemma 9, we can immediately obtain the PAC guarantee.
Theorem 10 (PAC guarantee of R EP -UCB) .
By interacting with the environment for a number of steps at most
N log( N/δ )
[Equation]
with probability 1 − δ , we can ensure V π (cid:63) P (cid:63) ,r − V
ˆ π P (cid:63) ,r ≤ (cid:15).
Proof.
From Lemma 9 and Lemma 22, when N is
[Equation]
with probability 1 − δ , we can ensure
[Equation]
With probability 1 − δ , we need ln(1 /δ ) interactions with the environment to get one triple ( s, a, s (cid:48) ) from d πP (cid:63) ( s, a ) .
Thus, the total sample complexity is O ( N (1 − γ ) − 1 ln( N/δ )) .
Next, we provide an important lemma to ensure the concentration of the bonus term.
The version for ﬁxed φ is proved in Zanette et al.
(2021a, Lemma 39).
Here, we take a union bound over the whole feature φ ∈ Φ .
Recall
[Equation]
The extension to the ofﬂine case is again straightforward.
Recall the deﬁnition of d ¯ π ( s, a ) in the beginning of Appendix.
Lemma 11 (Concentration of the bonus term) .
Set λ n = Θ( d ln( n | Φ | /δ )) for any n .
Deﬁne Σ
[Equation]
n ,φ = n E
[Equation]
With probability 1 − δ , we have
[Equation]
Next, we show our key lemma leveraging low-rank MDP structures.
For any g ∈ S ×
A → R , This lemma shows that E ( s,a ) ∼ d π ˆ Pn { g ( s, a ) } can be upper-bounded using E (˜ s, ˜ a ) ∼ d π ˆ Pn (cid:107) ˆ φ n (˜ s, ˜ a ) (cid:107) Σ − 1 ρn, ˆ φn as long as we have the convergence guarantee for E ( s,a ) ∼ ρ n (cid:2) g 2 ( s, a ) (cid:3) .
Lemma 12 (One-step back inequality for the learned model) .
Take any g ∈ S ×
A → R s.t.
(cid:107) g (cid:107) ∞ ≤ B .
We condition on the event where the MLE guarantee (15) :
[Equation]
holds.
Then, for any policy π , we have
[Equation]
Recall
[Equation]
Proof.
First, we have an equality:
[Equation]
The second term in (8) is upper-bounded by
[Equation]
−
[Equation]
Next we consider the ﬁrst term in (8).
By CS inequality, we have
[Equation]
Then, [Equation] a
[Equation]
n ( s ) (cid:62)
[Equation]
n
[Equation]
[Equation] (Use the assumption (cid:107) (cid:80) a π ( a | s ) g ( s, a ) (cid:107) ∞ ≤ B and (cid:82) (cid:107) ˆ µ n ( s ) h ( s ) d ( s ) (cid:107) 2 ≤ √ d for any h : S → [0 , 1] .)
[Equation]
Jensen
[Equation]
(Importance sampling)
[Equation] [Equation] [Equation] (cid:3)(cid:27) + λ n
[Equation]
nζ
[Equation]
Then, the ﬁnal statement is immediately concluded.
Here, we show a similar lemma as Lemma 12.
The difference is we aim for calculating E ( s,a ) ∼ d πP(cid:63) { g ( s, a ) } instead of E ( s,a ) ∼ d π ˆ Pn { g ( s, a ) } .
For any g ∈ S ×
A → R , this lemma shows that E ( s,a ) ∼ d πP(cid:63) { g ( s, a ) } can be upper-bounded using E (˜ s, ˜ a ) ∼ d πP(cid:63) (cid:107) φ (cid:63) (˜ s, ˜ a ) (cid:107) Σ − 1 ρn, ˆ φ(cid:63) as long as we have the convergence guarantee for E ( s,a ) ∼ ρ n (cid:2) g 2 ( s, a ) (cid:3) .
Note comparing to Lemma 12, this is not a probabilistic statement.
Note that (cid:107) φ (cid:63) ( s, a ) (cid:107) Σ − 1 ρn,φ(cid:63) is the usual elliptical potential function under the ﬁxed representation φ (cid:63) .
Lemma 13 (One-step back inequality for the true model ) .
Take any g ∈ S ×
A → R such that (cid:107) g (cid:107) ∞ ≤ B .
Then,
[Equation]
Recall Σ
[Equation]
Proof.
First [Equation] we have
[Equation]
The second term in (9) is upper-bounded by
[Equation]
By CS inequality, the ﬁrst term in (9) is further bounded as follows:
[Equation]
Here, we have
[Equation]
≤ [Equation] [Equation]
Therefore
[Equation]
Finally
[Equation]
we have
[Equation]
This concludes the proof.
C P ROOF OF THE THEORETICAL PROPERTY OF R EP -LCB
This section provides the detailed proofs for our results in the ofﬂine setting.
Below we ﬁrst prove that V π ˆ P,r − ˆ b is an almost pessimistic estimator of V πP (cid:63) ,r .
Lemma 14 (Almost Pessimism at the Initial Distribution) .
Let ω = max a,s 1 /π b ( a | s ) .
Set
[Equation]
With probability 1 − δ , for any policy π , we have
[Equation]
n .
Proof.
We deﬁne
[Equation]
where λ
= O
( d ln( |M| /δ )) .
In this proof, letting f ( s, a ) = (cid:107) ˆ P ( · | s, a ) − P (cid:63) ( · | s, a ) (cid:107) 1 , we condition on the events:
[Equation]
where ζ = O (ln( |M| /δ ) /n ) .
From the ofﬂine version of Lemma 17 and Lemma 11 6 , this event happens with probability 1 − δ .
Then, from simulation lemma (Lemma 20),
[Equation]
Here, we have
[Equation]
noting (cid:107) V πP (cid:63) ,r (cid:107) ∞ ≤ 1 .
This is further bounded by Lemma 16:
[Equation]
where
[Equation]
Here, we use f ( s, a ) ≤ 2 in Lemma 16 and E ( s,a ) ∼ ρ
[ f 2 ( s, a )] ≤ ζ .
Thus,
[Equation]
Going back to the simulation lemma
[Equation]
we have
[Equation]
This concludes the proof
[Equation]
With the above lemma
[Equation]
now we can proceed to prove the main theorem.
Theorem 15 (PAC guarantee of R EP -LCB) .
Set the parameters as in Lemma 14.
With probability 1 − δ [Equation] for any comparator policy π including history-dependent non-Markovian policies, we have
[Equation]
where C (cid:63)π is the relative condition number under φ (cid:63) :
[Equation]
Proof.
In this proof, letting f ( s, a ) = (cid:107) ˆ P ( · | s, a )
[Equation]
P
[Equation]
· | s
[Equation]
a
[Equation]
[-( · | s, a ) (cid:107) 
-]
1 we condition on the events:
[Equation]
From Lemma 11 and Lemma 17, this event happens with probability 1 − δ .
For any policy π , we have
[Equation]
Recall f
[Equation]
s [Equation] From the second line to the third line, note though ˆ π is the argmax over Markoovian polices, ˆ π is also the argmax over all history-dependent polices.
In the last line, we use a simulation lemma 20, which is tailored to a time-inhomogeneous policy.
We here use (cid:107) V
π ˆ P,r − ˆ b (cid:107) ∞ ≤ 2 / ((1 − γ )) .
noting (cid:107) ˆ b (cid:107) ∞
=
O (1) .
We further calculate the ﬁrst term (a).
Considering the ofﬂine version of Lemma 13 and noting (cid:107) ˆ b [Equation] ∞ ≤ 2 , we have
[Equation]
From (12), we have
[Equation]
] ≤ d.
Thus,
[Equation]
Second
[Equation]
we further calculate the second term (b).
Considering the ofﬂine version of Lemma 13 and noting f 2 ( s, a ) is upper-bounded by 4 ,
[Equation]
γ
[Equation]
In the ﬁnal line, recall (cid:112) ω { nγζ } + γλd + γnζ ≤ α.
Finally, by combining the calculation of the ﬁrst term (a) and the second term (b), we have
[Equation]
Now, we use the fact E
[Equation]
s,
[Equation]
a
[Equation]
d πP(cid:63)
[Equation]
φ
[Equation]
˜ s
[Equation]
a )
[Equation]
is upper-bounded as
[Equation]
Lemma
[Equation]
From (13))
Finally [Equation] we have
[Equation]
out [Equation] dominating terms)
[Equation]
n
[Equation]
The lemma below is a key technical lemma for our proof.
It shows that one can relate the expected value of any function f ( s, a ) with respect to d π ˆ P (i.e., inside the learned model ˆ P ) to the potential function with respect to d π ˆ P , i.e., E (˜ s, ˜ a ) ∼ d π ˆ P (cid:107) ˆ φ (˜ s, ˜ a ) (cid:107) Σ − 1 ρ, ˆ φ .
Pairing ˆ φ and ˆ P is important since ˆ P is the low-rank transition model deﬁned using ˆ φ .
As we have seen in the above analysis, when using the lemma below, we instantiate f ( s, a ) :
= (cid:107) [Equation] P ( [Equation]
s,
[Equation]
P (cid:63) (
[Equation]
| s,
[Equation]
(cid:107) 1 .
Lemma 16 (One-step back inequality for the learned model in ofﬂine setting) .
Take any f ⊂ S ×
A → R s.t.
(cid:107) f (cid:107) ∞ ≤ B .
We condition on the event where the MLE guarantee holds:
[Equation]
Then, letting ω = max s,a (1 /π b ( a | s )) , for any policy π , we have
[Equation]
E
[Equation]
Proof.
First, we have an equality:
[Equation]
The second term in (14) is upper-bounded by
[Equation]
Next we consider the ﬁrst term in (14).
By CS inequality [Equation] we have
[Equation]
Then,
[Equation]
Use the assumption (cid:107)
[Equation]
a ) (cid:107)
[Equation]
and (cid:107) (cid:82) ˆ µ (
[Equation]
√
[Equation]
for h : S
[Equation]
[Equation]
[Equation]
MLE
[Equation]
(cid:107) E
[Equation]
Finally, the ﬁrst term in (14) is upper-bounded by
[Equation]
(Deﬁnition of ρ )
The ﬁnal statement is immediately concluded.
D A UXILIARY LEMMAS
First, we present the MLE guarantee.
Regarding the proof, refer to Agarwal et al.
(2020b, Theorem 21).
Note ˆ P n and ¯ π n are the quantities appearing in the proposed online algorithm.
We can also immediately obtain the statement to the ofﬂine case.
Lemma 17 (MLE guarantee) .
For a ﬁxed episode n , with probability 1 − δ ,
[Equation]
n . As a straightforward corollary, with probability 1 − δ , ∀ n ∈ N + , n
[Equation]
The following is a standard inequality to prove regret bounds for linear models.
Refer to Agarwal et al.
(2020a, Lemma G.2.)
Lemma 18.
Consider the following process.
For n = 1 , · · · , N , M n = M n − 1 + G n with M 0 = λ
I and G n being a positive semideﬁnite matrix with eigenvalues upper-bounded by 1 .
We have that:
[Equation]
Lemma 19 (Potential function lemma) .
Suppose Tr( G n ) ≤ B 2 .
[Equation]
Proof.
Let σ 1 , · · · , σ d be the set of singular values of M N recalling M N is a positive semideﬁnite matrix.
Then, by the AM-GM inequality,
[Equation]
Since we have (cid:80) i σ i =
Tr( M N ) ≤ dλ 0 + NB [Equation] the statement is concluded.
Lemma 20 (Simulation lemma)
[Equation] Given two MDPs ( P (cid:48) , r + b ) and ( P, r ) [Equation] for any policy π
[Equation]
we have:
[Equation]
and
[Equation]
Proof.
We use
[Equation]
, π )]
[Equation]
f ( s, a )]]
Then,
[Equation]
Similarly,
[Equation]
The following lemma is used to deal with the distribution shift in the ofﬂine setting.
For the proof, refer to Chang et al.
Lemma 21 (Distribution shift lemma) .
Consider any policy π and state-action distribution ρ , and any representation φ (cid:63) , we have:
[Equation]
This is some auxiliary lemma to convert the ﬁnite sample error bound into the sample complexity.
Lemma 22 (Conversion of ﬁnite sample error bounds into sample complexities) .
By taking
[Equation]
It satisﬁes
[Equation]
where c is a constant independent of a 1 , a 2 , a 3 .
Proof.
We ﬁrst have
[Equation]
Here, we use
[Equation]
Then, we prove when N = 1 /(cid:15) 2 × ln 2 (1 + 1 /(cid:15)
[Equation] [Equation]
This is proved by
[Equation]
From the third line to the fourth line, we use ln( x ) ≤ 0 .
5( x 1 / 2 − 1) for x > 0 .
Then, the ﬁnal statement is concluded.
E M ORE COMPARISON TO X IE ET AL . (2021)
We brieﬂy explain the guarantee when we use Algorithm 1 (Xie et al., 2021).
For a given reward r , we ﬁrst deﬁne a new feature class Φ + r .
Deﬁnition 23 (Augmented feature) .
Let φ =
[ φ 1 , · · · , φ d ] .
[Equation]
Then, we set
[Equation]
where c is some suitable constant.
Given the hypothesis class F for the Q-function, we can run
Algorithm 1 in (Xie et al., 2021).
We denote the output policy as ˆ π .
We check two assumptions to ensure the algorithm works.
The ﬁrst assumption is realizability.
This is satisﬁed since for any policy π ∈ Π ( Π is the class of all Markovian polices), we have Q πP (cid:63) ,r ∈ F .
The second assumption is completeness.
This is also satisﬁed since T πP (cid:63) ,r F ⊂ F for any policy π ∈ Π where T π P (cid:63) ,
r is a Bellman-operator s.t.
[Equation]
where we denote f ( s, π )
=
E a ∼ π ( s ) f ( s, a ) .
Then, by invoking their Corollary 5, we have
Theorem 24 (PAC bound based on Xie et al.
With probability 1 − δ ,
[Equation]
where
[Equation]
We compare the above result with our result in Theorem 6.
First, since C † r includes r and all possible features in Φ , this partial coverage condition is stronger than ours (recall our partial coverage condition is only related to the true representation φ (cid:63) ), and we always have C (cid:63)π ≤ C † π,r .
Secondly, the dependence on n is much worse.
Third, it is unclear whether the learned policy can compete against any history-dependent policy.
Recall in Theorem 6, we show that our algorithm can compete with any history-dependent policies.