Title: Offline Multitask Representation Learning for Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Multitask learning

representation Learning

offline RL

provably efficient

suboptimality gap

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

Offline Multitask Representation Learning for Reinforcement Learning

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

1 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, 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

theoretically studied the online multitask representation learning for RL where the agent is allowed to interact with multiple source tasks to learn the shared representation (Cheng et al., 2022; Agarwal et al., 2023; Sam et al., 2024), there is currently no theoretical understanding on the effectiveness of multitask RL in the offline setting. This is crucial as in many practical scenarios (Kumar et al., 2022; Yoo et al., 2022; Lin et al., 2022), it is not feasible to interact with the different task environments in an online manner.

Moreover, when the tasks share the same representation, offline multitask representation learning can serve as a launchpad for effectively solving many other downstream tasks (Kumar et al., 2023). Consider the problem of learning to control robotic arms where we may already have offline datasets from different pick-and-place tasks in a kitchen such as the Bridge Dataset (Ebert et al., 2022). From this one can consider many possible downstream RL tasks where representation learned from these offline datasets can be beneficial. For example, one may consider solving a new pick-and-place task with different previously unseen objects in either an online or offline manner. Alternatively, one may consider a downstream reward-free RL (Jin et al., 2020a) task where the agent would first gather additional novel and diverse data without a pre-specified reward function and afterward, when provided with any reward function (e.g. slightly different target placing spot for the picked object), would be asked to provide a good policy without additional interaction.

In this work, we study the provable benefits of offline multi-task representation learning for RL in which the learner is only given access to pre-collected data from different source tasks which are modeled by low-rank MDPs (Agarwal et al., 2020b) with a shared (yet unknown) representation.

Our contributions. We develop a new offline multitask reinforcement learning algorithm that enables sample efficient representation learning in low-rank MDPs (Agarwal et al., 2020b) and further provide improved sample complexity to the downstream learning. In summary, our main contributions are:

We propose a new offline multitask representation learning algorithm called Multitask Offline Representation Learning (MORL) under low-rank MDPs. MORL represents a standard training procedure in modern machine learning, by pooling the data from all source tasks to learn a shared representation of the dynamics via maximum likelihood estimation oracle.

We prove that, MORL can learn a near-accurate model, and, when combined with the pessimism principle, find a near-optimal policy for each of the source tasks T in the average sense, more sample-efficiently, than learning each task in isolation. To our knowledge, this is the first theoretical result demonstrating the benefit of representation learning in offline multitask RL.

We then show theoretical benefits of using the learned representation from MORL in downstream reward-free RL Jin 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.

2 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

We define the action-value function for a given state-action pair ( s, a ) under policy π at step h as

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

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 function g : 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.

2.1 Offline Multitask RL with Downstream Learning

In offline multitask RL upstream learning, the agent is provided with an offline dataset collected from T source tasks, where the reward functions { r 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 h , a h ) , µ ( ∗ ,t ) h ( s ′ h ) ⟩ 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 dataset D = ⋃ 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, 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.

Algorithm 1 Multitask Offline Representation Learning (MORL) 1: Input: Dataset D = { ( s ( i,t ) h , a ( i,t ) h , r ( i,t ) h , s ( i,t ) h +1 } i ∈ [ n ] ,h ∈ [ H ] ,t ∈ [ T ] , Regularizer λ , Parameter α , Models { ( ϕ, µ ) : ϕ ∈ Ψ , µ ∈ Ψ } 2: for h = 1 , . . . , H do 3: Learn ( ̂ ϕ h , ̂ µ (1) h , . . . , ̂ µ ( T ) h ) via MLE ( ⋃ t ∈ [ T ] D ( t ) h ) as (3.1) 4: end for 5: for t = 1 , . . . , T do 6: Update estimated transitioned kernels ̂ P ( t ) h ( · | · , · ) as (3.2), empirical covariance matrix ̂ Σ ( t ) h as (3.3) and penalty term ̂ b ( t ) h as (3.4) 7: 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 ) } τ ∈ [ N off ] and D off were collected using some behavior policy ρ .

3 Upstream Offline Multitask Representation Learning

In this section, we introduce our algorithm Multitask Offline Reinforcement Learning (MORL) designed for upstream offline multitask RL in low-rank MDPs and describe its theoretical properties.

3.1 Algorithm Design

The details of the algorithm MORL is depicted in Algorithm 1. The agent passes all input offline data to estimate low-rank components ̂ ϕ h , ̂ µ (1) h , . . . , ̂ µ ( T ) h simultaneously via the Maximum Likelihood Estimation (MLE) oracle MLE ( ⋃ t ∈ [ T ] D ( t ) h ) on the joint distribution defined as follows:

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 , ̂ µ ( t ) h :

Using the representation estimator ̂ ϕ h , we set the empirical covariance matrix ̂ Σ ( t ) h, ̂ ϕ for task t as

Using both ̂ ϕ h and ̂ Σ ( t ) h, ̂ ϕ , we construct a lower confidence bound penalty term as follows:

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 .

3.2 Theoretical Result on Upstream Task

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, 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 :

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 ) = max s,a d π t P ( ∗ ,t ) h ( s, a ) /d π b t P ( ∗ ,t ) h ( s, a ) (Chen and Jiang, 2022). The relative condition number C ∗ t,h ( π t , π b t ) is always bounded by the concentrability coefficient (Uehara and Sun, 2021). In addition, the relative condition number is computed under the averaging over the state, actions, which could be much smaller than max s,a d π t P ( ∗ ,t ) h ( s, a ) /d π b t P ( ∗ ,t ) h ( s, a ) since the latter will be very large as long as any s, a pair gives large d t P ( ∗ ,t ) h ( s, a ) /d t P ( ∗ ,t ) h ( s, a ) ratio. Therefore, our C ∗ t could be much smaller compared to the concentrability coefficient, especially in large-scale MDPs (e.g. continuous state space). Moreover, in our definition of relative condition number, we use the unknown true representation ϕ ∗ . Finally, we generalize single task relative condition number to multitask setting by defining C ∗ , by simply taking the maximum over the single-task relative

π π b condition numbers.

Now we describe our main theorem.

Theorem 3.3. (a) Under Assumption 3.1, with probability at least 1 -δ , for any step h ∈ [ H ] , we have

where ̂ ϕ, ̂ P (1) , . . . , ̂ P ( T ) be the output of Algorithm 1.

(b) In addition, in Algorithm 1, if we set α = √ 2 nωζ n + λd , λ = cd log( | Φ || Ψ | T nH/δ ) with ζ n := 2 log(2 | Φ || Ψ | T nH/δ ) n and 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

where { ̂ π t } t ∈ [ T ] is the output of Algorithm 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 use ˜ O ( log | Φ | Tϵ 2 + log | Ψ | ϵ 2 ) samples per task, yet for the independent learning to use

