🤖 machine learning

WorldDiT: A Unified Diffusion Architecture for World and Action Modeling

O WorldDiT introduz uma arquitetura de transformador de difusão unificada que gera simultaneamente ações contínuas e prevê quadros visuais futuros, alcançando o estado da arte em múltiplos benchmarks de manipulação robótica sem depender de grandes modelos de visão-linguagem pré-treinados.

Autores originais: Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra

Publicado 2026-07-28

📖 7 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Sen Wang, R. Gnana Praveen, Bidhan Roy, Marcos Villagra

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine um mundo onde os robôs não apenas seguem uma lista rígida de instruções, mas realmente "sonham" sobre o que acontecerá a seguir. Esta é a fronteira do aprendizado robótico, um campo que tenta ensinar máquinas a se moverem através do nosso mundo real e caótico. Por muito tempo, os maiores avanços vieram de uma receita específica: pegue um "cérebro" massivo e pré-treinado que já sabe tudo sobre linguagem e imagens (como uma enciclopédia superinteligente que também consegue ver) e anexe um braço de robô a ele. Funciona, mas é como tentar ensinar um filhote de cachorro a buscar algo primeiro treinando-o para ser um professor; é pesado, caro e difícil dizer se o robô está aprendendo a tarefa ou apenas copiando as notas do professor. A grande questão que os cientistas estão fazendo é: Podemos construir um cérebro de robô que seja menor, mais leve e aprenda a se mover e a prever o futuro ao mesmo tempo, sem precisar daquela enciclopédia gigante e pré-carregada?

Apresentamos o WorldDiT, um novo cérebro de robô projetado por pesquisadores da Bagel Labs. Pense no WorldDiT não como um estudante copiando um livro didático, mas como um jovem artista aprendendo a pintar ao observar um filme do seu próprio futuro. Em vez de depender de um gigante "Modelo de Visão-Linguagem" pré-treinado (uma IA massiva que sabe descrever imagens e textos) para lhe dizer o que fazer, o WorldDiT usa um motor único e unificado que faz duas coisas simultaneamente. Primeiro, ele descobre os próximos movimentos do braço do robô. Segundo, ele tenta adivinhar o que a câmera verá alguns segundos à frente. É como um jogador de xadrez que não apenas planeja sua próxima jogada, mas também visualiza o tabuleiro três turnos à frente para garantir que seu plano faça sentido.

Os pesquisadores treinaram este sistema em um mundo simulado chamado LIBERO, onde os robôs têm que realizar tarefas como empilhar blocos ou mover objetos. Eles descobriram que o WorldDiT, que é surpreendentemente pequeno (menos de 400 milhões de parâmetros, comparado aos bilhões usados por outros métodos de ponta), consegue prever a visão futura da câmera e as ações do robô com uma precisão incrível. De fato, quando o testaram, o robô teve sucesso em cerca de 95% de suas tarefas em quatro tipos diferentes de desafios. A parte mais emocionante? Ele fez isso sem o "backbone" pesado e pré-treinado que a maioria dos outros robôs bem-sucedidos utiliza. Provou que você não precisa de um cérebro gigante e pré-carregado para ser um bom robô; você só precisa de um cérebro que aprenda a ver o futuro enquanto se move.

O Robô que "Sonha"

Então, como o WorldDiT realmente funciona? Imagine que você está tentando aprender a fazer malabarismo. A maioria dos robôs de hoje são como estudantes que memorizaram mil vídeos de mestres do malabarismo e estão tentando imitá-los exatamente. O WorldDiT é diferente. É como um estudante de malabarismo que está vendado, mas possui uma bola de cristal.

Aqui está o truque: Cada vez que o WorldDiT planeja um movimento, ele também tenta "sonhar" o que a câmera verá no futuro. Ele observa o estado atual do robô, a instrução que recebeu (como "pegue o bloco vermelho") e os últimos segundos de vídeo. Então, ele se pergunta duas coisas ao mesmo tempo:

"O que o robô deve fazer a seguir?"
"Se eu fizer isso, como será a imagem um momento depois?"

Os pesquisadores chamam isso de transformer de difusão. Se isso soa como jargão de ficção científica, pense nisso como um processo de "redução de ruído" (denoising). Imagine uma foto que está coberta por estática de ruído. O WorldDiT começa com um palpite completamente aleatório e ruidoso sobre qual deve ser o próximo movimento do robô e como deve ser a imagem futura. Então, passo a passo, ele limpa esse ruído, refinando seu palpite até encontrar um caminho claro e suave para o objetivo. Ele faz isso tanto para as ações do robô quanto para as imagens futuras.

A magia acontece porque essas duas tarefas se ajudam. Ao tentar prever a imagem futura, o robô é forçado a entender a física do mundo. Se ele planeja empurrar um bloco, mas seu "sonho" do futuro mostra o bloco flutuando no ar, ele sabe que seu plano está errado. Ele tem que ajustar sua ação até que o "sonho" corresponda à realidade. Isso cria um ciclo de feedback onde o robô aprende as regras do mundo apenas tentando prever o que vem a seguir.

