Lyapunov-Guided Self-Alignment: Test-Time Adaptation for Offline Safe Reinforcement Learning

1. Probleemstelling

Offline Reinforcement Learning (RL)-agenten, getraind op statische datasets (bijv. D4RL, DSRL), falen vaak bij implementatie in real-world omgevingen door distributieveranderingen. Zelfs als een agent goed presteert op trainingsdata, kunnen kleine verschillen in dynamiek of staatsovergangen leiden tot onveilig gedrag (bijv. botsingen, het overschrijden van kostendrempels).

Bestaande oplossingen staan voor aanzienlijke beperkingen:

• Op geloof gebaseerde adaptatie: Vereist hertraining vanaf nul om onzekerheid te modelleren.
• Beperkte/Conservatieve RL: Offer vaak prestaties (opbrengst) op of vereist fijne afstelling met extra kostensignalen.
• Adaptatie op testtijd: De meeste huidige methoden kunnen een vooraf getrainde offline agent niet aanpassen aan veiligheidsbeperkingen bij implementatie zonder parameterupdates of menselijk toezicht.

De kernvraag die wordt behandeld is: Kan een offline RL-agent op testtijd zelf worden uitgelijnd voor veiligheid, uitsluitend gebruikmakend van zijn vooraf getrainde distributie, zonder hertraining?

2. Methodologie: SAS (Self-Alignment for Safety)

De auteurs stellen SAS voor, een transformer-gebaseerd kader dat adaptatie op testtijd mogelijk maakt via Lyapunov-geleide zelfuitlijning. De methode werkt zonder modelparameters bij te werken, en vertrouwt in plaats daarvan op in-context learning en Bayesiaanse inferentie.

Kerncomponenten

1. Lyapunov Dichtheidsmodel (LDM) voor Offline RL:

   • De auteurs herformuleren Lyapunov-stabiliteit met behulp van bezettingsmaten afgeleid van offline data.
   • Ze definiëren een energie-functie $E(s, a) = -\log \hat{\rho}_\pi(s, a)$, waarbij $\hat{\rho}_\pi$ de modelgebaseerde bezettingsmaat is geschat uit verbeelde trajecten.
   • Een Lyapunov-surrrogaat $G_{SAS}(s, a)$ wordt geconstrueerd om ervoor te zorgen dat de agent binnen een controle-invariante verzameling ($R_{SAS}^G$) blijft. Deze verzameling garandeert dat trajecten geen gebieden met lage dichtheid (onveilig/buiten distributie) binnengaan.
   • $G_{SAS}$ wordt gedefinieerd als de minimale energie die wordt aangetroffen langs elk traject, waardoor het beleid voorkomt dat het instort in onveilige, lage-probabiliteitsgebieden.

2. Probabilistische Inferentie Visie:

   • Het paper interpreteert de transformer-gebaseerde RL-agent als een Hiërarchisch RL-systeem waarbij het hoog-niveau beleid latente vaardigheden ($z_t$) selecteert en het laag-niveau beleid acties uitvoert.
   • In-context learning wordt geframed als Bayesiaanse inferentie over deze latente vaardigheden. Door een prompt (een reeks staat-actieparen) te verstrekken, leidt het model de meest waarschijnlijke hoog-niveau beleid af die consistent is met die prompt.
   • Veiligheidsvoorwaarden worden gecodeerd als observeerbare indicatoren ($U_t$ en $V_t$):
     • $U_t$: Geeft aan of het huidige staat-actiepaar binnen het veilige dichtheidsgebied valt ($G_{SAS} > 0$).
     • $V_t$: Geeft aan of de Lyapunov-functie niet-toenemend is ($G_{SAS}(s_t, a_t) \geq G_{SAS}(s_{t+1}, a_{t+1})$).

3. Het SAS-algoritme (Algoritme 1):
   Het proces omvat twee lussen van "verbeelding" en selectie:

   • Stap 1 (Voorwaarde $U_t$): Vanuit een vaste beginstaat verbeeldt de transformer $N$ trajecten. Het selecteert het trajectsegment dat de maximale energie (worst-case risico) minimaliseert om een initiële prompt te vormen.
   • Stap 2 (Voorwaarde $V_t$): Met behulp van de initiële prompt verbeeldt de transformer $M$ nieuwe trajecten. Het selecteert het traject dat de cumulatieve Lyapunov-afname maximaliseert (som van $V_t$-indicatoren).
   • Stap 3 (Implementatie): Het geselecteerde veilige segment wordt gebruikt als in-context prompt ($p_{1:K}$) voor de transformer. De agent baseert zijn acties vervolgens op deze prompt en de huidige staat, waardoor het vooraf getrainde beleid effectief wordt "gestuurd" naar veilig gedrag zonder gradiëntupdates.

