Multivariate Distributional Reinforcement Learning Using Sliced Divergences

Resumo Técnico: Aprendizado por Reforço Distribucional Multivariado Usando Divergências Fatiadas

Declaração do Problema

O Aprendizado por Reforço Distribucional (DRL) modela a distribuição completa do retorno em vez de apenas sua esperança, oferecendo vantagens teóricas e empíricas sobre o RL clássico baseado em esperança. No entanto, estender o DRL para configurações multivariadas (onde os retornos são vetores de $d$ dimensões, $d > 1$) permanece um desafio significativo.

As abordagens existentes enfrentam três limitações primárias:

1. Intratabilidade Computacional: Métricas comuns como a distância de Wasserstein sofrem com altos custos computacionais em altas dimensões (tipicamente $O(n^3 \log n)$ para solvers de transporte ótimo) e taxas estatísticas ruins que dependem da dimensão ambiente.
2. Falta de Garantias de Contração: Embora o DRL escalar possua resultados de contração bem estabelecidos sob o operador de Bellman, configurações multivariadas com desconto matricial geral (onde o fator de desconto é uma matriz dependente de estado-ação $\Gamma(s, a) \in \mathbb{R}^{d \times d}$) carecem de garantias rigorosas de contração. Métodos existentes frequentemente falham em contrair sob descontos anisotrópicos (não uniformes).
3. Viés de Gradiente em Treinamento Estocástico: Muitas divergências, particularmente a de Wasserstein, falham na propriedade de gradiente de amostra não enviesado (U) quando o alvo de Bellman é instanciado a partir de um único sucessor amostrado (aprendizado TD padrão). Isso leva a gradientes enviesados e correspondência distribucional subótima.

Metodologia: Aprendizado por Reforço Distribucional Fatiado (SDRL)

Os autores introduzem o Aprendizado por Reforço Distribucional Fatiado (SDRL), um framework que eleva divergências unidimensionais tratáveis para distribuições de retorno multivariadas via projeções aleatórias (fatiamento).

Mecanismo Central

Em vez de calcular uma divergência de alta dimensão diretamente, o SDRL projeta as distribuições multivariadas em direções aleatórias $\theta \in S^{d-1}$, calcula uma divergência base $\Delta$ em uma dimensão e agrega os resultados.

