Résumé technique : Utiliser l'incertitude de la récompense pour induire des comportements diversifiés en apprentissage par renforcement

Énoncé du problème

L'apprentissage par renforcement (RL) classique cherche généralement une politique déterministe qui maximise la somme attendue d'une récompense scalaire. Bien que cette hypothèse garantisse l'existence d'une politique déterministe optimale, elle présente des limitations significatives dans les applications modernes telles que le réglage fin (fine-tuning) de modèles de langage de grande taille (LLM) et la découverte scientifique. Dans ces domaines, la convergence vers un comportement unique est souvent indésirable. Par exemple, préserver la diversité des politiques pré-entraînées est crucial pour les tâches créatives et le raisonnement de type « chaîne de pensée » (chain-of-thought). De plus, lorsque les modèles de récompense sont des approximations imparfaites des véritables préférences sous-jacentes, la sur-optimisation vers une politique déterministe unique peut entraîner une détérioration de la qualité de la génération (reward hacking).

Les remèdes existants, tels que la régularisation de l'entropie ou les bonus de diversité, tentent d'imposer la stochasticité mais souffrent de défauts critiques :

1. Compromis fragiles : La régularisation de l'entropie impose un compromis où l'augmentation de la stochasticité se fait au détriment de la récompense attendue.
2. Ordre incorrect des politiques : Les bonus de diversité entraînent souvent des classements indésirables où des politiques sous-optimales avec une grande diversité sont préférées à des politiques ayant des récompenses moyennes plus élevées.
3. Inexpressivité : De nombreux objectifs, y compris le gradient de politique classique et les rendements scalaires attendus (ESR), possèdent intrinsèquement des politiques optimales déterministes, échouant ainsi à induire la diversité nécessaire.
4. Défis d'optimisation : Les formulations de rendement scalaire attendu (SER) manquent souvent d'estimateurs de gradient non biaisés et basés sur des échantillons pour les objectifs non linéaires.

Les auteurs soutiennent que la diversité doit être comprise non pas comme une contrainte artificielle, mais comme la réponse rationnelle à l'incertitude de la fonction de récompense. Lorsque la fonction de récompense n'est pas parfaitement connue (en raison d'ambiguïtés de préférences ou de modèles de récompense imparfaits), s'engager vers une action unique peut être sous-optimal.

Méthodologie : ROSA

Le papier propose une reformulation fondamentale de l'objectif de RL en remplaçant la récompense scalaire par une distribution sur les fonctions de récompense et en appliquant un objectif non linéaire sur des ensembles d'actions. Ce cadre est nommé Randomized Objectives, Set Actions (ROSA).

Le critère ROSA+Max

En se concentrant sur le cadre des bandits contextuels, les auteurs définissent l'objectif principal, ROSA+Max. Soit $\rho$ une distribution sur les fonctions de récompense $R$, et soit $\mathbf{Y} = (Y_i)_{i=1}^n$ un multiensemble de $n$ actions échantillonnées i.i.d. à partir d'une politique $\pi_\theta$. L'objectif est :

$$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]$$

Intuitivement, ce critère agrège la performance en prenant la récompense maximale atteinte au sein d'un ensemble d'actions échantillonnées pour chaque réalisation de la fonction de récompense échantillonnée. Cela encourage la politique à garantir qu'au moins une action dans l'ensemble performe bien pour chaque réalisation possible de la fonction de récompense, plutôt que d'optimiser la récompense moyenne.

Estimation du gradient

Les auteurs dérivent un estimateur de gradient de politique sans biais, non biaisé et à variance réduite pour cet objectif. Étant donné $m$ fonctions de récompense échantillonnées et $n$ actions échantillonnées, l'estimateur du gradient est :

$$\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]$$

Cet estimateur attribue le crédit uniquement à l'action qui atteint la récompense maximale pour un échantillon spécifique de la fonction de récompense, le second terme servant de variable de contrôle pour la réduction de la variance.

Généralisation

Le cadre se généralise à une famille de fonctions d'ensemble $f$. Les auteurs prouvent que pour des fonctions de récompense binaires, toute fonction d'ensemble $f$ qui est strictement croissante et strictement concave par rapport à la somme des récompenses (par exemple, Softmax) produit une politique optimale maximalement diversifiée.

