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_cc_2024_Conference
Published: 2024-05-14 07:46:35

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.

Main paper: 
 introduction: 
 Recent advances in offline reinforcement learning (RL) (Levine et al., 2020) have opened up possibilities for training policies for real-world problems using pre-collected datasets, such as 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 et 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 on single-task problems, there are many practical scenarios where multiple tasks are correlated and it is beneficial to learn multiple tasks jointly by utilizing all of the data available (Kalashnikov et al., 2018; Yu et al., 2021 Yu et al., , 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 a shared low-dimensional representation function among related tasks and then using a 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 et al., 2022) , the theoretical understanding of it is still in its early stages (Brunskill and Li, 2013; Calandriello et al., 2014; Arora et al., 2020; D'Eramo et al., 2020; Hu et al., 2021; Lu et al., 2021; Müller and Pacchiano, 2022) . Although some works • We then show theoretical benefits of using the learned representation from MORL in downstream reward-free RL Jin et al. (2020a) ; Wang et al. (2020) . In particular, we show that, to guarantee an ϵ-suboptimal policy for uniformly over any reward function, our algorithm requires at most O( H 4 d 3 ϵ 2 ) episodes during the exploration phase where d is the dimension of the feature and H 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 a factor of O(HdK), where K is the cardinality of the action space. In addition, as a complementary result, we show that using the learned representation from MORL improves the suboptimality gap bound in both offline and online downstream task. 
preliminary: 
 Episodic MDP. We consider an episodic discrete-time Markov Decision Process (MDP), denoted by M = (S, A, H, P, r), where S is the state space, A is the action space with cardinality K, H is the finite episode length, P = {P h } H h=1 are the state transition probability distributions with P h : S × A → ∆(S), and r = {r h } H h=1 are the deterministic reward functions with r h : S × A → [0, 1]. Following prior work (Jiang et al., 2017; Sun et al., 2019) , we assume that the initial state s 1 is fixed for each episode. A policy π is a collection of H functions {π h : S → A} h∈ [H] where π h (s) is the action that the agent takes at state s and at the h-th step in the episode. Given a starting state s h , s h ′ ∼ (P, π) denotes a state sampled by executing policy π under the transition model P for h ′ -h steps and E (s h ,a h )∼(P,π) [•] denotes the expectation over states s h ∼ (P, π) and actions a h ∼ π. Moreover, for each h ∈ [H], we define the value function under policy π when starting from an arbitrary state s h = s at the h-th time step as V π h,P,r (s) = E (s h ′ ,a h ′ )∼(P,π)H h ′ =h r h ′ (s h ′ , a h ′ )|s h = s .We define the action-value function for a given state-action pair (s, a) under policy π at step h as Q π h,P,r (s, a) = E (s h ′ ,a h ′ )∼(P,π)H h ′ =h r h ′ (s h ′ , a h ′ )|s h = s, a h = a .Defining (P h f )(s, a) = E s ′ ∼P (•|s,a) [f (s ′ )] for any function f : S → R, we write the Bellman equation associated with a policy π as Q π h,P,r (s, a) = (r h + P h V π h+1,P,r )(s, a), V π h,P,r (s) = Q π h,P,r (s, π h (s)), V π H+1,P,r (s) = 0. (2.1) Since the MDP begins with the same initial state s 1 , for simplicity, we use V π P,r to denote V π 1,P,r (s 1 ). Another useful concept is the notion of occupancy measure of a policy π at time step h under transition kernel P . Specifically, we use d π P h (s, a) to denote the marginal probability of encountering the state-action pair (s, a) at time step h when executing policy π under MDP with transition kernel P . Finally, we denote U(S) and U(S, A) as the uniform distribution over S and S × 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 * h : S × A → ∆(S) admits a low-rank decomposition with dimension d ∈ N if there exists two unknown embedding functions ϕ * h : S ×A → R d and µ * h : S → R d such that for all s, s ′ ∈ S and a ∈ A, P * h (s ′ | s, a) = ⟨ϕ * h (s, a), µ * h (s ′ )⟩. Without loss of generality, we assume ∥ϕ * h (s, a)∥ 2 ≤ 1 for all (s, a) ∈ S × A and for any functiong : S → [0, 1], ∥ µ * h (s)g(s)ds∥ 2 ≤ √ d.We remark that the upper bounds on the norm of ϕ * and µ * are just for normalization. As the function class Φ for ϕ * can be a 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 is assumed that the true representation ϕ * is known to the agent a priori.In offline multitask RL upstream learning, the agent is provided with an offline dataset collected from T source tasks, where the reward functions {r t } t∈[T ] are assumed to be known. Each task t ∈ [T ] is associated with a low-rank MDP M t = (S, A, H, P t , r t ). Here, all T tasks are identical except for (1) their true transition model P ( * ,t) , which admits a low-rank decomposition with dimension d: P( * ,t) h (s ′ h | s h , a h ) = ⟨ϕ * h (s ′ ), µ ( * ,t) h(s, a)⟩ for all h ∈ [H], t ∈ [T ], and (2) their reward r t h . While the tasks may differ in µ ( * ,t) h and r t h , we emphasize that all tasks share the same feature function ϕ * h . We have access to offline datasetD = t∈[T ],h∈[H] D (t) h , where D (t) h = {(s (i,t) h , a (i,t) h , r (i,t) h , s (i,t) h+1 } i∈[n] with s (i,t) h+1 ∼ P ( * ,t) h (• | s (i,t) h , a (i,t) h ) and D (t)h was collected using a fixed behavior policy π b t . In the upstream learning stage, the goal is to find a near-optimal policy and a near-accurate model for any task t ∈ [T ] and any reward function {r t } t∈[T ] through the use of offline dataset D and provide a 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 π b 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 π b t globally covers the the occupancy measure under any possible policies, i.e., the concentrability ratio satisfies, max π,s,a d π P ( * ,t) h (s, a)/d π b t P ( * ,t) h (s, a) < ∞. Instead, we make a partial coverage assumption and our suboptimality bound scales with the relative condition number (Agarwal et al., 2020a (Agarwal et al., , 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 RL (Uehara and Sun, 2021; Uehara et al., 2022) and generalize it to offline multitask setting. Get policy π t = argmax π V P (t) ,r t -b (t) 8: end for 9: Output: ϕ, P (1) , . . . , P (T ) , π 1 , . . . , π T In downstream learning stage, a new target task T + 1 with a low-rank transition kernel P ( * ,T +1) and the same S, A and H is assigned to the agent. The transition kernel P ( * ,T +1) shares the same representation ϕ * with the T upstream tasks, but has a 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 Jin et al. (2020a) , the agent first interacts with the new task environment without accessing the reward function in the exploration phase for up to K RFE episodes. Afterwards, it is provided with a reward function r = {r h } H h=1 and asked to output an ϵ-optimal policy π in the planning phase. We define the sample complexity to be the number of episodes K RFE required in the exploration phase to output an ϵoptimal policy π in the planning phase for any given reward function r.In the offline and online setting the downstream task T + 1 is already assigned with an unknown reward function r T +1 and the goal is to find a 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 an offline dataset D off = h∈[H] D h , where D h = {(s τ h , a τ h , r τ h , s τ h+1 )} τ ∈[Noff] and D off were collected using some behavior policy ρ. 
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.The details of the algorithm MORL is depicted in Algorithm 1. The agent passes all input offline data to estimate low-rank components ϕ h , µh , . . . , µ(T ) h simultaneously via the Maximum Likelihood Estimation (MLE) oracle M LE t∈[T ] D (t) hon the joint distribution defined as follows:EQUATIONh , . . . , µ(T ) h = argmax ϕ h ∈Φ, µ (1) h ,...,µ (T ) h ∈Ψ n i=1 T t=1 log ϕ h (s (i,t) h , a (i,t) h ), µ t h (s (i,t) h+1 . (3.1)The MLE oracle in (3.1) is the offline multitask counterpart to the celebrated MLE oracle in the online multitask RL (Agarwal et al., 2020b; Cheng et al., 2022; Agarwal et al., 2023) . The MLE oracle can be reasonably approximated in practice whenever optimizing over Φ and Ψ is feasible through proper parameterization such as by neural network. For each task t, we obtain the estimated transition kernel P (t) at each step h using the learned embeddings ϕ h , µh :P (t) h (s ′ | s, a) = ⟨ ϕ h (s, a), µ (t) h (s ′ )⟩. (3.2)Using the representation estimator ϕ h , we set the empirical covariance matrixΣ (t) h, ϕfor task t asΣ (t) h, ϕ = n i=1 ϕ h (s (i,t) h , a (i,t) h ) ϕ h (s (i,t) h , a (i,t) h ) ⊤ + λI. (3.3) Using both ϕ h and Σ (t) h, ϕ, we construct a lower confidence bound penalty term as follows:b (t) h (s h , a h ) = min α∥ ϕ h (s h , a h )∥ ( Σ (t) h, ϕ ) -1 , 1 , (3.4)where α is a pre-determined parameter.Finally, for each task t, with the learned model P (t) and the reward r t -b (t) , we do planning to get policy π t .To facilitate the model selection task using the joint MLE oracle in (3.1), we posit a 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 {Φ, Ψ} that contains the true model, i.e., for any h∈ [H], t ∈ [T ], the embeddings ϕ * h ∈ Φ, µ ( * ,t) ∈ Ψ.For normalization, we assume that for any ϕ ∈ Φ, ∥ϕ(s, a)∥ 2 ≤ 1 and for any µ ∈ Ψ and any function g :S :→ [0, 1], ∥ µ h (s)g(s)ds∥ 2 ≤ √ d.For simplicity, we assume that the cardinality of the function classes Φ and Ψ are finite.Next, we define the multitask relative condition number C * , which is a natural extension of the standard relative condition number (Agarwal et al., 2020a (Agarwal et al., , 2021;; Uehara et al., 2022) .Definition 3.2 (Multi-task relative condition number). For task t and time step h, we define C * t,h (π t , π b t ) as the relative condition number under ϕ * h :C * t,h (π t , π b t ) := sup x∈R d x ⊤ E (s h ,a h )∼(P ( * ,t) ,πt) [ϕ * h (s h , a h )ϕ * h (s h , a h ) ⊤ ]x x ⊤ E (s h ,a h )∼(P ( * ,t) ,π b t ) [ϕ * h (s h , a h )ϕ * h (s h , a h ) ⊤ ]x . (3.5) We define C * t := max h∈[H] C * t,h (π t , π b t ) and C * := max t∈[T ] C * t .Intuitively, C * t,h (π t , π b t ) defined in (3.5) measures the deviation between a comparator policy π t and the behavior policy π b t at time step h. When tabular MDP is considered (i.e., ϕ * h is a one-hot encoding vector), this relative condition number reduces to the density ratio based single-policy concentrability coefficient, C * t,h,∞ (π t , π b t ) = maxP (t) h (• | s h , a h ) -P ( * ,t) h (• | s h , a h ) T V ≤ 2 log(2|Φ||Ψ| T nH/δ) nT ,(3.6) where ϕ, P (1) , . . . , P (T ) be the output of Algorithm 1.(b) In addition, in Algorithm 1, if we set α = √ 2nωζ n + λd, λ = cd log(|Φ||Ψ| T nH/δ) with ζ n := 2 log(2|Φ||Ψ| T nH/δ) nand c being a constant, where we assume that ω := max t max s,a (1/π b t (a | s)) < ∞, then under Assumption 3.1, with probability at least 1 -δ, we have 1 Theorem 3.3 (a) shows a 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)). In particular, to obtain an ϵ-suboptimal transition kernels, it suffices for the joint learning to useT T t=1 V πt P ( * ,t) ,r t -V πt P ( * ,t) ,r t ≤ ωαdH C * n +2dH 2 λC * n +ωH 2 dC * ζ n T +α d n +2H ωζ n T , (3.7) where { π t } t∈[T ] is the output of Algorithm 1.O log |Φ| T ϵ 2 + log |Ψ| ϵ 2samples per task, yet for the independent learning to useO log |Φ| ϵ 2 + log |Ψ| ϵ 2samples per task. This benefit naturally comes from the inductive bias that all the source tasks share a representation in Φ which can be learned more accurately with the aggregated data pooled from all the source tasks. Equation (3.7) in Theorem 3.3 further shows that, for all tasks on average, we can uniformly compete with any set of comparator policies {π t } t∈[T ] satisfying the partial coverage through C * < ∞. In particular, if the optimal policy π * t is covered by the offline data for all t ∈ [T ], then the output { π t } t∈[T ] of Algorithm 1 is able to compete against it on average as well. Remark 3.4. We note that in order for the bound to hold in Theorem 3.3, Algorithm 1 requires the knowledge of ω as it is required to set the value of α which is used in the lower confidence bound penalty term defined in (3.4). However, in practice, we expect that α can be treated as a 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 one-step back lemma (Lemma G.1), we first develop a new upper bound on model estimation error for each task which encapsulates the benefit of joint offline MLE model estimation over single-task offline learning. We then use the following lemma to show near pessimism in the average sense. Lemma 3.5. For any policy π t and reward r t , we have, with probability 1 -δ1 T T t=1 V πt P (t) ,r t -b (t) -V πt P ( * ,t) ,r t ≤ H ωζ n /T ,whereζ n := 2 log(2|Φ||Ψ| T nH/δ) 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 ∥ϕ* h (s h , a h )∥ (Σ h,π b t ,ϕ * ) -1 where Σ h,π b t ,ϕ * = nE (s h ,a h )∼(P ( * ,t) ,π b t ) [ϕ * ϕ * ⊤ ]+ λI using a one-step back inequality (Lemma G.1) and a distribution shift lemma (Lemma O.2).4 Downstream RL: Reward-free Exploration, Offline RL and Online RLP 1 (s ′ |s, a) = ⟨ϕ 1 (s, a), µ 1 (s ′ )⟩ and P 2 (s ′ |s, a) = ⟨ϕ 2 (s, a), µ 2 (s ′ )⟩ in the model class Φ × Ψ, we have, ∥P 1 (•|s, a) -P 2 (•|s, a)∥ TV ≤ C R E (s,a)∼U (S,A) ∥P 1 (•|s, a) -P 2 (•|s, a)∥ TV , for all (s, a) ∈ S × A and h ∈ [H] where C R is an absolute constant.4. The transition kernel of task T + 1, P ( * ,T +1) can be ξ-approximated by a linear combination of the T source upstream tasks, i.e. there exist T unknown coefficients c 1 , . . . , c T ≥ 0 such that T t=1 c t ≤ C L and ξ ≥ 0 such that for all (s, a) ∈ S × A and h ∈ [H], we have∥P ( * ,T +1) (•|s, a) - T t=1 c t P ( * ,t) (•|s, a)∥ TV ≤ ξ.Here, ξ is called the linear combination misspecification.The first point in Assumption 4.1 ensures that the behavior policy π b t for each task t can reach any state in S at any time step with a positive probability, an assumption that is previously used in (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.The third point in Assumption 4.1 ensures that for each source task t, the point-wise TV error between the learned estimated transition kernel P (t) and the true transition kernel P ( * ,t) is bounded by the population-level TV error. 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 of the target task P ( * ,T +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 performance because the downstream task utilizes the estimated feature from the upstream. We use the following notion of ϵ-approximate linear MDP to provide a guarantee for the estimated feature. Definition 4.2 (ϵ-approximate linear MDP (Jin et al., 2020b; Cheng et al., 2022) ). For any ϵ > 0, we say that MDP M = (S, A, H, P, r) is an ϵ-approximate linear MDP with a feature map ϕ h :S × A :→ R d , if for any h ∈ [H], there exist d unknown (signed) measures µ h = (µ (1) h , . . . , µ (d) h ) over S such that for any (s, a) ∈ S × A, we have ∥P h (•|s, a) -⟨ϕ h (s, a), µ h (•)⟩∥ TV ≤ ϵ.(4.1) Any ϕ satisfying (4.1) is called an ϵ-approximate feature map of M.The next lemma shows that the learned feature ϕ 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 1 is a ξ down -approximate feature for MDP M T +1 whereξ down = ξ + C L C R ν κ 2T log(2|Φ||Ψ| T nH/δ) n, i.e. there exist a time-dependent unknown (signed) measure µ * over S such that for any (s, a) ∈ S × A, we have∥P ( * ,T +1) h (•|s, a) -⟨ ϕ h (s, a), µ * h (•)⟩∥ TV ≤ ξ down . Furthermore, for any g : S → [0, 1], we have ∥ µ * h (s)g(s)ds∥ 2 ≤ C L √ d. 
relationship between upstream and downstream mdps: 
 In order for us to theoretically study the downstream RL tasks where we would use the learned feature ϕ 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 1. For each upstream source task t with transition kernel P ( * ,t) , the behavior policy π b t is such thatmin s∈S P (π b t ,t) h (s) ≥ κ, where P (π b t ,t) h (•) : S →R is the marginalized state occupancy measure over S using policy π b t at time step h. 2. The state space S is compact and has finite measure 1/ν, and the induced uniform probability density function is f (s) = ν, s ∈ S.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 in our analysis. Another difference with Wang et al. (2020) is that in the exploration phase, like Chen et al. (2022b) , we construct more aggressive reward function to avoid overly-conservative 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 4.1, after collecting K RFE trajectories during the exploration phase in Algorithm 2, with probability at least 1 -δ, the output of Algorithm 3, policy π satisfiesE s1∼µ [V * 1 (s 1 , r) -V π 1 (s 1 , r)] ≤ c ′ d 3 H 4 log(dK RFE H/δ)/K RFE + 6H 2 ξ down . (4.2)If the linear combination misspecification error ξ in Assumption 4.1 satisfies O( d 3 /K RFE ) and the number of trajectories in the offline dataset for each upstream task is at least O(T K RFE /d 3 ), then, provided K RFE is at least O(H 4 d 3 log(dHδ -1 ϵ -1 )/ϵ2 ), with probability 1 -δ, the policy π will be an ϵ-optimal policy for any given reward during the planning phase.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) ). 2 Cheng et al. (2023) further shows that the dependence of 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, as we see 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 O(HdK) compared to that of RAFFLE. Moreover, compared to standard linear MDP, where the true representation {ϕ h } H h=1 is known (Jin et al., 2020b) , the suboptimality gap in (4.2) contains an additional term H 2 ξ down which is due to the upstream misspecification error ξ down . When ξ down is small enough, our resulting sample complexity of O( H 4 d 3 ϵ 2 ) matches the reward-free exploration sample complexity for linear mixture MDPs (Chen et al., 2022b) O( H 4 d 3 ϵ 2 ) Multi-taskTable 1 : Sample complexities of different approaches to learning an ϵ-optimal policy for the rewardfree RL setting with low-rank MDPs.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 r T +1 in the downstream task T + 1 is linear with respect to the unknown feature ϕ * : S × A → R d . We emphasize that, unlike Cheng et al. (2022) , we assume the reward function r T +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 et 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 et al., 2021; Yin and Wang, 2021) . Assumption 4.5 (Feature coverage). There exists an absolute constant κ ρ such that for all h ∈ [H] andϕ h ∈ Φ h , λ min (E ρ [ϕ h (s h , a h )ϕ h (s h , a h ) ⊤ |s 1 = s]) ≥ κ ρ .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λ d = 1, β = O(Hd √ ι + H dN off ξ down ), where ι = log(HdN off max(ξ down , 1)/δ), with probability at least 1-δ, the suboptimality gap of Algorithm 4 is at mostV π * P ( * ,T +1) ,r (s) -V π P ( * ,T +1) ,r (s) ≤ 2H 2 ξ down + 2β H h=1 E π * ϕ h (s h , a h ) Λ -1 h s 1 = s . (4.3)Additionally if Assumption 4.5 holds, and the sample size satisfies N off ≥ 40/κ ρ • log(4dH/δ), then with probability 1 -δ, we have,V π * P ( * ,T +1) ,r (s) -V π P ( * ,T +1) ,r (s) ≤ O κ -1/2 ρ H 2 d log(HdNoff max(ξdown, 1)/δ) Noff + κ -1/2 ρ H 2 d 1/2 ξdown .(4.4)Online RL. For downstream online RL task, where the agent is allowed to interact with the new task MDP M T +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 π be the uniform mixture of π 1 , . . . , π Non in Algorithm 5. Under Assumption 4.1, setting λ = 1,β n = O(Hd ι n (δ) + H √ dnξ down + C L √ Hd), where ι n = log(Hdn max(ξ down , 1)/δ), with probability 1 -δ, the suboptimality gap of Algorithm 5 satisfiesV * P ( * ,T +1) ,r -V π P ( * ,T +1) ,r ≤ O(H 2 d 3/2 N -1/2on + H 2 dξ down ). 
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 a policy from a 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 et 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 single-task offline RL, including model-based approach (Rashidinejad et al., 2021; Uehara and Sun, 2022; Jin et al., 2021; Yu et al., 2020; Xie et al., 2021b; Uehara 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. 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 to perform offline multitask pretraining of a world model which is then finetuned on a downstream target task. Hu et al. (2023) proposes a provably efficient 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 (Uehara 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, we provide results in downstream offline and online RL setting as a complementary result. Moreover, unlike us, Cheng et al. (2022) assumes that the reward-function is known in the downstream task, which is a fairly strong assumption. Somewhat in a contrived manner, in Cheng et al. (2022) , the reward-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 to be within [0, 1]. We do not make this assumption for fair comparison to literature on reward-free RL for low-rank MDPs.Comparison to Concurrent Work. In a 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-wise 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. 
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 nearoptimal policy of a 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.  
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: It is a theory paper.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses 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 not use 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: Reward-Free Exploration (RFE). In reward-free RL setting, the agent does not have access to a reward function during exploration phase. However, the agent must propose a 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 in 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 et al., 2021; Ménard et al., 2021) and linear function approximation (Wang et al., 2020; Zanette et al., 2020; Zhang et al., 2021; Hu et al., 2022; Wagenmaker et al., 2022) K RFE ← c K • d 3 H 4 log(dHδ -1 ε -1 )/ε 2 for some c K > 0 3: β ← C L H √ d + dH log(dK RFE H max(ξ down , 1)/δ) + Hξ down √ dK RFE 4: for k = 1, . . . , K RFE do 5: Q k H+1 (•, •) ← 0, V k H+1 (•) ← 0 6: for h = H, H -1, . . . , 1 do 7: Λ k h = k-1 τ =1 ϕ h (s τ h , a τ h ) ϕ h (s τ h , a τ h ) ⊤ + I d 8: u k h (•, •) ← β ϕ(•, •) ⊤ (Λ k h ) -1 ϕ(•, •) 9:Define the exploration-driven reward functionr k h (•, •) ← u k h (•, •) 10: w k h = (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) V k h+1 (s τ h+1 ) 11: Q k h (•, •) = min{ ϕ h (•, •) ⊤ w k h + r k h (•, •) + u k h (•, •), H} 12: V k h (•) = max a∈A Q k h (•, a)13: Take actionπ k h (•) = argmax a∈A Q k h (•,a k h = π k h (s k h ) and observe s k h+1 ∼ P ( * ,T +1) h (s k h , a k h ) 18:end for 19: end for 20: Output:D RFE ← {(s k h , a k h )} (k,h)∈[K]×[H]Algorithm 3 Downstream Reward-Free Exploration: Planning Phase1: Input: Feature ϕ, dataset D RFE = {(s k h , a k h )} k∈[KRFE],h∈[H] , reward functions r = {r h } h∈[H] 2: Initialization: Q H+1 (•, •) ← 0 and V H+1 (•) ← 0 3: for h = H, H -1, . . . , 1 do 4: Λ h = KRFE τ =1 ϕ h (s τ h , a τ h ) ϕ h (s τ h , a τ h ) ⊤ + I d 5: Let u h (•, •) ← min β ϕ(•, •) ⊤ (Λ h ) -1 ϕ(•, •), H 6: w h ← Λ -1 h KRFE τ =1 ϕ h (s τ h , a τ h ) V h+1 (s τ h+1 ) 7: Q h (•, •) ← min{ ϕ h (•, •) ⊤ w h + r h (•, •) + u h (•, •), H} and V h (•) ← max a∈A Q h (•, a) 8: π h (•) ← argmax a∈A Q h (•, a) 9: end for 10: Output: π = { π h } h∈[H]B.2 Algorithm for Downstream Offline RL Algorithm 4 Pessimistic Value Iteration (PEVI) with Approximate Feature (Jin et al., 2021) 1:Input: Feature ϕ, dataset D down = {(s τ h , a τ h , r τ h , s τ h+1 )} τ ∈[Noff],h∈[H] , parameters λ d , β, ξ down 2: Initialization: V H+1 = 0 3: for h = H, H -1, . . . , 1 do 4: Λ h = Noff τ =1 ϕ h (s τ h , a τ h ) ϕ h (s τ h , a τ h ) ⊤ + λ d I d 5: w h = Λ -1 h ( Noff τ =1 ϕ h (s τ h , a τ h ) • (r τ h + V h+1 (s τ h+1 ))) 6: Γ h (•, •) = Hξ down + β[ ϕ ⊤ h Λ -1 h ϕ h (•, •)] 1/2 7: Q h (•, •) = min{ ϕ h (•, •) ⊤ w h -Γ h (•, •), H -h + 1} + 8: π h (•|•) = argmax π h ⟨ Q h (•, •), π h (•|•)⟩ A 9: V h (•) = ⟨ Q h (•, •), π h (•|•)⟩ A 10: end for 11: Output: { π h } H h=1 B.3 Algorithm for Downstream Online RL Algorithm 5 LSVI-UCB with Approximate Feature (Jin et al., 2020b) 1: Input: Feature ϕ, parameters λ d , β n , ξ down 2: for n = 1, . . . , N on do 3:Receive the initial state s n 1 = s 1 4:for h = H, H -1, . . . , 1 do 5: Λ n h = n-1 τ =1 ϕ h (s τ h , a τ h ) ϕ h (s τ h , a τ h ) ⊤ + λ d I d 6: w n h = (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h )[r h (s τ h , a τ h ) + V n h+1 (s τ h+1 )] 7: Q n h (•, •) = min{ ϕ h (•, •) ⊤ w n h + β n ∥ ϕ(•, •)∥ (Λ n h ) -1 , H -h + 1} + 8: V n h (•) = max a∈A Q n h (•, a) 9:end for 10:π n h (•) = argmax a∈A Q n h (•, a)11:for h = 1, . . . , H do 12:Take action a n h = π n (s n h ) and observe s n h+1 13:end for 14: end for 15: Output: π 1 , . . . , π Non and π where π is the uniform mixture of π 1 , . . . , π NonWe summarize frequently used notations in the following list.f (t) h (s, a) ∥ P (t) h (• | s, a) -P (⋆,t) h (• | s, a)∥ T V ζ n 2 log(2|Φ||Ψ| T nH/δ) n ζ (t) h E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 ζ (t) 1 E (s1,a1)∼(P ( * ,t) ,π b t ) f (t) 1 (s 1 , a 1 ) 2 α (t) h nωζ (t) h + λd + nζ (t) h-1 α √ 2nωζ n + λd Σ h,π b t ,ϕ nE (s h ,a h )∼(P ( * ,t) ,π b t ) [ϕϕ ⊤ ] + λI Σ (t) h,ϕ n i=1 ϕ h (s (i,t) h , a (i,t) h )ϕ h (s (i,t) h , a (i,t) h ) ⊤ + λI b (t) h (s h , a h ) min α∥ ϕ h (s h , a h )∥ Σ (t) h, ϕ , 1 ω t max s,a (1/π b t (a | s)) ω max t ω t ξ down ξ + C L C R ν κ 2T log(2|Φ||Ψ| T nH/δ) n µ * (•) T t=1 c t µ (t) (•) w * h s ′ µ * (s ′ ) V h+1 (s ′ )ds ′ For any h ∈ [H], we define P ( * ,T +1) h (•|s, a) = ⟨ϕ * h (s, a), µ ( * ,T +1) h (•)⟩, P h (•|s, a) = ⟨ ϕ h (s, a), µ * h (•)⟩.Given a reward function r, for any function f : S → R and h ∈ [H], we define the transition operators and their corresponding Bellman operators as follows(P ( * ,T +1) h f )(s, a) = s ′ ⟨ϕ * h (s, a), µ ( * ,T +1) h (s ′ )⟩f (s ′ )ds ′ , (B h f )(s, a) = r h (s, a) + (P ( * ,T +1) h f )(s, a) (P h f )(s, a) = s ′ ⟨ ϕ h (s, a), µ * h (s ′ )⟩f (s ′ )ds ′ , (B h f )(s, a) = r h (s, a) + (P h f )(s, a)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.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.1. For any task t, policy π t , and reward r t , we have for all h ≥ 2,E (s h ,a h )∼( P (t) ,πt) f (t) h (s h , a h ) ≤ E (s h-1 ,a h-1 )∼( P (t) ,πt) α (t) h ϕ h-1 (s h-1 , a h-1 ) (Σ h-1,π b t , ϕ ) -1, and for h = 1, we haveE a1∼πt f (t) 1 (s 1 , a 1 ) ≤ ω t ζ (t) 1 ,whereω t = max s,a (1/π b t (a | s)).In the next lemma, we prove that V πt P (t) ,r t -b (t) is an almost pessimistic estimator of V πt P ( * ,t) ,r t in the average sense. Lemma D.2 (Restatement of Lemma 3.5). For any policy π t and reward r t , we have, with probability 1 -δ 1 T T t=1 V πt P (t) ,r t -b (t) -V πt P ( * ,t) ,r t ≤ H ωζ n /T ,P (t) h (• | s h , a h ) -P ( * ,t) h (• | s h , a h ) T V ≤ 2 log(2|Φ||Ψ| T nH/δ) nT ,(D.1) where ϕ, P (1) , . . . , P (T ) be the output of Algorithm 1.Algorithm 1, if we set α = √ 2nωζ n + λd, λ = cd log(|Φ||Ψ| T nH/δ) with ζ n := 2 log(2|Φ||Ψ| T nH/δ) nand c being a constant term, where we assume that ω := max t max s,a (1/π b t (a | s)) < ∞, then under Assumption 3.1, with probability at least 1 -δ, we have1 T T t=1 V πt P ( * ,t) ,r t -V πt P ( * ,t) ,r t ≤ ωαdH C * n +2dH 2 λC * n +ωH 2 dC * ζ n T +α d n +2H ωζ n T , (D.2) where { π t } t∈[T ] is the output of Algorithm 1.Proof of Theorem D.3. As in Lemma D.2, we condition on the events:T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 ≤ ζ n , where ζ n := 2 log(2|Φ||Ψ| T nH/δ) n , (D.3) and ∀ϕ ∈ Φ : ∥ϕ h (s, a)∥ ( Σ (t) h,ϕ ) -1 = Θ ∥ϕ h (s, a)∥ (Σ h,π b t ,ϕ ) -1 . (D.4)From Lemma F.3 and Lemma G.2, this event happens with probability 1 -δ. Conditioning on this event, we have1 T T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) P (t) h (• | s h , a h ) -P ( * ,t) h (• | s h , a h ) T V = 1 T T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) (i) ≤ 1 T T T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 (ii) ≤ 1 T T T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 (iii) ≤ ζ n T = 2 log(2|Φ||Ψ| T nH/δ) nT ,where (i) follows from Cauchy-Schwarz inequality, (ii) follows from Jensen's inequality and (iii) 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 π 1 , . . . , π T , we haveT t=1 V πt P ( * ,t) ,r t -V πt P ( * ,t) ,r t = T t=1 V πt P ( * ,t) ,r t -V πt P (t) ,r t -b (t) + V πt P (t) ,r t -b (t) -V πt P ( * ,t) ,r t (i) ≤ T t=1 V πt P ( * ,t) ,r t -V πt P (t) ,r t -b (t) + V πt P (t) ,r t -b (t) -V πt P ( * ,t) ,r t (ii) ≤ T t=1 H h=1 E (s h ,a h )∼(P ( * ,t) ,πt) b (t) h (s h , a h ) + (P ( * ,t) h -P (t) h )V πt h+1, P (t) ,r t -b (t) (s h , a h ) + T t=1 V πt P (t) ,r t -b (t) -V πt P ( * ,t) ,r t (iii) ≤ T t=1 H h=1 E (s h ,a h )∼(P ( * ,t) ,πt) b (t) h (s h , a h ) + H T t=1 H h=1 E (s h ,a h )∼(P ( * ,t) ,πt) f (t) h (s h , a h ) + H ωT ζ n , (D.5)where (i) follows from the observation that π t is the argmax over all Markovian policies as well as all history-dependent policies for P (t) , (ii) follows from the simulation lemma, Lemma O.1, (iii) follows from the observation that V πt h, P (t) ,r t -b (t) ≤ H and Lemma D.2. Now, using Lemma G.1 (with setting P = P ( * ,t) and ϕ = ϕ * ) and noting that | b(t) h | ∞ ≤ 1, for h ≥ 2, we have E (s h ,a h )∼(P ( * ,t) ,πt) b (t) h (s h , a h ) ≤ E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) ∥ϕ * h-1 (s h-1 , a h-1 )∥ (Σ h-1,π b t ,ϕ * ) -1 × nωE (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s,a) 2 +λd (D.6)From (D.4), we havenE (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s, a) 2 ≤ nE (s h ,a h )∼(P (⋆,t) ,π b t ) min α 2 ∥ ϕ h (s h , a h )∥ 2 (Σ h,π b t , ϕ ) -1 , 1 ≤ nE (s h ,a h )∼(P (⋆,t) ,π b t ) α 2 ∥ ϕ h (s h , a h )∥ 2 (Σ h,π b t , ϕ ) -1 ≤ α 2 Tr nE (s h ,a h )∼(P (⋆,t) ,π b t ) [ ϕ h ϕ ⊤ h ]{nE (s h ,a h )∼(P ( * ,t) ,π b t ) [ ϕ h ϕ ⊤ h ] + λI} -1 ≤ α 2 Tr {nE (s h ,a h )∼(P (⋆,t) ,π b t ) [ ϕ h ϕ ⊤ h ] + λI}{nE (s h ,a h )∼(P ( * ,t) ,π b t ) [ ϕ h ϕ ⊤ h ] + λI} -1 = α 2 Tr[I d ] = α 2 d (D.7)Next, we upper boundE (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) ∥ϕ * h-1 (s h-1 , a h-1 )∥ 2 (Σ h-1,π bt ,ϕ * ) -1 as the followingE (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) ∥ϕ * h-1 (s h-1 , a h-1 )∥ 2 (Σ h-1,π b t ,ϕ * ) -1 (i) ≤ C * t,h (π t , π b t )E (s h-1 ,a h-1 )∼(P ( * ,t) ,π b t ) ∥ϕ * h-1 (s h-1 , a h-1 )∥ 2 (Σ h-1,π b t ,ϕ * ) -1 (ii) ≤ dC * t,h (π t , π b t ) n ≤ dC * n , (D.8)where (i) follows from Lemma O.2 and (ii) follows from similar steps as in (D.7).Combining (D.6), (D.7) and (D.8), we getT t=1 E (s h ,a h )∼(P ( * ,t) ,πt) b (t) h (s h , a h ) ≤ T t=1 E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) ∥ϕ * h-1 (s h-1 , a h-1 )∥ (Σ h-1,π b t ,ϕ * ) -1 × nωE (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s,a) 2 +λd ≤ T t=1 E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) ∥ϕ * h-1 (s h-1 , a h-1 )∥ 2 (Σ h-1,π b t ,ϕ * ) -1 × E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) nωE (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s, a) 2 + λd ≤ T t=1 dC * n E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) nωE (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s, a) 2 + λd ≤ dC * n T t=1 nωE (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) E (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s, a) 2 + T √ λd ≤ dC * n T t=1 nω 2 E (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s, a) 2 + T √ λd = ω √ dC * T t=1 E (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s, a) 2 + T d λC * n ≤ ω √ dC * T T t=1 E (s h ,a h )∼(P (⋆,t) ,π b t ) b (t) h (s, a) 2 + T d λC * n ≤ ω √ dC * T T α 2 d n + T d λC * n = ωαT d C * n + T d λC * n . (D.9)Following similar steps as in (D.7), we can further show thatE (s1,a1)∼(P ( * ,t) ,πt) b (t) 1 (s 1 , a 1 ) ≤ E (s1,a1)∼(P ( * ,t) ,πt) b (t) 1 (s 1 , a 1 ) 2 ≤ α d n . Now noting |f (t) h | ∞ ≤ 1, for h ≥ 2,we get T t=1 E (s h ,a h )∼(P ( * ,t) ,πt) f (t) h (s h , a h ) ≤ T t=1 E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) ∥ϕ * h-1 (s h-1 , a h-1 )∥ (Σ h-1,π b t ,ϕ * ) -1 × nωE (s h ,a h )∼(P (⋆,t) ,π b t ) f (t) h (s h ,a h ) 2 +λd ≤ T t=1 E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) ∥ϕ * h-1 (s h-1 , a h-1 )∥ 2 (Σ h-1,π b t ,ϕ * ) -1 × E (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) nωE (s h ,a h )∼(P (⋆,t) ,π b t ) f (t) h (s h , a h ) 2 + λd ≤ dC * n T t=1 nωE (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) E (s h ,a h )∼(P (⋆,t) ,π b t ) f (t) h (s h , a h ) 2 + λd ≤ dC * n T t=1 nωE (s h-1 ,a h-1 )∼(P ( * ,t) ,πt) E (s h ,a h )∼(P (⋆,t) ,π b t ) f (t) h (s h , a h ) 2 + √ λd ≤ dC * n T t=1 nω 2 E (s h ,a h )∼(P (⋆,t) ,π b t ) f (t) h (s h , a h ) 2 + T d λC * n = ω √ dC * T t=1 E (s h ,a h )∼(P (⋆,t) ,π b t ) f (t) h (s h , a h ) 2 + T d λC * n ≤ ω √ dC * T T t=1 E (s h ,a h )∼(P (⋆,t) ,π b t ) f (t) h (s h , a h ) 2 + T d λC * n ≤ ω dC * T ζ n + T d λC * n (D.10) Further, note that, T t=1 E (s1,a1)∼(P ( * ,t) ,πt) f (t) 1 (s 1 , a 1 ) ≤ T t=1 E (s1,a1)∼(P ( * ,t) ,πt) f (t) 1 (s 1 , a 1 ) 2 ≤ T t=1 ωE (s1,a1)∼(P ( * ,t) ,π b t ) f (t) 1 (s 1 , a 1 ) 2 = √ ω T t=1 ζ (t) 1 ≤ ωT ζ n ,where the last inequality follows from the Cauchy-Schwarz inequality and (D.3).Finally, from (D.5) we getT t=1 V πt P ( * ,t) ,r t -V πt P ( * ,t) ,r t ≤ T t=1 H h=1 E (s h ,a h )∼(P ( * ,t) ,πt) b (t) h (s h , a h ) + H T t=1 H h=1 E (s h ,a h )∼(P ( * ,t) ,πt) f (t) h (s h , a h ) + H ωT ζ n ≤ H h=2 T t=1 E (s h ,a h )∼(P ( * ,t) ,πt) b (t) h (s h , a h ) + H H h=2 T t=1 E (s h ,a h )∼(P ( * ,t) ,πt) f (t) h (s h , a h ) + T t=1 E (s1,a1)∼(P ( * ,t) ,πt) b (t) 1 (s 1 , a 1 ) + H T t=1 E (s1,a1)∼(P ( * ,t) ,πt) f (t) 1 (s 1 , a 1 ) + H ωT ζ n ≤ HωαT d C * n + HT d λC * n + H 2 ω dC * T ζ n + H 2 T d λC * n + αT d n + 2H ωT ζ n ≤ HωαT d C * n + 2H 2 T d λC * n + H 2 ω dC * T ζ n + αT d n + 2H ωT ζ n . (D.11) So, we have 1 T T t=1 V πt P ( * ,t) ,r t -V πt P ( * ,t) ,r t ≤ Hωαd C * n + 2H 2 d λC * n + H 2 ω dC * ζ n T + α d n + 2H ωζ n T .(D.12)In this section, we provide the proofs of the lemmas that we used in the proof of Theorem D.3.Proof of Lemma D.1. For h = 1,E (s1,a1)∼( P (t) ,πt) f (t) 1 (s 1 , a 1 ) ≤ E (s1,a1)∼( P (t) ,πt) f (t) 1 (s 1 , a 1 ) 2 ≤ ω t E (s1,a1)∼(P ( * ,t) ,π b t ) f (t) 1 (s 1 , a 1 ) 2 = ω t ζ (t) 1where the first inequality follows from Jensen's inequality and the second inequality follows from importance sampling. Denotingζ (t) h = E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t)h (s, a) 2 , for h ≥ 2, we haveE s h ∼( P (t) ,π t ) a h ∼π t f (t) h (s h , a h ) (i) ≤ E (s h-1 ,a h-1 )∼( P (t) ,πt) ϕ h-1 (s h-1 , a h-1 ) (Σ h-1,π b t , ϕ ) -1 × nωtE (s h ,a h )∼(P (⋆,t) ,π b t ) [f (t) h (s h ,a h ) 2 ]+λd+nE (s h-1 ,a h-1 )∼(P (⋆,t) ,π b t ) f (t) h-1 (s h-1 ,a h-1 ) 2 (ii) = E (s h-1 ,a h-1 )∼( P (t) ,πt) nω t ζ (t) h + λd + nζ (t) h-1 ϕ h-1 (s h-1 , a h-1 ) (Σ h-1,π b t , ϕ ) -1 = E (s h-1 ,a h-1 )∼( P (t) ,πt) α (t) h ϕ h-1 (s h-1 , a h-1 ) (Σ h-1,π b t , ϕ ) -1where (i) follows from Lemma G.1 and |f(t) h (s h , a h )| ≤ 1, (ii) uses notations defined in Appendix C. E.2 Proof of Lemma D.2Proof of Lemma D.2. We condition on the events:T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 ≤ ζ n , where ζ n := 2 log(2|Φ||Ψ| T nH/δ) n , (E.1) and ∀ϕ ∈ Φ : ∥ϕ h (s, a)∥ ( Σ (t) h,ϕ ) -1 = Θ ∥ϕ h (s, a)∥ (Σ h,π b t ,ϕ ) -1 . (E.2)From Lemma F.3 and Lemma G.2, this event happens with probability 1 -δ. Conditioning on this event, we haveα (t) h = nω t ζ (t) h + λd + nζ (t) h-1 ≤ 2nωζ n + λd = α (E.3)We haveT t=1 V πt P (t) ,r t -b (t) -V πt P ( * ,t) ,r t (i) = T t=1 H h=1 E (s h ,a h )∼( P (t) ,πt) -b (t) h (s h , a h ) + (P ( * ,t) h -P (t) h )V πt h+1,P ( * ,t) ,r t (s h , a h ) (ii) ≤ H T t=1 H h=1 E (s h ,a h )∼( P (t) ,πt) -b (t) h (s h , a h ) + f (t) h (s h , a h ) (iii) ≤ H T t=1 H h=2 E (s h-1 ,a h-1 )∼( P (t) ,πt) min α (t) h ϕ h-1 (s h-1 , a h-1 ) (Σ h-1,π b t , ϕ ) -1 , 1 + H T t=1 ωζ (t) 1 + H T t=1 H h=1 E (s h ,a h )∼( P (t) ,πt) -b (t) h (s h , a h ) (iv) ≤ H T t=1 H h=2 E (s h-1 ,a h-1 )∼( P (t) ,πt) min α (t) h ϕ h-1 (s h-1 , a h-1 ) (Σ h-1,π b t , ϕ ) -1 , 1 + H ωT ζ n + H T t=1 H h=1 E (s h ,a h )∼( P (t) ,πt) -b (t) h (s h , a h ) (v) ≲ H T t=1 H h=2 E (s h-1 ,a h-1 )∼( P (t) ,πt) min α ϕ h-1 (s h-1 , a h-1 ) (Σ h-1,π b t , ϕ ) -1 , 1 + H ωT ζ n + H T t=1 H h=1 E (s h ,a h )∼( P (t) ,πt) -min α ϕ h (s h , a h ) (Σ h,π b t , ϕ ) -1 , 1 (vi) ≤ H ωT ζ n + H T t=1E (s1,a1)∼( P (t) ,πt) -min α ϕ 1 (s 1 , a 1 )(Σ 1,π b t , ϕ ) -1 , 1 ≤ H ωT ζ n .where (i) follows from Lemma O.1, (ii) follows from the observation V πt P ( * ,t) ,r t ≤ H, (iii) follows from Lemma D.1, (iv) follows from Cauchy-Schwarz inequality and the fact thatT t=1 ζ (t) h ≤ ζ n , (v) follows from (E.2).Consider a sequential conditional probability estimation setting with an instance space X and target space Y and with a conditional probability density p(y | x) = f * (x, y). We consider a function class F : (X × Y) → R for modeling the condition distribution f * , and we further assume that the realizability condition holds i.e. f * ∈ F. We are given a dataset D := {(x i , y i )} n i=1 , wherex i ∼ D i = D i (x 1:i-1 , y 1:i-1 ) and y i ∼ p(• | x i ). Let D ′ denote a tangent sequence {(x ′ i , y ′ i )} n i=1 where x ′ i ∼ D i (x 1:i-1 , y 1:i-1 ) and y ′ i ∼ p(• | x ′ i ).Note that here x ′ i 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 a dataset of n examples, and let D ′ be a tangent sequence. Let L(f, D) = n i=1 ℓ(f, (x i , y i )) be any function that decomposes additively across examples where ℓ is any function, and let f (D) be any estimator taking as input random variable D and with range F. Then Agarwal et al. (2020b) ). For any two conditional probability densities f 1 , f 2 and any distribution D ∈ ∆(X ) we haveE D exp L( f (D), D) -log E D ′ exp(L( f (D), D ′ )) -log |F| ≤ 1. Lemma F.2 (Lemma 25 ofE x∼D ∥f 1 (x, •) -f 2 (x, •)∥ 2 T V ≤ -2 log E x∼D,y∼f2(• | x) exp - 1 2 log(f 2 (x, y)/f 1 (x, y)) .Lemma F.3 (Multitask offline MLE guarantee). Given δ ∈ (0, 1), consider the transition kernels learned in Algorithm 1. For any n, h with probability at least 1 -δ/2, we haveT t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 ≤ ζ n , where ζ n := 2 log(2|Φ||Ψ| T nH/δ) n . (F.1)Proof of Lemma F.3. Let f (D) denote empirical maximum likelihood estimator:f (D) := argmax f ∈F (xi,yi)∈D log f (x i , y i )We combine Lemma F.1 with the Chernoff method to obtain the following exponential tail bound: with probability 1 -δ, we have-log E D ′ exp(L( f (D), D ′ )) ≤ -L( f (D), D) + log |F| + log(1/δ). (F.2) Now, we set L(f, D) = n i=1 -1 2 log(f * (x i , y i )/f (x i , y i )) where D = {x i , y i } n i=1 is a dataset and D ′ = {x ′ i , y ′ i } n i=1 is a tangent sequence. In the multitask offline RL setting, let x = {(s t h , a t h )} T t=1 , y = {s t h+1 } T t=1 and f (x, y) = T t=1 P t h (s t h+1 | s t h , a t h ). Then, the dataset D h can be decomposed into D h = ∪ T t=1 D (t) h where D (t) h = {(s (i,t) h , a (i,t) h , s (i,t) h+1 } i∈[n] . Similarly, D ′ h = ∪ T t=1 (D ′ h ) (t), and(D t h ) i := (D t h ) i ((s t h , a t h ) 1:i-1 , (s t h+1 ) 1:i-1). Thus, in the multitask offline RL setting, we have the cardinality |F| = |Φ||Ψ| T . With this choice, the right hand side of (F.2) isn i=1 1 2 log(f * (x i , y i )/ f (x i , y i )) + log |F| + log(1/δ) ≤ log |F| + log(1/δ) = log(|Φ||Ψ| T /δ), (F.3) where the first inequality follows because f is the empirical maximum likelihood estimator and the realizability assumption. The equality follows because |F| = |Φ||Ψ| T . On the other hand, the left hand side of (F.2) is-log E D ′ h exp n i=1 - 1 2 log f * (x ′ i , y ′ i ) f (x ′ i , y ′ i ) D h (i) = -log E D ′ h exp n i=1 - 1 2 log T t=1 P ( * ,t) h (s (i,t) h+1 | s (i,t) h , a (i,t) h ) P (t) h (s (i,t) h+1 | s (i,t) h , a (i,t) h ) D h (ii) = - T t=1 log E (D ′ h ) (t) exp n i=1 - 1 2 log P ( * ,t) h (s (i,t) h+1 | s (i,t) h , a (i,t) h ) P (t) h (s (i,t) h+1 | s (i,t) h , a (i,t) h ) D h (iii) = - T t=1 n i=1 log E (D (t) h )i exp - 1 2 log P ( * ,t) h (s (i,t) h+1 | s (i,t) h , a (i,t) h ) P (t) h (s (i,t) h+1 | s (i,t) h , a (i,t) h ) (iv) ≥ T t=1 1 2 n i=1 E (s h ,a h )∼(D t h )i P (t) (• | s h , a h ) -P ( * ,t) (• | s h , a h ) 2 T V (v) = n 2 T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 , (F.4)where (i) follows from the above definition of f (x, y), (ii) follows because the data of T tasks are independent conditional on D h , (iii) follows because P (t) is independent of the dataset (D ′ h ) (t) and from the definition of D ′ h , (iv) follows from Lemma F.2, and (v) follows because in task T , the data is collected using behavior policy π b t . Combining (F.2), (F.3), (F.4), we getn 2 T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 ≤ log(|Φ||Ψ| T /δ) (F.5)Using union bound, we obtain that for any h ∈ [H] and n with probability at least 1 -δ/2, it holds thatT t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) f (t) h (s, a) 2 ≤ 2 log(2|Φ||Ψ| T nH/δ) n .This completes the proof.G.1 One-step back lemmaThe following one-step back lemma is a key technical lemma for our proof. One-step back lemma for offline setting was first introduced in Uehara et al. ( 2022) for infinite-horizon stationary MDP. Our lemma extends their result to finite-horizon non-stationary MDP for offline setting. For any function g ∈ S × A → R, policy π and transition kernel P , the lemma shows that we can relate the expected value E (s h ,a h )∼(P,π) [g(s h , a h )] to the potential functionE (s h-1 ,a h-1 )∼(P,π) ∥ϕ h-1 (s h-1 , a h-1 )∥ (Σ h-1,π b t ,ϕ ) -1 .Lemma G.1 (One-step back inequality for non-stationary finite-horizon MDP in offline setting). For each task t ∈ [T ], let P ∈ { P (t) , P ( * ,t) } with embedding ϕ ∈ { ϕ, ϕ * } and µ be an MDP model, and Σ h,π b t ,ϕ = nE (s h ,a h )∼(P ( * ,t) ,π b t ) [ϕϕ ⊤ ] + λI be the covariance matrix following the behavior policy π b t under the true environment P ( * ,t) . Denote the total variation distance between P ( * ,t) and P at time step h by f t (s h , a h ). Take any g ∈ S × A → R such that ∥g∥ ∞ ≤ B. Then, letting ω = max s,a (1/π b t (a | s)) for all h ≥ 2, and for any policy π, we haveE (s h ,a h )∼(P,π) [g(s h , a h )] ≤ E (s h-1 ,a h-1 )∼(P,π) ∥ϕ h-1 (s h-1 , a h-1 )∥ (Σ h-1,π b t ,ϕ ) -1 × nωE (s h ,a h )∼(P (⋆,t) ,π b t ) [g 2 (s h ,a h )]+λdB 2 +nB 2 E (s h-1 ,a h-1 )∼(P (⋆,t) ,π b t ) [f t (s h-1 ,a h-1 ) 2 ] .Proof of Lemma G.1. First, we haveE (s h ,a h )∼(P,π) [g(s h , a h )] = E (s h-1 ,a h-1 )∼(P,π) s h a h g(s h , a h )π(a h | s h )⟨ϕ h-1 (s h-1 , a h-1 ), µ h-1 (s h )⟩ds h ≤ E (s h-1 ,a h-1 )∼(P,π)   ∥ϕh-1(sh-1, a h-1 )∥ (Σ h-1,π b t ,ϕ ) -1 a h g(s h , a h )π(a h | s h )µ h-1 (s h )ds h Σ h-1,π b t ,ϕ    ,where the inequality follows from Cauchy-Schwarz inequality.Then,a h g(s h , a h )π(a h | s h )µ h-1 (s h )ds h 2 Σ h-1,π b t ,ϕ = a h g(s h , a h )π(a h | s h )µ h-1 (s h )ds h ⊤ nE s h-1 ∼P ( * ,t) a h-1 ∼π b [ϕϕ ⊤ ] + λI a h g(s h , a h )π(a h | s h )µ h-1 (s h )ds h (i) ≤ nE (s h-1 ,a h-1 )∼(P ( * ,t) ,π b )   a h g(s h , a h )π(a h | s h )µ h-1 (s h ) ⊤ ϕ(s h-1 , a h-1 )ds h 2   + B 2 λd = nE (s h-1 ,a h-1 )∼(P ( * ,t) ,π b ) E s h ∼P (• | s h-1 ,a h-1 ) a h ∼π [g(s h , a h ) 2 ] + B 2 λd (ii) ≤ nE (s h-1 ,a h-1 )∼(P ( * ,t) ,π b ) E s h ∼P ( * ,t) a h ∼π [g(s h , a h ) 2 ] + B 2 λd + nB 2 E (s h-1 ,a h-1 )∼(P ( * ,t) ,π b ) [f t (s h-1 , a h-1 ) 2 ] (iii) ≤ nωE (s h ,a h )∼(P ( * ,t) ,π b ) g(s h , a h ) 2 + B 2 λd + nB 2 E (s h-1 ,a h-1 )∼(P ( * ,t) ,π b ) [f t (s h-1 , a h-1 ) 2 ]where (i) follows from the assumption ∥g∥ ∞ ≤ B and for any functionh : S → [0, 1], ∥ µ h (s)h(s)ds∥ 2 ≤ √ d, (ii)follows from the definition of f t (s h , a h ) which is the total variation distance between P * and P at time step h, and finally (iii) follows from importance sampling. This completes the proof.EQUATIONh is equal to Σ h,π b t , ϕ in expectation. We now provide an important lemma to ensure the concentration of the penalty term. The version for fixed ϕ is proved in Zanette et al. (2021a) . Here, we take a union bound over the whole feature ϕ ∈ Φ, number of total tasks T , horizon H and cardinality n of each offline dataset from individual tasks.Lemma G.2 (Concentration of the penalty term). Fix δ ∈ (0, 1) and set λ = O(d log(2nT H|Φ|/δ)) for any n. With probability at least 1 -δ/2, we have that for any n ∈ N, h ∈ [H], t ∈ [T ] and ϕ ∈ Φ,β 1 ∥ϕ h (s, a)∥ (Σ h,π b t ,ϕ ) -1 ≤ ∥ϕ h (s, a)∥ ( Σ (t) h ) -1 ≤ β 2 ∥ϕ h (s, a)∥ (Σ h,π b t ,ϕ ) -1 ,where β 1 and β 2 are some absolute constants.We first restate Lemma 4.3.Lemma H.1. Under Assumption 4.1, the output ϕ of Algorithm 1 is a ξ down -approximate feature for MDP M T +1 whereξ down = ξ + C L C R ν κ 2T log(2|Φ||Ψ| T nH/δ) n, i.e. there exist a time-dependent unknown (signed) measure µ * over S such that for any (s, a) ∈ S × A, we have∥P ( * ,T +1) h (•|s, a) -⟨ ϕ h (s, a), µ * h (•)⟩∥ TV ≤ ξ down .Furthermore, for any g : S → [0, 1], we have∥ µ * h (s)g(s)ds∥ 2 ≤ C L √ d.The following proof is motivated from the proof of Lemma 1 in Cheng et al. (2022) .Proof of Lemma H.1. For all (s, a) ∈ S × A, h ∈ [H] and for any t ∈ [T ] we have ≤ T t=1 C R E (s h ,a h )∼U (S,A) ∥ P (t) h (•|s h , a h ) -P ( * ,t) h (•|s h , a h )∥ TV (ii) ≤ C R ν κ T t=1 E (s h ,a h )∼(P ( * ,t) ,π b t ) ∥ P (t) h (•|s h , a h ) -P ( * ,t) h (•|s h , a h )∥ TV (iii) ≤ C R ν κ 2T log(2|Φ||Ψ| T nH/δ) n (H.1)where (i), (ii) follows from Assumption 4.1 and (iii) follows from Theorem 3.3.Defining µ * (•) = T t=1 c t µ (t) (•), we have ∥P ( * ,T +1) h (•|s, a) -⟨ ϕ h (s, a), µ * h (•)⟩∥ TV = ∥P ( * ,T +1) h (•|s, a) -⟨ ϕ h (s, a), T t=1 c t µ (t) (•)⟩∥ TV = ∥P ( * ,T +1) h (•|s, a) - T t-1 c t P (t) h (•|s, a)∥ TV ≤ ∥P ( * ,T +1) h (•|s, a) - T t-1 c t P ( * ,t) h (•|s, a)∥ TV + T t=1 c t ∥P ( * ,t) h (•|s, a) -P (t) h (•|s, a)∥ TV (i) ≤ ξ + C L C R ν κ 2T log(2|Φ||Ψ| T nH/δ) n ,where For any h ∈ [H], we defineP ( * ,T +1) h (•|s, a) = ⟨ϕ * h (s, a), µ ( * ,T +1) h (•)⟩, P h (•|s, a) = ⟨ ϕ h (s, a), µ * h (•)⟩.Given a reward function r (as is provided in the planning phase of reward-free RL setting), for any function f : S → R and h ∈ [H], we define the transition operators as follows(P ( * ,T +1) h f )(s, a, r) = s ′ ⟨ϕ * h (s, a), µ ( * ,T +1) h (s ′ )⟩f (s ′ )ds ′ , (P h f )(s, a, r) = s ′ ⟨ ϕ h (s, a), µ * h (s ′ )⟩f (s ′ )ds ′ .When no reward function is provided as is the case in the exploration phase of reward-free RL setting, we simply omit r from the above operator notation.In this section for notational simplicity we denote V π h,P ( * ,T +1) ,r (s) and Q π h,P ( * ,T +1) ,r (s, a) by V π h (s, r) and Q π h (s, a, r) respectively where r is reward function provided in the planning phase of downstream reward-free RL task. We similarly denote the optimal value function and action-value function under reward function r as V * h (s, r) and Q * h (s, a, r) respectively. We also introduce the truncated optimal value function V * h (s, r) in the planning phase, which is recursively defined from step H + 1 to step 1. Compared to the definition of standard optimal value function V * h (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 I.1 (Truncated Optimal Value Function). We introduce the truncated optimal value function V * h (s, r) which is recursively defined from step H + 1 to step 1:V * H+1 (s, r) = 0, ∀s ∈ S Q * h (s, a, r) = r h (s, a) + P ( * ,T +1) h V * h+1 (s, a, r), ∀(s, a) ∈ S × A V * h (s, r) = min max a∈A r h (s, a) + P ( * ,T +1) h V * h+1 (s, a, r) , H , ∀s ∈ S, h ∈ [H].We can similarly define V π h (s, r) and Q π h (s, a, r).Now we state the supporting lemmas that are used in the proof of Theorem 4.4. The proof of these lemmas are provided in Appendix J.The following lemma shows that the linear weight w k h in Algorithm 2 is bounded. Lemma I.2 (Bounds on Weights in Algorithm 2). For any h ∈ [H], the weight w k h in Algorithm 2 satisfiesw k h 2 ≤ H √ dK.Lemma I.3. Let E be the event that for all (k, h)∈ [K RFE ] × [H], k-1 τ =1 ϕ τ h V k h+1 (s τ h+1 ) -P ( * ,T +1) h V k h+1 (s τ h , a τ h ) (Λ k h ) -1 ≲ dH log dK RFE H max(ξ down , 1) δ .Then Pr[E] ≥ 1 -δ/8. Lemma I.4. With probability 1 -δ/8, we have for all (s, a) ∈ S × A,ϕ h (s, a) ⊤ w k h -P ( * ,T +1) h V k h+1 (s, a) ≲ β∥ ϕ h (s, a)∥ (Λ k h ) -1 + Hξ down Lemma I.5. With probability 1 -δ/4, for all (h, k) ∈ [H] × [K RFE ],and any s ∈ S, we haveV * h (s, r k ) ≲ V k h (s) + H(H -h + 1)ξ down and KRFE k=1 V k 1 (s k 1 ) ≤ c d 3 H 4 K RFE log(dK RFE H/δ) + H 2 K RFE ξ down ,where c > 0 is a constant. Lemma I.6. With probability 1 -δ/2, for the function u h (•, •) defined in Line 5 of Algorithm 3, we haveE s∼µ V * 1 (s, u) ≤ c ′ d 3 H 4 log(dK RFE H/δ)/K RFE + 2H 2 ξ down . Lemma I.7.With probability 1 -δ/2, for any reward function which is linear with respect to the unknown feature ϕ * : S × A → R d , for all (s, a) ∈ S × A and h ∈ [H], we haveQ * h (s, a, r)-H(H -H +1)ξ down ≤ Q h (s, a) ≤ r h (s, a)+P ( * ,T +1) h V h+1 (s, a)+2u h (s, a)+Hξ down I.2 Proof of Theorem 4.4We first restate Theorem 4.4 Theorem I.8. Under Assumption 4.1, after collecting K RFE trajectories during the exploration phase in Algorithm 2, with probability at least 1 -δ, the output of Algorithm 3, policy π satisfiesE s1∼µ [V * 1 (s 1 , r) -V π 1 (s 1 , r)] ≤ c ′ d 3 H 4 log(dK RFE H/δ)/K RFE + 6H 2 ξ down . (I.1)If the linear combination misspecification error ξ in Assumption 4.1 satisfies O( d 3 /K RFE ) and the number of trajectories in the offline dataset for each upstream task is at least O(T K RFE /d 3 ), then, provided K RFE is at least O(d 3 H 4 log(dHδ -1 ϵ -1 )/ϵ 2 ), with probability 1 -δ, the policy π will be an ϵ-optimal policy for any given reward during the planning phase.Proof of Theorem I.8. We condition on the events defined in Lemma I.6 and Lemma I.7 which, by union bound, hold with probability at least 1 -δ. By Lemma I.7, for any s ∈ S, we haveV 1 (s) = max a∈A Q 1 (s, a) ≥ max a∈A Q * 1 (s, a, r) -H 2 ξ down = V * 1 (s, r) -H 2 ξ down .This impliesE s1∼µ [V * 1 (s 1 , r) -V π 1 (s 1 , r)] ≤ E s1∼µ [ V 1 (s 1 ) -V π 1 (s 1 , r)] + H 2 ξ down ,where π is the policy returned by Algorithm 3.Observe that, using Lemma I.7, we haveE s1∼µ [ V 1 (s 1 ) -V π 1 (s 1 , r)] = E s1∼µ [ Q 1 (s 1 , π 1 (s 1 )) -Q π 1 (s 1 , π 1 (s 1 ), r)] ≤ E s1∼µ [r 1 (s 1 , π 1 (s 1 )) + P ( * ,T +1) 1 V 2 (s 1 , π 1 (s 1 ), r) + 2u 1 (s 1 , π 1 (s 1 )) + Hξ down -r 1 (s 1 , π 1 (s 1 )) -P ( * ,T +1) 1 V π 2 (s 1 , π 1 (s 1 ), r)] = E s1∼µ,s2∼P ( * ,T +1) 1 (•|s1,π1(s1)) [ V 2 (s 2 ) -V π 2 (s 2 , r) + 2u 1 (s 1 , π 1 (s 1 ))] + Hξ down ≤ . . . ≤ 2E s∼µ [V π 1 (s, u)] + H 2 ξ down .Moreover, note that 0 ≤ V h (s) ≤ H and 0 ≤ V π h (s, r) ≤ H as 0 ≤ r(s, a) ≤ 1. Thus, we would always have V h (s) -V π h (s, r) ≤ H. Along with the previous derivation, this implies,E s1∼µ [ V 1 (s 1 ) -V π 1 (s 1 , r)] ≤ 2E s∼µ [ V π 1 (s, u)] + H 2 ξ down .By definition of V * 1 (s, u), we further haveE s∼µ [ V π 1 (s, u)] ≤ E s∼µ [ V * 1 (s, u)]. By Lemma I.6, we haveE s∼µ V * 1 (s, u) ≤ c ′ d 3 H 4 log(dK RFE H/δ)/K RFE + 2H 2 ξ down .So, we have,E s1∼µ [V * 1 (s 1 , r) -V π 1 (s 1 , r)] ≤ E s1∼µ [ V 1 (s 1 ) -V π 1 (s 1 , r)] + H 2 ξ down ≤ 2E s∼µ [ V π 1 (s, u)] + 2H 2 ξ down ≤ 2E s∼µ [ V * 1 (s, u)] + 2H 2 ξ down ≤ 2c ′ d 3 H 4 log(dK RFE H/δ)/K RFE + 6H 2 ξ down . (I.2)Recall the definition of ξ down from Lemma 4.3, that is,ξ down = ξ + C L C R ν κ 2T log(2|Φ||Ψ| T nH/δ) n .If the linear combination misspecification error ξ in Assumption 4.1 satisfies O( d 3 /K RFE ) and the number of trajectories in the offline dataset for each task in the upstream stage is at least O(T K RFE /d 3 ), then the first term in Equation (I.2) dominates the second term 6H 2 ξ down . Then, by taking K RFE = c K d 3 H 4 log(dHδ -1 ϵ -1 )/ϵ 2 for a sufficiently large constant c K > 0, we haveE s1∼µ [V * 1 (s 1 , r) -V π 1 (s 1 , r)] ≤ 2c ′ d 3 H 4 log(dK RFE H/δ)/K RFE + 6H 2 ξ down ≤ ϵ.This completes the proof.J.1 Proof of Lemma I.2Proof of Lemma I.2. We havew k h = (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) V k h+1 (s τ h+1 ) ≤ √ k k-1 τ =1 V k h+1 (s τ h+1 ) ϕ h (s τ h , a τ h ) 2 (Λ k h ) -1 1/2 ≤ K RFE • H • k-1 τ =1 ϕ h (s τ h , a τ h ) 2 (Λ k h ) -1 1/2 ≤ H dK RFEwhere the first inequality follows from Lemma O.5 and the fact that the largest eigenvalue of (Λ k h ) -1 is at most 1, second inequality follows from the fact that | V k h+1 (s)| ≤ H for all s ∈ S and the last inequality follows from Lemma O.4.Proof of Lemma I.3. The proof is similar to 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 (k, h) ∈ [K RFE ] × [H], by Lemma I.2, we have ∥ w k h ∥ 2 ≤ H √ dK. Moreover, we have r k h (•, •) = u k h (•, •) and hence we haver k h (•, •) + u k h (•, •) = 2β ϕ(•, •) ⊤ (Λ k h ) -1 ϕ(•, •).Thus, our value function V k h+1 is of the formV (•) := min max a∈A ϕ(•, a) ⊤ w + β ϕ(•, a) ⊤ (Λ) -1 ϕ(•, a), H ,for some Λ ∈ R d×d , and w ∈ R d which matches the value function class defined in Lemma O.9. Moreover, by construction, the minimum eigenvalue of Λ k h is lower bounded by 1. Combining Lemma O.7 and Lemma O.9, we have for any fixed ε > 0 thatk-1 τ =1 ϕ τ h V k h+1 (s τ h+1 ) -P ( * ,T +1) h V k h+1 (s τ h , a τ h ) 2 (Λ k h ) -1 ≤ 4H 2 d 2 log(k + 1) + d log 1 + 8H √ dk ε + d 2 log 1 + 32 √ dβ 2 ε 2 + log 8 δ + 8k 2 ε 2We set the hyperparameterβ = C L H √ d + dH log(dK RFE H max(ξ down , 1)/δ) + Hξ down √ dK RFE . Finally, picking ε = dH/k, we have k-1 τ =1 ϕ τ h V k h+1 (s τ h+1 ) -P ( * ,T +1) h V k h+1 (s τ h , a τ h ) (Λ k h ) -1 ≲ dH log dK RFE H max(ξ down , 1) δ ,which concludes the proof.J.3 Proof of Lemma I.4Proof of Lemma I.4. We condition on the event E defined in Lemma I.3, which holds with probability at least 1 -δ/8. We definew k h = s ′ V k h+1 (s ′ ) µ * (s ′ )ds ′ where µ * is defined in the proof of Lemma 4.3. Note that by Lemma 4.3, we have ∥ w k h ∥ ≤ C L H √ d. Now, ϕ h (s, a) ⊤ w k h -P ( * ,T +1) h V k h+1 (s, a) = ϕ h (s, a) ⊤ w k h -ϕ h (s, a) ⊤ w k h + ϕ h (s, a) ⊤ w k h -P ( * ,T +1) h V k h+1 (s, a) ≤ ϕ h (s, a) ⊤ w k h -ϕ h (s, a) ⊤ w k h + P h V k h+1 (s, a) -P ( * ,T +1) h V k h+1 (s, a) ≤ ϕ h (s, a) ⊤ w k h -ϕ h (s, a) ⊤ w k h + Hξ down , (J.1)where the last inequality follows from Lemma 4.3 and | V k h+1 | ≤ H. The first term in (J.1) can be written as,ϕ h (s, a) ⊤ w k h -ϕ h (s, a) ⊤ w k h = ϕ h (s, a) ⊤ (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) V k h+1 (s τ h+1 ) -ϕ h (s, a) ⊤ (Λ k h ) -1 (Λ k h ) w k h = ϕ h (s, a) ⊤ (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) V k h+1 (s τ h+1 ) -w k h - k-1 τ =1 ϕ h (s τ h , a τ h )P h V k h+1 = -ϕ h (s, a) ⊤ (Λ k h ) -1 w k h (a) + ϕ h (s, a) ⊤ (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) V k h+1 (s τ h+1 ) -P ( * ,T +1) h V k h+1 (s τ h+1 ) (b) + ϕ h (s, a) ⊤ (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) P ( * ,T +1) h -P h V k h+1 (s τ h+1 ) (c) (J.2)We now bound (a), (b), (c) in (J.2) individually.Term (a). We have,-ϕ h (s, a) ⊤ (Λ k h ) -1 w k h ≤ ϕ h (s, a) (Λ k h ) -1 ∥ w k h ∥ (Λ k h ) -1 ≤ ∥ w k h ∥ 2 ϕ h (s, a) (Λ k h ) -1 ≤ C L H √ d ϕ h (s, a) (Λ k h ) -1 . (J.3)Term (b). Using Cauchy-Schwarz inequality and the definition of the event E from Lemma I.3, we haveϕ h (s, a) ⊤ (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) V k h+1 (s τ h+1 ) -P ( * ,T +1) h V k h+1 (s τ h+1 ) ≤ ϕ h (s, a) (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) V k h+1 (s τ h+1 ) -P ( * ,T +1) h V k h+1 (s τ h+1 ) (Λ k h ) -1 ≲ dH log dK RFE H max(ξ down , 1) δ ϕ h (s, a) (Λ k h ) -1 . (J.4)Term (c). We haveϕ h (s, a) ⊤ (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) P ( * ,T +1) h -P h V k h+1 (s τ h+1 ) (i) ≤ ϕ h (s, a) ⊤ (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) Hξ down = k-1 τ =1 ϕ h (s, a) ⊤ (Λ k h ) -1 ϕ h (s τ h , a τ h ) Hξ down (ii) ≤ k-1 τ =1 ϕ h (s, a) 2 (Λ k h ) -1 k-1 τ =1 ϕ h (s τ h , a τ h ) 2 (Λ k h ) -1 Hξ down (iii) ≤ Hξ down √ dk ϕ h (s, a) (Λ k h ) -1 , (J.5)where (i) follows from Lemma 4.3, (ii) follows from Cauchy-Schwarz inequality and (iii) follows from Lemma O.4. Substituting (J.3), (J.4), (J.5), into (J.2), and denotingβ = C L H √ d + dH log dKRFEH max(ξdown,1) δ + Hξ down √ dk we get ϕ h (s, a) ⊤ w k h -ϕ h (s, a) ⊤ w k h ≲ β ϕ h (s, a) (Λ k h ) -1Putting everything together in (J.1), we getϕ h (s, a) ⊤ w k h -P ( * ,T +1) h V k h+1 (s, a) ≲ β ϕ h (s, a) (Λ k h ) -1 + Hξ down , which concludes the proof.Proof of Lemma I.5. We prove the first part using backward induction on h. For h = H, we haveV * H (s, r k ) = min max a∈A r k H (s, a) + P ( * ,T +1) H V * H+1 (s, a, r k ) , H = min max a∈A r k H (s, a), H ≤ min max a∈A r k H (s, a) + ϕ H (s, a) ⊤ w k H + β∥ ϕ H (s, a)∥ (Λ k H ) -1 , H = V k H (s) ≤ V k H (s) + H(H -H + 1)ξ down .Suppose for some h + 1 ∈ [H], it holds that for all s ∈ S,V * h+1 (s, r k ) ≤ V k h+1 (s) + H(H -h)ξ down .Then,V * h (s, r k ) = min max a∈A r k h (s, a) + P ( * ,T +1) h V * h+1 (s, a, r k ) , H ≤ max a∈A r k h (s, a) + P ( * ,T +1) h V * h+1 (s, a, r k ) ≤ max a∈A r k h (s, a) + P ( * ,T +1) h V k h+1 (s, a) + H(H -h)ξ down ≲ max a∈A r k h (s, a) + ϕ h (s, a) ⊤ w k h + β∥ ϕ h (s, a)∥ (Λ k h ) -1 + Hξ down + H(H -h)ξ down = max a∈A r k h (s, a) + ϕ h (s, a) ⊤ w k h + β∥ ϕ h (s, a)∥ (Λ k h ) -1 + H(H -h + 1)ξ down ,where the last inequality follows from Lemma I.4.Thus, we haveV * h (s, r k ) ≲ min max a∈A r k h (s, a) + ϕ h (s, a) ⊤ w k h + β∥ ϕ h (s, a)∥ (Λ k h ) -1 , H + H(H -h + 1)ξ down = V k h (s) + H(H -h + 1)ξ down ,as desired. This completes the first part of the proof. Now we prove the second part of the proof. For all (k, h)∈ [K RFE ] × [H -1],we denote,ξ k h = P ( * ,T +1) h V k h+1 (s k h , a k h ) -V k h+1 (s k h+1 ).Conditioned on E from Lemma I.3 where Pr[E] ≥ 1 -δ/8, KRFE k=1 V k 1 (s k 1 ) ≤ KRFE k=1 r k 1 (s k 1 , a k 1 ) + ϕ 1 (s k 1 , a k 1 ) ⊤ w k 1 + β∥ ϕ 1 (s k 1 , a k 1 )∥ (Λ k 1 ) -1 = KRFE k=1 ϕ 1 (s k 1 , a k 1 ) ⊤ w k 1 + 2β∥ ϕ 1 (s k 1 , a k 1 )∥ (Λ k 1 ) -1 ≲ KRFE k=1 P ( * ,T +1) 1 V k 2 (s k 1 , a k 1 ) + 3β∥ ϕ 1 (s k 1 , a k 1 )∥ (Λ k 1 ) -1 + Hξ down = KRFE k=1 ξ k 1 + V k 2 (s k 2 ) + 3β∥ ϕ 1 (s k 1 , a k 1 )∥ (Λ k 1 ) -1 + HK RFE ξ down ≤ • • • ≤ KRFE k=1 H-1 h=1 ξ k h + KRFE k=1 H h=1 3β∥ ϕ h (s k h , a k h )∥ (Λ k h ) -1 + H 2 K RFE ξ down ,where the second inequality follows from Lemma I.4.h ∈ [H -1], {ξ k h } KRFE k=1 is a martingale difference sequence with |ξ k h | ≤ H.We define the event E ′ to be the event thatKRFE k=1 H-1 h=1 ξ k h ≤ c ′ H 2 K RFE log(K RFE H/δ).By Azuma-Hoeffding inequality, we have Pr[E ′ ] ≥ 1 -δ/8. By Cauchy-Schwarz inequality, we haveKRFE k=1 H h=1 ∥ ϕ h (s k h , a k h )∥ (Λ k h ) -1 ≤ K RFE H KRFE k=1 H h=1 ϕ h (s k h , a k h ) ⊤ (Λ k h ) -1 ϕ h (s k h , a k h ).Using Lemma D.2 of Jin et al. (2020b) , we haveKRFE k=1 H h=1 ϕ h (s k h , a k h ) ⊤ (Λ k h ) -1 ϕ h (s k h , a k h ) ≤ 2dH log(K RFE ).Conditioned on E and E where Pr(E ∩ E ′ ) ≥ 1 -δ/4, we haveKRFE k=1 V k 1 (s k 1 ) ≤ c ′ H 2 K RFE log(K RFE H/δ) + 3β K RFE H • 2dH log(K RFE ) + H 2 K RFE ξ down ≤ c d 3 H 4 K RFE log(dK RFE H/δ) + H 2 K RFE ξ down ,which completes the proof.J.5 Proof of Lemma I.6Proof of Lemma I.6.Denote ∆ k = V * 1 (s k 1 , r k ) -E s∼µ [ V * 1 (s, r k )].Note that r k depends only on the data collected during the first k -1 episodes. Thus, {∆ k } KRFE k=1 is a martingale difference sequence. Moreover, we have |∆ k | ≤ H. Using Azuma-Hoeffding inequality, with probability 1 -δ/4, we haveKRFE k=1 ∆ k ≤ c 1 H K RFE log(1/δ),where c 1 > 0 is an absolute constant. Therefore, we haveE s∼µ KRFE k=1 V * 1 (s, r k ) ≤ KRFE k=1 V * 1 (s k 1 , r k ) + c 1 H K RFE log(1/δ). Now, notice that for all k ∈ [K RFE ], Λ h ⪰ Λ k h . Thus for all (k, h) ∈ [K RFE ] × [H], we have r k h (•, •) ≥ u h (•, •), which implies V * 1 (•, u h ) ≤ V * 1 (•, r k h ).Using Lemma I.5 and union bound, we have with probability 1 -δ/2,E s∼µ V * 1 (s, u) ≤ E s∼µ KRFE k=1 V * 1 (s, r k )/K RFE ≤ 1 K RFE KRFE k=1 V * 1 (s k 1 , r k ) + c 1 H log(1/δ) K RFE ≤ 1 K RFE KRFE k=1 V k 1 (s k 1 ) + H 2 ξ down + c 1 H log(1/δ) K RFE = 1 K RFE KRFE k=1 V k 1 (s k 1 ) + H 2 ξ down + c 1 H log(1/δ) K RFE ≤ 1 K RFE c d 3 H 4 K RFE log(dK RFE H/δ) + H 2 K RFE ξ down + H 2 ξ down + c 1 H log(1/δ) K RFE = c d 3 H 4 log(dK RFE H/δ) K RFE + c 1 H log(1/δ) K RFE + 2H 2 ξ down ≤ c ′ d 3 H 4 log(dK RFE H/δ) K RFE + 2H 2 ξ down ,for some absolute constant c ′ > 0. This completes the proof.J.6 Proof of Lemma I.7Proof of Lemma I.7. Following the same argument as in the proof of Lemma I.4, for all h ∈ [H] and (s, a) ∈ S × A, with probability 1 -δ/4, we haveϕ h (s, a) ⊤ w h -P ( * ,T +1) h V h+1 (s, a) ≲ β∥ ϕ h (s, a)∥ (Λ h ) -1 + Hξ downThus, for all h ∈ [H] and (s, a) ∈ S × A, we haveQ h (s, a) ≤ ϕ h (s, a) ⊤ w h + r h (s, a) + u h (s, a) ≲ r h (s, a) + P ( * ,T +1) h V h+1 (s, a) + 2β∥ ϕ h (s, a)∥ (Λ h ) -1 + Hξ down . Since, Q h (s, a) ≤ H and u h (•, •) = min β ϕ(•, •) ⊤ (Λ h ) -1 ϕ(•, •), H , we have, Q h (s, a) ≤ r h (s, a) + P ( * ,T +1) h V h+1 (s, a) + 2u h (s, a) + Hξ down .This completes the first part of the proof. Now, using induction, we prove that for all h ∈ [H] and (s, a) ∈ S × A, we haveQ * h (s, a, r) - H(H -h + 1)ξ down ≤ Q h (s, a).When h = H + 1, the claim is trivially true. Suppose for some h ∈ [H], we have, for all (s, a) ∈ S × A,Q * h+1 (s, a, r) -H(H -h)ξ down ≤ Q h (s, a). Note that, Q h (s, a) = min ϕ h (s, a) ⊤ w h + r h (s, a) + u h (s, a), H Since Q * h (s, a, r) ≤ H and u h (•, •) = min β ϕ(•, •) ⊤ (Λ h ) -1 ϕ(•, •), H , it suffices to show that Q * h (s, a, r) ≤ ϕ h (s, a) ⊤ w h + r h (s, a) + β∥ ϕ h (s, a)∥ (Λ h ) -1 + H(H -h + 1)ξ down .Applying the max operator on both side of the inductive hypothesis, we getV * h+1 (s, r) ≤ V h+1 + H(H -h)ξ down . Now, Q * h (s, a, r) = r h (s, a) + P ( * ,T +1) h V * h+1 (s, a, r) ≤ r h (s, a) + P ( * ,T +1) h V h+1 (s, a) + H(H -h)ξ down ≲ ϕ h (s, a) ⊤ w h + β∥ ϕ h (s, a)∥ (Λ h ) -1 + Hξ down + H(H -h)ξ down + r h (s, a) = ϕ h (s, a) ⊤ w h + r h (s, a) + β∥ ϕ h (s, a)∥ (Λ h ) -1 + H(H -h + 1)ξ down .This completes the proof.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.The following lemma shows that the linear weight w h in Algorithm 4 is bounded. Lemma K.1 (Bounds on Weights in Algorithm 4). For any h ∈ [H], the weight w h in Algorithm 4 satisfiesw h 2 ≤ H dN off /λ d .Next, we present our main concentration lemma for this section that upper-bounds the stochastic noise in regression.Lemma K.2. Setting λ d = 1, β(δ) = c β (Hd ι(δ) + H dN off ξ down ), where ι(δ) = log(HdN off max(ξ down , 1)/δ), with probability at least 1 -δ, for all h ∈ [H], we haveNoff τ =1 ϕ h (s τ h , a τ h ) (P ( * ,T +1) h V h+1 )(s τ h , a τ h ) -V h+1 (s τ h+1 ) Λ -1 h ≲ Hd √ ι.Recall that, we define the Bellman operator B h as(B h f )(s, a) = r h (s, a)+(P ( * ,T +1) h f )(s, a) for any f : S × A → R. In the next lemma, we denote the Bellman estimate ( B h V h+1 )(•, •) = ϕ h (•, •) ⊤ w h .This lemma provides an upper bound for the Bellman update error|(B h V h+1 -B h V h+1 )(s, a)|and characterizes the impact of the misspecification of the representation taken from the upstream learning.Lemma K.3 (Bound on Bellman update error). Set λ d = 1, β = c β (Hd √ ι + H dN off ξ down ), where ι = log(HdN off max(ξ down , 1)/δ). Define the following eventE(δ) = (B h V h+1 -B h V h+1 )(s, a) ≤ β ϕ h (s, a) Λ -1 h + Hξ down , ∀h ∈ [H] and ∀(s, a) ∈ S × A .Then under Assumption 4.1, we have P(E(δ)) ≥ 1 -δ. Definition K.4 (Model prediction error). For all h ∈ [H], we define the model prediction error as,l h (s, a) = (B h V h+1 )(s, a) -Q h (s, a).The following lemma decomposes the suboptimality gap into the summation of uncertainty metric of each step.Lemma K.5. Let { π} H h=1 be the output of Algorithm 4. Conditioned on the event E(δ) defined in Lemma K.3, for any h ∈ [H] and (s, a) ∈ S × A, we have0 ≤ l h (s, a) ≤ 2Γ h (s, a),where Γ h (s, a) = β ϕ h (s, a) Λ -1 h + Hξ down . Furthermore, we haveV π * P ( * ,T +1) ,r (s) -V π P ( * ,T +1) ,r (s) ≤ 2 H h=1 E π * [Γ h (s h , a h )|s 1 = s] K.2 Proof of Theorem 4.6We first restate Theorem 4.6. Theorem K.6. Under Assumption 4.1, settingλ d = 1, β = O(Hd √ ι + H dN off ξ down ), where ι = log(HdN off max(ξ down , 1)/δ), with probability at least 1-δ, the suboptimality gap of Algorithm 4 is at mostV π * P ( * ,T +1) ,r (s) -V π P ( * ,T +1) ,r (s) ≤ 2H 2 ξ down + 2β H h=1 E π * ϕ h (s h , a h ) Λ -1 h s 1 = s . (K.1)Additionally if Assumption 4.5 holds, and the sample size satisfies N off ≥ 40/κ ρ • log(4dH/δ), then with probability 1 -δ, we have, Proof of Theorem K.6. Letting Γ h = β∥ ϕ h (s, a)∥ Λ -1 h + Hξ down in Lemma K.5, with probability at least 1 -δ, we haveEQUATION≤ O κ -1/2 ρ H 2 d 1/2 ξ down + κ -1/2 ρ H 2 d log(HdN off max(ξ down , 1)/δ) N off . (K.V π * P ( * ,T +1) ,r (s) -V π P ( * ,T +1) ,r (s) ≤ 2 H h=1 E π * [Γ h (s h , a h )|s 1 = s] ≤ 2H 2 ξ down + 2β H h=1 E π * ϕ h (s h , a h ) Λ -1 h s 1 = sThis 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 et al. (2021) , if N off ≥ 40/κ ρ • log(4dH/δ), then with probability 1 -δ/2, for any h ∈ [H] and (s, a) ∈ S × A, we haveϕ h (s, a) Λ -1 h ≤ 2 κ ρ • 1 √ N off .Setting β(δ/2) = c β (Hd ι(δ/2) + H √ dN off ξ down ), with probability 1 -δ/2, we haveV π * P ( * ,T +1) ,r (s) -V π P ( * ,T +1) ,r (s) ≤ 2H 2 ξ down + 2β H h=1 E π * ϕ h (s h , a h ) Λ -1 h s 1 = s .Using union bound, we have the following bound with probability at least 1 -δ V π * P ( * ,T +1) ,r (s) -V π P ( * ,T +1) ,r (s)≤ 2H Hξ down + β(δ/2) 2 κ ρ • 1 √ N off = O κ -1/2 ρ H 2 d 1/2 ξ down + κ -1/2 ρ H 2 d log(HdN off max(ξ down , 1)/δ) N off . (K.3)L Proof of Supporting Lemmas in Appendix KIn 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.1Proof of Lemma K.1. We havew h = Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h )(r τ h + V h+1 (s τ h+1 )) ≤ N off λ d Noff τ =1 (r τ h + V h (s τ h+1 )) ϕ h (s τ h , a τ h ) 2 Λ -1 h 1/2 ≤ H N off λ d Noff τ =1 ϕ h (s τ h , a τ h ) 2 Λ -1 h 1/2 ≤ H dN off λ d ,where the first inequality follows from Lemma O.5 and the fact that the largest eigenvalue of Λ -1 h is at most 1/λ d , second inequality follows from the fact that r τ h ∈ [0, 1] and | V h+1 (s)| ≤ H -1 for all s ∈ S and the last inequality follows from Lemma O.4.Proof of Lemma K.2. The value function V h+1 has the parametric form ofV (•) = min max a∈A w ⊤ ϕ(•, a) -β ϕ(•, a) ⊤ Λ -1 ϕ(•, a), H -h + 1 ,where w ∈ R d and positive definite matrix Λ is such that its minimum eigenvalue satisfies λ min (Λ) ≥ λ d . From Lemma K.1, we have ∥ w h ∥ ≤ H dN off /λ d . Thus, applying Lemma O.4 and Lemma O.9, we have, for any fixed ε > 0 and for all h ∈ [H], with probability at least 1 -δ,Noff τ =1 ϕ h (s τ h , a τ h ) (P ( * ,T +1) h V h+1 )(s τ h , a τ h ) -V h+1 (s τ h+1 ) 2 Λ -1 h ≤ 4H 2 d 2 log N off + λ d λ d + log N ε δ + 8N 2 off ε 2 λ d ≤ 4H 2 d 2 log N off + λ d λ d + d log 1 + 4H √ dN off ε √ λ d + d 2 log 1 + 8 √ dβ 2 λ d ε 2 + log 1 δ + 8N 2 off ε 2 λ d ≤ 4H 2 d 2 log N off + λ d λ d + d log 1 + 4H √ dN off ε √ λ d + d 2 log 1 + 8 √ dβ 2 λ d ε 2 + log 1 δ + 8N 2 off ε 2 λ d . (L.1) Setting ε = dH/N off , λ d = 1, β = c β (Hd √ ι + H √ dN off ξ down ), where ι = log(HdN off max(ξ down , 1)/δ), we can further upper bound (L.1) by4H 2 d 2 log(1 + N off ) + d log(1 + 4d -1/2 N 3/2 off ) + d 2 log(1 + 8d -3/2 H -2 N 2 off β 2 ) + log(1/δ) + 8d 2 H 2 ≲ H 2 d log N off + H 2 d log(d -1/2 N 3/2 off ) + H 2 log(1/δ) + H 2 d 2 log(Hd 1/2 ιN 3 off ξ down ) + d 2 H 2 ≲ H 2 d 2 ι Therefore, we have Noff τ =1 ϕ h (s τ h , a τ h ) (P ( * ,T +1) h V h+1 )(s τ h , a τ h ) -V h+1 (s τ h+1 ) Λ -1 h ≲ Hd √ ι. L.3 Proof of Lemma K.3 Proof of Lemma K.3. For h ∈ [H], we define w * h = s ′ µ * (s ′ ) V h+1 (s ′ )ds ′ . Then we have (B h V h+1 -B h V h+1 )(s, a) = (B h V h+1 -B h V h+1 + B h V h+1 -B h V h+1 )(s, a) ≤ (B h V h+1 -B h V h+1 )(s, a) + (B h V h+1 -B h V h+1 )(s, a) = (P ( * ,T +1) h V h+1 -P h V h+1 )(s, a) + (P h V h+1 )(s, a) -ϕ h (s, a) ⊤ w h (i) ≤ H P ( * ,T +1) h (•|s, a) -⟨ ϕ h (s, a), µ * h )(•)⟩ TV + s ′ ⟨ ϕ h (s, a), µ * h (s ′ )⟩ V h+1 (s ′ )ds ′ -ϕ h (s, a) ⊤ w h (ii) ≤ Hξ down + ϕ h (s, a) ⊤ ( w * h -w) , (L.2)where (i) follows from the fact that | V h+1 (s)| ≤ H for all s ∈ S and (ii) follows from Lemma 4.3.We now decompose the second term in (L.2). Recall thatΛ h = Noff τ =1 ϕ h (s τ h , a τ h ) ϕ h (s τ h , a τ h ) ⊤ +λ d I d and w h = Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) V h+1 (s τ h+1 ). Then, we have ϕ h (s, a) ⊤ ( w * h -w) = ϕ h (s, a) ⊤ Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) ϕ h (s τ h , a τ h ) ⊤ + λ d I d w * h - Noff τ =1 ϕ h (s τ h , a τ h ) V h+1 (s τ h+1 ) = λ d ϕ h (s, a) ⊤ Λ -1 h w * h (a) + ϕ h (s, a) ⊤ Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) (P ( * ,T +1) h V h+1 )(s τ h , a τ h ) -V h+1 (s τ h+1 ) (b) + ϕ h (s, a) ⊤ Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) (P h V h+1 -P ( * ,T +1) h V h+1 )(s τ h , a τ h ) (c) (L.3)We now provide an upper bound for each of the terms in (L.3).Term (a). We haveλ d ϕ h (s, a) ⊤ Λ -1 h w * h (i) ≤ λ d ϕ h (s, a) Λ -1 h ∥ w * h ∥ Λ -1 h (ii) ≤ λ d ϕ h (s, a) Λ -1 h ∥ w * h ∥ 2 (iii) ≤ λ d Hd ϕ h (s, a) Λ -1 h , (L.4)where (i) follows from Cauchy-Schwarz inequality, (ii) follows from the fact that the largest eigenvalue of Λ -1 h is at most 1/λ d and (iii) follows from Assumption 3.1 and | V h+1 (s)| ≤ H for all s ∈ S.Term (b). We haveϕ h (s, a) ⊤ Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) (P ( * ,T +1) h V h+1 )(s τ h , a τ h ) -V h+1 (s τ h+1 ) ≤ ϕ h (s, a) Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) (P ( * ,T +1) h V h+1 )(s τ h , a τ h ) -V h+1 (s τ h+1 ) Λ -1 h ≲ Hd √ ι ϕ h (s, a) Λ -1 h (L.5)where the first inequality comes from Cauchy Schwarz inequality and the last inequality follows from Lemma K.2.Term (c). We haveϕ h (s, a) ⊤ Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) (P h V h+1 -P ( * ,T +1) h V h+1 )(s τ h , a τ h ) ≤ ϕ h (s, a) ⊤ Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) • Hξ down = Noff τ =1 ϕ h (s, a) ⊤ Λ -1 h ϕ h (s τ h , a τ h ) • Hξ down ≤ Noff τ =1 ϕ h (s, a) 2 Λ -1 h Noff τ =1 ϕ h (s τ h , a τ h ) 2 Λ -1 h • Hξ down ≤ Hξ down dN off ϕ h (s, a) Λ -1 h , (L.6)where the first inequality follows from | V h+1 (s)| ≤ H for all s ∈ S and from Lemma 4.3, the second inequality follows from Cauchy-Schwarz inequality and the last inequality follows from Lemma O.4.Settingλ d = 1, β = c β (Hd √ ι + H √ dN off ξ down ), where ι = log(HdN off max(ξ down , 1)/δ) and combining (L.2) to (L.6), we get for any h ∈ [H] and for any (s, a) ∈ S × A, with probability at least 1 -δ,(B h V h+1 -B h V h+1 )(s, a) ≤ β ϕ h (s, a) Λ -1 h + Hξ down .This completes the proof.L.4 Proof of Lemma K.5Proof of Lemma K.5. Recall thatQ h (•, •) = min{ ϕ h (•, •) ⊤ w h -Γ h (•, •), H -h + 1} + If ϕ h (s, a) ⊤ w h -Γ h (s, a) ≤ 0, Q h (s, a) = 0. Then, we have l h (s, a) = (B h V h+1 )(s, a) - Q h (s, a) = (B h V h+1 )(s, a) > 0. If ϕ h (s, a) ⊤ w h -Γ h (s, a) > 0, we have l h (s, a) = (B h V h+1 )(s, a) -Q h (s, a) ≥ (B h V h+1 )(s, a) -ϕ h (s, a) ⊤ w h + Γ h (s, a) = (B h V h+1 )(s, a) -( B h V h+1 )(s, a) + Γ h (s, a) ≥ 0,where the last inequality follows from conditioning on the event E(δ). Thus, we have l h (s, a) ≥ 0 for any h ∈ [H] and (s, a) ∈ S × A.We now show that l h (s, a) ≤ 2Γ h (s, a). Observe thatϕ h (s, a) ⊤ w h -Γ h (s, a) = ( B h V h+1 )(s, a) -Γ h (s, a) ≤ (B h V h+1 )(s, a) ≤ H -h + 1,where the first inequality follows from the conditioning on the event E(δ) and the last inequality follows from the fact that r h ∈ [0, 1] andV h+1 ∈ [0, H -h]. Now, Q h (s, a) = min{ ϕ h (s, a) ⊤ w h -Γ h (s, a), H -h + 1} + = max{ ϕ h (s, a) ⊤ w h -Γ h (s, a), 0} ≥ ϕ h (s, a) ⊤ w h -Γ h (s, a) = ( B h V h+1 )(s, a) -Γ h (s, a).Now, from the definition of l h , we havel h (s, a) = (B h V h+1 )(s, a) -Q h (s, a) ≤ (B h V h+1 )(s, a) -( B h V h+1 )(s, a) + Γ h (s, a) ≤ 2Γ h (s, a),where the last inequality follows from the conditioning on the event E(δ).Finally, we obtainEQUATION≤ -H h=1 E π [l h (s h , a h )|s 1 = s] + H h=1 E π * [l h (s h , a h )|s 1 = s] ≤ 2 H h=1 E π * [Γ h (s h , a h )|s 1 = s]where the first inequality follows from Lemma O.3 and definition of π, and the last inequality follows because with probability at least 1 -δ, we have 0 ≤ l h (s, a) ≤ 2Γ h (s, a) for any h ∈ [H] and (s, a) ∈ S × A.In this section for notational simplicity we denote V π h,P ( * ,T +1) ,r T +1 (s) and Q π h,P ( * ,T +1) ,r T +1 (s, a) by V π h (s) and Q π h (s, a) respectively. First, we state the supporting lemmas that are used in the proof of Theorem 4.7. The proof of these lemmas are provided in Appendix N.The following concentration lemma for online RL upper-bounds the stochastic noise in regression. The proof is omitted since it is quite similar to the one of Lemma K.2.Lemma M.1. Setting λ d = 1, β n = c β (Hd ι n (δ) + H √ dnξ down + C L √ Hd), where ι n = log(Hdn max(ξ down , 1)/δ), with probability at least 1 -δ/2, for all h ∈ [H] and any n ∈ [N ], we haven-1 τ =1 ϕ h (s τ h , a τ h ) (P ( * ,T +1) h V n h+1 )(s τ h , a τ h ) -V n h+1 (s τ h+1 ) Λ -1 h ≲ Hd √ ι n .The next lemma recursively bounds the difference between the value function maintained in Algorithm 5 (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 error term is upper-bounded by the bonus term with high probability. Lemma M.2. There exists an absolute constant c β such that forβ n = c β (Hd ι n (δ)+H √ dnξ down + C L √ Hd), where ι n = log(Hdn max(ξ down , 1)/δ), and for any policy π, with probability at least 1 -δ/2, for any (s, a) ∈ S × A, n ∈ [N on ] and h ∈ [H], we haveϕ h (s, a) ⊤ w n h -Q π h (s, a) = P ( * ,T +1) h ( V n h+1 -V π h+1 )(s, a) + ∆ n h (s, a), for some ∆ n h (s, a) that satisfies ∥∆ n h (s, a)∥ ≤ β n ∥ ϕ h (s, a)∥ (Λ n h ) -1 + 2Hξ down .We now prove optimism of the estimated value function in the following lemma. Lemma M.3 (Optimism of value function). With probability at least 1 -δ/2, for any (s, a) ∈ S × A, h ∈ [H] and n ∈ [N on ], we haveQ n h (s, a) ≥ Q * h (s, a) -2H(H -h + 1)ξ down .Lemma M.2 also easily transforms a recursive formula for the value function differenceδ n h = V n h (s n h ) -V π n h (s n h ).The following lemma will be useful in proving the regret bound for the online downstream task.Lemma M.4 (Recursive formula). Let δ n h = V n h (s n h ) -V π n h (s n h ) and ξ n h+1 = E[δ n h+1 |s n h , a n h ] -δ n h+1 .Then for any (n, h) ∈ [N on , H] with probability at least 1 -δ/2, we haveδ n h ≤ δ n h+1 + ξ n h+1 + β n ϕ h (s n h , a n h ) (Λ n h ) -1 + 2Hξ down .M.2 Proof of Theorem 4.7We are now ready to prove the main theorem in this section. We first restate Theorem 4. ) improving the suboptimality gap bound of REP-UCB (Uehara et al., 2022) 3 by an order of O(H 3/2 K √ d) 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 byEQUATIONi) ≤ Non n=1 V n 1 -V π n P ( * ,T +1) ,r + 2H 2 N on ξ down (ii) ≤ Non n=1 H h=1 ξ n h + β n ϕ h (s n h , a n h ) (Λ n h ) -1 + 2Hξ down + 2H 2 N on ξ down ≤ Non n=1 H h=1 ξ n h (a) + Non n=1 H h=1 β n ϕ h (s n h , a n h ) (Λ n h ) -1 (b)+4H Dividing both sides by N on , we getEQUATIONV * P ( * ,T +1) ,r -V π P ( * ,T +1) ,r ≤ O(H 2 dξ down + H 2 d 3/2 N -1/2 on ).This completes the proof.In this section, we provide the proofs of the lemmas that we used in the proof of Theorem 4.7.N.1 Proof of Lemma M.2Proof of Lemma M.2. For any policy π, we definew π h = V π h+1 (s ′ ) µ * (s ′ )ds ′ . Note that ϕ h (s, a) ⊤ w π h = P h V π h+1 (s, a) and by Lemma 4.3, we have ∥w π h ∥ ≤ C L √ d.Now, we derive the followingϕ h (s, a) ⊤ w n h -Q π h (s, a) = ϕ h (s, a) ⊤ w n h -ϕ h (s, a) ⊤ w π h + ϕ h (s, a) ⊤ w π h -Q π h (s, a) ≤ ϕ h (s, a) ⊤ w n h -ϕ h (s, a) ⊤ w π h + ϕ h (s, a) ⊤ w π h -Q π h (s, a) (N.1)Using Lemma 4.3, we bound the second term asQ π h (s, a) -ϕ h (s, a) ⊤ w π h = P ( * ,T +1) h V π h+1 (s, a) -P h V π h+1 (s, a) ≤ Hξ down , (N.2)where we used the observation |V π h+1 | ≤ H. Now, the first term in (N.1) can be bounded by Term (a). We have,EQUATION-λ d ϕ h (s, a) ⊤ (Λ n h ) -1 w π h ≤ ϕ h (s, a) (Λ n h ) -1 ∥λ d w π h ∥ (Λ n h ) -1 ≤ λ d ∥w π h ∥ 2 ϕ h (s, a) (Λ n h ) -1 ≤ C L λ d Hd ϕ h (s, a) (Λ n h ) -1 , (N.4)where the first inequality follows from Cauchy-Schwarz inequality, the second inequality follows from the fact that the largest eigenvalue of (Λ n h ) -1 is at most 1/λ d and the last inequality follows from Assumption 3.1 and |V π h+1 (s)| ≤ H for all s ∈ S. Term (b). Using Cauchy-Schwarz inequality and Lemma M.1, we haveϕ h (s, a) ⊤ (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h ) V n h+1 (s τ h+1 ) -P ( * ,T +1) h V n h+1 (s τ h+1 ) ≤ ϕ h (s, a) (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h ) V n h+1 (s τ h+1 ) -P ( * ,T +1) h V n h+1 (s τ h+1 ) (Λ n h ) -1 ≲ Hd √ ι n ϕ h (s, a) (Λ n h ) -1 .(N.5)Term (c). We haveϕ h (s, a) ⊤ (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h )P h ( V n h+1 -V π h+1 )(s τ h , a τ h ) ≤ ϕ h (s, a) ⊤ (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h ) ϕ h (s τ h , a τ h ) ⊤ ( V n h+1 -V π h+1 )(s ′ ) µ * h (s ′ )ds ′ = ϕ h (s, a) ⊤ (Λ n h ) -1 (Λ n h -λ d I) ( V n h+1 -V π h+1 )(s ′ ) µ * h (s ′ )ds ′ ≤ ϕ h (s, a) ⊤ ( V n h+1 -V π h+1 )(s ′ ) µ * h (s ′ )ds ′ + λ d ϕ h (s, a) ⊤ (Λ n h ) -1 ( V n h+1 -V π h+1 )(s ′ ) µ * h (s ′ )ds ′ (i) ≤ P h ( V n h+1 -V π h+1 )(s, a) + C L λ d Hd ϕ h (s, a) (Λ n h ) -1 (ii) ≤ P ( * ,T +1) h ( V n h+1 -V π h+1 )(s, a) + Hξ down + C L λ d Hd ϕ h (s, a) (Λ n h ) -1 , (N.6)where (i) follows from similar steps as in (N.4) and (ii) follows from Lemma 4.3.We haveϕ h (s, a) ⊤ (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h ) P ( * ,T +1) h -P h V n h+1 (s τ h+1 ) (i) ≤ ϕ h (s, a) ⊤ (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h ) • Hξ down = n-1 τ =1 ϕ h (s, a) ⊤ (Λ n h ) -1 ϕ h (s τ h , a τ h ) • Hξ down (ii) ≤ n-1 τ =1 ϕ h (s, a) 2 (Λ n h ) -1 n-1 τ =1 ϕ h (s τ h , a τ h ) 2 (Λ n h ) -1 • Hξ down (iii) ≤ Hξ down √ dn ϕ h (s, a) (Λ n h ) -1 , (N.7)where (i) follows from Lemma 4.3 and | V n h+1 (s)| ≤ H for all s ∈ S, (ii) follows from Cauchy-Schwarz inequality and (iii) follows from Lemma O.4. Substituting (N.4), (N.5), (N.6), (N.7) into (N.3), and setting λ d = 1 we getϕ h (s, a) ⊤ w n h -ϕ h (s, a) ⊤ w π h ≲ c β Hd ι n (δ) + H √ dnξ down + C L √ Hd ϕ h (s, a) (Λ n h ) -1 + P ( * ,T +1) h ( V n h+1 -V π h+1 )(s, a) + Hξ down = β n ϕ h (s, a) (Λ n h ) -1 + P ( * ,T +1) h ( V n h+1 -V π h+1 )(s, a) + Hξ down . (N.8)Combining (N.8), (N.2) in (N.1) completes the proof.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ϕ H (s, a) w n H -Q π H (s, a) = P ( * ,T +1) H ( V n H+1 -V π H+1 )(s, a) + ∆ n H (s, a) ≤ β n ϕ H (s, a) (Λ n H ) -1 + 2Hξ down .Thus, we haveQ n H (s, a) = min ϕ H (s, a) ⊤ w n H + β n ϕ H (s, a) (Λ n H ) -1 , H -h + 1 ≥ Q * H (s, a) -2Hξ down .Now, suppose the statement holds true at step h + 1 and consider the step h. Using Lemma M.2, we haveϕ h (s, a) ⊤ w n h -Q * h (s, a) = ∆ n h (s, a) + P ( * ,T +1) h ( V n h+1 -V * h+1 )(s, a) ≥ -β n ϕ h (s, a) (Λ n h ) -1 -2Hξ down -2H(H -h)ξ down = -β n ϕ h (s, a) (Λ n h ) -1 -2H(H -h + 1)ξ down . Therefore, Q n h (s, a) = min{ ϕ h (s, a) ⊤ w n h + β n ∥ ϕ(s, a)∥ (Λ n h ) -1 , H -h + 1} + ≥ Q * h (s, a) -2H(H -h + 1)ξ down , which completes the proof. N.3 Proof of Lemma M.4Proof of Lemma M.4. By Lemma M.2, for any (s, a) ∈ S × A, h ∈ [H] and n ∈ [N on ], with probability at least 1 -δ/2, we haveQ n h (s, a) -Q π n h (s, a) ≤ ∆ n h (s, a) + P ( * ,T +1) h ( V n h+1 -V π n h+1 )(s, a) ≤ β n ϕ h (s, a) (Λ n h ) -1 + 2Hξ down + P ( * ,T +1) h ( V n h+1 -V π n h+1 )(s, a).And finally, by definition of π n in Algorithm 5, we have π n (s n h ) = a n h = argmax a∈A Q n h (s h , a). This implies Q n h (s n h , a n h ) -Q π n h (s n h , a n h ) = V n h (s n h ) -V π n h (s n h ) = δ n h . Thus, we have δ n h ≤ δ n h+1 + ξ n h+1 + β n ϕ h (s, a) (Λ n h ) -1 + 2Hξ down .O.1 Miscellaneous LemmasThe following lemma measures the difference between two value functions under two MDPs and reward functions. Here, we use a shorthand notationP h V h+1 (s h , a h ) = E s∼P h (•|s h ,a h ) [V (s)].Lemma O.1 (Simulation lemma (Dann et al., 2017) ). Consider two MDPs with transition kernels P 1 and P 2 , and reward function r 1 and r 2 respectively. Given a policy π, we have,V π h,P1,r1 (s h ) -V π h,P2,r2 (s h ) = H h ′ =h E s h ′ ∼(P 2 ,π) a h ′ ∼π r 1 (s h ′ , a h ′ ) -r 2 (s h ′ , a h ′ ) + (P 1,h ′ -P 2,h ′ )V π h ′ +1,P1,r1 (s h ′ , a h ′ ) | s h = H h ′ =h E s h ′ ∼(P 1 ,π) a h ′ ∼π r 1 (s h ′ , a h ′ ) -r 2 (s h ′ , a h ′ ) + (P 1,h ′ -P 2,h ′ )V π h ′ +1,P2,r2 (s h ′ , a h ′ ) | s hWe use the following lemma to deal with distribution shift in offline RL setting. We use the following lemma to bound the suboptimality in downstream offline RL.Lemma O.3 (Decomposition of Suboptimality, Lemma 3.1 in (Jin et al., 2021) ). Let π = { π h } H h=1 be the policy such that V h (s) = ⟨ Q h (s, •), π h (•|s)⟩ A and for each step h ∈ [H], define the model evaluation error as l h (s, a) = (B h V h+1 )(s, a) -Q h (s, a). For any π and s ∈ S, we have V π * 1 (s) -V π 1 (s) = - H h=1 E π [l h (ϕ ⊤ i (Λ h ) -1 ϕ i ≤ d.Lemma O.5 (Lemma D.5 in (Ishfaq et al., 2021) ). Let A ∈ R d×d be a positive definite matrix where its largest eigenvalue λ max (A) ≤ λ. Let x 1 , . . . , x k be k vectors in R d . Then it holds that Let {s i } ∞ i=1 be a stochastic process on state space S with corresponding filtration {F i } ∞ i=1 . Let {ϕ i } ∞ i=1 be an R d -valued stochastic process where ϕ i ∈ F i-1 , and∥ϕ i ∥ ≤ 1. Let Λ k = λI + k i=1 ϕ i ϕ ⊤ i .Then for any δ > 0, with probability at least 1 -δ, for all k ≥ 0, and any V ∈ V with sup s∈S |V (s)| ≤ H, we havek i=1 ϕ i V (s i ) -E[V (s i ) | F i-1 ] 2 Λ -1 k ≤ 4H 2 d 2 log k + λ λ + log N ε δ + 8k 2 ϵ 2 λ ,where N ε is the ε-covering number of V with respect to the distance dist(V, V ′ ) = sup s∈S |V (s) -V ′ (s)|.Lemma O.8 (Covering number of Euclidean ball, Vershynin (2018) ). For any ε > 0, the ε-covering number, N ε , of the Euclidean ball of radius B > 0 in R d satisfiesN ε ≤ 1 + 2B ε d ≤ 3B ε d .• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It 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 a 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 of range (e.g. negative error rates). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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: It is a theoretical paper and thus we do not forsee any immediate societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).We rescale the result in(Cheng et al., 2023) by a factor of H 2 as we do not assume the sum of rewards to be within [0, 1].We convert the 1/(1 -γ) horizon dependence in REP-UCBUehara et al. (2022) to H. We further rescale their suboptimality gap by a factor of H 2 as we do not assume the sum of rewards to be within [0, 1]. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
1. 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.
1. The contributions to downstream online and offline RL are not clear.

Minor issues:
1. Equation (3.2) should be $P(s'|s,a)$... where $s'$ is missing.
questions: questions: 1. In line 295 in Section 4.3, why do you assume that reward function $r^{T+1}$ is unknown?
1. Could you provide any experimental results to support the claims in the paper?
limitations: limitations: The paper lacks a limitations section.
paper_summary: paper_summary: This paper provides a theoretical analysis of representation learning in Multi-Task Offline RL.
Specifically 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.
The 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.
They provide a suboptimality bound, based on bounding the TV distance between estimated and true transition kernels.
The authors also investigate the application of this acquired representation in down-stream online, offline and reward-free tasks.
strength_weakness: strength_weakness: **I don't work in theory, please see the below more as minor comments.**

 * [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.

 * 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.
Similarly 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.

 * 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.

* 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.

* 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.

* The checklist instruction part was meant to be removed before submission, i.e. the instructions in L1253,1254
questions: 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: limitations: Assumptions are discussed sufficiently, so I have no concerns about additional limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. The empirical validation through experiments is limited, lacking comparisons with other state-of-the-art algorithms to support the theoretical claims.
2. The theoretical proofs need additional explanations or visual aids, with a high-level overview before detailed proofs to improve accessibility for a broader audience.
questions: questions: 1. How does the performance of MORL scale with the number of tasks and the size of the state and action spaces?
2. Are there any practical considerations or limitations in applying MORL to real-world datasets that were not covered in the paper?
limitations: 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.
paper_summary: paper_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.
The paper assumes that the new task shares the same representation as the upstream tasks making up the offline data set.
The paper provides theoretical results but does not provide any empirical evaluation.
Rather 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.
strength_weakness: strength_weakness: 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.
I have put a confidence of 2 to reflect his.

Whilst this paper addresses an important and interesting problem, I did find it hard to follow.
This is likely due to the fact I do not research theoretical aspects of RL.
Nevertheless, I do think the paper should be digestible to a wider audience than it currently is.
In particular, the authors could do a better job of communicating intuition for the different equations that are introduced.
I will now detail some suggestions with the aim of helping the authors communicate their work more clearly.

At 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.
This would really help make this paper an easier read.
Similarly in Section 4.
What is the main takeaway from this section?

Equation 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.
I would advise the authors to state this reasoning when they introduce Equation 3.4.

I 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\\}$


## Minor comments and typos 
- 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.
- Equation 2.1 - What is $\mathbb{P}$? This is not defined until way later in the paper.
- Line 42 - "We develop new" should be "We develop a new"
- 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.
- Sections shouldn't lead straight into sub-sections, e.g. Section 2 to 2.1 and Section 4 to 4.1
    - You should add a few sentences explaining what you are introducing in this section.
- Line 137 - "was" should be "were"
- Line 152 - "we construct lower confidence" should be "we construct a lower confidence"
- Line 159 - "can" should be "has"
- Line 177 - "compared to concentrability coefficient" should be "compared to the concentrability coefficient"
- Line 215 - "sequel" is the wrong word here. It sounds like you mean the next paper...
- Line 233 - "task" should be "tasks"
- Line 233 - "assume" should be "assumes"
- Line 233 - "for reward-free" should be "for the reward-free"
- 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.
- Line 299 - "such" should be "this"
- Line 343 - "study the multitask RL" should be "study multitask RL"
- Line 347 - "setting" should be "settings"
questions: questions: - Why is $\phi$ shared and $\mu$ is not? This should be made clear in the text.
- What is $\lambda$ in Equation 3.3?
- Why are equations not referenced as Equation X but instead as (X)? You could be referring to an equation/section etc...
limitations: limitations: I do not think the limitations are adequately discussed.
The 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?

Comments: title: Reply to reviewer fKn9
Reply to reviewer Qe4K
Official Comment by Reviewer fKn9
Thank you!
comment: Dear reviewer fKn9,

Thank you for spending time checking our rebuttal. We promise to incorporate the suggestions you made ( a proof roadmap and visual aids) in the final version, and we authors are still available here to answer your questions until tomorrow in case you have any final questions.

Best, Authors
Dear reviewer Qe4K,

Thank you for checking our rebuttal. We authors are still available here to answer your questions until tomorrow in case you have any last-minute questions.

Best, Authors
Thanks for the authors' response. My concerns are addressed in the rebuttal. I have raised my score to 5 accordingly.
Thank you for your detailed rebuttal and for addressing the points I raised in my review. I appreciate your clarification on the scalability of the MORL algorithm, particularly that the sub-optimality gap depends on the feature dimension rather than the state and action space sizes. I have also reviewed the other feedback and acknowledge the broader concerns about the paper's clarity. Your plan to include a proof roadmap and visual aids in the revised version is a positive step that will improve readability. I believe your responses have effectively addressed my concerns, and while I will maintain my current score, I look forward to the improvements in the final version.
We thank the reviewer for their positive review and increasing the score.
We thank the reviewer for their suggestions on how to make the discussed limitation more easily findable. We will incorporate these suggestions in our next version.

Once again thank you so much for your time in reviewing the paper and providing meaningful suggestions on how to improve the paper.
Thank you for answering my questions. After reading the rebuttal, the other reviews, and the other rebuttal comments, I have decided to maintain my score of 5. There seems to be broad agreement amongst reviewers that this paper lacks intuition and is hard to follow. I suggest the authors take on board the reviews and rework their paper. With that said, I do think this paper has potential and I am happy for it to be published if all other reviewers believe it should. However, in its current form, I do not feel like I can give a score higher than 5.

Thank you for explaining where the limitations are discussed. As both Reviewer **Qe4K** and myself overlooked these limitations, I suggest you help the reader find them, for example, name the paragraph with \paragraph{Limitations}. I'd also note that it's pretty common practice to discuss the limitations of your work in the discussion or conclusion with a clearly labelled paragraph.
I would like to thank the authors for their additional explanation, especially of the assumptions made in the problem setting and their consequences.

I will raise my score from 6 to 7, as this paper seems useful and novel to me.
However, I am keeping my confidence at 1, as I am unfortunately not very familiar with this area.