Riepilogo Tecnico: PriorZero

Enunciato del Problema

L'integrazione di Modelli Linguistici di Grande Dimensione (LLM) in agenti di Apprendimento per Rinforzo (RL) offre una via promettente per sfruttare la ricca conoscenza del mondo nella risoluzione di compiti decisionali a lungo orizzonte e con ricompense sparse. Tuttavia, gli approcci esistenti sono ostacolati da un fondamentale disallineamento prior-dinamiche.

I paradigmi attuali affrontano specifici colli di bottiglia strutturali:

1. Politica Ingenua (P1): L'uso degli LLM come politiche zero-shot porta a decisioni miopi perché la probabilità linguistica non equivale al ritorno dell'ambiente, creando un "divario tra conoscenza e azione".
2. Affinamento Ingenuo con RL (P2): L'affinamento end-to-end degli LLM tramite RL soffre di assegnazione del credito ad alta varianza e instabilità di ottimizzazione, specialmente in ambienti con ricompense sparse.
3. Modellazione delle Dinamiche (P3): L'uso degli LLM per prevedere le transizioni di stato fallisce a causa di un disallineamento di granularità; gli LLM sono pre-addestrati sulla generazione di testo a livello di token, mentre i modelli di dinamiche latenti richiedono la previsione delle transizioni a livello stato-azione.
4. Encoder di Testo (P4): L'uso di feature LLM congelate come rappresentazioni statiche limita l'adattamento specifico al compito, imponendo un limite alle prestazioni.

La sfida centrale consiste nel colmare il divario tra le prior semantiche statiche degli LLM e i modelli di transizione dinamici e specifici dell'ambiente richiesti per una pianificazione efficace.

Metodologia: Framework PriorZero

PriorZero propone un framework unificato che disaccoppia l'LLM dal modello del mondo, integrandoli attraverso un design di rollout e addestramento disaccoppiato. Il sistema colloca un modello del mondo apprendibile tra l'LLM e l'ambiente, permettendo all'LLM di fornire guida semantica mentre il modello del mondo gestisce la modellazione delle transizioni, la stima del valore e la pianificazione multi-step.

1. Iniezione della Prior Radice (Fase di Rollout)

Durante la fase di rollout, PriorZero utilizza un Modello del Mondo Latente (basato su UniZero) per eseguire una Ricerca ad Albero Monte Carlo (MCTS) in uno spazio latente.

• Meccanismo: Un nuovo meccanismo di Iniezione della Prior Radice fonde la prior semantica di azione dell'LLM esclusivamente al nodo radice dell'MCTS.
• Implementazione: L'LLM genera un'analisi della situazione (Chain-of-Thought) e una distribuzione di azioni ($\pi_{LLM}$) basata sul contesto testuale corrente. Questa viene fusa con la politica del modello del mondo ($\pi_{WM}$) alla radice:
  $$P_{root}(a | z_t) = (1 - \alpha) \pi_{WM}(a | z_t) + \alpha \pi_{LLM}(a | C_t)$$
• Vincolo: I nodi non-radice nell'albero di ricerca si affidano interamente alle dinamiche latenti del modello del mondo. Questo previene il "disallineamento di granularità" del Paradigma P3, poiché l'LLM non è mai costretto a prevedere transizioni di stato per stati latenti immaginati.
• Vantaggio: Ciò guida la ricerca verso azioni semanticamente promettenti, impedendo all'MCTS di cadere in traiettorie prive di significato (ad esempio, cicli di ricompensa) preservando al contempo la capacità di previsione profonda del modello del mondo.

2. Affinamento Alternato per Rinforzo (Fase di Addestramento)

Per affrontare il fallimento dell'assegnazione del credito del Paradigma P2, PriorZero impiega una strategia di Affinamento Alternato che disaccoppia l'apprendimento del modello del mondo dall'adattamento dell'LLM.

• Programma: L'addestramento segue un programma asimmetrico (ad esempio, 10 aggiornamenti del modello del mondo per ogni 1 aggiornamento dell'LLM).
• Apprendimento del Modello del Mondo: Il modello del mondo viene aggiornato sui dati di interazione per affinare le sue previsioni di dinamiche, valore e politica.
• Adattamento dell'LLM: Invece di utilizzare ritorni on-policy ad alta varianza, l'LLM viene affinato utilizzando segnali di vantaggio a bassa varianza derivati dalle stime di valore del modello del mondo.
  • Il vantaggio è calcolato utilizzando un target TD bootstrap a $n$-step: $A(z_t, a_t) = Q_n(z_t, a_t) - v_\theta(z_t)$.
  • L'LLM viene aggiornato tramite un obiettivo PPO clipato, con uno schema di pesatura specifico che applica il peso completo ai token di azione e un peso ridotto ($w_{cot}=0.1$) ai token di ragionamento Chain-of-Thought per mantenere la stabilità.
