25年1月来自 PI 公司和伯克利分校、斯坦福一起撰写的论文“FAST: Efficient Action Tokenization for Vision-Language-Action Models”。

自回归序列模型,例如基于 Transformer 的视觉语言动作 (VLA) 策略,可以非常有效地捕捉复杂且可泛化的机器人行为。然而,这样的模型要求选择连续动作信号的 token 化,这决定模型预测的离散符号如何映射到连续的机器人动作。当前基于简单的按维度、按时间步长分 bin 方案的机器人动作 token 化方法,在从高频机器人数据中学习灵巧技能时通常表现不佳。为了应对这一挑战,基于离散余弦变换,提出一种基于压缩的机器人动作 token 化方案。该 token 化方法,频率空间动作序列 token 化 (FAST),能够训练自回归 VLA ,完成标准离散化方法完全无法完成的高度灵巧和高频任务。基于 FAST,发布通用机器人动作 token 化器 FAST+,该token化器在 1M 条真实机器人动作轨迹上进行训练。它可以用作各种机器人动作序列的黑盒子 token 化器,具有不同的动作空间和控制频率。最后,当与 π0 VLA 结合使用时,该方法可以扩展到对 10,000 小时的机器人数据进行训练,并匹配扩散 VLA 的性能,同时将训练时间缩短多达 5 倍。

FAST 如图所示:一种通过时间序列压缩对机器人动作轨迹进行 token 化的简单而有效方法,FAST 能够训练自回归 VLA,解决复杂的灵巧操作任务并广泛应用于新场景。

请添加图片描述


大型、高容量的 Transformer 模型可以非常有效地从头开始捕捉复杂且可泛化的机器人行为 [8, 69, 51, 6, 20, 62],也可以使用在互联网规模的图像文本语料库上预训练的下一个 token 预测模型 [10, 39, 63, 7, 65]。然而,这些模型需要选择连续动作信号的 token 化,这决定模型预测的离散符号如何映射到连续的机器人动作 [64, 34, 41, 12]。众所周知,良好的 token 化选择对序列模型的性能至关重要 [55, 57]。先前的此类机器人策略通常使用基于每个维度、每个时间步长分 bin 方案的朴素 token 化策略 [9, 10, 39]。在学习具有高频控制的灵巧技能时,此类方法表现不佳(如图右)。在预测未来动作序列(即动作“块”)时,时间步骤之间的相关性是简单 token 化策略面临的主要挑战,这在高频控制中很常见。高度相关的动作 token,会降低自回归 VLA 中使用的下一个 token 预测目标有效性。直观地说,在这种情况下,通常可以通过简单的映射(例如简单地复制最近的动作 tokens )来实现较低的 token 预测损失,从而使模型处于较差的局部最优状态。

请添加图片描述

语言、文本和音频的 token 化。Token 化是现代基于 Transformer 自回归序列模型训练流程的关键组成部分,token 化方法的选择会对模型训练和下游性能产生重大影响 [55]。虽然有多项研究探索直接对比特流进行操作的“无 token 化”语言模型的训练 [28, 53],但当今大多数语言模型在训练之前都依赖于文本 token 化阶段。一种常见的方法是字节对编码 [27, 55],它通过将频繁出现的 token 序列合并为新 token 来压缩输入文本。对于图像,学习到的压缩方案提供一种有效的方法:输入图像可以表示为由预训练视觉编码器 [44] 生成的“软 tokens”,并且可以使用矢量量化自动编码器 [22, 59] 实现完全自回归图像输入输出。类似的方法可以扩展到视频领域 [66]。在音频生成和语音合成中,它们共享动作预测的时间序列结构,最先进的模型通常使用频域频谱图 [29] 或使用学习的矢量量化器 [68] 对时间序列音频数据进行编码。

视觉-语言-动作模型。最近,多项研究开发通用机器人策略 [9, 51, 6, 10, 20, 39, 62, 11],这些策略在越来越大的机器人学习数据集上进行训练 [52, 38, 60, 24, 47, 35]。一种很有前途的训练通用策略方法是视觉-语言-动作模型(VLA;[10, 17, 39, 67, 7, 63, 73, 71, 13, 11])。VLA 对在互联网规模的图像和文本数据上预训练的视觉语言模型进行微调,以进行机器人控制。这样做有多个好处:使用具有数十亿个参数的大型视觉语言模型主干,可以为策略提供拟合大型机器人数据集所需的表达能力。重用在互联网规模数据集上预训练的权重还可以提高 VLA 遵循各种语言命令和泛化到新目标和场景背景等的能力 [10, 39, 67, 63, 36]。当今大多数 VLA 模型仅限于相当简单的低频控制任务,尤其是使用最常见自回归 VLA 设计的模型 [10, 39]。这是这些模型采用的动作 token化 方案直接结果,这使得对灵巧任务的训练具有挑战性。

