LeCun團隊新作:在世界模型中導航

机器之心發表於2024-12-07
現實世界版的 Genie-2?

最近,世界模型(World Models)似乎成為了 AI 領域最熱門的研究方向。

World Labs(李飛飛)谷歌 DeepMind 接連發布自己的世界模型研究之後,Meta FAIR 的 Yann LeCun 團隊也加入了戰場,也在同一周之內釋出了導航世界模型(Navigation World Models/NWM)。

我們知道,Yann LeCun 一邊在不斷唱衰當前主導 LLM 領域的自迴歸正規化,同時也一直是世界模型的「鼓吹者」。上月中旬,該團隊就已經發布了一篇世界模型相關研究成果,但那篇研究涉及的模型規模不大,環境也較為簡單,參閱機器之心報導《LeCun 的世界模型初步實現!基於預訓練視覺特徵,看一眼任務就能零樣本規劃》。

而這一次,LeCun 團隊釋出的 NWM 看起來能適應更復雜的環境了,並且與 World Labs 和 DeepMind 的世界模型一樣,也能基於單張影像生成連續一致的影片。只是 LeCun 團隊的這個世界模型更加強調世界模型的導航能力,其能夠在已知環境中按照軌跡行進,也能在未知環境中自己尋找前進道路,還能執行路徑規劃。不過整體而言,與能單圖生世界的 DeepMind Genie 2 相比,NWM 的單圖生影片還是要稍遜一些。
圖片
  • 論文標題:Navigation World Models
  • 論文地址:https://arxiv.org/pdf/2412.03572v1
  • 專案地址:https://www.amirbar.net/nwm/

從其專案網站的演示影片看,NWM 的效果很不錯,能夠基於單張真實照片執行相當好的導航操作。只能說,世界模型,也開始捲起來了。

NWM 效果演示

在深入瞭解 NWM 的技術細節之前,我們先來看看它的實際表現如何。

首先是在已知環境中按照軌跡行進的能力。NWM 能夠基於單張輸入幀和給定的輸入動作合成影片,這個過程是自迴歸式的。另需說明,在這裡,模型已經已經訓練階段看過了這個環境,但軌跡是全新的。可以看到,不管是室內環境還是室外環境,NWM 都具有相當不錯的場景理解表現。
圖片
圖片
NWM 也能在未知環境中導航:它不僅適用於已知環境,對於訓練中從未見過的單張輸入影像,模型也可以根據給定的輸入動作自迴歸式地預測後續幀。
圖片
圖片
下面是與其它模型的對比情況,可以看到,NWM 在保證合成影片的一致性和穩定性方面以及動作的執行效果方面都更加出色。
圖片
另外,該團隊也研究了使用 NWM 和外部導航策略 NoMaD 來執行規劃。具體來說,就是讓 NoMaD 給出軌跡,再讓 NWM 來進行排名 —— 後者會生成軌跡影片並選出其中得分最高的軌跡。
圖片
圖片
整體而言,LeCun 團隊的這項 NWM 研究做出了以下貢獻:

  • 提出了導航世界模型和一種全新的條件擴散 Transformer(CDiT);相比於標準 DiT,其能高效地擴充套件到 1B 引數,同時計算需求還小得多。
  • 使用來自不同機器人智慧體的影片和導航動作對 CDiT 進行了訓練,透過獨立地或與外部導航策略一起模擬導航規劃而實現規劃,從而取得了當前最先進的視覺導航效能。
  • 透過在 Ego4D 等無動作和無獎勵的影片資料上訓練 NWM,使其能在未曾見過的環境中取得更好的影片預測和生成效能。

導航世界模型

NWM 的數學描述

下面先來看看 NWM 的公式描述。直觀地說,NWM 是一個接收當前世界狀態(例如,對影像的觀察)和導航操作(描述物體移動到哪裡以及如何旋轉)的模型。然後,該模型根據智慧體的視角生成下一個世界狀態。

本文給出了一個第一人稱的影片資料集,其包含智慧體導航動作圖片,其中 圖片是影像,a_i = (u, ϕ) 是由平移引數 圖片給出的導航命令,控制向前 / 向後和左右運動,以及導航旋轉角 圖片。a_i 的導航動作可以被完全觀察到。