˜ O ( log | Φ | ϵ 2 + log | Ψ | ϵ 2 ) samples 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 -δ

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 ,ϕ ∗ = n E ( 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 RL

4.1 Relationship between upstream and downstream MDPs

In order for us to theoretically study the downstream RL tasks where we would use the learned feature ̂ ϕ from the upstream tasks, first we need to make certain connection between the upstream

and downstream MDPs. Naturally, we would have to resort to some assumptions on the transition kernels to make such connections. Next, we describe these assumptions, which we largely adopt from Cheng et al. (2022).

Assumption 4.1. We make the following assumptions

For each upstream source task t with transition kernel P ( ∗ ,t ) , the behavior policy π b t is such that min 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 .

The state space S is compact and has finite measure 1 /ν , and the induced uniform probability density function is f ( s ) = ν , s ∈ S .

For any two models P 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,

for all ( s, a ) ∈ S × A and h ∈ [ H ] where C R is an absolute constant.

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

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

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 ν κ √ 2 T 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

Furthermore, for any g : S → [0 , 1] , we have ∥ ∫ ̂ µ ∗ h ( s ) g ( s ) ds ∥ 2 ≤ C L d .

4.2 Downstream Reward-Free RL

Our goal in this part is to investigate the statistical efficiency of reward-free RL in low-rank MDP while having access to offline datasets from the upstream tasks.

Our algorithm for the reward-free setting is presented in Algorithm 2 (exploration phase) and Algorithm 3 (planning phase) which is built on the procedure of optimistic learning as Wang et al. (2020); Zhang et al. (2021). While having similar design principle as in Wang et al. (2020), our algorithm differs from them due to the misspecification of representation from the upstream task. Thus, the upstream learning error affects the learning accuracy and downstream suboptimality gap and we need to account for that 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 π satisfies

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 ( TK 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) achieves a sample complexity of ˜ O ( H 22 d 7 K 9 ϵ 10 ) whereas MOFFLE (Modi et al., 2024) achieves a sample complexity of ˜ O ( H 7 d 11 K 14 min { ϵ 2 η,η 5 } ) , where η is a reachability probability to all states. More recently, Cheng et al. (2023) proposed RAFFLE which has the best-known sample complexity of ˜ O ( H 5 d 4 K ϵ 2 ) . 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), which is worse off by only ˜ O ( d ) compared to the best known sample complexity of ˜ O ( H 4 d 2 ϵ 2 ) (Hu et al., 2022) for linear MDP.

4.3 Downstream Offline and Online RL

For completeness, we also consider downstream offline and online RL which was previously studied in Cheng et al. (2022) to show the effectiveness of our offline representation. In both cases, we assume that the reward function 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 most

Additionally if Assumption 4.5 holds, and the sample size satisfies N off ≥ 40 /κ ρ · log(4 dH/δ ) , then with probability 1 -δ , we have,

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 , . . . , π N on 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 satisfies

5 Related Work

Offline Reinforcement Learning. Offline RL (Ernst et al., 2005; Riedmiller, 2005; Lange et al., 2012; Levine et al., 2020) studies the problem of learning 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. Agarwal et al. (2020b) initiates the study of low-rank MDPs. Uehara et al. (2022) proposed model-based algorithms for both online and offline RL, while Modi et al. (2024) put forward

a model-free algorithm for low-rank MDPs. Moreover, Du et al. (2019); Misra et al. (2020); Zhang et al. (2022) studied block MDPs, which is a special case of low-rank MDPs.

Offline Data Sharing in RL. There has been several empirical works that investigated the benefits of using offline datasets from multiple tasks to accelerate downstream learning (Eysenbach et al., 2020; Kalashnikov et al., 2021; Mitchell et al., 2021; Yu et al., 2021; Yoo et al., 2022). Yu et al. (2021) show that selectively sharing data between tasks can be helpful for offline multitask learning. For instance, earlier studies have investigated the development of data sharing strategies through human domain knowledge (Kalashnikov et al., 2021), inverse RL (Reddy et al., 2019; Eysenbach et al., 2020; Li et al., 2020), and estimated Q-values (Yu et al., 2021). More recently, Xu et al. (2023) uses offline dataset from diverse tasks 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.

6 Conclusion

In this paper, we theoretically study multitask RL in the offline setting. We show that offline multitask representation learning is provably more sample efficient than learning each task individually. We further show the benefit of employing the learned representation from the upstream to learn a 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.

Acknowledgments

R. Arora's and T. Nguyen-Tang's research was supported, in part, by the NSF CAREER award IIS1943251. M. Yin and M. Wang acknowledge the support by NSF IIS-2107304, NSF CPS-2312093, and ONR 1006977. The authors would like to thank Yingbin Liang and Yu-Xiang Wang for their helpful discussions.

References

Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International Conference on Machine Learning , pages 22-31. PMLR, 2017. (p. 19.)

Alekh Agarwal, Mikael Henaff, Sham Kakade, and Wen Sun. PC-PG: Policy cover directed exploration for provable policy gradient learning. Advances in Neural Information Processing Systems , 33:13399-13412, 2020a. (pp. 3 and 5.)

Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. FLAMBE: Structural complexity and representation learning of low rank MDPs. Advances in Neural Information Processing Systems , 33:20095-20107, 2020b. (pp. 2, 3, 4, 5, 8, 9, 10, 19, 29, and 53.)

Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. The Journal of Machine Learning Research , 22(1):4431-4506, 2021. (pp. 3 and 5.)

Alekh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang, Mengdi Wang, and Xuezhou Zhang. Provable benefits of representational transfer in reinforcement learning. In The Thirty Sixth Annual Conference on Learning Theory , pages 2114-2187. PMLR, 2023. (pp. 2, 4, and 10.)

András Antos, Csaba Szepesvári, and Rémi Munos. Learning near-optimal policies with Bellmanresidual minimization based fitted policy iteration and a single sample path. Machine Learning , 71:89-129, 2008. (p. 3.)

Sanjeev Arora, Simon Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi. Provable representation learning for imitation learning via bi-level optimization. In International Conference on Machine Learning , pages 367-376. PMLR, 2020. (p. 1.)

Kai Arulkumaran, Dylan R Ashley, Jürgen Schmidhuber, and Rupesh K Srivastava. All you need is supervised learning: From imitation learning to meta-RL with upside down RL. arXiv preprint arXiv:2202.11960 , 2022. (p. 1.)

Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. In The Tenth International Conference on Learning Representations , 2022. (p. 9.)

Jonathan Baxter. A model of inductive bias learning. Journal of Artificial Intelligence Research , 12: 149-198, 2000. (p. 1.)

Avinandan Bose, Simon Shaolei Du, and Maryam Fazel. Offline multi-task transfer RL with representational penalization. arXiv preprint arXiv:2402.12570 , 2024. (p. 10.)

Emma Brunskill and Lihong Li. Sample complexity of multi-task reinforcement learning. In Proceedings of the Twenty-Ninth Conference on Uncertainty in Artificial Intelligence, UAI 2013, Bellevue, WA, USA, August 11-15, 2013 . AUAI Press, 2013. (p. 1.)

Daniele Calandriello, Alessandro Lazaric, and Marcello Restelli. Sparse multi-task reinforcement learning. Advances in Neural Information Processing Systems , 27, 2014. (p. 1.)

Rich Caruana. Multitask learning. Machine Learning , 28:41-75, 1997. (p. 1.)

Jonathan Chang, Masatoshi Uehara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift in imitation learning via offline data with partial coverage. Advances in Neural Information Processing Systems , 34:965-979, 2021. (p. 53.)

Jinglin Chen and Nan Jiang. Offline reinforcement learning under value and density-ratio realizability: the power of gaps. In Uncertainty in Artificial Intelligence , pages 378-388. PMLR, 2022. (p. 5.)

