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

Authors: xuehui  yu
mhairi  dunion
xin  li
stefano v albrecht
𝑐  𝜋
!  𝜏
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:58:06

Abstract: Meta-Reinforcement Learning (Meta-RL) agents can struggle to operate across tasks with varying environmental features that require different optimal skills (i.e., different modes of behaviour). Using context encoders based on contrastive learning to enhance the generalisability of Meta-RL agents is now widely studied but faces challenges such as the requirement for a large sample size, also referred to as the $\log$-$K$ curse. To improve RL generalisation to different tasks, we first introduce Skill-aware Mutual Information (SaMI), an optimisation objective that aids in distinguishing context embeddings according to skills, thereby equipping RL agents with the ability to identify and execute different skills across tasks. We then propose Skill-aware Noise Contrastive Estimation (SaNCE), a $K$-sample estimator used to optimise the SaMI objective. We provide a framework for equipping an RL agent with SaNCE in practice and conduct experimental validation on modified MuJoCo and Panda-gym benchmarks. We empirically find that RL agents that learn by maximising SaMI achieve substantially improved zero-shot generalisation to unseen tasks. Additionally, the context encoder trained with SaNCE demonstrates greater robustness to a reduction in the number of available samples, thus possessing the potential to overcome the $\log$-$K$ curse.

Main paper: 
 introduction: 
 Reinforcement Learning (RL) agents often learn policies that do not generalise across tasks in which the environmental features and optimal skills are different [des Combes et al., 2018 , Garcin et al., 2024] . Consider a set of cube-moving tasks where an agent is required to move a cube to a goal position on a table (Figure 1 ). These tasks become challenging if environmental features, such as table friction, vary between tasks. When facing an unknown environment, the agent have different transition dynamics caused by underlying environmental features, hence optimal skills are different across tasks. 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. Recent advances in Meta-Reinforcement Learning (Meta-RL) [Lee et al., 2020 , Agarwal et al., 2021 , Mu et al., 2022 , Dunion et al., 2023b ,a, McInroe et al., 2024] enable agents to understand environmental features by inferring context embeddings from a small amount of exploration, and to train a policy conditioned on the context embedding to generalise to novel tasks.In recent years, unsupervised contrastive learning algorithms have been shown to learn context embeddings that perform remarkably well on generalisation tasks [Clavera et al., 2019a , Lee et al., 2020] . In particular, some Meta-RL algorithms [Fu et al., 2021 , Wang et al., 2021 , Li et al., 2021 , Sang et al., 2022] train context encoders by maximising InfoNCE [Oord et al., 2019] , which has yielded vital insights into contrastive learning through the lens of mutual information (MI) analysis. The InfoNCE objective can be interpreted as a K-sample lower bound on the MI between trajectories and context embeddings. Despite significant advances, integrating contrastive learning with Meta-RL poses several unresolved challenges, of which two are particularly relevant to this research: (i) Existing context encoders based on contrastive learning do not distinguish tasks that require different skills. Many prior algorithms only pull embeddings of the same tasks together and push those of different tasks apart. However, for example, a series of cube-moving tasks with high friction may only require a Pick&Place skill (picking the cube off the table and placing it at the goal position), making further differentiation unnecessary. (ii) Existing K-sample MI estimators are sensitive to the sample size K (i.e., the log-K curse) [Poole et al., 2019] . The low sample efficiency of RL [Franke et al., 2021] and the sample limitations in zero-shot generalisation make collecting a substantial quantity of samples often impractical [Arora et al., 2019, Nozawa and Sato, 2021] . The effectiveness of K-sample MI estimators breaks down with a finite sample size and leads to a significant performance drop in downstream RL tasks [Mnih and Teh, 2012, Guo et al., 2022] .To enhance RL generalisation across different tasks, we propose that the context embeddings should optimise downstream tasks and indicate whether the current skill remains optimal or requires further exploration. This also reduces the necessary sample size by focusing solely on skill-related information. In this work, (1) we introduce Skill-aware Mutual Information (SaMI), a generalised form of MI objective between context embeddings, skills, and trajectories, designed to address issue (i). We provide a theoretical proof showing that by introducing skills as a third variable into the MI of context embeddings and trajectories, the resulting SaMI is smaller and easier to optimise. Furthermore, (2) we propose a data-efficient K-sample estimator, Skill-aware Noise Contrastive Estimation (SaNCE) to optimise SaMI, effectively addressing issue (ii). Additionally, (3) we propose a practical skill-aware trajectory sampling strategy that shows how to sample positive and negative examples without relying on any prior skill distribution. In that way, Meta-RL agents autonomously acquire a set of skills applicable to many tasks, with these skills emerging solely from the SaMI learning objective and data.We demonstrate empirically in MuJoCo [Todorov et al., 2012] and Panda-gym [Gallouédec et al., 2021] that SaMI enhances the zero-shot generalisation capabilities of two Meta-RL algorithms [Yu et al., 2020 , Fu et al., 2021] by achieving higher returns and success rates on previously unseen tasks, ranging from moderate to extreme difficulty. Visualisation of the learned context embeddings reveals distinct clusters corresponding to different skills, suggesting that the SaMI learning objective enables the context encoder to capture skill-related information from trajectories and incentivise Meta-RL agents to acquire a diverse set of skills. Moreover, SaNCE enables Meta-RL algorithms to use smaller sample spaces while achieving improved downstream control performance, indicating their potential to overcome the log-K curse. 
related works: 
 Meta-RL. By conditioning on an effective context embedding, Meta-RL policies can zero-shot generalise to new tasks with a small amount of exploration Kirk et al. [2023] . Existing algorithms can be categorised into three types based on different context embeddings. In the first category, the context embedding is learned by minimising the downstream RL loss [Rakelly et al., 2019 , Yu et al., 2020] . PEARL [Rakelly et al., 2019] learns probabilistic context embeddings by recovering the value function. Multi-task SAC with task embeddings as inputs to policies (TESAC) [Hausman et al., 2018 , Yu et al., 2020] parameterises the learned policies through a shared embedding space, aiming to maximise the average returns across tasks. However, the update signals from the RL loss are stochastic and weak, and may not capture the similarity relations among tasks [Fu et al., Figure 2 : A policy π conditioned on a fixed context embedding c is defined as a skill π(•|c) (shortened as π c ). The policy π conditioned on a fixed c alters the state of the environment in a consistent way, thereby exhibiting a mode of skill. The skill π(•|c 1 ) moves the cube on the table in trajectory τ + c1 and is referred to as the Push skill; correspondingly, the Pick&Place skill π(•|c 2 ) takes the cube off the table and places it in the goal position in the trajectory τ + c2 .2021]. The second category involves learning context embeddings through dynamics prediction [Lee et al., 2020 , Zhou et al., 2019] , which can make the context embeddings noisy, as they may model irrelevant dependencies and overlook task-specific information [Fu et al., 2021] . The third category employs contrastive learning [Fu et al., 2021 , Wang et al., 2021 , Li et al., 2021 , Sang et al., 2022] , achieving significant improvements in context learning. However, these methods overlook the similarity of skills between different tasks, thus failing to achieve effective zero-shot generalisation by executing different skills. Our improvements build upon this third category by distinguishing context embeddings according to different optimal skills.Contrastive learning. Contrastive learning has been applied to RL due to its significant momentum in representation learning in recent years, attributed to its superior effectiveness [Tishby and Zaslavsky, 2015 , Hjelm et al., 2019 , Dunion and Albrecht, 2024] , ease of implementation [Oord et al., 2019] , and strong theoretical connection to MI estimation [Poole et al., 2019] . MI is often estimated using InfoNCE [Oord et al., 2019] that has gained recent attention due to its lower variance [Song and Ermon, 2020] and superior performance in downstream tasks. However, InfoNCE may underestimate the true MI when the sample size K is finite. To address this limitation, CCM [Fu et al., 2021] uses a large number of samples for maximising InfoNCE. DOMINO [Mu et al., 2022] reduces the true MI by introducing an independence assumption; however, this results in biased estimates. We focus on proposing an unbiased alternative MI objective and a more data-efficient K-sample estimator tailored for downstream RL tasks, which, to our knowledge, have not been addressed in previous research. 
