Technische Samenvatting: SAPO (Single-Rollout Autoregressive Policy Optimization)

1. Probleemstelling

Agentic Reinforcement Learning (RL) is een cruciale post-training fase geworden voor Large Language Models (LLM's), waardoor ze langdurige interactieve taken kunnen uitvoeren zoals redeneren, het gebruik van tools en webnavigatie. Hoewel bestaande methoden succes hebben geboekt, kampen ze met drie primaire beperkingen:

1. Gebrek aan expliciete waarde-generalisatie: Critic-vrije, groep-relatieve methoden (bijv. GRPO) schatten voordelen (advantages) in door beloningen te normaliseren over meerdere gesamplede trajecten. Deze aanpak faalt in het leren van een gegeneraliseerde waardefunctie die specifieke staten of prompts onafhankelijk van de groep kan beoordelen, wat leidt tot zwakke temporele credit assignment.
2. Advantage Collapse: In taken met een lange horizon en ijle (sparse) feedback is het gebruikelijk dat alle rollouts in een groep identieke of bijna identieke beloningen ontvangen. Dit veroorzaakt dat de genormaliseerde advantages verdwijnen, wat resulteert in regio's met een nul-gradiënt of "advantage collapse", wat het leerproces stagneert.
3. Inefficiënte Trade-offs: Groep-relatieve methoden vereisen een kostbare afweging tussen het sampling-budget en prestaties. Kleine groepen leveren ruisende baselines op, terwijl grote groepen de stabiliteit verbeteren ten koste van enorme computationele kosten en synchronisatiebarrières. Omgekeerd biedt traditionele Proximal Policy Optimization (PPO) met een geleerde critic weliswaë expliciete waarde-lering, maar brengt het aanzienlijke geheugen- en computationele overhead met zich mee door de noodzaak van een aparte critic-model op policy-schaal.

Huidige pogingen om deze problemen te verzachten (bijv. DAPO, ReMax, VinePPO) blijven ofwel afhankelijk van meerdere rollouts, behouden grove credit assignment, of vereisen aparte waardenetwerken die parameters dupliceren. Er is behoefte aan een framework dat efficiënt single-rollout leren ondersteunt met expliciete temporele waarde-estimatie zonder een aparte critic te onderhouden.

2. Methodologie: SAPO

De auteurs stellen Single-Rollout Autoregressive Policy Optimization (SAPO) voor, een framework met een laag geheugen- en compute-verbruik dat policy- en waarde-lering verenigt binnen één autoregressieve backbone.

Kernarchitectuur: Autoregressieve Actor-Critic

SAPO maakt gebruik van de causale structuur van LLM's om de policy, de staatswaarde ($V$) en de actiewaarde ($Q$) binnen een enkele modelstroom te representeren.

• Causale Grenzen: De methode identificeert twee duidelijke causale grenzen binnen een beurt van een agent:
  1. Pre-action boundary: De staat-context ($c^s_t$) vóór de actie wordt gegenereerd.
  2. Post-action boundary: De staat-context geconcateneerd met de gegenereerde actie ($c^{s,a}_t$).
• Gedeelde Backbone: Hetzelfde transformer-model genereert de actie-tokens tussen deze grenzen.
• Two-Token Waarde-basis: In plaats van een aparte waarde-head, reserveert SAPO twee specifieke vocabulaire-tokens ($w^+$ en $w^-$) exclusief als waarde-basis. Het model leest de logits van deze twee tokens bij de respectieve causale grenzen om genormaliseerde scalaire waarden te berekenen:
  • $V_\theta(s_t)$ wordt gelezen bij de pre-action boundary.
  • $Q_\theta(s_t, a_t)$ wordt gelezen bij de post-action boundary.
  • De waarde wordt afgeleid van het verschil in logits, afgeknipt naar een genormaliseerd bereik van $[-1, 1]$, en geschaald door de maximale mogelijke return ($R_{max}$).
• Policy Separatie: De gereserveerde waarde-tokens worden uit de actie-vocabulaire uitgesloten tijdens de policy-generatie en waarschijnlijkheidsnormalisatie, zodat ze de actie-distributie niet verstoren.

Single-Rollout Advantage Estimatie

In tegen tegenstelling tot groep-relatieve methoden, samplet SAPO slechts één traject per taak.

• Traject-niveau GAE: De methode gebruikt een Generalized Advantage Estimator (GAE) die vertraagde terminale beloningen via het traject terugverplaatst door de geleerde staatswaarden ($V$).
• Targets:
  • De Staatswaarde wordt getraind richting de $\lambda$-return (het propageren van langetermijn-uitkomstinformatie).
  • De Actiewaarde wordt getraind richting een on-policy SARSA-target (het funderen van de post-action boundary in onmiddellijke consequenties).
• Batch Normalisatie: Om de training te stabiliseren zonder groep-statistieken, worden turn-level advantages genormaliseerd met behulp van batch-statistieken ($\mu_B, \sigma_B$) over de huidige trainingsbatch voordat ze naar alle tokens in de actie worden uitgezonden.

Gezamenlijke Optimalisatie

SAPO optimaliseert een verenigde objectieve functie in één enkele backward pass:

1. Token-niveau Policy Loss: Een afgeknipte (clipped) PPO-stijl loss met behulp van de batch-genormaliseerde advantages.
2. Waarde-losses: Afgeknipte regressie-losses voor zowel de $V$- als de $Q$-objectieven, gemiddeld over turns in plaats van tokens om ervoor te zorgen dat elke beslissing een gelijke weging krijgt, ongeacht de responslengte.
3. Regularisatie: Bevat KL-divergentie tegen een referentie-policy en entropie-regularisatie.

3. Belangrijkste Bijdragen

1. Verenigde Architectuur: Introductie van een single-rollout autoregressieve actor-critic architectuur die policy, staatswaarde en actiewaarde gezamenlijk representeert en traint via causale-boundary readouts, waardoor de noodzaak voor een apart critic-netwerk vervalt.
2. Nieuwe Estimator: Ontwikkeling van een traject-niveau advantage estimator die $\lambda$-returns combineert met batch-normalisatie, wat expliciete turn-level credit assignment mogelijk maakt zonder afhankelijk te zijn van groep-relatieve sampling.
3. Efficiëntie en Stabiliteit: Demonstratie dat dit verenigde framework een gunstige balans bereikt tussen prestaties, trainingsstabiliteit, geheugenefficiëntie en runtime-efficiëntie op langdurige agentic taken.

4. Experimentele Resultaten

De auteurs evalueerden SAPO op de ALFWorld (embodied reasoning) en WebShop (webnavigatie) benchmarks met behulp van Qwen2.5-1.5B en Qwen2.5-7B backbone-modellen.

• Prestaties: SAPO presteerde consequent beter dan standaard PPO en GRPO.
  • Op ALFWorld met Qwen2.5-1.5B behaalde SAPO 90,1% geaggregeerde succes, een verbetering van +35,7 en +17,3 procentpunten ten opzichte van respectievelijk PPO en GRPO.
  • Op WebShop behaalde SAPO een score-verbetering van +8,4 punten en een toename in succespercentage van +12,2 punten ten opzichte van PPO.
  • Over de experimenten heen (zoals samengevat in de Abstract), traint SAPO stabiel en presteert het gemiddeld +15,1 en +12,1 procentpunten beter dan PPO en GRPO, respectievelijk.
  • Met het 7B-model behaalde SAPO 94,0% ALFWorld succes en een WebShop-score van 88,6 met 82,4% succes, waarmee het zowel PPO als GRPO als recente varianten op alle drie de geaggregeerde metrieken verslaat.
• Efficiëntie:
  • Geheugen: SAPO elimineert de geheugenvoetafdruk van een apart critic-model (dat doorga's policy-schaal is).
  • Runtime: SAPO verminderde de per-iteratie runtime met 33,2% vergeleken met PPO. De primaire besparingen kwamen voort uit het verwijderen van de aparte waarde-inferentie en critic-optimalisatiestappen, waardoor de totale tijd per iteratie op ALFWorld daalde van 451,2s naar 301,4s.
• Stabiliteit: De methode trainde stabiel vanaf een enkele rollout per prompt, waarbij de advantage collapse-problemen die bij groep-relatieve methoden worden gezien onder ijle feedback, werden vermeden.

5. Betekenis en Claims

Het artikel claimt dat SAPO de spanning oplost tussen de expliciete waarde-generalisatie van actor-critic algoritmen en de geheugen-/computationele efficiëntie van critic-vrije, groep-relatieve optimalisatie.

• Architecturale Uitlijning: De auteurs beargumenteren dat de informatieflow die vereist is voor actor-critic leren natuurlijk aansluit bij de causale volgorde van taalgeneratie. Door deze uitlijning te benutten, maakt SAPO efficiënte waarde-gebaseerde optimalisatie mogelijk binnen een enkele causale model.
• Schaalbaarheid: Het framework demonstreert dat expliciet waarde-leren geen gedubliceerde modellen of kostbare sampling vereist, wat een levensvatbaar pad biedt naar efficiënte langdurige agentic RL.
• Generalisatie: De resultaten suggereren dat de architecturale uitlijning tussen autoregressieve modellering en actor-critic leren de afhankelijkheid van taakspecifieke dense reward engineering kan verminderen, aangezien de methode goed presteert met ijle (sparse) uitkomstbeloningen.

Het artikel concludeert dat SAPO een belangrijke stap voorwaarts vertegenwoordigt in het schaalbaarder en praktischer maken van agentic RL voor complexe, multi-turn interacties.