Provably Efficient Multi-Task Reinforcement Learning with Model Transfer

We have a set of $M$ MDPs $\mathinner{\left\{\mathcal{M}_{p}=(H,{\mathcal{S}},\mathcal{A},p_{0},\mathbb{P}_{p},r_{p})\right\}}_{p=1}^{M}$, each associated with a player $p\in[M]$. Each MDP $\mathcal{M}_{p}$ is regarded as a task. The MDPs share the same episode length $H\in\mathbb{N}_{+}$, finite state space ${\mathcal{S}}$, finite action space $\mathcal{A}$, and initial state distribution $p_{0}\in\Delta({\mathcal{S}})$. Let $\bot$ be a default terminal state that is not contained in ${\mathcal{S}}$. The transition probabilities $\mathbb{P}_{p}\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}\times\mathcal{A}\to\Delta({\mathcal{S}}\cup\mathinner{\left\{\bot\right\}})$ and reward distributions $r_{p}\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}\times\mathcal{A}\to\Delta([0,1])$ of the players are not necessarily identical. We assume that the MDPs are layered¹¹1This is a standard assumption (see, e.g., [44]). It is worth noting that any episodic MDP (with possibly nonstationary transition and reward) can be converted to a layered MDP with stationary transition and reward, with the state space size being $H$ times the size of the original state space., in that the state space ${\mathcal{S}}$ can be partitioned into disjoint subsets $({\mathcal{S}}_{h})_{h=1}^{H}$, where $p_{0}$ is supported on ${\mathcal{S}}_{1}$, and for every $p\in[M]$, $h\in[H]$, and every $s\in{\mathcal{S}}_{h},a\in\mathcal{A}$, $\mathbb{P}_{p}(\cdot\mid s,a)$ is supported on ${\mathcal{S}}_{h+1}$; here, we define ${\mathcal{S}}_{H+1}=\mathinner{\left\{\bot\right\}}$. We denote by $S\mathrel{\mathop{\ordinarycolon}}=\left|{\mathcal{S}}\right|$ the size of the state space, and $A\mathrel{\mathop{\ordinarycolon}}=\left|\mathcal{A}\right|$ the size of the action space.

The interaction process between the players and the environment is as follows: at the beginning, both $(r_{p})_{p=1}^{M}$ and $(\mathbb{P}_{p})_{p=1}^{M}$ are unknown to the players. For each episode $k\in[K]$, conditioned on the interaction history up to episode $k-1$, each player $p\in[M]$ independently interacts with its respective MDP $\mathcal{M}_{p}$; specifically, player $p$ starts at state $s_{1,p}^{k}\sim p_{0}$, and at every step (layer) $h\in[H]$, it chooses action $a_{h,p}^{k}$, transitions to next state $s_{h+1,p}^{k}\sim\mathbb{P}_{p}(\cdot\mid s_{h,p}^{k},a_{h,p}^{k})$ and receives a stochastic immediate reward $r_{h,p}^{k}\sim r_{p}(\cdot\mid s_{h,p}^{k},a_{h,p}^{k})$; after all players have finished their $k$-th episode, they can communicate and share their interaction history. The goal of the players is to maximize their expected collective reward $\mathbb{E}\left[\sum_{k=1}^{K}\sum_{p=1}^{M}\sum_{h=1}^{H}r_{h,p}^{k}\right]$.

A deterministic, history-independent policy $\pi$ is a mapping from ${\mathcal{S}}$ to $\mathcal{A}$, which can be used by a player to make decisions in its respective MDP. For player $p$ and step $h$, we use $V_{h,p}^{\pi}\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}_{h}\to[0,H]$ and $Q_{h,p}^{\pi}\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}_{h}\times\mathcal{A}\to[0,H]$ to denote its respective value and action-value functions, respectively. They satisfy the following recurrence known as the Bellman equation:

