[EAI-023] FAST: Efficient Action Tokenization for Vision-Language-Action Models

Paper Card

论文标题：FAST: Efficient Action Tokenization for Vision-Language-Action Models
论文作者：Karl Pertsch, Kyle Stachowicz, Brian Ichter, Danny Driess, Suraj Nair, Quan Vuong, Oier Mees, Chelsea Finn, Sergey Levine
论文链接：https://arxiv.org/abs/2501.09747
论文出处：/
论文被引：/
项目主页：https://www.pi.website/research/fast

Abstract

基于Transformer的视觉-语言-动作（VLA）策略等自回归序列模型，可以非常有效地捕捉复杂且可泛化的机器人行为。但是，此类模型对连续动作信号的进行 tokenization（离散标记化），这决定了模型预测的离散token （标记）如何映射到连续的机器人动作（action）。当前基于简单逐维度、逐时间步长分箱方案的机器人动作离散标记化方法，在从高频机器人数据中学习灵巧技能时，通常表现不佳。为了解决这一挑战，本文提出了一种基于离散余弦变换的基于压缩的机器人动作离散标记化方案——Frequency-space Action Sequence Tokenization (FAST)，能够为高度灵巧和高频的任务训练自回归VLA，而标准离散化方法在此类任务中完全失效。基于FAST，发布了一个通用的机器人动作离散标记化器 FAST+，它在100万个真实的机器人动作轨迹上训练。它可以用作各种机器人动作序列（具有不同的动作空间和控制频率）的黑盒分词器（tokenizer）。当与 VLA 模型 pi0 结合使用时，可以扩展到在1万小时的机器人数据上进行训练，并与扩散VLA模型的性能相当，但训练时间减少了5倍。

Summary

研究背景

性能好的tokenizer对序列模型的性能至关重要。当前机器人策略通常使用基于每个维度、每个时间步长分箱方案的简单 tokenization 方法，这类方法在学习具有高频控制的灵巧技能时的表现不好（本文中涉及的测试全部失败）。当预测未来动作序列时，简单的tokenizer难以捕获各个时间步之间的相关性，高度相关的action token降低了自回归VLA模型对 next token prediction 建模的有效性。

方法介绍

在这里插入图片描述

从第一性原理出发，开发新的 action tokenizer。关键想法：受到llama中使用的 BPE 编码方法的启发，在模型训练之前压缩机器人动作信号，减少连续token之间的相关性。考虑到机器人动作是连续信号，因此采用离散余弦变换编码，由此产生的 tokenization 方法称为 Frequency-space Action Sequence Tokenization (FAST)。在 DROID 数据集上高效训练 VLA 模型，在未见的真实环境通过自然语言提示进行零样本评测。
在这里插入图片描述

构建了 FAST+ 通用的机器人动作 tokenizer，在100万条真机数据上训练。可以对各种机器人动作序列进行分词，报错单臂机器人、双臂机器人和移动机器人。当pi0模型结合FAST进行训练时，比原来的结合扩散思想的模型训练时间缩短了 5 倍，但性能相当。

前置知识

问题描述：目标是训练策略 $π(a_{1:H}|o)$ ，将观测结果 $o$ 映射到未来机器人动作序列 $a_{1:H}$ 。假设策略输出动作块（Action Chunk），即序列长度为 $H$ 动作，这使得更容易产生时间上一致的动作并减少复合误差。动作分词（Action Tokenization）的目标是定义一个映射 $\mathcal{T}_{a}:a_{1:H} → [T_1,…,T_n]$ ，从维度为 $|\mathcal{A}|$ 的连续动作 $a_{1:H}$ 序列映射到来自大小为 $|\mathcal{V}|$ 的词汇表中的 $n$ 个离散token序列 $T∈|\mathcal{V}|$ 。动作序列之间的token数量 n 可能不同，就像相同长度的句子可能被离散化为可变数量的文本token一样。

基于分箱的动作token化：动作token化最常用的方法是简单的分箱离散化。对于给定的动作 a，这种方法独立地离散化每个维度，将训练数据集中值的范围划分为 $N$ 个均匀的箱，最常用的是 $N = 256$ 。对于D维的动作序列 $a_{1:H}$ ，此token化方案将应用于每个时间步，从而产生最终的token序列 $\mathcal{T}a(a1:H)=[T_{1,1},…,T_{1,D},…,T_{H,1},…,T_{H,D}]$ 。对于高频机器人数据，这种token化方案并非最优：很容易为每个动作片段生成数百个 token，使得难以训练并且推理速度慢。

Tokenization 对 VLA 模型训练的影响

在这里插入图片描述

创建了一个简单的合成时间序列数据集，其目标是预测一个插值四个随机生成点的三次曲线，如图3所示。这个问题反映了高频动作片段上训练的策略面临的挑战，即策略必须预测一系列连续的动作。训练一个小型自回归Transformer进行实验，序列时间步H从25到800，以模拟不同频率收集的动作数据。Navie方法表示将动作序列中的每一个元素进行256bins分箱操作。

