WOMBET: World Model-based Experience Transfer for Robust and Sample-efficient Reinforcement Learning

Probleemstelling

Versterkingsleren (RL) in de robotica wordt vaak beperkt door de hoge kosten en risico's van dataverzameling in de echte wereld. Bestaande methoden kampen met specifieke tekortkomingen:

• Online RL: Vereist veel interactie en is vaak onveilig of inefficiënt.
• Offline RL: Kan veilig zijn, maar degradeert vaak bij beperkte of mismatchende data en vereist vaak dure expert-demonstraties.
• Offline-naar-online RL: Combineert pre-training met fine-tuning, maar gaat er vaak ten onrechte van uit dat er al een hoogwaardig offline dataset beschikbaar is, zonder te adresseren hoe deze data veilig gegenereerd kan worden.
• Model-based RL (MBRL): Kan data genereren via een wereldmodel, maar lijdt vaak aan modelbias (overoptimisme) en produceert trajecten met lage diversiteit door exploitatie.

Het kernprobleem is het gebrek aan een framework dat betrouwbare data genereert vanuit een bron-taak en deze veilig overdraagt naar een doeltaak, zonder de beperkingen van statische datasets of onbeheerde modelbias.

Methodologie: WOMBET

WOMBET (World Model-based Experience Transfer) is een unificerend framework dat wereldmodelgebaseerde datageneratie koppelt aan online fine-tuning via een iteratief proces van model- en beleidsevolutie. Het proces verloopt in vier hoofdfasen:

1. World Model-based Data Generatie (Bron-taak):

   • Er wordt een probabilistisch ensemble van wereldmodellen ($\hat{P}$) getraind op de bron-taak.
   • Er wordt gebruikgemaakt van Model Predictive Control (MPC) met een onzekerheidsstraf (uncertainty-penalized planning). De doelstelling is om een ondergrens (lower bound) op de ware opbrengst te maximaliseren door onzekerheid te straffen: $\tilde{r}(s, a) = r(s, a) - \lambda u(s, a)$.
   • Dit zorgt voor conservatieve planning die hoge onzekerheidsgebieden vermijdt.

2. Dubbel-Criteria Filtering (Reliable Data Curation):

   • Niet alle gegenereerde rollouts worden gebruikt. WOMBET filtert trajecten ($\tau$) op basis van twee criteria:
     1. Hoge opbrengst: $J(\tau) \geq J_{th}$.
     2. Lage epistemische onzekerheid: De gemiddelde onzekerheid $\bar{u}(\tau)$ moet onder een drempel $u_{th}$ liggen.
   • Alleen trajecten die aan beide criteria voldoen, worden toegevoegd aan het offline dataset $D_S$. Deze data wordt vervolgens hernaamd met de beloning van de doeltaak ($r_T$).

3. Adaptieve Training (Doeltaak):

   • Het beleid wordt geoptimaliseerd in de doeltaak met een mix van offline data ($D_S$) en online data ($D_T$).
   • Adaptieve Sampling: De mix-ratio $\alpha$ (tussen offline en online data) wordt dynamisch aangepast op basis van de TD-fout (Temporal Difference error) van de critic. Bij hoge fouten (onderfitting) wordt meer gebruikgemaakt van stabiele offline data; bij lagere fouten verschuift de focus naar online data voor adaptatie.
   • Implicit Regularisatie: Om extrapolatiefouten (extrapolation error) te voorkomen bij Out-of-Distribution (OOD) inputs, worden twee mechanismen gebruikt:
     • LayerNorm: Beperkt de grootte van de Q-waarden.
     • Ensemble Minimum: Het gebruik van het minimum van een ensemble van critics (in plaats van het gemiddelde) introduceert een inherente pessimisme die overoptimisme onderdrukt.

4. Iteratieve Verfijning:

   • Het wereldmodel wordt continu bijgewerkt met de geaggregeerde dataset $D = D_S \cup D_T$, waardoor de voorspellingsnauwkeurigheid en de onzekerheidsschattingen in de loop van de tijd verbeteren.

Belangrijkste Bijdragen

1. Gekoppelde Datageneratie en Transfer: In plaats van uit te gaan van een statisch offline dataset, genereert WOMBET zelf betrouwbare prior-data via onzekerheidsbewuste planning en gebruikt deze direct voor transfer.
2. Betrouwbare Data-Curatie en Adaptieve Transfer:
   • Introductie van een dubbel-criteria filter (hoge opbrengst + lage onzekerheid) om modelbias te onderdrukken.
   • Een adaptieve sampling-strategie die de balans tussen bias en variance regelt op basis van de critic's prestatie.
3. Theoretische en Empirische Validatie:
   • Bewijs dat onzekerheidsgestrafde planning een ondergrens maximaliseert op de ware opbrengst.
   • Afleiding van een finite-sample fout-decompositie die distributiemismatch en benaderingsfouten kwantificeert.
   • Experimenten tonen aan dat WOMBET superieur is aan sterke baselines op continue controle-benchmarks.

Resultaten

De experimenten zijn uitgevoerd op diverse MuJoCo-benchmarks:

• Sample Efficiency: WOMBET leert aanzienlijk sneller dan online RL-baselines (SAC, PPO, TD3) en bereikt hogere asymptotische opbrengsten met minder dan de helft van het interactiebudget.
• Noodzaak van Fine-tuning: In tegenstelling tot offline-only methoden (zoals BC, IQL, CQL), die faalden bij significante verschuivingen tussen bron- en doeltaak, slaagt WOMBET door middel van online fine-tuning om zich aan te passen aan nieuwe beloningen en staten.
• Transfer Effectiviteit: WOMBET presteert beter dan bestaande offline-naar-online methoden, zelfs wanneer alle methoden starten met exact hetzelfde gegenereerde dataset. Dit bevestigt dat de verbetering voortkomt uit het geïntegreerde transfermechanisme (filtering + adaptieve sampling).
• Component Analyse: Ablatiestudies tonen aan dat zowel de adaptieve sampling als de dubbel-criteria filtering essentieel zijn. Varianten zonder filtering of met een vaste mix-ratio leiden tot lagere prestaties of instabiel leren.

Betekenis en Conclusie

WOMBET biedt een oplossing voor het fundamentele probleem van datagebrek in robotica-RL. Door de exploitatiekracht van model-based RL (via wereldmodellen) te combineren met de exploratiecapaciteit van model-free RL, creëert het een robuust en sample-efficiënt leerproces.

De belangrijkste innovatie is de verschuiving van het "gebruik van bestaande data" naar het "veilig genereren en selecteren van data" voor transfer. Dit maakt het mogelijk om complexe taken te leren met minder risico en minder interactie met de fysieke wereld, wat een belangrijke stap is richting de praktische toepasbaarheid van RL in de robotica.