RoboTTT——面向机器人策略的上下文扩展:将TTT集成至VLA中以推理时建立记忆信息,从而将视觉-运动上下文扩展到 8K 个时间步

前言

如原论文的摘要所述,近期的机器人基础模型通常只利用单步或极短历史的视觉-运动(visuomotor)上下文

而Test-Time-Training Robot Policies(RoboTTT),在不增加推理时延的前提下,将视觉-运动上下文扩展到 8K 个时间步

  1. RoboTTT 的核心思想是,将 Test-Time Training 融入到机器人基础模型(例如 Vision-Language-Action 策略)中,从而得到一个序列模型,其循环状态由"快权重"(fast weights)构成------这些参数在训练和推理过程中都通过梯度下降进行更新,将历史压缩到权重空间中,并在长上下文条件下检索相关的上下文信息

  2. 为扩展训练时的上下文长度,该方法结合了"序列动作强制"(sequence action forcing)与"截断式时间反向传播"(truncated backpropagation through time)

    在具有挑战性的真实机器人操作任务中,RoboTTT 相比单步上下文基线方法将整体性能提升了 87%,并且能够完整完成一个持续五分钟、包含十个阶段的装配任务,而所有基线方法均未能完成

第一部分 RoboTTT: Context Scaling for Robot Policies

1.1 引言与预备知识

1.1.1 引言

如原论文所述

  1. 目前,大多数最先进的机器人基础模型只在单步或短历史的视觉-运动上下文中运行(5; 7; 8; 29; 36; 37;41; 47; 51; 60; 67; 75; 76)
  2. 相比之下,上下文长度已经成为大语言模型的一个重要扩展维度(9; 16;27; 46; 53; 71)
    尽管机器人中的更长期推理可以交由外部记忆库完成(50; 69),但长时视觉-运动上下文对于以下能力仍然十分关键,例如:
    从人类视频示范中进行一次性、基于上下文的模仿学习(18)
    在机器人自身的部署历史基础上即时改进(39)
    以及在多阶段、长时程任务上实现更强的闭环控制性能

由此产生一个自然问题:如何构建能够从任意长上下文中学习并加以利用的视觉-运动策略?

为了回答这个问题,来自1 NVIDIA、2 Stanford University、3 The University of Texas at Austin的研究者提出了测试时训练机器人策略(RoboTTT)

这是一种机器人模型和训练方法,将视觉运动上下文扩展到8K 时间步(称为RoboTTT-8K),相比当前最先进的机器人基础模型(5; 51;75)在不增加推理延迟的情况下提升了三个数量级

在这种上下文长度下,RoboTTT 展现出新的机器人能力

  1. 通过长上下文条件化,它能够从单个上下文中的人类视频示范中进行一次性模仿,在10 次试验中成功6次,而基线方法则完全失败

    它还展现出即时策略改进能力,比未针对这一能力进行训练的同一模型性能提高了36 %

  2. 在外部扰动下,它在83 % 的试验中取得成功,而最佳短上下文基线的成功率为53 %
    在多阶段、长时程任务中,它相较于单步上下文基线将整体任务性能提高了87 %,并完整完成了一个持续超过五分钟、跨越十个阶段的装配任务,而没有任何基线方法能够做到这一点

  3. 最后,作者声称,他们首次表明,扩展预训练上下文长度能带来稳定的闭环性能提升:RoboTTT-8K 在任务完成得分上比使用1K 时间步上下文预训练的同一模型高出63 %,并且比最佳短上下文基线高出57 %
    这表明上下文长度可以作为机器人基础模型的一个新的扩展维度

在其核心,RoboTTT 将测试时训练(Test-Time Training, TTT)(64; 78) 集成到机器人基础模型中,例如视觉-语言-动作(Vision-Language-Action, VLA)策略(51)

RoboTTT 是一个序列模型,其循环状态由快速权重(58) 构成:不同于在推理阶段被冻结的慢速权重,++快速权重在训练和推理过程中都通过梯度下降进行更新++。该设计解决了长上下文视觉运动策略的三个挑战:以足够的容量编码长历史、利用条件上下文(12)、以及在上下文长度变化时保持推理成本不变

  1. 首先,由快速权重参数化的快速模型(例如,一个MLP)提供了其容量大于循环神经网络的向量值状态

  2. 其次,在部署期间训练快速模型,可以在机器人密集且重复的观测与动作流中保留显著特征并丢弃冗余特征

  3. 再次,随时间传播快速权重能够保持推理成本恒定,而即使使用 KV cache,Transformer 的推理成本仍会随历史长度增长

    更为关键的是,通过在部署期间进行学习,RoboTTT 实现了新的上下文内自适应与策略改进形式
    例如,它可以以人类演示新任务配置的视频作为条件,实现一次性模仿。通过 DAggerDistillation(一种将 DAgger 风格 (57) 的"失败到纠正"映射蒸馏到快速权重中的训练过程),它能够在运行过程中即时改进

为扩展训练时的上下文长度,作者的方法将 sequenceaction forcing 与截断反向传播通过时间(TBPTT)相结合,使上下文得以增长而不会增加GPU 内存占用

1.1.2 预备知识

测试时训练机制测试时训练(TTT)(64; 78) 引入了***++在训练和推理阶段都会更新的快速权重++ ,用于动态建模上下文信息(相当于建立记忆信息)***

