一句话定位

AdaWorld(HKUST / Harvard / UMass Amherst / MIT-IBM Watson AI Lab,ICML 2025)从无标注视频里自监督抽取连续潜动作(latent action),并以此为统一条件预训练一个自回归扩散世界模型:抽出的潜动作与具体画面内容解耦,使模型能在全新环境里零训练迁移已演示的动作,或仅用约 100 个动作样本 + 几百步微调就学会新的动作空间,在动作迁移、仿真保真度、视觉规划三项任务上全面超过动作无关预训练、光流条件、离散潜动作(Genie 式)三类基线。

背景与定位

世界模型的主流范式是从预训练视频生成模型出发(如以 Stable Video Diffusion 为基座的 Vista、GameGen-X 等),但要获得精确的动作可控性仍需大量动作标注和昂贵的训练,且新环境往往有不同的动作格式,导致每次适配都要重新付出高训练成本。为绕开人工动作标注,一条路线是用伪标签(VPT)标注视频;另一条路线是从视频里无监督抽取”潜动作”作为条件——但 Genie、IGOR、Moto 等工作大多聚焦离散潜动作 + 行为克隆目标,离散化带来表达歧义,且离散码本一旦训完数量就固定,难以灵活组合或迁移到未见过的细粒度动作。

AdaWorld 的核心主张是把潜动作压缩进连续潜空间:通过 β-VAE 信息瓶颈从相邻两帧里挤出”最本质的变化”,得到的潜动作具有上下文不变性(context-invariant)——同一个动作在不同场景里提取出的潜动作彼此接近,因而可以跨环境直接复用、平均、插值、聚类,构造出新的可控接口。模型骨干沿用 Vista 与 Jafar(flairox)的实现思路,并从 Stable Video Diffusion 权重初始化世界模型。论文将其定位为”世界模型预训练新范式”:把动作信息在预训练阶段就注入进来,而不是像现有方法那样只在动作无关视频上预训练、把动作可控性留到下游昂贵微调去解决。

模型架构

范式:两阶段——① 潜动作自编码器(信息瓶颈)无监督抽取动作;② 以抽出的潜动作为条件,预训练自回归视频扩散世界模型。

1)潜动作自编码器(Latent Action Autoencoder)

  • 基于 Transformer,500M 参数,16 层编码器块 + 16 层解码器块,通道数 1024,16 个注意力头;潜动作维度 32
  • 编码器:把相邻两帧 f_{t:t+1} 切成 16×16 image patch,投影为 patch embedding 并展平,与两个可学习 token a_{t:t+1} 拼接;施加正弦位置编码(ViT 式);用时空 Transformer(空间注意力看单帧内全部 token,时间注意力只看两帧同一空间位置的两个 token,并在时间注意力里加旋转位置编码表示因果关系)聚合两帧间的动态;取 a_{t+1} 位置的输出预测潜动作后验 (μ, σ),按标准 VAE 采样得到潜动作 ã。
  • 解码器:空间 Transformer,把采样出的 ã 附加到前一帧 f_t 上,在像素空间预测后一帧 f_{t+1}。
  • 目标函数是 β-VAE(预测重建项 + β·KL 项),β = 2×10⁻⁴,用于在”潜动作表达力”与”与上下文解耦的程度”之间取得平衡(若用标准 VAE 的 β=1,潜动作会因为后验约束过强而难以表达多样的帧间变化)。

2)自回归世界模型

  • 基于 3D UNet(沿用 Stable Video Diffusion (SVD) 架构,EDM 扩散框架初始化),1.5B 可训练参数,短期记忆窗口 K=6 帧历史帧。
  • 与原版 SVD(一次性生成整段视频片段)不同,AdaWorld 每次只对一帧加噪去噪,实现帧级(frame-level)而非片段级控制,交互粒度更细。
  • 条件注入:潜动作 ã 与 SVD 原有的 timestep embedding、CLIP image embedding 拼接;记忆中最后一帧作为 SVD 的条件图像;历史帧经 SVD 的图像编码器编码后与待预测帧的噪声 latent 沿通道拼接;训练时历史帧数量在 1~6 之间随机采样(记忆长度可变)。
  • 为缓解长程漂移(long-term drift),训练时对历史帧施加噪声增强(0.0~0.7 区间,间隔 0.1 随机采样)——论文发现即便推理时不加噪声,这一增强仍能有效改善长期漂移。
  • 两个模型默认输入分辨率均为 256×256

数据

规模:预训练数据总计约 2000M(20 亿)帧,来自 4 个公开数据集 + 2 个自动生成的游戏平台,共 1016 个环境(1000 个 Gym Retro + 16 个 Procgen)。

数据组成与配比(论文 Table 5,按视频数加权)

