Parameter-Efficient Distributional RL via Normalizing Flows and a Geometry-Aware Cramér Surrogate

Resumo Técnico: NFDRL

Declaração do Problema
O Aprendizado por Reforço Distribucional (DistRL) visa modelar a distribuição completa do retorno em vez de apenas o retorno esperado, fornecendo sinais mais ricos para incerteza, risco e variabilidade. No entanto, os métodos existentes enfrentam um compromisso entre expressividade e eficiência de parâmetros.

• Métodos categóricos (ex.: C51): Baseiam-se em suportes discretos fixos. Embora eficazes, carecem de parcimônia; a contagem de parâmetros escala linearmente com a resolução desejada (número de átomos). Eles também impõem limites rígidos aos retornos, limitando a aplicabilidade em domínios com recompensas dinâmicas ou de cauda pesada.
• Métodos de quantis (ex.: QR-DQN, IQN): Aprendem localizações de quantis para lidar com retornos ilimitados. No entanto, aproximam distribuições como misturas de deltas de Dirac (funções degrau). Essa representação pode ser ineficiente para distribuições suaves, multimodais ou de cauda pesada, frequentemente exigindo muitos quantis para aproximar densidades simples e falhando em capturar formas distribucionais de alta granularidade.
• Métodos de Densidade Contínua: Embora modelar funções de densidade de probabilidade (PDFs) contínuas ofereça suporte adaptativo e suavidade, abordagens anteriores lutaram para definir um objetivo de treinamento que seja simultaneamente teoricamente sólido (garantindo contração do operador de Bellman) e praticamente otimizável (fornecendo gradientes de amostra não viesados).

Metodologia: NFDRL
Os autores propõem o NFDRL, uma arquitetura DistRL eficiente em parâmetros que modela distribuições de retorno usando fluxos normais contínuos.

1. Arquitetura:

   • O modelo utiliza um fluxo normal condicional para mapear amostras de ruído base $z \sim U$ para resultados de retorno $y = f_\theta(z)$.
   • O fluxo é construído via uma Função de Distribuição Acumulada (CDF) de Misto Gaussiano. Uma rede neural $h_\theta$ mapeia um estado $x$ para os parâmetros de uma CDF de mistura para cada ação.
   • Para abordar a natureza limitada das CDFs ($[0, 1]$), uma transformação afim adicional reescala a saída para uma faixa dinâmica $[-G_{max}, G_{max}]$, permitindo que o modelo se adapte ao suporte de retornos do mundo real sem limites fixos.
   • A distribuição alvo é construída compondo o fluxo previsto com a função de bootstrap ($r + \gamma y$), tratada como uma camada de fluxo invertível para preservar a normalização via fórmula de mudança de variáveis.

2. Objetivo de Treinamento (Surrogado de Cramér Consciente de Geometria):

   • A divergência KL padrão não é adequada para suportes não sobrepostos, e a minimização exata de Wasserstein sofre de gradientes viesados.
   • Os autores introduzem uma distância inspirada em Cramér e consciente de geometria definida sobre massas de probabilidade em vez de densidades brutas.
   • A função de perda é definida como:
     $$D(\eta, \hat{T}\eta) = \left( \frac{1}{N^2} \sum_{i,j} (w_i - v_i)^2 \cdot |y_i - y_j| \right)^{1/2}$$
     onde $w_i$ e $v_i$ são massas de probabilidade nos pontos de suporte $y_i$.
   • Para lidar com o desalinhamento entre os suportes previsto e alvo (que diferem devido à etapa de bootstrap), o método emprega Estimação de Densidade por Kernel (KDE) para avaliar densidades em suportes compartilhados antes de calcular a perda.

Principais Contribuições

1. Eficiência de Parâmetros: Ao contrário de baselines categóricas onde o tamanho do modelo cresce com a resolução, o NFDRL mantém uma pegada de parâmetros constante independentemente da resolução efetiva da distribuição, oferecendo resolução efetiva infinita e suporte dinâmico.
2. Solvência Teórica: Os autores provam que sua distância baseada em massa é uma métrica verdadeira e que o operador de Bellman distribucional associado é uma contração de $\sqrt{\gamma}$. Isso garante convergência para um ponto fixo único.
3. Gradientes Não Viesados: O objetivo proposto admite gradientes de amostra não viesados, uma propriedade tipicamente ausente em métodos DistRL baseados em PDF anteriores (diferentemente de Wasserstein, que frequentemente requer aproximações viesadas).
4. Expressividade: A natureza contínua do fluxo permite que o modelo capture paisagens de retorno complexas e multimodais que métodos baseados em quantis frequentemente borram ou falham em resolver.

Resultados

• MDPs de Brinquedo: Em ambientes controlados, o NFDRL recupera com sucesso paisagens de retorno ricas e multimodais (ex.: distribuições bimodais) com PDFs suaves, enquanto as baselines de quantis (IQN) produziram distribuições ruidosas e borradas que obscureciam a modalidade.
• Benchmark Atari-5: O NFDRL alcança desempenho competitivo com baselines categóricas (C51) e aproxima-se do desempenho dos principais métodos de quantis (IQN).
  • No benchmark Atari-5, o NFDRL (variante surrogate) alcançou uma pontuação média normalizada por humanos de 418, superando o C51 (330) e o DQN (189), embora fique ligeiramente atrás do IQN (525).
  • Notavelmente, o NFDRL-S superou o IQN em jogos específicos como Double Dunk e QBert*.
• Eficiência: O NFDRL iguala o desempenho do C51 com significativamente menos parâmetros. Enquanto a contagem de parâmetros do C51 escala linearmente com o número de átomos, o NFDRL mantém um tamanho fixo, igualando a capacidade do C51 com apenas 11 átomos.

Significado e Alegações
O artigo posiciona o NFDRL não como um substituto universal para agentes fortes baseados em quantis como o IQN no Atari, mas como uma alternativa contínua e parcimoniosa com compromissos distintos. Seu significado principal reside em oferecer um framework teoricamente rigoroso (contração e gradientes não viesados) para distribuições de retorno contínuas com uma pegada de parâmetros constante. Isso é particularmente valioso em cenários onde decisões a jusante dependem da forma detalhada da distribuição de retorno (ex.: multimodalidade, risco de cauda) em vez de apenas do valor esperado. Os autores reconhecem limitações, incluindo maior variância de treinamento, convergência mais lenta em comparação com algumas baselines e o custo computacional da etapa de alinhamento KDE, vendo esses aspectos como áreas para refinamento futuro.