实验结果表明，分箱方案的模型在低采样频率的条件下预测效果较好，但是随着采样频率增加，预测误差急剧增加。为什么呢？因为自回归模型的训练目标是下一个token预测，因此，他们的学习信号在给定 $T_{1:i-1}$ 的情况下与 $T_i$ 的边际信息内容成正比。分箱方案随着采样频率的增加，边际信息接近于零：对于平滑信号，随着时间步长的缩短，每个时间步长的变化成比例地缩小。这极大地减慢了训练收敛速度，并且难以你和复杂的高频数据集。例如，OpenVLA在低频的 BridgeV2 和 RT-1 数据集上运行良好，但是对于高频的 DROID 数据集表现不佳。这说明为机器人动作设计更好的分词器的重要性。

通过时间序列压缩实现高效的动作分词器

在这里插入图片描述

为了解决前述提到的高频动作轨迹中的冗余会导致每个action token的边际信息量低进而导致训练性能差的问题，需要一种能够将高度冗余的动作信号压缩减少成少量高信息量token的动作离散化方法。

本文使用基于离散余弦变换（DCT）的压缩算法构建FAST。DCT是一种频域变换，它将连续信号表示为各种频率的余弦元素之和。低频捕获信号的整体形状，而高频份量反映尖锐的跳变（sharp jumps）。

图4说明了FAST从原始机器人动作到action token的变换步骤。首先对输入动作进行归一化，然后对每个动作维度应用DCT，为了压缩信号，忽略不重要的系数，得到量化后的稀疏的矩阵。然后将矩阵展平为一个一维整数向量，然后训练一个BPE分词器将其无损压缩成稠密的token。
在这里插入图片描述

通用机器人动作分词器

FAST 中唯一需要学习的组件是 BPE 编码器的词汇表，这个词汇表需要针对分词器应用的每个新数据集进行循例那，虽然只需要几分钟，但也增加了使用FAST的难度。因此，使用100万个1s的动作块训练了一个通用的机器人动作分词器。已经开源并合入到Transformers库了。
在这里插入图片描述

训练所需要的数据集：数据集有多种动作空间：联合空间、末端执行器世界坐标系和末端执行器相机坐标系，以确保所得分词器的通用性。Open X-Embodiment、DROID和Bridge V2则以其原始形式包含在内。在分词之前，所有动作都填充到32维，以适应不同维度的动作空间。在这里插入图片描述

消融实验

模型基线：pi0、OpenVLA

实验目标：验证FAST分词器+自回归VLA模型的有效性

评估任务

在这里插入图片描述

如图5所示，包含7个评估任务（6个真实机器人任务，1个模拟任务），旨在测试VLA在高度灵巧的任务（例如折叠衣物）和泛化任务（例如在未见环境中进行0样本桌面操作）上的性能。

Libero：在Libero [43]模拟基准套件上进行测试。测量了Libero-Spatial、Libero-Object、Libero-Goal和Libero-10的平均性能。
餐桌清理 [7] (20 Hz)：一台UR5单臂机器人需要清理桌子，将12个物体分类到垃圾桶（用于垃圾）和塑料容器（用于盘子、碗、杯子和餐具）中。此任务需要精确抓取各种物体。
折叠T恤 [7] (50 Hz)：一套双臂ARX机器人需要在一个静止的桌面上折叠各种衬衫。在任务开始时，衬衫平放在桌子上。成功完成此任务需要精确的抓取和移动才能折叠衬衫。
杂货装袋 [7] (20 Hz)：一台UR5单臂机器人需要将 7 个物体从桌子上装入杂货袋中，注意不要弄倒或撕破袋子。此任务需要拾取各种各样的物体并小心地将它们放入袋中。
从烤面包机中取出吐司 [7] (50 Hz)：一台双臂Trossen Viper-X机器人需要从烤面包机中取出两片面包并将它们放在盘子上。此任务需要精确地抓取和放置面包片。
衣物折叠 [7] (50 Hz)：一台双臂ARX机器人需要从篮子里取出衬衫和短裤，将它们平放在桌子上，然后折叠并堆叠起来。这是我们测试中最灵巧的任务。它需要精确的抓取，动态的动作来使衣物平整，在衣物缠结时进行重试和纠正，以及将折叠好的衣物精确地放置在现有的衣物堆上。报告了单个服装物品的成功率。
零样本DROID桌面操作 [38] (15 Hz)：测试了一个在完整DROID数据集上训练的策略，该策略涵盖各种桌面操作任务，例如拾取和放置物体、擦拭、打开和关闭抽屉等。在未见的环境中测试该策略，该环境具有新的桌子设置、背景、新颖的物体、视角和桌子高度。这是第一次在完全未见的环境中对DROID策略进行“零样本”评估，无需协同训练或微调，只需使用自然语言提示预训练模型即可。

