[论文解读] Masked Autoencoding for Scalable and Generalizable Decision Making

Fangchen Liu, Hao Liu+2 more|arXiv (Cornell University)|Nov 23, 2022

Reinforcement Learning in Robotics被引用 4

一句话总结

本文提出了一种名为掩码决策预测（Masked Decision Prediction, MaskDP）的自监督预训练方法，用于强化学习。该方法通过在状态-动作轨迹上应用掩码自编码，学习通用且可扩展的智能体。通过随机掩码状态和动作并重建它们，MaskDP 实现了在目标到达和离线强化学习任务上的零样本迁移，其在零样本设置下的表现优于先前方法，并且在模型规模增大时展现出强大的可扩展性。

ABSTRACT

We are interested in learning scalable agents for reinforcement learning that can learn from large-scale, diverse sequential data similar to current large vision and language models. To this end, this paper presents masked decision prediction (MaskDP), a simple and scalable self-supervised pretraining method for reinforcement learning (RL) and behavioral cloning (BC). In our MaskDP approach, we employ a masked autoencoder (MAE) to state-action trajectories, wherein we randomly mask state and action tokens and reconstruct the missing data. By doing so, the model is required to infer masked-out states and actions and extract information about dynamics. We find that masking different proportions of the input sequence significantly helps with learning a better model that generalizes well to multiple downstream tasks. In our empirical study, we find that a MaskDP model gains the capability of zero-shot transfer to new BC tasks, such as single and multiple goal reaching, and it can zero-shot infer skills from a few example transitions. In addition, MaskDP transfers well to offline RL and shows promising scaling behavior w.r.t. to model size. It is amenable to data-efficient finetuning, achieving competitive results with prior methods based on autoregressive pretraining.

研究动机与目标

• 开发一种可扩展的自监督预训练方法，用于强化学习，使其能够泛化到多种下游任务。
• 实现在不进行任务特定微调的情况下，对目标到达和离线强化学习任务的零样本迁移。
• 探索非自回归、掩码标记预测在序列决策制定中的有效性，受视觉和语言模型的启发。
• 研究在预训练期间改变掩码比率如何影响泛化能力和下游性能。
• 证明 MaskDP 能够从多样化、未标注的序列数据中学习，并在性能上与自回归预训练方法相媲美。

提出的方法

• MaskDP 通过随机掩码部分状态和动作标记，对状态-动作轨迹应用掩码自编码。
• 模型使用基于 Transformer 的架构重建被掩码的标记，从同一训练信号中同时学习前向和逆向动力学。
• 在预训练期间使用多种掩码比率（15%、35%、75%、95%），每条序列均以随机比率采样，以提升泛化能力。
• 模型在多样化、未标注的序列数据上进行预训练，然后使用少量数据对下游任务进行微调。
• 训练期间使用总损失，结合被掩码和未被掩码标记的重建，实证表明其收敛速度优于仅使用掩码损失的方法。
• 该方法在零样本目标到达、序列目标任务和离线强化学习基准上进行评估，并与自回归基线方法（如 Goal-GPT）进行性能比较。

实验结果

研究问题

• RQ1 在状态-动作轨迹上应用掩码自编码是否能够实现在多样化下游强化学习任务上的零样本迁移？
• RQ2 在预训练期间改变掩码比率如何影响泛化能力和下游性能？
• RQ3 非自回归、掩码预测方法是否在离线强化学习和目标到达任务中优于自回归预训练？
• RQ4 MaskDP 是否能随着模型规模有效扩展，且在更大架构中保持性能增益？
• RQ5 在数据效率和零样本迁移方面，MaskDP 与自回归预训练方法（如 Goal-GPT）相比如何？

主要发现

• 在单目标到达任务中，MaskDP 的性能优于或匹配了从零开始训练和基于预训练的方法，即使在零样本设置下也表现出色。
• 在序列多目标到达任务中，MaskDP 显著优于迭代基线方法，成功规划出无需闭环执行的多步轨迹。
• MaskDP 在离线强化学习任务中展现出强大的迁移能力，其性能与专门设计的离线强化学习方法相当，并在 Walker 域上取得了最先进结果。
• 该模型表现出强大的可扩展性：从 MaskDP 预训练微调的更大模型始终优于更小模型，而 Goal-GPT 中性能增益则微乎其微。
• 使用混合掩码比率策略（15%、35%、75%、95%）显著优于固定比率基线，其中 95% 的掩码比率在实现有意义重建方面最为有效。
• 使用总损失（掩码 + 未掩码）进行训练，即使在干净的专家数据上，其收敛速度也快于仅使用掩码损失的方法。