用于 VLA 训练的动作表征。先前的研究已经探索用于训练机器人策略(包括 VLA)的各种动作参数化。一项工作使用“语义”动作表示,如语言子任务 [21, 2, 4] 或关键点 [50, 32, 25, 19]。这类方法通常可以从少数例子中学习,甚至可以在没有任何机器人示例的情况下零样本执行任务 [50, 32, 25],但需要手工设计的低级控制器来执行任务,这限制了它们的泛化性。另一种方法是直接训练 VLA,在给定图像和语言指令输入的情况下输出低级机器人控制命令。最常见的设计直接将动作嵌入到离散 token 中,可以使用标准自回归序列模型生成,就像任何流行的视觉语言模型一样。现有的方法使用简单的每维度、每时间步长分 bin 方案将连续的机器人动作映射到离散动作token [9, 10, 39]。该方案难以扩展到高频机器人控制任务。许多工作还提出 token 化的替代方案,例如使用回归头或引入新的权重进行扩散解码 [20, 7, 41, 63]。另一组相关工作探索矢量量化动作表示 [41, 3, 49]。这种方法训练一个矢量量化的编码器-解码器网络,该网络的重建质量可能对超参选择和结构很敏感 [66]。这些方法在粗糙、低保真度重建任务中表现良好,但在需要细粒度控制的高频任务中表现不佳。

本文提出的频率空间动作序列 token 化 (FAST),能够通过简单的下一个 token 预测(见上图左)训练自回归 VLA 策略,用于高度灵巧和高频率的任务,而标准离散化方法则完全失败。此外,FAST 首次实现在最近推出的 DROID 数据集 [38] 上进行高效的 VLA 训练。

目标是训练策略 π(a_1:H|o),将观察 o 映射到未来机器人动作序列 a_1:H。假设策略输出一个“动作块”[69, 40],即 H 动作序列 [15, 7, 69],这使得生成时间一致的动作变得更容易,并减少复合误差。动作 token 化的目标是定义一个映射 T_a : a_1:H → [T_1,…,T_n],从维度为 |A| 的连续动作序列 a_1:H 到大小为 |V| 词汇表中的 n 个离散 token T 序列。请注意,动作序列之间 token 数 n 可能不同,就像长度相同的句子可以 token 为可变数量的文本token 一样。

最常用的动作 token 化方法是简单的分 bin 离散化方案 [8、10、39、72、56]。对于给定的动作 a,此方法独立地离散化每个维度,将训练数据集中的值范围划分为 N 个均匀的 bins,最常用的 N = 256。对于 D 维动作序列 a_1:H,此 token 化方案将应用于每个时间步,产生最终的 token 序列 T_a(a_1:H) = [T_1,1, …,T_1,D, …,T_H,1, …,T_H,D]。对于高频机器人数据,此 token 化方案不是最优的:它很容易为每个动作块产生数百个 tokens ,这使训练变得具有挑战性并导致推理速度变慢。


下面介绍 token 化如何影响 VLA训练。

从一个简单的教学示例开始。创建一个合成的时间序列数据集,其目标是预测插入四个随机生成点的三次样条(见下图底部)。这个小问题反映在高频动作块上训练策略所面临的挑战,这些策略必须在给定一些条件信息的情况下预测一系列连续动作。用以前 VLA 策略中使用的简单 token 化方案对目标序列进行 token 化,该方案将序列中的每个元素分别离散化为 256 个 bins 中的一个。然后,训练一个小型的自回归 Transformer 策略来预测给定条件点的 token 信号。对目标信号的不同采样率重复此实验,从每个序列 25 到 800 个时间步长,而不更改基础数据集。这模拟在不同频率收集的动作数据上训练自回归策略。下图顶部(“简单 naive”)显示以不同频率训练的自回归模型平均预测 MSE。具有分 bin 的 token 化模型在低采样率下实现良好的预测性能(即低 MSE)。但随着采样率的增加,预测误差急剧增加,直到最终模型只是复制第一个动作,如图左下方的定性可视化所示。请注意,这个问题不能归咎于数据本身:底层数据分布的复杂性不会改变,期望具有相同容量、经过相同步数训练的模型在所有采样率下都能实现可比的性能。那么发生了什么?