Jinglin Chen, Aditya Modi, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. On the statistical efficiency of reward-free exploration in non-linear RL. Advances in Neural Information Processing Systems , 35:20960-20973, 2022a. (p. 19.)

Xiaoyu Chen, Jiachen Hu, Lin F Yang, and Liwei Wang. Near-optimal reward-free exploration for linear mixture MDPs with plug-in solver. In The Tenth International Conference on Learning Representations , 2022b. (p. 8.)

Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang. Provable benefit of multitask representation learning in reinforcement learning. Advances in Neural Information Processing Systems , 35:31741-31754, 2022. (pp. 2, 4, 5, 7, 9, 10, and 32.)

Yuan Cheng, Ruiquan Huang, Jing Yang, and Yingbin Liang. Improved sample complexity for rewardfree reinforcement learning under low-rank MDPs. In The Eleventh International Conference on Learning Representations , 2023. (pp. 2, 8, and 19.)

Christoph Dann, Tor Lattimore, and Emma Brunskill. Unifying PAC and regret: Uniform PAC bounds for episodic reinforcement learning. Advances in Neural Information Processing Systems , 30, 2017. (p. 52.)

Leandro M De Lima and Renato A Krohling. Discovering an aid policy to minimize student evasion using offline reinforcement learning. In 2021 International Joint Conference on Neural Networks (IJCNN) , pages 1-8. IEEE, 2021. (p. 1.)

Carlo D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, Jan Peters, et al. Sharing knowledge in multi-task deep reinforcement learning. In 8th International Conference on Learning Representations, 2020, Addis Ababa, Ethiopia, April 26-30, 2020 , pages 1-11, 2020. (p. 1.)

Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient RL with rich observations via latent state decoding. In International Conference on Machine Learning , pages 1665-1674. PMLR, 2019. (p. 10.)

Frederik Ebert, Yanlai Yang, Karl Schmeckpeper, Bernadette Bucher, Georgios Georgakis, Kostas Daniilidis, Chelsea Finn, and Sergey Levine. Bridge data: Boosting generalization of robotic skills with cross-domain datasets. Proceedings of Robotics: Science and System , 2022. (p. 2.)

Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research , 6, 2005. (pp. 9 and 19.)

Ben Eysenbach, Xinyang Geng, Sergey Levine, and Russ R Salakhutdinov. Rewriting history with inverse RL: Hindsight inference for policy improvement. Advances in Neural Information Processing Systems , 33:14783-14795, 2020. (p. 10.)

Kamyar Ghasemipour, Shixiang Shane Gu, and Ofir Nachum. Why so pessimistic? Estimating uncertainties for offline RL through ensembles, and why their independence matters. Advances in Neural Information Processing Systems , 35:18267-18281, 2022. (p. 9.)

Arthur Guez, Robert D Vincent, Massimo Avoli, and Joelle Pineau. Adaptive treatment of epilepsy via batch-mode reinforcement learning. In AAAI , volume 8, pages 1671-1678, 2008. (p. 1.)

Hao Hu, Yiqin Yang, Qianchuan Zhao, and Chongjie Zhang. The provable benefits of unsupervised data sharing for offline reinforcement learning. In The Eleventh International Conference on Learning Representations , 2023. (p. 10.)

Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear RL. In International Conference on Machine Learning , pages 4349-4358. PMLR, 2021. (pp. 1 and 3.)

Pihe Hu, Yu Chen, and Longbo Huang. Towards minimax optimal reward-free reinforcement learning in linear MDPs. In The Eleventh International Conference on Learning Representations , 2022. (pp. 8 and 19.)

Haque Ishfaq, Qiwen Cui, Viet Nguyen, Alex Ayoub, Zhuoran Yang, Zhaoran Wang, Doina Precup, and Lin Yang. Randomized exploration in reinforcement learning with general value function approximation. In International Conference on Machine Learning , pages 4607-4616. PMLR, 2021. (p. 53.)

Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard. Way off-policy batch deep reinforcement learning of implicit human preferences in dialog. arXiv preprint arXiv:1907.00456 , 2019. (p. 1.)

Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contextual decision processes with low Bellman rank are PAC-learnable. In International Conference on Machine Learning , pages 1704-1713. PMLR, 2017. (pp. 2 and 3.)

Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward-free exploration for reinforcement learning. In International Conference on Machine Learning , pages 4870-4879. PMLR, 2020a. (pp. 2, 4, and 19.)

Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory , pages 2137-2143. PMLR, 2020b. (pp. 3, 7, 8, 9, 20, 36, 40, 53, and 54.)

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline RL? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021. (pp. 9, 20, 43, and 53.)

Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning , pages 651-673. PMLR, 2018. (p. 1.)

Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic reinforcement learning at scale. In Conference on Robot Learning . PMLR, 2021. (pp. 1 and 10.)

Emilie Kaufmann, Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Edouard Leurent, and Michal Valko. Adaptive reward-free exploration. In Algorithmic Learning Theory , pages 865-891. PMLR, 2021. (p. 19.)

Taylor W Killian, Haoran Zhang, Jayakumar Subramanian, Mehdi Fatemi, and Marzyeh Ghassemi. An empirical study of representation learning for reinforcement learning in healthcare. arXiv preprint arXiv:2011.11235 , 2020. (p. 1.)

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:1179-1191, 2020. (p. 9.)

Aviral Kumar, Anikait Singh, Frederik Ebert, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pretraining for robots: Leveraging diverse multitask data via offline reinforcement learning. In NeurIPS 2022 Foundation Models for Decision Making Workshop , 2022. (p. 2.)

Aviral Kumar, Anikait Singh, Frederik Ebert, Mitsuhiko Nakamoto, Yanlai Yang, Chelsea Finn, and Sergey Levine. Pre-training for robots: Offline RL enables learning new tasks from a handful of trials. Proceedings of Robotics: Science and System XIX , 2023. (p. 2.)

Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the-art , pages 45-73. Springer, 2012. (p. 9.)

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020. (pp. 1 and 9.)

Alexander Li, Lerrel Pinto, and Pieter Abbeel. Generalized hindsight for reinforcement learning. Advances in Neural Information Processing Systems , 33:7754-7767, 2020. (p. 10.)

Qinjie Lin, Han Liu, and Biswa Sengupta. Switch trajectory transformer with distributional value approximation for multi-task reinforcement learning. arXiv preprint arXiv:2203.07413 , 2022. (p. 2.)

Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch off-policy reinforcement learning without great exploration. Advances in Neural Information Processing Systems , 33:1264-1274, 2020. (p. 9.)

Rui Lu, Gao Huang, and Simon S Du. On the power of multitask representation learning in linear MDP. arXiv preprint arXiv:2106.08053 , 2021. (p. 1.)

Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Emilie Kaufmann, Edouard Leurent, and Michal Valko. Fast active learning for pure exploration in reinforcement learning. In International Conference on Machine Learning , pages 7599-7608. PMLR, 2021. (p. 19.)

Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In International Conference on Machine Learning , pages 6961-6971. PMLR, 2020. (p. 10.)

Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline metareinforcement learning with advantage weighting. In International Conference on Machine Learning , pages 7780-7791. PMLR, 2021. (p. 10.)

Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. Model-free representation learning and exploration in low-rank MDPs. Journal of Machine Learning Research , 25(6):1-76, 2024. (pp. 8, 9, and 19.)

