ScoRe-Flow: Complete Distributional Control via Score-Based Reinforcement Learning for Flow Matching

Resumo Técnico: ScoRe-Flow

1. O Problema

As políticas baseadas em Flow Matching (FM) tornaram-se uma espinha dorsal eficiente para o controle robótico, oferecendo geração de ações rápida e expressiva, superando os modelos de difusão tradicionais em termos de velocidade de amostragem (requerendo 4 a 10 vezes menos passos de integração). No entanto, a maioria das políticas FM é treinada apenas via Aprendizado por Imitação (IL), o que as limita ao desempenho das demonstrações humanas (subótimas) e as torna suscetíveis a mudanças de covariância.

Para superar essas limitações, o Aprendizado por Reforço (RL) é necessário para o ajuste fino (fine-tuning). Contudo, aplicar RL em políticas FM apresenta um desafio fundamental:

• As dinâmicas FM originais são determinísticas (equações diferenciais ordinárias - ODEs), o que faz com que as transições colapsem para medidas de Dirac. Isso torna as verossimilhanças das ações mal definidas, impedindo a aplicação direta de métodos de gradiente de política (como PPO).
• Métodos recentes (ex: ReinFlow) convertem o fluxo determinístico em uma Equação Diferencial Estocástica (SDE) injetando ruído aprendível. No entanto, essa abordagem controla apenas a variância (ruído), perturbando a posição sem modular a dinâmica de fluxo (o "drift"). Isso é ineficiente quando os dados de demonstração já fornecem um prior forte, pois o ruído apenas adiciona perturbações aleatórias em vez de guiar a exploração de forma geométrica.

2. Metodologia: ScoRe-Flow

O ScoRe-Flow (Score-based Reinforcement Learning for Flow Matching) propõe um método de ajuste fino via RL que alcança o controle distribucional completo, desacoplando o controle da média e da variância das transições estocásticas.

Principais Componentes Técnicos:

1. Dualidade Score-Velocidade em Fluxo Linear:

   • O método explora a relação conhecida entre a função score (gradiente da densidade logarítmica, $\nabla \log \rho_t$) e o campo de velocidade ($v_\theta$) em caminhos de fluxo linear.
   • Diferente de modelos de difusão que exigem redes auxiliares para estimar o score, o ScoRe-Flow deriva uma expressão de fechamento (closed-form) para o score diretamente a partir do campo de velocidade pré-treinado:
     $$s_t(a) = \frac{t \cdot v_\theta(t, a, s) - a}{1 - t}$$
   • Isso permite o controle do drift (tendência) sem adicionar parâmetros ou redes neurais extras.

2. Modulação de Drift Baseada em Score:

   • Em vez de apenas injetar ruído, o ScoRe-Flow modula o termo de deriva (drift) da SDE usando o score. O score aponta para regiões de alta densidade de probabilidade, guiando a exploração de volta para a variedade de dados (data manifold) quando a trajetória se desvia.
   • Para garantir estabilidade numérica (já que o score diverge quando $t \to 1$), o método aplica um decaimento temporal rígido no multiplicador do score: $\alpha_{scaled}(t) = (1-t) \cdot \alpha_\psi(t)$.

3. Controle Desacoplado de Média e Variância:

   • O método introduz dois componentes aprendíveis:
     • Um agendador de score ($\alpha_\psi$) que modula a correção do drift (controle da média/direção).
     • Um preditor de variância ($\sigma_\phi$) que controla independentemente a magnitude da exploração estocástica.
   • A equação SDE resultante é:
     $$da_t = [v_\theta + \alpha_{scaled}(t) \cdot s_t(a_t)] dt + \sigma_\phi(t, a_t, s) dW_t$$
   • Essa desacoplagem permite que a política mantenha uma correção direcional forte enquanto reduz o ruído de exploração conforme o treinamento avança, algo impossível em métodos baseados apenas em ruído ou apenas em score.

4. Otimização:

   • O método utiliza o algoritmo PPO (Proximal Policy Optimization) para otimizar a verossimilhança da trajetória, ajustando conjuntamente o campo de velocidade base ($v_\theta$), o agendador de score e o preditor de variância.

3. Contribuições Principais

• Observação Teórica: Identificar que a conversão SDE baseada apenas em ruído não modula a dinâmica do fluxo e que a função score pode ser calculada em fechamento a partir do campo de velocidade FM, permitindo controle de drift sem redes auxiliares.
• Algoritmo Novo (ScoRe-Flow): Proposta de um método que combina modulação de drift baseada em score com predição de variância aprendida, alcançando controle desacoplado sobre a média e a variância das transições estocásticas.
• Eficiência Computacional: Elimina a necessidade de redes auxiliares para estimar o score, adicionando sobrecarga computacional negligenciável (apenas ~800 parâmetros para o agendador).

4. Resultados Experimentais

O método foi avaliado em benchmarks de locomoção (D4RL) e manipulação (Robomimic e Franka Kitchen), comparado com o estado da arte (ReinFlow, DPPO, FQL).

• Convergência Mais Rápida: No D4RL, o ScoRe-Flow alcançou uma convergência 2,4 vezes mais rápida que o ReinFlow (o SOTA baseado em fluxo) com o mesmo número de passos de denoising ($K=4$).
• Desempenho Superior:
  • Locomoção (D4RL): Obteve o melhor desempenho médio, superando o ReinFlow e o DPPO (baseado em difusão), especialmente na tarefa complexa Humanoid-v3 (5100 pontos vs 5041 do ReinFlow).
  • Manipulação (Robomimic): Alcançou taxas de sucesso superiores em todas as tarefas, com uma melhoria de até 5,4% na taxa de sucesso em tarefas de manipulação visual.
  • Franka Kitchen: Demonstrou eficácia em tarefas de múltiplos passos com recompensas esparsas, completando mais subtarefas que os baselines.
• Eficiência de Tempo: Devido à natureza baseada em fluxo (menos passos de integração), o ScoRe-Flow é estruturalmente mais rápido que métodos baseados em difusão (como DPPO), com uma vantagem de tempo de parede de até 21,9x em média.

5. Significado e Impacto

O ScoRe-Flow representa um avanço significativo na aplicação de RL online em políticas generativas baseadas em fluxo. Ao demonstrar que é possível modular a dinâmica do fluxo (drift) de forma eficiente e geometricamente fundamentada usando a dualidade score-velocidade, o trabalho resolve o problema de instabilidade e ineficiência na exploração de políticas FM.

Isso permite que sistemas de IA corporificada (embodied AI) aprendam comportamentos complexos e multi-modais que superam as demonstrações humanas, mantendo a velocidade de inferência necessária para controle em tempo real. O método oferece um novo paradigma para o ajuste fino de políticas generativas, onde a direção da exploração e a sua magnitude são controladas de forma independente e adaptativa.