Near-optimal Representation Learning for Linear Bandits and Linear RL

TL;DR

This paper introduces a sample-efficient algorithm for multi-task linear bandits and RL that leverages shared low-dimensional representations, significantly improving regret bounds and demonstrating near-optimality.

Contribution

The paper develops the first theoretical algorithm for multi-task representation learning in linear bandits and RL, achieving near-optimal regret bounds.

Findings

Achieves regret of (M\u221a{d}kT + dkMT) in multi-task linear bandits.

Provides a lower bound showing near-optimality when d > M.

Extends results to multi-task episodic RL with linear value functions.

Abstract

This paper studies representation learning for multi-task linear bandits and multi-task episodic RL with linear value function approximation. We first consider the setting where we play $M$ linear bandits with dimension $d$ concurrently, and these bandits share a common $k$-dimensional linear representation so that $k\ll d$ and $k \ll M$. We propose a sample-efficient algorithm, MTLR-OFUL, which leverages the shared representation to achieve $\tilde{O}(M\sqrt{dkT} + d\sqrt{kMT} )$ regret, with $T$ being the number of total steps. Our regret significantly improves upon the baseline $\tilde{O}(Md\sqrt{T})$ achieved by solving each task independently. We further develop a lower bound that shows our regret is near-optimal when $d > M$. Furthermore, we extend the algorithm and analysis to multi-task episodic RL with linear value function approximation under low inherent Bellman error \citep{zanette2020learning}. To the best of our knowledge, this is the first theoretical result that characterizes the benefits of multi-task representation learning for exploration in RL with function approximation.