TransMASK: Masked State Representation through Learned Transformation

Dit paper introduceert TransMASK, een zelftoezichtmethode die een leerbare masker transformeert om irrelevante omgevingsvariabelen te filteren en zo robuustere robotbeleid te trainen voor betere generalisatie in nieuwe omgevingen.

De kern

TransMASK: De Kunst van het "Niet Kijken" voor Robots

Stel je voor dat je een robot wilt leren om een groene blok van de tafel te pakken en in het midden te leggen. Jij, als mens, kijkt alleen naar de blok en waar je hem neer wilt zetten. Je negeert automatisch de kleur van de tafel, de stofjes in de lucht, of de rommel op de achtergrond. Voor jou zijn dat details die er niet toe doen.

Maar een robot is niet zo slim als jij. Als je een robot leert met video's, ziet hij alles. Hij ziet de houtnerf van de tafel, de schaduw van een vaas, en de kleur van de muur. Als je de robot nu op een marmeren tafel zet in plaats van een houten tafel, raakt hij in paniek. Hij denkt: "Oh nee! De wereld is veranderd! Ik weet niet meer wat ik moet doen!" Hij faalt, omdat hij te veel naar de verkeerde dingen heeft gekeken.

Dit is het probleem dat de onderzoekers van TransMASK willen oplossen.

Het Probleem: De Robot die te veel luistert

Stel je voor dat de robot een student is die een proef maakt. De leraar (de mens) geeft aanwijzingen: "Pak de blok." Maar de student (de robot) luistert ook naar het geluid van de klok, de geur van de koffie en de kleur van de gordijnen. Als de leraar later in een andere klas lesgeeft met andere gordijnen, raakt de student in de war omdat hij dacht dat de gordijnen belangrijk waren voor het antwoord.

De robot leert dus een "kwetsbaar" gedrag. Hij maakt verbindingen met dingen die eigenlijk niets met de taak te maken hebben.

De Oplossing: TransMASK (De Slimme Filter)

De onderzoekers hebben een nieuwe methode bedacht die ze TransMASK noemen. Je kunt dit zien als een slimme bril of een versterker die de robot opzet.

In plaats van dat we de robot vertellen wat hij moet negeren (wat heel moeilijk is), laten we de robot leren wat hij moet negeren, terwijl hij de taak uitvoert.

Hoe werkt dat?

1. De Leren Bril: De robot krijgt een "masker" (een soort digitale bril) dat hij zelf moet leren instellen.
2. De Oefening: De robot probeert de taak te doen (de blok pakken). Hij kijkt door zijn bril naar de wereld.
3. De Feedback: Als de robot de blok mist, krijgt hij een signaal: "Je hebt verkeerd gekeken!"
4. Het Magische Moment: De robot kijkt terug naar zijn bril. Hij ziet dat hij te veel naar de achtergrond heeft gekeken en te weinig naar de blok. Dus, hij draait aan de knoppen van zijn bril. Hij maakt de lens voor de achtergrond dicht (zoals een sluiter die dichtgaat) en de lens voor de blok helder.

Na veel oefening heeft de robot een bril die automatisch alles wat niet belangrijk is, dichttint. Hij ziet alleen nog de blok en zijn eigen hand. De rest is wazig of onzichtbaar.

Waarom is dit zo slim?

Meerdere andere methoden proberen dit ook, maar vaak op een moeilijke manier:

• Sommige proberen de robot te trainen met duizenden verschillende achtergronden (zoals een student die 1000 verschillende examens moet maken). Dat kost veel tijd en werkt niet altijd goed.
• Andere methoden proberen de robot te dwingen om alleen de belangrijkste dingen te onthouden, maar dan vergeet de robot soms hoe hij de taak moet doen.

TransMASK is anders. Het is als een natuurlijk leerproces.
De robot hoeft niet te weten wat belangrijk is. Hij hoeft alleen maar te proberen de taak goed te doen. De wiskunde achter de schermen zorgt ervoor dat de robot vanzelf merkt: "Hé, als ik naar die achtergrondklok kijk, helpt dat me niet om de blok te pakken. Als ik naar de blok kijk, wel."

