Generalization in offline RL: The structure is more important than the amount of pessimism

Titel: Generalisatie in Offline RL: De Structuur is Belangrijker dan de Mate van Pessimisme

Probleemstelling
Offline Reinforcement Learning (RL) streeft naar het leren van optimale policies vanuit vaste, statische datasets zonder verdere interactie met de omgeving. Om overschattingsbias op out-of-distribution (OOD) acties te beperken, maken de meeste state-of-the-art methoden gebruik van pessimistische waarde-lering of conservatisme. Recente literatuur suggereert echter dat overmatig conservatief zijn de generalisatie naar nieuwe scenario's kan hinderen of het 'trajectory stitching' kan belemmeren. Hoewel sommige benaderingen pleiten voor "mild" conservatisme om de generalisatie te verbeteren, blijft het onduidelijk of de graad van pessimisme of de structurele eigenschappen van de pessimistische waardefunctie de primaire determinant is van zero-shot policy transfer (ZSPT) prestaties. Specifiek, in omgevingen met onderliggende symmetrieën (bijv. rotationele invariantie), is het onbekend of een hoogst pessimistische maar symmetrische waardefunctie beter generaliseert dan een mild pessimistische maar asymmetrische waardefunctie.

Methodologie
De auteurs onderzoeken deze vraag door een combinatie van theoretische analyse en empirische validatie binnen de context van Contextual MDP's (CMDP's) onder het Generalization-through-Invariance ZSPT (GTI-ZSPT) framework.

1. Theoretisch Framework:

   • De auteurs formaliseren generalisatie als het vermogen van een agent om de correcte symmetriegroep $G$ van de optimale oplossing te leren, ondanks het feit dat alleen een subgroep $B \leq G$ in de trainingsdata wordt waargenomen.
   • Ze maken gebruik van Q-waarde distillatie als een proxy voor offline RL-mechanismen. Dit isoleert het effect van het leren van een specifieke pessimistische waardefunctie van de complexiteiten van het RL-algoritme zelf.
   • De analyse gaat uit van een oneindig brede neurale netwerk limiet, waarbij gebruik wordt gemaakt van de Neural Tangent Kernel (NTK) en resultaten van Gerken & Kessel (2024) met betrekking tot de equivariantie van deep ensembles.
   • Er worden twee belangrijke stellingen bewezen:
     • Stelling 1: Als een pessimistische waardefunctie de onderliggende symmetriegroep $B$ van de trainingsdata respecteert, kan de agent een optimale generalisatieprestatie bereiken in de testomgeving, ongeacht hoe groot het niveau van pessimisme ($\eta_{max}$) is.
     • Stelling 2: Daarentegen bestaan er instanties waarbij een mild pessimistische waardefunctie die de symmetriestructuur schendt ($\hat{Q}_{asym}$), gegarandeerd slechter zal generaliseren dan een overdreven pessimistische maar symmetrische waardefunctie ($\hat{Q}_{sym}$).

2. Empirische Validatie:

   • Omgeving: Een rotationeel symmetrische "Rotational Reacher" continuous control omgeving. Trainingscontexten dekken een $C_4$ subgroep (90° rotaties), terwijl de testcontexten de volledige $SO(2)$ groep dekken.
   • Algoritmen: De studie evalueert Implicit Q-Learning (IQL) en Conservative Q-Learning (CQL).
   • Data Augmentatie (DA) Strategieën: Het artikel vergelijkt standaard DA (trainen op geaugmenteerde datasets) met Data Augmentation Consistency (DAC) regularisatie. DAC dwingt symmetrie af door een loss-term toe te voegen die het verschil minimaliseert tussen de outputs (of latente representaties) van het netwerk voor originele en geaugmenteerde inputs.
   • Hypothese: De auteurs stellen dat DAC superieur is voor offline RL omdat het expliciet de structuur (symmetrie) van de waardefunctie/policy afdwingt, waarbij prioriteit wordt gegeven aan symmetrie boven de nauwkeurigheid van de waardeschattingen, wat overeenkomt met hun theoretische bevindingen.

Belangrijkste Bijdragen

