Technische Samenvatting: CAST (Credit Assignment from Solver Teachers)

Probleemstelling

Het trainen van Large Language Models (LLMs) om te fungeren als generalistische agenten in scenario's met een lange horizon, zoals klassieke spellen (Sokoban, Minesweeper, Rush Hour), staat voor een kritieke flessenhals: sparse reward-signalen (ijle beloningssignalen). Hoewel Reinforcement Learning met Verifieerbare Beloningen (RLVR) succesvol is geweest in taken met één enkele beurt, leunt het op resultaatbeloningen aan het einde van de rit (winst/verlies) die geen informatie bieden over welke specifieke tussenliggende beslissingen hebben geleid tot succes of falen. Dit creëert een fundamenteel probleem van credit-assignment (toeschrijving van krediet), waarbij elke token in een traject dezelfde grove feedback ontvangt, wat het vermogen van het model om te leren over planning op lange termijn, doelgerichte exploratie en foutherstel belemmert. Bestaande pogingen om deze signalen te verdichten via kostbare zoekprocessen, geleerde process-reward modellen of vergelijkingen tussen trajecten, lijden vaak aan afruil tussen computationele kosten, kwaliteit van supervisie en de betrouwbaarheid van het signaal.

Methodologie: CAST

De auteurs stellen CAST (Credit Assignment from Solver Teachers) voor, een framework dat domeinspecifieke spel-solvers gebruikt om fijnmazige, per beurt gebaseerde supervisie te bieden aan LLM-agenten.

Kernmechanisme

1. Solver als Turn-Level Teacher: In plaats van een solver alleen te gebruiken om expert-trajecten te genereren (wat het leren beperkt tot door experts bezochte staten), gebruikt CAST de solver om de on-policy acties van de LLM in realtime te evalueren.
2. Constructie van de Solver Advantage:
   • De solver berekent een cost-to-go waarde, $N(s)$, die de minimale stappen vertegenwoordigt die nodig zijn om te winnen vanuit staat $s$.
   • Voor een actie $a_t$ die door de LLM wordt ondernomen bij de overgang van $s_t$ naar $s_{t+1}$, berekent de methode de verandering in de cost-to-go.
   • Een shifted solver advantage wordt gedefinieerd als:
     $$\tilde{A}^{\pi_{Solver}}(s_t, a_t) = N(s_t) - \mathbb{E}[N(s_{t+1})]$$
   • Deze waarde is positief als de actie de afstand tot het doel verkleint (vooruitgang), nul als deze neutraal is, en negatief als het de afstand vergroot of tot een doodlopend pad leidt.
3. Signal Shaping: Om instabiliteit tijdens de training te voorkomen die wordt veroorzaakt door heavy-tailed waarden (bijv. grote straffen voor doodlopende paden) en schaalverschillen tussen spellen, past CAST twee transformaties toe:
   • Asinh Compressie: $g(x) = \text{asinh}(x)$, wat kleine waarden lineair behoudt terwijl extreme uitschieters logaritmisch wordt gecomprimeerd.
   • Batch-level RMS Normalisatie: Schaalt het signaal door de wortel van de gemiddelde kwadratische grootte (RMS) binnen een batch om consistente schalen te garanderen, zonder het gemiddelde af te trekken (om de semantische betekenis van nul als "geen vooruitgang" te behouden).
4. Integratie met RLVR: De gevormde solver advantage wordt toegevoegd aan de standaard traject-niveau resultaat-advantage (bijv. uit GRPO) om een gecombineerd advantage-signaal te vormen:
   $$\hat{A}_{i,t} = \hat{A}^{outcome}_i + \alpha \cdot h(\tilde{A}^{\pi_{Solver}}_{i,t})$$
   Dit gecombineerde signaal wordt gebruikt in de clipped surrogate objective om het beleid (policy) bij te werken.

Theoretisch Inzicht: Logit-vrije On-Policy Distillatie

Het paper stelt een theoretische equivalentie vast tussen het maximaliseren van de solver advantage en on-policy distillatie (OPD) van een soft-optimale teacher.

