機(jī)器人視覺控制新范式!ByteDance Research新算法實(shí)現(xiàn)通過性能SOTA
世界模型(World Model)作為近年來機(jī)器學(xué)習(xí)和強(qiáng)化學(xué)習(xí)的研究熱點(diǎn),通過建立智能體對(duì)其所處環(huán)境的一種內(nèi)部表征和模擬,能夠加強(qiáng)智能體對(duì)于世界的理解,進(jìn)而更好地進(jìn)行規(guī)劃和決策。在強(qiáng)化學(xué)習(xí)領(lǐng)域中,世界模型通常被建模為一個(gè)神經(jīng)網(wǎng)絡(luò),通過歷史狀態(tài)和動(dòng)作,預(yù)測(cè)未來可能出現(xiàn)的狀態(tài)。其中,Dreamer 算法在多種模擬環(huán)境的成功表現(xiàn)讓我們看到了世界模型優(yōu)秀的表征和泛化能力。如果將世界模型應(yīng)用于復(fù)雜真實(shí)場(chǎng)景,是否能夠?qū)崿F(xiàn)更好的控制決策呢?
對(duì)此,ByteDance Research 研究團(tuán)隊(duì)成功將世界模型應(yīng)用于四足機(jī)器人視覺控制領(lǐng)域,提出了基于世界模型的感知算法 WMP(World Model-based Perception),WMP 通過在模擬器中學(xué)習(xí)世界模型和策略,其中世界模型通過歷史感知信息(包括視覺感知和本體感知)預(yù)測(cè)未來的感知,策略以世界模型提取的特征作為輸入,輸出具體控制動(dòng)作。
WMP 將模擬器中訓(xùn)練的世界模型和策略 Zero-Shot 遷移到宇樹 A1 機(jī)器人進(jìn)行驗(yàn)證,在多種環(huán)境下取得了出色的成績(jī),達(dá)到了目前為止 A1 機(jī)器人 SOTA 的通過性能。同時(shí),使用模擬數(shù)據(jù)訓(xùn)練的世界模型可以準(zhǔn)確預(yù)測(cè)真實(shí)軌跡,展示出卓越的泛化性能,有望成為一種機(jī)器人控制的新范式。
- 項(xiàng)目主頁:https://wmp-loco.github.io/
- 論文地址:https://arxiv.org/abs/2409.16784
研究背景
近年來,強(qiáng)化學(xué)習(xí) (RL) 通過在物理模擬器中訓(xùn)練策略,然后將其轉(zhuǎn)移到現(xiàn)實(shí)世界(Sim-to-Real transfer),在足式機(jī)器人的運(yùn)動(dòng)控制領(lǐng)域被廣泛應(yīng)用。在足式機(jī)器人的運(yùn)動(dòng)控制中,視覺圖像信息對(duì)于諸如越障等復(fù)雜環(huán)境是不可或缺的。
憑借強(qiáng)化學(xué)習(xí)的獎(jiǎng)勵(lì)信號(hào)從長(zhǎng)序列、高維的視覺信息中學(xué)習(xí)策略難度極大。為了更好地引入視覺信息,傳統(tǒng)特權(quán)學(xué)習(xí)(privileged learning)框架首先學(xué)習(xí)一個(gè)教師策略,教師策略的輸入包含只能在模擬器中得到的特權(quán)信息,如高度圖掃標(biāo)點(diǎn)(heightmap scandots),各類障礙物的幾何參數(shù)等。之后再學(xué)習(xí)一個(gè)以深度圖序列為輸入的學(xué)生策略以模仿教師策略的動(dòng)作。
由于特權(quán)學(xué)習(xí)兩階段的訓(xùn)練模式,學(xué)生策略的性能往往落后于教師策略;并且特權(quán)信息的選擇需要人工設(shè)計(jì),較為繁瑣,同時(shí)易受到具體環(huán)境的限制。如下圖所示,使用 scandots 作為特權(quán)信息無法處理需要精確距離的環(huán)境以及存在空中障礙物的環(huán)境。
相比之下,動(dòng)物能夠在沒有特權(quán)信息的情況下僅憑借視覺感官信息通過各種非結(jié)構(gòu)化的地形,并且在不熟悉的環(huán)境中依然能夠做出合理的決策。認(rèn)知科學(xué)的一種解釋是動(dòng)物會(huì)利用腦中建立的心智模型(mental model)對(duì)外部環(huán)境進(jìn)行理解和預(yù)測(cè)以便更好地理解和應(yīng)對(duì)周圍的環(huán)境,從而做出合理的動(dòng)作和決策。
ByteDance Research 的機(jī)器人研究團(tuán)隊(duì)在此啟發(fā)下設(shè)計(jì)了一種更加通用的機(jī)器人控制框架 WMP(World Model-based Perception)。WMP 通過構(gòu)建世界模型來處理復(fù)雜的感知信息,并將世界模型提取的環(huán)境信息輸入給策略,解決了特權(quán)學(xué)習(xí)中特權(quán)信息難以設(shè)計(jì)的局限性。通過訓(xùn)練的世界模型和策略可以直接遷移到真實(shí)環(huán)境的 A1 機(jī)器人上,在多種復(fù)雜任務(wù)中達(dá)到了目前該領(lǐng)域的領(lǐng)先水平效果。例如,在世界模型的幫助下,A1 機(jī)器人可以跳過 85cm 的間隙,跳上 55cm 的高臺(tái),穿過 22cm 高的橋洞。這些結(jié)果證明了世界模型對(duì)于決策的正向作用,為之后世界模型在機(jī)器人等領(lǐng)域的研究提供了重要參考。
方法
WMP 采用經(jīng)典的 RSSM 框架作為世界模型的結(jié)構(gòu),RSSM 包括編碼模塊 encoder,解碼模塊 decoder,以及循環(huán)模塊 recurrent model。encoder 將感知信息以及循環(huán)狀態(tài)編碼為一個(gè)隨機(jī)變量,decoder 通過循環(huán)狀態(tài)和隨機(jī)變量恢復(fù)出原始的感知信息,而循環(huán)模塊則通過循環(huán)狀態(tài)、隨機(jī)變量以及動(dòng)作序列預(yù)測(cè)下一個(gè)循環(huán)狀態(tài)。為了滿足真機(jī)運(yùn)行的算力要求,WMP 將世界模型的運(yùn)行頻率設(shè)定為策略運(yùn)行頻率的 k 分之一。由于世界模型主要處理更高層級(jí)的信息,較低的控制頻率同樣能滿足底層控制的需求,這與人體大腦和小腦展現(xiàn)出的不同的控制頻率情況有一定的相似之處。
一個(gè)訓(xùn)練有素的世界模型的循環(huán)狀態(tài)包含足夠多的信息從而預(yù)測(cè)未來的時(shí)間步,也有助于策略執(zhí)行動(dòng)作。因此在 WMP 框架中,策略會(huì)接受來自世界模型的循環(huán)狀態(tài)作為輸入。并使用強(qiáng)化學(xué)習(xí)算法 PPO 進(jìn)行訓(xùn)練。此外,策略的訓(xùn)練和世界模型的訓(xùn)練使用模擬數(shù)據(jù)同步進(jìn)行,簡(jiǎn)化了特權(quán)學(xué)習(xí)中的兩階段訓(xùn)練。訓(xùn)練后的策略和世界模型可以無需微調(diào)直接遷移到真實(shí)機(jī)器人設(shè)備。
實(shí)驗(yàn)結(jié)果
模擬實(shí)驗(yàn):
WMP 算法使用 Issacgym 模擬器構(gòu)建的 6 種地形上進(jìn)行訓(xùn)練:Slope、Stair、Gap、Climb、Crawl、Tilt。由于 scandot 特權(quán)信息的局限性,使用特權(quán)學(xué)習(xí)訓(xùn)練的 Student baseline 只使用前四個(gè)地形進(jìn)行訓(xùn)練。在模擬器的定量對(duì)比實(shí)驗(yàn)中,WMP 在絕大多數(shù)任務(wù)中獲得了比 Baseline 更高的回報(bào)獎(jiǎng)勵(lì)以及更小的速度追蹤誤差。
真機(jī)實(shí)驗(yàn):
在真機(jī)實(shí)驗(yàn)中,WMP 繼承了模擬器中的優(yōu)秀表現(xiàn),相比 baseline 能以更高的成功率通過更難的地形,并且在室內(nèi)和室外環(huán)境中表現(xiàn)保持一致,進(jìn)一步展現(xiàn)出世界模型優(yōu)秀的泛化能力。
一鏡到底視頻:
驗(yàn)證實(shí)驗(yàn):
使用模擬數(shù)據(jù)訓(xùn)練需要考慮的一個(gè)問題是世界模型對(duì)真實(shí)軌跡預(yù)測(cè)的準(zhǔn)確性如何。驗(yàn)證實(shí)驗(yàn)表明,世界模型對(duì)于真實(shí)軌跡的圖像能給出準(zhǔn)確的預(yù)測(cè),尤其是對(duì)于關(guān)鍵的部分。例如,世界模型對(duì)橋洞障礙物整體形狀的預(yù)測(cè)存在偏差,但對(duì)機(jī)器人需通過的窄縫的位置角度的預(yù)測(cè)十分準(zhǔn)確。這驗(yàn)證了世界模型有利于解構(gòu)和提取不同域中的關(guān)鍵要素,從而有助于模擬到真實(shí)的泛化。
總結(jié)
本研究提出了一種新的運(yùn)動(dòng)控制框架,通過構(gòu)筑的世界模型來處理視覺信息和輔助決策,在四足機(jī)器人運(yùn)動(dòng)控制領(lǐng)域取得了不錯(cuò)的效果。WMP 揭示了世界模型在 Sim2Real 以及機(jī)器人控制領(lǐng)域的巨大潛力,為之后世界模型在現(xiàn)實(shí)世界更廣泛的應(yīng)用提供了樣例和寶貴的經(jīng)驗(yàn)。