请添加图片描述

要了解 token 化方案如何影响学习表现,需要研究学习目标本身。从根本上说,自回归模型被训练为在给定所有先前 token 的情况下预测下一个 token 。因此,它们的学习信号与给定 T_1:i−1 的 T_i 边际信息内容成正比。至关重要的是,当使用简单的按时间步长 token 化方案时,随着训练信号的控制频率增加,这种边际信息趋近于零:对于平滑信号,随着时间步长变短,每个时间步长的变化会成比例减少。这大大减慢训练期间的收敛速度,并且可能使拟合复杂的高频数据集变得具有挑战性。事实上,在先前的工作中已经观察到这样的挑战。例如,OpenVLA 在低频 BridgeV2 和 RT-1 数据集上运行良好,但在拟合高频 DROID 数据集时却遇到困难 [39]。

该案例研究结果强调为机器人动作设计更好 token 化方案的重要性。


为了解决上面的问题,需要一种 token 化方法,将高度冗余的动作信号压缩为较少数量的高信息量token。

通过离散余弦变换进行时间序列压缩

关于有效压缩连续时间序列的研究有很多,从信号变换到频域进行压缩的方法 [18, 1, 61] 到学习压缩方法,例如基于矢量量化 [59, 48]。一个关键结论是,任何足够有效的压缩方法,当应用于动作目标时,都适合提高 VLA 模型的训练速度。在实践中,有一些考虑因素可能导致仍然偏爱某些压缩算法,而不是其他算法,例如,训练 token 器的复杂性,以及它在 token化和去 token 操作中的效率如何。

这里使用基于离散余弦变换(DCT)[1] 的压缩算法。与基于矢量量化的学习压缩方法相比,基于 DCT 的压缩是一种分析方法,因此非常简单和快速。

FAST token 化算法

采用离散余弦变换设计 FAST,这是一种快速有效的机器人动作 token 化方法。如图详细介绍从原始机器人动作到动作 token 的步骤。首先对输入动作进行规范化,使得训练数据集中每个动作维度的第 1 分位数和第 99 分位数的值映射到范围 [-1, . . . , 1]。此初始规范化步骤有助于将数据带入指定范围,还可以使具有不同动作尺度的跨实体数据集 token 化更容易。用分位数对偶尔出现在大型机器人数据集中的异常动作保持稳健。数据规范化后,将离散余弦变换分别应用于每个动作维度。为了压缩 DCT 转换后的信号,可以简单地忽略不重要的系数,通过缩放和舍入操作来实现,其中缩放系数是一个超参,它在 token 化操作的损耗和压缩率之间进行权衡。

请添加图片描述

舍入运算之后,DCT 系数矩阵通常是稀疏的,大多数元素为零,并且每个动作维度仅剩余少数几个重要系数。要真正实现压缩,必须将这个稀疏矩阵转换为密集 token 序列。将矩阵展平为一维整数向量,首先引入所有低频分量来交错动作维度,并训练字节对编码 (BPE) token 化器 [27] 将其无损压缩为密集动作 token。BPE 步骤“压缩”零值分量并合并在动作维度频繁出现的系数组合。选择 BPE 来压缩 DCT 矩阵,因为存在许多有效的实现,并且它可以生成固定大小的输出词汇表,可以轻松地集成到现有视觉语言模型词汇表中以进行 VLA 训练。可以使用其他无损压缩算法,如 Huffman 编码 [33] 或 Lempel-Ziv 方法 [75](gzip 压缩方法的基础算法)。

请注意,在 BPE 编码之前展平 |A|×H DCT 系数矩阵的顺序会对策略训练产生重大影响。有两种选择:列优先展平,即首先连接每个维度的最低频率分量,或行优先展平,即首先连接单个动作维度的所有频率分量。选择前者,因为在自回归预测期间首先预测表征输出序列整体形状的低频分量,会导致更稳定的策略推出。

