SegDAC: Visual Generalization in Reinforcement Learning via Dynamic Object Tokens

SegDAC is een nieuwe actor-critic-architectuur die visuele generalisatie in versterkend leren verbetert door dynamische, variabele object-tokens te gebruiken die zijn afgeleid van tekst-gebaseerde segmentatie, waardoor het aanzienlijk beter presteert dan bestaande methoden onder diverse visuele verstoringen.

De kern

🤖 De Robot die "Aandacht" leert, niet alleen "Kijken"

Stel je voor dat je een robot wilt leren om een blokje op een tafel te schuiven. Je leert de robot door hem duizenden keren te laten kijken naar foto's van de tafel.

Het probleem:
De meeste robots zijn als een kind dat alleen leert door te staren naar een foto. Als je de foto een beetje verandert – bijvoorbeeld door de lichten donkerder te maken, de muur een andere kleur te geven, of de robot een andere hoek te laten bekijken – raakt de robot in paniek. Hij denkt: "Oh nee, dit is een nieuwe wereld! Ik weet niet meer wat ik moet doen!" Hij is te afhankelijk van de exacte pixels (de kleine kleurtjes op het scherm) in plaats van de objecten zelf.

De oude oplossing:
Vroeger probeerden robot-onderzoekers dit op te lossen door de robot duizenden verschillende versies van dezelfde foto te laten zien (met verschillende filters, roterend, gekleurd). Dit is alsof je een kind duizenden foto's van een hond laat zien, van alle kanten en in alle weersomstandigheden, om het te leren herkennen. Het werkt soms, maar het is inefficiënt en de robot leert nog steeds niet echt wat een hond is, hij leert alleen hoe hij op een hond moet lijken.

💡 De nieuwe aanpak: SegDAC

De auteurs van dit paper hebben een slimme nieuwe manier bedacht, genaamd SegDAC. In plaats van naar de hele foto te kijken, leren ze de robot om de foto te ontleden in losse onderdelen.

Hier is hoe het werkt, stap voor stap, met een analogie:

1. De "Taal van de Objecten" (Segmentatie)

Stel je voor dat de robot een foto krijgt. In plaats van de hele foto als één grote, rommelige puzzel te zien, gebruikt SegDAC een slimme "bril" (een vooraf getrainde visuele AI) om de foto te snijden in losse stukjes.

• De robot ziet nu niet meer "een grijs vlak met een rode stip", maar hij ziet: "Hier is de robotarm, hier is het blokje, en hier is de achtergrond."
• De truc: De robot krijgt een lijstje met deze losse stukjes. Als er vandaag 3 objecten zijn en morgen 5, maakt dat niet uit. De robot is flexibel.

2. De "Dynamische Teamvergadering" (Transformer)

Nu heeft de robot een lijstje met objecten. Hoe beslist hij wat hij moet doen?

• Oude robots: Kijken naar een vast aantal "vakjes" (slots). Als er meer objecten zijn dan vakjes, gooien ze er een paar weg. Als er minder zijn, vullen ze de rest met "lege lucht". Dit is alsof je een vergadering houdt waarbij je altijd 10 stoelen hebt, ook al komen er maar 3 mensen. De lege stoelen verwarren de vergadering.
• SegDAC: Gebruikt een Transformer (dezelfde technologie als ChatGPT). Dit is als een vergadering waar je precies zoveel stoelen neerzet als er mensen zijn.
  • Als er 3 objecten zijn, zijn er 3 stoelen.
  • Als er 10 zijn, zijn er 10 stoelen.
  • De robot kan dan met al zijn objecten "praten" en beslissen: "Oké, de arm is hier, het blokje is daar, en de tafel is blauw. Ik moet de arm naar het blokje bewegen."

3. De "Ruimtelijke GPS" (Positie-codering)

Een groot probleem bij het losmaken van objecten is: "Waar zit dat blokje precies?" Als je alleen zegt "er is een blokje", weet je niet of het links of rechts is.

• SegDAC plakt een GPS-coördinaat op elk object-token. De robot weet dus niet alleen wat het is, maar ook waar het is, zelfs als de achtergrond verandert.

🚀 Waarom is dit zo geweldig?

De onderzoekers hebben hun robot getest op 8 verschillende taken (zoals blokjes schuiven, pinnen in gaten steken) en ze hebben de robot blootgesteld aan 12 soorten veranderingen:

