LAR-MoE: Latent-Aligned Routing for Mixture of Experts in Robotic Imitation Learning

Dit paper introduceert LAR-MoE, een tweestapskader voor imitatielearning dat ongelabelde demonstraties gebruikt om een latente ruimte te leren die de expert-routing reguleert, waardoor robots effectief kunnen leren van heterogene taken zonder handmatige vaardigheidsannotaties.

De kern

Stel je voor dat je een robot wilt leren om een complexe taak uit te voeren, zoals het vastpakken en terugtrekken van darmweefsel tijdens een operatie. Dit is niet zoiets als "gooi de bal naar de doos". Het is een reeks van verschillende stappen: eerst zoeken, dan vastpakken, dan wachten, dan voorzichtig trekken, en tenslotte de spanning vasthouden.

Als je een robot gewoon laat kijken naar duizenden video's van chirurgen, neigt de robot vaak tot het "gemiddelde" nemen. Hij probeert alles tegelijk te doen, wat resulteert in een rommelige, onzeker beweging. Het is alsof je een kok vraagt om een soep te maken, maar hij probeert tegelijkertijd te bakken, te koken en te braden, waardoor het eindresultaat een smakeloze soep wordt.

Hier komt LAR-MoE om de hoek kijken. Het is een slimme manier om robots te leren, en het werkt als een super-georganiseerd team van specialisten.

De Grote Idee: Een Team van Specialisten

In plaats van één grote, algemene "brein" voor de robot te bouwen, maken de onderzoekers een Mixture of Experts (MoE). Denk hierbij niet aan één super-intelligente persoon, maar aan een orkest.

• In een orkest heb je violisten, drummers en fluitisten.
• Als er een rustig stukje muziek is, spelen de violisten.
• Als het tempo stijgt, komen de drums erbij.
• Ze werken samen, maar elk speelt zijn eigen rol op het juiste moment.

Bij robots betekent dit: in plaats van één algoritme dat alles probeert, hebben we een team van kleine, gespecialiseerde netwerken (de "experts"). De ene is goed in "vastpakken", de andere in "wachten" en weer een andere in "trekken".

Het Probleem: Wie doet wat?

Het grote probleem bij dit soort teams is: wie bepaalt wie er mag spelen?
In de oude methoden moesten mensen handmatig zeggen: "Op seconde 10 moet de 'trek-expert' aan de slag." Dit is echter lastig, duur en in de echte wereld (zoals in een ziekenhuis) vaak onmogelijk om perfect te labelen.

De Oplossing: LAR-MoE (De "Geheime Code")

LAR-MoE lost dit op met een slimme, twee-stappen methode die geen menselijke labels nodig heeft. Het is alsof de robot zelf de muziekleer ontdekt zonder dat iemand hem de noten heeft gegeven.

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

Stel je voor dat de robot eerst een tijdmachine bouwt.

• De robot kijkt naar een foto van de huidige situatie (bijv. een schaar die naar een darm wijst).
• Vervolgens probeert hij te raden: "Wat gaat er straks gebeuren?"
• Hij leert een geheime code (een 'latent space') die de link legt tussen wat hij nu ziet en wat er later gebeurt.
• Dit is als een student die naar een meester kijkt en probeert te voorspellen wat de meester als volgende gaat doen, puur op basis van de beweging. De robot leert zo de "structuur" van de taak zonder dat iemand zegt: "Dit is nu fase 1".

Stap 2: De "Regisseur" (Routing)

Nu komt het slimme deel. De robot heeft nu die geheime code. Hij gebruikt deze code als een regisseur voor zijn team van experts.

• Als de code zegt: "We zijn in een situatie die lijkt op 'vastpakken'", stuurt de regisseur het signaal naar de 'vastpak-expert'.
• Als de code zegt: "We moeten nu wachten", gaat het signaal naar de 'wacht-expert'.

De onderzoekers hebben een slimme regel toegevoegd: de regisseur mag niet willekeurig kiezen. Hij moet zijn keuzes baseren op de geheime code die in stap 1 is geleerd. Dit zorgt ervoor dat de experts niet gaan "slapen" (niet gebruikt worden) of allemaal hetzelfde gaan doen (wat "expert collapse" wordt genoemd). Ze blijven gespecialiseerd.