• Onder de aanname dat een sterke solver fungeert als een soft-optimale policy waarbij de actiekansen proportioneel zijn aan $\exp(Q/\tau)$, is de solver advantage wiskundig proportioneel aan de log-waarschijnlijkheid van de actie onder de policy van de solver.
• Kernbijdrage: Dit impliceert dat de scalaire solver advantage inherent de voorkeursdistributie van de teacher codeert. Bijgevolg bereikt CAST distillatie zonder de noodzaak van teacher logits (de volledige actieverdeling), door enkel te vertrouwen op scalaire waarden. Dit stelt de student-LLM in staat om efficiënt te leren van de waarde-functie van de solver.

Belangrijkste Bijdragen

1. Solver-afgeleide Credit: De introductie van een methode om veranderingen in de state-waarde van een spel-solver om te zetten in fijnmazige, per beurt gebaseerde advantages voor LLM-training.
2. Logit-vrije Distillatie: Een theoretisch bewijs dat het maximaliseren van solver advantages equivalent is aan on-policy distillatie, wat efficiënte training mogelijk maakt met enkel scalaire waarden in plaats van volledige teacher-distributies.
3. Stabiele Signal Shaping: Het gebruik van asinh-compressie en batch-level RMS-normalisatie om extreme waarden en variaties in schaal tussen domeinen te beheersen, wat zorgt voor een stabiele integratie in RLVR.
4. Praktische Bruikbaarheid: Demonstratie dat exacte solvers verwaarloosbare overhead toevoegen en dat geleerde waarde-netwerken (bijv. DQN) effectieve benaderingen kunnen dienen wanneer exacte solvers niet beschikbaar zijn.

Experimentele Resultaten

De methode werd geëvalueerd op Sokoban, Minesweeper en Rush Hour, met zero-shot transfer naar ALFWorld en WebShop.

• In-Domain Prestaties: CAST presteerde beter dan alle getrainde baselines (GRPO, GSPO, DAPO, GiGPO) en training-vrije prompting-methoden (inclusclusief sterke closed-source modellen zoals Claude-Opus-4.6) in alle drie de spellen, zowel in in-domain als in ongeziene moeilijkheidsgraden.
  • Gemiddeld behaalde CAST een succespercentage van 62,1% (Avg@4) op in-domain taken en 28,4% op ongeziene moeilijkheidsgraden, waarmee het de backbone DAPO (44,7% ID, 18,7% Unseen) aanzienlijk versloeg.
• Sample Efficiëntie: CAST bereikte de piek validatieprestatie van DAPO in 1,7× tot 2,0× minder trainingsstappen over de drie spellen.
• Zero-Shot Generalisatie: Zonder verdere fine-tuning vertoonden de CAST-agenten een transfer naar ALFWorld en WebShop, waarbij ze de hoogste gemiddelde zero-shot prestatie (30,3% totaal) behaalden van alle methoden, waarmee ze de eerstvolgende beste baseline met 5,6 punten versloegen.
• Ablatie-studies:
  • De gewicht van de solver-advantage ($\alpha$) is cruciaal; $\alpha=0,1$ bood de beste balans tussen procesbegeleiding en outcome-objectieven.
  • Zowel asinh-compressie als RMS-normalisatie waren essentieel; het verwijderen van een van beide verslechterde de prestaties, waarbij het verwijderen van asinh de grootste daling veroorzaakte vanwege gradiënt-instabiliteit door extreme waarden.
• Geleerde Solvers: Het vervangen van de exacte solver door een geleerd DQN-waardenetwerk behield de meeste voordelen qua prestaties, wat de toepasbaarheid van de methode buiten spellen met exacte solvers bewijst.

Betekenis en Claims

Het paper claimt dat CAST het fundamentele probleem van credit-assignment in langdurige RLVR aanpakt door een schaalbaar, fijnmazig signaal te bieden dat afgeleid is van domein-solvers. Door de equivalentie met logit-vrije on-policy distillatie te bewijzen, biedt de methode een praktische weg om LLM's af te stemmen op expert waarde-functies zonder de computationele last van het genereren van volledige actieverdelingen. De resultaten suggereren dat betrouwbare staat-evaluatie—of dit nu via exacte solvers of geleerde waarde-functies gebeurt—een robuuste route biedt van ijle uitkomsten naar dichte, actiegerichte feedback, waardoor LLM's complexe, meerstaps besluitvormingstaken efficiënter kunnen beheersen dan alleen op uitkomsten gebaseerde benaderingen.