• Andere camera-hoeken.
• Andere lichtkleuren.
• Andere texturen (bijv. een houten tafel wordt een glazen tafel).
• Zelfs "semantische" verwarring (bijv. de tafel is precies even rood als het blokje, zodat ze samensmelten).

Het resultaat:

• De oude robots (die op pixels keken) vielen vaak volledig uit elkaar. Ze haalden soms maar 10% van hun normale score.
• SegDAC bleef kalm. Omdat hij naar de objecten keek en niet naar de pixelkleurtjes, kon hij de taak nog steeds uitvoeren, zelfs als de wereld er totaal anders uitzag.
• De prestatie: SegDAC was tot 88% beter dan de beste bestaande methoden op de moeilijkste niveaus.

🏆 De "Superkracht" van SegDAC

Het mooiste is dat SegDAC dit doet zonder dat hij duizenden extra foto's nodig heeft om te leren (data-augmentatie).

• Oude methode: "Laten we de robot 10.000 keer een foto laten zien met verschillende filters, zodat hij het leert." (Traag en veel werk).
• SegDAC: "Laten we de robot leren om de wereld te zien als losse objecten." (Snel en slim).

Het is alsof je iemand leert autorijden.

• De oude methode is: "Rijd 10.000 keer door regen, sneeuw, zon en mist, zodat je het leert."
• SegDAC is: "Leer de bestuurder om de weg, de auto's en de verkeersborden te herkennen, ongeacht het weer."

Conclusie

SegDAC is een doorbraak omdat het robots leert om abstracter te denken. Ze kijken niet meer naar de "ruis" (de achtergrond, het licht), maar focussen puur op de essentie (de objecten en hun relaties). Hierdoor worden robots veel robuuster en kunnen ze makkelijker worden ingezet in de echte wereld, waar dingen nooit precies hetzelfde zijn als in de training.

Kortom: SegDAC leert robots om te kijken met hun "verstand" in plaats van alleen met hun "ogen".

Technische samenvatting

Probleemstelling

Visuele versterkende leer (RL) agenten die zijn getraind op pixel-observaties, kampen vaak met ernstige generalisatieproblemen wanneer de visuele omstandigheden tijdens de testfase veranderen. Kleine wijzigingen in achtergrondtexturen, belichting of objectkleuren kunnen leiden tot grote prestatiedalingen, zelfs als de onderliggende taakstructuur onveranderd blijft.

Hoewel object-gerichte representaties een veelbelovend alternatief zijn, hebben bestaande methoden belangrijke beperkingen:

1. Ze gebruiken vaak vaste grootte "slot"-representaties (een vooraf bepaald aantal objecten), wat niet past bij dynamische scènes.
2. Ze vereisen beeldreconstructie als trainingsdoel, wat de representaties kan laten afwijken naar visuele precisie in plaats van taakrelevantie.
3. Ze hebben auxiliary losses (hulpverliezen) of uitgebreide data-augmentatie nodig om object-decompositie te leren.

Het blijft onduidelijk hoe men RL-beleid direct kan leren vanuit object-niveau input zonder deze beperkingen, vooral wanneer het aantal en de identiteit van objecten per tijdstip veranderen.

Methodologie: SegDAC

De auteurs stellen SegDAC (Segmentation-Driven Actor-Critic) voor, een architectuur die werkt met een variabele set van object-token embeddings. Het systeem bestaat uit drie hoofdfasen:

1. Van Pixels naar Object Maskers (Segmentatie):

   • In plaats van pixel-augmentatie of prompt-vrije segmentatie (wat traag is), gebruikt SegDAC een tekst-gebaseerde aanpak.
   • Een bevroren open-vocabulary detector (YOLO-World) genereert bounding boxes op basis van een korte lijst van conceptwoorden (bijv. "robot", "blok", "achtergrond").
   • Een bevroren semantische segmentatiemodel (EfficientViT-SAM) genereert maskers binnen deze boxes.
   • Een lichtgewicht morphologische nabewerking (openen en sluiten) verwijdert ruis en vult gaten zonder de latentheid significant te verhogen.
   • Het resultaat is een variabele set van $N$ objectmaskers per frame.

2. Van Maskers naar Contextuele Object Tokens:

   • Voor elk masker worden de patch-embeddings van een bevroren Vision Transformer (ViT) encoder geselecteerd die overlappen met het masker.
   • Door global average pooling over deze patches wordt één compacte vector (token) gegenereerd per object.
   • Belangrijk: Omdat de ViT-encoder al globale context via self-attention heeft verwerkt, bevatten deze tokens zowel lokale details als scène-context, zonder extra berekening.
   • Deze tokens worden direct in de replay buffer opgeslagen, wat het opnieuw uitvoeren van de encoder tijdens training voorkomt en de sample-efficiëntie verhoogt.

