# Conditional Mutual Information for Disentangled Representations in Reinforcement Learning

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

## Key facts

- Cited by 27 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 49 references that are in the map.
- First posted on arXiv 23 May 2023.

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

- [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)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations](https://mlsanity.com/paper/1811.12359/) (2018)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Disentangling by Factorising](https://mlsanity.com/paper/1802.05983/) (2018)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Captum: A unified and generic model interpretability library for PyTorch](https://mlsanity.com/paper/2009.07896/) (2020)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Understanding disentangling in β-VAE](https://mlsanity.com/paper/1804.03599/) (2018)
- [Unsupervised Feature Extraction by Time-Contrastive Learning and Nonlinear ICA](https://mlsanity.com/paper/1605.06336/) (2016)
- [Weakly-Supervised Disentanglement Without Compromises](https://mlsanity.com/paper/2002.02886/) (2020)
- [DARLA: Improving Zero-Shot Transfer in Reinforcement Learning](https://mlsanity.com/paper/1707.08475/) (2017)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [On Disentangled Representations Learned From Correlated Data](https://mlsanity.com/paper/2006.07886/) (2020)
- [Weakly Supervised Disentanglement with Guarantees](https://mlsanity.com/paper/1910.09772/) (2019)
- [Deep Reinforcement and InfoMax Learning](https://mlsanity.com/paper/2006.07217/) (2020)
- [Conditional independence testing based on a nearest-neighbor estimator of conditional mutual information](https://mlsanity.com/paper/1709.01447/) (2017)
- [Disentangling Identifiable Features from Noisy Data with Structured Nonlinear ICA](https://mlsanity.com/paper/2106.09620/) (2021)
- [Causal Curiosity: RL Agents Discovering Self-supervised Experiments for Causal Representation Learning](https://mlsanity.com/paper/2010.03110/) (2020)
- [Nonlinear Independent Component Analysis for Principled Disentanglement in Unsupervised Deep Learning](https://mlsanity.com/paper/2303.16535/) (2023)
- [Learning Markov State Abstractions for Deep Reinforcement Learning](https://mlsanity.com/paper/2106.04379/) (2021)
- [CCMI : Classifier based Conditional Mutual Information Estimation](https://mlsanity.com/paper/1906.01824/) (2019)
- [Decoupled Reinforcement Learning to Stabilise Intrinsically-Motivated Exploration](https://mlsanity.com/paper/2107.08966/) (2021)
- [Explicit Explore-Exploit Algorithms in Continuous State Spaces](https://mlsanity.com/paper/1911.00617/) (2019)

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

- [emg2pose: A Large and Diverse Benchmark for Surface Electromyographic Hand Pose Estimation](https://mlsanity.com/paper/2412.02725/) (2024)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)
- [DRED: Zero-Shot Transfer in Reinforcement Learning via Data-Regularised Environment Design](https://mlsanity.com/paper/2402.03479/) (2024)
- [Multi-view Disentanglement for Reinforcement Learning with Multiple Cameras](https://mlsanity.com/paper/2404.14064/) (2024)
- [Studying the Interplay Between the Actor and Critic Representations in Reinforcement Learning](https://mlsanity.com/paper/2503.06343/) (2025)
- [DEAR: Disentangled Environment and Agent Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2407.00633/) (2024)
- [Disentangled World Models: Learning to Transfer Semantic Knowledge from Distracting Videos for Reinforcement Learning](https://mlsanity.com/paper/2503.08751/) (2025)
- [Enhancing Tactile-based Reinforcement Learning for Robotic Control](https://mlsanity.com/paper/2510.21609/) (2025)
- [Range, not Independence, Drives Modularity in Biologically Inspired Representations](https://mlsanity.com/paper/2410.06232/) (2024)
- [Imagined Autocurricula](https://mlsanity.com/paper/2509.13341/) (2025)
- [From Pixels to Factors: Learning Independently Controllable State Variables for Reinforcement Learning](https://mlsanity.com/paper/2510.02484/) (2025)
- [Learning Disentangled Representations for Perceptual Point Cloud Quality Assessment via Mutual Information Minimization](https://mlsanity.com/paper/2411.07936/) (2024)
- [Entropy Regularized Task Representation Learning for Offline Meta-Reinforcement Learning](https://mlsanity.com/paper/2412.14834/) (2024)
- [PvP: Data-Efficient Humanoid Robot Learning with Proprioceptive-Privileged Contrastive Representations](https://mlsanity.com/paper/2512.13093/) (2025)
- [Contextual Pre-planning on Reward Machine Abstractions for Enhanced Transfer in Deep Reinforcement Learning](https://mlsanity.com/paper/2307.05209/) (2023)
- [Skill-aware Mutual Information Optimisation for Generalisation in Reinforcement Learning](https://mlsanity.com/paper/2406.04815/) (2024)
- [How the level sampling process impacts zero-shot generalisation in deep reinforcement learning](https://mlsanity.com/paper/2310.03494/) (2023)
- [Highway Graph to Accelerate Reinforcement Learning](https://mlsanity.com/paper/2405.11727/) (2024)
- [Maximum Total Correlation Reinforcement Learning](https://mlsanity.com/paper/2505.16734/) (2025)
- [PTSM: Physiology-aware and Task-invariant Spatio-temporal Modeling for Cross-Subject EEG Decoding](https://mlsanity.com/paper/2508.11357/) (2025)
- [Synthetic Data is Sufficient for Zero-Shot Visual Generalization from Offline Data](https://mlsanity.com/paper/2508.12356/) (2025)
- [Balanced Multimodal Learning: An Unidirectional Dynamic Interaction Perspective](https://mlsanity.com/paper/2509.02281/) (2025)
- [Inter-Agent Relative Representations for Multi-Agent Option Discovery](https://mlsanity.com/paper/2512.24827/) (2025)
- [Sample-wise Constrained Learning via a Sequential Penalty Approach with Applications in Image Processing](https://mlsanity.com/paper/2601.16812/) (2026)
- [Conflict-Aware Additive Guidance for Flow Models under Compositional Rewards](https://mlsanity.com/paper/2605.20758/) (2026)
- [Coordinated Disentanglement with Iterative Mode Discovery Under Hidden Correlations](https://mlsanity.com/paper/2607.17264/) (2026)
- [Divide-and-Remember: Recursive Action-Relevant Memory for Long-Horizon VLA Policies](https://mlsanity.com/paper/2610.00982/) (2026)

## Abstract (arXiv)

Reinforcement Learning (RL) environments can produce training data with spurious correlations between features due to the amount of training data or its limited feature coverage. This can lead to RL agents encoding these misleading correlations in their latent representation, preventing the agent from generalising if the correlation changes within the environment or when deployed in the real world. Disentangled representations can improve robustness, but existing disentanglement techniques that minimise mutual information between features require independent features, thus they cannot disentangle correlated features. We propose an auxiliary task for RL algorithms that learns a disentangled representation of high-dimensional observations with correlated features by minimising the conditional mutual information between features in the representation. We demonstrate experimentally, using continuous control tasks, that our approach improves generalisation under correlation shifts, as well as improving the training performance of RL algorithms in the presence of correlated features.