The Role of Representation Transfer in Multitask Imitation Learning

We thank reviewer Fs7F for their feedback. Thank you for emphasizing the signficance and rigor of our work.

W1: Thank you for pointing this out. Our statement here refers to the theoretical result that we derived. However, as suggested by you and reviewer Eicc, we aim to provide extra analysis on a new task with image-based observations shortly.

W2: Thank you for pointing out the clarity issue in our statement. Our sentence aims to describe that Arora et al. provided a statistical bound that holds in expectation---while this is an interesting bound, in practice we are given a particular target task. Practitioners may want to know the sample complexity of the particular transfer which is described by a high-probability bound.

W3: Thank you for the constructive feedback. We will include an extra ablation on larger number of demonstrations shortly during the discussion period.

Finally, is there any specific improvements that we can make to increase the score? Thank you reviewer Fs7F for your feedback again!

Thank you reviewer Fs7F for suggesting to run an extra ablation on BC. We have updated our manuscript to include this result on page 29, figures 11 and 12. Our experiment includes running BC on 10 target tasks per environment, each with 5 random seeds. We normalize the returns as each target task may have different performance range.

As expected, as we increase the amount of data to 4x and 8x, BC performs very closely to the expert with very low variance. While one may question the usage of multitask imitation learning, we note that the premise of pre-training using other source data is to reduce the amount of target data required. Recall that figures 3 and 4 on page 7 already demonstrates that generally we have achieved this.

We thank you again for the suggestion and we are happy to continue our discussions on this work.

We thank reviewer F18A for their feedback. Thank you for emphasizing the importance of our work.

W1: Assumptions and limitations of the current theoretical results: We thank you for emphasizing the interpreting consequences of the assumptions made. Due to space limitation we would like to use the space to highlight our theoretical and experimental findings, and deferred the exact assumptions and discussions on the limitations in the appendix. However, we have modified our main paper to include the high-level assumptions made and the intuitions. See paragraph before theorem 1 on page 4.

W2: Thank you for pointing out the clarity issue. We have updated the plot in the manuscript.

Finally, is there any specific improvements that we can make to increase the score? Thank you reviewer F18A for your feedback again!

We thank reviewer Eicc for their feedback. Thank you for emphasizing the potential significance of our work.

W1: Thank you for pointing this out. We agree that a task in extra environments will make the paper more impactful. As suggested by you and reviewer Fs7F, we aim to provide extra analysis on a new task with image-based observations shortly.

W2: Both Spearman and Kendall correlations aim to find monotonic correlations. In the main manuscript we briefly discussed that they are both positively correlated, albeit rather weak in some cases. Notice that frozen lake, cartpole, and pendulum tasks are the environments with weaker correlations. Our suspicion is that each variant of these environments are very sensitive to the policy's actions. In some cases, the policy will need to make consecutively accurate actions to perform well. For example, each variant of frozen lake differs in its dynamics and initial and goal positions. Consequently the learner policy in the target task may vary significantly in actions, when compared to the learned policies in the source tasks. This relates to the last sentence where in discrete action space, there is a possibility that we can simply permute the action space and have low task diversity even if the policy can learn perfectly well.

W3: Thank you for pointing out the clarity issue. We have updated the plot in the manuscript.

W4: We have fixed the typo in the manuscript. Thank you for pointing that out.

Q1: Thank you for the interesting question. Based on our theory, so long as the expert policy shares representation with the true representation, the pretraining will not stagnate or deteriorate the improvement. Consequently, we can likely use MMD or other divergences if the sample complexity improves and can be easily tied to the transfer risk, as suggested by reviewer StV9. We find that there has been work showing that some objectives work and some do not [1, 2, 3].

Q2: Yes, we believe that is the case based on the theory. Unless we have the degenerate case where both the representation and the task mappings are tabular, intuitively the learned representation should be capturing shared dynamics information relevant across all tasks. This is because the pretraining loss is induced by the dynamics and the policies of the source tasks. Consequently, this information should improve policy robustness to the covariate shift problem.

Q3: We have not yet considered inverse RL (IRL) methods, however we believe this work opens a new avenue to analyze various combinations of supervised IL and IRL approaches.

Finally, is there any specific improvements that we can make to increase the score? Thank you reviewer Eicc for your feedback again!

References:
[1] Yang, M., & Nachum, O. (2021, July). Representation matters: Offline pretraining for sequential decision making. In International Conference on Machine Learning (pp. 11784-11794). PMLR.
[2] Nachum, O., & Yang, M. (2021). Provable representation learning for imitation with contrastive fourier features. Advances in Neural Information Processing Systems, 34, 30100-30112.
[3] Kumar, A., Singh, A., Ebert, F., Nakamoto, M., Yang, Y., Finn, C., & Levine, S. (2022). Pre-training for robots: Offline rl enables learning new tasks from a handful of trials. arXiv preprint arXiv:2210.05178.

