Flex-$\pi$: A Multi-Stream World-Action Model with Compute Flexibility

Problema

Modelos de mundo-ação (WAMs) surgiram como uma alternativa poderosa aos modelos de visão-linguagem-ação (VLAs) ao prever conjuntamente observações e ações futuras, aprendendo assim representações espaço-temporais robustas e herdando priors de geração de vídeo em larga escala. No entanto, os WAMs generalistas atuais preveem quase exclusivamente latentes de imagens RGB futuras. Esses latentes, treinados puramente para reconstrução de pixels, capturam detalhes de aparência, mas frequentemente falham em se alinhar com a geometria 3D explícita e a semântica de nível de objeto necessária para a manipulação robótica precisa.

Supervisionar WAMs diretamente em geometria e semântica normalmente acarreta três custos significativos:

1. Sensores Adicionais: Exigindo novos hardwares de modalidades (ex: sensores de profundidade).
2. Novo Pré-treinamento: Necessitando o treinamento de novos priors visuais do zero.
3. Latência de Inferência: Desacelerando a implantação devido ao processamento multimodal complexo.

O artigo questiona: Como podemos amplificar os pontos fortes dos WAMs (previsão conjunta e fortes priors visuais) para incluir geometria 3D e semântica sem esses sacrifícios?

Metodologia: FLEX-π

Os autores introduzem o FLEX-π, um WAM de 6 bilhões de parâmetros que prevê observações RGB futuras, pointmaps 3D e semântica DINO centrada em objetos simultaneamente. A inovação central é alcançar este "almoço grátis" de fundamentação multimodal sem novos sensores ou pré-treinamento.

1. Espaço Latente Compartilhado via VAE Congelado

O FLEX-π utiliza um único VAE (Variational Autoencoder) congelado de um modelo de geração de vídeo pré-treinado (Wan-2.2).

• Codificação RGB: O VAE codifica imagens RGB como de costume.
• Codificação de Pointmap 3D: Surpreendentemente, os autores descobriram que este mesmo VAE congelado, treinado apenas em pixels RGB, pode codificar e decodificar pointmaps 3D quase sem perdas. Isso permite que a geometria 3D seja projetada no exato mesmo espaço latente que o RGB sem treinar um novo codificador.
• Codificação Semântica: Semânticas centradas em objetos são derivadas usando um codificador DINOv3 congelado.

2. Backbone de Mistura de Transformers (MoT) de Múltiplos Fluxos

O modelo incorpora três fluxos visuais (RGB, Pointmap, DINO) em um fluxo de tokens de espaço latente compartilhado.

• Arquitetura: Um backbone de Mistura de Transformers (5 bilhões de parâmetros) processa esses fluxos. Ele utiliza um backbone de transformer visual compartilhado, mas mantém redes de projeção feedforward independentes para cada modalidade.
• Especialista de Ação: Um Especialista de Ação separado e mais estreito (Transformer de 1 bilhão de parâmetros) realiza atenção cruzada (cross-attention) aos fluxos visuais para gerar blocos de ação (action chunks).
• Geração Conjunta Causal: Tokens de ação atendem às observações atuais e realizam atenção cruzada com os tokens visuais futuros ativos. Isso permite que a geração de ação se beneficie da representação evolutiva do futuro pelo modelo.

3. Treinamento Flexível via Dropout de Fluxo e Forçamento de Cross-Modalidade

Para permitir que um único checkpoint suporte qualquer combinação de entradas e saídas na inferência, o FLEX-π emprega duas estratégias de treinamento principais:

• Dropout de Fluxo Visual: Durante o treinamento, os fluxos de entrada são descartados aleatoriamente (com probabilidade de 0,5), garantindo que o modelo possa operar com observações parciais (ex: apenas RGB).
• Forçamento de Cross-Modalidade: Crucialmente, mesmo que um fluxo seja descartado da entrada, o modelo ainda é forçado a realizar o denoising e prever esse fluxo na saída. Por exemplo, se o pointmap 3D estiver ausente da entrada, o modelo deve sintetizar o pointmap futuro com base nos inputs RGB e DINO. Isso força o modelo a internalizar uma representação unificada onde aparência, geometria e semântica são mutuamente preditivas.

4. Flexibilidade de Inferência

