Resumo Técnico: DataMIL – Seleção de Dados para Aprendizado por Imitação de Robôs com Datamodels

1. Declaração do Problema

A comunidade de robótica mudou recentemente o foco para o treinamento de políticas de fundação generalistas usando conjuntos de dados massivos e diversos (ex: Open-X Embodiment). Embora essas políticas alcancem um forte desempenho médio em tarefas variadas, elas frequentemente apresentam desempenho inferior em tarefas individuais e especializadas, necessitando de ajuste fino (fine-tuning) com novos dados específicos da tarefa.

O desafio central abordado por este trabalho é a seleção de dados: identificar quais subconjuntos de grandes conjuntos de dados prévios, quando combinados com demonstrações limitadas e específicas da tarefa, produzem as melhores políticas especializadas.

• Limitações dos Métodos Atuais: A seleção ingênua baseada em noções humanas de qualidade (ex: similaridade semântica, similaridade visual ou proximidade estado-ação) frequentemente falha. Esses heurísticos assumem que os dados mais "similares" são os mais "úteis", ignorando o impacto real de um ponto de dado no desempenho da política downstream.
• A Barreira de Otimização: Idealmente, um selecionaria os dados testando exaustivamente subconjuntos através do retreinamento e avaliando a política resultante. No entanto, na robótica, a avaliação da política requer execuções (rollouts) no mundo real, o que é demorado, perigoso e computacionalmente inviável para uma busca em larga escala.
• A Lacuna: Estruturas de atribuição de dados existentes, como datamodels (que preveem o desempenho do modelo com base nos dados de treinamento), tiveram sucesso em NLP e Visão Computacional, mas não foram diretamente aplicáveis à robótica devido à natureza não diferenciável e dispendiosa da avaliação de políticas no mundo real.

2. Metodologia: DataMIL

Os autores introduzem o DataMIL (Datamodels for Imitation Learning), um framework que estende o paradigma de datamodels para a robótica. O DataMIL raciocina sobre a seleção de dados de forma ponta a ponta (end-to-end), utilizando a própria política para identificar pontos de dados que melhoram o desempenho, em vez de depender de heurísticas estáticas.

2.1 Estrutura Central

O DataMIL trata o algoritmo de aprendizado de política como uma "caixa preta" e treina um estimador (um datamodel) para prever o desempenho de uma política treinada em um subconjunto específico de dados.

• Objetivo: Encontrar um subconjunto $D' \subset D$ que maximize a métrica alvo $M(A(D'))$, onde $A$ é o algoritmo de aprendizado e $M$ é a métrica de desempenho.
• Formulação do Datamodel: Em vez de treinar a política, o DataMIL aprende uma função $\hat{f}(D') \approx M(A(D'))$. Para datamodels lineares, isso se reduz a atribuir uma pontuação escalar $\tau(z_i)$ a cada ponto de dado $z_i$, onde o desempenho previsto é a soma das pontuações para o subconjunto selecionado.

2.2 Adaptações Chave para Robótica

Para tornar os datamodels tratáveis para a robótica, os autores introduzem três modificações críticas:

1. Função de Perda Substituta (Evitando Rollouts):
   Execuções no mundo real são não diferenciáveis e caras. O DataMIL substitui a métrica real de taxa de sucesso por uma métrica substituta (proxy) $\tilde{M}$, definida como a perda negativa de Behavior Cloning (BC) em um conjunto mantido à parte de demonstrações alvo:
   $$\tilde{M}(\pi, D_{target}) = \frac{1}{|D_{target}|} \sum_{(s,a) \in D_{target}} -L_{BC}(\pi(s), a)$$
   Esta métrica é totalmente diferenciável e não requer interação com o ambiente, permitindo o uso de métodos de estimativa baseados em gradiente.

2. Estimadores de Datamodel:
   O artigo emprega dois métodos para estimar as pontuações de influência $\tau(z_i)$:

• Estimador de Regressão: Amostra subconjuntos aleatórios de dados, treina políticas, avalia a métrica substituta e ajusta um modelo linear para prever o desempenho com base na presença dos dados. Este método é preciso, mas computacionalmente pesado.
• Estimador Baseado em Metagradiente: Utiliza funções de influência e metagradientes para calcular a derivada da métrica alvo em relação aos pesos dos dados. Este método é mais eficiente para modelos grandes (como o Octo) e evita a necessidade de retreinar políticas para cada subconjunto.

3. Agrupamento (Clustering) e Mitigação de Desvio de Distribuição:

• Agrupamento: Para reduzir o ruído na estimativa de influência, pares individuais de estado-ação são agrupados em clusters temporais (sub-trajetórias ou trajetórias completas). A granularidade ideal do cluster depende do tamanho do conjunto de dados; conjuntos de dados maiores (como o OXE) beneficiam-se da agregação ao nível de trajetória.
• Desvio de Distribuição: Para evitar estimativas ruins causadas por desvios entre os domínios prévio e alvo, uma pequena fração de dados da tarefa alvo é incluída durante a fase de estimativa do datamodel para alinhar o processo de aprendizado com o domínio alvo.

