剛剛，LeCun團隊讓世界模型學(xué)會持續(xù)學(xué)習(xí)！

henry 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

世界模型，也能持續(xù)學(xué)習(xí)了！

剛剛，紐約大學(xué)聯(lián)合LeCun初創(chuàng)AMI帶來JEPA系列的最新成果——

AdaJEPA

與過去在預(yù)訓(xùn)練結(jié)束后就凍結(jié)參數(shù)的世界模型不同，AdaJEPA能夠在與環(huán)境交互中，基于測試時自適應(yīng)（Test-Time Adaptation, TTA），實時調(diào)整世界模型的編碼器和預(yù)測器參數(shù)，從而實現(xiàn)持續(xù)學(xué)習(xí)。

具體而言，AdaJEPA通過計劃、執(zhí)行、觀測、更新、再規(guī)劃的閉環(huán)，在每次交互中只執(zhí)行MPC規(guī)劃出的第一段動作，然后把真實觀察到的下一幀狀態(tài)，當(dāng)成自監(jiān)督信號來更新世界模型。

由此，在下一輪規(guī)劃時用的就不再是剛部署時那個凍結(jié)模型，而是已經(jīng)被當(dāng)前環(huán)境“校準(zhǔn)”過的模型。

這個思路有點像經(jīng)典強化學(xué)習(xí)里的Dyna

模型不是一次訓(xùn)練完就結(jié)束，而是在真實交互中不斷修正自己對世界的理解。

實驗結(jié)果也表明，無論是在分布內(nèi)環(huán)境，還是面對各類分布外偏移，AdaJEPA的規(guī)劃成功率都明顯優(yōu)于固定世界模型。

這是怎么做到的？

計劃、執(zhí)行、觀測、更新、再規(guī)劃

一直以來，基于JEPA路線的隱空間世界模型，都有一個默認(rèn)前提，就是模型訓(xùn)練完，就凍結(jié)參數(shù)。

整個流程大致是這樣：

模型先在離線軌跡上學(xué)習(xí)把高維圖像壓進latent space，然后再在這個隱空間里預(yù)測未來。

到了測試階段，MPC（Model Predictive Control，模型預(yù)測控制）就會調(diào)用這個凍結(jié)的世界模型，在隱空間里向前滾動“想象”未來，優(yōu)化出一串動作，再把第一步動作拿到真實環(huán)境里執(zhí)行。

（注：MPC的核心思想是每次只往前預(yù)測一小段，算出一串動作，但先只執(zhí)行第一步。等真實環(huán)境反饋回來，再重新預(yù)測、重新規(guī)劃）

但問題在于，環(huán)境一變，凍結(jié)世界模型就容易失準(zhǔn)。

當(dāng)系統(tǒng)面臨測試時分布偏移（Test-time Distribution Shift）時，在latent space里看起來能到達(dá)目標(biāo)的動作，落到真實環(huán)境里，可能一步都不對。

更麻煩的是，MPC本來就靠短時域滾動規(guī)劃，單步誤差往后滾幾步，就會被放大。

為解決這一問題，論文提出AdaJEPA框架。

它的核心判斷是：世界模型不該訓(xùn)練完就固定在那里。它應(yīng)該像真正部署中的智能體一樣，一邊行動，一邊用新經(jīng)驗校準(zhǔn)自己。

具體來說，AdaJEPA的循環(huán)可以分為四步：

• ：模型先把當(dāng)前觀測編碼成latent state，然后用當(dāng)前世界模型進行MPC，在隱空間里向前滾動預(yù)測，找出一串最接近目標(biāo)狀態(tài)的動作。
• 執(zhí)行：模型不會一次性執(zhí)行完整動作序列，而是只執(zhí)行第一段動作。隨后，真實環(huán)境返回下一幀觀測。
• 更新：把這次真實狀態(tài)轉(zhuǎn)移存進在線緩存區(qū)。AdaJEPA再讓模型根據(jù)觀測和動作預(yù)測下一步latent state，并和真實狀態(tài)預(yù)測編碼出的latent state對齊。預(yù)測錯在哪里，梯度就從哪里回來。
• 再規(guī)劃：更新后的世界模型立刻進入下一輪MPC。論文里默認(rèn)只更新視覺編碼器和預(yù)測器的最后幾層，每次重規(guī)劃只做1步梯度下降。