• Vantaggio: Ciò crea un ciclo chiuso in cui il modello del mondo fornisce una supervisione stabile e granulare per l'LLM, permettendo alle prior dell'LLM di co-evolvere con le dinamiche dell'ambiente senza l'instabilità dell'ottimizzazione end-to-end.

Contributi Chiave

1. Analisi dei Paradigmi: Il documento identifica il disallineamento prior-dinamiche come l'ostacolo centrale nel trasferire la conoscenza degli LLM al RL a lungo orizzonte, categorizzando gli approcci esistenti (P1–P4) e i loro specifici modi di fallimento.
2. Framework PriorZero: Un'architettura innovativa che sinergizza le prior degli LLM con la pianificazione del modello del mondo attraverso due meccanismi fondamentali:
   • Iniezione della Prior Radice: Inietta guida semantica alla radice dell'MCTS per facilitare l'esplorazione senza contaminare le dinamiche latenti.
   • Affinamento Alternato per Rinforzo: Utilizza le stime di valore del modello del mondo per costruire segnali di assegnazione del credito stabili e a bassa varianza per l'adattamento dell'LLM.
3. Validazione Empirica: Esperimenti completi su Jericho (giochi di avventura testuali) e BabyAI (esecuzione di istruzioni in grid-world) dimostrano che PriorZero migliora costantemente l'efficienza esplorativa e le prestazioni asintotiche rispetto alle linee di base.

Risultati Sperimentali

• Benchmark Jericho: PriorZero supera UniZero (una forte linea di base del modello del mondo senza prior LLM) su quattro giochi distinti (Detective, Acorncourt, Zork1, Omniquest).
  • In Zork1 (lungo orizzonte, ricompense sparse), PriorZero ottiene ritorni asintotici più elevati e una convergenza più rapida.
  • In Detective, PriorZero supera UniZero dopo i guadagni iniziali, dimostrando il valore dell'adattamento delle prior alle dinamiche specifiche del compito.
• Benchmark BabyAI: Su 18 livelli del grid-world di BabyAI, PriorZero raggiunge plateau di ritorno più alti (0.82 contro 0.79 per UniZero) e mostra guadagni significativi in compiti composizionali (ad esempio, SynthLoc), dove la prior dell'LLM aiuta a scomporre gli obiettivi in sottocompiti.
• Studi di Ablazione:
  • Programma Alternato: Cruciale per la stabilità; le varianti non alternate soffrono di "corruzione della conoscenza" a causa di segnali di valore rumorosi.
  • Iniezione Solo Radice: Essenziale per impedire all'LLM di sovrascrivere le dinamiche del modello del mondo durante la previsione.
  • Peso della Loss CoT: È richiesto un peso bilanciato ($w_{cot}=0.1$); rimuovere CoT causa instabilità numerica, mentre un peso completo porta a un apprendimento volatile.
  • MCTS: Disabilitare l'MCTS causa il collasso delle prestazioni, confermando la necessità di target migliorati dalla pianificazione.

Significato e Affermazioni

Il documento afferma che PriorZero stabilisce un framework promettente per il processo decisionale potenziato dagli LLM risolvendo i colli di bottiglia strutturali del trasferimento delle prior. Il suo significato risiede in:

• Disaccoppiamento: Separazione riuscita del ruolo dell'LLM (fornire prior semantiche) dal ruolo del modello del mondo (apprendere dinamiche e pianificare), evitando così il disallineamento di granularità e i fallimenti di assegnazione del credito dei metodi precedenti.
• Stabilità: Abilitazione di un affinamento stabile degli LLM in ambienti a lungo orizzonte e con ricompense sparse sfruttando il modello del mondo come critico a bassa varianza.
• Scalabilità: Dimostrazione che il framework scala con la capacità dell'LLM (migliorando le prestazioni con modelli da 7B in compiti complessi) mantenendo al contempo la robustezza attraverso il programma di addestramento alternato.

Gli autori posizionano PriorZero non come un sostituto per i modelli del mondo o gli LLM, ma come un'integrazione sinergica in cui l'LLM focalizza la ricerca e il modello del mondo affina le prior dell'LLM attraverso l'interazione continua.