# Skill-aware Mutual Information Optimisation for Generalisation in Reinforcement Learning

arXiv: https://arxiv.org/abs/2406.04815 · HTML: https://mlsanity.com/paper/2406.04815/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 39 of its 57 references that are in the map.
- First posted on arXiv 7 Jun 2024.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 39)

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [Learning deep representations by mutual information estimation and maximization](https://mlsanity.com/paper/1808.06670/) (2018)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Diversity is All You Need: Learning Skills without a Reward Function](https://mlsanity.com/paper/1802.06070/) (2018)
- [Learning Representations by Maximizing Mutual Information Across Views](https://mlsanity.com/paper/1906.00910/) (2019)
- [Contrastive Learning with Hard Negative Samples](https://mlsanity.com/paper/2010.04592/) (2020)
- [On Variational Bounds of Mutual Information](https://mlsanity.com/paper/1905.06922/) (2019)
- [A Theoretical Analysis of Contrastive Unsupervised Representation Learning](https://mlsanity.com/paper/1902.09229/) (2019)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learning](https://mlsanity.com/paper/1803.11347/) (2018)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Understanding the Limitations of Variational Mutual Information Estimators](https://mlsanity.com/paper/1910.06222/) (2019)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2005.06800/) (2020)
- [panda-gym: Open-source goal-conditioned environments for robotic learning](https://mlsanity.com/paper/2106.13687/) (2021)
- [Skill-based Model-based Reinforcement Learning](https://mlsanity.com/paper/2207.07560/) (2022)
- [Understanding Negative Samples in Instance Discriminative Self-supervised Representation Learning](https://mlsanity.com/paper/2102.06866/) (2021)
- [Environment Probing Interaction Policies](https://mlsanity.com/paper/1907.11740/) (2019)
- [Trajectory-wise Multiple Choice Learning for Dynamics Generalization in Reinforcement Learning](https://mlsanity.com/paper/2010.13303/) (2020)
- [Sample-Efficient Automated Deep Reinforcement Learning](https://mlsanity.com/paper/2009.01555/) (2020)
- [Towards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning](https://mlsanity.com/paper/2009.13891/) (2020)
- [Conditional Mutual Information for Disentangled Representations in Reinforcement Learning](https://mlsanity.com/paper/2305.14133/) (2023)

## Built on by (most cited first, 1 of 1)

- [Learning Long-Horizon Robot Manipulation Skills via Privileged Action](https://mlsanity.com/paper/2502.15442/) (2025)

## Abstract (arXiv)

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.