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

By Masatoshi Uehara, Xuezhou Zhang, Wen Sun. 2021

142 citations recorded by the source corpus.

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), 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 $\widetilde{O}( A^9 d^7 / (\epsilon^{10} (1-\gamma)^{22}))$ for FLAMBE to $\widetilde{O}( A^2 d^4 / (\epsilon^2 (1-\gamma)^{5}) )$ with $d$ being the rank of the transition matrix (or dimension of the ground truth representation), $A$ being the number of actions, and $\gamma$ being the discounted factor. Notably, REP-UCB is simpler than FLAMBE, as it directly balances the interplay between representation learning, exploration, and exploitation, while FLAMBE is an explore-then-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 distribution.

Subfields: Artificial Intelligence, Management Science and Operations Research

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Management Science and Operations Research

Citations in the corpus

Towards Continual Reinforcement Learning: A Review and Perspectives (2020, background): “Additionally, assumptions about low-rank structure enable a basis for provable generalization while allowing for more complexity than simpler Block MDP models (Agarwal et al., 2020; Uehara et al., 2021; Zhang et al., 2021a).”
Provably efficient representation selection in Low-rank Markov Decision Processes: from online to offline RL (2021, background): “Other follow-up works include the partial coverage [Uehara and Sun, 2021] in general function approximation and the statistical barriers for offline RL [Foster et al., 2021].”