AI Research 技术调研

HPT: Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

14分钟阅读

一句话定位

MIT CSAIL + FAIR（Lirui Wang、Kaiming He 等）提出的跨本体策略预训练框架：给每个本体（embodiment）配一个小型「stem」tokenizer，把各不相同的本体感受（proprioception）+ 视觉输入对齐成固定 16 个 token 的短序列，喂进一个所有本体共享的 Transformer trunk，用行为克隆在 52 个异构机器人数据集上监督预训练，首次系统展示了机器人策略预训练随数据 / 模型 / 算力缩放（scaling）的行为；trunk 冻结后只重训 stem/head（约占参数 2%）即可迁移到全新本体，下游任务成功率提升 20%+。NeurIPS 2024 Spotlight。

背景与定位

机器人学习长期受困于 heterogeneity（异构性）：不同机器人是物理上不同的硬件本体，自由度、末端执行器、控制器、动作空间各异（proprioception 异构），相机数量 / 安装位置 / 视觉外观也天差地别（vision 异构）。以往做法是针对单一本体、单一任务采集数据训练，昂贵且过拟合。

同期主流跨本体路线（rt-1、RT-2-X、Octo、openvla、Open-X Embodiment）走的是统一数据格式 / 统一动作空间 + 语言条件的 VLA 范式，且大多只预训练视觉主干（R3M、VC-1、Voltron 只训 vision encoder），把 proprioception 事后（post-hoc）在迁移时才加进去。HPT 的差异化定位：

  • 不强行同质化：不把所有本体拉平成统一数据形式，而是承认异构、用「per-embodiment tokenizer」来吸收异构——思路直接借鉴多模态对齐（Flamingo、Perceiver、ImageBind、LLaVA），把「不同本体」类比成「不同模态」。
  • proprioception 是一等公民：视觉与本体感受在预训练阶段就联合训练，而非迁移时补丁。
  • 范式命名：Heterogeneous Pre-training（异构预训练）+ 类 Mixture-of-Experts 的路由（每个本体一个 expert stem/head，router 由本体决定），共享 trunk 累积所有本体梯度。
  • 核心科学问题：机器人策略预训练是否像 NLP/CV 一样存在 scaling behavior？HPT 把规模推到 50+ 数据源、1B trunk 参数来回答。

模型架构

模块化为三段 Stem（本体专属）→ Trunk（全局共享）→ Head（任务专属）（Figure 2）。对某个本体，一对 stem/head 被激活（开关式路由），trunk 恒定共享、参数量与本体 / 任务数无关。

Stem（本体专属 tokenizer，参数极小以防过拟合）

  • Proprioception tokenizer：对本体 k，把维度 d^k_p 的本体感受序列（如末端位姿 7 维）先用一个 MLP 投影到隐维 d，加 正弦位置编码，再用 attention 让 N_p=16 个可学习 token 去 attend 状态特征，输出 16 个 d 维 token。隐维 d 取值范围 128–1024。
  • Vision tokenizer：对任意相机图像序列（多视角视频）H×W×3，先用冻结的预训练视觉编码器抽特征（默认 ResNet18，图像 resize 到 224×224，取 global pooling 前的 7×7 特征），flatten 后同样用 attention 把它映射到 N_v=16 个可学习 token。多视角时每个视角单独一个 projector MLP 再拼接。
  • 各模态分别处理后拼接，加 modality embedding + 正弦位置编码，作为 trunk 输入序列。关键设计：无论输入序列多长，token 数固定为每模态 16 个（对比 Octo 用 256 个视觉 token），保证固定 context length 的缩放 / 推理收益。
  • Stem 只有「一个 MLP + 一个 attention 层」的少量参数。

Trunk（共享 Transformer，预训练主体）

  • 标准 Transformer，隐维 d，输出序列长度 L = 输入序列长度，输出 token 做 pooling 得到该 observation 的综合特征。
  • 五档规模（Table 1，参数只算 trunk）：
型号DepthWidthAttn HeadsTrunk 参数
HPT-Small1612883.1M
HPT-Base16256812.6M
HPT-Large16512850.5M
HPT-XLarge3276816226.8M
HPT-Huge801024161.1B

（发布的 checkpoint 另有 HPT-Base-Lang 50.6M）

Head（任务专属动作解码器）

  • 取 trunk pooled 特征映射到各数据集动作空间 A，可为任意架构（MLP / Transformer / diffusion policy——架构与 head 解耦）。
  • 动作按数据集统计 element-wise 归一化到 [-1,1]。固定 动作 horizon = 8、观测 horizon = 4，训练时沿时间维随机 mask 以适配不同下游 horizon。
  • 迁移到新本体时 head 重新初始化。