Waarom is dit zo cool?

1. Geen dure labels nodig: Je hoeft niet urenlang video's te bekijken en te zeggen "nu is het moment om te trekken". De robot leert dit zelf uit de beelden.
2. Efficiëntie: Het systeem is klein (slechts 150 miljoen parameters), maar presteert beter dan gigantische modellen met miljarden parameters. Het is als een klein, slank team dat harder werkt dan een overbevolkte, trage fabriek.
3. Werkt in de echte wereld: De robot werd getest op een kunstmatige darm (een "phantom") en daarna zonder enige extra training op echte, dode varkensdarmen. Het team van experts wist zich aan te passen aan de nieuwe, echte situatie.

De Conclusie

LAR-MoE is als het geven van een robot een intuïtie voor hoe een taak in elkaar zit. In plaats van hem een stap-voor-stap instructieboekje te geven, laat je hem de "flow" van de beweging ontdekken. Vervolgens organiseer je zijn interne team van specialisten zo dat ze precies op het juiste moment ingrijpen.

Het resultaat? Een robot die soepel, veilig en slim complexe chirurgische taken uitvoert, alsof hij jarenlang ervaring heeft, terwijl hij eigenlijk gewoon heeft geleerd om naar de toekomst te kijken en de juiste specialist te bellen.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het artikel "LAR-MoE: Latent-Aligned Routing for Mixture of Experts in Robotic Imitation Learning" in het Nederlands.

Probleemstelling

Imitatielearning (IL) stelt robots in staat om vaardigheden te leren uit demonstraties. Echter, het implementeren van één beleid (policy) voor taken met heterogene dynamieken (zoals in chirurgische taken) blijft een uitdaging. Bestaande methoden neigen naar het "middelen" van verschillende gedragsmodi in plaats van zich te specialiseren, wat leidt tot suboptimale prestaties bij complexe taken.

Hoewel Mixture-of-Experts (MoE)-architecturen een oplossing bieden door gespecialiseerde subnetwerken te activeren, vereisen deze doorgaans expliciete, handmatig gedefinieerde taakfasen of supervisie voor het routeren van experts. In domeinen zoals robotchirurgie zijn dergelijke gelabelde datasets schaars en duur. Bestaande MoE-methoden kampen bovendien vaak met expert collapse (waarbij één expert alle taken overneemt) en inefficiënt parametergebruik.

Methodologie: LAR-MoE

De auteurs introduceren LAR-MoE (Latent-Aligned Routing for Mixture of Experts), een tweestapskader dat onbeheerde vaardigheidsontdekking ontkoppelt van beleidslearning.

1. Pre-training: Onbeheerde Latente Ruimte Leren

• Student-Teacher Co-training: Er wordt een gezamenlijke latente representatie geleerd tussen waarnemingen (beelden) en toekomstige actie-chunks.
  • De Teacher (gevoed met beelden én acties) probeert de actie-chunk te reconstrueren.
  • De Student (gevoed met alleen beelden) probeert de latente vector van de teacher te voorspellen.
• Doel: Deze onbeheerde strategie leert een latente ruimte die de onderliggende taakstructuur en gedragsmodi vastlegt zonder expliciete fase-annotaties.

2. Post-training: Latent-Aligned Routing

• Architectuur: Het beleid bestaat uit een visuele en taalkundige encoder, gevolgd door $N$ compacte actie-experts (geïmplementeerd als Transformer-decoders).
• Routing Mechanisme: De ingevroren student-model voorspelt de latente vector $\hat{z}_t$. Een soft-gating mechanisme (softmax) bepaalt de activatie van de experts.
• Regularisatie (Kerninnovatie): Om expert collapse te voorkomen en specialisatie te stimuleren, wordt de routing gedwongen om de structuur van de geleerde latente ruimte te volgen. Dit wordt bereikt via een Distance Consistency Loss ($L_{DC}$):
  • De afstand tussen de voorspelde latente vectoren (van de student) wordt vergeleken met de afstand tussen de expert-selectie-verdelingen.
  • Als twee waarnemingen dicht bij elkaar liggen in de latente ruimte (soortgelijke taken), moeten ze ook door dezelfde experts worden behandeld.
