[Paper Review] Sample Efficient Reinforcement Learning via Low-Rank Matrix Estimation

Devavrat Shah, Dogyoon Song+2 more|arXiv (Cornell University)|Jun 11, 2020

Reinforcement Learning in Robotics4 references12 citations

TL;DR

This paper proposes a sample-efficient reinforcement learning algorithm that leverages low-rank matrix estimation to learn the optimal Q-function in continuous state and action spaces. By representing the Q-function via its spectral decomposition and exploiting its low-rank structure, the method achieves a sample complexity of $\widetilde{O}(\epsilon^{-(\max(d_1,d_2)+2)})$, offering an exponential improvement over classical non-parametric bounds when the optimal Q-function has low rank.

ABSTRACT

We consider the question of learning $Q$-function in a sample efficient manner for reinforcement learning with continuous state and action spaces under a generative model. If $Q$-function is Lipschitz continuous, then the minimal sample complexity for estimating $ε$-optimal $Q$-function is known to scale as $Ω(\frac{1}{ε^{d_1+d_2 +2}})$ per classical non-parametric learning theory, where $d_1$ and $d_2$ denote the dimensions of the state and action spaces respectively. The $Q$-function, when viewed as a kernel, induces a Hilbert-Schmidt operator and hence possesses square-summable spectrum. This motivates us to consider a parametric class of $Q$-functions parameterized by its "rank" $r$, which contains all Lipschitz $Q$-functions as $r o \infty$. As our key contribution, we develop a simple, iterative learning algorithm that finds $ε$-optimal $Q$-function with sample complexity of $\widetilde{O}(\frac{1}{ε^{\max(d_1, d_2)+2}})$ when the optimal $Q$-function has low rank $r$ and the discounting factor $γ$ is below a certain threshold. Thus, this provides an exponential improvement in sample complexity. To enable our result, we develop a novel Matrix Estimation algorithm that faithfully estimates an unknown low-rank matrix in the $\ell_\infty$ sense even in the presence of arbitrary bounded noise, which might be of interest in its own right. Empirical results on several stochastic control tasks confirm the efficacy of our "low-rank" algorithms.

Motivation & Objective

• Address the curse of dimensionality in reinforcement learning for continuous state and action spaces.
• Develop a data-efficient learning method that does not require prior knowledge of low-dimensional features or representations.
• Provide a universal parametric representation of the Q-function based on its spectral properties to exploit low-rank structure.
• Achieve provably improved sample complexity for learning $\epsilon$-optimal Q-functions when the true Q-function has low rank.
• Design a novel matrix estimation algorithm that ensures $\ell_\infty$-accurate recovery of low-rank matrices under bounded noise.

Proposed method

• Represent the optimal Q-function as a kernel-induced Hilbert-Schmidt operator with a generalized singular value decomposition: $Q^*(s,a) = \sum_{i=1}^\infty \sigma_i f_i(s)g_i(a)$.
• Parameterize the Q-function by its rank $r$, forming a parametric class that includes all Lipschitz continuous Q-functions as $r \to \infty$.
• Develop a novel matrix estimation algorithm that estimates a low-rank matrix in the $\ell_\infty$-norm under arbitrary bounded noise.
• Use iterative sampling from a generative model to construct a matrix of Q-values at sampled state-action pairs.
• Apply low-rank matrix estimation (e.g., via nuclear norm or soft-thresholding) to recover the underlying Q-function structure.
• Integrate the matrix estimation step into a reinforcement learning loop to iteratively refine the Q-function estimate with minimal samples.

Experimental results

Research questions

• RQ1 Can a universal representation of the Q-function be developed that enables data-efficient learning when the Q-function has low-dimensional structure?
• RQ2 What is the minimal sample complexity required to learn an $\epsilon$-optimal Q-function when the optimal Q-function is low-rank?
• RQ3 Can a matrix estimation algorithm achieve $\ell_\infty$-accurate recovery of a low-rank matrix under arbitrary bounded noise?
• RQ4 How does the proposed method compare to existing approaches in terms of sample efficiency and error convergence on continuous control tasks?
• RQ5 Does the theoretical sample complexity improvement hold under varying discount factors $\gamma$?

Key findings

• The proposed algorithm achieves a sample complexity of $\widetilde{O}(\epsilon^{-(\max(d_1,d_2)+2)})$ for learning an $\epsilon$-optimal Q-function when the optimal Q-function has low rank $r$ and $\gamma$ is below a threshold.
• This sample complexity represents an exponential improvement over the classical minimax lower bound of $\Omega(\epsilon^{-(d_1+d_2+2)})$ for general Lipschitz Q-functions.
• Empirical results on Mountain Car, Double Integrator, Cart-Pole, and Acrobot tasks show significant gains in sample efficiency, with faster convergence to low $\ell_\infty$ and mean errors.
• The method outperforms baseline matrix estimation techniques like Soft-Impute and Nuclear Norm in both sample complexity and error reduction.
• An ablation study with $\gamma = 0.5$ on the Inverted Pendulum task confirms faster convergence, validating the theoretical dependence on $\gamma$.
• The novel matrix estimation algorithm successfully recovers low-rank matrices in $\ell_\infty$-norm even under arbitrary bounded noise, a result of independent interest.