We thank reviewer StV9 for their feedback. We appreciate that you find this work to be potentially large.

W1: We thank the reviewer for their comment on the limitation of our task-diversity metric. We note that our task-diversity metric is based on the observations that (1) the source tasks likely capture the target task information if there is a source policy that performs well in the target task, and (2) the KL-divergence is asymmetrical which is a desirable property in task transfer. However, we agree that there can be other metrics, as you have mentioned in the later points. We aim to consider these metrics in future work.

W2: Thank you for pointing this out, we will include a comparison against Bhattacharyya distance shortly.

W3: Thank you for initiating this discussion. While we do believe this may be out of scope of our paper, we are happy to discuss this. The concern with discrepancy measure is that usually it relates only one source task to one target task. Perhaps what can happen in this case is to consider the source task distribution as the distribution of mixture policies in various MDPs.

W4: Perhaps we misunderstood the comment, could you please rephrase it? The (empirical) evaluations did not involve any complexity measure, rather we focused on minimizing the loss through KL-divergence.

Q1: Thank you for the question. This can be answered through the pendulum example. Suppose the pendulum can vary in its maximum torque. The lower-torque variant includes less policies that can keep the link upright while the higher-torque variant includes more policies. As a result, the lower-torque variant can be used in the higher-torque environment (arguably without finetuning) to keep the link upright, while this fails in the other case.

Q2: Thank you for the interesting question. Generally, we do not believe that Rademacher complexity is the best way to quantify neural network complexity. In fact many analyses consider Gaussian complexities due to its mathematical properties, as demonstrated in prior work outlined in our paper. Another approach for analyzing neural networks can be using algorithmic stability along with metric entropy [3,4]. It is also common to consider norm-based complexity which is used to upper bound Rademacher complexity [1,5].

Q3-Q5: They are all very great suggestions that we believe can be considered generally. We have not considered discrepancy measure but it might be interesting, as you have discussed in W3. Our goal is to directly quantify the representations learned from the source tasks and the target task. The difficulty may lie in including the true shared representation to the empirical representations from the source tasks and the target task, and finally connecting it to the imitation gap. We can alternatively consider the MMD of the policies, taking the expectation over their respective occupancy measures. Localized Rademacher complexity may give us $O(1/M)$ rate in some situations, as you have mentioned. But, it may still be difficult to improve upon $O(1/\sqrt{NT})$ generally, as we have discussed with reviewer Ss6u.

Finally, thank you again for your insightful feedback. Please let us know if we have properly addressed your questions and what we can do to encourage you to increase our score.

References:
[1] Liang, T., Poggio, T., Rakhlin, A., & Stokes, J. (2019, April). Fisher-rao metric, geometry, and complexity of neural networks. In The 22nd international conference on artificial intelligence and statistics (pp. 888-896). PMLR. [2] Bartlett, P. L., Montanari, A., & Rakhlin, A. (2021). Deep learning: a statistical viewpoint. Acta numerica, 30, 87-201.
[3] Bousquet, O., & Elisseeff, A. (2000). Algorithmic stability and generalization performance. Advances in Neural Information Processing Systems, 13.
[4] Li, Y., Ildiz, M. E., Papailiopoulos, D., & Oymak, S. (2023). Transformers as algorithms: Generalization and stability in in-context learning. [5] Neyshabur, B., Tomioka, R., & Srebro, N. (2015, June). Norm-based capacity control in neural networks. In Conference on learning theory (pp. 1376-1401). PMLR.

We thank reviewer StV9 for suggesting another metric. We have included a metric using the Bhattacharyya distance, replacing the KL-divergence. As a preliminary step we computed the metrics only for the discrete environments. We can see that Bhattacharyya distance may negatively correlate to the imitation performance in Frozen Lake and have weaker correlation in Walker compared to KL. We plan to investigate further in the future different task diversity metrics and how their theoretical properties may correlate to the imitation performance.

EDIT: We have included this table (table 6) on page 32, with explanation on pages 30 and 31. we further included the L2 representation norm as suggested by reviewer kWs2

Pearson

Spearman

Kendall

We thank reviewer StV9 for the continued discussions, we appreciate the amount of effort you have put in so far.