preliminaries: 
 Reinforcement learning. In Meta-RL, we assume an environment is a distribution ξ(e) of tasks e (e.g. uniform in our experiments). Each task e ∼ ξ(e) has a similar structure that corresponds to a Markov Decision Process (MDP) [Puterman, 2014] , defined by M e = (S, A, R, P e , γ), with a state space S, an action space A, a reward function R(s t , a t ) where s t ∈ S and a t ∈ A, state transition dynamics P e (s t+1 |s t , a t ), and a discount factor γ ∈ [0, 1). In order to address the problem of zeroshot generalisation, we consider the transition dynamics P e (s t+1 |s t , a t ) vary across tasks e ∼ ξ(e) according to multiple environmental features e = {e 0 , e 1 , ..., e N } that are not included in states s and can be continuous random variables, such as mass and friction, or discrete random variables, such as the cube's material. For instance, in a cube-moving environment (Figure 1 ), an agent has different tasks that are defined by different environmental features (e.g., mass and friction). The Meta-RL agent's goal is to learn a generalisable policy π that is robust to such dynamic changes. Specifically, given a set of training tasks e sampled from ξ train (e), we aim to learn a policy that can maximise the discounted returns, arg max π E e∼ξtrain(e) [ ∞ t=0 γ t R(s t , a t )|a t ∼ π(a t |s t ), s t+1 ∼ P e (s t+1 |s t , a t )], and can produce accurate control for unseen test tasks sampled from ξ test (e).Contrastive learning. In Meta-RL, the context encoder ψ(c|τ c,0:t ) first takes the trajectory τ c,0:t = {s 0 , a 0 , r 0 , ..., s t } from the current episode as input and compresses it into a context embedding c [Fu et al., 2021] . Then, the policy π, conditioned on context embedding c, consumes the current state s t and outputs the action a t . As a key component, the quality of context embedding c can affect algorithms' performance significantly. MI is an effective measure of embedding quality [Goldfeld et al., 2019] , hence we focus on a context encoder that optimises the InfoNCE objective I InfoNCE (x; y), which is a K-sample estimator and lower bound of the MI I(x; y) [Oord et al., 2019] . Given a query x and a set Y = {y 1 , ..., y K } of K random samples containing one positive sample y 1 and K -1 negative samples from the distribution p(y), I InfoNCE (x; y) is obtained by comparing pairs sampled from the joint distribution x, y 1 ∼ p(x, y) to pairs x, y k built using a set of negative examples y 2:K :I InfoNCE (x; y|ψ, K) = E log f ψ (x, y 1 ) 1 K K k=1 f ψ (x, y k ).(1)InfoNCE constructs a formal lower bound on the MI, i.e., I InfoNCE (x; y|ψ, K) ≤ I(x; y) [Guo et al., 2022 , Chen et al., 2021] . Given two inputs x and y, their embedding similarity is f ψ (x, y) = e ψ(x) ⊤ •ψ(y)/β , where ψ is the context encoder that projects x and y into the context embedding space, the dot product is used to calculate the similarity score between ψ(x), ψ(y) pairs [Wu et al., 2018 , He et al., 2020] , and β is a temperature hyperparameter that controls the sensitivity of the product. Some previous Meta-RL methods [Lee et al., 2020 , Mu et al., 2022] learn a context embedding c by maximising I InfoNCE (c; τ c |ψ, K) between the context c embedded from a trajectory in the current task, and the historical trajectories τ c under the same environmental features setting.4 Skill-aware mutual information optimisation for Meta-RL 4.1 The log-K curse of K-sample MI estimatorsIn this section, we provide a theoretical analysis of the challenge inherent in learning a K-sample estimator for MI, commonly referred to as the log-K curse. Based on this theoretical analysis, we give insights to overcome this challenge. Given that we focus on the generalisation of RL, we only consider cases with a finite sample size of K. If a context encoder ψ in Equation ( 1) has sufficient training epochs, then I InfoNCE (x; y|ψ, K) ≈ log K [Mnih and Teh, 2012, Guo et al., 2022] . Hence, the MI we can optimise is bottlenecked by the number of available samples, formally expressed as: Lemma 1 Learning a context encoder ψ with a Ksample estimator and finite sample size K, we have I InfoNCE (x; y|ψ, K) ≤ log K≤I(x; y), when x ̸⊥ ⊥ y (see proof in Appendix A).We do not consider the case when x ⊥ ⊥ y, i.e., log K ≥ I(x; y) = 0 (∀K ≥ 1), because a Meta-RL agent learns a context encoder by maximising MI between trajectories τ c and context embeddings c, which are not independent (as shown in Figure 2 ). While an unbounded sample size K for learning effective context embeddings is theoretically feasible and assumed in many studies Mu et al. [2022] , Lee et al. [2020] , it is often impractical in practice.  
skill-aware mutual information: a smaller ground-truth mi: 
 We aim for our MI learning objective to incentivise agents to acquire a diverse set of skills, enabling them to generalise effectively across tasks. To start with, we define skills [Eysenbach et al., 2018] :Definition 1 (Skills) A policy π conditioned on a fixed context embedding c is defined as a skill π(•|c), abbreviated as π c . If a skill π c is conditioned on a state s t , we can sample actions a t ∼ π(•|c, s t ). After sampling actions from π c at consecutive timesteps, we obtain a trajectory τ c,t:t T = {s t , a t , r t , s t+1 , . . . , s t+T , a t+T , r t+T } which demonstrates a consistent mode of behaviour.After a limited amount of exploration, an agent should be able to infer the task (i.e., environmental features e = e 0 , e 1 , . . . , e N ) and adapt accordingly within the current episode. The context em- bedding should encompass skill-related information, guiding the policy on when to explore new skills or switch between existing ones. We propose that the context encoder ψ should be trained by maximising the MI between the context embedding c, skills π c , and trajectories τ c . To this end, we propose a novel MI optimisation objective, Skill-aware Mutual Information (SaMI), defined as:I SaMI (c; π c ; τ c ) = I(c; τ c ) -I(c; τ c |π c ).(2)SaMI is defined according to interaction information [McGill, 1954] , serving as a generalisation of MI for three variables {c, π c , τ c }. Although we cannot evaluate p(c, π c , τ c ) directly, we approximate it by Monte-Carlo sampling, using K samples from p(c, π c , τ c ). As illustrated in Figure 3 , a context encoder ψ trained with the objective of maximising I SaMI (c; π c ; τ c ) converges more quickly, asI SaMI (c; π c ; τ c ) ≤ I(c; τ c ) (see proof in Appendix B). By focusing more on skill-related information, I SaM I enables agents to autonomously discover a diverse range of skills for handling multiple tasks.4.3 Skill-aware noise contrastive estimation: a tighter K-sample estimator Despite InfoNCE's success as a K-sample estimator for approximating MI [Laskin et al., 2020 , Eysenbach et al., 2022] , its learning efficiency plunges due to limited numerical precision, which is called the log-K curse, i.e., I InfoNCE ≤ log K ≤ I SaMI [Chen et al., 2021] (see proof in Appendix B).When K → +∞, we can expect I InfoNCE ≈ log K ≈ I SaMI [Guo et al., 2022] . However, increasing K is too expensive, especially in complex environments with enormous negative sample space. To address this, we propose a novel K-sample estimator that requires a significantly smaller sample size K ≪ +∞. First, we define K * : Definition 2 (K * ) K * = |c| • |π c | • M isEQUATIONwhere f ψ (c 1 , π c1 , τ c1 ) = e ψ(τc 1 ) ⊤ •ψ * (τc 1 )/β . The query c 1 = ψ(τ c1 ) is generated by the context encoder ψ. For training stability, we use a momentum encoder ψ * to produce the positive and negative embeddings. SaNCE significantly reduces the required sample space size K * SaNCE by sampling trajectories τ c based on different skills π c (Figure 4 (b)) in task e 1 , so that K * SaNCE = |c| • |π c | • M = |π c1 | • M ≤ K * InfoNCE (|c| = |c 1 | = 1).In this section, we propose a practical trajectory sampling method. Methods focusing on skill diversity often rely heavily on accurately defining and identifying individual skills [Eysenbach et al., 2018] . Some of these methods require a prior skill distribution, which is often inaccessible [Shi et al., 2022] , and it is impractical to enumerate all possible skills that we hope the model to learn. Besides, 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. Consequently, we do not directly teach any of these skills or assume any prior skill distribution. Diverse skills naturally emerge from the incentives of the SaMI learning objective in a multi-task setting, driven by the inherent need to develop generalisable skills. For example, in high-friction tasks, the agent must acquire the Pick&Place skill to avoid large frictional forces, whereas in high-mass tasks, the agent must learn the Push skill since it cannot lift the cube. In each task, we only identify whether the skills are optimal; in this way, under a multi-task setting, the agent will acquire a set of general skills that are applicable to many tasks.In a given task e, positive skills π + c are defined as optimal skills achieving highest return t+T i=t R(s i , a i ), whereas negative skills π - c are those that result in lower returns. As a result, the positive sample τ + c consists of trajectories generated by the skills with the highest ranked returns, while the negative samples correspond to those with the lowest returns. This straightforward approach of selecting positive samples based on the ranked highest return effectively aligns with positive skills and mitigates the challenge of hard negative examples [Robinson et al., 2021] . The SaNCE loss is then minimised to bring the context embeddings of the highest return trajectories closer while distancing those of negative trajectories. By the end of training, the top-ranked trajectories in the ranked replay buffer correspond to positive samples τ + c with high returns, while the lower-ranked trajectories represent negative samples τ - c with low returns. However, at the beginning of training, it is likely that all trajectories have low returns. Therefore, our SaNCE loss is a soft variant of the K-sample SaNCE:L SaNCE = -max ||ψ(τ + c ), ψ(τ - c )|| L2 , 1 • I SaNCE (4)where ||•|| L2 represents the Euclidean distance [Tabak, 2014] . Figure 5 provides a practical framework of SaNCE, with a cube-moving example task e 1 under high friction. In task e 1 , the positive skill π + c1 is the Pick&Place skill, which is used to generate queries ψ(τ + c1 ) and positive embeddings ψ * (τ + c1 ); after executing Push skill we get negative samples τ - c1 and negative embeddings ψ * (τ - c1 ). 
experiments: 
 Our experiments aim to answer the following questions: (1) Does optimising SaMI lead to increased returns during training and zero-shot generalisation (see Table 1 and 2)?; (2) Does SaMI help the RL agents to be versatile and embody multiple skills (see Figure 6 )?;(3) Can SaNCE overcome the log-K curse in sample-limited scenarios (see Table 1 and 2, and Section 5.4)?Modified benchmarks with multiple environmental features.1 We evaluate our method on two benchmarks, Panda-gym [Gallouédec et al., 2021] and MuJoCo [Todorov et al., 2012] Baselines. The loss function of the context encoder in all algorithms consists of two key components: the RL loss and the contrastive loss. The RL loss L RL , which is the same across all methods, corresponds to the RL value function loss. Our primary comparison focuses on the contrastive loss, taking the form of either SaNCE, InfoNCE, or no contrastive loss. Accordingly, we select baselines based on contrastive loss: CCM [Fu et al., 2021] , which utilises InfoNCE, and TESAC [Yu et al., 2020] , which relies solely on the RL loss, allowing assessment of the context encoder without contrastive loss. Since CCM and TESAC use an RNN encoder, we also include PEARL [Rakelly et al., 2019] , which utilises an MLP context encoder and follows the similar RL loss. Additionally, Appendix G includes comparisons with DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020] , using the same environmental setup in the MuJoCo benchmark.Our method.2 We use Soft Actor-Critic (SAC) [Haarnoja et al., 2018] Task description. Our modified Pandagym benchmark involves a robot arm control task using the Franka Emika Panda [Gallouédec et al., 2021] , where the robot moves a cube to a target position. Unlike previous works, we simultaneously modify multiple environmental features (cube mass and table friction) that characterise the transition dynamics, and the robot can flexibly execute different skills (Push and Pick&Place) for different tasks. This environment requires high skill diversity; for instance, the agent must use Pick&Place in high-friction tasks and Push in high-mass tasks. were used to visualise executed skills. Figure 6 shows that SaCCM agents learned the Push skill for large cube masses (30 Kg, 10 Kg) and Pick&Place for smaller masses, while CCM showed no clear skill grouping (Figure 16 in Appendix F.1). Overall, SaMI incentivises agents to autonomously learn diverse skills, enhancing generalisation across a wider range of tasks. Specifically, through the cycle of effective exploration, context inference, and adaptation, diverse skills emerge solely from the data. Further visualisation results are in Appendix F.Task description. We extended the modified MuJoCo benchmark introduced in DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020] . It contains ten typical robotic control environments based on the MuJoCo physics engine [Todorov et al., 2012] . Hopper, Walker, Half-cheetah, Ant, HumanoidStandup, and SlimHumanoid are influenced by continuous environmental features (i.e., mass, damping) that affect transition dynamics. Crippled Ant, Crippled Hopper, Crippled Walker, and PEARL 7524±2455 3355±2555 1984±356 7899±2532 4377±2563 2965±1426 TESAC 7747±1772 4355±1530 2581±407 9908±1561 5929±1971 3041±912 CCM 8136±557 5476±803 2519±682 10317±1137 6233±1869 3098±821 SaTESAC 8675±752 5840±676 3632±404 * 10389±1031 8387±1291 * 4280±485 * SaCCM 8361±586 5779±691 3481±332 * 10496±951 8235±1212 4824±839 *Crippled Half-cheetah are more challenging due to the addition of discrete environmental features (i.e., randomly crippled leg joints), requiring agents to master different skills (e.g., switching from running to crawling after a leg is crippled).Results and skill analysis. 1 and 2 , RL algorithms using SaNCE achieve better or comparable performance with significantly fewer negative samples (K) than InfoNCE. This suggests SaNCE effectively addresses the log-K curse, and the SaMI objective helps the contrastive context encoder extract critical information for downstream RL tasks.The number of negative samples K is influenced by two hyperparameters: buffer size, which determines the negative sample space, and contrastive batch size, which controls the number of samples used to train the contrastive context encoder per update.We analysed these hyperparameters further, and as shown in Figure 7 , reductions in buffer and contrastive batch size do not significantly impact the average return for SaCCM and SaTE-SAC, which maintain state-of-the-art performance with small buffers and batch sizes. The results in Table 2 correspond to a buffer size of 100,000 and a batch size of 12. Results across all environments (refer to Appendix E.2) show that SaNCE exhibits low sensitivity to K, highlighting its potential to overcome the log-K curse.(b) (a)  
conclusion and future work: 
 Zero-shot generalisation has been a longstanding challenge concerning Meta-RL agents, with skill diversity and sample efficiency being key to generalising to previously unseen environments. In this paper, we proposed Skill-aware Mutual Information (SaMI) to learn context embeddings for zero-shot generalisation in downstream RL tasks, and Skill-aware Noise Contrastive Estimation (SaNCE) to optimise SaMI and overcome the log-K curse, along with a practical skill-aware trajectory sampling strategy. Experimental results showed that RL algorithms equipped with SaMI achieved state-ofthe-art performance in MuJoCo and Panda-gym benchmarks, particularly in zero-shot generalisation within more complex environments. During the zero-shot generalisation, when faced with an unseen task, SaMI assists agents in exploring effectively, inferring context, and rapidly adapting their skills within the current episode. SaNCE's optimisation uses a significantly smaller negative sample space than baselines, and our analysis on buffer and contrastive batch sizes demonstrated its effectiveness in addressing the log-K curse.Given that environmental features are often interdependent, such as a cube's material correlating with friction and mass, SaMI does not introduce independence assumptions like DOMINO [Mu et al., 2022] . Therefore, future work will focus on verifying and enhancing SaMI's potential in more complex tasks where environmental features are correlated [Dunion et al., 2023a ]. This will contribute to our ultimate goal: developing a generalist and versatile agent capable of working across multiple tasks and even real-world tasks in the near future. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All assets used in our study are open-source, and we have acknowledged the authors' copyrights in the References.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: We provide detailed descriptions of the assets we will release in Appendix D and on our open-source homepage. During the double-blind review phase, we have anonymised the URLs.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: The paper does not involve crowdsourcing nor research with human subjects Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Given a query x and a set Y = {y 1 , . . . , y K } of K random samples, containing one positive sample y 1 and K -1 negative samples drawn from the distribution p(y), a K-sample InfoNCE estimator is obtained by comparing pairs sampled from the joint distribution (x, y 1 ) ∼ p(x, y) with pairs (x, y k ), constructed using the set of negative examples y 2:K . InfoNCE compares the positive pairs (x, y 1 ) with the negative pairs (x, y k ), where y k ∼ y 2:K , as follows:EQUATIONStep 1. Let us prove that the K-sample InfoNCE estimator is upper-bounded by log K. According to Mu et al. [2022] ,f ψ (x,y1) K k=1 f ψ (x,y k ) = f ψ (x,y1) f ψ (x,y1)+ K k=2 f ψ (x,y k ) ≤ 1.So we have:EQUATIONHence, we have I InfoNCE (x; y|ψ, K) ≤ log K.Step 2. We have the I(x; y) ≥ I InfoNCE (x; y|ψ, K) according to:Proposition Step 3. In this research, the context encoder ψ in f ψ (x, y) is implemented using an RNN to approximate p(y|x) p(y) [Oord et al., 2019] . With a sufficiently powerful deep learning model for ψ and a finite sample size K, such that I(x; y) ≥ log K, we can reasonably expect that I InfoNCE ≈ log K after a few training epochs. Therefore, during training, when K ≪ +∞, we always have I(x; y) ≥ log K.Proof. See Chen et al. [2021] for more detailed proof.Step 4. Let us prove that the K-sample InfoNCE bound is asymptotically tight. The specific choice of the context encoder ψ influences the tightness of the K-sample NCE bound. InfoNCE [Oord et al., 2019] sets f ψ (x, y) ∝ p (y|x) p(y) to model a density ratio that preserves the MI between x and y, where ∝ stands for 'proportional to' (i.e., up to a multiplicative constant). Substitutingf ψ (x, y) = f * ψ (x, y) = p(y|x) p(y)into InfoNCE, we obtain:I InfoNCE (x; y|ψ, K) = E log f * ψ (x, y 1 ) K k=1 f * ψ (x, y k ) + log K = -E log 1 + p(y) p(y|x) K k=2 p(y k |x) p(y k ) + log K ≈ -E log 1 + p(y) p(y|x) (K -1)E y k ∼p(y) p(y k |x) p(y k ) + log K = -E log 1 + p(y 1 ) p(y 1 |x) (K -1) + log K ≈ -E log p(y) p(y|x) -log(K -1) + log K = I(x; y) -log(K -1) + log K (7)Now taking K → +∞, the last two terms cancel out.Putting it together. Combining I(x; y) ≥ log K with Proposition 1 and Equation ( 6), we have Lemma 1:EQUATION) Moreover, according to Equation ( 7), as the sample size K → +∞, the K-sample InfoNCE bound becomes sharp and approaches the true MI I(x; y), i.e., I InfoNCE (x; y|ψ, K) ≈ log K ≈ I(x; y).Step 1. According to Lemma 1, we have 3 ).I InfoNCE (c; π c ; τ c |ψ, K) ≤ log K ≤ I SaMI (c; π c ; τ c ) (shown in FigureStep 2. Let us prove that SaNCE is a K-sample SaNCE estimator and is upper bounded by log K.Since Mu et al., 2022] , we have:f ψ (c,πc,τ + c ) f ψ (c,πc,τ + c )+ K k=2 f ψ (c,πc,τ - c,k ) ≤ 1 [I SaNCE (c; π c ; τ c |ψ, K) = E p(c1,πc 1 ,τ + c 1 )p(τ - c 1 ,2:K ) log K • f ψ (c 1 , π c1 , τ + c1 ) f ψ (c 1 , π c1 , τ + c1 ) + K k=2 f ψ (c 1 , π c1 , τ - c1,k ) = E p(c1,πc 1 ) E p(τ - c 1 ,2:K ) log K • f ψ (c 1 , π c1 , τ + c1 ) f ψ (c 1 , π c1 , τ + c1 ) + K k=2 f ψ (c 1 , π c1 , τ - c1,k ) ≤ log K (9) Thus, we obtain I SaNCE (c; π c ; τ c |ψ, K) ≤ log K, similar to Equation (6).Step 3. With the definition of K * , we can prove thatI InfoNCE (c; π c ; τ c |ψ, K) ≤ I SaNCE (c; π c ; τ c |ψ, K)with the same sample size K. In task e 1 , SaNCE obtains positive and negative samples from the current task e 1 . Since the variable c = c 1 is constant, we have:EQUATION)The required sample size isK * SaNCE = |c 1 | • |π| • M = |π| • M . As K → K * SaNCE , we have I SaNCE (c; π c ; τ c |ψ, K) ≈ I SaMI (c; π c ; τ c ).Correspondingly, for InfoNCE, in the current task e 1 with context embedding c 1 , positive samples are trajectories τ 1 generated after executing the skill π 1 in task e 1 , while negative samples are trajectories {τ - c2 , ...} from other tasks {e 2 , ...}. Under the definition of K * , we have:EQUATIONwhereK * InfoNCE = |c| • |π| • M ≈ |c| • K * SaNCE .In real-world robotic control tasks, the sample space size increases significantly due to multiple environmental features e = {e 0 , e 1 , . . . , e N }. The number of different tasks |c| grows exponentially due to the permutations and combinations of the N environmental features. When the current task e 1 has context embedding c 1 , the c 2:|c| refer to the context embeddings for the other tasks. As K → K * InfoNCE , we have I InfoNCE (c; π c ; τ c |ψ, K) ≈ I SaMI (c; π c ; τ c ). Thus, during the training process, I InfoNCE (c; π c ; τ c |ψ, K) ≤ I SaNCE (c; π c ; τ c |ψ, K) with the same sample size K.Step 4. According to the definition of SaMI in Equation ( 2), we have I SaMI (c; π c ; τ c ) ≤ I(c; τ c ), as illustrated using Venn diagrams in Figure 8 (a) and (b). SaMI is formulated based on interaction information [McGill, 1954] , which aims to capture the relationships among multivariate variables by quantifying the amount of information (redundancy or synergy) shared among three variables. Interaction information has been less extensively studied, partly due to its challenging interpretation from both information theory and neuroscience perspectives, as it can be either positive or negative [Li, 2022] . In this section, we illustrate the sample size of I Sa+InfoNCE (c; π c ; τ c |ψ, K). Sa+InfoNCE incorporates SaNCE into InfoNCE, using positive samples τ + c1 from task e 1 after executing skill π + c1 , and negative samples are trajectories τ - c 1:K from executing skills π - c 1:K in tasks e 1:K , respectively. Therefore, this approach is equivalent to first observing the variable c and then observing the variable π c , i.e., sampling from the distribution p(π c , τ c |c)p(c). We have:EQUATION)It should be noted that such a combination increases the size of the negative sample space, i.e.,K * Sa+InfoNCE = |c| i=1 |π - ci | + |π + c1 | • M ≥ K * SaNCE .The misaligned bars in Figure 4 (c) illustrate that negative sample spaces may vary across tasks. This variation arises because we define negative samples as trajectories with low returns, making the size of the negative sample space influenced by sampling randomness. With the same number K of samples, I Sa+InfoNCE is less precise and looser than I SaNCE . Therefore, a trade-off between sample diversity and the precision of the K-sample estimator is required. We modified the original Pick&Place task in Panda-gym [Gallouédec et al., 2021] by setting the z dimension (i.e., the desired height) of the cube's goal position to 03 and maintaining the freedom of the grippers4 , allowing the agent to explore whether it should push or grasp the cube. Skills in this benchmark are defined as:• Pick&Place skill: This skill specifically refers to the agent using the gripper to grasp the cube, lift it off the table, and place it in the goal position. We determine the Pick&Place skill by detecting no contact points between the table and the cube, two contact points between the robot's end effector and the cube, and the cube's height being greater than half its width.• Push skill: This skill involves the agent moving the cube on the table to the goal position, either by dragging or sliding it. We confirm the Push skill by detecting that the cube's height equals half its width.• Other skills: Any behaviour modes other than Pick&Place and Push are classified as other skills.Some elements in the RL framework are defined as follows: We extended the modified MuJoCo benchmark introduced in DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020] . In our extension, four new environments were added (Walker, Crippled Hopper, Crippled Walker, HumanoidStandup), compared to the original benchmark. Additionally, in our experiments, we used a different task set design (Table 3 ) than those used in the DOMINO and CaDM papers. For Hopper, Walker, Half-Cheetah, Ant, HumanoidStandup, and SlimHumanoid, we used the MuJoCo physics engine environments and implemented settings from Clavera et al. [2019b] and Seo et al. [2020] , scaling the mass of every rigid link by a scale factor m and the damping of every joint by a scale factor d. For Crippled Ant, Crippled Hopper, Crippled Walker, and Crippled Half-Cheetah, we used the implementation available from Seo et al. [2020] , scaled the mass of each rigid link by a factor of m, scaled the damping of each joint by a factor of d, and randomly selected joints to be uncontrollable (i.e., masking the corresponding actions with 0). Generalisation performance is measured in two different regimes: moderate and extreme, where the moderate regime draws environmental features from a range closer to the training range compared to the extreme regime.Our settings for training, extreme, and moderate test tasks are provided in Table 3 . In this section, we provide the implementation details for SaMI. We present the pseudo-code for using SaNCE during meta-training and meta-testing in Algorithms 1 and 2. Our codebase is built on top of the publicly released implementation of Stable Baselines3 by Raffin et al. [2021] and the implementation of InfoNCE by Oord et al. [2019] . A public, open-source implementation of SaMI is available at https://github.com/uoe-agents/SaMI.Base algorithm. We use SAC [Haarnoja et al., 2018] for the downstream evaluation of the learned context embedding. SAC is an off-policy actor-critic method that leverages the maximum entropy framework for soft policy iteration. At each iteration, SAC performs soft policy evaluation and improvement steps. We use the same SAC implementation across all baselines and other methods.During the meta-training phase, we trained agents for 1.6 million timesteps in each environment on the Panda-gym and MuJoCo benchmarks. For meta-testing, we evaluated 100 episodes in each environment, with tasks randomly sampled from the moderate and extreme task sets.Encoder architecture. In our method, the context encoder ψ is modelled as a Long Short-Term Memory (LSTM) network that produces a 128-dimensional hidden state vector, which is subsequently processed through a single-layer feed-forward network to generate a 6-dimensional context embedding. We aim for the agent to complete the three steps of "explore effectively, infer, adapt" within an episode. Therefore, we initialise the hidden state and cell state of the LSTM to zero at the start of each episode. Both the actor and critic use the same context encoder to embed trajectories. For contrastive learning, SaNCE utilises a momentum encoder ψ * to generate positive and negative context embeddings [Laskin et al., 2020 , He et al., 2020 ]. Formally, denoting the parameters of ψ as θ ψ and those of ψ * as θ ψ * , we update θ ψ * as follows:EQUATION)Here m ∈ [0, 1) is a soft-update rate. Only the parameters θ ψ are updated by back-propagation. The momentum update in Equation ( 13) makes θ ψ * evolve more smoothly by having them slowly track the θ ψ with m ≪ 1 (e.g., m = 0.05 in this research). This means that the target values are constrained to change slowly, greatly improving the stability of learning.Hyperparameters. A full list of hyperparameters is displayed in Table 4 .Hardware. For each experiment run we use a single NVIDIA Volta V100 GPU with 32GB memory and a single CPU. for each task e n do 5:for Roll-out time steps do Roll-out policy π cn (a t |s t , c n ) and accumulate transition (s t , a t , r t , s t+1 ); Add trajectory τ cn = {s 0 , a 0 , r 0 , s 1 , r 1 , ..., s T , a θ ψ * ← m • θ ψ + (1 -m) • θ ψ * ; 20:end for 21:end for 22: end while Algorithm 2 SaNCE Meta-testing Require: Batch of training tasks {e n } n=1,...,N from ξ test (e); 1: while not done do 2:for each task e n do 3:for each episode do Roll-out policy π cn (a t |s t , c n ) and accumulate transition (s t , a t , r t , s t+1 ); end for 10: end while E Additional resultsWhile past work has optimised hyperparameters to balance the contrastive loss coefficient α relative to the RL objective [Jaderberg et al., 2016 , Bachman et al., 2019] , we use both the contrastive and RL objectives with equal weight, setting α = 1.0 for the MuJoCo benchmark and α = 0.01 for the Panda-gym benchmark. Additionally, we analyse the effect of the loss coefficient α for CCM, SaTESAC, and SaCCM in the MuJoCo (Figure 12 ) and Panda-gym (Figure 11 ) benchmarks. Buffer sizes are 400,000, 100,000, 10,000, and 1,000. This section presents the heatmap results and further analysis of TESAC, CCM, and SaTESAC on the Panda-gym benchmark. From the heatmap results of SaTESAC and SaCCM (Figure 6 ), we observe that, with higher cube masses (30 and 10 kg), the agent executed the Push skill (indicated by the clustered points within the yellow bounding box in Figure 1(a) ). At lower masses, the agent executed the Pick&Place skill.In contrast, as shown in Figures 17(e-f ), CCM predominantly learned the Pick&Place skill, resulting in a drop in success rates for tasks with mass = 30, as the agent could not lift the cube off the table using the Pick&Place skill, as depicted in Figure 17(d) . The visualisation of the context embedding (Figure 16 ) did not reveal clear skill grouping across different tasks.Finally, TESAC primarily mastered the Push skill. The Push skill is relatively simpler to learn than the Pick&Place skill, as it does not require the agent to manipulate its fingers to pick up cubes. Consequently, TESAC's success rate notably decreased in environments with higher friction.SaMI enables RL agents to be versatile and embody multiple skills. Additionally, video demos (available at https://github.com/uoe-agents/SaMI) provide a better demonstration of the different skills. In this section, we give a brief comparison between our methods and methods from DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020] in the MuJoCo benchmark because we are using the exact same environmental setting (shown in Table 5 ).DOMINO [Mu et al., 2022] is based on the InfoNCE K-sample estimator. Their implementation, PPO+DOMINO, is a model-free RL algorithm with a pre-trained context encoder. This encoder reduces the demand for large contrastive batch sizes during training by decoupling representation learning for each modality, simplifying tasks while leveraging shared information. However, a pre-trained encoder necessitates a large sample volume, with DOMINO training PPO agents for 5 million timesteps on the MuJoCo benchmark. In contrast, SaTESAC and SaCCM, trained for 1.6 million timesteps without pre-trained encoders, achieve considerably higher average returns across four environments (Table 6 ). Therefore, it is crucial to focus on extracting MI in contrastive learning that directly optimises downstream tasks, integrating rather than segregating representation learning from task performance.CaDM [Lee et al., 2020] proposes a context-aware dynamics model adaptable to changes in dynamics. Specifically, they utilise contrastive learning to learn context embeddings, and then predict the next state conditioned on them. We copy the numerical results of PPO+CaDM, Vanilla+CaDM, and PE-TS+CaDM from CaDM [Lee et al., 2020] et al., 2020] : prediction models are more effective when the transition function changes across tasks. Therefore, we consider that a model-based approach to SaMI could be an interesting extension for future work. • The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).38th Conference on Neural Information Processing Systems (NeurIPS 2024).Our modified benchmarks are open-sourced at https://github.com/uoe-agents/Skill-aware-Panda-gymOur code, video demos and experimental data are available at https://github.com/uoe-agents/SaMIIf z is not equal to 0, the Pick&Place skill is always required to solve the tasks.In the original Push task, the grippers are blocked to ensure the agent can only push cubes. However, this restriction prevents the agent from learning Pick&Place skills, leading to "unpushable" failure in Figure1(b). 


Decision: The submission clearly presents a novel method with sufficient empirical evidence for its efficacy. There were some concerns raised in the reviews, but most (if not all) of them were adequately addressed by the authors during the rebuttal.

As such, I am recommending an acceptance.

Reviews: review: paper_summary: paper_summary: The paper presents an interesting way to deal with the "log k curse", that is novel while being sensible. The results seem promising, even though their presentation needs improving, and I believe it can serve as foundation for many future works to build upon. Overall I think the idea and execution are solid, but the paper as an artifact needs work, as I'll point out in detail in the Weaknesses section.
strength_weakness: strength_weakness: I believe the results, and more specifically their presentation is a serious weakness of the paper. In Panda-Gym the authors claim their method "boosts success rates by an average of 20.23% in moderate tests and18.36% in extreme tests", but in fact when taking into account the standard deviation of theirs and previous method performances we *cannot* ascertain whether there was any improvement whatsoever on the moderate and extreme tests. This also points to an abuse of the use of the blue background to "indicates that our method surpasses PEARL, TESAC and CCM". 

Moreover in the MuJoCo results while we in general do see more statistically significant results for Test (Moderate) and Test (extreme) there are still a few misuses of the blue background, and also we in general don't see statistically significant improvement in the training setting, which leads to the question of why do we see such different behaviours in Panda-Gym (non-statistically significant improvement in Moderate and Extreme, but improvement in Training) and MuJoCo (the opposite), a question that would have been good if the authors had explored a little more.

