Provable Benefit of Multitask Representation Learning in Reinforcement Learning
Authors: Yuan Cheng
Songtao Feng
Jing Yang
Hong Zhang
Yingbin Liang
Venue: NeurIPS 2022 Accept (2022)
Keywords: Multitask learning
representation Learning
offline RL
suboptimality gap.

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

Provable Benefit of Multitask Representation Learning in Reinforcement Learning
Yuan Cheng
Songtao Feng The Ohio State University feng. 13w99335sssss. edu
University of Science and Technology of China cy160mail.. uste. edu.
Jing Yang
Hong Zhang University of Science and Technology of China zhangh@ustc. edu. cn
The Pennsylvania State University Lungging@psu..m edu
Yingbin Liang
The Ohio State University
1 Lingg- 8898osu. edu
Abstract
As representation learning becomes a powerful technique to reduce sample com- plexity in reinforcement learning (RL) in practice, theoretical understanding of its advantage is still limited. In this paper, we theoretically characterize the benefit of representation learning under the low-rank Markov decision process (MDP) model. We first study multitask low-rank RL (as upstream training), where all tasks share a common representation, and propose new multitask reward-free algorithm called REFUEL. REFUEL learns both the transition kernel and the near-optimal policy for each task, and outputs well-learned representation for downstream tasks. Our result demonstrates that multitask representation learning is provably more sample-efficient than learning each task individually, as long as the total number of tasks is above certain threshold. We then study the downstream RL both online and offline settings, where the agent assigned with new task sharing the same representation the upstream tasks. For both online and offline settings, we develop sample-efficient algorithm, and show that it finds a near- optimal policy with the suboptimality gap bounded by the sum of the estimation error of the learned representation in upstream and a vanishing term as the number of downstream samples becomes large. Our downstream results of online and offline RL further capture the benefit of employing the learned representation from upstream opposed to learning the representation of the low-rank model directly. To the best of our knowledge, this is the first theoretical study that characterizes the benefit of representation learning in exploration-base reward-free multitask RL for both upstream and downstream tasks.
1 Introduction
Multitask representation learning has arisen as a popular and powerful framework to deal with data-scarce environment in supervised learning, bandits and reinforcement learning (RL) recently (Maurer et al. 2016 Du et al. 20200 Tripuraneni et al. 202I: Yang et al. 202T Cella et al.) 2022 Lu et al.] 2021 Arora et al. 20200 Hu al. [202)) The general multitask representation learning framework consists of both upstream multiple source tasks learning and downstream target task
learning. In RL applications, in upstream, the agent is assigned with T source tasks sharing
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

common representation and the goal is to learn near-optimal policy for each task via learning the representation. downstream, with the help of the learned representation, the agent aims to find near-optimal policy of new task that shares the same representation as the source tasks.
While representation learning has achieved great success in supervised learning (Du et al. 2020)
Tripuraneni et al.] [202T, Maurer et al. [0016] Kong et al., (2020) and multi-armed bandits (MAB)
problems (Yang et al. [0.]2221]] Qin et al. 2022 Cella et al. (0022), most works in multitask RL mainly
focus on empirical algorithms (Sodhani et al. [202 [ Arall
al., 2022 Teh et al. 2017) with
limited theoretical works (Arora et al. [2020] Hu et al. [022]] Brunskill and Li 2013]] Müller and
Pacchiano 2022 Calandriello et al. 2015 Lu et al. [021 Eramo et al. (0200). Generally speaking,
D'Errmo
there are two main challenges for multitask representation learning in RL. First, differently from supervised learning with static datasets, in MAB and RL, representation learning and data collection are intertwined with each other. The agent has to choose an exploration policy in each iteration based on the representation learned in the past. Such iterative process introduces temporal dependence to the collected data. Note that none of the aforementioned theoretical multitask RI. studies took the influence of sequential exploration and temporal dependence in data into consideration. Second, data distribution mismatch between the target task and the source tasks may occur. Additional assumption is usually required in order to ensure the representation learned from sources task will benefit the target task (Lu et al. 202T).
In this work, we aim to unveil the fundamental impact of multitask representation learning on RL by answering the following two open questions:
First, in the upstream, is it possible to design exploration-basee multitask representation learning algorithms for RL to achieve efficiency gain compared with single-task RL after taking the interactive data collection process and temporal data dependency into account?
Second, can the learned representation from upstream bring efficiency gain the downstream learning?
We focus on multitask representation learning under low-rank MDPs (Agarwal et al. 2020), where the transition kernel admits low-rank decomposition that naturally involves representation learning. answers the aforementioned questions affirmatively, as summarized below.
Main contributions:
First, we propose new multitask representation learning algorithm called REFUEL under low- rank MDPs in upstream. REFUEL features joint modeling learning, pseudo cumulative value function construction for exploration, joint exploration policy learning and unique termination criterion. These novel designs enable REFUEL to leverage the shared representation across source tasks to improve the learning efficiency.
We show that REFUEL can 1) find a near-optimal policy and 2) learn near-accurate model for each of T source tasks to certain average accuracy Cu. with a sample complexity of = +
H3 TEL
H²K3 where H, K, d denote the horizon steps, dT
the cardinality of action space and the low-rank dimension, respectively. Compared with state- of-the-art single-task representation in low-rank MDPs (Uehara et al.) 2022). REFUEL achieves reduced sample complexity as long as T' = Ω (K) ). which shows the benefit of multitask learning provably. To the best of our knowledge, this is the first theoretical result demonstrating the benefit of representation learning in exploration-based reward-free multitask RL.
Finally, we apply the learned representation from REFUEL downstream online and of- fline learning problems. In offline RL, we propose new sample efficient algorithm with the suboptimality gap bounded by Ö(HVaaaaw + H max(d, vP) /Noot), where Nofr is the number of samples of offline dataset and p measures the model complexity of reward. In online RL, we propose an online algorithm with the suboptimality gap bounded by δ HelEawwn + H max vp)). where Non is the number of iterations in the on- line algorithm. In both the offline and online settings, the first term captures the approximation error where Êdown measures the quality of the learned representation from upstream, and is negligible when the upstream learning is sufficiently accurate, and the second term improves the dependency on H and d compared with the state of the art (Uehara et al.) 2022) for low-rank MDPs without the knowledge of the representation, showing the benefit of employing the learned representation from upstream.

---

