Technische Samenvatting: Het gebruik van beloningsonzekerheid om divers gedrag te induceren in Reinforcement Learning

Probleemstelling

Klassieke reinforcement learning (RL) streeft doorgaans naar een deterministisch beleid dat de verwachte som van een scalaire beloning maximaliseert. Hoewel deze aanname de existentie van een optimaal deterministisch beleid garandeert, brengt dit significante beperkingen met zich mee in moderne toepassingen zoals het fine-tunen van large language models (LLM's) en wetenschappelijke ontdekkingen. In deze domeinen is convergentie naar een enkel gedrag vaak onwenselijk. In deze domeinen is het behoud van de diversiteit van vooraf getrainde beleidsregels bijvoorbeeld cruciaal voor creatieve taken en chain-of-thought redeneren. Bovendien, wanneer beloningsmodellen imperfecte proxies zijn voor de werkelijke onderliggende voorkeuren, kan overoptimalisatie naar een enkel deterministisch beleid leiden tot de verslechtering van de kwaliteit van generatie (reward hacking).

Bestaande remedies, zoals entropie-regularisatie of diversiteitsbonussen, proberen stochasticiteit af te dwingen maar lijden aan kritieke nadelen:

1. Fragiele Afwegingen: Entropie-regularisatie dwingt een afweging af waarbij verhoogde stochasticiteit ten koste gaat van de verwachte beloningen.
2. Incorrecte Beleidsordening: Diversiteitsbonussen resulteren vaak in ongewenste rangschikkingen waarbij suboptimale beleidsregels met een hoge diversiteit worden verkozen boven beleidsregels met hogere gemiddelde beloningen.
3. Inexpressiviteit: Veel doelstellingen, inclusief vanilla policy gradient en expected scalarized returns (ESR), bezitten inherent deterministische optimale beleidsregels, waardoor ze niet in staat zijn om noodzakelijke diversiteit te induceren.
4. Optimalisatieproblemen: Formuleringen van scalarized expected return (SER) missen vaak ongebiaste, op samples gebaseerde gradiënt-schatters voor niet-lineaire doelstellingen.

De auteurs stellen dat diversiteit niet begrepen moet worden als een kunstmatige beperking, maar als de rationele reactie op onzekerheid in de beloningsfunctie. Wanneer de beloningsfunctie niet perfect bekend is (door ambigue voorkeuren of imperfecte beloningsmodellen), kan het vastleggen op een enkele actie suboptimaal zijn.

Methodologie: ROSA

Het artikel stelt een fundamentele herformulering voor van de RL-doelstelling door de scalaire beloning te vervangen door een distributie over beloningsfuncties en een niet-lineaire doelstelling over verzamelingen van acties toe te passen. Dit raamwerk wordt Randomized Objectives, Set Actions (ROSA) genoemd.

De ROSA+Max Criterium

Bij het focussen op de contextuele bandit-setting definiëren de auteurs de primaire doelstelling, ROSA+Max. Laat $\rho$ een distributie zijn over beloningsfuncties $R$, en laat $\mathbf{Y} = (Y_i)_{i=1}^n$ een multiset zijn van $n$ acties die i.i.d. zijn gesampled uit een beleid $\pi_\theta$. De doelstelling is:

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

Intuïtief aggregeert dit criterium prestaties door de maximale beloning te nemen die bereikt wordt binnen een set gesamplede acties voor elke gesamplede beloningsfunctie, in plaats van te optimaliseren voor de gemiddelde beloning. Dit moedigt het beleid aan om te zorgen dat ten minste één actie in de set goed presteert voor elke mogelijke realisatie van de beloningsfunctie.

Gradiëntschatting

De auteurs leiden een principiële, ongebiasse en variantverlagende policy gradient schatter af voor deze doelstelling. Gegeven $m$ gesamplede beloningsfuncties en $n$ gesamplede acties, is de gradiëntschatter:

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

Deze schatter wijst effectief krediet toe aan alleen de actie die de maximale beloning bereikt voor een specifieke sample van de beloningsfunctie, waarbij de tweede term dient als een control variate voor variantvermindering.

Generalisatie

Het raamwerk generaliseert naar een familie van setfuncties $f$. De auteurs bewijzen dat voor binaire beloningsfuncties, elke setfunctie $f$ die strikt toenemend en strikt concave is met betrekking tot de som van beloningen (bijv. Softmax), een maximaal divers optimaal beleid oplevert.

Belangrijkste Bijdragen

1. ROSA Raamwerk: Introductie van een nieuwe familie van RL-doelstellingen die gekalibreerde diversiteitscontrole bieden via een distributie over beloningsfuncties.
2. Theoretische Analyse:
   • Optimaal Beleid: Bewezen dat voor een uniforme distributie over binaire beloningsfuncties, het unieke optimale beleid van ROSA+Max het beleid is dat gelijke waarschijnlijkheidsmassa toekent aan alle verschillende optimale acties.
   • Controleerbare Diversiteit: Afgeleid een formule die laat zien hoe niet-uniforme beloningsfunctiedistributies specifieke, controleerbare stochastische beleidsregels induceren, waardoor de kans op het samplen van specifieke correcte antwoorden kan worden afgestemd via de gewichten van de beloningsfunctiedistributie.
   • Landschap-analyse: Gedemonstreerd dat ROSA+Max de "platte" landschappen van entropie-regularisatie en de "herordening"-problemen van diversiteitsbonussen vermijdt, door een gradiëntrichting te bieden die gelijktijdig beloning en diversiteit verhoogt.
3. Empirische Validatie: Uitgebreide experimenten in toy-domeinen en grootschalige settings (inclusief LLM post-training) die aantonen dat ROSA divers gedrag induceert zonder de verwachte beloning op te offeren, waarbij het beter presteert dan standaard policy gradient en diversiteitsbonus-baselines in scenario's met conflicterende voorkeuren en onbekende grondwaarheid-beloningen.

Experimentele Resultaten

Het artikel evalueert ROSA in drie primaire scenario's:

1. Conflicterende Voorkeuren: In een taak met twee tegenstrijdige beloningsfuncties (bijv. pariteit versus anti-pariteit), faalt vanilla policy gradient door beloningsannulering. ROSA leert succesvol een beleid dat beide voorkeuren balanceert, wat diverse geldige outputs genereert.
2. Meerdere Correcte Antwoorden: In wiskundige redeneertaken waar meerdere antwoordlengtes geldig zijn, samplen met ROSA-getrainde beleidsregels uit alle gespecificeerde lengtevoorkeuren, terwijl standaard policy gradient convergeert naar een enkele lengtebereik.
3. Onbekende Grondwaarheid-Beloningen: In settings met ruisige beloningsmodellen (gesimuleerd door een ensemble van judges met variërende nauwkeurigheid), demonstreert ROSA superieure robuustheid. Het vangt de onzekerheid in de beloningsfunctie op door de waarschijnlijkheidsmassa over beloningsmodi te verdelen in verhouding tot hun gewichten, terwijl entropie-geregulariseerde PG ofwel overfit naar de meest waarschijnlijke beloning of te breed exploreert naar suboptimale acties.

In out-of-distribution evaluaties (AIME 2025), bereikten modellen getraind met ROSA op ruisige MATH-data hogere pass@k rates vergeleken met standaard PG en pessimistische ensemble-baselines.

Betekenis en Claims

Het artikel claimt dat ROSA een robuust en theoretisch gefundeerd alternatief biedt voor complexe RL-taken waar traditionele formuleringen er niet in slagen het gewenste bereik van agentgedrag te induceren. Door diversiteit te herformuleren als een rationele reactie op beloningsonzekerheid, lost het raamwerk de spanning tussen beloningsmaximalisatie en beleidsentropie op.

De auteurs benadrukken dat ROSA:

• Zowel vanilla policy gradient als recente action-set benaderingen (zoals best-of-$n$) natuurlijk generaliseert.
• Een mechanisme biedt om diversiteit te controleren via de distributie van de beloningsfunctie in plaats van fragiele hyperparameter-afstemming van entropiecoëfficiënten.
• Ervoor zorgt dat de optimale oplossing een stochastisch beleid is dat de distributie van mogelijke beloningsfuncties reflecteert, wat de instorting naar een enkel deterministisch gedrag voorkomt, zelfs wanneer het beloningsmodel imperfect is.

Uiteindelijk biedt dit werk een principieel fundament voor het trainen van stochastische beleidsregels die even divers en aanpasbaar zijn als de omgevingen waarin zij zich bevinden, met name in domeinen zoals LLM fine-tuning en wetenschappelijke ontdekkingen waar meerdere geldige oplossingen bestaan.