Offline Multitask Representation Learning for Reinforcement Learning

Summary

This paper investigates how representation learning in offline multitask low-rank RL can improve sample complexity when using the learned representations in downstream reward-free RL, offline and online RL settings. The paper assumes that the new task shares the same representation as the upstream tasks making up the offline data set. The paper provides theoretical results but does not provide any empirical evaluation. Rather intuitively, their result shows that learning a representation from the offline multitask data set can improve sample complexity in new downstream tasks when compared to learning the new downstream task from scratch.

Strengths

Representation learning and large-scale offline data sets have shown much promise for learning policies. As such, investigating the theory of learning representations from multi-task offline data seems an important research direction. I found the related work section very beneficial for pulling different pieces of the paper together. This really helped me understand things.

Weaknesses

I should start off by saying that I have little experience with the theoretical aspects of representation learning for RL and there were many aspects of this paper I had to take as a given without fully understanding. I have put a confidence of 2 to reflect his. Whilst this paper addresses an important and interesting problem, I did find it hard to follow. This is likely due to the fact I do not research theoretical aspects of RL. Nevertheless, I do think the paper should be digestible to a wider audience than it currently is. In particular, the authors could do a better job of communicating intuition for the different equations that are introduced. I will now detail some suggestions with the aim of helping the authors communicate their work more clearly. At the end of Section 3, it was not clear to me what the main result is. I would advise the authors to summarise the main result at the end of the section. This would really help make this paper an easier read. Similarly in Section 4. What is the main takeaway from this section? Equation 3.4 is introduced rather abruptly. That is, I did not understand its purpose until later in the paper when the authors stated that they wanted to improve sample complexity by improving exploration. I would advise the authors to state this reasoning when they introduce Equation 3.4. I am confused by the notation for $h \in [H]$, $t \in [T]$ and $i \in [N]$. Aren't these discrete numbers? If so, the authors should use curly braces instead, e.g. $h \in \\{0,\ldots,H\\}$ ## Minor comments and typos - Line 61 - $H$ and $d$ aren't defined yet so this doesn't make much sense. I found the contributions had way too many technical details which does not make any sense to the reader given what they've read so far. Consider shortening these and removing unnecessary technical details. - Equation 2.1 - What is $\mathbb{P}$? This is not defined until way later in the paper. - Line 42 - "We develop new" should be "We develop a new" - Line 69 - MDPs are discrete time by definition so it is confusing to state you consider discrete-time MDPs. It almost implied MDPs are not usually discrete time. - Sections shouldn't lead straight into sub-sections, e.g. Section 2 to 2.1 and Section 4 to 4.1 - You should add a few sentences explaining what you are introducing in this section. - Line 137 - "was" should be "were" - Line 152 - "we construct lower confidence" should be "we construct a lower confidence" - Line 159 - "can" should be "has" - Line 177 - "compared to concentrability coefficient" should be "compared to the concentrability coefficient" - Line 215 - "sequel" is the wrong word here. It sounds like you mean the next paper... - Line 233 - "task" should be "tasks" - Line 233 - "assume" should be "assumes" - Line 233 - "for reward-free" should be "for the reward-free" - Table 1 - It is not clear from the caption that this is for the reward-free setting. This should be made clear in the caption. - Line 299 - "such" should be "this" - Line 343 - "study the multitask RL" should be "study multitask RL" - Line 347 - "setting" should be "settings"

Questions

- Why is $\phi$ shared and $\mu$ is not? This should be made clear in the text. - What is $\lambda$ in Equation 3.3? - Why are equations not referenced as Equation X but instead as (X)? You could be referring to an equation/section etc...

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

I do not think the limitations are adequately discussed. The proofs utilise a lot of assumptions and the paper would benefit from a clear paragraph demonstrating the limitations associated with these assumptions. What are the practical implications of these assumptions?

Summary

This paper introduces the Multitask Offline Representation Learning (MORL) algorithm, which aims to enhance sample efficiency in offline multitask reinforcement learning (RL). By learning a shared representation from pre-collected datasets of different tasks, modeled by low-rank Markov Decision Processes (MDPs), the authors demonstrate the theoretical benefits of this method for various downstream RL tasks, including reward-free, offline, and online scenarios. The paper provides extensive theoretical analysis and introduces specific data-coverage assumptions, presenting the first theoretical results showing the advantages of representation learning in offline multitask RL.

