Resumo Técnico: PLUME

Declaração do Problema

A manipulação destre de múltiplos dedos é altamente sensível a parâmetros físicos, como formas de objetos, poses e coeficientes de fricção. Esses parâmetros são frequentemente incertos no momento da implantação, levando a dinâmicas de sistema desconhecidas. Embora a simulação permita a coleta de dados em larga escala com parâmetros conhecidos, políticas treinadas via randomização de domínio padrão costumam falhar ao generalizar para tarefas precisas (ex: girar uma chave de fenda) porque as estratégias de manipulação ideais podem precisar mudar com base em valores específicos de parâmetros. Além disso, coletar dados de alta qualidade para tarefas complexas de múltiplos dedos é desafiador, resultando frequentemente em conjuntos de dados heterogêneos contendo trajetórias subótimas ou falhas. Métodos de modelagem de mundo existentes tipicamente dependem de históricos de horizonte fixo para contabilizar implicitamente variações de dinâmica, o que é insuficiente para tarefas precisas sensíveis a tais variações. Adicionalmente, métodos de aprendizado por reforço offline geralmente não focam no problema conjunto de identificação de sistema e aprendizado de política, enquanto métodos de identificação de sistema online frequentemente exigem engenharia de recompensa complexa e não suportam planejamento.

Metodologia

Os autores propõem o PLUME (Probabilistic Latent Unified world Modeling and parameter Estimation), um framework que aprende conjuntamente um modelo de mundo e estima uma crença sobre parâmetros físicos latentes.

Arquitetura Central

O PLUME consiste em três componentes principais treinados de ponta a ponta:

1. Aprendizado de Representação de Parâmetros: Um Auto-Encoder de Wasserstein (WAE) com backbone Transformer aprende uma representação latente unificada ($z_t$) de variáveis heterogêneas, incluindo parâmetros físicos (ex: fricção, forma) e recompensas. O codificador tokeniza essas entradas e utiliza pooling de atenção cruzada para produzir um embedding latente. Uma priori normal padrão é imposta para facilitar a inicialização razoável das crenças de parâmetros.
2. Modelagem de Histórico: Um codificador Mamba-2 comprime o histórico da trajetória (ações passadas e observações) em um vetor de contexto ($c_t$), retendo um contexto histórico útil que suposições de Markov simples poderiam perder, especialmente com observações ruidosas ou partículas limitadas.
3. Flow Matching Conjunto: Um modelo de flow matching serve como o modelo de mundo unificado. Ele aprende a estimar parâmetros latentes e modelar a dinâmica do sistema condicionada a esses parâmetros. O modelo utiliza um backbone transformer e uma máscara auxiliar para lidar com diferentes distribuições condicionais:
   • Estimativa de Parâmetros: $p(Z_{t+1} | Z_t, y_t, u_{t+1}, y_{t+1}, c_t)$
   • Geração de Rollout: $p(u_{t+1}, y_{t+1}, Z_{t+1} | Z_t, y_t, c_t)$
   • Dinâmica Inversa: $p(u_{t+1} | Z_t, y_t, Z_{t+1}, y_{t+1}, c_t)$

Treinamento e Inferência

• Treinamento: O modelo é treinado em datasets offline contendo trajetórias com variados valores de parâmetros. A função de perda combina uma perda de reconstrução para parâmetros, uma regularização de Discrepância de Média Máxima (MMD) para alinhar o espaço latente com uma normal padrão, e a perda de flow matching. Crucialmente, os gradientes da perda de flow matching não são retropropagados para o codificador de parâmetros para evitar que o modelo otimize seu próprio alvo à custa da capacidade representacional.
• Planejamento (Implantação): Na implantação, o PLUME mantém uma crença sobre parâmetros latentes representada como um conjunto de partículas ($Z_t$). Ele emprega uma abordagem de amostragem de trajetória de horizonte recuante:
  1. Rollout: O modelo amostra autoregressivamente múltiplas trajetórias condicionadas à observação atual, histórico e crença de partículas.
  2. Pontuação: As trajetórias são pontuadas com base em uma combinação de recompensas decodificadas e verossimilhanças de transição estimadas. A verossimilhança é estimada usando um modelo de dinâmica inversa; verossimilhanças baixas penalizam trajetórias de alta recompensa que exibem dinâmicas irreais (detecção de Out-of-Distribution).
  3. Execução e Atualização: A primeira ação da trajetória de maior pontuação é executada. A crença latente é então atualizada online usando a nova observação, permitindo que o modelo de mundo se adapte à verdadeira dinâmica sem necessidade de re-treinamento ou fine-tuning.