where we use the convention that $V_{H+1,p}^{\pi}(\bot)=0$, and for $f\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}_{h+1}\to\mathbb{R}$, $(\mathbb{P}_{p}f)(s,a)\mathrel{\mathop{\ordinarycolon}}=\sum_{s^{\prime}\in{\mathcal{S}}_{h+1}}\mathbb{P}_{p}(s^{\prime}\mid s,a)f(s^{\prime})$, and $R_{p}(s,a)\mathrel{\mathop{\ordinarycolon}}=\mathbb{E}_{\hat{r}\sim r_{p}(\cdot\mid s,a)}\left[\hat{r}\right]$ is the expected immediate reward of player $p$. For player $p$ and policy $\pi$, denote by $V_{0,p}^{\pi}=\mathbb{E}_{s_{1}\sim p_{0}}\left[V_{1,p}^{\pi}(s_{1})\right]$ its expected reward.

For player $p$, we also define its optimal value function $V_{h,p}^{\star}\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}_{h}\to[0,H]$ and the optimal action-value function $Q_{h,p}^{\star}\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}_{h}\times\mathcal{A}\to[0,H]$ using the Bellman optimality equation:

where we again use the convention that $V_{H+1,p}^{\star}(\bot)=0$. For player $p$, denote by $V_{0,p}^{\star}=\mathbb{E}_{s_{1}\sim p_{0}}\left[V_{1,p}^{\star}(s_{1})\right]$ its optimal expected reward.

Given a policy $\pi$, as $V_{h,p}^{\pi}$ for different $h$’s are only defined in the respective layer ${\mathcal{S}}_{h}$, we “collate” the value functions $(V_{h,p}^{\pi})_{h=1}^{H}$ and obtain a single value function $V_{p}^{\pi}\mathrel{\mathop{\ordinarycolon}}{\mathcal{S}}\cup\mathinner{\left\{\bot\right\}}\to\mathbb{R}$. Formally, for every $h\in[H+1]$ and $s\in{\mathcal{S}}_{h}$,

We define $Q_{p}^{\pi},V_{p}^{\star},Q_{p}^{\star}$ similarly. For player $p$, given its optimal action value function $Q_{p}^{\star}$, any of its greedy policies $\pi_{p}^{\star}(s)\in\mathop{\mathrm{argmax}}_{a\in\mathcal{A}}Q_{p}^{\star}(s,a)$ is optimal with respect to $\mathcal{M}_{p}$.

For player $p$, we define the suboptimality gap of state-action pair $(s,a)$ as $\mathrm{gap}_{p}(s,a)=V_{p}^{\star}(s)-Q_{p}^{\star}(s,a)$. We define the mininum suboptimality gap of player $p$ as $\mathrm{gap}_{p,\min}=\min_{(s,a)\mathrel{\mathop{\ordinarycolon}}\mathrm{gap}_{p}(s,a)>0}\mathrm{gap}_{p}(s,a)$, and the minimum suboptimality gap over all players as $\mathrm{gap}_{\min}=\min_{p\in[M]}\mathrm{gap}_{p,\min}$. For player $p\in[M]$, define $Z_{p,\mathrm{opt}}\mathrel{\mathop{\ordinarycolon}}=\mathinner{\left\{(s,a)\mathrel{\mathop{\ordinarycolon}}\mathrm{gap}_{p}(s,a)=0\right\}}$ as the set of optimal state-action pairs with respect to $p$.

We measure the performance of the players using their collective regret, i.e., over a total of $K$ episodes, how much extra reward they would have collected in expectation if they were executing their respective optimal policies from the beginning. Formally, suppose for each episode $k$, player $p$ executes policy $\pi^{k}(p)$, then the collective regret of the players is defined as:

