# Temporal Disentanglement of Representations for Improved Generalisation in Reinforcement Learning

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

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 46 references that are in the map.
- First posted on arXiv 12 Jul 2022.

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [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)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations](https://mlsanity.com/paper/1811.12359/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (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)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [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)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Understanding disentangling in β-VAE](https://mlsanity.com/paper/1804.03599/) (2018)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [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)
- [Closing the Sim-to-Real Loop: Adapting Simulation Randomization with Real World Experience](https://mlsanity.com/paper/1810.05687/) (2018)
- [A Study on Overfitting in Deep Reinforcement Learning](https://mlsanity.com/paper/1804.06893/) (2018)
- [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)
- [Self-Supervised Policy Adaptation during Deployment](https://mlsanity.com/paper/2007.04309/) (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)
- [Pretraining Representations for Data-Efficient Reinforcement Learning](https://mlsanity.com/paper/2106.04799/) (2021)
- [Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation](https://mlsanity.com/paper/2107.00644/) (2021)
- [The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels](https://mlsanity.com/paper/2101.02722/) (2021)

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

- [Conditional Mutual Information for Disentangled Representations in Reinforcement Learning](https://mlsanity.com/paper/2305.14133/) (2023)
- [A Survey of State Representation Learning for Deep Reinforcement Learning](https://mlsanity.com/paper/2506.17518/) (2025)
- [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)
- [Learning to Play Atari in a World of Tokens](https://mlsanity.com/paper/2406.01361/) (2024)
- [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)
- [Learning Generalizable Agents via Saliency-Guided Features Decorrelation](https://mlsanity.com/paper/2310.05086/) (2023)
- [Zero-Shot Generalization of Vision-Based RL Without Data Augmentation](https://mlsanity.com/paper/2410.07441/) (2024)
- [Deep Reinforcement Learning for Multi-Agent Interaction](https://mlsanity.com/paper/2208.01769/) (2022)
- [From Pixels to Factors: Learning Independently Controllable State Variables for Reinforcement Learning](https://mlsanity.com/paper/2510.02484/) (2025)
- [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)
- [Salience-Invariant Consistent Policy Learning for Generalization in Visual Reinforcement Learning](https://mlsanity.com/paper/2502.08336/) (2025)
- [Highway Graph to Accelerate Reinforcement Learning](https://mlsanity.com/paper/2405.11727/) (2024)
- [MInCo: Mitigating Information Conflicts in Distracted Visual Model-based Reinforcement Learning](https://mlsanity.com/paper/2504.04164/) (2025)

## Abstract (arXiv)

Reinforcement Learning (RL) agents are often unable to generalise well to environment variations in the state space that were not observed during training. This issue is especially problematic for image-based RL, where a change in just one variable, such as the background colour, can change many pixels in the image. The changed pixels can lead to drastic changes in the agent's latent representation of the image, causing the learned policy to fail. To learn more robust representations, we introduce TEmporal Disentanglement (TED), a self-supervised auxiliary task that leads to disentangled image representations exploiting the sequential nature of RL observations. We find empirically that RL algorithms utilising TED as an auxiliary task adapt more quickly to changes in environment variables with continued training compared to state-of-the-art representation learning methods. Since TED enforces a disentangled structure of the representation, our experiments also show that policies trained with TED generalise better to unseen values of variables irrelevant to the task (e.g. background colour) as well as unseen values of variables that affect the optimal policy (e.g. goal positions).