ImageWAM — 图像编辑即世界动作模型

动机：视频生成 WAM 的三大缺陷

先厘清概念梯队：VLA(vision-language-action) 直接从"观测+指令"预测动作；WAM(world action model) 在中间插入一步"视觉想象"，先想象未来、再据此产生动作(reason-before-act)。主流 WAM 用视频生成来做这步想象。

论文指出，用视频生成当 WAM backbone 有三个结构性缺陷：

关键洞察：机器人操作 = 指令引导的视觉变换

作者把操作任务重新表述为：给定当前图 + 语言指令，产生一张"任务完成后的目标图"。这恰好就是图像编辑(text-guided image editing)做的事——按指令修改源图、保留无关内容。于是他们只建模端点帧：

ô_edit 是一张 source-conditioned 的单帧，概括了"当前观测在该任务指令下应发生的视觉变换"。

整体架构(图2)

❄ 冻结(VLM/多模态理解模块)　🔥 可训(diffusion 编辑分支 + 动作专家)

三个组件

① Image Editing Backbone(编辑分支)

基于现成图像编辑模型(OmniGen2 / Ovis-U1 / FLUX.2)。输入当前观测 o_t + 指令 l，本职是去噪生成"编辑后目标帧"。ImageWAM 复用它去噪过程中每层 transformer 的 KV cache 作为条件，而不只是拿它解码一张图。

② Action Expert(动作专家)

一个 flow-matching 的 DiT(见 ActionDiTOmnigen2)。把动作 chunk 编码成 token，通过 joint attention 读取编辑分支的 KV cache，预测动作速度场。

③ 冻结的 VLM 理解模块

编辑模型里负责编码指令/视觉上下文的 VLM 部分保持冻结，提供稳定的语言-视觉条件；只训练 diffusion 生成分支和动作专家。

核心机制：KV Cache 当作"世界-动作上下文"

这是全文最关键、也最容易误解的一步。请分三层理解。

第一层：什么是这里的 KV cache

编辑分支是一个多层 transformer。去噪时，视觉 latent 先与指令(经冻结 VLM)交互，然后逐层做 attention。每一层 ℓ 都会算出 Key/Value 张量。ImageWAM 把这些逐层 KV 收集起来：

τ 是随机采样的编辑去噪时间步。因为这份 cache 是在"视觉 latent 已经和指令充分交互之后"算出来的，所以它已经编码了"任务条件下的视觉变换信息"，而无需把最终目标图解码出来。

第二层：动作专家怎么用它(joint attention)

关键实现在 mot.py(Mixture-of-Transformers)。编辑专家和动作专家逐层并行前进，在每层的 attention 里，动作 token 的 Query 去 attend 拼接了编辑分支 K/V 的联合序列：

也就是说，动作 token 不是"看一张图"，而是"直接钻进编辑模型的中间推理里，读它每一层关于'该怎么变'的内部表征"。这就是论文说的 transfers the image editing model's internal reasoning process to robot control。

第三层：Fast-WAM 变体对照

为公平对比，作者还实现了 Fast-WAM 式的视频 WAM 变体：未来视频 token 只在训练时用于 video co-training，推理时移除，动作专家仅用"当前观测+指令"产生的 KV。这给出一个"和 ImageWAM 一样没有未来 token"的视频 WAM 基线。ImageWAM 在此基础上进一步把"视频"换成"图像编辑"。

训练：两个 flow-matching 目标联合优化

① 图像编辑目标 L_img

让编辑分支学会预测任务相关的端点帧。目标帧 o_t+H+1 经 VAE 编码成 latent z*；采噪声 ε_z、图像流时间 r，构造插值 latent，diffusion 分支预测速度场(flow matching)：

这个目标保住编辑分支"预测任务相关未来视觉状态"的能力，也逼它的 KV cache 编码有用的变换信息。

② 动作 flow matching 目标 L_act

