Exploring Conditions for Diffusion models in Robotic Control

1. Problema e Motivação

O trabalho aborda o desafio de adaptar modelos de difusão pré-treinados (especificamente modelos de texto-para-imagem como o Stable Diffusion) para tarefas de controle robótico (imitação de aprendizado).

• Limitação das Representações Atuais: As abordagens atuais utilizam representações visuais pré-treinadas (como CLIP ou VC-1) que são "frozen" (congeladas) durante o aprendizado da política. Isso torna as representações agnósticas à tarefa, ou seja, elas não se adaptam às necessidades específicas de cada tarefa de controle, exigindo inspeção manual para escolher a melhor representação para cada cenário.
• Falha das Condições de Texto: Em outras tarefas de visão computacional, o uso de prompts de texto (descrições) como condições para modelos de difusão tem sido altamente eficaz. No entanto, os autores descobriram que aplicar prompts de texto diretamente em ambientes de controle robótico (gerados por modelos de linguagem-vision) resulta em ganhos mínimos ou até mesmo em queda de desempenho.
• Causa da Falha: A análise revela uma lacuna de domínio (domain gap). Modelos de difusão são treinados em imagens da web (objetos comuns, cenas estáticas), enquanto o controle robótico envolve agentes especializados, dinâmicas de vídeo e a necessidade de granularidade visual fina (ex: distinguir partes específicas de um braço robótico em movimento). O modelo de difusão falha em associar corretamente palavras de texto a regiões específicas da imagem no contexto de controle, gerando mapas de atenção ruidosos.

2. Metodologia: ORCA

Para superar essas limitações, os autores propõem o ORCA (uma sigla que não é explicitamente definida no texto, mas refere-se ao framework de condicionamento). A abordagem visa obter representações visuais adaptáveis à tarefa sem fine-tuning (ajuste fino) do próprio modelo de difusão.

O framework introduz dois tipos de prompts aprendíveis (learnable prompts) que condicionam o modelo de difusão:

1. Prompts de Tarefa (Task Prompts):

   • Substituem os prompts de texto fixos.
   • São parâmetros aprendíveis que são otimizados durante o aprendizado da política de controle (imitação).
   • Objetivo: Aprender implicitamente a focar nas regiões ou objetos críticos para a tarefa específica (ex: o botão em uma tarefa de "pressionar botão" ou o braço robótico em uma tarefa de "corrida de guepardo"), evitando a alucinação de texto.

2. Prompts Visuais (Visual Prompts):

   • Para capturar detalhes dinâmicos e de alta granularidade de cada quadro do vídeo, o método incorpora informações visuais diretas.
   • Utiliza um codificador de visão (ex: DINOv2) para extrair representações densas (não globais) da imagem observada.
   • Essas representações são projetadas e concatenadas como prompts visuais para o modelo de difusão.
   • Isso permite que o modelo de difusão "veja" o estado atual do ambiente quadro a quadro, capturando movimentos e interações finas que o texto não consegue descrever com precisão.

Objetivo de Aprendizado:
Os prompts (tarefa e visual) são aprendidos de ponta a ponta (end-to-end) usando a função de perda de Clonagem de Comportamento (Behavior Cloning Loss). O modelo de difusão permanece congelado; apenas os prompts e a rede de política são atualizados.

3. Contribuições Chave

1. Descoberta da Ineficácia do Texto: Demonstração empírica e análise de atenção cruzada mostrando que prompts de texto gerados por IA são inadequados para condicionamento em controle robótico devido à falta de alinhamento semântico preciso no domínio de simulação/robótica.
2. Proposta do ORCA: Um framework simples e eficaz que substitui condições de texto por prompts de tarefa aprendíveis e prompts visuais densos, permitindo que o modelo de difusão se adapte dinamicamente a cada quadro e tarefa.
3. Eficiência: O método evita o ajuste fino completo do modelo de difusão (que levaria a overfitting em dados limitados de imitação), mantendo o modelo base congelado e aprendendo apenas um pequeno número de parâmetros (os prompts).
4. Análise de Atenção: Visualizações mostram que os prompts aprendidos conseguem focar corretamente em regiões relevantes (ex: pernas do agente, ferramentas), superando o ruído observado com condições de texto.

4. Resultados Experimentais

O ORCA foi avaliado em três benchmarks padrão de aprendizado por imitação: DeepMind Control (DMC), MetaWorld e Adroit.

• Desempenho Geral: O ORCA alcançou o estado da arte (SOTA) em todos os 12 tarefas testadas, superando significativamente baselines como CLIP, VC-1, SCR (que usa difusão sem condições) e métodos adaptados de visão (CoOp, TADP).
  • Exemplo no MetaWorld: Média de sucesso de 95.2% (ORCA) vs. 93.1% (TADP) e 90.4% (SCR).
  • Exemplo no Adroit: Média de sucesso de 65.3% (ORCA) vs. 58.0% (SCR) e 47.3% (VC-1).
• Comparação com Fine-Tuning: O artigo demonstra que o fine-tuning completo de modelos de difusão ou encoders visuais em dados limitados de robótica leva a uma degradação severa de desempenho (ex: queda de >80% no SCR com fine-tuning completo). O ORCA mantém a generalização do modelo pré-treinado enquanto se adapta à tarefa.
• Análise de Componentes: Experimentos de ablação confirmam que tanto os prompts de tarefa quanto os visuais são essenciais. O uso combinado supera o uso individual, indicando que a tarefa precisa de contexto global (prompt de tarefa) e detalhes locais (prompt visual).
• Eficiência Computacional: Embora o modelo de difusão seja grande, o número de parâmetros aprendíveis é baixo (~10.6M), tornando o método viável comparado ao custo de treinar modelos do zero ou fazer fine-tuning massivo.

5. Significado e Conclusão

O trabalho estabelece que condicionar modelos de difusão para robótica requer uma abordagem diferente da visão computacional tradicional. A simples tradução de tarefas para texto não funciona devido à natureza dinâmica e específica do controle robótico.

O ORCA demonstra que é possível alavancar a riqueza de representações de modelos de difusão pré-treinados em grandes conjuntos de dados da web para robótica, desde que se utilize mecanismos de condicionamento que sejam sensíveis ao quadro (frame-sensitive) e adaptáveis à tarefa. Isso abre caminho para o uso de modelos generativos como extratores de características poderosos e adaptáveis em agentes robóticos, sem a necessidade de treinar modelos do zero ou realizar ajustes finos custosos e propensos a overfitting.