西交大提出Fast LeWorldModel：動(dòng)作前綴并行預(yù)測(cè)動(dòng)態(tài)估計(jì)加速4倍

本文第一作者為西安交通大學(xué)碩士生高云天，通訊作者為西安交通大學(xué)教授許翔宇，其研究方向涵蓋世界模型、三維視覺(jué)與具身智能（個(gè)人主頁(yè)：https://xuxy09.github.io/）

在視覺(jué)規(guī)劃與具身智能中，“世界模型” 被認(rèn)為是智能體走向通用決策能力的核心組件：在真正執(zhí)行動(dòng)作之前，先在潛在空間中 “想象未來(lái)”，再選擇最優(yōu)行為。

但在視覺(jué)規(guī)劃里，這個(gè) “想象” 過(guò)程往往很慢

以 LeWorldModel（LeWM）為例，它在規(guī)劃時(shí)有一個(gè)重要瓶頸：每評(píng)估一條候選動(dòng)作序列，模型都要一步步自回歸 rollout。也就是說(shuō)，LeWM 先預(yù)測(cè)下一步 latent，再把預(yù)測(cè)出的 latent 輸入 dynamics model，繼續(xù)預(yù)測(cè)下一步：

這種方式有兩個(gè)問(wèn)題：一是規(guī)劃慢，CEM 需要反復(fù)評(píng)估大量候選動(dòng)作序列；二是誤差會(huì)沿 imagined trajectory 累積，早期預(yù)測(cè)偏一點(diǎn)，后面可能越滾越偏。

針對(duì)這一瓶頸，西安交通大學(xué)研究團(tuán)隊(duì)提出了Fast LeWorldModel（Fast-LeWM），試圖從根本上改變世界模型的預(yù)測(cè)方式：從 step-by-step rollout 變成 trajectory-level parallel prediction。

• 論文標(biāo)題：Fast LeWorldModel
• 作者：Yuntian Gao, Xiangyu Xu
• 單位：西安交通大學(xué)
• 論文鏈接：https://arxiv.org/pdf/2606.26217
• 項(xiàng)目主頁(yè)：https://fast-lewm.github.io/
• 代碼：https://github.com/Yuntian-Gao/Fast-LeWorldModel

它的核心思想非常直接：不再用一步轉(zhuǎn)移模型反復(fù) rollout，而是把一段動(dòng)作序列的不同前綴作為預(yù)測(cè)單元，直接并行預(yù)測(cè)執(zhí)行這些動(dòng)作前綴后到達(dá)的未來(lái)潛變量，并且通過(guò)密集的監(jiān)督迫使模型學(xué)會(huì)狀態(tài)隨著不同動(dòng)作序列的演化過(guò)程，而非狀態(tài)的單步轉(zhuǎn)移。

換句話說(shuō)，模型不再問(wèn)：“執(zhí)行下一個(gè)動(dòng)作后會(huì)怎樣？”，而是直接問(wèn)：“執(zhí)行 1 個(gè)、 2 個(gè)、…… H 個(gè)動(dòng)作后，分別會(huì)到達(dá)什么狀態(tài)？”

實(shí)驗(yàn)顯示，在與 LeWM 相同的規(guī)劃協(xié)議下，F(xiàn)ast-LeWM 將平均成功率從 85.8% 提升到 90.5%；加入自一致性約束后進(jìn)一步提升到92.0%。同時(shí)模型的 rollout 中的動(dòng)態(tài)模塊耗時(shí)從31.4s 降至 8.0s，完整 CEM 求解時(shí)間從54.4s 降至 28.3s

Fast-LeWM 的 pipeline

Fast-LeWM 的方法由三部分組成。

第一步，視覺(jué)編碼器把當(dāng)前觀測(cè)和未來(lái)觀測(cè)映射到 latent space：

第二步是Action-Prefix Encoder。它把候選動(dòng)作序列通過(guò) causal Transformer 編碼成一組 prefix tokens，每個(gè) token 對(duì)應(yīng)一個(gè)不同長(zhǎng)度的動(dòng)作前綴：

其中，第 k 個(gè) prefix token 只包含前 k 個(gè)動(dòng)作的信息：

考慮到同樣的動(dòng)作序列在不同初始位置、物體狀態(tài)和接觸關(guān)系下對(duì)于動(dòng)作會(huì)產(chǎn)生不同的后果，實(shí)際實(shí)現(xiàn)中，F(xiàn)ast-LeWM 還會(huì)把當(dāng)前 latent z_t 映射成 state token，放在動(dòng)作 token 序列最前面，為動(dòng)作的編碼提供上下文信息。

第三步，Parallel Latent Predictor 使用當(dāng)前 latent 和全部 prefix token，一次性輸出所有未來(lái) latent

訓(xùn)練時(shí)，F(xiàn)ast-LeWM 對(duì)每一個(gè)前綴的預(yù)測(cè)都施加監(jiān)督，而不只是監(jiān)督最終狀態(tài)：

最終目標(biāo)保留 SIGReg 防坍塌正則：

這也是 Fast-LeWM 區(qū)別于 LeWM 的關(guān)鍵：模型不只學(xué)習(xí)狀態(tài)的局部變化，還要學(xué)習(xí)動(dòng)作逐步累積時(shí)狀態(tài)如何連續(xù)變化。