Robert Müller and Aldo Pacchiano. Meta learning MDPs with linear transition models. In International Conference on Artificial Intelligence and Statistics , pages 5928-5948. PMLR, 2022. (p. 1.)

Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research , 9(5), 2008. (p. 3.)

Thanh Nguyen-Tang and Raman Arora. VIPer: Provably efficient algorithm for offline RL with neural function approximation. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=WOquZTLCBO1 . (p. 9.)

Thanh Nguyen-Tang and Raman Arora. On the statistical complexity of offline decision-making. In Forty-first International Conference on Machine Learning , 2024a. URL https://openreview. net/forum?id=dYDPcx78tm . (p. 9.)

Thanh Nguyen-Tang and Raman Arora. On sample-efficient offline reinforcement learning: Data diversity, posterior sampling and beyond. Advances in Neural Information Processing Systems , 36, 2024b. (p. 9.)

Thanh Nguyen-Tang, Sunil Gupta, Hung Tran-The, and Svetha Venkatesh. On sample complexity of offline reinforcement learning with deep reLU networks in besov spaces. Transactions on Machine Learning Research , 2022. ISSN 2835-8856. URL https://openreview.net/forum? id=LdEm0umNcv . (p. 9.)

Thanh Nguyen-Tang, Ming Yin, Sunil Gupta, Svetha Venkatesh, and Raman Arora. On instancedependent bounds for offline reinforcement learning with linear function approximation. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 37, pages 9310-9318, 2023. (p. 9.)

Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. In Learning for Dynamics and Control , pages 1154-1168. PMLR, 2021. (p. 1.)

Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline reinforcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems , 34:11702-11716, 2021. (p. 9.)

Siddharth Reddy, Anca D Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards. In The International Conference on Learning Representations , 2019. (p. 10.)

Shideh Rezaeifar, Robert Dadashi, Nino Vieillard, Léonard Hussenot, Olivier Bachem, Olivier Pietquin, and Matthieu Geist. Offline reinforcement learning as anti-exploration. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 8106-8114, 2022. (p. 9.)

Martin Riedmiller. Neural fitted Q iteration-first experiences with a data efficient neural reinforcement learning method. In Machine Learning: ECML 2005: 16th European Conference on Machine Learning, Porto, Portugal, October 3-7, 2005. Proceedings 16 , pages 317-328. Springer, 2005. (p. 9.)

Tyler Sam, Yudong Chen, and Christina Lee Yu. The limits of transfer reinforcement learning with latent low-rank structure, 2024. URL https://arxiv.org/abs/2410.21601 . (p. 2.)

Susan M Shortreed, Eric Laber, Daniel J Lizotte, T Scott Stroup, Joelle Pineau, and Susan A Murphy. Informing sequential clinical decision-making through reinforcement learning: an empirical study. Machine learning , 84(1-2):109, 2011. (p. 1.)

Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with contextbased representations. In International Conference on Machine Learning , pages 9767-9779. PMLR, 2021. (p. 1.)

Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based RL in contextual decision processes: PAC bounds and exponential improvements over model-free approaches. In Conference on learning theory , pages 2898-2933. PMLR, 2019. (p. 2.)

Yee Teh, Victor Bapst, Wojciech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. Advances in Neural Information Processing Systems , 30, 2017. (p. 1.)

Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. In International Conference on Learning Representations , 2021. (pp. 3 and 5.)

Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. In International Conference on Learning Representations , 2022. (p. 9.)

Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline RL in low-rank MDPs. In International Conference on Learning Representations , 2022. (pp. 3, 5, 9, 10, 31, 43, and 48.)

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018. (p. 53.)

Andrew J Wagenmaker, Yifang Chen, Max Simchowitz, Simon Du, and Kevin Jamieson. Reward-free RL is no harder than reward-aware RL in linear Markov decision processes. In International Conference on Machine Learning , pages 22430-22456. PMLR, 2022. (p. 19.)

Lu Wang, Wei Zhang, Xiaofeng He, and Hongyuan Zha. Supervised reinforcement learning with recurrent neural network for dynamic treatment recommendation. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining , pages 2447-2456, 2018. (p. 1.)

Ruosong Wang, Simon S Du, Lin Yang, and Russ R Salakhutdinov. On reward-free reinforcement learning with linear function approximation. Advances in Neural Information Processing Systems , 33:17816-17826, 2020. (pp. 2, 8, and 19.)

Ruosong Wang, Dean P. Foster, and Sham M. Kakade. What are the statistical limits of offline RL with linear function approximation? In 9th International Conference on Learning Representations, ICLR 2021, Austria, May 3-7, 2021 , 2021. (p. 9.)

Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang, Ruslan Salakhutdinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. In International Conference on Machine Learning . PMLR, 2021. (p. 9.)

Annie Xie and Chelsea Finn. Lifelong robotic reinforcement learning by retaining experiences. In Conference on Lifelong Learning Agents , pages 838-855. PMLR, 2022. (p. 1.)

Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistent pessimism for offline reinforcement learning. Advances in Neural Information Processing Systems , 34:6683-6694, 2021a. (p. 9.)

Tengyang Xie, Nan Jiang, Huan Wang, Caiming Xiong, and Yu Bai. Policy finetuning: Bridging sample-efficient offline and online reinforcement learning. Advances in Neural Information Processing Systems , 34, 2021b. (p. 9.)

Yifan Xu, Nicklas Hansen, Zirui Wang, Yung-Chieh Chan, Hao Su, and Zhuowen Tu. On the feasibility of cross-task transfer with model-based reinforcement learning. In The Eleventh International Conference on Learning Representations , 2023. (p. 10.)

Yuling Yan, Gen Li, Yuxin Chen, and Jianqing Fan. The efficacy of pessimism in asynchronous Q-learning. IEEE Transactions on Information Theory , 2023. (p. 9.)

Ming Yin and Yu-Xiang Wang. Towards instance-optimal offline reinforcement learning with pessimism. Advances in Neural Information Processing Systems , 34:4065-4078, 2021. (p. 9.)

Ming Yin, Yu Bai, and Yu-Xiang Wang. Near-optimal provable uniform convergence in offline policy evaluation for reinforcement learning. In International Conference on Artificial Intelligence and Statistics , pages 1567-1575. PMLR, 2021. (pp. 7 and 9.)

Ming Yin, Yaqi Duan, Mengdi Wang, and Yu-Xiang Wang. Near-optimal offline reinforcement learning with linear representation: Leveraging variance information with pessimism. International Conference on Learning Representations , 2022. (p. 9.)

Minjong Yoo, Sangwoo Cho, and Honguk Woo. Skills regularized task decomposition for multitask offline reinforcement learning. Advances in Neural Information Processing Systems , 35: 37432-37444, 2022. (pp. 2 and 10.)

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems , 33:14129-14142, 2020. (p. 9.)

Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Sergey Levine, and Chelsea Finn. Conservative data sharing for multi-task offline reinforcement learning. Advances in Neural Information Processing Systems , 34:11501-11516, 2021. (pp. 1 and 10.)

Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Chelsea Finn, and Sergey Levine. How to leverage unlabeled data in offline reinforcement learning. In International Conference on Machine Learning , pages 25611-25635. PMLR, 2022. (p. 1.)

Andrea Zanette, Alessandro Lazaric, Mykel J Kochenderfer, and Emma Brunskill. Provably efficient reward-agnostic navigation with linear value iteration. Advances in Neural Information Processing Systems , 33:11756-11766, 2020. (p. 19.)