Throughout the paper, we use [N] to denote set (I....,NT . for N E N. use || l2 to denote the lg norm of vector x, use | to denote the Frobenius norm of matrix X, use A (AA to denote the probability simplex over set A, and use ((A) to denote uniform sampling over A. given
Notations:
|A| < co. Furthermore, for any symmetric positive definite matrix Σ, we let ||>||> == Vx Er.
2 Related Work
Multitask bandits, multitask RL and meta-RL: The benefit of multitask learning in linear bandits
has been investigated in Yang al. (2021 D: Qin al. (2022): ) Cella et al. (2022) : Azizi et al. (2022):
] Deshmukh et al. (2017); Cella and Pontil (2021); Hu et al. (2021). p. For multitask RL, Arora et al.
(2020) showed that representation learning can reduce sample complexity for imitation learning.
Calandricllo et al. (013)) studied multitask learning under linear MDPs with known representation.
Hu al. (2021) studied multi-task RL with low inherent Bellman error (Zanette et al. 2020) with known representation. Assuming all tasks are sampled from an unknown distribution, Brunskill and Li 2013) studied the benefit of multitask RL, when each task is sampled independently from a distribution over finite set of MDPs while while|Müller Müller and Pacchiano| (2022) focused on meta-RL for linear mixture model (Yang and Wang 2020). Considering all tasks share a common representation, D' D'Eramo et al. (2020) showed the benefit in the convergence rate on value iteration and Lu et al. (2021) proved the sample efficiency gain of multitask RL under low-rank MDPs. None of the above works took the influence of sequential exploration and temporal dependence in data into consideration, where Lu et al. (2021) assumed a generative model is accessible and D Eramo et al. (2020) assumed that the dataset given in prior. In our study, we design reward-free. exploration for multitask RL.
Single task RL under low-rank MDPs: Low-rank MDPs are first studied by byAgarwal et al. (2020)
Uehara et al. (0022) studied low-rank MDPs for both online and offline RL. Modi et al. (2021) )
proposed a model-free algorithm MOFFLE under low- www-nonegative-ee -rank MDPs. Du et al. (2019);
Misra et al. 2020); Zhang et al. (2022) studied block MDP which is special case of low-rank MDPs.
Offine RL under linear and low-rank MDPs: Most of existing works on offline RL focus on single-
task learning. among which Jin et al. (202T) studied offline RL under linear MDPs and proposed an
algorithm with provable suboptimality gap, Uehara et al. (2022) studied offline RL under low-rank MDPs and proposed an algorithm REP-LCB that leverages pessimism, and Yin et al. (2022): Xiong
et al.) 2022) studied offline RL. leveraging variance information.
3 Problem formulation
3.1 Episodic MDP
Consider an episodic non non-stationary Markov Decision Process (MDP) with finite horizon H. denoted by M= (S, AA, H, P. P.r). where S is the state space (possibly infinite), A is the finite action space with cardinality K, H is the number of steps in each episode, P = {Pr ]e(HH] is the collection of {rikne Hi is the collection of deterministic reward functions with Tn : S x + [0, 1]. The initial state s is assumed to be fixed for each episode and the sum of the reward is normalized with M Th E [0, 1].
transition measures with PA : S x A + ASS),
[TR as sequence of mappings where Th S + A(A). We use (als)) to denote the probability of selecting action in state step h. Given a starting state 8h, we use Sn ~ (P, x) to denote a state sampled by executing policy ा under the transition model P for h - h steps and (xx,,x)) [:] to denote the expectation over states 8 (P. x)) and actions a; ~ n. Given policy ι and a state s at step h, the (state) value function is defined
Denote policy
as (s) = T (e E##=, (hr(sse, 84 = s) Similarly, the (state-)action value function for a given state-action pair (s) under policy rx(s, a) + (P% V S: R. For simplicity, we omit subscript h for h = 1. Since the MDP begins with the same initial state 81, to simplify the notation, we use V to denote (TT,,,(si) without any ambiguity. Let P* denote the transition kernel of the underlying environment. Given reward function r, there always exists an optimal policy x² that yields the optimal value function Vp.,, = sup, Vp. . abbreviated as Vp....
step h is defined Ph. ==
a), where we define (s, a) == E. (-)a.a)) [f(')] for any function
This paper focuses on low-rank MDPs Jiang et al., 2017 Agarwal et al. 2020 defined as follows.
3

---

Definition 1 (Low-rank MDPs), A transition kernel P. 8 S A - 4(A) admits a low-rank
decomposition with dimension d E N if there exists two embedding functions 0f :S x A + R and M : S + Rd such that
Px |s, a) = (6% (k())) Vs, E
Without loss of generality, we assume ||ø (s, a)ll2 S 1 for all (s, a) E S x A and for any function + [o: 1). [Iʃ P7 (s)g(())) S vd. An MDP is low-rank MDP with dimension d if for [M] its transition kernel P% admits a low rank decomposition with dimension d. Let 6- = {0% hiv M) and "² (1% he) be the embeddings for P', where o" is also called representation in RL literature.
3.2 Reward-fre multitask RL and downstream learning
We consider a multitask RL consisting of reward-free upstream and reward-known downstream.
In reward-free multitask upstream learning, the agent is expected to conduct reward-free ex- ploration over T source tasks, where all reward functions fr' hEerT are unknown. Each task t E [T] is associated with a low-rank MDP M = (S, H. Pt, r"). All T tasks are identical except for (a) their transition model Pt, which admits low-rank decomposition with dimension d: P'(s'|s, a) (8'),, (((,) (s, a)) for all & E [7] and (b) their reward rf. Namely, all tasks share the same feature function o", but may differ in (((+)) as well as reward r'. The goal of upstream learning to find a near-optimal policy and near-accurat, model for any task / [7]] and any reward function 1 E [T'] via sufficient reward-free exploration, and output well-learned representation for downstream task. For each / (T), we use p(*.)) to denote the true transition kernel of task t.
In downstream learning, the agent is assigned with new target task T associated with low-rank MDP M = (S, A, H, pT+1 , The task shares the same S, A. H and ø" with the T upstream tasks, but has task-specific u! (+++)) We also assume that the reward for the new task TT+ is given to the agent. Here, the agent is assumed to take the learned representation function ø during the upstream training, and then interacts with the new task environment MT+1 in an offline or online manner to find near-optimal policy for the new task. By utilizing the representation learned in the upstream, the agent is expected to achieve better sample efficiency in the downstream. Finally, we use ((+I++) to denote the true transition kernel of task T
4 Upstream reward-free multitask RL
In this section, present our proposed algorithm for upstream reward-fire -ree multitask RL in low-rank MDPs and characterize its theoretical performance.
Since it is impossible to learn model in polynomial time if there is assumption on features de and (h, we first adopt the following conventional assumption (Agarwal et al., 2020).
Assumption 1. (Realizability). A learning agent can access to a model class { V/) that contains
the true model, namely, for any E [H], [T]. the embeddings Φ) E 22,1)
(=,t) E V.
While we assume the cardinality of function classes to be finite for simplicity, extensions to infinite classes with bounded statistical complexity are not difficult (Sun al. 2019 Agarwal al. 2020
4.1 Algorithm design
We first describe our proposed algorithm REFUEL depicted in Algorithm
In each iteration n, for each task / and time step h, the agent executes the exploration policy ". #-1 followed by two uniformly chosen actions to collect trajectories for H episodes. Note that n, #-1 is designed the previous iteration. Then, to utilize the common feature of data collected from different source tasks, the agent processes the data jointly as following steps.
Joint MLE-based Modeling Learning: The agent passes all previously collected data to estimate the low-rank components 0% ,7th
(n.1)) ,...,7%h ((n.)) simultaneously via the MLE oracle
MLE (Uber Dfx.0) (line 8 on the joint distribution defined as follows:
(7() ((n.)) Ph ,..., Ph ((n.T))
Σ log II
(n.t.)) (n.tA))) a (8 (mii)ii))) (I) (on ,a
arg max dn w "....

---

Algorithm REFUEL (REward Free mUltitask rEpresentation Learning)
Input: Regularizer An. parameter Ön, Cn.B3 Ó, eus Models {(u, Φ) : V)
2: Initialize To(-|s) (s) to be uniform, set D ((u,)) for h E [H] t [T].
3: for = 1,..,, , N₂ do
4: for h = I...., H do
5: for I,.... do
Under MDP M', use 7i n- to roll into sh-1 uniformly choose (n.th)) ((.t)) and 0h-1 ,4
6:
enter into s).
hh+1
(n,t,h)
and collect data of sfmath) .] ((...h)) ,...,,h, s (ott.h) (math)) (mac))
,441
7: Add the triple (oh (att)) .h
(n,c)) (n.1))) to the dataset D(n.t)
D ((1-1,))
18411
[(5) (m (n,t,))
-h411 (..
"h
8: Learn ((n . Ph (n.1)
....,,hh" p(n11)
via MIE (Urem D. as Equation 1b.
9: for I..... T do
Update transition kernels pl = (3) (1, )). p2f (-)), empirical covariance matrix (7( as Equation
10:
2 and exploration-dre -driven reward (ln.t)(.,-) as Equation
Set Pseudo Cumulative Value Function PCV (p(n.), B/m-.) Tr; T) as Equation 44
Get exploration policy "T..... T arg max T1 1,, PCV (p(n.)), Byri) Tei TT
if 2PCV (pini), Blm-tt mT; + +2KKT - Teu then
13:
Terminate and set n. n and output : ((m,), p()) p(n..)), ..... Output: T 3(mm), p(1) p(m..)) p(T) p(n..))
...... P(T) = T).
14:
The MLE oracle above generalizes the MLE oracle commonly adopted for single-task RL in the literature (Agarwal et al. 2020; Uehara et al. 2022). In practice, MLE oracle can be efficiently implemented if the model classes 4 V are properly parameterized such as by neural networks.
Next, for each task the agent uses the learned embeddings 2)m) (a,f)
Ph to update the estimated transition kernel p(n,t) at each step h as ((m))) (s'ls, a) = (0) (s, a), Pip (n..)) (s')).
Pseudo Cumulative Value Function Construction for Exploration: The agent first uses the representation estimator
((n) to update the empirical covariance matrix ((m.)) as
((m) (-(.1)) E oh (s), (l (
(r,Ah+)) (...+++))) ()))) (s) (T,A411)
,4 aa ()) + A
(2)
(1.th1))) the tuple collected at the 7-th iteration, (h 1)-th episode, where {s/ (1,,,1111)) 0h ./+1
and step h. Then, the agent uses both (60) and (f(m.)) to provide an exploration driven reward as
(lm (sh, an) min min ān o (sk. an)
(0 ,
where B and ãn are pre-determined parameters. To integrate information of learned models and exploration-drivee driven rewards for T' tasks, define Pseudo Cumulative Value Function (PCV) as
PCV (pin.)) ((na) me T) EML
T
E
(PCm)) mm))
(st, ]
(4)
Joint Policy Learning for Exploration: The agent then learns T policies via optimization over the PCV, and uses these policies as the exploration policy for each task in the next iteration. Since PCV measures the cumulative model estimation eтror over all tasks, these policies will explore the state-action space where the model estimation has large uncertainty on average of T tasks. Here we adopt an optimization oracle for line 12 in Algorithm
Termination: REFUEL equipped with termination condition and outputs the current estimated model and representation during iterations if the PCV plus certain minor term is below threshold, which suffices to guarantee that for each task t, the returned model p(t) from REFUEL reaches an average proximity to the true model p(+.))
5