动作表示：默认对连续值直接回归（为通用性）。消融里也试过 rt-1 式的离散化 token 动作头（learned position encoding 作 action token + conv1D 回归头），但最终选连续回归。

数据

两档设置（Table 2）：

设置数据集数轨迹数样本数Batch
Default2716k5M256
Scaled52~200k–270k155M2048

Scaled 的 52 个数据集构成（Table 4）：41 个来自 Open-X Embodiment（含 Droid，正文另处也写「a subset of 42 datasets」）+ 7 个仿真 + 3 个人类视频 + 1 个野外部署机器人 = 52。四类异构域：

  • 真机遥操作：Open-X Embodiment 子集（各机构不同任务 / 环境 / 本体），默认每数据集最多取 1000 条轨迹。
  • 仿真（7 个）：跨 Drake、Mujoco（Metaworld / Robomimic）、Isaac Sim、PyBullet、Sapien、Flex，图像输入 + 专家演示。
  • 人类视频（surrogate 动作标注）：EPIC-Kitchen 用检测框归一化 2D 中心当 proprioception、与下一帧之差当 action，共 2000 条 video clip（最大长 500）；PoCo 用手部 3D 位置 + ICP 抽的 6-DoF 位姿当 action。（正文 4.3 处口径较小：EPIC + PoCo 共 300 条；附录 A.1 给的是 EPIC 2000 条。）
  • 野外部署机器人：FrodoBot-2k（野外游戏驾驶的移动机器人），前置相机，动作为线速度 / 角速度，proprioception 用 IMU，取 150 条轨迹、每条 >500 步。

采样策略：多任务 balanced weighted sampling——按每个数据集大小的平方根归一化作采样概率（inverse-exp 温度），防止大数据集主导一个 epoch。默认对各本体数据集用均匀分布采样。

数据处理刻意最小化：不做 Octo / OpenVLA 那样的动作 / 观测空间对齐与清洗，只做归一化（作者也自评 data cleanup/filtering 会有帮助但本工作欠探索）。视觉特征离线预抽后存盘。总数据集磁盘约 10 TB。

训练方法

  • 目标函数：跨 K 个异构数据集的行为克隆（BC），对归一化动作标签与预测用 Huber loss（δ=0.1）。总参数 θ = 各本体 stem/head ∪ 单一共享 trunk。
  • 两条数据缩放轴：单数据集轨迹量 M_k、数据集总数 K。
  • 类 MoE 的更新频率差异：预训练时每次迭代只更新 trunk；各本体 stem/head 按训练 batch 采样到才更新。分布式下每个 stem/head 由对应本体数据训练，trunk 累积所有 worker 的梯度。
  • 优化器：AdamW，weight decay 0.05，base lr 2e-4，cosine schedule + warmup + dropout，lr 按 batch size 比例缩放。作者称预训练对 lr schedule / 优化器 / 验证集选择都相当稳定。
  • 默认预训练算例：HPT-Small trunk（3.17M），27 数据集、每集 ≤1000 轨迹、共 16k 轨迹，batch 256、80k 迭代 ≈ latent 空间 0.65B token / 感知 token 空间 5B token（horizon 相关）。
  • 缩放实验协议：数据缩放让「更大模型 + 每上一个数量级轨迹数就 batch 翻倍」协同（否则 HPT-S 固定迭代会在每集 ~1000 轨迹处早早 plateau）；模型缩放从 1M→1B、batch 256→2048（每上一个数量级翻倍）、用 170k 轨迹；深 / 宽缩放无显著差异。评测用最后一步在固定验证集（默认 27 集，每集 ≤200 轨迹）上的平均验证 loss（借鉴 NLP 用 perplexity 度量预训练进度）。
  • 迁移学习：给新本体新数据集，冻结 trunk，按新本体的 proprioception / action 维度重新初始化 stem 和 head，用与预训练相同（或替代如 diffusion）的目标微调；可训练参数（stem+head）可低至全模型 2%，激活参数 <3 MB。

