LeCun團隊新作:在世界模型中導航
現(xiàn)實世界版的 Genie-2?
最近,世界模型(World Models)似乎成為了 AI 領域最熱門的研究方向。
繼 World Labs(李飛飛)、谷歌 DeepMind 接連發(fā)布自己的世界模型研究之后,Meta FAIR 的 Yann LeCun 團隊也加入了戰(zhàn)場,也在同一周之內(nèi)發(fā)布了導航世界模型(Navigation World Models/NWM)。
我們知道,Yann LeCun 一邊在不斷唱衰當前主導 LLM 領域的自回歸范式,同時也一直是世界模型的「鼓吹者」。上月中旬,該團隊就已經(jīng)發(fā)布了一篇世界模型相關研究成果,但那篇研究涉及的模型規(guī)模不大,環(huán)境也較為簡單,參閱機器之心報道《LeCun 的世界模型初步實現(xiàn)!基于預訓練視覺特征,看一眼任務就能零樣本規(guī)劃》。
而這一次,LeCun 團隊發(fā)布的 NWM 看起來能適應更復雜的環(huán)境了,并且與 World Labs 和 DeepMind 的世界模型一樣,也能基于單張圖像生成連續(xù)一致的視頻。只是 LeCun 團隊的這個世界模型更加強調(diào)世界模型的導航能力,其能夠在已知環(huán)境中按照軌跡行進,也能在未知環(huán)境中自己尋找前進道路,還能執(zhí)行路徑規(guī)劃。不過整體而言,與能單圖生世界的 DeepMind Genie 2 相比,NWM 的單圖生視頻還是要稍遜一些。
- 論文標題:Navigation World Models
- 論文地址:https://arxiv.org/pdf/2412.03572v1
- 項目地址:https://www.amirbar.net/nwm/
從其項目網(wǎng)站的演示視頻看,NWM 的效果很不錯,能夠基于單張真實照片執(zhí)行相當好的導航操作。只能說,世界模型,也開始卷起來了。
NWM 效果演示
在深入了解 NWM 的技術(shù)細節(jié)之前,我們先來看看它的實際表現(xiàn)如何。
首先是在已知環(huán)境中按照軌跡行進的能力。NWM 能夠基于單張輸入幀和給定的輸入動作合成視頻,這個過程是自回歸式的。另需說明,在這里,模型已經(jīng)已經(jīng)訓練階段看過了這個環(huán)境,但軌跡是全新的??梢钥吹剑还苁鞘覂?nèi)環(huán)境還是室外環(huán)境,NWM 都具有相當不錯的場景理解表現(xiàn)。
NWM 也能在未知環(huán)境中導航:它不僅適用于已知環(huán)境,對于訓練中從未見過的單張輸入圖像,模型也可以根據(jù)給定的輸入動作自回歸式地預測后續(xù)幀。
下面是與其它模型的對比情況,可以看到,NWM 在保證合成視頻的一致性和穩(wěn)定性方面以及動作的執(zhí)行效果方面都更加出色。
另外,該團隊也研究了使用 NWM 和外部導航策略 NoMaD 來執(zhí)行規(guī)劃。具體來說,就是讓 NoMaD 給出軌跡,再讓 NWM 來進行排名 —— 后者會生成軌跡視頻并選出其中得分最高的軌跡。
整體而言,LeCun 團隊的這項 NWM 研究做出了以下貢獻:
- 提出了導航世界模型和一種全新的條件擴散 Transformer(CDiT);相比于標準 DiT,其能高效地擴展到 1B 參數(shù),同時計算需求還小得多。
- 使用來自不同機器人智能體的視頻和導航動作對 CDiT 進行了訓練,通過獨立地或與外部導航策略一起模擬導航規(guī)劃而實現(xiàn)規(guī)劃,從而取得了當前最先進的視覺導航性能。
- 通過在 Ego4D 等無動作和無獎勵的視頻數(shù)據(jù)上訓練 NWM,使其能在未曾見過的環(huán)境中取得更好的視頻預測和生成性能。
導航世界模型
NWM 的數(shù)學描述
下面先來看看 NWM 的公式描述。直觀地說,NWM 是一個接收當前世界狀態(tài)(例如,對圖像的觀察)和導航操作(描述物體移動到哪里以及如何旋轉(zhuǎn))的模型。然后,該模型根據(jù)智能體的視角生成下一個世界狀態(tài)。
本文給出了一個第一人稱的視頻數(shù)據(jù)集,其包含智能體導航動作,其中
是圖像,a_i = (u, ?) 是由平移參數(shù)
給出的導航命令,控制向前 / 向后和左右運動,以及導航旋轉(zhuǎn)角
。a_i 的導航動作可以被完全觀察到。
目標是學習一個世界模型 F,即從先前的潛在觀察 s_τ 和動作 a_τ 隨機映射到未來的潛在狀態(tài)表示 s_(t+1 ):
由于此公式簡單易懂,因此它可以自然地跨環(huán)境共享,并輕松擴展到更復雜的動作空間,例如控制機械臂。
公式 1 模擬了動作,但無法控制時間動態(tài)(temporal dynamics)。因此,作者用時移輸入 k ∈ [T_min, T_max] 擴展此公式,設置,因此現(xiàn)在 a_τ 指定時間變化 k,用于確定模型應向未來(或過去)移動多少步。因此,給定當前狀態(tài) s_τ ,可以隨機選擇 k, token 化相應的視頻幀。然后可以將導航動作近似為從時間 τ 到 τ + k 的總和:
上述公式既可以學習導航動作,也可以學習時間動態(tài)。實際上,本文允許時間偏移最多 ±16 秒。
擴散 Transformer 作為世界模型
條件擴散 Transformer 架構(gòu)。本文使用的架構(gòu)是一個時間自回歸 transformer 模型,該模型利用高效的 CDiT 塊(見圖 2)。
CDiT 通過將第一個注意力塊中的注意力限制在正在去噪的目標幀中的 token 上,實現(xiàn)了在時間上高效的自回歸建模。為了對過去幀中的 token 進行條件處理,本文還整合了一個交叉注意力層,然后,交叉注意力通過跳躍連接層將表示情境化。
使用世界模型進行導航規(guī)劃
接下來,文章描述了如何使用經(jīng)過訓練的 NWM 來規(guī)劃導航軌跡。直觀地說,如果世界模型熟悉某個環(huán)境,可以用它來模擬導航軌跡,并選擇那些能夠達到目標的軌跡。在未知的、分布外的環(huán)境中,長期規(guī)劃可能依賴于想象力。
形式上,給定潛在編碼 s_0 和導航目標 s^?,目標是尋找動作序列 (a_0, ..., a_T),以最大化到達 s^? 的可能性。
定義能量函數(shù),使得最小化能量與最大化未歸一化的感知相似度得分相對應, 并遵循關于狀態(tài)和動作的潛在約束。
相似度的計算方法是,使用預訓練的 VAE 解碼器將 s^? 和 s_T 解碼為像素,然后測量感知相似度。
那么問題就簡化為尋找最小化該能量函數(shù)的動作:
該目標可被重新表述成一個模型預測控制(MPC)問題,并且可使用交叉熵方法(Cross-Entropy Method)來優(yōu)化它。
導航軌跡排名方法。假設已有一個導航策略 Π(a|s_0, s^?),可使用 NWM 來對采樣得到的軌跡進行排名。這里,該團隊的使用了一種 SOTA 的導航策略 NoMaD 來執(zhí)行機器人導航。在排名時,會從 Π 給出的多個樣本中選出能量最低的那個。
實驗結(jié)果
下面來看看 NWM 在實驗中實際表現(xiàn)。
首先,數(shù)據(jù)集方面,該團隊使用了 TartanDrive、RECON 和 HuRoN。NWM 可以獲取機器人的位置和角度數(shù)據(jù),然后推斷在當前位置的相關動作。
評估指標包括絕對軌跡誤差 (ATE)和相對姿態(tài)誤差 (RPE)。對比基線包括 DIAMOND、GNM 和 NoMaD。
消融實驗
模型在已知環(huán)境 RECON 上對驗證集軌跡對單步 4 秒未來預測進行評估。研究人員通過測量 LPIPS、DreamSim 和 PSNR 來評估相對于地面真實框架的性能。圖 3 中提供了定性示例:
模型大小和 CDiT。研究人員將 CDiT 與標準 DiT(其中所有上下文標記都作為輸入)進行比較。其中假設,對于導航已知環(huán)境,模型的容量是最重要的,圖 5 中的結(jié)果表明,CDiT 確實在具有多達 1B 個參數(shù)的模型中表現(xiàn)更好,同時消耗的 FLOP 不到 ×2。令人驚訝的是,即使參數(shù)數(shù)量相同(例如,CDiT-L 與 DiT-XL 相比),CDiT 也可以快 4 倍,并且表現(xiàn)更好。
目標數(shù)量。在給定固定上下文的情況下訓練具有可變目標狀態(tài)數(shù)量的模型,將目標數(shù)量從 1 更改為 4。每個目標都是在當前狀態(tài)的 ±16 秒窗口內(nèi)隨機選擇的。表 1 中報告的結(jié)果表明,使用 4 個目標可顯著提高所有指標的預測性能。
上下文大小。研究人員在訓練模型的同時將條件幀的數(shù)量從 1 變?yōu)?4(見表 1)。不出所料,更多的上下文帶來了幫助,而對于較短的上下文,模型通常會「迷失方向」,導致預測不佳。
時間和動作條件。研究人員同時使用時間和動作條件訓練模型,并測試每個輸入對預測性能的貢獻程度。結(jié)果包含在表 1 中。研究人員發(fā)現(xiàn),使用時間運行模型只會導致性能不佳,而不使用時間條件也會導致性能略有下降。這證實了兩種輸入對模型都有好處。
視頻預測與合成
這里評估的是模型遵從真實動作和預測未來狀態(tài)的能力。
以第一張圖像和上下文幀為條件,該模型需要根據(jù) ground truth 動作,以自回歸方式預測下一個狀態(tài),并給每個預測提供反饋。
通過比較在 1、2、4、8 和 16 秒的 ground truth 圖像,再得出在 RECON 數(shù)據(jù)集上的 FID 和 LPIPS 值,可以對這些預測結(jié)果進行比較。
圖 4 展示了在 4 FPS 和 1 FPS 幀率下,NWM 與 DIAMOND 的性能情況??梢悦黠@看到,NWM 的預測準確度比 DIAMOND 好得多。
一開始的時候,NWM 1 FPS 的表現(xiàn)更好,但 8 秒之后,它就會因為累積誤差和上下文損失而被 4 FPS 版本超過。
生成質(zhì)量。為了評估視頻質(zhì)量,該團隊以 4 FPS 的速度自回歸預測生成了一些 16 秒長的視頻,同時這是基于 ground truth 動作的。然后,再使用 FVD 評估生成視頻的質(zhì)量,并與 DIAMOND 進行比較。圖 6 中的結(jié)果表明 NWM 輸出的視頻質(zhì)量更高。
使用 NWM 執(zhí)行規(guī)劃
接下來的實驗衡量了 NWM 執(zhí)行導航的能力。
獨立規(guī)劃。實驗表明,這個世界模型可以有效地獨立執(zhí)行目標導向的導航。
基于過去的觀察和目標圖像,NWM 可以使用交叉熵方法找到一條軌跡,同時盡可能降低預測圖像和目標圖像之間的 LPIPS 相似度,實驗結(jié)果見下表 2,可以看到 NWM 的規(guī)劃能力足以比肩 SOTA 策略。
帶約束條件的規(guī)劃。在使用 NWM 進行規(guī)劃時,還可以指定約束條件,比如要求智能體走直線或只轉(zhuǎn)彎一次。
表 3 的結(jié)果表明,NWM 可以在滿足約束的同時進行有效規(guī)劃,并且規(guī)劃性能變化不大。
下圖 9 中包含了左右優(yōu)先約束下的規(guī)劃軌跡案例。
使用導航世界模型進行排序。NWM 可以增強目標條件導航中已有的導航策略。研究者根據(jù)過去觀察結(jié)果和目標圖像對 NoMaD 進行條件化,采樣了 n ∈ {16,32} 條軌跡,其中每條軌跡長度為 8,并通過使用 NWM 來自回歸地遵循動作以對這些軌跡進行評估。
最后,研究者通過測量與目標圖像的 LPIPS 相似性來對每條軌跡的最終預測結(jié)果進行排序,結(jié)果如下圖 7 所示。他們還在上表 2 中報告了 ATE 和 RPE,發(fā)現(xiàn)對軌跡進行排序可以產(chǎn)生 SOTA 導航性能,并且采樣的軌跡越多結(jié)果越好。
泛化到未知環(huán)境的能力
研究者嘗試添加未標注的數(shù)據(jù),并詢問 NWM 是否可以使用想象力在新環(huán)境中做出預測。他們在所有域內(nèi)數(shù)據(jù)集以及來自 Ego4D 的未標注視頻子數(shù)據(jù)集上訓練了一個模型,并且只能訪問時移操作。
研究者訓練了一個 CDiT-XL 模型,并在 Go Stanford 數(shù)據(jù)集以及其他隨機圖像上對該模型進行了測試。結(jié)果如下表 4 所示,可以發(fā)現(xiàn),在未標注數(shù)據(jù)上進行訓練可以顯著提升各項視頻預測結(jié)果,包括提高生成質(zhì)量。
研究者在下圖 8 中提供了一些定性案例。相較于域內(nèi)(上圖 3),模型崩潰得更快并且在生成想象環(huán)境的遍歷時還會產(chǎn)生幻覺路徑。
更多實驗細節(jié)請參閱原論文。