---

The joint learning and decision making among all tasks involved in the four main components differ- entiate REFUEL from single-task RL algorithms, and enable it to exploit the shared representation across the source tasks to improve the sample efficiency.
4.2 Theoretical results on sample complexity
In this section, we first establish the sample complexity upper bound for REFUEL and then compare our multitask result with that of single-task RI.
Theorem 1. For any fixed 5 E (0. | min }1), set Xn Od o ( Sn Ollog nH//)//// on = O(VK log( + |IV/nH// + A,d'TT) and B - O
Let の. P(1) ..... P(T) be the output of Algorithm 1 Then, under Assumption
with probability at L
least 1 - 8. for any policy Te of task E [T] and any step h E [H], we have
+ Ei (x.) aa).
[|| (-|sh an) - ph(,.0) (+|s/, and
[rvl
ccw cu-
(5)
Further, for any task given any reward r², let क. arg max pi Then with probubility at
least 1 - 6, + Σt Vieior -Vpi..o, Vei S San eu-
Meanwhile, the number of trajectories collected by each task is at most (H K K + H²
T
Equation (5) in Theorem I Lindicates that the estimated transition kernels of the T MDPs meet the required accuracy on average. Theorem T further indicates that with the data collected during such reward-free exploration, if every task is further given any reward function, near-optimal policies can be found for all tasks on average well.
For the sample complexity bound in Equation (6). the first term is inversely proportional T, which captures sample-benefit for each task due to learning the 7 models jointly. The second term is related to the number of samples to guarantee the concentration of empirical covariance matrix (n.t) in order to identify desirable exploration policies. Such concentration needs to be satisfied for each task t independently and can not be implemented jointly, which makes the second term independent with T'. The third term originally arises from the model estimation error shift among T tasks when implementing joint MLE-based model estimation. In the algorithm, this error shift is then transferred to B defined in Equation (3. which provides an explicit upper bound for the exploration-driven
reward (i).t) Note that (/)))) consists the PCV to guide exploration, which finally contributes to convergence of REFUEL and affects its sample complexity.
Compared with the state-of-the-art sample complexity that scales in Ō ( [H'dk² <
when performing T' single-task RL independently (Uchara et al., (022), REFUEL achieves lower sample complexity when T 5 k ) and it reduces the single-task sample complexity in (Uchara et 2022) by an order of O min H22 TT, H2 K, d dK, dT
K
Technically, the joint learning feature of REFUEL requires new analytical techniques. We outline the main steps of the proof of Theorem I heere and defer the detailed proof to Appendix AA] Step 1: We develop new upper bound on model estimation error for each task, which captures the advantage of joint MILE model estimation over single-task learning, as shown in Proposition [ Step 2: We establish the PCV as an uncertainty measure that captures the difference between the estimated and ground truth models. This justifies using such a PCV as a guidance for further exploration, as shown in Proposition 2 Step 3: We show that the summation of the PCVs over all iterations is sublinear with respect to the total number N,, of iterations, as shown in Proposition which further implies polynomial efficiency of REFUEL in learning the models.
Downstream offline and online RL
In downstream RL, the agent is given a new target task T + under MDP MT+1 and aims to find a near-optimal policy for it. The downstream task is assumed to share the same representation
We convert their results in discounted setting to our episodic MDP setting by replacing 1/(1 -7) = ((H). And similarly in following sections, we convert results in other settings to the same setting with us for a fair comparason.

---

as upstream tasks, and hence the agent can adopt the representation learned from the upstream to expedite its learning. Since the agent does not know the task-specific embedding function j u++1 the agent is allowed to interact with MDP M in an online RI. setting or have access to an offline dataset Daown = (15 07 The + 1
Nott . in an offline RL setting, which is rolled out from some
behavior policy p.
In this section, first introduce the connections between upstream and downstream MDPs in Section[5.]] We then provide algorithms for the downstream task and characterize their suboptimality gap for offline RL in Section 5.2 and for online RL in Section 5.3 respectively.
5.1 Connections between upstream and downstream MDPs
In downstream RL tasks, the agent will adopt the feature Ø learned from upstream Algorithm | To ensure that the feature learned in upstream can still work well in downstream tasks, upstream and downstream MDPs should have certain connections. We next elaborate some reasonable assumptions on transition kernels to build such connections.
The following reachability assumption is common in relevant literature (Modi et al. 2021
Assumption 2 (Reachability). For each source task in upstream, whose transition kernel is p(+.t),
there exists a policy of such that min,es P/A (p( (2) 2 Ru,, where P/ (.): : :S - IR is the density (s)
function over using policy qं to roll into state & at timestep h. Assumption 3. The state space S is compact.
Combining Assumption 3 with Assumption 2, it can be shown that there exists uniform distribution on with the density function f(s) = u. where 1/U is the measure of S. Then, we denote ((()) as the uniform distribution over S and (AA as the uniform distribution over S x A.
Assumption 4. For any two different models in the model class {(0, })}. say Pl (s'Is, a) = (s (8,a), .= (s')) and (2 = (d) a), (s8)), there exists a constant Cn such that for all (s, E S A and h E [M].
[- a) P- a)llrv I CRE (1sh,, a4)) P- arll (7)
For normalization, we assume that for any @ E 4, 11, a)ll I and for any a E v and any function g : S + [0, 1]. [I (s(s)s)))ds) S Vd.
Assumption 4 ensures that for each source task t any (s, a) E S A, the total variation distance between the learned (-Is, a) and true transition kernels (-ls, a) will not be too large when the expectation of the total variation distance over the entire state action space is small.
Assumption 5. For the underlying MDP of task T + 1, we assume that the reward function 7T is chosen from a function class R with E-covering number Nxx (E) for any given and (x((E) Ô (2) The true transition kerne P(+,T+1) can be g-approximated by a linear combination of T source tasks, iee there exist T (unknown) coefficients CI...., CT E [0, Ct. such that
v(s, a) S X A, h E [H),
p((+T+)))sss - Σ e, (l)s))((sssss)))))
rv S.
(8)
Assumption 5 establishes the underlying connection between upstream source tasks and the down- stream target task. The assumption on N: (R) is more general than the common assumption on reward under linear MDP and can be easily achieved in practice. For example, if the reward is linear with respect to an unknown feature Φ S x A + RR,, i.e., there exists E TR and ||0]| < R,
r(s, a) = (ō(s. a), f), then (R) = Ō ( (1+ + 2²)") (see Lemma D.5 in Jin et al. (0200), which satisfies the assumption.
Since the downstream task applies the estimated feature from the upstream, the downstream perfor- mance is highly affected by the upstream estimation accuracy of the feature. Hence, we first introduce the definition of c-approximate linear MDP, and provide a guarantee for the estimated feature.
Definition 2 (e-Approximate linear MDP). For any >> 0, MDP M = (S.A. H. P. P,r) is an e- approximate linear MDP with time-dependent feature map & : S x A - Rdd if for any h E (H, there exist a time-dependent unknown (signed) measures μ over S such that for any (s, a)
A.
(a (),))))),)(())l S.
(9)
7

