Paper: The Role of Representation Transfer in Multitask Imitation Learning (ICLR-2024)

Reviewer comment (Writing):
In the theoretical contributions paragraph on page one, what do you mean by the second sentence? Is that from the Arora et al. paper? If so, please say that it refers to that because it sounds off.

Author response:
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.

Reviewer comment (Experiments):
For the BC baseline, it seems like an easy one to compare and in the SM you have BC with 2|D|. Would it not be appropriate to also show that similar comparison with the main text experiments to show what if BC had more |D| to what your method has?

Author response:
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.

Reviewer comment (Writing):
It would be better to state clearly in the main paper about what assumptions are made in the paper and discuss about the limitations.

Author response:
**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.

Reviewer comment (Presentation):
It is difficult to read Figure 1-4. The lines are hard to distinguish from one another.

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

Reviewer comment (Experiments):
While the story told in Figures 1-4 are clear, it would have been nice to see some evaluations on at least one multitask environment such as FrankaKitchen or Metaworld. While the current results on Mujoco support the claim, the environments are relatively simple. The results of this paper would be of interest to a much larger community if the same compelling results are shown on just one of the above environments.

Author response:
**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.

Reviewer comment (Evaluation):
It would have been nice to see some more interpretation on the results of Spearman and Kendall correlations on the bottom of page 8.

Author response:
**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.

Reviewer comment (Presentation):
Due to the boldness of the lines and overlap, it is a bit challenging to tell the differences between the blue, yellow, green, and red lines. I might suggest using different shapes or decreasing the boldness of the lines.

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

Reviewer comment (Writing):
On (ii) towards the bottom half of page 7, "let" should be capitalized in "let \hat{r}_t be the average rewar of the expert...".

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

Reviewer comment (Theory):
Is there any demand on the optimality of the source task demonstrations? For example, would a large batch of suboptimal demonstrations actually deteriorate or stagnate improvement given the proposed method?

Author response:
**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].

Reviewer comment (Experiments):
I'm curious whether increasing the amount of source task data would also make the policy more robust to covariate shift, since it theoretically should have a larger state-space coverage. Or when transition dynamics are stochastic.

Author response:
**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.

Reviewer comment (Experiments):
Have the authors attempted to generalize their framework to imitation learning algorithms beyond behavioral cloning, such as IRL methods?

Author response:
**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.

Reviewer comment (Novelty):
Overall, I found the theoretical statements to be fairly simple extensions of known results by Tripuraneni et al. and Ross & Bagnell. In essence, by focusing only on behavioral cloning, the authors are able to almost entirely ignore the sequential nature of the imitation problem and apply the standard analysis for multi-task supervised learning. Then, once they have a bound on the KL divergence between the learner and the expert, they can apply the well-known upper bounds for behavioral cloning (https://www.cs.cmu.edu/~sross1/publications/Ross-AIStats11-NoRegret.pdf) to get an overall policy performance guarantee. So, I really didn't get much out of the theorems they proved.

Author response:
**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).

Reviewer comment (Theory):
There's a few pieces of odd terminology throughout the paper. First, instead of "policy error", people usually use "performance difference" or "imitation gap" (https://arxiv.org/abs/2103.03236). Also, I think you're missing a $H^2$ or $\frac{1}{(1-\gamma)^2}$ in the first equation in the paper? Second, when people talk about "number of demonstrations" (i.e. $|\mathcal{D}|$ in the paper), they usually mean the number of whole trajectories rather than the number of state-action tuples (as you seem to use it in Table 1). Do you mind re-naming this? I got super confused for a while by why one would need millions of samples for BC to work on a Mujoco task. Can you also clarify whether $N$ and $M$ are measured in terms of trajectories or state-action pairs?

Author response:
**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.

Reviewer comment (Experiments):
Once you divide the numbers in table 1 by the horizon of the problem (1000 for Mujoco), you realize that they're attempting to learn based on effectively 1-2 demonstrations. This is a somewhat absurdly small amount of data (usually people do ~25 demos for Mujoco tasks). So, while the experimental results make sense to me, I think it is somewhat important to note that they are under fairly contrived settings.

Author response:
**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.)