Andrea Zanette, Ching-An Cheng, and Alekh Agarwal. Cautiously optimistic policy optimization and exploration with linear function approximation. In Conference on Learning Theory , pages 4473-4525. PMLR, 2021a. (p. 32.)

Andrea Zanette, Martin J Wainwright, and Emma Brunskill. Provable benefits of actor-critic methods for offline reinforcement learning. Advances in Neural Information Processing Systems , 34: 13626-13640, 2021b. (p. 9.)

Xianyuan Zhan, Haoran Xu, Yue Zhang, Xiangyu Zhu, Honglei Yin, and Yu Zheng. Deepthermal: Combustion optimization for thermal power generating units using offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 4680-4688, 2022. (p. 1.)

Weitong Zhang, Dongruo Zhou, and Quanquan Gu. Reward-free model-based reinforcement learning with linear function approximation. Advances in Neural Information Processing Systems , 34, 2021. (pp. 8 and 19.)

Xuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun. Efficient reinforcement learning in block MDPs: A model-free representation learning approach. In International Conference on Machine Learning , pages 26517-26547. PMLR, 2022. (p. 10.)

Contents

A Additional Related Work

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) settings. Agarwal et al. (2020b); Modi et al. (2024); Chen et al. (2022a); Cheng et al. (2023) study reward-free RL in low-rank MDPs which is particularly interesting as here representation learning is interwined with reward-free exploration.

B Omitted Algorithms

B.1 Algorithms for Downstream Reward-Free RL

Algorithm 3 Downstream Reward-Free Exploration: Planning Phase

1: Input: Feature ̂ ϕ , dataset D RFE = { ( s k h , a k h ) } k ∈ [ K RFE ] ,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

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 ) } τ ∈ [ N off ] ,h ∈ [ H ] , parameters λ d , β, ξ down

2: Initialization: ̂ V H +1 = 0

3: for h = H,H -1 , . . . , 1 do

4: Λ h = ∑ N off τ =1 ̂ ϕ h ( s τ h , a τ h ) ̂ ϕ h ( s τ h , a τ h ) ⊤ + λ d I d

6: Γ h ( · , · ) = Hξ down + β [ ̂ ϕ ⊤ h Λ -1 h ̂ ϕ h ( · , · )] 1 / 2

5: ̂ w h = Λ -1 h ( ∑ N off τ =1 ̂ ϕ h ( s τ h , a τ h ) · ( r τ h + ̂ V h +1 ( s τ h +1 )))

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 , . . . , π N on and ˜ π where ˜ π is the uniform mixture of π 1 , . . . , π N on

C Notations

We summarize frequently used notations in the following list.

For any h ∈ [ H ] , we define

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

D Proof of Multitask Offline Representation Learning

We first state the supporting lemmas that are used in the proof of Theorem 3.3. The proof of these lemmas are provided in Appendix E.

D.1 Supporting Lemmas

In the following lemma, we provide an upper bound for the model estimation error for each task that captures the advantage of joint MLE model estimation over single-task learning.

Lemma D.1. For any task t , policy π t , and reward r t , we have for all h ≥ 2 ,

and for h = 1 , we have

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 -δ

D.2 Proof of Theorem 3.3

We first restate Theorem 3.3.

Theorem D.3 (Restatement of Theorem 3.3) . Under Assumption 3.1, with probability at least 1 -δ , for any step h ∈ [ H ] , we have

where ̂ ϕ, ̂ P (1) , . . . , ̂ P ( T ) be the output of Algorithm 1.

In addition, in Algorithm 1, if we set α = √ 2 nωζ n + λd , λ = cd log( | Φ || Ψ | T nH/δ ) with ζ n := 2 log(2 | Φ || Ψ | T nH/δ ) n and 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 have

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:

and

From Lemma F.3 and Lemma G.2, this event happens with probability 1 -δ . Conditioning on this event, we have

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 have

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

From (D.4), we have

Next, we upper bound E ( s h -1 ,a h -1 ) ∼ ( P ( ∗ ,t ) ,π t ) [ ∥ ϕ ∗ h -1 ( s h -1 , a h -1 ) ∥ 2 (Σ h -1 ,π b t ,ϕ ∗ ) -1 ] as the following

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 get

Following similar steps as in (D.7), we can further show that

Now noting | f ( t ) h | ∞ ≤ 1 , for h ≥ 2 ,we get

Further, note that,

where the last inequality follows from the Cauchy-Schwarz inequality and (D.3).

Finally, from (D.5) we get

So, we have

E Proof of Supporting Lemmas in Appendix D

In this section, we provide the proofs of the lemmas that we used in the proof of Theorem D.3.

E.1 Proof of Lemma D.1

Proof of Lemma D.1. For h = 1 ,

where 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 have

where ( 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.2

Proof of Lemma D.2. We condition on the events:

and

From Lemma F.3 and Lemma G.2, this event happens with probability 1 -δ . Conditioning on this event, we have

We have

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 that ∑ T t =1 ζ ( t ) h ≤ ζ n , ( v ) follows from (E.2).

F Multitask Offline MLE

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 , where x 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

Lemma F.2 (Lemma 25 of Agarwal et al. (2020b)) . For any two conditional probability densities f 1 , f 2 and any distribution D ∈ ∆( X ) we have

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 have

Proof of Lemma F.3. Let ̂ f ( D ) denote empirical maximum likelihood estimator:

We combine Lemma F.1 with the Chernoff method to obtain the following exponential tail bound: with probability 1 -δ , we have

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) is

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

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 get

Using union bound, we obtain that for any h ∈ [ H ] and n with probability at least 1 -δ/ 2 , it holds that

This completes the proof.

G One-Step Back Lemma and Concentration of Penalty Term

G.1 One-step back lemma

