Feedback World Model Enables Precise Guidance of Diffusion Policy

Paper: arXiv:2605.15705v1 Project: Feedback World Model project page

1. Motivation (研究动机)

这篇论文解决的是一个很具体的部署问题：robot diffusion policy 可以用 world model 在 test time 评估候选动作的后果，但这个 world model 在 OOD 初始状态或长 horizon 执行中会变成静态 open-loop predictor；一旦第一步预测偏离真实轨迹，后续 guidance 会继续基于错误 latent，反而把 policy 往错误方向推。现有做法通常有两类：一类在训练阶段把 world model 的动态、偏好或 reward 当作额外 supervision；另一类在推理阶段用 world model 对候选 action sequence 做 planning/guidance。

这两类方法共同假设 world model 的预测足够可靠，但机器人部署时非常容易遇到训练分布外的 robot pose、object configuration 或 visual observation。论文强调，单纯扩大数据、换更大的 video/world model、或收集额外 rollout 都成本高，而且在 data-limited manipulation 场景里不一定覆盖最终 OOD 测试分布。本文的核心观察是：机器人执行动作后天然会看到真实下一状态 $o_{t+1}$。

这个 observation 直接暴露了预测状态和真实状态之间的 mismatch，但现有 test-time world model guidance 很少把这个误差写回 world model 的内部预测状态。也就是说，deployment 本身产生的 feedback signal 没有被用来闭环校正预测。研究价值在于：如果 world model 能在 inference time 用真实 observation 做轻量校正，就可以在不增加训练数据、不更新参数、不依赖更大模型的情况下，让 diffusion policy 的 guidance 更可靠。

这对 real-world robot learning 很重要，因为真实机器人 rollout 昂贵，OOD 初始姿态又常见；一个只改推理环节的 feedback 机制更容易部署。

2. Idea (核心思想)

论文的核心思想是把 world model 从“每一步只根据当前 observation/action 做静态预测”的 open-loop model，改成带 feedback state 的 latent observer：系统维护一个辅助状态 ar z_t 表示 world model 根据历史 executed actions 认为当前应该处在的 latent；真实 observation 编码得到 $z_t$ 后，误差 e_t=z_t-ar z_t 会作为 corrective signal 加到下一步 latent velocity prediction 上。

这个设计与普通 world-model guidance 的关键区别在于：普通方法只用 ${\hat{z}}_{t+1}=f_{\theta }(z_t,A_t)$ 构造到 expert latent 的距离；本文用 online residual 修正后得到 $z_{t+1}^{\mathrm{fb}}$，再把它用于 diffusion score guidance。它不是重新训练 world model，而是在 deployment 的每个 environment step 后更新一个轻量 latent feedback state。第二个创新是 action-aware guidance。

普通 latent MSE 会把背景、光照等 action-irrelevant latent 维度也纳入距离，可能让 diffusion policy 优化不可控因素；本文用 counterfactual action variance 估计每个 latent dimension 对 action 的敏感度，用 $w_j^{(\beta )}$ 对 guidance energy 加权，让 guidance 更集中在机器人动作可控、和任务进度相关的 latent 维度上。

3. Method (方法)

3.1 Overall framework：Feedback-Guided Policy

Figure 1 解读：左侧环境执行 clean action $A_{t-1}^{(0)}$ 后返回 observation $o_t$，encoder $\psi$ 得到当前 latent $z_t$；中间的 guided diffusion policy 在低噪声 denoising step 对候选动作 $A_t^{(\tau )}$ 做 world-model guidance；下方橙色模块是 Feedback World Model，它维护 feedback state ${\bar{z}}_t$ 并用 residual 修正 next-state latent；绿色模块是 Action-Aware Energy，它从 expert set 中找到 nearest expert latent $z_{i^{\ast }}^E$，并用 controllability weights 强化 action-relevant 维度。执行 $A_t^{(0)}$ 后，新 observation 再更新 ${\bar{z}}_{t+1}$，形成外层闭环。

