Technisches Resümee: Nutzung von Belohnungsunsicherheit zur Induktion diverser Verhaltensweisen in Reinforcement Learning

Problemstellung

Klassisches Reinforcement Learning (RL) strebt typischerweise eine deterministische Policy an, die die erwartete Summe einer skalaren Belohnung maximiert. Während diese Annahme die Existenz einer optimalen deterministischen Policy garantiert, stellt sie in modernen Anwendungen wie dem Fine-Tuning von Large Language Models (LLMs) oder der wissenschaftlichen Entdeckung erhebliche Einschränkungen dar. In diesen Domänen ist die Konvergenz zu einem einzelnen Verhalten oft unerwünscht. Beispielsweise ist die Erhaltung der Diversität vortrainierter Policies entscheidend für kreative Aufgaben und das Chain-of-Thought-Reasoning. Darüber hinaus kann die Überoptimierung auf eine einzige deterministische Policy, wenn Belohnungsmodelle unvollkommene Proxies für die wahren zugrunde liegenden Präferenzen sind, zu einer Verschlechterung der Generationsqualität führen (Reward Hacking).

Bestehende Abhilfemaßnahmen, wie Entropie-Regularisierung oder Diversitäts-Boni, versuchen, Stochastizität zu erzwingen, leiden jedoch unter kritischen Nachteilen:

1. Fragile Trade-offs: Die Entropie-Regularisierung erzwingt einen Kompromiss, bei dem erhöhte Stochastizität auf Kosten der erwarteten Belohnungen geht.
2. Falsche Policy-Rangfolge: Diversitäts-Boni führen oft zu unerwünschten Rangfolgen, in denen suboptimale, aber hochgradig diverse Policies gegenüber Policies mit höheren durchschnittlichen Belohnungen bevorzugt werden.
3. Mangelnde Ausdrucksstärke: Viele Zielsetzungen, einschließlich des klassischen Policy Gradient und der erwarteten skalierten Erträge (Expected Scalarized Returns, ESR), besitzen inhärent deterministische optimale Policies und versagen daher bei der Induktion notwendiger Diversität.
4. Optimierungsprobleme: Formulierungen des skalierten erwarteten Ertrags (Scalarized Expected Return, SER) verfügen oft über keine unverzerrten, probenbasierten Gradientenschätzer für nicht-lineare Zielsetzungen.

Die Autoren argumentieren, dass Diversität nicht als künstliche Einschränkung, sondern als die rationale Antwort auf die Unsicherheit in der Belohnungsfunktion verstanden werden sollte. Wenn die Belohnungsfunktion nicht perfekt bekannt ist (aufgrund ambivalenter Präferenzen oder unvollkommener Belohnungsmodelle), kann das Festhalten an einer einzigen Aktion suboptimal sein.

Methodik: ROSA

Das Paper schlägt eine grundlegende Umformulierung der RL-Zielsetzung vor, indem es die skalare Belohnung durch eine Verteilung über Belohnungsfunktionen ersetzt und eine nicht-lineare Zielsetzung über Aktionsmengen anwendet. Dieses Framework wird als Randomized Objectives, Set Actions (ROSA) bezeichnet.

Das ROSA+Max Kriterium

Mit Fokus auf das Kontext-Banditen-Setting definieren die Autoren die primäre Zielsetzung, ROSA+Max. Sei $\rho$ eine Verteilung über Belohnungsfunktionen $R$, und sei $\mathbf{Y} = (Y_i)_{i=1}^n$ eine Multimenge von $n$ Aktionen, die i.i.d. aus einer Policy $\pi_\theta$ gezogen wurden. Die Zielsetzung ist:

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

Intuitiv aggregiert dieses Kriterium die Leistung, indem es die maximale Belohnung innerhalb einer Menge von Stichproben-Aktionen für jede Stichprobe der Belohnungsfunktion bildet. Dies ermutigt die Policy sicherzustellen, dass mindestens eine Aktion in der Menge für jede mögliche Realisierung der Belohnungsfunktion gut abschneidet, anstatt nur die durchschnittliche Belohnung zu optimieren.

Gradientenschätzung

Die Autoren leiten einen fundierten, unverzerrten und varianzreduzierten Policy-Gradient-Schätzer für diese Zielsetzung ab. Gegeben seien $m$ Stichproben von Belohnungsfunktionen und $n$ Stichproben von Aktionen, ist der Gradientenschätzer:

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

Dieser Schätzer weist die Kreditwürdigkeit effektiv nur der Aktion zu, die die maximale Belohnung für eine spezifische Stichprobe der Belohnungsfunktion erzielt, wobei der zweite Term als Kontrollvariate zur Varianzreduktion dient.

Verallgemeinerung

Das Framework generalisiert auf eine Familie von Mengenfunktionen $f$. Die Autoren beweisen, dass für binäre Belohnungsfunktionen jede Mengenfunktion $f$, die streng monoton steigend und streng konkav bezüglich der Summe der Belohnungen ist (z. B. Softmax), eine maximal diverse optimale Policy liefert.