动作 chunk a* 采噪声 ε_a、动作流时间 s，构造插值，动作专家conditioned 在编辑 cache C_edit^τ 上预测动作速度场：

注意两个不同的"时间"：s 是动作 flow-matching 时间，τ 是图像编辑去噪时间步。训练时随机采 τ，让动作专家见到去噪不同阶段的 cache。

推理：一步编辑 forward，不解码图

推理时既不做完整未来视频生成，也不解码完整编辑图。固定一个编辑去噪时间步 τ*，只跑一次编辑分支 forward 取 cache，然后动作专家据此去噪产生动作：

实验设置

关键差异：ImageWAM 不用额外具身预训练(P.T.=✗)，只在下游 benchmark 演示上训练，却超越许多需要预训练的 VLA/WAM。评测 benchmark：LIBERO、LIBERO-Plus、RoboTwin 2.0，外加真机 Dobot XTrainer 双臂 4 任务。

评测协议矩阵

RoboTwin 2.0(表1)

LIBERO(表2)与 LIBERO-Plus(表3)

LIBERO 四套件平均：ImageWAM 98.4(π0.5 96.9，与视频 WAM 相当)。LIBERO-Plus(更强分布偏移，含 Camera/Robot/Language/Light/Background/Noise/Layout 七类扰动)三变体平均成功率：

真机(表4)与效率(表5)

真机四任务(叠三碗/叠毛巾/开抽屉存马克笔/挂杯)ImageWAM 全部最优，T2(叠毛巾，可形变物体)增益最大(+9 vs FastWAM)。

消融要点(4.4)

• Q1 能换编辑模型吗？能。OmniGen2/Ovis-U1/FLUX.2 都超过 FastWAM，越强的编辑 backbone 策略越鲁棒 → 不绑定特定 edit 模型。
• Q2 为何不用统一理解-生成模型？理解要语义抽象、生成要空间细节，深层冲突；解耦(冻结理解+只训生成分支)优于 UniVLA/BagelVLA。
• Q3 编辑 backbone 更大更好吗？FLUX.2 4B→9B 平均 83.1→85.2，主要提升在 Robot/Language/Background/Layout；但 Camera/Light/Noise 非单调 → 规模有益但依赖 cache 与扰动类型的对齐。

训练/评测资源配置 🟢论文核对

来源：PDF §5.2 训练细节 + 表8(通用超参)/表9(数据集配置)/表10(训练成本)。

训练成本(表10)

† 有效 per-GPU batch，含 3 步梯度累积。LIBERO 训 10 epoch、RoboTwin 5 epoch、真机 10 epoch。

复现评测：FLUX.2 4B × LIBERO 全量

🟢实测复现 在 ICL（8× NVIDIA H20, 96GB/卡）上，用官方 release checkpoint ImageWAM-FLUX.2-4B-LIBERO 完成全量 LIBERO 4 套件评测。总成功率 97.9%，论文报告 98.4%，差异 0.5pp 在统计波动范围内，复现成功。

评测结果总览

套件级成功率（可视化对比论文）

逐任务热力图（40 任务 × 25 trials）

颜色深度 = 成功率（绿=100%，黄=92-96%，橙=84-88%）

libero_spatial（96.8%）

libero_object（98.4%）

libero_goal（98.0%）

libero_10（98.4%）

hover 查看具体数值。8 个非满分任务（共 21/1000 失败 trial），最低 84%（spatial task5: "pick up bowl on ramekin"——需精准定位堆叠物体）

复现 vs 论文对比表

评测配置

部署踩坑与经验

复现命令

# SSH 连接开发机
ssh -p 443 root@10.82.96.110

# 环境准备
cd /robot/bionic-control-data/czy/projects/ImageWAM
source .venv/bin/activate
unset ALL_PROXY HTTP_PROXY HTTPS_PROXY http_proxy https_proxy all_proxy
export MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa
export PYTHONPATH=$PWD/third_party/LIBERO

