Shift Before You Learn: Enabling Low-Rank Representations in Reinforcement Learning
cs.LG
Submitted: 2025-09-05
Updated: 2026-08-28
Comments: 63 pages, 11 figures. Accepted to NeurIPS 2025 (Spotlight). This version includes updated numerical experiments following a correction to the bootstrapping procedure used for TD estimation
Code: https://github.com/stestoKTH/shift-SM
License: http://creativecommons.org/licenses/by/4.0/
The gist: Low-rank structure is a common implicit assumption in many modern reinforcement learning (RL) algorithms.
Terminology
Abstract
Low-rank structure is a common implicit assumption in many modern reinforcement learning (RL) algorithms. For instance, reward-free and goal-conditioned RL methods often presume that the successor measure admits a low-rank representation. In this work, we challenge this assumption by first remarking that the successor measure itself is not approximately low-rank. Instead, we demonstrate that a low-rank structure naturally emerges in the shifted successor measure, which captures the system dynamics after bypassing a few initial transitions. We provide finite-sample performance guarantees for the entry-wise estimation of a low-rank approximation of the shifted successor measure from sampled entries. Our analysis reveals that both the approximation and estimation errors are primarily governed by a newly introduced quantitity: the spectral recoverability of the corresponding matrix. To bound this parameter, we derive a new class of functional inequalities for Markov chains that we call Type II Poincaré inequalities and from which we can quantify the amount of shift needed for effective low-rank approximation and estimation. This analysis shows in particular that the required shift depends on decay of the high-order singular values of the shifted successor measure and is hence typically small in practice. Additionally, we establish a connection between the necessary shift and the local mixing properties of the underlying dynamical system, which provides a natural way of selecting the shift. Finally, we validate our theoretical findings with experiments, and demonstrate that shifting the successor measure indeed leads to improved performance in goal-conditioned RL.
Sources
- Learning Successor States and Goal-Dependent Values: A Mathematical Viewpoint
- Spectral gap of nonreversible Markov chains
- D4RL: Datasets for Deep Data-Driven Reinforcement Learning
- Deep Successor Reinforcement Learning
- Matrix Estimation for Offline Reinforcement Learning with Low-Rank Structure
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks