Skill-aware Mutual Information Optimisation for Zero-shot Generalisation in Reinforcement Learning
Authors: Xuehui Yu
Mhairi Dunion
Xin Li
Stefano V Albrecht
Venue: NeurIPS 2024 poster (2024)
Keywords: contrastive learning
reinforcement learning
meta reinforcement learning
zero-shot generalisation
Primary area: reinforcement_learning

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.

Skill-aware Mutual Information Optimisation for Generalisation in Reinforcement Learning
Xuehui Yu!
Mhairi Dunion
Xin Li Stefano V. Albrecht
Harbin Institute of Technology University of Edinburgh
(yuxuehui (2ss03169))sttu hit edu. en (mhairi. dunion, s. albrecht]@ed ac. ac.uk
Abstract
Meta-Reinforcement Learning (Meta-RL) agents 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 the requirement for large sample size, also referred to as the log-K 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 --ampl 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 reduction in the number of available samples, thus possessing the potential to overcome the log-K curse.
Introduction
(a)
(b)
Reinforcement Learning (RL) agents often learn policies that do not generalise across tasks
Unpushable (High Friction)
Pushable (Low Friction)
in which the environmental fea- tures and optimal skills are dif- ferent [des Combes et al. 22188] Garcin al.) 2024). Consider a set of cube-moving tasks where an agent required to move a cube to goal position on table (Figure M] These tasks become challenging if environmental fea- tures, such as table friction, vary between tasks. When facing an unknown environment, the agent needs to explore effectively, understand the environment, and adjust its behaviour accordingly within an episode. For instance, the agent tries to push a cube across a table covered by tablecloth and finds it "unpushable," it should infer that the table friction is relatively high and adapt by lifting
Goal Position
Environmental features
Mass, Friclion, Material,
Eection
Lifable (Low Mass)
Unliftable (High Mass)
Figure 1: (a) In cube-moving environment, tasks are defined according to different environmental features. (b) Different tasks have different transition dynamics caused by underlying environ- mental features, hence optimal skills are different across tasks.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the cube to avoid friction, rather than continuing to push. Recent advances in Reinffrerrrrem
Learning (Meta- RL) [ee al. 2020 Agarwal al. 2021 Mu al. 2022 Dunion et al., 2023b
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. [202]]
al. [022]
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-RI. 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 Pick& 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. 202I] and the sample limitations in zero-shot generalisation make collecting 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 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 fur- ther 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), generalised form of MI objective between context embeddings, skills, and trajectories, designed to address issue (i). We provide 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 K-sample estimator, Skill-aware Noise Contrastive Estimation (SaNCE) to optimise SaMI, effectively addressing issue (ii). Additionally, (3) we propose 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 of skills applicable 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 al. 2020 Fu et al. T by achieving higher returns and success rates 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 RR agents 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.
2 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. [2033 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 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 shared embedding space, aiming to maximise the average returns across tasks. However, the update signals from the RI. loss are stochastic and weak, and may not capture the similarity relations among tasks [Fu et al.)

---

at
+++
af+7
FE+T
Tt (t))
Push
a (c2)
C
at
r
SE
2i+7
aC+t
TE+T
Tc2
Pick&Place
Figure 2: A policy conditioned on fixed context embedding e is defined skill x(- (shortened as πT). The policy conditioned on a fixed alters the state of the environment in consistent way, thereby exhibiting a mode of skill. The skill π (-lc1) moves the cube on the table in trajectory T2) and is referred to as the Push skill; correspondingly, the Pick&Place skill x x--(02) takes the cube off the table and places it in the goal position in the trajectory T
2021]. The second category involves learning context embeddings through dynamics prediction
[Lee al., 2020 Zhou et al. 2019) which can make the context embeddings noisy, as they may
model irrelevant dependencies and overlook task- specific information [uu 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. [019 Dunion and Albrecht 2024]. ease of implementation [Oord et al. [019]] and strong theoretical connection MI estimation [Poole et al., 2019]. MI is often estimated using InfoNCE FOrd 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 [uu et al. 2021] uses large number of samples for maximising InfoNCE. DOMINO [Mu et al., 2022] reduces the true MI by introducing independence assumption; however, this results in biased estimates. We focus on proposing an unbiased alternative MI objective and 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. Meta-RI,, we assume an environment a distribution E(e) of tasks e (e.g. uniform in our experiments). Each task ~ {(e) has similar structure that corresponds to Markov Decision Process (MDP) [Puterman 2014), defined by M, (S, A, R, Pe. 7), with a state space S, an action space A, reward function R(s, dz) where s, S and a E A. state transition dynamics (s dt), and a discount factor 7 E [0, 1). In order to address the problem of zero- shot generalisation, we consider the transition dynamics P. (s | az) vary across tasks ~ E(e) according to multiple environmental features e = fe²; e²_ ...., e 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 's material. For instance, in a cube- srbbinnovigg environment (Figure [], 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 sampled from Errin (e),
aim to learn a policy that can maximise the discounted returns, arg max, ş Ee-gmm(() (e) 2 7 R(s,, az) (a) ~ (Tarls)); 8241 ~ Pe (8 a;)], and can produce accurate control for unseen test tasks sampled from (ue((e).
Contrastive learning. In Meta-RL, the context encoder (c()TT.))) first takes the trajectory Te.0:
(8o, ag To. ..., } from the current episode as input and compresses it into a context embedding e [Fu et al.] 202T]. Then, the policy π, conditioned on context embedding c, consumes the current state 8, and outputs the action a, 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 context encoder that optimises the InfoNCE objective JIntONCE y), which is -ample estimator and lower bound of the MI [(x; y) (Oord et al. 2019). Given a query and set Y {uI,-.. ..... UK² of K random samples containing one positive sample y1 and K - negative samples from the distribution p(y), IIttNCC y) is obtained by comparing pairs sampled
from the joint distribution x. V1 ~ p(r, y) to pairs æ, Uk built using a set of negative examples V2:K:
fo v1)
I (z; v.d, K) E log
E fe yk) K
1
InfoNCE constructs a formal lower bound on the MI, i.e., IINNNCu((; yoo, K) S. [(::y) [Guo et al.
2022
Chen al., 2021). Given two inputs and y. their embedding similarity is fz y)
e: (at)))
-v(M/p, where v is the context encoder that projects x and into the context embedding space.
the dot product is used to calculate the similarity score between ()) pairs [Wu et al. 2018
He et al. 2020], and /8 is a temperature hyperparameter that controls the sensitivity of the product.
Some previous Meta-RL methods [Lee et al. 2020 Mu et al. [022] learn context embedding c by maximising /InoNCC (c; Talth, K) between the context с embedded from a trajectory in the current task, and the historical trajectories Te under the same environmental features setting.
4
Skill-aware mutual information optimisation for Meta-RL
4.1 The log- K curse of K -ssmplee MI estimators
In this section, provide theoretical analysis of the challenge inherent in learning K K-sample estimator for MI, commonly referred to as the log- curse. Based on this theoretical analysis, we give insights to overcome this challenge. Given that focus the generalisation of RL, we only consider cases with a finite sample size of K. If a context encoder tb in Equation (I) has sufficient training epochs, then InfoNCC (#; y/ K) log K [Mnih and Teh 2012 Guo al. 2022]. Hence, the MI we can optimise is bottlenecked by the number of available samples, formally expressed as: Lemma 1 Learning context encoder v0 with a K- sample estimator and finite sample size K, we have Inffocee( (2;; y) b) K) < log K (x((x;)) y), when + y (see proof in Appendix[A].
Mutual Informatiin
Isaiis Re ke)
logk
We do not consider the case when x 11 y. i.e., log K 2 ((x;y) = (vK 1), because a Meta- RL agent learns a context encoder by maximising MI between trajectories To and context embeddings which are not independent (as shown in Figure
Sanca (e
2.. While an unbounded sample size K for learning
effective context embeddings is theoretically feasi-
ble and assumed in many studies Mu al.) 2022]]
Training Epochs
Lee et | 2020], it is often impractical in practice. Therefore, building on Lemma [ we derive two key
Figure 3: AntoNCE(es ). with finite sample size of K, is a loose lower bound of [[c; Ta) and leads to lower performance embeddings. Isami (c; mei Ta) is a lower ground-truth MI, and TSNNCE (cy Maj Ta) is a tighter lower bound.
insights with limited samples: (1) generalising the
current MI objective to be smaller than I(a;; y) (see
Section 4.2p ; (2) developing a K K-sample estimator
tighter than InfONCE (see Section 4.3)
4.2 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 (Skills) A policy π conditioned on a fixed context embedding c is defined as a skill (ccc), abbreviated as ne. If a skill Tee is conditioned a state st. can sample actions a (-1c,,8 8t). After sampling actions from ne at consecutive timesteps, we obtain a trajectory Te,titr = {se, 0t, Tt, 8141,..., StTT, Ot+T, l14r) which demonstrates consistent mode of behaviour.
After a limited amount of exploration, an agent should be able to infer the task (i.e., environmental features = ee,e! ,...,, N and adapt accordingly within the current episode. The context em-

---

Tc
R2
Rc
F
Fc2
f
T2
fG
Positive Sample Space
Fcz
F22
Td
F22
Negative Sample Space
E2²
Fiy
Not Sample Space
Tc
12
Observed Variable
(a) InfoNCE
(b) SaNCE
(c) Sa+InfoNCE
Figure 4: A comparison of sample spaces for task Positive samples To, or T2i are always from current task é1. For SaNCE, in a task e₁ with embedding cb. the positive skill T conditions on ck and generates positive trajectories T and the negative skill Tea generates negative trajectories Tep The top graphs show the relationship between e, Te and Tc-
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 eb should be trained by maximising the MI between the context embedding c, skills Tes and trajectories Te- To this end, propose a novel MI optimisation objective, Skill-aware Mutual Information (SaMI), defined as:
[saMi(c; To) = [c; - (c; Tel
(2)
SaMI defined according interaction information [McGill 1954). serving as a generalisation of MI for three variables {c. Te, Je). Although we cannot evaluate p(c, Ne, Te) directly, we approximate by Monte-Carlo sampling, using K samples from p(c, TT. Te). As illustrated in Figure context encoder &b trained with the objective of maximising /saMi Ne; T2) converges more quickly, as (s; nel TetTe) I [c;; Te) (see proof in Appendix B). By focusing more on skill-related information, ISOMI enables agents to autonomously discover diverse range of skills for handling multiple tasks.
4.3 Skill-aware noise contrastive estimation: a tighter K K-sample estimator
Despite InfoNCE's success as 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., J S log K S SaMI [Chen al. 202T (see proof Appendix B). When K + +00, we can expect IINONCE: R log K = Isamı [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 K-sample estimator that requires significantly smaller sample size K < +o0. First, we define K*:
Definition (K*) lcl |mrl M is defined as the number of trajectories in the replay buffer (i.e., the sample space), in which (c) represents the number of different context embeddings [Tel represents the number of different skills ne, and M is natural number.
To ensure that JINANCE is a tight bound of ISaMI, we require that JINNNCE log K Rs IsaMI when K + K" Under the definition of K", the replay buffer can be divided according to the different context embeddings e and skills Tee (i.e., observing context embeddings c and skills π.). In real- real-world robotic control tasks, the sample space size significantly increases due multiple environmental features e = {e e²,.... N J. Taking the sample space of InfoNCE as an example (Figure a)), in the current task e₁ with context embedding C1, positive samples are trajectories Te, generated after executing the skill ne, in task ej, and negative samples are trajectories {Te.....} from other tasks {e...... The permutations and combinations of N environmental features lead to exponential growth in task number |c), which in turn results an increase of sample space K = |c] |T| -M.
We introduce tight K K-sample estimator, Skill-aware Noise Contrastive Estimation (SaNCE). which is used approximate (c n.ं Te) with K KA < KintaNCE For SaNCE, both positive samples - and negative samples ग. are sampled from the current tasks e but are generated by executing positive skills ग.. and negative skills गे respectively. Here, positive skill is intuitively defined by whether optimal for the current task e, with more formal definition provided in Section 4.4 For instance, in cube-moving task under large friction setting, the agent executes
5

---

Negutive Ta (with mús Zird R()
Negathre embedeings gr asd)
Ramked X5e Ris
Push
Momentum Context Encodar V
samplex sh. (wiic max 2 R(5,5,))
SANCE Lass
Postlive embueddings
v'(r())
ESINCE
Positive Samples 12.
Raplay Buffer for oube-moving look with high frictio
(with max Mi. RGr
Contoxt Encoder
Quertes
Pick&Placa
Reinforcement Learning
RL
CRL
Figure 5: A practical framework for using SaNCE in the meta-training phase. During meta-training, we sample trajectories from the replay buffer for off-policy training. Queries are generated by context encoder v. which is updated with gradients from both the SaNCE loss CsaNCE and the RL loss CRI- Negative/Positive embeddings are encoded by momentum context encoder 1f", which is driven by a momentum update with the encoder a). During meta-testing, the meta-traimg context encoder e embeds the current trajectory, and the RL policy takes the embedding as input together with the state for adaptation within an episode.
a skill ग. after several iterations of learning, and obtains corresponding trajectories 7. where the cubes leave the table surface. This indicates that the skill क is Pick&Place and other skills xę may include Push or Flip (flipping the cube to the goal position), with corresponding trajectories Tæ where the cube remains stationary or rolls on the table. Formally, we can optimise the K-sample lower bound ISACCE to approximate Isami:
[saNCE(;; Toi Tolb,,
K- (c1, Teg T log (e, Teg v(t) + Ek= (et, Tea Ter
=I p(e it Iol (
S Isaai (15 Tei
where (c1 Tey Tep (F - (T.))/8 The query G ( generated by the context encoder v. For training stability, we use momentum encoder x/" to produce the positive and negative embeddings. SaNCE significantly reduces the required sample space size KŠANCE by sampling trajectories based different skills Te (Figure 4b) in task e1, so that KAANCE
|c| [el .M
| M S KIANNC (cc = |cil 1). Therefore, [saNCE satisfies Lemma 2
Lemma 2 With context encoder ip and finite sample size K, we have InfoNCE (c; n.; Tel db K) S
[sawe[(s; Tei Tei Tal b K) < log K < IsaMr(C; Te) S [[( (see proof in Appendix B
SaNCE can be used alone or combined with other optimisation objectives train context encoders in Meta-RL algorithms. For instance, integrating SaNCE with InfoNCE diversifies the negative sample
space, with Sa+InfaNC:: = E | + | M. The sample space for /SatiiiaNCE is depicted in Figure (e) and further analysed in detail in Appendix C
4.4 Skill-aware trajectory sampling strategy
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 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 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, only identify whether the skills
slight difference in states can

---

1.0
0.1
0.9
0.1
0,1
0.8
Fitction values
1.0
1.0
0.9
0.1
Frietion values
0.2
0.2
1.0
Inition values
0.8
0.0
0.6
5.0
0.9
0.7
8.3
0.1
10,0 5.0
0.2
0.2
0.2
0.0
10.0 5.0
0.8
0.7
0.4
0.0
10.0
0.9
0.6
0.1
0.8
0.8
0.7
0.0
30
3000
30,0
5
10
30
5
10
30
5
30
30
0.0
Mass values
Mess valuee
Mass values
(a) UMAP visualisation
(b) Success rate
[c] Push probability
(d) Pick&Place probability
Figure 6: (a) UMAP visualisation of context embeddings for the SaCCM in the Panda- gym envi- ronment, with points in the yellow box representing the Push skill in high-mass tasks. Heatmap of (b) success rate, (c) Push skill probability, and (d) Pick&Place skill probability for SaCCM. In large-mass scenarios, the Push skill is more likely be executed than Pick&Place.
optimal; in this way, under a multi-task setting, the agent will acquire of general skills that are applicable to many tasks.
In given task e, positive skills oo are defined as optimal skills achieving highest return Eist a;), whereas negarive skills aे are those that result in lower returns. As a result, the positive sample Tt 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 then minimised 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 7 with high returns, while the owwer-ranked trajectories represent negative samples T, with low returns. However, at the beginning of training, it is likely that all trajectories have low returns. Therefore, our SaNCE loss is soft variant of the K -ssamll SaNCE:
LsaNCE: = max (IIv(T="), (T. 1l1,, ISaNCE:
where ||-||z2 represents the Euclidean distance Tabak 2014). Figure provides Provides practical framework of SaNCE, with a cube-moving example task é1 under high friction. In task e1, the positive skill क.. the Pick&Place skill, which is used to generate queries ()(+)t ) and positive embeddings v0" (r after executing Push skill we get negative samples Tāj and negative embeddings 1h" (TTT)-
5 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 and and Section 5,4
5.1 Experimental setup
Modified benchmarks with multiple environmental features We evaluate our method on two benchmarks, Panda-gym [Gallouèdec et al., 2021] and MuJoCo Todorov al. 2012 (details in Sections 5.2 and 5.3) The benchmarks are modified to be influenced by multiple environmental features, which are sampled the start of each episode during meta-training and meta-testing. During meta-training we uniform- uniform-randomly select a combination of environmental features from a training task set. At test time, we evaluate each algorithm unseen tasks with environmental features outside the training range. Generalisation performance is measured in two different regimes: moderate and extreme. The moderate regime draws environmental features from closer range to the training range compared to the extreme. Our results report the mean and standard deviation of models trained over five seeds in both training and test tasks. Further details are available in Appendix D
OOu modified benchmarks open- sourced at http://ggihubbcoomuueeeeeettttte Paw.eent///kllllaaaaaaaaaaaaga Panda-gym
7

---

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 LRL. which is the same across all methods, corresponds to the RI. 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 202T) which utilises InfoNCE, and TESAC (Yu cct 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 (ct 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 al. 2020 using the same environmental setup in the MuJoCo benchmark.
Our method We use Soft Actor-Critic (SAC) (Haarnoja et al. 2018 as the base RL algorithm.
H
training agents for 1.6 million timesteps in each environment (details Appendix (D.3). SaNCE is a simple objective based on MI that can be used to train any context encoder. We integrate SaNCE into two RL. algorithms: (1) SaTESAC is TESAC with SaNCE, which uses SaNCE for contrastive learning, using |c| times smaller sample space (Figure 4b)); (2) SaCCM is CCM with SaNCE, where the contrastive learning combines InfoNCE and SaNCE, as shown in Figure
5.2 Panda-gym
Table 1:
Comparison of success rate + standard devia- tion with baselines in Panda-gym (over seeds). Bold text signifies the highest average return. * next to the number means that the algorithm with SaMI has statisti- cally significant improvement over the same algorithm without SaMI. All significance claims based on paried t-tests with significance threshold of p 0.05.
Task description. Our modified Panda- gym benchmark involves a robot arm con- trol task using the Franka Emika Panda [Galouédee et al. 202T, where the robot moves cube to target position. Unlike previous works, we simultaneously mod- ify multiple environmental features (cube mass and table friction) that characterise the transition dynamies, and the robot can flexibly execute different skills (Push and Pick&Place) for different tasks. This en- vironment requires high skill diversity; for instance, the agent must use Pick&Place in high-friction tasks and Push in high-mass tasks.
Training
Test (moderate)
Test (extreme)
PEARL TESAC CCM
0.42.10.10 19
0.10:10.06
0.11:10.05
0.50:10.22
0.31:10.20
0:22:10.21
0.
0.49:1:.23
0.29:1:028
SaTESAC
0.92:10.04
0.56110.24
0.37:10.34
SaCCM
0.93:10.05'
0.57:10.26*
0.36:10.35*
Results and skill analysis. As shown in Table I SaTESAC and SaCCM achieve superior general- isation performance compared to PEARL, TESAC, and CCM, with smaller sample space. The t-test results in Table [ that SaMI significantly improves success rates across training, moderate, and extreme test sets at a 0.05 significance level. Video demos show that agents equipped with SaMI acquired multiple skills (Push, Pick&Place) to handle various tasks. When faced with an unknown task, the agents explore by attempting to lift the cube, infer the context, and adjust their skills accordingly within the episode. We visualised the context embeddings using UMAP (McInnes] et al. 2020] (F) and Appendix F and t-SNE [Van der Maaten and Hinton| 2008 (Appendix D. plotting the final step from 100 tests per task. Skills were identified through contact points between the end effector, cube, and table (see Appendix D for more details), and hcatmaps [Waskom] 202T] were used to visualise exccuted skills. Figure 6 shows that SaCCM agents Iearned 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.T.p. 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
5.3 MuJoCo
Task description. We extended the modified MuloCo benchmark introduced in DOMINO Muu ctt al., [202]] and CaDM [Le 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.c., mass, damping) that affect transition dynamics. Crippled Ant, Crippled Hopper, Crippled Walker, and
"Orr code, video demos and experimental data are available at https//giitbbbcom/ooc-aeaaa
8

---

Table 2: Comparison of average return + standard deviation with baselines in modified MuJoCo benchmark (over 5 seeds). Bold number signifies the highest return. * next to the number means that the algorithm with SaMI has statistically significant improvement over the same algorithm without SaMI. All significance claims based on t-tests with significance threshold of p < 0.05.
Crippled Ant
Crippled Half-cheetah
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
PEARL TESAC
1682-173
996.1.21
888
1-93
69%.15488
7462
2139=900
1952-::40
1048:1124
3967:1-955
874:19011
84641849
CCM
2361=114
2047
1527=301
3481::488
821#575
873#914
SaTESAC
2638::406
2379::528
2131=1132"
4528=11922
1143:1:664
1540=10944
SaCCM
2355
2310::314
2007:1-88*
4478:11131
1007:1568
1027:1882
Ant
Half-chectah
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme))
PEARL
51531.581
3873:1.235
3802:1.09
5802.1.73
219011970
1346.1692
TESAC
6789+-451
4705:+-279
4108==69
6298:22310
3173::1210
1159:
CCM
6901
5179:1902
4900::696
6955-1.888
392
1325:1269
SaTESAC
73141:545
5513.16648
4940.15311
7430 1026
4058.1890
1780.1 102"
SaCCM
7378:5339
5717-1-488
5215=377
7154:1965
3849:1689
1926:1-218'
SlimHumanoid
HumunoidStanduu
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
PEARL
6947:1.3541
2697:1:674
2018:1907
95456:1 13445
63242:4 13546
(N224::1546
TESAC
8437.11198
1989:1130)
37601:00
15838411 14455
153944. 15046
74220.1
CCM
7996::1907
5784-1:31
1887
146480-+33345
154601#1629
94991 15258
SaTESAC
102161:1620
7883
61231:1403*
178142.: 10081*
168337 12123
133335 1. 24607"
SaCCM
9312+705
7430= 1587
6473--2001
179930=19338
181033=14626
141750==27426
Hopper
Crippled Hopper
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
PEARL
934-12222
799-4-998
2387-1656
456:1-335
TESAC
1492:159
14991:35
1459:172
3575
3298 1:551
122:1:161
CCM
1484=54
1446::64
1452::58
3455++301
3409::239
1009:4289
SaTESAC
1502=20
14534:39
1447:114
3391=84
3262:1166 166
1839:1130
SaCCM
1462.145
14621:14 14
1451:167
3449 1.491:103
3390:1.211
205911220
Walker
Crippled Walker
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
PEARL
7524-12455
25555
1984
1899112532
4377.1
2965:1: 1426
TESAC
7747:11772
1335
2581::407
1908
5929-+1997
3041
CCM
813611557
5476.1803
251911682
10317:11137
6233.1:
3098.1821
SATESAC
8675-1752
5840-1676
3632:11404
10389-1103
8387 12911
4280-1448*
SaCCM
83611-586
5779:1691
3481:1332*
10496:951
8235:t: 1212
4824:1839'
Crippled Half-cheetah are more challenging due 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 crippled).
Results and skill analysis. Table 2 shows the average return of our method and baselines on training and test tasks. SaTESAC and SaCCM achieved higher returns most tasks, except for Ant, Half-Cheetah, and Hopper, where only single skill was needed. For instance, the Hopper robot learned hop forward, adapting to different mass values. When environments become complex and require diverse skills for different tasks (Crippled Ant, Crippled Hopper, Crippled Salf-Cheetaah, SlimHumanoid, HumanoidStandupp and Crippled Walker), SaNCE brings significant improvements. For example, when the Ant robot has or 4 legs available, it learns to roll to generalise across varying mass and damping. In more challenging zero-shot settings, when only 2 legs are available, the ant robot can no longer roll and it adapts by walking using its healthy legs. This aligns with the results Table 2, where SaMI significantly improved performance in extreme test sets. summary, i) SaMI helps the RL agents to be versatile and embody multiple skills; ii) SaMI leads to increased returns

---

during training and zero-shot generalisation, especially in environments that require different skills.
Our video demos and visualisation results in Appendix F.2 show different skills in all environments.
5.4 Analysis of the log-K curse in sample-linited scenarios
(a)
This section analyses whether SaNCE can overcome the log- K curse. During training, environmental features are sampled at the start of each episode, requiring the context encoder to learn the context embedding distribution across multiple tasks. Since InfoNCE samples negative examples from all tasks and SaNCE samples from the current task, SaNCE's negative sample space is |c] times smaller than InfoNCE's. For instance, in the SlimHumanoid environment, where both mass and damping have five values, InfoNCH's 's sampling space can be 25 times larger than SaNCE's. shown in Tables [ 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.
12000
10000
Average Return
8000
4000
TESAC
CCM
2000
SaTESAC
SaCCM
1000
10000
100000
400000
(b)
Buffer Size
11000
Average Retur
10000
The number of negative samples K is influenced by two hyper- parameters: 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 SaTEE- 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 batch size of 12. Results across all environments (refer to Appendix E..) show that SaNCE exhibits low sensitivity to K, highlighting its potential to overcome the log- K curse.
8000
7000
CCM
6000
SaTESAC
SaCCM
5000
512
Contrastive Batch Size
Figure 7: Effect of (a) buffer size (TESAC, CCM, SaTESAC, SaCCM) and (b) contrastive batch size (CCM. SaTESAC, SaCCM) in the SlimHumanoid environment.
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 practical skill-aware trajectory sampling strategy. Experimental results showed that RL algorithms equipped with SaMI achieved state-of- the-art performance MuloCo 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 significantly smaller negative sample space
within the current episode. SaNCE's ' optimisation 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 cube's 's material correlating with friction and mass, SaMI does not introduce independence assumptions like DOMINO [Mu jet 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.
10

---

References
Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint arXiv:2101.05265 2021.
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saun- shi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229 2019.
Philip Bachman, Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. Advances in neural information processing systems, 32, 2019.
Junya Chen, Zhe Gan, Xuan Li, Qing Guo, Liqun Chen, Shuyang Gao, Tagyoung Chung, Yi Xu, Belinda Zeng. Wenlian Lu, et al. Simpler, faster, stronger: Breaking the log-k curse on contrastive 152, 2021.
learners with flatnce. arXiv preprint
Ignasi Clavera, Anusha Nagabandi, Simin Liu, Ronald S. Fearing, Picter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt dynamic, real-world environments through meta-reinforeement learning. In International Conference on Learning Representations, 2019a.
Ignasi Clavera, Anusha Nagabandi, Simin Liu, Ronald S. Fearing. Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning adapt dynamic, real-world environments through meta-reinforcemmen learning. In International Conference on Learning Representations, 2019b.
Remi Tachet des Combes, Philip Bachman, and Harm van Seijen. Learning invariances for policy generalization, 2018. URL htttps://ppenre/// http://pporrrrees.e.eeeorrorooo http://oorummdd-BBBBBR--
Mhairi Dunion and Stefano V Albrecht. Multi-view disentanglement for reinforcement learning with multiple cameras. In Reinforcement Learning Conference, 2024.
Mhairi Dunion, Trevor McInroe, Kevin Sebastian Luck, Josiah P. Hanna, and Stefano V Albrecht. Conditional mutual information for disentangled representations in reinforcement learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023a.
Mhairi Dunion, Trevor Mclnroe, Kevin Sebastian Luck, Josiah P. Hanna, and Stefano V Albrecht. Temporal disentanglement of representations for improved generalisation in reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023b.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine Diversity is all you need: Learning skills without a reward function, 2018. URL httppp://sapi...tttttt hwt. abs///bs//// 1802. 06070
Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. Advances in Neural Information Processing Systems, 2022.
Jörg K.H. Franke, Gregor Koehler, André Biedenkapp, and Frank Hutter. Sample-efficien automated deep reinforcement learning. In International Conference on Learning Representations, 2021.
Haotian Fu, Hongyao Tang, Jianye Hao, Chen Chen, Xidong Feng, Dong Li, and Wulong Liu. Towards effective context for meta-reinforcement learning: approach based on contrastive learning. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 7457-7465, 2021.
Quentin Gallouédec, Nicolas Cazin, Emmanuel Dellandréa, and Liming Chen. panda-gym: Open- Source Goal-Conditioned Environments for Robotic Learning. 4th Robot Learning Workshop: Self-Supervised and Lifelong Learning at NeurIPS, 2021.
Samuel Garcin, James Doran, Shangmin Guo, Christopher G. Lucas, and Stefano V. Albrecht. DRED: Zero-shot transfer reinforcement learning via data-regularised environment design. In International Conference Machine Learning (ICML), 2024.
Ziv Goldfeld, Ewout van den Berg, Kristjan Greenewald, Igor Melnyk, Nam Nguyen, Brian Kingsbury, and Yury Polyanskiy. Estimating information flow deep neural networks, 2019. URL httpss
//arxiv. org/ abs/ /10///110...0 05728
11

---

Qing Guo, Junya Chen, Dong Wang, Yuewei Yang, Xinwei Deng, Jing Huang, Larry Carin, Fan Li, and Chenyang Tao. Tight mutual information estimation with contrastive fenchel-legendre optimization. Advances in Neural Information Processing Systems, 35:28319-28334 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. In International Conference on Learning Representations, 2018.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 9729-9738, 2020.
R Devon Hjelm, Alex Fedorov, Samuel
Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In International Conference on Learning Representations, 2019.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks, 2016.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76: 201-264, 2023.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International conference on machine learning, pages 5639-5650 PMLR, 2020.
Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. International Conference on Machine Learning, pages 5757-5766. PMLR, 2020.
Lanqing Li, Yuanhao Huang, Mingzhe Chen, Siteng Luo, Dijun Luo, and Junzhou Huang. Provably improved context-based offline meta-rl with attention and contrastive learning. arXiv preprint 10Xiv22102.10774, 2021.
Qiang Li. Functional connectivity inference from fmri data using multivariate information measures. Neural Networks, 146:85-97 97, 2022.
William McGill. Multivariate information transmission. Transactions of the IRE Professional Group on Information Theory, 4(4):93-111, 1954.
Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction, 2020. URL httph:///axxx http://///xxxx.or//ab/// httg//bs/18002/ 03426
Trevor McInroe, Lukas Schäfer, and Stefano V. Albrecht. Multi-horizon representations with hierarchical forward models for reinforcement learning. Transactions on Machine Learning Research (TMLR), 2024.
Andriy Mnih and Yee Whye Teh. A fast and simple algorithm for training neural probabilistic language models. arXiv preprint arXiv: 120n.6266.6 2012.
Yao Mu, Yuzheng Zhuang, Fei Ni, Bin Wang. Jianyu Chen, Jianye HAO, and Ping Luo. DOMINO: Decomposed mutual information optimization for generalized context in meta-reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
Leland Gerson Neuberg. Causality: models, reasoning, and inference, by judea pearl, cambridge university press, 2000. Econometric Theory, 1964))675-685 2003.
12

---

Kento Nozawa and Issei Sato. Understanding negative samples in instance discriminative self- supervised representation learning. Advances Neural Information Processing Systems, 34: 5784-5797, 2021.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03745 2019.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning, pages 5171- 5180. PMLR, 2019.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL http://jmlz www.ooggggaeeeoo...oooooo pttppppapas//////2// 1364. html
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off oflicy meta-reinforcemmnt learning via probabilistic context variables. International conference on machine learning, pages 5331-5344 PMLR, 2019.
John A Rice and John A Rice. Mathematical stafistics and data analysis, volume 371. Thom- son/Brooks/Cole Belmont, CA, 2007.
Joshua Robinson, Ching Yao Chuang, Suvrit Sra, and Stefanie Jegelka. Contrastive learning with hard negative samples, 2021. URL htttts/////xxxi wttp://gg/a....o..///// 04592
Tong Sang, Hongyao Tang, Yi Ma, Jianye Hao, Yan Zheng, Zhaopeng Meng, Boyan Li, and Zhen Wang. Pandr: Fast adaptation to new environments from offline experiences via decoupling policy and environment representations, 2022.
Younggyo Seo, Kimin Lee, Ignasi Clavera Gilaberte, Thanard Kurutach, Jinwoo Shin, and Pieter Abbeel. Trajectory-wise multiple choice learning for dynamics generalization in reinforcement learning. Advances in Neural Information Processing Systems, 13:12968-1299 2020.
Lucy Xiaoyang Shi, Joseph J. Lim, and Youngwoon Lee. Skill-based model basee reinforcement learning. Conference on Robot Learning, 2022.
Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information estimators. In International Conference on Learning Representations, 2020.
John Tabak. Geometry: the language of space and form. Infobase Publishing, 2014.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw), pages 1-5. IEEE, 2015.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033. IEEE, 2012.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008
Bernie Wang, Simon Xu, Kurt Keutzer, Yang Gao, and Bichen Wu. Improving context-based meta-reinforcemment learning with self-supervised trajectory contrastive learning, 2021.
Michael L. Waskom. seabor:: statistical data visualization. Journal of Open Source Software, 6(60): 3021, 2021. doi: 02211005joss.0022.. URL https: org/10 21105 03021
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non- parametric instance discrimination. Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3733-3742, 2018
13

---

Tianhe Yu. Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-11000 1100. PMLR, 2020.
Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. In International Conference on Learning Representations, 2019.
14

---

A Proof of Lemma 1
Given a query z and a set Y f....... yK K random samples, containing one positive sample y1 and K negative samples drawn from the distribution p(y), K s-sample InfoNCE estimator is obtained by comparing pairs sampled from the joint distribution (ær, 31 p(æ; y) with pairs (z. yk:) constructed using the set of negative examples y2:k- InfoNCE compares the positive pairs (æ,, y1) with the negative pairs (I, yk), where yk ~ Y2:K, as follows:
InfoNCE vlv, Ep(a,a) (plaae)
log
(5)
R k= fI 4k))
Step 1. Let us prove that the K K-sample InfoNCE estimator is upper-bounder by log K. According to fa(
Mu et al. 2022],
f0(2.01)
S So we have:
Σ E fo(maen)) (o(mm)
(a(e,11)
InfoNCE( (2; ulv,, K) Bp(aa))))(
log
K EE (1(,,11)
-p((n)))
Ep(rax)) log
E fo Mkb
S log K
(6)
Hence, we have (2 K) S log K. Step 2. We have the / (x; InnfaNCE w/w,, K) according to:
Proposition 1 Poole et al. 2019 A K K-sample estimator is an asymptotically tight lower bound on the MI, i.e.,
I(x;y) 2 InfoNce(Z: vl,,,, K), lim InjoNCE(a; ylv, I(x;))
++00
Proof. See Poole al. 2019 for neat proof of how the multi-sample estimator (e.g., InfoNCE) lower bounds MI.
Step 3. In this research, the context encoder vb in fo y) is implemented using RNN to p(ylm) (Oord et 2019). With sufficiently powerful deep learning model for y and
approximate plu) finite sample size K, such that (x:y) log K, we can reasonably expect that IIfONCE: F log K after a few training epochs. Therefore, during training, when K < +oo, we always have [(x; u) 2 log K.
Proof. See Chen et al. 2021 for more detailed proof.
Step 4. Let prove that the K-sample InfoNCE bound is asymptotically tight. The specific choice of the context encoder & influences the tightness of the K K-sample NCE bound. InfoNCE
OOrdd et al. 2019] sets (z, y) x p(1)
P(ylæ) to model density ratio that preserves the MI between and y. where stands for "proportional to' (i.e., up multiplicative constant). Substituting
15

---

fx f2 = p(yl=)
into InfoNCE, we obtain:
p(a)
(a v1)
(F: K) = E log
+ log K
2k=1 fof (2, vn.)
p(v)
K
p(yalz)
]
-E
log
1 14
Σ
+ log K
p(y](x)
p(vx)
k=
plu)
(K 1)B plunk (2:)
)
log
log K
p(y(x))
p(yk)
p(y1 1z) (K' (5-11) 11)] + log K
p(v))
-E
log 1
log p(v)
E
log( K 1) log K
p(y)z)
I [(;
v)
log( K + log K
Now taking K- +00, the last two terms cancel out.
Putting it together. Combining [x;;y) 2 log K with Proposition and Equation 66. we have Lemma
IINoNNCE b/, K) I log K I yx:
(8)
Moreover, according to Equation 77) as the sample size K + +o0, the K K-sample InfoNCE bound becomes sharp and approaches the true MI I(x; y)::). i.e., [ItONCe ulv. K) log K I(x v):
B Proof for Lemma 2
Step According to Lemma we have InfoNCe (c: Mei Teli, K) I log I IsaMi(e; nei Te) (shown
in Figure 3).
Step Let prove that SaNCE a K K-sample SaNCE estimator and is upper bounded by log K. have:
Since fule. +t) doleeertt S 1 [Mu et al. 2022
IsaNcE (c: K)
K fc Teis Tot)
=E pler tet 1pl 2
log
felei, To, Tot) + EK fulc1, Tea To)
K R.folch, Te Te) Tet)
=Ep(er
E Bplr log
(ol
(TL Teiy Tet) Ek (c1, Tes To.k
< <log
(9)
Thus, we obtain ISaNCE Teli,, K) log K, similar to Equation 6
Step 3. With the definition of K", we can prove that IIfENCE (c; mei To b K) S ISaNCE Tei mat Talus, K) with the same sample size K. In task e1, SaNCE obtains positive and negative samples from the current task e1. Since the variable c = c1 is constant, we have:
TsaNce (c; Tei Tolib, K)
K (c Te,, Ter,TTh Tc))
E p(c mey net 1n(7 log
+ sepee, TT
KŠANCE' Splc1, Tey Ter rot
< E Epporiee mey ret 1pl
log fol Te, Tot E
KyNen for (CI, Ter toi
)
KÊANCE fe(
= iplt aiNGUNEe
log
(ci is constant.)
(os Tet) + Ekonoo fos Torn)
2 log KEANCE.
(10)
16

---

The required sample size is SaNCE ccl M = M. As K - SaNCE, we have Isacc (2 Mat Tel (,, (o, K) = IsaMI(C; mei Te). Correspondingly, for InfoNCE, in the current task e with context embedding e1, positive samples are trajectories generated after executing the skill m in task 61, while negatíve samples are trajectories (Tez.-..) from other tasks (e2,.... Under the definition of K*, we have:
InNNNE Tei Te)
K ful
D a
log
føl (c, Te: Tei) + fulc, The Toale1k))
Ki felc, Tes Ted)
Telesrerte 1n(7 (n(Te K.Niace
log
KinNNe fa Tes Toales))) fulc Tes Toi) 2h
2 log KIMONNEE
(11)
where K iatENCE |c| |T| M |c| SANCE- In real-world robotic control tasks, the sample space size increases significantly due to multiple environmental features e {c), ...... e.....,,................... The number of different tasks |c] grows exponentially due to the permutations and combinations of the N environmental features. When the current task 61 has context embedding C1, the C2:1e] refer to the context embeddings for the other tasks. As K + KitONNE: have /InfOCC mei Tol (6, K) (sMMi(;; Thi To). Thus, during the training process, InfoNCE Tei Tolwi, K) S ISaNCE(C) Tei To Taleh, K) with the same sample size K.
Step 4. According to the definition of SaMI in Equation (2). we have IsaMi nef Te) S [c: Te), as illustrated using Venn diagrams in Figure 8 (a) and (b). SaMI is formulated based interaction information [MGilll 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 negative
Li 2022
(cc)
[(c; T.)
H(c)
((t,)
(Tr)
(
π
IsaMi T2)
(a)
(b)
(c)
Figure 8: Venn diagrams illustrating (a) mutual information [[c; Te), (b) interaction information IsaMi(C; Te; Te), and (c) the MDP graph the context embedding c, skill Tes and trajectory Te, which
represents a common-caus, structure [Neuberg] 2003
In the Meta-RL setting, the MDP (causal) graph structure, shown in Figure 3 (c), illustrates that the context embedding skill n., and trajectory т. form common-cause structure, where acts as a shared cause influencing both the skill т. and the trajectory тe- Consequently. T( Te c) < [TT,, Te). Therefore, IsaMi (c; me; Te) 0 is guaranteed, making interpretable in real-world robotic control tasks.
Putting it together. Thus, we establish Lemma 2 we always have JINfNNCE (c; T Ta V. K)
IsaNce (G; mas rely, K) S log < TsaMI(;; mej Te) [CCTTa) as shown Figure 3 while learn-
ing skirl-aware context encoder v with the SaNCE estimator. Since K KS,NCE Isacce mat Tel K) serves as a much tighter lower bound for the true IsaMi [C nei Te) than
K
InfoNCE (c; Maj To (,, K).
C Sample size of ISa+InfNNCE.
In this section, we illustrate the sample size of Isa-IntaNCC ne Telv), K). Sa+InfoNCE, incorporates SaNCE into InfoNCE, using positive samples - from task ej after executing skill ntt and negative
17

---

samples are trajectories Toixx from executing skills TeL in tasks 61:K, respectively. Therefore, this approach is equivalent first observing the variable and then observing the variable Tes ie.. sampling from the distribution p( Telc)(()). We have:
Isa-InooNC Ths Tells, K)
K fule, Ttj
= =E ple)p(m2 1c)p (221 log
folc, To; 52t) + fulc, T2:/c/,k 22/0,1
KSa-inNNNE (e Tet Tot)
E
log
plc)p(ad, rcj lc (T
folc, set) > 2k=2
Suttancc fulc, T 2/cl, 22cc),
2: 2.
log SatiateNCE
(12)
should be noted that such combination increases the size of the negative sample space, i.e.. |xt) M 2 KaNCCE 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, /SaInfoNCE is less precise and looser than IS,NCE. Therefore, trade-off between sample diversity and the precision of the K sample estimator is required.
K Sa+InfNNC: (E
D Environmental setup
1.0 0.1
Training Tasks
Friction values
Moderate Test Tasks
5.0
30.0 10.0
Extreme Test Tasks
30
Mass values
(a) Panda-gym
(b) Task setting
Figure 9: (a) Modified Panda-gym benchmarks, (b) the training tasks, moderate test tasks, and extreme test tasks. The moderate test task setting involves combinatorial interpolation, while the extreme test task setting includes unseen ranges of environmental features and represents an extrapolation.
D.1 Modified Panda-gym
We modified the original Pick& task in Panda-gym Gallouédec et 2021] by setting the dimension (i.e., the desired height) of the cube's goal position to 0 and maintaining the freedom of the grippers allowing the agent explore whether it should push grasp the cube. Skills in this benchmark are defined as:
Plac& 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 s end effector and the cube, and the cube's 's height being greater than half its width.
IIf z not equal to 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 Plck&Place skills, leading "unpushable" failure in Figure (b).
18

---

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 's height equals half its width.
Other skills: Any behaviour modes other than Pick& and Push are classified as other skills.
Some elements in the RL framework are defined as follows:
State space: We use feature vectors that contain the cube's position (3 dimensions), cube rotation dimensions), cube velocity dimensions). cube angular velocity (3 dimensions). end-effector position (3 dimensions), end-effector velocity (3 dimensions), gripper width (1 dimension), desired goal (3 dimensions), and achieved goal (3 dimensions). Environmental features are not included in states.
Action space: The action space has 4 dimensions; the first three dimensions represent changes in the end-effeecto's position, and the last dimension represents the change in the gripper's width.
During training, randomly select combination of environmental features from training set by sampling combinations from the following sets: mass
and friction E (0.1, 1.0, 5.0, 10.0); mass E {1.0, 5.0, 10.0} and friction 0.1. At test time, evaluate each algorithm tasks from the moderate test setting, where mass E (5.0, 10.0] and friction E {1.0, 5.0, 10.0) (shown in Figure b)), and all tasks from the extreme test setting: mass 30.0 and friction E {0.1, 1.0, 10.0,30.0); 30.0); mass E {1.0, 5.0, 10.0, 30.0) and friction = 30.0 (shown in Figure ()).
all
D.2 Modified MuJoCo
We extended the modified MuJoCo benchmark introduced in DOMINO [Mu al. 2022 and CaDM [ee et al., 2020]. In our extension, four new environments were added (Walker, Crippled Hopper, Crippled Walker, HumanoidStandupp) compared to the original benchmark. Additionally, in our experiments, we used a different task set design (Table than those used in the DOMINO and CaDM papers. For Hopper, Walker, Half-Cheetah, Ant, HumanoidStandup. and SlimHumanoid, we used the MuloCo physics engine environments and implemented settings from Clavera et al. (2019b) and Seo al. (2020), scaling the mass of every rigid link by scale factor m and the damping every joint by scale factor d. For Crippled Ant, Crippled Hopper, Crippled Walker, and Crippled Half-Cheethe we used the implementation available from [Soo et al. [2020], scaled the mass of each rigid link by 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 provided in Table
Ant
Half-cheetah
SlimHumanoid
Hopper
Walker
Crippled
Crippled Half-cheetah
Humanold Standup
Crippled Hopper
Crippled Walker
Figure 10: Ten environments in modified MuJoCo benchmark.
19

---

Episode Length
1000
1000
1000
1000
1000
1000
2000
2000
2000
2000
Chippled Legs (R,, R2) E (0, 2, 4,5) (8) * Rd)
Crippled Joints (R,, Re) (0. 1,2,3,4, (R2 Ra)
T. E (0.20, 0040, 1.0, 1800, 4,00)
SU, 4.00)
(0.20, 0. 1.0, 1.80, 4.00)
d (0.20 0 40, L.00 4.00)
m E (0.20, 0 4040 1.0, 180, 4.00)
de (0.20, D. 40, 1.60, L., 4.00)
E (0.20, 0.40, 1.00, 1.80, 4.00
de (0.20, 0.40, 1,00, 1.80, 4.0)
Crippled Legs R. e (0. 1.2]
(0.40, 0.50, 1.70, 1.80)
d {0.40, 0.50 1.70, 1800)
TH E (0.40, 0 (.50 1.70, 180, 4.00)
de (0.40, 0.50, 1.70, 4.00).
E (0.20, 0. 40. 160, 180, 4.00).
de (0.20, n. 40, LOO, L80. 4.00)
m E (0:20, 40, 160, L.00, 4.00).
Ls0, 4.00).
de (n.2 n. 40, Crippled Joints (left leg) E (3. (3,4,8
m E {0.20. 0.40, 100, 1.80)
de [0.20, 0.40, 1.60 1.80)
m E (0:20, 0:40, 1.60, 1.80)
de (0.20, 0.40, 1.00, 1.80)
Test (Extrame)
de (0.20 0 40,
Table 3: Environmental features used for MuloCo benchmark.
Crippled Joints (back leg) Ri e (0, 1,2)
Crippled Joints (right leg) E 10,1,2)
Crippled Legs R; e {0, 1,2,3) 3)
1 E (0-40, 0.50, 1.50, Lo0)
d' (0. 41,4 0.50, 1,50, 1600)
E (0.20 (0.20, 0. 0.40, 160, 1800) d (0.20, 0 404 6600 L.S0)
m E (0-40, 0.50, 1500, 1.00) d' (0. (0.40, 0 0050 5,50 1.60)
E (0 40, 0.50, 1,50, L,60)
de (0. 40, 0.50, 1.50, 1.00)
C (0.60, 0,70, 1,50, 1. 6,60) dE (0.60, 0.70, 1.50, 1.0))
T E (0.60, 0.70, 150, 6600
dE (n.60, 0.70, 1,00 L.03.
E 40, 0.50, 1,50, 1. 6.60) dE (n. 40 0.50, 1.50, Lo0))
m E (0-40, 0.50, 1.50) (.60) de (n. 1n, 0.50, L50, Lon).
m E (n. 0.50, 1,50, L.60)
d (0 40, 0.50, 1550. 1.60)
m E (0. 40, 0.50, 1.50, 6.60)
d (0. 40, 0.50, 1.50, 1.00)
Test (Moderate)
Crippled Logs Ri 0 or R, E (0, 1,
Crippled Joints (front leg) R; (3,4.5)
Crippled Joints (right leg) = (0,1, 1
m (0.75, 0.85, I.0, 1.15, 1.25)
d e (0.75, 0.83, 1.0. 1.15, 1.25) 1.25)
(0.80, 0. 10,00 1.0, 1.15, 1.25) E (0.80, 0.90, 1.0, 1.15, 1.25)
i (0.8) (0.80, 0 10,0, 1.0, 1,15, 1.25) E (0.80, 1., L.0 1.16, .
(0.75, 0185, I.0, 1.15, ..25)
E (0.75, 1.83,100. 1.15, 1.25)
E (0.75, 85, 1.0, 1.15, 1.25)
d E (0.75, 0.85, 1.0, 1.15, 1.25)
(0.40, 40, 0.50, .50, 1.60)
c (0-40, 0.50 50, 1.50, (60))
m E (0.75, 1.0, 1.25)
d'c (0.73,1.0) 1.0. 1,25)
E (0.75, 1.0, 1.25)
de (0.75,1.0) 1.0, 1.25)
(0.75. 1.0, 1. ..2
dE (0.75, 1.0, ..25)
12 E [0.75, L.U, 1.255
de (n.75, L.0, 25)
Training
Crippled Half-cheetah
Crippled Hopper
HumamoidStandup
Crippled Walker
Half-cheetah
Crippled Ant
Hopper
Slimluumani
Walker
Ant
20

---

Table Hyperparameters used in the Panda-gy gym and MuJoCo benchmarks, Most hyperparameter values remain unchanged across tasks, except for the contrastive batch size and the SaNCE loss coefficient.
Hyperparameter
Value
Replay buffer size Contrastive batch size SaNCF loss coefficient a Context emhedding dimension
100,000
MuJoCo 12. Panda-gym 256
MuJoCo 1.0, Panda-gym 0.01
6
Hidden state dimension Learning rate (actor, critic and encoder)
128
le-3
Training frequency (actor, critic and encoder)
128
Gradient steps
16
Momentum context encoder ri soft-update rate
0.05
SAC target soft-update rate
critic 0.01, actor 0.05
SAC hatch size
256
Discount factor
0.99
Optimizer
Adam
D.3 Implementation details
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 and [ Our codebase built
on top of the publicly released implementation of Stable Baselines3 by Raflin et al. [022]] and the
implementation of InfoNCE by Oord al. (2019). public, open-source implementation of SaMI
available at https://gihhub.ooouoeeeeee aeeee-ggnt//Sa....
Base algorithm. We SAC (Haarnoja 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 each environment on the Panda-gym 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-Tear Memory (LSTM) network that produces 128- 228-dimensional hidden state vector, which is subse- quently processed through a single-layer feed-forward network to generate a 6-dimensional context embedding. We aim for the agent 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 the same context encoder to embed trajectories. For contrastive learning, SaNCE utilises a momentum encoder 1f" to generate positive and negative context embeddings [Laskin al. 2020 He et al. 2020). Formally, denoting the parameters of ei as 0, and those of s² as Ug., we update Ug.. as tollows:
Bg- m. As + (1 (1-m)) m) 8p..
m.
(13)
Here m E [0, 1) is a soft-update rate. Only the parameters 0g are updated by back- pack-propagation. The momentum update in Equation 133 makes 0g- evolve more smoothly by having them slowly track the 6e 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 displayed in Table
Hardware. For each experiment run we use a single NVIDIA Volta V100 GPU with 32GB memory and single CPU.
21

---

Algorithm SaNCE Meta-training
Require: Batch of training tasks {en from (train(e), soft-update rate mi;
1:
Initialize RL replay buffer BRI. encoder replay buffer Benci
2: Initialize parameters d/ for context encoder, x/" for momentum context encoder and @ for the off-policy SAC:
3: while not done do
4: for each task e,, do for Roll-out time steps do
5:
6:
for time step maximum episode length T do
7:
Update context embedding Cv. <b(Cnl .u..
8:
Roll-out policy TT.,, (0r|s; Cn) and accumulate transition (St. aes Pty St+1); end for
9:
Add trajectory Te,, 88, do, ro, 81, f1, mn 8T, OT, rT) to replay buffer Bi and B B end for
11:
12:
end for
13:
for each training step do
14:
for each task e, do
15:
Sample RL batch fre. B
16:
Sample a positive sample T for generating query with highest return, positive samples
{T= } and negative samples {- for encoding positive and negative embeddings; Update with RL loss LRL.
17:
18:
Update v with SaNCE loss LSaNCE and RL loss LRL:
19:
By- +- m Ag + (1 m) 0x=;
20:
end for
21:
end for
22:
end while
Algorithm SaNCE Meta-testing
Require: Batch of training tasks {en]n= from Stest(e):
1: while not done do
2: for each task e,, do
3: for each episode do
4:
for time step < maximum episode length T do
5:
Update context embedding Cn vb(cal .oe)
6:
Roll-out policy т.,, (oxl St, Cx) and accumulate transition (st, Ce, Tt, 8t+1 1 end for end for
7:
8:
9:
end for
10:
end while
22

---

E
Additional results
E.1 Balance contrastive and RL updates: loss coefficient
While past work has optimised hyperparameters to balance the contrastive loss coefficient a relative the RL objective [Jaderberg et al. 2016 Bachman al. 2019]. we use both the contrastive and RL objectives with equal weight, setting = 1.0 for the MuloCo benchmark and a 0.01 for the Panda-gym benchmark, Additionally, we analyse the effect of the loss coefficient a for CCM, SaTESAC, and SaCCM in the MuloCo (Figure 12 and Panda-gym (Figure I1) benchmarks.
0.6
CDM
CCM
SeTTEAAC
3.4
SaCCM
Succes Rate
0.5
Success Rate
0.3
0.3
0.2
0.1
0.1
0.1
8.05
0.01
1.0
0.1
0.05
0.01
1.0
0.05
0.03
loss coefficient
coefficient
loss coefficient or (c) Test tasks (extreme)
(a) Training tasks
(b) Test tasks (moderate)
Figure 11: Loss coefficient a analysis of Panda-gym benchmark in training and test (moderate and extreme) tasks.
8000
6000
2008
6000
5000
Average Return
Aver age Return
Awerage Return
8000
5000
4000
4000
6000
3000
2000
2000
2000
CCM
1000
Sei TTTSAA
1000
SaTESAC
S:IESAC
SaCEM
SNCCM
S10CC
1.0 loss coefficient d (a) Half-cheetah
10.0
coefficient
(b) Crippled Half-cheetah
(c) SlimHHumnnidd
1750
800
1600
SATESAC SaCCI
700
Average Return
Kiet ege Belune
600
Awerage Return
1,000
500
1000
400
800
750
300
900
500
200
400
100
SuCCM
SaCOM
contficient
confficien
contficient
(d) Ant
(c) Crippled Ant
(f) Hopper
Figure 12: Loss coefficient analysis of MuJoCo benchmark in training tasks.
23

---

E.2 Result of log-K curse analysis
E.2.1
Buffer size
7000
TESAC
CCM
12000
6000
SATESAC
Average Return
10000
5000
Jveraae
5000
Average
4000
6000
8000
TESAC
2000
4000
3000
2000
2006
SaCEM
1000
SHCCM
1000
10000
100000
400000
1000
10000
100000
400000
1000
10000
100000
400000
Buffer Size
Buffer Size
Buffer Size
(4)
Crippied
2000
TESAC
1600
1/70
CCM
1200
DCM
1300
Aversge Return
1000
SaTESAC.
SaCCM
Average Raturn
1400
1250
1300
1000
Average
200
1200
730
600
TESAC
400
CCM
250
200
900
SaTESAC
3000
10000
Buffer Size
100000
430000
1000
10000
Duffer Size
100000 400000
800
Buffer Size
Ad)
(6) Crippled
Hopper
Figure 13: Comparison of different buffer sizes in the MuJoCo benchmark on training tasks (averaged over seeds). Buffer sizes are 400,000, 100,000, 10,000, and 1,000.
E.2.2
Contrastive batch size
7500
Raturn
Averoge Return
Average Rehurn
7000
8500
5000
4500
3500
CCM
8000
CCM
CCM
5000
SaTESAC
SaTESAC
SWTESAC
SacCM
SaOCM
5000
Contrastive Batch
Contrastive Batch Size
Contrastive Batch Size
[8] Half-cheetah
(b) Crippled Har-cheesshh
SinHHmmnnodd
900
1200
Avarage Retur
Average Retum
800
1520
1000
700
2500
Juarage
1480
890
3460
600
CCM SITESAC
1440
SaCOM
SNCCM
1420
SeCCM
128
512
16
128
512
Contrestive Batch Size
Contrastive Batch Size
Contrastive Batchh Size
(4) Crippled Ari
( Hopper
Figure 14: Comparison different contrastive batch sizes in the MuJoCo benchmark on training tasks (averaged over 5 seeds). Contrastive batch sizes are 512, 128, 16, and 8.
24

---

Further skill analysis
E1
Panda-gym
F1111 Visualisation of context embedding
We visualise the context embedding using UMAP McInnes et al. 2020 (Figure ??) and t-SNE Van der Maaten and Hinton 2008 (Figure I6). When the mass of the cube is high (30 kg and 10 kg), the agent learned the Push skill (indicated by the yellow bounding box in Figure T [a)), whereas with lower masses, the agent learned the Plac& skill. However, as shown in Figure 16tb).
CCM did not display clear skill grouping. This indicates that SaMI extracts high-quality skill-related information from the trajectories and enables agents autonomously discover diverse range of skills for handling multiple tasks.
TESAC
Bccart MASELLA
wert
(c) SaTESAC
(d) SaCOM
Figure 15: UMAP visualisation of context embeddings extracted from trajectories collected in the Panda-gym environments.
-10
TESAC
(b) CCM
-10
(6) SaTEBAC
(d) SaCCM
Figure 16: t-SNE visualisation of context embeddings extracted from trajectories collected in the Panda-gym environments.
25

---

F11.2
Heatmap of Panda-gym benchmark
0.7
0.7
0.7
0.7
0.1
0.9
0.8
0.7
1.0
0.2
0.1
0.2
0.2
0.B
Friction values
1.0
0.2
1.0
0.8
0.7
Friction values
0.1
0.1
0.1
0.3
5.0
0.1
Friction
5.0
5.0
10.0
0.7
0.4
0.2
0.1
10.0
0.5
0.6
0.6
0.8
10.0
0.1
0.1
0.1
0.2
0.6
0.4
0.2
0.1
30.0
0.6
0.6
0.6
0.7
30.0
0.0
0.0
0.0
0.0
10
30
5
10
30
0.0
values
Mass
(a) TESAC Success Rate
(b) TESAC Push Rate
(c) TESAC Plck&Place Rate
0.1
0.6
0.3
0.1
0.1
0.0
0.1
0.8
0.9
0.6
0.0
1.0
0.9
0.8
0.5
0.2
friction values
1.0
0.0
0.0
0.2
0.1
Friction values
1.0
0.8
0.8
0.5
0.0
5.0
0.9
0.7
0.3
0.1
10.0 5.0
0.0
0.1
0.3
0.2
5.0
0.8
0.8
0.5
0.0
0.4
10.9
0.9
0.1
0.1
0.1
0.3
0.8
0.7
0.5
0.0
0.2
30.0
30,0
30,0
30
1
5
10
30
10
30
Mass values
Mass values
Mass values
(d) CCM Success Rate
(e)
CCM Push Rate
(f) CCM Pick&Place Rate
1.0
0.1
0.1
0.3
1.0
0.1
0.8
0.8
0.5
0.0
1.0
0.9
0.9
0.7
0.2
Friction values
0.2
0.2
0.3
0.9
1.0
Friction
5.0
0.9
0.6
0.3
0.1
0.1
0.1
0.3
0.7
Friction
5.0
0.8
0.7
0.4
0.0
0.4
30.0 10.0
0.9
0.6
0.3
0.1
0.1
0.1
0.3
0.7
10.0
0.8
0.7
0.4
0.0
0.2
0.9
0.5
0.2
30.0
0.1
0.1
0.2
0.8
30,0
10
30
10
30
10
30
Mass values
Mass values
Mass values
(g) SaTESAC Success Rate
(h) SaTESAC Push Rate
(i) SaTESAC Pick&Place Rate
Figure 17: Heatmap of success rates and learned skills of SaCCM. For each grid, the (i, location shows the probability of the skills executed over 100 evaluations with (mass i, frictionn j).
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 T), 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 Figure Ta)). At lower masses, the agent executed the Pick& skill.
In contrast, as shown in Figures e-f), CCM predominantly learned the Pick& Place skill, resulting in 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, depicted in Figure 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 does not require the agent to manipulate its fingers to pick up cubes. Consequently, TESAC's success notably decreased in environments with higher friction.
26

---

F.2 MuJoCo
SaMI enables RL agents be versatile and embody multiple skills. Additionally, video demos (available at https://gihuub.oo/uue-gggee provide better demonstration of the different skills.
OCM
BaTESAC
BaOCM
(4)
0 COM
(P)
I-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in the Crippled Calf-Cheetah environment. "cripple_see set" refers to the index of the crippled joint. The figure shows the context embeddings of tasks in three moderate test settings. Combined with the video demos for skill analysis, the Crippled Half-Cheetal robot executed three distinet forward running skills.
Figure 18:
(0) PCA TESAC
COM
(0) PGA to SiTesii
(h)
Figure 19: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in the Crippled Ant environment. "cripple_ set" refers the index of the crippled joint. When or 4 legs are available, the Ant robot (trained with SaCCM) rolls to adapt to varying mass and damping. However, with only 2 adjacent legs during zero-shot generalisation, it switches to walking. 2 opposite legs are available, the Ant can still roll but eventually tips over. Please refer to the video 2
27

---

BNE
(6) UNAP for TESAC
0 UMNP for COM
02 UMAP for SaTEBAC
(1) UMMP for SaOCM
Figure 20: t-SNE and UMAP visualisations of context embeddings extracted from trajectories environment. During zero-shot generalisation, SaCCM demonstrates
collected in the different skills for running forwards at various speeds, as well as skills for doing flips and faceplanting.
DINSNE for TESAC
(J)- for CCM
1-NEE for BaTEBAC
(d) H-ME for SaOCM
TESTAC
UMAP SOTESC
LMAD SWOCM
Figure 21: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in Ant environment. The Ant robot learned a single skill, rolling, to adapt to different mass and damping values.
(4JUMAP for TEBAC
1 UMAP for COM
UNUU for SATEBAC
Figure 22: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in SlimHumanoid environment. The Humanoid Robot crawls the ground using one elbow. When the damping is relatively high (damping=1.8), the Humanoid Robot can crawl forward stably, but when the damping is low (damping=0.4)), it tends to roll.
28

---

(6) UMAP for TESAG
01 UMAP for COM
(0) UMAP for SATESAC
Figure 23: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in HumanoidStanduu environment. SaCCM and SaTESAC learned sitting posture that makes easier stand up, allowing it to generalise well when mass and damping change.
INSE for TESAC
0011-NE CCM
DONE SOTESAO
0111 for BaCCM
TEBAC
Figure 24: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in Hopper environment. Combined with the video demos 0 for skill analysis, the plots for SaCCM and SaTESAC show skills: 1) when the mass is low, the Hopper hops in an upright posture; 2) when the mass is higher, the Hopper hops forward on the floor.
(0)--SME TESAC
SSME SATESAC
10NE SaCOMI
(e) UISAP TESAC
(3 for COM
LIMG for Sai LSAC
UM for SaOOOM
Figure 25: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in Crippled Hopper environment. "cripple_see set" refers the index of the crippled joint. SaCCM and SaTESAC learned to take a big hop forward at the beginning (i.e., effective exploration) and then switch different skills based on environmental feedback.
29

---

IAJ--NNE
(BIL-BNE
(0):BBE
AM
(5) UIMP fO TEBAC
LINAP COM
UMAP SATESAC
UAAAP for SWOCM
Figure 26: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected Walker environment. The Walker Iearned a single skill, hopping forward on the floor through both right and left legs.
IRJ--BNE for TESAC
(bIL-NNE COM
(ONE BUTESRO
001-SNE For BOCOM
(9)
(1) UIAMAPPN for COM
(g) UMAP SASTAAAA
(HJUMAP RWOOM
Figure 27: t-SNE and UMAP visualisations of context embeddings extracted from trajectories collected in Crippled Walker environment. "cripple_sse set" refers to the index of the crippled joint. The Crippled Walker (trained with SaTESAC and SaCCM) learned to hop forward using the right leg (cripple_ set:0 and cripple_sede set:2) in the training and moderate test tasks and switched to hopping forward using the left leg (cripple_sete) the extreme test tasks.
Table 5: Environmental features used for MuloCo benchmark from DOMINO and CaDM.
Training
Test (Moderate)
Test (Extrame)
Episode Length
0,65, 1.0, 1.25) de (0.75, 0.85, 1.0, 1.1,
(0.10,0.50 1.60)
(0.20, 40, 1.60, 1.80) d (0,20, 0.40, 1.60, 1.80)
Half-chectah
d (0.40, 0.500 1.60)
m. (0.75,0085, 1.0, 1.15,
in [0.40,0.500 1.50, 1.60)
C 00.20,0.40,
0.40, 1.60, 1.80]
d e (0.75, 0.85, 1.15, 1.25)
d (0.40, 0.50, 1.60)
d (0.20, 0. 40, 1.60, 1.80)
i. [0.80, 1.0, 1.25)
in 0.70,, 1.60)
C 00.00,0.50, 1.80 0,50,
SlimHHumooodd
d e (0.80, 1.0, 1.15
d (0.60,0.77 0.70, 1.60)
de (0.10, 0,50,,00,.1000 1.80)
m (0:75, 0.85, 1.0, 1.25)
i (0.40, 0.50, 1.50, 1.60)
e (0.20,0.40,
0.40,
1.80
Crippled Half-chestah
C [0.75, 0.85, 1.0, 1.15, 1.25) Crippled Joints: (0, 1, 3)
C (0.40,0.35 0.50, 1.50, 1.60]
[0.20, 0.40, 1.60, 1.80)
1000
Crippled Joints: (4.5)
Crippled Joints: (1.5)
30

---

Table 6: Comparison of average return standard deviation with baselines in MuJoCo bench-
mark (over 5 seeds). The bold text signifies the highest average return. The numerical results for PPO+DOMINO are copied from Mu et al. [0022]; the numerical results for PPO+CaDM, Vanilla+CaDMM, and PE-TS+CADM are copied from Lee et al. (020).
Half cheetali
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
PPO+DOMINO
227:186
216::52
2472:4803
1034::476
PPO+CaDM
268.617774
228.8:1-8.4
199.2:::2.1
2652.0:1133.4
1224.2:::30.0
1021.11676.6
MaDI
1831,0=#13.77
1315.7:445.5
131.4:::13.5
3536,5-:641.7
1556.14:260.6
1264.5=228.7 7
PE-TS*CaDM
2848.4:661.9
2121.0:::00.
1200.7=21.88
$264.0=1374MM
7087.2+
4661.8:783.9
SaTESAC
908:65
640:111
532-:88
7430=
4058:890
1780s 102
SaCCM
928=141
635-94
555+888
7154=1965
3849:6689
1926#2188
SlimHumanoid
Crippled Half-Cheetah
Training
(moderate)
(extreme)
Training
(moderate)
(extreme)
PPO+DOMINO
7825=1256
5258::103
2503:1658
1326:2:49
PPO+CaDM
10455.0:: 1004.9
4975.74: 1305.7
3015.1+:15088
2356.6 6:1.244.3
1454.0 0:::462..6
1025.0#:296.2
VanIllaCaDDM
1758.2:1459.1
1274.911374.0
1487.9:1339.0
1375.331:290.6
966.9.189.4
PE-TS-CDD
1371.9:140000
903.7:1:43.99
814.5:1.74.8
3294.9:1733.9
2618.7:1647,1
1294.2:1-144.9
SaTESAC
10216:1102 1620
7886:1:203
6123:::1403
5169:1730
21841:592
1628:2281
SaCCM
9312:4705
7580-11587
6473:1-2001
5709:1744
2795:1446
2115:::666
31

---

G A comparison with DOMINO and CaDM
In this section, we give brief comparison between our methods and methods from DOMINO Mu et al. 2022] and CaDM [Lee et al. 2020 in the MuloCo benchmark because we are using the exact same environmental setting (shown in T5)
DOMINO [Mu et al. 2022 is based on the InfoNCE K K-sample estimator. Their implementation, PPO+DOMINO, is 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, pre-trained encoder necessitates a large sample volume, with DOMINO training PPO agents for 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 Therefore, it 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 learn context embeddings, and then predict the next state conditioned on them. We copy the numerical results of PPO+CaDM, VanillaCCaDM, and PE-TS+CaDM from CaDM [Lee et al. 2020) their environmental setting is identical ours, where PPO+CaDM is a model-free RL algorithm, while Vanilla+CaDM and PE-TS+CaDM are model- based. The model-free RL approach, PPO+CaDM, trained for 5 million timesteps on the MuJoCo benchmark. As shown in Table 6 SaTESAC and SaCCM significantly outperform PPO+CaDM. The model-based RL algorithms, Vanilla+CaDM and PE-TS+CaDM, require 2 million timesteps for learning in model-based setups, compared to our fewer samples (i.e., million timesteps). In the Ant environment, Vanilla+CaDM and PE-TS+CaDM achieve higher returns than SaTESAC and SaCCM; similarly, in the Half-cheetah environment, PE-TS+CaDM outperforms them. Results in the SlimHumanoid and Crippled Half-cheetah environments show that skill-aware context embeddings are notably effective. An insight here is that our method outperforms the model-free CaCM approach, but not the model-based one. This consistent with what is empirically found in CaDM [Lee 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.
32

---

Table 7: The p-value of the statistical hypothesis tests (paried t-tests) for comparing the effectiveness of SaMI in MuJoCo benchmark (over 5 seeds). * next to the number means that the algorithm with SaMT has statistically significant improvement over the same algorithm without SaMI at a significance level of 0.05. The "SaTESAC-TESSAC row indicates the p-value for the return improvement brought by SaMI to AESAC; the "SaCCM-CCM" row indicates the p-value for the return improvement brought by SaMI to CCM.
Crippled Ant
Crippled Half-cheetah
Training
Test (moderatee)
Test (extreme)
Training
Test (moderate)
Test (extreme)
SaTESAC-TESAC SaCCM-CCM
0.121
0.109
9.54E-07 0.008"
0.154
0.0024*
0.0889
0.913
0.108
0.04"
0.307
0.106
Ant
Half-cheetah
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
SATESAC-TESAC S&CCM-CCM
0.136
0.034
0.021
0.346
0.224
0.004
0.138
0.275
0.163
0.73
0.791
0.005*
SlimHumanoid
HumanoidStandup
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
SATESAC-TESAC
0.139
0.456
0,006
0.037
0.129
0,003*
SaCCM-CCM
0.113
0.059
0.008
0.048
0.027
0.01
Hopper
Crippled Hopper
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
SATESAC-TESAC
0.747
0.089
0.707
0.459
0.69
0.088
SaCCM-CCM
0.52
0.599
0.969
0.967
0.897
0.0002
Walker
Crippled Walker
Training
Test (moderate)
Test (extreme)
Training
Test (moderate)
Test (extreme)
SaTESAC-TESAN
0.312
0.082
0.003>
0.223
0.048*
0.028*
SaCCM-CCM
0.55
0.541
0.022
0.794
0.079
0.011
Table 8: The p-value of the statistical hypothesis tests (paried t-tests) for comparing the effectiveness of SaMI in Panda-gym benchmark (over 5 seeds). * next to the number means that the algorithm with SaMI has statistically significant improvement over the same algorithm without SaMI at a significance level of 0.05. The "SaTESAC-TESAG row indicates the p-value for the return improvement brought by SaMI to TESAC; the "SaCCM-CCM" row indicates the p-value for the return improvement brought by SaMI to CCM.
Training
Test (moderate)
Test (extreme)
SaTESAC-TESAC
0.000260*
0.000160*
0.001310*
SaCCM-CCM
0.000230*
0.002190*
0.001390*
H Statistical hypothesis tests (paried t-tests)
We used a t-test Rice and Rice, 2007 to conduct a statistical hypothesis test to determine whether SaMI brought statistically significant improvement. we reported the p-value of the t-test MuloCo (Table 7 and Panda-gym (Table 8) benchmarks. next to the number is used to indicate that the algorithm with SaMI has statistically significant improvement over the same algorithm without SaMI at significance level of 0.05. From Table 7 and g SaMI brings significant improvement on the extreme test set in which the RL agent needs to execute diverse skills. The statistically significant test aligns with our results in the skill analysis (i.e., video demos, t-SNE and UMAP visualisation). In complex environments that require high skill diversity from the RL agent, the statistically significant improvement and higher returns/success rates brought by SaMI are evident.
33

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our abstract and introduction clearly state the claims made, including the contributions made in the paper and important assumptions and limitations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalise to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations and future work of this work in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
34

---

Justification:
Proofs are provided in Appendix A and
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The code will be open-sourced and the implementation details for both the environment setup and algorithms in the Appendix D
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
35

---

Answer: [Yes]
Justification: We will open-source our modified benchmarks (i.e., data) and code after the double-blind review phase ends.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https /mips.ccc/pulic//uuue guwde//ddee/e/e Cot//ooodeeeiiseeeeeeee wwwpiimmsssioooooooooooooooi for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have provided experimental details in Appendix Djincluding hyperparame- ters, how they were chosen, type of optimizer, etc.
Guidelines:but
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the mean and standard deviation of our environmental results over five seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence 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.)
36

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative enrr rates).
show in tables
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have provided sufficient information the computer resources (type of compute workers, memory, time of execution) Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://ppp.//.oooooooo/////////a
Answer: [Yes]
Justification: Our research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We have fully considered the potential societal impacts of research. Our work is conducted in simulated environments, and there is no societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
37

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 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 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).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 usage guidelines 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 best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: All assets used in our study are open-source, and have acknowledged the authors' copyrights in the References.
Guidelines:
The answer NA means that the paper does 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 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.
38

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta
latasets.papersstt
has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
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 release assets.
Researchers should communicate the details the dataset/code/mooell as part 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included 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.
15. 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 crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
39

---

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.
40


Reviews:
[{"review_id": "hI5vS7IT5k", "reviewer": "Reviewer_cm3M", "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.", "questions": "+ Do the authors have any hypothesis on why the generalization behaviour is different when using SaNCE in Panda-Gym vs MuJoCo?\n+ 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\n+ 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": "The paper is mostly about addressing a limitation of previous methods, and I do not believe it has any potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents an interesting novel technique that is intuitive without being trivial. There are experiments across 2 suites of tasks and solid mathematical work, and the overall idea is flexbible enough for future work to build upon.", "weaknesses": "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\". \n\nMoreover 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.\n\nFinally 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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.\n\n**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,\n\n*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.\n\n*2) MuJoCo:* In the MuJoCo environment, we found:\ni) 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\";\nii) 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. \n\nTo 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).\n\nTo 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. \n\n**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. \n\n**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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oMsuom87tj", "reviewer": "Reviewer_hnsz", "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.", "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?\n  - Similarly, I don't understand Figure 6—K is on the x-axis, but $\\log K$ is drawn as a horizontal line.\n\n- 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))$.\n\n  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)$\n\n- How is the \"momentum encoder\" parameterized in this setting?\n\n- Is there significance to the bars not lining up in Figure 2 (c)?\n\n- The authors may wish to discuss related work on hard negative mining", "limitations": "Addressed", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "- The log K curse is an important practical issue when using contrastive infoNCE losses in RL settings with many tasks—often the classification becomes trivial without hard negatives. Approaches like those proposed in this paper will be important to scaling contrastive learning methods for control.\n- Empirical results show strong benefits against the baselines tested in many environments.", "weaknesses": "- There seem to be issues with the mathematical formulation (see questions)\n- I found the presentation of the method to be confusing, in particular regarding how the learned representations were actually used for meta learning", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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: \n1) 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).\n2) 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).\n\nRegarding 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:\n\n1) 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.\n\n2) 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.\n\nWe 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):\nLemma 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)\n\n**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: \n$$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)} $$\nBesides, 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)$.\n\n**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).\n\n**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.\"\n\n**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).”", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jQOzJNLxRE", "reviewer": "Reviewer_TjCk", "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.", "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.\n\n- 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": "Overall I found this a good paper with minor issues/unclear points (see Weaknesses/Questions).", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- relevant topic: Generalising RL agents across tasks is an important topic that will enhance the field of RL\n- well-motivated: The example in the introduction in addition to Fig. 1 is helping the reader to visually understand the considered problem\n- The paper clearly states the contributions", "weaknesses": "only minor things such as an algorithm box would be useful to understand the learning procedure", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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\": \n\n\"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.\n\nIn 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.\"\n\n**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.”\nThis 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.\n\nAdditionally, 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.”", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L5WZ6xJPvm", "reviewer": "Reviewer_h8bu", "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.", "questions": "### Q1. Scalability to Complex Environments\n\nHow 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?\n\n### Q2. Variability in Skill Definitions\n\nHow 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": "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.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "### S1. Novel Problem Formulation\n\nThe paper introduces a novel and relevant problem formulation by addressing the need for distinguishing context embeddings according to skills in Meta-RL. This approach targets the challenge of generalizing across tasks with varying optimal skills, which is crucial for effective RL in diverse environments.\n\n### S2. Empirical Validation\n\nThe proposed methods, SaMI and SaNCE, are thoroughly validated through experiments on modified MuJoCo and Panda-gym benchmarks. The results demonstrate substantial improvements in zero-shot generalization to unseen tasks, showcasing the practical applicability and effectiveness of the methods.\n\n### S3. Clear Presentation\n\nThe paper is well-structured, with clear explanations of the proposed methods and their benefits. The use of figures, such as visualizations of context embeddings and success rate comparisons, effectively supports the presentation of results. The theoretical proofs and experimental details provided in the appendices further enhance the clarity and comprehensiveness of the paper.", "weaknesses": "### W1. Scalability Concerns\n\nThe 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.\n\n### W2. Dependency on Skill Definitions\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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:\n1) 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.\n2) 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.\nWe 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.\n\nAs 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.\n\nWe 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.\n\n**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\": \n\n\"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.\n\nIn 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.\"\n\nAdditionally, 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Lri3rg9VoP", "reviewer": "Reviewer_Qb7x", "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.", "questions": "- For how many episodes does the agent interact with the environment at meta-training and testing times?\n- How does the encoder accumulate across multiple episodes of interaction?\n- Is there a way to use the proposed method for adaptation within an episode?", "limitations": "Limitations are not discussed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper is clearly written.\n- The proposed method is well motivated.\n- The proposed benchmark is interesting.", "weaknesses": "- 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.\n- Due to the above issue, it is inaccurate to say that SaNCE can be integrated with any Meta-RL algorithm.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**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.\n\n**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.\n\nFor 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.\n\nAdditionally, 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.\n\n**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).\n\n**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.\n\n**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.\n\n**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. \n\nFinally, 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thank you very much for the reviewer's feedback. We appreciate the time and effort that the reviewers dedicated to providing feedback on our manuscript, and are grateful for the insightful comments and valuable improvements to our paper. We have addressed each of the reviewers' comments individually. Please find video demos in our Anonymous GitHub repository (https://anonymous.4open.science/r/SaMI). We also attach a one-page PDF file which contains all of the Tables and Figures mentioned in our responses to the reviewers. You can also find our updated anonymous manuscript in our Anonymous GitHub repository (named "SaMI.pdf").

Meta-review: 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.
Decision: Accept (poster)