Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization
For visual reinforcement learning agents, this method improves generalization to unseen environments by filtering out task-irrelevant features.
T2RD decouples visual observations into task-relevant and task-irrelevant representations via self-supervised learning, achieving state-of-the-art generalization and sample efficiency on DeepMind Control Suite and robotic manipulation tasks.
Visual Reinforcement Learning (VRL) has achieved considerable success in solving control tasks. However, generalizing learned policies to new environments remains a major challenge, as agents often overfit to task-irrelevant features in the training environment. To solve this problem, we introduce the concept of decoupling observations into task-relevant and task-irrelevant representations. Building on this idea, we propose a self-supervised Task-Relevant Representation Decoupling (T2RD) algorithm for VRL. This algorithm consists of three components: task-relevant representation consistency, cross-reconstruction, and cross-dynamic prediction. The first two components achieve the decoupling of content and style features, but the resulting content representations are not necessarily task-relevant. To further refine task-relevant features from content representations, we design the third component that introduces dynamic prediction. T2RD achieves State-Of-The-Art (SOTA) generalization performance and sample efficiency in the DeepMind Control Suite and Robotic Manipulation tasks.