Sintesi Tecnica: Utilizzare l'Incertezza della Ricompensa per Indurre un Comportamento Diversificato nel Reinforcement Learning

Problematica

Il reinforcement learning (RL) classico cerca tipicamente una politica deterministica che massimizzi la somma attesa di una ricompensa scalare. Sebbene tale assunzione garantisca l'esistenza di una politica deterministica ottimale, essa presenta limitazioni significative in applicazioni moderne come il fine-tuning dei modelli linguistici di grandi dimensioni (LLM) e la scoperta scientifica. In questi domini, la convergenza verso un singolo comportamento è spesso indesiderabile. Ad esempio, preservare la diversità delle politiche pre-addestrate è cruciale per compiti creativi e per il ragionamento di tipo "chain-of-thought". Inoltre, quando i modelli di ricompensa sono proxy imperfetti delle vere preferenze sottostanti, l'ottimizzazione eccessiva verso una singola politica deterministica può portare al deterioramento della qualità della generazione (reward hacking).

I rimedi esistenti, come la regolarizzazione dell'entropia o i bonus di diversità, tentano di imporre la stochasticità ma soffrono di critici svantaggi:

1. Compromessi Fragili: La regolarizzazione dell'entropia impone un compromesso in cui l'aumento della stochasticità avviene a scapito delle ricompense attese.
2. Ordinamento Errato delle Politiche: I bonus di diversità spesso causano classificazioni indesiderabili in cui vengono preferite politiche sub-ottimali con alta diversità rispetto a politiche con ricompense medie più elevate.
3. Inespressività: Molti obiettivi, inclusi il vanilla policy gradient e l'expected scalarized returns (ESR), possiedono intrinsecamente politiche ottime deterministiche, fallendo nell'indurre la necessaria diversità.
4. Sfide di Ottimizzazione: Le formulazioni di scalarized expected return (SER) spesso mancano di stimatori del gradiente non distorti e basati su campioni per obiettivi non lineari.

Gli autori sostengono che la diversità debba essere intesa non come un vincolo artificiale, ma come la risposta razionale all'incertezza nella funzione di ricompensa. Quando la funzione di ricompensa non è perfettamente nota (a causa di preferenze ambigue o modelli di ricompensa imperfetti), impegnarsi in un'unica azione può essere sub-ottimale.

Metodologia: ROSA

Il documento propone una riformulazione fondamentale dell'obiettivo di RL sostituendo la ricompensa scalare con una distribuzione su funzioni di ricompensa e applicando un obiettivo non lineare su insiemi di azioni. Questo framework è denominato Randomized Objectives, Set Actions (ROSA).

Il Criterio ROSA+Max

Concentrandosi sul contesto dei bandit contestuali, gli autori definiscono l'obiettivo primario, ROSA+Max. Sia $\rho$ una distribuzione sulle funzioni di ricompensa $R$, e sia $\mathbf{Y} = (Y_i)_{i=1}^n$ un multiset di $n$ azioni campionate i.i.d. da una politica $\pi_\theta$. L'obiettivo è:

$$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, questo criterio aggrega le prestazioni prendendo la ricompensa massima raggiunta all'interno di un insieme di azioni campionate per ogni istanza campionata della funzione di ricompensa. Ciò incoraggia la politica a garantire che almeno un' azione nell'insieme funzioni bene per ogni possibile realizzazione della funzione di ricompensa, piuttosto che ottimizzare per la ricompensa media.

Stima del Gradiente

Gli autori derivano uno stimatore del gradiente della politica, non distorto e con riduzione della varianza, per questo obiettivo. Dato un numero $m$ di funzioni di ricompensa campionate e un numero $n$ di azioni campionate, lo stimatore del 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]$$

Questo stimatore assegna il merito solo all'azione che raggiunge la ricompensa massima per una specifica istanza della funzione di ricompensa, con il secondo termine che funge da variabile di controllo per la riduzione della varianza.

Generalizzazione

Il framework si generalizza a una famiglia di funzioni di insieme $f$. Gli autori dimostrano che per funzioni di ricompensa binarie, qualsiasi funzione di insieme $f$ che sia strettamente crescente e strettamente concava rispetto alla somma delle ricompense (ad esempio, Softmax) produce una politica ottimale massimamente diversificata.