De robot leert dus een statisch masker: een vaste regel die zegt "Kijk altijd naar de blok, nooit naar de tafelkleur". Dit werkt zelfs als de robot later in een heel andere kamer komt. Omdat hij de tafelkleur al heeft uitgeschakeld, maakt het niet uit of de tafel nu van hout, marmer of plastic is. Hij ziet alleen nog de blok.

De Resultaten

In hun experimenten hebben ze robots getest in een virtuele wereld en in de echte wereld.

• Zonder TransMASK: De robot faalde vaak als de achtergrond veranderde (bijvoorbeeld van hout naar marmer).
• Met TransMASK: De robot was veel sterker. Hij faalde veel minder vaak, zelfs als de omgeving totaal anders was. Hij kon de taak doen alsof hij er al jaren ervaring mee had, omdat hij zich alleen richtte op de essentie.

Samenvattend

TransMASK is een manier om robots te leren om niet te kijken naar de rommel, maar alleen naar wat er echt toe doet. Het is alsof je een robot een "focus-regelaar" geeft die automatisch de ruis uitschakelt. Hierdoor worden robots veel robuuster en kunnen ze hun vaardigheden beter toepassen in nieuwe, onbekende situaties.

Het is de kunst van het robotleren: niet alles zien, maar wel het juiste zien.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "TransMASK: Masked State Representation through Learned Transformation" in het Nederlands.

Probleemstelling

In het domein van imitatielearning (nabootslearning) leren robots taken door te kijken naar demonstraties van een menselijke expert. Een fundamenteel probleem is dat robots vaak overgevoelig worden voor irrelevante kenmerken in hun waarnemingen (zoals de kleur van de tafel, achtergrondobjecten of belichting), terwijl de menselijke expert zich uitsluitend richt op de voor de taak relevante elementen (zoals de positie van het object en de robotgrijper).

Wanneer een beleid (policy) dat is getraind op deze "volle" toestand wordt ingezet in een nieuwe omgeving met andere achtergronden of texturen, faalt het vaak omdat het spurious correlaties (toevallige samenhangen) heeft geleerd. Bestaande methoden om dit op te lossen, zoals data-augmentatie of het gebruik van Vision-Language Models (VLMs), hebben beperkingen: ze kunnen de prestaties binnen de oorspronkelijke distributie verminderen, vereisen vaak extra labels, of leiden tot catastrofale vergetelheid van eerder geleerde features. Andere benaderingen, zoals het Information Bottleneck (IB) principe, lijden vaak onder slecht gestelde optimalisatieproblemen waarbij de latente representatie instort tot een actie-representatie in plaats van een zinvolle staat-representatie.

Methodologie: TransMASK

De auteurs stellen TransMASK voor, een zelftoezicht (self-supervised) methode die een masker leert om irrelevante staten uit de waarneming te filteren, zonder de bestaande imitatielearning-structuur of de verliesfunctie (loss function) te hoeven aanpassen.

Kernconcepten:

1. Gedissocieerde Toestand: De methode gaat ervan uit dat de toestand $s$ kan worden opgesplitst in relevante elementen ($\mu$, zoals objectposities) en irrelevante ruis ($\eta$, zoals achtergrondkleur).
2. Jacobian als Proxy: De auteurs redeneren dat de gradiënten van de imitatielearning-optimalisatie informatie bevatten over causaliteit. Omdat de expert alleen reageert op relevante elementen, zullen de gradiënten voor deze elementen groot zijn, terwijl ze voor irrelevante elementen (die de expert negeert) dicht bij nul liggen. De grootte van de Jacobiaan van het expert-beleid fungeert dus als een proxy voor causale relevantie.
3. Het Masker (M): TransMASK introduceert een leerbare transformatiematrix $M$ (een masker) die de waargenomen toestand $s$ transformeert naar een latente representatie $z = Ms$.
   • $M$ is een $n \times n$ matrix.
   • Tijdens het trainen worden de gewichten van $M$ bijgewerkt via de gradiënten van de imitatielearning-loss (bijv. MSE tussen voorspelde en expert-acties).
   • De kolommen van $M$ die corresponderen met irrelevante staten convergeren naar nul, terwijl kolommen voor relevante staten een hoge magnitude behouden.
