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

TL;DR

SCPL tackles generalization in visual reinforcement learning by enforcing salience-guided, task-relevant representations and consistent decisions across original and perturbed observations. It combines a value consistency module, a dynamics module, and a policy consistency module with a KL constraint to achieve zero-shot generalization. The approach is supported by theoretical insights linking policy consistency to generalization and is validated across DMC-GB, Robotic Manipulation, and CARLA, where SCPL surpasses state-of-the-art baselines by substantial margins. The work advances practical generalization of visual RL and highlights policy-level consistency as a key factor, with future work aimed at adaptive saliency mechanisms.

Abstract

Generalizing policies to unseen scenarios remains a critical challenge in visual reinforcement learning, where agents often overfit to the specific visual observations of the training environment. In unseen environments, distracting pixels may lead agents to extract representations containing task-irrelevant information. As a result, agents may deviate from the optimal behaviors learned during training, thereby hindering visual generalization.To address this issue, we propose the Salience-Invariant Consistent Policy Learning (SCPL) algorithm, an efficient framework for zero-shot generalization. Our approach introduces a novel value consistency module alongside a dynamics module to effectively capture task-relevant representations. The value consistency module, guided by saliency, ensures the agent focuses on task-relevant pixels in both original and perturbed observations, while the dynamics module uses augmented data to help the encoder capture dynamic- and reward-relevant representations. Additionally, our theoretical analysis highlights the importance of policy consistency for generalization. To strengthen this, we introduce a policy consistency module with a KL divergence constraint to maintain consistent policies across original and perturbed observations.Extensive experiments on the DMC-GB, Robotic Manipulation, and CARLA benchmarks demonstrate that SCPL significantly outperforms state-of-the-art methods in terms of generalization. Notably, SCPL achieves average performance improvements of 14\%, 39\%, and 69\% in the challenging DMC video hard setting, the Robotic hard setting, and the CARLA benchmark, respectively.Project Page: https://sites.google.com/view/scpl-rl.

Figures (17)

• Figure 1: (Left) Saliency masked map of SVEA, SGQN, and SCPL (ours), which shows the attention regions of value functions on the DMC-GB benchmark. (Middle) The KL divergence of action distribution between training and test environments on DMC-GB, where our method holds the smallest KL divergence. (Right) Contribution overview of SCPL, which aims to improve visual generalization by achieving task-relevant representations and consistent and superior decisions.
• Figure 2: Overview of SCPL. The value consistency module is trained using the original and augmented observations $s$ and $s_{\alpha}$, along with their saliency attribute maps $\hat{s}$ and $\hat{s}_{\alpha}$. The dynamics module aids the encoder $f_{\theta}$ in acquiring task-relevant representations, while the policy consistency module introduces a constraint to maintain consistency in action distributions.
• Figure 3: The generation of saliency attribute masked maps.
• Figure 4: The performance of SAC, SVEA, SGQN, and SCPL in Video hard setting. SCPL (red line) shows better generalization.
• Figure 5: Saliency attribute maps for SAC, SVEA, SGQN, and SCPL in Training and Video hard setting. In observations of each task, the first column is the original observation, and the second column is the perturbed observation.