• Aanvullende Regularisatie: Er wordt gebruikgemaakt van entropie-regularisatie (voor specialisatie) en groep-sparse regularisatie (voor stabiliteit).

Belangrijkste Bijdragen

1. Onbeheerde Co-training Strategie: Een methode om een beschrijvende latente ruimte te leren die de relatie tussen visuele waarnemingen en toekomstige bewegingstrajecties vastlegt zonder supervisie.
2. LAR-MoE Architectuur: Een nieuw framework met een "latent-alignment" regularisatiestrategie. Dit verankert de zachte expert-routing aan de structuur van de latente ruimte, wat expert collapse voorkomt en de parameter-efficiëntie aanzienlijk verhoogt.
3. Validatie zonder Annotaties: Het bewijs dat routingstructuren kunnen worden geleerd puur uit de uitlijning van waarnemingen en toekomstige bewegingen, gevalideerd op zowel simulatie- als hardware-chirurgische taken.

Resultaten

1. Simulatie Benchmark (LIBERO)

• Op de LIBERO-benchmark bereikt LAR-MoE een gemiddelde succesratio van 95,2%.
• Dit wordt bereikt met slechts 150M parameters, wat aanzienlijk minder is dan grote VLA-modellen (bijv. $\pi_0.5$ met ~3,5B parameters).
• Het model presteert vergelijkbaar met of beter dan veel grotere modellen, wat de efficiëntie van de MoE-architectuur onderstreept.
• Ablatiestudies tonen aan dat zowel het invriezen van de student-encoder als de latent-alignment regularisatie essentieel zijn voor de prestaties.

2. Hardware Experimenten (Chirurgische Darmgreep en Retractie)

• Taak: Een complexe chirurgische taak met vijf fasen (aangeven, grijpen, wachten, rekken, spanning houden).
• Data: Getraind op slechts 120 demonstraties zonder enige fase-annotatie.
• Prestatie: LAR-MoE bereikte een succesratio die vergelijkbaar is met een gesuperviseerde MoE-baseline (die wel annotaties gebruikte), maar zonder de kosten van labeling.
• Zero-shot Transfer: Het model werd getest op ex vivo varkensdarmweefsel zonder extra training en behaalde een succesratio van 45% (9/20 rollouts), wat aantoont dat het model generaliseert naar realistisch weefsel met andere visuele en mechanische eigenschappen.

3. Interpretatie van Expert Activatie

• De analyse toont aan dat de experts zich automatisch specialiseren in specifieke taakfasen (bijv. benadering, grijpen, vasthouden).
• De activatiepatronen van de experts corresponderen sterk met handmatig geannoteerde taakfasen door chirurgen, ondanks dat het model nooit deze labels heeft gezien.
• Er is sprake van ruimtelijke consistentie: experts specialiseren zich in specifieke gebieden van de taakruimte, wat bijdraagt aan de robuustheid.

Betekenis en Impact

LAR-MoE biedt een principieel alternatief voor gesuperviseerde vaardigheidsdecompositie in robotica. Door de routing te baseren op een onbeheerd geleerde latente ruimte, elimineert het de noodzaak voor dure en tijdrovende handmatige annotaties van taakfasen.

Dit is vooral waardevol in domeinen zoals chirurgische robotica, waar data schaars is en annotaties complex zijn. De methode maakt het mogelijk om gestructureerde, gespecialiseerde expert-systemen te bouwen die efficiënter zijn (minder parameters) en beter generaliseren naar nieuwe omgevingen (zoals ex vivo weefsel) dan traditionele end-to-end benaderingen. De resultaten suggereren dat impliciete taakfasen kunnen worden geleerd uit visuele en kinematische data, wat de weg vrijmaakt voor robuustere en schaalbare robotleersystemen.