3. Transformer Actor-Critic voor Dynamische Tokens:

   • De actor en critic zijn gebaseerd op een Transformer-decoder.
   • Variabele lengte verwerking: In plaats van padding tot een vaste lengte, worden alle tokens van een batch samengevoegd tot één "packed" sequence met een attention-mask. Dit zorgt voor efficiëntie en schaalbaarheid zonder een harde bovengrens op het aantal objecten.
   • Segment Positieve Encoding: Elke object-token krijgt een leerbaar positie-encoding gebaseerd op de bounding-box coördinaten. Dit is cruciaal omdat de ViT-features alleen niet voldoende ruimtelijke gronding bieden voor precisie-taken.
   • Modality Embeddings: Tokens worden onderscheiden van proprioceptieve input en query-tokens.
   • De critic gebruikt een actie-geconditioneerde query om Q-waarden te voorspellen voor specifieke acties.

Het hele systeem wordt getraind met de standaard SAC (Soft Actor-Critic) loss, zonder reconstructieverliezen, hulpverliezen of data-augmentatie.

Kernbijdragen

1. Een transformer-based Actor-Critic die stabiele, model-vrije beleidsregels leert van een variabele set object-tokens, robuust tegen natuurlijke variatie in aantal en identiteit, zonder reconstructie of data-augmentatie.
2. Een methode voor contextuele object-tokens die bevroren pretrained visuele modellen gebruikt met segmentatie-positieve encoding, wat ruimtelijke gronding behoudt zonder ground-truth maskers of finetuning.
3. Empirische evaluatie op 8 ManiSkill3 manipulatietaakken met 12 soorten visuele perturbaties, waarbij SegDAC de state-of-the-art (SOTA) methoden overtreft.
4. Een nieuw visueel generalisatie-benchmark op ManiSkill3 met drie moeilijkheidsniveaus (Eenvoudig, Middel, Moeilijk) en een semantische taxonomie voor scene-entiteiten.

Resultaten

SegDAC werd geëvalueerd tegen zes baselines (waaronder DrQ-v2, SAM-G, MaDi, SADA) over 8 manipulatietaakken.

• Visuele Generalisatie: SegDAC presteert aanzienlijk beter dan bestaande methoden onder visuele perturbaties:
  • +15% verbetering op "Eenvoudige" settings.
  • +66% verbetering op "Middel" settings.
  • +88% verbetering op de "Moeilijkste" settings (waar andere methoden vaak volledig falen).
• Sample Efficiency: SegDAC bereikt een sample-efficiëntie die gelijkwaardig is aan DrQ-v2 (de SOTA voor sample-efficiëntie), terwijl DrQ-v2 zelf slecht generaliseert onder visuele veranderingen. Dit doorbreekt de gebruikelijke trade-off tussen generalisatie en sample-efficiëntie.
• Robuustheid: De methode is robuust tegen variatie in het aantal gedetecteerde segmenten (bijv. door occlusie). De agent faalt "gracefully" (gestructureerd) in plaats van chaotisch gedrag te vertonen wanneer perturbaties optreden.
• Ablatie Studies: Het verwijderen van de "segment positional encoding" of het forceren van een vast aantal tokens leidt tot significante prestatiedalingen, wat aantoont dat beide componenten essentieel zijn.

Betekenis en Impact

Dit paper is significant omdat het een brug slaat tussen object-gerichte representaties en praktische, efficiënte RL-training.

• Efficiëntie: Door te werken met een compacte set van object-tokens in plaats van pixels, en door bevroren encoders te gebruiken, wordt de rekentijd drastisch verlaagd, waardoor online RL op hoge resolutie (512x512) haalbaar wordt op één GPU.
• Generalisatie: Het bewijst dat het redeneren over objecten in plaats van pixels een gunstige inductieve bias biedt voor zowel leren als generalisatie.
• Praktische Toepassing: De methode vereist geen handmatige data-augmentatie of complexe reconstructiedoelen, wat het eenvoudiger en robuuster maakt voor toepassing in complexe, dynamische omgevingen.

Kortom, SegDAC toont aan dat het direct leren van dynamische object-tokens via een transformer-architectuur de huidige beperkingen van visuele RL voorbijstijgt, met name in scenario's met grote visuele variatie.