Offline Multitask Representation Learning for Reinforcement Learning

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

Thank you so much for your suggestion on how to improve the readability for a broader audience. We will definitely add more intuition for each of the different equations that are presented in the main paper.

Main take-away for Section 3

The main take-away for the main result (Theorem 3.3) in Section 3 is highlighted in L187-193. Namely, in Theorem 3.3, Equation (3.6) provides the statistical scaling rate for the average accuracy of the estimated transition kernels of the $T$ MDPs. Equation (3.7) implies that if the optimal policy $\pi_t^*$ is covered by the offline data for all $t \in$ {1, 2, ..., T} as characterized by multitask relative condition number $C^*$, then the output policy of MORL is able to compete against it on average as well.

Main take-away for Section 4

On high-level the main take-away for section 4 is that representation learning in upstream multitask offline dataset can help in improving the sample complexity, statistical rate of sub-optimality gap and regret bound in reward-free, offline and online RL respectively.

The main take-away for Theorem 4.4 is provided in Remark 4.5. Due to space constraint in the main paper, the main take-away for Theorem 4.7 is provided in Remark K.7 (L1042-L1046) in Appendix K. Similarly, the main take-away for Theorem 4.8 is provided in Remark M.6 (L1140-L1145) in Appendix M.

Other comments

Thanks for suggesting to put the purpose of introducing Equation 3.4 earlier. We will do it in our next version. Essentially, Equation 3.4 is used for inducing pessimism in the learned policy in Algorithm 1. The notation $[N]$ indicates the set {1, 2, ..., N}. Even though it’s a standard notation, we will clarify it in our next version. Thanks for the other minor comments and pointing out some typos. We will correct and update them.

Questions

1. By assuming $\mu$ is not shared across tasks, we ensure that the transition kernel of each task is different from one another. However, for theoretical analysis it is imperative to make an assumption that will connect all the MDP tasks in some manner. This is why we assume that the feature $\phi$ is shared across tasks.

2. Here $\lambda$ is a regularization factor in Equation 3.3.

3. Thanks for your suggestion regarding equations as Equation X instead of as (X). We will update this in our next version as per your suggestion.

Limitations

We have already discussed the assumptions made at length in L230-L241. We will add further discussion on how these assumptions can be limiting.

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

• As discussed in Additional Related Work in Appendix A under Offline Data Sharing in RL, there have been numerous empirical works that studied the benefit of using offline datasets from multiple tasks to accelerate downstream learning. The aim of this paper is to provide a framework for investigating this approach from a theoretical lens. While empirical validation is limited in this paper, we hope that the theoretical analysis presented in this work will provide insights for algorithmic design principles that can be used to design empirically better performing algorithms compared to current state-of-the-art algorithms.

• Thanks a lot for your suggestion regarding improving the accessibility of the proofs for a broader audience. We will add a proof roadmap section along with a visual proof roadmap in our next version. Again, thanks a lot for the great suggestion!

Questions

1. As can be observed from Equation 3.7, the sub-optimality gap (in average sense) of MORL does not depend on the size of the state and action space. Instead it depends on $d$ - the dimension of the feature $\phi$. Moreover, the average accuracy of the estimated transition kernels of the $T$ MDPs scales at the rate $\tilde{O}(\sqrt{\frac{\log |\Phi|}{nT} + \frac{\log|\Psi|}{n}})$.

2. In our setting we assumed that the different tasks share the feature representation $\phi$ under low-rank MDP structure. However, in most real-world datasets this might be difficult to satisfy exactly. It would be an interesting future work to study how we can extend this work to the case where the features are not exactly shared but instead shared with some perturbed approximation.

Finally, we are glad that the reviewer felt that we have adequately addressed the limitations of the work in terms of the assumptions used for theoretical analysis. As per the suggestion of the reviewer, we will add more discussion on potential practical limitations of MORL along with how to possibly mitigate them in our next version.

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

• Thanks for your suggestion regarding [89]. A preprint version of our results appeared publicly around at the same time of the preprint of [89]. We will provide a more detailed discussion of it and move the discussion from the Appendix to the main paper.

• A conservative reward $r-b$ allows us to show near-pessimism in the average sense as depicted in Lemma 3.5. This is possible through a concentration argument for the penalty term $\hat{b}_h^{(t)}$ as defined in Equation (3.4). Under linearity assumption as in linear MDP, the representation $\phi^*$ is assumed to be known a priori. This allows one to use linear regression analysis to derive point-wise model uncertainty quantification and that can serve as a conservative penalty. However, this is not the case under low-rank assumption. It would be interesting to see how we can interpret each term in Equation 3.7. But as the reviewer mentioned, it is unclear how to do it in a meaningful way. Our motivation for representing Equation 3.7 in this way was to provide a fine-grained view for each term.

• We considered non-stationary reward and transition functions in our work as it is a standard practice in literature that assumes episodic MDP with finite horizon that have either low-rank MDP or linear MDP structure. Without non-stationarity, the dependency on $H$ would change in the resulting bounds. Indeed, with stationary reward and transition function one can achieve a better regret bound in terms of $H$ dependency.

• Thanks for pointing out potential confusion with the acronym MORL. We will try to come up with an alternative acronym in the next version.

• We totally understand that it can be challenging to follow the notations sometimes as there are so many moving parts involved that add up in the form of subscript and superscript. We tried our best to keep things concise and clear as much as possible. We would be happy to incorporate any suggestion that the reviewer might have to make the notations easier to follow.

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?

We thank the reviewer for their valuable time and effort in providing detailed feedback on our work. We hope our response will fully address all of the reviewer's points.

Points raised as weaknesses

1. To our knowledge, there is only one other concurrent offline multitask RL theory work [89]. We provided detailed comparison with that work in Appendix A: Additional Related Work. However, as described in our discussion the theoretical results of [89] are not directly comparable to our theoretical results. Table 1 only compares with single-task methods as there is no existing work that studies reward-free RL in multi-task settings. The goal of Table 1 is to show that performing offline multitask representation learning in the upstream task, can help improve sample complexity of downstream reward-free tasks compared to single task counterparts. For offline and online downstream tasks we provided comparison with other algorithms in Remark K.7 (Appendix K.2) and Remark M.6 (Appendix M.2). We provided them in the appendix due to space constraints in the main paper.

2. As we mentioned in the Introduction section where we highlighted our contributions, for the downstream part, our core contribution is in the reward-free setting. We provided results for offline and online downstream settings as complementary results.

Questions

1. In Section 4.3, for downstream offline and online RL tasks, we assume the reward function $r^{T+1}$ to be unknown. We made this assumption to make our result more general. Note that, knowing the reward function $r^{T+1}$ would have made the downstream task easier to solve.

2. As discussed in Additional Related Work in Appendix A under Offline Data Sharing in RL, there have been numerous empirical works that studied the benefit of using offline datasets from multiple tasks to accelerate downstream learning. The aim of this paper is to provide a framework for investigating this approach from a theoretical lens. While empirical validation is limited in this paper, we hope that the theoretical analysis presented in this work will provide insights for algorithmic design principles that can be used to design empirically better performing algorithms compared to current state-of-the-art algorithms.

Limitations

As all other reviewers have acknowledged, we have addressed limitations of our works, especially through a thorough discussion of the assumptions made for the theoretical analysis. Moreover, as per the suggestion of the Reviewer fKn9, we will add more discussion on potential practical limitations of MORL along with how to possibly mitigate them in our next version.

We hope we have addressed all of your questions/concerns. If you have any further questions, we would be more than happy to answer them and if you don’t, would you kindly consider increasing your score?