PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations

1. Formulação do Problema

Os atuais modelos Visão-Linguagem-Ação (VLA) avançaram o controle robótico ao aproveitar Modelos Visão-Linguagem (VLMs) como espinhas dorsais perceptivas, tipicamente utilizando uma arquitetura de sistema duplo (Sistema 2 para raciocínio, Sistema 1 para controle). No entanto, as abordagens existentes sofrem de uma limitação fundamental:

• Estático vs. Temporal: Elas enquadram predominantemente o pré-treinamento como clonagem de comportamento supervisionada (BC), focando em correlações semânticas visuais-linguísticas estáticas. Elas carecem de uma compreensão intrínseca da alcançabilidade do objetivo—a probabilidade quantitativa de que um par estado-ação atual levará a um objetivo especificado por linguagem ao longo do tempo.
• Gargalos de Estimativa de Valor: Tentativas de incorporar aprendizado de valor frequentemente dependem de redes de valor separadas, exigem anotações manuais de recompensa ou utilizam aprendizado de Diferença Temporal (TD), que sofre de erros de bootstrap e alta variância em tarefas de longo horizonte.
• Lacuna de Generalização: Sem consciência temporal explícita, esses modelos lutam com planejamento de longo horizonte, instruções novas em zero-shot e recuperação de intervenções humanas, pois não conseguem distinguir entre estados "fáceis-mas-errôneos" e "difíceis-mas-alcançáveis".

2. Metodologia: PRTS

Os autores propõem o PRTS, um modelo fundamental VLA que reformula o pré-treinamento através de Aprendizado por Reforço Condicionado a Objetivos (GCRL) utilizando Aprendizado por Reforço Contrastivo (CRL).

A. Conceito Central: RL Contrastivo para VLA

Em vez de prever ações diretamente ou usar uma rede de valor separada, o PRTS aprende um espaço de incorporação unificado onde o produto interno de uma incorporação estado-ação ($\phi(s, a)$) e uma incorporação de objetivo ($\psi(l)$) aproxima a ocupação de objetivo com desconto logarítmico.

• Base Teórica: Baseia-se na equivalência entre a função Q para recompensas de alcance de objetivo e a probabilidade de ocupação condicional (Lema 1).
• Objetivo: O modelo aprende a discriminar entre pares estado-ação que podem alcançar um objetivo de linguagem específico (amostras positivas) e aqueles que não podem (amostras negativas).

B. Inovações Técnicas Chave

1. RL Contrastivo Condicionado a Linguagem:

   • Diferentemente do CRL padrão que amostra estados futuros geometricamente, o PRTS trata a instrução de linguagem como o objetivo fixo para toda uma trajetória.
   • Problema de Múltiplos Positivos: Como um único objetivo de linguagem se aplica a todos os passos de tempo em uma trajetória, um lote contém muitos pares positivos para um objetivo.
   • Pesagem Temporal: Para recuperar a distribuição de amostragem geométrica do CRL padrão, os autores introduzem um esquema de pesagem temporal. Pares estado-ação mais próximos do final de uma trajetória (mais próximos do objetivo) são ponderados mais alto ($\gamma^{T-t}$). Isso permite que o modelo aprenda a estrutura temporal do alcance de objetivo sem rótulos de recompensa explícitos.
   • Objetivos Bidirecionais:
     • Estado-Ação $\to$ Linguagem: Alinha $(s, a)$ com sua instrução específica.
     • Linguagem $\to$ Estado-Ação: Força o modelo a ranquear estados-ações dentro de uma trajetória por sua proximidade ao objetivo, codificando o progresso temporal.

2. Arquitetura de Passada Única (Sobrecarga Negligenciável):

   • O PRTS anexa dois blocos de tokens auxiliares à sequência de entrada: <CRL_action> e <CRL_goal>.
   • Máscara Causal Consciente de Função: Uma máscara de atenção personalizada (implementada via um kernel CuTe-FlashAttention fundido) garante isolamento de informação:
     • Os tokens <CRL_action> atendem apenas à visão/propriocepção (não à linguagem).
     • Os tokens <CRL_goal> atendem apenas à instrução.
     • Os tokens de ação autoregressivos padrão mantêm a atenção causal padrão.
   • Resultado: O modelo produz logits de ação discretos, incorporações estado-ação e incorporações de objetivo em uma única passada, incorrendo em sobrecarga computacional negligenciável em comparação com a clonagem de comportamento padrão.

