Latent Wasserstein Adversarial Imitation Learning

Dit paper introduceert Latent Wasserstein Adversarial Imitation Learning (LWAIL), een nieuw framework dat door gebruik te maken van een dynamische latente ruimte en slechts één expert-episode in staat is om expert-niveau prestaties te bereiken zonder toegang tot expert-acties.

De kern

Het Grote Probleem: Leren zonder de "Handreiking"

Stel je voor dat je een robot wilt leren om te dansen. De traditionele manier (Reinforcement Learning) is alsof je de robot laat vallen, laat vallen, en elke keer dat hij struikelt, je hem een flinke duw geeft (een negatieve beloning) en als hij goed staat, een snoepje (een positieve beloning). Dit kost enorm veel tijd en energie.

Een slimmere manier is Imitatie Learning: je laat de robot gewoon kijken hoe een expert (een echte danser) het doet. Maar hier zit een addertje onder het gras:

1. Vaak hebben we alleen video's van de danser (we zien waar hij staat, maar niet hoe hij zijn spieren spant).
2. We hebben vaak maar één of twee video's van die expert.
3. De robot moet leren van deze beperkte informatie.

De Oude Methode: "De Euclidische Valstrik"

Hoe meet een robot of hij goed doet? Hij vergelijkt zijn eigen bewegingen met die van de expert.
Stel je voor dat de expert op punt A staat en dan naar punt B springt. De robot staat op punt C.

• De oude methode (Euclidische afstand) meet de afstand alsof je een rechte lijn trekt met een liniaal door de lucht.
• Het probleem: In de echte wereld (en in een video) is de lucht vaak vol obstakels. Misschien is punt C wel heel dicht bij punt B in de lucht (korte lijn), maar als de robot daarheen gaat, botst hij tegen een muur of valt hij in een gat. De "rechte lijn" is dus een slechte maatstaf voor wat er echt gebeurt.

De oude methoden gebruiken deze simpele "liniaal-meting", waardoor de robot vaak de verkeerde weg op gaat, zelfs als hij dicht bij de expert lijkt te zijn.

De Nieuwe Oplossing: LWAIL (De "Intuïtieve Kaart")