目標是學習一個世界模型 F,即從先前的潛在觀察 s_τ 和動作 a_τ 隨機對映到未來的潛在狀態表示 s_(t+1 ):
圖片
由於此公式簡單易懂,因此它可以自然地跨環境共享,並輕鬆擴充套件到更復雜的動作空間,例如控制機械臂。

公式 1 模擬了動作,但無法控制時間動態(temporal dynamics)。因此,作者用時移輸入 k ∈ [T_min, T_max] 擴充套件此公式,設定圖片,因此現在 a_τ 指定時間變化 k,用於確定模型應向未來(或過去)移動多少步。因此,給定當前狀態 s_τ ,可以隨機選擇 k, token 化相應的影片幀。然後可以將導航動作近似為從時間 τ 到 τ + k 的總和:
圖片
上述公式既可以學習導航動作,也可以學習時間動態。實際上,本文允許時間偏移最多 ±16 秒。

擴散 Transformer 作為世界模型

條件擴散 Transformer 架構。本文使用的架構是一個時間自迴歸 transformer 模型,該模型利用高效的 CDiT 塊(見圖 2)。
圖片
CDiT 透過將第一個注意力塊中的注意力限制在正在去噪的目標幀中的 token 上,實現了在時間上高效的自迴歸建模。為了對過去幀中的 token 進行條件處理,本文還整合了一個交叉注意力層,然後,交叉注意力透過跳躍連線層將表示情境化。

使用世界模型進行導航規劃

接下來,文章描述瞭如何使用經過訓練的 NWM 來規劃導航軌跡。直觀地說,如果世界模型熟悉某個環境,可以用它來模擬導航軌跡,並選擇那些能夠達到目標的軌跡。在未知的、分佈外的環境中,長期規劃可能依賴於想象力。

形式上,給定潛在編碼 s_0 和導航目標 s^∗,目標是尋找動作序列 (a_0, ..., a_T),以最大化到達 s^∗ 的可能性。

定義能量函式圖片,使得最小化能量與最大化未歸一化的感知相似度得分相對應, 並遵循關於狀態和動作的潛在約束。
圖片
相似度的計算方法是,使用預訓練的 VAE 解碼器將 s^∗ 和 s_T 解碼為畫素,然後測量感知相似度。

那麼問題就簡化為尋找最小化該能量函式的動作:
圖片
該目標可被重新表述成一個模型預測控制(MPC)問題,並且可使用交叉熵方法(Cross-Entropy Method)來最佳化它。

導航軌跡排名方法。假設已有一個導航策略 Π(a|s_0, s^∗),可使用 NWM 來對取樣得到的軌跡進行排名。這裡,該團隊的使用了一種 SOTA 的導航策略 NoMaD 來執行機器人導航。在排名時,會從 Π 給出的多個樣本中選出能量最低的那個。

實驗結果

下面來看看 NWM 在實驗中實際表現。

首先,資料集方面,該團隊使用了 TartanDrive、RECON 和 HuRoN。NWM 可以獲取機器人的位置和角度資料,然後推斷在當前位置的相關動作。

評估指標包括絕對軌跡誤差 (ATE)和相對姿態誤差 (RPE)。對比基線包括 DIAMOND、GNM 和 NoMaD。

消融實驗

模型在已知環境 RECON 上對驗證集軌跡對單步 4 秒未來預測進行評估。研究人員透過測量 LPIPS、DreamSim 和 PSNR 來評估相對於地面真實框架的效能。圖 3 中提供了定性示例:
圖片
模型大小和 CDiT。研究人員將 CDiT 與標準 DiT(其中所有上下文標記都作為輸入)進行比較。其中假設,對於導航已知環境,模型的容量是最重要的,圖 5 中的結果表明,CDiT 確實在具有多達 1B 個引數的模型中表現更好,同時消耗的 FLOP 不到 ×2。令人驚訝的是,即使引數數量相同(例如,CDiT-L 與 DiT-XL 相比),CDiT 也可以快 4 倍,並且表現更好。
圖片
目標數量。在給定固定上下文的情況下訓練具有可變目標狀態數量的模型,將目標數量從 1 更改為 4。每個目標都是在當前狀態的 ±16 秒視窗內隨機選擇的。表 1 中報告的結果表明,使用 4 個目標可顯著提高所有指標的預測效能。

