Provable Benefits of Multi-task RL under Non-Markovian Decision Making Processes
Authors: Ruiquan Huang
Yuan Cheng
Jing Yang
Vincent Tan
Yingbin Liang
Venue: ICLR 2024 poster (2024)
Keywords: Reinforcement learning
multi-task learning
bracketing number
predictive state representation
POMDP
Primary area: reinforcement learning

Abstract: In multi-task reinforcement learning (RL) under Markov decision processes (MDPs), the presence of shared latent structures among multiple MDPs has been shown to yield significant benefits to the sample efficiency compared to single-task RL. In this paper, we investigate whether such a benefit can extend to more general sequential decision making problems such as predictive state representations (PSRs). The main challenge here is that the large and complex model space makes it hard to identify what types of common latent structure of multi-task PSRs can reduce the model complexity and improve sample efficiency.
To this end, we posit a  joint model class for tasks and use the notion of $\eta$-bracketing number to quantify its complexity; this number also serves as a general metric  to capture the similarity of tasks and thus determines the benefit of multi-task over single-task RL. We first study  upstream multi-task learning over PSRs, in which all tasks share the same observation and action spaces. We propose a provably efficient algorithm  UMT-PSR for finding near-optimal policies for all PSRs, and demonstrate that the advantage of multi-task learning manifests if the joint model class of PSRs has a smaller $\eta$-bracketing number compared to that of individual single-task learning. We further investigate downstream learning, in which the agent needs to learn a new target task that shares some commonalities with the upstream tasks via a similarity constraint. By exploiting the learned PSRs from the upstream, we develop a sample-efficient algorithm that provably finds a near-optimal policy. 
Upon specialization to some examples with small $\eta$-bracketing numbers, our results further highlight the benefit compared to directly learning a single-task PSR.

Published as a conference paper at ICLR 2024
PROVABLE BENEFITS OF MULTI-TASK RL UNDER NON- MARKOVIAN DECISION MAKING PROCESSES
Ruiquan Huang"
Yuan Cheng' :
Jing Yang!
Vincent Tan
Yingbin Liang
ABSTRACT
In multi-task reinforcement learning (RL) under Markov decision processes (MDPs), the presence of shared latent structures among multiple MDPs has been shown to yield significant benefits to the sample efficiency compared to single-task RL. In this paper, we investigate whether such benefit can extend to more gen- eral sequential decision making problems such as predictive state representations (PSRs). The main challenge here is that the large and complex model space makes it hard to identify what types of common latent structure of multi-task PSRs can reduce the model complexity and improve sample efficiency. To this end, we posit a joint model class for tasks and use the notion of 1-bracketing number to quantify complexity; this number also serves as general metric to capture the similarity of tasks and thus determines the benefit of multi-task over single-task RL. We first study upstream multi-task learning over PSRs, in which all tasks share the same observation and action spaces. We propose provably efficient algorithm UMT-PSR for finding near-optimal policies for all PSRs, and demonstrate that the advantage of multi-task learning manifests if the joint model class of PSRs has smaller bybbrckktiing number compared to that of individual single-task learning. We further investigate downstream learning, in which the agent needs to learn a new target task that shares some commonalities with the upstream tasks via a de-
similarity constraint. By exploiting the learned PSRs from the upstream, velop samplc-efficient algorithm that provably finds a near-optimal policy. Upon specialization some examples with small y-bracketing numbers, our results further highlight the benefit compared to directly learning a single-task PSR.
INTRODUCTION
Multi-task sequential decision making, or multi-task reinforcement learning (MTRL) is subfield of reinforcement learning (RL) that extends the learning process across multiple tasks. Many real-world applications can be modeled by MTRL. For instance, in robotics and autonomous driving, different types of robots and vehicles in a shared environment can have different observational capabilities based on their sensors and learning goals. Other applications include personalized healthcare, weather forecasting across different regions, and manufacturing quality control on different types of products. The fundamental idea behind MTRL is to leverage the inherent similarities among a set tasks in order to improve the overall learning efficiency and performance. For Markov decision processes
(MDPs), a line of works (Pathak et al. 2017]] Tang et al. 20177 Oord et al.)
20188 Laskin et al. Lu et al. 2021 Cheng et al. 2022 Agarwal al.) 2022 Pacchiano et al. 2022) have explored
2020
multi-task representation learning and shown its benefit both practically and theoretically.
However, it is still open question whether such benefit can extend to more general sequential decision making problems, even in partially observable MDPs (POMDPs), let alone more general predictive state representations (PSRs). In this context, it is even unclear:
'Equal contribution.
1 Penn State University, State College, PA 16801, USA. {zzh514, yangiingnn @psu. edu National University of Singapore, 119077, Singapore 6 Ohio State University, Columbus, OH 43210, USA. 1iang. 8ttgg888900 08w.0su...uue... edu.
yuan. www.geu.nus......o.o mus..nuuuuuu vtangnus, edu..s

---

Published as a conference paper at ICLR 2024
When can latent similarity structure encompassed by multiple PSRs be potentially beneficial?
The challenges mainly emanate from two aspects. First, the large and complex model space makes hard to identify what rypes of common latent structure of multi-task PSRs can reduce the model complexity. The non-Markovian property of these problems implies that the sufficient statistics or belief about the current environmental state encompasses all the observations and actions from past interactions with the environment. This dramatically increases the statistical complexity. Even for finite observation space and action space, model complexity can be exponentially large in the number of observations and actions. Such complex parameter space makes it difficult to identify what types of latent similarity structure multi-task PSRs reduce the model complexity. Second, reduced model complexity does not necessarily result in benefit in statistical efficiency gain of RL. In RL, model learning and data collection are intertwined. The agent has to choose an exploration policy in each iteration based on the model learned in the past. Such iterative process introduces temporal dependence to the collected data, which makes the analysis of multi-task PSRs complicated.
In this paper, we answer the question above with upstream multi-task learning and downstream transfer learning. We summarize our contributions below.
1. To deal with the first challenge, we propose a unified approach to characterize the effect of task similarity model complexity by introducing the notion of the y-bracketing number for the joint model space of multiple tasks. Regardless of whether the concrete form of task similarity is implicit explicit, desirable task similarity should contribute reduce the n7bbrccktinng number compared to that without similarity structures. This significantly generalizes existing studies of multi-task MDPs that considered only specific task similarity structures.
2. We deal with the second challenge in both upstream and downstream learning. For the former, propose a novel multi-task PSRs algorithm called UMT-PSR, which features a pairwise additive distance-based based optimistic planning and exploration well confidence set construction based on the bracketing number of the joint model class. We then prove that if the bracketing number of the multi-task model class normalized by the number of tasks is lower than that of single task, UMT-PSR benefits from multi-task learning with these novel designs. We then provide several specific multi-task POMDP/PSR examples with low bracketing number to demonstrate that UMT-PSR is often more efficient than single-task learning.
3. We further employ the upstream learning to downstream learning by connecting upstream and downstream models via similarity constraints. We show that the downstream learning can identify a near-accurate model and find a near-optimal policy. Upon specialization to the examples used to elucidate the 7--backetinng numbers, our downstream results further highlight the benefit in comparison to directly learning parameters of PSRs without upstream information. Our analysis here features a novel technique of using Rényi divergence to measure the approximation error which guarantees the sub-optimality bound without requiring the realizability condition.
Our work is the first theoretical study that characterizes the benefits of multi-task RL with PSRs/POMDPs over its single-task counterpart.
RELATED WORK
Multitask representation learning and transfer learning have been extensively studied in RL, particularly under MDPs. Arora et al. (2020) demonstrated that representation learning can reduce sample complexity for imitation learning. Hu et al. (2021) analyzed MTRL with low inherent Bellman error Zanette et al., 2020) and known representation. IZhang & Wang (202T) studied multi-task learning under similar transition kernels. In contrast, Brunskill Li (2013) studied the benefit of MTRL when each task is independently sampled from distribution over finite set of MDPs. Recent studies have also considered the case where all tasks share common representation, including D' 'ramo et al. (2020) which demonstrated the convergence rate benefit on value iteration, Lu and Lu et al. (2021) which proved the sample efficiency gain MTRL under low-rank MDPs. Some recent work further took the impact of sequential exploration and temporal dependence in data into account. Considering sequential exploration with shared unknown representation, Cheng et al. (2022): ; Agarwal et al. (2022) studied reward free MTRL under low-rank MDPs as upstream learning and applied the learned representation from upstream to downstream RL. Pacchiano et al. (2022) focused on common low-dimensional linear representation and investigated MTRL under linearly-factored
MTRL under MDPs:
2

---

