Resumo Técnico: BayesFP: Estimativa de Posterior para Políticas Baseadas em Fluxo via Amostragem de Feynman-Kac

Declaração do Problema

Robôs operando em ambientes do mundo real devem gerar trajetórias que sejam suficientemente expressivas para capturar comportamentos multimodais de especialistas e simultaneamente reativas o suficiente para satisfazer restrições de segurança e objetivos específicos de tarefas. Um desafio crítico surge quando essas restrições (ex: desvio de obstáculos, limites articulares, custos de energia) são desconhecidas no momento do treinamento ou mudam entre as implantações.

Abordagens existentes para geração de trajetórias com restrições geralmente se dividem em duas categorias:

1. Filtragem/Otimização pós-hoc: Métodos que retificam rollouts não restritos usando planejadores baseados em gradiente ou filtros de segurança. Estes frequentemente desacoplam a geração da correção, levando a trajetórias fora da distribuição (out-of-distribution), mínimos locais e alta sobrecarga computacional.
2. Orientação generativa (Generative guidance): Métodos que modificam o processo de amostragem diretamente (ex: orientação baseada em classificador, projeção, dinâmica refletida). Embora eficazes, estes são tipicamente modificações heurísticas das dinâmicas de amostragem. Eles carecem de uma relação fundamentada com uma distribuição posterior bem definida que respeite conjuntamente a distribuição de dados aprendida e as restrições de tempo de inferência.

A questão central abordada é: Dada uma política de difusão ou de correspondência de fluxo (flow-matching) pré-treinada, como podemos amostrar trajetórias que respeitem um problema de otimização com restrições especificado apenas no tempo de inferência, permanecendo fiéis à distribuição de demonstração aprendida?

Metodologia

Os autores formulam a geração de trajetórias com restrições como amostragem de posterior Bayesiano. A distribuição de demonstração aprendida $p_{data}(x)$ serve como um prior, e uma função de custo $J(x)$ codificando as restrições de tempo de inferência serve como uma verossimilhança (likelihood). O objetivo é amostrar do posterior:
$$p(x | O = 1) \propto p_{data}(x) \exp(\beta J(x))$$
onde $O=1$ indica otimalidade e $\beta < 0$ controla a força da inclinação induzida pelo custo.

Para amostrar deste posterior sem retreinar a política base, o artigo utiliza o framework do Corretor de Feynman-Kac (FKC). As principais contribuições metodológicas são:

1. Dinâmica Ponderada de Feynman-Kac

Em vez de adicionar heuristicamente gradientes de custo ao score ou campo de velocidade da política, os autores derivam uma SDE (Equação Diferencial Estocástica) ponderada. Esta SDE governa a evolução de uma população de partículas onde cada partícula carrega um peso logarítmico $w_t$. As dinâmicas são derivadas da PDE de Fokker-Planck da distribuição alvo com inclinação de Gibbs.

• Correção de Drift: O termo de drift inclui uma correção consciente do custo proporcional a $\nabla J(x)$.
• Evolução de Pesos: Os pesos acumulam um termo de reponderação derivado da discrepância entre a evolução da dinâmica simulada e a densidade alvo. Isso garante que a população de partículas seja um estimador assintoticamente não enviesado do posterior, mesmo que a correção de drift seja imperfeita.
• Reamostragem: Para mitigar a degenerescência de pesos em horizontes longos, o método emprega técnicas de Monte Carlo Sequencial (SMC), especificamente a reamostragem sistemática sobre um intervalo de tempo ativo.

2. Extensão para Correspondência de Fluxo Determinística (Flow-Matching)

O framework FKC foi originalmente formulado para processos de difusão estocástica. Os autores o estendem para políticas de correspondência de fluxo determinísticas (que resolvem uma ODE $dx = -v_t(x)dt$) recodificando o fluxo como uma SDE de preservação de marginais.

• Seguindo Singh e Fischer [16], eles introduzem um termo de difusão auxiliar $\tilde{\sigma}(t)$ e uma correspondente correção de drift baseada em score.
• Crucialmente, para fluxos retificados com uma fonte Gaussiana, o score $\nabla \log q_t(x)$ pode ser imputado em forma fechada a partir do campo de velocidade aprendido $v_t(x)$, eliminando a necessidade de uma rede de score separada.
• Esta construção permite que o framework FKC seja aplicado uniformemente tanto a políticas de difusão quanto de correspondência de fluxo.

3. Garantias Teóricas

