Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets

Autori: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · RSS 2025

Licenze

• Open access
• Articolo: CC BY 4.0

Tutto ciò che segue si basa sull'articolo completo secondo la sua licenza CC BY 4.0; la figura è riportata citando gli autori.

In breve

Un unico diffusion transformer con passi temporali di rumore separati per le azioni e le immagini future funge da policy, modello di dinamica o predittore video, e può essere addestrato su video senza azioni.

Perché è importante

Un solo modello può essere addestrato sia su dati robotici etichettati con le azioni sia su video senza azioni, molto più facili da ottenere. Nei test su robot reale su cinque compiti con una configurazione Franka in stile DROID, UWM preaddestrato ha ottenuto il successo più alto in ogni compito, e il co-addestramento su video aggiuntivi lo ha aumentato ulteriormente, anche in presenza di distrazioni visive.

Problema

L'apprendimento per imitazione per grandi modelli robotici dipende da dimostrazioni esperte di alta qualità, costose e lente da raccogliere tramite teleoperazione. I video che mostrano la dinamica del mondo reale abbondano, ma la maggior parte degli attuali metodi di apprendimento per imitazione non può usarli perché privi di etichette delle azioni, e il behavior cloning non cattura esplicitamente la dinamica temporale presente nelle traiettorie.

Soluzione proposta

UWM accoppia un processo di diffusione delle azioni e un processo di diffusione delle immagini future in un unico Transformer, condizionato dall'osservazione corrente, con un passo temporale di diffusione indipendente per ciascuna modalità. Impostare il passo temporale di una modalità a T la marginalizza e impostarlo a 0 condiziona su di essa, così lo stesso modello campiona una policy p(a|o), una dinamica diretta p(o′|o,a), una dinamica inversa p(a|o,o′) o una predizione video p(o′|o). I video senza azioni sono usati in addestramento fissando il passo temporale delle azioni a T e riempiendo le azioni mancanti con rumore gaussiano.

Idee chiave

• I passi temporali di diffusione indipendenti per modalità agiscono come un mascheramento parziale, così un'unica rete copre le distribuzioni marginali e condizionate sulle azioni e sui fotogrammi futuri.
• A produrre il guadagno è la predizione dei fotogrammi futuri, non la semplice ricostruzione del fotogramma corrente: 0,86/0,76 di successo su Stack-Bowls/Block-Cabinet contro 0,70/0,66 con la ricostruzione del fotogramma corrente e 0,48/0,60 con Diffusion Policy.
• I register token (8 per impostazione predefinita) aiutano a scambiare informazioni tra token di azione e di immagine, e il condizionamento AdaLN ha superato la cross-attention nell'ablazione su LIBERO.
• UWM e Diffusion Policy rendono più o meno allo stesso modo quando sono addestrati da zero; UWM guadagna di più dal preaddestramento.
• Il co-addestramento su video robotici in-domain ha aiutato più del co-addestramento su video di attività umane da Internet (Kinetics-400 e Something-Something v2), anche se entrambi hanno battuto i soli dati robotici.

Risultati principali

• Robot reale, preaddestrato su 2.000 traiettorie DROID, successo in distribuzione: UWM 0,86, 0,76, 0,78, 0,82 e 0,60 su Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel e Rice-Cooker, contro 0,48, 0,60, 0,52, 0,64, 0,35 per Diffusion Policy e 0,66, 0,66, 0,60, 0,66, 0,40 per GR1.
• Il co-addestramento su altre 2.000 traiettorie DROID usate come video senza azioni ha portato il successo in distribuzione di UWM a 0,92, 0,84, 0,86, 0,86 e 0,65, e il successo fuori distribuzione (primi quattro compiti) da 0,76/0,60/0,78/0,64 a 0,84/0,72/0,84/0,76.
• Scene OOD categorizzate (illuminazione, sfondo, disordine): UWM co-addestrato 21/30 e 15/30 su Stack-Bowls e Block-Cabinet, contro 15/30 e 8/30 per UWM preaddestrato e 12/30 e 6/30 per Diffusion Policy.
• LIBERO (preaddestramento su LIBERO-90, fine-tuning su cinque compiti di LIBERO-10 con spostamenti di distribuzione): successo medio di 0,79 ± 0,11 per UWM contro 0,71 ± 0,12 per DP, 0,57 ± 0,19 per PAD e 0,58 ± 0,14 per GR1.
• Modalità di dinamica inversa che insegue 50 traiettorie LIBERO entro la lunghezza della traiettoria: 0,73 e 0,45 di successo su Book-Caddy e Soup-Cheese, contro 0,37 e 0,33 per la modalità policy.
• Addestrare un UWM su DROID per 100K passi richiede 24 ore su 4 GPU NVIDIA A100.

