AI Research 技术调研

VQ-BeT: Behavior Generation with Latent Actions

16分钟阅读

一句话定位

NYU/首尔大学（Seungjae Lee、Yibin Wang、Haritheja Etukuru、H. Jin Kim、Nur Muhammad Mahi Shafiullah、Lerrel Pinto）提出的 VQ-BeT：用 Residual VQ-VAE 把连续动作（或动作块）量化成一组分层离散码（primary + secondary code），再由 GPT 式 Transformer 自回归/并行地从观测预测这些码、外加一个连续 offset 头找回精度，单次前向即可生成多模态动作；在七个仿真环境（含 nuScenes 自动驾驶）与真机厨房任务上持平或超过 BeT、Diffusion Policy，且比 Diffusion Policy 快 5 倍（仿真）/ 25 倍（真机，Hello Robot 板载 CPU）。ICML 2024 Spotlight，arXiv 2403.03181（2024-03-05）。

背景与定位

行为生成（behavior generation）的核心难点是动作分布本身的多模态性：同一观测下人类演示可能有多种同样合理的完成方式。Behavior Transformer (BeT) 用 k-means 把连续动作聚类成离散 bin 再配合连续 offset 来捕捉多模态，但 k-means 基于 ℓ2 度量空间，难以扩展到高维动作空间或长时序动作块（action chunking），且聚类不可微、缺乏梯度信息。同期的 Diffusion Policy 与 BESO 用去噪扩散头建模多模态，表达力强，但需要多步迭代去噪，推理慢。

VQ-BeT 的定位是用可学习、可微的层级向量量化（Residual VQ-VAE，源自 Zeghidour et al. 2021 的 SoundStream/RVQGAN 一脉）替换 BeT 里的 k-means 分箱器，同时保留 BeT「离散码 + 连续 offset」的两头设计与单次前向推理的速度优势，试图同时吃到 Diffusion Policy 的表达力和 BeT 的速度。论文在操作、locomotion、自动驾驶、真机四类任务上验证，并做条件（goal-conditioned）与非条件两种任务形式化。

模型架构

两阶段流水线（Figure 2）：Stage 1 训练 Residual VQ-VAE 动作量化器；Stage 2 冻结量化器、训练 GPT 式 Transformer 做码预测。

  • Stage 1：Residual VQ-VAE 动作分词器——编码器 φ 把动作（或长度 n>1 的动作块）a_{t:t+n} 映射到潜向量 x=φ(a)；x 依次通过 N_q 层残差量化：第 1 层码本（primary code，做粗粒度聚类）量化后取残差，再送入第 2 层码本（secondary code，做细粒度修正），量化结果 z_q(x) 为各层码向量之和；解码器 ψ 把 z_q(x) 重构回动作。全部实验固定用 N_q=2 层残差 VQ 即够用，承诺损失权重 λ_commit=1 恒定。码本更新用滑动平均（EMA）而非梯度反传（沿用 Islam et al. 2022 / Mazzaglia et al. 2022 的做法）。VQ-VAE 训练损失 = L1 重构损失 + VQ 目标项（stop-gradient 双向逼近，标准 VQ-VAE commitment loss 形式）。

  • Stage 2：GPT 式 Transformer 码预测器（基于 minGPT 实现）——输入观测序列 o_{t-h:t}（非条件）或额外拼接未来/目标观测 o_{N-g:N}（goal-conditioned），预测每层码本上的类别分布 ζ_code^i，用 Focal loss 加权：L_code = L_focal(ζ_code^{i=1}) + β·L_focal(ζ_code^{i>1})，即给 primary/secondary code 不同权重 β（因为二者承担的角色不对称）。

  • Offset 头：额外一个连续 offset 头 ζ_offset(o_t)，对量化后的粗动作做 L1 修正（ζ_offset 预测 a_{t:t+n} 与量化动作之差），补足精度损失。总损失 L_VQ-BeT = L_code + L_offset。消融显示这个 offset 头对最终精度非常关键。

  • 动作分块（action chunking）：可选，n>1 时一次预测多步动作；但消融发现（4.6 节）在多数环境里加 chunking 反而降低性能——VQ-BeT 本身推理够快，真机实时控制不需要 chunking 换速度。

  • 码预测方式：多数环境（Kitchen/Ant/BlockPush/UR3/PushT）非自回归并行预测 primary+secondary code；仅 nuScenes 与真机环境启用自回归码预测（先出 primary code 再条件预测 secondary code），且消融显示自回归预测在真实世界环境中更重要。

  • 视觉编码器：图像输入环境（Image Kitchen、Image PushT）用 ResNet18；真机环境沿用 Dobb·E 的 HPR 编码器并在训练中微调；状态输入环境（Ant/BlockPush/UR3/NuScenes）无独立视觉编码器。

  • 各环境具体配置（Table 13，7 个环境列）：

    超参KitchenAntBlockPushUR3PushTNuScenesReal-world
    Obs window size10100310516
    Goal window size（条件任务）101031051–
    预测动作序列长度11110561
    自回归码预测否否否否否是是
    β（Eq.4）0.10.60.10.10.10.10.5
    MinGPT 层数/头数/embed dim6/6/1206/6/1204/4/726/6/1206/6/1206/6/1206/6/120
    VQ-VAE latent dim512512256512512512512
    VQ-VAE 码本大小（每层 k）161081616108/10/16

    码本大小 816、N_q=2 意味着实际码组合数（模式数上限）= k² = **64256**。附录 B.1 额外测试把组合数扩大 10×~250×（Table 12），发现 VQ-BeT 对码本尺寸总体稳健：Kitchen（条件）环境组合数扩大 256 倍后，总体码预测准确率仅降到原来的 0.08 倍，但 primary code 单独的预测准确率仍保持原来的 0.8 倍，最终任务表现只降了约 4.5%（3.78→3.61）——说明 VQ-BeT 会自动把 primary code 承担主要模式区分、把 secondary code 的权重压低来应对过大的码本。