该 token 化流水线中的所有操作都很容易逆转,从而可以快速解码预测的动作。token 化器只有两个超参数:舍入之前应用于 DCT 系数的比例,以及 BPE 压缩步骤的词汇量。这两个参数都不太敏感,在所有单数据集 token 化实验中使用相同的值(四舍五入为 10,BPE 词汇表大小为 1024)。这与依赖矢量量化的端到端学习压缩模块形成对比 [59]。这样的网络通常训练起来很繁琐,需要仔细选择特定于数据集的超参数才能实现良好的重建 [66, 48]。实验表明,基于 DCT 的 token 化方法比基于 VQ 的方法训练出性能更高的策略,同时明显更简单、更易于调整。

通过上面的玩具示例,通过经验证明基于 DCT token 化的好处。在 DCT 压缩的目标 token 上训练自回归模型可以在很宽的采样频率范围内实现持续较低的预测误差。如下算法 1 中简要总结 token 化方法:

请添加图片描述

通用机器人动作 token 化器

Token 化器唯一学习的组件是 BPE 编码器的词汇表,需要针对 token 化器应用到每个新数据集进行训练。虽然这个学习过程很快(通常只需几分钟),但它增加了使用 FAST token化的额外摩擦。因此,目标是训练一个通用动作 token 化器,它可以对来自任何机器人的机器人动作块进行编码。为此,使用上面描述的流水线在一个大型、跨实体机器人动作数据集上训练 token 化器,该数据集由来自单臂、双手和移动操作机器人的大约一百万个 1 秒动作块组成,具有关节和末端执行器控制动作空间和各种控制频率。经过训练后,通用动作 token化 器 FAST+ 可以作为黑盒子 token 化器应用于来自任何机器人设置的 1 秒动作序列。实验评估表明,它与针对单个数据集进行调整的 token 化器相比具有竞争力。

代码发布。在便捷的 HuggingFace AutoProcessor 类中发布了预训练的通用动作 token 化器 FAST +,这使得只需三行代码即可轻松地将 token 化器应用于任何新的机器人动作块:

请添加图片描述

为了获得最佳压缩效果,建议通过 quantile 规范化将输入动作归一化到范围 [−1,…,1],并一次 token 化 1 秒的动作。本文模块还可以轻松地在给定训练的动作块数据集上训练一个新的快速 FAST 化器:

请添加图片描述


实验中,测试 FAST 用两个 VLA 主干:π0 [7] 和 OpenVLA [39]。将 FAST 与替代动作 token 化方案进行比较,并消融关键设计决策。然后,将使用 FAST token化训练的 π0 模型与最先进的 π0 流匹配(扩散)VLA 进行比较,并测试使用 FAST 将自回归 VLA 训练扩展到大型、跨具身数据集(包含 10,000 小时的灵巧机器人操作数据)。

策略实施。用流行的 VLA 主干测试自回归 VLA 训练的不同 token 化方案。对于大多数实验,用 π0 [7],这是一种基于 PaliGemma-3B [5] 的 VLA。还使用基于 Prismatic 7B [37] 构建的 OpenVLA [39] 进行测试。在训练过程中,对 1 秒动作块进行 token 化,并按照先前的 VLA [10, 39] 用生成的动作 token 覆盖 VLM 词汇表中使用最少的tokens。对 VLA 模型进行微调,进行机器人动作预测,而无需冻结权重。

对于 DROID 数据集的训练,根据单个第三人称视角和手腕摄像头视角来设定策略。由于 DROID 每个 episode 提供两个外部摄像头视角,因此在训练期间随机抽样第三人称视角。同样,DROID 为每个训练集提供三个自然语言注释,在训练期间对它们进行随机化。不使用摄像头标定信息。因此,训练后的策略可以在开箱即用的新视点上进行测试,而无需标定。用关节速度和绝对夹持器位置动作空间,并训练策略以预测 15 步动作块(在推理时执行 8 或 15 步块开环)。应用轻度数据管理:只对标记为“成功”的episodes(75k episodes)进行训练,并在训练期间过滤掉所有具有全零动作的空闲时间步(通常是遥控操作员在数据收集期间重置 VR 控制器位置的时间步)。除此之外,在完整数据集上进行训练效果很好,但通过更仔细的管理,性能可能会有所提升。训练三个 epochs 的策略(240k 次迭代 @ 256 批次大小),对于正在使用的 3B 参数 VLA,其在 8 x H100 GPU 上大约需要 4 天时间。

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