Grounding Bodily Awareness in Visual Representations for Efficient Policy Learning

Dit artikel introduceert ICon, een contrastieve leermethode voor Vision Transformers die agent-specifieke en omgevings-specifieke tokens scheidt om robuustere visuele representaties te creëren die de beleidsleer en overdraagbaarheid van robotmanipulatie taken verbeteren.

De kern

Stel je voor dat je een robot wilt leren om een kopje thee te pakken. Je geeft de robot een camera en laat hem kijken. Maar hier zit een probleem: de robot ziet niet alleen het kopje, maar ook zijn eigen arm, de tafel, de achtergrondmuur en misschien zelfs een vlieg die voorbijvliegt.

Voor een mens is het heel makkelijk om te weten: "Dat is mijn arm, dat is het kopje." Maar voor een computer is dat lastig. De computer ziet alleen een wirwar van pixels. Als de robot te veel aandacht besteedt aan de muur of de vlieg, vergeet hij misschien hoe zijn eigen arm beweegt. Dit maakt het leren van nieuwe taken traag en onstabiel.

De auteurs van dit paper, Junlin Wang en Zhiyun Lin, hebben een slimme oplossing bedacht die ze ICon noemen. Laten we uitleggen hoe dit werkt met een paar creatieve vergelijkingen.

1. Het Probleem: De "Verwarde Kunstgalerij"

Stel je voor dat de robotkijken als een kunstgalerij is waar alle schilderijen (de beelden) door elkaar liggen. De robot moet een opdracht uitvoeren, maar hij kan niet goed onderscheiden wat "hemzelf" is en wat "de wereld" is. Hij probeert te leren, maar hij raakt in de war over wie de schilder is en wie het schilderij.

In de oude methoden probeerden ze dit op te lossen door de robot te dwingen om het beeld opnieuw te tekenen (reconstructie). Dit is alsof je een kunststudent dwingt om elke keer dat hij naar een schilderij kijkt, een perfecte kopie te maken. Dat kost veel tijd en energie, en soms vergeten ze de echte opdracht (de robotarm bewegen) omdat ze te druk zijn met het kopiëren.

2. De Oplossing: ICon (De "Scheidingslijn")

De auteurs hebben een nieuwe methode bedacht genaamd Inter-token Contrast (ICon). In plaats van de robot te dwingen om alles na te tekenen, geven ze hem een heel simpel spelletje: "Vind de gelijkenissen en maak onderscheid."

Hier is hoe het werkt, stap voor stap:

• De Puzzelstukjes (Tokens): De robot kijkt niet naar één groot beeld, maar breekt het beeld op in duizenden kleine puzzelstukjes (dit noemen ze 'tokens' in de computerwereld).
• Het Spel: De robot krijgt de opdracht: "Kijk naar al je puzzelstukjes. Als een stukje van mijn arm is, moet het lijken op de andere stukjes van mijn arm. Als een stukje van de muur is, moet het lijken op de andere stukjes van de muur. Maar: stukjes van mijn arm mogen er nooit op lijken als stukjes van de muur!"
• Het Resultaat: Door dit spel te spelen, leert de robot vanzelf een mentale scheidslijn te trekken. Hij leert een "lichaamsbewustzijn" (proprioceptie) te ontwikkelen zonder dat hij zijn eigen lichaam hoeft te tekenen. Hij leert gewoon: "Dit is mijn arm, en dat is de rest van de wereld."

3. De Slimme Trucs (De "Versterkers")

Om dit spelletje nog beter te laten werken, gebruiken ze twee slimme trucjes:

• De "Verstrooiings-Truc" (Farthest Point Sampling):
  Stel je voor dat je een groep mensen moet selecteren om een foto te maken. Als je willekeurig kiest, kun je per ongeluk 10 mensen uit dezelfde hoek van de kamer kiezen. Dat geeft een onvolledig beeld.
  De auteurs gebruiken een methode die ze Farthest Point Sampling noemen. Dit is alsof je zegt: "Kies iemand, en kies dan de volgende persoon die zo ver mogelijk van de eerste staat, en de volgende die weer ver weg staat."
  Hierdoor kiezen ze voorbeelden van de arm die over de hele arm verspreid zijn (bijv. de schouder, de elleboog en de hand), in plaats van alleen de hand. Dit zorgt voor een veel vollediger en sterker begrip van hoe de robot eruitziet.

