Grounding Bodily Awareness in Visual Representations for Efficient Policy Learning

Este artigo apresenta o ICon, um método de aprendizado contrastivo em nível de token para Vision Transformers que separa representações específicas do agente das do ambiente, gerando representações visuais centradas no corpo que melhoram o aprendizado e a transferência de políticas em tarefas de manipulação robótica.

O essencial

Imagine que você está tentando ensinar um robô a fazer tarefas domésticas, como fechar uma porta de micro-ondas ou pegar um copo. O robô tem "olhos" (câmeras), mas ele tem um problema: ele vê o mundo todo misturado. Para ele, a mesa, a cadeira, a parede e o próprio braço do robô são apenas um grande emaranhado de pixels coloridos.

O robô precisa aprender a distinguir: "O que é EU (meu corpo) e o que é O MUNDO (o ambiente ao redor)". Sem essa distinção, ele fica confuso e aprende a tarefa de forma lenta e ineficiente.

Este artigo apresenta uma solução inteligente chamada ICon (Contraste Inter-token). Vamos explicar como funciona usando analogias do dia a dia:

1. O Problema: O Robô "Cego" para o Próprio Corpo

Pense em um aluno tentando aprender a andar de bicicleta. Se ele olhar apenas para a estrada e ignorar completamente onde estão seus pés e o guidão, ele vai cair. Da mesma forma, os robôs muitas vezes ignoram o próprio corpo nas imagens porque o foco do treinamento é apenas "fazer a tarefa". Eles acabam filtrando as informações sobre o próprio braço como se fossem "ruído" desnecessário.

2. A Solução: O "Espelho" de Contraste (ICon)

Os autores criaram um método para ensinar o robô a olhar para a imagem e dizer: "Ok, esses pixels são o meu braço, e aqueles pixels são a mesa".

Eles usam uma tecnologia chamada Vision Transformer (que é como um cérebro que divide a imagem em pequenos pedaços, chamados de "tokens", como se fossem peças de um quebra-cabeça).

O ICon funciona como um jogo de "Semelhança e Diferença":

• A Regra: "Tudo que é parte do robô deve parecer com tudo que é parte do robô. Tudo que é o ambiente deve parecer com o ambiente. Mas o robô NÃO deve parecer com o ambiente."
• A Analogia: Imagine que você está em uma festa lotada. O ICon é como um professor que diz: "Agrupem-se! Todos que usam camisa vermelha (o robô) fiquem juntos. Todos que usam camisa azul (o ambiente) fiquem juntos. Mas não deixem a camisa vermelha se misturar com a azul." Isso cria uma separação clara na mente do robô.

3. As Truques Inteligentes (Tecnologia Simplificada)

Para fazer isso funcionar bem, o ICon usa dois truques criativos:

• Amostragem por Ponto Mais Longe (FPS):
  Imagine que você quer escolher representantes de uma sala para um grupo de discussão. Se você escolher aleatoriamente, pode acabar pegando 5 pessoas sentadas no mesmo canto. O ICon usa um método chamado "Farthest Point Sampling". É como se o professor dissesse: "Escolham pessoas que estejam o mais longe possível umas das outras na sala". Isso garante que o robô veja o braço inteiro (da ponta dos dedos até o ombro) e não apenas um pedacinho, capturando a estrutura completa do corpo.

• Contraste em Múltiplos Níveis:
  O cérebro do robô (o Transformer) tem várias camadas. As camadas iniciais veem formas simples (bordas), e as camadas finais veem conceitos complexos. O ICon aplica essa regra de "separar corpo e ambiente" em todas as camadas, não apenas no final. É como ensinar a criança a distinguir "eu" e "outro" desde que ela vê uma linha simples até quando ela entende uma cena complexa.

4. O Resultado: Robôs que Aprendem Mais Rápido e Se Adaptam

Os testes mostraram que robôs usando o ICon:

1. Aprendem mais rápido: Eles conseguem realizar tarefas (como abrir caixas ou fechar gavetas) com mais sucesso do que os robôs sem essa ajuda.
2. São mais estáveis: O treinamento não oscila tanto; o robô não "esquece" o que aprendeu facilmente.
3. Transferem conhecimento: Se você treina um robô (o "Robô A") e depois quer ensinar a mesma tarefa para um robô diferente (o "Robô B", que tem um braço de formato diferente), o robô que usou o ICon aprende muito mais rápido. Como ele já entendeu o conceito de "meu corpo vs. o mundo", ele se adapta facilmente a um novo corpo.

Resumo em uma Frase

O ICon é como dar ao robô um espelho interno que o ajuda a separar visualmente o que é ele mesmo do que é o mundo ao redor, tornando-o mais inteligente, rápido e capaz de aprender novas tarefas com menos esforço.

Resumo técnico

1. O Problema

O aprendizado de políticas de manipulação robótica baseadas em visão enfrenta um desafio fundamental: a complexidade das dinâmicas corporais envolvidas na execução de ações. Embora a visão seja crucial para a percepção do ambiente e do próprio corpo (conhecida como propriocepção visual), extrair informações conscientes do corpo a partir de imagens de alta dimensão em frameworks de aprendizado end-to-end é difícil.

Nesses frameworks, onde os codificadores visuais e as redes de políticas são otimizados conjuntamente, o modelo tende a convergir para soluções que filtram inadvertidamente sinais visuais relacionados ao corpo do agente, focando apenas em tarefas irrelevantes para a execução da ação. Métodos anteriores que tentam resolver isso usando funções de perda de reconstrução (como autoencodificadores para reconstruir máscaras do agente) podem comprometer a estabilidade do treinamento e a performance da política.