De auteurs van dit paper (Siqi Yang en collega's) hebben een slimme truc bedacht. Ze zeggen: "Laten we niet met een liniaal meten, maar met een kaart die de struikelblokken al kent."

Dit is hoe hun methode, LWAIL, werkt, in drie stappen:

Stap 1: De "Droomtraining" (Pre-training)

Voordat de robot de dans van de expert gaat leren, laten we hem eerst een beetje rondhobbelen in het donker met willekeurige bewegingen.

• De Analogie: Stel je voor dat je een kind in een nieuw park laat rennen, zonder dat het doel is om te winnen, maar gewoon om de grond te voelen. Het kind leert: "Als ik hier spring, land ik op gras. Als ik daar spring, val ik in een modderpoel."
• De Techniek: Ze gebruiken een klein beetje willekeurige data (slechte kwaliteit) om een ICVF (Intention Conditioned Value Function) te trainen. Dit is een soort "intuïtie" die leert hoe de wereld werkt. Het leert een geheime kaart (latent space) waar de afstanden niet worden gemeten in meters, maar in "hoe moeilijk is het om hier naartoe te komen?".
  • In deze geheime kaart is punt C (de valkuil) plotseling heel ver weg van punt B (de expert), zelfs als ze er visueel dichtbij uitzien.

Stap 2: De "Vergelijkingspartij" (Imitation)

Nu de robot die geheime kaart heeft, gaan we de echte dans leren.

• We nemen de video van de expert.
• We kijken niet meer naar de ruwe coördinaten, maar naar de positie op de geheime kaart.
• We gebruiken een wiskundige methode (Wasserstein-afstand) om de robot te dwingen zijn bewegingen op die kaart zo dicht mogelijk bij die van de expert te houden.
• Omdat de kaart de "valkuilen" en "struikelblokken" al kent, weet de robot precies welke weg hij moet nemen om de expert te imiteren, zonder in de modder te vallen.

Stap 3: De "Scheidsrechter" (Adversarial Learning)

Tijdens het leren speelt er een tweede robot (een scheidsrechter) mee.

• Deze scheidsrechter probeert te raden: "Is dit een beweging van de expert, of van mijn vriend die net begint?"
• De dansende robot probeert de scheidsrechter te misleiden door zo goed mogelijk te dansen.
• Door dit spelletje te spelen, wordt de dansende robot steeds beter, zelfs als hij maar één video van de expert heeft gezien.

Waarom is dit zo speciaal?

1. Weinig data nodig: Normaal gesproken heb je honderden video's nodig. Met LWAIL volstaat één video van de expert.
2. Geen acties nodig: De robot hoeft niet te weten welke spieren de expert heeft aangespannen; hij hoeft alleen te kijken waar de expert is.
3. Robuust: Als de robot per ongeluk een beetje uit de toon raakt (bijvoorbeeld door ruis of een ongelukje), weet hij dankzij de "geheime kaart" hoe hij zichzelf weer terug naar de goede weg moet sturen.

Samenvattend

Stel je voor dat je een auto wilt leren rijden door alleen naar een video van een coureur te kijken, zonder dat je het stuur of de pedalen kunt zien.

• Oude methode: De auto probeert de coördinaten van de coureur na te bootsen. Omdat de weg vol gaten zit, rijdt de auto er recht in.
• LWAIL-methode: De auto krijgt eerst een "gevoel" voor de weg (de geheime kaart) door even wat rond te rijden. Daarna kijkt hij naar de coureur, maar niet op de kaart van de weg, maar op de kaart van de "veilige routes". Hierdoor rijdt hij moeiteloos en veilig, zelfs als hij maar één keer heeft gekeken.

Kortom: LWAIL leert robots niet alleen wat er gebeurt, maar vooral hoe de wereld in elkaar zit, zodat ze met heel weinig voorbeelden expert kunnen worden.

Technische samenvatting

Probleemstelling

Imitatie Learning (IL) stelt agenten in staat om expertgedrag na te bootsen door te leren van demonstraties. Traditionele methoden hebben echter vaak grote hoeveelheden demonstraties nodig, inclusief de acties van de expert, wat in de praktijk vaak onbeschikbaar of duur is.

• Observation-only (LfO): Er is een verschuiving naar "Imitation Learning from Observations" (LfO), waarbij alleen staten (zonder acties) van de expert worden gebruikt.
• Data-efficiëntie: Zelfs staten-only demonstraties zijn kostbaar. Het is cruciaal om methoden te ontwikkelen die leren van een minimaal aantal expert-trajectoires (bijv. slechts één).
• Beperkingen van bestaande methoden:
  • f-divergenties: Methoden die gebruikmaken van f-divergenties (zoals KL- of JS-divergentie) vereisen dat de verdelingen van de agent en de expert op dezelfde steunset (support set) liggen. Dit leidt tot theoretische beperkingen en numerieke fouten, vooral als de data van lage kwaliteit is.
  • Wasserstein-afstand: Hoewel de Wasserstein-afstand (Earth Mover's Distance) populair is geworden omdat deze geen steunset-overlap vereist, gebruiken bestaande methoden vaak de Kantorovich-Rubinstein (KR) dual. Een groot nadeel hiervan is dat de onderliggende metriek (de "ground metric") tussen staten vaak beperkt blijft tot de Euclidische afstand.
  • Dynamica-ignorantie: De Euclidische afstand in de ruwe statenruimte faalt om de dynamica van de omgeving te vangen. Twee staten kunnen numeriek dicht bij elkaar liggen in de Euclidische ruimte, maar dynamisch gezien volledig onbereikbaar voor elkaar zijn (zie Figuur 1 in het paper). Dit misleidt het leerproces.

Methodologie: LWAIL

De auteurs stellen Latent Wasserstein Adversarial Imitation Learning (LWAIL) voor, een tweestapsproces dat de beperkingen van de Euclidische metriek oplost door een dynamica-bewuste latent space te gebruiken.

1. Pre-training Stage (Dynamica-bewuste Embedding)

In plaats van de ruwe staten direct te gebruiken, leert LWAIL een betere representatie van de statenruimte:

• Data: Er wordt gebruikgemaakt van een zeer kleine hoeveelheid (ongeveer 1% van de online data) ongestructureerde, lage-kwaliteit data (bijv. willekeurige rollouts) die alleen statenparen $(s, s')$ bevatten.
• ICVF (Intention Conditioned Value Function): De auteurs trainen een ICVF-model op deze data. ICVF leert een waardefunctie $V(s, s^+, z)$ die de kans schat om een toekomstige staat $s^+$ te bereiken vanuit $s$ met een specifieke intentie (doel) $z$.
• Embedding: Uit deze waardefunctie wordt een statische embedding $\phi(s)$ afgeleid. Deze embedding codeert de dynamica-bewuste bereikbaarheid tussen staten. In deze latent space vertegenwoordigt de Euclidische afstand tussen $\phi(s)$ en $\phi(s')$ de werkelijke "kosten" om van $s$ naar $s'$ te gaan binnen de dynamica van de omgeving, in plaats van de fysieke afstand.

2. Imitation Stage (Adversarial Learning)

• Vastgevroren Embedding: De geleerde embedding $\phi$ wordt vastgevroren (frozen).
• Wasserstein Matching: De agent en de discriminator (discriminator $f$) opereren nu in de latent space. Het doel is om de 1-Wasserstein-afstand te minimaliseren tussen de bezettingsverdeling van de expert en die van de agent, maar dan gemeten in de latent space:
  $$\min_{\pi} \max_{\|f\|_L \le 1} \left( \mathbb{E}_{(s,s') \sim d^\pi_{ss}} [f(\phi(s), \phi(s'))] - \mathbb{E}_{(s,s') \sim d^E_{ss}} [f(\phi(s), \phi(s'))] \right)$$
• Beloning: De discriminator genereert een pseudo-beloning voor de agent (gebaseerd op TD3 als downstream RL-algoritme), waarbij de beloning wordt berekend als $\sigma(-f(\phi(s), \phi(s')))$.

Kernbijdragen

1. Dynamica-bewuste Ground Metric: Het paper toont aan dat de ICVF-geleerde latent space een effectieve "ground metric" biedt voor de Wasserstein-afstand, zelfs wanneer deze wordt getraind op een klein aantal willekeurige, staten-only data. Dit lost het fundamentele probleem op van de Euclidische metriek in eerdere KR-dualiteit-gebaseerde methoden.
2. Data-efficiëntie: De methode bereikt expert-niveau prestaties met slechts één staten-only expert-trajectoire.
3. Theoretische onderbouwing: De auteurs bewijzen (Theorema 3.1) dat in een bijna-deterministische MDP de bezettingsverdeling van statenparen een lineaire combinatie is van de ICVF-embedding, wat het gebruik van deze embedding voor optimalisatie rechtvaardigt.
4. Empirische Superioriteit: LWAIL overtreft bestaande baselines (zowel f-divergentie als eerdere Wasserstein-methoden) op diverse taken.

Resultaten

De methode is getest op meerdere omgevingen uit de D4RL-benchmark (MuJoCo en Maze2D):

• MuJoCo (Hopper, HalfCheetah, Walker2D, Ant): LWAIL behaalde consistent de beste resultaten, vaak dicht bij of zelfs beter dan de expert-scores, zelfs met slechts één expert-trajectoire. Het overtrof sterke concurrenten zoals WDAIL, IQlearn, en OPOLO.
• Robuustheid tegen Ruis: In experimenten met gestoorde startposities (Gaussian noise) in Maze2D, faalde de methode zonder ICVF-embedding catastrofale, terwijl LWAIL met ICVF stabiel bleef presteren. Dit bevestigt dat de embedding de dynamica van de omgeving beter begrijpt.
• Ablatie Studies:
  • Het gebruik van ICVF-embeddings presteerde significant beter dan andere contrastieve leermethoden (zoals CURL of PW-DICE).
  • De methode is robuust tegen verschillende niveaus van omgevingsruis.
  • Het werkt goed met zeer kleine offline datasets (slechts 10k willekeurige transities) voor het pre-trainen van de ICVF.

Betekenis en Impact

LWAIL biedt een belangrijke doorbraak in het veld van Imitation Learning door twee kritieke beperkingen te overwinnen:

1. Afhankelijkheid van expert-acties: Het vereist geen expert-acties, alleen staten.
2. Data-honger: Het kan leren van extreem weinig expert-data (slechts één trajectoire).

Door de introductie van een dynamica-bewuste metriek via ICVF, maakt LWAIL het mogelijk om de geometrische beperkingen van de Euclidische afstand in de KR-dualiteit te omzeilen. Dit resulteert in een efficiëntere, stabielere en nauwkeurigere imitatie, wat vooral waardevol is voor toepassingen in de robotica waar expert-demonstraties schaars of moeilijk te verkrijgen zijn. De methode democratiseert de toegang tot Reinforcement Learning door de afhankelijkheid van kostbare, hoogwaardige annotaties te verminderen.