LaGO: Latent Action Guidance for Online Reinforcement Learning

Technische Samenvatting: LaGO – Latent Action Guidance voor Online Reinforcement Learning

Probleemstelling
Hoewel Large Language Models (LLM's) aanzienlijke potentie hebben getoond in planning en sequentiële besluitvorming, vertrouwen bestaande benaderingen vaak op het gebruik van deze modellen als directe controllers, expliciete high-level planners of wereldmodel-simulatoren. Deze methoden vereisen dat LLM's precieze acties, gedetailleerde plannen of nauwkeurige toekomstige trajecten genereren. In de praktijk is deze vereiste echter vaak te streng; de capaciteiten van LLM's variëren aanzienlijk, en zelfs kleine voorspellingsfouten kunnen over tijd accumuleren, wat leidt tot slechte controleprestaties in reinforcement learning (RL)-omgevingen die een precieze uitvoering en langetermijnredenering vereisen. Bijgevolg blijft het direct vertrouwen op een LLM als primaire controller onbetrouwbaar, met name in gevectoriseerde omgevingen waar observaties en acties niet van nature in taal of visuele vorm zijn gerepresenteerd.

Methodologie
Om deze uitdagingen aan te pakken, stellen de auteurs LaGO (Latent Action Guidance for Online Reinforcement Learning) voor, een tweestaps-framework dat een vooraf getraind LLM niet gebruikt als een expliciete controller, maar als een bron van "zachte" sturing via latente actie-priors.

1. Fase 1: Offline Prior Learning
   Het framework begint met het trainen van een latente beleidsprior vanuit een offline expert-dataset ($D$). Een vooraf getrainde LLM-backbone ($F_\phi$) wordt gebruikt, waarbij niet-tekstuele staat-actie-sequenties uit demonstraties worden omgezet in latente representaties via leerbare projectiemodules ($E_s$ en $E_a$). Deze embeddings, gecombineerd met een tekstuele taakbeschrijving ($x$), worden in de LLM gevoed. Een nieuw toegevoegde actie-kop ($G_\theta$) wordt vervolgens getraind via supervised fine-tuning om de volgende actie ($a_{t+1}$) te voorspellen op basis van het traject-prefix ($\tau_t$) en de taakbeschrijving. Het doel is om de negatieve log-likelihood van de expert-data te minimaliseren, waardoor een latente beleidsprior $\pi_K(a_{t+1} | x, \tau_t)$ ontstaat die een grove gedragsmatige prior vastlegt in plaats van een optimale controller.

2. Fase 2: Online RL met Latent Guidance
   In de tweede fase wordt de getrainde latente prior $\pi_K$ bevroren en gebruikt om een online RL-beleid $\pi_\psi(a|s)$ (bijv. PPO of SAC) te sturen. De RL-agent interacteert met de omgeving om rollouts te verzamelen, maar de beleidsoptimalisatie wordt geregulariseerd richting de latente prior. De totale trainingsverliesfunctie is gedefinieerd als:
   $$L_{online} = L_{PG} + \beta L_{prior}$$
   waarbij $L_{PG}$ de standaard policy gradient-doelstelling is en $\beta$ de invloed van de prior reguleert.

   • Voor discrete actieruimtes is de prior-loss de negatieve log-likelihood van de door de prior voorspelde actie.
   • Voor continue actieruimtes wordt de prior-loss geïmplementeerd als een cross-entropy term tussen de distributie van de prior en het huidige beleid, wat dient als een equivalent optimalisatiesignaal voor KL-divergentie regularisatie (tot de entropie van de vaste prior).

Dit ontwerp stelt het RL-beleid in staat om te profiteren van de structurele kennis die in de LLM is gecodeerd, zonder dat de LLM tijdens interactie exacte beslissingen hoeft te nemen.

Belangrijkste Bijdragen

• Nieuw Framework: LaGO introduceert een methode om de kennis van een vooraf getraind LLM te benutten als een latente gedragsmatige prior voor online RL, wat de strikte nauwkeurigheidseisen voor het gebruik van LLM's als directe controllers versoepelt.
• Tweestaps-architectuur: De scheiding tussen offline prior learning (via supervised fine-tuning op expert-data) en online beleidsoptimalisatie (via regularisatie) maakt het systeem in staat om LLM-kennis te benutten, zelfs wanneer de prior niet nauwkeurig genoeg is om als zelfstandige agent te fungeren.
• Generalisatie: Het framework is ontworpen om te werken in zowel discrete als continue controle-domeinen door de latente prior-regularisatieterm aan te passen aan de specifieke actieruimte.

Experimentele Resultaten
De auteurs hebben LaGO geëvalueerd op twee benchmarks: CLEVR-Robot (discrete controle) en Meta-World (continue controle). De experimenten vergeleken LaGO met "Vanilla PPO" over vier taakcategorieën: Real, Rephrase, Easy en Hard.

• Prestatiewinst: LaGO presteerde consistent beter dan Vanilla PPO op zowel beloning als succespercentage.
  • Op CLEVR-Robot steeg het gemiddelde succespercentage van 15,1% naar 27,2%.
  • Op Meta-World steeg het gemiddelde succespercentage van 2,7% naar 15,2%.
• Generalisatie: Verbeteringen werden waargenomen in alle taakcategorieën, inclusocief "Easy" en "Hard" ongeziene taken, wat suggereert dat de prior algemene sturing biedt in plaats van louter memorisatie.
• Impact van LLM-kwaliteit: Een ablatie-studie die een Llama-2-7B prior vergeleek met een TinyLlama-1.1B prior op Meta-World, toonde aan dat sterkere vooraf getrainde modellen effectievere latente priors opleveren. LaGO met Llama-2-7B behaalde consistent de hoogste prestaties, terwijl de TinyLlama prior soms onderpresteerde ten opzichte van Vanilla PPO, wat het belang van de kwaliteit van het fundamentele model benadrukt.

Betekenis en Claims
Het artikel claimt dat LaGO een eenvoudig en praktisch framework biedt voor het integreren van vooraf getrainde LLM-kennis in reinforcement learning. De primaire betekenis ligt in het aantonen dat LLM's de online besluitvorming en planning kunnen verbeteren, zelfs wanneer ze niet nauwkeurig genoeg zijn om als directe controllers te dienen. Door de LLM te behandelen als een bron van zachte, latente sturing in plaats van een expliciete planner, mitigeert het framework de risico's van foutaccumulatie, terwijl het profiteert van de rijke inductieve biases en wereldkennis verworven tijdens de pretraining.

De auteurs concluderen dat de kennis en generalisatievermogens gecodeerd in vooraf getrainde LLM's kunnen dienen als sterke priors voor planning en RL. Ze suggereren verder dat naarmate foundation modellen verbeteren, de effectiviteit van deze sturing van nature zal toenemen, wat een schaalbaar pad biedt voor het benutten van LLM's in complexe besluitvormingstaken zonder dat daarvoor nauwkeurige actiegeneratie tijdens interactie vereist is.