Distributional Inverse Reinforcement Learning

Resumo Técnico: Distributional Inverse Reinforcement Learning (DistIRL)

1. O Problema

O Aprendizado por Reforço Inverso (IRL) tradicional visa inferir a função de recompensa e a política de um especialista a partir de trajetórias observadas. No entanto, a maioria dos métodos existentes (como MaxEntIRL e variantes Bayesianas) assume que a função de recompensa é determinística ou recupera apenas uma estimativa pontual (média).

Essa suposição é limitante em cenários do mundo real onde:

• Os sinais de recompensa são inerentemente estocásticos (ex.: incerteza em manipulação robótica de objetos frágeis, variabilidade na atividade de neurônios dopaminérgicos).
• A política do especialista pode ser sensível ao risco (aversão ou busca de risco), o que não pode ser capturado apenas pela maximização do retorno esperado (média).
• Métodos existentes falham em recuperar a distribuição completa da recompensa, ignorando momentos de ordem superior (variância, assimetria, caudas), essenciais para comportamentos robustos e seguros.

O desafio central é: como aprender a distribuição completa da recompensa e uma política consciente de risco diretamente de demonstrações offline, sem interação com o ambiente?

2. Metodologia

Os autores propõem o DistIRL, um framework que modela conjuntamente a incerteza sobre as funções de recompensa e as distribuições completas de retorno. A abordagem é dividida em três componentes principais:

A. Formulação Baseada em Dominância Estocástica de Primeira Ordem (FSD)
Em vez de apenas igualar os retornos esperados (média), o DistIRL busca alinhar as distribuições completas de retorno do agente e do especialista.

• Utiliza o critério de Dominância Estocástica de Primeira Ordem (FSD). Se a distribuição de retorno do especialista $Z_E$ domina a do agente $Z_\pi$ ($Z_E \succeq_{FSD} Z_\pi$), significa que o especialista tem uma probabilidade maior de obter retornos altos em qualquer nível de corte.
• O objetivo de aprendizado da recompensa minimiza a violação do FSD, definida como a integral da diferença positiva entre as funções de distribuição acumulada (CDF):
  $$\mathcal{L}_{FSD} = \int_{-\infty}^{\infty} [F_{Z_E}(z) - F_{Z_\pi}(z)]_+ dz$$
• Isso é tratado como uma função de energia em um Modelo Baseado em Energia (EBM), permitindo inferir a distribuição posterior da recompensa via Inferência Variacional.

B. Aprendizado de Política Consciente de Risco (Risk-Aware)
Para otimizar a política, o método substitui a função indicadora não diferenciável (que marca violações de FSD) por Medidas de Risco de Distorção (DRM).

• O objetivo da política torna-se maximizar uma medida de risco $M_\xi$ do retorno, onde $\xi$ é uma função de distorção que pondera diferentes partes da distribuição (ex.: focar em caudas inferiores para aversão ao risco).
• A função de perda da política é:
  $$\mathcal{L}_\pi(\phi) = \int_0^1 F^{-1}_{Z_\pi}(v) d\tilde{\xi}(v) + H(\pi_\phi)$$
  onde $F^{-1}$ é a função quantílica, $\tilde{\xi}$ é a função de distorção dual e $H$ é a entropia.

C. Algoritmo Prático
O algoritmo opera em um loop alternado (offline):

1. Modelagem da Recompensa: Parametriza a distribuição de recompensa $q_\phi(r|s,a)$ (ex.: distribuição normal assimétrica de Azzalini) e atualiza os parâmetros $\phi$ minimizando a perda de FSD baseada em amostragem de Monte Carlo.
2. Avaliação do Crítico (Critic): Aprende a distribuição de retorno $Z_\pi$ usando Regressão Quantílica (Quantile Regression) com perda Huber Quantílica, estimando os quantis da distribuição de retorno.
3. Atualização da Política: Atualiza a política $\pi_\phi$ maximizando a medida de risco (DRM) derivada da distribuição de retorno aprendida.

3. Principais Contribuições

1. Aprendizado de Distribuição de Recompensa: É o primeiro trabalho a recuperar a distribuição completa da função de recompensa (não apenas parâmetros de uma distribuição fixa) de forma principial no IRL offline, capturando momentos de ordem superior.
2. Política Consciente de Distribuição: Estende o IRL para aprender políticas que são sensíveis ao risco, alinhando-se com comportamentos que evitam ou buscam variabilidade, não apenas maximizando a média.
3. Análise Teórica: Fornece uma prova de convergência do algoritmo, demonstrando uma complexidade de iteração de $O(\epsilon^{-2})$ para atingir um ponto estacionário de primeira ordem.
4. Validação Empírica Abrangente:
   • Gridworld: Demonstra a capacidade de inferir variância e média corretamente em ambientes com recompensas estocásticas.
   • Dados Neurocomportamentais (Ratos): Aplica o método a dados reais de dopamina. Mostra que o DistIRL consegue recuperar a distribuição assimétrica dos sinais de dopamina a partir do comportamento, algo que métodos determinísticos falham em fazer.
   • MuJoCo (Robótica): Alcança desempenho state-of-the-art em tarefas de controle offline com recompensas estocásticas e penalidades de segurança, superando baselines como BIRL, ValueDICE e BC.

4. Resultados Chave

• Gridworld: O DistIRL recuperou com precisão a média e a variância das recompensas, enquanto o BIRL (Bayesian IRL) recuperou a média, mas falhou em capturar a variância, produzindo estimativas espúrias.
• Neurociência: Ao analisar o comportamento espontâneo de ratos e os sinais de dopamina, o modelo S-DistIRL (com distribuição normal assimétrica) alcançou a maior correlação (Pearson ~0.3) entre a recompensa inferida e a dopamina real, e a menor distância Wasserstein-1, indicando que a estrutura completa da distribuição de recompensa é crucial para modelar sinais neurológicos.
• MuJoCo (D4RL): Em tarefas com penalidades estocásticas (ex.: falhas catastróficas em alta velocidade), o DistIRL superou consistentemente métodos baseados em MaxEntIRL e aprendizado por imitação (BC), alcançando retornos próximos aos do especialista em ambientes de risco.
• Ablação: Estudos mostraram que a combinação de recompensas distribucionais com perda FSD é superior a métodos que apenas igualam a média, especialmente para políticas avessas ao risco.

5. Significado e Impacto

O DistIRL representa um avanço fundamental ao reconhecer que, em muitos domínios reais (robótica, neurociência, finanças), a incerteza não é apenas ruído, mas uma característica estrutural que define o comportamento.

• Segurança e Robustez: Permite aprender políticas que evitam cenários de baixa probabilidade, mas alto impacto (caudas da distribuição), essencial para robótica autônoma.
• Interpretabilidade Biológica: Oferece uma nova ferramenta para decodificar estados motivacionais internos e substratos neurais (como a variabilidade da dopamina) apenas observando o comportamento, sem necessidade de medições diretas.
• Generalização: O framework é geral o suficiente para se reduzir ao IRL determinístico padrão quando necessário, mas oferece uma capacidade expressiva superior quando a estocasticidade é relevante.

Em suma, o paper estabelece um novo paradigma para o IRL offline, movendo-se da estimativa de pontos (médias) para a modelagem de distribuições completas, permitindo uma imitação mais fiel e segura de comportamentos complexos.