Skill-aware Mutual Information Optimisation for Zero-shot Generalisation in Reinforcement Learning

We are very grateful for the reviewer's professional review work on our paper. We would like to express our sincere appreciation for the reviewer's recognition of our work, especially regarding the writing and the motivation behind our work. We also greatly appreciate the acknowledgement of our theoretical and empirical evidence. Please let us know if further clarification is required.

Re Weakness (1) and (2): Regarding the reviewer's question about whether SaNCE can be combined with any meta-RL algorithm (e.g., $\text{RL}^2$ and MAML), SaNCE is a simple objective based on mutual information that can be used to train any context encoder as defined in Section 4.2 “Skill-aware noise contrastive estimation: a tighter $K$-sample estimator” in our paper and enable contextual RL agents to autonomously discover skills. We have now specified the types of algorithms where SaNCE can be used to make the scope clearer in the new manuscript, specifically in Section 5.1 “Experimental Setup” under the “Our Methods” paragraph.

Re Weakness (3): Please note that the modified MuJoCo environment is not entirely new. We extended the modified MuJoCo benchmark introduced in DOMINO (https://arxiv.org/abs/2210.04209) and CaDM (https://arxiv.org/abs/2005.06800). Hence, in our extension, there are four new-added environments (Walker, Crippled Hopper, Crippled Walker, HumanoidStandup) compared with the original benchmark. Additionally, in our experiments, we used a different task set design than in the DOMINO and CaDM papers. Specifically, we evaluated RL algorithms in out-of-distribution tasks which have more extreme environmental features.

For benchmarking, we referred to the DOMINO paper and chose PEARL as a baseline. Our baseline CCM (https://arxiv.org/pdf/2009.13891) is also a zero-shot meta-RL algorithm, and in the CCM paper, CaDM and PEARL were used as baselines for comparison. TESAC, mentioned in the paper (https://arxiv.org/pdf/1910.10897) as a multi-task soft actor-critic with an off-policy version of task embeddings, was used for comparison with the reviewer's recommended $\text{RL}^2$ and MAML. Therefore, we believe our choice of benchmarking aligns with the reviewer's expectations.

Additionally, we directly compared our zero-shot meta-RL algorithms with DOMINO and CaDM using exactly the same task set design. We present the comparison results in Appendix G "A Comparison with DOMINO and CaDM." Since we cannot replicate their pre-trained encoder and prediction networks, we have only compared our results with the experimental results from their original papers. The results show that our model-free methods SaCCM and SaTESAC significantly outperform the model-free DOMINO and are slightly inferior to the model-based CaDM.

Re Weakness (4): In the updated manuscript, we have added explanations on how the proposed method works during meta-training and meta-testing times. First, SaNCE is an objective for contrastive learning that can be used to train context encoders. Therefore, we have added details in Section 3, "Preliminaries," under the "Contrastive learning" subsection, on how the context encoder is trained and works during the meta-training and meta-testing phases: "In a Meta-RL setting, the context encoder $\psi(c|\tau_{0:t})$ first takes the trajectory $\tau_{0:t}= \{s_{0}, a_{0}, r_{0},...,s_{t}\}$ from the current episode as input and compresses it into a context embedding $c$. Then, the policy $\pi$, conditioned on context embedding $c$, consumes the current state $s_t$, outputs the action $a_t$. The policy $\pi$ conditioned on a fixed $c$ alters the state of the environment in a consistent way, thereby exhibiting a mode of skill." Furthermore, in Figure 5, we updated the caption and gave a more illustrative description of how the context encoder works during meta-training and meta-testing times. Besides, to help readers understand the learning procedure, we provide pseudo-code for how to use SaNCE during meta-training and meta-testing in Algorithms 1 and 2 in Appendix D.3 (you can find them in the pdf file we submitted in the global rebuttal).

Re Q(1) how many episodes does the agent interact with the environment: We have detailed this in Appendix D.3 "Implementation details" under the “Base algorithm” section in the original manuscript. In the Meta-training phase, we trained agents for 1.6 million timesteps in each environment on the Panda-gym and MuJoCo benchmarks. The maximum episode length is listed in Appendix Table 3 (can be found in the pdf file we submitted in the global rebuttal). For meta-testing, we tested 100 episodes in each environment, with tasks randomly sampled from the moderate and extreme task sets. Considering this might be of interest to the readers, we have added this detail in the main paper's Section 5.1 “Experimental Setup” under the “Our Method” paragraph.

Re Q(2) How does the encoder accumulate across multiple episodes of interaction?: Our proposed method needs to adapt within an episode. Therefore, at the beginning of each episode, the encoder is reinitialised and does not accumulate across multiple episodes. Specifically, we use an LSTM as the encoder, and we initialise the hidden state and cell state to zero at the start of each episode. We have added the aforementioned content in the “Encoder Architecture” paragraph of Appendix D.3 "Implementation Details" in the updated manuscript.

Re Q(3) Is there a way to use the proposed method for adaptation within an episode?: Our proposed method adapts within an episode. We have added descriptions in Section 3, “Preliminaries,” under the “Contrastive learning” subsection and in Figure 5 to make this clearer.

Finally, we have provided a demonstration video, "SaMI.MP4," in our Anonymous GitHub repository (https://anonymous.4open.science/r/SaMI) to better illustrate that our method helps the agent complete the three steps of "explore effectively, infer, adapt" within an episode.

Thank you very much for your feedback, especially your comments on the scalability of the method and variability in skill definitions. We are also grateful for the reviewer's recognition of the novelty and contributions of our work, as well as the acknowledgement of our empirical validation and presentation. Please let us know if further clarification is required.

Re Q1. Scalability to Complex Environments: To further explore the full potential and limitations of SaMI, we have added 4 new environments in the MuJoCo benchmark (Walker, Crippled Walker, Crippled Hopper, and HumanoidStandup), adopted broader test tasks (with more extreme unseen mass and damping values in testing tasks, shown in Table 3 in the pdf file we submitted in the global rebuttal) and conducted a more comprehensive analysis (i.e., video demos in our Anonymous GitHub repository (https://anonymous.4open.science/r/SaMI) to demonstrate the different skills learned by various algorithms. Through additional experiments, we have gained a clearer understanding of SaMI's scalability to more complex, large-scale environments:

1. In more complex MuJoCo environments (Crippled Ant, Crippled Hopper, Crippled Half-Cheetah, SlimHumanoid, HumanoidStandup, and Crippled Walker) as well as the Panda-gym benchmark, SaMI helps the RL agents to be versatile and embody multiple skills, and leads to increased returns/success rates during training and zero-shot generalisation. These environments require the agent to master different skills. When the environmental features (e.g., cube mass, table friction) vary between tasks, the agent can use different skills to work across multiple tasks.
2. In the MuJoCo Ant, Half-Cheetah and Hopper environments, we observed that when the environment requires only a single skill, the improvement is not as significant. However, SaMI requires far fewer samples. We believe that with the 10 environments in MuJoCo and Panda-gym, we can conclude that SaMI demonstrates significant advantages in complex, large-scale environments. It enables the agent to become more versatile, thereby performing well in these complex settings.

As for the limitations, we discussed that we did not make independence assumptions about environmental features in the Discussion section. Therefore, there is potential to apply our approach to more complex environments where environmental features interact with each other (e.g., the material of the cube affecting both its mass and friction). This will be explored in our future work. We also hope this study inspires others to focus on scalability in complex environments without relying on restrictive assumptions.

We have also added another limitation in Section 5.3 “MuJoCo” under the “Results and Skill Analysis” paragraph, noting that the RL agent tends to learn a single skill to work across multiple tasks in some environments. For example, in the Hopper environment, we found that to adapt to different mass values, the TESAC/CCM/SaTESAC/SaCCM policy tends to learn only one skill, which is the Hopper hopping forward on the floor. As a result, we can see from Table 2 (in the pdf file we submitted in the global rebuttal) that the returns obtained by the four algorithms are very similar across the training tasks, moderate test tasks, and extreme test task settings. In this case, our conclusion is that "even though SaNCE uses fewer samples, it does not degrade RL performance when only one skill is required." Additionally, further discussion and research are needed to determine when multiple skills are required and what improvements multiple skills can bring beyond generalisation.

Re Q2. Variability in Skill Definitions: To more clearly explain the distinctiveness of skills, we have added a description in Section 4.4 "Skill-aware trajectory sampling strategy":

"Methods that focus on skill diversity often rely heavily on accurately defining and identifying skills (https://arxiv.org/abs/1802.06070), with some requiring a prior skill distribution that is often inaccessible (https://arxiv.org/abs/2207.07560). This variability in skill definitions across tasks can affect the robustness and generalisability of these methods. For these reasons, our approach does not use such skill definitions and priors for specific environments or tasks. In this section, we give the distinctiveness of skills and propose a practical trajectory sampling method.

In this study, we believe that distinctiveness of skills is inherently difficult to achieve — a slight difference in states can make two skills distinguishable, and not necessarily in a semantically meaningful way. Instead, we should focus on whether the skills acquired by the agent can complete the task. For example, in high-friction tasks, the agent must acquire the Pick&Place skill to avoid large frictional forces, while in high-mass tasks, the agent must learn the Push skill since it cannot lift the cube. In that way, without skills definition in a semantically meaningful way, we only need to train an agent on a set of tasks, and it will autonomously discover diverse skills to work across multiple tasks."

Additionally, based on the reviewer's comments, we have added video demos (https://anonymous.4open.science/r/SaMI) in our experiments to help readers better understand how RL agents complete different tasks through different skills. Meanwhile, using the t-SNE and PCA plots provided in the original manuscript, we can understand the distinctiveness of the acquired skills in a semantically meaningful way. For example, in the videos, we can see that SaCCM has learned the Push, Pick&Place, Drag and Slide skills in Panda-gym.

Thank you very much for your recognition of our work. We fully agree with the reviewer that generalising RL agents across tasks is an important topic that will enhance the field of RL, and we are very grateful for your positive feedback on it. We appreciate your acknowledgement of our presentation and our contribution statement. Based on the reviewer's comments, to help readers understand the learning procedure, we provide pseudo-code for how to use SaNCE during meta-training and meta-testing in Algorithms 1 and 2 in Appendix D.3 "Implementation Details" (you can find them in the pdf file we submitted in the global rebuttal). Please let us know if further clarification is required.

Re Q(1) “How prone is this to declare very similar samples, i.e. similar behaviours, as positive and negative samples? For example, the return for sample 1 might be high and the return for sample 2 is low, but they are in essence executing the same behaviour”: In Section 5.3, we define samples (i.e., trajectories) with high returns as positive samples, and the skill that generates this positive sample is defined as a positive skill. In this way, if sample 1 and sample 2 do not have similar returns (i.e., sample 1 has a high return, and sample 2 has a low return), then their corresponding skills are also not similar (i.e., skill 1 is a positive skill, and skill 2 is a negative skill). Note that the similarity of skills mentioned here is not semantic similarity; it does not mean that all skills that exhibit "pick the cube off the table and place it to the goal position" are considered similar skills. Furthermore, the reviewer's concern involves the issue of the distinctiveness of skills. To more clearly explain the distinctiveness of skills, we have added a description in Section 4.4 "Skill-aware trajectory sampling strategy":

"Methods that focus on skill diversity often rely heavily on accurately defining and identifying skills (https://arxiv.org/abs/1802.06070), with some requiring a prior skill distribution that is often inaccessible (https://arxiv.org/abs/2207.07560). This variability in skill definitions across tasks can affect the robustness and generalisability of these methods. For these reasons, our approach does not use such skill definitions and priors for specific environments or tasks. In this section, we give the distinctiveness of skills and propose a practical trajectory sampling method.

In this study, we believe that distinctiveness of skills is inherently difficult to achieve — a slight difference in states can make two skills distinguishable, and not necessarily in a semantically meaningful way. Instead, we should focus on whether the skills acquired by the agent can complete the task. For example, in high-friction tasks, the agent must acquire the Pick&Place skill to avoid large frictional forces, while in high-mass tasks, the agent must learn the Push skill since it cannot lift the cube. In that way, without skills definition in a semantically meaningful way, we only need to train an agent on a set of tasks, and it will autonomously discover diverse skills to work across multiple tasks."

Re Q(2): “How is the sampling of features exactly done and what exactly is the training set as mentioned in Section 6.1, line 246? Does this mean the RL agent is building upon existing offline data?”: The original text in line 246 is “During training, we uniform-randomly select a combination of environmental features from a training set.” This sentence means that during meta-training, at the beginning of each episode, we sample a task from the training task set to interact with for one episode. Since tasks are defined by environmental features, this means selecting a combination of environmental features from a training task set at the start of each episode. To ensure it is clear to readers that "training set" refers to data collected during training rather than offline data, we have revised the original text to: “During meta-training, we uniform-randomly select a combination of environmental features from a training task set.” Similarly, during meta-testing, at the beginning of each episode, we select a combination of environmental features from the moderate/extreme task set.

Additionally, in response to the reviewer's comment, “Could the authors elaborate on the training procedure?”, we have added the description of the meta-training and meta-testing processes in the new manuscript. We have added descriptions in Section 3, “Preliminaries,” under the “Contrastive learning” subsection and in Figure 5. We also provided an example in the first paragraph of the introduction section illustrating the need for the agent to adapt within an episode: “When facing an unknown environment, the agent needs to explore effectively, understand the environment, and adjust its behaviour accordingly within an episode. For instance, if the agent tries to push a cube across a table covered by a tablecloth and finds it ‘unpushable,’ it should infer that the table friction is relatively high and adapt by lifting the cube to avoid friction, rather than continuing to push.”

Thank you for the reviewer's feedback and for acknowledging our empirical results and recognising the importance of our research contributions to scaling contrastive learning methods for control. We hope our clarifications, corrections and analysis can convince the reviewer to reconsider their score. Please let us know if further clarification is required.

Re Q(1) the statement in Lemma 1 and typo in Figure 6: Thank you for pointing out the typo in Figure 6 in the original paper, we’ve fixed it in the updated paper, including:

1. The x-axis should be "training epochs" instead of $K$. The main idea that Figure 6 aims to present is that as the training progresses, the lower bound of $I(x;y)$, $I_\text{SaNCE}$ (yellow line), gradually approaches $I(x;y)$ (dark blue line).
2. Considering the importance of Figure 6 to Lemma 1, we have moved it to the main text (so it's Figure 3 now, you can find it in the one-page pdf).

Regarding the reviewer's question, "In some cases, it seems the mutual information between $x$ and $y$ can be less than $\log K$, regardless of the parameterization used." We have revised the statement of Lemma 1 in Section 4.1 to ensure we have considered all cases. We believe the following two changes will address the reviewer's concern:

1. We do not consider the case when $x \perp y$, i.e. $\log K > I(x;y) = 0$ ($\forall K \geq 1$), because a meta-RL agent learns a context encoder by maximising MI between trajectories $\tau_c$ and context embeddings $c$, which are not independent according to the MDP graph in Figure 2. We also added the aforementioned content in Section 4.1. More specifically, in Section 4.2, we define skills as follows: “A policy $\pi$ conditioned on a fixed context embedding $c$ is defined as a skill $\pi(\cdot|c)$, abbreviated as $\pi_c$. If a skill $\pi_c$ is conditioned on a state $s_t$, we can sample actions $a_t \sim \pi(\cdot|c,s_t)$. After sampling actions from $\pi_c$ at consecutive timesteps, we obtain a trajectory $\tau_{c}=\{s_t, a_t, r_t, s_{t+1}, \ldots, s_{t+T}, a_{t+T}, r_{t+T}\}$ which demonstrates a consistent mode of behaviour.” Therefore, trajectories $\tau_c$ and context embeddings $c$ are not independent.

2. Given that we focus on generalisation, we consider cases with a finite sample size of $K$, hence the true $I(x;y) \geq \log K$. This is because we want to learn a context encoder to compress helpful information even with finite samples, which is crucial for generalisation. The key to achieving good compression is having a reliable measure of compression, and MI is a good measure of compression (https://arxiv.org/pdf/1810.05728). $K$-sample estimators like NCE or InfoNCE were originally proposed to estimate MI from finite samples. Therefore, we believe it is important to consider whether the $K$-sample estimator can approximate MI with finite samples, i.e., we should consider cases where the true $I(x;y) \geq \log K$. We also added the above statement in Section 4.1 on page 4 of the updated manuscript.

We have included the restrictions of "finite sample size $K$" and the independence of $x$ and $y$ in Lemma 1 to ensure that the Lemma covers all cases (as shown on page 4 of the updated manuscript): Lemma 1: Learning a context encoder $\psi$ with a $K$-sample estimator and finite sample size $K$, we always have ${I}_{\text{InfoNCE}}(x;y|\psi,K)$ $\leq$ $\log K$ $\leq$ $I(x;y)$, when $x \not \perp y$. (see proof in Appendix A)

Re Q(2) definition of SaMI in eq.2: Thank you for pointing out our typo "$p(\tau_c|c,\pi_c)/p(\tau_c)$" in eq.2. We have corrected it, and it now aligns with the formula in Appendix B "Proof for Lemma 2". We have corrected eq. 2 to: $I_{\text{SaMI}}(c;\pi_c;\tau_c) = E_{p(c,\pi_c,\tau_c)} \log \ \frac{p(c,\pi_c,\tau_c)}{p(c)p(\pi_c)p(\tau_c)}$ Besides, we gave proof in Appendix B of the original manuscript: $I_{\text{SaMI}}(c;\pi_c;\tau_c) = I(c;\tau_c) - I(c;\tau_c|\pi_c)$, hence $I_{\text{SaMI}}(c;\pi_c;\tau_c) \leq I(c;\tau_c)$.

Re Q(3) How is the "momentum encoder" parameterized?: In Appendix D.3 "Implementation Details" of the original manuscript, the "momentum encoder" is updated $\theta_{\psi^*}$ by $\theta_{\psi^*} \leftarrow m \cdot \theta_{\psi} + (1-m) \cdot \theta_{\psi^*}$. This is a similar setup to existing research (e.g. https://arxiv.org/abs/1911.05722). The momentum update makes $\theta_{\psi^*}$ evolve more smoothly by having it slowly track $\theta_{\psi}$ with $m \ll 1$ (e.g., $m=0.05$ in this research).

Re Q(4) Is there significance to the bars not lining up in Figure 2(c)?: Yes, the bars not lining up in Figure 2 (c) (Figure 4 in the updated manuscript) illustrate that the negative sample space may differ across tasks. To address the reviewer's concern, we have added an explanation in Appendix C, "Sample Size of $I_{\text{Sa+InfoNCE}}$,": "Figure 4 illustrates the differences in the size of the negative sample space, in which the negative sample space may vary across tasks (as shown by the non-lining-up bars in Figure 4(c)). This is because we define negative samples as trajectories with low returns, so the size of the negative sample space is influenced by sampling randomness."

Re Q(5) related work on hard negative mining: We thank the reviewer for pointing out this related work, such methods provide insight into defining positive/negative samples. We have carefully considered this suggestion, but our method does not make significant contributions in terms of sample hardness. In the third paragraph of Section 4.4 “Skill-aware trajectory sampling strategy,” we briefly added insight we gain from hard negative samples: “Positive samples are generated by the optimal skill for the current task, while lower return samples are classified as negative. This polarised definition helps the model select the optimal skill from among many skills with varying returns and avoids the issue of hard negative examples during training (https://arxiv.org/abs/2010.04592).”

Thank you for your response. The second statement is the correct definition $I_{SaMI}(c;\pi_c;\tau_c)= I(c;\tau_c)-I(c;\tau_c|\pi_c)$, i.e., the 3-variable interaction information. So we are optimising the interaction information $I_{SaMI}(c;\pi_c;\tau_c)$, defined as $I_{SaMI}(c;\pi_c;\tau_c) = I(c;\tau_c)-I(c;\tau_c|\pi_c)$. Please kindly check the updated manuscript for the updated eq. (2) at https://anonymous.4open.science/r/SaMI/SaMI.pdf.

So the entire story of the version changes is as follows:

1. The SaMI formula $I_{SaMI}(c;\pi_c;\tau_c)$, defined based on interaction information, has not changed, and the proof (in Appendix A and B of our original manuscript) based on this formula still holds.

2. The only change made was to the detailed version of the formula with the log probabilities, i.e., $E_{p(c,\pi_c,\tau_c)} \log \ \frac{p(\tau_c|c,\pi_c)}{p(\tau_c)}$. This was corrected due to a typo pointed out by the reviewer, but it was not essential for the theory and has been removed from the paper. Since we cannot evaluate $p(c,\pi_c,\tau_c)$ directly, we approximate it by Monte Carlo sampling using $K$ samples from $p(c,\pi_c,\tau_c)$ (i.e., trajectories from environments). Therefore, extensive discussion about the details of the distribution $p(c,\pi_c,\tau_c)$ is not required.

3. We ensure that our framework effectively maximises $I_{SaMI}$ through the $K$-sample estimator $I_{SaNCE}$ proposed in Section 4.3 of the original manuscript, and the skill-aware trajectory sampling strategy proposed in Section 4.4 of the original manuscript.

Please refer to Section 4.1 of our original manuscript for our motivation. Since maximising MI $I(c;\tau_c)$ faces the $\log K$ curse, we propose $I_{\text{SaMI}}(c;\pi_c;\tau_c)$, which is a ground-truth MI that is smaller than $I(c;\tau_c)$. We state in our Lemma 1 that $I_{\text{InfoNCE}}(x;y|\psi, K) \leq \log K \leq I(x;y)$; please refer to Appendix A of our original manuscript for the proof of Lemma 1.

We have presented the motivation for this study in Section 4.1 of our original manuscript: “Good compression is crucial for generalisation, and compressing valuable information from a limited number of $K$ samples requires MI as an effective measure of compression (https://arxiv.org/abs/1810.05728). We derive three key insights when learning a context encoder with finite sample size: (1) focus on a ground-truth MI that is smaller than $I(x;y)$; (2) develop a $K$-sample estimator tighter than the $I_{\text{InfoNCE}}$; (3) increase sample quantity $K$, however, this is usually impractical. A meta-RL agent learns a context encoder by maximizing MI between trajectories $\tau_c$ and context embeddings $c$. Driven by insight (1), we introduce Skill-aware Mutual Information (SaMI) in Section 4.2, designed to enhance the zero-shot generalization of downstream RL tasks. Corresponding to insight (2), we propose Skill-aware Noise Contrastive Estimation (SaNCE) to maximize SaMI with finite samples in Section 4.3. Finally, Section 4.4 demonstrates how to equip a Meta-RL agent with SaNCE in practice.”

Besides, there are many formulas (or generalised versions of these formulas) for calculating the mutual information between variables $c,\pi_c,\tau_c$. We aim to address the $\log K$ curse. Therefore, we proposed $I_{SaMI}(c;\pi_c;\tau_c)$, which aligns with our motivation: $I_{SaMI}$ serves as a ground-truth MI that is smaller than $I(c;\tau_c)$. Hence, we maximise the interaction information $I_{SaMI}$. This allows for faster convergence in sample-limited tasks, thereby facilitating zero-shot generalisation. Achieving this is crucial for applying RL-based robots to sample-limited real-world scenarios and enabling rapid adaptation, which is one of the significant bottlenecks.

We believe there are various bounds of mutual information, lower bounds and upper bounds (https://arxiv.org/pdf/1905.06922). For example, the total correlation pointed out by the reviewer has great potential for calculating mutual information among multiple variables. However, it does not align with our motivation. Since $TC \geq I_{SaMI}$, it is not suitable for zero-shot generalisation scenarios.

Thank you for the reviewer’s comments, especially the helpful comments on the presentation of experimental results. We appreciate your recognition of the novelty of our work and the potential of our method. We have made revisions to enhance the clarity and presentation of our experimental results based on your suggestions. Please let us know if further clarification is required.

Re statistical significance tests: We have added a t-test to conduct a statistical hypothesis test to determine whether SaMI brought a statistically significant improvement over the respective base RL algorithm without SaMI. i) In the updated manuscript, we marked with an asterisk (*) in Table 1 and Table 2 where the algorithm with SaMI shows statistically significant improvement over the same algorithm without SaMI (you can also find them in the pdf file we submitted in the global rebuttal). The t-test results indicate that SaMI brings significant improvement to the extreme test set. This aligns with our conclusions, that SaMI leads to increased returns during training and zero-shot generalisation, especially in environments that require different skills. We have added this information in Section 5.3. ii) We also report the p-values of the t-tests in Appendix H in the updated manuscript.

Re Q(1) whether the generalisation behaviour is different for SaNCE in Panda-Gym vs MuJoCo: The results from Panda-Gym and MuJoCo are consistent as they both show: i) SaMI helps the RL agents to be versatile and embody multiple skills; ii) SaMI leads to increased performance during training and zero-shot generalisation, especially in environments that require different skills. Specifically,

1) Panda-Gym: The RL agent must equip itself with different skills to achieve a high success rate across multiple tasks because different environmental features require different skills. For example, tasks with high friction require the Pick&Place skill (picking the cube off the table and placing it in the goal position), while tasks with high mass require the Push skill (pushing the cube to move it across the table). Therefore, in the Panda-Gym benchmark, we see that multiple skills are needed in training tasks, moderate test tasks, and extreme test tasks, so SaNCE brings improvement across all three.

2) MuJoCo: In the MuJoCo environment, we found: i) In some environments (Ant, Hopper, and Half-Cheetah), the RL agent only needs to learn a single skill to generalise across different tasks. For example, in the Hopper environment, we found that to adapt to different mass values, the TESAC/CCM/SaTESAC/SaCCM policy tends to learn only one skill, which is the Hopper hopping forward on the floor. As a result, we can see from Table 2 that the returns obtained by the four algorithms are very similar across the training tasks, moderate test tasks, and extreme test task settings. In this case, our conclusion is that "even though SaNCE uses fewer samples, it does not degrade RL performance"; ii) When the environment (Crippled Ant, Crippled Hopper, Crippled Half-Cheetah, SlimHumanoid, HumanoidStandup, and Crippled Walker) requires different skills to complete tasks, SaNCE brings significant improvements. For example, in the Crippled Ant environment, when it has 3 or 4 legs available, the Crippled Ant Robot learns to roll to adapt to varying mass and damping in the training tasks and moderate tasks. However, during zero-shot generalisation in extreme tasks, when only 2 legs are available, the Ant Robot can no longer roll. Instead, it adapts by walking using its two legs. Therefore, we can see that SaNCE brings significant improvement, especially during zero-shot generalisation.

To better showcase generalisation behaviour, we have added 4 new environments in the MuJoCo benchmark (Walker, Crippled Walker, Crippled Hopper, and HumanoidStandup), and adopted broader test tasks (with more extreme unseen mass and damping values in testing tasks, shown in Table 3 in the pdf file we submitted in the global rebuttal). This allows us to better demonstrate SaNCE's zero-shot generalisation ability across more environments and less familiar testing tasks. Additionally, to better present our results, we have utilised videos on GitHub to clearly demonstrate how SaNCE helps the RL algorithm learn different skills and perform zero-shot generalisation (https://anonymous.4open.science/r/SaMI).

To help readers better understand the generalisation behaviour, we have added the aforementioned descriptions in Section 5.2 "Panda-gym" and Section 5.3 "MuJoCo" under the "Result and Skill Analysis" paragraph.

Re Q(2) and (3) t-SNE plots and claims on improvements: We have updated the t-SNE plots in Figure 6 in the one-page pdf. The original plot represented 100 trajectories for each task, where each point in the t-SNE plot represents the context embedding of the final time step of each trajectory, encapsulating the skill information of the entire trajectory. However, when randomly initialising the cube's starting position, if the cube is very close to the target position, the agent only needs to slightly move the cube to complete the task. In this case, the embeddings for all 4 tasks cluster together. Additionally, there are failure cases across all 4 tasks where the agent attempts to move the cube, but the cube remains still, leading to clustering of the embeddings. Therefore, the 2 clear clusters containing all tasks represent these two scenarios, with one cluster on the lower left and another on the lower right. To better illustrate the clustering of embeddings based on different skills, we have removed the aforementioned two scenarios from the 100 tests for each task. In the revised t-SNE plot, we can clearly see two distinct clusters.

Regarding the presentation of results: To better present our experimental results, we have made the following improvements: 1) We’ve deleted the blue background. 2) We have removed the average improvement percentage.