The following one-step back lemma is 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 function E ( 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 ,ϕ = n E ( 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 have

Proof of Lemma G.1. First, we have

where the inequality follows from Cauchy-Schwarz inequality.

Then,

where ( i ) follows from the assumption ∥ g ∥ ∞ ≤ B and for any function h : 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.

G.2 Concentration of penalty term

Recall that Σ h,π b t ,ϕ = n E ( s h ,a h ) ∼ ( P ( ∗ ,t ) ,π b t ) [ ϕϕ ⊤ ] + λI . Thus, ̂ Σ ( t ) h 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(2 nTH | Φ | /δ )) for any n . With probability at least 1 -δ/ 2 , we have that for any n ∈ N , h ∈ [ H ] , t ∈ [ T ] and ϕ ∈ Φ ,

where β 1 and β 2 are some absolute constants.

H Proof of Lemma 4.3: Approximate Feature for New Task

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 ν κ √ 2 T 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

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

where ( i ) , ( ii ) follows from Assumption 4.1 and ( iii ) follows from Theorem 3.3.

Defining T

where ( i ) follows from Assumption 4.1, (H.1) and the fact that c t ∈ [0 , C L ] for all t ∈ [ T ] . Moreover, by normalization, for any g : S → [0 , 1] , we get

where the last inequality follows from Assumption 3.1.

I Proof for Downstream Reward-Free RL

For any h ∈ [ H ] , we define

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

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:

We can similarly define ˜ V π h ( s, r ) and ˜ Q π h ( s, a, r ) .

I.1 Supporting Lemmas

Now we state the supporting lemmas that are used in the proof of Theorem 4.4. The proof of these lemmas are provided 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 satisfies

Lemma I.3. Let E be the event that for all ( k, h ) ∈ [ K RFE ] × [ H ] ,

Then Pr [ E ] ≥ 1 -δ/ 8 .

Lemma I.4. With probability 1 -δ/ 8 , we have for all ( s, a ) ∈ S × A ,

Lemma I.5. With probability 1 -δ/ 4 , for all ( h, k ) ∈ [ H ] × [ K RFE ] , and any s ∈ S , we have

and

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 have

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 have

I.2 Proof of Theorem 4.4

We 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 π satisfies

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 ( TK 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 have

This implies

where π is the policy returned by Algorithm 3.

Observe that, using Lemma I.7, we have

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,

By definition of ˜ V ∗ 1 ( s, u ) , we further have E s ∼ µ [ ˜ V π 1 ( s, u )] ≤ E s ∼ µ [ ˜ V ∗ 1 ( s, u )] . By Lemma I.6, we have

So, we have,

.

Recall the definition of ξ down from Lemma 4.3, that is, ξ down = ξ + C L C R ν κ √ 2 T 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 ( TK RFE /d 3 ) , then the first term in Equation (I.2) dominates the second term 6 H 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 have

This completes the proof.

J Proof of Supporting Lemmas in Appendix I

J.1 Proof of Lemma I.2

Proof of Lemma I.2. We have

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

J.2 Proof of Lemma I.3

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 have

Thus, our value function ̂ V k h +1 is of the form

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 that

We 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.4

Proof of Lemma I.4. We condition on the event E defined in Lemma I.3, which holds with probability at least 1 -δ/ 8 . We define ˜ w 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,

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,

We now bound ( a ) , ( b ) , ( c ) in (J.2) individually.

Term (a). We have,

Term (b). Using Cauchy-Schwarz inequality and the definition of the event E from Lemma I.3, we have

Term (c). We have

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 ( dK RFE H max( ξ down , 1) δ ) + Hξ down √ dk we get

Putting everything together in (J.1), we get

J.4 Proof of Lemma I.5

Proof of Lemma I.5. We prove the first part using backward induction on h . For h = H , we have

Suppose for some h +1 ∈ [ H ] , it holds that for all s ∈ S ,

Then,

where the last inequality follows from Lemma I.4.

Thus, we have

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,

Conditioned on E from Lemma I.3 where Pr [ E ] ≥ 1 -δ/ 8 ,

where the second inequality follows from Lemma I.4.

Note that for each h ∈ [ H -1] , { ξ k h } K RFE k =1 is a martingale difference sequence with | ξ k h | ≤ H . We define the event E ′ to be the event that

By Azuma-Hoeffding inequality, we have Pr [ E ′ ] ≥ 1 -δ/ 8 .

By Cauchy-Schwarz inequality, we have

Using Lemma D.2 of Jin et al. (2020b), we have

Conditioned on E and E where Pr ( E ∩ E ′ ) ≥ 1 -δ/ 4 , we have

which completes the proof.

J.5 Proof of Lemma I.6

Proof 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 } K RFE k =1 is a martingale difference sequence. Moreover, we have | ∆ k | ≤ H . Using Azuma-Hoeffding inequality, with probability 1 -δ/ 4 , we have

where c 1 > 0 is an absolute constant. Therefore, we have

Now, notice that for all k ∈ [ K RFE ] , Λ h ⪰ Λ k h . Thus for all ( k, h ) ∈ [ K RFE ] × [ H ] , we have

which implies

Using Lemma I.5 and union bound, we have with probability 1 -δ/ 2 ,

for some absolute constant c ′ > 0 . This completes the proof.

J.6 Proof of Lemma I.7

Proof 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

Thus, for all h ∈ [ H ] and ( s, a ) ∈ S × A , we have

Since, ̂ Q h ( s, a ) ≤ H and u h ( · , · ) = min { β √ ̂ ϕ ( · , · ) ⊤ (Λ h ) -1 ̂ ϕ ( · , · ) , H } , we have,

This completes the first part of the proof.

Now, using induction, we prove that for all h ∈ [ H ] and ( s, a ) ∈ S × A , we have Q ∗ 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 ,

Note that,

Since Q ∗ h ( s, a, r ) ≤ H and u h ( · , · ) = min { β √ ̂ ϕ ( · , · ) ⊤ (Λ h ) -1 ̂ ϕ ( · , · ) , H } , it suffices to show that

Now,

This completes the proof.

K Proof for Downstream Offline RL

First, we state the supporting lemmas that are used in the proof of Theorem 4.6. The proof of these lemmas are provided in Appendix L.

K.1 Supporting Lemmas

The following lemma shows that the linear weight ̂ 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 satisfies

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 have

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 event

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,

The following lemma decomposes the suboptimality gap into the summation of uncertainty metric of each step.

Applying the max operator on both side of the inductive hypothesis, we get

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 have

where Γ h ( s, a ) = β ∥ ∥̂ ϕ h ( s, a ) ∥ ∥ Λ -1 h + Hξ down. Furthermore, we have

K.2 Proof of Theorem 4.6

We 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 most

Additionally if Assumption 4.5 holds, and the sample size satisfies N off ≥ 40 /κ ρ · log(4 dH/δ ) , then with probability 1 -δ , we have,

Remark K.7 . Compared to Theorem 4.4 in Jin et al. (2021), the suboptimality gap in (K.1) has an additional term 2 H 2 ξ down. When the linear misspecification error ξ = ˜ O ( √ d/N off ) and the number of trajectories n in each upstream offline task dataset satisfies n = ˜ O ( TN off d ) , the RHS of (K.2) is dominated by ˜ O ( N -1 / 2 off H 2 d ) improving the suboptimality gap bound of REP-LCB (Uehara et al., 2022) by an order of ˜ O ( Hd ) under low-rank MDP with unknown representation.

Proof of Theorem K.6. Letting Γ h = β ∥ ̂ ϕ h ( s, a ) ∥ Λ -1 h + Hξ down in Lemma K.5, with probability at least 1 -δ , we have

This finishes the first part of the proof. Now we will provide the suboptimality bound under the feature coverage assumption in Assumption 4.5. From Appendix B.4 of Jin et al. (2021), if N off ≥ 40 /κ ρ · log(4 dH/δ ) , then with probability 1 -δ/ 2 , for any h ∈ [ H ] and ( s, a ) ∈ S × A , we have

Setting β ( δ/ 2) = c β ( Hd √ ι ( δ/ 2) + H √ dN off ξ down ) , with probability 1 -δ/ 2 , we have

Using union bound, we have the following bound with probability at least 1 -δ

L Proof of Supporting Lemmas in Appendix K

In this section, we provide the proofs of the lemmas that we used in the proof of Theorem 4.6.

L.1 Proof of Lemma K.1

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.

L.2 Proof of Lemma K.2

Proof of Lemma K.2. The value function ̂ V h +1 has the parametric form of

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 -δ ,

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) by

Therefore, we have

L.3 Proof of Lemma K.3

Proof of Lemma K.3. For h ∈ [ H ] , we define w ∗ h = ∫ ′ µ ∗ ( s ′ ) ̂ V h +1 ( s ′ ) ds ′ . Then we have