# 完整评测（默认参数 = 4套件 × 25 trials × 8 GPUs）
bash scripts/flux2/run_eval_flux2_libero.sh

# Smoke 快速验证（2 分钟）
NUM_GPUS=2 TASK_SUITE_NAMES='[libero_spatial]' NUM_TRIALS=2 \
  bash scripts/flux2/run_eval_flux2_libero.sh

真机训练规划：从采集数据到部署

🟢实践指南 本节基于代码仓库的训练流程分析 + 论文真机实验配置，给出从真机数据到 ImageWAM 部署的完整路径。

数据要求：标准 Episode 即可

是的，标准的 state-action episode + 相机图像 + 语言指令就是全部。不需要特殊标注、不需要人类视频、不需要仿真对。

每条 Episode 需要包含

两个组件分别用什么数据训练？

ImageWAM 内部有两个可训组件，它们联合训练（同一份数据，一个 loss 函数同时优化两者）：

单臂 vs 双臂配置

相机配置与图像处理

推荐配置

ImageWAM 对图像方向敏感——确保送入模型的图像是正立的（桌面在下）。训练和推理必须一致。

训练流程（完整步骤）

Phase 0: 数据准备

将真机采集的 episodes 转为 LeRobot v2/v3 格式（parquet + mp4）。每条 episode 必须有：

• 连续的 RGB 帧序列（视频）+ 对应的 action/state 数值序列
• 一条描述任务的语言指令
• 推荐 ≥50 episodes/task（论文真机每任务约 50-100 条）

Phase 1: 预计算（自动化）

# 1. 预计算 Qwen3 文本 embedding cache（只需跑一次，加速后续训练）
GPU_PER_NODE=8 TASK_TYPE=real FLUX2_VARIANT=4b PRECOMPUTE_QWEN3_CACHE=true \
  bash scripts/flux2/run_train_flux2_klein_imagewam.sh

# 2. ActionDiT 初始化权重（自动从 FLUX.2 骨干提取，无需手动）
#    脚本检测 ACTION_INIT 不存在时自动调用 preprocess_action_dit_flux2.py

Phase 2: 训练

# 环境变量
export DATA_ROOT="/path/to/your/real_data"
export FLUX2_SRC="./third_party/flux2"
export FLUX2_AE_MODEL_PATH="./checkpoints/flux2/FLUX.2-dev/ae.safetensors"

# 启动训练（8×H20，ZeRO-1，~18h）
GPU_PER_NODE=8 \
TASK_TYPE=real \
FLUX2_VARIANT=4b \
bash scripts/flux2/run_train_flux2_klein_imagewam.sh

Phase 3: 评测 & 部署

# 使用训练产出的 checkpoint
export CKPT_PATH="./runs/train/YYYYMMDD_HHMMSS/checkpoints/checkpoint_best.pt"
export DATASET_STATS_PATH="./runs/train/YYYYMMDD_HHMMSS/dataset_stats.json"

# 推理延迟: ~263ms/step（单卡 A6000），满足实时控制
# action_horizon=16, replan_steps=12（每 12 步 replan 一次）

训练超参数推荐

数据量 vs 预期效果

方案选择决策树

与其他方案对比：为什么选 ImageWAM

部署踩坑预警（基于我们的复现经验）

执行 Checklist

三种"生成模型服务于动作"的范式对比

"预测未来图像/视频到底有没有用？"是当前 WAM 领域最核心的争论。ImageWAM、DiT4DiT、DynaGuide 分别给出了三种不同但互补的答案。

核心主张对比

它们在说同一件事的不同方面

"预测图像没必要"——三者怎么看？

表面上 DiT4DiT 说"预测图像没必要"而 ImageWAM 还用了图像编辑模型，看似矛盾。但实际上：

效果与效率对比

适用场景差异

启发：生成模型在机器人中的三种用法

代码印证:机制在仓库哪里