• Fatiamento Uniforme: Médias a divergência base sobre $L$ direções aleatórias:
  $$S\Delta_p^p(\mu, \nu) = \int_{S^{d-1}} \Delta^p((P_\theta)_\#\mu, (P_\theta)_\#\nu) d\sigma(\theta)$$
• Max-Fatiamento (MSDRL): Otimiza sobre a direção mais discriminativa para obter garantias de contração mais fortes para descontos matriciais gerais:
  $$MS\Delta(\mu, \nu) = \sup_{\theta \in S^{d-1}} \Delta((P_\theta)_\#\mu, (P_\theta)_\#\nu)$$

O framework suporta uma ampla classe de divergências base, especificamente analisando Wasserstein ($W_p$), distância de Cramér ($C_2$) e Maximum Mean Discrepancy (MMD).

Contribuições Teóricas

O artigo estabelece fundamentos teóricos rigorosos para o SDRL:

1. Propriedades Métricas: Prova que tanto os levantamentos fatiados uniformes quanto os maximizados preservam os axiomas métricos se a divergência base for uma métrica.
2. Garantias de Contração:
   • Fatiamento Uniforme: Prova a contração de Bellman sob desconto escalar compartilhado ($\Gamma = \gamma I$) para qualquer divergência base que satisfaça não-expansão de translação, contração de escala e $p$-convexidade de mistura.
   • Max-Fatiamento: Introduz um novo resultado de contração para desconto matricial denso geral (atualizações anisotrópicas). Prova que o max-fatiamento gera contração sob a norma de operador da matriz de desconto, um cenário onde o fatiamento uniforme e o MMD padrão frequentemente falham em contrair.
3. Complexidade de Amostra: Demonstra que o fatiamento uniforme herda taxas de convergência independentes da dimensão da divergência base de 1D, evitando a maldição da dimensionalidade. As taxas de max-fatiamento escalam polinomialmente com a dimensão ($O(d \log n / n)$), o que é significativamente melhor que o OT exato de alta dimensão.
4. Análise de Viés de Gradiente:
   • Fatiamento Uniforme: Preserva a propriedade de gradiente de amostra não enviesado (U) se a divergência base a satisfizer. Isso torna o Cramér Fatiado e o MMD Fatiado compatíveis com o bootstrapping TD de um único amostrador.
   • Max-Fatiamento: Falha na propriedade (U) devido ao viés de seleção introduzido pelo passo de maximização. A direção selecionada depende do lote de amostra específico, levando a gradientes enviesados em configurações TD de amostra única padrão.

Resultados Experimentais

Os autores avaliam o SDRL em três benchmarks: uma cadeia MDP tabular, um ambiente de labirinto baseado em pixels e um subconjunto de jogos de Atari.

1. Ambiente de Cadeia (Avaliação de Política):

   • Compara TD de uma amostra (bootstrapping padrão) vs. TD quase-exato (construção explícita de mistura).
   • Descoberta: Objetivos que satisfazem a propriedade (U) (ex: Cramér Fatiado, MMD Fatiado) alcançam correspondência distribucional precisa no regime de uma amostra. Objetivos que violam (U) (ex: Wasserstein Fatiado, variantes de Max-Fatiamento) degradam significativamente, a menos que o alvo de mistura completo seja usado.
   • Max-Fatiamento: Desempenha bem no caso degenerado unidimensional (onde o viés de seleção desaparece), mas falha em configurações multivariadas gerais sob TD de uma amostra.

2. Labirinto Baseado em Pixels (Avaliação de Política):

   • Avalia a precisão distribucional usando a distância Wasserstein-2 empírica entre os retornos de Monte Carlo e a previsão do crítico.
   • Descoberta: Cramér Fatiado e MMD Fatiado performam em nível de seus equivalentes não fatiados e superam significativamente o Wasserstein Fatiado no regime TD de uma amostra, confirmando a importância da propriedade (U).

3. Jogos de Atari (Controle):

   • Avalia o desempenho de controle em recompensas multidimensionais decompostas.
   • Descoberta: Embora objetivos que satisfazem (U) forneçam estimativas de retorno precisas, o Wasserstein-2 Fatiado surpreendentemente alcança um forte desempenho de controle, apesar da correspondência distribucional pobre e gradientes enviesados. Isso sugere que, para controle, a estimativa precisa da esperança do retorno pode ser mais crítica do que a fidelidade da distribuição completa, mesmo que o objetivo distribucional seja enviesado.

Significância e Alegações

O artigo afirma fornecer um framework principiado para estender o DRL para configurações multivariadas, mantendo a tratabilidade computacional e garantias teóricas.

• Recomendação Prática: Os autores identificam o Cramér Fatiado como uma escolha padrão forte para aprendizado distribucional multivariado. É computacionalmente eficiente ($O(n \log n)$), satisfaz a propriedade de gradiente não enviesado (permitindo treinamento TD padrão) e oferece garantias de contração no cenário isotrópico padrão.
• Insight Teórico: O trabalho esclarece os compromissos entre diferentes escolhas de divergência. Destaca que, embora o Max-Fatiamento resolva o problema de contração para descontos matriciais gerais, ele introduz um viés de gradiente que o torna inadequado para o TD de uma amostra padrão sem modificações.
• Identificação de Lacunas: O artigo conclui que uma divergência que seja simultaneamente provadamente contraível sob descontos anisotrópicos gerais (apenas no sentido de norma), satisfaça a propriedade de gradiente de amostra não enviesado e ofereça estimativa tratável, permanece uma questão em aberto.

Em resumo, o SDRL oferece uma abordagem escalável e teoricamente fundamentada para o DRL multivariado, com o Cramér Fatiado emergindo como a escolha prática mais robusta para pipelines de aprendizado por reforço padrão.