2.3 Pipeline de Treinamento

Uma vez que as pontuações de influência são estimadas:

1. Os $x\%$ superiores de dados prévios (maior influência positiva) são selecionados para formar $D_{sel}$.
2. A política final é treinada via co-treinamento nos dados alvo $D_{target}$ e nos dados selecionados $D_{sel}$. Em cada etapa, um lote (batch) é amostrado de $D_{target}$ com probabilidade $\alpha$ e de $D_{sel}$ com probabilidade $1-\alpha$.

3. Principais Contribuições

• Extensão de Datamodels para Robótica: O DataMIL é o primeiro framework a aplicar com sucesso datamodels ao aprendizado por imitação robótica, superando a barreira das execuções caras no mundo real através de uma perda substituta diferenciável.
• Consciência de Desempenho Ponta a Ponta: Ao contrário dos métodos heurísticos (similaridade visual, de movimento ou de estado), o DataMIL otimiza diretamente o impacto no sucesso da política, permitindo que selecione dados que melhoram o desempenho e descarte dados que o degradam, mesmo que os dados pareçam visualmente similares.
• Estimadores Escaláveis: A introdução do estimador baseado em metagradiente permite que o framework escale para políticas grandes e complexas (ex: Octo) e conjuntos de dados massivos (ex: OXE), onde abordagens baseadas em regressão seriam computacionalmente proibitivas.
• Seleção Cross-Embodiment e Multi-Tarefa: O método demonstra a capacidade de selecionar dados úteis para tarefas e implementações robóticas (embodiments) não presentes no conjunto de dados prévio, e de curar conjuntos de dados que suportam múltiplos objetivos downstream simultaneamente.

4. Resultados Experimentais

Os autores validam o DataMIL através de mais de 60 tarefas em simulação e ambientes reais:

• MetaWorld (50 Tarefas): O DataMIL alcançou um aumento de 10% no desempenho sobre os baselines de última geração (Behavior Retrieval, Flow Retrieval, STRAP). Ele filtrou com sucesso demonstrações ruidosas e subótimas de dados de exploração autônoma, enquanto os baselines baseados em similaridade frequentemente recuperavam dados irrelevantes ou de baixa qualidade.
• LIBERO (10 Tarefas): Usando a política Transformer Octo, o DataMIL superou consistentemente os baselines em diversas tarefas de longo horizonte. Enquanto os baselines de similaridade visual performaram bem em algumas tarefas específicas, o DataMIL proporcionou o desempenho médio mais robusto.
• Open-X Embodiment (OXE) (Mundo Real):
  • Franka-Ball & Franka-Pouch: O DataMIL selecionou dados que melhoraram significativamente as taxas de sucesso em comparação com a seleção aleatória e os baselines heurísticos.
  • Tiago-Sink: Em um cenário desafiador de cross-embodiment (selecionando dados para um robô Tiago de um dataset que não contém dados de Tiago), o DataMIL identificou com sucesso demonstrações relevantes de outros robôs (ex: RT-1, BC-Z) que compartilhavam a "essência" da tarefa (manipulação de bancada), enquanto os baselines falharam em generalizar.
  • Droid-Multitask: O DataMIL curou um único conjunto de dados que melhorou o desempenho em três tarefas distintas simultaneamente, superando os baselines em taxa de sucesso média.

Achados Qualitativos:

• O DataMIL seleciona uma mistura diversificada de conjuntos de dados em vez de depender excessivamente de uma única fonte (ao contrário dos baselines, que frequentemente recuperam mais de 80% de um único dataset).
• Ele distingue efetivamente entre dados úteis e prejudiciais, mesmo quando parecem visualmente idênticos (ex: mesmo estado, mas distribuições de ação diferentes), um modo de falha comum na recuperação baseada em similaridade.

5. Significância e Alegações

O artigo afirma que a seleção de dados orientada pelo desempenho e ponta a ponta é crítica para desbloquear o potencial de grandes conjuntos de dados prévios na robótica. Os autores argumentam que, embora os modelos de fundação forneçam uma base sólida, sua especialização depende fortemente da qualidade dos dados de ajuste fino.

• Alegações Modestas sobre Escalabilidade: Os autores reconhecem que estimar datamodels continua sendo computacionalmente caro (várias vezes o custo de treinar uma política em todos os dados), embora a abordagem de metagradiente mitigue isso. Eles observam que o ajuste de hiperparâmetros (ex: tamanho do cluster, razão de seleção) atualmente carece de uma intuição teórica forte.
• Escopo: O trabalho foca primariamente na especialização de tarefa única, com o cenário Droid-Multitask servindo como um passo preliminar em direção à curadoria multi-tarefa.
• Conclusão: O DataMIL demonstra que ir além dos heurísticos definidos por humanos para uma seleção baseada em dados e consciente do modelo gera ganhos consistentes nas taxas de sucesso, particularmente em cenários complexos, heterogêneos e do mundo real, onde os métodos de recuperação tradicionais falham.