预备定义中，observation context $O_t$ 先被编码为 latent： $z_t=\psi (O_t),$ latent world model 预测 candidate action sequence $A_t$ 后的下一 latent： ${\hat{z}}_{t+1}=f_{\theta }(z_t,A_t).$ Diffusion policy 学习条件 action distribution 的 score： $s_{\varphi }(A_t^{(\tau )},O_t,\ell ,\tau )\approx {\nabla }_{A_t^{(\tau )}}\log p_{\tau }(A_t^{(\tau )}\mid O_t,\ell ),$ 并可在推理时加入 guidance term： $\tilde{s}(A_t^{(\tau )},O_t,\ell ,\tau )=s_{\varphi }(A_t^{(\tau )},O_t,\ell ,\tau )+\lambda g(A_t^{(\tau )},O_t).$ 标准 world-model guidance 会把预测 latent 匹配到 expert latent memory ${\mathcal{Z}}^E=\{z_i^E{\}}_{i=1}^N$ 中最近的状态： $i^{\star }=\arg {\min }_i{\Vert {\hat{z}}_{t+1}(A_t^{(\tau )})-z_i^E\Vert }_2^2,$ 并定义 energy： $E(A_t^{(\tau )})={\Vert {\hat{z}}_{t+1}(A_t^{(\tau )})-z_{i^{\star }}^E\Vert }_2^2.$ 最终 guided score 写成： $\tilde{s}=s_{\varphi }-\lambda {\nabla }_{A_t^{(\tau )}}E(A_t^{(\tau )}).$ 本文指出这个标准形式有两个失败点：一是 ${\hat{z}}_{t+1}$ 在 OOD 状态下会 drift；二是 uniform latent matching 会把不可控 latent 维度也纳入 guidance。Feedback world model 和 action-aware energy 分别对应这两个问题。

3.2 Feedback World Model：用真实 observation 校正 latent dynamics

