Resumo Técnico: FlowMPC

Definição do Problema

O Flow Matching (FM) emergiu como um método poderoso para o aprendizado de comportamento (behavior cloning) em espaços de ação multimodais, particularmente para manipulação robótica, onde múltiplas trajetórias de ação distintas podem resolver a mesma tarefa. No entanto, o FM é fundamentalmente um método de aprendizado por imitação treinado para corresponder ao comportamento do especialista, em vez de maximizar diretamente o retorno esperado. Consequentemente, as políticas de FM são limitadas pela qualidade e cobertura dos dados de demonstração. Elas frequentemente têm dificuldade em recuperar-se de pequenos erros, erros compostos ou estados não bem representados no conjunto de dados do especialista.

Embora trabalhos recentes tenham tentado abordar isso refinando as políticas de FM usando aprendizado por reforço (RL), tais abordagens exigem a modificação do próprio objetivo de treinamento do FM. Este artigo investiga uma questão diferente: Pode um modelo de mundo aprendido melhorar uma política de FM sem alterar o objetivo de treinamento do FM? O objetivo é determinar se o planejamento baseado em modelo pode complementar políticas de imitação baseadas em fluxo durante o teste.

Metodologia: Estrutura FlowMPC

A estrutura proposta, FlowMPC, integra uma política de FM aprendida por imitação com um modelo de mundo aprendido dentro de um loop de planejamento Model Predictive Path Integral (MPPI). A ideia central é usar a política de FM como um mecanismo de proposta estruturado para gerar sequências de ações candidatas, que são então avaliadas e refinadas pelo modelo de mundo.

1. Paradigma de Treinamento

O pipeline de treinamento consiste em cinco estágios:

1. Coleta de Dados do Especialista: Um agente SAC com informações de estado privilegiadas gera demonstrações de especialistas.
2. Coleta de Trajetórias: Trajetórias completas (RGB, propriocepção, ações, recompensas) são coletadas.
3. Treinamento da Política de FM: Uma política de FM é treinada nessas trajetórias de especialistas para aprender um prior de ação multimodal.
4. Pré-processamento Visual: Trajetórias de especialistas são recolhidas novamente com observações RGB pré-processadas pelo codificador visual da política de FM para produzir latentes visuais compactos ($z_{viz,t}$).
5. Treinamento do Modelo de Mundo: O modelo de mundo é treinado online usando um buffer de replay pré-carregado com dados de especialistas, utilizando os latentes visuais pré-processados combinados com o estado proprioceptivo ($x_t = [s_t, z_{viz,t}]$).

2. Adaptações Arquiteturais ao TD-MPC2

A estrutura baseia-se no TD-MPC2, mas substitui seu ator aprendido pela política de FM aprendida por imitação. Várias modificações críticas foram feitas para acomodar este híbrido:

• Preditor de Ação de Behavior-Cloning ($\pi_{BC}$): Uma rede dedicada mapeia o estado atual para uma ação. Isso é usado para o bootstrapping de valor terminal durante os rollouts de planejamento para garantir que a função de valor avalie ações dentro da distribuição, estabilizando o processo de planejamento.
• Interface de Representação Unificada: O modelo de mundo consome o latente visual do FM diretamente concatenado com o estado proprioceptivo, em vez de aprender codificadores separados. Isso permite que o modelo de mundo aproveite os recursos visuais congelados da política de FM.
• Objetivos Simplificados: Dada a estrutura de recompensa binária das tarefas, as complexas recompensas transformadas por symlog discretizadas do TD-MPC2 padrão foram substituídas por entropia cruzada binária para recompensas e erro quadrático médio padrão para regressão de Q-valor.
• Duas Variantes de Modelo de Mundo:
  • Variante de Espaço de Estado: Prediz diretamente no espaço de estado concatenado e latente visual. Usada para a tarefa PickCube ($H=1$).
  • Variante de Espaço Latente: Segue o padrão do TD-MPC2 com uma representação latente aprendida e rollouts de múltiplos passos. Usada para a tarefa mais complexa PickSingleYCB ($H=3$).

3. Inferência e Planejamento

No teste, o planejador opera da seguinte forma:

1. Proposta: Dada a observação atual, o planejador amostra $K$ trajetórias de ação candidatas. Um subconjunto ($K_\pi$) é proposto pela política de FM, enquanto o restante é amostrado aleatoriamente.
2. Rollout & Pontuação: Cada trajetória é executada (rolled out) no modelo de mundo aprendido. As trajetórias são pontuadas somando as recompensas previstas e uma estimativa de valor terminal ($\hat{Q}$) derivada do preditor de behavior-cloning.
3. Refinamento: O MPPI repondera e refina as trajetórias com base nessas pontuações. A primeira ação da sequência otimizada é executada.
4. Eficiência: A codificação visual do FM é computada uma única vez por observação e reutilizada em todas as amostras e etapas de integração.

Resultados Principais

A abordagem foi avaliada em duas tarefas de manipulação do ManiSkill: PickCube-v1 e PickSingleYCB-v1.

• Sucesso de Fim de Episódio: Os ganhos mais significativos foram observados na manutenção do sucesso até o final do episódio.
  • PickCube: Melhorou de 93,14% (apenas FM) para 97,44% (FlowMPC).
  • PickSingleYCB: Melhorou de 56,81% (apenas FM) para 66,41% (FlowMPC).
• Sucesso a Qualquer Momento (Anytime Success): Melhorias também foram vistas em atingir um estado de sucesso em qualquer ponto, embora os ganhos tenham sido menores.
  • PickCube: 95,78% $\to$ 98,68%.
  • PickSingleYCB: 68,77% $\to$ 69,78%.

Os resultados indicam que o modelo de mundo é particularmente eficaz em corrigir pequenos erros de ação e manter o controle estável próximo à conclusão da tarefa, em vez de apenas ajudar o agente a alcançar um estado de sucesso transitório.

Significância e Alegações

O artigo postula que o planejamento baseado em modelo de mundo pode complementar efetivamente as políticas de imitação baseadas em fluxo sem exigir um novo objetivo de treinamento de FM.

• Abordagem Híbrida: O trabalho demonstra uma via híbrida entre aprendizado por imitação e RL baseado em modelo. Utiliza a imitação para aprender um prior de ação multimodal forte e o planejamento para refinar as decisões no tempo de inferência.
• Sem Necessidade de Retreinamento: A melhoria provém de uma melhor utilização da política no tempo de teste, não do retreinamento da política com um novo objetivo de RL.
• Escopo Modesto: Os autores declaram explicitamente que os resultados são limitados a duas tarefas e que o design ideal do modelo de mundo (espaço de estado vs. espaço latente, horizonte de planejamento) variou entre as tarefas. Eles não alegam uma configuração universalmente ótima, mas sugerem que a combinação de políticas generativas expressivas e planejamento de antecipação (lookahead) é uma direção promissora para a melhoria robusta da política.
• Linha de Base de Desempenho: O método superou uma forte linha de base de FM (que já estava próxima do desempenho do TD-MPC2 em uma tarefa), sugerindo que o planejador "afia" uma política forte em vez de meramente resgatar uma política fraca.