# Provable Multi-Task Reinforcement Learning: A Representation Learning Framework with Low Rank Rewards

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 16 of its 21 references that are in the map.
- First posted on arXiv 4 Apr 2026.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

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

- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [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)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [Provable Meta-Learning of Linear Representations](https://mlsanity.com/paper/2002.11684/) (2020)
- [Multi-task Deep Reinforcement Learning with PopArt](https://mlsanity.com/paper/1809.04474/) (2018)
- [Spectral Methods for Data Science: A Statistical Perspective](https://mlsanity.com/paper/2012.08496/) (2020)
- [Sample Complexity of Multi-task Reinforcement Learning](https://mlsanity.com/paper/1309.6821/) (2013)
- [Impact of Representation Learning in Linear Bandits](https://mlsanity.com/paper/2010.06531/) (2020)
- [Reward-Free RL is No Harder Than Reward-Aware RL in Linear Markov Decision Processes](https://mlsanity.com/paper/2201.11206/) (2022)
- [Near-optimal Representation Learning for Linear Bandits and Linear RL](https://mlsanity.com/paper/2102.04132/) (2021)
- [Provable Benefit of Multitask Representation Learning in Reinforcement Learning](https://mlsanity.com/paper/2206.05900/) (2022)
- [Nearly Minimax Optimal Reward-free Reinforcement Learning](https://mlsanity.com/paper/2010.05901/) (2020)
- [Offline Multitask Representation Learning for Reinforcement Learning](https://mlsanity.com/paper/2403.11574/) (2024)
- [Multi-User Reinforcement Learning with Low Rank Rewards](https://mlsanity.com/paper/2210.05355/) (2022)
- [Accelerating Multi-Task Temporal Difference Learning under Low-Rank Representation](https://mlsanity.com/paper/2503.02030/) (2025)

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

- [Scalable Multi-Task Inverse Reinforcement Learning](https://mlsanity.com/paper/2610.00758/) (2026)

## Abstract (arXiv)

Multi-task representation learning (MTRL) is an approach that learns shared latent representations across related tasks, facilitating collaborative learning that improves the overall learning efficiency. This paper studies MTRL for multi-task reinforcement learning (RL), where multiple tasks have the same state-action space and transition probabilities, but different rewards. We consider T linear Markov Decision Processes (MDPs) where the reward functions and transition dynamics admit linear feature embeddings of dimension d. The relatedness among the tasks is captured by a low-rank structure on the reward matrices. Learning shared representations across multiple RL tasks is challenging due to the complex and policy-dependent nature of data that leads to a temporal progression of error. Our approach adopts a reward-free reinforcement learning framework to first learn a data-collection policy. This policy then informs an exploration strategy for estimating the unknown reward matrices. Importantly, the data collected under this well-designed policy enable accurate estimation, which ultimately supports the learning of an near-optimal policy. Unlike existing approaches that rely on restrictive assumptions such as Gaussian features, incoherence conditions, or access to optimal solutions, we propose a low-rank matrix estimation method that operates under more general feature distributions encountered in RL settings. Theoretical analysis establishes that accurate low-rank matrix recovery is achievable under these relaxed assumptions, and we characterize the relationship between representation error and sample complexity. Leveraging the learned representation, we construct near-optimal policies and prove a regret bound. Experimental results demonstrate that our method effectively learns robust shared representations and task dynamics from finite data.