3. Pipeline de Treinamento em Duas Etapas:

   • Pré-treinamento: Otimiza conjuntamente a espinha dorsal VLM na clonagem de comportamento (previsão do próximo token), tarefas auxiliares (Robot-QA) e as perdas contrastivas bidirecionais. Nenhum especialista em ação é anexado ainda.
   • Pós-treinamento: Um Especialista de Ação de Correspondência de Fluxo (baseado em Diffusion Transformer) é anexado e ajustado finamente em dados de domínio específico. A espinha dorsal pré-treinada fornece os priores de alcançabilidade de objetivo.

3. Contribuições Principais

1. Reformulação do Pré-treinamento VLA: Primeiro a escalar RL Contrastivo sem recompensa e condicionado a linguagem diretamente para a fase de pré-treinamento VLA, eliminando a necessidade de anotações manuais de recompensa ou redes de valor separadas.
2. Consciência Intrínseca de Alcançabilidade de Objetivo: As representações aprendidas codificam quantitativamente a probabilidade de alcançar um objetivo, fechando a lacuna entre raciocínio semântico e progresso temporal da tarefa.
3. Arquitetura Eficiente: A "Máscara Causal Consciente de Função" e o kernel CuTe personalizado permitem supervisão densa de alcançabilidade de objetivo com sobrecarga computacional próxima de zero, tornando-a escalável para modelos grandes (4B parâmetros) e conjuntos de dados massivos (167B tokens).
4. Framework Unificado: Unifica raciocínio semântico, previsão de ação discreta e estimativa de valor dentro de uma única arquitetura VLM.

4. Resultados Experimentais

O PRTS foi pré-treinado em 167,8 bilhões de tokens de conjuntos de dados diversos (AgiBotWorld, RoboMind, Open X-Embodiment, etc.) e avaliado em benchmarks de simulação e do mundo real.

• Benchmarks de Simulação (LIBERO, LIBERO-Plus, LIBERO-Pro, SimplerEnv):

  • Desempenho SOTA: O PRTS alcança resultados state-of-the-art, igualando ou superando bases fortes (ex: $\pi^0.5$, ABot-M0) enquanto usa significativamente menos computação de pós-treinamento (ex: 1/8 do orçamento de $\pi^0.5$ no LIBERO).
  • Tarefas de Longo Horizonte: Ganhos significativos no conjunto "Long" do LIBERO (96,6% vs. 94,8% para bases com orçamento correspondente), provando a eficácia da consciência temporal.
  • Robustez: Supera as bases na generalização zero-shot sob deslocamentos de distribuição (câmera, estado do robô, iluminação, disposição de objetos) e perturbações de Tarefa/Posição no LIBERO-Pro.
  • Instrução Zero-Shot: Melhoria massiva em instruções novas (31,5% vs. 0,8% para $\pi^0.5$ no eixo "Tarefa" do LIBERO-Pro).

• Avaliação no Mundo Real (RealMan Dual-Arm & Flexiv Single-Arm):

  • Manipulação Complexa: Alcançou 95,9% de taxa de sucesso média em 11 tarefas de dois braços e 90,0% em 3 tarefas de braço único de alta granularidade.
  • Recuperação de Intervenção Humana: O PRTS demonstrou robustez superior quando humanos interromperam tarefas (ex: removendo objetos durante a execução). Diferentemente de bases que falharam ou reproduziram trajetórias obsoletas, o PRTS reavaliou com sucesso o estado e replanejou para alcançar o objetivo.
  • Generalização: Mantive altas taxas de sucesso sob mudanças de iluminação, identidade de objetos e disposição espacial.

• Visualização de Valor:

  • O produto interno $\phi(s, a)^\top \psi(l)$ mostrou-se fortemente correlacionado com o progresso da tarefa, aumentando à medida que o robô se aproxima do objetivo e permanecendo baixo para instruções incompatíveis, confirmando que o modelo aprende uma função de valor condicionada a objetivo válida.

5. Significado

O PRTS representa uma mudança de paradigma nos modelos fundamentais robóticos:

• Da Imitação ao Raciocínio: Avança além da simples clonagem de comportamento (imitação) para aprender consciência de alcançabilidade de objetivo, permitindo que o robô entenda quão perto está de um objetivo e quão viável é um caminho.
• Escalabilidade: Demonstra que objetivos contrastivos, anteriormente limitados a domínios menores de tomada de decisão, podem escalar para espinhas dorsais VLA grandes (4B+ parâmetros) e conjuntos de dados massivos sem custo proibitivo.
• Inteligência de Propósito Geral: Ao injetar priores temporais e conscientes de objetivo durante o pré-treinamento, o PRTS alcança generalização zero-shot e robustez notáveis, sugerindo um caminho para inteligência robótica verdadeiramente de propósito geral e orientada a objetivos, capaz de lidar com ambientes reais complexos, de longo horizonte e dinâmicos.