Flexible Multitask Learning with Factorized Diffusion Policy

1. Declaração do Problema

O artigo aborda os desafios do aprendizado por imitação multitarefa em robótica, especificamente a dificuldade de modelar distribuições de ação altamente multimodais e diversas.

• Limitações de Modelos Monolíticos: Políticas de rede única tradicionais (por exemplo, Política de Difusão padrão) lutam para generalizar entre tarefas diversas, frequentemente subajustando distribuições complexas ou falhando em representar modos comportamentais distintos (por exemplo, diferentes estratégias de preensão para objetos diferentes).
• Limitações de Modelos Modulares Existente: Embora arquiteturas de Mistura de Especialistas (MoE) tentem decompor tarefas, elas sofrem com:
  • Instabilidade de Treinamento: Dificuldade em equilibrar a utilização dos especialistas.
  • Problemas de Roteamento: A seleção discreta de especialistas pode levar a um "desequilíbrio de roteamento" (alguns especialistas nunca são usados, outros são superutilizados).
  • Baixa Interpretabilidade: Os especialistas frequentemente têm papéis sobrepostos ou pouco claros, e as habilidades podem estar distribuídas entre camadas em vez de serem desentrelaçadas.
  • Esquecimento Catastrófico: Adaptar-se a novas tarefas frequentemente exige o re-treinamento de todo o modelo, arriscando a perda de habilidades previamente aprendidas.

2. Metodologia: Política de Difusão Fatorizada (FDP)

Os autores propõem a FDP, uma arquitetura de política modular que fatoriza a distribuição de ação em uma composição de modelos de difusão especializados.

A. Arquitetura Central

• Fatorização: Em vez de uma única política $p(a_t|o_t)$, a FDP modela a distribuição de ação como um produto ponderado de $N$ distribuições componentes:
  $$p(a_t|o_t) \propto \prod_{i} p_i(a_t|o_t)^{w_{t,i}}$$
  Aqui, cada $p_i$ é um modelo de difusão especializado (especialista) capturando um modo comportamental distinto (sub-habilidade), e $w_{t,i}$ são pesos dependentes da observação.
• Amostragem Composicional:
  • Roteador: Um MLP leve recebe a observação $o_t$ e produz pesos contínuos $\{w_{t,i}\}$ para todos os componentes.
  • Agregação de Pontuação: Diferentemente do MoE discreto (que seleciona um especialista), a FDP agrega as estimativas de pontuação (gradientes da probabilidade logarítmica) de todos os componentes:
    $$\nabla \log p(a_t|o_t) \approx \sum_{i} w_{t,i} \nabla \log p_i(a_t|o_t)$$
  • Processo de Remoção de Ruído: O passo iterativo de remoção de ruído utiliza essa soma ponderada de pontuações para guiar a geração da trajetória de ação. Isso corresponde à amostragem a partir do produto de distribuições, enquadrandose efetivamente a geração de ação como um problema de satisfação de restrições.

B. Treinamento e Adaptação

• Treinamento Conjunto: Todos os componentes de difusão e o roteador são treinados de ponta a ponta usando uma perda de Erro Quadrático Médio (MSE) na previsão de pontuação agregada.
• Adaptação Eficiente (Upcycling): Para adaptar-se a uma nova tarefa sem re-treinar todo o modelo:
  1. Um novo componente de difusão é introduzido.
  2. Ele é inicializado por upcycling (cópia de pesos) a partir de um componente existente.
  3. Apenas o novo componente e o roteador são ajustados finamente; todos os componentes anteriores são congelados.
  4. Isso previne o esquecimento catastrófico e reduz significativamente o número de parâmetros treináveis.

3. Contribuições Principais

