Self-Consistent Model-based Adaptation for Visual Reinforcement Learning

TL;DR

This paper introduces SCMA, a model-based method that enhances visual reinforcement learning robustness against distractions by unsupervised denoising, improving generalization and sample efficiency without altering policies.

Contribution

The paper proposes a novel self-consistent, model-based adaptation method for visual RL that does not require policy fine-tuning and is optimized via an unsupervised distribution matching objective.

Findings

SCMA improves performance across various visual distractions.

It enhances sample efficiency in visual RL tasks.

Effective on multiple benchmarks and real robot data.

Abstract

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.