论文把 learned transition 改写成 latent velocity： $v_t(A_t^{(\tau )})\triangleq \frac{f_{\theta }(z_t,A_t^{(\tau )})-z_t}{\delta t}=\frac{{\hat{z}}_{t+1}-z_t}{\delta t}.$ 除了真实 observation latent $z_t$，方法维护一个 auxiliary feedback state ${\bar{z}}_t$，初始化为 $z_0$。${\bar{z}}_t$ 表示“按照 world model 和过去实际执行动作推进后，模型内部认为当前应该在哪里”。二者的差： $e_t=z_t-{\bar{z}}_t$ 就是累积预测误差。关键更新式为： $\{\begin{array}{rl}& {\hat{v}}_t(A_t^{(\tau )})=v_t(A_t^{(\tau )})+L{e}_t,\\ & z_{t+1}^{\mathrm{fb}}(A_t^{(\tau )})=z_t+\delta t\cdot {\hat{v}}_t(A_t^{(\tau )}),\end{array}$ 其中 $L$ 是 positive-definite feedback gain。直觉上，如果模型内部 belief ${\bar{z}}_t$ 落后于或偏离真实 $z_t$，误差 $e_t$ 会作为一个 velocity correction 直接修正所有 candidate action 的 predicted next latent。这样 denoising 过程看到的不是 open-loop prediction，而是经过真实执行反馈校准后的 $z_{t+1}^{\mathrm{fb}}$。

执行 clean action $A_t^{(0)}$ 后，feedback state 用同一个 corrected velocity 推进： ${\bar{z}}_{t+1}={\bar{z}}_t+\delta t\cdot {\hat{v}}_t(A_t^{(0)}).$ 环境返回 $o_{t+1}$ 后，编码得到 $z_{t+1}$，下一步 residual 为： $e_{t+1}=z_{t+1}-{\bar{z}}_{t+1}.$ 这个 alignment 很重要：$A_t^{(0)}$ 先推进内部 belief 到 ${\bar{z}}_{t+1}$，真实 observation 再形成下一轮 correction。这样 feedback signal 不需要更新 $f_{\theta }$ 的参数，只更新推理时的 latent state。

论文还给出 convergence guarantee。若 true latent velocity 与 learned latent velocity 的 residual bounded by $\gamma$，则： ${\lim }_{t\to \infty }\Vert z_t-{\bar{z}}_t\Vert \le \frac{\gamma }{{\lambda }_{\min }(L)}.$ scalar gain $L=lI$ 时界为 $\gamma /l$。这说明 world model 的误差大小由 learned dynamics 的 residual 决定，而 feedback gain 决定真实 observation 抑制累积误差的强度；只要 residual bounded，feedback state 不会无限漂移。

3.3 Action-aware guidance：只强调 action-controllable latent 维度

Uniform latent MSE 的问题是 representation 往往混合了 robot-object geometry、背景、光照等因素。机器人动作能控制前者，却很难控制后者；如果 guidance 强迫 action 去匹配不可控维度，会带来无效甚至干扰性的 gradient。本文用 learned world model 的 counterfactual action variance 估计第 $j$ 个 latent 维度的 controllability： $w_j=\frac{1}{M}{\sum }_{i=1}^M{\mathrm{Var}}_{a\sim {\mathcal{A}}_{\mathrm{demo}}}\left([f_{\theta }(z_i,a){]}_j\right),j\in \{1,\dots ,D\}.$ 其中 $\{z_i{\}}_{i=1}^M$ 是 demonstration trajectory 中均匀采样的 observations，variance 在 demo action pool ${\mathcal{A}}_{\mathrm{demo}}$ 上计算。论文实现中用 $M=200$ 个 state、每个 state 采样 32 个 action。

为了保持 energy 尺度，先 normalize 再做 soft interpolation： ${\bar{w}}_j=\frac{w_j}{\frac{1}{D}{\sum }_{k=1}^D{w}_k+ϵ},w_j^{(\beta )}=1+\beta ({\bar{w}}_j-1),$ 其中 $ϵ=10^{-8}$，$\beta \in [0,1]$ 控制 controllability modulation 强度。最终 guidance loss 为： $E_{\mathrm{ctrl}}(A_t^{(\tau )})={\sum }_{j=1}^D{w}_j^{(\beta )}(z_{t+1,j}^{\mathrm{fb}}(A_t^{(\tau )})-{\hat{z}}_{i^{\ast },j}^E{)}^2.$ 梯度 ${\nabla }_A{E}_{\mathrm{ctrl}}$ 替代普通 ${\nabla }_{A}E$ 加入 diffusion score，因此 action-sensitive dimensions 对 denoising guidance 贡献更大，action-invariant dimensions 被 down-weight。

3.4 Complete inference pipeline 与伪代码

论文完整 pipeline 分为 offline preprocessing 和 online deployment。Offline 阶段把 expert demonstrations 编码成 expert latent memory ${\mathcal{Z}}^E$，并计算 action-aware weights $\{w_j^{(\beta )}\}$；online 阶段在 diffusion denoising 的最后 ${\tau }_g$ 个 low-noise steps 才启用 world-model guidance，因为此时 action sequence 已经有较明确任务结构，world model prediction 更有意义。

def feedback_guided_policy_step(policy_score, encoder, world_model, expert_latents,
                                ctrl_weights, obs_t, zbar_t, action_noise,
                                feedback_gain_L, guide_lambda, tau_g, T):
    z_t = encoder(obs_t)
    e_t = z_t - zbar_t
    A_tau = action_noise
 
    for tau in range(T, 0, -1):
        score = policy_score(A_tau, obs_t, tau)
        if tau <= tau_g:
            z_next_open = world_model(z_t, A_tau)
            v_open = (z_next_open - z_t) / delta_t
            v_fb = v_open + feedback_gain_L @ e_t
            z_next_fb = z_t + delta_t * v_fb
 
            i_star = nearest_neighbor(z_next_fb, expert_latents)
            z_expert = expert_latents[i_star]
            e_ctrl = (ctrl_weights * (z_next_fb - z_expert).pow(2)).sum()
            score = score - guide_lambda * grad(e_ctrl, A_tau)
 
        A_tau = diffusion_denoise(A_tau, score, tau)
 
    A_clean = A_tau
    execute_first_action_chunk(A_clean)
    zbar_next = zbar_t + delta_t * v_fb_for_clean_action(world_model, z_t, A_clean, e_t, feedback_gain_L)
    return A_clean, zbar_next

代码搜索未找到开源实现。检索范围包括 arXiv 页面、项目页、GitHub repository search（paper title / “Feedback World Model” / “Diffusion Policy” / authors），仅发现项目页仓库 Lorenzo-0-0/Feedback_World_Model，其页面标注 “Code release coming soon”，仓库内容为 README.md 和 index.html，没有训练、推理或模型实现。因此本节伪代码依据论文 Algorithm 1 和公式整理，不提供 released-code mapping table，也不设置 github_ref。

4. Experimental Setup (实验设置)

4.1 Benchmarks 与 OOD 设置

Figure 2 解读：论文在两个模拟 benchmark 上评估。Robomimic 使用 Square、Tool-Hang、Transport 三个 manipulation tasks；LIBERO-Plus 使用 LIBERO-10 中四个 single-task settings，重点测试 Robot Initial States 扰动类别。所有模拟实验都保持任务目标和场景语义不变，只扰动 robot initial state，从而构造 robot-pose-induced OOD。

Robomimic 设置：Square、Tool-Hang、Transport 三个任务，diffusion policy 使用 official proficient-human demonstrations 的 20% 训练，输入包括 RGB observations 和 proprioception。Square / Tool-Hang 是 single-arm，Transport 是 dual-arm；visual observations 使用标准 dual-view setup，如 agent/shoulder view 和 wrist-mounted camera。评估时只 perturb robot initial joint configuration，保持 object arrangement、camera placement、task goal 和 scene semantics 不变。

LIBERO-Plus 设置：选择四个 LIBERO-10 single tasks，并在 Robot Initial States category 下测试。具体任务为：

Alias	Scene	Language instruction	# variants
Kitchen-S3	KITCHEN_SCENE3	turn on the stove and put the moka pot on it	44
Kitchen-S6	KITCHEN_SCENE6	put the yellow and white mug in the microwave and close it	42
Kitchen-S8	KITCHEN_SCENE8	put both moka pots on the stove	34
LR-S6	LIVING_ROOM_SCENE6	put the white mug on the plate and put the chocolate pudding to the right of the plate	42

LIBERO policies 使用 agent-view camera、eye-in-hand camera、proprioceptive state 和 task language embedding；每个 single task 都训练 dedicated diffusion policy 和对应 lightweight world model。

Figure A1 解读：Square OOD 设置只扰动 robot arm joints。左侧是 in-distribution demonstration start，后续列展示不同 OOD starts；上下两行分别对应 agentview camera 和 front-view camera，说明视觉域没有随机化，变化主要来自 robot initial pose。

Figure A2 解读：Transport 是 dual-arm task，joint perturbations 分别施加到两只手臂。上下两行对应 shouldercamera0 和 robot-0 wrist camera；这类初始姿态变化会让 policy/world model 遇到训练示范中很少出现的 robot configuration。

Figure A3 解读：Tool-Hang 的 OOD starts 也遵循同样 protocol，上下两行分别是 side view 和 front view。它展示了本文 OOD 的边界：不是改变物体语义或背景，而是测试 policy 对 robot initial-state shift 的鲁棒性。

4.2 Real-world tasks、模型与 baselines

Figure 3 解读：真实机器人实验使用 R1 Lite 的右臂和右夹爪。Peach-P&P 要把桃子从桌面抓起并放入目标箱；Drawer-Open 要定位柔性把手、接触并拉开抽屉。右侧柱状图显示每个 setting 做 20 个 rollout episodes：in-distribution 下 base 已较强，ours 从 0.90/0.70 提升到 0.95/0.75；OOD 下提升更明显，Peach-P&P 从 0.40 到 0.80，Drawer-Open 从 0.20 到 0.70。

真实任务数据规模：Peach-P&P 收集 50 条 expert trajectories；Drawer-Open 收集 65 条 expert trajectories。两者都从 scratch 训练 diffusion policy 和 world model；测试时把右臂 reset 到训练分布外的 abnormal initial poses，workspace layout 和 task goal 不变。

模型实现细节：world model 是 task-specific lightweight latent dynamics model，不是 large pretrained video prediction model 或 VLA-scale world model。它用 deployed diffusion policy 的 visual encoder 提取每个 RGB view 的 visual features，再融合 proprioception encoder 和 action encoder，最后用 Transformer predictor 预测 next-step latent state。

训练目标是一阶 latent prediction，不重建 pixels，也不生成 future videos。部署时 world model frozen；feedback 只维护和校正 online latent state estimate。Baselines 包括：Base diffusion policy（不使用 world model guidance）、Std.

Guid.（使用 frozen base world model guidance，无 feedback correction）、Rollout Guid.（使用 rollout-enhanced world model）、Mixed BC（expert subset + all collected rollouts）、Filtered BC（expert subset + successful rollouts）。

Rollout 数据只从 original in-distribution validation environments 收集，不来自最终 OOD evaluation setting，用来检验“额外 rollout/retraining 是否能替代 feedback”。论文没有公开完整训练脚本或 config，因此具体 GPU type/count、learning rate、batch size、training steps 无法从 released code 验证。论文正文/appendix 明确给出的训练与实现数字包括：Robomimic 使用 20% official proficient-human demonstrations；

Peach-P&P / Drawer-Open 分别 50 / 65 expert trajectories；action-aware weights 用 $M=200$ states、每个 state 32 actions；真实任务每个 setting 20 rollout episodes；diffusion policy 使用 UNet-based architecture、AdamW、cosine learning-rate schedule、EMA、DDPM denoising，但未给出所有 hyperparameter 数值。

5. Experimental Results (实验结果)

5.1 Feedback world model 降低 OOD prediction error

Figure 4 解读：灰色是 Base WM，蓝色是 Rollout-Enhanced WM，橙色是 Feedback WM，绿色折线为 reduction。Robomimic Square 上 Feedback WM 把 latent prediction MSE 从约 0.016 降到 0.0003，reduction 73.3%；Tool-Hang / Transport 的 reduction 为 36.0% / 35.8%。LIBERO-Plus 上 Kitchen-S3 / Kitchen-S6 / Kitchen-S8 的 reduction 为 14.5% / 16.1% / 21.7%。真实任务图按 $\times 0.1$ 显示，Peach-P&P 从 0.0013 到 0.0008（65.4%），Drawer-Open 从 0.0022 到约 0.0005（76.4%）。这说明 feedback correction 在无需额外训练数据或参数更新的情况下显著降低 OOD prediction error。

Figure 5 解读：作者把 Base WM、Feedback WM 和 ground-truth latent trajectory 投影到 3-D PCA view（PC1、PC2、timestep）。Base WM 随 episode 逐步偏离真实 manifold，而 Feedback WM 的轨迹更贴近 ground truth，并能在早期 transient 后恢复。这补充说明 Figure 4 的 MSE 改善不是单点统计，而是在线 rollout 过程中持续抑制 drift。

5.2 Action-aware guidance 的必要性

Figure 6 解读：左侧 sorted controllability weights $w_j$ 呈 long-tail，Robomimic 三个任务中 top 5% latent dimensions 贡献约 20%—25% 总 controllability weight；右侧 Transport 的 variance map 显示部分维度跨 observations 都保持 action-responsive。这支持把 $W$ offline frozen 并在 inference 中使用：latent representation 确实存在 action anisotropy，不同维度对 action 的响应强弱差异很大。

Ablation 结果如下：

Method	Square	Transport	Tool-Hang
Base	0.36	0.34	0.27
+ Feedback WM	0.40	0.54	0.32
+ Action-aware	0.46	0.60	0.36

这个 ablation 说明 prediction 更准只是第一步：如果 guidance energy 仍然平均处理所有 latent dimensions，policy 仍会被不可控/弱相关维度干扰。加入 action-aware weighting 后，Square 从 0.40 到 0.46，Transport 从 0.54 到 0.60，Tool-Hang 从 0.32 到 0.36。

5.3 OOD policy success rate

模拟 OOD manipulation tasks 的 success rate 表如下：

Benchmark	Task	Base	Mixed BC	Filtered BC	Std. Guid.	Rollout Guid.	Ours
Robomimic	Square	0.36	0.08	0.12	0.40	0.44	0.46
Robomimic	Tool-Hang	0.27	0.22	0.22	0.22	0.31	0.36
Robomimic	Transport	0.34	0.26	0.36	0.46	0.52	0.60
LIBERO-Plus	Kitchen-S3	0.55	0.57	0.70	0.61	0.50	0.68
LIBERO-Plus	Kitchen-S6	0.45	0.48	0.54	0.48	0.50	0.55
LIBERO-Plus	Kitchen-S8	0.59	0.42	0.35	0.62	0.62	0.65
LIBERO-Plus	LR-S6	0.26	0.24	0.24	0.21	0.24	0.36
Average	-	0.40	0.32	0.36	0.43	0.45	0.52

整体看，Ours 的平均 success rate 为 0.52，高于 Base 0.40、Std. Guid. 0.43、Rollout Guid. 0.45；相对 Base 是 30% improvement。Filtered BC 在 Kitchen-S3 上达到 0.70，略高于 Ours 0.68，但平均只有 0.36；Mixed BC 平均 0.32，说明简单加入 rollout 甚至可能引入失败轨迹、biased state distribution 或覆盖不足，不能稳定解决最终 OOD generalization。

真实任务的结果更直接：in-distribution 下 Peach-P&P 从 90% 到 95%，Drawer-Open 从 70% 到 75%；OOD 下 Peach-P&P 从 40% 到 80%，Drawer-Open 从 20% 到 70%。这表明 feedback correction + action-aware energy 对 robot initial-state shift 特别有效，因为它针对的是执行过程中 prediction drift 与 action-relevance 两个推理时瓶颈。

主要限制：论文没有公开 released implementation 或完整训练 config，因此无法核对实际代码中 feedback gain $L$、guidance strength $\lambda$、$\beta$、${\tau }_g$、diffusion steps、batch size、LR 等数值；实验 OOD 主要是 robot initial-state perturbation，尚未证明对 object layout、visual domain、language ambiguity 或更长 horizon real-world tasks 同样稳定。