Infra（训练 / 推理工程）

  • 训练硬件：8× 到 128× NVIDIA V100（跨数量级算例）；训练时长半天到 1 个月。RAM 需求 <50 GB，数据盘 ~10 TB。
  • 算力 / token 跨度（Table 5 / 附录 A）：模型 0.5M–1B 参数；全模态数据 token ~32M–5B；单 batch token 0.03M–2M；单迭代 FLOPs 0.03–31 GFLOPs；80k 迭代（batch 512）近似 epoch 数 200→2。作者明确自评：token / FLOPs 规模只到 LLM 训练的中等量级，仍未充分收敛，异构预训练收敛慢。
  • 推理 / 边缘：迁移后闭环推理在旧款 RTX 3070 上——正文 5.1 给 HPT-base 47 Hz / HPT-XL 19 Hz；附录 A.2 给（含全部预处理 + encoder 时间）HPT-base 39 Hz / HPT-xlarge 33 Hz；A100 约快 3–4 倍。真机遥操作控制频率 10 Hz。真机微调在单张 RTX 2080Ti 上 batch 256、20000 迭代 ≈ 4 小时。

评测 benchmark

预训练用平均验证 loss（无绝对数值表，仅缩放曲线 Fig 5–8：数据 / 模型 / epoch/batch 三轴均可持续降 loss 直到 plateau；加仿真 + 人类视频联合预训练在大 embodiment gap 下仍成立）。

仿真迁移：Metaworld、Robomimic、Fleet-Tools、Simpler 四类 benchmark。用 HPT-Small trunk 迁移，5 次独立训练取平均；对比 No-Trunk / From-Scratch / Pretrained-Frozen / Pretrained-Finetuned（HPT-B、HPT-XL）。HPT-B→HPT-XL 迁移成功率随规模改善。Simpler benchmark（Google EDR 本体，Close Drawer / Move Near / Pick Coke Can，>300 episodes）上与 Octo、RT-1-X、RT-2-X 对比（HPT-base backbone，用官方 baseline 结果）。

真机迁移：两个本体、pet care + assembly 四任务（Sweep Leftover / Fill Water / Scoop Food / Switch Insertion），均不在预训练数据中；1 vs 2 相机、relative vs absolute pose 两种动作空间；每任务 ~100 demo、15 trial。HPT 稳定优于 No-Trunk / From-Scratch（Fill-Water 的 From-Scratch 用了 SOTA diffusion policy 架构做对比）。

Sweep Leftover 消融（Table 3，成功率 %）——凸显视觉-only 预训练与「联合预训练 proprioception」的差距：

方法Success (%)
From Scratch No Prop.26.7 ± 3.3
From Scratch43.3 ± 3.8
R3M50.0 ± 3.0
Voltron46.7 ± 3.8
VC-153.3 ± 2.6
No Prop. Finetuned63.3 ± 2.6
HPT-B Finetuned70.0 ± 3.0
HPT-XL Finetuned76.7 ± 3.3

总体：微调后策略在多个仿真 / 真机 unseen 任务上成功率提升 20%+。

创新点与影响

  • 异构预训练范式：不同本体 = 不同模态，用 per-embodiment stem tokenizer 对齐到固定 16-token 短序列 + 共享 Transformer trunk，摆脱 VLA 路线「统一动作 / 观测空间 + 语言条件」的前置约束，能一并吸收真机 / 仿真 / 人类视频 / 部署机器人四类异构域。
  • 首次系统展示机器人策略预训练的 scaling：数据（52 集、270k 轨迹）、模型（trunk 到 1B）、算力三轴联合缩放可持续降验证 loss，且预训练 loss 与下游成功率有相关性——把「策略学习也有 scaling law」这件事做实。
  • proprioception 与 vision 联合预训练（而非事后补丁），在 Sweep Leftover 上对比 R3M/VC-1/Voltron 等视觉-only 方法明显占优。
  • 极轻量迁移：冻结 trunk、只重训 ~2% 参数（<3 MB）即可上新本体，head 可插 diffusion policy 等任意架构，工程上「把可缩放 Transformer 塞进你策略的中间、别从头训」。
  • 开源：代码（liruiw/HPT + liruiw/lerobot 的 HF 版实现）、五档权重（HPT-Small/Base/Large/XLarge + Base-Lang）全公开。
  • 作者自述局限：balanced mixture 的本体划分过于简单、数据过滤 / 质量把控欠探索；仅监督学习目标、token/FLOPs 只到 LLM 中等量级未充分收敛，异构预训练收敛慢；仿真 + 真机评测都限于固定本体的短程操作，可能限制高容量模型的收益；学到的策略可靠性仍不高（多数任务成功率 <90%），失败模式常见 over/undershoot（浇水 / 舀食物欠到位）。

原始链接

一手源存档（sources/）

关系图谱

反向链接