具身智能,是人工智能(AI)行业的下一个浪潮。如何有效训练 Transformers 模型来控制具身机器人,是当前亟需要解决的难题,尤其是对于更复杂、需要精确和高频控制的精巧技能,现有的视觉-语言-动作(VLA)模型几乎失效。尽管扩散或流匹配通常表现得更好,但扩散需要更长的训练时间。
那么,如何在保持灵巧性和精确性的同时,快速训练 Transformers 进行机器人控制呢?使用一个好的 tokenizer 对有效的大规模训练至关重要。
今天,具身智能初创公司 Physical Intelligence 推出了一种专门为动作设计的新 tokenizer——FAST。
据介绍,FAST 的灵感来自于 JPEG 图像的连续压缩方法,它能够处理标准分箱离散化无法应对的高频精巧任务,并达到与流匹配或扩散相似的灵巧程度,同时训练速度提高 5 倍。通过像处理语言一样用离散 token 表示动作,FAST 提高了从互联网规模预训练的迁移能力,并改善了语言指令跟随。通过自然语言命令的提示,他们首次在 DROID 数据集上训练出能够在全新环境中执行一系列操作任务的策略。
为了促进更强大机器人基础模型的研究,他们发布了一个在 100 万个真实机器人动作序列上训练过的 FAST tokenizer 的通用变体。
通过 FAST,他们开发了一种高效的机器人动作 tokenization 方法,从而能够无缝连接机器人技术与自回归 Transformer 训练管道。
图|FAST 是一个动作 tokenizer,能够通过简单的下一个 token 预测,在高度精巧的任务上训练通用策略。
实验表明,这种自回归策略使人们能够使用简单的方法解决一些迄今为止最具挑战性的机器人任务,同时训练速度远快于现有模型。与此同时,FAST 展示了对当前通用策略训练管道进行小幅改动如何对训练效率和性能产生重大影响,这表明可能还有许多其他改动可以改善策略训练。
Physical Intelligence 团队表示,他们将发布在 100 万个真实机器人动作序列上训练过的 FAST tokenizer 版本。这样,任何人都可以使用 FAST 训练策略,只需三行代码,就可以把动作 token 化:
有关 tokenizer 的更多信息以及如何在自己的数据上训练 FAST tokenizer,详见 https://huggingface.co/KarlP/fast。
FAST:专为动作设计的 tokenizer
FAST 通过在训练前对原始动作块进行压缩,改进了简单的分箱方法。它可以显著提高在精巧机器人数据上进行策略训练和推理的效率。具体来说,该 tokenization 方法依赖于离散余弦变换(DCT),后者是一种常用于信号压缩的技术,比如 JPEG 或 MP3 编解码器。Physical Intelligence 团队将 DCT 与字节对编码(BPE)相结合,后者是一种常用于训练大语言模型(LLM)的压缩算法。两者结合后,原始动作块被压缩成少量密度的动作 token,通常每个块包含 30 到 60 个 token,比之前的动作 tokenization 方法压缩了 10 倍。
图|FAST tokenizer 使用离散余弦变换(DCT)对动作序列进行压缩,生成了一个密集的压缩动作 token 序列。
此外,Physical Intelligence 团队还将 FAST 与 π0 模型结合,进行了测试。
测试结果显示,与之前局限于简单操作任务的离散化 VLA 模型不同,FAST 支持在高精巧任务上训练自回归 Transformer 策略,例如折叠衣物、清理桌子和打包购物袋。与此同时,FAST 训练的速度比之前的模型快多达 5 倍。下方视频展示了通过 FAST 策略可以解决的一些任务。
此外,他们利用 FAST 在最近发布的 DROID 数据集上训练出了第一个通用策略,其能够在新环境中将其泛化到各种指令。DROID 是一个包含各种机器人操纵任务的开源数据集,由来自世界各地的机器人研究人员历时两年收集而成。该数据集包含从大学建筑到真实家庭的各种场景和任务,但迄今为止,还没有一种方法能够在完整数据集上训练出通用策略,使其能够在新环境中零样本执行语言指令。
与加州大学伯克利分校、斯坦福大学和华盛顿大学合作进行的测试结果显示,该策略能够在所有测试的环境中直接执行简单的操作任务。如下方视频所示:
即使策略在某项任务上失败了,它通常也会做出直观的尝试来解决问题(见下方视频)。这为展现了一个未来的可能性——在未来,可以像使用语言模型一样,直接下载并使用通用机器人策略。
π0-FAST:第一个自回归通用策略
他们还使用 FAST tokenizer 训练了 π0-FAST,这是他们的第一个自回归通用策略。
π0-FAST 建立在 π0 模型基础上,并使用相同的模型骨干和训练数据集。π0-FAST 能够解决与标准基于扩散的 π0 模型相同的复杂和精巧任务,但由于使用了简单的自回归离散化方法,它的训练速度快了 5 倍。在对比实验中,标准离散化方法无法解决实验中的任何精巧任务。
图|使用 FAST 进行训练非常高效。通用策略 π0-FAST 的训练速度比原始的 π0 模型快 5 倍,并取得了相似的性能。
当然,他们的模型也并非完美。他们表示,当前模型的一个显著缺点是推理速度较慢:π0-FAST 的自回归解码明显慢于 π0 中使用的流匹配解码方法。虽然加速自回归 VLA 的推理仍是一个有待解决的问题,但在其他领域(如语言建模)中,关于自回归 Transformer 模型快速推理的研究已有丰富的成果,这些研究可以为 VLA 的解决方案提供启示。