[Paper Review] Provably Efficient Representation Selection in Low-rank Markov Decision Processes: From Online to Offline RL

Weitong Zhang, Jiafan He+3 more|arXiv (Cornell University)|Jun 22, 2021

Advanced Bandit Algorithms Research41 references4 citations

TL;DR

This paper proposes ReLEX, a provably efficient algorithm for representation selection in low-rank Markov Decision Processes, enabling adaptive use of multiple representations in both online and offline reinforcement learning. It achieves improved regret bounds in online RL and constant sample complexity in offline RL under coverage assumptions, outperforming single-representation methods.

ABSTRACT

The success of deep reinforcement learning (DRL) lies in its ability to learn a representation that is well-suited for the exploration and exploitation task. To understand how the choice of representation can improve the efficiency of reinforcement learning (RL), we study representation selection for a class of low-rank Markov Decision Processes (MDPs) where the transition kernel can be represented in a bilinear form. We propose an efficient algorithm, called ReLEX, for representation learning in both online and offline RL. Specifically, we show that the online version of ReLEX, called ReLEX-UCB, always performs no worse than the state-of-the-art algorithm without representation selection, and achieves a strictly better constant regret if the representation function class has a "coverage" property over the entire state-action space. For the offline counterpart, ReLEX-LCB, we show that the algorithm can find the optimal policy if the representation class can cover the state-action space and achieves gap-dependent sample complexity. This is the first result with constant sample complexity for representation learning in offline RL.

Motivation & Objective

• To address the challenge of selecting optimal representations across different state-action pairs in reinforcement learning, especially when multiple representations are available.
• To improve sample efficiency in both online and offline RL by dynamically selecting the best representation for each state-action pair.
• To establish theoretical guarantees for representation selection in low-rank MDPs, particularly under bilinear transition kernels.
• To close the gap in offline RL representation learning by providing the first constant sample complexity result with representation selection.
• To demonstrate empirically that representation selection significantly enhances performance compared to fixed representations.

Proposed method

• ReLEX employs a representation selection mechanism that chooses the representation minimizing optimistic Q-values in online RL (ReLEX-UCB) and maximizing pessimistic Q-values in offline RL (ReLEX-LCB).
• The algorithm uses a bilinear form to model the transition kernel as $ \mathbb{P}(s'\mid s,a) = \bm{\phi}(s,a)^\top \mathbf{M}^* \bm{\psi}(s') $, enabling low-rank structure exploitation.
• For online RL, ReLEX-UCB uses upper confidence bounds to balance exploration and representation selection, with regret bounds improving when the representation class has good coverage.
• For offline RL, ReLEX-LCB applies pessimism to handle distribution shift and ensures optimal policy identification under coverage assumptions.
• The method relies on matrix concentration and self-normalized martingale inequalities to control uncertainty in representation estimation.
• It introduces a novel analysis framework to bound the representation selection error and derive gap-dependent and constant sample complexity results.

Experimental results

Research questions

• RQ1 Can dynamic representation selection across state-action pairs improve sample efficiency in online and offline reinforcement learning?
• RQ2 Does representation selection lead to strictly better regret bounds than fixed-representation methods in online RL?
• RQ3 Can offline RL algorithms achieve constant sample complexity when using representation selection under coverage assumptions?
• RQ4 How does representation coverage affect the performance and generalization of RL algorithms in low-rank MDPs?
• RQ5 Can the proposed algorithm ReLEX outperform single-representation baselines in both online and offline settings?

Key findings

• ReLEX-UCB achieves a regret bound that is no worse than state-of-the-art algorithms without representation selection, and strictly better when the representation class has full coverage over the state-action space.
• ReLEX-LCB identifies the optimal policy in offline RL with gap-dependent sample complexity, and achieves constant sample complexity under a coverage assumption on the behavior policy.
• The constant sample complexity in offline RL is the first such result in representation learning, marking a significant theoretical advance.
• Empirical evaluations confirm that ReLEX outperforms any single representation in both online and offline settings across various MDPs.
• The algorithm's performance gain is most pronounced when the representation class includes a function that covers the optimal policy's state-action distribution.
• Theoretical analysis shows that the representation selection mechanism effectively reduces uncertainty, leading to tighter generalization bounds.