Os Resultados: Cérebro Pequeno, Grande Sucesso

Os pesquisadores testaram o WorldDiT em quatro diferentes suítes de simulação (pense nelas como quatro pistas de obstáculos diferentes: Espacial, Objeto, Objetivo e Longa). Eles o compararam com outros 24 métodos famosos de aprendizado robótico.

Aqui está o ponto crucial: o WorldDiT é minúsculo. Ele possui cerca de 399 milhões de parâmetros (os "neurônios" em seu cérebro). Muitos dos outros robôs de alto desempenho utilizam modelos com 1 bilhão ou até 10 bilhões de parâmetros. Geralmente, cérebros maiores significam melhor desempenho. Mas o WorldDiB quebrou essa regra.

Nestas simulações, o WorldDiT alcançou uma taxa de sucesso de 94,9% em média em todos os quatro cursos.

Espacial: 98,0% de sucesso
Objeto: 97,0% de sucesso
Objetivo: 92,8% de sucesso
Longa: 91,8% de sucesso

Ao plotar esses resultados em um gráfico, o WorldDiT situa-se na "fronteira de Pareto". Em termos simples, isso significa que ele é o robô mais eficiente da sala. Se você quer um robô que seja tão inteligente quanto o WorldDiT, mas que use um cérebro menor, você não encontrará um. Se você quer um robô tão pequeno quanto o WorldDiT, mas mais inteligente, também não encontrará. Todos os outros robôs que tiveram um desempenho superior a 94,9% eram significativamente maiores e mais pesados.

Por Que Isso Importa

O artigo argumenta explicitamente contra a ideia de que você deve usar um enorme Modelo de Visão-Linguagem (VLA) pré-treinado para obter um bom controle robótico. Por um tempo, o campo pensou: "Precisamos de um cérebro gigante para entender o mundo". O WorldDiT diz: "Não necessariamente".

Os pesquisadores mostraram que, ao acoplar a geração de ação (movimento) com a modelagem de mundo (previsão do futuro), uma arquitetura única e unificada pode aprender a controlar um robô de forma muito eficaz. Eles não apenas sugeriram isso; eles mediram isso em milhares de episódios simulados. O robô não apenas "teve sorte"; ele consistentemente superou modelos muito maiores.

No entanto, há uma ressalva. Esses resultados vêm de simulações (jogos de computador onde os robôs vivem). O artigo observa que, embora o desempenho seja forte, este é um ponto de partida para estudos futuros. Ainda não sabemos se este robô que "sonha" funcionará perfeitamente em uma cozinha real e bagunçada, com gravidade real e pisos reais escorregadios. Mas os resultados da simulação são um forte indício de que poderemos construir robôs menores, mais baratos e mais eficientes no futuro — robôs que aprendem imaginando o futuro, em vez de apenas memorizar o passado.

No fim das contas, o WorldDiT sugere que o segredo para um robô inteligente não é apenas ter uma enorme biblioteca de conhecimento; é ter um cérebro que possa projetar o futuro em sua mente para tomar melhores decisões agora.

Resumo Técnico: WorldDiT

Declaração do Problema

Avanços recentes no aprendizado de políticas robóticas frequentemente dependem da anexação da geração de ações a grandes modelos de visão-linguagem pré-treinados (VLMs). Embora essa abordagem aproveite a percepção ampla e a compreensão de linguagem, ela obscurece as contribuições específicas da escala e da arquitetura do modelo. Quando uma política contém vários bilhões de parâmetros, torna-se difícil discernir se o forte desempenho de controle advém do design da ação, do backbone pré-treinado ou de sua combinação. Os autores levantam a questão: Pode uma arquitetura de transformer de difusão unificada combinar a geração de ação contínua com um objetivo auxiliar de predição visual futura para alcançar um forte desempenho de controle sem depender de um grande VLM pré-treinado como o backbone de ação?

Metodologia: WorldDiT

O WorldDiT é uma arquitetura de transformer de difusão unificada projetada para acoplar a geração de ações com a modelagem de mundo visual. Diferente dos modelos de mundo-ação autorregressivos que geram tokens de ação sequencialmente usando um grande VLM, o WorldDiT emprega um único backbone de Diffusion Transformer (DiT) compartilhado para modelar tanto ações robóticas contínuas quanto patches RGB normalizados selecionados de frames de câmera futuros.

Arquitetura e Codificação

O sistema utiliza uma abordagem híbrida de "congelado" e "treinável":

Codificadores: Codificadores visuais "congelados" (MAE para imagens, CLIP para texto) e um codificador de estado do robô treinável (codificando separadamente os estados do braço e da garra) processam as entradas.
Tokenização: O sistema constrói tokens multimodais consistindo de:
- 34 tokens visuais (34 patches das câmeras primária e do pulso).
- 1 token de linguagem (da instrução).
- 1 token de estado (estado do robô).
- Isso resulta em 36 tokens de contexto limpos por slot temporal.
Backbone: O núcleo é um backbone DiT compartilhado com profundidade de 4, tamanho oculto de 1024, 16 cabeças de atenção e 4 tokens de registro. Ele utiliza AdaLN-Zero por token e atenção estritamente bloco-causal e segura para ação, garantindo que as consultas de ação não possam ler alvos RGB ruidosos.

