一句话定位

ATM(Any-point Trajectory Model)把「从视频中学控制」的中间表征从像素预测换成任意 2D 点的未来轨迹:先在无动作视频上预训练一个 track transformer,预测画面上任意采样点在未来若干步的位置,再用这些预测轨迹作为稠密子目标去训练一个只需极少量带动作标注演示的视觉运动策略。在 LIBERO 130 个语言条件操作任务上,平均成功率 63% 对比此前最强视频预训练基线 37%,提升超 80%;并进一步展示了人类视频、跨机器人本体视频的零动作标注迁移能力。RSS 2024 论文。

背景与定位

论文要解决的问题:机器人演示数据(action-labeled trajectories)采集昂贵——RT-1 收集 130K 条轨迹耗时 17 个月——而无动作标注的视频(人类操作视频、其他机器人视频)便宜且量大,但缺少动作标签使其难以直接用于模仿学习。此前两条路线各有缺陷:一是学习自监督特征表示(如 R3M、VIP),只描述当前时刻状态,丢失了视频里的转移动态信息;二是像素级视频预测模型(如 UniPi、AVDC/Ko et al.)显式建模转移动态,但把物理运动和纹理光照等无关因素耦合在一起,容易产生幻觉和不真实的未来预测,且训练推理计算量大,高推理成本又迫使这类方法采用开环执行,策略鲁棒性差。

ATM 的立意是用一种介于两者之间的结构化表征——2D 点轨迹(tracks/particles)——来解耦运动信息与纹理光照等无关像素信息,天然带有物体恒常性(object permanence)等归纳偏置,同时保持轻量、可闭环推理。方法直接建立在近年 Tracking Any Point (TAP) 进展之上,用现成的点跟踪器 CoTracker(Karaev et al. 2023)为视频生成自监督轨迹标注。与同期 RoboTAP(Vecerik et al. 2023)不同——RoboTAP 不学轨迹模型,而是在测试时直接跑 tracker 做视觉伺服(需要把任务切分成多阶段、逐阶段预测目标点位置)——ATM 训练一个可预测未来轨迹的模型,框架更简单、适用场景更广。实验主战场是 LIBERO 基准(130 个语言条件操作任务),并对比 Diffusion Policy 验证通用性。

模型架构

ATM 分两阶段、两个独立模型:

(1) Track Transformer(轨迹预测模型,视频预训练阶段) 输入:当前帧图像 o_t、该帧上任意一组 2D 查询点的位置 pt = {p_t,k}{k=1}^K、语言指令 ℓ;输出:这些点未来 H 步的轨迹 p_{t:t+H} ∈ ℝ^{H×K×2}(仅建模相机坐标系下的 2D 轨迹,不假设标定相机或深度可用,便于扩展到更多样的视频源)。

  • 多模态 masked prediction:轨迹 token(每个点单独编码成一个 token,未来位置在编码前被 mask 掉)、语言 token(预训练 BERT 编码器)、图像 patch token(随机 mask 50% patch,参考 MAE 做重建作为辅助任务)三路输入共享 embedding 空间,一起送入一个大 transformer,解码出未来轨迹;辅助任务是重建被 mask 的图像 patch。
  • 训练时用启发式过滤掉背景静止点:先在帧上采样 n×n 网格点、跟踪整段视频得到初始轨迹 τ∈ℝ^{n²×T×2},按位置方差过滤掉基本不动的点,再在过滤后位置附近重采样、用 tracker 生成最终轨迹标注,避免静态相机下大部分随机点落在背景、训练信号稀疏。
  • 轨迹长度(预测 horizon)默认 H=16 步;消融显示 4/8/16/32/64 步中 16 步平均最优,更长轨迹在 LIBERO-Long/Goal 上因噪声增大反而掉点。

(2) Track-guided Policy(轨迹引导策略,模仿学习阶段) 策略 π(a_t | o_t, p_{t:t+H}) 输入当前观测和(冻结的 track transformer 预测出的)未来轨迹,输出动作。测试时不再用训练阶段的方差过滤,而是固定用网格采样的 32 个点(每视角)喂给策略——ATM 对输入点集合具有置换不变性且对采样分布鲁棒,训练和推理可以用不同的点采样方案。

  • 策略骨架沿用 LIBERO/ViLT 一路的 ViT-T 三段式:空间编码(各视角 track token + 图像 patch token + 一个 spatial CLS token 做自注意力,取 CLS token 作为该时刻表征)→ 时序解码(本体感受信息映射到同一 embedding 空间,与 spatial CLS token、action CLS token 沿时间轴做因果自注意力)→ 动作头(每个时间步的 action CLS token 与当前时刻的重构轨迹拼接后过 MLP 回归动作,MSE loss)。
  • 轨迹在策略中早融合 + 晚融合两处注入:早融合(送入 spatial transformer)+ 晚融合(拼到 action head 前);消融显示去掉晚融合对性能损伤最大,说明未编码的原始轨迹本身就是易于解读的目标信号。
  • 有意思的发现:轨迹已把「做什么」的高层任务目标分解成逐步细粒度子目标,策略阶段不再需要语言指令作为任务指定——相当于把策略学习问题退化为一个逆动力学(inverse dynamics)问题。
  • ATM 是模型无关的条件信号:论文额外把预测轨迹作为附加条件接入 Diffusion Policy(ATM Diffusion Policy),验证了跨策略架构的通用性。

