跳转至

Rethinking Causal Action Tokenization with Conditional Annealing in Flow Matching

会议: NeurIPS 2026
arXiv: 2609.35469
论文: https://arxiv.org/abs/2609.35469
代码: https://chenyuzhangx.github.io/CATok/
领域: 机器人
关键词: 动作分词器, 流匹配, 自回归VLA, 条件退火, 具身智能

一句话总结

CATok 将连续动作分词重构为基于条件退火的因果流匹配过程,使离散 Token 具备自粗到细的层次因果语义,在保证极高动作重建质量的同时极大提升了自回归 VLA 模型的下游训练与控制性能。

研究背景与动机

视觉-语言-动作(VLA)模型在机器人通用操纵策略中展现出广阔的前景。当前主流 VLA 路线主要分化为两类:一类采用纯自回归架构,使用离散 Token 预测动作,与底层大语言模型的训练基础设施完全兼容;另一类则是混合架构,在大模型后外挂扩散或流匹配动作头。混合头虽然具备很高的连续控制精度,但割裂了动作生成与语言骨干的统一自回归建模,丢失了自回归预训练的统一扩展优势;而现有纯自回归 VLA 的瓶颈在于缺乏合理的动作离散化策略。

现有的动作分词器基本停留在单纯的“数据压缩”视角,导致生成的 Token 语义与自回归骨干的因果预测目标严重错位。简单的一维均匀分箱(如 OpenVLA、RT-2)产生序列长度膨胀且忽略动作维度相关性;基于频域变换(FAST)的方法输出变长编码且缺乏时序因果对应;基于残差量化(FASTer)或随机掩码(OAT)的方法虽然得到固定长度序列,但各个 Token 位置缺乏明确的语义物理对应与层级生成语义。这些方法只编码了“需要重建什么动作”,却忽略了“动作应当如何被因果生成”。

生成扩散与流匹配模型的时间反演轨迹天然蕴含抽象层次:高噪声阶段刻画动作的宏观全局结构,低噪声阶段逐步填充局部微观轨迹与关节残差。如果动作 Token 能与这一生成阶梯对齐,那么前序 Token 便能承载全局动作意图,后续 Token 细化微观控制指令。核心 idea:将动作分词重构为流匹配的条件退火生成过程,通过逐步掩码前序 Token 迫使其分别对齐流轨迹各阶段的信息增量,在离散隐空间建立自粗到细的因果动作 Token 表征。

方法详解

整体框架

CATok 的核心由动作分词器(Tokenizer)和自回归 VLA 集成(VLA-CATok)两部分组成。在分词阶段,输入的连续动作块首先由双流动作编码器通过 2D CNN 和双模态跨注意力提取压缩隐状态;紧接着通过降维瓶颈矢量量化器映射至离散码本;最后基于 MMDiT 的流匹配解码器在条件退火机制下,根据不同时间步提供的有效 Token 子集重建动作。在策略执行阶段,VLA 骨干模型自回归预测因果 Token,并经由冻结的流匹配解码器单步生成执行动作,天然实现语言语义与底层执行的知识绝缘。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["连续动作块 𝒜<br/>(H × d_a)"] --> B["双流动作编码器<br/>2D-CNN 提取与 MMDiT 跨模态查询交互"]
    B --> C["瓶颈矢量量化器<br/>线性降维与码本最近邻量化"]
    C --> D["条件退火流匹配解码器<br/>κ(t) 动态掩码构建粗到细因果阶梯"]
    D --> E["频域平滑与多任务联合优化<br/>时域流匹配 + 频域 DCT + VQ 承诺损失"]
    E --> F["VLA-CATok 自回归集成<br/>离散瓶颈知识绝缘与单步 ODE 动作生成"]

关键设计

