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

前言大模型真真切切的各方面推动了具身模型的发展包括我司七月(视频号七月具身july)在2023年起 做了一年多的大模型应用开发真是值得的比如之前本博客中解读过的TTT这不 便用在了具身模型上具体而言如原论文的摘要所述近期的机器人基础模型通常只利用单步或极短历史的视觉-运动visuomotor上下文而Test-Time-Training Robot PoliciesRoboTTT在不增加推理时延的前提下将视觉-运动上下文扩展到 8K 个时间步RoboTTT 的核心思想是将 Test-Time Training 融入到机器人基础模型例如 Vision-Language-Action 策略中从而得到一个序列模型其循环状态由“快权重”fast weights构成——这些参数在训练和推理过程中都通过梯度下降进行更新将历史压缩到权重空间中并在长上下文条件下检索相关的上下文信息为扩展训练时的上下文长度该方法结合了“序列动作强制”sequence action forcing与“截断式时间反向传播”truncated backpropagation through time在具有挑战性的真实机器人操作任务中RoboTTT 相比单步上下文基线方法将整体性能提升了 87%并且能够完整完成一个持续五分钟、包含十个阶段的装配任务而所有基线方法均未能完成第一部分 RoboTTT: Context Scaling for Robot Policies1.1 引言与预备知识1.1.1 引言如原论文所述目前大多数最先进的机器人基础模型只在单步或短历史的视觉-运动上下文中运行5; 7; 8; 29; 36; 37;41; 47; 51; 60; 67; 75; 76相比之下上下文长度已经成为大语言模型的一个重要扩展维度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 展现出新的机器人能力通过长上下文条件化它能够从单个上下文中的人类视频示范中进行一次性模仿在10 次试验中成功6次而基线方法则完全失败它还展现出即时策略改进能力比未针对这一能力进行训练的同一模型性能提高了36 %在外部扰动下它在83 % 的试验中取得成功而最佳短上下文基线的成功率为53 %在多阶段、长时程任务中它相较于单步上下文基线将整体任务性能提高了87 %并完整完成了一个持续超过五分钟、跨越十个阶段的装配任务而没有任何基线方法能够做到这一点最后作者声称他们首次表明扩展预训练上下文长度能带来稳定的闭环性能提升RoboTTT-8K 在任务完成得分上比使用1K 时间步上下文预训练的同一模型高出63 %并且比最佳短上下文基线高出57 %这表明上下文长度可以作为机器人基础模型的一个新的扩展维度在其核心RoboTTT 将测试时训练Test-Time Training, TTT(64; 78) 集成到机器人基础模型中例如视觉-语言-动作Vision-Language-Action, VLA策略(51)RoboTTT 是一个序列模型其循环状态由快速权重(58) 构成不同于在推理阶段被冻结的慢速权重快速权重在训练和推理过程中都通过梯度下降进行更新。该设计解决了长上下文视觉运动策略的三个挑战以足够的容量编码长历史、利用条件上下文(12)、以及在上下文长度变化时保持推理成本不变首先由快速权重参数化的快速模型例如一个MLP提供了其容量大于循环神经网络的向量值状态其次在部署期间训练快速模型可以在机器人密集且重复的观测与动作流中保留显著特征并丢弃冗余特征再次随时间传播快速权重能够保持推理成本恒定而即使使用 KV cacheTransformer 的推理成本仍会随历史长度增长更为关键的是通过在部署期间进行学习RoboTTT 实现了新的上下文内自适应与策略改进形式例如它可以以人类演示新任务配置的视频作为条件实现一次性模仿。通过 DAgger Distillation一种将 DAgger 风格 (57) 的“失败到纠正”映射蒸馏到快速权重中的训练过程它能够在运行过程中即时改进为扩展训练时的上下文长度作者的方法将 sequenceaction forcing 与截断反向传播通过时间TBPTT相结合使上下文得以增长而不会增加GPU 内存占用1.1.2 预备知识测试时训练机制测试时训练TTT(64; 78) 引入了在训练和推理阶段都会更新的快速权重用于动态建模上下文信息(相当于建立记忆信息)这与慢速权重即模型参数形成对比后者(慢速权重)只在训练期间更新并在推理期间保持冻结形式化地考虑一个由d 维标记组成的序列X以及其由投影矩阵诱导得到的查询、键和值序列Q, K, V 其中表示时间步t 时的投影快速权重W 参数化了一个小型神经网络例如一个线性层或MLP代表着在时间步t快速权重被更新用于将K 与其对应的值投影V 关联起来通过其中通常是均方误差且表示可学习的学习率随后在应用步骤中更新后的快速权重计算输出其方式为这种”先更新再应用” 的操作在训练和推理阶段都会发生直观地说更新步骤将上下文信息编码到快速模型的参数空间中而应用步骤则为下游预测从中进行提取投影矩阵和快速权重初始化是通过外层任务损失来学习的从而使历史机制针对当前任务得到优化在推理阶段TTT 会将先前处理过的token 压缩到其快速权重中而标准的完全注意力机制则在内存中保留所有先前的keys 和values并在每一步对它们进行注意力计算机器人序列模型在这项工作中作者考虑在其回滚历史上进行条件化的机器人策略也称为机器人序列模型(22; 32; 56; 67)具体来说机器人轨迹由图像、本体感受和动作块三元组组成为简化起见作者省略语言模态。即学习在时间步t 之前的历史以及当前观测上进行条件化的策略。对于长上下文策略目标是扩展上下文长度1.2 测试时训练的机器人策略接下来将首先描述模型架构以及如何将 TTT 整合到现代机器人基础模型中接着给出一套训练配方将序列动作强制sequence action forcing与截断时间反向传播truncated backpropagation through timeTBPTT相结合以扩展训练时的上下文长度随后说明 RoboTTT 如何实现长上下文条件建模从而解锁新的上下文内自适应与策略改进形式基于上下文的人类视频示范的一次性模仿one-shot imitation以及DAgger Distillation——一种元学习方法通过蒸馏 DAgger 风格 (57) 的纠正过程将次优的机器人动作映射到人类纠正到其快速权重中从而教会策略在推理过程中即时自我改进1.2.1 模型架构RoboTTT 将 TTT 层集成到机器人基础模型中例如 VLA 模型同时保持与底层架构的兼容性因此可以应用于广泛种类的骨干网络。在本文中作者在 GR00T N1.7 (51) 之上实例化 RoboTTT首先对于测试时训练用于机器人动作如图2 所示『TTT 层被添加在 DiT 动作头中的注意力层之后注意力在每个时间步内进行运算而 TTT 层则跨时间步进行运算。训练阶段使用带有序列动作强制的序列流匹配损失对每个动作块独立采样噪声水平。推理从学习得到的初始化开始在每次观察时更新快速权重并将其向前传播』RosoTTT 由视觉-语言模型VLM骨干和带有TTT 层的Diffusion Transformer (DiT) (52) 动作头组成。在时间步t 它预测一个H 步的动作块且在自注意力和交叉注意力层之后加入TTT 层使得注意力处理单步信息而TTT 层处理跨时间维度的信息具体而言RosoTTT 的DiT 的输入是一个跨越T 个时间步的机器人轨迹其中是由VLM 输出的视觉-语言VLtoken是编码后的本体感知token是加噪的动作token而是在每个时间步前置的N 个学习到的寄存器token (14; 30)它们与所有其他token 进行注意力交互注意力层在单步token和上操作并与该时间步的VL token进行交叉注意力————这点跟GR00T的做法 是类似的因为GR00T也是类似处理其处理机器人的本体感知状态和(噪声化的)动作然后与来自Eagle-2 VLM主干网络的图像和文本token进行交叉注意力计算以输出去噪后的电机动作————————每个时间步的注意力输出随后在时间维度上进行拼接并传入TTT 层以进行快速权重更新定义为公式1 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 投影矩阵和快权重初始化作为模型参数的一部分进行学习具体来说给定一个训练序列在内循环中对其运行TTT并使用快权重损失LFW在每个时间步计算外部任务损失并在平均损失上优化整个模型通过这种方式投影矩阵可以直接从外部任务的梯度中学习而则通过梯度的梯度(21; 65) 进行元学习从而使快权重更新适应机器人轨迹具体而言他们的数据集包含N 条轨迹每条轨迹是由语言、图像、本体感受和动作块元组序列组成其中重新引入在整条轨迹中共享的语言指令每个训练序列是一条完整轨迹或在最大上下文长度限制下的一段连续子轨迹。将逐步流匹配目标记为则给定快速权重初始化时一条轨迹上的序列损失为其中是进入时间步的快速权重状态并在TTT 层内部(公式1)更新为随后每一步优化都会对进行梯度更新同时更新普通模型权重和快速权重的初始化首先对于序列动作强制RoBoTTT 使用针对动作的流匹配目标进行训练DiT 动作头学习对去噪其中是流匹配时间步是采样的噪声在序列训练中作者发现有必要为序列中的每个动作块独立采样噪声水平然后将这一技术称为序列动作强制如果没有这一技术训练会变得不稳定这可能是因为在整个序列上共享一个噪声水平全序列扩散会使整个序列在学习难度上要么均匀简单低噪声要么均匀困难高噪声这与Chen 等人(10) 的发现相呼应总结来说将DiT 动作头记作将timestep- 元组记作作者通过最小化带有序列动作强制的RoboTTT 进行训练其中每个都是独立采样的采样方式为其次对于截断的时间反向传播在长序列上使用完整的时间反向传播BPTT进行训练时需要为每个时间步存储激活值因此GPU 内存会随着序列长度增长作者则采用截断的时间反向传播TBPTT将输入序列划分为多个片段并且梯度只在每个片段内部传播图4梯度在片段边界处被截断快速权重会被延续从而使 TTT 能够在整个序列上持续进行关键在于快速权重会在片段之间被保留从而使TTT 在整个序列上持续进行而它们的梯度会在片段边界处被截断GPU 内存因此由片段长度而非总序列长度决定从而在固定内存预算下允许任意长的训练上下文需要注意的是快速权重初始化仍然会通过第一个片段接收梯度其更新直接来源于最后对于推理如图2 所示RooDTTT 从学习到的初始化开始每次展开在当前观测上更新快速权重并将其传播到下一个时间步。在每个时间步使用步去噪生成动作块1.2.3 从上下文中高效学习RoboTTT 将快速权重更新与慢速权重更新解耦通过在选定的时间步上对 flow-matching 损失进行掩蔽这些时间步只作为纯上下文使用只更新快速权重而不提供模仿目标。这样的灵活性使RoboTTT 能够从异质的上下文中学习例如人类的视频演示或者机器人自身的不理想 rollout首先对于从上下文视频演示中进行模仿作者将人类视频演示序列与执行相同任务的机器人轨迹进行配对视频序列只更新快速权重其流匹配损失被掩蔽而动作损失是在以更新后的快速权重为条件的机器人轨迹上计算的在这种配对数据上进行训练后模型学会从上下文视频中提取任务信息在测试时只需对一个从未见过的任务配置的人类视频进行条件化就可以实现一次性模仿其次对于DAgger 蒸馏考虑按照DAgger (57) 收集有人类纠正的机器人rollout每当机器人犯错时人类操作员就通过纠正动作进行干预从而得到一条轨迹其中每一个执行的动作块要么是机器人动作要么是人类纠正这种交错的rollout 形成了一种自然的在线策略改进模式标准DAgger 仅在有人类纠正的数据上进行微调并丢弃次优的机器人动作然而恰恰是这些动作揭示了每个纠正所应对的失败RoboTTT 则是同时使用两者但承担不对称的角色在序列训练过程中fast weights 会基于完整的交互历史进行更新不仅包含已经执行的人类纠正还包括那些次优的机器人动作本身而流匹配损失则仅在有人类纠正的部分上进行掩蔽计算作者将这一过程称为 DAgger 蒸馏图 6。这种不对称性——将失败作为上下文、将纠正作为目标——正是人类从失败到纠正的映射被“蒸馏”进快速权重的方式模型学习在遇到失败时产生纠正而不是在脱离上下文的情况下单纯模仿纠正从而更充分地利用同一批采集到的数据作者将其视为机器人领域中算法蒸馏39的一种具体实例由人工干预引发的改进过程被蒸馏进策略的快速权重自适应之中在测试阶段模型会在线执行此类纠正而无需人工干预它自身做出的纠正同样会进入历史轨迹并被吸收到快速权重中方式与训练期间人类纠正被吸收的过程完全相同正如作者在原论文第 4 节所展示的那样DAgger 蒸馏DAggerDistillation相比标准的 DAgger 训练能够实现更强的失败恢复能力和更高的任务性能// 待更