2. Metodologia: Inter-token Contrast (ICon)

Os autores propõem o ICon (Inter-token Contrast), um método de aprendizado contrastivo aplicado às representações em nível de token de Vision Transformers (ViTs). O objetivo é criar representações visuais centradas no agente, incorporando vieses indutivos específicos do corpo.

Componentes Principais:

• Máscaras em Nível de Token:
  • A imagem é dividida em patches (tokens) pelo ViT.
  • Uma máscara binária do agente (gerada por modelos de segmentação como o SAM) é mapeada para o nível de tokens. Se a proporção de pixels do agente em um patch exceder um limiar $\beta$, o token é classificado como "agente"; caso contrário, como "ambiente".
• Perda Contrastiva Inter-token:
  • O método aplica aprendizado contrastivo (baseado na perda InfoNCE) diretamente nos tokens do ViT.
  • Lógica: Recursos (features) pertencentes à mesma classe (agente ou ambiente) devem ser similares, enquanto recursos de classes diferentes devem ser dissimilares.
  • São calculadas consultas (queries) médias para o agente e para o ambiente.
  • Amostragem de Pontos Mais Distantes (Farthest Point Sampling - FPS): Adaptada para o domínio 2D, a FPS é usada para selecionar tokens-chave (keys) que estejam espacialmente bem distribuídos, garantindo que a amostragem capture a diversidade estrutural do agente e do ambiente, evitando agrupamentos locais.
• Contraste Multi-nível (MLC):
  • Em vez de aplicar a perda apenas na camada final, o ICon aplica o contraste em múltiplas camadas do encoder do ViT.
  • As perdas de cada camada são somadas com pesos exponenciais, dando mais importância às camadas mais profundas (que contêm informações semânticas) para garantir um desenovelamento (disentanglement) completo entre agente e ambiente.
• Integração com Políticas:
  • O ICon é integrado como um objetivo auxiliar à função de perda de uma política de Diffusion Policy (algoritmo de aprendizado por imitação).
  • A perda total é: $L = L_{diffusion} + \lambda L_{ICon}$.

3. Contribuições Chave

1. Novo Paradigma de Representação: Propõe explicitamente o desacoplamento de características específicas do agente e do ambiente no espaço de características do ViT, promovendo a "consciência corporal" visual.
2. Adaptação da FPS para 2D: Introduz a técnica de Farthest Point Sampling no domínio de imagens 2D para garantir uma amostragem diversificada e representativa de tokens para o aprendizado contrastivo.
3. Arquitetura Multi-nível: Desenvolve uma estratégia de fusão de perdas contrastivas em múltiplas camadas do transformer para um desenovelamento mais robusto.
4. Estabilidade de Treinamento: Demonstra que, diferentemente de métodos de reconstrução, o ICon melhora a performance sem sacrificar a estabilidade do treinamento end-to-end.

4. Resultados Experimentais

Os autores avaliaram o ICon em 8 tarefas de manipulação distribuídas em 3 robôs diferentes (Franka Panda, Kinova Gen3, KUKA IIWA) e dois benchmarks de simulação (RLBench e Robosuite).

• Desempenho Geral: O ICon, quando acoplado ao Diffusion Policy (ICon-Diff-C e ICon-Diff-T), superou consistentemente as políticas base (Diff-C, Diff-T) e o método de reconstrução (Crossway-Diff-C) na maioria das tarefas.
  • Exemplo: No RLBench, houve melhorias absolutas de até 21,3% na tarefa "Open Box" e 13,3% na tarefa "Close Microwave" em comparação com a base.
• Transferência de Política (Few-shot): O método facilitou a transferência de políticas entre robôs com morfologias diferentes. Políticas pré-treinadas no Franka e ajustadas (fine-tuned) com poucos dados no Kinova ou IIWA mantiveram ou melhoraram a performance em comparação com a base.
• Estabilidade de Treinamento: O ICon demonstrou maior estabilidade durante o treinamento, mantendo uma taxa de sucesso média mais alta e consistente em comparação com métodos que usam perda de reconstrução, que apresentaram grandes oscilações entre o desempenho máximo e médio.
• Estudo Ablativo:
  • O limiar de máscara $\beta = 0.5$ foi o mais eficaz.
  • A combinação de FPS e Contraste Multi-nível (MLC) foi essencial; a remoção de qualquer um deles causou queda significativa no desempenho.

5. Significado e Impacto

O trabalho demonstra que incorporar vieses indutivos relacionados ao corpo do agente diretamente nas representações visuais aprendidas é uma estratégia superior para o aprendizado de políticas robóticas.

• Eficiência: Permite que robôs aprendam representações estruturadas agente-ambiente a partir de pixels brutos de forma mais eficiente.
• Generalização: A capacidade de transferir políticas entre robôs diferentes sugere que o ICon aprende características invariantes à morfologia, focando na interação dinâmica entre o corpo e o ambiente.
• Estabilidade: Oferece uma alternativa mais robusta aos métodos de reconstrução, que muitas vezes instabilizam o treinamento em cenários complexos.

Limitações: O método ainda apresenta sobrecarga computacional devido ao processo de amostragem FPS durante a propagação direta e foi testado apenas em simulação, sem validação em cenários do mundo real devido a limitações de hardware.

Em resumo, o ICon representa um avanço significativo na forma como robôs percebem a si mesmos através da visão, melhorando tanto o aprendizado quanto a transferência de habilidades de manipulação.