Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization
2026-07-01 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address the challenge that visual reinforcement learning agents struggle to apply what they learned in one environment to new ones because they focus on irrelevant details. They propose a method called Task-Relevant Representation Decoupling (T2RD), which separates important task information from unimportant background features. Their approach uses three techniques to ensure the agent focuses on useful features for the task, improving how well it generalizes and learns efficiently. They tested T2RD on control and robotic tasks, where it showed strong performance.
Visual Reinforcement LearningGeneralizationTask-Relevant RepresentationSelf-Supervised LearningRepresentation DecouplingDynamic PredictionCross-ReconstructionDeepMind Control SuiteRobotic ManipulationSample Efficiency
Authors
Jinwen Wang, Youfang Lin, Xiaobo Hu, Qian Xu, Shuo Wang, Zhuo Chen, Kai Lv
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.