4. Normalisatie: Om te voorkomen dat de waarden van het masker onbeperkt groeien, wordt een normalisatielaag (zoals softmax of sparsemax) per rij toegepast. Dit dwingt het masker om een harde selectie te maken (waarden dicht bij 0 of 1).
5. Integratie: TransMASK is modulair en kan worden toegevoegd aan bestaande frameworks (zoals Diffusion Policies of MLP's) zonder extra labels of wijzigingen in de trainingspipeline. Het masker is statisch en onafhankelijk van de input, in tegenstelling tot dynamische attention-mechanismen.

Belangrijkste Bijdragen

1. Analyse van Bestaande Methodes: De auteurs tonen aan dat bestaande methoden voor staterepresentatie (zoals VAEs en IB) vaak vastlopen in lokale minima, instabiel zijn of leiden tot representatie-instorting (collapse) waarbij de latente ruimte geen zinvolle toestand meer bevat.
2. TransMASK Framework: Ze ontwikkelen een nieuwe methode die gebruikmaakt van de inherente gradiënten van de imitatielearning-taak om een masker te leren dat relevante features isoleert. Dit vereist geen extra supervisie.
3. Experimentele Validatie: Ze vergelijken TransMASK met state-of-the-art baselines (BC, VAE, CLASS, VINN) in zowel gesimuleerde als real-world omgevingen.
4. Robuustheid: Ze demonstreren dat TransMASK leidt tot beleidsregels die aanzienlijk robuuster zijn tegen distributieveranderingen (distribution shifts) in de omgeving.

Resultaten

De auteurs testen de methode in twee scenario's: gesimuleerde omgevingen (Panda-Gym) en een real-world UR10 robot.

• Gesimuleerde Experimenten:

  • Taken: Picken, Duwen en Roteren van een Rubik's kubus in een omgeving met afleidende objecten (distractors) en variërende achtergronden (houten vs. marmeren tafel).
  • Privileged State: Wanneer de robot directe toegang heeft tot objectposities, overtreft TransMASK de baselines met tot 15% hogere succespercentages in de trainingsdistributie (ID) en 9% in de buiten-distributie (OOD).
  • Beeldwaarneming: Zelfs met hoge-dimensionale beeldinvoer (via ResNet-encoders) en segmentatiemaskers, behaalt TransMASK betere resultaten dan VAE, CLASS en VINN, vooral in OOD-scenario's.
  • Masks Visualisatie: De geleerde maskers tonen duidelijk aan dat de methoden de gewichten voor relevante objecten (doelobject, robot, doellocatie) hoog houden, terwijl gewichten voor achtergrond en distractors naar nul worden gedrukt.

• Real-World Experimenten:

  • Taken: Picken, Stapelen van LEGO en Scheppen uit een kom.
  • De robot werd getraind op een houten tafel en getest op een tafel bedekt met een witte laken (verandering in achtergrond, schaduwen en reflecties).
  • TransMASK presteerde consistent beter dan alle baselines, inclusief in de OOD-situatie. Hoewel CLASS en VINN soms trainden op een mix van ID- en OOD-data (wat hen een voordeel gaf), presteerde TransMASK, dat alleen op ID-data was getraind, toch superieur of gelijkwaardig, wat de effectiviteit van het masker benadrukt.

Betekenis en Conclusie

TransMASK biedt een elegante oplossing voor het probleem van "causal confusion" in imitatielearning. Door gebruik te maken van de gradiënten van de primaire taak om een statisch masker te leren, kunnen robots beleidsregels ontwikkelen die zich uitsluitend richten op de structuur van de taak en niet op toevallige eigenschappen van de omgeving.

De belangrijkste implicaties zijn:

• Generalisatie: Robots kunnen taken succesvol uitvoeren in nieuwe omgevingen zonder extra trainingsdata of handmatige labeling.
• Efficiëntie: De methode vereist geen complexe architecturale wijzigingen of extra verliesfuncties, wat het makkelijk toepasbaar maakt op bestaande systemen (zoals Diffusion Policies).
• Toekomstige Richting: Hoewel de methode afhankelijk is van een zekere mate van ontkoppeling (disentanglement) van de input, biedt het een sterke basis voor robuuste robotica. Toekomstig werk richt zich op het uitbreiden naar Reinforcement Learning en het bieden van theoretische garanties voor de convergentie van het masker.

Kortom, TransMASK stelt robots in staat om te leren wat echt belangrijk is voor een taak, waardoor ze minder kwetsbaar zijn voor veranderingen in de wereld om hen heen.