Resumo Técnico: Utilizando a Incerteza da Recompensa para Induzir Comportamento Diverso em Aprendizado por Reforço

Declaração do Problema

O aprendizado por reforço (RL) clássico busca tipicamente uma política determinística que maximiza a soma esperada de uma recompensa escalar. Embora essa premissa garanta a existência de uma política determinística ótima, ela apresenta limitações significativas em aplicações modernas, como o ajuste fino de grandes modelos de linguagem (LLMs) e a descoberta científica. Nesses domínios, a convergência para um único comportamento é frequentemente indesejável. Por exemplo, preservar a diversidade das políticas pré-treinadas é crucial para tarefas criativas e raciocínio de cadeia de pensamento (chain-of-thought). Além disso, quando os modelos de recompensa são proxies imperfeitos para as verdadeiras preferências subjacentes, a otimização excessiva em direção a uma única política determinística pode levar à deterioração da qualidade da geração (reward hacking).

Remédios existentes, como a regularização de entropia ou bônus de diversidade, tentam impor estocasticidade, mas sofrem de desvantagens críticas:

1. Compromissos Frágeis: A regularização de entropia força um compromisso onde o aumento da estocasticidade ocorre às custas da redução das recompensas esperadas.
2. Ordenação Incorreta de Políticas: Bônus de diversidade frequentemente resultam em classificações indesejáveis, onde políticas subótimas com alta diversidade são preferidas em relação a políticas com recompensas médias mais elevadas.
3. Inexpressividade: Muitos objetivos, incluindo o gradiente de política vanilla e o retorno escalarizado esperado (ESR), possuem inerentemente políticas determinísticas ótimas, falhando em induzir a diversidade necessária.
4. Desafios de Otimização: Formulações de retorno escalarizado (SER) frequentemente carecem de estimadores de gradiente baseados em amostras e não enviesados para objetivos não lineares.

Os autores argumentam que a diversidade deve ser entendida não como uma restrição artificial, mas como a resposta racional à incerteza na função de recompensa. Quando a função de recompensa não é perfeitamente conhecida (devido a preferências ambíguas ou modelos de recompensa imperfeitos), comprometer-se com uma única ação pode ser subótimo.

Metodologia: ROSA

O artigo propõe uma reformulação fundamental do objetivo de RL ao substituir a recompensa escalar por uma distribuição sobre funções de recompensa e aplicar um objetivo não linear sobre conjuntos de ações. Este framework é denominado Randomized Objectives, Set Actions (ROSA).

O Critério ROSA+Max

Focando no cenário de bandit contextual, os autores definem o objetivo primário, ROSA+Max. Seja $\rho$ uma distribuição sobre funções de recompensa $R$, e seja $\mathbf{Y} = (Y_i)_{i=1}^n$ um multiconjunto de $n$ ações amostradas i.i.d. de uma política $\pi_\theta$. O objetivo é:

$$J_{\text{ROSA+Max}}(\pi_\theta) = \mathbb{E}_{X \sim \mu, \mathbf{Y} \sim \pi_\theta} \left[ \mathbb{E}_{R \sim \rho} \left[ \max_{1 \le i \le n} R(X, Y_i) \right] \right]$$

Intuitivamente, este critério agrega o desempenho ao tomar a recompensa máxima alcançada dentro de um conjunto de ações amostradas para cada amostra da função de recompensa, em vez de otimizar para a recompensa média. Isso incentiva a política a garantir que pelo menos uma ação no conjunto performe bem para cada realização possível da função de recompensa.

Estimativa de Gradiente

Os autores derivam um estimador de gradiente de política baseado em política, não enviesado e com redução de variância para este objetivo. Dada uma amostra de $m$ funções de recompensa e $n$ ações amostradas, o estimador de gradiente é:

$$\hat{g}(\theta) = \frac{1}{m} \sum_{k=1}^m \left[ \sum_{i=1}^n \left( \max_{1 \le j \le n} R_k(X, Y_j) - \max_{1 \le j \le n, j \ne i} R_k(X, Y_j) \right) \nabla_\theta \log \pi_\theta(Y_i|X) \right]$$

Este estimador atribui crédito efetivamente apenas à ação que alcança a recompensa máxima para uma amostra específica da função de recompensa, com o segundo termo servindo como uma covariável de controle para redução de variância.