Q1: This is a good point, we have included a pendulum example in the updated draft, please see page 5, figure 1. Generally, the result agrees with the intuition we have provided in the previous response, where the expert policy from the lower-torque variant is more robust to the higher-torque variants, but not the other way around. We hope that this addition would provide better intuition on the asymmetry of policy transfer.

Q3-Q5: Thank you for pointing out the pitfalls of KL-divergence in the continuous setting. We agree that in practice the continuous setting is more difficult, and KL-divergence requires distribution realizability---often intractable to compute. Even in remark 3 (now moved to appendix A on page 15), our suggestion is to discretize the continuous space at the cost of the approximation error. We agree that this perhaps may be a limitation in practice, but we note that this limitation is very common in various areas of machine learning.

We once again thank reviewer StV9 for their insightful discussions. We hope they can reconsider our score after our adjustments.

We thank reviewer StV9 for their feedback. We hope our response addresses your questions and concerns.

W1: Thank you for your question. Our statement on environment asymmetry is a motivation for why we cannot simply transfer one policy from one environment to another. Indeed, there can be cases where you can simply transfer but this is not the case generally, another concrete example can be found under section 3.1, example 3 of Hanneke & Kpotufe, 2020. Consequently, we want to ensure our metric to be able to capture asymmetry when we swap order of the environments.

We also thank you for pointing the paragraph out, we notice that the reference was incorrect and is now corrected to be Hanneke & Kpotufe, 2020 and have pointed the readers to the example for more intuition.

W3: Thank you for raising this question. This depends on how MMD will be used. Suppose we compare the MMD of the representations, then it is not immediately clear because we will need to consider what is the distribution of the learned representation. On the other hand, when computing MMD of the source policies and target policy, then we agree that we can try using MMD. We note that, however, one consideration to make is what kernel to use when we compute the empirical estimation of MMD.

W4: Thank you for indicating the possible confusion. The intuition comes from definition 4 on page 14, where $\sigma$-diversity is defined to be upper bounded by the ratio of the worst-case representation difference (definition 3) and the task-average representation difference (definition 2). Both representation differences in this case are actually KL-divergences due to the fact that we are doing behavioural cloning. Consequently, we approximate this ratio via our approximate KL metric. More details can be found in appendix B on page 15.

We once again thank reviewer StV9 for their thoughtful discussions!

We thank reviewer Fs7F for their insightful feedback. We believe we have addressed your concerns and we are very thrilled to continue the thoughtful discussion with you.

W1: We thank reviewer Fs7F for noting the simplicity of our theoretical proofs. We agree that our theoretical result is a synthesis of known results and is straightforward. Our novelty is not a new proving technique. Instead, our novelty lies in bridging and connecting learning theories in multitask learning and imitation learning. As noted by many other reviewers (StV9, Eicc, F18A, Fs7F, Ss6u), this result does not appear to be commonly known and can have significant impact in the community. Furthermore, we note that our result is an improvement of existing results by leveraging the properties of the log loss. With any policy parameterization that is Lipschitz, we can improve the sample complexity by $\log(NT)$ when compared to using Gaussian complexity. We finally argue that a simpler proof is often preferred since it can be easily understandable, which appears to be the case as stated by many reviewers (StV9, kWs2, Eicc, F18A).

W2: We thank reviewer Fs7F for raising an undefined term in the paper. Based on your feedback, we have modified our terminology to imitation gap.

Regarding the horizon term, indeed we have neglected the horizon term. Due to Xu et al, 2020. this is an unavoidable term and we thus treat it as a constant. We further note that we also neglected other constants (e.g. $\lvert \mathcal{A} \rvert$) for presentation purposes.

Finally, thank you for raising this concern of the term "demonstration". We emphasize that we treat each demonstration as a state-action pair (i.e. a transition) and we have explicitly defined this in section 2.2. We wish to keep the term demonstration as transitions can be referred to any policies, and terms such as "expert transitions" are cumbersome from the reading perspective. However, we are open to suggestions on what can be a better term.

W3: Thank you for pointing out that the number of trajectories is small for MuJoCo tasks in many papers [1,2]. We believe that this "contrived setting" is preferrable in practice---practitioners aim to minimize the number of demonstrations as much as possible since it can be costly (e.g. human-operating time, computational cost, etc.)

W4: We acknowledge reviewer kWs2's comment on the task-diversity metric. We note that this is an attempt to defining a task-diversity metric that may positively correlate with the imitation performance. Indeed, we do not provide any theoretical guarantees on the metric and the conditions in appendix B are very restrictive. Intuitively we hope to have the irreducible error to be sufficiently small such that the ratio can be estimated by our proposed metric. We consequently experimentally showed that this metric can be strongly correlated to the imitation performance. Furthermore, the metric still captures the asymmetry of task transfer which is a desirable property.

