Resumo Técnico: Otimização de Política One-step MeanFlow Baseada em Score (SOM)

1. Formulação do Problema

O aprendizado por reforço (RL) em espaços de ação contínuos tem adotado cada vez mais modelos generativos, como difusão e correspondência de fluxo, para superar as limitações de expressividade de políticas Gaussianas unimodais. Esses modelos podem capturar distribuições de ação complexas e multimodais, permitindo exploração mais rica e comportamentos ótimos diversos. No entanto, um gargalo crítico permanece: políticas padrão de difusão e correspondência de fluxo exigem desruído iterativo e multi-etapa para gerar uma única ação. Isso introduz uma sobrecarga computacional substancial no tempo de inferência, o que é particularmente prejudicial no RL online, onde a inferência da política é continuamente intercalada com a coleta de dados.

Embora o MeanFlow tenha sido proposto como uma solução para reduzir a amostragem a uma única avaliação de rede aprendendo um campo de velocidade médio, sua aplicação direta ao RL online é impedida por um requisito fundamental: a construção do campo de velocidade alvo tipicamente necessita de amostras da distribuição alvo (a política ótima). No RL online, essas amostras de verdade fundamental não estão disponíveis, pois o agente deve aprender a distribuição alvo simultaneamente à política. Abordagens existentes que tentam contornar isso (por exemplo, amostragem Best-of-N) frequentemente lutam com a escalabilidade à medida que a dimensionalidade da ação aumenta.

2. Metodologia

Os autores propõem a Otimização de Política One-step MeanFlow Baseada em Score (SOM), um algoritmo ator-crítico que permite o treinamento de MeanFlow em RL totalmente online sem exigir amostras explícitas da distribuição alvo. A inovação central reside na construção do campo de velocidade alvo diretamente a partir da função Q via estimativa de score.

2.1 Função de Energia Dependente do Tempo

O método trata a função Q negativa como a energia da distribuição alvo. Sob um framework de RL regularizado por entropia, a política ótima segue uma distribuição de Boltzmann:
$$\pi^*(a|s) \propto \exp(\alpha Q(s, a))$$
onde $\alpha$ é um parâmetro de temperatura. Os autores definem uma função de energia dependente do tempo $Q_t(s, a_t)$ que preserva essa estrutura de Boltzmann ao longo dos passos de tempo de difusão $t \in [0, 1]$. O score da distribuição em qualquer passo de tempo $t$ é derivado como:
$$\nabla_{a_t} \log p_t(a_t|s) = \nabla_{a_t} Q_t(s, a_t)$$
Crucialmente, embora calcular $Q_t$ em forma fechada seja intratável, o método requer apenas seu gradiente, que pode ser aproximado.

2.2 Estimativa de Score via Adaptação iDEM

Para estimar o score $\nabla_{a_t} Q_t(s, a_t)$ sem amostras alvo, o SOM adapta o estimador Iterated Denoising Energy Matching (iDEM). Em vez de amostrar da distribuição alvo verdadeira, o método amostra ações candidatas a partir de uma proposta Gaussian centrada no estado ruidoso atual $a_t$. Essas candidatas são reponderadas com base em seus valores-Q para aproximar a expectativa posterior. O estimador de score resultante é:
$$\nabla_{a_t} Q_t(s, a_t) \approx \nabla_{a_t} \log \sum_{i=1}^K \exp(\alpha Q(s, a^{(i)}_{0|t}))$$
onde $a^{(i)}_{0|t}$ são amostras extraídas de uma distribuição Gaussian condicionada em $a_t$. Isso efetivamente viésa as atualizações de ação para regiões de alto valor, mantendo a exploração através do ruído.

2.3 Construção do Campo de Velocidade

O score estimado é integrado a uma Equação Diferencial Ordinária de Fluxo de Probabilidade (PF-ODE) para derivar o campo de velocidade alvo. Para uma EDE de Preservação de Variância (VP), o campo de velocidade é definido como:
$$v_t(a_t|s) = -\frac{1}{2}\beta(t) \left( a_t + \nabla_{a_t} Q_t(s, a_t) \right)$$
Este campo de velocidade transporta o ruído diretamente para modos de alto valor induzidos pela função Q atual. A política MeanFlow $u_\theta$ é então treinada para satisfazer a identidade MeanFlow usando este campo de alvo construído, minimizando a perda:
$$\mathcal{L}(\theta) = \mathbb{E} \left[ \| u_\theta(a_t, r, t, s) - \text{sg}(u^{\text{new}}_{\text{tgt}}) \|^2_2 \right]$$
onde $u^{\text{new}}_{\text{tgt}}$ é derivado da velocidade da PF-ODE.

