Offline Multitask Representation Learning for Reinforcement Learning
Authors: Haque Ishfaq
Thanh Nguyen-Tang
Songtao Feng
Raman Arora
Mengdi Wang
Ming Yin
Doina Precup
Venue: NeurIPS 2024 poster (2024)
Keywords: Multitask learning
representation Learning
offline RL
provably efficient
suboptimality gap
reward-free RL
Primary area: reinforcement_learning
TL;DR: We show provable benefit of offline multitask reinforcement learning for both upstream and downstream tasks when certain part of representation is shared among the tasks.

Abstract: We study offline multitask representation learning in reinforcement learning (RL), where a learner is provided with an offline dataset from different tasks that share a common representation and is asked to learn the shared representation. We theoretically investigate offline multitask low-rank RL, and propose a new algorithm called MORL for offline multitask representation learning. Furthermore, we examine downstream RL in reward-free, offline and online scenarios, where a new task is introduced to the agent that shares the same representation as the upstream offline tasks. Our theoretical results demonstrate the benefits of using the learned representation from the upstream offline task instead of directly learning the representation of the low-rank model.

Offline Multitask Representation Learning for
Reinforcement Learning
Haque Ishfaq" Mila, McGill University
Thanh Nguyen-Tang Johns Hopkins University nguyent@cs. jhu. edu
Songtao Feng
University of Florida sfeng10ufl.. edu
haque ishfaq@mail mwgill.le...
Raman Arora
Mengdi Wang Princeton University
Ming Yin" Princeton University my0049@princetton.cm edu
Johns Hopkins University arorašcs.jbl juu edu
mengdivepprnneeton. edu
Doina Precup" Mila, McGill University dprecupècs www.wmc.cccco..oo
Abstract
We study offline multitask representation learning in reinforcement learning (RL), where learner is provided with offline dataset from different tasks that share a common representation and is tasked to learn the shared representation. We theoretically investigate offline multitask low-rank RL, and propose new algo- rithm called MORL for offline multitask representation learning. Furthermore, examine downstream RL in reward-free, offline and online scenarios, where a new task is introduced to the agent that shares the same representation as the upstream offline tasks. Our theoretical results demonstrate the benefits of using the learned representation from the upstream offline task instead of directly learning the representation of the low-rank model.
Introduction
Recent advances in offline reinforcement learning (RL) (Levine et al., 2020) have opened up possi- bilities for training policies for real-world problems using pre-collected datasets, such robotics (Kalashnikov et al., 2018; Rafailov et al., 2021; Kalashnikov et al., 2021), natural language processing (Jaques et al., 2019), education (De Lima and Krohling, 2021), electricity supply (Zhan al., 2022) and healthcare (Guez et al., 2008; Shortreed et al., 2011: Wang et al., 2018; Killian et al., 2020). While most offline RL studies focused single-task problems, there many practical scenarios where multiple tasks are correlated and is beneficial to learn multiple tasks jointly by utilizing all of the data available (Kalashnikov et al., 2018; Yu et al., 2021, 2022; Xie and Finn, 2022). One popular approach in such cases is multitask representation learning, where the agent aims to tackle the problem by extracting shared oow-dimensional representation function among related tasks and then using simple function (e.g., linear) on top of this common representation to solve each task (Caruana, 1997; Baxter, 2000), Despite the empirical success of multitask representation learning, particularly in reinforcement learning for its efficacy in reducing the sample complexity (Teh et al., 2017; Sodhani et al., 2021; ; Arulkumaran al., 2022), the theoretical understanding of is still in its early stages (Brunskill and Li, 2013; Calandriello et al., 2014; Arora et al., 2020; D E'Eamoo al., 2020; Hu et al., 2021; Lu et al., 2021; Müller and Pacchiano, 2022). Although some works
'The corresponding authors.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

theoretically studied the online multitask representation learning for RL where the agent is allowed interact with multiple source tasks learn the shared representation (Cheng et al., 2022; Agarwal et al., 2023; Sam et al., 2024), there is currently no theoretical understanding on the effectiveness of multitask RL. in the offline setting. This crucial as many practical scenarios (Kumar et al., 2022; Yoo et al., 2022; Lin et al., 2022), it is not feasible to interact with the different task environments in an online manner.
Moreover, when the tasks share the same representation, offline multitask representation learning can serve as launchpad for effectively solving many other downstream tasks (Kumar et al., 2023). Consider the problem of learning to control robotic arms where we may already have offline datasets from different pick-and and--aacce tasks in a kitchen such as the Bridge Dataset (Ebert et al., 2022). From this one can consider many possible downstream RL tasks where representation learned from these offline datasets can be beneficial. For example, one may consider solving a new pick-and-place task with different previously unseen objects in either an online or offline manner. Alternatively, one may consider a downstream reward-free RI. (Jin et al., 2020a) task where the agent would first gather additional novel and diverse data without a pre-specified reward function and afterward, when provided with any reward function (e.g. slightly different target placing spot for the picked object), would be asked to provide good policy without additional interaction.
In this work, we study the provable benefits of offline multi-task, representation learning for RL in which the learner is only given access to pre-collected data from different source tasks which are modeled by low-rank MDPs (Agarwal et al., 2020b) with a shared (yet unknown) representation.
Our contributions. We develop a new offline multitask reinforcement learning algorithm that enables sample efficient representation learning in low-rank MDPs (Agarwal et al., 2020b) and further provide improved sample complexity to the downstream learning. In summary, our main contributions are:
We propose a new offline multitask representation learning algorithm called Multitask Offline Representation Learning (MORL) under low-rank MDPs. MORL represents a standard training procedure in modern machine learning, by pooling the data from all source tasks learn shared representation of the dynamics via maximum likelihood estimation oracle.
We prove that, MORL can Iearn near-accurate model, and, when combined with the pessimism principle, find a near-optimal policy for each of the source tasks T in the average sense, more sample-efficiently. than learning each task in isolation. To our knowledge, this is the first theoretical result demonstrating the benefit of representation learning in offline multitask RL.
We then show theoretical benefits of using the learned representation from MORL in downstream reward-free RL Jin al. (2020a); Wang al. (2020). In particular, we show that, guarantee an e-suboptimal policy for uniformly over any reward function, our algorithm requires at most Ō O(H+2) episodes during the exploration phase where d is the dimension of the feature and / is the planning horizon. This improves the best known sample complexity for the reward-free RL in low-rank MDP (Cheng et al., 2023) by factor of Ö (Hdd), where K is the cardinality of the action space. In addition,
complementary result, we show that using the learned representation from MORL improves the suboptimality gap bound in both offline and online downstream task.
2 Preliminary
Episodic MDP. We consider an episodic discrete-time Markov Decision Process (MDP), denoted
by M (S, A, H, r), where S is the state space, A is the action space with cardinality K. H is the finite episode length, P = HE are the state transition probability distributions with Pμ S x A + A(S), and r = { Following prior work (Jiang et al., 1= are the deterministic reward functions with rı : x A + [o, 1]. 2017; Sun et al., 2019), we assume that the initial state 8 is fixed for each episode. A policy x is a collection action that the agent takes at state and the h-th step in the episode. Given starting state Sh. Sy ~ (P. T)) denotes a state sampled by executing policy π under the transition model P for h' steps and E (s,.,a)>>(P (+) denotes the expectation over states Sh (P, T) and actions an Moreover, for each h E [H], we define the value function under policy æ when starting from an
H functions {mA : 3 A]ke[H] where T (s) is the

---

arbitrary state at the h-th time step as
H The (sw aN) =
Vi, (s) = [s (oanaa)) an (P(r))
We define the action-value function for a given state-action pair (s, a) under policy æ at step / as
H
Q ( = Blea (x)
The ar = 8, On = 0
h'=h
Defining (Pr)) a) = P(yy-P(is..))((s' for any function f S - R. we write the Bellman equation associated with policy ι as
(5, = (rh- + Pa (s, a), (s),,,,)) (T,,,, (s) 0. (2.1)
Since the MDP begins with the same initial state 81, for simplicity, we use V7,, to denote ($1) Another useful concept is the notion of occupancy measure of a policy и at time step h under transition kernel P. Specifically,
use d'p, (s, a) to denote the marginal probability of encountering the state-action pair (s, a) at time step h when exccuting policy under MDP with transition kernel P. Finally, we denote (S) and (S .A) as the uniform distribution over S and S x A respectively.
We study low-rank MDPs (Jiang et al., 2017; Agarwal et al., 2020b) defined as follows.
Definition 2.1 (Low-rank MDPs). A transition kernel P% S x A + ASS) admits a low-rank decomposition with dimension d N if there exists two unknown embedding functions 67 A Rd and MN S Idd such that for all s. E S and E A. (s s,a) (s, p (s)). Without loss of generality, we assume (1,%, a)ll < 1 for all (s,a) E S x A and for any function 9 > [0, 1). J pi (s) S vd.
We remark that the upper bounds on the norm of o² and μr¹ are just for normalization. As the function class + for o" can be non-linear, flexible function class, the low-rank MDP generalizes prior works with linear representations (Jin et al., 2020b; Hu et al., 2021) where it assumed that the true representation o² is known to the agent priori.
2.1 Offline Multitask RL with Downstream Learning
In offline multitask RL upstream learning, the agent is provided with an offline dataset collected from T' source tasks, where the reward functions {r" herry are assumed to be known. Each task t e [2'] is associated with a low-rank MDP M = (S,A, H, r'), Here, all I' tasks are identical except for (1) their true transition model p(=1), which admits a low-rank decomposition with dimension d: p(*,')(s)) (s%, Shy ax) (0% an), ( (*11) (s())) for all h E [H] t E [T], and (2) their (+.t) and rf. we emphasize that all tasks share the same feature function 0% We have access to offline dataset D = ULE(TT ((t), where D(!
reward rhe While the tasks may differ in Ph
.04 ,7h ,r(..) -8h41 [iem) with (1.t) ~ PA p(*.))(. (-1%). (1.2) "aaa and D/.") was collected using a fixed behavior policy n/. In the upstream learning stage, the goal is to find near-optimal policy and near-accurate model for any task / E [T] and any reward function {rr (raheerr through the use of offline dataset D and provide well-learned representation for the downstream task. In order to achieve bounded sample complexity in offline RL, we need additional coverage assumption on the behavior policy TÝ. One common coverage assumption is the global coverage assumption (Antos et al., 2008; Munos and Szepesvári, 2008), which assumes the occupancy measure under the behavior policy Tद globally covers the the occupancy measure under any possible policies, i.e., the concentrability ratio /d
((sh1.) (1.1)
(t.1) (6.0))
satisfies, maxx,..a di....) (s,,))d" p(-..) (s, a) < oc. Instead, we make a partial coverage assumption
and our suboptimality bound scales with the relative condition number (Agarwal et al., 2020a, 2021) instead of the global concentrability ratio, where the former can be substantially smaller than the latter. Under this assumption, we want to compete against any comparator policy covered by the offline data. In Section 3.2, we define the partial coverage condition using relative condition number, which was previously used in the context of single-task offline RI (Uehara and Sun, 2021; Uehara et al., 2022) and generalize it to offline multitask setting.
3

---

Algorithm Multitask Offline Representation Learning (MORL)
1: Input:
Dataset D {(0)
(1.0) (1) : Φ E v 1 e )
(1.0) (1,4) (6.4) an
(1.4) Th (HHt)eeN)))e(())xee))) Regularizer X, Parameter a, Models
2: for h = 1.
H do
3: Learn 1b. Ph ,....,,, P/ (1
via MLE (Ulerry Dk²) as (3.1)
4: end for
5: for 1 1...... do do
6: Update estimated transitioned kernels ((t) (.)) (3.2), empirical covariance matrix (0) as (3.3) and penalty term ((0) as (3.4)
h
2
7: Get policy ग. argmax, -800
8: end for
9: Output: P(1), PIT), $ fT
In downstream learning stage, a new target task T + 1 with a low-rank transition kernel P(+7++) and the same S, AA and is assigned the agent. The transition kernel p(+,++1) shares the same representation d* with the T upstream tasks, but has task-distinct (((+T+1), We consider three settings for downstream tasks reward-free, offline and online RL, where the agent needs to use the representation function @ learned during the upstream stage to interact with the new task environment.
In the reward-free setting, firstly proposed in Jín al. (2020a), the agent first interacts with the new task environment without accessing the reward function in the exploration phase for up to KREE episodes. Afterwards, is provided with reward function e-optimal policy π in the planning phase. We define the sample complexity to be the number of episodes KRFF required in the exploration phase output an e- optimal policy in the planning phase for any given reward function r.
(r)) and asked to output
In the offline and online setting the downstream task T + 1 is already assigned with unknown reward function T++ and the goal is find near-optimal policy for the new task. The agent is expected to expedite its downstream learning through using the representation learned from the offline upstream task. In the online setting, it is allowed to interact with the new task environment and in the offline setting, it is instead provided with offline dataset Dar Useun DR. where
Dr = ((T), 0h The 8 (++))re(Nal) and Dofr were collected using some behavior policy p.
Upstream Offline Multitask Representation Learning
In this section, we introduce our algorithm Multitask Offline Reinforcement Learning (MORL) designed for upstream offline multitask RL in low-rank MDPs and describe its theoretical properties.
3.1 Algorithm Design
The details of the algorithm MORL is depicted in Algorithm 1. The agent passes all input offline data
to estimate low-rank components &n. Ph ..... Ph simultaneously via the Maximum Likelihood
(()) (T)
Estimation (MLE) oracle MLE (Ukem D(.)) on the joint distribution defined as follows:
T
(Nn, Ah " ((T)
,...,,1?
argmax
ΣΣ log ((on(s) (6.0) (, ph((64))))
(3.1)
(m)
The MLE oracle (3.1) is the offline multitask counterpart to celebrated MLE oracle in online multitask RI. (Agarwal et al., 2020b; Cheng et al., 2022; Agarwal et al., 2023). The MIL oracle can be reasonably approximated in practice whenever optimizing over & and V is feasible through proper parameterization such by neural network. For each task f, we obtain the estimated

---

transition kernel p(o) at each step h using the learned embeddings 6. ( Ph h ( 8, (s, (s)))
(3.2)
Using the representation estimator the we set the empirical covariance matrix 500 for task as h.
500 (6.0) (n ,4 "
Σ
oo T + AL.
(3.3)
Using both 6n and (tt) h,8
construct lower confidence bound penalty term follows:
(sh. (sh. an) = min { (s).sss)))lss (o
(3.4)
where a is pre-determined parameter.
Finally, for each task I, with the learned model p(1) and the reward rf B(0, we do planning to get policy fe..
3.2 Theoretical Result on Upstream Task
To facilitate the model selection task using the joint MLE oracle in (3.1), posit realizability assumption which is standard in low-rank MDP literature (Agarwal et al., 2020b; Cheng et al., 2022).
Assumption 3.1 (Realizability). A learning agent has access to a model class {Ø, v that contains the true model, i.e., for any E [M]] TT. the embeddings Φ6 E 4, p(++t)) E v. For normalization, assume that for any ø कb [ a2 S 1 and for any μ E V and any function g S : [0, 1],
II pm (s)g())d=|12 I vd.
For simplicity, we assume that the cardinality of the function classes & and v are finite.
Next, we define the multitask relative condition number C*, which is natural extension of the standard relative condition number (Agarwal et al., 2020a, 2021; Uchara et al., 2022).
Definition 3.2 (Multi-task relative condition number). For task t and time step h, we define Ci (Te, m the relative condition number under oh:
(Lm...)<PPt az ( (sh, an Jof an) a Jx
(T) = sup
(3.5)
ZERd (( 10 (5h, an )on)) (sh, an) T
We define C maxRE(H) Ci, mL) and C* max C7.
Intuitively, Cū,, (Te, nT) defined in (3.5) measures the deviation between a comparator policy ne and the behavior policy of time step h. When tabular MDP is considered (i.e., of is one-hot encoding vector), this relative condition number reduces to the density ratio based single-policy concentrability
coefficient, 07,,,00 (Tt, m maxs,a (pp... (8, a) of" p.... (s, a) (Chen and Jiang, 2022). The relative
condition number (m,,, m)) is always bounded by the concentrability coefficient (Uehara and Sun, 2021). In addition, the relative condition number is computed under the averaging over the state, actions, which could be much smaller than max,.a d di (8, "d
p):
(s, a) since the latter p pt...) (s, a) ratio. Therefore,
p(-.e)
will be very large as long as any s, a pair gives large d p (s, a) /d our C could be much smaller compared to the concentrability coefficient, especially in large-scale MDPs (e.g. continuous state space). Moreover, in our definition of relative condition number, use the unknown true representation ø² Finally, we generalize single task relative condition number multitask setting by defining C*, by simply taking the maximum the single-task relative condition numbers.
Now describe main theorem.
5

---

Theorem 3.3. (a) Under Assumption 3.1, with probability at least 1 5, for any step h [H]
we
have
T
log (2 1 1 T nH//0)
E (o))a)) ((Plr.).)))
(
sh, an) (. an) rv S
T
nT
(3.6)
where 6() p(T) be the output of Algorithm 1.
(b) In addition, in Algorithm 1, if we set a with Sn = (2] w/ [(' and с being a constant, where we assume that w == max; max, a (1/T)(a| < co, then under Assumption 3.1, with probability at least 1 - 6, we have
V2nwGn + Ad, À = ed cdlo[ /
VTL V p(e pi swodtly
C +2d11
AC +w n
dC'Sa
d +ay -
WSn
(3.7)
where (TelrE(T] the output of Algorithm
Theorem 3.3 (a) shows potential benefit of a joint learning of the source-task transition kernels, as compared to independent learning. as measured by the task-average TV distance (defined in the LHS of Equation (3.6)). particular, to obtain an c-suboptimal transition kernels, it suffices for the joint learning to use Ö
log (+)
T22 + log| samples per task, yet for the independent learning to
δ |og.
[+1
log: |og:"" [w]
samples per task. This benefit naturally comes from the inductive bias that all
the source tasks share representation in I which can be learned more accurately with the aggregated data pooled from all the source tasks. Equation (3.7) Theorem 3.3 further shows that, for all tasks on average, we can uniformly compete with any set of comparator policies {Te hetry satisfying the partial coverage through C* < oo. particular, if the optimal policy मे is covered by the off data for all E [T], then the output (ETT) of Algorithm is able to compete against it on average as well.
Remark 3.4. We note that in order for the bound hold in Theorem 3.3, Algorithm I requires the knowledge of w as it is required to set the value of a which is used in the lower confidence bound penalty term defined in (3.4). However, practice, we expect that a can be treated hyperparameter that might be tuned using grid search.
Proof outline. Here, we highlight the key steps for the proof of Theorem 3.3. The detailed proof is deferred to Appendix D. Using offline multitask MLE lemma (Lemma F.3) and ste-step back lemma (Lemma G. 1), we first develop new upper bound on model estimation error for each task which single-task offline learning. We then use the following lemma to show near pessimism in the average sense.
encapsulates the benefit of joint offline MLE model estimation
Lemma 3.5. For any policy Ti and reward r², we have, with probability 1 -
VTL pt
V p
S H H.W.ST/T
T
2 log (2) // / (7/8)
where S..
n
The proof of Lemma 3.5 relies on simulation lemma and a concentration argument for the penalty term defined in (3.4). Finally, to obtain a result only depending on the relative condition number using the true representation ø² but not the learned feature σ, we translate the penalty defined with to the potential function (sp, an) an)ll(,,....>>>>>>> where En, = (B x²) ( 0 + Al using a one-step back inequality (Lemma G. and distribution shift lemma (Lemma O.2).
4 Downstream RL: Reward-free Exploration, Offline RL and Online RL
4.1 Relationship between upstream and downstream MDPs
In order for us to theoretically study the downstream RL tasks where we would use the learned feature 6 from the upstream tasks, first we need to make certain connection between the upstream

---

and downstream MDPs. Naturally, we would have to resort to some assumptions on the transition kernels to make such connections. Next, we describe these assumptions, which we largely adopt from Cheng et al. (2022).
Assumption 4.1. We make the following assumptions
For each upstream source task & with transition kernel p(*)), the behavior policy ni. is such that 1% p((((.0) (ofof:)())
(s) 2 K, where P p(a):"(-)) : + IR is the marginalized state occupancy measure over S using policy ni at time step
minges
2. The state space is compact and has finite measure 1/v, and the induced uniform probability density function is (s) = v, 8 S.
For any two models Pa (' a) - (0) (s- (), and (2 (62 a), (²())) in the model class + x V, we have, (-|s,, a)llTV I CnB( P' (-)s,
||P] (-1s,a)
for all (s, E x A and h E [H] where On is an absolute constant.
4. The transition kernel of task T 1, p(+T++) can be E-approximated by a linear combination of the source upstream tasks, i.e. there exist unknown coefficients cI...., or 2 such that Elz a S C2 and & 2 0 such that for all (s, a) E x A and h E [H], we have
T
Σ papp (1s, S g.
Here, & is called the linear combination misspecification.
The first point in Assumption 4.1 ensures that the behavior policy म. for each task / can reach any state in S at any time step with a positive probability, (Yin et al., 2021). Compared to Cheng et al. (2022), which assumes the existence of a policy with reachability property in each of the upstream online tasks, ours assumes reachability property for the behavior policies used to collect the upstream offline dataset.
assumption that is previously used in
The third point in Assumption 4.1 ensures that for each source task f, the point-wise TV eтror between
the learned estimated transition kernel p(t) and the true transition kernel p(++t) is bounded by the population-lewe TV eиror. This assumption is necessary to transfer the MLE error from the upstream source tasks to the downstream target task.
Finally, the fourth point in Assumption 4.1 connects the upstream source tasks with the downstream target task by assuming that the transition kernel the target task p(+I++1) can be approximated by a linear combination of transition kernels of T upstream source tasks.
The precision of the feature estimation in the upstream has a significant impact on the downstream task's 's performance because the downstream task utilizes the estimated feature from the upstream. We use the following notion of e-approximate linear MDP to provide guarantee for the estimated feature.
Definition 4.2 (e-approximate linear MDP (Jin et al., 2020b; Cheng al., 2022)). For any > 0,
we say that MDP M = (S,A, H, P.r) is an e-approximate linear MDP with a feature map Φn S A + Rd. if for any h [H] there exist d unknown (signed) measures μ4 (14) Hh over S such that for any (s, a) E S A, we have
(1) (d)
[[Ps a) (dn(s, un (-)) S Any & satisfying (4.1) is called a-apprxximate feature map of M.
(4.1)
The next lemma shows that the learned feature T from the upstream offline tasks can approximate the true feature in the new downstream task.
Lemma 4.3. Under Assumption 4.1, the output @ of Algorithm for MDP MT++ where Edown = E Cr.Cav v 2T (1 i.e, there exist a time-dependent
Edown abpr-xxipo feature
1 unknown (signed) measure ft" over S such that for any (s,a) E S x A, we have
(-]s. (6n(s, (R))) I Edown-
Furthermore, for any g: S + [0, 1]. we have [I (s) dssl S CCVdd.
7

---

4.2 Downstream Reward-Free RL
Our goal in this part is to investigate the statistical efficiency of reward-free RL in low-rank MDP while having access to offline datasets from the upstream tasks.
Our algorithm for the reward-free setting is presented in Algorithm 2 (exploration phase) and Algorithm 3 (planning phase) which is built on the procedure of optimistic learning as Wang et al. (2020); Zhang et al. (2021). While having similar design principle as in Wang et al. (2020), our algorithm differs from them due to the misspecification of representation from the upstream task. Thus, the upstream learning error affects the learning accuracy and downstream suboptimality gap and we need to account for that our analysis. Another difference with Wang et al. (2020) that in the exploration phase, like Chen et al. (2022b), we construct more aggressive reward function to avoid overly-conservativy exploration, which removes the extra dependency of sample complexity on episode length H. Below, we provide our main theorem for downstream reward-free RL task and defer the proof to Appendix I.
Theorem 4.4. Under Assumption 1. after collecting KEFE trajectories during the exploration phase in Algorithm 2. with probability at least 1 - 6, the output of Algorithm 3, policy a satisfies
Es
d loggd lor(dKureHHH))) / KRFE + Edown. - -VV r))
2down
(4.2)
If the linear combination misspecificatiion error & in Assumption 4.1 satisfies Ō( d3 KRFE) and the number of trajectories in the offline datasel for each upstream task is at least O K (eredd), then, provided KAFF at least O dd (o an c-optimal policy for any given reward during the planning phose.
with probability
the policy n will be
We compare the above result with other algorithms developed for the reward-free RL under low-rank MDPs and summarize the comparison in Table 1. FLAMBE (Agarwal et al., 2020b) achieves a sample complexity of Ō === k") whereas MOFFLE (Modi al., 2024) achieves a sample complexity of
H" d Kia " ), where is reachability probability to all states. More recently, Cheng et al. (2023)
proposed RAFFLE which has the best-known sample complexity of Ō ((H"g" ))2 Cheng al. (2023)
further shows that the dependence sample complexity on action space cardinality K is unavoidable when performing reward-free exploration in low-rank MDPs from which they conclude that it is strictly harder to find a near-optimal policy under low-rank MDPs than under linear MDPs. However, we from Theorem 4.4, by using estimated representation from the upstream offline datasets, we can avoid this dependence of sample complexity on K and overall improve the sample complexity by Ō(HdK) compared to that of RAFFLE. Moreover, compared to standard linear MDP, where the true representation {0r, { is known (Jin et al., 2020b), the suboptimality gap in (4.2) contains an additional term {aown which is due the upstream misspecification error Edown. When Édown is small enough, our resulting sample complexity of Ō( H*d") matches the reward-free exploration sample complexity for linear mixture MDPs (Chen et al., 2022b), which is worse off by only Ö(d)
compared to the best known sample complexity of Ō H& (Hu et al., 2022) for linear MDP.
Algorithm
Sample Complexity
Task
FLAMBE (Agarwal et al., 2020b)
0 H K"
Single task
.10
MOFFLE (Modi et al., 2024)
Ó H' &"K K
Single task
min (
RAFFLE (Cheng et al., 2023)
Ö H
Single task
This work (Algorithm 2 and Algorithm 3)
Ó H
Multi-task
Table 1:
Sample complexities of different approaches to learning an e-optimal policy for the reward- free RL setting with low-rank MDPs.
We rescale the result in (Cheng et al., 2023) by a factor of H² as we do not assume the sum of rewards to be within [0, 1].
8

---

4.3 Downstream Offline and Online RL
For completeness, we also consider downstream offline and online RL which was previously studied in Cheng et al. (2022) to show the effectiveness of our offline representation. In both cases, we assume that the reward function TT+ in the downstream task T is linear with respect to the unknown feature o² S x A -> IRd. We emphasize that, unlike Cheng et al. (2022), we assume the reward function --1 is unknown.
Offline RL. For downstream offline RL task, similar to Cheng et al. (2022), we use standard pessimistic value iteration algorithm (PEVI) (Jin al., 2021) with approximate feature learned from upstream task. We make the following data-coverage type of assumption which is standard in the study of offline RL (Xie et al., 2021b; Wang et al., 2021; Yin et al., 2021). Moreover, this assumption has been shown to be necessary for sample efficient offline RL for tabular and linear MDPs (Wang al., 2021; Yin and Wang, 2021).
Assumption 4.5 (Feature coverage). There exists an absolute constant Rip such that for all h E [H] and Φn, E D/. Amiu (E,lon an (s,, T s)) 2 Kpp
Next, we provide our result for the downstream offline RL task and defer the proof to Appendix K.
Theorem 4.6. Under Assumption 4.1, setting Xa 1, B O(Hdvi+ + H / (oooooo)))) where
== log( Hd max 1))8), with probability at least 1 - 5, the suboptimality gup of Algorithm 4
at most
H
VP.. (P v po 2HH Edown + 2
> E..
ET.
| and an)lla==18 s)
(4.3)
Additionally if Assumption 4.5 holds, and the sample size satisfies Nog 2 40/mpp looo (og(4dH//), then with probability 1 6, we have,
, (6) - V pt- (8) s 0 H²d log log(HdNog max(Eamm, 1)/5)
Vpi 1
khp 22 Edown
Nog
(4.4)
Online RL. For downstream online RL task, where the agent is allowed to interact with the new task MDP MI+1 for policy optimization, similar to Cheng et al. (2022), we use standard LSVI-UCB algorithm (Jin et al., 2020b) with approximate feature. We next provide our result for downstream online RL task and defer the proof to Appendix M.
Theorem 4.7. Let Te be the uniform mixture of π TI....,
T Now sumption 4.1, setting X 1, Bn O(Hd (o) + H Hn + CLVHd) where in log H dn max(Edownr 1)/5), with probability ô, the suboptimality gap of Algorithm satisfies
Non in Algorithm S. Under As-
Vpo.r+u, V ,1., S 8/2 =H2dooo)))
5 Related Work
Offline Reinforcement Learning. Offline RL (Ernst et al., 2005; Riedmiller, 2005; Lange et al., 2012; Levine et al., 2020) studies the problem of learning policy from static dataset without interacting with the environment. The key challenge in offline RL is the insufficient coverage of the dataset, due to the lack of exploration (Levine al., 2020; Liu et al., 2020). One prevalent approach to address this challenge is the pessimism principle to penalize the estimated value of the under-covered state-action pairs. There have been extensive studies on incorporating pessimism into the development of different approaches in angle-task offline RL, including model-based approach (Rashidinejad et al., 2021; I; Uchara and Sun, 2022; Jin et al., 2021; Yu et al., 2020; Xie et al., 2021b; Uchara et al., 2022; Yin et al., 2022), model-free approaches (Kumar et al., 2020; Wu et al., 2021; Bai et al., 2022; Ghasemipour et al., 2022; Yan et al., 2023; Nguyen- Tang et al., 2022, 2023; Nguyen-Tang and Arora, 2023), and policy-based approach (Rezaeifar et al., 2022; Xie et al., 2021a: Zanette et al., 2021b; Nguyen- Tang and Arora, 2024a,b). Our algorithm for upstream offline multitask RL is inspired by the uncertainty-based pessimism methods in single-task offline RL.
Low-rank MDPs. Agarwal et al. (2020b) initiates the study of low-rank MDPs. Uehara et al. (2022) proposed model-based algorithms for both online and offline RL, while Modi et al. (2024) put forward

---

a model-free algorithm for low-rank MDPs. Moreover, Du et al. (2019); Misra et al. (2020); Zhang al. (2022) studied block MDPs, which a special case of low-rank MDPs.
Offline Data Sharing in RL. There has been several empirical works that investigated the benefits of using offline datasets from multiple tasks to accelerate downstream learning (Eysenbach et al., 2020; Kalashnikov et al., 2021; Mitchell et al., 2021; Yu et al., 2021; Yoo et al., 2022). Yu et al. (2021) show that selectively sharing data between tasks can be helpful for offline multitask learning. For instance, earlier studies have investigated the development of data sharing strategies through human domain knowledge (Kalashnikov et al., 2021), inverse RL (Reddy et al., 2019; Eysenbach et al., 2020; Li et al., 2020), and estimated Q-values (Yu et al., 2021). More recently, Xu et al. (2023) uses offline dataset from diverse tasks perform offline multitask pretraining of a world model which is then finetumed on a downstream target task. Hu et al. (2023) proposes provably efficient self- self-supervised offline data-sharing algorithm for linear MDP. However, they assume access to reward-free data.
Comparison to Cheng et al. (2022). Closest to our work is Cheng et al. (2022) who studied online multitask RL. Their proposed REFUEL algorithm combines design principles from FLAMBE (Agarwal et al., 2020b) and REP-UCB (Uchara et al., 2022) and performs joint MLE based model learning while collecting data in an online manner. On the contrary, MORL first performs joint MLE based model learning using the offline dataset collected for each source task and then upon constructing penalty terms for each task, performs planning using pessimistic reward functions. While both works rely on an MLE oracle, first proposed in Agarwal et al. (2020b) for single-task RL, our proposed offline multitask MLE lemma (Lemma F.3) conveys fundamentally very different ideas compared to its online counterpart, Lemma 3 in Cheng et al. (2022). Lemma 3 in Cheng et al. (2022) says that when exploration policies for each upstream online source task is uniformly chosen, the summation of the estimation error of transition probability can be bounded with high probability. On the contrary, our Lemma F.3 states that when the offline datasets for each of the upstream offline source tasks are collected using respective behavior policies, the summation of the estimation error is bounded with high probability. For the downstream task, Cheng et al. (2022) studies only offline and online RL task, whereas our primary contribution in this part is in the study of downstream reward-free RL which has not been previously studied in the context of multitask representation learning. For completeness, complementary result. Moreover, unlike us, Cheng et al. (2022) assumes that the ceward-function is known in the downstream task, which is fairly strong assumption. Somewhat in contrived manner, in Cheng et al. (2022), the reward-functe function is further assumed to be general and not necessarily linear in the feature which complicates their downstream analysis. On the contrary, we assume that the reward function is linear with respect to the feature. Finally, Cheng et al. (2022) assumes that for each episode in any task MDP, the sum of reward is normalized be within [0, 1]. We do not make this assumption for fair comparison to literature on reward-free RL for low-rank MDPs.
provide results in downstream offline and online RL setting
Comparison to Concurrent Work. In concurrent work, Bose et al. (2024) studies representational transfer in offline low-rank RL. In the upstream task, similar to ours, Bose et al. (2024) also uses an offline MLE oracle. Compared to our Theorem 3.3, where we provide bound for average accuracy of the estimated transition kernels of the upstream source tasks, Bose et al. (2024) provides bound for the sum of the point-wiss errors in the transition dynamics averaged over the points in the source datasets. To bound the representational transfer error in the downstream target task, they introduce a notion of neighborhood occupancy density. Moreover, to connect the upstream tasks and the downstream target task, they make a pointwise linear span assumption from Agarwal et al. (2023). Finally, for the downstream target task, they only consider offline setting, whereas our primary focus and contribution is in the study of reward-free setting in the downstream task.
6 Conclusion
In this paper, we theoretically study multitask RL in the offline setting. We show that offline multitask representation learning is provably more sample efficient than learning each task individually. We further show the benefit of employing the learned representation from the upstream to learn a near- optimal policy of new downstream task, in reward-free, offline and online setting, that shares the same representation. We believe our work will open up many promising directions for future work, for example, studying the general function class representation learning in offline multitask setting.
10

---

Acknowledgments
R. Arora's and T. Sruyen-Tang's research was supported, in part, by the NSF CAREER award IIS- 1943251. M. Yin and M. Wang acknowledge the support by NSF IIS-2107304, NSF CPS-231 and ONR 1006977. The authors would like to thank Yingbin Liang and Yu-Xiang Wang for their helpful discussions.
References
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International Conference Machine Learning, pages 22-31. PMLR, 2017. (p. 19.)
Alekh Agarwal, Mikael Henaff, Sham Kakade, and Wen Sun. PC-PG: Policy cover directed explo- ration for provable policy gradient learning. Advances in Neural Information Processing Systems, 33:13399- 13412, 2020a. (pp. 3 and 5.)
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. FLAMBE: Structural com- plexity and representation learning of low rank MDPs. Advances in Neural Information Processing Systems, 133:0095--20111 2020b. (pp. 2, 3, 4, 5, 8, 9, 10, 19, 29, and 53.)
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. The Journal of Machine Learning Research, 22(1):4431 4506, 2021. (pp. and 5.)
Alekh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang. Mengdi Wang, and Xuezhou Zhang. Prov- able benefits of representational transfer in reinforcement learning. In The Thirty Sixth Annual Conference Learning Theory, pages 2114 2187. PMLR, 2023. (pp. 4, and 10.)
András Antos, Csaba Szepesvári, and Rémi Munos. Learning near-optimal policies with Bellman- residual minimization based fitted policy iteration and a single sample path. Machine Learning, 71:89-1299 129, 2008 (p. 3.)
Sanjeev Arora, Simon Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi. Provable representation learning for imitation learning via bi-level optimization. In International Conference on Machine Learning, pages 367-376. PMLR, 2020. (p. 1.)
Kai Arulkumaran, Dylan R Ashley, Jürgen Schmidhuber, and Rupesh K Srivastava. All you need is supervised learning: From imitation learning to meta-RL. with upside down RL. arXiv preprint arXiv: 1960, 2022. (p. 1.)
Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-drin driven offline reinforcement learning. In The Tenth International Conference on Learning Representations, 2022. (p. 9.)
Jonathan Baxter. A model of inductive bias learning. Journal of Artificial Intelligence Research, 12: 149-198, 2000. (p. 1.)
Avinandan Bose, Simon Shaolei Du, and Maryam Fazel. Offline multi-task transfer RL with representational penalization. arXiv preprint arXiv:2402.12570. 2024. (p. 10.)
Emma Brunskill and Lihong Li. Sample complexity of multi-task reinforcement learning. In Proceedings of the Twenty-Ninth Conference on Uncertainty in Artificial Intelligence, UAI 2013, Bellevue, WA, USA, August 11-15, 2013. AUAI Press, 2013. (p. 1.)
Daniele Calandriello, Alessandro Lazaric, and Marcello Restelli. Sparse multi-task reinforcement learning. Advances in Neural Information Processing Systems, 27, 2014. (p. 1.)
Rich Caruana. Multitask learning. Machine Learning, 28:41-75, 1997. (p. 1.)
Jonathan Chang, Masatoshi Uehara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift in imitation learning via offline data with partial coverage. Advances in Neural Information Processing Systems, 34:965-979, 2021. (p. 53.)
Jinglin Chen and Nan Jiang. Offline reinforcement learning under value and density-ratio realizability: the power of gaps. In Uncertainty in Artificial Intelligence, pages 378-388. PMLR, 2022. (p. 5.)
Jinglin Chen, Aditya Modi, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. On the statistical efficiency of reward-free exploration non-linear RI., Advances in Neural Information Processing Systems, 1-220000-209777 2022a. (p. 19.)

---

Xiaoyu Chen, Jiachen Hu, Lin F Yang, and Liwei Wang. Near-optimal reward free exploration for linear mixture MDPs with plug-in solver. In The Tenth International Conference on Learning Representations, 2022b. (p. 8.)
Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang, Provable benefit of multitask representation learning in reinforcement learning, Advances in Neural Information Processing Systems, 35:31741-31754, 2022. (pp. 2, 4, 5, 7, 9, 10, and 32.)
Yuan Cheng, Ruiquan Huang, Jing Yang, and Yingbin Liang. Improved sample complexity for reward- free reinforcement learning under low-rank MDPs. In The Eleventh International Conference on Learning Representations, 2023. (pp. 2, 8, and 19.)
Christoph Dann, Tor Lattimore, and Emma Brunskill. Unifying PAC and regret: Uniform PAC bounds for episodic reinforcement learning. Advances in Neural Information Processing Systems, 30, 2017. (p. 52.)
Leandro M De Lima and Renato A Krohling. Discovering an aid policy to minimize student evasion using offline reinforcement learning. In 2021 International Joint Conference on Neural Networks
(LJCNN), pages 1-8. IEEE, 2021. (p. 1.)
Carlo D' D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, Jan Peters, et al. Sharing knowledge in multi-task deep reinforcement learning. In 8th International Conference Learning Representations, 2020, Addis Ababa, Ethiopia, April 26-30, 2020, pages 1-11, 2020. (p. 1.)
Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal. Miroslav Dudik, and John Langford. Provably efficient RL with rich observations via latent state decoding. In International Conference on Machine Learning, pages 1665-1674, PMLR, 2019. (p. 10.)
Frederik Ebert, Yanlai Yang, Karl Schmeckpeper, Bernadette Bucher, Georgios Georgakis, Kostas Danilidiis, Chelsea Finn, and Sergey Levine. Bridge data: Boosting generalization of robotic skills with cross-domain datasets. Proceedings of Robotics: Science and System. 2022. (p. 2.)
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6, 2005. (pp. 9 and 19.)
Ben Eysenbach, Xinyang Geng, Sergey Levine, and Russ R Salakhutdinov. Rewriting history with inverse RL: Hindsight inference for policy improvement. Advances in Neural Information Processing Systems, 33: 14783-14795 2020. (p. 10.)
Kamyar Ghasemipour, Shixiang Shane Gu, and Ofir Nachum. Why so pessimistic? Estimating uncertainties for offline RL through ensembles, and why their independence matters. Advances in Neural Information Processing Systems, 1866-18818,,,, 2022. (p. 9.)
Arthur Guez, Robert D Vincent, Massimo Avoli, and Joelle Pineau. Adaptive treatment of epilepsy via mode reinforcement learning. In AAAI, volume 8, pages 1671 2008. (p. 1.)
Hao Hu, Yiqin Yang, Qianchuan Zhao, and Chongjie Zhang. The provable benefits of unsupervised data sharing for offline reinforcement learning. In The Eleventh International Conference on Learning Representationss 2023. (p. 10.)
Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear RL. In International Conference on Machine Learning, pages 4349-4358. PMLR, 2021. (pp. I and 3.)
Pihe Hu, Yu Chen, and Longbo Huang. Towards minimax optimal reward-free reinforcement learning in linear MDPs. In The Eleventh International Conference Learning Representations, 2022. (pp. 8 and 19.)
Haque Ishfaq, Qiwen Cui, Viet Nguyen, Alex Ayoub, Zhuoran Yang. Zhaoran Wang, Doina Precup, and Lin Yang. Randomized exploration in reinforcement learning with general value function approximation. In International Conference on Machine Learning, pages 4607-4616. PMLR,
2021. (p. 53.)
Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard. Way off- off-policy batch deep reinforcement learning of implicit human preferences in dialog. arXiv preprint 1rX 2019. (p. 1.)
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contex- tual decision processes with low Bellman rank are PAC-learnable. In International Conference on Machine Learning, pages 1704-1713. PMLR, 2017. (pp. 2 and 3.)
12

---

Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward- exploration for reinforcement learning. In International Conference on Machine Learning, pages 4870 1870-4879. PMLR, 2020a. (pp. 2, 4, and 19.)
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137-2143.
PMLR, 2020b. (pp. 3, 7, 8, 9, 20, 36, 40, 53, and 54.)
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline RL? In International Conference on Machine Learning, pages 5084-5096. PMLR, 2021. (pp. 9, 20, 43,
and 53.)
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning, pages 651 PMLR, 2018. (p. 1.)
Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic rein- forcement learning at scale. In Conference on Robot Learning. PMLR, 2021. (pp. 1 and 10.)
Emilie Kaufmann, Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Edouard Leurent, and Michal Valko. Adaptive reward-free exploration. In Algorithmic Learning Theory, pages 865-891. PMLR, 2021. (p. 19.)
Taylor W Killian, Haoran Zhang, Jayakumar Subramanian, Mehdi Fatemi, and Marzyeh Ghassemi. An empirical study of representation learning for reinforcement learning in healthcare. arXiv preprint arXiv:2011.1122: 2020. (p. 1.)
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. Advances Neural Information Processing Systems, 13:1179-119,, 2020.
(p. 9.)
Aviral Kumar, Anikait Singh, Frederik Ebert, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre-
training for robots: Leveraging diverse multitask data via offline reinforcement learning. In NeurlPS 2022 Foundation Models for Decision Making Workshop, 2022. (p. 2.)
Aviral Kumar, Anikait Singh, Frederik Ebert, Mitsuhiko Nakamoto, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre-training for robots: Offline RL enables learning new tasks from a handful of trials. Proceedings of Robotics: Science and System XIX, 2023. (p. 2.)
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the- pages 45-73. Springer, 2012. (p. 9.)
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv: 016433, 2020. (pp. and 9.)
Alexander Li, Lerrel Pinto, and Pieter Abbeel. Generalized hindsight for reinforcement learning. Advances in Neural Information Processing Systems, 33:7754-7767 2020. (p. 10.)
Qinjie Lin, Han Liu, and Biswa Sengupta. Switch trajectory transformer with distributional value approximation for multi-task reinforcement learning. arXiv preprint arXiv:2203.07413 2022. (p. 2.)
Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch off- off-policy reinforcement learning without great exploration. Advances in Neural Information Processing Systems, 33:1264412220 1274, 2020. (p. 9.)
Rui Lu, Gao Huang, and Simon S Du On the power of multitask representation learning in lincar MDP. arXiv preprint arXiv:2106. 100.0003,, 2021. (p. 1.)
Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Emilie Kaufmann, Edouard Leurent, and Michal Valko. Fast active learning for pure exploration in reinforcement learning. In Interna-
tional Conference on Machine Learning, pages 7599-7608. PMLR, 2021. (p. 19.)
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstrac- tion and provably efficient rich-observation reinforcement learning. In International Conference on Machine Learning. pages 6961-6971 PMLR, 2020. (p. 10.)
Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pages 7780-7791. PMLR, 2021. (p. 10.)
13

---

Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. Model-free representation learning and exploration in low-rank MDPs. Journal of Machine Learning Research, 25(6):1-76, 2024. (pp. 8, 9, and 19.)
Robert Müller and Aldo Pacchiano. Meta learning MDPs with linear transition models. In Interna- tional Conference on Artificial Intelligence and Statistics, pages 5928-5948. PMLR, 2022. (p- 1.)
Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research, 9(5), 2008. (p. 3.)
Thanh Nguyen-Tang and Raman Arora. VIPer: Provably efficient algorithm for offline RL with neural function approximation. In The Eleventh International Conference on Learning Representations, 2023. URL https: /ttp///eppeeeeeeee//e/ee/e wwt/ peet/forummid--Woooe oroporrrddd-00000TTC000..... (p. 9.)
Thanh gguyen-Tang and Raman Arora. On the statistical complexity of offiine decision-making. In Forty-first International Conference on Machine Learning, 2024a. URL https://
(p. 9.)
Thanh Nguyen-Tang and Raman Arora. On sample-efficient offline reinforcement learning: Data diversity, posterior sampling and beyond. Advances in Neural Information Processing Systems, 36, 2024b. (p. 9.)
Thanh Nguyen- Tang, Sunil Gupta, Hung Tran-The, and Svetha Venkatesh. On sample complexity of offline reinforcement learning with deep reLU networks in besov spaces. Transactions on Machine Learning Revearch, 2022. ISSN 2835-8856. URL id=LdEmOumNeev. (p. 9.)
Thanh Ngugen-Tang, Ming Yin, Sunil Gupta, Svetha Venkatesh, and Raman Arora. On instance- dependent bounds for offline reinforcement learning with linear function approximation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 9310-9318, 2023. (p. 9.)
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. In Learning for Dynamics and Control, pages 1154-1168. PMLR, 2021. (p. 1.)
Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline rein- forcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems, 34:11772--171 2021. (p. 9.)
Siddharth Reddy, Anca D Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. In The International Conference on Learning Representations, 2019.
(p. 10.)
Shideh Rezaeifar, Robert Dadashi, Nino Vieillard, Léonard Hussenot, Olivier Bachem, Olivier Pietquin, and Matthieu Geist. Offline reinforcement learning as anti-exploration. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8106 81 $114, 2022. (p. 9.)
Martin Riedmiller. Neural fitted Q iteration first experiences with a data efficient neural reinforcement learning method. In Machine Learning: ECML 2005: 16th European Conference on Machine Learning, Porto, Portugal, October 3-7, 2005. Proceedings 16, pages 317-328. Springer, 2005. (p. 9.)
Tyler Sam, Yudong Chen, and Christina Lee Yu. The limits of transfer reinforcement learning with latent low-rank structure, 2024. URL https: http://xxxvvvrrree wrgg/bss//21ar1 21601. (p. 2.)
Susan M Shortreed, Eric Laber, Daniel J Lizotte, T Scott Stroup, Joelle Pincau, and Susan A Murphy, Informing sequential clinical decision-making through reinforcement learning; empirical study, Machine learning, 1099,)) 2011. (p. 1.)
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pages 9767-9779. PMLR, 2021. (p. 1.)
Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. 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, 2019. (p. 2.)
14

---

Yee Tch, Victor Bapst, Wojciech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. Advances in Neural Information Processing Systems, 30, 2017. (p. 1.)
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. International Conference on Learning Representations, 2021. (pp. and 5.)
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. In International Conference on Learning Representations, 2022. (p. 9.)
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline RL in low-rank MDPs. In International Conference on Learning Representations, 2022. (pp. 3, 5, 9, 10, 31, 43, and 48.)
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018, (p. 53.)
Andrew Wagenmaker, Yifang Chen, Max Simchowitz, Simon Du, and Kevin Jamieson. Reward- free RL is no harder than reward-aware RL in linear Markov decision processes. In International Conference on Machine Learning, pages 22430-22456. PMLR, 2022. (p. 19.)
Lu Wang, Wei Zhang, Xiaofeng He, and Hongyuan Zha. Supervised reinforcement learning with recurrent neural network for dynamic treatment recommendation. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pages 2447-2456 2018. (p. I.)
Ruosong Wang, Simon S Du, Lin Yang, and Russ R Salakhutdinov. On reward-free. reinforcement learning with linear function approximation. Advances in Neural Information Processing Systems, 33:17816-11 17726667726,, 2020. (pp. 2, 8, and 19.)
Ruosong Wang, Dean P. Foster, and Sham M. Kakade. What the statistical limits of offline RL with linear function approximation? In 9th International Conference on Learning Representations, ICLR 2021, Austria, May 3-7, 2021, 2021. (p. 9.)
Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang, Ruslan Salakhutdinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. In International Conference on Machine Learning. PMLR, 2021. (p. 9.)
Annie Xie and Chelsea Finn. Lifelong robotic reinforcement learning by retaining experiences. In Conference on Lifelong Learning Agents, pages 838-855. PMLR, 2022. (p. 1.)
Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistent pessimism for offline reinforcement learning. Advances in Neural Information Processing Systems, 34:6683-6694, 2021a. (p. 9.)
Tengyang Xie, Nan Jiang, Huan Wang, Caiming Xiong, and Yu Bai. Policy finetuning: Bridging sample-efficient offline and online reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021b. (p. 9.)
Yifan Xu, Nicklas Hansen, Zirui Wang, Yung-Chich Chan, Hao Su, and Zhuowen Tu. On the feasibility of cross-task transfer with model-based reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023, (p. 10.)
Yuling Yan, Gen Li, Yuxin Chen, and Jianqing Fan. The efficacy of pessimism in asynchronous Q-learning. JEEE Transactions on Information Theory, 2023. (p. 9.)
Ming Yin and Yu-Xiang Wang. Towards instance-optimal offline reinforcement learning with pessimism. Advances in Neural Information Processing Systems, 34:4065 4078, 2021. (p. 9.)
Ming Yin, Yu Bai, and Yu-Xiang Wang. Near-optimal provable uniform convergence in offline policy evaluation for reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pages 1567-1575. PMLR, 2021. (pp. 7 and 9.)
Ming Yin, Yaqi Duan, Mengdi Wang, and Yu-Xiang Wang. Near-optimal offline reinforcement learning with linear representation: Leveraging variance information with pessimism. International Conference on Learning Representations, 2022. (p.9, 9.)
Minjong Yoo, Sangwoo Cho, and Honguk Woo. Skills regularized task decomposition for multi- task offline reinforcement learning. Advances Neural Information Processing Systems, 35: 37432-37444 2022. (pp. 2 and 10.)
15

---

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 13:114129114143 2020. (p. 9.)
Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Sergey Levine, and Chelsea Finn. Conservative data sharing for multi-task offline reinforcement learning. Advances in Neural Information Processing Systems, 34:11501-11516, 2021. (pp. 1 and 10.)
Taahe Yu. Aviral Kumar, Yevgen Chebotar, Karol Hausman, Chelsea Finn, and Sergey Levine.
How to leverage unlabeled data in offline reinforcement learning. In International Conference Machine Learning, pages 2561 1-25635. PMLR, 2022. (p. 1.)
Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Provably efficient reward-agnostic navigation with linear value iteration. Advances in Neural Information Processing
Systems, 13:11756 2020. (p. 19.)
Andrea Zanette, Ching-An Cheng, and Alekh Agarwal. Cautiously optimistic policy optimization and exploration with linear function approximation. In Conference on Learning Theory, pages 4473 4425-4525.
4525. PMLR, 2021a. (p. 32.)
Andrea Zanette, Martin Wainwright, and Emma Brunskill. Provable benefits of actor-critic methods for offline reinforcement learning. Advances in Neural Information Processing Systems, 34:
13626-13640, 2021b. (p. 9.)
Xianyuan Zhan, Haoran Xu, Yue Zhang, Xiangyu Zhu, Honglei Yin, and Yu Zheng. Deepthermal: Combustion optimization for thermal power generating units using offline reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelligence, volume 36, pages 4680-4688
2022. (p. 1.)
Weitong Zhang, Dongruo Zhou, and Quanquan Gu. Reward-free model-based reinforcement learning with linear function approximation. Advances in Neural Information Processing Systems, 34, 2021.
(pp. and 19.)
Xuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun. Efficient reinforcement learning in block MDPs: A model-free representation learning approach. In International Conference on Machine Learning, pages 16-27-26 PMLR, 2022. (p. 10.)
16

---

Contents
1
Introduction
2
Preliminary
2.1
Offline Multitask RL with Downstream Learning
3
Upstream Offline Multitask Representation Learning
3.1
Algorithm Design
3.2
Theoretical Result on Upstream Task
4
Downstream RL: Reward-free Exploration, Offline RL and Online RL
4.1
Relationship between upstream and downstream MDPs
4.2
Downstream Reward- RL
4.3
Downstream Offline and Online RL
5
Related Work
6
Conclusion
10
Additional Related Work
19
Omitted Algorithms
19
B.1 Algorithms for Downstream Reward-Free RI.
19
B.2 Algorithm for Downstream Offline RL
20
B.3 Algorithm for Downstream Online RL
20
C Notations
21
D Proof of Multitask Offline Representation Learning
22
D.I Supporting Lemmas
22
D.2 Proof of Theorem 3.3
22
Proof of Supporting Lemmas Appendix D
27
E.1 Proof of Lemma D.I
................................
27
E.2 Proof of Lemma D.2
28
Multitask Offliine MLE
29
One-Step Back Lemma and Concentration of Penalty Term
31
G.1 One-step back lemma
G.2 Concentration of penalty term
32
H Proof of Lemma 4.3: Approximate Feature for New Task
32
17

---

I Proof for Downstream Reward-Free RL
33
1.1 Supporting Lemmas
34
1.2 Proof of Theorem 4.4
35
J Proof of Supporting Lemmas in Appendix
36
J.1 Proof of Lemma 1.2
36
J.2 Proof of Lemma 1.3
36
J.3 Proof of Lemma 1.4
37
J.4 Proof of Lemma 1.5
38
J.5 Proof of Lemma 1.6
40
J.6 Proof of Lemma 1.7
41
K
Proof for Downstream Offline RL
42
K.I Supporting Lemmas
42
K.2 Proof of Theorem 4.6
43
Proof of Supporting Lemmas in Appendix K
44
L.1 Proof of Lemma K.1
44
L.2 Proof of Lemma K.2
44
L.3 Proof of Lemma K.3
45
L.4 Proof of Lemma K.5
46
M Proof for Downstream Online RL
47
M.1 Supporting Lemmas
47
M.2 Proof of Theorem 4.7
48
N Proof of Supporting Lemmas in Appendix M
49
N.I Proof of Lemma M.2
N.2 Proof of Lemma M.3
51
N.3 Proof of Lemma M.4
52
Auxiliary Lemmas
52
0.1 Miscellaneous Lemmas
52
0.2 Inequalities for summations..
53
0.3 Covering numbers and self-normalized processes
18

---

A Additional Related Work
Reward-Free Exploration (RFE). In reward-free RL setting, the agent does have access to reward function during exploration phase. However, the agent must propose near-optimal policy for an arbitrary reward function revealed only after the initial exploration phase. This setting is particularly relevant when there are multiple reward functions of interest (Achiam et al., 2017) or the batch RL setting (Ernst et al., 2005). In recent years, reward-free RL has been extensively studied in both tabular (Jin et al., 2020a; Kaufmann al., 2021; Ménard et al., 2021) and linear function approximation (Wang et al., 2020; Zanette et al., 2020; Zhang et al., 2021; Hu al., 2022; Wagenmaker et al., 2022) settings. Agarwal al. (2020b); Modi al. (2024); Chen et al. (2022a); Cheng et al. (2023) study reward-firee RL in low-rank MDPs which is particularly interesting as here representation learning is interwined with reward-free exploration.
B Omitted Algorithms
B.1 Algorithms for Downstream Reward-Free RL
Algorithm Downstream Reward-Frees Exploration: Exploration Phase
1: Input: Feature G. Failure probability >> 0 and target accuracy & > 0
2: KREE ск- d H log b
for some ск > 0
3: B x- Cr, Hd dH /
max(Caown: 1)/5) + Edown v KRKE
4: for mork=,.... KRFE do
5:
Qir+1(1.')
0, V
0
6:
for h = H, H
do
7:
A/ = Σー= (5 (s TT + Ia
( + 3v o(. )T (1)()))
8:
9:
Define the exploration-driver reward function rf(;))) (.
o (A -1 Σ on (8% a V
11:
(+,)) min m (. TT + + (,),),
(-) = max,e.A 2 (-
13:
(.) argmax, A (,,a)
end for
15: Receive initial state s
for h I...... do
17: Take action o n (s and observe s4 p(*,T+1)
(st at
18: end for
19:
end for
Output: DRFE {( (RA)))IIIX
19

---

Algorithm 3 Downstream Reward-Free Exploration: Planning Phase
1: Input: Feature o, dataset DRFE {(s%, (e (KKr))))e reward reward functions fri (re)
2: Initialization: 20+1('+' + and V+++(.) t
3: for h H, H ,,..., do
4: An KREE ( o (5h a7) La
5: Let (n.))) t min BV 06 (An)
Th + An Exer
KEEE ( of of (1+2(5411)
6:
(n(:)) f min (in (:.) T + ra(ls)) + un(-)), H and Pa(+) + maxoed a)
7:
8:
(.) f argmax ax(.aa)
9: end for
10: Output: ã = (
B.2 Algorithm for Downstream Offline RL
Algorithm 4 Pessimistic Value Iteration (PEVI) with Approximate Feature (Jin et al., 2021)
1: Input: Feature o, dataset Doown {s5 dh The sh+1) (ar)))reeeee re (N())),,e)() parameters Ad, 8, Edown
2: Initialization: V
3: for h H, H -,...,,1 do
Natl (s/ at (2 (8), af) AaTa 4: A₃ EmT
5: T6 An (Ete (5 0 (ri, + (+++())111))
6:
Ta(::)) = HEdwwn B|00, A (
7:
Qn(:,)) mintØn( (. T wh - -Tn(s,) h 13+
8:
(..)) argmax,, m1 (1)())); (7(1)))mm((1))))
9:
V.(.) (,, x- (-1-))a
end for
11:
Output: TTh 1
B.3 Algorithm for Downstream Online RL
Algorithm 5 LSVI-UCB with Approximate Feature (Jin et al., 2020b)
1: Input: Feature 6, parameters Ad, Bng Edown
2: for 1, ... N.. do
3: Receive the initial state s = 81 1,..., do
4: for h H, H
5:
AA Σー= (3 o (s TF + Aala
6:
Wi = (AR) ΣT= on (n a (ra a7) + V (8h+1))
7:
QR(..) = min(on(,)) T
oh + Bn (86. ))(ag))), H 13+
8:
(-)()) max.e.4 (.)aa
9:
end for
nH argmax,e A Q (.,a)
11:
for h = 1...., H do
Take action ai. n (s%) and observe s
13:
end for
end for
15:
Output: n. ..... No. and मे where ã is the uniform mixture of x1 ..... Non
20

---

C
Notations
We summarize frequently used notations in the following list.
f (s
P. - 220 nH/82
(-1s,)))2
Sn
E ((P...))(P))000)))
Sh (s,
(1)
</
E ((P.0)))))))
A (81,
(0)
V (nwch
(1) Ad Adn+ n
(1)
O'
a
/2nwin +Ad
n (m).a)) ())))) ldp I + A
5(0)
-i=1 @n (s ah
(f.t) (1it)
)on (on(s, (5/11)) a
y + AI
2 an)
min
allon (5h, an) 15:00
max (a/m((aa))))))))))))))))({)))(
w
max, We
Edown
++ C.Cru
27 1og 1 nH/6
TC (.)
Σi (.)
4%
(s
For any h E [H], we define
pl p (-1s, =
(.)),
Pn(.)s,)) = (2, a), 77 (-))
Given a reward function r, for any function f : S + R and E [H] we define the transition operators and their corresoondigg Bellman operators follows
(P(+ (s, (0)(s, (T.TT)(()))))) (..T))))))).)))))))))))))))) (s)))f(s))d''
a) + (p(t
(Pa))ss,a)= = / (6n(s, a) (1) (Brf)(s, = (s(s a) + (P
21

---

D Proof of Multitask Offline Representation Learning
We first state the supporting lemmas that are used in the proof of Theorem 3.3. The proof of these lemmas are provided in Appendix E.
D.1 Supporting Lemmas
In the following lemma, we provide an upper bound for the model estimation error for each task that captures the advantage of joint MLE model estimation over single-task learning.
Lemma D.I. For any task I, policy nu, and reward rt, we have for all 2 2,
E (o),,)))(BB)), (sRr an) & (e00-1.00-1>>PPPP.
(1)
1 (8n-11 (o
and 1, we have
Burnnt [1 (s) s Vwrsi"
where wt max
In the lemma, prove that TT pt, an almost pessimistic estimator of V pis in the
average sense.
Lemma D.2 (Restatement of Lemma 3.5). For any policy and reward r", we have, with probubility
1 1
2 2ogg 22/ TH//
V (t Bi Vpcr.), ,+
V
H H.W/w/// where Cn
71
D.2 Proof of Theorem 3.3
We first restate Theorem 3.3.
Theorem D.3 (Restatement of Theorem 3.3), Under Assumption 3.1, with probability at least for any step h [H), we have
[|| 1 (-1sh, an) (.
sh, (-1shhann) [
2 2ogg
E (gg))))
nT
(D.1)
where P(),)..... ((T) be the output of Algorithm 1.
In addition. in Algorithu if we set a with So 2 2lug(2]0) 2) nHH6) and being constant where we assume that max, max,, (L/T|(a|) (s)) oo, then under Assumption 3.1, with probability at least ô, we have
V2nwGn Xd, X
ed ed ( | /8)
1 T
Σ
Vū vpth.o,, S wodH V - V
XC*
+WH
dC T
d tag
- +2H V
WCn
T
n
T
(D.2)
where ( is the output of Algorithm 1.
Proof of Theorem D.3. As Lemma D.2, we condition on the events:
T
where Go = 2 2 T nH/8)
E ( (.
(s, a)2 s Sn.
(D.3)
and
VΦ E : (on = -111
(D.4)
22

---

From Lemma F.3 and Lemma G.2, this event happens with probability - 6. Conditioning on this event, we have
| (-1sh:
(lshaan)) an) an) [Tv]
E
(n on (P()
T
T
E
S/ an )-p0))(PPp00000
(
1
T
[1h (ss.aa)
T
T
(9)
T
(s )~ (s)
T
E
(H6)
nT
where (i) follows from Cauchy-Schwarz inequality, (fi) follows from Jensen's 's inequality and (fii follows from (D.3). This completes the first part of the proof.
Conditioning on the event in (D.3) and (D.4), for any set of policies Ti..... TT, we have
v pt
V pr
Vi+
+ V
Vū. p(
(
T
V
+ Rr
p()) - V
pi
(8)
B ((PC..)))))))
(6t (sh, and (P/(tt)) p((t) V
V
b00 (sh, an)
1=1
V 1p
(80)
T
H
T H
E (ssta)))(P m)) (sh. (s2,
+ H
Σ
E
(s (P(s))))
[1h an)
1=1
h=1
t=
h=1
H WTTG.
(D.5)
where (i) follows from the observation that 7t is the argmax over all Markovian policies as well as all
policies for (ot), (ii) follows from the simulation lemma, Lemma 0.1, (tii
follows from the observation that m. A & H and Lemma D.2.
Now, using Lemma G.1 (with setting P= p(+.t) and Φ³) and noting that (8 1, for h 2 we have
8% (sh, an)
B(ex.0n)
| (sh-1, ah -1X
< E ((en-nos->>>pP) m
V /nwk (ocpoo
B (s.a) +>d
(D.6)
23

---

From (D.4), we have
nB, (e (s)
(n min a² (ol an)llss p }
S (P00)))))PPPP00000
a 22l1on (sn, an) lis
S 2
(on (xPPr (ono% + -11
2a² Tr
कo XT) (s..))))) (1 + -13-11
= aa2
= add
(D.7)
Next, we upper bound Eee oh- ()---p/t)
107h1(0h-2 an -)ls)) a )-1
as the follow-
ing
llo 1 1000-.0
an -))l 17
C. (t (2 | ( a. -
(10)
d 7b)
n
dC"
(D.8)
n1
where (1) follows from Lemma O.2 and (ii) follows from similar steps as in (D.7).
Combining (D.6), (D.7) and (D.8), get
24

---

T
B(0) (sh: an)
E (n
T
E (sn 1)
(9-- an-
R.wE (on (ara) +Ad
Eisn (=(
|10%_ (8h-1, )-1 x
(1 (s, a)] + dd
(8h- 1)-
dC
(oB (1 (s, a)) Ad
R (100-1000)))000000...))
n2
dC
Blsnnan 3// (s,a)) + +TVAd
nWB (sn- 150 (m(()))))
d
n B (s (s, a)] +TVAd
A0* (oo [ (s. (5. a)] + +Tdg
wv
n
AC ( (s.a) + +dd )-P(..)))
S dv
T
Am1
Ta
XC*
SwvdC.T
Tdy
n
C*
xC*
=waTdy
+ Tdy
(D.9)
n
n
Following similar steps as in (D.7), we can further show that
8 (s1, an)
8800 (sL, on)
$ say d
25

---

Now noting 100 S 1, for h 2 we get
E (1s.0n))(P(00000)) (Sh
|10% 190-1 imf 1X
E3 (B 10
1 (on +Xd
T
B 110011000))))0000.00000
(9h-,, 2np x
(o 1h (Sh, (oh, + Ad
E (op laak (P(())))))))))
do
E ((o..o))))())) (ons + Ad nwBB aaa (x(p)))))
dC
nwlE (o 1000000)0000
( ax))) vAd
n
dC
2 (pe.. (11) (sn, ande] + +dd
A0
n
n
AC E ((pt..0)) [ (sh, an) + +Tdy
wv d0
t=1
SwvdC=
T
(en (pt(n)))) (sn, + +Tdy
AC*
n
AC*
SwVdCOTG. + Tdl.
(D.10)
n
Further, note that,
T
[+ ( (1)
E to ((Pt)))))))))
WI (o) (.B0,,.a)))ppp,,,pp)) (s))
1=1
V
((00)
$ /TT...
where the last inequality follows from the Cauchy- Schhhy- incquality and (D.3).
26

---

Finally, from (D.5) we get
T
pl
T
(sh, an) H
[Th (sh.
1=1
1=1 A=1
H / wTG
H
110 (sn, B )(phi)00)))))
[1 (on,
H
B ((Ph0.)))))
2= t=1
h=2 r=
1
T
T
E6 ((
(1n, +
18
( (s1, + 11 vWTG..
C
AC +HHWVdC.T,, H²Td
AC
+aTI /= +2H VWTG.
HwaTdy
+ +HTd C² C*+
n
n
&HwaTdy
=22112Td1
XCC*
d +HH2 + +211 v wTTG.
(D.11)
n
n
So, have
1
C* S Hweady _
AC*
dC
WCn
Vil pi.
VB: p(
2H2 dy
HPW
an
2H /
T
T
T
(D.12)
E Proof of Supporting Lemmas in Appendix
In this section, we provide the proofs of the lemmas that used in the proof of Theorem D.3.
E.1 Proof of Lemma D.1
Proof of Lemma D.I. For h = 1,
(11,))))((L),,)) (1(0 (s). s
E ( and (()
wrB ( [1 (11
27

---

where the first inequality follows from Jensen's 's inequality and the second inequality follows from importance sampling. Denoting Sh = E (os.aa))(PP.00.
[1% (s, a)² for 2 2, we have
E B,n-(p(1).s [h (sn, an)
%A m
| (5h-1,
E (.h-
nwr (o 1 (sn.an))) - Si0 ((n--1 (
(2))
E
(bh an (P((),
nuussh + Ad nc (1)
n
1 (8 18h-
(
E 11)-
On
(1) 1
where (1) follows from Lemma G.1 and 15% (sh, an) S 1, (it) uses notations defined in Appendix C.
E.2 Proof of Lemma D.2
Proof of Lemma D.2. We condition on the events:
2 [V (
E (n ((Pl..))))) (s, s Sa. where Gn
(E.1)
and
Vp E +: [1) = (16n(())))l
(E.2)
From Lemma F.3 and Lemma G.2, this event happens with probability - 8. Conditioning on this event, we have
(6) a
= (wis Xd + nch-1
(t)
(1)
2rwwnn + Ad
(E.3)
28

---

We have
T
V
Vp pl
pe Bie)
T
H
[-0%]] (sh: an) (P/h -P/-)) V po.,, (Sh, an)
E (
Am
(8))
(sk,
(sh, an) + f/ (sh, an)
(eh (ehh . (P(t),
1=1
((0)
T
H
(sn 1.0h ()(
min
Oh | j 18h 9n-)) (
1=1
h=2
H
[ (sh.
H
H
E
(on (.t)).))PP))..
A=1
(80)
T
H
(1)
E
(Na-1 (P(0),)
min
ag
1=1 h=
T H1
+ H WTT + H
E
1b/ (sh, (sh,
A=1
E
(oh aa ()W)))
min on 10hh1 an-1) 1
(=1 h=2
T
H
min a ( and an (S
+ H WVWTC +
E
(08)
H HWWT + H
E (s (P(00,
min
I H. .HVWTG....
where (i) follows from Lemma O.1, (ii) follows from the observation V pl S H, (iii follows
from Lemma D.I, (iv) follows from Cauchy- Sehhwa- inequality and the fact that Σ Sh Gn. (v) follows from (F.2).
(0)
F Multitask Offline MLEE
Consider sequential conditional probability estimation setting with an instance space A' and target space y and with a conditional probability density /=y fz y). We consider function class F : (X' x }) + RR for modeling the condition distribution f", and we further assume that the realizability condition holds i.e. f* E F. We are given a dataset == (( y: ] where J, DD = D1(11=11, V1-i-1) and y/s p(- v..). Let D' denote a tangent sequence {(r,, v/) ) where x ~ Di
and y p(- lx.). Note that here x depends on the original sequence, and so the tangent sequence is independent conditional on D.
Lemma F.1 (Lemma 24 of Agarwal et al. (2020b)). Let D be datasel of examples, and let D' be a langent sequence. Let T(J, D) E e(S, (1;, y1)) be any function that decomposes additively across examples where E is any function, and let f(D) be any estimator taking as input random variable D and with range F. Then
exp (L(F(D), D) - log EDr (exp(L(((D), log Fi)] 1
ED
Lemma F2 (Lemma 25 of Agarwal et al. (2020b)). For any two conditional probability densities fi, f2 and any distribution D E ((X) we have
E Ex~Dllfi :) - -1l S -2 De- exp (-1 (ogg(f((..w)//(((..))))
29

---

Lemma F.3 (Multitask offline MLE guarantee). Given E (0, 1), consider the transition kernels
learned in Algorithm For any n, with probability at least 1 5/2. we have
B a.....)) [12] (s, a)²] SCn
2 2
v T nH
where Gn
(F.1)
Proof of Lemma F.3. Let f(D) denote empirical maximum likelihood estimator:
F(D)
argmax
Σ
log f y6)
JEF
(R..B))DD
We combine Lemma F.1 with the Chernoff method obtain the following exponential tail bound: with probability 1 - 6, we have
llg D (exp(L(F())) -(P())
[F| log(L/0).
(F22)
Now, we set L(S, Xi - (z; f(r, Mk)) where D = {a;, yi 3= is dataset and D' (1l, 14) is tangent sequence. In the multitask offline RI. setting, let = ((/, )
1 and [(x, v) = II (1 h+ 1 1 d.)) Then, the dataset D,, can be decomposed into Dj = -T_,, where D = ((5 (6.) (8.)) (6,)) (6.2) 1
ag $b411 (iej)- Similarly, D/, = (D() and
Thus, in the multitask offline RL setting, we have the cardinality | | With this choice, the right hand side of (F.2) is
(Ph)+ == (D)) of 11
Σ log(f"(x; (2; 9/)/ (z) v:)) log FF log(1/6) S log |F| + log(1/0) log( |og((
(F.)
where the first inequality follows because î is the empirical maximum likelihood estimator and the realizability assumption. The equality follows because |F] hand side of (F.2) is
| On the other hand, the left
n
(ddr 211
V1) ))
log EDy
exp
log
DT.
n
pl
(i.1)
(1A) (i.t)
log Epg
exp
2
log
II
p(0)
hti
(i.1)
(1,4)
(1,4)
Dn
(1.0)
log E ((
exp
log
P
(2.0)
Dp
i=1
p(+.)) (1.1) (3.0 (5.0)
(132)
1
h h+t
ah
log (D(
exp
log
p(0) ($ 18h
(1.2) (1.0)
(..))
a
(Kt)
T
(1..h.0000 and - p(+.t(). sh, an) 2 1
E eh (en
(. sh,
2
1=1
B ) (s,
(F.4)
where (i) follows from the above definition of f(z, y), (ii) follows because the data of T' tasks are independent conditional on Da, (iii) follows because P(() is independent of the dataset (D))(*) and from the definition D/,- (iv) follows from Lemma F.2, and (v) follows because task T, the data collected using behavior policy T..
Combining (F2), (F.3), (F.4), we get
T
7 2
(0s,.an)(((-00)))) a}]] < log(L
(F.)
30

---

Using union bound, we obtain that for any h E [H] and with probability at least 6/2, it holds that
T R (r (s. (x, 1
2 log( [ nH/6
This completes the proof.
G
One-Step Back Lemma and Concentration of Penalty Term
G.1 One-step back lemma
The following one-step back lemma is key technical lemma for proof. One-step back lemma for offline setting was first introduced in Uehara et al. (2022) for infinite-horizon sta- tionary MDP. Our lemma extends their result finite-horizon non-stationary MDP for offline setting. For any function ! E S x A + R. policy т and transition kernel P, the lemma shows that we can relate the expected value ((s)))n) az)] to the potential function #4-1)ll(,, 1xp..)
(x) || 11
Lemma G.1 (One-step back inequality for non-stationary finite-horizon MDP in offline setting). For each task E (T). let P (P(0, plx00) with embedding (Ø, d²) } and p be MDP model, (nn.s...n))((pt... x1)00
and E
XI be the covariance matrix following the behavior policy n. under the true environment P(+.)) Denote the total variation distance between p(r.) and P time step by (Sh, an). Take any g S x A such that l1glloc S B. Then, letting max,, (1/ (a [))ss]](())()))))) for all h 2, and for any policy have
(P on)l S Been 1en--1nn----// [1 (sn 1, au 11-1)1l ing 3-1
(o y (ne . (on mn)) AdB3 4nB4
(n(p(tit) ( (f 1:0h 13a
Proof of Lemma G.I. First, have
Bier ((e))a)) an))
Egsn an (T(a (30 an 1), phh (8n dsa
/
<Ben-110s->><PPm 1onn-1 an-1))l
gssn, an an))T(an st)an (on)dss)
where the inequality follows from Cauchy-Schware inequality.
31

---

Then.
[ g(sn, ( (s))un Asa
{ V g(sn: an)t(an 1n))n-1(sa)) )dsh (ddd { nB - 100 [00]] 1 1+
{
{
/
glsn, an (aan sn)uu (sp Jdsn
4/
(
27
Ea(s (a) sn)mn-1(SN) (00nnn((00(((0n01111//00 ah-h-1; E/Zo
n.
BB
E Pea-P(lon (g(sn an) + B²Ad
(5t)
S 1p )-
[g(sh, ah + A32Ad + nB² + // 14
(181 I nwlB (s1.0)))
[g(sh, a
+ B²Ad nB² B Boo (0
where (i) follows from the assumption l| I B and for any function S - [0. 1]. I pn (s))(s)dssl I vd, (ii) follows from the definition of (' as) which the total vari- ation distance between p and P at time step h, and finally (iéi) follows from importance sampling. This completes the proof.
G.2 Concentration of penalty term
Recall that 2 = n. (as, (a(tp)))))PP))ppp.... xx?) 100 ] >T. Thus, 5:(0) is equal to " in expectation. We now provide important lemma to ensure the concentration of the penalty term. The version for fixed ø is proved Zanette et al. (2021a). Here, we take a union bound over the whole feature & E & number of total tasks T. horizon H and cardinality of each offline dataset from individual tasks.
Lemma G.2 (Concentration of the penalty term). Fix ô E (0, 1) and set À = (nn for any n. With probability at least 1 1- 6/2, we have that any N, [H], t E [T] and ø E 4
B1 Il (8, a) 1 (o < (1 (ss)) S Ba 1lob (8, a) (1)
where 3n and 32 are some absolute constants.
H Proof of Lemma 4.3: Approximate Feature for New Task
We first restate Lemma 4.3.
Lemma H.1. Under Assumption 4.1, the output é of Algorithm 1 is approxpipo feature for MDP MT+1 where Edown E C.C.w, { (2I [] nH/8) i.e. there exist a time-dependent unknown (signed) measure ço- over such that for any (s, E A, we have
())ss. S Edown-
Furthermore, for any g S + [0, 11. we have [l T4 (s)g(s)dsll < Cr.vd.
The following proof is motivated from the proof of Lemma I in Cheng et al. (2022).
32

---

Proof of Lemma H.1. For all (s, ES S x h [H] and for any t E [T] we have
10 (1s, P (-1s,))lTVv (1s,,))
(-1s, (-1s,
max =EA
C (L (/(0) (118h; (+ lnl
(1) CRU
T
(Pt.. (1/ (.18n; un) (. (1 an)llTv
K
(10) CRU
27 (
(H.1)
n where (1), (ii) follows from Assumption 4.1 and (iii) follows from Theorem 3.3.
Defining ()) Ei= (t). have
ppt h (.
(5,(s, a), Σ aa)l
= -11p(07t1)(1ss0)
Σ
( (-Is
T
< (-Is,a)
Σ
(-1s, alllv + p(++t) h
Σ cill ( - p (-|s, 1T
55+ CLCRY
2T
where (i) follows from Assumption 4.1, (H.1) 1) and the fact that E [0, Cz. for all t E [T]. Moreover,
by normalization, for any g S > [0, 1]. we get
|/ p [mmolssssds
|/
Σ
Mh 11)
(s)g(s)ds
2
I CiVd,
where the last inequality follows from Assumption 3.1.
I Proof for Downstream Reward-Free RL
For any h E [H], we define
pl 1(+ (-1s,)))))))))) = (0i a), () (,.T+1())),
Pa(.1s, = (on(s, a), (7))
Given reward function (as is provided in the planning phase of reward-free RL setting), for any function f S > IR and h E [H], we define the transition operators as follows
(1(+, (s, - - (+ (4) (+T+1)
(Pnf)(s,a.)) - (all:saar)-/(cttt ((1
33

---

When no reward function is provided as is the case in the exploration phase of reward-frre RL setting, we simply omit r from the above operator notation.
In this section for notational simplicity we denote V pt, (s) and Q P.. n,r (s, a) by rs, and Q% (s, a.r) respectively where is reward function provided in the planning phase of downstream reward-free free RL task. We similarly denote the optimal value function and action-valus function under reward function as VE (s, r) and Q a, respectively.
We also introduce the truncated optimal value function Vi (s, in the planning phase, which is recursively defined from step H + to step 1. Compared to the definition of standard optimal value function V (s, r). the main difference is that we take minimization over the value function and H in each step in this definition. We provide the formal definition as follows.
Definition 1.1 (Truncated Optimal Value Function).
introduce the truncated optimal value function Vi rs which recursively defined from step H + to step 1:
(8, 0, As ES
(it1(s, v(s,a) (%(s,,, = - a) + PA p( V
(8, = min { max (r(()))) + p (ir(,,,) H). As S aEA
We can similarly define V. (s, and (s,
I.1 Supporting Lemmas
Now we state the supporting lemmas that are used in the proof of Theorem 4.4. The proof of these lemmas are provided Appendix J.
The following lemma shows that the linear weight wf in Algorithm 2 bounded.
Lemma 1.2 (Bounds on Weights in Algorithm 2). For any h E [H]. the weight wh in Algorithm satisfies
1 S HVVKK.
Lemma L.3. Let & be the event that for all hk, E [KREE] (H),
dKKrrH max(Ednon: 1)
Σ 0% (V442(5441) ph. pl (5% a4)
S dll y log
(A/)-1
Then Pr(E] 2 8/8.
Lemma I.4. With probability 8/8, we have for all (s, E S x A.
(n(s. a) wh p. i (s, a) S (s, a))( Hdowwn
Lemma 1.5. With probability - 6/4, for all (h, k) E [H] [KREE), and any E S, we have (s, (s) + H(H h 1)
and
Kers.
/v d3 logdd H KRFEEdow,
where e > is a constant.
Lemma L... With probability have
6/2, for the function uy (g(c.) defined Line 5 of Algorithm we
Enng Vi (s) u) < c' /d d H log(d (og H H/8) KEFE + 2H² Edom-
2
Lemma L.7. With probability 1 8/2, 2 for any reward function which is linear with respect to the unknown feature 0² x A Rdd for all (s,a) E S x A and h E [H), we have
Q (HH H Edown (s a) (n(s (*,TT+ (*,++1
V (s 42n a)
H
34

---

1.2 Proof of Theorem 4.4
We first restate Theorem 4.4
Theorem L88. Under Assumption 4. I. after collecting KAFE trajectories during the exploration phase in Algorithm 2. with probability at least 1 - 6, the output of Algorithm 3, policy ι satisfies
Es ~ (B - - (si, r)] S HoggdKrrrt KRFF + 6HH Edown.
(1.1)
If the linear combination misspecificctiion error & in Assumption 4.1 satisfies Ö / KRFE) and the number of trajectories in the offline dataset for each upstream task at least ( K KReE/ d²),, then, provided KEFF is at least (dd H
with probability 1 8, the policy a will be an e-optimal policy for any given reward during the planning phase.
Proof of Theorem 1.8. We condition on the events defined Lemma 1.6 and Lemma 1.7 which, by union bound, hold with probability least 1 1 8. By Lemma 1.7, for any S, have
(2(s) max (1( a) 2 max Qi(s, a, H2 Cdown aed
aEA
VV H Edown-
This implies
E [ (s1, (s1, Esis (s1,)) + H Edowns where a is the policy returned by Algorithm 3. Observe that, using Lemma 1.7, we have
E ((1) rs)
Ear (11. (1
E (r t (s1)) p(- PI ((s))))))) + 2u (11, Hdown - (11,
- p( Vo m. (s)),
E s
(e))) [2((2) (82, + (x1: + Edown
2E s (s, u)] + H2 Saown.
Moreover, note that 0 (i()s) s H and 0 < VT H as 0 r(s,a) < 1. Thus, would always have (s))) Vi (s, & H. Along with the previous derivation, this implies,
Eg Eg [V( (s 2E. 2E (V (s, w)] + H² Edown.
By definition of Vit(s, u), we further have Es IVT u)] s Esr (V, have
(s, u)). By Lemma I.6,
(s, S c' lo H/8)/ KREE 2H2 2
So, we have,
Pa)na(V1((1177) Vsl,))) < Exj ~p (1())
(81) - (iT(s, + 2 Sdown
< 2F s V7T (s, u)] + 2H2 Edown
2
< 2F pn [7 w)) + 2H Edown
< 2c' v H log d H/8)//////K + 6H Edom-
(I.2)
35

---

Recall the definition of Édown from Lemma 4.3, that is, {down = C₁ Cr,Cav / /
If the linear combination misspecification error in Assumption 4.1 satisfies Olv d KRFE) and the number trajectories in the offline dataset for each task in the upstream stage at least ÖT KREE ) then the first term in Equation (1.2) dominates the second term Sdown- Then, by taking KRFE cgd H log(dH0- ///e? for a sufficiently large constant eg 0, we have
B (Vi r) (1" r) & 2d /d9 H
/ 6HH Sdown S 6.
This completes the proof.
J Proof of Supporting Lemmas in Appendix I
J.1 Proof of Lemma 1.2
Proof of Lemma 1.2. We have
(A Σ on (5%, a Vi (
--1
1/2
I Vi
1 (
ooh am) (3)
S v KREE H T an) an)
S V d.KFFE
where the first inequality follows from Lemma 0.5 and the fact that the largest eigenvalue of (A is at most 1, second inequality follows from the fact that 177 h+1 (s)l I H for all S and the last inequality follows from Lemma O.4.
J.2 Proof Lemma 1.3
Proof of Lemma I.3. The proof is similar that of Lemma B.3 in (Jin et al., 2020b) with the major difference being the usage of approximate feature map ((: and different reward function at different episodes. We provide the full outline of the proof for completeness.
For all (ki, h) [KReF] x [M]] by Lemma I.2, we have || S Hvdk. Moreover, we have
r = (., and hence we have
*...)) + -B
Thus, value function V is of the form
V = min { max "(, TT Tw ()) }
[max
for some A Rdddd and Ed which matches the value function class defined in Lemma O.9. Moreover, by construction, the minimum eigenvalue of A is lower bounded by 1. Combining Lemma 0.7 and Lemma O.9, we have for any fixed > 0 that
Σ 34 (7/42())44)= p(a.T+1) Dh
k+1 141 (8%%
(A))-1
d
8/1 dk
(
32 /2 1
8
< 4
log(k: + 1) dlog
1+
+ dd log
+ log
+
8
2
g2
36

---

We set the hyperparameter 3 = CLHVd dH Vlog( RdK H max(Edown, 1)/6) + H Hdown v /dKKFE: Finally, picking E dH/kk k. we have
ddKFFll H max 1)
Σ Th (Ph+2(0h+1)
D 10 aN)
S ddd log
(Af)
which concludes the proof.
J.3 Proof of Lemma 1.4
Proof of Lemma 1.4. We condition on the event & defined in Lemma I.3, which holds with probability
at least 1 - 8/8. We define wh = J. (s')""" (s')ds' where Ø is defined in the proof of Lemma 4.3. Note that by Lemma 4.3, we have || Cv..vd. Now,
Φn (8, a) T th ph. p(+.T+1) Pk (s.
(5n(s,a) aaTTT TTO t 6n a) w + a) p/
h+1 (s,
(s) aa_aaa + (s a)
Wp P pl. (()
h+1 (s,
(x, ( t En (8, a)
wh
Hddwun:
(I.1)
where the last inequality follows from Lemma 4.3 and Whttll $ H. The first term in (I.1) can be written as,
(n, a) T (8, aa
k-1
= (8, a)TT( (AA))
Σ
(n a7 V/t - 2n(s, a) T()()) -
= =),(s, (e(s,,)T(AA) ((A/)
Σ Bn af) V/ 2 wht Σ Sn Pn PnV
-- (8, (A
Gz (s, a) (ls,))((A))) (AA)
Σ On (sh an)
[ (8h+ p(
h+ (sh+1)
(b)
+ (asss))((A))
[
Σ
(n a7) (pf -Ph) Vitz((hhz2)
(J.2)
(c)
We now bound (a), (b), (c) in (J.2) individually.
Term (a). We have,
- -x ((A))((A))))))))))
< 1on(s, a) I1 (
| @l (s,
Cr. H 2n (s. (s. willant> 1-
(J.3)
37

---

Term (b). Using Cauchy-Schwarz inequality and the definition of the event & from Lemma I.3, we have
a) (11)))
Σ Th (sh (h)
V/ (8h+1)) Pk.
(x
h+1 (6b+1))
s (
|
Sn w/h 1 P/.
h+1 (3
(a8)
d'KFFE H max 1)
S dH /
log
||ōn(s, (1)((s)))l(a
(J.4)
8
Term (c). We have
k-1
(x (AN)
Σ
Φn (8% an) (p)
Dh (h) - Pn.
T=1
Σ on (sh ant HHEaom
()) (A
(s a) (AA) -G (sh ahl H Edown
68)
(s a )1As )
Σ 1 "Wllass)1)
H Edown
H Edown v dk Sn (s, ())
(J.5)
where (i) follows from Lemma 4.3, (di) follows from Cauchy-Schwarr inequality and (iii) follows from Lemma O.4.
Substituting (J.3), (J.4),
(J.5) into (I.2), and denoting 8 = C.H.Vd +
+ Edown / dk we get dH, log d'Kams HI (nvmi))) o
(s, a) T Twnh 2n(s, T wh S (
Putting everything together in (J.1), we get
f0l a) क pl. p(, (4 & (1) + H Edown:
which concludes the proof.
J.4 Proof Lemma I.5
Proof of Lemma I.5. We prove the first part using backward induction on h. For h H, we have
rs
min
max
{ (ri (s. + P/ p(c.T+1) + Ve " (*) H}
aEA
-min
max (s 0). }} aEA
<min
mmx
{ L (s, + Tn a) 1 + fl
all (
aEd
V (s)
S D (s) + H(H - H + 1)
38

---

Suppose for some h + 1 E [H], it holds that for all E S,
(i, S (s) + H(H - &donn
Then,
(3, min
max { // ( + Ph p 1 (3,0, H
aEd
(/ (5) a) V (8,0,
max aEA
max atA
( a) - (s, a)
H(H h Edoma
max acA
{1h((,, a) (s) T
2} + H H(H 1t)
(1%(ss,) (3, क + Bl ((a)))))) } + H(H - h + Edwwwn:
max aed
where the last inequality follows from Lemma 1.4.
Thus, we have
(s, x min max { (r a) + a) wA Bl) () } H H(H h h 1)
aEA
(s) H(H + 1dowwa
as desired. This completes the first part of the proof.
Now we prove the second part of the proof. For all (ke, h) E [KRFE] x [H 1], we denote,
ch Pl
p
i (s ah)) V (sh+
Conditioned on & from Lemma I.3 where Pr[E] 2 - 88/8,
Kgr
KEY
(+
(1 at + (s) कt + Bll (f
Kun
(oi(s), at) Tat (5t
k=1
KRsE
(pl (z + 3 ( + H Edown
Kun
+ (s + (()la)))))))))) H KRFESAWn
Kum H-1
Kamn H
3 3 Ma) + H2 KRrE&oow
h=1
k=hht
where the second inequality follows from Lemma 1.4.
KESE is martingale difference sequence with |ch| < H. We Note that for each h E [H] - 1], { define the event 8 be the event that
Kunn #-1
Σ
ch
S c2H KREE log H/8).
39

---

By Azuma- Hoefdaig inequality, we have Pr(E"] 2 6/8.
By Sauchy-Schwand inequality, we have
KirE
H
KirE H
[ (a)) 1 S
KREEH H
( h:
(
1= 1=1
Using Lemma D.2 of Jin et al. (2020b), we have
( 12 aBt
2d.H log (KRFE).
Conditioned on & and & where Pr(E (E') 2 1 8/4, have
KREE
(s & d & KRFE log KREe + 33 KRECH. 2dH (og())))) + H² KRFESdoWn
S Scv CoV/HHKRRFE log(d + 2 KREFSdown
which completes the proof.
J.5 Proof of Lemma I.6
Proof of Lemma 1.6. Denote A = (st r*) - Eoww (s, Note that n² depends only on Kam is martingale difference sequence.
the data collected during the first k 1 episodes. Thus, { Moreover, we have A < H. Using Zzuma-Hoeffding inequality, with probability 1 - /4 have
Jk=1
KAFE
S ea Hy KRIF log(1/5)
where C₁ > is an absolute constant. Therefore, we have
Kam
Kum
Egngg
E
Σ
Vi(s, < Σ (= + H. KRIE log(1/5).
h=1
Now, notice that for all l: [KRrE], An 2 ΛB. Thus for all (k, h) [KRFE] x [H], we have
( (n)))
which implies
Vi(s.u) S (.().
40

---

Using Lemma 1,5 and union bound, we have with probability 1 1 5//2.
Kum Vi (s)r /KRFF
Bawn (V, w)] s Esnu
Kum
log(1/6) KREE
+ H
KREE
k=1
KAFE
log(1/5)
(sf) + H22 (down) + t
KRFE
KRFE
k=1
KREE
log(1/6)
(sF) 112 Edown cill
KRFF
KRFF
logí (16) KRFF
(ev d3 H KRFE Ro H/0) + H² KRFECdown) +
H2 Sdown H
KRFF
d H log(dKKnr
log(1/6) KRFB
coH
22H2 +2H2Edowu
KRFE
H4 log(d H/0)
<
+ Edowa:
KRFE
for some absolute constant d'
This completes the proof.
J.6 Proof of Lemma 1.7
Proof of Lemma 1.7. Following the same argument as in the proof of Lemma 1.4, for all E [H] and (s,a) E S x A, with probability 5/4, have
an wn p( Vn+ (s, $11o,ss,
Hdown
Thus, for all h E [H] and (s, E x A. we have
(n) S (n(s, (a(s,a)TTin Twn + rr((
+ (s)
am((sa)) pl p(x.T+11) Ph+1 D(
+ H Hdownn-
Since, (s) & and un(: (.) = min B
we have,
(x a) S (n a) + pk p(+ 1 + (8, a) + Hddwon-
This completes the first part the proof.
Now, using induction, we prove that for all h E [M] and (s,a) E S A, we have (,,,,,,))
H(H 1) & Qu(s,a), a).
When h H + the claim is trivially true. Suppose for some [H], we have, for (s, a) S A,
(h+1(8, H(H Edown @n(s,a).
Note that,
Qu(s,a)= min { ( a) Ton, (s( an) H}
Since (s, a) S and un(..) min (VO()))
it suffices to show that
41

---

(s, S a Sn (s. + Bl (s.)(aa)) H(H h 1)wwn Applying the max operator on both side of the inductive hypothesis, we get
(8, S Vn+1+ + H(H Edownn
Now,
(s. = ()) + Pk. pl
V (8,
S r a) + ph. p (s: + H(H hdow
5 0, a) _ wn + Bl (1)l + H -((H hd) + (s(,aa)
= (s, a)
wh (, + (3, + 11 - hh+ 1)
This completes the proof.
K Proof for Downstream Offline RL
First, we state the supporting lemmas that are used in the proof of Theorem 4.6. The proof of these lemmas are provided in Appendix L.
K.1 Supporting Lemmas
The following lemma shows that the linear weight ūy in Algorithm 4 is bounded.
Lemma K.1 (Bounds on Weights in Algorithm 4). For any h E [H], the weight Sn in Algorithm satisfies
| H H dNNtt/D
Next, present main concentration lemma for this section that upper-bounds the stochastic noise in regression.
Lemma K.2. Setting Ad 1, B(8) = ca ((()) + H H where c())
log( H dN max(Edown, 1)/8), with probability at least - for all h E [H], we have
No
Σ
on (1 a7)
- V
S Hdvi.
A
Recall that, we define the Bellman operator B, as (B,,f)ss,) = (p f), a) forany f: S x A >> RR. In the next lemma, we denote the Bellman estimate (TB, 1) (n(:.)) wh.
This lemma provides upper bound for the Bellman update error [B, V++1 B, V (s, a)l and characterizes the impact of the misspecification of the representation taken from the upstream learning.
Lemma K.3 (Bound on Bellman update error). Set Aa B c (HddHHH dN offfedmmm) where H ma 1)/8). Define the following event
{(τ) {
(BB, P (s, a) - { a) + HEdoom Yh E [H] and V(s, ES an
Then under Assumption 4.1, we have (P( 2 8. Definition K.4 (Model prediction emror). For all h [H]. define the model prediction as,
(n a) (B, ( (s. a) (s.
The following lemma decomposes the suboptimality gap into the summation of uncertainty metric of each step.
42

---

Th=] be the output of Algorithm 4. Conditioned on the event E() defined in Lemma K.3. any h [H] and (s, a) E x A, we have
+ Hl Édown in Lemma K.5, with probability at Lemma K.5. Let {T
(s, (s,
where (n(s,a) = 8 B1 (s, + Hdown Furthermore, we have
H
V 11
V , 2 Σ
Ex Tn(sn,aa)le
hml
K.2 Proof of Theorem 4.6
We first restate Theorem 4.6.
Theorem K.6. Under Assumption 4.1, setting Xa 1, O(HdVi+ + (dNoffoooon) where
1 == log( HdN off max 1)/8), with probability at least 5,,, the suboptimality gap of Algorithm 4
at most
VP (s) Vpc T+n, (8) 2H2 Edown + 2B
E B (iensssnnnlla =s] (K.1)
s).
Additionally if Asssumptiin 4.5 holds, and the sample size satisfies Not z 40/mzp log(AdH//0) then with probability 5, we have,
V ((,,)+),, V p1),
5o
Kpp 1/2 //2 Edown Kx -1/2//1 log( log(HdNNo max(Edown 1)/8)
(K.2)
Nowf
Remark K.7. Compared to Theorem 4.4 Jin et al. (2021), the suboptimality gap in (K.1) has an additional term 2dow- When the linear misspecification error Ç = O 0 d/Nott and the number of trajectories in each upstream offline task dataset satisfies Ō d the RHS of (K.2) is dominated by O(N D(N-1/22H'dd) Hdd improving the suboptimality gap bound of REP-LCB (Uehara et al., 2022) by an order of O(Hd) under low-rank MDP with unknown representation.
Proof of Theorem K.6. Letting I' least 1 8, we have
H , (s) Va (p <
V
Σ
E [ (Sh, an)ssi =
h
< 2H Edown + 2,
Σ Eg 1=1
[ and an) 11 =
This finishes the first part of the proof. Now we will provide the suboptimality bound under the feature coverage assumption in Assumption 4.5. From Appendix B.4 of Jin al. (2021),
Non 2 40/K, (og(4d//)) then with probability 1 - 8/2, for any h [H] and (s, E S x A, we have
2 Kp
16, a)
S V
/Noff
Setting ((6/2) cg(Hd)/(6/22) H H /dNo with probability 1 -8/2, we have
H
4x"" p1. (8) -V (8) 5 211 Edown + 2,3
Er
12 (8h, an (h)l
Am
43

---

Using union bound, we have the following bound with probability at least 1 -
Vp. ((o V ((...+1),,5))
2
$ 2H Hddww + +8(6/2)
Kp
VNoit
(0(1) kpp 1/2 Edown
= =0
log( HdN max 12/0)
k/ H2 d
Rip
(K.3)
Nott
L Proof of Supporting Lemmas in Appendix K
In this section, we provide the proofs of the lemmas that we used in the proof of Theorem 4.6.
L.1 Proof of Lemma K.1
Proof of Lemma K.I. We have
wA
A
(s (/T ((+1(5++1))
Nai
(16% + [/))),/))/ .) 1/2
Nan
X
Nonr
Naw
| (o. -7)1
< H
da
dN
Ad
where the first inequality follows from Lemma O.5 and the fact that the largest eigenvalue of Az at most 1/ Ad, second inequality follows from the fact that r% E [0, 1] and [V)] s H - for all E S and the last inequality follows from Lemma 0.4.
L.2 Proof of Lemma K.2
Proof of Lemma K.2. The value function Vn41 has the parametric form of
V(-) = min {
max T () B A),())())) aEA
1).
where w E _d and positive definite matrix A is such that its minimum eigenvalue satisfies Amin (A) 2 Aa- From Lemma K.I, we have || || < H we have, for any fixed > and for H.vd ddNa//aa Thus, applying Lemma O.4 and Lemma 0.9, [M]. with probability least 1 6,
Nat
Σ
(%h an)
(142)
Ph+z (
AA
d
Nair + Ad
No
8N2
2
log
+ log
+
Ad
8
dd
Nor + Ar
/ d. off
8v
1
8 8/2 N
< 4H²
log
+ dlog
1+
+ &olog log
1+
+ log
+
Ad
/v
AgeZ
8
Ad
dd
Nott + Ad As
4/1 /
8VdB²
1
8.V2 S
2
log
+ d log
1 1
d &log log
1+
+ +log
2
EV Ax
AgE
8
Ad
(L.
44

---

Setting
dH/Nou, Xd 1, 3 = Cadii + H log H d. max 1)/6), we can further upper bound (L.1) by
HydNontSawnll
where
4H²
2 log(1 + Nott + doog 4d N22 dl log(1 + 8d 3/2 2 22)) + log(1/8) + +82212
off
dog Nott H2 dog 1/2 (/2 H log H dl log Hdd 12,3 offodowa) H
Therefore, have
Nai Σ
(sh a7)
( (5t1)(5h, a7) (
(PP(
2 Hdvi.
Ax
L.3 Proof of Lemma K.3
Proof of Lemma K.3. For h E [H), we define Wj = (5) Then we have
[TB, Vh+1 BB (8, a)l
(BA Vn+1 Bb Vn + BA Vn+1 Bn (s)
(B,, Pa+1 - B, 1 a) + (B, P By (
(P/f--T 1/ Ph 1 h+ (s, 0)) + (P - Gh
H| H ((+,T)))1)( (-1s,a) (on a) (M)) + +1ʃ (1 - -3, a) Tn
H Haown 10, (8. a) (wh -
(1.2)
where (i) follows from the fact that [x)1(s)) < H for all S and (ii) follows from Lemma 4.3.
We now decompose the second term in (L.2). Recall that A₁ = >m Nat
Nat @5 105.of%000 0n (5%, a7) AaTa
and wh = Ar² Ene ( of Vr+) (8). Then, we have Gn (8, a) (wū w)
= (s, a) TA=]' {
Non
-EE
Nit
On (5% at) (h
Φn (s) oj) on ah) + sala 16%
Nar
(aon(s, a) An wh
on a) An
Σ
2n (5h an) (oek Vn+1) (ntz((htts)
(0)
Nut
2n (8, a) TA,,
A
Σ (nh
(PPPar
ah)
ah)
(L.3)
(6)
We now provide upper bound for each of the terms in (L.3).
Term (a). We have
a) N (1) (s, 11
(2l) a) a,- |,,
(w))
Vatld 2n(0 ,))l,
(L.4)
45

---

where (i) follows from Cauchy-Schward inequality, (ii) follows from the fact that the largest eigenvalue of A is most 1/da and (dit) follows from Assumption 3.1 and [V++((s]] - H for all
8 E S.
Term (b). We have
Nat
(5 a) A A
To (0% ah)
[1P) (7) (on10)
[1P/
Nuat
Σ
2n (5% an)
(pl V. 11+1) a - P
2 (s.
(L.S)
where the first inequality comes from Cauchy Schwarz inequality and the last inequality follows from Lemma K.2.
Term (c). We have
Na
as TTAT An
Σ
2n (sh a/)
(P (Ph Ph +(+
1 0
6n att) ||
(5 0) T An
H Edown
Nai
Σ (s. a) 1n añ)
Edrwn
Now
(s, a) ||== )
Non
Σ
[on(5T,
H Edowa
< H Edown v dN 5n (8, a))l
(L.6)
where the first inequality follows from (s) S H for all E S and from Lemma 4.3, the second inequality follows from Cauchy-Schee Schwarz inequality and the last inequality follows from Lemma O.4.
Jog(HdN off max(Esown: 1)/d) and combining (L.2) to (L.6), we get for any [H] and for any (s, a) E x A, with probability at least 1 - 8,
Setting A 1. B ( H ddNoofddown) where
BB 2 - (a s 36 1) + Edownn-
This completes the proof.
L.4 Proof of Lemma K.S
Proof of Lemma K.5. Recall that
Qu((y) min(fn(:.)Tan T.. - - h
If (s, T Qn(s,)) = (Br ( > 0.
Th - a)) I 0, Qn(s))) = 0. Then, have (s) (BB
(s) T wi - (n(s,))= > 0, we have
(s, a)
(BnVn+1)(s,)) - Qu(s,a)
> (Ba - an(s,a) wn (x
(Br (BADnt))(ssa) (BrP + (n(),,))
2 0.
46

---

where the last inequality follows from conditioning on the event (().. Thus, we have (s 2 0 for any h [H] and (s, a) S A.
We now show that 44 (8, a) - (s, a). Observe that
Th 2n a) TT= - Ta(s,a) - (
- (a)
< (())())))8ssa)) -
where the first inequality follows from the conditioning the event E()) and the last inequality follows from the fact that rh E [0, 1] and Vn+1 [0, 11 hl.
Now,
an(s,a) minfōn a) T w. (n(s. h 11+
= max a) Th - ())
2 as T6 (s.
= ((+))t))ss F (8.
Now, from the definition of 4, have
(s a) = (B, 1) -)x(s)))
S BB, (n+)))((, (B, + (s)
S (s,aa,
where the last inequality follows from the conditioning the event {
Finally, we obtain
V p
V (1..))))..s))
H
Ealla (sh, an)(s1 (shh = =
E (s an)(si s]
h=1
h=1
Σ E. [T an) = s]
where the first inequality follows from Lemma O.3 and definition of म. and the last inequality follows because with probability at least - we have S (,, a) I 21 (s, a) for any [ and
(s,a) E S A.
M Proof for Downstream Online RL
this section for notational simplicity we denote Vi pt 1,, and Q7, pt..z (1: by Vr (s) and 0% (s, a) respectively.
First, we state the supporting lemmas that used in the proof of Theorem 4.7. The proof of these lemmas are provided in Appendix N.
M.1 Supporting Lemmas
The following concentration lemma for online RL upper-bounds the stochastic noise in regression. The proof is omitted since it quite similar to the one of Lemma K.2.
Lemma M.I. Setting Ad = 1, 3n = ca ( + H HVdnEaao + CL. CLVHdd, where in = log( H dn max(Edown; 1)/5), with probability at least 1 8/2, all h [H] and any [N], have
Σ Zu (sh, a7) [Ip (pf-, nn (1h, - D (
S Hdd/in
A
47

---

The next lemma recursively bounds the difference between the value function maintained in Al- gorithm (with-out bonus) and the true value function of any policy æ. We provide a bound for this difference using their expected difference at next step, plus an error term. This eroor term is upper-bounded by the bonus term with high probability.
Lemma M.2. There exists an absolute constant Cg such that for Bn (Hd H dnLdown
CL. where in = log( dn m 1)/5), and for any policy n,, with probability at least 1- 6/2, for any (s. S A. [Non] and h [H. we have
(n(s,0) 0) T T Q% (8, a) = Ph p(= (V - (s, a) + (8),
for some A a) that satisfies (s, a)ll S 9o, Ballon (Ag)) + 2H
We now prove optimism of the estimated value function in the following lemma. Lemma M.3 (Optimism of value function). With probability at least - 6/2, for any (s, (s,0) EXA h [H] and [No,), we have
(s, a) 2 a) h++ 1)
Lemma M.2 also easily transforms recursive formula for the value function difference 8% V (5%) VA" (sK). The following lemma will be useful in proving the regret bound for the online downstream task.
Lemma M.4 (Recursive formula). Let 0% = (s - VV"" (5%) and 5h-1 = 1 ahl 0% h+1
Then for any (n, h) E [Non, H] with probability at least - 5/2, we have
5% S 8h41 + ch+ + 8m 10 (5h a (1()(())>= 2H
M.2 Proof of Theorem 4.7
We are now ready prove the main theorem in this section. We first restate Theorem 4.7.
Theorem M.5. Let ā be the uniform mixture of π ....., Non in Algorithm 5. Under As- HLV. where in
H sumption 4.1, setting Ad 1, 3., (( HV log H dn max(Edowmn 1)/5), with probability 1 - 6, the suboptimality gap of Algorithm 5 satisfies
Vp P V P I O(II2dEáow + 112 3/2 N 1/2).
(M.1)
Remark M.6. As we use estimated representation (on 11=1 (o1 from the upstream tasks, we get an extra term dd2d[down in the suboptimality gap above. When the linear misspecification error E
D(yd/Nort and the of trajectories in each upstream offline task dataset satisfies Ö -
the RHS of (M.1) is dominated by ((F22 1/2 improving the suboptimality gap bound of REP-UCB (Uchara et al., 2022) by an order of (1H3/22 K vd) under low-rank MDP with unknown representation which attests to the benefit of upstream representation learning.
Proof of Theorem M.5. We first bound the cumulative regret by
No
Vi y - V VPC..TH..))
No
(Vr
2on
N
en + Bn (5 a;) 1la + Satwn] + 2H² NonEdown
1=1
N,.
H
Non
H
3m
(5 an (ag)
-onHoNonooow
(M.2)
(b)
horizon dependence in REP-UCB Uehara et al. (2022) to H. We further rescale their suboptimality gap by factor of H² as we do not assume the sum of rewards be within [0, 1).
'We convert the 1/(1 7)
48

---

where (1) follows from Lemma M.3 and (ii) follows from Lemma M.4.
Na.H Note that in term (a), fSh I==1,h=1 is a martingale difference sequence with [5%
S 2. By Azuma-
Hoeffding inequality, we have, with probability at least 1 8/4,
Nai
ER
S
8. 11 log(8/0)
(M.3)
For term (b), we have
Nen
H
3n @n (sh, aB (aa)a
(1)
32
Tn (sh oB>ll(n)>~
200 H22 Non + H2N2 de? ) v 2dg(1 + Non/(dA))
k=1
H 20 H2 H2 Non + H2N dedde d 4d Non
(188) 2v 2c 2 / L log Non +
v log Non):
(M.4)
where (i) follows from Cauchy-Schwarr inequality, (ii) follows from Lemma O.6, and (iit) follows from the inequality that /I + y S v / for all x, y
Combining (M.2), (M.3) and (M.4) we get
Na
V
Vaa Vp S ddN + 112 2 Non).
Dividing both sides by Non. we get
VOTT,, V PO-TTY, I &HH deown 2223/22
This completes the proof.
N Proof of Supporting Lemmas in Appendix M
In this section, we provide the proofs of the lemmas that used in the proof of Theorem 4.7.
N.1 Proof of Lemma M.2
Proof of Lemma M.2. For any policy n, we define W (s" Note that Tn (8, a) T wH PAVT. (s,a) and by Lemma 4.3, we have ||, S CLVd.
Now, we derive the following
(3, a) 16 (s)
(s, a) T Th (s (s, T wH + (s, a)
T w (s,
< (0n a) Twp (s. wx.) + (s, w/ - (s,
(N.I)
Using Lemma 4.3, we bound the second term as (s, (n(s,a) a)T "wil p -P
< H&aom.
(N.2)
49

---

where we used the observation | S H. Now, the first term in (N.1) be bounded by
(n(s,)) 1h ())ssa)
wh
n
=5(ss a (AR)
Σ
Sn ab Vn (5%+1)) (5
wx
R-1 Aoo
= (s,s, ( (AR))
Σ
Sn ai) Vitt
Σ
(5% a Pr Vot1
= as) (Ag) wng + (s, (T,)))((A)) (AR)
Σ
(5h an)
Vi (8 h
h+ (8h+1)
(1)
+ ()
Σ (sh ah) Pn)P (ot
|
+ a (AA)
Σ
on (sh an) (p/ - Ph -Ph)
Diit
(N.3)
We now bound (a), (b), (c), (d) in (N.3) individually.
Term (a). We have,
- as) " & (ōn(s, a>)l Mowi
V 2l (s. a)ll(ay)
s &C.VAa (n a) (ln 17
(N.4)
where the first inequality follows from Bauchy-Schwarz inequality, the second inequality follows from the fact that the largest eigenvalue of (AT) is at most 1/Aa and the last inequality follows from Assumption 3.1 and I [s)] < H for all E S.
Term (b). Using Cauchy-Schwarz inequality and Lemma M. we have
(a(s)))T(AB)) (Ag))1 Σ Sn(sh (h)
V ( ((..T+1) pr h (671()
$ (1
Σ (s and
V
1() m (8h+
(ag)-1
$ Hdd/i (
(N.5)
50

---

Term (c). We have
(5, a) (AA) (AA))
2n Pn(( (itt) ah)
@n ah (shh / ( VI+I (s'
(s) (AA)
an(s (AA) (AA; -at) / (Vitt -6
@ / (Mitz (s + |20ōn(0.0) (AA)) / - (s)
Ph(V ( a)l + CCLV (ns
(1) P (1)) ( - (*,))s,, + HSauwn Ct. Cn.VaaHdllnn (s,,)a))))))
(N.6)
where (1) follows from similar steps as in (N.4) and (ii) follows from Lemma 4.3.
Term (d). We have
(1 0) ( (AA)
Σ
on (sh, ai)
(p/ttT=)) Ph Dit (oh2
Σ On (sh: añ} }1
Φ (5 a) T(A))
H
(s (AA) (A%) (sh, an)l Hddown
I (5, (1), .)
Σ | a 1oll)
Sdown
Edown dn Sn ())(
(N.7)
where (i) follows from Lemma 4.3 and 1i (s) S H for all s e S, (ii) follows from Cauchy- Schwarz inequality and (iit) follows from Lemma O.4.
Substituting (N.4), (N.5), (N.6), (N.7) into (N.3), and setting du = we get
(s TH (n Tw $ cal Hdd (HdV (0) + H H Vdngawn + Hdd (sn ())))))))))))
+ p(( -((.T-111 h+ (o, + H&
Bn (n a) (Ag) + p(*.Tt1((m.. (s, + Haown-
(N.8)
Combining (N.8), (N.2) in (N.1) completes the proof.
N.2 Proof of Lemma M.3
Proof of Lemma M.3. We use backward induction for our proof. First, we prove the base case, at the last step H. By Lemma M.2, we have
(s Qir(s,a)
11(..4+1) (vn+i (V (ss,) (7,
< Bol (1 a)) + SHown
51

---

Thus, we have
@h(,,a) min {ōn( a) wH + 3m Bn (()))))))))))) 1, - - hh+1
(s, - 2/ 5down.
Now, suppose the statement holds true step h 1 and consider the step h. Using Lemma M.2, have
(s. a) 8h (7) = (s + P/ 1 (Vh+1 - Vi+t)(s,
Sn (o,(s.))))
2HHdown -2H(H - h))
= Sn 30. -2H(H - h + 1dddwn.
Therefore,
(s, = min Tw, + B,llo(s, (all@(ssa)llaa))-..1 -
()) -H h + E)own
which completes the proof.
N.3 Proof of Lemma M.4
Proof of Lemma M.4. By Lemma M.2, for any (s,a) SxAne [H] and E [Non]. with probability at least 1 - 6/2, we have
(s) - QT (s, < (x) + p(t
(V/+1 V (s,
-
S Ball (s 2HH p(t (h+1 as)
And finally, by definition of Λ" in Algorithm 5, we have (S)) = aj argmax 2% a)).
This implies Q7 a Q (5h ag) V7) VAr (0%) = 8% Thus, we have
5% h + Eh (5, 22H
O
Auxiliary Lemmas
O.1 Miscellaneous Lemmas
The following lemma measures the difference between two value functions under two MDPs and reward functions. Here, we use a shorthand notation PA V.+11 on) = (BggPalllenn,a (V))..
Lemma O.1 (Simulation lemma (Dann et al., 2017)). Consider two MDPs with transition kernels P and P2, and reward function T and ra respectively. Given a policy n. we have,
H
V V. (NN) - V...a.ro (SN)
(x(ppm [ri(sN, aw) - r2 (8h1, aw)
=
h
A
(PIN - P2N VW (Shr, aw)) sn
E - mu
[ri(sn, aw) r2 (2 any)
+ (PI, P2n)) War ss]]
We use the following lemma to deal with distribution shift in offline RL setting.
52

---

Lemma 0.2 (Distribution shift lemma (Chang et al., 2021)), Consider two distributions p1 E AS A) and P2 ASS A), and a feature mapping S A R Denote C suPPER Bnawp, o(s.a)0(s.a)"= Then for any positive definite matrix A, we have 5.s.po0ls.))0ss00
E. 0a () A S CEs,a-p20ss,0 a) Ad(s,a).
We the following lemma to bound the suboptimality downstream offline RL. Lemma 0.3 (Decomposition Suboptimality, Lemma 3.1 in (Jin et al., 2021)). Let ( J be the policy such that (s) (Qn ) (-s)))) and for each step h E [H] define the model evaluation error as (s, a) = (B, Qu(s, a). For any & and s E S, we have
H
H
VP (8)
Σ
Ee (sk. = o|]
Σ
Bx U1 (sh, an)[s1 =
H
E ((Q (+1s)) - (-(-18n)) =s].
h=1
0.2 Inequalities for summations
Lemma 0.4 (Lemma D.I in Jin et al. (2020b)). Let A₁ = AT Ei &io,, where Φs E RRd and
O. Then it holds that
Σ
o (An) 6. S d.
1=1
Lemma 0.5 (Lemmaa D.5 in (Ishfaq et al., 2021)). Let A E Rdxd be a positive definite matrix where its largest eigenvalue Amax S X. Let Φ1, Itk be vectors RR. Then holds that
1/2
Iz
Ak
Σ
| 11
Lemma 0.6 (Lemma G.2 (Agarwal al., 2020b)). Consider a sequence of semidefinite matrices Xi...., .......N paxd with Tr(Xx) < for all [N]. Define Mo Al and Mn Mn-1 Xn.
Then
N Σ Tr(X,, 2og (M))
det (Nx)
s 2d + N/((d)).
0.3 Covering numbers and self-normalized processes
Lemma 0.7 (Lemma D.4 in Jin et al. (2020b)). Let { be a stochastic process on state space
100 be an Rdd -valued stochastic process where di E F, and [ S 1. Let A₂ = > Ee= T. Then for any ō > O, with probubiliyy at least 1 for all O, and any V V with sup V(s)) S H, we have
with corresponding filtration {Fr) Let {0
E(V(s) Σ (VV)
EV 4H² d log
k + +) + log No + A
8k262
A
where Ng is the e-covering number of V with respect the distance dist(V, V, V sup.es (V()) V'(s).
Lemma 0.8 (Covering number of Euclidean ball, Vershynin (2018) ). For any è >> O, the s-covering number, No. of the Euclidean ball of radius > 0 in I2d satisfies
N₂ < 1+
E )
2B d
3B
53

---

Lemma 0.9 (e-covering number (Jin et al., 2020b)), Let V denote class of functions mapping from E with the following parametric form
v()) min
max w o( + 8 - aEA
where the parameters (w, A) satisfy ||||| - L. E [0, B) and the minimum eigenvalue satisfies Amin (A) 2 X. Assume [| a)]l < for all (s, a) pairs, and let N₂ be the e-covering number of V with respect to the distance dist( (V, V*) = sup, (V()) - V*()). Then,
log.N; v dlog(1 + 4L/E) + of log(1 + /3 32(()))))
54

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and the introduction are reflected in the paper's main contribution.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the assumptions made in order for our theorems to hold true and diseuss the limitations of those assumptions.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
55

---

Answer: [Yes]
Justification: We provide detailed list of assumptions along with their interpretation. The complete proofs are provided in the appendix.
Guidelines:
The answer NA means that the paper does include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [NA]
Justification:
We do not have any experiment in this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. .., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
56

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
Mibllic/guides/CdeStttiii oooosooooosooioooioooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: //nips.coom cttpsscccppulic//ug guides/ /CodeSubmi ssi sisoiiyyi for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: There is no experiment in the paper as is theoretical nature.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
57

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification: The paper does not include any experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: pettps:/nneurs http///nurrrp..peeee hccppbblic/tthhcccdddeeeele..l.e
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
a theoretical paper and thus we do not forsee any immediate societal
impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
58

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
is a theory paper.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
59

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Not applicable
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
There was no crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: There was no study subject.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
60


Reviews:
[{"review_id": "eF168mgVBc", "reviewer": "Reviewer_Qe4K", "summary": "This paper proposes a representation learning algorithm for offline multitask reinforcement learning. The proposed algorithm, MORL, is designed for offline multitask RL in low-rank MDPs. The learned representation is examined in downstream RL in reward-free, offline, and online scenarios.", "questions": "1. In line 295 in Section 4.3, why do you assume that reward function $r^{T+1}$ is unknown?\n1. Could you provide any experimental results to support the claims in the paper?", "limitations": "The paper lacks a limitations section.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Detailed theoretical analysis of the proposed algorithm.\n1. Better sample complexity than existing algorithms.", "weaknesses": "I do not conduct research in theoretical offline RL, so I cannot provide an accurate evaluation for this paper. I will try to provide some comments based on my understanding.\n1. This paper lacks a comparison with other multitask RL algorithms. Table 1 only compares it with single-task methods. Section 4.3 does not provide any comparison with other algorithms.\n1. The contributions to downstream online and offline RL are not clear.\n\nMinor issues:\n1. Equation (3.2) should be $P(s'|s,a)$... where $s'$ is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points. \n\n### Points raised as weaknesses\n\n1. To our knowledge, there is only one other concurrent offline multitask RL theory work [89]. We provided detailed comparison with that work in Appendix A: Additional Related Work. However, as described in our discussion the theoretical results of [89] are not directly comparable to our theoretical results. Table 1 only compares with single-task methods as there is no existing work that studies reward-free RL in multi-task settings. The goal of Table 1 is to show that performing offline multitask representation learning in the upstream task, can help improve sample complexity of downstream reward-free tasks compared to single task counterparts. For offline and online downstream tasks we provided comparison with other algorithms in Remark K.7 (Appendix K.2) and Remark M.6 (Appendix M.2). We provided them in the appendix due to space constraints in the main paper. \n\n2. As we mentioned in the Introduction section where we highlighted our contributions, for the downstream part, our core contribution is in the reward-free setting. We provided results for offline and online downstream settings as complementary results. \n\n### Questions\n1. In Section 4.3, for downstream offline and online RL tasks, we assume the reward function $r^{T+1}$ to be unknown. We made this assumption to make our result more general. Note that, knowing the reward function $r^{T+1}$ would have made the downstream task easier to solve.\n\n2. As discussed in Additional Related Work in Appendix A under **Offline Data Sharing in RL**, there have been numerous empirical works that studied the benefit of using offline datasets from multiple tasks to accelerate downstream learning. The aim of this paper is to provide a framework for investigating this approach from a theoretical lens. While empirical validation is limited in this paper, we hope that the theoretical analysis presented in this work will provide insights for algorithmic design principles that can be used to design empirically better performing algorithms compared to current state-of-the-art algorithms. \n\n\n### Limitations\n\nAs all other reviewers have acknowledged, we have addressed limitations of our works, especially through a thorough discussion of the assumptions made for the theoretical analysis. Moreover, as per the suggestion of the Reviewer fKn9, we will add more discussion on potential practical limitations of MORL along with how to possibly mitigate them in our next version.\n \n\nWe hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?", "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": "LVo7QNanJz", "reviewer": "Reviewer_Z4xa", "summary": "This paper provides a theoretical analysis of representation learning in Multi-Task Offline RL.\nSpecifically they consider a setting in which the transition kernels have a low-rank decomposition $P(s'|s,a)=\\langle \\phi(s'),\\mu(s,a)\\rangle$, such that $\\mu(s,a)$ depends on the task but $\\phi(s')$ is shared by all tasks.\nThe authors propose an algorithm for this setting that first learns a transition model and then does planning on the model with a low-confidence-penalized reward.\nThey provide a suboptimality bound, based on bounding the TV distance between estimated and true transition kernels.\nThe authors also investigate the application of this acquired representation in down-stream online, offline and reward-free tasks.", "questions": "* Why was a nonstationary setting analyzed, and how did this affect the analysis? Is it possible to achieve a better regret by considering a stationary reward and transition function?", "limitations": "Assumptions are discussed sufficiently, so I have no concerns about additional limitations.", "rating": 7, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is very well written, assumptions that are made are explained and well justified.\n * The exploration of down-stream applications of Multi-Task Offline RL is of large practical interest.", "weaknesses": "**I don't work in theory, please see the below more as minor comments.**\n\n * [89] analyzes a very similar setting but is only mentioned in the appendix as concurrent work. Strictly speaking, [89] was uploaded ot arxiv three months before the deadline, so it is not considered concurrent as per NeurIPS guidelines. Either way, it should probably be discussed more prominently and not just in the Appendix.\n\n * It would have been nice to specify which improvements in the bound are results of which decisions. For example, a conservative reward $r-b$ in planning vs the normal reward $r$, or how exactly using the low-rank assumption rather than a linearity assumption affects the analysis.\nSimilarly it would be nice to provide an interpretation for the different terms in (3.7), but I'm also not sure if that is really possible to do in a meaningful way.\n\n * The paper treats a setting with non-stationary reward and transition functions. It's is not discussed why the paper treats this setting and what effect it has on the analysis.\n\n* The approach name and acronym MORL (Multitask Offline Representation Learning) is a somewhat confusing choice, as MORL is frequently used to denote Multi-Objective RL.\n\n* The notation is somewhat confusing at times with $t$ representing tasks, $h$ representing timesteps in the MDP, $\\tau$ indexing transitions in the downstream offline dataset and $i$ indexing samples in the upstream dataset. The usage of non-stationary policies, rewards,  and transition functions also does not help to reduce subscripts.\n\n* The checklist instruction part was meant to be removed before submission, i.e. the instructions in L1253,1254", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points. \n\n* Thanks for your suggestion regarding [89]. A preprint version of our results appeared publicly around at the same time of the preprint of [89]. We will provide a more detailed discussion of it and move the discussion from the Appendix to the main paper. \n\n*  A conservative reward $r-b$ allows us to show near-pessimism in the average sense as depicted in Lemma 3.5. This is possible through a concentration argument for the penalty term $\\hat{b}_h^{(t)}$ as defined in Equation (3.4). Under linearity assumption as in linear MDP, the representation $\\phi^*$ is assumed to be known a priori. This allows one to use linear regression analysis to derive point-wise model uncertainty quantification and that can serve as a conservative penalty. However, this is not the case under low-rank assumption. It would be interesting to see how we can interpret each term in Equation 3.7. But as the reviewer mentioned, it is unclear how to do it in a meaningful way. Our motivation for representing Equation 3.7 in this way was to provide a fine-grained view for each term. \n\n*  We considered non-stationary reward and transition functions in our work as it is a standard practice in literature that assumes episodic MDP with finite horizon that have either low-rank MDP or linear MDP structure. Without non-stationarity, the dependency on $H$ would change in the resulting bounds.  Indeed, with stationary reward and transition function one can achieve a better regret bound in terms of $H$ dependency.\n\n*  Thanks for pointing out potential confusion with the acronym MORL. We will try to come up with an alternative acronym in the next version. \n\n*  We totally understand that it can be challenging to follow the notations sometimes as there are so many moving parts involved that add up in the form of subscript and superscript. We tried our best to keep things concise and clear as much as possible. We would be happy to incorporate any suggestion that the reviewer might have to make the notations easier to follow. \n\n\nWe hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?", "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": "1GeH8FtgX6", "reviewer": "Reviewer_fKn9", "summary": "This paper introduces the Multitask Offline Representation Learning (MORL) algorithm, which aims to enhance sample efficiency in offline multitask reinforcement learning (RL). By learning a shared representation from pre-collected datasets of different tasks, modeled by low-rank Markov Decision Processes (MDPs), the authors demonstrate the theoretical benefits of this method for various downstream RL tasks, including reward-free, offline, and online scenarios. The paper provides extensive theoretical analysis and introduces specific data-coverage assumptions, presenting the first theoretical results showing the advantages of representation learning in offline multitask RL.", "questions": "1. How does the performance of MORL scale with the number of tasks and the size of the state and action spaces?\n2. Are there any practical considerations or limitations in applying MORL to real-world datasets that were not covered in the paper?", "limitations": "The authors have adequately addressed the limitations of their work, particularly in terms of the assumptions required for the theoretical analysis. However, more discussion on the potential practical limitations and how they can be mitigated would be beneficial.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper addresses the novel problem of offline multitask representation learning in RL with a new algorithm (MORL) that incorporates innovative techniques such as joint model learning via MLE and penalty functions for pessimism. The approach is well-grounded in theory and fills a relevant gap in the literature.\n2. The submission is technically sound with rigorous theoretical analysis supporting the proposed algorithm. Detailed proofs of the main theoretical results and thorough explanations of the algorithm's components are provided.\n3. The paper is clearly written and well-organized, with a logical progression from problem introduction to algorithm presentation and theoretical results. The definitions and assumptions are clearly stated, and the proofs are detailed and comprehensible.\n4. The results are important as they provide the first theoretical demonstration of the benefits of multitask representation learning in offline RL. The approach has potential applications in domains where collecting data online is infeasible, making it highly relevant for real-world problems.", "weaknesses": "1. The empirical validation through experiments is limited, lacking comparisons with other state-of-the-art algorithms to support the theoretical claims.\n2. The theoretical proofs need additional explanations or visual aids, with a high-level overview before detailed proofs to improve accessibility for a broader audience.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points. \n\n*  As discussed in Additional Related Work in Appendix A under **Offline Data Sharing in RL**, there have been numerous empirical works that studied the benefit of using offline datasets from multiple tasks to accelerate downstream learning. The aim of this paper is to provide a framework for investigating this approach from a theoretical lens. While empirical validation is limited in this paper, we hope that the theoretical analysis presented in this work will provide insights for algorithmic design principles that can be used to design empirically better performing algorithms compared to current state-of-the-art algorithms. \n\n*  Thanks a lot for your suggestion regarding improving the accessibility of the proofs for a broader audience. We will add a proof roadmap section along with a visual proof roadmap in our next version. Again, thanks a lot for the great suggestion!\n\n### Questions\n\n1.  As can be observed from Equation 3.7, the sub-optimality gap (in average sense) of MORL does not depend on the size of the state and action space. Instead it depends on $d$ - the dimension of the feature $\\phi$. Moreover, the average accuracy of the estimated transition kernels of the $T$ MDPs scales at the rate $\\tilde{O}(\\sqrt{\\frac{\\log |\\Phi|}{nT} + \\frac{\\log|\\Psi|}{n}})$.\n\n2. In our setting we assumed that the different tasks share the feature representation $\\phi$ under low-rank MDP structure. However, in most real-world datasets this might be difficult to satisfy exactly. It would be an interesting future work to study how we can extend this work to the case where the features are not exactly shared but instead shared with some perturbed approximation.\n\nFinally, we are glad that the reviewer felt that we have adequately addressed the limitations of the work in terms of the assumptions used for theoretical analysis. As per the suggestion of the reviewer, we will add more discussion on potential practical limitations of MORL along with how to possibly mitigate them in our next version.\n\nWe hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?", "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": "7OvT8my0aX", "reviewer": "Reviewer_oHwh", "summary": "This paper investigates how representation learning in offline multitask low-rank RL can improve sample complexity when using the learned representations in downstream reward-free  RL, offline and online RL settings.\nThe paper assumes that the new task shares the same representation as the upstream tasks making up the offline data set.\nThe paper provides theoretical results but does not provide any empirical evaluation.\nRather intuitively, their result shows that learning a representation from the offline multitask data set can improve sample complexity in new downstream tasks when compared to learning the new downstream task from scratch.", "questions": "- Why is $\\phi$ shared and $\\mu$ is not? This should be made clear in the text.\n- What is $\\lambda$ in Equation 3.3?\n- Why are equations not referenced as Equation X but instead as (X)? You could be referring to an equation/section etc...", "limitations": "I do not think the limitations are adequately discussed.\nThe proofs utilise a lot of assumptions and the paper would benefit from a clear paragraph demonstrating the limitations associated with these assumptions. What are the practical implications of these assumptions?", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Representation learning and large-scale offline data sets have shown much promise for learning policies.\nAs such, investigating the theory of learning representations from multi-task offline data seems an important research direction.\nI found the related work section very beneficial for pulling different pieces of the paper together. This really helped me understand things.", "weaknesses": "I should start off by saying that I have little experience with the theoretical aspects of representation learning for RL and there were many aspects of this paper I had to take as a given without fully understanding.\nI have put a confidence of 2 to reflect his.\n\nWhilst this paper addresses an important and interesting problem, I did find it hard to follow.\nThis is likely due to the fact I do not research theoretical aspects of RL.\nNevertheless, I do think the paper should be digestible to a wider audience than it currently is.\nIn particular, the authors could do a better job of communicating intuition for the different equations that are introduced.\nI will now detail some suggestions with the aim of helping the authors communicate their work more clearly.\n\nAt the end of Section 3, it was not clear to me what the main result is. I would advise the authors to summarise the main result at the end of the section.\nThis would really help make this paper an easier read.\nSimilarly in Section 4.\nWhat is the main takeaway from this section?\n\nEquation 3.4 is introduced rather abruptly. That is, I did not understand its purpose until later in the paper when the authors stated that they wanted to improve sample complexity by improving exploration.\nI would advise the authors to state this reasoning when they introduce Equation 3.4.\n\nI am confused by the notation for $h \\in [H]$, $t \\in [T]$ and $i \\in [N]$. Aren't these discrete numbers? If so, the authors should use curly braces instead, e.g. $h \\in \\\\{0,\\ldots,H\\\\}$\n\n\n## Minor comments and typos \n- Line 61 - $H$ and $d$ aren't defined yet so this doesn't make much sense. I found the contributions had way too many technical details which does not make any sense to the reader given what they've read so far. Consider shortening these and removing unnecessary technical details.\n- Equation 2.1 - What is $\\mathbb{P}$? This is not defined until way later in the paper.\n- Line 42 - \"We develop new\" should be \"We develop a new\"\n- Line 69 - MDPs are discrete time by definition so it is confusing to state you consider discrete-time MDPs. It almost implied MDPs are not usually discrete time.\n- Sections shouldn't lead straight into sub-sections, e.g. Section 2 to 2.1 and Section 4 to 4.1\n    - You should add a few sentences explaining what you are introducing in this section.\n- Line 137 - \"was\" should be \"were\"\n- Line 152 - \"we construct lower confidence\" should be \"we construct a lower confidence\"\n- Line 159 - \"can\" should be \"has\"\n- Line 177 - \"compared to concentrability coefficient\" should be \"compared to the concentrability coefficient\"\n- Line 215 - \"sequel\" is the wrong word here. It sounds like you mean the next paper...\n- Line 233 - \"task\" should be \"tasks\"\n- Line 233 - \"assume\" should be \"assumes\"\n- Line 233 - \"for reward-free\" should be \"for the reward-free\"\n- Table 1 - It is not clear from the caption that this is for the reward-free setting. This should be made clear in the caption.\n- Line 299 - \"such\" should be \"this\"\n- Line 343 - \"study the multitask RL\" should be \"study multitask RL\"\n- Line 347 - \"setting\" should be \"settings\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points. \n\nThank you so much for your suggestion on how to improve the readability for a broader audience. We will definitely add more intuition for each of the different equations that are presented in the main paper. \n\n### Main take-away for Section 3\nThe main take-away for the main result (Theorem 3.3) in Section 3 is highlighted in L187-193. Namely, in Theorem 3.3, Equation (3.6) provides the statistical scaling rate for the average accuracy of the estimated transition kernels of the $T$ MDPs. Equation (3.7) implies that  if the optimal policy $\\pi_t^*$ is covered by the offline data for all $t \\in$ {1, 2, ..., T} as characterized by multitask relative condition number $C^*$, then the output policy of MORL is able to compete against it on average as well. \n\n### Main take-away for Section 4\nOn high-level the main take-away for section 4 is that representation learning in upstream multitask offline dataset can help in improving the sample complexity, statistical rate of sub-optimality gap and regret bound in reward-free, offline and online RL respectively. \n\nThe main take-away for Theorem 4.4 is provided in Remark 4.5. Due to space constraint in the main paper, the main take-away for Theorem 4.7 is provided in Remark K.7 (L1042-L1046) in Appendix K. Similarly, the main take-away for Theorem 4.8 is provided in Remark M.6 (L1140-L1145) in Appendix M. \n\n### Other comments\nThanks for suggesting to put the purpose of introducing Equation 3.4 earlier. We will do it in our next version. Essentially, Equation 3.4 is used for inducing pessimism in the learned policy in Algorithm 1. The notation $[N]$ indicates the set {1, 2, ..., N}. Even though it’s a standard notation, we will clarify it in our next version. Thanks for the other minor comments and pointing out some typos. We will correct and update them.\n\n### Questions\n1. By assuming $\\mu$ is not shared across tasks, we ensure that the transition kernel of each task is different from one another. However, for theoretical analysis it is imperative to make an assumption that will connect all the MDP tasks in some manner. This is why we assume that the feature $\\phi$ is shared across tasks.\n\n2. Here $\\lambda$ is a regularization factor in Equation 3.3. \n\n3. Thanks for your suggestion regarding equations as Equation X instead of as (X). We will update this in our next version as per your suggestion.\n\n### Limitations\n\nWe have already discussed the assumptions made at length in L230-L241. We will add further discussion on how these assumptions can be limiting. \n\nWe hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?", "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:

Meta-review: Although there has been much empirical work in representation learning for offline multitask RL, there has been relatively little theoretical work. This paper helps close this gap by developing a novel theoretical framework to address this problem, specifically focusing on low-rank MDPs. The paper overall seems technically solid, makes good contributions to the field, and is likely to have a solid impact. Its framework is well developed and supported by extensive proofs. As pointed out in the reviews, there are two shortcomings of this paper that would improve it — one semi-major and one minor: 1) the semi-major issue is that the paper lacks empirical validation of the theoretical ideas. Adding this would help connect the work to the well-established and complementary empirical literature on the same topic, and help establish some practical implications of the developed theory. 2) As a minor point, the paper would benefit from more intuitive explanations and illustrations to help broaden its appeal beyond the pure theory audience.
Decision: Accept (poster)