cs.LGApr 4, 2026

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

Authors: Yaoze Guo, Shana Moothedath

Organizations: Electrical and Computer Engineering, Iowa State University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

Fig. 1 : Results for Experiment 1: We set d=100,T=100,r=2,∣\pazocalS∣=1000,∣\pazocalA∣=10.d=100,T=100,r=2,|\pazocal{S}|=1000,|\pazocal{A}|=10. Random policy baseline replaces stage 2 in Algorithm 1 with a random policy that chooses each action uniformly. MoM baseline replaces stage 3 in Algorithm 1 with an MoM estimator. Our proposed approach outperforms both the baselines.
Fig. 2 : Results for Experiment 1: Figure 2(a) presents regret plots for Experiment 1 of the proposed algorithm against the three baseline approaches. Results for Experiment 2: Figure 2(b) and 2(c) present results for the grid maze environment. The five goal points for five tasks are (1,1), (2,2), (3,3), (4,4), (5,5), respectively. Figure 2(b) presents subspace distance vs. number of samples KK . Fig. 2(c) presents the estimation error vs. number of samples KK . Estimation error is defined as ∥Θ^h−Θh⋆∥F/∥Θh⋆∥F\|\widehat{\Theta}_{h}-\Theta^{\star}_{h}\|_{\text{F}}/\|\Theta^{\star}_{h}\|_{\text{F}} .
More in this topic

Multi-Task RL

Multi-task reinforcement learning trains agents to solve related tasks by sharing policies, representations, or experience.

See all 86 papers

Explore similar work

CardsList
May 1, 2026

Near-optimal and Efficient First-Order Algorithm for Multi-Task Learning with Shared Linear Representation

Multi-task learning (MTL) has emerged as a pivotal paradigm in machine learning by leveraging shared structures across multiple related tasks. Despite its empirical success, the development of likelihood-based efficiently solvable algorithms--even for shared linear representations--remains largely underdeveloped, primarily due to the non-convex structure intrinsic to matrix factorization. This paper introduces a first-order algorithm that jointly learns a shared representation and task-specific parameters, with guaranteed efficiency. Notably, it converges in O~(1)\widetilde{\mathcal{O}}(1) iterations and attains a near-optimal estimation error of O~(dk/(TN))\widetilde{\mathcal{O}}(dk/(TN)), improving over existing likelihood-based methods by a factor of kk, where dd, kk, TT, NN denote input dimension, representation dimension, task count, and samples per task, respectively. Our results justify that likelihood-based first-order methods can efficiently solve the MTL problem.

Shihong Ding, Fangyu Du, Cong Fang

State Key Lab of General AI, School of Intelligence Science and Technology, Peking University, China · School of Mathematical Sciences, Peking University, China · Institute for Artificial Intelligence, Peking University, China

Shared Representation LearningRepresentation LearningMulti-Task LearningMulti-Task OptimizationMatrix Factorization
AskPDF
Jun 4, 2026

Representation Learning Enables Scalable Multitask Deep Reinforcement Learning

Scaling reinforcement learning (RL) to diverse multitask settings remains a central challenge. While recent advances in model-based RL achieve strong performance, they rely on planning and complex training pipelines, making it unclear which components are essential for scalability. We revisit this question and argue that the primary driver of scalable multitask RL is not model-based control, but representation learning. In particular, we show that combining predictive, model-based representations with high-capacity value function approximation is sufficient to achieve strong performance, even without planning. We evaluate a simple model-free algorithm, MR.Q, coupled with auxiliary predictive objectives into a scalable actor-critic architecture. This approach outperforms a recent world-model-based method and a range of deep RL baselines across a diverse suite of multitask continuous control tasks, while significantly reducing computational overhead and improving wall-clock efficiency. We observe consistent improvements with increased model capacity and show through ablations that predictive representation learning is critical for performance.

Johan Obando-Ceron, Lu Li, Scott Fujimoto +3

Mila – Québec AI Institute · Université de Montréal · McGill University +2

Reinforcement LearningRepresentation Learning for RLContinuous ControlMulti-Task RLModel-Free RL
AskPDF
Apr 1, 2026

RepUCB: Representation Learning-Based UCB for Heterogeneous Multi-Task Linear Bandits

Multi-task representation learning exploits the shared structure among related tasks by learning a common latent representation, thereby improving sample efficiency. This paper introduces a novel approach to multi-task representation learning in heterogeneous linear bandits. We consider TT concurrent heterogeneous linear bandit tasks, each with feature dimension dd, whose reward parameters share a common latent representation of dimension r≪min⁡{d,T}r \ll \min\{d, T\}, capturing the underlying task relatedness. We propose RepUCB, a novel Upper Confidence Bound (UCB) algorithm that leverages shared low-rank representations to enhance learning in a sample-efficient manner. Our algorithm first collects data through an exploration phase, estimates the shared representation, and then performs UCB-based learning on our proposed confidence set. We provide theoretical guarantees for the confidence set and prove that the unknown reward parameters lie within the confidence set with high probability. We derive cumulative regret bound and show that the proposed approach achieves O~(drNT)\widetilde{O}(\sqrt{drNT}), a significant improvement over solving the TT tasks independently, resulting in a regret of O~(dTN)\widetilde{O}(dT\sqrt{N}). We performed numerical simulations to validate the performance of our algorithm for different problem sizes and compared with baseline algorithms.

Jiabin Lin, Shana Moothedath

Department of Electrical and Computer Engineering, Iowa State University, USA

Multi-Armed BanditsMulti-Task LearningUCB Algorithms
AskPDF