• De "Meerdere Lagen" (Multi-Level Contrast):
  Een robotbrein (een 'Vision Transformer') heeft verschillende lagen, net als een schoolkinderen die van klein naar groot gaan. De onderste lagen zien alleen lijntjes en kleuren, de bovenste lagen zien vormen en objecten.
  Normaal gesproken kijken ze pas naar het einde van de les. Maar ICon zegt: "Kijk naar elke les, van de kleuterschool tot de middelbare school." Door op alle niveaus te oefenen met het scheiden van "ik" en "wereld", wordt de robot veel slimmer en sneller.

4. Wat levert dit op?

De resultaten zijn indrukwekkend:

• Sneller leren: Robots met ICon leren nieuwe taken (zoals een lade openen of een doos sluiten) veel sneller dan robots zonder deze methode.
• Beter overdragen: Als je een robot hebt getraind om een taak te doen met een Franka-arm, en je wilt dat hij het doet met een Kinova-arm (een ander model), werkt het veel beter met ICon. Het is alsof de robot het concept van "mijn arm bewegen" heeft geleerd, in plaats van alleen de specifieke vorm van die ene arm.
• Stabiel leren: De training is rustiger. De robot raakt niet in paniek of blijft hangen in de training, wat vaak gebeurt bij de oude methoden.

Samenvatting

Kortom, dit paper zegt: "Laten we robots niet dwingen om de wereld na te tekenen. Laten we ze in plaats daarvan een spelletje laten spelen waarbij ze leren hun eigen lichaam te onderscheiden van de rest van de wereld."

Door dit te doen, krijgen robots een soort "innerlijk gevoel" voor hun eigen lichaam, zelfs als ze alleen maar naar een camera kijken. Dit maakt ze slimmer, sneller en beter in staat om taken uit te voeren in een chaotische wereld.

Technische samenvatting

Probleemstelling

Het leren van effectieve visuele representaties voor robotmanipulatie blijft een fundamentele uitdaging, voornamelijk vanwege de complexe lichaamsdynamiek die betrokken is bij het uitvoeren van acties. Hoewel visuele proprioceptie (het vermogen om het eigen lichaam via visie waar te nemen) essentieel is voor flexibele actieplanning, is het extraheren van lichaamsbewuste informatie uit hoge-dimensionale beelden moeilijk, vooral in end-to-end leerframeworks.

In bestaande benaderingen worden visuele encoders en beleidsnetwerken (policy networks) gezamenlijk geoptimaliseerd met één doel. Dit leidt vaak tot convergentie naar bottlenecks waarbij task-irrelevante signalen, waaronder visuele cues over het lichaam van de agent, onbedoeld worden gefilterd. Bestaande methoden die proberen dit op te lossen door reconstructieverliezen (bijv. het reconstrueren van RGB-beelden of agent-maskers) toe te voegen, kunnen de trainingsstabiliteit ondermijnen. De kernvraag is: hoe kunnen we ontkoppelde representaties van agent en omgeving afleiden zonder de prestaties of stabiliteit te schaden?

Methodologie: Inter-token Contrast (ICon)

De auteurs stellen Inter-token Contrast (ICon) voor, een contrastieve leermethode die specifiek is ontworpen voor de token-niveau representaties van Vision Transformers (ViTs). Het doel is om agent-specifieke en omgevings-specifieke features in de feature space te scheiden, waardoor een agent-gerichte visuele representatie ontstaat met ingebouwde inductieve biases voor het lichaam.

De kerncomponenten van de methode zijn:

1. Token-niveau Agent Maskers:

   • Een ViT verdeelt een beeld in patches (tokens).
   • Een segmentatiemodel genereert een binaire masker van de robot in het beeld.
   • Dit pixel-niveau masker wordt omgezet naar een token-niveau masker ($M_{token}$). Een token wordt als "agent-dominant" gemarkeerd als het percentage agent-pixels in de bijbehorende patch een drempelwaarde $\beta$ overschrijdt.

2. Inter-token Contrastief Verlies:

   • De features worden gescheiden in agent-specifieke en omgevings-specifieke sets.
   • Er worden queries ($q_a$ voor agent, $q_e$ voor omgeving) berekend door het middelen van de corresponderende features.
   • Voor de keys wordt Farthest Point Sampling (FPS) toegepast (oorspronkelijk uit point cloud processing, hier aangepast voor 2D). In plaats van willekeurig te sample, selecteert FPS punten die ruimtelijk zo ver mogelijk van elkaar verwijderd zijn. Dit zorgt voor een diverse en representatieve selectie van features die de structuur van zowel de agent als de omgeving goed vastleggen.
   • Er wordt een symmetrisch InfoNCE-verlies berekend: agent-features worden dicht bij elkaar getrokken en ver weg geduwd van omgevings-features (en vice versa).

