Abstract
Reinforcement Learning (RL) environments can produce training data withspurious correlations between features due to the amount of training data orits limited feature coverage. This can lead to RL agents encoding thesemisleading correlations in their latent representation, preventing the agentfrom generalising if the correlation changes within the environment or whendeployed in the real world. Disentangled representations can improverobustness, but existing disentanglement techniques that minimise mutualinformation between features require independent features, thus they cannotdisentangle correlated features. We propose an auxiliary task for RL algorithmsthat learns a disentangled representation of high-dimensional observations withcorrelated features by minimising the conditional mutual information betweenfeatures in the representation. We demonstrate experimentally, using continuouscontrol tasks, that our approach improves generalisation under correlationshifts, as well as improving the training performance of RL algorithms in thepresence of correlated features.