Sparse Planning in Visual World Models via Cost Gradients¶

会议: NeurIPS 2026
arXiv: 2610.10274
论文: https://arxiv.org/abs/2610.10274
代码: https://ycxuyingchen.github.io/costgrad/
领域: 其他
关键词: 视觉世界模型, 稀疏规划, 成本梯度, 动作条件调制, 交叉熵方法 (CEM)

一句话总结¶

提出无需重训且目标条件化的 Token 选择器 CostGrad，通过规划成本对输入 Token 的梯度范数筛选关键空间位置，在 50% 稀疏度下使视觉世界模型的 CEM 规划速度提升 2.6 至 5 倍，且在多个连续控制基准上持平或超越全量 Token 性能。

研究背景与动机¶

基于 Token 的视觉世界模型（如 DINO-WM、JEPA-WMs）通过预训练视觉编码器将高维观测映射为二维空间特征网格，并在隐空间中预测未来动态，为复杂连续控制任务提供了强大的表征能力。然而，在测试时执行动作规划（如交叉熵方法 CEM）面临着极其高昂的计算代价：规划器需要在较长的时间视界内对数百条候选动作序列进行多轮迭代采样与自回归推演，在 $16 \times 16$ 的 Token 网格与标准搜索宽度下，单步环境规划就需要执行数千万次 Token 级的前向计算。实际上，在绝大多数控制场景中，智能体、受操纵物体与目标仅占据画面的一小部分，大面积的背景 Token 消耗了海量计算资源却对下游决策贡献微弱。

直觉上的加速思路是复用世界模型自身的预测误差或注意力权重来裁剪 Token。但“预测不等于规划”：模型基于动态预测损失挑选出的区域往往对应于视觉不可预测性，而非控制相关性。例如在推块任务中，受控推杆和物体在动作作用下运动确定且易于预测，预测误差极低，但在规划中却是决定成败的核心要素；反之，背景中的细微扰动或光影变化可能带来虚高的预测误差。此外，自注意力权重受训练损失的欠定性影响，在同性能的不同训练种子间极不稳定，无法为动作搜索提供可靠且一致的敏感区域指引。

本文的切入角度是直接从下游控制任务的规划目标出发，用规划成本本身指导 Token 的选择。核心 idea：利用单次零动作探测的反向传播计算规划成本对各个输入 Token 的梯度范数，直接筛选出对当前目标最敏感的关键空间位置，在固定且无需重训的世界模型上实现计算量骤减且性能无损的稀疏隐空间规划。

方法详解¶

整体框架¶

CostGrad 是一种完全免重训（training-free）且目标条件化（goal-conditioned）的测试时空间 Token 稀疏选择机制。在每个环境决策步开始时，系统输入当前观测帧编码与目标图像编码。首先，通过一次全量 Token 的前向预测与反向传播，计算规划成本对最新观测帧各空间 Token 的梯度范数；随后，提取梯度范数最大的前 $K$ 个关键 Token 索引集合，并在随后的整个 CEM 动作搜索迭代过程中仅保留该子集进行多步自回归展开与代价评估。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入观测与目标特征<br/>z, z_g ∈ R^(256×384)"] --> B["目标条件规划成本梯度选择<br/>零动作探测与反向传播获取 Top-K 敏感度"]
    B --> C["稀疏隐空间 CEM 动作搜索<br/>仅在子集 S 上展开 H 步预测与评估"]
    C --> D["动作注入通路与局部调制<br/>AdaLN 保持无上下文漂移的动作响应"]
    D --> E["最优动作输出与单步执行<br/>执行首步动作并在下个观测动态重选"]

关键设计¶

