Provable Benefit of Multitask Representation Learning in Reinforcement Learning
As representation learning becomes a powerful technique to reduce sample
complexity in reinforcement learning (RL) in practice, theoretical
understanding of its advantage is still limited. In this paper, we
theoretically characterize the benefit of representation learning under the
low-rank Markov decision process (MDP) model. We first study multitask low-rank
RL (as upstream training), where all tasks share a common representation, and
propose a new multitask reward-free algorithm called REFUEL. REFUEL learns both
the transition kernel and the near-optimal policy for each task, and outputs a
well-learned representation for downstream tasks. Our result demonstrates that
multitask representation learning is provably more sample-efficient than
learning each task individually, as long as the total number of tasks is above
a certain threshold. We then study the downstream RL in both online and offline
settings, where the agent is assigned with a new task sharing the same
representation as the upstream tasks. For both online and offline settings, we
develop a sample-efficient algorithm, and show that it finds a near-optimal
policy with the suboptimality gap bounded by the sum of the estimation error of
the learned representation in upstream and a vanishing term as the number of
downstream samples becomes large. Our downstream results of online and offline
RL further capture the benefit of employing the learned representation from
upstream as opposed to learning the representation of the low-rank model
directly. To the best of our knowledge, this is the first theoretical study
that characterizes the benefit of representation learning in exploration-based
reward-free multitask RL for both upstream and downstream tasks.