Reviewer comment (Evaluation):
While the idea of a metric to capture the effectiveness of multi-task IL data for learning a transferrable representation is interesting, I found the ideas in Appendix B to be a bit sloppy and the empirical reported correlations in Tables 2/3 to be fairly low.

Author response:
**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.

Reviewer comment (Theory):
Most analysis of imitation learning doesn't have to make assumptions about the optimality of the expert policy. In Footnote 1, you note that you do this. Is this actually important for any of your analysis?

Author response:
**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.

Reviewer comment (Novelty):
I might add some more citations for multi-task imitation learning outside of behavioral cloning. While they are clearly different than your work, it would be good to add in some references and discuss the differences: https://arxiv.org/pdf/1805.12573.pdf, https://arxiv.org/pdf/1909.09314.pdf, https://arxiv.org/pdf/1805.08882.pdf, https://arxiv.org/pdf/2309.00711.pdf. You might also want to cite some work on representation learning for sequential decision making (e.g. https://arxiv.org/abs/2207.08229).

Author response:
**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.

Reviewer comment (Theory):
I think it would be helpful if you could add in a comparable statement to Theorem 1 for single-task BC. You could then give ranges of T and N under which you'd have a meaningful difference in upper bounds between multi-task and single-task BC. The sharpest analysis I know for IL under a deterministic expert assumption is in https://arxiv.org/pdf/2205.15397.pdf -- you might be able to just copy some of their theorems.

Author response:
**Q1:** Thank you for the suggestion. We believe your provided reference paper, Xu et al., 2020, and Ross and Bagnell, 2011 are all great references for indicating the BC sample complexity. In fact, our proof of Theorem 3 also works with minimal changes.
Consider the ERM $\min_{f \in \mathcal{F}, \phi \in \Phi} R_{test}(f, \phi)$ and follow similar proof techniques, then we have that w.p. $1 - \delta$,
$$
R_{test}(f, \phi) - R_{test}(f^*, \phi^*) \leq 4 \mathfrak{R}_M (\ell \circ \mathcal{F} \circ \Phi) + 2B \sqrt{\frac{\log (2/\delta)}{2M}}.
$$

Perhaps through this bound we can easily see that the first term will dominate RHS and will require larger $M$ when we have an expressive model such as a neural network.

Reviewer comment (Evaluation):
Do you have any hypothesis for why, in Figure 2, things look quite bad for CartPole? Is it perhaps because of the kinds of environment modifications you were considering?

Author response:
**Q3:** We find cartpole as a very sensitive environment generally. If you consider figure 10 in the appendix on page 27, you can see that BC performance does not increase after very small number of demonstrations on the default environment. We suspect that some of the source tasks include these difficult variants, thereby impeding the performance of MTBC.

Reviewer comment (Experiments):
In Figures 3 and 4, you're giving the learner quite a bit of source data so perhaps the performance gap is already quite close. If you have the compute resources, could you ablate these results across smaller values of $N$?

Author response:
**Q4:** Thank you for the suggestion, we are currently gathering the results for this experiment and will provide them along with other extra experimental results.

Reviewer comment (Experiments):
Generally, could you be more specific about the ranges over which you varied environment parameters to generate the multi-task data (e.g. what link sizes for Walker)?

Author response:
**Q5:** Yes, we are happy to share this:

Frozen lake:
1. We vary the initial and goal positions uniformly (disallowing equal positions)
1. We sample uniformly the probability of not slipping $p_1 \sim U[0, 1]$, then sample uniformly the probability of slipping to one adjacent cell $p_2 \sim U[0, 1 - p_1]$, and assign the remaining mass $1 - p_1 - p_2$. This probability is assigned to all state-action pairs.

Cartpole:
1. Joint stiffness: $U[0, 5]$

Pendulum:
1. Pendulum max torque: $U[0.01, 4]$

Cheetah:
1. Front and back thigh joints:
    1. Damping: $U[4, 8]$
    1. Stiffness: $U[230, 250]$
1. Front and back feet:
    1. Link size: $U[0.04, 0.05]$ by $U[0.065, 0.075]$

