NeurIPS 2022

摘要

Randomly masking and predicting word tokens has been a successful approach in pre-training language models for a variety of downstream tasks. In this work, we observe that the same idea also applies naturally to sequential decision making, where many well-studied tasks like behavior cloning, offline reinforcement learning, inverse dynamics, and waypoint conditioning correspond to different sequence maskings over a sequence of states, actions, and returns. We introduce the Uni[MASK] framework, which provides a unified way to specify models which can be trained on many different sequential decision making tasks. We show that a single Uni[MASK] model is often capable of carrying out many tasks with performance similar to or better than single-task models. Additionally, after fine-tuning, our Uni[MASK] models consistently outperform comparable single-task models.

随机屏蔽和预测单词标记一直是为各种下游任务预训练语言模型的一个成功方法。在这项工作中，我们观察到同样的想法也自然地适用于顺序决策，其中许多研究良好的任务，如行为克隆、离线强化学习、逆向动力学和航线调度，都对应于一连串的状态、动作和回报的不同序列屏蔽。我们介绍了Uni[MASK]框架，它提供了一种统一的方式在许多不同的序列决策任务上训练模型。我们表明，一个单一的Uni[MASK]模型通常能够执行许多任务，其性能类似于或优于单一任务模型。此外，经过微调后，我们的Uni[MASK]模型的表现优于可比的单任务模型。

研究动机

• NLP中常使用Masked language modeling，其也可以用于序列决策中。

主要贡献

• 提出了一个新框架Uni[mask]，该框架将序列决策问题统一为不同的掩码方案
• 演示了随机掩码方案如何训练单个多任务模型，该模型可用于行为克隆（Behavior Cloning），奖励条件（reward-conditioning），动态建模（dynamics modeling）和更多现成的框架。
• 测试了多任务训练如何影响单任务性能，并展示如何通过随机masking训练的模型经过微调超越单任务模型
• 展示了在开发选择Uni [bask]体系结构时获得的见解，以改善其他state-of-the-art方法

方法描述

• 建模：$\tau = \{(s_0,a_0,p_0),\dots,(s_T,a_T,p_T)\}$

  • $s$：状态；$a$：动作
  • $p$：属性token（可选），一般使用return-to-go(RTG)，也可以是需满足的特定环境条件，agent风格或性能

• Tasks as Masking Schemes

  • 原文：formulate tasks in sequential decision problems as input masking schemes.
  • 这里的任务指不同的masking方案，而非寻常意义的task

• Model Architecture & Training Regimes

  • Model Architecture：使用BERT模型

    • 将state、action、property stack到一个向量中，而非独立的token

      • 个人小疑问：如下图对应的s2和a2对应的token是啥样的？对于s2其他位置都为0？对于a2是只取action对应的下标数值？

    • 使用传统的位置编码

      • trajectory片段只是用位置编码（state、action）
      • RTG使用timestep编码，标志当前所处的timestep
    • Reward-conditioned：只feed context window第一步的return-to-go（Appendix D指出只输入第一个RTG token的效果比每一步都输入效果更好）