1. Arquitetura Modular Inovadora: Introdução da FDP, que utiliza agregação contínua de pontuação em vez de seleção discreta de especialistas. Isso garante treinamento estável, previne desequilíbrio de roteamento e promove especialização clara entre componentes.
2. Interpretação Probabilística: Fundamentada em Modelos Baseados em Energia (EBM) e correspondência de pontuação, fornecendo uma formulação principista onde a política composta representa a interseção de restrições comportamentais.
3. Estratégia de Adaptação Escalável: Proposição de um mecanismo de "ajuste seletivo" usando upcycling e congelamento, permitindo expansão eficiente de habilidades com dados mínimos e sem esquecimento catastrófico.
4. Validação Empírica: Demonstrou desempenho superior sobre bases fortes (DP Monolítico, Política de Difusão Esparsa e MoDE) tanto em ambientes de simulação quanto do mundo real.

4. Resultados Experimentais

Os autores avaliaram a FDP no MetaWorld, RLBench, LIBERO e em manipulação robótica do mundo real (braço UR5e).

• Desempenho Multitarefa:

  • A FDP superou consistentemente a Política de Difusão (DP) monolítica e bases modulares (SDP, MoDE) em mais de 30 tarefas.
  • MetaWorld: A FDP alcançou 74,8% de sucesso médio versus 70,8% (DP) e 69,8% (SDP).
  • RLBench: A FDP alcançou 63,9% de sucesso médio versus 45,6% (DP) e 42,1% (SDP).
  • Mundo Real: A FDP mostrou taxas de sucesso mais altas em tarefas complexas (por exemplo, pendurar canecas, pegar cubos coloridos) onde as bases falharam devido a sobreajuste ou incapacidade de capturar distribuições multimodais.

• Adaptação e Transferência de Tarefas:

  • Eficiência de Dados: Ao adaptar-se a novas tarefas com dados limitados (10–25 demonstrações), a FDP usando a estratégia "Adicionar Novo Módulo" superou o ajuste fino de parâmetros completos e o ajuste fino parcial.
  • Aprendizado Contínuo: Em um benchmark de aprendizado contínuo de 12 tarefas (LIBERO), a FDP manteve alto desempenho à medida que novos especialistas eram adicionados, enquanto outros métodos degradavam.
  • Retenção de Conhecimento: Congelar componentes antigos preservou efetivamente habilidades anteriores. O uso de um buffer de replay mínimo (5 demonstrações/tarefa) permitiu que o modelo retivesse quase 100% do desempenho original após a adaptação.

• Análise e Escalonamento:

  • Escalonamento de Componentes: O desempenho melhorou à medida que o número de componentes aumentava de 2 para 4, depois estabilizou.
  • Escalonamento de Tarefas: A lacuna de desempenho entre a FDP e modelos monolíticos aumentou à medida que o número de tarefas crescia, destacando a escalabilidade da FDP.
  • Interpretabilidade: A visualização mostrou que componentes individuais se especializaram em sub-habilidades distintas (por exemplo, alinhamento, aproximação, preensão) em vez de níveis de ruído ou camadas arbitrárias.
  • Velocidade de Inferência: A FDP suporta "reconstrução parcial" (usando apenas os top-$k$ componentes), alcançando um aceleração de 2x na inferência com apenas uma queda de ~19% no desempenho.

5. Significado

• Aprendizado Robótico: A FDP oferece uma solução prática para implantar robôs em ambientes dinâmicos onde devem aprender novas habilidades rapidamente sem esquecer as antigas.
• Estabilidade: Ao substituir o roteamento discreto pela composição contínua de pontuação, a FDP resolve os problemas de instabilidade de treinamento que afligiram políticas de difusão baseadas em MoE.
• Interpretabilidade: A fatorização permite que pesquisadores inspecionem e compreendam sub-habilidades específicas (por exemplo, "preensão" versus "aproximação") dentro da política, auxiliando na depuração e verificação de segurança.
• Direções Futuras: O framework abre caminhos para designs de módulos heterogêneos (mistura de arquiteturas) e sistemas de aprendizado vitalício onde robôs acumulam continuamente uma biblioteca de habilidades especializadas.

Em resumo, a FDP representa um avanço significativo no aprendizado por imitação robótica, combinando o poder generativo de modelos de difusão com uma arquitetura modular robusta, interpretável e adaptável.