Published as a conference paper at ICLR 2024
MDPs. Lu et al. (0022) explored MTRL with general function approximation. There also exist few papers studying multi-task POMDPs from theoretical (Li et al. 2009) and practical (Omidshatiei et al. 2017) side. Note that all the above studies considered specific common model structures shared among tasks, whereas our paper proposes a unified way to characterize the similarity among tasks. Further, none of the existing studies considered model-based multi-task PSRs, which is the focus of our paper.
Single-task RL with PSRs and general sequential decision making problems: A general decision making framework PSR (Littman & Sutton 200T) was proposed to generalize MDPs and POMDPs. Since then, various approaches have been studied to make the problem tractable with polynomial sample efficiency. These methods include spectral type of techniques (Boots et al. 2011] Heiny et al.
2015 Jiang et al., 2018 Zhang et al. 2022), methods based on optimistic planning and maximum
log-likelihood estimators together with confidence set-based design (Zhan et al. 2022 Liu et al.
2022). the bonus- basud approaches (Huang et al. 2023), value-based actor-critic approaches (Uehara et al. 2022). posterior sampling methods (Zhong et al. (0022). Chen et al. (2022) further improved
the sample efficiency for previous work including OMLE (Liu et al. 2022), MOPS
(Agarwal &
Zhang 2022). and E2D (Foster et al. 2021..
3 PRELIMINARIES
Notations. For any positive integer N, we use [M] to denote the set .N,...,,N, For any vector æ, the i-th coordinate of æ is represented as [zr]:. For set X, the Cartesian product of N copies of X denoted by X For probability distributions and Q supported on a countable set X, the total variation distance between them is Drv (IP, Q) ΣS [P(x>) Q(x)), and the Rényi divergence of order for 1, between them is D (IP, Q) a log Esp (dLP (dLPP/)))))) [dQ)=
3.1 THE NON-MARKOVIAN DECISION MAKING PROBLEM
We consider episodic decision making process, which is generally non-Markovian, with an observation space O and finite action space A. We assume that the process is episodic and each episode contains H steps. i.e., with horizon H. At each step, the evolution of the process is controlled by an underlying distribution IP, where Fo,, oo..... -h-1, 01...... 0/-1) is the probability of visiting on at step given that the learning agent has observed 0g E O and taken action a, E A for previous steps t E [h- 1]. And the learning agent receives a reward at each episode determined by the reward function R (O x A) + [0, 1]. We denote such a process compactly as P = (O, H, IP, R). For each step h, we denote historical trajectory T (01, a,,..., On, an). the set of all possible historical trajectories as Hr = (O x A)", the future trajectory as Wh. := (oh+1. 0h+1. 0H, a and the set of all possible future trajectories as Ωr (O x A) =-=
The agent interacts with the environment in each episode as follows. At step 1, a fixed initial observation 01 is drawn. At each step h [H) due to the non-Markovian nature, the action selection and environment transitions are based on whole history information. Specifically, the agent can choose an action dn based on the history Th-1 and the current observation on with strategy (probability) (a Th-1, 1, on). We denote such a strategy as a policy, and collect the policies over H steps into (mh 1= and denote the set of all feasible policies II. Then the environment takes a transition to oh+1 based on P(o.+1|TT)) The episode terminates after H steps.
For any historical trajectory The we further divided it into T (oi,.... o%.) and T (ai..... ag) which is observation and action sequences contained in The respectively. Similar to The for the future trajectories wh, we denote wi. as the observation sequence in wh. and w/t as the action sequence in wh. For simplicity, we write (TT).) π xaa (on, 7h-1) ... . (Ta1 (o1) to denote the probability of choosing the sequence of actions T given the observations Tğ under the policy π. We denote pa as the distribution of the trajectories induced by the policy æ under the dynamics P. The value function of a policy π under IP and the reward R is denoted by V = Erg=r= [R(HH)] The primary learning goal is to find an c-optimll policy म, which is one that satisfies max, Ve,RR VE, S.
Given that addressing general docision-making problem entails an exponentially large sample complexity in the worst case, this paper focuses on the low-rank class of problems as in Zhan et al. (2022); Liu al. (2022); Chen al. (2022). Before formal definition of the low-rank problem,
3

---

Published as a conference paper at ICLR 2024
introduce the dynamics matrix D,, E R[Ha|x|9al for each h, where use Th E H, and wh E S2, to index the rows and columns of the matrix D,, respectively, and the entry at the Ty, row and w/ column of D, equals to the conditional probability P(wg, T (k.w))
Definition 1 (Rank-r sequential decision making problem). A sequential decision making problem is rank r for any h, the model dynamics matrix 1 has rank at most v.
As result, for each h, the probability of observing w can be represented by linear combination of probabilities on of future trajectories known to the agent called core tests Qn = ( dh C
On where dA 2 r. Specifically, there exist functions m 2t - IRd,, 1/ : Hi - Rdda such that (i) the value of the ε-th coordinate of (T.) equals to the conditional probability (( T -laa 7 on
(ah Th), where of, and af, to denote the observation and the action sequences of q and (ii) for any
Wh E Sh, Th Hi. the conditional probability can be factorized as
P( TRTE w()) m(w) (T(T))
(1)
Predictive State Representation. Following from Theorem C.I in Liu et al. (2022), given core tests {Qh any low rank decision making problem admits (self-consistent) predictive state repre- sentation (PSR) 0 = [(db,, MR)) Ma such that Eq. 1 can be reparameterized by O. Mathematically, For any h E [ Th HH,, Wh Elpy
m (w))) T o M (o) ...
44+1),
(TT)) = MA aA) ... ...MM (01. 01 11)
and [oox,a))oxx h++ M, (on, a,)) oT. For ease of the presentation, assume th is known
The following assumption is standard in the literature (Liu al. 2022 Chen et al. 2022).
Assumption 1 (7- conditoeed PSR). We assume any PSR 0 = {(o,. M, ) considered in this paper cywell-condittonedd for some & i.e.
Vh [H], max, m Ea Enn,n, T(Wn|TA]|m(as Txl 4
(2)
In the following context, use Py indicate the model determined by the PSR 0. For simplicity, we denote V as Vē_R- Moreover, let 2 - {af )d be the action sequence set from core tests which is constructed by climinating any repeated action sequence. The set O is also known as the core action sequence set. The set of all rank-r and --welll-ooniitoneed PSRs is denoted by θ.
3.2 UPSTREAM MULTI-TASK LEARNING
In upstream multi-task learning, the agent needs to solve N low-rank decision making problems (also known as source tasks) at the same time instead of only one single problem (task). The of N source tasks is denoted by (P,, > where P,, = (O, A. H, To; Rn), and 6% E e In other words, all N tasks are identical except for their model parameters 0; {(o) n. M, Ji= and reward functions Rn, Morcover, we denote the model class multi-tase task PSRs as O (the subscript stands for upstream), subset of N
")
The goal of the upstream learning consists of two parts: (i) Finding near-optimal policies for all N tasks on average. Mathematically, given accuracy level c, the set of N policies that produced by the algorithm { ...., # should satisfy 2 E (max, V V Ve,, S e; (ii) Characterizing the theoretical benefit of multi-task PSRs learning in terms of the sample complexity. compared to learning each task individually.
3.3 BRACKETING NUMBER OF JOINT PARAMETER SPACE
One critical factor that affects the efficiency of multi task learning compared to separate task learning is the presence of shared latent structure among the tasks, which yields reduced model space in
'The sample complexity of learning % if is unknown is relatively small compared to the learning of the other parameters (Liu et 2022).
For simplicity, assume all tasks have the same rank and y, but have different core test sets. The extension different ranks and 7 's is straightforward.
4

---

Published as a conference paper at ICLR 2024
multi-task PSRs learning, compared to separately learning single tasks the Cartesian product of N model spaces (see Figure I for an illustration in 2 dimensions). Consequently, this reduction in model complexity can ultimately lead to improved sample efficiency. Unlike the specific shared model structures among multiple tasks that the previous works studied, such as shared representation in Cheng al. (2022) and similar transition kernels in Zhang & Wang (202T), here we focus a general shared model space and use the notion of the y-bracketing number to quantify the complexity of the joint model space. Such a notion plays a central role in capturing the benefit of multi-task PSR learning over single-task learning.
We start with a domain X' and a single task function class F, in which each element f X - Rt,, For the multi-task case, the function class is subset Fu of J.N.
D2
model class P of task
Definition 2 (17- Br-Baacketing number of vector-valued function class F., w.r.t. II-. Given two vector-valued functions 1 and
(p(r.1) plr)2) (p(m11))p()11)
X + R.Y, the bracket [1. gl is the set of all functions f fe F. satisfying 1 S gf An n-bracket is a bracket [L, g)
with l|g -1 minimum number of t--backkes needed to cover F
II| < 17. The bracketing number Nr |H-ID )) is the
model class P task
In this paper. we are interested in the bracketing number of the joint model space, i.c., distribution spaces over (O x A)" parameterized by u... For simplicity, we use (O,.) to de- note the yy-bracketing number of {(Po...... Poy E O.))
(a) Independent learning of each task where the joint model class p?
Fo S pa
w.r.t. the to policy weighted norm I1-11.. where the too policy weighted norm between two vector-valued functions
model class
........9... } defined on (Ox A) 1= 14,..., |N3 and g (
of task
is equal to [gg - - = max max;,e11 Σn (4(H) (((T))(T,(T)) As we will show later, a lower by-bracketing number of the joint model space results in lower sample structures and hence their joint model space will have much lower yy-bracketing number compared the product of model spaces (i.e, treating the model of each task separately). We provide several such examples of Man-Morovvinn decision processes in Section 4.3 We provide more examples of MDPs with their 17- brackkkeiing numbers in Appx. F Notably, there can be much richer scenarios beyond these examples.
(n(t)t) ( pl
complexity in multi-task PSR learning.
In practice, it is common tasks share certain common model
model class P of task
(b) Joint learning of tasks with shared la- tent model structure. The joint model class is Fu, strict subset of p2.
Figure I: Reduction in y yy-brackeing number when F,, S. p2
3.4 DOWNSTREAM TRANSFER LEARNING
In downstream learning, the agent is assigned with a new target task P. (O,,A,H,P, Ro). where 0j E O, which shares some similarities with source tasks to benefit from upstream learning. Here, we capture the shared structure between upstream and downstream tasks via the similarity constraint C(fg. ..... 0%) < where N R=4, na E N. The similarity constraint establishes the relationship between the downstream target task and the upstream source tasks. Hence, the downstream model class is given by θy (0p θ C(0,, 0f 0%) S 0). We note that the similarity constraint is general enough to capture various relationships between upstream and downstream tasks. For example, Cheng al. (0022) consider the case when the downstream task shares the same representation as upstream tasks, which is equivalent to assuming [C(eo..... Oy)],, [] - d0l12. 2. where n E [N] and ((*.a) is the representation of task n. However, the similarity constraint allows much richer beyond the above example, for example, downstream tasks can have
similar but not the same representations as the upstream, or may share only some representation
features, but not of them.
Wee write that two vectors a, b satisfy S b if a is coordinate-wise nonnegative. We say that collection of sets ..., S, cover a set S if St
5

---

Published as a conference paper at ICLR 2024
The goal of the downstream learning is to find near-optimal policy, by exploiting the constraint similarity with upstream tasks and utilizing upstream knowledge to achieve better sample efficiency compared with learning without upstream knowledge.
UPSTREAM LEARNING OVER MUSTI-TASK PSRs
We present our upstream algorithm in Section 4.1
characterize its theoretical performance in Sec- tion 4.2 and present examples to validate the benefit of upstream multi-task learning Scction 4.3
We use bold symbol to represent the multi-task parameters or policy. Specifically, 0 = (01.... xn), and (TI....., TN). ). We define QA max, max, 10% where 2% is the core action sequence set of task n at step The policy, denoted by 1/ (T,, 7²), takes x at the initial - steps and switches to n' from the fi-th step. Lastly, uxx represents the uniform distribution over the set X.
4.1 ALGORITHM:: UPSTREAM MULKI-TASK PSRS (UMT-PSR)
We provide the pseudo-code of our upstream multi-task algorithm called Upstream Multi- Task PSRs (UMT-PSR) in Algorithm
This iterative algorithm consists of three main steps as follows.
Algorithm Upstream Multi- Task PSRs (UMT-PSR)
Input: B₁ = model class, estimation margin 3(N), maximum iteration number K.
2: for k = 1...., K" do
3: Set nh = arg maX maxe,0.cB EB. EmE(n] Drv (Fz,, P..")
4: for n, h E [N] [H] do
5: Use un to collect data TH
k.k.
6: end for
7:
Construct Bk+1
E e
log
p 9,
(T
(TH
>
max g
Σ
log
( B
NE(R] WE[N]
AE( nE(N.
8: end for
9: Output: Any 0 E BK+1, and a greedy multi-task policy π arg max EMEINy
V Rn.Ra
Pairwise additive distance based multi-task planning (Line 3): To promote joint planning among tasks, a natural choice to measure the distance between two multi-task models is the distance between the two product distributions Poj x... x PAN and pat x.. XPEN However, such "distance between product distributions" is not sufficient to guarantee the accuracy of the individual models of each task, which is needed in the analysis of the of the individual value functions. Hence, we propose to the "pairwise additive distance" for our planning, defined D (0, & EnE[N] Drv Pat
N.K) More specifically, at each iteration k. UMT-PSR seleets a multi-tasi policy π = (T 1,k that maximizes the largest pairwise additive distance maxø. or D, (0, O') within the confidence set B (which will be specified later). An important property of B is that it contains the true model 0 with high probability. Using this property, the largest pairwise additive distance serves an optimistic value of the uncertainty Dr (0',0 0) for any multi-task model E B
Multi-task exploration (Line 5): Building upon the planning policy Th for each task n and each step h, UMT-PSR executes the policy n.k k for first /: - steps, and then uniformly selects an action sequence in A x Q for the following H H-h+) steps. In particular, at step h, UMT-PSR uniformly takes an action in A, and then uniformly chooses core action sequence a, such that regardless of what the observation sequence is, UMT-PSR always plays the action in the sampled core action sequence. In summary, for each (n, h) E [N] x [M] UMT-PSR adopts the policy vs Vn(T"" ni,k
Ax
to collect sample trajectory TH
n.k.h We abbreviate un ". "Ax Uh

---

Published as a conference paper at ICLR 2024
Confidence set construction via bracketing number of joint model class (Line 7): Given the sampled trajectories, UMT-PSR calls a maximum likelihood estimation oracle to construct the multi- task confidence set. A novel element here is the use of the bracketing number of the joint model class to characterize estimation margin ((M), which is upper bound of the gap between the maximum log-likelihood within ,, and the log-likelihood of the true model. Such a design provides a unified way for any MTRL problem and avoids individual design for each problem in case-by-case manner.
4.2 MAIN THEORETICAL RESULT
The following theorem characterizes the guarantee of the model estimation and the sample complexity to find near-optimal multi-task policy.
Theorem 1. Under Assumption 1 for any fixed 6 >> O, let θ. be the multi-task parameter space, 3( = ei (log KHN + log log.NN, (Q.).. where e > 0 and n) S KHN- Then with probability at least
1 - 6, UMT-PSR finds multi-task model 0 = (0i,..., 0N) such that
N Ex max DTV (P'5" oz
Ô
QA
B K
In addition, if K earlAl HD( for large enough 02 > UMT-PSR produces multi-task policy
π ,...,,,), ) such that the average sub-optiiaiiy gap is at most c, i.e.
V RR -o -V. Sc (max,en
[ shows that with the sample complexity Ō( 00144193 112g(0) UMT-PSR identifies an e-optimal multi-task policy. As comparison, the best
Benefits of multi-task learning: Theorem
known sample complexity of a single-task PSR RL is given by O r1A122
=**, in Chen et al. 2022 122
7222
where ((1) = ( H²) scales the logarithm of the bracketing number of a single-task PSR with rank r. is clear that as long as ((N) < (,) then UMT-PSR enjoys multi-task benefit in the sample complexity. In Section 4.3] we will provide several example multi-task POMDPs/PSRs to illustrate that such a condition can be satisfied broadly.
Next, we make a few comparisons concerning ((N), known sample complexity given in Chen et al. 2022 (ii If none of tasks share any similarity, i.e. O n = N we have 3(N) N NB and the sample complexity does not exhibit any benefit compared to learning the tasks separately. This coincides with the intuition that in the worst case, multi-task learning is not required. (iii) The benefits of multi-task learning
(i) N = 1, Theorem matches the best
more evident when g(N) /N decreases. An extreme example is that when all tasks also share the same dynamics, leading to ((N) (()), In this case, multi-tast taski learning reduces to the batch setting and as the batch size increases, the iteration number decreases linearly in N.
4.3 IMPORTANT EXAMPLES OF MULTI-TASK PSRS
As shown in Section 3.2 and Theorem I for multi-tast models with low brbracketing number, i.e., satisfying ((N) < N J UMT-PSR exhibits better sample complexity than single-task learning. In this section, we provide example multi-task POMDPs and PSRs and show that their bybracketingg number satisfies the condition. Detailed proof's for these examples can be found in Appendix [..2
We consider tabular POMDPs, which is a classic subclass of PSRs. Specifically, the dynamics in POMDPs consist of H transition distributions (T', Axxxx - [0, 11 and H emission distributions (Or, S x O + [o. 1] where S is a finite state space. The states capture the entire system information, but are not directly observable. In POMDPs, at cach step h, the current system state is sh, the agent observes on with probability (o, ss)) Then, if the agent takes an action a, based on previous observations Op..... O1 and actions a d1, the system state transits to $ with probability Tn (s (8h, a.). We use the notation Po (O.A, H, S, T,, O, R) to represent a POMDP instance. Note that the tuple (S, T, O) in POMDPs determine the general dynamics P in PSRs. If all tasks share the same state, observation, and action spaces, then Pp. (O, A, H. S, T 0", R) represents the model of task n.
Muli-task POMDPs.
7

---

Published as a conference paper at ICLR 2024
Example 1 (Multi-task POMDP with common transition kernels). All tasks (i.., all POMDPs) share the same transition kernel, i.e., there exists a set of transition distributions {Ty such that T% =T;, for all n E [N] and h E [.]. The emission distributions can be different. Such scenario arises if the agent observes the same environment from from different angles and hence receives different observations. Then, g(N) is at most ((H((S?)AI) + |S||O|
[|[][]) N) log H|01 A whereas the single lask ((1) is given by ((H(IST)A) + [8|[0]) log (IO)) Clearly, 3(N) < N2
"
Multi-task PSRs: We next provide two example multi-task PSRs, in which tasks do not share common model parameters. In these examples, the similarities among tasks could alternatively be established via implicit connections and correlations in latent spaces, which reduce the complexity of the joint model class, hence the estimation margin and the sample complexity of algorithms significantly compared with separately learning each single task.
Example 2 (Multi-tas PSR with perturbed models). Suppose there exist latent base task Pb. and finite noisy perturbation space A. Each task n E [N] is a noisy perturbation of the latent base task and can be parameterized into two parts: the base task plus task-specified noise term. Specifically, for each step h E [H] and task n E [N], any (o,a) E O X A, we have
MM an) (o a1.) + + an), AR E.
Such a multi-task PSR satisfies that 3(M) = 0 + HN log A whereas 3(1) for a single task is given by Ö ) Clearly, 3(M) < ( holds if log |A] < ((AHH, which can be easily satisfied for small-size perturbation environments. Hence, the multi-task PSR benefis from a significantly reduced sample complexity compared to single-task learning.
Example 3 (Multi-task PSR: Linear combination of core tasks). Suppose that the multi-task PSR lies in the linear span of m core tasks (P1..... .. P.m.. J. Specifically. for each task n E [N], there exists
coefficient vector ax = (aff' ... ( E R. s.I. for any E {H) and (on, an)) EEO x A.
67, an) = Ei= oj 0) (oh,
Oj MMM My an) = If M; (on, an).
For regularization, we assume 0 S of for all [mm] and E [N], and E of 1 for all n E [N). It can be shown that g( 0(m{)~100|A||H N)). whereas 3(L) = 012100)11111)) Clearly, B(N) Ng(() holds if m S min(N, ( H which satisfied in practice.
DOWNSTREAM LEARNING FOR PSRs
In downstream learning, the agent is assigned new task Po (O,A, H, Ro), where Oj E Oj- and 6ğ is defined in Section 3.4 As explained in Section 3.4 upstream and downstream tasks are connected via the similarity constraint C(O,, .i 0%) 0. Therefore, the agent can use the estimated model parameter 0,.... By in the upstream construct an empirical candidate model
class for the downstream task as Ó: 10,- E C/((0, 0 0........ 0x) 0). Then for downstream learning, we adopt the standard OMLE (Liu et al. 2022 Chen et al. 2022) for the model class Qy.
The sample complexity of downstream learning will be determined by the bracketing number of Óg. which is nearly the same as that of the ground truth Ojj. Since the similarity constraint will significantly reduces the complexity of the model parameter space, the bracketing number of Qg should be much smaller than that of the original parameter space O. In this way, the downstream can benefit from the upstream learning with reduced sample complexity. In the following subsections. we first characterize the performance guarantee for downstream learning in terms of the bracketing number of éz and then show that the similarity constraint reduces the bracketing number for the examples given Section 4.3
5.1 THEORETICAL GUARANTEE FOR DOWNSTREAM LEARNING
One main challenge in the downstream learning is that the true model may not lie in Qg. To handle this, we employ Rényi divergence to measure the "distance" from the model class to the true model as follows, mainly because its unique advantage under the MLE oracle: the Rényi divergence of order a with a > 1 serves as an upper bound on the TV distance and the KL divergence, and thus has more robust performance.
8

---

Published as a conference paper at ICLR 2024
Definition 3. Fix 1. The approximation enror of Q under a-Rényi divergence is defined as (6 e. max o 00) 1P
a-1 Theorem 2. Fix a > 1. Let 60 = ea. (Qg), 8 = co No (6g) H + { (-iepzo) + 1) log KH
for some large cu, where n S KB Under Assumption with probability at least 6, the output of Algorithm satisfies that
max Drv (60 - 00
Q4
VED)
Benefits of downstream transfer learning: Theorem shows that when 6o <²//4, with sample (A! ). OMLE identifies an c-optimal policy for the downstream task. As comparison, the best known sample complexity for single-task PSR RL without transfer learning is Ō( Or"aaLA where 3 = (log (+)) (Chen et al. 2022). It is clear that as long as So < B, then downstream learning enjoys transfer benefit in the sample complexity.
complexity at most Ö 2
Notably, in the realizable case when 6o 0, i.e. 0g E θy: we must have 30 = NI, S B, since éø O In the non-realizable case when 6o > 0, compared to the realizable case, the estimation error of &g has an additive factor of O(v + vI/(K(e - 1))) after hiding system parameters. We remark that this factor shrinks if the approximation error of éμ decreases and the order of Rényi divergence grows, which coincide with the intuition.
5.2 EXAMPLES DOWNSTREAM LEARNING TASKS
We revisit the examples presented in upstream multi-task learning, specifically Examples | to and subsequently extend their application downstream tasks under the realizable setting. With the prior knowledge obtained from upstream learning, these examples exhibit reduced r--bracketing number, and hence benefit in the sample efficiency. Detailed proofs are Appx.
Example (Multi task POMDP with Common transition kernels). Suppose T is the output from UMT-PSR. In this case, the downstream θ is constructed by combining and all possible emission distributions. Then SP = (())) However, for POMDP without prior knowledge, 3 Ó(H(IS))AI + [S](OO)). Clearly, 3o S 3. indicating the benefit of downstream learning.
For PSRs without prior knowledge, we have 3PSR 0(1-(0)1A
Example 2 (Multi-task PSR with perturbed models). The downstream task Po is noisy perturbation of base task Pb. Specifically, for each step h E [H), any (0,a) O x A. we have
dh M. (on, a..) MA (i a.,) AR an). E
PH
Then, Bu = Ö(H log ) which is much lower than 3PSR iflog A & O(P))O
Example 3 (Multi-task PSR: Linear combination of core tasks). The downstream task Pu lies in the linear span of L upstream tasks (e.g. the firs L source tasks). Specifically, there exists coefficient
vector a = (a), ,aa)) E RRL s.l. for any lt E [H] and (on, ax) E EO x A,
Mo, an) = E/ ( (o), For regularization, we assume S for is much smaller than 3PSR if L 5 N A1)
1i = Ete af
for all E [T., and El= at = 1. Then Go = Ô(HH), which
CONCLUSION
In this paper, we study multi-task learning on general non-markovian low-rank decision making problems. Given that all tasks share the same observation and action spaces, using the approach of PSRs, we theoretically characterize that multi-task learning presents benefit over single-task learning if the joint model class of PSRs has smaller 77 brackeetinng number. We also provide specific example multi-task PSRs with small 17- bracketing numbers. Then, with prior knowledge from the upstream, show that downstream learning is more efficient than learning from scratch.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The work of R. Huang and J. Yang was supported in part by the U.S. National Science Foundation under the grants CNS-1956276 CNS-2003131 and CNS-2030026. The work of Y. Liang was supported in part by the U.S. National Science Foundation under the grants ECCS- 3860. DMS- 2134145, and CNS-21 12477... The work of Y. Cheng and V. Tan was supported by the Singapore Ministry of Education Academic Research Fund (AcRF) Tier under grant number A-8000423-00-00 and the Singapore Ministry of Education (AcRF) Tier 1 under grant number A-80001 89-01-00.
IMPACT STATEMENT
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
REFERENCES
Alekh Agarwal and Tong Zhang. Model-based rl with optimistic posterior sampling: Structural conditions and sample complexity. Advances in Neural Information Processing Systems, 35: 33294, 2022.
Alekh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang, Mengdi Wang, and Xuezhou Zhang. Provable benefits of representational transfer in reinforcement learning. arXiv preprint 140551, 2022.
Sanjeev Arora, Simon S. Du, Sham M. Kakade, Yuping Luo, and Nikunj Saunshi. Provable rep- resentation learning for imitation learning via bi-level optimization. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Leurning Research, pp. 367-376. PMLR, 2020. URL http: http:///rcceeddnns.. mm www.ppressvvve 3/arora20a.....mo htt.oraa00a.hmmll
Byron Boots, Sajid M Siddiqi, and Geoffrey J Gordon.
Closing the learning-planne planning loop with predictive state representations. The International Journal of Robotics Research, 30(7)9954-966, 2011.
Emma Brunskill and Lihong Li. Sample complexity of multi-task reinforcement learning. In Ann E.
Nicholson and Padhraic Smyth (eds.), Proceedings of the Twenty-Ning -inth Conference on Uncer- tainty in Artificial Intelligence, UAI 2013, Bellevue, WA, USA, August 1 11-15, 2013. AUAI Press, jsp ?mmnu=
2013. URL httph:://ssssssses p wwtpt/tttggg//aaiaiiiiaeeeeeeeet weciiispllyyiiii www.laattteeeeeeee 1 2 2ttpeadddddd1311dededdeeeeee 1d=29
Fan Chen, Yu Bai, and Song Mei. Partially observable rl with b-stability: Unified structural condition and sharp sample-efficient algorithms. arXiv preprint 1iXXv:2209.1499000 2022.
Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang. Provable benefit of multitask representation learning in reinforcement learning. In Advances Neural Information Processing Systems, November 28- December 9, 2022, New Orleans, USA, volume 35, 2022.
Carlo D' D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Sharing knowl- edge in multi-task deep reinforcement learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL
net wtw.:e/wwwww...eeeoooeeeddddeeee dd
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487, 2021.
Ahmed Hefny, Carlton Downey, and Geoffrey J Gordon. Supervised learning for dynamical system learning. Advances in neural information processing systems, 28, 2015.
Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear RL. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event,
10

---

Published as a conference paper at ICLR 2024
volume 139 of Proceedings of Machine Learning Research, pp. 4349 4358. PMLR, 2021. URL http://proceeddnngs .mw..innsm.i......mmm.mm......e. wwl..
httm
Ruiquan Huang, Yingbin Liang, and Jing Yang. Provably efficient ucb-type algorithms for learning predictive state representations. arXiv preprini arXiv:2307.004055 2023.
Nan Jiang, Alex Kulesza, and Satinder Singh. Completing state representations using spectral learning. Advances in Neural Information Processing Systems, 31, 2018.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning, pp. 5639-5650. PMLR, 2020.
Hui Li, Xuejun Liao, and Lawrence Carin. Multi-task reinforcement learning in partially observable stochastic environments. Journal of Machine Learning Research, 10(5), 2009.
Michael Littman and Richard S Sutton. Predictive representations of state. Advances in neural information processing systems, 14, 2001.
Qinghua Liu, Praneeth Netrapalli, Csaba Szepesvari, and Chi Jin. Optimistic mle-a a generic model-based algorithm for partially observable sequential decision making. arXiv preprint arXiv:2209. 14997, 2022.
Rui Lu, Gao Huang, and Simon S Du. On the power of multitask representation learning in linear mdp. arXiv preprint wrXiv:2106.08055 2021.
Rui Lu, Andrew Zhao, Simon S Du, and Gao Huang. Provable general function class representation learning in multitask bandits and mdp. Advances in Neural Information Processing Systems, 35: 11507-11519, 2022.
Shayegan Omidshafiei, Jason Pazis, Christopher Amato, Jonathan P How, and John Vian. Deep decen- tralized multi-task multi-agent reinforcement learning under partial observability. International Conference on Machine Learming, pp. 2681 16900. PMLR, 2017.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 187..0180..0374, 2018
Aldo Pacchiano, Ofir Nachum, Nilseh Tripuraneni, and Peter Bartlett. Joint representation training in sequential tasks with shared structure. arXiv preprint arXiv:2206. /2441, 2022.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference machine learning, pp. 2778-2787. PMLR, 2017.
Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, OpenAl Xi Chen, Yan Duan, John Schulman, Filip DeTurck, and Pieter Abbeel. # exploration: A study of count-based exploration for deep reinforcement learning. Advances in neural information processing systems, 30, 2017.
Masatoshi Uehara, Ayush Sekhari, Jason D Lee, Nathan Kallus, and Wen Sun. Provably efficient re- inforcement learning in partially observable dynamical systems. arXiv preprint 22Xi:22206.12020, 2022.
Tim Van Erven and Peter Harremos. Rényi divergence and kullback-leibler divergence. IEEE Transactions on Information Theory, (0(7):3777-882) 2014.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint 1rXiv:1011.3027, 2010.
Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Learning near optimal policies with low inherent bellman error. In International Conference on Machine Learning, pp. 10978-10989, PMLR, 2020.
Wenhao Zhan, Masatoshi Uchara, Wen Sun, and Jason D Lee. Pac reinforcement learning for predictive state representations. arXiv preprint arXiv:2207.05738, 2022.
11

---

Published as a conference paper at ICLR 2024
Chicheng Zhang and Zhi Wang. Provably efficient multi-task reinforcement learning with model transfer. Advances in Neural Information Processing Systems, 34:19771-19783, 2021.
Zhi Zhang, Zhuoran Yang, Han Liu, Pratap Tokekar, and Furong Huang. Reinforcement learning under a multi-agent predictive state representation model: Method and theory. In The Tenth International Conference on Learning Representations (ICLR 2022), 2022.
Han Zhong, Wei Xiong, Sirui Zheng, Liwei Wang, Zhaoran Wang, Zhuoran Yang, and Tong Zhang. posterior sampling framework for interactive decision making. arXiv preprint arXiv:2211.01933 2022.
12

---

Published as a conference paper at ICLR 2024
Supplementary Materials
A MULTI-TASK MLE ANALYSIS
The following lemma shows that the true model lies in the confidence set with high probability. Proposition (Confidence Set), For all k E [K], and for any 0 = (01,...,00 ......,,)) E θ let S HH(NKH), g(N) = e log (K) for some e 2 0. With probability at least 5, for any k E [K], we have
ΣΣΣD RSN ISk ASH
log
(TRth)) (TH
ΣΣΣ DSN ISRASH
log
-B(N). (
(4)
Consider of p-brackets, denoted by O,, that covers θ For any we can find two measures in , parameterized by 0 and 0 such that
Proof.
Vn, TH. rT, Vn, Drv Th
P (TH) Po (TH) 2 Pe. (TH),
(5)
(TH), Po. (TH) S n7.
Note that the above two inequalities imply that Enp 0n
(TH) S 1 + EnH
2. (TH) S 1 + n..
Then, we have
pog On
(TH
E
exp
log
1=1
Pp 0
(TH
R=
(TH
E
exp
log
E
[=
0
(TH
O
(TH To
k- |
P
(TH m.t.
(TH knn
exp
log
On
12
1=1
h
6
(TH
h
h-1
p
(T
$
exp
Σ
log
1+
NKH
12
1=1
(TH
NKH
Se.
NKH
Therefore, by Chernoff type bound, we have
FR
(7
p mi an
exp
Σ
log
(T
(T
(T
P
log
log(1/8)
e8.
1/8
(T
WEN
0%
Taking union bound over all k0, E O x [K] and rescaling 8. have for any O O
(T7.5.h))
IP
Vo v E [K].
log
20g(eKKN,
$ 8.
P 0:
(T
The proof is finished by noting that 0 is optimistic measure (see Equation
13

---

Published as a conference paper at ICLR 2024
The following lemma establishes the relationship between the Hellinger- -qquiree--quurred distance and the difference log likelihood functions between true parameters and any possible parameters from the model class.
Proposition 2. Let 17 11/ g(d) probability at least - 8. we have, for any 0 (01,..., .N. E O the following inequality holds.
clog (N O (KK/() for some O. Then, with
(TH
(TB.
Σ
D
Σ
log
+ B(
IP'h " m. 8,
NSRASS USN
ASA uSN
(TH
Proof. By the definition of 1y-bracket, for any multi-task parameter 0, find within finite set of r--brackets such that Σ-H iP (TH) 3. (TT) S Then, for any and we have
[1 On Po;
1 Σ
p
(TH) P (TH) 0x on
TH
1 Σ
p (TH PP (TH)
[ (TH) PT (TH)
P (TH)
0.,
Σ 7H
(TH)
(TH) + Σ
P7 (TH) - 1p (TH)
P (TH)
8n
8,
(it)
1PE D
(TH)
log
Pg. (TH)
P (TH)
TH
1g-Pop (-)
PP (TH)
8.
TM
(1i) log
(TH)
S
E
+ VT,
TH
(-)
P
(TH)
where (i) follows from fact that Va S V a + Vbb. (ii) follows from the fact that x S log x for 2 0 (the first term) and the Cauchy-Schwand inequality (the second term), and (Mii follows from the definition of bracket.
Then, for any fixed 0n we have
n (T (n..b
1 (TH) 0_
exp
ΣE ISA ASH
log
Σ
log
ASN
Tp T
N xEI TH
(+)
0:
(TH)
(7H)
InSS Hesh l
(TH)
p p mit
(7H)
E
In IeSk I
(TT)
1.
Hence, by laking union bound over the finite of rr-bacceets and & [K], with probability least 1 - 8, we have for any & E [K]
ΣΣΣN
(TP
T
14

---

Published as a conference paper at ICLR 2024
(TH)
ΣΣΣ
log
NKH
#SN I 4SH
TH TH
( V
Pd
(TH)
(4)
NKH 2
ΣΣΣ
log
(TH
+ +log
(N)((O)
USN ASH
Ph
Dn
(TH
8
(86) 1 + ΣEE 2
P
(THA.))
+ log (N)
log F On
BSN BSNISK ISBBSH
(TH
where (i) follows from Equation 6 (it) follows from Equation union bound, and (dit) follows from that r7 S H²²
7 the Chernoff's method and the
The proof is finished by noting that 0, is an optimistic measure.
B PROPERTIES PSRS
First, for any model {dh, Mo an)), we have the following identity
Mд an (Th-1) Po(oo/TT-))(((
(8)
The following proposition is adapted from Lemma C.3 in Liu (2022) and Proposition in Huang et al. 2023).
Proposition (TV-distance < Estimation error). For any task n E [N], policy n, and any two parameters 0,6 l' O we have
H
DTV (TT P5)
m'(w (M, an) - MA a ((T-)))
h=1 TH
C PROOFS FOR UPSTREAM LEARNING: PROOF OF THEOREM I
In this section, we first prove two lemmas, and then provide the proof for Theorem
First, by the algorithm design and the construction of the confidence set, we have the following estimation guarantee.
Lemma (Estimation Guarantee in Upstream Learning). Let r 1/(N*K2H*), g(d)
e log (N).NHH8))) for some 2 0. With probability at least 1 - 8. for any k and any
θ (0 ... 0n)) Br. we have
ΣΣΣ ASN ISK-IASH
D
PPA 6
The
pog d-
S 2g(N)
Proof. The proof follows directly by combining Proposition and Proposition] 2 and the optimality of the confidence set B
Then, we show that estimation error can be upper bounded by the norm of prediction features.
By Lemma G.3 Liu al. 20222. for any task n, and step h. we can find projection A E 1dda-1xr such that
(6) 10 (Th-1) (T)-)) AA(A)) (Th-)), (1i) : $1
15

---

Published as a conference paper at ICLR 2024
Lemma 2. Let Ab Rd be the projector satisfying Equation ( Fix &o [K]. For any Hh (0 ...,, ( By and any multi-task task policy π (Z1,...,, Tn); we have
( 0z
[I(AA)) (Ta (Ta)))tmn((Tam)))) (UE =]
Drv
E
Th-1
USN
BSN A-
where
U. XI (AR)
Σ
(10 (Th- (T=-1)="" ((AA)))
Ack 7h 7a-
Proof. By Proposition 3, we have
DTv
m mm/k(w)))
(M (o, an) M,, and v (Th-)) T (TH)
For ease of presentation, we fix a task index n. Index Th-1 by i, wh-1 1 by j. Denote (A%) 1. (Tn-1)
by u.f m" (wn)) (Mr. (on, "on an) TT (wn (wh-- 1/7 by (wfja)
Then, we have
m n (wh)
(M an) M (on. M (o))
30th (mn-(() (a(TH)
n.k (wh) (M, "on, an) MM an)) gm (nn- 1
Poy (Tn-1)
TW
Iita (i)
(wihr) T2
(x:7)T Σ sgn (( (Tp)) (a?))
(18)
Bi
| |0f -1
Σ wft sgn (ww))))) x??) xz?
where (i) follows from the property of the projection Ag and definition of the prediction feature T and (ii) is due to the Cauchy's inequality.
Fix an index in. We aim to analyze
E, WIi.. sgn((w)),.) TT%%))lo
We have
Σ 7), sgn((ggnnn sgn((w))to) (T)) Tay)
UT
Σ
= A
Σ sgn( ((w)) T2ry 2 Juijio
sga((w))))) Tx,0 (wig(t))) TT
bck
smpth
0%
Ia
For the first term I,, we have
VT. V
Σ sgn (10f)))) T (o (w Tox
=ER": |
Σ n.e (wh-1 1600
i mk (wa)
(Mr (on, an) -Mh""on (on,an) ARt
=ER":||=|[p==
16

---

Published as a conference paper at ICLR 2024
V
max
Σ
nn (wh- (uo)
| (wn. 1) A 25
>ER":||a||a-
+
max
Σ
Ta (wh-ww |m
mi (wn)) M (on, Af
=
Wh-1
( max
Af + Y
max
Σ
T. Ta op. 10) |M' (on, an A
7 zER E
wER =ER*: 1 08
(1)) 2v ArQA
where (i) follows from Assumption and (ii) follows from the property of AE, stated before Lemma 2
For the second term 12, we have
Σ Tr (6 (=*((wn) (Nr an)) &n ( 2
12
tcck Th-1n-ppr ox
wh-1
Σ ne (wh- - (ō0) wh-1
m (wr) M (on, an) (v)* - Gr (m) (7h-1))
t<ck
Th- 1 Pa
2
Σ Tn (wh- 1600) (in=)((w))) (
(M (on, and (nn (.m.)))-M*M(0 MMn anon,an)v: an) (Tam))))(((
+
-
(a)
(7))) - (Tn-1)
E tek Th-1 Parit
- Fez (o (Th (/11)))))
+ Σ Tn(anlon i0)
Oh
Tn. (mn-())T - 30 (Th
E
tck Th no
27
E
Σ oh.as
Tn (an lon (an(anon i0) | (o
Poo (on|Ta-1 Th (T))
72
i<k
122
where (a) is due to Equation (8) and Assumption
Recall that the L-th coordinate a prediction feature (TT-1) is the conditional probability of core test of - Hence, for the term 121, we have
27
T21
E
P n. oy (o 11 -
n.e Pay (on (0n-1170-11 Poz (on (on-117a 177-
Po; (o 100-21/m-1,,10 a-
p. Th 1
<
<Q²
E
Eawa oing
Σ
Th- Th-1 Pann
P "eni a%
= Q2
E
D2
(wn-1//Th-1)
Po:
(wh-1 (wn-117a-1)
17

---

Published as a conference paper at ICLR 2024
(4)
E
E
Dj
P
(wn-11TA-1),3
P),
(wh (wn-1|m(-2)
-0-2.0h-1~--
QAIADD
TTa).P% (TH)
02
where (a) and (b) follow from Lemma 7 In addition, we can bound T22 as follows.
2
Σ A8h
Tr (anlon, 7h-1) [Fog (onlmn-1) - Po; (onlTn-1) ) 118m (Th)lli
122 5 2
E
Th. 1 1-
Th (anlooo Pp= (on (onl 1)) (1 ( -55m)) -m00)) Σ
+ 2
E
Th Th-1r-pani
Oh an
(1)
2Q2
E
Drv
(Pog (on
Th pupi
"ogin P (walTn), Po: (wnl)T))
+ +22
A
E
E
E
Div
Th- in-
P (1) an min
(0)
22% |A|
E
E
D ( Pos(onlma (ooolmnn1)) Po. (onl (on(Ta-1)
The 11--2:1-11000 P
pun Og
I
+ +2Q
E
E D (poo (walTn), P
"ogin (wn (TA)
Th-1 on Pap (-(mn-)) anmas
(186) S 2Q |A|
[
P 6% (TH), (TH)
+ D
(TH), P
(TH)
where (i) follows from that the coordinate of & takes the value on the conditional probability over core test, (ii) and (iii) follow from Lemma 7
Substituting the upper bounds of /21 and 122 into 12, we obtain that
1 < 72
6Q% |A|
Σ
D
(TH), - (mH)) + D 04
(Ie07" 0%
(TH), (TH) )
Denote D
(TH),
(TH)
by Gi,.. Therefore,
DTV
Iotn Ge
C^ + Σ Cin-1 + Σ
Shn ((A)))t (Ta-1) I
tek
te
Summing over we have
DTv
[p gz
pat
C2 + Es + Σ Bn 5 (Th (7-
eck
tm (Th (Ta--)))))))))) 2
NHCA + ΣΣΣ
ΣΣ
Ey BnLnne
T
h tCk
E BTa-IN [|((A]) fm (To 1)1l(0ew)
18

---

Published as a conference paper at ICLR 2024
where (i) is due to the estimation guarantee Lemma
Theorem 3 (Restatement of Theorem D. For any fixed 8 2 0. let @ be the multi-task parameter space, B(N) = ci (log KHN + log (e)), where e₁ 2 0 and 1j S KHN- Then, under Assumption
with probability at least 5, UMT-PSR finds multi-task model 0 (01....,0r .On) such that
QA
HH A
Σ max DTV 6, Po;
SO
(10)
K
In addition, if K car LIQ2 N
for some C2, UMT-PSR produces a multi-task policy π (#) ...., #N) such that the average sub-optimality gap is at most i.e.
1
max V R.
V BB, Sc.
(11)
N
Proof. Note that 0 B& for all & E [K 1]. Therefore, by Lemma
we have
K
Σ
Drv
p pat
0
(
DTv (pT"ii pan.. 0%
max A NEN @k a.ek...
2
Σ
max DIV 8% E
h:
8QAV A 3(N
ΣΣ NSNASH RSH
E pprit ((()))t (T) (Ta-1)112 (nn-))le
KSK ?
(1)
8QAV K ((
ΣΣΣ
SSKNSNASI ASH
(10) SQAV KBBB
VNNH log(1 +rK/A),
where (i) follows from the fact that 0, B₁ for all & E [K]. (ii) is due to the Cauchy's 's inequality, and (iis) follows from Lemma
Hence,
Sō CAVTIAINHHINN log rK/A)
Σ
pinn Drv
p e. 6
YVK
D DOWNSTREAM LEARNING: PROOF OF THEOREM 2
In this section, we first provide the full algorithm of OMLE. Then, we prove a new estimation guarantee under the presence of approximation error ég. Finally, we provide the proof of Theorem 2
19

---

Published as a conference paper at ICLR 2024
D.1 OPTIMISTIC HODEL-BASEL ALGORITHM
In this section, we provide the full algorithm of OMLE (Liu et al. 2022 Chen et al. 2022) given model class and an estimation margin 60. for the completeness of the paper.
First, OMLE seeks exploration policy n² that maximizes the largest total variation distance between any two model parameters 0 and (' within a confidence set B Then, OMLE uses policies adapted from n² collect data, Finally, using the collected sample trajectories, OMLE constructs confidence set which includes the true model parameter. The pseudo code is provided in Algorithm
Algorithm 2 Downstream multi-task PSR (OMLE)
1: Input: B Ojj, estimation margin Sg.
2: for k
L.... K Down do
3:
arg max max Drv (P)) sEl A.0'
4: for [H] do
5: Use v/ to collect data 7
6: end for
7:
Construct confidence set B
Bx+1 -60= {0 E Σ
θ: tsk
log p B
(TH
max z
Σ
Σ
log pic (F - 30 nBa
tck
8: end for
9:
Output: Any Op E p 1 and greedy policy To arg max
Vx Ro
D.2 ESTIMATION GUARANTEE OF OMLE
(P Recall that co ca. (Qg) min n, z max, D ( 0g P%,) is the approximation error of the model
class ég. this section, let 0% arg ming g max, DR. (177,, ).,))
The following lemma is from Proposition B.1 in Liu (2022).
Lemma 3. Let KH-
With probability at least for any 0g Qg, we have
t.b (T
ΣΣ h
log
S log(N,((6))) + log cK 5
ICR
I g
(TH
L.k
Then, show that the log-likelihood of model 0 is sufficiently large. Lemma 4. With probability at least -- 8, we have
ΣΣ
log
(TH
2.
S CoKH + [co log
t.hy
a
tsk
(TH
Proof. By the definition of 0%, we have
a-1
p (TH) Por gro
log Eng.
c0-
(12)
a
(TH)
If 6o 0. then the proof is trivial. In the following, mainly consider the when ép 0.
20

---

Published as a conference paper at ICLR 2024
By the Markov's inequality, for any R, we have
P
f.le
HTH
log
I
P
tcc
TH
p (Thah)
a-1
P
z (0-11)
p Li (TH
Aic h
Pa Patti 1.h
1.h
a-1
pa kih
(TH.
a-1
ee (a- 1
(TH
n k
par (The.t)
(TH
[<k
(6) 5 c ta- 1) (ee
(a- 1)(z-KHoo))
where (i) follows from Equation
12
By choosing x = 60 K H + a-1 log( K/6) and taking union bound over k. we conclude that, with probability least 1 8,
Σ LCc
log
pr (T
f.l I KH +
1
log
p 0.0
(T
h
Combining Lemma and Lemma following bound holds.
we immediately obtain that with probability at least 1 - 8/2, the
ΣΣ
log po 00
(TH
max Σ Σ log pet 8
(T
[ogg,N, (6 + log
4eK
1 Len log
4K
8 + KH +
where 7) S KH-
Setting 80 log No, (+t) + log deK Eo [soz0)
a-I log 3. we conclude that 0 E Bg for all
KE [K]. Based on this fact, we have the following estimation guarantee.
Lemma 5. With probability at least 8. for any k
[K] and du BR, have
ΣΣ Dēy
[pus 80
p.r 230
Bck h
Proof. We follow the same argument as in Proposition except setting N = 1 and O -g.. Then, we obtain that, with probability at least - 6/2, the following inequality holds.
E B
2. + log
(O,
ΣΣ t<b
D
purt 0
log
tSk
Ao
(T
Since 0, 0% E BR. by the optimality of Bg, we further have
ΣΣ
D
p
8o
21

---

Published as a conference paper at ICLR 2024
p (5tth) %
(N
log
+ Bo + log
P 0
(TH
5
t5k
6 H + (c0
log + 30 + log (O
K
a
S 2.0-
where (i) is due to Lemma
D.3 PROOF OF THEOREM 2
Theorem 4 (Restatement of Theorem D Fix a 1. Let En @6 3o = Ollg KH..
6
log log,N((Ot) + H ((ouz0) where n S Kи- Under Assumption with probability at least
a-1
1 6, the output of Algorithm 2
satisfies that
JIIT max Drv (PT,
Q4
H K
VED
(13)
Proof. First, we follow the proof in Lemma 2 except setting N 1. We obtain that
QalA1
DTv
pst
(TH).
Po (TH)
Cx +
E
Sth
88
i<k
h
0 QAV A D
[I(AP)] (aaittt.*'mn (000 11)
E I
where
Ar LA
C^
3030
X
rQUIAl
(14)
G., = D
(TR r.
),
p
(TH
UhA AI
30.0
(T-)))))
Bog
and (i) is due to the estimation guarantee. Therefore.
KDIV
[.T Po5
I KDTV
T. 0000
KDrv
p Pa 000 0g
Σ max DTv
(Rzz. 1
+ Veo
0 6B EB
< max DTv
(p-* oB
+ K To
BREB
Q Al( + KH - log( (K/())))
HTHK log(1 + rK/A) + KVoo
22

---

Published as a conference paper at ICLR 2024
E BRACKETING NUMBERS OF EXAMPLES AND MISSING PROOFS IN SECTION 4.3 AND SECTION 5.2
In the section, we present bracketing numbers of examples and missing Proofs in Section 4.3 and Section 5.2 Because 0% = E (on, (on.aa) JEOxAA MA ax) for each E [ - 1]. 0% can be decided by @H and {M
1/f h==1- For simplicity, this section, we reparameterize the PSRs parameters h--1hhi }. Without loss of generality, from Theorem C.1 and C.7 in Liu al. (2022).
0 H [MA)]]]
we assume for any (o, a) EO x A, MM a) E IRxT,, and the rank of M2 a) is T.
E.1 BRACKETING NUMBERS OF BASIC SINGLE-TASK PSRS
For completeness, in this subsection, we first present the analysis of bracketing number of basic single-task PSRs.
Lemma 6 (Bracketing number of single-task PSRs). Let θ be the collection of PSR parameters of all rank-r sequential decision making problems with obseration space O action space and horizon H. Then we have
OIIA ))-
log, log,N, (((
Proof. We assume the is known and ||vol12 S VA By Corollary C.8 from Liu et al. (002)) the
PSR model class have following form:
10: 0 ton, (M,) [M, a)ll S for (o, E EO 1 s 11}.
Denote GG the 8-cover of w.r.t fo -oom with 8 [, for some large O. Mathemati-
cally, for any = {6H, (M, e O, there exists θ (0H, [Mr θs. such that for any (o, a) O x A,
pn T2| S 8,
(ec(MM((,a)) Vec(Ma(o, I
Next, show that θ₂ can constitute ry-bracket for O. For any policy æ,, we have
Σ
[P((T)) - (P(TH)) x ((TH)
x))(
H
Σ
Σ
my mi (wn)
(Ma an) MM (n)))))))) un ( x x((Tn)
T hm1
Σ
Σ
Im (wm) 112
MMi(n, an) - Ma an)
( I x T(TH)
T=E(OxA)) h=1
Σ
Σ
(Mh. (.n(o,.a)-M(((oona)) an) VIAl" T((H)
THE
Σ
Σ
(Ma(on) Mn an) V IA x <T(T)) (15)
TH THE(OxA)) h=1
H
Σ Σ
(Vec(MM an)) (oc(Ma(oh,ah)) /1 VIAI" x x(TH)
T)E
< <
Succ assumption does not influence the order of bracketing number since the model complexity related to T does dominate.
23

---

Published as a conference paper at ICLR 2024
where (i) follows from the property of operation norm of matrix, (it) follows from the fact that (wn)llllle I ι la MM (oH oH)ll2.... M (0h+1, 0h+1)ll2 S and ( I
@1)lla l1vlll2 I VIAlM, and (ii) follows from the relation- ship between the operation norm induced by l2 norm and &oo norm.
-, 1112212..... [M
By Lemma 10
(1
1 2
r4 IHOIAA
r45 II (3
8
"
and
[O||A|| ))
log 05 00==1100011
7
which equals to the log uy-backeting number.
E.2 BRACKETING NUMBER OF UPSTREAM EXAMPLES
Stample[[Muutti-tattss POMDP with same transition kernels): The multi-task parameter space is { (TTAa Sh. Oh...., ON Th. E RlSii|s), Of [ol|xIAl, vi [M]] ME(H),EAA
Note that the value of each coordinate of these matrices are probabilities, and thus are bounded within [0,1]. Therefore, the 7ybrackkeingg number this case is (( + ||O[|S] log H|
Example 2 Multi PSR with perturbed models): Suppose there exist a latent base task Pb. and noisy perturbation space A. Each task [MV is noisy perturbation of the latent base task and can be parameterized into two parts: the base task plus a task-specified noise term. Specifically, for each step h E [H] and task E [N], any (o,a) O x A, we have
M (on, as) M as) + A (on, an),
A E.
Such multi-task PSR satisfies that ((N) < Ollog KHN PIOIAA log [O1IA| + HN log A " whereas ((1) for single task is given by ((+2 log |o|A|| H(N - 1) log [)). Clearly,
g(N) < (1) holds if log A < (( which can be easily satisfied for low perturbation environments. In such case, the multi-task PSR benefits from significantly reduced sample complexity compared to single-task learning.
Proof of Example 2 Suppose there exist a latent base task model space:
( = (Øn, (M. Ma a)lla - (or a) E O x A,1 S 11..
θ = (0:0=
A base task model is selected: gb {0h, (M)) ). The parameters of each task in multi-task PSR models are as follows:
O" = 00:0= (0b, (M) AN AR e A, 11 $1}
where (
Bdaxds for any h E [H]. and E [N]. and is the noisy perturbation space with finite cardinality.
A :
Let G be the d-cover of @ w.r.t -nomm with d goाLāear. From Lemma 6 have 6
(""ILAD"" 9
2r+r *HH1)IA
For each n E [N], denote
6 A
0: 0 (7) [M ARH!.)): (pir. (M.)..)) E E A}.
24

---

Published as a conference paper at ICLR 2024
Obviously, [0%] (3y7 +
x A | Then denote the multi-task 8-cover
as O₂ {
G Following from Lemma 6 for any task n 2 1, Oy can constitute an n-bracket for θ", so O can constitute an 77-backet for θ By noticing that each @g has common part, and the changing part is only related A. The corresponding v-bracketing
number is (3)F(IO1)a)EE) r+r2 (OJOIAA
X A JN,, and the log n-bracketing number at most
"
1111111 ( log LOTLA + H(N - 1) log ||)
Example 3 Multi-task PSRs: Linear combination of core tasks): Suppose that the multi-task PSR lies in the linear span of m core tasks, i.c., there exist a set of core tasks indexed by {1,2....,m? m] such that each PSR can be represented as a linear combination of those m core tasks, Specifically, for each task n E [N]] there exists a coefficient vector a = (af, ... , E R s.t. for any h E [H]
and (Oh, an) E O x A.
67 az) = Ei= of (on, (n), (y, oA) = -i= M (on, ax).
For regularization, assume 0 - of I for all [m] and n E [N], and Σ/= of for all [N]. It can be shown that ((N) = 0m( H2 + N) log - ), whereas 3(2) for single task is given by 11 log ||| Clearly, 3(M) < (1( holds if m S min
which is satisfied in practice.
Proof of Example model class
Denote the core tasks model class Q @0.1 x ... x O0.m and the multi-task x... x ON,, where for any E [N]. O is defined as:
H
m Σ
T"
P af PH
of M
: a a="= (af (af..... a R"; (on (M Qua
h=1
For any 2 0, we first consider the 5-cover of the model class of core tasks 6 - x 50.m
where for each base tasks E [m]. 6) is d-cover for Om. Similar to proof of Lemma 6 10 (3vF x =1 and | (3 4)²² [0
Then we consider the cover for multi-task parameter space. Denote C as a 8-cover for the unit ball in R. w.r.t. & norm. Matherattialll, for any n E [N] and vector ax E IR", there exists an
â" E C such that la" - a S 6. In addition, the cardinality |C = )m
Define the task model class θ5 T x GN where for any n [N], 63 is defined
H
m Σ
6
aj
Q HH
ai
M
ā" (ā
2m T E C.; (on: (M. E 00.
h=1
We next show that G₂ is bracket of O, By definition, for any model of base task [m]: 0 (dH (M) E Ol, there exists
g (z)z: (M/ 1 (M) E 6% such that for any (o, 0 A,
oo Jy sb, ve.
Vee (o, (o.)))
5.
Then, for any 0¹ E O α² E Rm there exist Ō E Q and " C's such that
of
af
a
H
1=1
m
0
111
+
Σ
0
8
11
1=1
1=1
25

---

Published as a conference paper at ICLR 2024
7
2
0f
- 26
1=1
1=1
and for any (0%, aA) E O x A
m
m
Σ
0
M
01 M
m
[of'l
M/h
Mi
+
of
af'
M
0i Vr
Vec(MM) Vec(M))))
+
a7
M
1=1
1=1
vr + Σ Vrlaß - a7'l
H=1
-=1
Vrd + /r ||a a" s 2/r5.
Similar the analysis in the proof of Lemma specifically, Equation 15 O can constitute an n-bracket for O' with 5 2
The cardinality of the cover of multi-task model class is [osl {)=*#|)]]44] (3)m..
+ |
(3vr x
conclusion, the log brbrcketeing number is O (12111001A{m log( PHIOIIA| + m.V log( PHOIAA
E.3 BRACKETING NUMBER OF DOWNSTREAM EXAMPLES
Example Note that the model class for downstream learning is θ = (hetn) Thus, we immediately obtaint that No, (θx) O(H OHO)) log [0|(S)
PSR with perturbed models): Similar to the upstream tasks, the downstream task is also noisy perturbation of the latent base task. Specifically, for each step h E [H] any (o, a) E O x A, we have
b = oh M (o,, a.) = (,, a.) + (o an). A
(16)
The log nrbbackeeinng number is at most H log A
Proof. Suppose the estimated model parameter of the base task is 0 = (o) (M") the downstream task model parameters satisfy Equation (16), the empirical candidate model class can be characterized as
h }. Because
: (0: 0 (#) (M
Th: MA, (oh, an)) M, (on, an) Di an). (on, an) O A: A)
If Tb is given, then the candidate model class is decided by A. Then for any 1j > the brackettnn number is A
Example Example3[Multi- Multii-ask PSRs: Linear combination of core tasks): Suppose the downstream task also lies in the linear span of core tasks same the upstream. Moreover, assume the upstream tasks are diverse enough to span the whole core tasks space. As result, the downstream task can be represented
lincar combination of subset of the upstream tasks. Specifically, there exists constant I. satisfying m < L < N and a coefficient vector a (af, ... ,,,)) E TR s.t. for any h [H] and (Oh, on) O x A.
6 Σ of Oi
( Ets afM) (on,
(17)
26

---

Published as a conference paper at ICLR 2024
For regularization, we assume I af 1 for all I E [L], and Σl= af 1. It can be shown that 50- = O(LH log( log( ))) whereas g(1) for learning without prior information is given by
0 log ) Clearly, 8o << 3(1) holds if m I m ( which is satisfied in practice.
Proof. Suppose for the upstream task / E [4]. the estimated model parameter is 0: f ( Because the downstream task model parameters satisty Equation 17 the empiri- cal candidate model class can be characterized
θ = : (0r (M of = E a a
M an)
0j
i M (on, (on,
an), (on, E A: E
Then we consider the cover for Gg. Denote C as d-cover for the unit ball in R2 w.r.t. & norm.
Mathematically, for any vector a R², there exists an â E C such that ||a" - a²1|, s 8. In addition, the cardinality |C₁¹ (3)
Define
O 10:0- (0H:
- (M, 1/ 04
0f
(on, an) = Σ ai M (on, an), an) = }
Then, for any 0² Ô with a E RL, there exist E G with â C such that
L
07
af
af 11 s $I1a² -6 S
1=1
1=1
and for any (O%, an) 0 x A
1
1.
Σ
0i
M
(on, an)
Σ
ai
(o,, an)
1=1
Σ
0
afl
M (z) /1 svilla"- s vro.
Similar to the analysis in the proof of Lemma
specifically, Equation 15 Oy constitute
b-bbacket for Gg with .-<
The cardinality of the cover multi-task model class is jōyl = |C| = ( In conclusion, the log
br-backkeinng number is (LH log( ((OlIA)) )
F EXAMPLES OF MULTI-TASK MDPS FROM PREVIOUS WORK
To demonstrate that our framework encompasses multi-task learning under MDPs, we provide several examples of MDPs from previous work in this subsection. Suppose S is the state space, there exist N source tasks, and p(+.a) S x A S > IR is the true transition kernel of task n.
Cheng et al. 2022) studied multi-task learning under low-rank MDPs in which the transition kernel p(..) has d dimension low-rank decomposition into embedding functions g(() S A
27

---

Published as a conference paper at ICLR 2024
Fu C
model class
of task 2 p
|o(n1)), p(+.a)
p1
Pminl (1, (t)),)))
no
model class of task P
Figure 2: Supplementary illustration of joint model class in two dimensions for MDPs. A concrete example of tasks with similar transition kernels, i.e., any point (Po, Pl) the joint class Fu satisfies (ax(s.)) Drv (P"(sss, a))s S e for small positive constant
Rd (((,))))) S + Id as p(x,") a, (6(*)) (s, (s,a), 1 (ll+ (s')) for all (s, (s,a,5' S S x for each task n. In this setting, Cheng et al. (2022) assume the N source tasks share common representations d(()) and for each task n, (5 E Φ (((x)) E V for finite model class (Φ, V) Consequently, the n-bracketing number for the multi-task low-rank MDPs model class is at most if (
2 w which is much smaller than the one of the individual single-task with O(H) | w M
Zhang & Wang (2021) studied multi-task learning under tabular MDPs with an assumption that for any two tasks 1, n2 [NV. it holds that (a Drv (p(( a) (() a)) s for some small 6 > 0. Then the multi-task model class much smaller than the model class of N individual single-task (see Figure 2 for an illustration when N 2). Conse- quently, the n--backeting number for the multi-task low-rank MDPs model class is at most ((H]o S + (N 1) log M(A|Sk))
)). This is smaller than that of each individual single-task, which is O(HN S|3 log H|Al|S if and N 2 1. "
"
G AUXILLARY LEMMAS
The following lemma characterizes the relationship between the total variation distance and the Hellinger-squaned distance. Note that the result for probability measures has been proved Lemma H.I in Zhong al. (2022). Since we consider more general bounded measures, we provide the full proof for completeness.
Lemma 7. Given two bounded measures and Q defined on the set X, let PP| = P(x) and [Q] = Exex (x). We have
D (P.Q) S 4(IP)+ (P.Q))))))
In addition, if PrIx. Qrix are two conditional distributions over random variable Y. and Px, Pr Qx Qyix Q are the joint distributions when follows the distributions and Q, respectively, we have
E X~P
x()Xxxx(11),,. Xxxxx)) & SDx Qx.x).
Lemma 8. Suppose IP and Q are two probability distributions. For any 1, have the following inequality.
(Iv Q) D
Da.a (IP, Q)
28

---

Published as a conference paper at ICLR 2024
Proof. By Pinsker's inequality, have
DPv Q S 2 DiL Q)
By Theorem 5 from Van Erven & Harremos 2014), we have
DRc Q) lim Da.o (P,Q)= S inf Da.o (I, Q).
atl
021
Lemma 9 (Elliptical potential lemma)) ). For any sequence of vectors X {a1...., .n.... C R let Up Al + Eick JaTk T where x is positive constant, and B > 0 is a real number. If the rank of X is most r, then, have
min [1osslle BJ 1 (1+ B)rlog(1 + K/A),
K
Σ
min ||welllz VB S v(1 B)rk log(1 + K/A).
Proof. Note that the second inequality is an immediate result from the first inequality by the Cauchy's inequality. Hence, it suffices to prove the first inequality. To this end, we have
min 2 1, B}
B)
log [ ||
K
(1 + B)
log (1 + trace ((U++1 - 11-1-UU)U0 UF
k=1
log (1 trace (Uz (Uk+1 U2) UZ 112)
(1 + B)
k=1
S1 B)
log det [Ia + Up U-!/2 (U++1 - Uz)
k=1
K
det (Uk+1)
(1 B)
Σ
log
det(Uk)
k=1
det(UK+1)
(1 + B) log
det(Ui)
K
(1 log det 1+
Σ
kml
(1)) < (1 + Byr log(1 (1 + K/A),
where (i) follows because x < (1 + B) log(1 x) if O < x B, and (ii) follows because rank( x)) S..
To compute the bracketing number of mulit-task model class, we first require a basic result on the covering number of Euclidean ball as follows. Proof of the lemma can be found in Lemma 5.2 in Vershynin (010)).
Lemma 10 (Covering Number of Euclidean Ball), For any e > O, the e-covering number of the Euclidean ball in dd with radius Rt > 0 is upper bounded by (1 + 2R/e)d.
29


Reviews:
[{"review_id": "Mgo45oxBbo", "reviewer": "Reviewer_iJpQ", "summary": "This paper investigates the (upper-bound) benefits of multitask learning in PSR environments vs per-task single-task learning. The setting is episodic RL and the proposed algorithm constructs and refines a confidence set of candidate environment parameters, and in each iteration uses those parameters to compute the data collection policy.\n\nBesides introducing the RL algorithm for multitask learning, the paper gives performance bounds for multitask learning and transfer to downstream tasks. It also explores the bounds in specific PSR settings, compared to bounds from using separate RL methods for each task. The comparisons highlight the advantages of the multitask approach.\n\nThe key algorithmic idea is to maintain a joint confidence set for the potential per-task environments, and the key technique used in the bound is considering a covering number of this confidence set. The log of this covering number grows slower or sometimes much slower with the number of tasks, than the sum of log-covering numbers of separate, per-task confidence sets.", "questions": "There is another family of PSRs that I would like to suggest as an example (from personal experience) and they can be relevant for the current and future work:\n* A set of tasks that is not observed uniformly. We still want uniformly good behavior on all the tasks, but the algorithm can only sample the tasks from a distribution, rather than go over each of them one by one in each iteration. This is highly relevant to how we train RL agents in POMDPs with procedural generation of the initial state, because procedural generation gives very coarse control over the resulting initial state distribution. Here, I am seeing each initial state in the support of the distribution as a task. The downstream learning is also interesting for this.\n* A single PSR with block structure in the dynamics. This is like the example above, but the multiple tasks are not explicitly recognized as such.\n\nSome things I would like to see in the downstream regime:\n* What impact does more and more training on the upstream tasks have on the zero-shot performance on the downstream task?\n* What impact does more and more training on the upstream tasks have on the speed of learning on the downstream task? It would be a surprising and interesting find if some amount of \"pre-training\" upstream would actually improve the rate of convergence of the downstream. I guess it's more likely that the guarantee would be like \"if you want to train for a given budget X downstream, then you can get good rates if you train for Y amount of experience upstream.\"\n\nAt a higher level, not as a criticism to the paper, though, I find the overall setting a bit odd. The proposed algorithm does not have any sequential interaction with the environment. Instead it runs the policy in the tasks for collecting data and updates its confidence set. What I find odd is therefore that the tasks can be so hard intra-episode that there is nothing we can do by adapting as we act, and we might as well pick policies, deploy them, and update. I guess the setting also does not quite apply to the kinds of environments that would be \"easy\" and where intra-episode adaptation could help improve performance.\n\nI liked the fact that the results allow us to recover the batch setting when we train on N copies of the same task.\n\nI am somewhat confused about what $\\pi(\\omega_h)$ means in Eq. 2, considering that $\\pi$ for the end of the episode $\\omega$ depends on what happened in the beginning of the episode ($\\tau$). So how does $\\tau$ factor into Eq. 2?\n\nIt would also be nice to understand the relationship between r from Definition 1 and multiple tasks. Considering the correspondence between block-dynamics and multiple tasks I mentioned above (that is, multiple tasks can be put together into a single POMDP that samples the task as part of the initial state), what is the relationship between r and N? r is arguably harder to scrutinize than N and the shared structure between tasks, so maybe it's possible to get rid of r as a proxy for multiple tasks and formalize everything in multitask terms?\n\nTypos:\n* I was a bit confused reading algorithm one because it refers to quantities that are defined after it is shown ($\\nu^{\\pi_n,k}$).\n* In Example 3, the P for the core tasks seem to be on the wrong font type?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well written and clear. Overall I am happy with this paper. It is clearly written and easy to follow. My reading of the contributions is that they offer a better understanding of multitask PSR problems more than offering a solution to multitask learning in PSRs. This is mostly because the PSRs are assumed to be given and the algorithm uses components that are useful for a theoretical study but tricky to set up in practice (Though it would be great to hear from the authors with details if they disagree.) To me the main takeaway is that, through studying upper-bounds, we can say that multitask PSRs are easier to learn (jointly), with interesting examples of shared structures across tasks that accelerate learning.", "weaknesses": "I do not see any issues with the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jRnO4E8d6Z", "reviewer": "Reviewer_t71h", "summary": "This paper studies the problem of multi-task reinforcement learning under non-Markovian decision making process. By assuming the multi-tasks share the same action and observation spaces, and the models are from a certain parameter class, the sample complexity of learning an averaged optimal multi-task policy (using UMT-PSR, an proposed algorithm) is given in Theorem 1. The complexity is related to the complexity of the parameter class, which is measured by $\\\\eta$-bracketing numbers. This result shows the benefit of multi-task learning when compared with learning tasks separately. For the downstream class learning, by adopting OMLE, the sample complexity is given in Theorem 2, which is related to the complexity of the downstream model class (constructed by upstream learning) that can be reduced by previous upstream learning. The authors also instantiate their generic framework on three examples.", "questions": "1. can these results generalize beyond low-rank problems? (maybe not low-rank, but some other structures)\n2. The results hold for finite action and observation spaces. Can they be generalized to general infinite spaces? (maybe using function approximation or other techniques)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: Studies the combination of non-Markovian process and multi-task RL, which is a relatively unexplored topic. \nQuality: provide generic frameworks together with concrete examples, which shows the applicability of this theoretical analysis.\nClarity: the writing is smooth, the ideas and intuitions are also clear.", "weaknesses": "The downstream learning seems to be only applying previous results on a smaller downstream model class, without further new ideas.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "heZBbggFF6", "reviewer": "Reviewer_W5uL", "summary": "This theoretical paper discusses the problem of multi-task reinforcement learning in the setting of low-rank and well-conditioned PSRs. It analyzes two situations when there is structure in the multiple tasks: upstream learning and downstream.", "questions": "1. Why is the first goal of upstream learning finding near-optimal policies for all $N$ tasks _on average_, as opposed to, say, finding near-optimal policies for all $N$ tasks, i.e. \n$$\n\\begin{align*}\nmax_{n \\in [N]} \\max_\\pi \\left(V_{\\theta_n^*,\\,R_n}^{\\pi} - V_{\\theta_n^*,\\,R_n}^{\\bar{\\pi}^n}\\right) \\le \\epsilon.\n\\end{align*}\n$$\n2. Why use $\\lVert \\cdot \\rVert_\\infty^p$ as the norm? Is it even a norm (i.e., satisfies the conditions required)?\n3. In the calculation of $\\eta$-bracketing number of $\\{(\\mathbb{P} _ {\\theta_1}, \\ldots, \\mathbb{P} _ {\\theta_N}) : \\mathbf{\\theta} \\in \\mathbf{\\Theta} _ u\\}$, what is the domain of the functions? Is it the $\\sigma-$algebra over $(\\mathcal{O} \\times \\mathcal{A})^H,$ which is the domain of the distributions? Consider a simpler calculation: how to calculate the $\\eta$-bracketing number for $\\{\\mathbb{P} _ {\\theta} : \\theta \\in \\Theta\\}$? Now $\\mathbb{P} _ {\\theta}$ is a probability measure which is defined over some $\\sigma$-algebra $\\mathscr{S}$. If it is contained in some $\\eta$-bracket $[\\mathbb{A}, \\mathbb{B}]$, then we must have $\\mathbb{A}(S) \\leq \\mathbb{P} _ \\theta(S) \\le \\mathbb{B}(S)$ for every $S \\in \\mathscr{S}$. But this would imply (it might require measures to be regular, I am not sure) that $\\mathbb{A} = \\mathbb{P} _ \\theta = \\mathbb{B}$. So the $\\eta$-bracketing number for $\\{\\mathbb{P} _ {\\theta} : \\theta \\in \\Theta\\}$ becomes $|\\Theta|$. I am assuming that this is not what the authors had in mind. Could you please clarify the calculation? The calculations in Appendix E are assuming the observation and action spaces are finite.\n4. I do not understand the discussion of pairwise additive distance based multi-task planning. Why is a distance between product distributions not sufficient as opposed to what the paper uses? Also, do the authors realize that\n   $$\n  \\sum _ {n \\in [N]}\\mathtt{D _ {TV}}(\\mathbb{P} _ {\\theta_n}, \\mathbb{P} _ {\\theta'_n}) = \\mathtt{D _ {TV}}(\\mathbb{P} _ {\\theta_1} \\otimes \\cdots \\otimes \\mathbb{P} _ {\\theta_N}, \\mathbb{P} _ {\\theta'_1} \\otimes \\cdots \\otimes \\mathbb{P} _ {\\theta'_N}),\n  $$\n   a divergence (not distance) over product distribution?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Strengths of this paper are that it tackles a challenging problem, introduces a nice formulation for the problem and showcases interesting theoretical results on when multi-task learning is beneficial as opposed to multiple single-task learning.", "weaknesses": "The biggest weakness of this paper is that it crams too much content in the main paper, and does not use the appendix to explicate it. This rushed discussion makes the job of reading the paper more difficult than it needs to be. For example, consider the assumptions of rank-$r$ and $\\gamma$-well-conditioned PSRs. These two assumptions are present throughout the paper, but they never get their deserved attention. A mere half a page of terse definitions for them might be justified in the main paper due to page number constraints, but it is hard to justify why they were never given due discussion in the appendix. This discussion should discuss the intuitive meaning of these assumptions, examples of PSRs which satisfy the assumptions (otherwise questions like, is the set $\\Theta$ even non-empty, surface), examples of PSRs which don't satisfy the assumptions, what fails in the proofs if each assumption is relaxed, etc. As for other examples, the notation $\\phi_h$ and $\\mathbf{M} _ h$ is never even defined, the norm $\\lVert \\cdot \\rVert_\\infty^p$  is \"pulled out of a hat\", etc.\n\nThe paper does a poor job at literature review. For example, it states that \"none of the existing studies considered multi-task POMDPs/PSRs\", but consider _Multi-task Reinforcement Learning in Partially Observable Stochastic Environment_ by Li, Liao and Carin (JMLR 2009) or _Deep Decentralized Multi-task Multi-Agent Reinforcement Learning under Partial Observability_ by Omidshafiei, Pazis, Amato, How and Vian (ICML 2017).\n\nWhile not necessarily a weakness, it would have been nice to demonstrate the usefulness of the theory developed in the paper on some simple experiments. As long as I am asking for things it would be nice to calculate the computational complexity of implementing the algorithms. Note that I am not at all expecting these additions to this paper -- the paper is already very terse as it is.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8kN8lZlPsX", "reviewer": "Reviewer_qpJY", "summary": "This paper studies the benefits of multi-task learning in low-rank PSRs when tasks share similar latent structures. It proposes a measurement for the similarity of $N$ tasks called the $\\eta$-bracketing number, which is shown to be small in several standard classes. Examples include multi-task POMDPs sharing the state space, action space, and transition kernel, and multi-task low-rank PSRs with similar core set observable matrices. The algorithm proposed in the paper leverages the $\\eta$-bracketing number to find the optimistic exploration policies, and follows the idea of MLE to build confidence set for models. It is proved to gain benefits when the average $\\eta$-bracketing number over $N$ tasks is smaller than the $\\eta$-bracketing number for a single task in terms of sample complexity to find the optimal policies of all tasks. Additionally, the paper also gives a downstream algorithm that improves the sample complexity to identify a new target task based on the similarity between target task and the $N$ original tasks.", "questions": "It seems that the $\\eta$-bracketing number is used to build the optimism of optimistic policies with uniform convergence over the bracketing set. This is essentially (and also used as) the covering number of the model class in terms of $\\|\\cdot\\|^{\\mathrm{p}}_{\\infty}$ norm. Why use bracketing number as the name instead of calling it the covering number directly?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The setting is an extension of multi-task learning of MDPs and bandits, which is novel in the literature of multi-task decision making. \n\n2. It provides a key measurement of similarity in the multi-task setting, the $\\eta$-bracketing number, to identify the effectiveness of the multi-task transfer learning in low-rank PSRs.", "weaknesses": "1. The results of the paper is interesting in the viewpoint of setting and techniques, but not surprising given many previous works on multi-task reinforcement learning. The standard tool to establish the benefits of shared structure of multiple tasks is the reduced covering number of the joint model class or value class. For example, the common low-rank assumption in the linear setting essentially reduce the log covering number of the function class from $nm$ to $nk + km$, where $n, m, k$ denotes the ambient dimension, the number of tasks, and the rank that is small. This work studies a more complicated setting, but the $\\eta$-bracketing number is essentially some type of covering number over the joint model class (see the questions below). As long as this key property is identified, the remaining task is to follow OMLE to perform optimistic planning in the joint model space. \n\n2. The generality of the theorems in the paper allows the various common structure with different $\\eta$-bracketing number of the joint model class. Several examples are already explained in the paper such as the multi-task observable POMDP sharing the common transition kernel. The generic upstream algorithm is highly computational inefficient in building the confidence set and find the optimistic policy. Therefore, an important question is how the optimization steps in the algorithm look like in a specific setting (e.g., the multi-task POMDPs). This helps to evaluate the effectiveness of the upstream algorithms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1LcHLvrJYK", "reviewer": "Reviewer_FV3A", "summary": "This paper studied multi-task reinforcement learning (MTRL) in complex environments like partially observable MDPs and predictive state representations. The authors identified two main challenges: 1. Identifying Beneficial Common Latent Structures: The large and complex model space of multi-task predictive state representations (PSRs) made it difficult to identify types of common latent structures that could reduce model complexity. 2. Intertwining of Model Learning and Data Collection: in RL, model learning and data collection are intertwined, creating temporal dependencies in the collected data. This complicates the analysis of multi-task PSRs, making it challenging to gauge the benefits of reduced model complexity in terms of statistical efficiency gains in RL. To solve these challenges, the authors introduced the η-bracketing number to quantify model complexity and task similarity. They developed the UMT-PSR algorithm for efficient upstream multi-task learning and addressed downstream transfer learning by leveraging similarities with previously learned tasks. Their contributions include a new complexity metric, the innovative UMT-PSR algorithm, and techniques for enhanced downstream learning, marking a pioneering theoretical exploration of multi-task RL's benefits over single-task approaches in complex environments.", "questions": "To AC: I do not have the expertise in this research area to review the strengths and weaknesses of the paper. Please lower the weight of my review.", "limitations": "", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "To AC: I do not have the expertise in this research area to review the strengths and weaknesses of the paper. Please lower the weight of my review.", "weaknesses": "To AC: I do not have the expertise in this research area to review the strengths and weaknesses of the paper. Please lower the weight of my review.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)