这与慢速权重(即模型参数)形成对比,后者(慢速权重)只在训练期间更新,并在推理期间保持冻结

  1. 形式化地,考虑一个由d 维标记组成的序列X,以及其由投影矩阵诱导得到的查询、键和值序列Q, K, V ,其中表示时间步t 时的投影

    快速权重W 参数化了一个小型神经网络,例如一个线性层或MLP
    代表着在时间步t,快速权重被更新,用于将K 与其对应的值投影V 关联起来 ,通过

    其中通常是均方误差,且 表示(可学习的)学习率

  2. 随后,在应用步骤中,更新后的快速权重计算输出,其方式为

    这种" 先更新再应用 " 的操作在训练和推理阶段都会发生
    直观地说,更新步骤将上下文信息编码到快速模型 的参数空间中,而应用步骤则为下游预测从中进行提取
    投影矩阵和快速权重初始化 是通过外层任务损失来学习的,从而使历史机制针对当前任务得到优化

    在推理阶段,++TTT 会将先前处理过的token 压缩到其快速权重中++,而标准的完全注意力机制则在内存中保留所有先前的keys 和values,并在每一步对它们进行注意力计算

机器人序列模型在这项工作中,作者考虑在其回滚历史上进行条件化的机器人策略,也称为机器人序列模型(22; 32; 56; 67)

具体来说,机器人轨迹由图像、本体感受 和动作块 三元组组成;为简化起见,作者省略语言模态。即学习在时间步t 之前的历史以及当前观测上进行条件化的策略。对于长上下文策略,目标是扩展上下文长度

1.2 测试时训练的机器人策略

接下来,将

  1. 首先描述模型架构,以及如何将 TTT 整合到现代机器人基础模型中
    接着给出一套训练配方,将序列动作强制(sequence action forcing)与截断时间反向传播(truncated backpropagation through time,TBPTT)相结合,以扩展训练时的上下文长度
  2. 随后,说明 RoboTTT 如何实现长上下文条件建模,从而解锁新的上下文内自适应与策略改进形式:基于上下文的人类视频示范的一次性模仿(one-shot imitation),以及 DAgger Distillation------一种元学习方法,通过蒸馏 DAgger 风格 (57) 的纠正过程(将次优的机器人动作映射到人类纠正)到其快速权重中,从而教会策略在推理过程中即时自我改进

1.2.1 模型架构

RoboTTT 将 TTT 层集成到机器人基础模型中,例如 VLA 模型,同时保持与底层架构的兼容性,因此可以应用于广泛种类的骨干网络。在本文中,作者在 GR00T N1.7 (51) 之上实例化 RoboTTT

首先,对于测试时训练用于机器人动作

如图2 所示『TTT 层被添加在 DiT 动作头中的注意力层之后:注意力在每个时间步内进行运算,而 TTT 层则跨时间步进行运算。训练阶段使用带有序列动作强制的序列流匹配损失,对每个动作块独立采样噪声水平。推理从学习得到的初始化 开始,在每次观察时更新快速权重并将其向前传播

  1. RosoTTT 由视觉-语言模型(VLM)骨干和带有TTT 层的Diffusion Transformer (DiT) (52) 动作头组成。在时间步t ,它预测一个H 步的动作块

    且在自注意力和交叉注意力层之后加入TTT 层,使得注意力处理单步信息,而TTT 层处理跨时间维度的信息

  2. 具体而言,RosoTTT 的DiT 的输入是一个跨越T 个时间步的机器人轨迹,

    其中
    是由VLM 输出的视觉-语言(VL)token
    是编码后的本体感知token
    是加噪的动作token
    是在每个时间步前置的N 个学习到的寄存器token (14; 30),它们与所有其他token 进行注意力交互

    注意力层在单步token 上操作,并与该时间步的VL token 进行交叉注意力

    这点,跟GR00T的做法 是类似的,因为GR00T也是类似处理:其处理机器人的本体感知状态和(噪声化的)动作,然后与来自Eagle-2 VLM主干网络的图像文本token进行交叉注意力计算,以输出去噪后的电机动作

    每个时间步的注意力输出随后在时间维度上进行拼接

    并++传入TTT 层++
    以进行快速权重更新(Eq. 1)


    输出(Eq. 2)

    为了计算效率,作者避免将VL token 直接传入TTT 层,而是依赖数量较少的(N = 16) 个寄存器token R 在时间上携带VL 信息

其次,对于用于保留预训练能力的门控机制

为了保留预训练VLA 模型的知识,RoboTTT 从基础模型权重初始化,并采用一种学习得到的tanh 门控机制(1),在训练开始时保持TTT 的贡献较小

具体来说,对于每一层DiT,学习,初始化为接近零(0.001),并通过以下方式对TTT 的贡献进行门控

其中 是来自公式2 **的TTT 层输出, 是注意力层输出。通过这种方式,++RoboTTT 学会调整TTT 层的贡献++,而不会压倒预训练模型的计算

1.2.2 RoboTTT 序列训练

作者在机器人轨迹序列上训练RoboTTT,使得快权重在每个训练序列内部得到更新,同时学习合适的快权重初始化以及其更新动态。TTT 投影矩阵和快权重初始化 作为模型参数的一部分进行学习

  1. 具体来说,给定一个训练序列,我们在内循环中对其运行TTT,并使用快权重损失LFW(),在每个时间步计算外部任务损失,并在平均损失上优化整个模型
  2. 通过这种方式,投影矩阵可以直接从外部任务的梯度中学习,而 则通过梯度的梯度(21; 65) 进行元学习,从而使快权重更新适应机器人轨迹

// 待更