Contributi Chiave

1. Framework ROSA: Introduzione di una nuova famiglia di obiettivi di RL che forniscono un controllo della diversità calibrato attraverso una distribuzione sulle funzioni di ricompensa.
2. Analisi Teorica:
   • Politica Ottimale: Dimostrato che per una distribuzione uniforme su funzioni di ricompensa binarie, l'unica politica ottimale di ROSA+Max è quella che assegna una massa di probabilità uguale a tutte le azioni ottime distinte.
   • Diversità Controllabile: Derivata una formula che mostra come distribuzioni non uniformi delle funzioni di ricompensa inducano politiche stocastiche specifiche e controllabili, permettendo di regolare la probabilità di campionare risposte corrette specifiche tramite i pesi della distribuzione della ricompensa.
   • Analisi del Paesaggio (Landscape): Dimostrato che ROSA+Max evita i paesaggi "piatti" della regolarizzazione dell'entropia e i problemi di "riordinamento" dei bonus di diversità, fornendo una direzione del gradiente che aumenta simultaneamente la ricompensa e la diversità.
3. Validazione Empirica: Esperimenti estesi in domini di gioco (toy domains) e contesti su larga scala (incluso il post-training di LLM) dimostrano che ROSA induce comportamenti diversificati senza sacrificare la ricompensa attesa, superando le basi standard del policy gradient e dei bonus di diversità in scenari con preferenze contrastanti e ricompense reali sconosciute.

Risultati Sperimentali

Il documento valuta ROSA in tre scenari principali:

1. Preferenze Conflittuali: In un compito con due funzioni di ricompensa contraddittorie (ad esempio, parità vs anti-parità), il vanilla policy gradient fallisce a causa della cancellazione della ricompensa. ROSA impara con successo una politica che bilancia entrambe le preferenze, generando output validi e diversificati.
2. Molteplici Risposte Corrette: In compiti di ragionamento matematico dove lunghezze di risposta multiple sono valide, le politiche addestrate con ROSA campionano da tutte le lunghezze preferite, mentre il policy gradient standard collassa verso un singolo intervallo di lunghezza.
3. Ricompense Ground-Truth Sconosciute: In contesti con modelli di ricompensa rumorosi (simulati da un ensemble di giudici con diversa accuratezza), ROSA dimostra una robustezza superiore. Cattura l'incertezza nella funzione di ricompensa distribuendo la massa di probabilità tra le mode della ricompensa in proporzione ai loro pesi, mentre la politica regolarizzata dall'entropia o il PG tende a fare overfitting sulla ricompensa più probabile o a esplorare troppo ampiamente verso azioni sub-ottimali.

Nelle valutazioni out-of-distribution (AIME 2025), i modelli addestrati con ROSA su dati MATH rumorosi hanno ottenuto tassi di pass@k più elevati rispetto al PG standard e alle baseline pessimistiche di ensemble.

Significato e Rivendicazioni

Il documento afferma che ROSA offre un'alternativa robusta e teoricamente fondata per compiti di RL complessi in cui le formulazioni tradizionali falliscono nell'indurre la desiderata ampiezza del comportamento dell'agente. Riformulando la diversità come una risposta razionale all'incertezza della ricompensa, il framework risolve la tensione tra massimizzazione della ricompensa ed entropia della politica.

Gli autori sottolineano che ROSA:

• Generalizza naturalmente sia il vanilla policy gradient che i recenti approcci basati su insiemi di azioni (come il best-of-$n$).
• Fornisce un meccanismo per controllare la diversità attraverso la distribuzione della funzione di ricompensa anziché attraverso il fragile tuning degli iperparametri del coefficiente di entropia.
• Garantisce che la soluzione ottimale sia una politica stocastica che riflette la distribuzione delle possibili funzioni di ricompensa, impedendo il collasso verso un singolo comportamento deterministico anche quando il modello di ricompensa è imperfetto.

In definitiva, questo lavoro fornisce una base rigorosa per l'addestramento di politiche stocastiche che siano tanto diverse quanto gli ambienti che abitano, particolarmente in domini come il fine-tuning di LLM e la scoperta scientifica dove esistono molteplici soluzioni valide.