Planning 時(shí)：

基于動(dòng)作前綴的快速 rollout

論文這種動(dòng)作前綴的設(shè)計(jì)還額外帶來(lái)了一個(gè)可選的 self-consistency scoring：模型一方面可以直接從長(zhǎng)度為 H 的動(dòng)作前綴預(yù)測(cè)終點(diǎn)；另一方面也可以先預(yù)測(cè)一個(gè)中間 latent，再?gòu)闹虚g latent 繼續(xù)預(yù)測(cè)剩余時(shí)域。兩種終點(diǎn)預(yù)測(cè)之間的差異被作為一致性懲罰項(xiàng)：

其中，β 控制 self-consistency 項(xiàng)的權(quán)重。當(dāng) β=0，F(xiàn)ast-LeWM 退化為只使用 goal distance 的 CEM 打分；當(dāng) β>0 時(shí)，CEM 可以選擇那些在不同 prefix 分解下預(yù)測(cè)結(jié)果一致的候選動(dòng)作序列，進(jìn)一步提升規(guī)劃穩(wěn)定性。

成功率提升，

規(guī)劃時(shí)間近乎減半

實(shí)驗(yàn)沿用 LeWM 的 goal-conditioned latent planning 協(xié)議，在 Two-Room、Reacher、PushT、OGBench-Cube 四個(gè)環(huán)境上評(píng)測(cè):

結(jié)果顯示，F(xiàn)ast-LeWM 在四個(gè)任務(wù)上的平均成功率從 LeWM 的85.8%提升到90.5%；加入動(dòng)作前綴預(yù)測(cè)帶來(lái)的額外 Self-Consistency 規(guī)劃項(xiàng)后進(jìn)一步達(dá)到92.0%

效率提升更明顯。在相同 CEM budget、單張 NVIDIA 4090 上，F(xiàn)ast-LeWM 的 dynamics time 從31.4s降到8.0s，加速約4 倍，其中包含動(dòng)作編碼和 predictor 預(yù)測(cè)時(shí)間。完整 CEM solve time 從54.4s降到28.3s，減少48.0%

另外開(kāi)環(huán)情況下，F(xiàn)ast-LeWM 想象未來(lái)時(shí)的 latent 誤差和誤差隨 Horizon 的增長(zhǎng)率也更小：

消融實(shí)驗(yàn)：

不是簡(jiǎn)單把動(dòng)作塊變長(zhǎng)就行

作者進(jìn)一步通過(guò)消融實(shí)驗(yàn)驗(yàn)證了 Fast-LeWM 各個(gè)組件的作用。首先一個(gè)看似直接的加速方式是：把 LeWM 的動(dòng)作塊變長(zhǎng)，讓一次 transition 覆蓋更長(zhǎng)時(shí)間。作者構(gòu)造了 Long-Action LeWM，將原本 action encoding 從 5 個(gè) primitive actions 改為 25 個(gè) primitive actions, 結(jié)果效果并不好，Terminal-only Fast-LeWM 只監(jiān)督最終 latent，不監(jiān)督中間 prefix latent，表現(xiàn)優(yōu)于 Long-Action LeWM，但仍低于完整模型。這說(shuō)明 action prefix 本身已經(jīng)是更有效的長(zhǎng)時(shí)域表示，但 dense prefix supervision 對(duì)學(xué)習(xí)連續(xù)狀態(tài)演化仍然關(guān)鍵。

作者還發(fā)現(xiàn)，去掉 state token 后，模型在多個(gè)任務(wù)上性能下降。這進(jìn)一步說(shuō)明了動(dòng)作編碼需要提供有效的上下文信息。

總結(jié)

Fast-LeWM 針對(duì)世界模型在規(guī)劃階段的關(guān)鍵瓶頸，提出了 action-prefix prediction 機(jī)制，將傳統(tǒng)的一步自回歸 rollout 改為并行多時(shí)域潛變量預(yù)測(cè)。

在相同 LeWM 評(píng)測(cè)協(xié)議下，F(xiàn)ast-LeWM 將平均規(guī)劃成功率從 85.8% 提升到 90.5%，加入 self-consistency 后達(dá)到 92.0%；同時(shí)將動(dòng)態(tài)模塊耗時(shí)從 31.4s 降到 8.0s，將完整 CEM 求解時(shí)間從 54.4s 降到 28.3s。

更深層意義：

世界模型的瓶頸不在 “模型”，而在 “接口”

這項(xiàng)工作的核心啟示并不只是加速，更本質(zhì)的是它表明，對(duì)于面向規(guī)劃的世界模型而言，動(dòng)態(tài)模型的接口設(shè)計(jì)本身可能與表征學(xué)習(xí)目標(biāo)同樣重要。相比一步步預(yù)測(cè) “下一個(gè) latent”，直接預(yù)測(cè)動(dòng)作前綴導(dǎo)致的多時(shí)域未來(lái)狀態(tài)，或許是讓視覺(jué)世界模型走向高效規(guī)劃的一條更直接的路徑。這可能意味著，世界模型正在從 “逐步想象未來(lái)”，走向 “并行生成未來(lái)”。

視頻鏈接：https://mp.weixin.qq.com/s/JIvkoo1pKkzup4Q6RdTLyQ