Zentrale Beiträge

1. ROSA Framework: Einführung einer neuen Familie von RL-Zielsetzungen, die eine kalibrierte Diversitätskontrolle durch eine Verteilung über Belohnungsfunktionen ermöglichen.
2. Theoretische Analyse:
   • Optimale Policy: Es wurde bewiesen, dass für eine Gleichverteilung über binäre Belohnungsfunktionen die eindeutige optimale Policy von ROSA+Max diejenige ist, die die gleiche Wahrscheinlichkeitsmasse allen distinkten optimalen Aktionen zuweist.
   • Kontrollierbare Diversität: Eine Formel wurde hergeleitet, die zeigt, wie nicht-gleichmäßige Verteilungen der Belohnungsfunktionen spezifische, kontrollierbare stochastische Policies induzieren, wodurch die Wahrscheinlichkeit der Stichprobenziehung spezifischer korrekter Antworten über die Gewichte der Belohnungsverteilung gesteuert werden kann.
   • Landschaftsanalyse: Es wurde demonstriert, dass ROSA+Max die "flachen" Landschaften der Entropie-Regularisierung und die "Neu-Rangfolge"-Probleme von Diversitäts-Boni vermeidet, indem es eine Gradientenrichtung bereitstellt, die gleichzeitig Belohnung und Diversität erhöht.
3. Empirische Validierung: Umfangreiche Experimente in Spielumgebungen und groß angelegten Settings (einschließlich LLM Post-Training) zeigen, dass ROSA diverse Verhaltensweisen induziert, ohne die erwartete Belohnung zu opfern, und dabei Standard-Policy-Gradient- und Diversitäts-Bonus-Baselines in Szenarien mit widersprüchlichen Präferenzen und unbekannten Ground-Truth-Belohnungen übertrifft.

Experimentelle Ergebnisse

Das Paper evaluiert ROSA in drei primären Szenarien:

1. Konfliktierende Präferenzen: In einer Aufgabe mit zwei widersprüchlichen Belohnungsfunktionen (z. B. Parität vs. Anti-Parität) scheitert der klassische Policy Gradient aufgrund von Belohnungskompensation. ROSA lernt erfolgreich eine Policy, die beide Präferenzen ausbalanciert und diverse valide Outputs generiert.
2. Multiple korrekte Antworten: In mathematischen Denkaufgaben, bei denen mehrere Antwortlängen gültig sind, sampelt eine mit ROSA trainierte Policy aus allen spezifizierten Längenpräferenzen, während ein Standard-Policy-Gradient zu einer einzigen Längenspanne kollabiert.
3. Unbekannte Ground-Truth-Belohnungen: In Settings mit verrauschten Belohnungsmodellen (simuliert durch ein Ensemble von Judges mit variierender Genauigkeit) zeigt ROSA eine überlegene Robustheit. Es erfasst die Unsicherheit der Belohnungsfunktion, indem es die Wahrscheinlichkeitsmasse proportional zu den Gewichten über die Belohnungsmoden verteilt, während eine entropie-regularisierte PG entweder auf die wahrscheinlichste Belohnung überoptimiert oder zu breit in suboptimale Aktionen exploriert.

In Out-of-Distribution-Evaluierungen (AIME 2025) erreichten Modelle, die mit ROSA auf verrauschten MATH-Daten trainiert wurden, höhere pass@k-Raten im Vergleich zu Standard-PG und pessimistischen Ensemble-Baselines.

Bedeutung und Behauptungen

Das Paper behauptet, dass ROSA eine robuste und theoretisch fundierte Alternative für komplexe RL-Aufgaben bietet, bei denen traditionelle Formulierungen die gewünschte Breite des Agentenverhaltens nicht induzieren können. Durch die Umformulierung von Diversität als rationale Antwort auf Belohnungsunsicherheit löst das Framework das Spannungsverhältnis zwischen Belohnungsmaximierung und Policy-Entropie auf.

Die Autoren betonen, dass ROSA:

• Sowohl den klassischen Policy Gradient als auch jüngere Ansätze für Aktionsmengen (wie Best-of-$n$) natürlich generalisiert.
• Einen Mechanismus bietet, um Diversität durch die Verteilung der Belohnungsfunktionen statt durch das fragile Hyperparameter-Tuning von Entropiekoeffizienten zu steuern.
• Sicherstellt, dass die optimale Lösung eine stochastische Policy ist, welche die Verteilung möglicher Belohnungsfunktionen widerspiegelt, was einen Kollaps zu einem einzigen deterministischen Verhalten verhindert, selbst wenn das Belohnungsmodell unvollkommen ist.

Letztlich bietet die Arbeit eine fundierte Basis für das Training stochastischer Policies, die so divers und anpassungsfähig sind wie die Umgebungen, in denen sie agieren – insbesondere in Domänen wie dem LLM-Fine-Tuning und der wissenschaftlichen Entdeckung, in denen multiple valide Lösungen existieren.