Finally the interpretation of the t-SNE results seem quite problematic, where 2 clear clusters, one on the lower left and another on the lower right, containing context embeddings of the 4 kinds of settings for Panda-Gym seem to be ignored, but an artificial cluster is suggested by using a yellow bounding box. I consider such cluster artificial as  many of the points on the upper part of the bounding seem closer to points outside it then to points inside it, and it's always good to remember that for t-SNE local distances are meaningful, but global ones are not.
questions: questions: + Do the authors have any hypothesis on why the generalization behaviour is different when using SaNCE in Panda-Gym vs MuJoCo?
+ I believe fixing the plots and claims on improvements is fundamental for the paper to be a good research artifact and recommend the authors do so
+ Given what was said about the t-SNE analysis I believe it either should simply be removed, or a different analysis using clustering methods in higher dimensions should be used instead.
limitations: limitations: The paper is mostly about addressing a limitation of previous methods, and I do not believe it has any potential negative societal impact.
paper_summary: paper_summary: The paper proposes an alternative approach to infoNCE for learning contrastive task representations by using the structure of skills in a meta learning setting. In doing so, the algorithm samples more negatives within the same task as the positive (but with lower reward), in a procedure that somewhat resembles hard negative mining. A benefit to this procedure is it is less sensitive to the choice of infoNCE batch size K, which in large multi-task settings is often only able to approximate mutual information for infeasibly large values. Empirical results show this alternative contrastive loss improves performance on RL benchmarks.
strength_weakness: strength_weakness: - There seem to be issues with the mathematical formulation (see questions)
- I found the presentation of the method to be confusing, in particular regarding how the learned representations were actually used for meta learning
questions: questions: - I'm confused by the statement in Lemma 1. The RHS, $\log K  \le I(x;y)$ only depends on the data distribution. In some cases, it seems the mutual information between $x$ and $y$ can be less than $\log K$, regardless of the parameterization used. Is this a contradiction?
  - Similarly, I don't understand Figure 6—K is on the x-axis, but $\log K$ is drawn as a horizontal line.

- From eq. 2, it seems the definition of the SaMI $I_{\mathrm{SaMI}}(c ; \pi_c ; \tau_c)=\mathbb{E}_{p(c, \pi_c, \tau_c)} \{\log \frac{p(\tau_c \mid c, \pi_c)}{p(\tau_c)}\}$ is equivalent to the mutual information $I(\tau_c;(c,\pi_c))$.

  But by the chain rule we know $I(\tau_c;(c,\pi_c))=I(\tau_c;c)+I(\tau_c;\pi_c\mid{}c)$ which contradicts the next statement $I_{\mathrm{SaMI}}(c ; \pi_c ; \tau_c) \leq I(c ; \tau_c)$

- How is the "momentum encoder" parameterized in this setting?

- Is there significance to the bars not lining up in Figure 2 (c)?

