PolicyFlow: Policy Optimization with Continuous Normalizing Flow in Reinforcement Learning

Resumo Técnico: PolicyFlow

Declaração do Problema

A Otimização de Política Próxima (PPO) é um algoritmo de aprendizado por reforço (RL) on-policy padrão, conhecido por sua estabilidade e simplicidade. No entanto, ela tipicamente depende de distribuições de política gaussianas, que têm dificuldade em representar distribuições de ações complexas, multimodais ou altamente assimétricas. Embora modelos generativos expressivos, como Fluxos de Normalização Contínua (CNFs) e modelos de difusão, possam capturar tais distribuições complexas, integrar-os ao PPO apresenta desafios significativos.

O principal gargalo é a necessidade de avaliação de verossimilhança para calcular as razões de importância ($\pi(a|s)/\hat{\pi}(a|s)$) para o objetivo de superfície do PPO. Para CNFs, a avaliação exata da verossimilhança exige a simulação de toda a trajetória da Equação Diferencial Ordinária (ODE) e a retropropagação por caminho (path-wise backpropagation). Este processo é computacionalmente caro, intensivo em memória e frequentemente instável numericamente devido a gradientes explosivos ou evanescentes. Tentativas existentes de aplicar atualizações do tipo PPO a políticas baseadas em fluxos (ex: FPO, DPPO) sofrem de viés de estimativa assimétrico, exigem grandes tamanhos de lote ou degradam o desempenho ao treinar do zero devido à dificuldade da exploração fora da variedade (off-manifold exploration). Além disso, técnicas padrão de regularização de entropia são difíceis de aplicar a políticas baseadas em fluxo porque o cálculo da entropia da política explicitamente requer integração de divergência dispendiosa.

Metodologia

Os autores propõem o PolicyFlow, um algoritmo de RL on-policy que integra políticas CNF expressivas com objetivos do tipo PPO sem exigir a avaliação da verossimilhança de trajetória completa. O método consiste em dois componentes principais:

1. Aproximação da Razão de Importância via Variações do Campo de Velocidade

Em vez de simular a trajetória completa da ODE para calcular o deslocamento terminal entre a política atual $\phi_1$ e a política de referência $\hat{\phi}_1$, o PolicyFlow aproxima a razão de importância usando as variações do campo de velocidade ao longo de um caminho de interpolação linear simples.

• Representação da Política: A política gera ações integrando um fluxo $\phi_t(z; s)$ governado por um campo de velocidade neural $v_t$, adicionando ruído gaussiano no tempo terminal: $a = \phi_1(z; s) + n$.
• Aproximação: O deslocamento terminal $\delta\phi_1 = \phi_1 - \hat{\phi}_1$ é aproximado integrando a diferença nos campos de velocidade $\delta v_t = v_t - \hat{v}_t$ ao longo de um caminho de interpolação linear $x_t = (1-t)z + t\hat{\phi}_1(z; s)$.
• Objetivo de Superfície: A razão de importância $\rho$ é estimada usando a razão de verossimilhança gaussiana das variações do campo de velocidade em vez das trajetórias de fluxo completas. Isso permite que o objetivo de treinamento seja estimado eficientemente ao longo do caminho de interpolação sem retropropagar através da simulação completa da ODE durante o passo de atualização. Análises teóricas indicam que isso introduz apenas um erro de primeira ordem na log-verossimilhança, que é naturalmente limitado pelo mecanismo de clipping do PPO.

2. Regularizador Browniano para Entropia

Para evitar o colapso de modo (mode collapse) e incentivar a exploração diversificada sem calcular explicitamente a entropia intratável do fluxo, os autores introduzem o Regularizador Browniano.