Limiti

• Il modello non apprende ancora da video umani su larga scala, quindi il divario di embodiment tra umano e robot non viene affrontato.
• Le predizioni della dinamica diretta possono contenere artefatti, il che può ridurne l'utilità per la pianificazione.
• Gli autori si aspettano ulteriori guadagni da una predizione video più densa, che questa versione non usa.
• I guadagni rispetto alle baseline nello scenario fuori distribuzione di LIBERO sono molto più piccoli che sul robot reale; gli autori suggeriscono che la dinamica della simulazione è più semplice.
• Rice-Cooker è stato valutato solo su 20 inizializzazioni vicine alla distribuzione e non in scenari OOD, e tutti i test su robot reale usano un'unica configurazione Franka.

Impatto sul settore

L'articolo ha 244 citazioni. Afferma che video e codice sono disponibili sulla pagina del progetto, che rimanda al repository del codice WEIRDLabUW/unified-world-model su GitHub.

Abstract

L'apprendimento per imitazione si è affermato come un approccio promettente per costruire robot generalisti. Tuttavia, scalare l'apprendimento per imitazione per grandi foundation model robotici resta difficile a causa della sua dipendenza da dimostrazioni esperte di alta qualità. Nel frattempo, sono facilmente disponibili grandi quantità di dati video che ritraggono un'ampia gamma di ambienti e comportamenti diversi. Questi dati sono una ricca fonte di informazioni sulla dinamica del mondo reale e sulle interazioni tra agente e ambiente. Sfruttare questi dati direttamente per l'apprendimento per imitazione si è però rivelato difficile per la mancanza delle annotazioni delle azioni richieste dalla maggior parte dei metodi attuali. In questo lavoro presentiamo Unified World Models (UWM), un framework che consente di sfruttare sia dati video sia dati di azione per l'apprendimento delle policy. In particolare, un UWM integra un processo di diffusione delle azioni e un processo di diffusione del video all'interno di un'architettura Transformer unificata, in cui passi temporali di diffusione indipendenti governano ciascuna modalità. Mostriamo che, semplicemente controllando ciascun passo temporale di diffusione, UWM può rappresentare in modo flessibile una policy, una dinamica diretta, una dinamica inversa e un generatore video. Attraverso esperimenti simulati e nel mondo reale, mostriamo che: (1) UWM consente un preaddestramento efficace su dataset robotici multi-task su larga scala con predizioni sia della dinamica sia delle azioni, ottenendo policy più generalizzabili e robuste rispetto all'apprendimento per imitazione, (2) UWM facilita in modo naturale l'apprendimento da dati video senza azioni attraverso il controllo indipendente dei passi temporali di diffusione specifici di ciascuna modalità, migliorando ulteriormente le prestazioni delle policy sottoposte a fine-tuning. I nostri risultati suggeriscono che UWM rappresenta un passo promettente verso lo sfruttamento di dataset grandi ed eterogenei per un apprendimento robotico scalabile, e offre una semplice unificazione tra i paradigmi, spesso separati, dell'apprendimento per imitazione e della modellazione del mondo. Video e codice sono disponibili su https://weirdlabuw.github.io/uwm/

Link

• arXiv 2504.02792: https://arxiv.org/abs/2504.02792
• DOI https://doi.org/10.15607/rss.2025.xxi.015: https://doi.org/https://doi.org/10.15607/rss.2025.xxi.015
• PDF: https://arxiv.org/pdf/2504.02792v1
• Codice: https://github.com/WEIRDLabUW/unified-world-model
• Pagina del progetto: https://weirdlabuw.github.io/uwm/