由此，AdaJEPA的循環(huán)不再只是傳統(tǒng) MPC 的：規(guī)劃，執(zhí)行，再規(guī)劃。

而是變成了：規(guī)劃，執(zhí)行，觀測，更新，再規(guī)劃。

世界模型也因此不再只是一個被動調(diào)用的“想象器”，而變成了一個會在部署過程中持續(xù)校準(zhǔn)自己的模塊。

實現(xiàn)細(xì)節(jié)

在實現(xiàn)上，AdaJEPA的底座依然是JEPA（Joint-Embedding Predictive Architectures），聯(lián)合嵌入預(yù)測架構(gòu)。

和傳統(tǒng)像素級預(yù)測世界模型不同，JEPA并不直接預(yù)測未來圖像，而是先把圖像壓進一個更緊湊的隱空間，只在latent space里預(yù)測未來狀態(tài)。

具體來說，整個模型由三個核心組件組成：

• 狀態(tài)編碼器，把當(dāng)前觀測編碼成隱狀態(tài)。
• 動作編碼器 ，把動作編碼成動作嵌入。
• 預(yù)測器，根據(jù)當(dāng)前隱狀態(tài)和動作嵌入，預(yù)測下一步隱狀態(tài)。

AdaJEPA的在線更新，就發(fā)生在這個隱空間里。

每次執(zhí)行動作后，系統(tǒng)都會把真實狀態(tài)轉(zhuǎn)移存入在線緩存區(qū)。這個緩存區(qū)不會無限增長，論文里默認(rèn)只保留最近N條轉(zhuǎn)移。

更新時，AdaJEPA會讓模型根據(jù)當(dāng)前觀測和動作預(yù)測下一時刻的隱狀態(tài)，再和真實下一幀觀測編碼出的隱狀態(tài)對齊。

為了防止在線更新把原本的表征空間拉崩，論文做了兩個限制：

一是對目標(biāo)表征使用stop-gradient；二是只更新少量參數(shù)。

實驗?zāi)J(rèn)只更新視覺編碼器和預(yù)測器的最后幾層，并且每次MPC重規(guī)劃只進行1步梯度下降。

所以，這并不是把整個世界模型在線重訓(xùn)一遍。

它更像是每走一步，就用剛剛從真實環(huán)境里獲得的新反饋，把世界模型往當(dāng)前環(huán)境上輕輕校準(zhǔn)一下。

實驗測試

為了驗證測試時自適應(yīng)能否讓凍結(jié)世界模型在環(huán)境變化中重新校準(zhǔn)自己，論文在PushT/PushObjPointMaze兩個基準(zhǔn)上做了實驗。

實驗表明，在 PushObj未見過形狀上，AdaJEPA幾乎把規(guī)劃成功率翻倍。

而在PointMaze的未見過布局中，提升同樣明顯。

其中，GD規(guī)劃成功率從53.3%提到 78.7%，CEM從49.3%提到70.7%。

更關(guān)鍵的是，這種在線更新帶來的額外延遲并不高。

論文默認(rèn)只更新視覺編碼器和預(yù)測器的最后幾層，每次MPC重規(guī)劃只做1步梯度下降，額外延遲只有 0.01到0.03秒。

也就是說，AdaJEPA并不是用很重的在線訓(xùn)練，強行換來更高成功率。

它更像是在原有世界模型的基礎(chǔ)上，加入了一個輕量的“部署時自我校準(zhǔn)”機制。

總體來看，這篇論文想說明的是：

世界模型不必訓(xùn)練完就被凍結(jié)，只要讓它在部署時利用真實交互結(jié)果做輕量更新，就已經(jīng)能顯著提升面對環(huán)境變化時的魯棒性。

作者介紹

最后，讓我們來簡單介紹一下這篇論文的作者們。

作者Ying Wang目前是紐約大學(xué)數(shù)據(jù)科學(xué)中心CILVR Lab 的博士生，研究方向是世界模型，導(dǎo)師是Mengye RenYann LeCun

另一位作者Oumayma Bounou目前是紐約大學(xué)博士后研究員，她的研究興趣集中在世界模型、控制和優(yōu)化。

目前，她正與LeCun合作研究世界模型。

此外，紐約大學(xué)計算機科學(xué)與數(shù)據(jù)科學(xué)助理教授Mengye Ren和圖靈獎獲得者Yann LeCun共同擔(dān)任指導(dǎo)作者。

參考鏈接[1]https://arxiv.org/pdf/2606.32026