• Conceito: Inspirado no movimento browniano e na equação do calor, este regularizador incentiva o campo de velocidade aprendido a seguir a dinâmica de um processo de aumento de entropia.
• Mecanismo: Ele aproveita a relação entre o campo de velocidade e a função de pontuação (score function $\nabla \log p_t$) em fluxos retificados. O regularizador penaliza desvios do campo de velocidade aprendido de um campo de vetores alvo derivado do score do fluxo de referência.
• Formulação: O termo do regularizador $J_{Reg}$ inclui um termo que minimiza a diferença entre o campo de velocidade aprendido e um termo de alinhamento alvo (derivado do campo de velocidade de referência e do estado), somado a um termo que incentiva a variância do ruído gaussiano injetado. Esta abordagem evita cálculos caros de log-verossimilhança e injeção de ruído heurística, fornecendo uma forma leve e fundamentada de manter a diversidade da trajetória.

Principais Contribuições

1. Integração Eficiente de CNF-PPO: O PolicyFlow permite o uso de políticas CNF expressivas em RL on-policy ao aproximar as razões de importância via variações do campo de velocidade, eliminando a necessidade de retropropagação por caminho custosa através de trajetórias completas de ODE.
2. Regularizador Browniano: Um regularizador de entropia novo e computacionalmente leve, inspirado no movimento browniano, que mitiga o colapso de modo e incentiva comportamentos diversos sem exigir o cálculo explícito de entropia.
3. Validação Empírica: O método é validado em diversos ambientes, demonstrando desempenho competitivo ou superior em comparação ao PPO Gaussiano e aos basais de fluxo de última geração (FPO, DPPO).

Resultados Experimentais

Os autores avaliaram o PolicyFlow nos benchmarks MultiGoal, MuJoCo Playground e IsaacLab:

• MultiGoal (Teste de Colapso de Modo): Em um ambiente com seis metas simétricas, o PPO padrão (Gaussiano) falhou em atingir todas as metas. Os basais FPO e DPPO colapsaram para alguns modos devido à falta de uma regularização de entropia eficaz. O PolicyFlow, com o regularizador browniano, capturou com sucesso uma distribuição multimodal equilibrada, atingindo todas as metas com probabilidade aproximadamente igual.
• MuJoço Playground: O PolicyFlow alcançou desempenho comparável ou superior ao FPO e DPPO em várias tarefas (ex: WalkerWalk, CheetahRun). Frequentemente convergiu mais rápido, indicando maior eficiência de amostragem.
• IsaacLab: O PolicyFlow foi comparado ao PPO em tarefas de robótica (locomoção, manipulação). Ele consistentemente igualou ou superou o desempenho assintótico do PPO.
• Eficiência Computacional: Embora o PolicyFlow envolva simulação de ODE para amostragem, o passo de treinamento (que evita a retropropagação total da ODE) aumentou o tempo de treinamento por iteração em menos de 50% em comparação ao PPO em tarefas padrão, e permaneceu abaixo de duas vezes o custo mesmo com o aumento das dimensões de incorporação (embedding).

Significância e Alegações

O artigo afirma que o PolicyFlow oferece uma estrutura prática e expressiva para RL on-policy. Ao desacoplar a necessidade de avaliação de verossimilhança de trajetória completa da atualização da política, ele preenche a lacuna entre o poder de representação dos fluxos de normalização contínua e a estabilidade do PPO.

Os autores enfatizam que:

• A aproximação das razões de importância é teoricamente limitada e suficiente para um treinamento estável dentro do regime de clipping do PPO.
• O regularizador browniano fornece uma alternativa fundamentada à injeção de ruído heurística ou à integração de divergência dispendiosa para manter a exploração.
• O PolicyFlow aborda as limitações específicas dos métodos de fluxo baseados em RL existentes (FPO, DPPO) em relação à estabilidade de otimização, custo computacional e viés, tornando as políticas generativas expressivas mais acessíveis para aprendizado por reforço online.

O trabalho sugere que o PolicyFlow serve como uma base para unir modelagem generativa e aprendizado por reforço, com potenciais direções futuras, incluindo o ajuste fino (fine-tuning) de políticas de flow-matching, extensão para RL offline e exploração de conexões com modelos de difusão.