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 T concurrent heterogeneous linear bandit tasks, each with feature dimension d, whose reward parameters share a common latent representation of dimension r ll mind, 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 widetildeO(√drNT), a significant improvement over solving the T tasks independently, resulting in a regret of widetildeO(dT√N). We performed numerical simulations to validate the performance of our algorithm for different problem sizes and compared with baseline algorithms.