A naive baseline for multi-task RL is to let each player run a separate RL algorithm without communication. For concreteness, we choose to let each player run the state-of-the-art Strong-Euler algorithm [36] (see also its precursor Euler [46]), which enjoys minimax gap-independent [3, 8] and gap-dependent regret guarantees, and we refer to this strategy as individual Strong-Euler. Specifically, as it is known that Strong-Euler has a regret of $\tilde{O}(\sqrt{H^{2}SAK}+H^{4}S^{2}A$), individual Strong-Euler has a collective regret of $\tilde{O}(M\sqrt{H^{2}SAK}+MH^{4}S^{2}A)$. In addition, by a union bound and summing up the gap-dependent regret guarantees of Strong-Euler for the $M$ MDPs altogether, it can be checked that with probability $1-\delta$, individual Strong-Euler has a collective regret of order²²2The originally-stated gap-dependent regret bound of Strong-Euler ([36], Corollary 2.1) uses a slightly different notion of suboptimality gap, which takes an extra minimum over all steps. A close examination of their proof shows that Strong-Euler has regret bound (2) in layered MDPs. See also Remark LABEL:rem:clipping in Appendix LABEL:subsec:conclude-reg-bounds.

Our goal is to design multi-task RL algorithms that can achieve collective regret strictly lower than this baseline in both gap-dependent and gap-independent fashions when the tasks are similar.

Throughout this paper, we will consider the following notion of similarity between MDPs in the multi-task episodic RL setting.

If the MDPs in $(\mathcal{M}_{p})_{p=1}^{M}$ are $0$-dissimilar, then they are identical by definition, and our interaction protocol degenerates to the concurrent RL protocol [34]. Our dissimilarity notion is complementary to those of [6, 16]: they require the MDPs to be either identical, or have well-separated parameters for at least one state-action pair; in contrast, our dissimilarity notion allows the MDPs to be nonidentical and arbitrarily close.

We have the following intuitive lemma that shows the closeness of optimal value functions of different MDPs, in terms of the dissimilarity parameter $\epsilon$:

For each player $p$, the construction of its value function bound estimates relies on empirical estimates on its transition probability and expected reward function. For both estimands, we use two estimators with complementary roles, which are at two different points of the bias-variance tradeoff spectrum: one estimator uses only the player’s own data (termed individual estimate), which has large variance; the other estimator uses the data collected by all players (termed aggregate estimate), which has lower variance but can easily be biased, as transition probabilities and reward distributions are heterogeneous. Such an algorithmic idea of “model transfer”, where one estimates model in one task using data collected from other tasks has appeared in prior works (e.g.,  [39]). Specifically, at the beginning of episode $k$, for every $h\in[H]$ and $(s,a)\in{\mathcal{S}}_{h}\times\mathcal{A}$, the algorithm has its empirical count of encountering $(s,a)$ for each player $p$, along with its total empirical count across all players, respectively:

The individual and aggregate estimates of immediate reward $R(s,a)$ are defined as:

Similarly, for every $h\in[H]$ and $(s,a,s^{\prime})\in{\mathcal{S}}_{h}\times\mathcal{A}\times{\mathcal{S}}_{h+1}$, we also define the individual and aggregate estimates of transition probability as:

If $n(s,a)=0$, we define $\hat{R}(s,a)\mathrel{\mathop{\ordinarycolon}}=0$ and $\hat{\mathbb{P}}(s^{\prime}\mid s,a)\mathrel{\mathop{\ordinarycolon}}=\frac{1}{\left|{\mathcal{S}}_{h+1}\right|}$; and if $n_{p}(s,a)=0$, we define $\hat{R}_{p}(s,a)\mathrel{\mathop{\ordinarycolon}}=0$ and $\hat{\mathbb{P}}_{p}(s^{\prime}\mid s,a)\mathrel{\mathop{\ordinarycolon}}=\frac{1}{\left|{\mathcal{S}}_{h+1}\right|}$. The counts and reward estimates can be maintained by Multi-task-Euler efficiently in an incremental manner.