Na implantação, os usuários podem trocar entre latência e performance selecionando quais fluxos gerar:

• Modo Apenas-Ação: O modelo prevê apenas ações (inferência mais rápida).
• Geração Conjunta Total: O modelo prevê ações mais RGB futuros, pointmaps e recursos DINO (performance mais alta).
• Geração Parcial: Qualquer subconjunto de fluxos pode ser gerado.

Principais Contribuições

1. Fundamentação Geométrica e Semântica Gratuita: Demonstrar que um VAE de vídeo pré-treinado em RGB pode codificar pointmaps 3D sem perdas, permitindo que WAMs sejam supervisionados em geometria 3D e semântica DINO sem novos sensores ou pré-treinamento.
2. Flexibilidade de Computação: Um único checkpoint treinado suporta qualquer subconjunto de fluxos de entrada/saída, permitindo que praticantes escolham seu ponto de operação na fronteira velocidade-performance no momento da implantação.
3. Forçamento de Cross-Modalidade: Um regime de treinamento que força o modelo a prever modalidades ausentes a partir das observadas, melhorando significativamente a representação interna e o desempenho da política.

Resultados Experimentais

Os autores avaliam o FLEX-π em benchmarks de simulação (RoboTwin, LIBERO, LIBERO-Plus) e tarefas de manipulação bimanual no mundo real com um robô YAM.

• Desempenho no Mundo Real: Em cinco tarefas manuais precisas e destras (incluindo "Self-Repair Gripper" e "Soft-Bag Zipping"), o FLEX-π supera fortemente os baselines (π0.5, ManiFlow, Fast-WAM) em 2 a 7 vezes nas taxas de sucesso.
  • Notavelmente, a variante apenas-ação do FLEX-π (que não gera futuros visuais na inferência) já supera todos os baselines e roda mais rápido que o π0.5.
  • A geração conjunta total aumenta ainda mais o desempenho, alcançando até 42,7% de melhoria absoluta na conclusão de tarefas em tarefas difíceis comparado ao baseline mais forte.
• Eficiência de Dados: O FLEX-π é significativamente mais eficiente em dados. No RoboTwin, ele supera os baselines por 1,9 a 4,5 vezes em regimes de baixos dados (50 demonstrações). Mesmo quando treinado com apenas 50% dos dados reais, o FLEX-π supera os baselines treinados com o conjunto de dados completo.
• Generalização (OOD): O FLEX-π demonstra robustez a mudanças fora da distribuição (objetos não vistos, desordem, distratores). Enquanto os baselines como ManiFlow e π0.5 sofrem quedas significativas de desempenho (25–26%) sob mudança de distribuição, o FLEX-π cai menos de 5%.
• Ablações:
  • Remover inputs de DINO ou Pointmap durante o treinamento reduz as taxas de sucesso significativamente (em 6,8% e 20%, respectivamente).
  • Remover o Forçamento de Cross-Modalidade prejudica as taxas de sucesso em 21%, confirmando que forçar o modelo a sintetizar modalidades ausentes é crítico para aprender uma representação compartilhada forte.

Significância e Alegações

O artigo afirma que o FLEX-π demonstra que um WAM pode ganhar fundamentação semântica e geométrica essencialmente de graça. Ao aproveitar o espaço latente de um modelo de vídeo pré-treinado e empregar o forçamento de cross-modalidade, o modelo alcança:

• Melhor Desempenho de Política: Taxas de sucesso significativamente maiores em tarefas complexas, ricas em contato e de longo horizonte.
• Sem Sobrecarga de Hardware: Não há necessidade de sensores adicionais além de RGB.
• Sem Sobrecarga de Pré-treinamento: Não há necessidade de treinar novos priors visuais.
• Flexibilidade de Implantação: A capacidade de rodar em um modo rápido de apenas-ação ou em um modo de geração conjunta de alta performance, usando o mesmo checkpoint.

Os autores concluem que, embora o FLEX-π seja altamente eficaz, ele ainda requer mais épocas de ajuste fino para convergir do que modelos mais simples, e seu modo de geração conjunta total é mais lento que VLAs com parâmetros comparáveis. Eles sugerem que melhorias futuras podem vir da combinação do FLEX-π com backbones de VLM mais fortes e acesso a conjuntos de dados robóticos ainda maiores.