1. 目标条件规划成本梯度选择：以控制目标取代动态预测误差作为 Token 筛选源 现有 Token 裁剪方法依赖于注意力矩阵分布或自监督预测残差，常常错误排除掉易于预测但决定控制成败的受控物体。CostGrad 直接将控制目标引入重要性评估：给定当前观测上下文编码 $\bm{z} \in \mathbb{R}^{P \times D}$（$P=256, D=384$）与目标编码 $\bm{z}_g$，设定默认的探测动作 $\bm{a}_0 = \bm{0}$（归一化动作空间均值），运行单步前向预测得到 $\hat{\bm{z}} = \text{predictor}(\bm{z}, \bm{a}_0)$。定义视觉规划损失为单步预测特征与目标特征的欧氏平方距离： $$ \mathcal{L}_{\text{plan}}(\bm{a}_0; \bm{z}, \bm{z}_g) = |\hat{\bm{z}} - \bm{z}_g|2^2 $$ 通过对模型计算图进行一次反向传播，提取各输入空间位置 $j$ 的梯度范数作为重要性评分： $$ s_j = |\nablaj} \mathcal{L}|_2, \quad j=1,\dots,P $$ 梯度的 }$L_2$ 范数刻画了一阶局部敏感度：即便某个 Token 自身特征预测已很接近目标，只要它在自注意力机制中为预测其他目标关键区域提供了不可或缺的上下文，其梯度范数依然显著。依据 $s_j$ 降序选取 Top-$K$ 个空间索引构成保留集合 $\mathcal{S}$。