• Theoretisch Inzicht: Het artikel stelt vast dat de structuur van pessimisme belangrijker is dan de omvang ervan voor generalisatie. Het bewijst dat optimale generalisatie haalbaar is met willekeurig hoog pessimisme, mits de waardefunctie de symmetrieën van de omgeving respecteert.
• Contraintuïtieve Bevinding: Het demonstreert dat mild pessimisme geen garantie is voor betere generalisatie; als het milde pessimisme de vereiste symmetrieën verbreekt, kan de prestatie willekeurig slecht zijn vergeleken met een sterk pessimistische maar symmetrische variant.
• Methodologische Voorstel: De auteurs stellen voor om data augmentatie toe te passen via een consistentie-loss tijdens policy extractie (specifiek voor de actor in IQL), in plaats van de standaardpraktijk van trainen op geaugmenteerde datasets. Deze aanpak dwingt de symmetrische structuur van de policy/waardefunctie direct af.
• Empirisch Bewijs: Experimenten op IQL en CQL tonen aan dat DAC-Output (het afdwingen van consistentie op de netwerkoutput) de beste generalisatieprestaties levert, waarmee zowel de no-augmentation baselines als de standaard augmented training (Aug-D) worden overtroffen.

Resultaten

• Q-Waarde Distillatie: Bij de theoretische validatie met de Rotational Reacher behield een symmetrische pessimistische target ($\hat{Q}^*_{sym}$) een bijna optimale testprestatie, zelfs wanneer het pessimisme met een factor tien toenam. In tegen afstelling zag een asymmetrische target ($\hat{Q}^*_{asym}$) dat de testprestaties significant verslechterden naarmate het pessimisme toenam.
• Offline RL Experimenten:
  • Voor zowel IQL als CQL leverde het toepassen van DAC-Output consequent de hoogste testreturns op voor expert, mixed en suboptimal datasets.
  • Standaard data augmentatie (Aug-D) verbeterde ten opzichte van de no-augmentation baseline, maar presteerde over het algemeen minder goed of gelijk aan de methoden met consistentie-loss.
  • Het toepassen van DA enkel op de critic (zonder de actor) in IQL toonde geen verbetering ten opzichte van de baseline, wat bevestigt dat het afdwingen van symmetrie op het policy-extractiecomponent cruciaal is.

Significantie en Claims
Het artikel claimt dat de heersende intuïtie — dat het verminderen van conservatisme noodzakelijk is voor generalisatie — incompleet is. In plaats daarvan betogen de auteurs dat het afdwingen van de correcte structurele symmetrieën de primaire drijfveer is voor succesvolle zero-shot transfer in offline RL.

De significantie van dit werk ligt in de verschuiving van de focus van "hoe conservatief moeten we zijn?" naar "respecteert ons conservatisme de symmetrieën van de omgeving?". De auteurs stellen dat in offline RL, waar dataset coverage vaak niet-symmetrische biases induceert, technieken zoals Data Augmentation Consistency (DAC) essentieel zijn. Door prioriteit te geven aan symmetrie boven de waarde-accuratesse tijdens policy extractie, stelt DAC agents in staat om optimaal te generaliseren, zelfs wanneer de onderliggende waardefunctie zeer pessimistisch is. Dit daagt de standaardpraktijk uit om data augmentatie enkel te gebruiken om de datasetgrootte voor reguliere training te vergroten, en suggereert in plaats daarvan dat het gebruikt moet worden als een regularisator om structurele invariantie af te dwingen.

Beperkingen
De auteurs erkennen enkele beperkingen:

• De theoretische bewijzen vertrouwen op de infinite width limit van neurale netwerken, hoewel empirische resultaten op eindige breedte netwerken de theorie ondersteunen.
• De analyse richt zich op intrinsieke groepsymmetrieën die consistent zijn met de datadistributie. Het is nog niet bewezen of deze conclusies ook gelden voor extrinsieke of inconsistente transformaties (bijv. willekeurige ruis) die enkel voor regularisatie worden gebruikt.
• De bevindingen zijn momenteel gevalideerd op de Rotational Reacher taak en twee specifieke algoritmen (IQL, CQL); bredere validatie over complexe omgevingen en algoritmen is nodig.