Principais Contribuições

1. Modelo de Mundo de Flow-Matching Unificado: Um modelo que estima parâmetros do sistema para condicionar rollouts do modelo de mundo, permitindo que uma única política se adapte a variações de dinâmica.
2. Representação Latente End-to-End: Um método para aprender um espaço latente unificado para parâmetros qualitativamente diferentes (fricção, forma) e recompensas, otimizado conjuntamente com o modelo de mundo para ser útil ao aprendizado de dinâmica.
3. Planejamento Ponderado por Verossimilhança: Uma estratégia de planejamento que utiliza rollouts do modelo de mundo e verossimilhanças de transição estimadas para otimizar ações robóticas, filtrando efetivamente planos de alta recompensa, porém irreais.

Resultados Experimentais

Os autores avaliaram o PLUME em quatro tarefas simuladas (girar chave de fenda, girar válvula, dar um toque em disco, levantar balde) e uma tarefa de hardware (girar chave de fenda).

• Baselines: Comparações foram feitas contra Unified World Model (UWM), Flow Q-Learning (FQL) e Decision Diffuser (DD).
• Desempenho:
  • Chave de fenda Simulada: O PLUME alcançou uma taxa de validade de 81,3%, superando significativamente o UWM (70,7%), FQL (21,3%) e DD (19,3%). Também reduziu a distância ao objetivo em 31% comparado ao melhor baseline.
  • Chave de fenda em Hardware: O PLUME alcançou uma taxa de validade de 93,3% (28/30 tentativas) com transferência zero-shot da simulação, superando o UWM (86,7%).
  • Outras Tarefas: O PLUME superou significativamente todos os baselines nas tarefas de girar válvula e toque em disco (medido pela distância ao objetivo) e foi o único método a ter sucesso consistente na tarefa de levantar o balde DexArt.
• Ablações: O artigo nota que a capacidade do PLUME de raciocinar sobre variáveis de estado parcialmente observáveis (posição, velocidade, orientação), além de parâmetros físicos, contribui para seu sucesso, como evidenciado pela estimativa precisa de recompensa.

Significância e Alegações

O artigo alega que, para tarefas de manipulação de múltiplos dedos precisas e sensíveis, estimar explicitamente e condicionar em representações de parâmetros do sistema leva a um desempenho superior comparado a métodos que aprendem políticas agnósticas a parâmetros ou modelam a variação implicitamente.

As principais alegações incluem:

• Transferência Zero-Shot: O PLUME pode implantar uma política treinada em simulação para tarefas de hardware (girar chave de fenda) sem fine-tuning, adaptando-se à verdadeira dinâmica através de inferência de parâmetros online.
• Robustez à Qualidade de Dados: O método pode aprender modelos úteis de datasets heterogêneos contendo trajetórias subótimas e falhas, aproveitando rótulos de recompensa para usar dados de falha para aprendizado de dinâmica enquanto evita ações problemáticas através de planejamento.
• Adaptabilidade: Ao contrário de métodos que exigem re-treinamento ou fine-tuning, o PLUME alinha seu modelo de mundo à verdadeira dinâmica na implantação ao atualizar sua crença sobre parâmetros, permitando que um único modelo lide com diversas condições físicas.

Os autores reconhecem limitações, especificamente a necessidade de valores de parâmetros reais durante o treinamento (disponíveis em simulação, mas nem sempre no mundo real) e a escala quadrática de memória dos transformers em relação ao número de partículas. Eles sugerem que trabalhos futuros poderiam envolver o fine-tuning de modelos pré-treinados em dados do mundo real para alinhar as prioris sem exigir rótulos de verdade fundamental (ground-truth).