where ( i ) follows from the fact that | ̂ V h +1 ( s ) | ≤ H for all s ∈ S and ( ii ) follows from Lemma 4.3. Wenowdecompose the second term in (L.2). Recall that Λ h = ∑ N off τ =1 ̂ ϕ h ( s τ h , a τ h ) ̂ ϕ h ( s τ h , a τ h ) ⊤ + λ d I d and w h = Λ -1 h ∑ N off τ =1 ̂ ϕ h ( s τ h , a τ h ) ̂ V h +1 ( s τ h +1 ) . Then, we have

We now provide an upper bound for each of the terms in (L.3).

Term (a). We have

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

where the first inequality comes from Cauchy Schwarz inequality and the last inequality follows from Lemma K.2.

Term (c). We have

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 -δ ,

This completes the proof.

L.4 Proof of Lemma K.5

Proof of Lemma K.5. Recall that

If ̂ ϕ h ( s, a ) ⊤ ̂ w h -Γ h ( s, a ) > 0 , we have

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

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] and ̂ V h +1 ∈ [0 , H -h ] . Now,

Now, from the definition of l h , we have

where the last inequality follows from the conditioning on the event E ( δ ) . Finally, we obtain

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 .

M Proof for Downstream Online RL

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.

M.1 Supporting Lemmas

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 have

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

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

for some ∆ n h ( s, a ) that satisfies ∥ ∆ n h ( s, a ) ∥ ≤ β n ∥ ̂ ϕ h ( s, a ) ∥ (Λ n h ) -1 +2 Hξ 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 have

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

M.2 Proof of Theorem 4.7

We are now ready to prove the main theorem in this section. We first restate Theorem 4.7.

Theorem M.5. Let ˜ π be the uniform mixture of π 1 , . . . , π N on in Algorithm 5. Under Assumption 4.1, setting λ d = 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 satisfies

Remark M.6 . As we use estimated representation { ̂ ϕ h } H h =1 from the upstream tasks, we get an extra term H 2 dξ down in the suboptimality gap above. When the linear misspecification error ξ = ˜ O ( √ d/N off ) and the of trajectories n in each upstream offline task dataset satisfies n = ˜ O ( TN on d ) , the RHS of (M.1) is dominated by ˜ O ( H 2 d 3 / 2 N -1 / 2 on ) 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 by

where ( i ) follows from Lemma M.3 and ( ii ) follows from Lemma M.4.

Note that in term ( a ) , { ξ n h } N on ,H n =1 ,h =1 is a martingale difference sequence with | ξ n h | ≤ 2 . By AzumaHoeffding inequality, we have, with probability at least 1 -δ/ 4 ,

For term ( b ) , we have

where ( i ) follows from Cauchy-Schwarz inequality, ( ii ) follows from Lemma O.6, and ( iii ) follows from the inequality that √ x + y ≤ √ x + √ y for all x, y ≥ 0 .

Combining (M.2), (M.3) and (M.4) we get

Dividing both sides by N on, we get

This completes the proof.

N Proof of Supporting Lemmas in Appendix M

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

Proof of Lemma M.2. For any policy π , we define w π 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

Using Lemma 4.3, we bound the second term as

π

where we used the observation | V h +1 | ≤ H . Now, the first term in (N.1) can be bounded by

We now bound ( a ) , ( b ) , ( c ) , ( d ) in (N.3) individually.

Term (a). We have,

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

Term (c). We have

where ( i ) follows from similar steps as in (N.4) and ( ii ) follows from Lemma 4.3. Term (d). We have

where ( i ) follows from Lemma 4.3 and | ̂ V n h +1 ( s ) | ≤ H for all s ∈ S , ( ii ) follows from CauchySchwarz 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

Combining (N.8), (N.2) in (N.1) completes the proof.

N.2 Proof of Lemma M.3

Proof of Lemma M.3. We use backward induction for our proof. First, we prove the base case, at the last step H . By Lemma M.2, we have

Thus, we have

Now, suppose the statement holds true at step h +1 and consider the step h . Using Lemma M.2, we have

Therefore,

which completes the proof.

N.3 Proof of Lemma M.4

Proof 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 have

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

O Auxiliary Lemmas

O.1 Miscellaneous Lemmas

The following lemma measures the difference between two value functions under two MDPs and reward functions. Here, we use a shorthand notation P 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,

We use the following lemma to deal with distribution shift in offline RL setting.

Lemma O.2 (Distribution shift lemma (Chang et al., 2021)) . Consider two distributions ρ 1 ∈ ∆( S × A ) and ρ 2 ∈ ∆( S × A ) , and a feature mapping ϕ : S × A → R d . Denote C := sup x ∈ R d x ⊤ E s,a ∼ ρ 1 ϕ ( s,a ) ϕ ( s,a ) ⊤ x x ⊤ E s,a ∼ ρ 2 ϕ ( s,a ) ϕ ( s,a ) ⊤ x . Then for any positive definite matrix Λ , we have

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

O.2 Inequalities for summations

Lemma O.4 (Lemma D.1 in Jin et al. (2020b)) . Let Λ h = λI + ∑ t i =1 ϕ i ϕ ⊤ i , where ϕ i ∈ R d and λ > 0 . Then it holds that

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

Lemma O.6 (Lemma G.2 in (Agarwal et al., 2020b)) . Consider a sequence of semidefinite matrices X 1 , . . . , X N ∈ R d × d with Tr( X n ) ≤ 1 for all n ∈ [ N ] . Define M 0 = λI and M n = M n -1 + X n . Then

O.3 Covering numbers and self-normalized processes

Lemma O.7 (Lemma D.4 in Jin et al. (2020b)) . 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 have

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 satisfies

Lemma O.9 ( ε -covering number (Jin et al., 2020b)) . Let V denote a class of functions mapping from S to R with the following parametric form

where the parameters ( w,β, Λ) satisfy ∥ w ∥ ≤ L , β ∈ [0 , B ] and the minimum eigenvalue satisfies λ min (Λ) ≥ λ . Assume ∥ ϕ ( s, a ) ∥ ≤ 1 for all ( s, a ) pairs, and let N ε be the ε -covering number of V with respect to the distance dist ( V, V ′ ) = sup s | V ( s ) -V ′ ( s ) | . Then,

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims made in the abstract and the introduction are reflected in the paper's main contribution.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the assumptions made in order for our theorems to hold true and discuss the limitations of those assumptions.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide detailed list of assumptions along with their interpretation. The complete proofs are provided in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [NA]

Justification: We do not have any experiment in this paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: There is no experiment in the paper as it is theoretical in nature.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

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

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