O artigo fornece uma garantia teórica (Corolário 4.2) de que, conforme os pesos de penalidade e a magnitude da temperatura inversa $|\beta|$ aumentam, o posterior de Gibbs se concentra arbitrariamente de forma apertada em trajetórias que são viáveis e próximas do ótimo. Isso implica que o método pode satisfazer restrições sem exigir filtros de segurança externos ou planejadores de backup.

Principais Contribuições

1. Formulação Fundamentada: O trabalho reformula a geração de trajetórias com restrições como inferência Bayesiana, oferecendo uma alternativa fundamentada às orientações heurísticas e projeções pós-hoc.
2. Amostrador Unificado: Instancia o posterior como uma SDE ponderada de Feynman-Kac, estendendo o framework de modelos de difusão estocástica para políticas de correspondência de fluxo determinísticas via recodificação de SDE de preservação de marginais.
3. Convergência Teórica: Prova que o posterior resultante se concentra em soluções viáveis e próximas do ótimo sob força de penalidade suficiente.
4. Validação Empírica: A abordagem é validada através de tarefas simuladas e do mundo real usando modelos pré-treinados (Diffusion Policy, GR00T-N1.6 e $\pi0.5$), demonstrando sucesso em planejar contornando obstáculos não convexos introduzidos no tempo de inferência.

Resultados Experimentais

Os autores avaliaram o BayesFP em vários benchmarks, comparando-o com baselines incluindo Orientação por Combinação Linear, CASF (Constraint-Aware Streaming Flow) e JM2D (Joint Model-based Model-free Diffusion).

• Tarefas de Simulação:
  • Diffusion Policy: Em tarefas RoboMimic (Can, Transport) com obstáculos cilíndricos adicionados, o BayesFP alcançou 0% de taxa de colisão comparado a 30% (Sem Orientação) e 2% (Combinação Linear).
  • GR00T-N1.6: Em tarefas LIBERO-Object com obstáculos cilíndricos e em formato de V não convexos, o BayesFP alcançou 84% e 75% de taxas de sucesso respectivamente, superando significativamente os baselines que tiveram dificuldades com geometrias não convexas.
  • $\pi0.5$: Em tarefas zero-shot RoboLab (ex: FoodPacking, BBQSauce), o BayesFP melhorou as taxas de sucesso (ex: 36% vs 28% para HammersInLeftBin) e reduziu as taxas de colisão em comparação com baselines de combinação linear.
• Experimentos no Mundo Real:
  • Usando um braço SO101 ajustado finamente com apenas 20 demonstrações, o BayesFP lidou com sucesso com uma tarefa de PickAndPlace com um obstáculo de caneca introduzido na inferência (reduzindo a colisão de 98% para 6%) e uma tarefa de BimodalMugSelection onde uma restrição de linguagem foi reforçada por uma restrição de custo para suprimir um modo indesejado (alcançando 100% de seleção do caneca correta vs 72% para o baseline).
• Ablações:
  • Aumentar o número de partículas ($K$) melhorou o desempenho, mas saturou em torno de $K=4$ para obstáculos simples.
  • Aumentar a força de orientação $|\beta|$ levou a um decaimento exponencial nas violações de restrição, consistente com as previsões teóricas.

Significância e Alegações

O artigo afirma que o BayesFP fornece um amostrador unificado, em tempo de inferência e sem necessidade de retreinamento para políticas de difusão e de fluxo. Sua principal significância reside em:

• Orientação Fundamentada: Ao contrário dos métodos de combinação linear que amostram de uma soma de distribuições, o BayesFP amostra do verdadeiro produto (posterior) do prior e da verossimilhança de custo, garantindo um tratamento matematicamente consistente das restrições.
• Lidando com Não Convexidade: Ao amostrar trajetórias inteiras do posterior em vez de aplicar correções de gradiente local, o método lida naturalmente com obstáculos não convexos e espaços de solução multimodais.
• Generalidade: O framework é agnóstico à arquitetura da política base (difusão ou fluxo) e à natureza específica das restrições, desde que possam ser codificadas em uma função de custo diferenciável.

Os autores reconhecem limitações, observando uma lentidão computacional (1.25–2$\times$) devido à natureza sequencial da construção da Função de Distância Assinada (SDF) e à sobrecarga da simulação de partículas, embora isso seja mitigado pela execução paralela em GPU. Eles também observam a exigência de que a função de custo seja diferenciável. Sugere-se como trabalho futuro estender o framework para políticas de fluxo de streaming para controle de baixa latência e incorporar representações de restrições mais ricas.