ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

Autores: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

Licenças

Acesso aberto
Artigo: CC BY 4.0

Tudo o que segue se baseia no artigo completo sob a licença CC BY 4.0.

Em resumo

O ReMix usa otimização distribucionalmente robusta por grupos sobre a perda excedente de uma política para aprender quanto amostrar de cada conjunto de dados robótico, superando os pesos uniformes e os pesos de especialistas do RT-X em robôs reais.

Por que importa

Equipes que fazem pré-treinamento com dados robóticos agrupados podem substituir pesos de amostragem ajustados à mão por pesos calculados a partir dos dados e, nos testes relatados, podem descartar 75 % dos dados com perda limitada de desempenho. Os pesos aprendidos também mostram quais conjuntos de dados devem ter peso reduzido e por quê, por exemplo os dados do Kuka coletados de forma autônoma, o que ajuda a decidir quais dados coletar ou manter.

Problema

Grandes conjuntos de dados robóticos como o Open X-Embodiment reúnem dados de diferentes robôs, câmeras, espaços de ações e controladores, mas as misturas usadas para treinar políticas generalistas como RT-X, Octo e OpenVLA foram escolhidas à mão. Os métodos de filtragem de dados de visão e linguagem se transferem mal, porque n-gramas não se aplicam e embeddings visuais ignoram as ações e a natureza sequencial dos episódios.

Solução proposta

O ReMix formula a curadoria como um problema min-max: uma política discreta minimiza sua perda excedente de clonagem de comportamento em relação a um modelo de referência, enquanto os pesos de domínio são atualizados por subida de gradiente exponenciado para maximizá-la, de modo que o peso vai para os domínios em que a política ainda pode melhorar. Para tornar as perdas comparáveis entre robôs, as ações são normalizadas com uma gaussiana por domínio e discretizadas em intervalos, e o checkpoint de referência é escolhido por parada antecipada antes que ele sofra overfitting em algum domínio. Os pesos médios ao longo do treinamento são então reutilizados para treinar políticas de difusão separadas ou para subamostrar os dados.

Principais ideias

Usar a perda excedente em relação a uma política de referência, em vez da perda bruta, reduz o peso tanto de domínios que já estão ajustados quanto de domínios difíceis demais de ajustar.
Diferenças de escala das ações distorcem a otimização robusta: em um teste simplificado com normalização por limites, 94,3 % do peso foi para um domínio de ações puramente aleatórias; com normalização gaussiana por domínio, os dados reais receberam 84,2 %.
Perdas contínuas L1/L2 sobre as ações são influenciadas por valores atípicos e ações multimodais, então discretizar as ações gera pesos mais confiáveis mesmo quando a política final é um modelo de difusão contínuo.
Os conjuntos de dados robóticos (cerca de 10k-100k demonstrações) são pequenos o bastante para o modelo de referência sofrer overfitting; usar um checkpoint 50K passos após a parada antecipada reduziu o desempenho em mais de 15 % em média.

Principais resultados

Na mistura de dados do RT-X, políticas treinadas com os pesos do ReMix superaram os pesos uniformes em 38 % e os pesos de especialistas humanos do RT-X em 32 % em média, em quatro tarefas reais com WidowX e duas com Franka; os pesos dos especialistas superaram os uniformes em 6 %.
Em subconjuntos de 25 % da mistura do RT-X, o ReMix manteve o desempenho, enquanto os pesos humanos caíram mais de 10 %; o ReMix manteve 72 % dos dados do UR5 e 12 % dos do Kuka, contra 30 % e 24 % dos pesos humanos.
O ReMix elevou o pequeno conjunto de dados Toto de 3,42 % (uniforme) para 16,3 % e reduziu o Kuka de 24,9 % para 12,1 %.
No conjunto de dados Bridge V2 completo não houve diferença notável entre o ReMix e os pesos uniformes; em subconjuntos de 25 % do Bridge, o ReMix superou os pesos uniformes e a baseline SSP baseada em embeddings de imagem.
Em uma tarefa do WidowX com um cubo e um objeto distrator, os pesos do ReMix tiveram sucesso em 7/10 contra 0/10 tanto para os pesos humanos quanto para os uniformes.
Em simulação (Robomimic NutAssemblySquare dividido por operador), o ReMix obteve 77/100 contra 53/100 com pesos uniformes em subamostragem de 50 %, e 59/100 contra 39/100 em 25 %.

Limitações

Os autores observam que a avaliação no mundo real abrangeu apenas dois braços, WidowX e Franka, e deixam mais embodiments para trabalhos futuros.
O ReMix aumenta o peso de conjuntos de dados com distribuições de ações anômalas e multimodais como o Toto, o que, segundo os autores, não é necessariamente desejável.
Calcular os pesos exige treinar duas vezes com todos os dados, uma para o modelo de referência e outra para a otimização robusta.
Testado no Bridge V2 e no subconjunto RT-X do OpenX, não no conjunto de dados OpenX completo de mais de 2M episódios. A seção 4.1 descreve o Bridge V2 como 50k demonstrações, enquanto o apêndice B.2 diz que foram usadas cerca de 45.000 trajetórias das 60K do conjunto de dados completo.
Políticas treinadas apenas com os dados de pré-treinamento não generalizaram zero-shot para as configurações dos autores, então toda a avaliação depende de cotreinamento com 25 demonstrações do próprio domínio por tarefa.

Impacto na área

Finalista do CoRL 2024 Outstanding Paper Award e citado 74 vezes. Os autores disponibilizam o código em github.com/jhejna/remix. O artigo também publica seus pesos de domínio aprendidos para a mistura do RT-X (tabela 2) e para os 32 domínios do Bridge V2 (tabela 5).

Resumo

Conjuntos de dados de aprendizado por imitação cada vez maiores estão sendo coletados com o objetivo de treinar modelos de fundação para robótica. No entanto, apesar de a seleção de dados ter sido de suma importância em visão e processamento de linguagem natural, poucos trabalhos em robótica questionaram com quais dados esses modelos deveriam de fato ser treinados. Neste trabalho, investigamos como ponderar diferentes subconjuntos ou “domínios” de conjuntos de dados robóticos para o pré-treinamento de modelos de fundação para robôs. Concretamente, usamos otimização distribucionalmente robusta (DRO) para maximizar o desempenho no pior caso em todos os possíveis domínios de aplicação posteriores. Nosso método, Re-Mix, enfrenta a ampla gama de desafios que surgem ao aplicar DRO a conjuntos de dados robóticos, incluindo a variabilidade dos espaços de ações e da dinâmica entre conjuntos de dados. O Re-Mix emprega parada antecipada, normalização de ações e discretização para contornar esses problemas. Por meio de ampla experimentação com o maior conjunto de dados de código aberto de manipulação robótica, o conjunto de dados Open X-Embodiment, demonstramos que a curadoria de dados pode ter um impacto desproporcional no desempenho posterior. Especificamente, os pesos de domínio aprendidos pelo Re-Mix superam os pesos uniformes em 38 % em média e os pesos escolhidos por humanos em 32 % nos conjuntos de dados usados para treinar políticas robóticas generalistas existentes, especificamente os modelos RT-X.

Open Source Robotics