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

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

Licenças

Acesso aberto
Artigo: CC BY 4.0

Tudo o que segue se baseia no artigo completo sob a licença CC BY 4.0; sua figura é mostrada com crédito aos autores.

Em resumo

Um único Diffusion Transformer com passos de tempo de ruído separados para ações e imagens futuras atua como política, modelo de dinâmica ou preditor de vídeo, e pode ser treinado com vídeo sem ações.

Por que importa

Um único modelo pode ser treinado tanto com dados robóticos rotulados com ações quanto com vídeo sem ações, que é muito mais fácil de obter. Em testes com robô real em cinco tarefas com uma montagem Franka no estilo DROID, o UWM pré-treinado teve o maior sucesso em todas as tarefas, e o cotreinamento com vídeos extras o elevou ainda mais, inclusive sob distrações visuais.

Problema

O aprendizado por imitação para grandes modelos robóticos depende de demonstrações especialistas de alta qualidade, que são caras e lentas de coletar por teleoperação. Vídeos que mostram a dinâmica do mundo real são abundantes, mas a maioria dos métodos atuais de aprendizado por imitação não consegue usá-los porque não têm rótulos de ação, e a clonagem de comportamento não captura explicitamente a dinâmica temporal presente nas trajetórias.

Solução proposta

O UWM acopla um processo de difusão de ações e um processo de difusão de imagens futuras em um único transformador, condicionado pela observação atual, com um passo de tempo de difusão independente para cada modalidade. Definir o passo de tempo de uma modalidade como T a marginaliza e defini-lo como 0 condiciona nela, de modo que o mesmo modelo amostra uma política p(a|o), dinâmica direta p(o′|o,a), dinâmica inversa p(a|o,o′) ou predição de vídeo p(o′|o). Vídeo sem ações é usado no treinamento fixando o passo de tempo das ações em T e preenchendo as ações ausentes com ruído gaussiano.

Principais ideias

Passos de tempo de difusão independentes por modalidade funcionam como mascaramento parcial, de modo que uma única rede cobre as distribuições marginais e condicionais sobre ações e quadros futuros.
O que gera o ganho é prever quadros futuros, e não apenas reconstruir o quadro atual: 0,86/0,76 de sucesso em Stack-Bowls/Block-Cabinet contra 0,70/0,66 com reconstrução do quadro atual e 0,48/0,60 com Diffusion Policy.
Os register tokens (8 por padrão) ajudam a trocar informação entre tokens de ação e de imagem, e o condicionamento AdaLN superou a atenção cruzada na ablação no LIBERO.
UWM e Diffusion Policy têm desempenho parecido quando treinados do zero; o UWM ganha mais com o pré-treinamento.
O cotreinamento com vídeos de robôs do próprio domínio ajudou mais que o cotreinamento com vídeos de atividade humana da Internet (Kinetics-400 e Something-Something v2), embora ambos tenham superado os dados de robô sozinhos.

Principais resultados

Robô real, pré-treinado com 2.000 trajetórias do DROID, sucesso dentro da distribuição: UWM 0,86, 0,76, 0,78, 0,82 e 0,60 em Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel e Rice-Cooker, contra 0,48, 0,60, 0,52, 0,64, 0,35 com Diffusion Policy e 0,66, 0,66, 0,60, 0,66, 0,40 com GR1.
O cotreinamento com mais 2.000 trajetórias do DROID usadas como vídeo sem ações elevou o sucesso do UWM dentro da distribuição para 0,92, 0,84, 0,86, 0,86 e 0,65, e o sucesso fora da distribuição (quatro primeiras tarefas) de 0,76/0,60/0,78/0,64 para 0,84/0,72/0,84/0,76.
Cenas OOD categorizadas (iluminação, fundo, desordem): UWM cotreinado 21/30 e 15/30 em Stack-Bowls e Block-Cabinet, contra 15/30 e 8/30 do UWM pré-treinado e 12/30 e 6/30 da Diffusion Policy.
LIBERO (pré-treino no LIBERO-90, ajuste fino em cinco tarefas do LIBERO-10 com mudanças de distribuição): sucesso médio de 0,79 ± 0,11 para o UWM contra 0,71 ± 0,12 para DP, 0,57 ± 0,19 para PAD e 0,58 ± 0,14 para GR1.
Modo de dinâmica inversa rastreando 50 trajetórias do LIBERO dentro do comprimento da trajetória: 0,73 e 0,45 de sucesso em Book-Caddy e Soup-Cheese, contra 0,37 e 0,33 no modo política.
Treinar um UWM no DROID por 100K passos leva 24 horas em 4 GPUs NVIDIA A100.

Limitações

O modelo ainda não aprende com vídeos humanos de grande escala, então a lacuna de incorporação entre humano e robô não é tratada.
As predições de dinâmica direta podem conter artefatos, o que pode reduzir sua utilidade para planejamento.
Os autores esperam ganhos adicionais com uma predição de vídeo mais densa, que esta versão não usa.
Os ganhos sobre as referências no cenário fora da distribuição do LIBERO são muito menores que no robô real; os autores sugerem que a dinâmica da simulação é mais simples.
Rice-Cooker foi avaliada apenas em 20 inicializações próximas da distribuição e não em cenários OOD, e todos os testes com robô real usam uma única montagem Franka.

Impacto na área

O artigo tem 244 citações. Ele afirma que vídeos e código estão disponíveis na página do projeto, que aponta para o repositório de código WEIRDLabUW/unified-world-model no GitHub.

Resumo

O aprendizado por imitação surgiu como uma abordagem promissora para construir robôs generalistas. No entanto, escalar o aprendizado por imitação para grandes modelos de fundação robóticos continua difícil devido à sua dependência de demonstrações especialistas de alta qualidade. Enquanto isso, grandes quantidades de dados de vídeo que retratam uma ampla gama de ambientes e comportamentos diversos estão prontamente disponíveis. Esses dados são uma fonte rica de informação sobre a dinâmica do mundo real e as interações entre agente e ambiente. Aproveitar esses dados diretamente para o aprendizado por imitação, porém, tem se mostrado difícil pela falta das anotações de ações exigidas pela maioria dos métodos atuais. Neste trabalho, apresentamos os Unified World Models (UWM), um framework que permite aproveitar tanto dados de vídeo quanto de ações para o aprendizado de políticas. Especificamente, um UWM integra um processo de difusão de ações e um processo de difusão de vídeo em uma arquitetura de transformador unificada, na qual passos de tempo de difusão independentes governam cada modalidade. Mostramos que, simplesmente controlando cada passo de tempo de difusão, o UWM pode representar de forma flexível uma política, uma dinâmica direta, uma dinâmica inversa e um gerador de vídeo. Por meio de experimentos simulados e no mundo real, mostramos que: (1) o UWM permite um pré-treinamento eficaz em conjuntos de dados robóticos multitarefa de grande escala com predições tanto de dinâmica quanto de ações, resultando em políticas mais generalizáveis e robustas que o aprendizado por imitação, (2) o UWM facilita naturalmente o aprendizado a partir de dados de vídeo sem ações por meio do controle independente dos passos de tempo de difusão específicos de cada modalidade, melhorando ainda mais o desempenho das políticas ajustadas. Nossos resultados sugerem que o UWM representa um passo promissor para aproveitar conjuntos de dados grandes e heterogêneos no aprendizado robótico escalável, e oferece uma unificação simples entre os paradigmas, muitas vezes separados, do aprendizado por imitação e da modelagem de mundo. Vídeos e código estão disponíveis em https://weirdlabuw.github.io/uwm/

Open Source Robotics