Generalização

O framework generaliza para uma família de funções de conjunto $f$. Os autores provam que, para funções de recompensa binárias, qualquer função de conjunto $f$ que seja estritamente crescente e estritamente côncava em relação à soma das recompensas (ex: Softmax) produz uma política ótima maximamente diversa.

Principais Contribuições

1. Framework ROSA: Introdução de uma nova família de objetivos de RL que fornecem controle de diversidade calibrado através de uma distribuição sobre funções de recompensa.
2. Análise Teórica:
   • Política Ótima: Provou-se que, para uma distribuição uniforme sobre funções de recompensa binárias, a política única ótima de ROSA+Max é aquela que atribui massa de probabilidade igual a todas as ações ótimas distintas.
   • Diversidade Controlável: Derivou-se uma fórmula mostrando como distribuições de funções de recompensa não uniformes induzem políticas estocásticas específicas e controláveis, permitindo que a probabilidade de amostrar respostas corretas específicas seja ajustada via pesos da distribuição de recompensa.
   • Análise de Paisagem: Demonstrou-se que o ROSA+Max evita as paisagens "planas" da regularização de entropia e os problemas de "reordenação" de bônus de diversidade, fornecendo uma direção de gradiente que aumenta simultaneamente a recompensa e a diversidade.
3. Validação Empírica: Experimentos extensos em domínios de brinquedo e configurações de larga escala (incluindo pós-treinamento de LLM) demonstrando que o ROSA induz comportamentos diversos sem sacrificar a recompensa esperada, superando baselines padrão de gradiente de política e bônus de diversidade em cenários com preferências conflitantes e recompensas de verdade fundamental desconhecidas.

Resultados Experimentais

O artigo avalia o ROSA em três cenários principais:

1. Preferências Conflitantes: Em uma tarefa com duas funções de recompensa contraditórias (ex: paridade vs. anti-paridade), o gradiente de política vanilla falha devido ao cancelamento de recompensa. O ROSA aprende com sucesso uma política que equilibra ambas as preferências, gerando saídas válidas diversas.
2. Múltiplas Respostas Corretas: Em tarefas de raciocínio matemático onde múltiplos comprimentos de resposta são válidos, as políticas treinadas com ROSA amostram de todas as preferências de comprimento especificadas, enquanto o gradiente de política padrão colapsa para um único intervalo de comprimento.
3. Recompensas de Verdade Fundamental Desconhecidas: Em configurações com modelos de recompensa ruidosos (simulados por um ensemble de juízes com precisão variável), o ROSA demonstra robustez superior. Ele captura a incerteza na função de recompensa distribuindo a massa de probabilidade entre os modos de recompensa proporcionalmente aos seus pesos, enquanto a entropia regularizada do PG ou sofre de overfitting para a recompensa mais provável ou explora demais ações subótimas.

Em avaliações fora de distribuição (AIME 2025), modelos treinados com ROSA em dados MATH ruidosos alcançaram taxas de pass@k superiores comparadas ao PG padrão e baselines de ensemble pessimistas.

Significância e Alegações

O artigo alega que o ROSA oferece uma alternativa robusta e teoricamente fundamentada para tarefas complexas de RL onde formulações tradicionais falham em induzir a amplitude desejada de comportamento do agente. Ao reenquadrar a diversidade como uma resposta racional à incerteza da recompensa, o framework resolve a tensão entre maximização de recompensa e entropia da política.

Os autores enfatizam que o ROSA:

• Generaliza naturalmente tanto o gradiente de política vanilla quanto abordagens recentes de conjuntos de ações (como o best-of-n).
• Fornece um mecanismo para controlar a diversidade através da distribuição da função de recompensa, em vez de um ajuste frágil de hiperparâmetros de coeficientes de entropia.
• Garante que a solução ótima seja uma política estocástica que reflete a distribuição das possíveis funções de recompensa, evitando o colapso para um comportamento determinístico único mesmo quando o modelo de recompensa é imperfeito.

Em última análise, o trabalho fornece uma base principista para treinar políticas estocásticas que sejam tão diversas e adaptáveis quanto os ambientes que habitam, particularmente em domínios como o ajuste fino de LLMs e a descoberta científica, onde múltiplas soluções válidas existem.