[Paper Review] Unsupervised Visual Attention and Invariance for Reinforcement Learning

Xudong Wang, Long Lian+1 more|arXiv (Cornell University)|Apr 7, 2021

Reinforcement Learning in Robotics62 references4 citations

TL;DR

This paper proposes Visual Attention and Invariance (VAI), an unsupervised method for reinforcement learning that learns to adapt visual observations by removing background clutter and enhancing task-relevant features through a learnable adapter module. VAI achieves state-of-the-art performance on DeepMind Control and DrawerWorld benchmarks with 15% higher average cumulative reward than prior methods, while being over 2x faster and using 40% less GPU memory at inference time.

ABSTRACT

Vision-based reinforcement learning (RL) is successful, but how to generalize it to unknown test environments remains challenging. Existing methods focus on training an RL policy that is universal to changing visual domains, whereas we focus on extracting visual foreground that is universal, feeding clean invariant vision to the RL policy learner. Our method is completely unsupervised, without manual annotations or access to environment internals. Given videos of actions in a training environment, we learn how to extract foregrounds with unsupervised keypoint detection, followed by unsupervised visual attention to automatically generate a foreground mask per video frame. We can then introduce artificial distractors and train a model to reconstruct the clean foreground mask from noisy observations. Only this learned model is needed during test to provide distraction-free visual input to the RL policy learner. Our Visual Attention and Invariance (VAI) method significantly outperforms the state-of-the-art on visual domain generalization, gaining 15 to 49% (61 to 229%) more cumulative rewards per episode on DeepMind Control (our DrawerWorld Manipulation) benchmarks. Our results demonstrate that it is not only possible to learn domain-invariant vision without any supervision, but freeing RL from visual distractions also makes the policy more focused and thus far better.

Motivation & Objective

• To improve the robustness of reinforcement learning agents to distractors in visual observations without requiring task-specific supervision.
• To reduce computational and memory costs during inference by avoiding back-propagation through the encoder at test time.
• To enable efficient, scalable adaptation to new environments using a lightweight adapter module that processes raw observations.
• To provide a visually interpretable and debuggable method for visual adaptation in real-world deployment scenarios.
• To achieve strong generalization across environments with varying backgrounds and distractors, including real-world-like settings such as DrawerWorld.

Proposed method

• The method uses a learnable adapter module that transforms raw pixel observations into a more invariant and attention-focused representation, removing background clutter and emphasizing task-relevant features.
• The adapter is trained in an unsupervised manner using contrastive learning on image patches, encouraging the model to preserve identity while discarding irrelevant background variations.
• The adapter operates in the observation space, not in the latent feature space, enabling direct visual interpretation and debugging of adapted inputs.
• The RL agent is trained end-to-end with the adapter using standard RL algorithms (e.g., SAC), with the adapter parameters updated via back-propagation during training but frozen at inference.
• The method does not require ground-truth keypoint annotations or external datasets, relying solely on self-supervised contrastive learning on image patches.
• Evaluation uses a fixed adapter across diverse test environments, with no fine-tuning or retraining on deployment data, ensuring generalization and low deployment cost.

Experimental results

Research questions

• RQ1 Can an unsupervised visual adaptation method improve RL agent performance on distractor-rich visual environments without requiring task-specific annotations?
• RQ2 How does the proposed adapter-based method compare to feature-space adaptation methods (e.g., PAD) in terms of inference speed and memory efficiency?
• RQ3 To what extent can the adapter generalize to unseen environments with different textures, lighting, and distractor configurations?
• RQ4 Can visual inspection of the adapter's output provide reliable, human-interpretable diagnostics for deployment readiness?
• RQ5 Does removing background clutter through unsupervised adaptation lead to more stable and robust policy learning in the presence of distractors?

Key findings

• VAI achieves a 15% average improvement in cumulative reward over the previous SOTA (PAD) across distractor-rich environments in DeepMind Control, with a mean return of 806 compared to PAD’s 698.
• On the Ball in Cup, catch task, VAI achieves 956±4 cumulative reward, outperforming PAD’s 545±173 by 420 points (77% relative gain).
• VAI is over 2x faster and uses 40% less GPU memory than PAD during evaluation, due to avoiding back-propagation through the encoder at test time.
• The method generalizes well to unseen environments: performance remains close to training-distribution levels, even with significant background changes.
• Visualizations confirm that the adapter effectively suppresses background clutter and enhances foreground task-relevant structures, except in cases with high-reflective surfaces (e.g., marble) that confuse the model.
• The adapter’s output is visually interpretable, enabling quick assessment of deployment suitability without requiring policy rollouts or reward functions.