@inproceedings{
zhang2024focus,
title={Focus On What Matters: Separated Models For Visual-Based {RL} Generalization},
author={Di Zhang and Bowen Lv and Hai Zhang and Feifan Yang and Junqiao Zhao and Hang Yu and Chang Huang and Hongtu Zhou and Chen Ye and changjun jiang},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=wz2KvvEk44}
}
A primary challenge for visual-based Reinforcement Learning (RL) is to generalize effectively across unseen environments. Although previous studies have explored different auxiliary tasks to enhance generalization, few adopt image reconstruction due to concerns about exacerbating overfitting to task-irrelevant features during training. Perceiving the pre-eminence of image reconstruction in representation learning, we propose SMG (\blue{S}eparated \blue{M}odels for \blue{G}eneralization), a novel approach that exploits image reconstruction for generalization. SMG introduces two model branches to extract task-relevant and task-irrelevant representations separately from visual observations via cooperatively reconstruction. Built upon this architecture, we further emphasize the importance of task-relevant features for generalization. Specifically, SMG incorporates two additional consistency losses to guide the agent's focus toward task-relevant areas across different scenarios, thereby achieving free from overfitting. Extensive experiments in DMC demonstrate the SOTA performance of SMG in generalization, particularly excelling in video-background settings. Evaluations on robotic manipulation tasks further confirm the robustness of SMG in real-world applications. Source code is available at \url{https://anonymous.4open.science/r/SMG/}.
Focus On What Matters: Separated Models For Visual-Based RL Generalization
Di Zhang Bowen Lv Hai Zhang Feifan Yang Junqiao Zhao Hang Yu Chang Huang Hongtu Zhou Chen Ye changjun jiang