Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization

2026-07-01

AI Summary

This work addresses the limited generalization of visual reinforcement learning policies, which often overfit to task-irrelevant features in observations. To mitigate this issue, the authors propose T2RD, a self-supervised algorithm that introduces, for the first time, a task-relevant representation disentanglement mechanism. T2RD isolates and refines task-critical information from visual inputs through three core components: task-relevant representation consistency, cross-reconstruction, and cross-dynamics prediction. Evaluated on the DeepMind Control Suite and robotic manipulation benchmarks, the method achieves state-of-the-art performance, significantly improving both generalization capability and sample efficiency compared to existing approaches.

Abstract

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.

Problem

• Visual Reinforcement Learning
• Generalization
• Task-Relevant Representation
• Overfitting
• Representation Decoupling

Innovation

• Task-Relevant Representation
• Representation Decoupling
• Visual Reinforcement Learning
• Self-Supervised Learning
• Generalization