![Uni[MASK]模型](https://darkdawn-typora-img.oss-cn-beijing.aliyuncs.com/img202211141440829.png)

• Training regimes

  • single-task：使用一种masking scheme训练任务

  • multi-task：使用multiple masking schemes训练单个模型，任意trajectory片段使用随机masking schemes训练

    • Intuition：单个模型可以胜任多任务，且可能超越single-task的性能，因为学习到了更丰富的环境表征

  • random-mask：使用fully randomized masking训练单个模型。对任意trajectory片段，随机采样masking概率$p_{mask} \in [0,1]$；使用该概率随机mask state/action；RTG使用0.5的概率随机mask

    • Intuition：随机训练可以使模型表现很好，无需在训练中指定特定的任务。模型可以比multi-task学习到更丰富的表征，因为需要理解环境的所有组成。

  • finetune：在特定任务上微调random-mask所学的模型

    • Intuition：微调可以使模型从指定的表征中进一步优化

• A Unified Model for Any Inference Task

  • 环境设置：4*4的网格全局可见，智能体需要先拿到钥匙开门，再移动到指定目标点；

    • 智能体和钥匙位置随机；
    • 智能体每向目标点靠近一步奖励+1，远离奖励-1，否则0；
    • 序列长度T=10

  • One Model to Rule Them All：适用于不同的序列决策范式以及mask方法

    • Behavioral cloning：使用专家数据进行前向预测
    • Goal-conditioned：属性包含最终到达的目标点
    • Reward-conditioned：属性包含RTG
    • Waypoint-conditioned：属性包含指定时间点到达指定目标点
    • Backwards inference：以最终状态为条件的历史序列（对动作状态向后进行采样）
    • global properties：其他对于trajectory的条件，例如是否经过指定的位置

• Future State Predictions：可以基于已知部分状态对轨迹状态进行预测，而不对动作进行预测；即模型准确掌握了环境动态和常规智能体行为

  • 个人理解：即模型对状态转移也进行了学习or建模
  • 举例：已知（t=0,3,6）的状态，预测其他时间智能体的位置状态

• Measuring Single-Task Performance

  • Validation loss：评估模型对于trajectory分布和transition的适应

  • 实验是模型在单（或多）task训练的模型在单一task上的适应性

    • （H1不完全成立）倒数第二行和对角线对比表明： random-mask在一半task性能上优于single-task，所以在多task上训练可以改善部分task的性能
    • （H1不完全成立）倒数第一行和对角线对比表明：random-mask with finetune的表现最好，所以先在多task上generally训练可以改善性能
    • （H2成立）倒数第二行和倒数第三行对比表明：随机mask比在特定task子集上训练效果更佳

理论分析

无

实验验证

• 实验设置

  • 环境：Mujoco-physics Maze2D
  • goal-conditioned by default：一个随机初始化的质点移动到随机的指定目标位置
  • 抹除速度信息，要求智能体根据上下文推断出速度信息
  • 专家数据：使用D4RL’s PD controller产生，添加噪音0均值0.5方差，产生200步的1000 trajectories

• 实验结果：

  • The value of pre-training and fine-tuning for Uni[MASK] models

    • （不符合H2）multi-task优于single-task，也优于random-mask
    • （符合H1）finetune取得了比single-task更好的效果
    • 其他：multi-task的finetune没有random-mask的fintune效果明显

  • Isolating the effect of GPT vs. BERT（Decision-GPT v.s. single-task Uni[MASK）

    • GPT在长context window的效果更好

其他思考

• 本文好像和多任务强化学习没啥关联，当然该框架可以用于训练多任务的模型，即实际意义的多任务再加上不同的masking方案。
• 本文的创新点个人感觉并不多，比较有意思的是他对不同的离线强化学习范式都进行了transformer的尝试，给我们的启示是在RL中尝试引入transformer的时候可以考虑不同的范式，甚至是引入新的范式结构，使其更适配于序列决策问题，同时本文从实验验证了random mask的可行性。
• 但是random mask还是存在一定的局限性，对于大部分的环境而言，状态信息量远大于动作信息量，使用tranformer如何构建世界模型还有待进一步探索，例如像dota、星际这样的大型游戏环境，应用transformer学习transition function的难度还是很大的。
• 另外对于RL的预训练模型，个人始终持保留意见，无论是泛化到新任务还是进一步训练特定任务，此时继续使用离线数据就没有发挥预训练模型的最大优势，因为无论什么任务都需要专家数据。当然本文中的预训练是针对不同的mask方案，那么随机mask就包含了所有可能的mask任务类型了。

原文链接：https://openreview.net/forum?id=GisHNaleWiA

OpenReview：https://openreview.net/forum?id=GisHNaleWiA

Code：https://github.com/micahcarroll/uniMASK