[论文解读] Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning

Zhecheng Yuan, Zhengrong Xue+5 more|arXiv (Cornell University)|Dec 17, 2022

Domain Adaptation and Few-Shot Learning被引用 13

一句话总结

本文提出 PIE-G，一种视觉强化学习框架，通过利用 ImageNet 预训练的 ResNet 编码器的现成早期层特征，实现对未见过的视觉环境的强零样本泛化能力。通过冻结编码器并使用持续更新的 BatchNorm，PIE-G 在具有挑战性的 DMControl 泛化基准上相比之前最先进方法实现了平均 55% 的性能提升，表明预训练表征可显著提升视觉强化学习中的样本效率与鲁棒性。

ABSTRACT

Learning generalizable policies that can adapt to unseen environments remains challenging in visual Reinforcement Learning (RL). Existing approaches try to acquire a robust representation via diversifying the appearances of in-domain observations for better generalization. Limited by the specific observations of the environment, these methods ignore the possibility of exploring diverse real-world image datasets. In this paper, we investigate how a visual RL agent would benefit from the off-the-shelf visual representations. Surprisingly, we find that the early layers in an ImageNet pre-trained ResNet model could provide rather generalizable representations for visual RL. Hence, we propose Pre-trained Image Encoder for Generalizable visual reinforcement learning (PIE-G), a simple yet effective framework that can generalize to the unseen visual scenarios in a zero-shot manner. Extensive experiments are conducted on DMControl Generalization Benchmark, DMControl Manipulation Tasks, Drawer World, and CARLA to verify the effectiveness of PIE-G. Empirical evidence suggests PIE-G improves sample efficiency and significantly outperforms previous state-of-the-art methods in terms of generalization performance. In particular, PIE-G boasts a 55% generalization performance gain on average in the challenging video background setting. Project Page: https://sites.google.com/view/pie-g/home.

研究动机与目标

• 解决视觉强化学习智能体在特定环境中过拟合并无法泛化到未见视觉场景的问题。
• 探究现成的预训练视觉表征是否可作为视觉强化学习的通用先验，从而在无需额外环境特定数据的情况下提升泛化能力。
• 识别预训练编码器中提升泛化能力和样本效率的最优组件，特别是哪些层和归一化策略。
• 在多种基准上验证该框架，包括模拟控制任务和真实驾驶模拟器，以证明其鲁棒性与可迁移性。

提出的方法

• 冻结 ImageNet 预训练的 ResNet 编码器的参数，并使用其早期层特征作为强化学习智能体的视觉表征。
• 在训练过程中应用持续更新的 Batch Normalization（BatchNorm），使模型能够在不微调编码器的情况下适应新的数据分布。
• 将预训练编码器用作固定特征提取器，避免对环境特定数据进行预训练或数据增强。
• 将提取的特征集成到标准强化学习训练流程中，实现策略的端到端训练，同时保持视觉编码器冻结。
• 在多个基准上使用零样本泛化协议评估该框架，并与最先进方法进行性能对比。
• 通过将 ImageNet 编码器替换为在 CLIP、Ego4D 和 MoCo-v2 上预训练的模型，测试该方法的鲁棒性，确认其在其他预训练表征上的可泛化性。

实验结果

研究问题

• RQ1 来自 ImageNet 的现成预训练视觉表征是否能提升视觉强化学习中的零样本泛化能力？
• RQ2 预训练 ResNet 编码器的哪些层能为视觉强化学习任务提供最具泛化能力的特征？
• RQ3 在冻结编码器中使用持续更新的 BatchNorm 是否相比固定统计量能提升泛化性能？
• RQ4 与需要环境特定预训练或数据增强的方法相比，预训练编码器的性能如何？
• RQ5 所提出的框架是否能在多样化基准上实现泛化，包括模拟控制任务和真实自动驾驶模拟器？

主要发现

• 在 DMControl 泛化基准最具挑战性的设置中，PIE-G 相比之前最先进方法实现了平均 55% 的性能提升。
• 使用冻结的 ImageNet 预训练 ResNet 的早期层特征显著提升了泛化能力，优于仅依赖环境特定数据增强的方法。
• 持续更新的 BatchNorm 对性能至关重要，因为它使模型能够在推理阶段适应分布外的视觉分布。
• 微调预训练编码器会导致性能下降，表明冻结表征在泛化方面更具鲁棒性。
• 来自其他数据集（如 CLIP 和 MoCo-v2）的预训练模型也取得了优异结果，证实收益源于预训练分布而非特定数据集。
• 该框架在多样化基准（包括 DMControl、Drawer World 和 CARLA）上均保持强劲性能，证明其广泛适用性与鲁棒性。