---

Any Φ satisfying Equation (9 is called an e-approximate feature map for M.
Let Édown = - C.CCTTes The following lemma shows that the feature & learned in upstream is a
Nu
Sdowu-aph ceprox-pppo feature map and can approximate the true feature in the new task.
Lemma 1. Under Assuumtions 2 to the output of Algorithm o is Edowa-approxine approximate feature for MDP M = A, H, TT.,T r. $1), i.e. there exist a time- denendenn unknown (signed) measures fi' over S such that for any (s.a) E CSX A. we have
I +1) (-18, a) - (on(s. P.. 0 I &dowwn-
(10)
Furthermore, for any g: S + [0.1]. |l. (s))s( S &LLVd
The proof of Lemma T presented in Appendix B
5.2 Downstream Offline RL
We present our pessimistic value iteration algorithm for downstream offline RL called DOFRL and defer detailed Algorithm 2 Appendix C Although our algorithm shares similar design principles as traditional algorithms for linear MDPs (Jin et al. 202 D. it differs from them significantly as deseribed in the following, due the misspecification of representation from upstream and the general rather than linear reward function adopted. For ease of exposition, we define Bellman ff) operator B,, (B, a) (s, a) + (((+)++)) f) a) for any f x A + R. The main body
of the algorithm consists of a backward iteration over steps. each iteration ht the agent executes the following main steps.
Approximate Bellman update with general reward function. We construct an estimated Bellman update B 1 based on the dataset Daown to approximate TB, V4+1. Here V 1 an estimated value function constructed in the previous iteration based Ddown. Since our reward function is more general and not necessarily linear in the feature, we choose to parameterized pf (s, T a) Tu (s, a) + (P(+,T++))) 1)++1)(s a) due to the linear structure of p(+,T+1) Although the ground truth representation is unknown, the agent is able to use the estimated representation
T +T+1)+ Vi+t1 instead of
learned from the upstream learning parameterize the linear term (P(+,T+3)V 1 (s, a). Then, the coefficient W can be estimated by solving the following regularized least squares problem:
T argmin À ||w l2 Mott w ToTTn (5% a7)) V
(11)
which has the closed-from solution Wn. = D on (n an) ( (8)+1) with An = A + ENat (8,,,, oo (8%, TT) We obtain an estimate (B, (3+))())) ra(:))) TH (n(:)))
Design pessimism for incorporating upstream misspecification of representation. Pessimism is realized by subtracting uncertainty metric Fn from Qn (line 5 The uncertainty metric F quantifies the Bellman update error (IB, Vn+1 B,, Vs+1) a)l at step h. Our selection of I', is in the form of Edown ( A (0,) 1/22, which captures both the approximation error (first term) and the estimation error in the Bellman update (second term). Both of the above errors are affected by the misspecification of representation from upstream learning. Intuitively, represents the effective number of samples the agent has explored along the & direction, and the penalty term B/ /mm represents the uncertainty along the & direction. By choosing a proper value for 3. we can prove that with high probability, the uncertainty metric l' always upper bounds the Bellman update error.
(or,A," ) Gx)-"
Select greedy policy. executee greedy policy म. to maximize Qt, which will be the output of the algorithm after the backward iteration over steps completes.
Then we make the following coverage assumption, which can be easily checked in practice. Assumption 6 (Feature coverage). There exists absolute constant Kp such that
Vh [)) Amin (Ex) 2 Kp: where Σ, E, (sp, an) (sh, an) T s).
We are ready to provide our main result in the following theorem and defer the proof to Appendix C
8

