J-Mac: Jacobian Matrix Meets Masked Contrastive Learning for Generalization in Reinforcement Learning.

Abstract

When applying Reinforcement Learning (RL) algorithms to vision-based tasks, the significant variations between training and actual working environments pose a challenge to their generalization capability. While previous methods can enhance the generalization of the base RL algorithm, they are often limited to cases with minor changes between training and working environments. In this paper, we propose an effective auxiliary task called Jacobian Matrix Meets Masked Contrastive Learning (J-Mac), which aims to enhance the base RL algorithm's generalization capability even when there are significant changes between training and working environments. Specifically, we learn the correlations between visual states via transition dynamic learning. Meanwhile, on this basis, we eliminate task-irrelevant features from the representation of the visual state via masked contrastive learning. Extensive experiments demonstrate that our approach significantly improves the generalization of various base RL algorithms, outperforming other state-of-the-art methods across different vision-based benchmarks.