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

Abstract

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 efficient 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 offline RL settings. For the online setting, operating with the same computational oracles used in FLAMBE(Agarwal et al., 2020b)--the state-of-art algorithm for learning representations in low-rank MDPs, we propose an algorithm REP-UCB-Upper Confidence Bound driven REPresentation learning for RL, which significantly improves the sample complexity from O(A 9 d 7 /( 10 (1with 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, REP-UCB is simpler than FLAMBE, as it directly balances the interplay between representation learning, exploration, and exploitation, while FLAMBE is an explorethen-commit style approach and has to perform reward-free exploration step-by-step forward in time. For the offline 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 offline data distribution.

Provable Reinforcement Learning with a Short-Term Memory

“…The second line of work concerns rich observation RL, where the observation space can be infinite and arbitrarily complex, in (for the most part) Markovian environments. These works provide structural conditions that permit sample efficient RL with function approximation Jiang et al (2017); Sun et al (2019); Jin et al (2021); Du et al (2021);Foster et al (2021) as well as algorithms that are provably efficient in some special cases (Du et al, 2019;Misra et al, 2020;Agarwal et al, 2020;Uehara et al, 2021). However, as we will see, these structural conditions are not satisfied in our POMDP model so these results do not directly apply.…”

Efficient Reinforcement Learning in Block MDPs: A Model-free Representation Learning Approach

“…Note that from here, our analysis significantly departs from the prior Block MDP works' analysis and the analysis of MOFFLE which rely on a reachability assumption. This part of our proof leverages some ideas in the analysis of a recent model-based representation learning algorithm REP-UCB (Uehara et al, 2021).…”

“…Low-rank MDPs Low-rank MDP is strictly more general than linear MDPs which assume representation is known a priori. There are several related papers come from the recent literature on provable representation learning for low-rank MDPs (Agarwal et al, 2020b;Modi et al, 2021;Uehara et al, 2021;Ren et al, 2021). Low-rank MDPs generalize Block MDPs, so these algorithms are applicable in our setting.…”

“…Model-free vs model-based Like HOMER (Misra et al, 2020), we only model φ via function approximation (hence model-free), while FLAMBE and Rep-UCB (Agarwal et al, 2020b;Uehara et al, 2021) additionally model omission distributions o(•|z) (i.e., model-based), which could be complex when observations are high dimensional.…”

“…Prior works (Dann et al, 2018;Du et al, 2019;Misra et al, 2020;Zhang et al, 2020;Sodhani et al, 2021) have motivated the Block MDP model through scenarios such as navigation tasks and image based robotics tasks where the observations can often be reasonably mapped to the latent physical location and states. We develop a new Sample Complexity Reward OLIVE (Jiang et al, 2017) |Z| 3 H 3 |A| 2 log |Φ| 2 Yes REP-UCB (Uehara et al, 2021) |Z| 4 H 5 |A| 2 ln(|Φ||Υ|) 2 Yes MOFFLE (Modi et al, 2021) |Z| 7 H 8 |A| 13 ln |Φ| min( 2 η min ,η 5 min )…”