[论文解读] Masked World Models for Visual Control

Younggyo Seo, Danijar Hafner+5 more|arXiv (Cornell University)|Jun 28, 2022

Advanced Vision and Imaging被引用 10

一句话总结

本文提出掩码世界模型（MWM），一种解耦的基于视觉模型的强化学习框架，通过视觉变换器训练卷积特征掩码自编码器以实现细粒度视觉表征学习，并使用独立的潜在动力学模型进行规划。通过引入辅助奖励预测以编码任务相关细节，MWM在50项Meta-World视觉操作任务上达到81.7%的成功率，显著优于先前SOTA基线的67.9%。

ABSTRACT

Visual model-based reinforcement learning (RL) has the potential to enable sample-efficient robot learning from visual observations. Yet the current approaches typically train a single model end-to-end for learning both visual representations and dynamics, making it difficult to accurately model the interaction between robots and small objects. In this work, we introduce a visual model-based RL framework that decouples visual representation learning and dynamics learning. Specifically, we train an autoencoder with convolutional layers and vision transformers (ViT) to reconstruct pixels given masked convolutional features, and learn a latent dynamics model that operates on the representations from the autoencoder. Moreover, to encode task-relevant information, we introduce an auxiliary reward prediction objective for the autoencoder. We continually update both autoencoder and dynamics model using online samples collected from environment interaction. We demonstrate that our decoupling approach achieves state-of-the-art performance on a variety of visual robotic tasks from Meta-world and RLBench, e.g., we achieve 81.7% success rate on 50 visual robotic manipulation tasks from Meta-world, while the baseline achieves 67.9%. Code is available on the project website: https://sites.google.com/view/mwm-rl.

研究动机与目标

• 为解决在视觉控制任务中准确建模机器人与小型物体之间交互的挑战。
• 克服端到端训练在视觉世界模型中所面临的表征质量与动力学准确性之间的权衡。
• 通过解耦表征与动力学学习，提升视觉世界模型强化学习中的样本效率与渐近性能。
• 探究使用卷积特征掩码的自监督表征学习是否能在捕捉细粒度视觉细节方面优于像素块掩码。
• 探究在表征学习过程中引入辅助奖励预测是否能增强视觉表征的任务相关性信息，从而提升下游策略学习效果。

提出的方法

• 使用基于视觉变换器的自编码器，通过掩码卷积特征而非掩码图像块来重建视觉观测。
• 使用独立的潜在动力学模型，基于自编码器提取的表征来预测未来状态。
• 在自编码器训练过程中引入辅助奖励预测头，以将任务相关性信息注入视觉表征。
• 通过环境交互中持续收集的在线经验，持续更新自编码器与动力学模型。
• 解耦训练过程：先使用重建损失与奖励预测损失更新自编码器，再使用冻结的自编码器特征更新动力学模型。
• 通过逐步增加自编码器预训练过程中的掩码比例，应用课程学习以提升稳定性和特征质量。

实验结果

研究问题

• RQ1 解耦视觉表征学习与动力学学习是否能提升视觉世界模型强化学习的性能？
• RQ2 与图像块掩码相比，使用卷积特征掩码是否能在视觉控制任务中更优地捕捉细粒度视觉细节？
• RQ3 在表征学习过程中引入辅助奖励预测是否能提升视觉特征质量，从而改善下游策略学习？
• RQ4 MWM在样本效率与最终性能方面与端到端世界模型及其他解耦基线相比如何？
• RQ5 在单一任务上预训练的表征在多大程度上能泛化到未见的操作任务？

主要发现

• MWM在Meta-World的50项视觉机器人操作任务上达到81.7%的成功率，显著优于先前SOTA基线的67.9%。
• 消融实验表明，加入奖励预测的MWM优于相同设置但无奖励预测的版本，证实了任务相关信号在表征学习中的优势。
• 在Push任务上预训练的冻结表征能良好泛化到未见任务（如Pick Place、Push Back、Drawer Open），性能可与从零开始训练MWM相媲美或更优。
• 状态回归分析表明，结合奖励预测训练的表征在预测本体感觉状态时误差更低，表明其特征质量更优。
• MWM在复杂任务（如Pick Place）上优于使用对比表示学习的DreamerV2，凸显了解耦式掩码特征方法的优势。
• 尽管在ImageNet分类任务上性能相近，掩码卷积特征在视觉控制任务中表现优于像素块掩码，凸显了任务相关的优越性。