# Provable General Function Class Representation Learning in Multitask Bandits and MDPs

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 44 references that are in the map.
- First posted on arXiv 31 May 2022.

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

- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [End-to-End Training of Deep Visuomotor Policies](https://mlsanity.com/paper/1504.00702/) (2015)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [Policy Distillation](https://mlsanity.com/paper/1511.06295/) (2015)
- [Contextual Decision Processes with Low Bellman Rank are PAC-Learnable](https://mlsanity.com/paper/1610.09512/) (2016)
- [Distral: Robust Multitask Reinforcement Learning](https://mlsanity.com/paper/1707.04175/) (2017)
- [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)
- [Taming the Monster: A Fast and Simple Algorithm for Contextual Bandits](https://mlsanity.com/paper/1402.0555/) (2014)
- [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)
- [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)
- [Multi-task Deep Reinforcement Learning with PopArt](https://mlsanity.com/paper/1809.04474/) (2018)
- [Impact of Representation Learning in Linear Bandits](https://mlsanity.com/paper/2010.06531/) (2020)
- [Near-optimal Representation Learning for Linear Bandits and Linear RL](https://mlsanity.com/paper/2102.04132/) (2021)
- [Provable Representation Learning for Imitation Learning via Bi-level Optimization](https://mlsanity.com/paper/2002.10544/) (2020)
- [Provable Model-based Nonlinear Bandit and Reinforcement Learning: Shelve Optimism, Embrace Virtual Curvature](https://mlsanity.com/paper/2102.04168/) (2021)
- [On the Power of Multitask Representation Learning in Linear MDP](https://mlsanity.com/paper/2106.08053/) (2021)
- [Hypothesis Transfer Learning via Transformation Functions](https://mlsanity.com/paper/1612.01020/) (2016)
- [Reinforcement Learning in Linear MDPs: Constant Regret and Representation Selection](https://mlsanity.com/paper/2110.14798/) (2021)
- [Nearly Minimax Algorithms for Linear Bandits with Shared Representation](https://mlsanity.com/paper/2203.15664/) (2022)
- [Learning to Control in Metric Space with Optimal Regret](https://mlsanity.com/paper/1905.01576/) (2019)
- [Model-based RL in Contextual Decision Processes: PAC bounds and Exponential Improvements over Model-free Approaches](https://mlsanity.com/paper/1811.08540/) (2018)

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

- [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)
- [Understanding, Predicting and Better Resolving Q-Value Divergence in Offline-RL](https://mlsanity.com/paper/2310.04411/) (2023)
- [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)
- [Offline Multi-task Transfer RL with Representational Penalization](https://mlsanity.com/paper/2402.12570/) (2024)
- [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)
- [Distributed Multi-Task Learning for Stochastic Bandits with Context Distribution and Stage-wise Constraints](https://mlsanity.com/paper/2401.11563/) (2024)
- [Towards Understanding the Benefit of Multitask Representation Learning in Decision Process](https://mlsanity.com/paper/2503.00345/) (2025)
- [Provable Benefits of Multi-task RL under Non-Markovian Decision Making Processes](https://mlsanity.com/paper/2310.13550/) (2023)
- [Multi-Task Representation Learning for Conservative Linear Bandits](https://mlsanity.com/paper/2605.12176/) (2026)

## Abstract (arXiv)

While multitask representation learning has become a popular approach in reinforcement learning (RL) to boost the sample efficiency, the theoretical understanding of why and how it works is still limited. Most previous analytical works could only assume that the representation function is already known to the agent or from linear function class, since analyzing general function class representation encounters non-trivial technical obstacles such as generalization guarantee, formulation of confidence bound in abstract function space, etc. However, linear-case analysis heavily relies on the particularity of linear function class, while real-world practice usually adopts general non-linear representation functions like neural networks. This significantly reduces its applicability. In this work, we extend the analysis to general function class representations. Specifically, we consider an agent playing M contextual bandits (or MDPs) concurrently and extracting a shared representation function φ from a specific function class Phi using our proposed Generalized Functional Upper Confidence Bound algorithm (GFUCB). We theoretically validate the benefit of multitask representation learning within general function class for bandits and linear MDP for the first time. Lastly, we conduct experiments to demonstrate the effectiveness of our algorithm with neural net representation.