Focus On What Matters: Separated Models For Visual-Based RL Generalization
This addresses generalization in visual RL for robotics and simulation, though it is incremental as it builds on existing representation learning techniques.
The paper tackles the challenge of visual-based RL generalization across unseen environments by proposing SMG, a method that uses separated models for task-relevant and task-irrelevant features via image reconstruction, achieving state-of-the-art performance in DMC benchmarks and robust results in robotic manipulation tasks.
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 (Separated Models for Generalization), 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.