上下文大小。研究人員在訓練模型的同時將條件幀的數量從 1 變為 4(見表 1)。不出所料,更多的上下文帶來了幫助,而對於較短的上下文,模型通常會「迷失方向」,導致預測不佳。

時間和動作條件。研究人員同時使用時間和動作條件訓練模型,並測試每個輸入對預測效能的貢獻程度。結果包含在表 1 中。研究人員發現,使用時間執行模型只會導致效能不佳,而不使用時間條件也會導致效能略有下降。這證實了兩種輸入對模型都有好處。
圖片
影片預測與合成

這裡評估的是模型遵從真實動作和預測未來狀態的能力。

以第一張影像和上下文幀為條件,該模型需要根據 ground truth 動作,以自迴歸方式預測下一個狀態,並給每個預測提供反饋。

透過比較在 1、2、4、8 和 16 秒的 ground truth 影像,再得出在 RECON 資料集上的 FID 和 LPIPS 值,可以對這些預測結果進行比較。

圖 4 展示了在 4 FPS 和 1 FPS 幀率下,NWM 與 DIAMOND 的效能情況。可以明顯看到,NWM 的預測準確度比 DIAMOND 好得多。
圖片
一開始的時候,NWM 1 FPS 的表現更好,但 8 秒之後,它就會因為累積誤差和上下文損失而被 4 FPS 版本超過。

生成質量。為了評估影片質量,該團隊以 4 FPS 的速度自迴歸預測生成了一些 16 秒長的影片,同時這是基於 ground truth 動作的。然後,再使用 FVD 評估生成影片的質量,並與 DIAMOND 進行比較。圖 6 中的結果表明 NWM 輸出的影片質量更高。
圖片
使用 NWM 執行規劃

接下來的實驗衡量了 NWM 執行導航的能力。

獨立規劃。實驗表明,這個世界模型可以有效地獨立執行目標導向的導航。

基於過去的觀察和目標影像,NWM 可以使用交叉熵方法找到一條軌跡,同時儘可能降低預測影像和目標影像之間的 LPIPS 相似度,實驗結果見下表 2,可以看到 NWM 的規劃能力足以比肩 SOTA 策略。
圖片
帶約束條件的規劃。在使用 NWM 進行規劃時,還可以指定約束條件,比如要求智慧體走直線或只轉彎一次。

表 3 的結果表明,NWM 可以在滿足約束的同時進行有效規劃,並且規劃效能變化不大。
圖片
下圖 9 中包含了左右優先約束下的規劃軌跡案例。
圖片
使用導航世界模型進行排序。NWM 可以增強目標條件導航中已有的導航策略。研究者根據過去觀察結果和目標影像對 NoMaD 進行條件化,取樣了 n ∈ {16,32} 條軌跡,其中每條軌跡長度為 8,並透過使用 NWM 來自迴歸地遵循動作以對這些軌跡進行評估。

最後,研究者透過測量與目標影像的 LPIPS 相似性來對每條軌跡的最終預測結果進行排序,結果如下圖 7 所示。他們還在上表 2 中報告了 ATE 和 RPE,發現對軌跡進行排序可以產生 SOTA 導航效能,並且取樣的軌跡越多結果越好。
圖片
泛化到未知環境的能力

研究者嘗試新增未標註的資料,並詢問 NWM 是否可以使用想象力在新環境中做出預測。他們在所有域內資料集以及來自 Ego4D 的未標註影片子資料集上訓練了一個模型,並且只能訪問時移操作。

研究者訓練了一個 CDiT-XL 模型,並在 Go Stanford 資料集以及其他隨機影像上對該模型進行了測試。結果如下表 4 所示,可以發現,在未標註資料上進行訓練可以顯著提升各項影片預測結果,包括提高生成質量。
圖片
研究者在下圖 8 中提供了一些定性案例。相較於域內(上圖 3),模型崩潰得更快並且在生成想象環境的遍歷時還會產生幻覺路徑。
圖片
更多實驗細節請參閱原論文。

相關文章