数据

纯模拟环境实验数据来自 LIBERO 基准:5 个任务套件(LIBERO-Spatial/Object/Goal/Long 各 10 任务,LIBERO-90 90 任务),每任务 50 条人类专家演示(通过 3Dconnexion SpaceMouse 遥操作采集)。默认设置下每任务全部 50 条演示都被剥离动作标签用作 track transformer 预训练的无动作视频(每个 10-任务套件累计 500 条视频),而策略训练阶段只用其中 20%(10 条/任务)保留动作标签的演示做模仿学习。观测含第三人称相机 + 腕部相机 RGB 图像(分辨率 128×128)与本体感受(关节/夹爪状态),语言指令过预训练 BERT 编码。动作空间 7 维(末端执行器平移、旋转、开合)。

真实世界实验(v3 新增):6-DOF UR5 机械臂,GELLO 遥操作系统采集,关节位置控制 + 夹爪状态动作空间,两个 RealSense 相机 RGB 观测(叠 2 帧补偿部分可观测,不喂本体感受以避免策略过拟合本体、忽视视觉输入);共采集 50 条带动作标注轨迹 + 250 条无动作视频(后者由前者去掉动作标签得到,5× 视频量)。

跨具身迁移实验(三组):

  • 人类→机器人(3 个真实任务:叠布料、放番茄关柜门、扫具+簸箕):每任务采 10 条机器人遥操作演示(带动作)+ 100 条人类直接操作视频(无动作)
  • 机器人→机器人(Franka→UR5 拾放任务):160 条 Franka 无动作视频 + 10 条 UR5 带动作演示

未见跨数据集/网络规模视频预训练(未如 EPIC-KITCHENS/Ego4D 大规模网络视频那样的大规模预训练,仅 in-domain 视频,论文自陈为局限)。

训练方法

两阶段流水线,完全解耦:

  1. Track Transformer 预训练:在无动作视频集 𝒯_o 上,用 CoTracker 生成的 ground-truth 轨迹作为监督,多模态 masked prediction 目标(轨迹回归 + 图像 patch 重建辅助任务),不做帧堆叠(frame stacking)以避免 copycat/causal confusion 问题;用验证 loss 最低的 checkpoint 作为最终模型,之后冻结。
  2. Track-guided Policy 训练:用少量带动作标注演示 𝒯_a,输入观测 + 冻结 track transformer 预测的轨迹,MSE 回归动作;因行为克隆验证 loss 不总可靠,保存最后一个 epoch 的 checkpoint 做在线 rollout 评测。

超参数(LIBERO 实验固定值):

  • Track Transformer:100 epoch,batch size 1024,AdamW,lr 1e-4,weight decay 1e-4,cosine 调度 + 5 epoch warmup,clip grad 10,点采样用方差过滤,点数 32,轨迹长度 16,轨迹 patch size 4,图像 mask ratio 0.5,数据增强 ColorJitter + RandomShift。
  • Track-guided Policy:100 epoch,batch size 512,AdamW,lr 5e-4,weight decay 1e-4,cosine 调度,clip grad 100,点采样用网格,点数 32,轨迹长度 16,帧堆叠 10,数据增强同上。

基线对照四种视频预训练范式:BC(无视频先验)、R3M-finetune(对比学习表征)、VPT(先学逆动力学模型给视频打伪动作标签再 BC)、UniPi/UniPi-Replan(文本条件视频扩散模型生成图像子目标 + 逆动力学动作预测,开环 vs. 每 8 步重规划)。

Infra(训练 / 推理工程)

  • 训练:4× A100 GPU,DeepSpeed 训练策略(两阶段模型均在此配置下训练,论文未进一步拆分单阶段用时)。
  • 推理计算与延迟对比(V100 GPU 单次生成,闭环 vs 开环):
