Technische Samenvatting: PriorZero

Probleemstelling

De integratie van Large Language Models (LLM's) in Reinforcement Learning (RL)-agenten biedt een veelbelovende weg om rijke wereldkennis te benutten voor het oplossen van beslissingsproblemen met lange horizon en schaarse beloningen. Echter, bestaande aanpakken worden gehinderd door een fundamentele mismatch tussen prior en dynamiek.

Huidige paradigma's staan voor specifieke structurele knelpunten:

1. Naïeve Policy (P1): Het gebruik van LLM's als zero-shot policies leidt tot kortzichtige beslissingen, omdat linguïstische waarschijnlijkheid niet gelijkstaat aan omgevingsopbrengst, wat een "kennis-doen-kloof" creëert.
2. Naïeve RL-finetuning (P2): End-to-end finetuning van LLM's via RL lijdt onder hoge variantie in credit assignment en optimalisatie-instabiliteit, vooral in omgevingen met schaarse beloningen.
3. Dynamiekmodellering (P3): Het gebruik van LLM's om toestandsovergangen te voorspellen faalt door een mismatch in granulariteit; LLM's zijn voorgetraind op token-niveau tekstgeneratie, terwijl latent-dynamiekmodellen overgangspredicties op state-action-niveau vereisen.
4. Tekstencoders (P4): Het gebruik van bevroren LLM-kenmerken als statische representaties beperkt taakspecifieke aanpassing en legt een prestatiedekplaat op.

De kernuitdaging is het overbruggen van de kloof tussen de statische, semantische priors van LLM's en de dynamische, omgevingspecifieke overgangsmodellen die nodig zijn voor effectieve planning.

Methodologie: PriorZero Framework

PriorZero stelt een unificerend framework voor dat de LLM ontkoppelt van het wereldmodel, en deze integreert via een ontkoppeld rollout-trainingsontwerp. Het systeem plaatst een leerbare wereldmodel tussen de LLM en de omgeving, waardoor de LLM semantische begeleiding kan bieden terwijl het wereldmodel overgangsmodellering, waarde-schatting en meerstapsplanning verzorgt.

1. Root-Prior Injectie (Rollout-fase)

Tijdens de rollout-fase maakt PriorZero gebruik van een Latent World Model (gebaseerd op UniZero) om Monte Carlo Tree Search (MCTS) uit te voeren in een latente ruimte.

• Mechanisme: Een nieuw Root-Prior Injection-mechanisme fuseert de semantische actieprior van de LLM uitsluitend op de rootnode van de MCTS.
• Implementatie: De LLM genereert een situatieanalyse (Chain-of-Thought) en een actieverdeling ($\pi_{LLM}$) op basis van de huidige tekstuele context. Dit wordt op de root gefuseerd met het beleid van het wereldmodel ($\pi_{WM}$):
  $$P_{root}(a | z_t) = (1 - \alpha) \pi_{WM}(a | z_t) + \alpha \pi_{LLM}(a | C_t)$$
• Beperking: Niet-rootnodes in de zoekboom vertrouwen volledig op de latent-dynamiek van het wereldmodel. Dit voorkomt de "granulariteitsmismatch" van Paradigma P3, omdat de LLM nooit wordt gedwongen om toestandsovergangen te voorspellen voor verbeelde latente toestanden.
• Voordeel: Dit leidt de zoektocht naar semantisch veelbelovende acties, voorkomt dat de MCTS in zinloze trajecten terechtkomt (bijvoorbeeld beloningslussen), en behoudt tegelijkertijd het diepe vooruitziend vermogen van het wereldmodel.

2. Alternerende Reinforcement Finetuning (Trainingsfase)

Om het credit assignment-falen van Paradigma P2 aan te pakken, hanteert PriorZero een Alternerende Finetuning-strategie die het leren van het wereldmodel ontkoppelt van de aanpassing van de LLM.

• Schema: Het trainingsschema is asymmetrisch (bijvoorbeeld 10 updates voor het wereldmodel voor elke 1 update voor de LLM).
• Leren van het Wereldmodel: Het wereldmodel wordt bijgewerkt op interactiegegevens om zijn dynamiek, waarde en beleidsvoorspellingen te verfijnen.
• Aanpassing van de LLM: In plaats van gebruik te maken van on-policy returns met hoge variantie, wordt de LLM gefinetuned met low-variance advantage-signalen die zijn afgeleid van de waardeschattingen van het wereldmodel.
  • De advantage wordt berekend met een $n$-step gebootstrapt TD-doel: $A(z_t, a_t) = Q_n(z_t, a_t) - v_\theta(z_t)$.
  • De LLM wordt bijgewerkt via een afgeknipt PPO-doel, met een specifiek wegingsschema dat volledige gewicht toekent aan actietokens en verlaagd gewicht ($w_{cot}=0.1$) aan Chain-of-Thought redeneertokens om stabiliteit te behouden.
• Voordeel: Dit creëert een gesloten lus waarbij het wereldmodel stabiele, fijnkorrelige supervisie biedt voor de LLM, waardoor de priors van de LLM kunnen mee-evolueren met de omgevingsdynamiek zonder de instabiliteit van end-to-end optimalisatie.

Belangrijkste Bijdragen

1. Analyse van Paradigma's: Het artikel identificeert de mismatch tussen prior en dynamiek als het centrale obstakel bij het overdragen van LLM-kennis naar lange-horizon RL, en categoriseert bestaande aanpakken (P1–P4) en hun specifieke faalmodi.
2. PriorZero Framework: Een nieuw architectuur dat LLM-priors en wereldmodelplanning synergetisch combineert via twee kernmechanismen:
   • Root-Prior Injectie: Injecteert semantische begeleiding op de MCTS-root om exploratie te faciliteren zonder de latent-dynamiek te vervuilen.
   • Alternerende Reinforcement Finetuning: Gebruikt waardeschattingen van het wereldmodel om stabiele, low-variance credit assignment-signalen te construeren voor LLM-aanpassing.
3. Empirische Validatie: Uitgebreide experimenten op Jericho (tekstgebaseerde avonturenspellen) en BabyAI (grid-world instructie-opvolging) tonen aan dat PriorZero consistent de exploratie-efficiëntie en asymptotische prestaties verbetert ten opzichte van baselines.

Experimentele Resultaten

• Jericho Benchmarks: PriorZero presteert beter dan UniZero (een sterke wereldmodel-baseline zonder LLM-priors) op vier verschillende spellen (Detective, Acorncourt, Zork1, Omniquest).
  • In Zork1 (lange horizon, schaarse beloningen) bereikt PriorZero hogere asymptotische returns en snellere convergentie.
  • In Detective overtreft PriorZero UniZero na initiële winst, wat de waarde aantoont van het aanpassen van priors aan taakspecifieke dynamiek.
• BabyAI Benchmarks: Op 18 niveaus van de BabyAI grid-world bereikt PriorZero hogere return-plateaus (0,82 versus 0,79 voor UniZero) en toont significante winst in compositieve taken (bijvoorbeeld SynthLoc), waarbij de LLM-prior helpt doelen op te splitsen in subtaken.
• Ablatiestudies:
  • Alternerend Schema: Cruciaal voor stabiliteit; niet-alternerende varianten lijden onder "kenniscorruptie" door ruisige waardesignalen.
  • Alleen-Root Injectie: Essentieel om te voorkomen dat de LLM de dynamiek van het wereldmodel overneemt tijdens het vooruitkijken.
  • CoT-verliesgewicht: Een gebalanceerd gewicht ($w_{cot}=0.1$) is vereist; het verwijderen van CoT veroorzaakt numerieke instabiliteit, terwijl volledig gewicht leidt tot volatiel leren.
  • MCTS: Het uitschakelen van MCTS leidt tot een ineenstorting van de prestaties, wat de noodzaak bevestigt van door planning verbeterde doelen.

Betekenis en Beweringen

Het artikel beweert dat PriorZero een veelbelovend framework vestigt voor besluitvorming met LLM-steun door de structurele knelpunten van prior-overdracht op te lossen. De betekenis hiervan ligt in:

• Ontkoppeling: Het succesvol scheiden van de rol van de LLM (leveren van semantische priors) van de rol van het wereldmodel (leren van dynamiek en plannen), waardoor de granulariteitsmismatch en credit assignment-falen van eerdere methoden worden vermeden.
• Stabiliteit: Het mogelijk maken van stabiele finetuning van LLM's in omgevingen met lange horizon en schaarse beloningen door het wereldmodel te benutten als een low-variance critic.
• Schalbaarheid: Het aantonen dat het framework schaalt met de capaciteit van de LLM (verbetering van prestaties met 7B-modellen in complexe taken) terwijl robuustheid wordt behouden via het alternerende trainingsschema.

De auteurs positioneren PriorZero niet als een vervanging voor wereldmodellen of LLM's, maar als een synergetische integratie waarbij de LLM de zoektocht focust en het wereldmodel de priors van de LLM verfijnt via continue interactie.