Strengths

1. The paper addresses the novel problem of offline multitask representation learning in RL with a new algorithm (MORL) that incorporates innovative techniques such as joint model learning via MLE and penalty functions for pessimism. The approach is well-grounded in theory and fills a relevant gap in the literature. 2. 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. 3. 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. 4. The results are important as they provide the first theoretical demonstration of the benefits of multitask representation learning in offline RL. The approach has potential applications in domains where collecting data online is infeasible, making it highly relevant for real-world problems.

Weaknesses

1. The empirical validation through experiments is limited, lacking comparisons with other state-of-the-art algorithms to support the theoretical claims. 2. The theoretical proofs need additional explanations or visual aids, with a high-level overview before detailed proofs to improve accessibility for a broader audience.

Questions

1. How does the performance of MORL scale with the number of tasks and the size of the state and action spaces? 2. Are there any practical considerations or limitations in applying MORL to real-world datasets that were not covered in the paper?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations of their work, particularly in terms of the assumptions required for the theoretical analysis. However, more discussion on the potential practical limitations and how they can be mitigated would be beneficial.

Summary

This paper provides a theoretical analysis of representation learning in Multi-Task Offline RL. Specifically they consider a setting in which the transition kernels have a low-rank decomposition $P(s'|s,a)=\langle \phi(s'),\mu(s,a)\rangle$, such that $\mu(s,a)$ depends on the task but $\phi(s')$ is shared by all tasks. The authors propose an algorithm for this setting that first learns a transition model and then does planning on the model with a low-confidence-penalized reward. They provide a suboptimality bound, based on bounding the TV distance between estimated and true transition kernels. The authors also investigate the application of this acquired representation in down-stream online, offline and reward-free tasks.

Strengths

* The paper is very well written, assumptions that are made are explained and well justified. * The exploration of down-stream applications of Multi-Task Offline RL is of large practical interest.

Weaknesses

**I don't work in theory, please see the below more as minor comments.** * [89] analyzes a very similar setting but is only mentioned in the appendix as concurrent work. Strictly speaking, [89] was uploaded ot arxiv three months before the deadline, so it is not considered concurrent as per NeurIPS guidelines. Either way, it should probably be discussed more prominently and not just in the Appendix. * It would have been nice to specify which improvements in the bound are results of which decisions. For example, a conservative reward $r-b$ in planning vs the normal reward $r$, or how exactly using the low-rank assumption rather than a linearity assumption affects the analysis. Similarly it would be nice to provide an interpretation for the different terms in (3.7), but I'm also not sure if that is really possible to do in a meaningful way. * The paper treats a setting with non-stationary reward and transition functions. It's is not discussed why the paper treats this setting and what effect it has on the analysis. * The approach name and acronym MORL (Multitask Offline Representation Learning) is a somewhat confusing choice, as MORL is frequently used to denote Multi-Objective RL. * The notation is somewhat confusing at times with $t$ representing tasks, $h$ representing timesteps in the MDP, $\tau$ indexing transitions in the downstream offline dataset and $i$ indexing samples in the upstream dataset. The usage of non-stationary policies, rewards, and transition functions also does not help to reduce subscripts. * The checklist instruction part was meant to be removed before submission, i.e. the instructions in L1253,1254

Questions

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

Rating

7

Confidence

1

Soundness

3

Presentation

3

Contribution

3

Limitations

Assumptions are discussed sufficiently, so I have no concerns about additional limitations.

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?

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

The paper lacks a limitations section.

Official Comment by Reviewer fKn9

Thank you for your detailed rebuttal and for addressing the points I raised in my review. I appreciate your clarification on the scalability of the MORL algorithm, particularly that the sub-optimality gap depends on the feature dimension rather than the state and action space sizes. I have also reviewed the other feedback and acknowledge the broader concerns about the paper's clarity. Your plan to include a proof roadmap and visual aids in the revised version is a positive step that will improve readability. I believe your responses have effectively addressed my concerns, and while I will maintain my current score, I look forward to the improvements in the final version.