For each player $p$, based on these model parameter estimates, Multi-task-Euler performs optimistic value iteration to compute the value function estimates for states at all layers (lines 1 to 1). For the terminal layer $H+1$, $V_{H+1}^{\star}(\bot)=0$ trivially, so nothing needs to be done. For earlier layers $h\in[H]$, Multi-task-Euler iteratively builds its value function estimates in a backward fashion. At the time of estimating values for layer $h$, the algorithm has already obtained optimal value estimates for layer $h+1$. Based on the Bellman optimality equation (1), Multi-task-Euler estimates $(Q_{p}^{\star}(s,a))_{s\in{\mathcal{S}}_{h},a\in\mathcal{A}}$ using model parameter estimates and its estimates of $(V_{p}^{\star}(s))_{s\in{\mathcal{S}}_{h+1}}$, i.e., $(\overline{V}_{p}(s))_{s\in{\mathcal{S}}_{h+1}}$ and $(\underline{V}_{p}(s))_{s\in{\mathcal{S}}_{h+1}}$ (lines 1 to 1).

Specifically, Multi-task-Euler constructs estimates of $Q_{p}^{\star}(s,a)$ for all $s\in{\mathcal{S}}_{h},a\in\mathcal{A}$ in two different ways. First, it uses the individual estimates of model of player $p$ to construct $\underline{\text{ind-}Q}_{p}$ and $\overline{\text{ind-}Q}_{p}$, upper and lower bound estimates of $Q_{p}^{\star}$ (lines 1 and 1); this construction is reminiscent of Euler and Strong-Euler [46, 36], in that if we were only to use $\underline{\text{ind-}Q}_{p}$ and $\overline{\text{ind-}Q}_{p}$ as our optimal action value function estimate $\overline{Q}_{p}$ and $\underline{Q}_{p}$, our algorithm becomes individual Strong-Euler. The individual value function estimates are key to establishing Multi-task-Euler’s fall-back guarantees, ensuring that it never performs worse than the individual Strong-Euler baseline. Second, it uses the aggregate estimate of model to construct $\underline{\text{agg-}Q}_{p}$ and $\overline{\text{agg-}Q}_{p}$, also upper and lower bound estimates of $Q_{p}^{\star}$ (lines 1 and 1); this construction is unique to the multitask learning setting, and is our new algorithmic contribution.

To ensure that $\overline{\text{agg-}Q}_{p}$ and $\overline{\text{ind-}Q}_{p}$ (resp. $\underline{\text{agg-}Q}_{p}$ and $\underline{\text{ind-}Q}_{p}$) are valid upper bounds (resp. lower bounds) of $Q_{p}^{\star}$, Multi-task-Euler adds bonus terms $\text{ind-}b_{p}(s,a)$ and $\text{agg-}b_{p}(s,a)$, respectively, in the optimistic value iteration process, to account for estimation error of the model estimates against the true models. Specifically, both bonus terms comprise three parts:

where

and $L(n)\eqsim\ln(\frac{MSAn}{\delta})$.

The bonus terms altogether ensures strong optimism [36], i.e.,

In short, strong optimism is a stronger form of optimism (the weaker requirement that for any $p$ and $(s,a)$, $\overline{Q}_{p}(s,a)\geq Q^{\star}_{p}(s,a)$ and $\overline{V}_{p}(s)\geq V^{\star}_{p}(s)$), which allows us to use the clipping lemma (Lemma B.6 of [36], see also Lemma LABEL:lem:clipping-main in Appendix LABEL:subsec:conclude-reg-bounds) to obtain sharp gap-dependent regret guarantees. The three parts in the bonus term serve for different purposes towards establishing (6):

1. 1.

   The first component accounts for the uncertainty in reward estimation: with probability $1-O(\delta)$, $\mathinner{\!\left\lvert\hat{R}_{p}(s,a)-R_{p}(s,a)\right\rvert}\leq b_{\mathrm{rw}}\mathinner{\left(n_{p}(s,a),0\right)}$, and $\mathinner{\!\left\lvert\hat{R}(s,a)-R_{p}(s,a)\right\rvert}\leq b_{\mathrm{rw}}\mathinner{\left(n(s,a),\epsilon\right)}$.

