Policy Contrastive Decoding for Robotic Foundation Models

1. Declaração do Problema

Modelos fundamentais robóticos (políticas robóticas generalistas) demonstraram grande potencial para manipulação flexível e de propósito geral. No entanto, os autores identificam uma falha crítica: esses modelos frequentemente aprendem correlações espúrias a partir de trajetórias de pré-treino. Em vez de focar em características de objetos relevantes para a tarefa, as políticas dependem de pistas de fundo irrelevantes (por exemplo, iluminação, texturas ou objetos de fundo específicos) para prever ações.

• Consequência: Quando o ambiente de implantação muda (por exemplo, alterações no fundo, ajustes de iluminação ou reposicionamento de objetos), o desempenho da política degrada-se significativamente porque os "atalhos" aprendidos deixam de ser válidos.
• Limitações Atuais: Soluções existentes frequentemente exigem retreinamento ou ajuste fino, o que é computacionalmente caro e requer acesso aos pesos do modelo. Existe uma falta de métodos livres de treinamento e plug-and-play para corrigir essas falhas no momento da inferência.

2. Metodologia: Decodificação Contrastiva de Políticas (PCD)

Os autores propõem a Decodificação Contrastiva de Políticas (PCD), um plugin livre de treinamento para o momento da inferência projetado para redirecionar a atenção de uma política robótica de características espúrias para pistas visuais relevantes aos objetos.

Mecanismo Central

A PCD opera contrastando as distribuições de probabilidade de ação geradas a partir de duas entradas visuais diferentes:

1. Observação Original ($o_i$): A entrada visual padrão.
2. Observação com Máscara de Objeto ($\hat{o}_i$): A mesma entrada visual onde os objetos-alvo relevantes para a tarefa são mascarados (removidos/inpainted).

O método computa uma nova distribuição de probabilidade de ação contrastiva ($\pi^*$) usando a seguinte fórmula:
$$\pi^*_\theta(a_i | \ell, o_i) = \frac{1}{C} \cdot \pi_\theta(a_i | \ell, o_i) \cdot \left( \frac{\pi_\theta(a_i | \ell, o_i)}{\pi_\theta(a_i | \ell, \hat{o}_i)} \right)^\alpha$$

• Lógica: Se a política depender de características espúrias (fundo), a distribuição de probabilidade será semelhante para ambas $o_i$ e $\hat{o}_i$, resultando em uma razão próxima de 1 (sem alteração). Se a política depender do objeto, a distribuição para $\hat{o}_i$ diferirá significativamente, amplificando a probabilidade da ação correta na saída final.
• Hiperparâmetro ($\alpha$): Controla a força do contraste. $\alpha=0$ recupera a política de base.

Componentes Técnicos Chave

Para tornar a PCD aplicável a diversas políticas robóticas, os autores introduzem duas estratégias específicas:

1. Rastreamento para Máscara (Track2Mask):

   • Objetivo: Gerar automaticamente observações com máscara de objeto para cada quadro em uma trajetória sem intervenção manual.
   • Processo:
     • Anotação: Objetos-alvo no quadro inicial são identificados usando prompts humanos (Ponto/Box) ou detectores de vocabulário aberto prontos para uso (por exemplo, Grounding DINO).
     • Rastreamento: O SAM2 (Segment Anything Model 2) rastreia e segmenta esses objetos nos quadros subsequentes.
     • Inpainting: Os objetos segmentados são inpainted (removidos) para criar a observação mascarada $\hat{o}_i$.

2. Modelagem Probabilística Baseada em KDE (KDE-PM):

   • Desafio: Políticas autoregressivas (por exemplo, OpenVLA) emitem probabilidades de token diretamente, mas políticas baseadas em difusão (por exemplo, Octo, $\pi_0$) geram ações através de um processo de remoção de ruído e não emitem nativamente uma distribuição de probabilidade $\pi_\theta(a|o)$.
   • Solução: Os autores utilizam Estimativa de Densidade de Kernel (KDE) para aproximar a distribuição de probabilidade de ação. Eles amostram $N$ ações candidatas da política de difusão e ajustam uma densidade de kernel gaussiano a essas amostras para estimar $\pi_\theta(a|o)$ e $\pi_\theta(a|\hat{o})$, permitindo que a fórmula de decodificação contrastiva funcione com modelos de difusão.

3. Contribuições Principais

1. Primeira Solução Livre de Treinamento: A PCD é o primeiro trabalho a propor uma abordagem plug-and-play e livre de treinamento para mitigar correlações espúrias em políticas robóticas durante a inferência.
2. Agnóstica ao Modelo: Funciona como um plugin caixa-preta compatível com políticas autoregressivas (OpenVLA) e baseadas em difusão (Octo, $\pi_0$) sem exigir acesso a pesos pré-treinados ou ajuste fino.
3. Robustez a Mudanças de Distribuição: Ao contrastar explicitamente entradas relevantes para o objeto versus irrelevantes para o objeto, a PCD força a política a focar em características causais, melhorando significativamente a generalização para cenários não vistos.

4. Resultados Experimentais

Os autores avaliaram a PCD em 15 tarefas diversas através de simulação (SIMPLER) e ambientes do mundo real usando três políticas de última geração: OpenVLA, Octo e $\pi_0$.

• Desempenho em Simulação:
  • OpenVLA: Taxa de sucesso melhorada em 50,6% em média.
  • Octo: Taxa de sucesso melhorada em 29,7% em média.
  • $\pi_0$: Taxa de sucesso melhorada em 8,9% em média (apesar de o $\pi_0$ já ser uma base forte).
• Desempenho no Mundo Real:
  • Em tarefas de manipulação do mundo real usando a política $\pi_0$, a PCD alcançou uma melhoria massiva de 108% na taxa de sucesso em comparação com a base.
  • O método provou ser robusto contra vários distratores (mudanças de iluminação, desordem de fundo, deslocamentos de posição de objetos).
• Estudos de Ablação:
  • O método é robusto a diferentes estratégias de detecção de objetos (Prompts humanos vs. Grounding DINO vs. YOLO World).
  • É robusto a diferentes algoritmos de inpainting (LaMa teve o melhor desempenho).
  • O hiperparâmetro $\alpha$ mostra um impacto positivo claro quando $>0$.

5. Significado e Limitações

• Significado: A PCD aborda uma fraqueza fundamental nos modelos fundamentais robóticos atuais (dependência de correlações espúrias) sem o alto custo de retreinamento. Demonstra que intervenções simples no momento da inferência podem melhorar drasticamente a generalização e a segurança na implantação no mundo real.
• Limitações:
  • Sobrecarga Computacional: A PCD requer uma segunda passagem forward (para a imagem mascarada) e rastreamento/inpainting de objetos, aumentando a latência de inferência em aproximadamente 24-100% (dependendo da política e paralelização). No entanto, os autores argumentam que isso é aceitável dado os ganhos massivos nas taxas de sucesso.
  • Escopo: Foca atualmente na correção no momento da inferência; trabalhos futuros visam abordar correlações espúrias durante a fase de treinamento.

Conclusão: O artigo apresenta um framework altamente eficaz e versátil (PCD) que melhora significativamente a confiabilidade e a generalização de modelos fundamentais robóticos, forçando-os a atender às características visuais corretas, fechando a lacuna entre o desempenho em simulação e a robustez no mundo real.