Contributions clés

1. Cadre ROSA : Introduction d'une nouvelle famille d'objectifs de RL qui fournissent un contrôle de la diversité calibré via une distribution sur les fonctions de récompense.
2. Analyse théorique :
   • Politique optimale : Prouvé que pour une distribution uniforme sur les fonctions de récompense binaires, l'unique politique optimale de ROSA+Max est celle qui assigne une masse de probabilité égale à toutes les actions optimales distinctes.
   • Diversité contrôlable : Dérivé une formule montrant comment les distributions de fonctions de récompense non uniformes induisent des politiques stochastiques spécifiques et contrôlables, permettant de régler la probabilité d'échantillonner des réponses correctes spécifiques via les poids de la distribution de récompense.
   • Analyse du paysage : Démontré que ROSA+Max évite les paysages « plats » de la régularisation de l'entropie et les problèmes de « réordonnancement » des bonus de diversité, fournissant une direction de gradient qui augmente simultanément la récompense et la diversité.
3. Validation empirique : Expériences approfondies dans des domaines de test et des contextes à grande échelle (incluant le post-entraînement de LLM) démontrant que ROSA induit des comportements diversifiés sans sacrifier la récompense attendue, surpassant les standards de gradient de politique et les bases de comparaison de bonus de diversité dans des scénaux de préférences conflictuelles et de récompenses de vérité terrain inconnues.

Résultats expérimentaux

Le papier évalue ROSA à travers trois scénarios principaux :

1. Préférences conflictuelles : Dans une tâche avec deux fonctions de récompense contradictoires (ex: parité vs anti-parité), le gradient de politique classique échoue en raison de l'annulation des récompenses. ROSA apprend avec succès une politique qui équilibre les deux préférences, générant des sorties valides et diverses.
2. Plusieurs bonnes réponses : Dans des tâches de raisonnement mathématique où plusieurs longueurs de réponse sont valides, les politiques entraînées avec ROSA échantillonnent parmi toutes les préférences de longueur spécifiées, tandis que le gradient de politique standard s'effondre vers une plage de longueur unique.
3. Récompenses de vérité terrain inconnues : Dans des contextes avec des modèles de récompense bruités (simulés par un ensemble de juges aux précisions variables), ROSA démontre une robustesse supérieure. Il capture l'incertitude de la fonction de récompense en distribuant la masse de probabilité à travers les modes de récompense proportionnellement à leurs poids, alors que la politique de gradient de politique régularisée par l'entropie sur-optimise soit vers la récompense la plus probable, soit explore trop largement vers des actions sous-optimales.

Dans les évaluations hors distribution (AIME 2025), les modèles entraînés avec ROSA sur des données MATH bruitées ont obtenu des taux pass@k plus élevés que le PG standard et les bases de comparaison d'ensembles pessimistes.

Signification et affirmations

Le papier affirme que ROSA offre une alternative robuste et théoriquement fondée pour des tâches de RL complexes où les formulations traditionnelles échouent à induire la largeur de comportement souhaitée des agents. En recadrant la diversité comme une réponse rationnelle à l'incertitude de la récompense, le cadre résout la tension entre la maximisation de la récompense et l'entropie de la politique.

Les auteurs soulignent que ROSA :

• Généralise naturellement le gradient de politique classique et les approches récentes par ensemble d'actions (comme le best-of-$n$).
• Fournit un mécanisme pour contrôler la diversité via la distribution de la fonction de récompense plutôt que par le réglage fragile des hyperparamètres de coefficients d'entropie.
• Garantit que la solution optimale est une politique stochastique qui reflète la distribution des fonctions de récompense possibles, empêissant l'effondrement vers un comportement déterministe unique même lorsque le modèle de récompense est imparfait.

En fin de compte, ce travail fournit un fondement principalisé pour entraîner des politiques stochastiques qui sont aussi diverses et adaptables que les environnements qu'elles habitent, particulièrement dans des domaines comme le réglage fin de LLM et la découverte scientifique où de multiples solutions valides existent.