3. Belangrijkste Bijdragen

1. Nieuwe Formulering: Een nieuwe definitie van Lyapunov-stabiliteit voor offline RL gebaseerd op bezettingsmaten, waardoor veiligheidsbeperkingen kunnen worden afgedwongen via datadichtheid in plaats van expliciete kostfuncties.
2. Theoretisch Kader: Een Hiërarchisch RL-interpretatie van transformer-gebaseerde RL via Bayesiaanse inferentie, waarbij prompting overeenkomt met het afleiden van latente hoog-niveau vaardigheden. Dit biedt een theoretische basis voor waarom in-context learning werkt voor veiligheid.
3. Adaptatiemethode op Testtijd (SAS): Een praktisch algoritme dat Lyapunov-geleide prompts genereert om het gedrag van de agent opnieuw uit te lijnen. Het garandeert veiligheid door de agent te beperken tot controle-invariante verzamelingen afgeleid van offline data.
4. Probabilistische Veiligheidsgarantie: De auteurs bieden een formele grens (Stelling 4) die aantoont dat de waarschijnlijkheid dat de agent de veilige verzameling verlaat, exponentieel afneemt met het aantal verbeelde rollouts ($N$) en opnieuw-verbeeldingsproeven ($M$).

4. Experimentele Resultaten

De methode werd geëvalueerd op Safety Gymnasium, MuJoCo en Bullet-Safety-Gym benchmarks.

• Prestaties vs. Veiligheid: SAS verlaagde consistent kosten en faalpercentages (tot 2x verbetering) terwijl het de opbrengst (beloning) handhaafde of verbeterde in vergelijking met baselines zoals Decision Transformer (DT), Constrained Decision Transformer (CDT) en diverse veilige RL-baselines (BC-Safe, COptiDICE, etc.).
• Generalisatie: SAS verbeterde de veiligheid zelfs voor ruggengraatmodellen (DT) die niet waren getraind met expliciete kostensignalen, wat aantoont dat de op dichtheid gebaseerde Lyapunov-gids fungeert als een impliciete veiligheidsproxy.
• Ablatiestudies:
  • Lyapunov Afname ($V_t$): Het verwijderen van de afnamevoorwaarde verhoogde de kosten aanzienlijk, wat de noodzaak ervan bewijst voor stabiliteit op lange termijn.
  • Prompt Selectie: Willekeurige prompting of prompting met maximale energie presteerde slechter dan de Lyapunov-geleide selectie, wat de doeltreffendheid van de specifieke selectiecriteria bevestigt.
  • Model Robuustheid: De methode toonde robuustheid over verschillende promptlengtes, aantallen verbeelde trajecten en ruggengraatmodelgroottes (van DT tot GPT-2).

5. Betekenis en Impact

• Brug tussen LLM en RL: Het paper past het concept van zelfuitlijning (populair in Large Language Models) succesvol toe op het domein van Veilig RL, waardoor agenten kunnen "redeneren" over veiligheid met behulp van hun eigen interne wereldmodellen.
• Implementatiehaalbaarheid: Door de behoefte aan hertraining of fijne afstelling te elimineren, biedt SAS een schaalbare oplossing voor het implementeren van offline RL-agenten in dynamische, onzekere omgevingen waar veiligheid kritiek is.
• Theoretische Strenge: Het gaat voorbij aan heuristische veiligheidscontroles door een probabilistische garantie te bieden dat de agent binnen een controle-invariante verzameling blijft, gebaseerd op Lyapunov-stabiliteitstheorie en Bayesiaanse inferentie.
• Efficiëntie: Hoewel het rekenkracht op inferentietijd vereist voor verbeelding (rollouts), vermijdt het de rekenkosten van volledige hertraining, waardoor het geschikt is voor scenario's waar data statisch is maar de omgeving niet.

Kortom, SAS vertegenwoordigt een significante stap voorwaarts in het robuust maken van offline RL-agenten voor distributieveranderingen door gebruik te maken van Lyapunov-geleide verbeelding om gedrag op testtijd zelf te corrigeren via in-context learning.