数据

仿真侧共 7 个环境（8 个变体，含图像/状态两种输入）+ 1 个自动驾驶数据集，均为已有基准复用而非自采数据：

  • Franka Kitchen（Gupta et al. 2019）：Franka Panda 机械臂，7 维动作空间，566 条人类演示；7 种可能子任务，每条轨迹按某种顺序完成其中 4 个；图像变体用 MuJoCo 渲染成 112×112。
  • PushT（Chi et al. 2023，Diffusion Policy 环境）：2 维末端执行器速度控制，206 条人类演示。
  • BlockPush（Florence et al. 2022）：1,000 条轨迹，由预编程 primitive（非人类）采集，在「推红/绿块到红/绿目标格」的 4 种（目标块, 推动顺序）组合间均匀分配。
  • UR3 BlockPush（Kim et al. 2022）：UR3 机械臂推两块到两个目标圆，每条演示天然多模态（先推哪块不定）；具体演示条数论文未披露。
  • Multimodal Ant（MuJoCo Ant，Brockman et al. 2016 OpenAI Gym）：locomotion 任务，演示轨迹以不同顺序到达地图四角目标；具体演示条数论文未披露。
  • nuScenes 自动驾驶（Caesar et al. 2020）：采用 GPT-Driver（Mao et al. 2023a）预处理好的、以物体为中心的数据集，684 个演示场景，任务是给定几帧观测预测车辆未来 6 个时间步的位置；驾驶方向目标在 rollout 时给定。
  • 真机厨房环境（Hello Robot: Stretch + Dobb·E 同款 Stick 采集工具）：每个单阶段任务采集 45 条演示，80% 训练 / 20% 验证；12 个真机任务（5 个单阶段、3 个两阶段、4 个长程任务，每个长程任务由 3~4 个子任务顺序拼接）。

无跨本体/跨环境混合训练——每个环境独立训练一套 VQ-BeT。

