Resumo Técnico: Cosmos Policy

Declaração do Problema

Avanços recentes em modelos de geração de vídeo demonstraram uma capacidade impressionante de capturar interações físicas complexas, causalidade temporal e evolução de cena. No entanto, adaptar esses modelos para o controle robótico tem exigido tradicionalmente pipelines de treinamento complexos de múltiplos estágios e a introdução de novos componentes arquiteturais (por exemplo, difusores de ação separados ou modelos de dinâmica inversa). Os modelos de vídeo-ação unificados existentes muitas vezes falham em aproveitar os poderosos priors espaço-temporais de modelos de vídeo pré-treinados em larga escala devido a designs customizados treinados do zero. Além disso, embora os modelos Vision-Language-Action (VLA) tenham se tornado backbones populares, eles são primariamente treinados em pares estáticos de imagem-texto, o que pode limitar sua compreensão da dinâmica temporal e da física implícita em comparação com modelos de vídeo.

O desafio central abordado por este trabalho é como adaptar efetivamente um modelo de vídeo fundacional grande e pré-treinado em uma política robótica através de um processo simples de estágio único, sem modificações arquiteturais, enquanto simultaneamente permite que a política realize planejamento baseado em modelo usando modelos de mundo aprendidos e funções de valor.

Metodologia

Arquitetura Cosmos Policy

Os autores apresentam o Cosmos Policy, uma política robótica ajustada (fine-tuned) a partir do modelo de fundação de vídeo NVIDIA Cosmos-Predict2-2B. A abordagem distingue-se pela ausência de mudanças arquiteturais; em vez disso, ela aproveita o processo de difusão latente nativo do modelo para lidar com entradas e saídas multimodais.

1. Injeção de Quadro Latente (Latent Frame Injection):
Para incorporar modalidades não suportadas originalmente pelo modelo de vídeo (propriocepção robótica, ações, valores de estado e câmeras de múltiplas visões), os autores propõem a Injeção de Quadro Latente.