1. 双流动作编码器:保留时空结构的多模态查询交互 传统动作编码器通常将多维动作块直接展平或做一维切块,忽略了动作维度(关节坐标、姿态、夹爪)与时间跨度之间的结构化局部相关性。CATok 引入双流架构:动作流首先对分位数归一化后的动作矩阵进行基于权重归一化的 2D 卷积,并融入可学习的 2D 位置编码以保留时空拓扑;查询流初始化 \(K\) 个可学习向量并赋予独立位置编码。双流送入对称 MMDiT 块进行双向跨注意力交互,动作流提供富结构上下文,查询流则将信息高效汇聚至紧凑向量序列中,为离散化提供高保真表征。

2. 瓶颈矢量量化器:解耦隐层维度与码本容量 标准 VQ-VAE 在高维隐空间中直接量化极易遭遇码本崩溃与利用率低下,且直接增加码本维度会导致优化困难。CATok 采用降维瓶颈机制,将维度为 \(D\) 的潜在隐向量通过线性映射 \(\psi: \mathbb{R}^D \rightarrow \mathbb{R}^{D'}\) 投射到低维空间(\(D' < D\)),在此空间中计算欧氏距离匹配离散索引: $\(c_k = \arg\min_{j \in \{1,\dots,|\mathcal{C}|\}} \|\mathbf{e}_j - \psi(\mathbf{q}_k)\|_2\)$ 随后通过反投影 \(\phi: \mathbb{R}^{D'} \rightarrow \mathbb{R}^D\) 恢复至原始维度。这一设计使模型能在保持 Transformer 骨干高表征容量的同时,维持密集、高效的离散量化码本利用率,并结合指数移动平均(EMA)与低频死码重置策略彻底根除码本塌陷。

3. 条件退火流匹配解码器:赋予 Token 自粗到细的因果生成阶梯 常规生成解码器在整个轨迹中对全部 Token 做无差别交叉注意力,无法迫使不同 Token 位置承担不同尺度的控制语义。CATok 创新性地将流匹配的时间反演过程与 Token 因果性绑定。在整流流(Rectified Flow)路径 \(x_t = (1-t)x_0 + tx_1\)(\(x_0 \sim \mathcal{N}(0, I), x_1 = \mathcal{A}\))上,设定条件退火调度函数: $\(\kappa(t) = \lfloor t K \rfloor, \quad t \in [0, 1]\)$ 在时间步 \(t\),流匹配网络仅以被遮蔽前 \(\kappa(t)\) 个 Token 后的剩余活跃子集 \(\hat{\mathcal{Q}}_{>\kappa(t)}\) 作为条件。当 \(t\) 较小(接近纯噪声)时,需要前置 Token 引导全局运动轮廓;当 \(t\) 逐渐增大时,前置 Token 被退火掩码,强迫靠后的 Token 专门负责修正残差与微观关节细节。各 Token 实际上编码了流生成轨迹在邻近阶段的信息增量,形成了与自回归模型完全契合的因果语义层级。

4. 频域平滑与多任务联合优化:压制高频震荡与保证长期平滑 纯时域速度均方/绝对值误差容易在动作轨迹中积累高频抖动,导致机器人在实际执行时发生不平稳颤动。CATok 在目标中联合引入离散余弦变换(DCT)的频域平滑正则项 \(\mathcal{L}_{\text{smooth}}\): $\(\mathcal{L}_{\text{smooth}} = \mathbb{E}_{t, x_0, x_1} \left\| \text{DCT}(x_1 - x_0) - \text{DCT}(v_\theta(x_t, t, \hat{\mathcal{Q}}_{>\kappa(t)})) \right\|_1\)$ 将总训练损失定义为时域流匹配损失 \(\mathcal{L}_{\text{FM}}\)、频域平滑损失 \(\mathcal{L}_{\text{smooth}}\) 与 VQ 承诺损失 \(\mathcal{L}_{\text{VQ}}\) 的加权和,既保证了时域动作的精确还原,又约束了频域谱能量分布,显著提升了长程轨迹的时序一致性。

5. VLA-CATok 自回归集成:离散瓶颈诱导的知识绝缘 在将分词器接入多模态 VLA 骨干(如 Qwen-VL 系列)时,混合架构直接反传动作回归梯度往往破坏预训练多模态语言模型的泛化表征。VLA-CATok 采用教师强制的交叉熵损失直接监督 Token 预测: $\(\mathcal{L}_{AR} = -\sum_{i=1}^K \log p_\theta(c_i \mid I_t, s_t, l, c_{<i})\)$ 离散码本构建了一个天然的信息绝缘屏障,底层的流匹配解码器在策略训练中完全冻结。策略在推理时自回归生成 \(K\) 个 Token,解码器单步解 ODE 即可输出连续动作块,实现了“上层语义高维推理”与“底层低维连续控制”的无缝解耦。

一个完整示例

以 WidowX 机械臂执行长程堆叠任务(StackGreenCubeOnYellowCube)为例: 1. 状态输入与 Token 预测:当前环境图像 \(I_t\)、语言指令 \(l=\) "stack green cube on yellow cube" 输入 Qwen3-VL 骨干,模型依序自回归输出 \(K=16\) 个离散索引 \(C=(c_1, \dots, c_{16})\)。 2. 因果语义映射:前序 Token(如 \(c_1 \sim c_4\))在码本中对应夹爪和机械臂末端的大尺度全局移动趋向;中间 Token(\(c_5 \sim c_{10}\))决定对齐与高度调整;尾部 Token(\(c_{11} \sim c_{16}\))精确调制夹爪开合度与接触微调。 3. 单步流匹配解码:各离散 Token 索引经查表恢复嵌入 \(\hat{\mathbf{Q}}\),送入冻结 MMDiT 流匹配解码器,求解整流 ODE(单步执行),瞬时输出长度 \(H=10\)、动作维度 \(d_a=7\) 的平滑连续动作矩阵,机械臂精准拾取并完成绿块在黄块上的稳定堆放。

实验关键数据

主实验

在三个主流仿真基准(LIBERO、SimplerEnv、RoboTwin 2.0)以及真实 Franka 机械臂平台上对比了 BIN、FAST、OAT 及 CATok:

评测基准 任务场景 / 指标 BIN FAST OAT CATok (本文) 相对提升 / 说明
LIBERO (40任务) Spatial (空间概括) 0.586 0.960 0.428 0.978 +1.8% 优于 FAST
Object (物体交互) 0.878 0.998 0.876 0.994 保持极高成功率
Goal (目标条件) 0.680 0.962 0.704 0.954 与最强基线持平
Long (长程操纵) 0.604 0.901 0.276 0.910 +0.9% 优于 FAST
LIBERO 平均 0.687 0.955 0.571 0.959 整体最优
SimplerEnv (WidowX) Spoon 0.542 0.500 0.417 0.458 稳健泛化
Carrot 0.333 0.375 0.208 0.417 +4.2% 优于 FAST
Stack (方块堆叠长程) 0.208 0.375 0.167 0.542 +16.7% 显著超越基线
Eggplant 0.708 0.625 0.667 0.542 复杂交互验证
SimplerEnv 平均 0.448 0.469 0.365 0.490 整体最优
RoboTwin 2.0 (双臂50任务) Clean (干净环境) 0.213 0.478 0.229 0.489 +1.1% 优于 FAST
Randomized (随机域) 0.221 0.478 0.233 0.531 +5.3% 强鲁棒性
真实 Franka 机器人 Pick-Spatial (空间泛化) - 0.450 - 0.650 +20.0% 成功率跃升
Pick-Color (颜色泛化) - 0.300 - 0.500 +20.0% 成功率跃升
Stack-Long (长程堆叠) - 0.275 - 0.475 +20.0% 成功率跃升

消融实验

在 LIBERO 基准上针对条件退火机制与解码器架构开展的受控消融实验:

配置 / 变体 Spatial Object Goal Long 平均成功率 说明与结论
A. CATok 完整模型 0.978 0.994 0.954 0.910 0.959 具备条件退火流匹配解码
B. 去除条件退火 (w/o Cond. Anneal) 0.920 0.980 0.918 0.834 0.914 掉点 4.5%,证实退火时序因果结构至关重要
C. VQ-VAE + MMDiT + 随机掩码 0.724 0.924 0.756 0.542 0.737 掉点 22.2%,单纯无序随机掩码无法诱导因果层级
D. 标准 VQ-VAE + MMDiT 0.614 0.796 0.296 0.240 0.486 掉点 47.3%,长程和复杂目标完全丧失自回归适应力

关键发现

  • 因果预测方向性得到信息熵验证:在对动作序列进行纯自回归预测时,CATok 的 Token 预测信息熵随序列位置后移单调递减;若逆序预测,熵值单调递增。相比之下,FAST 与均匀分箱的信息熵平坦无序,证实 CATok 确实诱导出了自前向后的严格因果依赖。
  • 训练收敛加速 2 倍:在 SimplerEnv 任务中,CATok 达到 FAST 最佳性能所需的训练步数减少约 50%,显著提升了 VLA 策略训练效率。
  • 优异的保真度与压缩率权衡:CATok 在 \(K=8\) 时即可实现 \(\text{VRR} \times \text{CR} = 16.85\) 的最优综合得分,高保真重建率(VRR)比肩无压缩分箱,而压缩倍率(CR)与端到端编解码延迟远优于 FAST。

亮点与洞察

  • 将扩散/流匹配的时间演化投射为自回归的 Token 因果序列:传统 VQ 无法内生因果次序,而扩散过程天然是从宏观轮廓向微观残差演进。通过将时间步与 Token 掩码结合,精妙地使自回归模型的时序展开与流生成的降噪阶梯完全重合。
  • 离散 Token 带来的天然知识绝缘机制:无须像混合扩散头架构那样引入复杂的人工注意力遮罩防止动作梯度污染语言语义;离散 Token 形成硬信息瓶颈,既保全了 VLM 的通用语言常识,又让冻结解码器专注低层动作生成。
  • 频域 DCT 损失与瓶颈 VQ 的结合技巧:在整流场预测中融入 DCT 谱损失有效遏制了自回归拼接引起的轨迹跳变,而低维线性瓶颈极大提升了多维离散码本的活性。

局限与展望

  • 真实硬件平台与动作维度覆盖有限:实验主要在单臂 Franka 及特定双臂场景展开,尚未在多类异构机器人(如移动双臂四足、灵巧手高维自由度操作)上系统测试跨本体泛化性。
  • 大尺度海量开源数据集验证仍显欠缺:分词器当前主要在 LIBERO、Bridge 及 RoboTwin 内部数据上训练,未来需进一步扩展至 Open X-Embodiment 等数百万级多源机器人混合数据集以验证缩放定律(Scaling Law)。

相关工作与启发

  • vs FAST / FASTer: FAST 使用频域 DCT 与 BPE 分词,产生变长编码且缺乏时序因果对应;FASTer 采用 RVQ 残差量化但目标局限于静态保真度重建。CATok 将时间演化机制引入分词,生成固定长度且严格因果对齐的粗到细语义 Token。
  • vs OAT (Ordered Action Tokenization): OAT 尝试使用嵌套随机丢弃来迫使前序 Token 保留更多信息,但其 Token 位置缺乏与物理过程明确绑定的语义锚点。CATok 借助流匹配的时间步调度,显式赋予每个位置明确的生成阶段物理意义,在长程控制任务上表现出绝对优势。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙利用流匹配时间轨迹诱导因果离散动作空间,构思非常独到。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖三大仿真评测、50个双臂任务、真实机械臂以及细致的因果干预和探针消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严密,问题定义与理论推导清晰,图表呈现扎实。
  • 价值: ⭐⭐⭐⭐⭐ 为纯自回归 VLA 模型的扩展提供了兼具离散兼容性与连续控制精度的坚实基石。