Temporal Disentanglement of Representations for Improved Generalisation in Reinforcement Learning
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).
Code (1)
Tasks
Disentanglementreinforcement-learningReinforcement LearningReinforcement Learning (RL)Representation LearningSimilar Papers 제목 키워드 기반
The role of Disentanglement in Generalisation
Combinatorial generalization — the ability to understand and produce novel combinations of familiar elements — is considered a core capacity of humans that modern AI systems currently struggle at. Recently, it has been s…
DisentanglementOut-of-Distribution GeneralizationSystematic GeneralizationSemi-supervised Meta-learning with Disentanglement for Domain-generalised Medical Image Segmentation
Generalising deep models to new data from new centres (termed here domains) remains a challenge. This is largely attributed to shifts in data statistics (domain shifts) between source and unseen domains. Recently, gradie…
DisentanglementImage SegmentationMedical Image SegmentationMeta-Learning+1The Role of Diverse Replay for Generalisation in Reinforcement Learning
In reinforcement learning (RL), key components of many algorithms are the exploration strategy and replay buffer. These strategies regulate what environment data is collected and trained on and have been extensively stud…
Diversityreinforcement-learningReinforcement LearningReinforcement Learning (RL)Unsupervised Model Selection for Variational Disentangled Representation Learning
Disentangled representations have recently been shown to improve fairness, data efficiency and generalisation in simple supervised and reinforcement learning tasks. To extend the benefits of disentangled representations …
AttributeDisentanglementFairnessmodel+3Conditional Mutual Information for Disentangled Representations in Reinforcement Learning
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…
continuous-controlContinuous ControlDisentanglementreinforcement-learning+2