• Mecanismo: Modalidades adicionais são codificadas como novos quadros latentes e intercaladas diretamente na sequência de difusão latente do modelo de vídeo.
• Estrutura da Sequência: Para uma configuração robótica com múltiplas câmeras e propriocepção, a sequência latente é estruturada como $(s, a, s', V(s'))$, onde:
  • $s$: Estado atual (propriocepção do robô + imagens de câmera atuais).
  • $a$: Action chunk (codificado como um quadro latente).
  • $s'$: Estado futuro (propriocepção futura + imagens de câmera futuras).
  • $V(s')$: Valor do estado futuro (recompensa cumulativa esperada, codificada como um quadro latente).
• Flexibilidade: Este método permite configurações robóticas arbitrárias (por exemplo, números variados de câmeras) simplesmente adicionando ou removendo quadros latentes na sequência, sem alterar os pesos ou a estrutura do modelo.

2. Objetivos de Treinamento Conjunto:
O modelo é ajustado em dados de demonstração robótica usando um objetivo de estágio único que aprende conjuntamente três funções:

• Política (Policy): Prediz ações, estados futuros e valores condicionados ao estado atual ($p(a, s', V(s') | s)$).
• Modelo de Mundo (World Model): Prediz estados futuros e valores condicionados ao estado e ação ($p(s', V(s') | s, a)$).
• Função de Valor (Value Function): Prediz valores futuros condicionados ao estado, ação e estado futuro ($p(V(s') | s, a, s')$)$ .
  Os dados de treinamento consistem em uma mistura de demonstrações de especialistas (50% do lote) e rollouts da política (50% do lote). Os dados de rollout são usados para treinar o modelo de mundo e a função de valor, permitindo que o sistema aprenda tanto com trajetórias bem-sucedidas quanto com falhas.

3. Planejamento com Busca Baseada em Modelo:
O Cosmos Policy suporta dois modos de implantação:

• Política Direta (Direct Policy): Utiliza decodificação paralela para gerar ações diretamente a partir do estado atual.
• Planejamento Baseado em Modelo (Model-Based Planning): Utiliza o modelo de mundo e a função de valor aprendidos para buscar ações ótimas.
  • Processo: O sistema amostra múltiplas propostas de ação, utiliza o "modelo de planejamento" (um checkpoint refinado treinado intensamente em dados de rollout) para prever os estados futuros e valores resultantes, e seleciona a ação que leva ao maior valor previsto.
  • Estratégia de Ensemble: Para lidar com distribuições multimodais e variância, o sistema realiza o ensemble das previsões consultando o modelo de mundo e a função de valor múltiplas vezes por ação, agregando os resultados via "média majoritária" para determinar a probabilidade de sucesso.

Contribuições Principais

1. Adaptação Simples de Modelos de Vídeo: O artigo demonstra que um modelo de vídeo pré-treinado de grande escala pode ser adaptado para uma política robótica de alto desempenho através de um único estágio de pós-treinamento, sem modificações arquiteturais, utilizando a injeção de quadro latente para codificar ações e valores.
2. Política, Modelo de Mundo e Função de Valor Unificados: A abordagem unifica o aprendizado de política, modelagem de mundo e estimativa de valor dentro de um único framework de difusão latente, aproveitando os priors pré-treinados do modelo para capturar distribuições de ação complexas e dinâmicas temporais.
3. Planejamento via Refinamento de Rollout: Os autores mostram que, ao coletar dados de rollout da política e realizar o ajuste fino dos componentes de valor e de mundo, o sistema pode realizar um planejamento baseado em modelo eficaz (best-of-N sampling) para melhorar as taxas de sucesso.
4. Desempenho de Estado da Arte: O método atinge novos marcos em simulação e manipulação no mundo real, superando políticas treinadas do zero, políticas baseadas em vídeo e VLAs ajustados.

Resultados Experimentais

Os autores avaliaram o Cosmos Policy em três benchmarks:

• LIBERO (Simulação): Alcançou uma taxa de sucesso média de 98,5%, superando fortes baselines incluindo Diffusion Policy, Dita, $\pi_0$, OpenVLA-OFT e CogVLA.
• RoboCasa (Simulação): Alcançou uma taxa de sucesso média de 67,1% usando apenas 50 demonstrações humanas por tarefa (significativamente menos que os mais de 300 usados por muitos competidores), superando métodos como GR00T-N1.5 e Video Policy.
• Tarefas ALOHA (Mundo Real): Em tarefas desafiadoras de manipulação bimanual (ex: dobrar camisas, manipular sacos tipo ziploc), o Cosmos Policy alcançou a maior pontuação média (93,6%).
  • Notavelmente, superou VLAs ajustados ($\pi_0.5$, OpenVLA-OFT+) em tarefas que exigem preensões de alta precisão e manipulação de alta multimodalidade de ação.
  • Impacto do Planejamento: Quando potencializado com planejamento baseado em modelo usando dados de rollout, a política teve um aumento médio de 12,5% nas pontuações de conclusão de tarefa nas tarefas mais difíceis do ALOHA.

Estudos de Ablação:

• Remover alvos de aprendizado auxiliares (predizer $s'$ e $V$ junto com $a$) resultou em uma queda de 1,5% na taxa de sucesso.
• Treinar o modelo do zero (sem os pesos de vídeo pré-treinados) resultou em uma queda de 3,9%, destacando a importância dos priors do modelo de vídeo.
• O planejamento baseado em modelo usando uma formulação $V(s')$ (que requer um modelo de mundo) superou uma variante model-free $Q(s, a)$, sugerindo o valor da dinâmica de ambiente aprendida.

Significância e Alegações

O artigo afirma que o Cosmos Policy representa um passo significativo à frente no aprendizado robótico ao demonstrar que os priors espaço-temporais de modelos de geração de vídeo de larga escala são uma fundação superior para controle de baixo nível em comparação com modelos VLA baseados em imagens estáticas.

As principais alegações incluem:

• Simplicidade e Eficiência: A capacidade de criar uma política robótica poderosa sem mudanças arquiteturais ou pipelines de treinamento de múltiplos estágios.
• Eficiência de Dados: O método alcança alto desempenho com menos demonstrações (50 vs. 300+) comparado a outros métodos de estado da arte, provavelmente devido à forte inicialização do modelo de fundação de vídeo.
• Capacidade de Planejamento: A arquitetura unificada suporta naturalmente o planejamento baseado em modelo, permitindo que o robô "imagine" estados futuros e valores para selecionar melhores ações, uma capacidade que melhora o desempenho em tarefas de alta precisão e multimodalidade onde políticas diretas têm dificuldade.
• Generalização: A abordagem mostra forte generalização tanto em cenários de distribuição interna (in-distribution) quanto de distribuição externa (out-of-distribution) em ambientes de simulação e mundo real.

Os autores reconhecem limitações, especificamente a velocidade de inferência do planejamento baseado em modelo (aproximadamente 5 segundos por action chunk) e a necessidade atual de dados substanciais de rollout para refinar o modelo de mundo, identificando estes como áreas para melhorias futuras.