方法计算量 (TFLOPS/次生成)单次生成耗时 (s)
ATM(闭环,每步生成一次轨迹)1.560.015
UniPi-Replan(开环,每 8 步重生成一次图像子目标)13.094.51
UniPi(开环,仅初始生成一次)39.298.14

ATM 的轨迹生成比像素级视频扩散便宜约 25×(对 UniPi)到 8×(对 UniPi-Replan),且单次耗时仅 15ms,是唯一在计算上可行的闭环方案——这也是 ATM 相对 UniPi 系基线的核心工程优势之一:可以逐步重新预测子目标而非依赖昂贵的开环视频生成。真实机器人部署的控制频率/端到端时延未披露。

评测 benchmark

LIBERO 主结果(20% 带动作标注演示,成功率 % ± 标准差,3 个随机种子):

方法SpatialObjectGoalLongLIBERO-90
BC-Full-Trainset(Oracle,100% 数据)71.83±3.7071.00±7.9776.33±1.3124.17±2.59-
BC(无视频先验)39.00±8.2051.83±13.5442.50±4.9516.67±3.6629.78±1.14
R3M-finetune49.17±3.7952.83±2.255.33±1.439.17±2.669.59±0.27
VPT37.83±4.2919.50±0.823.33±2.363.83±1.65-
UniPi69.17±3.7559.83±3.0111.83±2.025.83±2.08-
UniPi-Replan(单种子,仅 Object/Goal)-31.003.00--
ATM(本文)68.50±1.7868.00±6.1877.83±0.8239.33±15.8048.41±2.09

平均成功率 ATM 63% vs. 此前最强基线(BC/UniPi 各套件取高)37%,提升超 80%。ATM 在 LIBERO-Long(长时程)与 LIBERO-90(任务/物体/布局最多样)两个最难套件上优势最大;用仅 4% 演示的 ATM 就能追平 BC 用 20% 演示的表现,20% 演示的 ATM 逼近 BC 用全量数据的表现。

Diffusion Policy 通用性验证(同样 20% 数据):

方法SpatialObjectGoalLongLIBERO-90
Diffusion Policy67.67±1.2578.00±2.4535.00±3.7437.33±2.0533.85±1.71
ATM Diffusion Policy79.00±3.7481.00±2.4558.67±4.6444.00±6.3862.89±1.10

轨迹条件对 Diffusion Policy 同样带来一致提升,尤以 LIBERO-90(+29 点)、LIBERO-Goal(+23.7 点)明显。

消融:图像 mask(50% vs 0%)在 Object/Goal/Long 三个套件上有小幅正向收益(Long 39.33 vs 30.83);早融合+晚融合同时保留优于单独任一项,去掉晚融合损伤最大(如 Goal 从 77.83 掉到 5.33)。

真实世界 & 跨具身:UR5 Kitchen 5 任务,ATM 一致优于 BC 与其他视频预训练基线(具体每任务数值论文未列表,仅图示)。人类→机器人迁移量化结果(10 次 rollout 平均成功率):

方法遥操作演示人类视频叠布料放番茄关柜门扫具
BC0%10%30%
ATM(仅机器人视频)0%0%13%
ATM(+人类视频)63%63%60%

机器人→机器人(Franka→UR5 拾放):UR5-only ATM 47% → Franka+UR5 ATM 80%。

创新点与影响

  • 把「视频预训练做子目标生成」的中间表征从像素/图像换成结构化的任意点 2D 轨迹,用一次性、低成本的轨迹生成替代计算高昂、易产生幻觉的像素级视频扩散/预测,同时保留了显式建模转移动态(相对纯表征学习如 R3M/VIP 的优势)。
  • 证明轨迹子目标比图像子目标(UniPi 系)更利于长时程、多目标任务:闭环、低计算量使其可以逐步重新预测,而不必依赖昂贵的开环视频生成。
  • 用一个共同的、跨具身的 2D 点轨迹空间,展示了从人类视频、跨机器人本体视频向目标机器人策略迁移的可行性——用极少量目标域带动作演示(10 条)加大量跨域无动作视频即可显著提升成功率。
  • 验证了该轨迹条件对不同策略架构(transformer BC、Diffusion Policy)均通用,是一个可插拔的模块而非绑定特定策略骨架。
  • 论文自陈的局限(Limitations 章节):仅在 in-domain 视频上预训练,尚未验证对更大规模、更多样视频数据集(如互联网视频)的扩展;训练策略阶段仍依赖带动作标注的演示数据,未探索无演示的强化学习设定。

原始链接

一手源存档(sources/)