🟢源码核对 仓库 ~/projects/ImageWAM/,推荐从 FLUX.2 变体入手。核心在 src/imagewam/models/backbones/。

① MoT:编辑专家 + 动作专家的 joint attention

mot.py(Mixture-of-Transformers)是 KV 融合的落点。它校验所有 expert 层数/头数一致,逐层并行推进,在 _mixed_attention 里把动作 Query 与拼接后的 K/V 一起做 attention:

def _mixed_attention(self, q_cat, k_cat, v_cat, attention_mask, ...):
    # q_cat: 动作 token 的 Query
    # k_cat / v_cat: [K_action ; K_edit] / [V_action ; V_edit]  ← 编辑分支 KV 拼进来
    key_len = k_cat.shape[1]
    H, H_kv, D = self.num_heads, self.num_kv_heads, self.attn_head_dim
    ...
    out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, enable_gqa=enable_gqa)
    # 动作 token 通过 joint attention 读取编辑分支的逐层 KV

这正是式(5)-(9)里 "动作专家 conditioned 在 C_edit" 的实现:不是把 cache 拼成一个 context 向量,而是在每层 attention 里把编辑分支的 K/V 拼到动作 token 的注意力键值上。

② 动作专家:flow-matching DiT

action_dit_omnigen2.py 的 ActionDiTOmnigen2:动作经 action_encoder(Linear) 变 token,加 RoPE 与 timestep 调制,过若干 OmniGen2TransformerBlock,head 输出动作维度的速度场。它声明 block_protocol="omnigen2" 以便和编辑专家在 MoT 中对齐拼接。

class ActionDiTOmnigen2(nn.Module):
    block_protocol = "omnigen2"
    def __init__(self, action_dim, ..., num_layers=26, ...):
        self.action_encoder = nn.Linear(action_dim, hidden_dim)
        self.blocks = nn.ModuleList([OmniGen2TransformerBlock(...) for _ in range(num_layers)])
        self.head = LuminaLayerNormContinuous(..., out_dim=action_dim)  # 输出动作速度场

③ 编辑专家封装

omnigen2_video_expert.py 的 OmniGen2VideoExpert 包住 OmniGen2Transformer2DModel,暴露 blocks/num_heads/num_kv_heads/RoPE,供 MoT 逐层取 KV。仓库还提供 action_dit_flux2.py/ovis_u1_video_expert.py 等对应三个变体。

评价与启发

它把什么问题问清楚了

"世界模型是否必须=预测未来视频"是当前具身领域的一个隐含假设。ImageWAM 通过把 WAM 的中间态从"视频"退化到"单张编辑目标帧",再退化到"根本不解码、只取 KV cache",一步步逼问:reason-before-act 真正需要的是"想象的像素",还是"想象过程中的中间表征"?答案偏向后者。

方法论上的三个亮点

• 把生成任务当表征提取器:不要生成结果(图/视频),只要生成过程中的 KV——这是"生成即表征"思路在机器人上的一次干净落地。
• 解耦理解与生成:冻结 VLM 理解、只训 diffusion 生成分支,规避统一模型的目标冲突(消融 Q2 有实证)。
• 效率是一等公民:1/6 FLOPs、1/4 延迟且不掉点,对真机实时控制(263ms)有直接意义。

值得追问/存疑处

• τ* 的选择:推理固定单个去噪时间步取 cache,τ* 怎么选、对不同任务是否稳健,论文给的是经验设定,值得关注其敏感性。
• "不解码目标图"的可解释性代价:视频/图像 WAM 至少能可视化"想象",ImageWAM 的中间态是 KV,调试与失败归因更难(图4的 attention 可视化是一种补偿)。
• 依赖强编辑基座:最好结果来自 FLUX.2,意味着方法的上限被通用图像编辑模型的进步牵引——是优点(搭生成社区顺风车)也是依赖。

自测:6 题检验理解