类别数据源帧数配比
Robot DataOpen X-Embodiment170M30%
Human ActivityEgo4D330M1%
Human ActivitySomething-Something V27M3%
2D Video GameGym Retro1000M49%
2D Video GameProcgen Benchmark144M2%
3D RenderingMiraData200M14%
City WalkingMiraData120M1%

(7 项加总 ≈1971M 帧,配比加总 100%,与正文”约 2000M 帧”一致。)

采集与处理

  • Open X-Embodiment:保留第一/第三人称示教,人工剔除静止、非连续、低帧率、低分辨率的”劣质子集”。
  • Gym Retro:从网上人工搜集 1000 个 ROM 导入模拟器;每个环境随机动作采集 1M 条 transition
  • Procgen:16 个环境,划出 1000 个起始关卡做评测、9000 个关卡用于训练,全部用 “hard” 难度;每个环境采集 10M 条 transition
  • 偏置动作采样(biased action sampling):不同于 Genie 对动作均匀采样,AdaWorld 在一段时间内提高某个特定动作的采样概率、随后再切换,鼓励探索更长的场景轨迹,实验显示比均匀采样能生成明显更丰富多样的场景(论文 Fig.7)。
  • 所有视频统一按元信息降采样到约 10 Hz 训练;训练时对输入帧做随机亮度抖动增强,不同宽高比统一做中心裁剪避免 padding。
  • 动作标注:预训练阶段不需要显式动作标签,全部动作信息由潜动作自编码器无监督抽取;仅在下游”世界模型适配”实验里才使用真实动作标签(每个动作 100 个样本)做嵌入初始化。
  • 数据多样性消融(论文 Table 4):分别用 OpenX、Gym Retro、Retro+OpenX 三种组合训练潜动作自编码器 40K 步,在 Procgen 上评测解码器预测的 PSNR/LPIPS:OpenX 单独 25.51/0.318,Retro 单独 26.43/0.250,Retro+OpenX 混合 26.62/0.234——即便 OpenX 主要是真实机器人视频,混入它仍能提升在 2D 虚拟游戏上的泛化,说明数据多样性对潜动作泛化有正向帮助。

训练方法

目标函数:潜动作自编码器用 β-VAE 目标(重建 log-likelihood − β·KL,β=2×10⁻⁴);世界模型用扩散重建损失 ‖x₀ − x̂₀(x_t, t, c)‖²,条件 c 包含历史帧与潜动作。

两阶段流程

  1. 从头训练潜动作自编码器:200K 步,batch size 960,AdamW,学习率 2.5×10⁻⁵,weight decay 0.01。
  2. 用固定的潜动作编码器抽取条件,预训练自回归世界模型:80K 步,batch size 64,学习率 5×10⁻⁵,16 张 NVIDIA A100 GPU;余弦学习率调度 + 10K 步 warmup;训练期间使用 EMA(指数滑动平均)提升预测质量。

采样:默认 5 步采样、classifier-free guidance scale 1.05;历史帧推理时不加噪声,但仍以增强等级 0.1 对世界模型做条件(作者发现这样效果更好);额外使用 timestep shifting 策略提升生成质量。

基线对比的训练预算(区别于上面 80K 步的正式预训练配置):为在 Table 1/2/3 里公平比较”动作无关预训练 / 光流条件 / 离散潜动作(VQ-VAE,8 个离散码,仿 Genie)/ AdaWorld”四种变体,论文将四者都统一训练 50K 迭代(架构、数据完全对齐,仅条件形式不同)。

下游适配(两种范式,均无需重新设计动作格式)

  • 零训练动作迁移:给一段演示视频,用潜动作编码器抽取动作序列,直接在新场景初始帧上复用该序列自回归生成,无需任何训练。
  • 有限交互微调适配:新环境收集少量动作-视频对(每个动作 100 个样本),用潜动作编码器推断各自的潜动作并取平均作为该动作的嵌入初始化(continuous 潜空间使”平均”有意义),再用该初始化对整个模型做少量步数微调(三个仿真环境微调 800 步,视觉规划环境微调 500 步,batch size 32,学习率 5×10⁻⁵,预训练权重的学习率再打 0.1 折)。
  • 额外能力:动作组合(对两个观测到的潜动作取平均,得到融合两者功能的新动作);动作创建(对潜动作做 K-means 聚类,聚出任意数量的可控选项)。

视觉规划(MPC):基于 Cross-Entropy Method 的采样式 MPC——初始均匀分布采样 N=100 条长度 L=15 的动作序列,世界模型 rollout 后按预测观测与目标图像的余弦相似度算奖励,取 top K=10 更新分布,共 i=2 轮 CEM 迭代,执行最优序列的前 T=5 步,最多尝试 20 步判定成功/失败。