W5: Thank you for raising this. Our result does not require optimality of the expert policy as the learner policy only aims to imitate the expert. We have removed footnote 1 which hopefully resolves the confusion.

W6: Thank you for providing extra references for multi-task IRL methods. We have included them in the first paragraph of the related work. We hope this better covers the literature and demonstrates the differences in our work.

We thank reviewer kWs2 for the feedback. We have updated our manuscript to reflect the follow-up suggestions.

W2:

• We have included the suggested reference on page 2, under section 2.2.
• We have also included the horizon term in page 1, under Theoretical contributions.

Q1:

• We have added a remark (Remark 1) on page 4 to include the sample-complexity of BC, as well as comparing BC with MTIL on page 4.

Q2 & Q4:

• We have provided a sentence in the caption of figures 13-15 on pages 30 and 31, as well as including paragraphs under Inclusion of Target Data during Training Phase and Impact of Source Tasks and Source Data on pages 28 and 29.

Q5:

• We have included the parameter ranges and described how we sample them on pages 24 and 25.

Finally, we thank reviewer kWs2 for their constructive feedback. We invite reviewer kWs2 to comment on the other points as the feedback has been very helpful. We also would love reviewer kWs2 to once again reconsider their score.

We thank reviewer Ss6u for their relevant feedback. We hope our response encourages further discussions.

W1: We acknowledge reviewer Ss6u's concern on the structure of the shared representation. The benefit of our theoretical result is that the shared representation has a minimal structure (e.g. in $\mathbb{R}^d$ and bounded). The sample complexity bound holds with high probability. Indeed, we believe that imposing a particular structure on the representation will yield a potentially better result and will further enable practitioners to understand the problem. The state-estimation problems in robotics are great examples for imposing structures on the representation. For example, consider a image-based pendulum task where we have no true readings of the joint information. If we are given the joint information we can solve it using linear controllers. Thus our representation can simply be enforced to be in the same space.

W2: We thank reviewer Ss6u for highlighting the importance of lower bounds---we agree that a lower bound will provide significant insight to our result, specifically to understand whether our result is minimax optimal. We first emphasize that this result means that it is hopeful to obtain a near-expert policy as we increase the number of samples, at the rate of $O(1/\sqrt{NT} + 1/\sqrt{M})$. Regarding the lower bound, there are existing results on general Lipschitz and smooth loss functions (see table 1 of [1]). In the single-task BC setting, indeed the $\Omega(1/\sqrt{M})$ is tight and can perhaps be tightened to $\Omega(1/M)$ with specific function classes (i.e. $L^*$ shrinks to 0 thus removing the $\Omega(1/\sqrt{M})$ dependence). While this bound is not established on the multi-task setting, Our intuition is that following similar proof techniques we will arrive at similar $\Omega(1/\sqrt{NT} + 1/\sqrt{M})$ result. We thank reviewer Ss6u again for the very helpful comment and we plan to clearly address this in our future work due to time limitation of the rebuttal period.

Q1: We acknowledge reviewer Ss6u's comment on the potential vagueness of the shared representation's structure. Generally, the representation can be a transformation from a visual image to a latent space, texts to a latent space, or a stacked sequence of past robot information into a latent space---our theoretical result imposes minimal latent space structure (other than $\mathbb{R}^d$ for example). An intuitive example is a visual pendulum task where the input space is a stack of images, which allows the representation to capture positional, velocity, acceleration, and higher-order derivatives.

Q2: We thank reviewer Ss6u for pointing out the lack of constants, conditions and assumptions made in the main paper. We provide the technical details and exact theorem result with constants in appendix C with a high-level explanation on each assumptions made.

Regarding practicality, we believe we can address this through whether the assumptions are reasonable and whether the bound can be used in practice. The former is provided in appendix C.1. For the latter, we believe a concern is whether certain assumptions hold---if not, the bound is hopeless and is uninformative. However, albeit tedious, we can still validate these assumptions by increasing the complexity of the models. Suppose we make the correct assumptions, then there are multiple facets that are practical. First, we can expect the policy performance to shrink at a square-root rate, with respect to the source and target data, meaning that it is hopeful to obtain near-expert performance. Second, we know exactly how much target data to collect in order for the imitation gap to shrink below the desired value as it is independent of the task-diversity constant. Third, the fact that the representation is more complex, we are free to gather more source, perhaps even include the target data in the pretraining as suggested by reviewer kWs2. The benefit here is that it is possible for some task data to be cheaper to collect, thus we can focus on exploiting the cheaper tasks.