机器人动作分词器对比

在这里插入图片描述

使用1秒的动作片段。FAST分词器对所有数据集都实现了有效的压缩，高频动作上效果更显著，token为20-53。
在这里插入图片描述

先前工作中应用的naive分词方法难以在高频机器人数据上学习有效的策略。最高频的任务中尤为明显：餐桌整理 (20Hz) 和 T 恤折叠 (50Hz)。

在这里插入图片描述

FAST 分词技术实现了在 DROID 数据集上成功训练强大的通用策略，该策略可以通过自然语言提示，在未见过的环境中进行零样本评估，无需微调。所有先前的工作都没有显示零样本结果，而是完全专注于联合训练或微调评估。在三个大学的校园中对各种桌面操作任务进行测试，证明了策略的通用性（图 7）。无需额外训练，该策略能够熟练地执行简单的操作任务，例如在各种场景和摄像机视角下拾取和放置物体、打开和关闭橱柜以及打开水龙头。即使是不成功的尝试也表现出合理的行为，例如靠近微波炉和洗碗机门的把手，即使最终未能打开它们。

消融研究

回答两个问题：

FAST分词方法是否独立于底层VLA主干？
BPE压缩步骤有多重要？

在这里插入图片描述

为了回答第一个问题，在高频T恤折叠数据集上训练了一个OpenVLA策略，修改了OpenVLA模型代码以接受多个输入图像并预测1秒的动作块。结果表明，FAST能够显著提高OpenVLA的性能，使其能够有效地训练高频机器人操作数据。这表明，分词方法独立于底层模型主干，并且可以轻松应用于各种预训练的自回归Transformer模型。

在这里插入图片描述

在桌面整理和T恤折叠任务上消融了BPE编码步骤。结果表明，没有BPE编码的策略获得了更差的性能（但仍然优于朴素分词）。直观地说，DCT变换仍然将大部分信号信息集中在少数几个token中，从而改善了学习信号。然而，如果没有BPE，就会出现大量重复的值为0的token，这会稀释学习信号，并显著减慢推理速度，因为模型需要自回归地预测数百个动作token，最终导致策略性能下降。

自回归VLA与扩散VLA对比

在这里插入图片描述
图 9 所示，在小型数据集（Libero，折叠T恤；<50小时）上，两种VLA的性能相当。但是在像搬运桌子这样的大型数据集上，基于FAST的VLA收敛速度明显更快，在训练步骤比π0的扩散变体少3倍的情况下达到了高性能。使用FAST分词训练的自回归 π0 模型更严格地遵循语言指令：在DROID评估中，扩散π0模型经常忽略语言指令，导致分数较低。未来会继续研究扩散和自回归VLA的语言遵循能力。

自回归VLA的一个当前局限性在于其推理速度：π0使用扩散模型通常可以在NVIDIA 4090 GPU上在100毫秒内预测一秒钟的动作片段，但π0模型使用FAST分词需要大约750毫秒的每个片段推理时间，因为它必须执行更多自回归解码步骤（通常需要解码30-60个动作token，而扩散模型π0需要10个扩散步骤）并使用完整的20亿参数语言模型主干进行自回归解码（而扩散模型π0使用3亿参数的“动作专家”）。未来会继续研究离散token自回归Transformer模型的推理提速。

实验结论

本文介绍了 FAST，一种用于高频机器人控制数据的动作分词器。FAST使用离散余弦变换（DCT）和字节对编码（BPE）来压缩动作块，使得其具有更好的压缩效果。实验表明，与以前使用的简单动作离散化方法相比，FAST带来了显著的性能提升，并且优于基于矢量量化的更复杂的学习分词方法。

未来工作：
动作分词器。FAST是朝着通用机器人动作分词器迈出的重要一步，但仍有很多问题有待解决。在这项工作对静态机器人机械臂测试了 FAST。FAST+在其他机器人形态（如移动机器人、灵巧手和人形机器人）上具有良好的压缩能力。在这些平台上测试实际策略性能是未来工作的一个令人兴奋的方向。探索替代压缩方案，以及测试基于压缩的动作编码与非自回归解码方法（如扩散[7]）的组合，是未来研究的有趣方向。

VLA 架构。本文初步探索了两种主要类型的 VLA 架构（自回归和扩散解码 VLA）之间的权衡，但最佳 VLA 架构仍未确定。未来工作应仔细研究训练速度、语言基础能力和任一方法的表达能力之间的权衡。

推理速度。虽然π0-FAST 的整体性能与扩散π0匹配，但在推理时间上较慢。未来的工作应该探索加快自回归 VLA 模型推理速度的方法，以使它们能够解决高度动态的任务。