---

Theorem 2. Under Assumptions 2 to 5 for any 5 E (0, 1). if we set X = 1 and 3 = O (d
+ +odNontdown + p log No where = O (log (pdHNot muax[Edwn, 1)/5),, then with proba- bility at least 1 - 8, the suboptimality gap of Algorithm is at most
H
V D, VIOTIID S Edown 2
Σ Em-
[|
(12)
Ax
h=1
If Assumption O holds, and Nof 2 40kp log(AdH then the suboptimality gap is at most
Edown + Kip -1/2 0 Kip 1/2 H22
/ loglpdH Nott max (aox/,don.111pp
+ +Kpp
112 p plog Notl Notl
H
(13)
Nott
Further, if the linear combination misspecification enror in Assumption then with at most (HAK2TNN + H KPT2KT + stream, we have
5 satisfies & O(Va/ (No)).
trajectories collected in up-
d
V ,. PL.ZII, S (N, 11 max (d, vpl) Kip
(14)
To the best of our knowledge, Theorem 2 provides the first suboptimality gap for offline RL with approximation error in representation. The first result in Cheorem@charactenii gap if the offline data rolled out from a behavior policy P If the offline data further satisfies the feature coverage assumption,
the suboptimality
able to upper bound the suboptimality gap by three terms in Equation 13b. Such result can be further simplified as Ô Hdd + Noo max vp}
where the first term a constant error that arises from the approximation error in Bellman update. The second term diminishes as Noff grows to infinity, and it arises from the estimation error of wn regularized linear regression. The dependence on d and vp arises from the model complexity of the linear function approximation class when we seek concentration of the estimated value function.
Compared standard linear MDP with the knowledge of ground truth representation Din al. (2021), our suboptimality gap contains an additional upstream misspecification error (i.c.. Edown). ) The remaining term yields order of Ö(HNN 1/2) which matches that in Jin et al. 2021 if we specialize our reward to be linear with p Odd).
Theorem 2 also demonstrates that our downstream offline RL benefits from the learned represen- tation in upstream. As shown in Equation (14). if the linear combination misspecification error is small enough, and the number of trajectories collected in upstream learning is sufficient large,
then the suboptimality gap is dominated by the term Ö (N 1V-122 H max (d, vp)) which improves
Ō(N 1/2 =²dd²) of REP-LCB Uehara et al. 2022) under low-rank MDP with unknown representa- tion, with better dependence on Hl and d due to the benefit of upstream representation learning.
off
5.3 Downstream Online RL
We present our downstream online RL algorithm called DONRL and defer the detailed Algorithm in Appendix DD. where the agent is allowed to interact with the new task environment for policy optimization, and utilizes the learned feature from upstream to implement linear approximation for state-action value function with traditional UCB type of algorithms.
The algorithm consists of three iterations, one outer iteration and two inner backward iterations over steps. Construct optimistic action value function: In each inner backward iteration, the agent first constructs the empirical feature covariance matrices and estimates the weights Ø via the trajectories collected in previous iterations. Then, the empirical covariance matrices and the weights क are used to approximate Bellman update and construct an optimistic action value function (line 7 Learn optimistic policy: Next, in the forward inner iteration, the agent learns greedy optimistic policy and uses to explore and collect a new trajectory. After the outer iteration completes, the greedy policies are output by Algorithm 3
Algorithm differs from traditional UCB type of algorithms under linear MDP (i.e., LSVI-UCB in Jin et al. 2 as follows. First, Jin et al. 2020 assumes the true representation o" is known with
The result remains valid if offfine data is compliant with MDP M (Assumption 2.2 in Jin et al. 2021

---

small model misspecification error, whereas our representation is obtained from upstream, and hence the upstream learning error affects the learning accuracy here. Second, Algorithm is applicable to reward functions from general function class RR which do not need to be linear with respect to the representation. Both of the above differences will require additional techniques in our analysis of the performance guarantee.
To analyze our downstream online Algorithm we need to characterize how misspecification error of the learned representation from upstream affects downstream suboptimality gap. We further need handle the general reward rather than the linear reward. Both make our analysis different from that in Jin et al. (2020) which directly assumes an approximation of the representation and works with linear reward functions. We next provide our main result and defer the proof to Appendix [
Theorem 3. Under Assumptions to fix 8 E (0. If we set X = and Pn O (d Vndldown + tVploggn where in = log (2pdn max 1)/6). Let ūt be the uni- form mixture of ,',..,,,"" T Then with probability at least 1 1- 6, the suboptimality gap of Algorithm satisfies
V VIL.T V pr Ō ( HHddSown + Ha N max (d. vDt).
(15)
Furthermore, if the lineur combination misspecification error & in Assumption satisfies Ö(Vd/ then with most Ô HHKKTT H d
H2 trajectories
collected in upstream, Edown reduces to O(Va/ and the second term in Equation 15 nates that the the suboptimality gap is bounded as
domi-
Vpc.r413,r Vp V. = Ō [Hd max VDt).
(16)
We first explain the terms in suboptimality gap in Equation (15). The first term captures how the approximation error of using the estimated representation from upstream affects suboptimality gap, and becomes small if such the approximation error is well controlled. The second term arises from the estimation error of Wn in constructing the optimistic action value function and vanishes as the number Non of samples becomes large. The dependence on and vp arises from model complexity of the linear functiin approximation class when we seek a concentration of the estimated value function. Thus, if the learned representation in upstream approximates the ground truth well and Edown reduces Ö(Vd/ /0////oo.. ). then suboptimality gap is dominated by the second term Hd /2 /2 1/2 max(d, vpt) as in Equation (6). If we specialize rewards to be linear and Ö p (dd) this matches suboptimality gap δ (H2dd/No))) in línear MDP (Jin et al. 2020
Theorem 3 also demonstrates that our downstream online RL benefits from the learned representation from upstream. Compared to the state-of-the-arr work under the low-rank MDP with unknown repre-
sentation (Uehara al. 2022 which has the suboptimality gap being bounded as Ö( ((VHodt our downstream online RL has an improved suboptimality gap in terms of the dependence on H, d and K benefited from the upstream representation learning.
6 Conclusion
In this paper, showed that shared representation of low rank MDPs can be Iearned more efficiently and in a transferred manner. We provided upper bounds on the sample complexity of upstream wwward-free multitask RL, and the suboptimality gap of both downstream offline and online RL settings with the help of the learned representation from upstream. Our results show that multitask representation learning is provably more sample efficient than learning each task individually and capture the benefit of employing the learned representation from upstream to learn near-optimal policy of new task in downstream that shares the same representation. As future work, we note that our analysis of the downstream RL depends on some assumptions that connect the upstream and downstream MDPs. It is interesting to explore whether these assumptions can be further relaxed.
7 Acknowledgement
The work of Feng and Y. Liang was supported in part by the startup fund of the Ohio State University.
10

---

References
Agarwal, A., Kakade, S., Krishnamurthy, A., and Sun, W. (2020). Flambe: Structural complexity and representation learning of low rank mdps. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H., editors, Advances in Neural Information Processing Systems, volume 33, pages 20095-20107. Curran Associates, Inc.
Arora, S., Du, S. S., Kakade, S. M., Luo, Y., and Saunshi, N. (2020). Provable representation 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 Learning Research, pages 367-376. PMLR.
Arulkumaran, K., Ashley, D. R., Schmidhuber, J., and Srivastava, R. K. (2022). All you need is super- vised learning: From imitation learning meta-rl with upside down RL. CoRR, abs/2202. 1960.
Azizi, M. J.. Duong, T., Abbasi- Yadkori, Y., György, A., Vernade, C., and Ghavamzadeh, M. (2022). Non-stationary bandits and meta-learning with a small set of optimal arms. CoRR, abs/2202.1300
Barreto, A., Dabney, W., Munos, R., Hunt, J. J., Schaul, T., van Hasselt, H. P., and Silver, D. (2017). Successor features for transfer in reinforcement learning. Guyon, I., Luxburg, U. V., Bengio. S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
Brunskill, E. and Li, L. (2013). Sample complexity of multi-task reinforcement learning. In Nicholson, A. E. and Smyth, P., editors. Proceedings of the Twenty-Nintt Conference on Uncertainty in Artificial Intelligence, UAI 2013, Bellevue, WA, USA, August 11-15, 2013. AUAI Press.
Calandriello, D., Lazaric, A., and Restelli, M. (2015). Sparse multi-task reinforcement learning. Intelligenza Artificiale, 9(1):5-20.
Cella, L., Lounici, K., and Pontil, M. (2022). Multi-task representation learning with stochastic linear bandits. CoRR, abs/2202.10066
Cella, L. and Pontil, M. (2021). Multi-task and meta-learning with sparse linear bandits. In de Campos, C. P... Maathuis, M. H., and Quaeghebeur, E.,, editors, Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence, UAI 2021, Virtual Event, 27-30 July 2021, volume 161 of Proceedings of Machine Learning Research, pages 1692-1702. AUAI Press.
Dann, C., Lattimore, T., and Brunskill, E. (2017). Unifying pac and regret: Uniform pac bounds for episodic reinforcement learning. arXiv preprint arXiv: 1703.077
D'Eramo, C., Tateo, D., Bonarini, A., Restelli, M., and Peters, J. (2020). Sharing knowledge in multi- task deep reinforcement learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Deshmukh, A. A., Dogan, U. and Scott, C. (2017). Multi-task learning for contextual bandits. In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R., editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 4856.
Du, S., Krishnamurthy, A., Jiang, N., Agarwal, A., Dudik, M., and Langford, J. (2019). Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning, pages 1665-16744 PMLR.
Du, S. S.. Hu, W., Kakade, S. M., Lee, J. D., and Lei, Q. (2020). Few-shot learning via learning the representation, provably. arXiv preprint wrxiv:2002.09434.
Hu, J., Chen, X., Jin, C., Li, L., and Wang, L. (2021). Near-optimal representation learning for linear bandits and linear RL. In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages
PMLR.

---

Jiang, N., Krishnamurthy, A., Agarwal, A., Langford, J., and Schapire, R. E. (2017). Contextual decision processes with low bellman rank are pac-learnable. In International Conference on Machine Learning, pages 1704-1713. PMLR.
Jin, C., Yang, Z., Wang. Z., and Jordan, M. I. (2020). Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137-2143. PMLR.
Jin, Y., Yang, Z., and Wang. Z. (2021). Is pessimism provably efficient for offline ri? In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 5084-5096. PMLR.
Kong, W., Somani, R., Song, Z., Kakade, S. M., and Oh, S. (2020). Meta-learning for mixed linear regression. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5394 5404. PMLR.
Lehnert, L. and Littmann M. L. (2020). Successor features combine elements of model-free and model-based reinforcement learning. J. Mach. Learn. Res., 21:1:::1:19999 1-16::33,
Lu, R., Huang, G., and Du, S. S. (2021). On the power of multitask representation learning in linear mdp. arXiv preprint arXiv: 2tt:::206608000 08053.
Maurer, A., Pontil, M., and Jomera-Paredess B. (2016). The benefit of multitask representation learning. J. Mach. Learn. Res., 17:81:1-81:32.
Min, Y., Wang, T., Zhou, D., and Gu, Q. (2021). Variance-aware off-policy evaluation with linear function approximation. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang. and Vaughan, J. W., editors, Advances in Neural Information Processing Systems, volume 34, pages 7598-7610. Curran Associates, Inc.
Misra, D., Henaff, M., Krishnamurthy, A., and Langford, (2020). Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In International conference machine learning, pages 6961-6971. PMLR.
Modi, A., Chen, J., Krishnamurthy, A.. Jiang, N., and Agarwal, A. (2021). Model-free Representation Learning and Exploration in Low-rank MDPs. arXiv e-prints, page arXiv:2102.07035.
Müller, R. and Pacchiano, A. (2022). Meta learning mdps with linear transition models. CoRR, abs/2201.08732.
Qin, Y. Menara, T., Oymak, S., Ching, S., and Pasqualetti. F. (2022). Non-stationary representation learning in sequential linear bandits. CoRR, abs/2201.04800
Sodhani, S., Zhang, A., and Pineau, J. (2021). Multi-task reinforcement learning with context- based representations. In Meila, M. and Zhang. T., editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9767-9779 PMLR.
Sun, W., Jiang, N.. Krishnamurthy, A., Agarwal, A., and Langford, J. (2019). Model-based rl in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference on learning theory, pages 2898-2933, PMLR.
Teh, Y. W., Bapst, V., Czarnecki, W. M., Quan, J.. Kirkpatrick, J., Hadsell, R., Heess, N., and Pascanu, R. (2017). Distral: Robust multitask reinforcement learning. In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R., editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 4496-4500
Tripuraneni, N., Jin, C., and Jordan, M. I. (2021). Provable meta-learning of linear representations. In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages
PMLR.
12

---

Uehara, M., Zhang. X., and Sun, W. (2022). Representation learning for online and offline RL in low-rank MDPs. International Conference on Learning Representations.
Wang, R., Foster, D. P., and Kakade, M. (2021). What are the statistical limits of offline RL with linear function approximatin?? In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Xie, T., Cheng, C., Jiang, N., Mineiro, P., and Agarwal, A. (2021). Bellman-consistem pessimism for offline reinforcement learning. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 6683-6694.
Xiong. W., Zhong, H., Shi, C., Shen, C., Wang. L., and Zhang, T. (2022). Nearly Minimax Optimal Offline Reinforcement Learning with Linear Function Approximation: Single-Agent MDP and Markov Game. arXiv e-prints, page arXiv:2205.1531
Yang. J., Hu, W., Lee, J. D., and Du, S. S. (2021). Impact of representation learning in linear bandits. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Yang. L. and Wang, M. (2020). Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound. In Proceedings of the 37th International Conference Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 10746-10756, PMLR.
Yin, M., Duan, Y., Wang, M., and Wang, Y.-X. (2022). Near-optimal Offline Reinforcement Learning with Linear Representation: Leveraging Variance Information with Pessimism. arXiv e-prints, page wiXiv:223300500.
Yin, M. and Wang. Y.-X. (2021). Towards instance-optimal offline reinforcement learning with pessimism. In Ranzato, M., Beygelzimer, A.. Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems, volume 34, pages 4065-4078. Curran Associates, Inc.
Zanette, A., Lazaric, A., Kochenderfer, M., and Brunskill, E. (2020). Learning near optimal policies with low inherent bellman error. In International Conference on Machine Learning, pages 10978- 10989. PMLR.
Zhang, X., Song, Y., Uchara, M., Wang. M., Sun, W., and Agarwal, A. (2022). Efficient reinforce- ment learning in block mdps: model-free representation learning approach. arXiv preprint arXiv:2202.00063.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 6
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions all theoretical results? [Yes] See section and Section 5.1
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [N/A]
13

---

(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [N/A]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [N/A]
4. If you are using existing assets (c.g., code, data, models) or curating/releusing new assets...
(a) your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you' re using/curating? [N/A]
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "yI641lUz33y", "reviewer": "Reviewer_1pug", "summary": "This paper theoretically characterize the benefit of representation learning under the low-rank Markov decision process (MDP) model. The framework under consideration is a reward-free upstream and reward-known downstream task, where all tasks share the same feature function $\\phi∗$ but might differ in $\\mu$. The downstream offline RL, consists of a new task sharing the same representation as the\nupstream tasks given an offline dataset, and the goal is to find a near-optimal policy.\nThe authors propose a new theoretical algorithm `REFUEL` (REward Free mUltitask rEpresentation Learning) and theoretically show that a shared transferrable representation of low rank MDPs can be learned efficiently. They provide various theoretical results on representation uncertainty measure, sample complexity, and  upper bound on both the sample complexity and the suboptimality gap of downstream offline RL with the help of the learned representation.", "questions": "\n- In terms of the polynomial efficiency of REFUEL, in Proposition 3, it is shown that the summation of the PCVs over all iterations is sublinear with respect to the total number $N_u$ of iterations, assuming \nexploration policy is derived by an oracle. How does this results change in practice when this assumption of existance of an oracle is violated?\n\n- Is  there any theoretical connection between this framework to frameworks such as Successor Features (SFs) where the the representation is decomposed to the dynamics $\\phi$ and $w$ representing the task-specific features? In general I have not seen such connection to be made in the majority of related work in  low-rank MDP and Block MDP literature. Is there a difference between this setting of low-rank MDP which is stricly more general than linear MDPs which assume representation is known a priori ?", "limitations": "- For the analysis to hold, there is dependency on the assumptions that connect the upstream and downstream MDPs. In practice this can be very restrictive in the space of designing multiple tasks.\n- The analysis relies on the existance of an oracle for learning the joint distribution via MLE for the representation and during the policy exploration.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The authors proceed with the algorithm design as follows:\n\n### Joint MLE-based Modeling Learning:\nThe agent passes all previously collected data to estimate the low-rank components $\\phi^n, \\mu^{n,1} ... \\mu^{n,t}$ simultaneously via the MLE oracle on the joint distribution of these random variables. In practice In practice, the representations $\\Phi$, $\\Psi$ are properly parameterized and learned by neural networks. Next, for each task t, the agent uses the learned embeddings to update the estimated transition kernel $\\hat{P}^{(n,t)}$.\n\n### Pseudo Cumulative Value Function Construction for Exploration\nThe agent first uses the representation estimator $\\hat{\\phi}^n$ to update the empirical covariance matrix $\\hat{U}^{(n,t)}$. The Pseudo Cumulative Value Function (PCV) is then established as an uncertainty measure that captures the difference between the estimated and ground truth model.\n\nIn downstream RL, the agent is given a new target task $T + 1$ under MDP $M^{T+1}$ and aims to find a near-optimal policy for $i$. . The downstream task is assumed to share the same representation as\nthe upstream tasks, and hence the agent can adopt the representation learned from the upstream to\nexpedite its learning.\n\nFinally, the joint learning and decision making among all tasks involved in the four main components differ from single-task RL algorithms, and enable it to exploit the shared representation across the source tasks to improve the sample efficiency.\n\nTo ensure that the feature learned in upstream can still work well in downstream tasks, upstream and\ndownstream MDPs should have certain connections. `Assumption 2` (Reachability) provides the assumptions made  on transition kernels to build such connections. \n\n## Strengths\n- Good coverage of theoretical results and the overall problem of learning joint representation in the context of multi-task RL is important for the community.\n- Theoretical results on sample complexity : sample complexity upper bound for REFUEL and then compare multitask result with that of single-task RL.\n- The theoretical results throughout the paper seems sound.\n\n## Weaknesses\n- This paper only provides theoretical results without empirical results\n- There might be missing connection to related work to other MDP frameworks learning decomposed representations\n- The assumptions and upper bounds do not hold in the case where oracle representation is not available and the case where these need to be approximated using function approximation\n\n\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "xtlk02srB6E", "reviewer": "Reviewer_XpNw", "summary": "This paper studies the benefit of joint representation learning in the presence of multiple reinforcement learning tasks. The authors introduce an algorithm called REFUEL that in the presence of $T$ tasks with a shared representation it is able to learn an ensemble of almost optimal policies in a number of trajectory samples that would be necessary if training each of these $T$ tasks independently. The benefits are only seen when $T = \\Omega(\\frac{K}{d^5})$ where $K$ is the number of actions (although the authors do not explicitly mention this). \n\nThese results are in line with a flurry of results in the literature that show the benefit of learning multiple tasks simultaneously in the presence of joint structures. This fact is communicated in the submissions' related work section. The authors then ask the question of wether representation learning can be beneficial in the setting where the learner is to encounter a test task. To show results for this setting the authors introduce strong coverage and reachability assumptions, and also require to assume the point wise TV is upper bounded by a constant multiple of the expected TV of all states sampled under the uniform distribution as well as a `linear dependence' relationship between target and source tasks. Moreover, they also require the learned features during training to satisfy a minimum eigenvalue assumption (over an unspecified measure $\\rho$). Combined these assumptions allow the authors to derive guarantees for downstream offline RL. \n\n\n ", "questions": "I would like to hear some clarification from the authors about the weaknesses described above. ", "limitations": "This work does not pose any negative societal impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths. The setting of multitask representation learning is very interesting and relevant to the community. In particular, studying the problem of downstream offline RL and representational transfer is very exciting. The first part of this work provides expected results for the multi task setting and a well executed algorithm and set of results. Although not super surprising it is a welcome addition to the literature. \n\nWeaknesses: the paper has a couple of typos and undefined variables ($K$, $\\rho$). To date I don't know what $\\rho$ is in Assumption 6. Although the setup of section 5 is interesting, the results presented require a series of strong assumptions that make the reader believe these are there just to make the results work. It would be beneficial if some work had been put in place to stress test these assumptions (in the form of for example lower bound arguments) and make the reader know which of these are necessary or which ones are just a device to make the proofs go through without further effort. ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "TfpyC_yI3FN", "reviewer": "Reviewer_hbyz", "summary": "The paper proposes a representation learning algorithm for multi-task reinforcement learning under low-rank MDPs assumption (transition kernels have a low-rank decomposition in the form of $P^t(s’|s, a) = \\langle \\phi(s’), \\mu^t(s, a)\\rangle$ where the feature of the next state $s’$ is shared across tasks $t \\in [T]$). The algorithm adopts the step-by-step exploration approach from FLAMBE (Agarwal et al., 2020), uses the reward bonus from REP-UCB (Uehara et al., 2021), and extend their single-task analyses to the multi-task setting. The authors showed that jointly learning the representation for multiple tasks is more sample efficient than learning each individual task independently, as long as the number of tasks is large enough. The authors also show that the sample complexity of learning a new task in an offline setting using the learned representation can be reduced under mild assumptions on the similarity of the new task to the pre-training tasks.", "questions": "- It would be also interesting to see if the pre-trained representation would also benefit online learning of the new task $T+1$ (the paper studied offline learning of new task). Are there some challenges that might hinder such analysis?\n- It would be helpful to highlight more on how the proof techniques used in this work are different from the prior work at the end of Section 4.2. The authors listed the new technical pieces that were introduced but I did not have a good understanding on how much different they are from the prior methods. ", "limitations": "The authors have adequately addressed the limitations of their work.  ", "rating": 6, "confidence": 1, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The proposed representation learning for multi-task reinforcement learning setting is novel and the proposed algorithm is shown to provide provable sample complexity benefits when the number of tasks is high (with the number only dependent on the number of actions and the low-rank dimension). \n- Although I did not check the proof carefully, the theoretical results look reasonable. \n\nWeaknesses:\n- Algorithmic-wise the proposed algorithm has limited innovation — 1) for the online setting (Algorithm 1), most components are directly from two prior works (FLAMBE and REP-UCB) with relatively straightforward modifications to make them amenable to the multi-task setting; 2) for the offline setting (Algorithm 2), the algorithm is an one-step version of LSVI-UCB. \n- Assumption 4 in the downstream offline RL seems quite strong. It states that the point-wise (for a certain pair of $s, a$) TV distance of two transition kernels is bounded by the average (over the entire support of $S \\times A$) TV distance of these two kernels multiplied by a constant. If $C_R$ were to depend on $S$ and $A$, it could be as big as $S \\times A$, in which case would not give sample efficiency benefits in the end over the baseline REP-LCB.\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept