Representation Learning for Online and Offline RL in Low-rank MDPs

1 INTRODUCTION . When applying Reinforcement Learning ( RL ) to large-scale problems where data is complex and high-dimensional , learning effective transformations of the data , i.e. , representation learning , can often significantly improve the sample and computation efficiency of the RL procedure . Indeed , several empirical works have shown that leveraging representation learning techniques developed in supervised or unsupervised learning settings can accelerate the search for good decision-making strategies ( Silver et al. , 2018 ; Stooke et al. , 2021 ; Srinivas et al. , 2020 ; Yang & Nachum , 2021 ) . However , representation learning in RL is far more subtle than it is for non-sequential and noninteractive learning tasks ( e.g. , supervised learning ) . Prior works have shown that even if one is given the magic representation that exactly linearizes the optimal policy ( Du et al. , 2019b ) or the optimal value functions ( Wang et al. , 2020 ; Weisz et al. , 2021 ) , RL is still challenging ( i.e. , one may still need exponentially many samples to learn ) . This indicates that an effective representation that permits efficient RL needs to encode more information about the underlying Markov Decision Processes ( MDPs ) . Despite the recent empirical success of representation learning in RL , its statistical guarantee and theoretical properties remain under-investigated . In this work , we study the representation learning question under the low-rank MDP assumption . Concretely , a low-rank MDP assumes that the MDP transition matrix admits a low-rank factorization , i.e. , there exists two unknown mappings µ ( s′ ) , φ ( s , a ) , such that P ( s′|s , a ) = µ ( s′ ) > φ ( s , a ) for all s , a , s′ , where P ( s′|s , a ) is the probability of transiting to the next state s′ under the current state and action ( s , a ) . The representation φ in a low-rank MDP not only linearizes the optimal stateaction value function of the MDP ( Jin et al. , 2020a ) , but also linearizes the transition operator . A lowrankness assumption on large stochastic matrices is a common assumption and has enabled successful development of algorithms for real world applications such as movie recommendation systems ( Koren et al. , 2009 ) . We note that a low-rank MDP strictly generalizes the linear MDP model ( Yang & Wang , 2020 ; Jin et al. , 2020a ) which assumes φ is known a priori . The unknown representation φ makes learning in low-rank MDPs much more challenging than that in linear MDPs since one can no longer directly use linear function approximations . On the other hand , the fact that linear MDPs can be solved statistical and computational efficiently if φ is known a priori implies that if one could learn the representation of the low-rank MDP , one could then efficiently learn the optimal policy . Indeed , prior works have shown that learning in low-rank MDPs is statistically feasible ( Jiang et al. , 2017 ; Sun et al. , 2019 ; Du et al. , 2021 ) via leveraging rich function approximators . However , these algorithms are version space algorithms and are not computationally efficient . Recent work FLAMBE proposes an oracle-efficient algorithm1 that learns in low-rank MDPs with a polynomial sample complexity , where the computation oracle is Maximum Likelihood Estimation ( MLE ) . In this work , we follow the same setup from FLAMBE ( Agarwal et al. , 2020b ) , and propose a new algorithm — Upper Confidence Bound driven Representation Learning , Exploration and Exploitation ( REP-UCB ) , which can learn a near optimal policy for a low-rank MDP with a polynomial sample complexity and is oracle-efficient . Comparing to FLAMBE , our algorithm significantly improves the sample complexity from O ( d7A9/ ( 10 ( 1− γ ) 22 ) for FLAMBE to O ( d4A4/ ( 2 ( 1− γ ) 3 ) , where d is the rank of the transition matrix ( or dimension of the true representation ) , A is the number of actions , is the suboptimality gap and γ ∈ [ 0 , 1 ) is the discount factor in the MDP . Our algorithm is also arguably much simpler than FLAMBE : FLAMBE is an explore-then-commit algorithm , has to explore in a layer-by-layer forward way , and does not permit data sharing across different time steps . In contrast , REP-UCB carefully trades exploration versus exploitation by combining the reward signal and exploration bonus ( constructed using the latest learned representation ) , and enables data sharing across all time steps . Our sample complexity nearly matches the ones from those computationally inefficient algorithms ( Jiang et al. , 2017 ; Sun et al. , 2019 ; Du et al. , 2021 ) . We summarize the comparison with the prior works that study representation learning in Table 1 . In addition to the online exploration setting , we also show that our new techniques can be directly used for designing offline RL algorithms for low-rank MDPs under partial coverage . More specifically , we propose an algorithm REP-LCB—Lower Confidence Bound driven Reprepresentation Learning for offline RL , that given an offline dataset , can learn to compete against any policy ( including history-dependent policies ) as long as it is covered by the offline data where the coverage is measured using the relative condition number ( Agarwal et al. , 2021 ) associated with the ground 1The oracle generally refers to empirical risk minimization oracle employing non-bilevel optimization . We seek to design an algorithm that runs in polynomial time , with each oracle call counting as O ( 1 ) . This algorithmic framework has lead to many successful algorithms in the literature of contextual bandit ( Agarwal et al. , 2014 ; Dudı́k et al. , 2017 ; Foster & Rakhlin , 2020 ) . truth representation . Thus , our offline RL result generalizes prior offline RL works on linear MDPs ( Jin et al. , 2020b ; Zhang et al. , 2021b ) which assume representation is known a priori and use linear function approximation . Computation-wise , our approach uses one call to the MLE computation oracle , and hence is oracle-efficient . REP-LCB is the first oracle efficient offline algorithm for low-rank MDP enjoying the aforementioned statistical guarantee . See Section 2 for a more detailed comparison with the existing literature on representation learning in offline RL . Our contributions . We develop new representation learning RL algorithms that enable sample efficient learning in low-rank MDPs under both online and offline settings : 1 . In the online episodic learning setting , our new algorithm REP-UCB integrates representation learning , exploration , and exploitation together , and significantly improves the sample complexity of the prior state-of-art algorithm FLAMBE ; 2 . In the offline learning setting , we propose a natural concentrability coefficient ( i.e. , relative condition number under the true representation ) that captures the partial coverage condition in low-rank MDP , and our algorithm REP-LCB learns to compete against any policy ( including history-dependent ones ) under such a partial coverage condition . 2 RELATED WORK . Online Setting We list the comparison as follows , which is summarized in Table 1 . FLAMBE ( Agarwal et al. , 2020b ) was a state-of-the-art oracle-efficient algorithm for low-rank MDPs . In all parameters , the statistical complexity is much worse than REP-UCB . Our algorithm and FLAMBE operate under the same computation oracle . FLAMBE does not balance exploration and exploitation , and uses explore-then-committee style techniques ( i.e. , constructions of absorbing MDPs ( Brafman & Tennenholtz , 2002 ) ) which results in its worse sample complexity . With a more complex oracle , Moffle ( Modi et al. , 2021 ) is a model-free algorithm for low-rank MDPs , with two additional assumptions : ( 1 ) the transition has low non-negative rank ( nnr ) , and ( 2 ) reachability in latent states . The first assumption significantly restricts the scope of low-rank MDPs as there are matrices whose nnr is exponentially larger than the rank ( Agarwal et al. , 2020b ) . The sample complexity of Moffle can scale O ( d6|A|13/ ( 2η5 ( 1− γ ) 5 ) ) , where η is the reachability probability , and 1/η could be as large as nnr1/2 ( Proposition 4 in Agarwal et al . ( 2020b ) ) , which essentially means that Moffle has a polynomial dependence on the nnr . Thus , Moffle needs the nnr of the transition matrix to be small . OLIVE ( Jiang et al. , 2017 ) , Witness rank ( Sun et al. , 2019 ) and Bilinear-UCB ( Du et al. , 2021 ) , when specialized to low-rank MDPs , have slightly tighter dependence on d ( e.g. , O ( d2/ 2 ) ) . But these algorithms are computationally inefficient as they are version space algorithms . Dann et al . ( 2021 ) shows that with a policy class , solving a low-rank MDP can take Ω ( 2d ) samples . In this work , similar to Witness rank ( Sun et al. , 2019 ) and FLAMBE , we use function approximators to model the transition . Thus our positive result is not in contradiction to the result from Dann et al . ( 2021 ) . VALOR ( Dann et al. , 2018 ) , PCID ( Du et al. , 2019a ) , HOMER ( Misra et al. , 2020 ) , RegRL ( Foster et al. , 2020 ) , and the approach from Feng et al . ( 2020 ) are algorithms for block MDPs which is a more restricted setting than low-rank MDPs . These works require additional assumptions such as deterministic transitions ( Dann et al. , 2018 ) , reachability ( Misra et al. , 2020 ; Du et al. , 2019a ) , Bellman completion ( Foster et al. , 2020 ) , and strong unsupervised learning oracles ( Feng et al. , 2020 ) . Offline Setting We discuss related works in offline RL . Uehara & Sun ( 2021 ) obtained similar results for offline RL on low-rank MDPs . Though the sample complexity in their algorithm is slightly tighter , our algorithm is oracle-efficient , while the CPPO algorithm from Uehara & Sun ( 2021 ) is a version space algorithm . Xie et al . ( 2021 ) propose a ( general ) pessimistic model-free algorithm in the offline setting . We can also apply their algorithm to low-rank MDPs and show some finite-sample guarantee . However , it is unclear whether the final bounds in their results can be characterized by the relative condition number only using the true representation , and whether they can compete with history-dependent policies . Thus , our result is still considered superior on low-rank MDPs . The detail is given in Section E. In addition to the above two works , the pessimistic approach in offline RL has been extensively investigated . Empirically , it can work on simulation control tasks ( Kidambi et al. , 2020 ; Yu et al. , 2020 ; Kumar et al. , 2020 ; Liu et al. , 2020 ; Chang et al. , 2021 ) . On the theoretical side , pessimism allows us to obtain the PAC guarantee on various models when a comparator policy is covered by offline data in some forms ( Jin et al. , 2020b ; Rashidinejad et al. , 2021 ; Yin et al. , 2021 ; Zanette et al. , 2021b ; Zhang et al. , 2021b ; Chang et al. , 2021 ) . However , these algorithms and their analysis rely on a known representation . We emphasize the differences in Section 5 .

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