2.4 Detalhes de Implementação

• Normalização: O gradiente do score é normalizado para norma unitária e reescalado por um coeficiente ajustável $w$ para estabilizar o treinamento em diversas escalas de recompensa.
• Treinamento do Crítico: O crítico é atualizado usando um objetivo de Bellman distribucional com redes Q gêmeas.
• Best-of-N (BoN): Embora a contribuição principal seja a velocidade baseada em score, os autores incorporam amostragem BoN durante as atualizações do crítico para acelerar o treinamento, selecionando a candidata de maior Q entre múltiplas amostras para a estimativa do valor alvo.

3. Contribuições Principais

1. Identificação de um Gargalo: O artigo identifica a falta de amostras da distribuição alvo como a principal barreira para aplicar MeanFlow no RL online.
2. Construção de Velocidade Baseada em Score: Propõe um método inovador para construir o campo de velocidade alvo diretamente a partir de gradientes da função Q, contornando a necessidade de amostras alvo explícitas.
3. Algoritmo SOM: A introdução do SOM, um algoritmo ator-crítico que integra MeanFlow com RL online, alcançando aprendizado de política expressivo com amostragem de ação em passo único.
4. Validação Empírica: Demonstração de que o SOM alcança desempenho state-of-the-art em tarefas de locomoção, reduzindo significativamente o tempo de treinamento e inferência em comparação com baselines iterativas de difusão e correspondência de fluxo.

4. Resultados Experimentais

Os autores avaliaram o SOM em nove tarefas de controle contínuo MuJoCo v4 (incluindo Ant, HalfCheetah, Hopper, Humanoid, Walker2d, Swimmer, Pusher, Reacher e InvertedDoublePendulum).

• Desempenho: O SOM alcançou o retorno médio mais alto em quatro de cinco benchmarks de locomoção e o segundo mais alto em Hopper-v4. A melhoria foi mais pronunciada em Humanoid-v4 (alta dimensionalidade de ação), onde o SOM superou o segundo melhor método por uma margem clara, sugerindo que a abordagem baseada em gradiente escala favoravelmente com a dimensionalidade.
• Eficiência: Como mostrado na Tabela 2, o SOM alcança um tempo de inferência de passo único de 0,218 ms em HalfCheetah-v4, comparado a 0,6–1,5 ms para baselines iterativas (que exigem 10–100 passos de desruído). O tempo de treinamento também é substancialmente reduzido.
• Análise da Política: Visualizações (t-SNE e paisagens Q) indicam que o SOM concentra amostras de ação firmemente ao redor de modos de alto Q, enquanto métodos concorrentes de passo único (como MFP) tendem a colapsar para modos únicos ou distribuir-se difusamente.
• Robustez: Experimentos com funções Q ruidosas demonstraram que o estimador de score suavizado por Gaussian no SOM permanece robusto, efetivamente mediando erros grosseiros do crítico enquanto preserva a estrutura local perto de modos de alto valor.

5. Significado e Alegações

O artigo alega que o SOM resolve a tensão entre a expressividade de políticas generativas e as restrições computacionais do RL online. Ao derivar a velocidade alvo diretamente da função Q, o SOM permite que políticas MeanFlow sejam treinadas em configurações totalmente online onde amostras alvo estão indisponíveis.

Os autores enfatizam que o SOM alcança desempenho state-of-the-art com um único passo de geração, oferecendo uma solução prática para aplicações de controle em tempo real onde a latência de inferência é crítica. Eles notam que, embora o método seja atualmente limitado a tarefas de controle contínuo (MuJoCo) e dependa da qualidade do crítico aprendido, ele representa um passo significativo em direção à otimização de política eficiente e expressiva. O trabalho sugere que direções futuras incluem estender o SOM para configurações de dimensionalidade mais alta ou parcialmente observáveis (por exemplo, controle visual) e integrar críticos mais avançados conscientes de incerteza.