- The authors may wish to discuss related work on hard negative mining
limitations: limitations: Addressed
paper_summary: paper_summary: Learning generalizable skills across different tasks is desirable in Reinforcement Learning. Some methods embed task information into a context latent space which is then used to train policies. This paper proposes a new objective Skill-aware Mutual Information (SaMI) that incentivizes the distinction of context embeddings according to skills. If optimized, the agent can then execute a skill depending on the underlying context. In addition to the objective, the paper proposes a K-sample estimator, Skill-aware Noise Contrastive Estimation (SaNCE), that enables an agent to optimize the SaMI objective in a sample-efficient manner.
strength_weakness: strength_weakness: only minor things such as an algorithm box would be useful to understand the learning procedure
questions: questions: - Section 5.3 provides insights into how the positive and negative samples are determined. How prone is this to declare very similar samples, i.e. similar behaviors 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 behavior.

- Could the authors elaborate on the training procedure? 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?
limitations: limitations: Overall I found this a good paper with minor issues/unclear points (see Weaknesses/Questions).
paper_summary: paper_summary: This paper introduces Skill-aware Mutual Information (SaMI) and Skill-aware Noise Contrastive Estimation (SaNCE) to enhance zero-shot generalization in reinforcement learning (RL). The authors address the challenges faced by Meta-Reinforcement Learning (Meta-RL) agents in tasks requiring different optimal skills by using context encoders based on contrastive learning. SaMI distinguishes context embeddings according to skills, while SaNCE, a K-sample estimator, optimizes the SaMI objective, reducing the negative sample space and mitigating the log-K curse. The proposed methods are empirically validated on modified MuJoCo and Panda-gym benchmarks, demonstrating significant improvements in zero-shot generalization and robustness to sample size reductions.
strength_weakness: strength_weakness: ### W1. Scalability Concerns

The scalability of the proposed methods to more complex, large-scale environments is not thoroughly discussed. While the methods show promising results in the tested benchmarks, a broader analysis of their scalability and practical utility in more complex scenarios is needed to understand their full potential and limitations.

### W2. Dependency on Skill Definitions

The success of the proposed methods depends on the accurate definition and identification of skills. The paper does not address potential issues related to the variability in skill definitions across different environments and tasks, which could impact the robustness and generalizability of the methods.
questions: questions: ### Q1. Scalability to Complex Environments

How do the proposed methods scale to more complex, large-scale environments? Are there any specific challenges or limitations that need to be addressed for practical deployment in such scenarios?

### Q2. Variability in Skill Definitions

