Policy-Induced Self-Supervision Improves Representation Finetuning in Visual RL

TL;DR

This paper introduces a self-supervised learning method that enhances representation transfer in visual reinforcement learning by clustering policy-induced representations, improving robustness and learnability across tasks.

Contribution

It proposes a novel self-supervised objective that clusters representations based on induced policies, outperforming traditional methods in visual RL transfer learning.

Findings

Self-supervised clustering improves transfer performance.

Finetuning enhances robustness and learnability.

Pretrained bottom layers are task-agnostic, top layers are task-specific.

Abstract

We study how to transfer representations pretrained on source tasks to target tasks in visual percept based RL. We analyze two popular approaches: freezing or finetuning the pretrained representations. Empirical studies on a set of popular tasks reveal several properties of pretrained representations. First, finetuning is required even when pretrained representations perfectly capture the information required to solve the target task. Second, finetuned representations improve learnability and are more robust to noise. Third, pretrained bottom layers are task-agnostic and readily transferable to new tasks, while top layers encode task-specific information and require adaptation. Building on these insights, we propose a self-supervised objective that clusters representations according to the policy they induce, as opposed to traditional representation similarity measures which are policy-agnostic (e.g. Euclidean norm, cosine similarity). Together with freezing the bottom layers, this objective results in significantly better representation than frozen, finetuned, and self-supervised alternatives on a wide range of benchmarks.