On the Power of Multitask Representation Learning in Linear MDP

Abstract

While multitask representation learning has become a popular approach in reinforcement learning (RL), theoretical understanding of why and when it works remains limited. This paper presents analyses for the statistical benefit of multitask representation learning in linear Markov Decision Process (MDP) under a generative model. In this paper, we consider an agent to learn a representation function φ out of a function class Φ from T source tasks with N data per task, and then use the learned φ to reduce the required number of sample for a new task. We first discover a Least-Activated-Feature-Abundance (LAFA) criterion, denoted as κ, with which we prove that a straightforward least-square algorithm learns a policyHere H is the planning horizon,and n is the number of samples for the new task. In multitask RL, N T is usually sufficiently large and the bound is dominated by the second term. Thus the required n is O(κdH 4 ) for the sub-optimality to be close to zero, which is much smaller than O(C(Φ) 2 κdH 4 ) in the setting without multitask representation learning, whose sub-optimality gap is Õ H 2 κC(Φ) 2 d n . This theoretically explains the power of multitask representation learning in reducing sample complexity. Further, we note that to ensure high sample efficiency, the LAFA criterion κ should be small. In fact, κ varies widely in magnitude depending on the different sampling distribution for new task. This indicates adaptive sampling technique is important to make κ solely depend on d. Finally, we provide empirical results of a noisy grid-world environment to corroborate our theoretical findings.Preprint. Under review.

Representation Learning for Online and Offline RL in Low-rank MDPs

“…In contrast, REP-UCB carefully trades exploration versus exploitation by combining the reward signal and exploration bonus (constructed using the latest learned representation), and enables data sharing across all time steps. Our sample complexity nearly matches the ones from those computationally inefficient algorithms (Jiang et al, 2017;Sun et al, 2019;Du et al, 2021). We summarize the comparison with the prior works that study representation learning in Table 1.…”

“…OLIVE (Jiang et al, 2017), Witness rank (Sun et al, 2019) and Bilinear-UCB (Du et al, 2021), when specialized to low-rank MDPs, have slightly tighter dependence on d (e.g., O(d 2 / 2 )). But these algorithms are computationally inefficient as they are version space algorithms.…”

“…Finally, we mention works that tackle representation learning from different viewpoints and settings. Lu et al (2021) consider multitask representation learning under a generative model assumption. Hao et al (2021a) study the feature selection problem in sparse linear MDPs under the assumption that the offline data admits some form of full coverage condition.…”

“…On the other hand, the fact that linear MDPs can be solved statistical and computational efficiently if φ is known a priori implies that if one could learn the representation of the low-rank MDP, one could then efficiently learn the optimal policy. Indeed, prior works have shown that learning in low-rank MDPs is statistically feasible (Jiang et al, 2017;Sun et al, 2019;Du et al, 2021) via leveraging rich function approximators. However, these algorithms are version space algorithms and are not computationally efficient.…”

Joint Learning of Linear Time-Invariant Dynamical Systems

“…Recently, Du et al (2020) and Tripuraneni et al (2021) consider the multi-task regression problem and show that under a shared low-dimensional representation, improved statistical rates can be obtained if the intrinsic dimension is small. Further, Hu et al (2021) and Lu et al (2021) show similar results for linear contextual bandits and Markov decision processes.…”