How does the variability in skill definitions across different environments and tasks affect the proposed methods? Have the authors considered any strategies to mitigate the potential impact of inconsistent skill definitions on the robustness and generalizability of the methods?
limitations: limitations: The authors acknowledge the limitations related to the requirement for accurate skill definitions and the focus on specific benchmarks. However, a more detailed discussion on potential negative societal impacts and strategies to mitigate them would be beneficial.
paper_summary: paper_summary: Proposes a new contrastive learning objective for use in meta-RL with contextual policies. The new objective incorporates a notion of skills into the mutual information estimate. Theoretical and empirical evidence in support of the superiority of the proposed method is presented.
strength_weakness: strength_weakness: - Meta-RL is presented as the problem of learning a contextual RL agent and a context generator for solving the task. However, there are many other kinds of meta-RL algorithms that do not use explicit context representations at all. For some canonical examples, see [RL^2](https://arxiv.org/abs/1611.02779) and [MAML](https://arxiv.org/abs/1703.03400). It is not important to include a thorough survey of other types of meta-RL methods, but for the sake of accuracy, the full scope of meta-RL should be mentioned. Especially important is to note that the proposed method requires a posterior sampling style meta-RL approach, which is not optimal at test-time compared to directly optimizing the meta-RL objective via, e.g., training a sequence model as the policy.
- Due to the above issue, it is inaccurate to say that SaNCE can be integrated with any Meta-RL algorithm.
- A new benchmark environment should come with thorough benchmarking. A good selection of contrastive meta-RL methods are included, but it would be more convincing if at least some RL^2 style method with suitable architecture (e.g. [transformers](https://proceedings.mlr.press/v162/melo22a.html) or [hypernetworks](https://proceedings.neurips.cc/paper_files/paper/2023/hash/c3fa3a7d50b34732c6d08f6f66380d75-Abstract-Conference.html)) was also run as a baseline. The current benchmarking doesn't try zero-shot meta-RL algorithms at all.
- The paper isn't self-contained in terms of understanding how the proposed method works during meta-training and meta-testing times. See questions for details.
questions: questions: - For how many episodes does the agent interact with the environment at meta-training and testing times?
- How does the encoder accumulate across multiple episodes of interaction?
- Is there a way to use the proposed method for adaptation within an episode?
limitations: limitations: Limitations are not discussed.

Comments: title: Response to the official comment from Reviewer hnsz
Response to the official comment from Reviewer hnsz
Response
Response to the official comment from Reviewer cm3M
Acknowledgement
comment: 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 your response. In your first comment, you stated $I_{\mathrm{SaMI}}\left(c ; \pi_c ; \tau_c\right)=E_{p\left(c, \pi_c, \tau_c\right)} \log \frac{p\left(c, \pi_c, \tau_c\right)}{p(c) p\left(\pi_c\right) p\left(\tau_c\right)}$, while in the second you state $I_{\text {SaMI }}\left(c ; \pi_c ; \tau_c\right)=I\left(c ; \tau_c\right)-I\left(c ; \tau_c \mid \pi_c\right)$. The first definition is total correlation while the second is 3-variable interaction information. Is the second statement the correct definition then? I am still confused as to which form is optimized by your objective and the motivation for maximizing interaction information.
**Re “The new statement of eq. (2) seems to be equivalent to the definition of the total correlation”:** We thank the reviewer for pointing out this related work. The eq. (2) $I_{\text{SaMI}}(c;\pi_c;\tau_c)$ is the **interaction information**, and it differs from the **total correlation** in definition. Both interaction information and total correlation are generalisations of mutual information. Specifically, let's first consider the **total correlation** for three variables. For a given set of 3 random variables $\{x,y,z\}$, the total correlation $TC(x,y,z)$ is defined as the Kullback–Leibler divergence from the joint distribution $p(x,y,z)$ to the independent distribution $p(x)p(y)p(z)$. This divergence simplifies to the difference of entropies:
$$TC(x,y,z)=H(x)+H(y)+H(z)-H(x,y,z)$$
where $H(x)$ is the information entropy of variable $x$, and $H(x,y,z)$ is the joint entropy of the variable set $\{x,y,z\}$.

In contrast, the **interaction information** $I(x;y;z)$ for three variables $\{x,y,z\}$ is given by:
$$I(x;y;z)=(H(x)+H(y)+H(z)) - (H(x,y)+H(x,z)+H(y,z)) + H(x,y,z)$$
In Appendix B, “Proof for Lemma,” we refer to it as “interaction information,” drawing from references "Interaction Information for Causal Inference: The Case of Directed Triangle" (https://arxiv.org/pdf/1701.08868) and "Multivariate Information Transmission" (https://ieeexplore.ieee.org/document/1057469). There are several other names for Interaction Information, including the amount of information [1], information correlation [2], co-information [3], and simply mutual information [4].
To make eq. (2) clearer, we have moved the definition of Interaction Information into the main text. Therefore, we directly defined $I_{SaMI}$ using the Interaction Information definition:
$$I_{SaMI}(c;\pi_c;\tau_c)= I(c;\tau_c)-I(c;\tau_c|\pi_c)$$
Therefore, we proposed a lower bound $I_{\text{SaMI}}(c;\pi_c;\tau_c)$ to approximate $I(c;\tau_c)$. We believe that other generalisations of mutual information (i.e., MI estimators) could also be helpful in approximating $I(c;\tau_c)$.

**Re "the relationship between the updated eq. (2) definition of $I_{SaMI}$ and estimators such as Bai et al. (2023)":** TC estimators [5] is an interesting study, we fully agree with Bai et al.'s perspective [5] that when calculating mutual information among multiple variables, existing work often seeks ways to decompose interaction information, which leads to too much strong independence assumptions about the data distribution. For instance, as mentioned in [5], some studies make a strong assumption that all variables are conditionally independent [6].  Without such strong independence assumptions, TC estimator [5] is a method of decomposing interaction information and assumes that the relationships between variables follow a tree-like or line-like structure, thus applying tree-like or line-like decomposition to the interaction information, resulting in tree-like or line-like MI estimators (i.e., TC estimators). 

Our proposed SaMI approach does not make any assumptions about the data, and therefore does not decompose the interaction information. This is primarily because we focus on mutual information among just three variables in the context of the RL problem, which is significantly simpler than the scenarios addressed in [5]. Additionally, through extensive experiments, we have verified that maximising the mutual information $I_{SaMI}$ significantly enhances the zero-shot generalisation ability of reinforcement learning.

We thank the reviewer again for noting the typo in our original eq(2) and helping us to refine Lemma 1. We would be grateful if the reviewer could reevaluate their score in light of our given clarifications.


**References**

[1] Ting, Hu Kuo. "On the amount of information." Theory of Probability \& Its Applications 7.4 (1962): 439-447. 

[2] Wolf, David R. "The Generalization of Mutual Information as the Information between a Set of Variables: The Information Correlation Function Hierarchy and the Information Structure of Multi-Agent Systems." (2004).

[3] Bell, Anthony J. "The co-information lattice." Proceedings of the fifth international workshop on independent component analysis and blind signal separation: ICA. Vol. 2003. 2003.

[4] Yeung, Raymond W. "A new outlook on Shannon's information measures." IEEE transactions on information theory 37.3 (1991): 466-474.

[5] Bai, Ke, et al. "Estimating total correlation with mutual information estimators." International Conference on Artificial Intelligence and Statistics. PMLR, 2023.

[6] Poole, Ben, et al. "On variational bounds of mutual information." International Conference on Machine Learning. PMLR, 2019.
Thank you for your response. The new statement of eq. (2) seems to be equivalent to the definition of the *total correlation* (Watanabe, 1960). Could you comment on the relationship between the updated eq. (2) definition of $I_{\\mathrm{SaMI}}$ and existing information-theoretic quantities like total correlation (and estimators such as Bai et al. (2023))?


### References
Bai, K., Cheng, P., Hao, W., Henao, R., & Carin, L. (2023). Estimating total correlation with mutual information estimators. _Proceedings of The 26th International Conference on Artificial Intelligence and Statistics_, 2147–2164. [https://proceedings.mlr.press/v206/bai23a.html](https://proceedings.mlr.press/v206/bai23a.html)

Watanabe, S. (1960). Information theoretical analysis of multivariate correlation. _IBM Journal of Research and Development_, _4_, 66–82. IBM Journal of Research and Development. [https://doi.org/10.1147/rd.41.0066](https://doi.org/10.1147/rd.41.0066)
The new UMAP plot indeed shows the clustering you allude to. Given that you addressed my comments I'll increase my rating
I highly appreciate the authors' responses that clarify my questions. The responses confirm my initial assessment in recommending accepting this paper.
Thank you for your response. The table below provides the requested 5 values for the performances of SaTESAC and TESAC showing the average success rate ± standard deviation (on 8 extreme test tasks):

|             | seed 1 | seed 2  | seed 3 | seed 4  | seed 5 |
| -------- | ------- | ------- | ------- | ------- | ------- |
| TESAC  | 0.19±0.13 | 0.20±0.26 | 0.25±0.19 | 0.18±0.14 | 0.28±0.31 | 
| SaTESAC | 0.37±0.37 | 0.37±0.35 | 0.36±0.36 | 0.36±0.34 | 0.38±0.35 |

These 5 values in the table demonstrate that *the improvements made by SaTESAC over TESAC in extreme test tasks are significant*, which aligns with our paired t-test results. The paired t-test indicates that *the improvements made by SaTESAC over TESAC in all three settings (Training, Moderate, and Extreme) are statistically significant at a significance level of 0.05*. The paired t-test results are shown in the table below (you can also find these results in Appendix H of the updated paper at: https://anonymous.4open.science/r/SaMI/SaMI.pdf):

|                   | Training  | Moderate  | Extreme   |
|-------------------|-----------|-----------|-----------|
| t-statistic       | 12.180000 | 13.800000 | 8.020000  |
| p-value    | 0.000260 | 0.000160 | 0.001310 |

Please find all the results for TESAC/SaTESAC/CCM/SaCCM across multiple seeds used to generate Table 1 and Table 2 in our anonymous GitHub repository (https://anonymous.4open.science/r/SaMI/data/MuJoCo.xlsx and https://anonymous.4open.science/r/SaMI/data/Panda-gym.xls). Please refer to the "success rate \& t-test" sheet in the file “data/Panda-gym.xls” to find all experimental results related to Panda-gym. The content of the table above can also be found in the "Mean Success Rate-Multiple Envs" sheet.

Thank you for your comment on the projection method. We have added **UMAP** (https://arxiv.org/pdf/1802.03426) as an alternative projection method, which is similar to t-SNE but more efficient and tends to better preserve the global structure of the data than t-SNE (i.e., more clearly separates groups of similar categories from each other). We have replaced the t-SNE visualisation in Figure 6 with UMAP in the updated paper. In the new context embedding, two clear clusters are visible, corresponding to Push and Pick\&Place, respectively. Please find the updated Figure 6 at https://anonymous.4open.science/r/SaMI/data/figure_6.png

Finally, to better understand what the method is doing, combining the newly added UMAP visualisations (in Appendix F of the updated paper at: https://anonymous.4open.science/r/SaMI/SaMI.pdf) with Heatmap visualisations (provided in the original manuscript in Appendix F), and video demos (in our anonymous GitHub repository https://anonymous.4open.science/r/SaMI) would be helpful. We have provided skill analysis in the original manuscript in sections 5.2 and 5.3 under the “Results and skill analysis” paragraph, which you can also find on page 9 (Panda-gym environment) and page 10 (MuJoCo environment) of the updated manuscript (https://anonymous.4open.science/r/SaMI/SaMI.pdf).

Please let us know if further clarification is required.
Thank you for the response.

I don't think it is fair to describe CCM as a zero-shot meta-rl method as it explicitly uses a different policy for first collecting exploration episodes and then runs the exploitation policy given the context encoding.

My concerns have been mostly addressed so I'm raising my score.
I appreciate the authors for the detailed response, especially the video demos, which seem to be very effective.

Most of my concerns are addressed, therefore I'm raising my assessment.
Thank you for your response and the changes made to the paper.
Could you tell me the 5 values for the performances of SaTESAC and TESAC on the Panda-Gym extreme test setting? I'm somewhat confused by the t-test results.
I still do not understand how you can call the region inside the yellow box a cluster in T-SNE, as there are multiple points with x < -10 who look like they should be part of the so-called cluster, but which would be pretty bad for the push skill, which end up just making me doubt a little whether the method is truly doing what the authors claim it is doing.