一句话定位
IRASim(ByteDance Research/Seed 与港科大合作,2024-06 挂 arXiv、2025-07 改版为 ICCV 2025 收录版)是一个帧级动作条件的视频世界模型:给定初始/历史观测帧与一段精确到每一帧的机械臂动作轨迹(最多 7-DoF 连续动作),用一个**扩散 Transformer(DiT)**在 VAE 潜空间里逐帧生成”该轨迹执行后会发生什么”的高分辨率、长时程视频。核心发明是把轨迹条件从”整段轨迹编码成一个 embedding”(video-level,类似文生视频的文本条件)细化为”每个动作单独编码、逐帧注入 AdaLN”(frame-level),从而做到动作-帧的精确对齐。改版后的 v2 进一步把它从单纯的视频生成 benchmark 论文,扩展成”世界模型当模拟器/规划器”的论文:策略评估与真实 Mujoco 仿真器的成功率相关系数达到 0.99;作为基于模型规划(model-based planning)的动力学模型,把 Push-T 基准上一个普通扩散策略的 IoU 从 0.637 提升到 0.961。
背景与定位
IRASim 把自己放在”机器人操作世界模型”这条线上,与三类工作对照:
- 文本/语言条件的视频世界模型:unisim(UniSim)用文本+动作提示文生视频模型模拟真实世界交互,但只演示了 2D 平面上的机械臂移动;IRASim 直接吃真实机械臂的物理动作(3D 平移+旋转+夹爪),不经过语言中介。
- 隐动作(latent action)世界模型:Genie 从无标注互联网视频里学习隐式动作空间,让用户”扮演”隐动作与环境交互;IRASim 用的是真实、有物理意义的机器人动作(7-DoF),而非学出来的隐变量。
- 自回归 token 世界模型:ivideogpt 用条件 VQGAN 压缩 token 化 + GPT 式自回归 Transformer 做逐 token 预测;IRASim 走的是扩散 Transformer(DiT)+ 潜空间扩散路线,直接借鉴 Sora/DiT 的架构选择,并在 v2 的 RoboNet 设置上与 iVideoGPT、MaskViT 直接对比。
论文自称是”第一个在真实世界里建模复杂 7-DoF 连续机械臂动作视觉动力学的工作”。范式命名:trajectory-to-video / frame-level action-conditioned diffusion transformer world model。
模型架构
- 骨干:Diffusion Transformer(DiT),在 SDXL 预训练且冻结的 VAE 潜空间中做扩散,而非直接在像素空间。每帧潜表示
z^i = Enc(I^i)含P个 patch token(patch size 2),按时间顺序拼接得到N×P个 token,叠加空间位置编码与时间位置编码。 - 计算效率设计:标准 Transformer 对
N×P个 token 做全量 MHA 是平方复杂度;IRASim 采用空间-时间分离注意力(借鉴 Genie/Latte 一路的做法,代码实现直接改自 Latte)——空间块只在同一帧内的1×Ptoken 上做 MHA(建模帧内交互),时间块只在同一 patch 位置跨N帧的N×1token 上做 MHA(建模帧间交互),使长视频、高分辨率生成在计算上可行。 - 历史/初始帧条件:v1 只用单帧
I^1作条件,训练时只对第 2~N 帧的 token 加噪、只在这些帧上算 loss,初始帧 token 保持”干净”混入序列,靠注意力机制维持与初始帧的一致性;v2 把这一设计泛化为h帧历史观测I^{t-h:t}条件(RT-1/Bridge/Language-Table 用 1 帧历史预测 15 帧,RoboNet 沿用 iVideoGPT 设置用 2 帧历史预测 10 帧)。 - 轨迹条件——两种 AdaLN 变体(这是论文的核心贡献点):
- Video-Level 条件(Video-Ada):整段轨迹用一个线性层编码成单一 embedding,与扩散时间步 embedding 相加后,为每个 spatial/temporal 注意力块回归出 scale(γ,α)与 shift(β)参数——类比文生视频里”一段文本条件全片”。
- Frame-Level 条件(Frame-Ada,主推变体):每个动作单独经线性层编码为逐帧 embedding,与扩散时间步 embedding 相加,只用于该帧对应的 spatial 注意力块的 AdaLN;temporal 块仍共享与 video-level 相同的单一条件 embedding。这样每个生成帧都能精确对齐它对应的那个动作。
- 输出层:线性层输出噪声预测
ε̂ = ε_θ(x_t, t, c),仅预测均值、不预测协方差(作者实测去掉协方差预测能提升视频质量,是相对标准 DiT 的一个显式改动)。 - 推理:PNDM 采样器,50 步。
- 配置(论文默认报告的 IRASim-XL):28 层、hidden size 1152、16 头、patch size 2、输入通道 4(对应 VAE 潜通道数)、dropout 0.1,679M 可训练参数。缩放实验的四档模型:IRASim-S(12 层/384 隐层/6 头/33M)、IRASim-B(12 层/768/12 头/132M)、IRASim-L(24 层/1024/16 头/461M)、IRASim-XL(28 层/1152/16 头/679M)。
- 可生成分辨率最高 288×512,自回归 rollout 可生成 150+ 帧的长时程视频(上一段生成的最后一帧作为下一段的历史/初始帧)。
数据
IRASim Benchmark 建立在真实机械臂公开数据集之上(均为真实机器人采集,非仿真):
| 数据集 | 训练 episode/clip | 验证 episode/clip | 测试 episode/clip | 原始分辨率 | 训练分辨率 | 动作空间 |
|---|---|---|---|---|---|---|
| RT-1 | 82,069 / 2,314,893 | 2,167 / 4,810 | 2,167 / 4,799 | 256×320 | 256×320 | 7-DoF:6-DoF SE(3) 位姿 + 1-DoF 连续夹爪 |
| Bridge | 25,460 / 482,701 | 1,737 / 2,905 | 1,738 / 2,946 | 480×640 | 256×320(resize) | 同 RT-1,7-DoF |
| Language-Table | 170,256 / 1,483,133 | 4,446 / 5,119 | 4,562 / 5,243 | 360×640 | 288×512(resize) | 2-DoF:2D 平面平移 Δx,Δy |
| RoboNet(仅 v2) | 162,161 / 2,540,500 | — | 256 / 407 | — | 256×256 | 最多 5-DoF,7 种机械臂统一为 5 维 [Δx,Δy,Δz,Δγ,g](缺失维度补零) |
- 训练片段构造:连续滑窗切出 4 秒/16 帧的片段;验证/测试为降低评测成本,按 16 帧间隔采样。
- 所有数据集的臂动作统一转换为相对 delta 动作(非绝对位姿)。
- 长轨迹评测用完整 episode 做自回归 rollout,平均长度 RT-1 42.5 帧、Bridge 33.4 帧、Language-Table 23.7 帧。
- v2 新增的下游实验另用了三份数据:LIBERO 仿真基准(策略评估实验,专家演示 + 用训练后策略额外收集的”post-trained rollouts”,含成功与失败两类,一起训练 IRASim);Push-T 仿真基准(200 条专家演示 + 0/100/200/500/1000 条不等的 post-trained rollouts,做规模消融);“自建真实机器人数据集”(论文原话:规模与 RT-1 相近、动作空间相同)用于三项真实世界任务(关抽屉、把橘子放到绿色/红色盘子上)的规划实验。这三组实验因数据有限,都用预训练的 OpenSora 权重初始化 IRASim 再微调,而非从零训练。
- 2025-05 起数据集在 Hugging Face 发布镜像
huggingface.co/datasets/fangqi/IRASim(需下载分片后用仓库自带merge.sh合并再解压)。
训练方法
- 目标函数:标准 DDPM 噪声预测 L2 loss(
L_simple = ||ε_θ(x_t,t)-ε_t||²),仅在历史/初始帧之后的帧上计算;历史帧本身不加噪、不计损,靠注意力对齐一致性。 - 优化超参(IRASim Benchmark 主实验,从零训练):AdamW(β=0.9, 0.999)、恒定学习率 1e-4、batch size 64、梯度裁剪 0.1、weight decay 0、EMA 衰减 0.9999;主文明确写”训练 300k 步”(附录超参表里对应字段写的是”3000000”,与正文”300k”不一致,疑似论文表格排版错误,取正文口径)。所有主实验模型均从零训练。
- 作者发现 IRASim 训练非常稳定,即便不做梯度裁剪也未观察到 loss spike;而两个 U-Net 基线(VDM、LVDM)不加梯度裁剪时经常出现 loss spike。
- 消融:去掉协方差预测(只预测均值)经验上提升视频质量;历史帧数从 1 帧增到 2 帧(Bridge 上测试)PSNR 几乎不变(都在 25 左右),作者推测轨迹本身已隐含速度信息,多给历史帧收益有限。
- v2 下游实验的训练方式(策略评估 / 基于模型规划):数据有限,故用预训练 OpenSora 权重初始化再用 frame-level 条件方法微调;基于模型规划用简单的”排序(ranking)“算法——从策略采样 K 条候选轨迹,逐条用 IRASim rollout,用价值/代价函数给每条轨迹的终帧打分,执行分数最高(Push-T:训练一个 ResNet50 回归器预测 IoU 作价值函数;真实机器人:用 MSE 或 ResNet50 特征余弦相似度衡量终帧与目标图像的相似度作代价函数,取代价最低的前 5 条在真实世界执行取平均成功率)的一条。
- 论文明确把扩散蒸馏(加速推理)与用 IRASim 做动力学模型的强化学习策略提升列为未做的未来工作。
Infra(训练 / 推理工程)
- 训练算力(仅 v1/v2 共有的 IRASim Benchmark 三数据集主实验有披露;RoboNet 及 v2 新增的 OpenSora 初始化微调实验的算力未披露):
| 数据集 | 并发 GPU 数 | GPU-小时 | GPU 型号 |
|---|---|---|---|
| RT-1 | 32 | 2,381 | A800 (40GB) |
| Bridge | 32 | 2,371 | A800 (40GB) |
| Language-Table | 32 | 2,369 | A100 (80GB) |
- 并行策略、混合精度等细节:未披露。
- 推理:PNDM 50 步采样,在单张 A100 上生成一段 16 帧(约 4 秒)视频仅需约 30 秒、显存占用仅 8GB——高吞吐、省显存,但作者明确承认非实时(“is not real-time”)。
- 基线对比:像素空间扩散 VDM(40M 参数)比潜空间扩散 LVDM/IRASim 消耗更多计算,尽管参数量更小。
评测 benchmark
短轨迹视频生成质量(Table 1,RT-1/Bridge/Language-Table,与 VDM、LVDM、IRASim-Video-Ada 对比):
| 数据集 | 方法 | PSNR↑ | SSIM↑ | Latent L2↓ | FID↓ | FVD↓ |
|---|---|---|---|---|---|---|
| RT-1 | VDM | 13.762 | 0.554 | 0.4983 | 41.23 | 371.13 |
| LVDM | 25.041 | 0.815 | 0.2244 | 4.26 | 30.72 | |
| Video-Ada | 25.446 | 0.823 | 0.2191 | 4.34 | 29.27 | |
| Frame-Ada(本文) | 26.048 | 0.833 | 0.2099 | 5.60 | 25.58 | |
| Bridge | VDM | 18.520 | 0.741 | 0.3709 | 39.82 | 127.25 |
| LVDM | 23.546 | 0.810 | 0.2155 | 10.59 | 35.06 | |
| Video-Ada | 24.733 | 0.827 | 0.2021 | 10.30 | 23.03 | |
| Frame-Ada(本文) | 25.275 | 0.833 | 0.1947 | 10.51 | 20.91 | |
| Language-Table | VDM | 23.067 | 0.857 | 0.3204 | 64.63 | 136.56 |
| LVDM | 28.254 | 0.889 | 0.1704 | 6.85 | 24.34 | |
| Video-Ada | 23.893 | 0.859 | 0.2028 | 7.05 | 73.84 | |
| Frame-Ada(本文) | 28.818 | 0.888 | 0.1660 | 6.38 | 48.49 |
作者明确以 Latent L2 与 PSNR 为主评价指标(因轨迹到视频任务的输出变化空间远小于文生视频,FID/FVD 不够贴合人类偏好)。
RoboNet 数据集(Table 2,v2 新增,与非扩散基线对比):MaskViT PSNR 20.4/SSIM 67.1;iVideoGPT PSNR 23.8/SSIM 80.8;IRASim PSNR 24.6/SSIM 81.1(MaskViT/iVideoGPT 数字取自 iVideoGPT 论文)。
长轨迹视频生成(Table 3,Latent L2 / PSNR):
| 数据集 | LVDM | Video-Ada | Frame-Ada |
|---|---|---|---|
| RT-1 | 0.2567 / 23.573 | 0.2519 / 23.984 | 0.2408 / 24.615 |
| Bridge | 0.2534 / 21.792 | 0.2385 / 22.868 | 0.2306 / 23.260 |
| Language-Table | 0.1776 / 26.215 | 0.2112 / 22.551 | 0.1730 / 26.773 |
人类偏好评估:5 名参与者、每人对每个数据集 10 个测试视频、共 90 对两两比较,Frame-Ada 在三个数据集上都击败所有对比方法,结果与 Latent L2 排序一致。
规模效应:33M→679M 参数、增加训练步数,在三数据集上视频质量均单调提升(论文以图表形式给出趋势,未在正文列出具体分数表)。
策略评估(Table 4,LIBERO,与 ground-truth Mujoco 仿真器对照):4 个不同训练步数的扩散策略分别在两个”评估器”里各跑 50 次,成功率:
| 评估器 | 1 | 2 | 3 | 4 |
|---|---|---|---|---|
| Ground-Truth 仿真器 | 0.18 | 0.50 | 0.80 | 1.00 |
| IRASim | 0.28 | 0.48 | 0.74 | 0.96 |
两组评估结果 Pearson 相关系数 0.99。
Push-T 基于模型规划(Table 5):K 为采样轨迹数,P 为训练 IRASim 用的 post-trained rollout 数量,报告 100 次试验平均 IoU。K=1(即无规划的原始策略)恒为 0.637。P=1000、K=50 时 IRASim 达到 0.961;对比最新 SOTA GPC(generative predictive control)的 GPC-RANK(K=5 时 0.642、K=50 时 0.698)与 GPC-RANK+OPT(K=5 时 0.642、K=10 时 0.824、K=50(M 步梯度优化)时 0.882)——P≥200 时 IRASim 全面超过两个 GPC 变体。
真实机器人基于模型规划(Table 6):3 项任务(关抽屉 / 橘子放绿盘 / 橘子放红盘)成功率——随机基线 0.20/0.07/0.13;IRASim(ResNet50 代价) 0.60/0.73/0.60;IRASim(MSE 代价) 0.87/0.80/0.87。MSE 代价函数整体优于 ResNet50 特征相似度,但并非完美(偶尔选错预测视频导致任务失败)。
鲁棒性定性实验:输入一条”物理上不可能”的轨迹(命令机械臂穿过桌面向下移动),IRASim 生成的视频里机械臂仍停留在桌面上而非穿透,作者以此说明模型对物理规律有一定隐式理解。
创新点与影响
- 核心贡献:提出帧级动作条件(Frame-Ada)AdaLN 机制,把”整段轨迹一个 embedding”细化为”每个动作单独编码、逐帧调制 spatial attention 的 AdaLN 参数”,在所有评测数据集和短/长轨迹设置下一致优于 video-level 条件与两个 U-Net 扩散基线(VDM/LVDM)。
- 改变了什么:把机器人操作世界模型的输入从”语言”(UniSim)或”学出来的隐动作”(Genie)换成了真实、连续、7-DoF 的机器人物理动作,并证明了扩散 Transformer 路线在这一任务上优于同期的 U-Net 扩散与(RoboNet 上对比的)自回归 token 模型(iVideoGPT/MaskViT)。v2 改版进一步把工作从”视频生成质量”论文扩展为”世界模型当模拟器/规划器”论文:策略评估与真实仿真器高度相关(Pearson 0.99),且作为基于模型规划的动力学模型能显著提升策略表现,并展示了性能随采样轨迹数 K 与训练数据量 P 同步扩大而持续提升的”测试时扩展(test-time scaling)“趋势。
- 作者自陈局限:(1)和其他生成模型一样存在幻觉;(2)推理非实时,尽管吞吐高、显存占用低(单张 A100 上 30 秒生成 16 帧、仅需 8GB);(3)v1 版本不支持灵活的输入分辨率与动作空间,限制了对不同分辨率/动作维度机器人数据的利用;(4)v2 讨论里指出真实机器人规划实验中代价函数的准确性会反过来限制评测结果的可靠性,且规划器只是一个简单的采样式规划器,可与更先进的动作分块(action-chunking)策略结合以扩展任务范围。
- 未来工作方向:扩散蒸馏加速推理;把 IRASim 当动力学模型、结合强化学习进一步提升机器人策略(论文本身未做)。
原始链接
- 论文(arXiv,v2 为 ICCV 2025 收录版《IRASim: A Fine-Grained World Model for Robot Manipulation》,v1 为 2024-06 原始版《IRASim: Learning Interactive Real-Robot Action Simulators》):https://arxiv.org/abs/2406.14540
- PDF:https://arxiv.org/pdf/2406.14540
- 项目主页:https://gen-irasim.github.io/
- 代码(GitHub):https://github.com/bytedance/IRASim
- 数据集镜像(Hugging Face):https://huggingface.co/datasets/fangqi/IRASim
一手源存档(sources/)
- irasim—github-readme — GitHub README(安装、数据下载表、训练/评测脚本、引用信息)
- irasim—project-page — 项目主页正文(abstract、四组实验的定性展示说明)
- arXiv 2406.14540 原文(v1 + v2 全文通读,PDF 不入 git):https://arxiv.org/abs/2406.14540