Self-Consistent Model-based Adaptation for Visual Reinforcement Learning

Visual reinforcement learning agents typically face serious performance declines in real-world applications caused by visual distractions. Existing methods rely on fine-tuning the policy's representations with hand-crafted augmentations. In this work, we propose Self-Consistent Model-based Adaptation (SCMA), a novel method that fosters robust adaptation without modifying the policy. By transferring cluttered observations to clean ones with a denoising model, SCMA can mitigate distractions for various policies as a plug-and-play enhancement. To optimize the denoising model in an unsupervised manner, we derive an unsupervised distribution matching objective with a theoretical analysis of its optimality. We further present a practical algorithm to optimize the objective by estimating the distribution of clean observations with a pre-trained world model. Extensive experiments on multiple visual generalization benchmarks and real robot data demonstrate that SCMA effectively boosts performance across various distractions and exhibits better sample efficiency.

Authors

• Xinning Zhou
• Chengyang Ying
• Yao Feng
• Hang Su Dept. of CST, BNRist Center, Inst. for AI, Tsinghua-Bosch Joint Center for ML, Tsinghua University
• Jun Zhu Tsinghua University National Center for Protein Sciences (Beijing) Beijing Institute of Lifeomics

Keywords

• Machine Learning: ML: Model-based and model learning reinforcement learning
• Machine Learning: ML: Reinforcement learning
• Machine Learning: ML: Robustness