Title: Skill-aware Mutual Information Optimisation for Zero-shot Generalisation in Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: contrastive learning

reinforcement learning

meta reinforcement learning

zero-shot generalisation
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

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 S killa ware M utual I nformation (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 S killa ware N oise C ontrastive E stimation (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.

1 Introduction

Reinforcement Learning (RL) agents often learn policies that do not generalise across tasks in which the environmental features and optimal skills are different [des Combes et al., 2018, Garcin et al., 2024]. Consider a set of cube-moving tasks where an agent is required to move a cube to a goal position on a table (Figure 1). These tasks become challenging if environmental features, such as table friction, vary between tasks. When facing an unknown environment, the agent

Unliftable (High Mass)

Unpushable

(High Friction)

Pushable

(Low Friction)

Liftable (Low Mass)

Environmental features

Mass, Friction, Material, …

👉

Friction

(a)

(b)

needs to explore effectively, understand the environment, and adjust its behaviour accordingly within an episode. For instance, if the agent tries to push a cube across a table covered by a tablecloth and finds it 'unpushable,' it should infer that the table friction is relatively high and adapt by lifting

the cube to avoid friction, rather than continuing to push. Recent advances in Meta-Reinforcement Learning (Meta-RL) [Lee et al., 2020, Agarwal et al., 2021, Mu et al., 2022, Dunion et al., 2023b,a, McInroe et al., 2024] enable agents to understand environmental features by inferring context embeddings from a small amount of exploration, and to train a policy conditioned on the context embedding to generalise to novel tasks.

In recent years, unsupervised contrastive learning algorithms have been shown to learn context embeddings that perform remarkably well on generalisation tasks [Clavera et al., 2019a, Lee et al., 2020]. In particular, some Meta-RL algorithms [Fu et al., 2021, Wang et al., 2021, Li et al., 2021, Sang et al., 2022] train context encoders by maximising InfoNCE [Oord et al., 2019], which has yielded vital insights into contrastive learning through the lens of mutual information (MI) analysis. The InfoNCE objective can be interpreted as a K -sample lower bound on the MI between trajectories and context embeddings. Despite significant advances, integrating contrastive learning with Meta-RL poses several unresolved challenges, of which two are particularly relevant to this research: (i) Existing context encoders based on contrastive learning do not distinguish tasks that require different skills . Many prior algorithms only pull embeddings of the same tasks together and push those of different tasks apart. However, for example, a series of cube-moving tasks with high friction may only require a Pick&Place skill (picking the cube off the table and placing it at the goal position), making further differentiation unnecessary. (ii) Existing K -sample MI estimators are sensitive to the sample size K (i.e., the logK curse) [Poole et al., 2019]. The low sample efficiency of RL [Franke et al., 2021] and the sample limitations in zero-shot generalisation make collecting a substantial quantity of samples often impractical [Arora et al., 2019, Nozawa and Sato, 2021]. The effectiveness of K -sample MI estimators breaks down with a finite sample size and leads to a significant performance drop in downstream RL tasks [Mnih and Teh, 2012, Guo et al., 2022].

To enhance RL generalisation across different tasks, we propose that the context embeddings should optimise downstream tasks and indicate whether the current skill remains optimal or requires further exploration. This also reduces the necessary sample size by focusing solely on skill-related information. In this work, (1) we introduce Skill-aware Mutual Information (SaMI) , a generalised form of MI objective between context embeddings, skills, and trajectories, designed to address issue (i). We provide a theoretical proof showing that by introducing skills as a third variable into the MI of context embeddings and trajectories, the resulting SaMI is smaller and easier to optimise. Furthermore, (2) we propose a data-efficient K -sample estimator, Skill-aware Noise Contrastive Estimation (SaNCE) to optimise SaMI, effectively addressing issue (ii). Additionally, (3) we propose a practical skill-aware trajectory sampling strategy that shows how to sample positive and negative examples without relying on any prior skill distribution. In that way, Meta-RL agents autonomously acquire a set of skills applicable to many tasks, with these skills emerging solely from the SaMI learning objective and data.

We demonstrate empirically in MuJoCo [Todorov et al., 2012] and Panda-gym [Gallouédec et al., 2021] that SaMI enhances the zero-shot generalisation capabilities of two Meta-RL algorithms [Yu et al., 2020, Fu et al., 2021] by achieving higher returns and success rates on previously unseen tasks, ranging from moderate to extreme difficulty. Visualisation of the learned context embeddings reveals distinct clusters corresponding to different skills, suggesting that the SaMI learning objective enables the context encoder to capture skill-related information from trajectories and incentivise Meta-RL agents to acquire a diverse set of skills. Moreover, SaNCE enables Meta-RL algorithms to use smaller sample spaces while achieving improved downstream control performance, indicating their potential to overcome the log -K curse.

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. [2023]. Existing algorithms can be categorised into three types based on different context embeddings. In the first category, the context embedding is learned by minimising the downstream RL loss [Rakelly et al., 2019, Yu et al., 2020]. PEARL [Rakelly et al., 2019] learns probabilistic context embeddings by recovering the value function. Multi-task SAC with task embeddings as inputs to policies (TESAC) [Hausman et al., 2018, Yu et al., 2020] parameterises the learned policies through a shared embedding space, aiming to maximise the average returns across tasks. However, the update signals from the RL loss are stochastic and weak, and may not capture the similarity relations among tasks [Fu et al.,

𝑎

!

"

𝜋($ |𝑐

"

, 𝑠

!

)

𝑐

"

𝑠

!

𝜋($ |𝑐

#

, 𝑠

!

)

𝑐

#

𝑠

!$"

"

𝑠

!$%

"

𝑟

!$%

"

𝝉

𝒄

𝟏

$

𝑎

!

#

𝑟

!

#

𝑠

!$"

#

𝑠

!$%

#

𝑎

!$%

#

𝝉

𝒄

𝟐

$

…

…

…

𝑟

!

"

𝑎

!$%

"

𝑟

!$%

#

Push

Pick&Place

2021]. The second category involves learning context embeddings through dynamics prediction [Lee et al., 2020, Zhou et al., 2019], which can make the context embeddings noisy, as they may model irrelevant dependencies and overlook task-specific information [Fu et al., 2021]. The third category employs contrastive learning [Fu et al., 2021, Wang et al., 2021, Li et al., 2021, Sang et al., 2022], achieving significant improvements in context learning. However, these methods overlook the similarity of skills between different tasks, thus failing to achieve effective zero-shot generalisation by executing different skills. Our improvements build upon this third category by distinguishing context embeddings according to different optimal skills.

Contrastive learning. Contrastive learning has been applied to RL due to its significant momentum in representation learning in recent years, attributed to its superior effectiveness [Tishby and Zaslavsky, 2015, Hjelm et al., 2019, Dunion and Albrecht, 2024], ease of implementation [Oord et al., 2019], and strong theoretical connection to MI estimation [Poole et al., 2019]. MI is often estimated using InfoNCE [Oord et al., 2019] that has gained recent attention due to its lower variance [Song and Ermon, 2020] and superior performance in downstream tasks. However, InfoNCE may underestimate the true MI when the sample size K is finite. To address this limitation, CCM [Fu et al., 2021] uses a large number of samples for maximising InfoNCE. DOMINO [Mu et al., 2022] reduces the true MI by introducing an independence assumption; however, this results in biased estimates. We focus on proposing an unbiased alternative MI objective and a more data-efficient K -sample estimator tailored for downstream RL tasks, which, to our knowledge, have not been addressed in previous research.

3 Preliminaries

Reinforcement learning. In Meta-RL, we assume an environment is a distribution ξ ( e ) of tasks e (e.g. uniform in our experiments). Each task e ∼ ξ ( e ) has a similar structure that corresponds to a Markov Decision Process (MDP) [Puterman, 2014], defined by M e = ( S , A , R, P e , γ ) , with a state space S , an action space A , a reward function R ( s t , a t ) where s t ∈ S and a t ∈ A , state transition dynamics P e ( s t +1 | s t , a t ) , and a discount factor γ ∈ [0 , 1) . In order to address the problem of zeroshot generalisation, we consider the transition dynamics P e ( s t +1 | s t , a t ) vary across tasks e ∼ ξ ( e ) according to multiple environmental features e = { e 0 , e 1 , ..., e N } that are not included in states s and can be continuous random variables, such as mass and friction, or discrete random variables, such as the cube's material. For instance, in a cube-moving environment (Figure 1), an agent has different tasks that are defined by different environmental features (e.g., mass and friction). The Meta-RL agent's goal is to learn a generalisable policy π that is robust to such dynamic changes. Specifically, given a set of training tasks e sampled from ξ train ( e ) , we aim to learn a policy that can maximise the discounted returns, arg max π E e ∼ ξ train ( e ) [ ∑ ∞ t =0 γ t R ( s t , a t ) | a t ∼ π ( a t | s t ) , s t +1 ∼ P e ( s t +1 | s t , a t )] , and can produce accurate control for unseen test tasks sampled from ξ test ( e ) .

Contrastive learning. In Meta-RL, the context encoder ψ ( c | τ c, 0: t ) first takes the trajectory τ c, 0: t = { s 0 , a 0 , r 0 , ..., s t } from the current episode as input and compresses it into a context embedding c [Fu et al., 2021]. Then, the policy π , conditioned on context embedding c , consumes the current state s t and outputs the action a t . As a key component, the quality of context embedding c can affect algorithms' performance significantly. MI is an effective measure of embedding quality [Goldfeld

et al., 2019], hence we focus on a context encoder that optimises the InfoNCE objective I InfoNCE ( x ; y ) , which is a K -sample estimator and lower bound of the MI I ( x ; y ) [Oord et al., 2019]. Given a query x and a set Y = { y 1 , ..., y K } of K random samples containing one positive sample y 1 and K -1 negative samples from the distribution p ( y ) , I InfoNCE ( x ; y ) is obtained by comparing pairs sampled from the joint distribution x, y 1 ∼ p ( x, y ) to pairs x, y k built using a set of negative examples y 2: K :

InfoNCE constructs a formal lower bound on the MI, i.e., I InfoNCE ( x ; y | ψ, K ) ≤ I ( x ; y ) [Guo et al., 2022, Chen et al., 2021]. Given two inputs x and y , their embedding similarity is f ψ ( x, y ) = e ψ ( x ) ⊤ · ψ ( y ) /β , where ψ is the context encoder that projects x and y into the context embedding space, the dot product is used to calculate the similarity score between ψ ( x ) , ψ ( y ) pairs [Wu et al., 2018, He et al., 2020], and β is a temperature hyperparameter that controls the sensitivity of the product. Some previous Meta-RL methods [Lee et al., 2020, Mu et al., 2022] learn a context embedding c by maximising I InfoNCE ( c ; τ c | ψ, K ) between the context c embedded from a trajectory in the current task, and the historical trajectories τ c under the same environmental features setting.

4 Skill-aware mutual information optimisation for Meta-RL

4.1 The log -K curse of K -sample MI estimators

In this section, we provide a theoretical analysis of the challenge inherent in learning a K -sample estimator for MI, commonly referred to as the log -K curse. Based on this theoretical analysis, we give insights to overcome this challenge. Given that we focus on the generalisation of RL, we only consider cases with a finite sample size of K . If a context encoder ψ in Equation (1) has sufficient training epochs, then I InfoNCE ( x ; y | ψ, K ) ≈ log K [Mnih and Teh, 2012, Guo et al., 2022]. Hence, the MI we can optimise is bottlenecked by the number of available samples, formally expressed as:

Lemma 1 Learning a context encoder ψ with a K -sample estimator and finite sample size K , we have I InfoNCE ( x ; y | ψ, K ) ≤ log K ≤ I ( x ; y ) , when x ̸ ⊥ ⊥ y (see proof in Appendix A).

We do not consider the case when x ⊥ ⊥ y , i.e., log K ≥ I ( x ; y ) = 0 ( ∀ K ≥ 1 ), because a MetaRL agent learns a context encoder by maximising MI between trajectories τ c and context embeddings c , which are not independent (as shown in Figure 2). While an unbounded sample size K for learning effective context embeddings is theoretically feasible and assumed in many studies Mu et al. [2022], Lee et al. [2020], it is often impractical in practice. Therefore, building on Lemma 1, we derive two key insights with limited samples: (1) generalising the current MI objective to be smaller than I ( x ; y ) (see Section 4.2); (2) developing a K -sample estimator tighter than I InfoNCE (see Section 4.3).

𝑰(𝒄; 𝝉

𝒄

)

𝑰

𝑰𝒏𝒇𝒐𝑵𝑪𝑬

(𝒄; 𝝅

𝒄

; 𝝉

𝒄

)

𝑰

𝑺𝒂𝑴𝑰

(𝒄; 𝝅

𝒄

; 𝝉

𝒄

)

𝑰

𝑺𝒂𝑵𝑪𝑬

(𝒄; 𝝅

𝒄

; 𝝉

𝒄

)

Training Epochs

Mutual Information

𝒍𝒐𝒈𝑲

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 1 (Skills) A policy π conditioned on a fixed context embedding c is defined as a skill π ( ·| c ) , abbreviated as π c . If a skill π c is conditioned on a state s t , we can sample actions a t ∼ π ( ·| c, s t ) . After sampling actions from π c at consecutive timesteps, we obtain a trajectory τ c,t : t T = { s t , a t , r t , s t +1 , . . . , s t + T , a t + T , r t + T } which demonstrates a consistent mode of behaviour.

After a limited amount of exploration, an agent should be able to infer the task (i.e., environmental features e = e 0 , e 1 , . . . , e N ) and adapt accordingly within the current episode. The context em-

Positive Sample Space

Negative Sample Space

Not Sample Space

(a) InfoNCE

(b) SaNCE

(c) Sa+InfoNCE

𝑐

𝜋

!

𝜏

!

Observed Variable

𝜏

!

!

𝜏

!

"

𝜏

!

#

𝜏

!

$

𝜏

!

!

"

𝜏

!

"

𝜏

!

#

𝜏

!

$

𝜏

!

!

#

𝑐

𝜋

!

𝜏

!

𝑐

𝜋

!

𝜏

!

𝜏

!

!

"

𝜏

!

!

#

𝜏

!

"

"

𝜏

!

"

#

𝜏

!

#

"

𝜏

!

#

#

𝜏

!

$

"

𝜏

!

$

#

bedding should encompass skill-related information, guiding the policy on when to explore new skills or switch between existing ones. We propose that the context encoder ψ should be trained by maximising the MI between the context embedding c , skills π c , and trajectories τ c . To this end, we propose a novel MI optimisation objective, Skill-aware Mutual Information (SaMI) , defined as:

SaMI is defined according to interaction information [McGill, 1954], serving as a generalisation of MI for three variables { c, π c , τ c } . Although we cannot evaluate p ( c, π c , τ c ) directly, we approximate it by Monte-Carlo sampling, using K samples from p ( c, π c , τ c ) . As illustrated in Figure 3, a context encoder ψ trained with the objective of maximising I SaMI ( c ; π c ; τ c ) converges more quickly, as I SaMI ( c ; π c ; τ c ) ≤ I ( c ; τ c ) (see proof in Appendix B). By focusing more on skill-related information, I SaMI enables agents to autonomously discover a diverse range of skills for handling multiple tasks.

4.3 Skill-aware noise contrastive estimation: a tighter K -sample estimator

Despite InfoNCE's success as a K -sample estimator for approximating MI [Laskin et al., 2020, Eysenbach et al., 2022], its learning efficiency plunges due to limited numerical precision, which is called the log -K curse, i.e., I InfoNCE ≤ log K ≤ I SaMI [Chen et al., 2021] (see proof in Appendix B). When K → + ∞ , we can expect I InfoNCE ≈ log K ≈ I SaMI [Guo et al., 2022]. However, increasing K is too expensive, especially in complex environments with enormous negative sample space. To address this, we propose a novel K -sample estimator that requires a significantly smaller sample size K ≪ + ∞ . First, we define K ∗ :

Definition 2 ( K ∗ ) K ∗ = | c | · | π c | · M 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 c , | π c | represents the number of different skills π c , and M is a natural number.

To ensure that I InfoNCE is a tight bound of I SaMI, we require that I InfoNCE ≈ log K ≈ I SaMI when K → K ∗ . Under the definition of K ∗ , the replay buffer can be divided according to the different context embeddings c and skills π c (i.e., observing context embeddings c and skills π c ). In real-world robotic control tasks, the sample space size significantly increases due to multiple environmental features e = { e 0 , e 1 , ..., e N } . Taking the sample space of InfoNCE as an example (Figure 4(a)), in the current task e 1 with context embedding c 1 , positive samples are trajectories τ c 1 generated after executing the skill π c 1 in task e 1 , and negative samples are trajectories { τ c 2 , ... } from other tasks { e 2 , ... } . The permutations and combinations of N environmental features lead to an exponential growth in task number | c | , which in turn results in an increase of sample space K ∗ InfoNCE = | c | · | π | · M .

We introduce a tight K -sample estimator, Skill-aware Noise Contrastive Estimation (SaNCE) , which is used to approximate I SaMI ( c ; π c ; τ c ) with K ∗ SaNCE < K ∗ InfoNCE . For SaNCE, both positive samples τ + c 1 and negative samples τ -c 1 are sampled from the current tasks e 1 , but are generated by executing positive skills π + c 1 and negative skills π -c 1 , respectively. Here, a positive skill is intuitively defined by whether it is optimal for the current task e , with a more formal definition provided in Section 4.4. For instance, in a cube-moving task under a large friction setting, the agent executes

Context Encoder

𝜓

for a cube-moving task

𝑒

!

with high friction

Reinforcement

Learning

Negative samples 𝜏

"

!: #

,$:&

'

𝑤𝑖𝑡ℎ 𝑚𝑖𝑛 ∑

()*

*+,

𝑅(𝑠

(

, 𝑎

(

)

Ranked by

∑

()*

*+,

𝑅(𝑠

(

, 𝑎

(

)

Positive sample𝑠

𝜏

"

$

+

Momentum

Context Encoder

𝜓

∗

Positive embedding

s

𝜓

∗

(𝜏

"

$

+

)

Negative embeddings

𝜓

∗

(𝜏

"

!: #

,$:&

'

)

Quer𝑖𝑒𝑠 𝜓(𝜏

"

$

+

)

Pick&Place

Push

𝑤𝑖𝑡ℎ 𝑚𝑎𝑥 ∑

()*

*+,

𝑅(𝑠

(

, 𝑎

(

)

Positive Samples 𝜏

"

$

+

𝑤𝑖𝑡ℎ 𝑚𝑎𝑥 ∑

()*

*+,

𝑅(𝑠

(

, 𝑎

(

)

a skill π + c after several iterations of learning, and obtains corresponding trajectories τ + c where the cubes leave the table surface. This indicates that the skill π + c is Pick&Place and other skills π -c may include Push or Flip (flipping the cube to the goal position), with corresponding trajectories τ -c where the cube remains stationary or rolls on the table. Formally, we can optimise the K -sample lower bound I SaNCE to approximate I SaMI:

where f ψ ( c 1 , π c 1 , τ c 1 ) = e ψ ( τ c 1 ) ⊤ · ψ ∗ ( τ c 1 ) /β . The query c 1 = ψ ( τ c 1 ) is generated by the context encoder ψ . For training stability, we use a momentum encoder ψ ∗ to produce the positive and negative embeddings. SaNCE significantly reduces the required sample space size K ∗ SaNCE by sampling trajectories τ c based on different skills π c (Figure 4(b)) in task e 1 , so that K ∗ SaNCE = | c | · | π c | · M = | π c 1 | · M ≤ K ∗ InfoNCE ( | c | = | c 1 | = 1 ). Therefore, I SaNCE satisfies Lemma 2:

Lemma 2 With a context encoder ψ and finite sample size K , we have I InfoNCE ( c ; π c ; τ c | ψ, K ) ≤ I SaNCE ( c ; π c ; τ c | ψ, K ) ≤ log K ≤ I SaMI ( c ; π c ; τ c ) ≤ I ( c ; τ c ) . (see proof in Appendix B)

SaNCE can be used alone or combined with other optimisation objectives to train context encoders in Meta-RL algorithms. For instance, integrating SaNCE with InfoNCE diversifies the negative sample space, with K ∗ Sa+InfoNCE = ( ∑ | c | i =1 | π -c i | + | π + c 1 | ) · M . The sample space for I Sa+InfoNCE is depicted in Figure 4(c) 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 - a slight difference in states can make two skills distinguishable, and not necessarily in a semantically meaningful way. Consequently, we do not directly teach any of these skills or assume any prior skill distribution. Diverse skills naturally emerge from the incentives of the SaMI learning objective in a multi-task setting, driven by the inherent need to develop generalisable skills. For example, in high-friction tasks, the agent must acquire the Pick&Place skill to avoid large frictional forces, whereas in high-mass tasks, the agent must learn the Push skill since it cannot lift the cube. In each task, we only identify whether the skills

(b) Success rate

(d) Pick&Place probability

(c) Push probability

(a) UMAP visualisation

1.0

20

friction:0.1,mass:10

friction:0.1,mass:30

0.9

0.9

0.9

0.9

0.1

0.2

0.2

0.4

1.0

0.8

0.8

0.5

0.0

friction:10,mass:1

0.8

15

10

values

S

0

1.0

0.9

0.8

0.1

values

S

1.0

0.2

0.2

0.3

1.0

values

0

1

0.8

0.8

0.4

0.0

0.6

5

iction

0

5

0.9

0.7

0.3

0.1

iction

5.0

0.2

0.2

0.2

0.8

iction

5

0

0.8

0.7

0.4

0.0

0.4

0-

Fric

1

0.9

0.6

0.2

0.1

E

10.0

0.2

0.1

0.3

0.8

E

10

0.8

0.7

0.4

0.0

5

0.2

10

30

0.

0.9

0.5

0.2

0.1

30

0

0.1

0.1

0.2

0.8

30

0.

0.8

0.8

0.4

0.0

15

1

5

10

30

1

5

10

30

1

5

10

30

0.0

10

0

10

20

Mass values

Mass values

Mass values

are optimal; in this way, under a multi-task setting, the agent will acquire a set of general skills that are applicable to many tasks.

In a given task e , positive skills π + c are defined as optimal skills achieving highest return ∑ t + T i = t R ( s i , a i ) , whereas negative skills π -c are those that result in lower returns. As a result, the positive sample τ + c consists of trajectories generated by the skills with the highest ranked returns, while the negative samples correspond to those with the lowest returns. This straightforward approach of selecting positive samples based on the ranked highest return effectively aligns with positive skills and mitigates the challenge of hard negative examples [Robinson et al., 2021]. The SaNCE loss is then minimised to bring the context embeddings of the highest return trajectories closer while distancing those of negative trajectories. By the end of training, the top-ranked trajectories in the ranked replay buffer correspond to positive samples τ + c with high returns, while the lower-ranked trajectories represent negative samples τ -c with low returns. However, at the beginning of training, it is likely that all trajectories have low returns. Therefore, our SaNCE loss is a soft variant of the K -sample SaNCE:

where ||·|| L 2 represents the Euclidean distance [Tabak, 2014]. Figure 5 provides a practical framework of SaNCE, with a cube-moving example task e 1 under high friction. In task e 1 , the positive skill π + c 1 is the Pick&Place skill, which is used to generate queries ψ ( τ + c 1 ) and positive embeddings ψ ∗ ( τ + c 1 ) ; after executing Push skill we get negative samples τ -c 1 and negative embeddings ψ ∗ ( τ -c 1 ) .

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 1 and 2, and Section 5.4)?

5.1 Experimental setup

Modified benchmarks with multiple environmental features. 1 We evaluate our method on two benchmarks, Panda-gym [Gallouédec et al., 2021] and MuJoCo [Todorov et al., 2012] (details in Sections 5.2 and 5.3). The benchmarks are modified to be influenced by multiple environmental features, which are sampled at the start of each episode during meta-training and meta-testing. During meta-training, we uniform-randomly select a combination of environmental features from a training task set. At test time, we evaluate each algorithm in 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 a 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.

Baselines. The loss function of the context encoder in all algorithms consists of two key components: the RL loss and the contrastive loss. The RL loss L RL , which is the same across all methods, corresponds to the RL value function loss. Our primary comparison focuses on the contrastive loss, taking the form of either SaNCE, InfoNCE, or no contrastive loss. Accordingly, we select baselines based on contrastive loss: CCM [Fu et al., 2021], which utilises InfoNCE, and TESAC [Yu et al., 2020], which relies solely on the RL loss, allowing assessment of the context encoder without contrastive loss. Since CCM and TESAC use an RNN encoder, we also include PEARL [Rakelly et al., 2019], which utilises an MLP context encoder and follows the similar RL loss. Additionally, Appendix G includes comparisons with DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020], using the same environmental setup in the MuJoCo benchmark.

Our method. 2 We use Soft Actor-Critic (SAC) [Haarnoja et al., 2018] as the base RL algorithm, training agents for 1.6 million timesteps in each environment (details in 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 Meta-RL algorithms: (1) SaTESAC is TESAC with SaNCE, which uses SaNCE for contrastive learning, using a | c | times smaller sample space (Figure 4(b)); (2) SaCCM is CCM with SaNCE, where the contrastive learning combines InfoNCE and SaNCE, as shown in Figure 4(c).

5.2 Panda-gym

Task description. Our modified Pandagym benchmark involves a robot arm control task using the Franka Emika Panda [Gallouédec et al., 2021], where the robot moves a cube to a target position. Unlike previous works, we simultaneously modify multiple environmental features (cube mass and table friction) that characterise the transition dynamics, and the robot can flexibly execute different skills (Push and Pick&Place) for different tasks. This environment requires high skill diversity; for instance, the agent must use Pick&Place in high-friction tasks and Push in high-mass tasks.

Table 1: Comparison of success rate ± standard deviation with baselines in Panda-gym (over 5 seeds). Bold text signifies the highest average 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 paried t-tests with significance threshold of p < 0 . 05 .

Results and skill analysis. As shown in Table 1, SaTESAC and SaCCM achieve superior generalisation performance compared to PEARL, TESAC, and CCM, with a smaller sample space. The t-test results in Table 1 show that SaMI significantly improves success rates across training, moderate, and extreme test sets at a 0.05 significance level. Video demos 2 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] (Figure 6(a) and Appendix F) and t-SNE [Van der Maaten and Hinton, 2008] (Appendix F), 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 heatmaps [Waskom, 2021] were used to visualise executed skills. Figure 6 shows that SaCCM agents learned the Push skill for large cube masses (30 Kg, 10 Kg) and Pick&Place for smaller masses, while CCM showed no clear skill grouping (Figure 16 in Appendix F.1). Overall, SaMI incentivises agents to autonomously learn diverse skills, enhancing generalisation across a wider range of tasks. Specifically, through the cycle of effective exploration, context inference, and adaptation, diverse skills emerge solely from the data. Further visualisation results are in Appendix F.

5.3 MuJoCo

Task description. We extended the modified MuJoCo benchmark introduced in DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020]. It contains ten typical robotic control environments based on the MuJoCo physics engine [Todorov et al., 2012]. Hopper, Walker, Half-cheetah, Ant, HumanoidStandup, and SlimHumanoid are influenced by continuous environmental features (i.e., mass, damping) that affect transition dynamics. Crippled Ant, Crippled Hopper, Crippled Walker, and

Crippled Half-cheetah are more challenging due to the addition of discrete environmental features (i.e., randomly crippled leg joints), requiring agents to master different skills (e.g., switching from running to crawling after a leg is crippled).

Results and skill analysis. Table 2 shows the average return of our method and baselines on training and test tasks. SaTESAC and SaCCM achieved higher returns in most tasks, except for Ant, Half-Cheetah, and Hopper, where only a single skill was needed. For instance, the Hopper robot learned to hop forward, adapting to different mass values. When environments become complex and require diverse skills for different tasks (Crippled Ant, Crippled Hopper, Crippled Half-Cheetah, SlimHumanoid, HumanoidStandup, and Crippled Walker), SaNCE brings significant improvements. For example, when the Ant robot has 3 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 2 healthy legs. This aligns with the results in Table 2, where SaMI significantly improved performance in extreme test sets. In 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 2 and visualisation results in Appendix F.2 show different skills in all environments.

5.4 Analysis of the log -K curse in sample-limited scenarios

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, InfoNCE's sampling space can be 25 times larger than SaNCE's. As shown in Tables 1 and 2, RL algorithms using SaNCE achieve better or comparable performance with significantly fewer negative samples ( K ) than InfoNCE. This suggests SaNCE effectively addresses the log -K curse, and the SaMI objective helps the contrastive context encoder extract critical information for downstream RL tasks.

The number of negative samples K is influenced by two hyperparameters: buffer size , which determines the negative sample space, and contrastive batch size , which controls the number of samples used to train the contrastive context encoder per update. We analysed these hyperparameters further, and as shown in Figure 7, reductions in buffer and contrastive batch size do not significantly impact the average return for SaCCM and SaTESAC, which maintain state-of-the-art performance with small buffers and batch sizes. The results in Table 2 correspond to a buffer size of 100,000 and a batch size of 12. Results across all environments (refer to Appendix E.2) show that SaNCE exhibits low sensitivity to K , highlighting its potential to overcome the log -K curse.

6 Conclusion and future work

Zero-shot generalisation has been a longstanding challenge concerning Meta-RL agents, with skill diversity and sample efficiency being key to generalising to previously unseen environments. In this paper, we proposed Skill-aware Mutual Information (SaMI) to learn context embeddings for zero-shot generalisation in downstream RL tasks, and Skill-aware Noise Contrastive Estimation (SaNCE) to optimise SaMI and overcome the log -K curse, along with a practical skill-aware trajectory sampling strategy. Experimental results showed that RL algorithms equipped with SaMI achieved state-ofthe-art performance in MuJoCo and Panda-gym benchmarks, particularly in zero-shot generalisation within more complex environments. During the zero-shot generalisation, when faced with an unseen task, SaMI assists agents in exploring effectively, inferring context, and rapidly adapting their skills within the current episode. SaNCE's optimisation uses a significantly smaller negative sample space than baselines, and our analysis on buffer and contrastive batch sizes demonstrated its effectiveness in addressing the log -K curse.

Given that environmental features are often interdependent, such as a cube's material correlating with friction and mass, SaMI does not introduce independence assumptions like DOMINO [Mu et al., 2022]. Therefore, future work will focus on verifying and enhancing SaMI's potential in more complex tasks where environmental features are correlated [Dunion et al., 2023a]. This will contribute to our ultimate goal: developing a generalist and versatile agent capable of working across multiple tasks and even real-world tasks in the near future.

(b)

(a)

12000-

10000-

8000

R

e

6000

a

e

4000

TESAC

CCM

2000

SaTESAC

SaCCM

0

1000

10000

100000

400000

11000-

E

10000

R

9000

g

8000

7000

CCM

6000

SaTESAC

SaCCM

5000

8

16

128

512

ContrastiveBatchSize

References

Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc 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 Saunshi. A theoretical analysis of contrastive unsupervised representation learning. arXiv preprint arXiv:1902.09229 , 2019.

Philip Bachman, R 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 learners with flatnce. arXiv preprint arXiv:2107.01152 , 2021.

Ignasi Clavera, Anusha Nagabandi, Simin Liu, Ronald S. Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta-reinforcement 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 to adapt in dynamic, real-world environments through meta-reinforcement 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 https://openreview.net/forum?id=BJHRaK1PG .

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 McInroe, 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 https://arxiv.org/abs/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 , 35:35603-35620, 2022.

Jörg K.H. Franke, Gregor Koehler, André Biedenkapp, and Frank Hutter. Sample-efficient 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: an 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: OpenSource 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 in reinforcement learning via data-regularised environment design. In International Conference on Machine Learning (ICML) , 2024.

Ziv Goldfeld, Ewout van den Berg, Kristjan Greenewald, Igor Melnyk, Nam Nguyen, Brian Kingsbury, and Yury Polyanskiy. Estimating information flow in deep neural networks, 2019. URL https: //arxiv.org/abs/1810.05728 .

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 on computer vision and pattern recognition , pages 9729-9738, 2020.

R Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, 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. In 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 arXiv:2102.10774 , 2021.

Qiang Li. Functional connectivity inference from fmri data using multivariate information measures. Neural Networks , 146:85-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 https://arxiv.org/abs/1802.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:1206.6426 , 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 , 19(4):675-685, 2003.

Kento Nozawa and Issei Sato. Understanding negative samples in instance discriminative selfsupervised representation learning. Advances in 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.03748 , 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 51715180. 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-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 22(268):1-8, 2021. URL http://jmlr.org/papers/v22/20-1364.html .

Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning , pages 5331-5340. PMLR, 2019.

John A Rice and John A Rice. Mathematical statistics and data analysis , volume 371. Thomson/Brooks/Cole Belmont, CA, 2007.

Joshua Robinson, Ching-Yao Chuang, Suvrit Sra, and Stefanie Jegelka. Contrastive learning with hard negative samples, 2021. URL https://arxiv.org/abs/2010.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 , 33:12968-12979, 2020.

Lucy Xiaoyang Shi, Joseph J. Lim, and Youngwoon Lee. Skill-based model-based reinforcement learning. In 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-reinforcement learning with self-supervised trajectory contrastive learning, 2021.

Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software , 6(60): 3021, 2021. doi: 10.21105/joss.03021. URL https://doi.org/10.21105/joss.03021 .

Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via nonparametric instance discrimination. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 3733-3742, 2018.

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-1100. PMLR, 2020.

Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. In International Conference on Learning Representations , 2019.

A Proof of Lemma 1

Given a query x and a set Y = { y 1 , . . . , y K } of K random samples, containing one positive sample y 1 and K -1 negative samples drawn from the distribution p ( y ) , a K -sample InfoNCE estimator is obtained by comparing pairs sampled from the joint distribution ( x, y 1 ) ∼ p ( x, y ) with pairs ( x, y k ) , constructed using the set of negative examples y 2: K . InfoNCE compares the positive pairs ( x, y 1 ) with the negative pairs ( x, y k ) , where y k ∼ y 2: K , as follows:

Step 1. Let us prove that the K -sample InfoNCE estimator is upper-bounded by log K . According to Mu et al. [2022], f ψ ( x,y 1 ) ∑ K k =1 f ψ ( x,y k ) = f ψ ( x,y 1 ) f ψ ( x,y 1 )+ ∑ K k =2 f ψ ( x,y k ) ≤ 1 . So we have:

Hence, we have I InfoNCE ( x ; y | ψ, K ) ≤ log K .

Step 2. We have the I ( x ; y ) ≥ I InfoNCE ( x ; y | ψ, K ) according to:

Proposition 1 [Poole et al., 2019] A K -sample estimator is an asymptotically tight lower bound on the MI, i.e.,

Proof . See Poole et al. [2019] for a neat proof of how the multi-sample estimator (e.g., InfoNCE) lower bounds MI.

Step 3. In this research, the context encoder ψ in f ψ ( x, y ) is implemented using an RNN to approximate p ( y | x ) p ( y ) [Oord et al., 2019]. With a sufficiently powerful deep learning model for ψ and a finite sample size K , such that I ( x ; y ) ≥ log K , we can reasonably expect that I InfoNCE ≈ log K after a few training epochs. Therefore, during training, when K ≪ + ∞ , we always have I ( x ; y ) ≥ log K .

Proof . See Chen et al. [2021] for more detailed proof.

Step 4. Let us prove that the K -sample InfoNCE bound is asymptotically tight. The specific choice of the context encoder ψ influences the tightness of the K -sample NCE bound. InfoNCE [Oord et al., 2019] sets f ψ ( x, y ) ∝ p ( y | x ) p ( y ) to model a density ratio that preserves the MI between x and y , where ∝ stands for 'proportional to' (i.e., up to a multiplicative constant). Substituting

f ψ ( x, y ) = f ∗ ψ ( x, y ) = p ( y | x ) p ( y ) into InfoNCE, we obtain:

Now taking K → + ∞ , the last two terms cancel out.

Putting it together. Combining I ( x ; y ) ≥ log K with Proposition 1 and Equation (6), we have Lemma 1:

Moreover, according to Equation (7), as the sample size K → + ∞ , the K -sample InfoNCE bound becomes sharp and approaches the true MI I ( x ; y ) , i.e., I InfoNCE ( x ; y | ψ, K ) ≈ log K ≈ I ( x ; y ) .

B Proof for Lemma 2

Step 1. According to Lemma 1, we have I InfoNCE ( c ; π c ; τ c | ψ, K ) ≤ log K ≤ I SaMI ( c ; π c ; τ c ) (shown in Figure 3).

Step 2. Let us prove that SaNCE is a K -sample SaNCE estimator and is upper bounded by log K . Since f ψ ( c,π c ,τ + c ) f ψ ( c,π c ,τ + c )+ ∑ K k =2 f ψ ( c,π c ,τ -c,k ) ≤ 1 [Mu et al., 2022], we have:

I

(

c

;

π

c

;

τ

c

|

ψ, K

)

Thus, we obtain I SaNCE ( c ; π c ; τ c | ψ, K ) ≤ log K , similar to Equation (6).

Step 3. With the definition of K ∗ , we can prove that I InfoNCE ( c ; π c ; τ c | ψ, K ) ≤ I SaNCE ( c ; π c ; τ c | ψ, K ) with the same sample size K . In task e 1 , SaNCE obtains positive and negative samples from the current task e 1 . Since the variable c = c 1 is constant, we have:

The required sample size is K ∗ SaNCE = | c 1 | · | π | · M = | π | · M . As K → K ∗ SaNCE , we have I SaNCE ( c ; π c ; τ c | ψ, K ) ≈ I SaMI ( c ; π c ; τ c ) . Correspondingly, for InfoNCE, in the current task e 1 with context embedding c 1 , positive samples are trajectories τ 1 generated after executing the skill π 1 in task e 1 , while negative samples are trajectories { τ -c 2 , ... } from other tasks { e 2 , ... } . Under the definition of K ∗ , we have:

I

(

c

;

π

c

;

τ

c

|

ψ, K

)

where K ∗ InfoNCE = | c | · | π | · M ≈ | c | · K ∗ SaNCE . In real-world robotic control tasks, the sample space size increases significantly due to multiple environmental features e = { e 0 , e 1 , . . . , e N } . The number of different tasks | c | grows exponentially due to the permutations and combinations of the N environmental features. When the current task e 1 has context embedding c 1 , the c 2: | c | refer to the context embeddings for the other tasks. As K → K ∗ InfoNCE , we have I InfoNCE ( c ; π c ; τ c | ψ, K ) ≈ I SaMI ( c ; π c ; τ c ) . Thus, during the training process, I InfoNCE ( c ; π c ; τ c | ψ, K ) ≤ I SaNCE ( c ; π c ; τ c | ψ, K ) with the same sample size K .

Step 4. According to the definition of SaMI in Equation (2), we have I SaMI ( c ; π c ; τ c ) ≤ I ( c ; τ c ) , as illustrated using Venn diagrams in Figure 8 (a) and (b). SaMI is formulated based on interaction information [McGill, 1954], which aims to capture the relationships among multivariate variables by quantifying the amount of information (redundancy or synergy) shared among three variables. Interaction information has been less extensively studied, partly due to its challenging interpretation from both information theory and neuroscience perspectives, as it can be either positive or negative [Li, 2022].

𝐻 𝑐

𝐻 𝜏

!

𝐻 𝜏

!

𝐻 𝑐

𝐻 𝜋

!

(b)

(a)

(c)

𝐼(𝑐; 𝜏

!

)

𝐼

"#$%

(𝑐; 𝜋

!

; 𝜏

!

)

In the Meta-RL setting, the MDP (causal) graph structure, shown in Figure 8 (c), illustrates that the context embedding c , skill π c , and trajectory τ c form a common-cause structure, where c acts as a shared cause influencing both the skill π c and the trajectory τ c . Consequently, I ( π c , τ c | c ) < I ( π c , τ c ) . Therefore, I SaMI ( c ; π c ; τ c ) > 0 is guaranteed, making it interpretable in real-world robotic control tasks.

Putting it together. Thus, we establish Lemma 2: we always have I InfoNCE ( c ; π c ; τ c | ψ, K ) ≤ I SaNCE ( c ; π c ; τ c | ψ, K ) ≤ log K ≤ I SaMI ( c ; π c ; τ c ) ≤ I ( c ; τ c ) , as shown in Figure 3, while learning a skill-aware context encoder ψ with the SaNCE estimator. Since K ∗ SaNCE ≪ K ∗ InfoNCE , I SaNCE ( c ; π c ; τ c | ψ, K ) serves as a much tighter lower bound for the true I SaMI ( c ; π c ; τ c ) than I InfoNCE ( c ; π c ; τ c | ψ, K ) .

C Sample size of I Sa+InfoNCE

In this section, we illustrate the sample size of I Sa+InfoNCE ( c ; π c ; τ c | ψ, K ) . Sa+InfoNCE incorporates SaNCE into InfoNCE, using positive samples τ + c 1 from task e 1 after executing skill π + c 1 , and negative

samples are trajectories τ -c 1: K from executing skills π -c 1: K in tasks e 1: K , respectively. Therefore, this approach is equivalent to first observing the variable c and then observing the variable π c , i.e., sampling from the distribution p ( π c , τ c | c ) p ( c ) . We have:

It should be noted that such a combination increases the size of the negative sample space, i.e.,

K ∗ Sa+InfoNCE = ( ∑ | c | i =1 | π -c i | + | π + c 1 | ) · M ≥ K ∗ SaNCE . The misaligned bars in Figure 4 (c) illustrate that negative sample spaces may vary across tasks. This variation arises because we define negative samples as trajectories with low returns, making the size of the negative sample space influenced by sampling randomness. With the same number K of samples, I Sa+InfoNCE is less precise and looser than I SaNCE. Therefore, a trade-off between sample diversity and the precision of the K -sample estimator is required.

D Environmental setup

(a) Panda-gym

Moderate

Test Tasks

(b) Task setting

Training Tasks

Extreme Test Tasks

0.1

Friction values

30.0 10.05.01.0

1

5

10

30

Mass values

D.1 Modified Panda-gym

We modified the original Pick&Place task in Panda-gym [Gallouédec et al., 2021] by setting the z dimension (i.e., the desired height) of the cube's goal position to 0 3 and maintaining the freedom of the grippers 4 , allowing the agent to explore whether it should push or grasp the cube. Skills in this benchmark are defined as:

Pick&Place skill : This skill specifically refers to the agent using the gripper to grasp the cube, lift it off the table, and place it in the goal position. We determine the Pick&Place skill by detecting no contact points between the table and the cube, two contact points between the robot's end effector and the cube, and the cube's height being greater than half its width.

Push skill : This skill involves the agent moving the cube on the table to the goal position, either by dragging or sliding it. We confirm the Push skill by detecting that the cube's height equals half its width.

Other skills : Any behaviour modes other than Pick&Place and Push are classified as other skills.

Some elements in the RL framework are defined as follows:

State space : We use feature vectors that contain the cube's position (3 dimensions), cube rotation (3 dimensions), cube velocity (3 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-effector's position, and the last dimension represents the change in the gripper's width.

During training, we randomly select a combination of environmental features from a training set by sampling combinations from the following sets: mass = 1 . 0 and friction ∈ { 0 . 1 , 1 . 0 , 5 . 0 , 10 . 0 } ; mass ∈ { 1 . 0 , 5 . 0 , 10 . 0 } and friction = 0 . 1 . At test time, we evaluate each algorithm on all tasks from the moderate test setting, where mass ∈ { 5 . 0 , 10 . 0 } and friction ∈ { 1 . 0 , 5 . 0 , 10 . 0 } (shown in Figure 9(b)), and on all tasks from the extreme test setting: mass = 30 . 0 and friction ∈ { 0 . 1 , 1 . 0 , 5 . 0 , 10 . 0 , 30 . 0 } ; mass ∈ { 1 . 0 , 5 . 0 , 10 . 0 , 30 . 0 } and friction = 30 . 0 (shown in Figure 9(b)).

D.2 Modified MuJoCo

We extended the modified MuJoCo benchmark introduced in DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020]. In our extension, four new environments were added (Walker, Crippled Hopper, Crippled Walker, HumanoidStandup), compared to the original benchmark. Additionally, in our experiments, we used a different task set design (Table 3) than those used in the DOMINO and CaDM papers. For Hopper, Walker, Half-Cheetah, Ant, HumanoidStandup, and SlimHumanoid, we used the MuJoCo physics engine environments and implemented settings from Clavera et al. [2019b] and Seo et al. [2020], scaling the mass of every rigid link by a scale factor m and the damping of every joint by a scale factor d . For Crippled Ant, Crippled Hopper, Crippled Walker, and Crippled Half-Cheetah, we used the implementation available from Seo et al. [2020], scaled the mass of each rigid link by a factor of m , scaled the damping of each joint by a factor of d , and randomly selected joints to be uncontrollable (i.e., masking the corresponding actions with 0). Generalisation performance is measured in two different regimes: moderate and extreme, where the moderate regime draws environmental features from a range closer to the training range compared to the extreme regime. Our settings for training, extreme, and moderate test tasks are provided in Table 3.

Ant

Crippled Ant

Half-cheetah

Crippled

Half-cheetah

SlimHumanoid

Hopper

Humanoid

Standup

Walker

Crippled

Hopper

Crippled

Walker

̸

̸

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 1 and 2. Our codebase is built on top of the publicly released implementation of Stable Baselines3 by Raffin et al. [2021] and the implementation of InfoNCE by Oord et al. [2019]. A public, open-source implementation of SaMI is available at https://github.com/uoe-agents/SaMI.

Base algorithm. We use SAC [Haarnoja et al., 2018] for the downstream evaluation of the learned context embedding. SAC is an off-policy actor-critic method that leverages the maximum entropy framework for soft policy iteration. At each iteration, SAC performs soft policy evaluation and improvement steps. We use the same SAC implementation across all baselines and other methods. During the meta-training phase, we trained agents for 1.6 million timesteps in each environment on the Panda-gym and MuJoCo benchmarks. For meta-testing, we evaluated 100 episodes in each environment, with tasks randomly sampled from the moderate and extreme task sets.

Encoder architecture. In our method, the context encoder ψ is modelled as a Long Short-Term Memory (LSTM) network that produces a 128-dimensional hidden state vector, which is subsequently processed through a single-layer feed-forward network to generate a 6-dimensional context embedding. We aim for the agent to complete the three steps of "explore effectively, infer, adapt" within an episode. Therefore, we initialise the hidden state and cell state of the LSTM to zero at the start of each episode. Both the actor and critic use the same context encoder to embed trajectories. For contrastive learning, SaNCE utilises a momentum encoder ψ ∗ to generate positive and negative context embeddings [Laskin et al., 2020, He et al., 2020]. Formally, denoting the parameters of ψ as θ ψ and those of ψ ∗ as θ ψ ∗ , we update θ ψ ∗ as follows:

Here m ∈ [0 , 1) is a soft-update rate. Only the parameters θ ψ are updated by back-propagation. The momentum update in Equation (13) makes θ ψ ∗ evolve more smoothly by having them slowly track the θ ψ with m ≪ 1 (e.g., m = 0 . 05 in this research). This means that the target values are constrained to change slowly, greatly improving the stability of learning.

Hyperparameters. A full list of hyperparameters is displayed in Table 4.

Hardware. For each experiment run we use a single NVIDIA Volta V100 GPU with 32GB memory and a single CPU.

Algorithm 1 SaNCE Meta-training

Algorithm 2 SaNCE Meta-testing

E Additional results

E.1 Balance contrastive and RL updates: loss coefficient α

While past work has optimised hyperparameters to balance the contrastive loss coefficient α relative to the RL objective [Jaderberg et al., 2016, Bachman et al., 2019], we use both the contrastive and RL objectives with equal weight, setting α = 1 . 0 for the MuJoCo benchmark and α = 0 . 01 for the Panda-gym benchmark. Additionally, we analyse the effect of the loss coefficient α for CCM, SaTESAC, and SaCCM in the MuJoCo (Figure 12) and Panda-gym (Figure 11) benchmarks.

(a) Training tasks

(b) Test tasks (moderate)

(c) Test tasks (extreme)

0.6

0.5

CCM

CCM

CCM

SaTESAC

SaTESAC

SaTESAC

0.8

SaCCM

0.5

SaCCM

0.4

SaCCM

Success Rate

Rate

0.4

SuccessRate

0.6

0.3

SuccessR

0.3

0.4

0.2

0.2

0.2

0.1

0.1

0.0-

1.0

0.1

0.05

0.01

0.0-

1.0

0.1

0.05

0.01

0.0-

1.0

0.1

0.05

0.01

losscoefficientα

losscoefficientα

(a) Half-cheetah

(b) Crippled Half-cheetah

(c) SlimHumanoid

(d) Ant

(e) Crippled Ant

(f) Hopper

8000-

6000

7000

10000

6000-

5000

Return

Return

Return

8000

5000

4000

Average

6000

3000

3000

4000-

2000-

2000

1000-

CCM

SaTESAC

1000-

CCM

SaTESAC

2000-

CCM

SaTESAC

SaCCM

SaCCM

SaCCM

0

0.1

1.0

10.0

0

0.1

1.0

10.0

0

0.1

1.0

10.0

loss coefficient α

1750-

CCM

800-

1600

1500

SaTESAC

700

1400

SaCCM

Return

1250

Return

600

Return

1200-

Average

1000

AverageR

500

AverageR

1000

400

800

750

300

600-

500

200-

400

CCM

CCM

250

100

SaTESAC

200-

SaTESAC

SaCCM

SaCCM

0

0.1

1.0

10.0

0-

0.1

1.0

10.0

0

0.1

1.0

10.0

loss coefficient α

E.2 Result of log -K curse analysis

E.2.1 Buffer size

(a) Half-cheetah

(b) Crippled Half-cheetah

(c) SlimHumanoid

(d) Ant

(e) Crippled Ant

(f) Hopper

7000

TESAC

7000

6000

CCM

12000-

Return

SaTESAC

Average Return

10000

5000

SaCCM

5000

8000

4000

6000

4000

TESAC

3000

4000-

TESAC

3000

CCM

CCM

SaTESAC

2000

2000-

SaTESAC

2000-

SaCCM

1000

0-

SaCCM

1000

10000

100000400000

1000

10000100000400000

1000

10000100000400000

BufferSize

BufferSize

BufferSize

2000-

TESAC

TESAC

1600

1750-

CCM

1200

CCM

1500-

SaTESAC

SaCCM

Average Return

1000

SaTESAC

SaCCM

Average Return

1400

1250

1300

800

1200

750

600

1100

500

400

1000

CCM

TESAC

250-

200-

900

SaTESAC

°1000

800-

SaCCM

10000100000400000

BufferSize

1000

10000100000400000

BufferSize

1000

10000

BufferSize

100000400000

E.2.2 Contrastive batch size

(d) Ant

(e) Crippled Ant

(f) Hopper

(a) Half-cheetah

(b) Crippled Half-cheetah

(c) SlimHumanoid

8000

6000

11000

7500-

Return

7000-

Return

10000

6500

5000-

Average

9000

4500

8000-

7000

CCM

CCM

CCM

5000

SaCCM

SaTESAC

3500-

SaCCM

SaTESAC

6000

SaTESAC

SaCCM

8

16

128

512

8

16

128

512

5000

8

16

128

512

900-

1540

1200

Average Return

Return

800

Average Return

1520-

1000

700

1500-

1480-

800

1460-

600

CCM

SaTESAC

400-

CCM

SaTESAC

1440

1420-

CCM

SaTESAC

SaCCM

300

SaCCM

SaCCM

8

16

128

512

8

16

128

512

8

16

128

512

F Further skill analysis

F.1 Panda-gym

F.1.1 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 16). 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 1(a)), whereas with lower masses, the agent learned the Pick&Place skill. However, as shown in Figure 16(b), CCM did not display clear skill grouping. This indicates that SaMI extracts high-quality skill-related information from the trajectories and enables agents to autonomously discover a diverse range of skills for handling multiple tasks.

(b) CCM

(a) TESAC

(d) SaCCM

(c) SaTESAC

15

friction:0.1,mass:10

friction:0.1,mass:30

friction:10,mass:1

10

friction:0.1,mass:10

friction:0.1,mass:30

friction:10,mass:1

0

5

10

15

8

10

12

11

friction:0.1,mass:10

friction:0.1,mass:10

10

friction:0.1,mass:30

friction:10,mass:1

20

friction:0.1,mass:30

friction:10,mass:1

15

10

8

5

10

15

8

10

12

14

10

0

10

20

(d) SaCCM

(c) SaTESAC

(b) CCM

(a) TESAC

friction:0.1,mass:10

7.5

10

friction:0.1,mass:30

friction:10,mass:1

5.0

2.5

0.0

2.5

5.0

7.5

friction:0.1,mass:10

friction:0.1,mass:30

friction:10,mass:1

10

10.0

40

30

-20-10

10

20

30

15

20

-10

0

10

20

10

friction:10,mass:1

friction:0.1,mass:10

friction:0.1,mass:30

10

friction:0.1,mass:10

friction:0.1,mass:30

friction:10,mass:1

5

10

15

10

0

10

20

30

6

8

10

12

14

F.1.2 Heatmap of Panda-gym benchmark

(a) TESAC Success Rate

(c) TESAC Pick&Place Rate

(b) TESAC Push Rate

(d) CCM Success Rate

(f) CCM Pick&Place Rate

(e) CCM Push Rate

(g) SaTESAC Success Rate

(i) SaTESAC Pick&Place Rate

(h) SaTESAC Push Rate

1.0

0.7

0.7

0.7

0.7

0.9

0.8

0.7

1.0

0.2

0.1

0.2

0.2

0.8

Friction values

1.0

0.7

0.7

0.6

0.2

0.8

0.7

0.6

0.8

0.1

0.1

0.1

0.3

0.6

30.0 10.05.0

0.7

0.5

0.2

0.1

0.6

0.6

0.7

0.8

0.1

0.1

0.1

0.2

5

10.0

0.4

10.0

0.7

0.4

0.2

0.1

0.5

0.6

0.6

0.8

0.1

0.1

0.1

0.2

30.0

30.0

0.2

0.6

0.4

0.2

0.1

0.6

0.6

0.6

0.7

0.0

0.0

0.0

0.0

0.0

1

5

10

30

1

5

10

30

1

5

10

Mass values

Mass values

Massvalues

1.0

0.9

0.9

0.6

0.3

0.1

0.0

0.2

0.1

0.8

0.9

0.6

0.0

0.8

Friction values

1.0

0.9

0.8

0.5

0.2

0.0

0.0

0.2

0.1

0.8

0.8

0.5

0.0

0.6

10.05.0

0.9

0.7

0.3

0.1

0.0

0.1

0.3

0.2

5.0

0.8

0.8

0.5

0.0

5

10.0

10.0

0.4

0.9

0.6

0.3

0.1

0.1

0.1

0.3

0.3

0.8

0.7

0.5

0.0

30.0

30.0

30.0

-0.2

0.9

0.6

0.2

0.1

0.0

0.1

0.3

0.2

0.9

0.8

0.5

0.0

0.0

1

5

10

30

1

5

10

30

1

5

10

30

Mass values

Mass values

Mass values

1.0

0.9

0.9

0.9

0.8

0.2

0.2

0.3

1.0

0.8

0.8

0.5

0.0

0.8

Friction values

1.0

0.9

0.9

0.7

0.2

0.2

0.2

0.3

0.9

0.8

0.8

0.3

0.0

0.6

10.05.0

0.9

0.6

0.3

0.1

5.0

0.1

0.1

0.3

0.7

0.8

0.7

0.4

0.0

10.0

10.05

0.4

0.9

0.6

0.3

0.1

0.1

0.1

0.3

0.7

0.8

0.7

0.4

0.0

-0.2

30.0

30.0

30.0

0.9

0.5

0.2

0.1

0.1

0.1

0.2

0.8

0.8

0.7

0.4

0.0

0.0

1

5

10

30

1

5

10

30

1

5

10

30

Mass values

Mass values

Mass values

This section presents the heatmap results and further analysis of TESAC, CCM, and SaTESAC on the Panda-gym benchmark. From the heatmap results of SaTESAC and SaCCM (Figure 6), we observe that, with higher cube masses (30 and 10 kg), the agent executed the Push skill (indicated by the clustered points within the yellow bounding box in Figure 1(a)). At lower masses, the agent executed the Pick&Place skill.

In contrast, as shown in Figures 17(e-f), CCM predominantly learned the Pick&Place skill, resulting in a drop in success rates for tasks with mass = 30 , as the agent could not lift the cube off the table using the Pick&Place skill, as depicted in Figure 17(d). The visualisation of the context embedding (Figure 16) did not reveal clear skill grouping across different tasks.

Finally, TESAC primarily mastered the Push skill. The Push skill is relatively simpler to learn than the Pick&Place skill, as it does not require the agent to manipulate its fingers to pick up cubes. Consequently, TESAC's success rate notably decreased in environments with higher friction.

F.2 MuJoCo

SaMI enables RL agents to be versatile and embody multiple skills. Additionally, video demos (available at https://github.com/uoe-agents/SaMI) provide a better demonstration of the different skills.

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

30

40

cripple_set:0

cripple_set:1

cripple_set:2

20

30

10

20

10

10

0

0

10

10

10

20

20

20

30

-20

60

cripple_set:1

cripple_set:0

cripple_set:2

40

20

0

20

40

09

40

30

40

20

0

20

40

50

40

cripple_set:1

cripple_set:0

cripple_set:2

20

0

20

40

40

40

20

0

20

cripple_set:2

cripple_set:0

cripple_set:1

40

cripple_set:2

cripple_set:1

15

cripple_set:2

cripple_set:0

cripple_set:1

15.0

12.5

cripple_set:2

cripple_set:1

cripple_set:0

15

cripple_set:2

cripple_set:1

cripple_set:0

10

10.0

10

7.5

5.0 -

2.5

0.0

2.5

2.5

0.0

2.5

5.0

7.5

10.0

12.5

10

-5

OT

15

20

25

-5

10

5

10

15

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) PCA for TESAC

(f) PCA for CCM

(g) PCA for SaTESAC

(h) PCA for SaCCM

60

40

20

60

80

40

cripple_set:163

cripple_set:0

cripple_set:051

cripple_set:2

09

40

cripple_set:061

0qas"adduo

80

60

40

cripple_set:2

cripple_set:0&1

cripple_set:0

cripple_set:1&3

0

20

TC

20

40

-20

20

02-

60

cripple_set:2

40

40

-80

cripple_set:061

cripple_set:163

09

60

60

08-

60

4020

0

20

40

60

60

40

20

0

20

40

60

60

20

0

20

40

60

80

60

40

20

0

20

40

60

15

cripple_set:061

cripple_set:163

cripple_set:0

cripple_set:2

20

15

sadyo

cripple_set:163

cripple_set:2

25

20

cripple_set:163

cripple_set:0

cripple_set:061

cripple_set:2

20

cripple_set:2

cripple_set:163

cripple_set:0

cripple_set:061

10

10

15

10

10

0

10

5

-10

-10

10

5

10

15

20

25

10

5

5

10

15

20

25

-10

10

15

20

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

30

20

40

mass:0.2,damping:0.2

mass:1.8,damping:0.2

30

20

10

30

40

20

40

20

mass:0.2,damping:0.2

mass:0.2,damping:1.8

mass:1.8,damping:0.2

mass:1.8,damping:1.8

10

0

10

0

30 -

10

20

30

20

40

mass:0.2,damping:0.2

mass:0.2,damping:1.8

mass:1.8,damping:0.2

mass:1.8,damping:1.8

10

20

30

mass:1.8,damping:0.2

mass:0.2,damping:0.2

mass:0.2,damping:1.8

mass:1.8,damping:1.8

40

60

40

20

0

20

40

60

40

20

0

20

40

40

20

20

40

60

40

20

0

40

50

mass:1.8,damping:1.8

mass:1.8,damping:0.2

mass:0.2,damping:1.8

mass:0.2,damping:0.2

15

mass:0.2,damping:1.8

mass:1.8,damping:1.8

mass:0.2.damping:0.2

mass:1.8,damping:0.2

12.5

10.0

15

20

！

0mass:1.8,damping:1.8

mass:0.2,damping:0.2

mass:1.8,damping:0.2

mass:0.2,damping:1.8

10

7.5

10

5.0

2.5

- 0°0

mass:0.2,damping:0.2

5.0

2.5

0.02.55.07.510.0

12.5

5

0

5

10

15

2.5

5.0

mass:0.2,damping:1.8

mass:1.8,damping:0.2

mass:1.8,damping:1.8

5

0

5

10

15

5

10

15

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

40

30

20

mass:0.8,damping:1.0

40

mass:0.8,damping:1.0

40

30

mass:0.8,damping:1.0

30

20 -

10

20

10

0

Q

10

0

10

10

20

-20

10

20

30

-40

20

30

40

40

-30

-20-10

10

30

40

60

40

20

20

40

60

30

40

20

0

20

40

60

30

20

10

0

10

mass:0.8,damping:1.0

20

30

17.5

15.0

12.5

mass:0.8,damping:1.0

25

20

10

14

12

10.0

15

7.5

10

6

5.0

2.5

0

2

2.5

0.0

-10

mass:0.8,damping:1.0

mass:0.8,damping:1.0

0

mass:0.8,damping:1.0

-10

5

0

5

10

15

10

0

10

20

30

2.5

0.02.55.07.510.012.5

2.5

0.0

2.55.0

7.5

10.012.5

15.0

10

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

10

20

OE

mass:0.4,damping:1.8

mass:0.4,damping:0.4

mass:1.8,damping:1.8

20

40

30

mass:0.4,damping:0.4

mass:0.4,damping:1.8

mass:1.8,damping:1.8

30 -

20

40

mass:0.4,damping:0.4

mass:0.4,damping:1.8

mass:1.8,damping:1:8

20

00

10

0

0

10

10

20

60

=40

20

0

20

40

60

30

10

20

40

40

20

0

20

40

40

20

30

-40

-30

-20-10

10

20

3040

20

40

40

20

0

mass:0.4,damping:1.8

mass:0.4,damping:0.4

mass:1.8,damping:1.8

20

40

11

10

12

mass:0.4,damping:0.4

mass:1.8,damping:1.8

mass:0.4,damping:1.8

10

mass:0.4,damping:0.4

mass:0.4,damping:1.8

mass:1.8,damping:1.8 a

mass:1.8,damping:1.8

mass:0.4,damping:1.8

ng:0.4

mass:0.4,damping:0.4

mass:1.8,damping:1.8

mass:0.4,damping:1.8

8

10

12

6

8

10

12

8

10

12

14

8

10

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

20

30

mass:0.1,damping:0.1

mass:5.0,damping:0.1

mass:0.1,damping:1.0

30

20

20

40

mass:0.1,damping:0.1

mass:5.0,damping:0.1

mass:0.1,damping:1.0

mass:5.0,damping:1.0

20

40

30

mass:0.1,damping:1.0

mass:0.1,damping:0.1

mass:5.0,damping:0.1

mass:5.0,damping:1.0

10

10

10

10

0

0

0

10

10

10

10

20

30

60

40

20

0

20

40

20

30

40

mass:0.1,damping:0.1

mass:5.0,damping:1.0

mass:0.1,damping:1.0

mass:5.0,damping:0.1

40

20

海

0

20

20

OE-

40

=40

20

0

20

40

20

30

40

40

20

0

20

10.0

12.5

12

10

mass:0.1,damping:0.1

mass:5.0,damping:0.1

mass:5.0,damping:1.0

mass:0.1,damping:1.0

10

mass:0.1,.damping:1.0

mass:5.0,damping:0.1

mass:0.1,damping:0.1

mass:5.0,damping:1.0

7.5

10

5.0

2.5

2.5

0.0

mass:5.0,damping:0.1

mass:0.1,damping:0,1

mass:0.1,damping:1.0

mass:5.0,damping:1.0

。

mass:0.1,damping:0.1

mass:5.0,damping:0.1

mass:0.1,damping:1.0

mass:5.0,damping:1.0

6

8

10

12

14

2

6

8

10

12

4

6

10

2.5

0.0

2.55.07.510.012.5

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

40

20

20

O

40

20

20

60

40

mass:0.4,damping:1.8

mass:0.4,damping:0.4

mass:1.8,damping:1.8

20

20

40

60

mass:1.8,damping:1.8

mass:0.4,damping:0.4

mass:0.4,damping:1.8

40

mass:0.4,damping:0.4

mass:1.8,damping:1.8

mass:0.4,damping:1.8

60

40

mass:0.4,damping:0.4

mass:1.8,damping:1.8

mass:0.4,damping:1.8

20

40

60

-40-20

20

40

60

60

-4020

0

20

40

60

60

40

0

20

40

60

80

60

-40

-20

0

20

40

15.0

12.5

mass:1.8,damping:1.8

mass:0.4,damping:0.4

mass:0.4,damping:1.8

10 -

12.5

15.0

mass:0.4,damping:0.4

mass:1.8,damping:1.8

mass:0.4,damping:1.8

10

12

mass:0.4,damping:1.8

mass:0.4,damping:0.4

mass:1.8,damping:1.8

10.0

10.0

7.5

7.5

5.0

5.0

2.5

00

mass:0.4,damping:0.4

mass:0.4,damping:1.8

0.0

2.5

mass:1.8,damping:1.8

2.5

0.02.55.07.510.012.515.0

0

246810

1214

7.5-5.02.50.02.55.07.5

10.0

10

15

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

20

cripple_set:1

cripple_set:2

oasadd

20

cripple_set:0

cripple_set:1

cripple_set:2

20

cripple_set:2

cripple_set:0

cripple_set:1

90

cripple_set:0

cripple_set:1

cripple_set:2

10

10

10

10

0

10

10

10

10

20

20

20

20

10

0

10

20

40

20

0

20

40

40

30

20-10

0

10

20

20

10

0

10

20

30

10

cripple_set:1

cripple_set:2

15

cripple_set:0

cripple_set:1

cripple_set:2

12

cripple_set:1

cripple_set:0

cripple_set:2

12

cripple_set:1

cripple_set:2

cripple_set:0

10

10

10 -

5

0

0

5

10

15

-5

0

10

15

5

10

15

4

6

8

10

12

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

60

mass:0.1,damping:1.0

60

mass:0.1,damping:1.0

mass:0.1,damping:1.0

40

mass:1.0,damping:1.0

mass:2.0,damping:1.0

40

mass:2.0,damping:1.0

mass:1.0,damping:1.0

40

mass:1.0,damping:1.0

mass:2.0,damping:1.0

OE

20

20

20

10

20

20

20

40

30

mass:0.1,damping:1.0

mass:1.0,damping:1.0

40

60

40

20

20

40

40

40

20

0

20

40

60

60

=40

20

0

20

40

40

60

40

20

0

20

mass:2.0,damping:1.0

40

60

10

15

mass:2.0,damping:1.0

mass:0.1,damping:1.0

mass:1.0,damping:1.0

10

10

0T

10

5

-5

5

10

mass:1.0,damping:1.0

mass:2.0,damping:1.0

mass:0.1,damping:1.0

15

5

10

15

mass:0.1,damping:1.0

mass:1.0,damping:1.0

mass:2.0,damping:1.0

0.0

2.55.0

7.510.012.515.0

-5.0-2.50.02.55.07.510.012.515.0

mass:0.1,damping:1.0

mass:2.0,damping:1.0

mass:1.0,damping:1.0

(a) t-SNE for TESAC

(b) t-SNE for CCM

(c) t-SNE for SaTESAC

(d) t-SNE for SaCCM

(e) UMAP for TESAC

(f) UMAP for CCM

(g) UMAP for SaTESAC

(h) UMAP for SaCCM

cripple_set:0

40

cripple_set:0

14

cripple_set:2

cripple_set:4

40

cripple_set:4

12

10

8

20

20

40 -

20

4

6

8

10

12

14

60

60

cripple_set:0

cripple_set:4

cripple_set:2

40

20

0

20

40

60

40

cripple_set:0

cripple_set:2

60

cripple_set:4

-40-200

20

40

60

40

60

40

20

0

20

40

09

60

40

cripple_set:0

cripple_set:4

10

cripple_set:0

cripple_set:2

cripple_set:4

10

cripple_set:2

cripple_set:0

cripple_set:4

cripple_set:4

cripple_set:2

8

20

20

40

60

80

60

40

20

20

40

09

6

8

0

2

4

8

10

10

G A comparison with DOMINO and CaDM

In this section, we give a brief comparison between our methods and methods from DOMINO [Mu et al., 2022] and CaDM [Lee et al., 2020] in the MuJoCo benchmark because we are using the exact same environmental setting (shown in Table 5).

DOMINO [Mu et al., 2022] is based on the InfoNCE K -sample estimator. Their implementation, PPO+DOMINO, is a model-free RL algorithm with a pre-trained context encoder. This encoder reduces the demand for large contrastive batch sizes during training by decoupling representation learning for each modality, simplifying tasks while leveraging shared information. However, a pre-trained encoder necessitates a large sample volume, with DOMINO training PPO agents for 5 million timesteps on the MuJoCo benchmark. In contrast, SaTESAC and SaCCM, trained for 1.6 million timesteps without pre-trained encoders, achieve considerably higher average returns across four environments (Table 6). Therefore, it is crucial to focus on extracting MI in contrastive learning that directly optimises downstream tasks, integrating rather than segregating representation learning from task performance.

CaDM [Lee et al., 2020] proposes a context-aware dynamics model adaptable to changes in dynamics. Specifically, they utilise contrastive learning to learn context embeddings, and then predict the next state conditioned on them. We copy the numerical results of PPO+CaDM, Vanilla+CaDM, and PE-TS+CaDM from CaDM [Lee et al., 2020] as their environmental setting is identical to ours, where PPO+CaDM is a model-free RL algorithm, while Vanilla+CaDM and PE-TS+CaDM are modelbased. The model-free RL approach, PPO+CaDM, is 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 is consistent with what is empirically found in CaDM [Lee et al., 2020]: prediction models are more effective when the transition function changes across tasks. Therefore, we consider that a model-based approach to SaMI could be an interesting extension for future work.

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 a statistically significant improvement. we reported the p-value of the t-test in MuJoCo (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 a significance level of 0.05. From Table 7 and 8, 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.

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

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Proofs are provided in Appendix A and B.

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

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: The code will be open-sourced and all the implementation details for both the environment setup and algorithms are in the Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a 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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so '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 to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a 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 as possible in supplemental material (appended to the paper) is 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have provided experimental details in Appendix D including hyperparameters, 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 a 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 are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have provided sufficient information on the computer resources (type of compute workers, memory, time of execution) in 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 the individual experimental runs as well as 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 it 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 https://neurips.cc/public/EthicsGuidelines ?

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 not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 our 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 is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All assets used in our study are open-source, and we have acknowledged the authors' copyrights in the References.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide detailed descriptions of the assets we will release in Appendix D and on our open-source homepage. During the double-blind review phase, we have anonymised the URLs.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

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 crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nProposes 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.\n\nSTRENGTHS:\n- The paper is clearly written.\n- The proposed method is well motivated.\n- The proposed benchmark is interesting.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nLimitations are not discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n### 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.\n\nWEAKNESSES:\n### 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.\n\nQUESTIONS:\n### 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

comments: 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.

STRENGTHS:
- relevant topic: Generalising RL agents across tasks is an important topic that will enhance the field of RL
- well-motivated: The example in the introduction in addition to Fig. 1 is helping the reader to visually understand the considered problem
- The paper clearly states the contributions

WEAKNESSES:
only minor things such as an algorithm box would be useful to understand the learning procedure

QUESTIONS:
- Section 5.3 provides insights into how the positive and negative samples are determined. How prone is this to declare very similar samples, i.e. similar behaviors as positive and negative samples? For example, the return for sample 1 might be high and the return for sample 2 is low, but they are in essence executing the same behavior.

- Could the authors elaborate on the training procedure? How is the sampling of features exactly done and what exactly is the training set as mentioned in Section 6.1 line 246? Does this mean the RL agent is building upon existing offline data?

LIMITATIONS:
Overall I found this a good paper with minor issues/unclear points (see Weaknesses/Questions).
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nAddressed", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2024-07-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n+ 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.\n\nLIMITATIONS:\nThe paper is mostly about addressing a limitation of previous methods, and I do not believe it has any potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-25"}