# Provable Benefit of Multitask Representation Learning in Reinforcement Learning

arXiv: https://arxiv.org/abs/2206.05900 · HTML: https://mlsanity.com/paper/2206.05900/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 53 of its 75 references that are in the map.
- First posted on arXiv 13 Jun 2022.

## Builds on (most cited first, 30 of 53)

- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [Minimax Regret Bounds for Reinforcement Learning](https://mlsanity.com/paper/1703.05449/) (2017)
- [Is Q-learning Provably Efficient?](https://mlsanity.com/paper/1807.03765/) (2018)
- [Is Pessimism Provably Efficient for Offline RL?](https://mlsanity.com/paper/2012.15085/) (2020)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [Contextual Decision Processes with Low Bellman Rank are PAC-Learnable](https://mlsanity.com/paper/1610.09512/) (2016)
- [Bellman-consistent Pessimism for Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06926/) (2021)
- [Distral: Robust Multitask Reinforcement Learning](https://mlsanity.com/paper/1707.04175/) (2017)
- [FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs](https://mlsanity.com/paper/2006.10814/) (2020)
- [Bellman Eluder Dimension: New Rich Classes of RL Problems, and Sample-Efficient Algorithms](https://mlsanity.com/paper/2102.00815/) (2021)
- [Few-Shot Learning via Learning the Representation, Provably](https://mlsanity.com/paper/2002.09434/) (2020)
- [Provably Efficient Exploration in Policy Optimization](https://mlsanity.com/paper/1912.05830/) (2019)
- [Reinforcement Learning in Feature Space: Matrix Bandit, Kernels, and Regret Bound](https://mlsanity.com/paper/1905.10389/) (2019)
- [The Benefit of Multitask Representation Learning](https://mlsanity.com/paper/1505.06279/) (2015)
- [Learning Near Optimal Policies with Low Inherent Bellman Error](https://mlsanity.com/paper/2003.00153/) (2020)
- [Bilinear Classes: A Structural Framework for Provable Generalization in RL](https://mlsanity.com/paper/2103.10897/) (2021)
- [Unifying PAC and Regret: Uniform PAC Bounds for Episodic Reinforcement Learning](https://mlsanity.com/paper/1703.07710/) (2017)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Provably efficient RL with Rich Observations via Latent State Decoding](https://mlsanity.com/paper/1901.09018/) (2019)
- [Provable Meta-Learning of Linear Representations](https://mlsanity.com/paper/2002.11684/) (2020)
- [Policy Finetuning: Bridging Sample-Efficient Offline and Online Reinforcement Learning](https://mlsanity.com/paper/2106.04895/) (2021)
- [What are the Statistical Limits of Offline RL with Linear Function Approximation?](https://mlsanity.com/paper/2010.11895/) (2020)
- [Representation Learning for Online and Offline RL in Low-rank MDPs](https://mlsanity.com/paper/2110.04652/) (2021)
- [Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning](https://mlsanity.com/paper/1911.05815/) (2019)
- [Generalization and Exploration via Randomized Value Functions](https://mlsanity.com/paper/1402.0635/) (2014)
- [Optimism in Reinforcement Learning with Generalized Linear Function Approximation](https://mlsanity.com/paper/1912.04136/) (2019)
- [Provable Benefits of Actor-Critic Methods for Offline Reinforcement Learning](https://mlsanity.com/paper/2108.08812/) (2021)
- [Model-free Representation Learning and Exploration in Low-rank MDPs](https://mlsanity.com/paper/2102.07035/) (2021)
- [Towards Instance-Optimal Offline Reinforcement Learning with Pessimism](https://mlsanity.com/paper/2110.08695/) (2021)
- [Near-optimal Offline Reinforcement Learning with Linear Representation: Leveraging Variance Information with Pessimism](https://mlsanity.com/paper/2203.05804/) (2022)

## Built on by (most cited first, 29 of 29)

- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Multi-Task Imitation Learning for Linear Dynamical Systems](https://mlsanity.com/paper/2212.00186/) (2022)
- [Overcoming the Sim-to-Real Gap: Leveraging Simulation to Learn to Explore for Real-World RL](https://mlsanity.com/paper/2410.20254/) (2024)
- [Minimax Optimal and Computationally Efficient Algorithms for Distributionally Robust Offline Reinforcement Learning](https://mlsanity.com/paper/2403.09621/) (2024)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)
- [Improved Sample Complexity for Reward-free Reinforcement Learning under Low-rank MDPs](https://mlsanity.com/paper/2303.10859/) (2023)
- [Offline Multi-task Transfer RL with Representational Penalization](https://mlsanity.com/paper/2402.12570/) (2024)
- [Offline Multitask Representation Learning for Reinforcement Learning](https://mlsanity.com/paper/2403.11574/) (2024)
- [Sample-Efficient Linear Representation Learning from Non-IID Non-Isotropic Data](https://mlsanity.com/paper/2308.04428/) (2023)
- [QMP: Q-switch Mixture of Policies for Multi-Task Behavior Sharing](https://mlsanity.com/paper/2302.00671/) (2023)
- [Multi-task Representation Learning for Pure Exploration in Linear Bandits](https://mlsanity.com/paper/2302.04441/) (2023)
- [Joint Representation Training in Sequential Tasks with Shared Structure](https://mlsanity.com/paper/2206.12441/) (2022)
- [Task Diversity Shortens the ICL Plateau](https://mlsanity.com/paper/2410.05448/) (2024)
- [Learning Policy Committees for Effective Personalization in MDPs with Diverse Tasks](https://mlsanity.com/paper/2503.01885/) (2025)
- [Provable Pathways: Learning Multiple Tasks over Multiple Paths](https://mlsanity.com/paper/2303.04338/) (2023)
- [Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Tasks](https://mlsanity.com/paper/2502.02834/) (2025)
- [Provably Efficient Algorithm for Nonstationary Low-Rank MDPs](https://mlsanity.com/paper/2308.05471/) (2023)
- [Sample Efficient Myopic Exploration Through Multitask Reinforcement Learning with Diverse Tasks](https://mlsanity.com/paper/2403.01636/) (2024)
- [Pessimism Principle Can Be Effective: Towards a Framework for Zero-Shot Transfer Reinforcement Learning](https://mlsanity.com/paper/2505.18447/) (2025)
- [Multi-Asset Closed-Loop Reservoir Management Using Deep Reinforcement Learning](https://mlsanity.com/paper/2207.10376/) (2022)
- [The Limits of Transfer Reinforcement Learning with Latent Low-rank Structure](https://mlsanity.com/paper/2410.21601/) (2024)
- [Towards Understanding the Benefit of Multitask Representation Learning in Decision Process](https://mlsanity.com/paper/2503.00345/) (2025)
- [Provable Multi-Task Reinforcement Learning: A Representation Learning Framework with Low Rank Rewards](https://mlsanity.com/paper/2604.03891/) (2026)
- [Provable Benefits of Multi-task RL under Non-Markovian Decision Making Processes](https://mlsanity.com/paper/2310.13550/) (2023)
- [Generalisation in Multitask Fitted Q-Iteration and Offline Q-learning](https://mlsanity.com/paper/2512.20220/) (2025)
- [Robust Transfer Learning with Side Information](https://mlsanity.com/paper/2603.07921/) (2026)
- [Near-optimal and Efficient First-Order Algorithm for Multi-Task Learning with Shared Linear Representation](https://mlsanity.com/paper/2605.00473/) (2026)
- [BIFROST: Bridging Invariant Feature Representation for Observation-space Sim2Real Transfer](https://mlsanity.com/paper/2607.01410/) (2026)
- [Out-of-Distribution Generalisation with Sequence Models in Offline Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2609.03667/) (2026)

## Abstract (arXiv)

As representation learning becomes a powerful technique to reduce sample complexity in reinforcement learning (RL) in practice, theoretical understanding of its advantage is still limited. In this paper, we theoretically characterize the benefit of representation learning under the low-rank Markov decision process (MDP) model. We first study multitask low-rank RL (as upstream training), where all tasks share a common representation, and propose a new multitask reward-free algorithm called REFUEL. REFUEL learns both the transition kernel and the near-optimal policy for each task, and outputs a well-learned representation for downstream tasks. Our result demonstrates that multitask representation learning is provably more sample-efficient than learning each task individually, as long as the total number of tasks is above a certain threshold. We then study the downstream RL in both online and offline settings, where the agent is assigned with a new task sharing the same representation as the upstream tasks. For both online and offline settings, we develop a sample-efficient algorithm, and show that it finds a near-optimal policy with the suboptimality gap bounded by the sum of the estimation error of the learned representation in upstream and a vanishing term as the number of downstream samples becomes large. Our downstream results of online and offline RL further capture the benefit of employing the learned representation from upstream as opposed to learning the representation of the low-rank model directly. To the best of our knowledge, this is the first theoretical study that characterizes the benefit of representation learning in exploration-based reward-free multitask RL for both upstream and downstream tasks.
