# On the Power of Multitask Representation Learning in Linear MDP

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

## Key facts

- Cited by 30 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 42 references that are in the map.
- First posted on arXiv 15 Jun 2021.

## Builds on (most cited first, 23 of 23)

- [A Theoretical Analysis of Contrastive Unsupervised Representation Learning](https://mlsanity.com/paper/1902.09229/) (2019)
- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [Meta-Learning with Differentiable Convex Optimization](https://mlsanity.com/paper/1904.03758/) (2019)
- [Policy Distillation](https://mlsanity.com/paper/1511.06295/) (2015)
- [Meta-learning with differentiable closed-form solvers](https://mlsanity.com/paper/1805.08136/) (2018)
- [Distral: Robust Multitask Reinforcement Learning](https://mlsanity.com/paper/1707.04175/) (2017)
- [Few-Shot Learning via Learning the Representation, Provably](https://mlsanity.com/paper/2002.09434/) (2020)
- [The Benefit of Multitask Representation Learning](https://mlsanity.com/paper/1505.06279/) (2015)
- [On the Theory of Transfer Learning: The Importance of Task Diversity](https://mlsanity.com/paper/2006.11650/) (2020)
- [Adaptive Gradient-Based Meta-Learning Methods](https://mlsanity.com/paper/1906.02717/) (2019)
- [Actor-Mimic: Deep Multitask and Transfer Reinforcement Learning](https://mlsanity.com/paper/1511.06342/) (2015)
- [Provable Meta-Learning of Linear Representations](https://mlsanity.com/paper/2002.11684/) (2020)
- [DARLA: Improving Zero-Shot Transfer in Reinforcement Learning](https://mlsanity.com/paper/1707.08475/) (2017)
- [Multi-task Deep Reinforcement Learning with PopArt](https://mlsanity.com/paper/1809.04474/) (2018)
- [Online Meta-Learning](https://mlsanity.com/paper/1902.08438/) (2019)
- [Model-Based Reinforcement Learning with a Generative Model is Minimax Optimal](https://mlsanity.com/paper/1906.03804/) (2019)
- [Impact of Representation Learning in Linear Bandits](https://mlsanity.com/paper/2010.06531/) (2020)
- [Learning-to-Learn Stochastic Gradient Descent with Biased Regularization](https://mlsanity.com/paper/1903.10399/) (2019)
- [Provable Representation Learning for Imitation Learning via Bi-level Optimization](https://mlsanity.com/paper/2002.10544/) (2020)
- [Hypothesis Transfer Learning via Transformation Functions](https://mlsanity.com/paper/1612.01020/) (2016)
- [Regret Bounds for Lifelong Learning](https://mlsanity.com/paper/1610.08628/) (2016)
- [Incremental Learning-to-Learn with Statistical Guarantees](https://mlsanity.com/paper/1803.08089/) (2018)
- [Is Plug-in Solver Sample-Efficient for Feature-based Reinforcement Learning?](https://mlsanity.com/paper/2010.05673/) (2020)

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

- [Representation Learning for Online and Offline RL in Low-rank MDPs](https://mlsanity.com/paper/2110.04652/) (2021)
- [RLHF from Heterogeneous Feedback via Personalization and Preference Aggregation](https://mlsanity.com/paper/2405.00254/) (2024)
- [Provable Benefits of Representational Transfer in Reinforcement Learning](https://mlsanity.com/paper/2205.14571/) (2022)
- [Provable Benefit of Multitask Representation Learning in Reinforcement Learning](https://mlsanity.com/paper/2206.05900/) (2022)
- [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)
- [Understanding, Predicting and Better Resolving Q-Value Divergence in Offline-RL](https://mlsanity.com/paper/2310.04411/) (2023)
- [Reinforcement Learning in Linear MDPs: Constant Regret and Representation Selection](https://mlsanity.com/paper/2110.14798/) (2021)
- [Joint Learning of Linear Time-Invariant Dynamical Systems](https://mlsanity.com/paper/2112.10955/) (2021)
- [Nearly Minimax Algorithms for Linear Bandits with Shared Representation](https://mlsanity.com/paper/2203.15664/) (2022)
- [Provable and Efficient Continual Representation Learning](https://mlsanity.com/paper/2203.02026/) (2022)
- [Improved Sample Complexity for Reward-free Reinforcement Learning under Low-rank MDPs](https://mlsanity.com/paper/2303.10859/) (2023)
- [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)
- [Provable General Function Class Representation Learning in Multitask Bandits and MDPs](https://mlsanity.com/paper/2205.15701/) (2022)
- [On the Power of Pre-training for Generalization in RL: Provable Benefits and Hardness](https://mlsanity.com/paper/2210.10464/) (2022)
- [Multi-task Representation Learning for Pure Exploration in Linear Bandits](https://mlsanity.com/paper/2302.04441/) (2023)
- [Task Diversity Shortens the ICL Plateau](https://mlsanity.com/paper/2410.05448/) (2024)
- [Meta Learning MDPs with Linear Transition Models](https://mlsanity.com/paper/2201.08732/) (2022)
- [Provable Pathways: Learning Multiple Tasks over Multiple Paths](https://mlsanity.com/paper/2303.04338/) (2023)
- [Provably Efficient Offline Reinforcement Learning with Perturbed Data Sources](https://mlsanity.com/paper/2306.08364/) (2023)
- [Provable Zero-Shot Generalization in Offline Reinforcement Learning](https://mlsanity.com/paper/2503.07988/) (2025)
- [Sample Efficient Myopic Exploration Through Multitask Reinforcement Learning with Diverse Tasks](https://mlsanity.com/paper/2403.01636/) (2024)
- [Guarantees for Nonlinear Representation Learning: Non-identical Covariates, Dependent Data, Fewer Samples](https://mlsanity.com/paper/2410.11227/) (2024)
- [Uncertainty-aware transfer across tasks using hybrid model-based successor feature reinforcement learning](https://mlsanity.com/paper/2310.10818/) (2023)
- [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)
- [Accelerating Multi-Task Temporal Difference Learning under Low-Rank Representation](https://mlsanity.com/paper/2503.02030/) (2025)
- [Provable Benefits of Multi-task RL under Non-Markovian Decision Making Processes](https://mlsanity.com/paper/2310.13550/) (2023)
- [Generalized Linear Markov Decision Process](https://mlsanity.com/paper/2506.00818/) (2025)

## Abstract (arXiv)

While multitask representation learning has become a popular approach in reinforcement learning (RL), theoretical understanding of why and when it works remains limited. This paper presents analyses for the statistical benefit of multitask representation learning in linear Markov Decision Process (MDP) under a generative model. In this paper, we consider an agent to learn a representation function φ out of a function class Phi from T source tasks with N data per task, and then use the learned hatφ to reduce the required number of sample for a new task. We first discover a Least-Activated-Feature-Abundance (LAFA) criterion, denoted as kappa, with which we prove that a straightforward least-square algorithm learns a policy which is tildeO(H2√fracC(Phi)2 kappa dNT+frackappa dn) sub-optimal. Here H is the planning horizon, C(Phi) is Phi's complexity measure, d is the dimension of the representation (usually dll C(Phi)) and n is the number of samples for the new task. Thus the required n is O(kappa d H4) for the sub-optimality to be close to zero, which is much smaller than O(C(Phi)2kappa d H4) in the setting without multitask representation learning, whose sub-optimality gap is tildeO(H2√frackappa C(Phi)2dn). This theoretically explains the power of multitask representation learning in reducing sample complexity. Further, we note that to ensure high sample efficiency, the LAFA criterion kappa should be small. In fact, kappa varies widely in magnitude depending on the different sampling distribution for new task. This indicates adaptive sampling technique is important to make kappa solely depend on d. Finally, we provide empirical results of a noisy grid-world environment to corroborate our theoretical findings.