Infra(训练 / 推理工程)

  • 训练硬件:世界模型预训练用 16×NVIDIA A100;GPU-hours、并行策略(数据/模型并行细节)未披露
  • 精度:论文未明确说明混合精度设置(未披露 fp16/bf16 选择)。
  • 推理:默认 5 步扩散采样、CFG scale 1.05;论文在”局限性”中明确承认当前不支持实时(real-time)频率推理,具体 FPS / 延迟数字未披露;作者建议未来可引入蒸馏/加速采样技术(如 Diffusion Forcing、一致性蒸馏)来提速。
  • 边缘硬件 / 部署:未涉及,论文只做研究原型,无边缘端部署报告。

评测 benchmark

Table 1·动作迁移(LIBERO / SSv2,未见过的评测集,1300 对视频,FVD↓ / ECS↑)

方法LIBERO FVDLIBERO ECSSSv2 FVDSSv2 ECS
Act-agnostic1545.20.702847.20.592
Flow cond.1409.50.724702.80.611
Discrete cond.1504.50.700726.80.596
AdaWorld767.00.804473.40.639

Table 2·世界模型适配(Habitat / Minecraft / DMLab,100 样本/动作,800 步微调,PSNR↑ / LPIPS↓)

方法Habitat PSNRHabitat LPIPSMinecraft PSNRMinecraft LPIPSDMLab PSNRDMLab LPIPS
Act-agnostic20.340.45019.440.53220.960.386
Flow cond.22.490.37320.710.49222.220.357
Discrete cond.23.310.34221.330.46522.360.349
AdaWorld23.580.32721.590.45722.920.335

另在 Minecraft 上做了样本量/微调步数的扫描曲线(Fig.6):AdaWorld 在微调初期即优于其他方法,且提升速度明显更快。

Table 3·视觉规划(Procgen 4 环境×30 场景,成功率±标准误,5 个随机种子)

方法HeistJumperMazeCaveFlyer平均
Random19.33±1.80%22.00±1.02%41.33±2.22%22.00±1.02%26.17±0.52%
Act-agnostic20.67±1.45%20.67±1.00%39.33±1.17%23.33±0.75%26.00±0.20%
AdaWorld(不微调,仅换嵌入)38.67±0.82%68.00±0.92%41.33±1.11%31.33±1.02%44.83±0.28%
AdaWorld(微调)66.67±1.67%58.67±1.02%68.00±0.69%33.33±1.55%56.67±0.44%
Q-learning22.67±1.58%47.33±2.74%4.67±0.33%34.00±2.52%27.17±0.26%
Oracle(真实环境)86.67±1.29%77.33±1.09%84.67±1.19%74.00±1.63%80.67±0.43%

即便不做任何权重微调(仅用采样得到的平均潜动作替换嵌入),AdaWorld 平均成功率(44.83%)已超过微调过的动作无关基线(26.00%);微调后进一步达到 56.67%,远超经典无模型方法 Q-learning(27.17%),但仍明显低于用真实环境做规划的 Oracle 上界(80.67%)。

Table 4·数据多样性消融(见”数据”节);Table 7·动作可控性(ΔPSNR,K-means 聚出 4~12 个动作选项):AdaWorld 在 4/5/6/7/8/9/10/11/12 个选项下分别取得 5.67/5.15/7.28/8.23/6.26/7.32/6.07/6.68/6.53,控制力与固定 8 码本的离散基线(仅 8 个选项时 ΔPSNR=6.47,其余选项数下 N/A,因码本数量训完即固定)相当,但 AdaWorld 支持任意数量的动作选项自由定制。

创新点与影响

  • 首次证明”连续”潜动作可作为世界模型预训练阶段的统一条件:不同于 Genie 类工作聚焦离散潜动作 + 行为克隆,AdaWorld 把潜动作压缩进连续空间,换来跨上下文的可迁移性、可组合性(平均/插值/聚类)。
  • 提出两种低成本适配范式:零训练动作迁移(复用潜动作序列)与有限交互微调适配(用真实动作的平均潜动作初始化嵌入),使世界模型能用远少于传统方法的标注数据和训练步数(100 样本 + 几百步)掌握新动作空间。
  • 在动作迁移、仿真保真度、视觉规划三类任务上系统性优于三类基线(动作无关预训练、光流条件、离散潜动作条件),且不微调也能超过微调过的动作无关基线,说明预训练阶段注入动作信息本身就是更好的初始化。
  • 作者自述局限:① 不支持实时频率推理,需要未来引入蒸馏/加速采样;② 当 rollout 超出初始场景范围时难以生成全新内容(类似先前工作的通病),需要靠扩大模型和数据规模缓解;③ 极长程 rollout 的质量仍不足,留待未来工作解决(如结合 Diffusion Forcing 等技术)。

原始链接

一手源存档(sources/)