Joint Representation Training in Sequential Tasks with Shared Structure

TL;DR

This paper provides a theoretical analysis of joint representation training in multi-task reinforcement learning, demonstrating improved regret bounds and efficient algorithms leveraging shared low-dimensional structures.

Contribution

It introduces the Shared-MatrixRL algorithm for multitask RL with shared low-dimensional representations and proves regret bounds showing benefits over single-task approaches.

Findings

Regret bounds are improved from $O(PHd ext{ }\sqrt{NH})$ to $O((Hd ext{ }\sqrt{rP} + HP ext{ }\sqrt{rd}) ext{ }\sqrt{NH})$.

Shared low-dimensional representations lead to better learning efficiency in multitask RL.

Efficient algorithms are developed using quadratic programming reductions.

Abstract

Classical theory in reinforcement learning (RL) predominantly focuses on the single task setting, where an agent learns to solve a task through trial-and-error experience, given access to data only from that task. However, many recent empirical works have demonstrated the significant practical benefits of leveraging a joint representation trained across multiple, related tasks. In this work we theoretically analyze such a setting, formalizing the concept of task relatedness as a shared state-action representation that admits linear dynamics in all the tasks. We introduce the Shared-MatrixRL algorithm for the setting of Multitask MatrixRL. In the presence of $P$ episodic tasks of dimension $d$ sharing a joint $r \ll d$ low-dimensional representation, we show the regret on the the $P$ tasks can be improved from $O(PHd\sqrt{NH})$ to $O((Hd\sqrt{rP} + HP\sqrt{rd})\sqrt{NH})$ over $N$ episodes of horizon $H$. These gains coincide with those observed in other linear models in contextual bandits and RL. In contrast with previous work that have studied multi task RL in other function approximation models, we show that in the presence of bilinear optimization oracle and finite state action spaces there exists a computationally efficient algorithm for multitask MatrixRL via a reduction to quadratic programming. We also develop a simple technique to shave off a $\sqrt{H}$ factor from the regret upper bounds of some episodic linear problems.