2. 稀疏隐空间 CEM 动作搜索：固定空间子集自回归推演与动态重选机制 在筛选出空间位置子集 $\mathcal{S}$（例如保持比例 $K/P = 50\%$）后，该几何索引在当前决策步的整个 CEM 优化循环中保持固定，并同时约束上下文帧与目标图像。CEM 采样 300 条候选动作序列，在时间跨度 $H=6$ 内自回归调用世界模型，此时世界模型仅处理长度为 $K$ 的稀疏序列。由于 Transformer 注意力计算开销关于序列长度呈二次方增长（减半即节省 4 倍 FLOPs），MLP 呈线性增长（减半即节省 2 倍 FLOPs），计算密度大幅压缩。在视界终点，规划代价仅在保留的 $K$ 个空间位置上计算与目标特征的平均均方误差： $$ J_{\mathcal{S}}(A) = \frac{1}{KD} \sum_{j \in \mathcal{S}} |\hat{\bm{z}}{H,j}(A) - \bm{z}_g|_2^2 $$ 其中本体感觉项权重设为 }|_2^2 + \frac{\alpha}{d_p} |\hat{\bm{p}}_H(A) - \bm{p$\alpha = 0.1$。CEM 迭代优化动作分布并输出最终精英动作的首步执行。当环境迈入下一时间步产生新观测后，再次调用单步反向传播重新确定新的 $\mathcal{S}$，实现对动态场景的自适应跟踪。

3. 动作注入通路与漂移抑制：揭示调制架构决定稀疏规划成败的深层机制 稀疏规划的有效性不仅取决于挑选出的 Token 是否关乎控制，更取决于模型架构在 Token 被剥离后能否维持稳定的动作响应。研究发现，在以 AdaLN-Zero（自适应层归一化）作为动作注入机制的模型中，动作特征通过全连接层映射为逐层的缩放、平移与门控参数，直接作用于每个 Token 内部，不依赖于空间上下文的完整性。而在常见的输入通道拼接（Concat）架构中，动作特征作为额外通道拼接入输入后完全依赖自注意力网络传递。当剪切 Token 时，Concat 架构在保留位置上的动作效应分布发生了严重的空间重配比（动作通路漂移）。作者设计了动作效应相对散度指标： $$ p_i = \frac{|\bm{\delta}_i^{\mathcal{S}}|2^2}{\sum|}} |\bm{\delta}_j^{\mathcal{S}}|_2^2}, \quad q_i = \frac{|\bm{\delta}_i^{\mathcal{P}2^2}{\sum}} |\bm{\deltaj^{\mathcal{P}}|2^2}, \quad D(p \parallel q) = \sum} $$ 其中 }} p_i \log \frac{p_i}{q_i$\bm{\delta}_i^{\mathcal{U}} = f_i(\mathcal{U}, \bm{a}) - f_i(\mathcal{U}, \bm{0})$。在 AdaLN-Zero 上，CostGrad 的漂移量仅为随机选择的 $0.29\times$，有效保护了 CEM 候选动作的判别度；而在 Concat 上漂移量高达随机选择的 $1.33\times$，导致纯梯度选择失效。这一设计揭示出世界模型架构与稀疏选择器之间深刻的兼容性边界。

损失函数 / 训练策略¶

CostGrad 算法本身完全免训练（training-free），可即插即用部署于预训练的冻结视觉世界模型上。所依托的世界模型基于 JEPA-WMs 架构，采用冻结的 DINOv2 ViT-S/14 提取视觉表征（$16 \times 16 \times 384$），预测器为 6 层、16 个注意头的 Transformer（参数量约 17.6M），采用 RoPE 位置编码，在轨迹离线数据上以 AdamW 优化器（学习率 $5 \times 10^{-4}$，批大小 8）训练 20 个 epoch，最小化多步特征均方误差。测试时，CostGrad 仅在当前观测步引入一次全网格的零动作单步前向与反向传播（在 H100 GPU 上耗时仅 27.53 ms），相对随后的 30 轮 CEM 循环开销极低。

实验关键数据¶

主实验¶

在 DINO-WM / JEPA-WMs 的 4 个连续控制基准（PointMaze 迷宫导航、Wall 绕墙导航、PushT 机器人推块、MetaWorld Reach 机械臂末端控制）上评测规划成功率（%）。基线包括全量 Token（Full, 256 tokens）、基于最后一层注意力的选择（PredAttn）、基于自监督历史预测误差梯度的选择（PredGrad）、Token 合并算法（ToMe）以及均匀随机保留（Random）。所有稀疏方法保留比例均为 $50\%$（128 tokens）。

方法 (50% Tokens)	PointMaze (MZ)	Wall	PushT (PT)	MetaWorld (MW)	平均成功率
Full (100% Tokens)	88.2 ± 8.7	84.7 ± 2.6	62.4 ± 4.4	58.7 ± 1.8	73.5
CostGrad (Ours)	86.2 ± 6.1	92.0 ± 1.8	65.2 ± 6.5	65.6 ± 4.7	77.3
PredAttn	86.7 ± 9.7	67.3 ± 1.7	62.0 ± 8.0	44.8 ± 6.1	65.2
PredGrad	76.8 ± 9.2	67.2 ± 3.3	41.3 ± 5.9	51.4 ± 3.5	59.2
ToMe	74.4 ± 9.1	84.4 ± 2.6	51.4 ± 5.9	39.1 ± 4.1	62.3
Random	77.0 ± 9.4	67.0 ± 3.6	41.3 ± 5.9	46.9 ± 3.1	58.1

在 4 个基准中，CostGrad 在 3 个环境中超过了全量 Token 的规划成功率（Wall 提高 7.3 个百分点，PushT 提高 2.8 个百分点，MetaWorld 提高 6.9 个百分点），平均胜出 Full 达 3.8 个百分点，同时由于 Token 序列减半带来单步规划端到端 2.60 倍的真实墙钟加速（从 88.61s 降至 34.04s）。

消融实验¶

1. 动作注入机制与选择器兼容性对比（Wall 环境） 对比 AdaLN-Zero 与完全对齐的 Concat 预测器在 Full、CostGrad、Random 下的表现以及动作效应漂移度（KL 散度）：

预测器架构	Full 成功率	CostGrad 成功率	Random 成功率	CostGrad 相对 Random 增益	动作通路相对漂移比 (CG / Rand)
AdaLN-Zero	84.7 ± 2.6	92.0 ± 1.8	67.0 ± 3.6	+25.0 pp	(0.29 ± 0.03)×
Concat (Matched)	85.1 ± 5.2	80.0 ± 2.9	80.2 ± 3.8	-0.1 pp	(1.33 ± 0.17)×

在 Full 模式下两者性能相仿（84.7% vs 85.1%），但在稀疏模式下，AdaLN 配合 CostGrad 产生巨大的增益（+25.0 pp），而 Concat 架构下 CostGrad 甚至略微落后于随机选择，动作通路漂移比从 0.29 飙升至 1.33。

2. 算力匹配与极限加速（Wall 环境） 将 Token 稀疏（50%）与削减 CEM 迭代轮次（从 30 轮减半至 15 轮）相结合：

配置	规划耗时 / 步 (s)	规划成功率 (%)	相对 Full 加速比	说明
Full (256 tokens, 30 iters)	88.61	84.7 ± 2.6	1.0×	原始基准
Full (256 tokens, 15 iters)	44.30	84.4 ± 1.0	2.0×	仅降低搜索迭代
CostGrad (128 tokens, 30 iters)	34.04	92.0 ± 1.8	2.6×	仅 Token 稀疏
CostGrad (128 tokens, 15 iters)	17.02	92.0 ± 5.1	5.2×	双重加速，性能仍胜 Full (+7.3 pp)

关键发现¶

稀疏规划表现反超全量 Token 的本质：剔除无关背景和低敏感度 Token 后，CEM 目标代价函数中的无效噪声被过滤，目标函数的优化曲面更加陡峭清晰。在 Wall 任务中换用更容易受背景噪声干扰的 $L_1$ 规划代价时，CostGrad 相比 Full 的增益由 $L_2$ 下的 +7.3 pp 激增至 +22.9 pp（70.8% vs 47.9%）。
预测容易度与决策关键性的解耦：象限分析显示，PushT 中高达 31.9 个 Token（理论独立期望为 16 个）属于“容易预测但极其关键”区域（推杆与 T 块），这类区域预测误差极低但规划梯度范数极高。传统的自监督预测误差法会严重漏选此类目标，导致 PredGrad 在 PushT 上仅得 41.3% 的成功率。
选择信号稳定性显著优于注意力：不同随机种子训练的预测器在验证集损失相当，但注意力 Top-128 重合度（IoU）仅 0.44 ± 0.02，规划波动高达 50 pp；而 CostGrad 在相同模型间的 Token 选择 IoU > 0.82，规划成功率在各评估种子上保持高度一致。

亮点与洞察¶

从下游任务拉取关注度：巧妙地反向利用规划代价梯度作为无监督视觉世界模型的目标导向注意力机制，以仅仅一次前向+反向传播的开销，摆脱了以往模型重训或启发式注意力剪枝的困境。
揭示了调制架构（AdaLN）对 Token 裁剪的鲁棒性：明确指出在世界模型中，动作注入路径是影响稀疏推理成败的关键因素，逐层独立的自适应特征调制相较于自注意力通道拼接能显著抵御空间截断引起的语义漂移。
正交维度的双重加速空间：Token 空间维度压缩（减少前向计算复杂度）与 CEM 时间搜索维度压缩（减少采样轮数）高度互补，两者叠加在取得 5.2 倍实测加速的同时，规划成功率依然提升 7.3 个百分点。

局限与展望¶

单步探测动作的视界局限：当前 CostGrad 仅在 $t=0$ 执行单步零动作前向探测，难以捕捉需要多步动作级联交互后才显现关键性的深层依赖；扩展至长视界前向探索有望进一步提高复杂长程任务的精度，但需权衡计算开销。
仿真与物理真实场景的差距：目前的验证集中于仿真连续控制环境（PointMaze, Wall, PushT, MetaWorld），尚未在伴随传感器噪声、动态光照变化及物理标定误差的真实机器人平台展开验证。
对于 Concat 架构的修复仍未彻底解决：针对 Concat 类模型的动作通路漂移，采用随机锚点混合虽然缓解了分布漂移，但收益高度依赖具体检查点权重，未来仍需探索显式的稀疏正则训练以提升泛化适应性。

评分¶

新颖性: ⭐⭐⭐⭐☆ 提出利用规划目标反向梯度作为世界模型测试时空间 Token 动态选择器，并从动作通路漂移视角揭示了网络架构兼容性机理。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 大基准、多基线对比、消融实验、详实墙钟时间分解以及深入的动作漂移诊断。
写作质量: ⭐⭐⭐⭐⭐ 逻辑缜密，剖析清晰，概念与实证数据严谨自洽。
价值: ⭐⭐⭐⭐⭐ 为基于视觉世界模型的具身智能与连续控制高效在线规划提供了极具实用价值的落地范式。