Justification: The paper does not include any experiment.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification:

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

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

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Not applicable

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: There was no crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: There was no study subject.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates how representation learning in offline multitask low-rank RL can improve sample complexity when using the learned representations in downstream reward-free  RL, offline and online RL settings.\nThe paper assumes that the new task shares the same representation as the upstream tasks making up the offline data set.\nThe paper provides theoretical results but does not provide any empirical evaluation.\nRather intuitively, their result shows that learning a representation from the offline multitask data set can improve sample complexity in new downstream tasks when compared to learning the new downstream task from scratch.\n\nSTRENGTHS:\nRepresentation learning and large-scale offline data sets have shown much promise for learning policies.\nAs such, investigating the theory of learning representations from multi-task offline data seems an important research direction.\nI found the related work section very beneficial for pulling different pieces of the paper together. This really helped me understand things.\n\nWEAKNESSES:\nI should start off by saying that I have little experience with the theoretical aspects of representation learning for RL and there were many aspects of this paper I had to take as a given without fully understanding.\nI have put a confidence of 2 to reflect his.\n\nWhilst this paper addresses an important and interesting problem, I did find it hard to follow.\nThis is likely due to the fact I do not research theoretical aspects of RL.\nNevertheless, I do think the paper should be digestible to a wider audience than it currently is.\nIn particular, the authors could do a better job of communicating intuition for the different equations that are introduced.\nI will now detail some suggestions with the aim of helping the authors communicate their work more clearly.\n\nAt the end of Section 3, it was not clear to me what the main result is. I would advise the authors to summarise the main result at the end of the section.\nThis would really help make this paper an easier read.\nSimilarly in Section 4.\nWhat is the main takeaway from this section?\n\nEquation 3.4 is introduced rather abruptly. That is, I did not understand its purpose until later in the paper when the authors stated that they wanted to improve sample complexity by improving exploration.\nI would advise the authors to state this reasoning when they introduce Equation 3.4.\n\nI am confused by the notation for $h \\in [H]$, $t \\in [T]$ and $i \\in [N]$. Aren't these discrete numbers? If so, the authors should use curly braces instead, e.g. $h \\in \\\\{0,\\ldots,H\\\\}$\n\n\n## Minor comments and typos \n- Line 61 - $H$ and $d$ aren't defined yet so this doesn't make much sense. I found the contributions had way too many technical details which does not make any sense to the reader given what they've read so far. Consider shortening these and removing unnecessary technical details.\n- Equation 2.1 - What is $\\mathbb{P}$? This is not defined until way later in the paper.\n- Line 42 - \"We develop new\" should be \"We develop a new\"\n- Line 69 - MDPs are discrete time by definition so it is confusing to state you consider discrete-time MDPs. It almost implied MDPs are not usually discrete time.\n- Sections shouldn't lead straight into sub-sections, e.g. Section 2 to 2.1 and Section 4 to 4.1\n    - You should add a few sentences explaining what you are introducing in this section.\n- Line 137 - \"was\" should be \"were\"\n- Line 152 - \"we construct lower confidence\" should be \"we construct a lower confidence\"\n- Line 159 - \"can\" should be \"has\"\n- Line 177 - \"compared to concentrability coefficient\" should be \"compared to the concentrability coefficient\"\n- Line 215 - \"sequel\" is the wrong word here. It sounds like you mean the next paper...\n- Line 233 - \"task\" should be \"tasks\"\n- Line 233 - \"assume\" should be \"assumes\"\n- Line 233 - \"for reward-free\" should be \"for the reward-free\"\n- Table 1 - It is not clear from the caption that this is for the reward-free setting. This should be made clear in the caption.\n- Line 299 - \"such\" should be \"this\"\n- Line 343 - \"study the multitask RL\" should be \"study multitask RL\"\n- Line 347 - \"setting\" should be \"settings\"\n\nQUESTIONS:\n- Why is $\\phi$ shared and $\\mu$ is not? This should be made clear in the text.\n- What is $\\lambda$ in Equation 3.3?\n- Why are equations not referenced as Equation X but instead as (X)? You could be referring to an equation/section etc...\n\nLIMITATIONS:\nI do not think the limitations are adequately discussed.\nThe proofs utilise a lot of assumptions and the paper would benefit from a clear paragraph demonstrating the limitations associated with these assumptions. What are the practical implications of these assumptions?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. The paper addresses the novel problem of offline multitask representation learning in RL with a new algorithm (MORL) that incorporates innovative techniques such as joint model learning via MLE and penalty functions for pessimism. The approach is well-grounded in theory and fills a relevant gap in the literature.\n2. The submission is technically sound with rigorous theoretical analysis supporting the proposed algorithm. Detailed proofs of the main theoretical results and thorough explanations of the algorithm's components are provided.\n3. The paper is clearly written and well-organized, with a logical progression from problem introduction to algorithm presentation and theoretical results. The definitions and assumptions are clearly stated, and the proofs are detailed and comprehensible.\n4. The results are important as they provide the first theoretical demonstration of the benefits of multitask representation learning in offline RL. The approach has potential applications in domains where collecting data online is infeasible, making it highly relevant for real-world problems.\n\nWEAKNESSES:\n1. The empirical validation through experiments is limited, lacking comparisons with other state-of-the-art algorithms to support the theoretical claims.\n2. The theoretical proofs need additional explanations or visual aids, with a high-level overview before detailed proofs to improve accessibility for a broader audience.\n\nQUESTIONS:\n1. How does the performance of MORL scale with the number of tasks and the size of the state and action spaces?\n2. Are there any practical considerations or limitations in applying MORL to real-world datasets that were not covered in the paper?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a theoretical analysis of representation learning in Multi-Task Offline RL.\nSpecifically they consider a setting in which the transition kernels have a low-rank decomposition $P(s'|s,a)=\\langle \\phi(s'),\\mu(s,a)\\rangle$, such that $\\mu(s,a)$ depends on the task but $\\phi(s')$ is shared by all tasks.\nThe authors propose an algorithm for this setting that first learns a transition model and then does planning on the model with a low-confidence-penalized reward.\nThey provide a suboptimality bound, based on bounding the TV distance between estimated and true transition kernels.\nThe authors also investigate the application of this acquired representation in down-stream online, offline and reward-free tasks.\n\nSTRENGTHS:\n* The paper is very well written, assumptions that are made are explained and well justified.\n * The exploration of down-stream applications of Multi-Task Offline RL is of large practical interest.\n\nWEAKNESSES:\n**I don't work in theory, please see the below more as minor comments.**\n\n * [89] analyzes a very similar setting but is only mentioned in the appendix as concurrent work. Strictly speaking, [89] was uploaded ot arxiv three months before the deadline, so it is not considered concurrent as per NeurIPS guidelines. Either way, it should probably be discussed more prominently and not just in the Appendix.\n\n * It would have been nice to specify which improvements in the bound are results of which decisions. For example, a conservative reward $r-b$ in planning vs the normal reward $r$, or how exactly using the low-rank assumption rather than a linearity assumption affects the analysis.\nSimilarly it would be nice to provide an interpretation for the different terms in (3.7), but I'm also not sure if that is really possible to do in a meaningful way.\n\n * The paper treats a setting with non-stationary reward and transition functions. It's is not discussed why the paper treats this setting and what effect it has on the analysis.\n\n* The approach name and acronym MORL (Multitask Offline Representation Learning) is a somewhat confusing choice, as MORL is frequently used to denote Multi-Objective RL.\n\n* The notation is somewhat confusing at times with $t$ representing tasks, $h$ representing timesteps in the MDP, $\\tau$ indexing transitions in the downstream offline dataset and $i$ indexing samples in the upstream dataset. The usage of non-stationary policies, rewards,  and transition functions also does not help to reduce subscripts.\n\n* The checklist instruction part was meant to be removed before submission, i.e. the instructions in L1253,1254\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nAssumptions are discussed sufficiently, so I have no concerns about additional limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

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

STRENGTHS:
1. Detailed theoretical analysis of the proposed algorithm.
1. Better sample complexity than existing algorithms.

WEAKNESSES:
I do not conduct research in theoretical offline RL, so I cannot provide an accurate evaluation for this paper. I will try to provide some comments based on my understanding.
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:
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:
The paper lacks a limitations section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-20