3. Multi-Level Contrast (MLC):

   • In plaats van het contrast alleen op de laatste laag van de ViT toe te passen, wordt het verlies berekend over meerdere lagen van de encoder.
   • Een gewogen som van de verliezen per laag wordt gebruikt, waarbij diepere lagen (die semantische features bevatten) zwaarder wegen dan ondiepere lagen (die voornamelijk positionele informatie bevatten).

4. Integratie in Beleidsleren:

   • ICon fungeert als een hulddoel (auxiliary objective) naast de hoofdverliesfunctie (bijv. Diffusion Loss).
   • De totale trainingsdoelstelling is: $L = L_{diffusion} + \lambda L_{ICon}$.
   • Dit wordt getest in combinatie met Diffusion Policy, een state-of-the-art imitatieleeralgoritme.

Belangrijkste Bijdragen

• ICon Framework: Een nieuwe contrastieve leermethode die token-niveau features van ViTs gebruikt om agent- en omgevingsinformatie expliciet te ontkoppelen.
• Farthest Point Sampling (FPS) in 2D: De innovatieve toepassing van FPS voor het selecteren van keys in het beelddomein, wat zorgt voor een betere ruimtelijke dekking en representativiteit dan willekeurige sampling.
• Multi-Level Ontkoppeling: Een ontwerp dat contrastief leren toepast over meerdere lagen van de transformer, wat leidt tot een completere ontkoppeling van agent en omgeving.
• Verbeterde Trainingsstabiliteit: In tegenstelling tot reconstructie-gebaseerde methoden, behoudt ICon de trainingsstabiliteit van het end-to-end leerproces.

Resultaten

De auteurs hebben ICon geëvalueerd op 8 manipulatie-taken over 3 verschillende robots (Franka, Kinova, KUKA) in twee simulatiebenchmarks: RLBench en Robosuite.

• Prestatieverbetering: ICon geïntegreerd met Diffusion Policy (ICon-Diff-C en ICon-Diff-T) overtrof consistent de baselines (standaard Diffusion Policy en Crossway Diffusion) op bijna alle taken.
  • Bijvoorbeeld: In de "Open Box" taak was er een verbetering van 21,3% ten opzichte van de standaard CNN-basis.
  • In complexe taken zoals "Close Microwave" en "Put Rubbish in Bin" toonde ICon-Diff-T aanzienlijk betere resultaten dan de baselines.
• Transferability (Overdraagbaarheid): ICon faciliteerde few-shot transfer van beleid tussen robots met verschillende morfologieën (bijv. van Franka naar Kinova of IIWA). De prestaties van getransfereerde beleidslijnen waren hoger bij gebruik van ICon dan bij de baselines.
• Trainingsstabiliteit: Experimenten toonden aan dat ICon een hogere gemiddelde succesratio behoudt tijdens het trainen vergeleken met methoden die reconstructieverlies gebruiken. Dit suggereert dat ICon robuustere gedragingen leert zonder de training instabiel te maken.
• Ablatiestudies:
  • Een drempelwaarde $\beta = 0.5$ voor het masker bleek optimaal.
  • Het gebruik van FPS in plaats van willekeurige sampling was cruciaal voor de prestaties.
  • Het verwijderen van Multi-Level Contrast leidde tot een merkbare daling in prestaties.

Betekenis en Toekomstperspectief

Dit paper toont aan dat het expliciet inbouwen van "lichaamsbewustzijn" (bodily awareness) in visuele representaties via contrastief leren op token-niveau een krachtige methode is voor robotmanipulatie. Het lost het probleem op van het verliezen van lichaamsinformatie in end-to-end systemen zonder de stabiliteit te offeren.

De methodiek biedt een nieuwe richting voor het ontwerpen van visuele encoders die niet alleen de omgeving begrijpen, maar ook de eigen actoren binnen die omgeving. Hoewel de huidige beperkingen bestaan in de rekenkosten van FPS en het gebrek aan real-world validatie, vormt ICon een veelbelovende basis voor toekomstig werk gericht op zero-shot transfer en robuustheid in complexe, real-world omgevingen.