Walker:
1. Left and right feet:
    1. Link size: $U[0.04, 0.06$ by $U[0.09, 0.11]$
1. Left and right legs:
    1. Link size: $U[0.03, 0.05$ by $U[0.23, 0.27]$
1. Left and right thighs:
    1. Link size: $U[0.04, 0.06$ by $U[0.2, 0.25]$
1. Left and right knee joint:
    1. Range: $U[-140, -160]$

We hope you can reconsider our score after our response. Thank you reviewer kWs2 for your feedback again!

References:  
[1] Ho, J., & Ermon, S. (2016). Generative adversarial imitation learning. Advances in neural information processing systems, 29.  
[2] Kostrikov, I., Agrawal, K. K., Dwibedi, D., Levine, S., & Tompson, J. (2018, September). Discriminator-Actor-Critic: Addressing Sample Inefficiency and Reward Bias in Adversarial Imitation Learning. In International Conference on Learning Representations.

Reviewer comment (Writing):
Lack of Clarity: The paper does not sufficiently describe the "shared representation" it aims to learn. A more detailed exposition, possibly including visualizations or analysis of the learned representation, is needed.

Author response:
**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.

Reviewer comment (Theory):
Theoretical Bound Practicality: The paper presents an order bound on policy error but does not provide a comprehensive discussion on its tightness or practical applicability, leaving its usefulness in question.

Author response:
**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.

Reviewer comment (Writing):
Section 4 requires further detail on the nature of the source tasks for each target task investigated. The concept of "shared representation" is pivotal yet remains vague within the paper. Is this representation a transformation from a visual image to a latent space, or something else? A detailed analysis or visual depiction of this shared representation would greatly enhance the clarity, perhaps focusing on a single task as an example.

Author response:
**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.

Reviewer comment (Theory):
The paper introduces a theoretical bound but does not elucidate on its tightness or practical applicability. It is essential to quantify or provide conditions under which the bound holds with a fixed constant, thereby ensuring utility in policy improvement with additional data.

Author response:
**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.

Reviewer comment (Theory):
How do you define the “policy error” on Page 1? How can it be measured? Is the measurement done in the task reward, action space, or divergence perspective? Besides, you mentioned that f-divergence imitation learning states that a learned policy is minimizing the divergence between expert and learner trajectory. How does the policy error fit / contrast with such existing framework?

Author response:
**Q3:** We thank reviewer Ss6u for raising an undefined term in the paper. Based on reviewer kWs2's feedback, we have modified our terminology to imitation gap. The imitation gap is measured through the max-norm of the difference between the value functions of the learner and expert policy.

Reviewer comment (Novelty):
“However, current methods require thousands of demonstrations even in simple tasks (Mandlekar et al., 2022; Jang et al., 2021; Ablett et al., 2023)” => *Thousands* of demonstrations seems to misrepresent the SOTA of IL? Refer to https://www.roboticsproceedings.org/rss19/p009.pdf  https://medium.com/toyotaresearch/tris-robots-learn-new-skills-in-an-afternoon-here-s-how-2c30b1a8c573 https://deepmind.google/discover/blog/robocat-a-self-improving-robotic-agent/  Please clarify.

Author response:
**Q5:** We appreciate reviewer Ss6u for pointing out the potential misrepresentation of the sample complexity. We note that in our paper a demonstration corresponds to a state-action pair, known as the transition. On the other hand, the references use trajectories, a sequence of state-action pairs, as the definition of a demonstration. When taking into account the horizon/length of the trajectories, we will get similar number of transitions. For conciseness, we wish to keep the term "demonstration" as we have defined this explicitly in section 2.2.

Reviewer comment (Reproducibility):
For all experiments detailed in Section 4, please define the state and action space, including their dimensions.

Author response:
**Q6:** We have provided the environment description in appendix F.1. We will also describe the environments here for completeness:

*Discrete environments*
|                             | Frozen Lake | Pendulum | Cheetah | Walker |
|-----------------------------|-------------|----------|---------|--------|
| dim($\mathcal{S}$)          | 64          | 3        | 17      | 24     |
| $\lvert \mathcal{A} \rvert$ | 5           | 11       | 64      | 64     |

*Continuous environments*
|                    | Cartpole | Pendulum | Cheetah | Walker |
|--------------------|----------|----------|---------|--------|
| dim($\mathcal{S})$ | 5        | 3        | 17      | 24     |
| dim($\mathcal{A})$ | 1        | 1        | 6       | 6      |

Reviewer comment (Theory):
Can the derived bound be applied to low-dimensional imitation learning that does not learn a representation? Assuming that the learned representation is the identity matrix, can we extend the bound to low-dimensional state space? Is the result implying that, adding more data into pre-training, will result in smaller policy error? However, given what we saw in https://ieeexplore.ieee.org/abstract/document/10161474/ it seems that more data is not guaranteed to help imitation learning performance on the real robot. Can the author provide some insight?

Author response:
**Q7:** We thank reviewer Ss6u for bringing up the surprising result in real-life robotics experiments. After reading the referenced paper, we would like reviewer Ss6u to clarify the statement, "more data is not guaranteed to help IL performance on the real robot." In particular, section III-d of the paper states that:
```
Dataset: Our Real-ORL dataset consists of around 3000 trajectories and the characteristics of our dataset is shown in Table. II. To validate the quality of our datasets, we use the Top-K% portion of dataset and trained behavior cloning for each task (K ∈ {100, 90}), observing that using the full datasets allowed behavior cloning to achieve the best performance on all tasks. To avoid biases favoring task/algorithm, we froze the dataset ahead of time.
```
which seems to agree with our sample complexity bound.

Regarding the derived bound being applied to low-dimensional IL, our theory suggests that the learner will pay a constant factor of the state-space dimension, which is captured within the Rademacher complexity.

Reviewer comment (Theory):
Our result is due to the objective of behavioral cloning, where the method aims to minimize the Kullback–Leibler (KL) divergence between the expert and the learner (Ghasemipour et al., 2019; Xu et al., 2020). => Could you elaborate how the findings from Xu et al., 2020 is related to this statement?

Author response:
**Q8:** The result (Theorem 1) from Xu et al., 2020 indicates that the policy error (imitation gap) can be upper bounded by the KL-divergence between policy $\pi$ and expert $\pi^E$. In standard behavioural cloning, we often aim to maximize the log-likelihood of the policy based on data generated by the expert (i.e. $\max_\pi \mathbb{E}_{s \sim \rho^*} \log \pi(a \vert s)$). Consequently, this objective is equivalent to the log loss and corresponds to the KL-divergence when the expert is deterministic.

Reviewer comment (Theory):
Regarding Equation 4 and the "under some assumptions" qualifier, a more intuitive explanation of these assumptions and their impact on the model's generalizability would be essential.

Author response:
**Q9:** Thank you for raising the challenge in understanding Eq. 4 and the assumptions. As mentioned in the response of Q2, we provide the exact theorem result with constants in appendix C with a high-level explanation on each assumptions made. We have further made modifications to our draft to explain the high-level assumptions made and the intuition. See paragraph before theorem 1 on page 4.

Reviewer comment (Evaluation):
The paper uses only KL-divergence to measure task-diversity for source and target tasks.

Author response:
**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.

Reviewer comment (Theory):
In the context of bounding sample-complexity in transfer learning, particularly for evaluating the richness of representation classes, Rademacher Complexity may not be the most appropriate choice because it avoids accounting for the nuances in transfer learning on sequential tasks.

Author response:
**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.

Reviewer comment (Theory):
Utility of Rademacher Complexity: Measures the ability of a function class to fit random noise, providing a general sense of the capacity of the function class.

Author response:
**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.

Reviewer comment (Evaluation):
Given the specific requirements of transfer learning, which often involve understanding the relationship and distributional differences between source and target tasks, measures like Maximum Mean Discrepancy (MMD), Discrepancy Distance, and Task Similarity Measures may be more pertinent because they are more directly aligned with the challenges of assessing how well a learned representation from one task can be applied to another.

Author response:
**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.

Reviewer comment (Theory):
Gaussian complexity is somewhat similar to Rademacher Complexity but might be more suitable in certain contexts, especially those in which approximately Gaussian assumptions naturally occur. Therefore, the significance of the finding with Rademacher Complexity in practice is questioned.

Author response:
**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].