世界模型WHALE來了！

人類能夠在腦海中設(shè)想一個(gè)想象中的世界，以預(yù)測(cè)不同的動(dòng)作可能導(dǎo)致不同的結(jié)果。受人類智能這方面的啟發(fā)，世界模型被設(shè)計(jì)用于抽象化現(xiàn)實(shí)世界的動(dòng)態(tài)，并提供這種“如果……會(huì)怎樣”的預(yù)測(cè)。

因此，具身智能體可以與世界模型進(jìn)行交互，而不是直接與現(xiàn)實(shí)世界環(huán)境交互，以生成模擬數(shù)據(jù)，這些數(shù)據(jù)可以用于各種下游任務(wù)，包括反事實(shí)預(yù)測(cè)、離線策略評(píng)估、離線強(qiáng)化學(xué)習(xí)。

世界模型在具身環(huán)境的決策中起著至關(guān)重要的作用，使得在現(xiàn)實(shí)世界中成本高昂的探索成為可能。為了促進(jìn)有效的決策，世界模型必須具備強(qiáng)大的泛化能力，以支持分布外（OOD）區(qū)域的想象，并提供可靠的不確定性估計(jì)來評(píng)估模擬體驗(yàn)的可信度，這兩者都對(duì)之前的可擴(kuò)展方法提出了重大挑戰(zhàn)。

近日，來自南京大學(xué)、南棲仙策等機(jī)構(gòu)的研究者在論文中引入了WHALE（World models with beHavior-conditioning and retrAcing-rollout LEarning），這是一個(gè)用于學(xué)習(xí)可泛化世界模型的框架，由兩種可以與任何神經(jīng)網(wǎng)絡(luò)架構(gòu)普遍結(jié)合的關(guān)鍵技術(shù)組成。

在確定策略分布差異是泛化誤差的主要來源的基礎(chǔ)上，研究者引入了一種行為—條件（behavior-conditioning）技術(shù)來增強(qiáng)世界模型的泛化能力，該技術(shù)建立在策略條件模型學(xué)習(xí)的概念之上，旨在使模型能夠主動(dòng)適應(yīng)不同的行為，以減小分布偏移引起的外推誤差。此外，研究者還提出了一種簡單而有效的技術(shù)，稱為retracing-rollout，以便對(duì)模型想象進(jìn)行有效的不確定性估計(jì)，作為一種即插即用的解決方案，可以有效地應(yīng)用于各種實(shí)施任務(wù)中的末端執(zhí)行器姿態(tài)控制，而無需對(duì)訓(xùn)練過程進(jìn)行任何更改。

通過整合WHALE的這兩種技術(shù)，研究者提出了WHALE-ST，這是一種可擴(kuò)展的基于時(shí)空transformer的世界模型，旨在實(shí)現(xiàn)更有效的決策。研究者進(jìn)一步提出了WHALE-X，這是一個(gè)在970K機(jī)器人演示上預(yù)訓(xùn)練的414M參數(shù)世界模型。最后，研究者進(jìn)行了大量的實(shí)驗(yàn)，以證明WHALE-ST和WHALE-X在模擬和現(xiàn)實(shí)世界任務(wù)中的卓越可擴(kuò)展性與泛化性，突出了它們?cè)谠鰪?qiáng)決策方面的效果。

為了評(píng)估WHALE-X在實(shí)際物理環(huán)境中的泛化能力，研究團(tuán)隊(duì)在ARX5機(jī)器人上進(jìn)行了全面實(shí)驗(yàn)。與預(yù)訓(xùn)練數(shù)據(jù)不同，評(píng)估任務(wù)調(diào)整了攝像機(jī)角度和背景等，增加了對(duì)世界模型的挑戰(zhàn)。他們收集了每個(gè)任務(wù)60條軌跡的數(shù)據(jù)集用于微調(diào)，任務(wù)包括開箱、推盤、投球和移動(dòng)瓶子，還設(shè)計(jì)了多個(gè)模型從未接觸過的任務(wù)來測(cè)試模型的視覺、運(yùn)動(dòng)和任務(wù)泛化能力。

結(jié)果顯示，WHALE-X在真實(shí)世界中展現(xiàn)出明顯的優(yōu)勢(shì)：與沒有行為—條件的模型相比，WHALEX的一致性提高了63％，表明該機(jī)制顯著提升了OOD泛化能力；在97萬個(gè)樣本上進(jìn)行預(yù)訓(xùn)練的WHALE-X，比從零開始訓(xùn)練的模型具有更高的一致性，凸顯了大規(guī)模互聯(lián)網(wǎng)數(shù)據(jù)預(yù)訓(xùn)練的優(yōu)勢(shì)；增加模型參數(shù)能夠提升世界模型的泛化能力，WHALE-X-base（203M）動(dòng)態(tài)模型在三個(gè)未見任務(wù)中的一致性比率是77M版本的3倍。此外，視頻生成質(zhì)量與一致性的結(jié)果一致。通過行為—條件策略、大規(guī)模預(yù)訓(xùn)練數(shù)據(jù)集和擴(kuò)展模型參數(shù)，三種策略結(jié)合，顯著提高了模型的OOD泛化能力，尤其是在生成高質(zhì)量視頻方面。