Objetivo de Treinamento: Flow Matching

O WorldDiT é treinado usando flow matching. A janela de treinamento ( $N=10$ ) consiste em:

Contexto ( $C=3$ ): Três passos observados contendo linguagem, RGB multiview e estados.
Alvos ( $H=7$ ): Um chunk de ação de 7 passos e um patch RGB futuro normalizado.

O alvo visual futuro é construído a partir dos frames das câmeras primária e do pulso em um deslocamento temporal de $H$ . Cada frame de 224x224 é dividido em patches de 16x16; 64 patches espaçados uniformemente são retidos de cada câmera, normalizados por sua própria média e variância, e concatenados para formar um alvo de 128 tokens ( $Y_{rgb}$ ).

O modelo recebe tokens de contexto limpos e tokens de alvo corrompidos (patches de ação e RGB) com ruído Gaussiano. Ele prevê a velocidade do fluxo para um caminho reto do ruído até o alvo limpo. A perda total é uma soma ponderada da perda de velocidade da ação e da perda de velocidade do patch RGB:
$L_{total} = w_{action}L_{action}^{flow} + w_{rgb}L_{rgb}^{flow}$
Crucialmente, a predição do patch RGB serve como um sinal de supervisão auxiliar durante o treinamento, mas está ausente na inferência.

Inferência

Na implantação, o sistema opera em um horizonte recuante:

O backbone unificado recebe o histórico (linguagem, observações, estados) até o passo atual.
Nenhum alvo RGB futuro ou rótulos de ação são fornecidos.
O modelo gera um chunk de ação de 7 passos integrando numericamente o campo de velocidade de ação aprendido de um ruído Gaussiano ( $\tau=0$ ) para o estado limpo ( $\tau=1$ ) ao longo de 20 passos de Euler.
Os três primeiros passos da ação do chunk previsto são executados, o estado é observado e a janela é deslocada para replanejamento.

Principais Contribuições

Arquitetura Unificada: Um único transformer de difusão que aprende conjuntamente a geração de ação contínua e a predição de patch RGB futuro normalizado sem um grande backbone de VLM pré-treinado.
Design de Treinamento/Implantação: Uma configuração inovadora onde a predição de patch RGB fornece supervisão durante o treinamento, mas é removida do grafo de inferência, permitindo um caminho de implantação compacto, focado apenas em ação.
Baseline de Sub-Bilhão de Parâmetros: Uma configuração com aproximadamente 399 milhões de parâmetros totais (135 milhões treináveis) que estabelece um forte baseline para futuros estudos de escala.

Resultados

Os autores avaliaram o WorldDiT no benchmark de manipulação LIBERO, especificamente através de quatro suítes: Spatial, Object, Goal e Long (LIBERO-10).

Desempenho: O WorldDiT alcançou as seguintes taxas médias de sucesso:
- Spatial: 98.0%
- Object: 97.0%
- Goal: 92.8%
- Long: 91.8%
- Média Geral: 94.9%
Fronteira de Pareto: Quando plotado contra os parâmetros totais do modelo instanciado, o WorldDiT situa-se na fronteira de Pareto relatada. Entre os métodos que reportam resultados em todas as quatro suítes, cada método com uma taxa de sucesso média maior utiliza mais parâmetros totais. Inversamente, métodos com parâmetros similares ou menores reportam taxas de sucesso mais baixas.
Comparação: O WorldDi-T supera outros métodos que não utilizam grandes backbones de VLM pré-treinados (ex: DreamVLA, FlowVLA, Octo) e compete favoravelmente com métodos que utilizam, apesar de possuir significativamente menos parâmetros.

Significância e Alegações

O artigo afirma que o WorldDiT demonstra que um forte desempenho em benchmarks é alcançável sem colocar um grande modelo de visão-linguagem pré-treinado no backbone de ação. Ao acoplar a geração de ação com um objetivo de predição visual futura, a arquitetura alcança um tradeoff de eficiência de parâmetros elevado.

Os autores posicionam o WorldDiT como um "forte baseline de sub-bilhão de parâmetros" para futuros estudos de escala. Eles sugerem que este resultado apoia a modelagem unificada de mundo-e-ação como uma base compacta para escalonamento através de capacidade de modelo, diversidade de dados e configurações de implantação. O artigo nota modestamente que, embora os resultados caracterizem o tradeoff abaixo de um bilhão de parâmetros, trabalhos futuros são necessários para determinar se esse tradeoff persiste em escalas maiores e como o objetivo de RGB futuro normalizado especificamente molda o controle. Os autores também destacam o potencial de decompor a formulação de flow-matching compartilhada em especialistas treinados independentemente para hardware heterogêneo.