2. 2.

   The second component accounts for the uncertainty in estimating $(\mathbb{P}_{p}V_{p}^{\star})(s,a)$: with probability $1-O(\delta)$, $\mathinner{\!\left\lvert(\hat{\mathbb{P}}_{p}V_{p}^{\star})(s,a)-(\mathbb{P}_{p}V_{p}^{\star})(s,a)\right\rvert}\leq b_{\mathrm{prob}}\mathinner{\left(\hat{\mathbb{P}}_{p}(\cdot\mid s,a),n_{p}(s,a),\overline{V}_{p},\underline{V}_{p},0\right)}$ and $\mathinner{\!\left\lvert(\hat{\mathbb{P}}V_{p}^{\star})(s,a)-(\mathbb{P}_{p}V_{p}^{\star})(s,a)\right\rvert}\leq b_{\mathrm{prob}}\mathinner{\left(\hat{\mathbb{P}}(\cdot\mid s,a),n(s,a),\overline{V}_{p},\underline{V}_{p},\epsilon\right)}$.

3. 3.

   The third component accounts for the lower order terms for strong optimism: with probability $1-O(\delta)$, $\mathinner{\!\left\lvert(\hat{\mathbb{P}}_{p}-\mathbb{P}_{p})(\overline{V}_{p}-V_{p}^{\star})(s,a)\right\rvert}\leq b_{\mathrm{str}}\mathinner{\left(\hat{\mathbb{P}}_{p}(\cdot\mid s,a),n_{p}(s,a),\overline{V}_{p},\underline{V}_{p},0\right)}$, and $\mathinner{\!\left\lvert(\hat{\mathbb{P}}-\mathbb{P}_{p})(\overline{V}_{p}-V_{p}^{\star})(s,a)\right\rvert}\leq b_{\mathrm{str}}\mathinner{\left(\hat{\mathbb{P}}(\cdot\mid s,a),n(s,a),\overline{V}_{p},\underline{V}_{p},\epsilon\right)}$.

Based on the above concentration inequalities and the definitions of bonus terms, it can be shown inductively that, with probability $1-O(\delta)$, both $\overline{\text{agg-}Q}_{p}$ and $\overline{\text{ind-}Q}_{p}$ (resp. $\underline{\text{agg-}Q}_{p}$ and $\underline{\text{ind-}Q}_{p}$) are valid upper bounds (resp. lower bounds) of $Q_{p}^{\star}$.

Finally, observe that for any $(s,a)\in{\mathcal{S}}_{h}\times\mathcal{A}$, $Q_{p}^{\star}(s,a)$ has range $[0,H-h+1]$. By taking intersections of all confidence bounds of $Q_{p}^{\star}$ it has obtained, Multi-task-Euler constructs its final upper and lower bound estimates for $Q_{p}^{\star}(s,a)$, $\overline{Q}_{p}(s,a)$ and $\underline{Q}_{p}(s,a)$ respectively, for $(s,a)\in{\mathcal{S}}_{h}\times\mathcal{A}$ (line 1 to 1). Similar ideas on using data from multiple sources to construct confidence intervals and guide explorations have been used by [37, 43] for multi-task noncontextual and contextual bandits. Using the relationship between the optimal value $V_{p}^{\star}(s)$ and and optimal action values $\mathinner{\left\{Q_{p}^{\star}(s,a)\mathrel{\mathop{\ordinarycolon}}a\in\mathcal{A}\right\}}$, Multi-task-Euler also constructs upper and lower bound estimates for $V_{p}^{\star}(s)$, $\overline{V}_{p}(s)$ and $\underline{V}_{p}(s)$, respectively for $s\in{\mathcal{S}}_{h}$ (line 1).

At each episode $k$, for each player $p$, its optimal action-value function upper bound estimate $\overline{Q}_{p}$ induces a greedy policy $\pi^{k}(p)\mathrel{\mathop{\ordinarycolon}}s\mapsto\mathop{\mathrm{argmax}}_{a\in\mathcal{A}}\overline{Q}_{p}(s,a)$ (line 1); the player then executes this policy in this episode to collect a new trajectory and use this to update its individual model parameter estimates. After all players finish their episode $k$, the algorithm also updates its aggregate model parameter estimates (lines 1 to 1) using Equations (3), (4) and (5), and continues to the next episode.