训练方法

  • 两阶段串行训练：先离线预训练 Residual VQ-VAE 动作量化器（收敛后冻结），再训练 Transformer 码预测器 + offset 头，训练目标是标准行为克隆（在观测上做条件的最大似然/交叉熵，无 RL、无奖励信号）。
  • 码预测损失：Focal loss 分别作用于 primary code 与 secondary code，用 β 加权二者（Eq. 4）；offset 头用 L1 回归损失。
  • 训练轮数与学习率（按环境，Table 13）：Kitchen 1000 epoch / lr 5.5e-5；Ant 300 epoch / 5.5e-5；BlockPush 1500 epoch / 1e-4；UR3 300 epoch / 5.5e-5；PushT 2000 epoch / 5.5e-5；NuScenes 1000 epoch / 5.5e-5；真机 600 epoch / 3e-4。
  • 消融验证的关键设计（4.6 节，在 conditional Kitchen / unconditional Ant / nuScenes 三个环境上做）：
    • Residual VQ vs 单层 VQ：去掉残差层显著掉点（表达力不足）。
    • β=1（primary/secondary 等权重）：同样显著掉点，说明二者需要非对称加权。
    • 自回归码预测：在 Kitchen 上反而掉点，但在真实机器人环境里因果自回归预测是必要的（论文称这一现象”anomalous”但真机上验证一致）。
    • 动作分块（action chunking）：在能做的环境上尝试后性能反而不如不分块——因 VQ-BeT 本身够快，不需要用分块换速度。
    • Offset 头：去掉后精度明显下降，说明为达到完整动作保真度这一头必不可少。 （以上均为论文正文的方向性结论，Figure 5 为雷达图汇总，未给出各消融项的具体数值增量。）

Infra（训练 / 推理工程）

  • 训练硬件：论文正文与附录均未披露训练用 GPU 型号、卡数或总 GPU-hours。

  • 推理效率对比（Table 3，Kitchen 环境，单步/多步 forward 延迟；Diffusion Policy 用 10 步迭代去噪，对应其真实世界部署设置）：

    方法组 A：C-BeT / C-BESO / CFG-BESO / VQ-BeT方法组 B：BeT / DiffusionPolicy-C / DiffusionPolicy-T / VQ-BeT
    Single step22.6ms / 25.9ms / 41.7ms / 22.8ms13.2ms / 100.5ms / 98.6ms / 15.1ms
    Multi step✗ / ✗ / ✗ / 23.3ms✗ / 100.7ms / 98.6ms / 15.2ms

    （表头原文如此：左侧表头标”Unconditional”、模型列为 C-BeT/C-BESO/CFG-BESO；右侧表头标”Conditional”、模型列为 BeT/DiffusionPolicy-C/T——与正文 4.1 节对”unconditional 用 BeT/DiffusionPolicy""conditional 用 C-BeT/BESO”的描述刚好对调，本页按 arXiv HTML 与 PDF 双重核对后原样转录，不做修正。）VQ-BeT 单次前向即可出动作（无论单步或多步动作块），Diffusion Policy 系需 10 步迭代去噪，故耗时约为 VQ-BeT 的 4~7 倍。

  • 真机推理延迟（Table 7，RTX A4000 GPU 工作站 vs Hello Robot 板载 4 核 Intel CPU）：VQ-BeT 18.06ms（GPU）/ 207.25ms（CPU）；DiffusionPolicy-T 573.49ms（GPU）/ 5243.82ms（CPU）；BC w/ Depth 5.66ms / 87.28ms；BC 4.73ms / 83.28ms。真机部署用全闭环控制（非 receding-horizon 开环），因为低成本移动机械臂运动噪声大，开环执行仅 3 个时间步就会明显偏离分布；闭环控制下 VQ-BeT 相对 Diffusion Policy 的推理时间差距拉大到板载 CPU 上的 25×（207.25ms vs 5243.82ms）。

  • 精度/并行策略：未披露。

评测 benchmark

Table 1（goal-conditional，7 个仿真环境，数值越高越好）：VQ-BeT 在除 BlockPush 外全部任务上超过 GCBC / C-BeT / C-BESO / CFG-BESO：PushT Final IoU 0.39（次优 C-BESO 0.30）；Image PushT 0.10（次优 C-BeT 0.02）；Kitchen Goals 3.78/4（次优 C-BESO 3.75）；Image Kitchen 2.60/4（次优 C-BeT 2.41）；Multimodal Ant 1.72/2（次优 C-BeT 1.68）；UR3 BlockPush 1.94/2（并列最优）；BlockPush 0.87/1（略低于 C-BESO 0.93、CFG-BESO 0.88，与 C-BeT 持平）。

