Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization

Jinwen Wang, Youfang Lin, Xiaobo Hu, Qian Xu, Shuo Wang +2 more
7/1/2026
cs.LG

No official code found. Best implementation: Wjw124/T2RD (2 stars, MIT, last commit 2026-09-06).

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.

View on arXivView PDF

Code Implementations(1)

Wjw124/T2RDCommunity25%
MIT

Official code for Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization (T2RD), ACM TOMM 2026.

20Shell, PythonSep 6, 20261 months agoMIT
domain-generalizationmujocoreinforcement-learningself-supervised-learningt2rd+2 more

Similar Papers

Generalizable Task Representation Learning for Offline Meta-Reinforcement Learning with Data Limitations

87% match2 repos72023

DyGRO-VLA: Cross-Task Scaling of Vision-Language-Action Models via Dynamic Grouped Residual Optimization

85% match2026

Task-Aware Virtual Training: Enhancing Generalization in Meta-Reinforcement Learning for Out-of-Distribution Tasks

85% match2025

Knowledge Reutilization in Meta-Reinforcement Learning

84% match1 repo2026

Salience-Invariant Consistent Policy Learning for Generalization in Visual Reinforcement Learning

83% match2025

Temporal Disentanglement of Representations for Improved Generalisation in Reinforcement Learning

83% match1 repo132022

Cite this paper

@article{wang2026taskrelevant,
  title  = {Task-Relevant Representation Decoupling for Visual Reinforcement Learning Generalization},
  author = {Jinwen Wang and Youfang Lin and Xiaobo Hu and Qian Xu and Shuo Wang and Zhuo Chen and Kai Lv},
  year   = {2026},
  eprint = {2607.00796},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2607.00796v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2607.00796.svg)](https://codeofpaper.com/paper/2607.00796)

Discussion