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 representation learning in low rank MDPs in both online and offline setting. This paper proposes REP-UCB algorithm, which improves the sample complexity of FLAMBE significantly. In addition, REP-UCB is much simpler than FLAMBE.  For the offline setting, the REP-LCB algorithm also achieves polynomial sample complexity with partial coverage. In particular, the coverage is measured by the ground-truth feature.