Table 2（unconditional，数值越高越好）：VQ-BeT 在多数任务上超过或持平 BC / BeT / DiffPolicy-C / DiffPolicy-T：PushT Final IoU 0.78（次优 DiffPolicy-T 0.74）；Image PushT 0.68（次优 DiffPolicy-C 0.66）；Kitchen Goals 3.66/4（次优 DiffPolicy-T 3.44）；Multimodal Ant 3.22/4（次优 DiffPolicy-C 3.12）；UR3 BlockPush 1.84/2（次优持平 DiffPolicy-C/T ~1.82-1.83）；Image Kitchen 2.98（被 DiffPolicy-C 3.11 超过）；BlockPush 1.79（被 DiffPolicy-T 1.93 超过）。

行为熵（多模态捕获能力，4.3 节）：VQ-BeT 在除 Image Kitchen 外的全部环境上熵最高（被 DiffusionPolicy-T 反超）。

nuScenes 轨迹规划（Table 4，L2 误差/碰撞率越低越好）：VQ-BeT（仅部分信息输入）L2=0.73m、碰撞率=0.29%，在”partial information”模型里 L2 误差最低（优于 GPT-Driver 0.84/0.44%、扩散轨迹模型 0.96/0.49%），碰撞率也优于二者，且逼近使用完整场景信息的 Agent-Driver（0.74m/0.21%，full-info 最优）。

真机单阶段/两阶段任务（Table 5，共 50 次单阶段 + 30 次两阶段评测）：单阶段任务 VQ-BeT 总成功 47/50，DiffusionPolicy-T（改进版）45/50，BC w/ Depth 27/50，BC 29/50；两阶段任务 VQ-BeT 19/30 vs DiffusionPolicy-T 11/30，相对优势 73%（引言原话）。

真机长程任务（4 个任务，每个 3~4 个子任务顺序拼接，Table 6/Figure 6）：VQ-BeT 在任务终点的成功率至少是 DiffusionPolicy 的 3 倍（例如 Task1 终点 Close Drawer：6/10 vs 2/10；Task4 终点 Return to Home：6/10 vs 1/10），且优势在长程任务后段更明显。

码本尺寸消融（Table 12/附录 B.1）：码本组合数从默认的 64256 扩大到 10×250×，多数环境性能变化很小；仅 Kitchen（Conditional）性能随码本增大而下降，降幅约 4.5%（3.78→3.61）。

创新点与影响

  • 把 BeT 里不可微、难扩展的 k-means 动作分箱替换为可学习、可微的层级 Residual VQ-VAE（primary code 粗聚类 + secondary code 细修正 + 连续 offset 头找回精度），在保留 BeT 单次前向速度优势的同时把表达力提升到接近甚至超过 Diffusion Policy 的水平。
  • 论文用统一架构同时覆盖条件（goal-conditioned）与非条件行为生成两种任务形式化，并首次把这套动作量化+Transformer 预测的方案扩展到自动驾驶轨迹规划（nuScenes）与长程真机操作，证明其不局限于机器人操作单一场景。
  • 推理速度优势显著：仿真中相对 Diffusion Policy 快约 5×，真机闭环部署（受限于板载 CPU 算力）快 25×，这对依赖多步去噪、真机高频闭环控制吃紧的扩散策略是明确的工程可用性优势。
  • 作者自陈方向/局限：结论部分提出未来可探索方向——扩展到数量级更大的行为数据集（更多环境、本体、行为模式）；在不同本体间学习共享动作潜空间以实现策略”跨本体翻译”；用学到的离散动作空间加速真实世界 RL。论文本身未给出训练算力/GPU 配置，也未对 Figure 5 的五项消融给出具体数值增量（仅雷达图与方向性文字结论）。

原始链接

一手源存档（sources/）

关系图谱

反向链接