Object-Centric World Models from Few-Shot Annotations for Sample-Efficient Reinforcement Learning

Deze paper introduceert OC-STORM, een object-gecentreerd model-based reinforcement learning-framework dat met slechts enkele geannoteerde frames de sample-efficiëntie aanzienlijk verbetert door dynamiek en interacties van beslissingsrelevante objecten te modelleren in complexe visuele omgevingen.

De kern

Stel je voor dat je een videospel moet leren spelen, zoals Hollow Knight of een klassieke Atari-game. Je hebt een robot die dit spel moet leren. De oude manier om dit te doen is alsof je de robot een camera voor de ogen hangt en zegt: "Kijk naar al die pixels op het scherm en probeer te raden wat er gebeurt."

Het probleem? De robot wordt overspoeld door details. Hij ziet de achtergrond, de wolken, de muren en de kleine vijand die hij moet verslaan. Omdat hij probeert alles perfect na te tekenen, vergeet hij vaak de belangrijkste dingen: de vijand en de bal. Het is alsof je probeert een auto te repareren door te kijken naar de kleur van de lucht erboven; je ziet de auto wel, maar je mist de motor.

Dit paper introduceert een slimme nieuwe methode genaamd OC-STORM. Laten we het uitleggen met een paar creatieve metaforen.

1. Het Probleem: De "Pixel-Overload"

In de huidige wereld van kunstmatige intelligentie (AI) proberen robots vaak het hele scherm na te bouwen. Ze zeggen: "Ik zie een blauwe pixel hier, een rode daar." Maar in een drukke, dynamische game is de achtergrond vaak saai en groot, terwijl de belangrijke dingen (zoals een vijand die een aanval doet) klein en snel zijn.

De robot besteedt zijn geheugen en rekenkracht aan het perfect nabootsen van de achtergrond, en vergeet daardoor de kleine, cruciale details. Het is alsof je een schilderij probeert te kopiëren, maar je besteedt 99% van je tijd aan het schilderen van de lucht, waardoor je vergeet de figuur in het midden te tekenen.

2. De Oplossing: De "Slimme Assistent"

De auteurs van dit paper zeggen: "Wacht even, waarom kijken we niet naar de objecten zelf?"

Ze introduceren een systeem dat werkt met een voorgekookte "object-herkenner".

• De Analogie: Stel je voor dat je een nieuwe game speelt. In plaats van dat de robot alles zelf moet leren, geef je hem een paar seconden van het spel en zeg je: "Kijk, dit is de speler, dit is de vijand, en dit is de bal."
• De robot gebruikt een slimme, vooraf getrainde "bril" (een bestaande technologie genaamd Cutie of SAM2) die deze objecten herkent. Deze bril is al expert in het vinden van dingen in video's.
• De robot hoeft niet meer te raden wat een "vijand" is; de bril zegt: "Hier is de vijand, hier is de speler."

3. Hoe het Werkt: De "Droomwereld"

De robot bouwt een wereldmodel. Dit is een soort "droomwereld" in zijn hoofd.

• Oude manier: De robot droomt van het hele scherm (achtergrond + alles) en probeert te voorspellen wat er morgen gebeurt. Omdat de achtergrond saai is, verliest hij de focus.
• Nieuwe manier (OC-STORM): De robot droomt alleen van de belangrijke personages. Hij zegt: "Oké, de vijand beweegt naar links, de speler springt." Hij negeert de saaie achtergrond.

Door zich te focussen op de "hoofdrolspelers" (de objecten), leert de robot veel sneller hoe het spel werkt. Hij heeft minder ervaring nodig om te winnen.

4. Het Resultaat: Sneller Leren met Minder Hulp

Het mooie van deze methode is dat je de robot niet hoeft te vertellen alles wat er in het spel zit. Je geeft hem maar een paar voorbeelden (bijvoorbeeld 6 tot 12 beelden) waarin je aangeeft wat de belangrijke dingen zijn. Daarna doet de slimme bril de rest van het werk.

• Op de Atari-benchmark: De robot leert veel sneller dan de oude methoden.
• In Hollow Knight: Dit is een heel moeilijke game met veel visuele rommel. Hier blinkt de robot uit. Hij verslaat de "bosses" (de grote vijanden) veel efficiënter dan voorheen, omdat hij precies weet waar hij op moet letten.

Samenvattend in één zin

In plaats van dat de robot probeert het hele schilderij na te maken, geven we hem een bril die alleen de belangrijke personages ziet, zodat hij zich kan concentreren op het spelen van het spel in plaats van het bestuderen van de achtergrond.

Dit maakt het mogelijk voor robots om complexe taken veel sneller en met minder "trainingstijd" te leren, wat een enorme stap is voor het gebruik van AI in de echte wereld.

Technische samenvatting

Probleemstelling

Diepe versterkende leer (Deep Reinforcement Learning - RL) vanuit pixels heeft indrukwekkende resultaten geboekt, maar lijdt aan een fundamenteel gebrek aan sample-efficiëntie. Agents hebben vaak orders van grootte meer ervaring nodig dan mensen om een taak te beheersen.

• Beperking van Model-Based RL (MBRL): Bestaande MBRL-methoden leren een wereldmodel door de omgeving te voorspellen op basis van pixelreconstructie (vaak met $L_2$-verlies).
• Het Kernprobleem: Deze reconstructiedoelstellingen worden gedomineerd door grote, statische achtergronden. Hierdoor worden kleine, maar cruciale objecten (zoals een vijand of een bal in een spel) genegeerd. In complexe omgevingen, zoals het spel Hollow Knight, failt een standaard wereldmodel (bijv. STORM) om deze beslissingsrelevante elementen correct te modelleren, wat leidt tot slechte beleidsleer (policy learning).

Methodologie: OC-STORM

De auteurs introduceren OC-STORM, een object-gecentreerd (Object-Centric - OC) MBRL-framework dat wereldmodellen verrijkt met objectrepresentaties afgeleid van een vooraf getraind segmentatienetwerk. Het doel is om de modelcapaciteit te richten op semantisch betekenisvolle entiteiten zonder uitgebreide labeling of toegang tot interne toestandsinformatie.

De werking van het framework:

1. Few-Shot Annotatie: In plaats van elke frame te labelen, annoteren de gebruikers slechts een klein aantal frames (bijv. 6-12) om de belangrijkste objecten in het spel te markeren.
2. Object-Feature Extractie: Een bevroren, vooraf getraind video-segmentatiemodel (zoals Cutie of SAM2) wordt gebruikt om compacte feature-vectoren te extraheren voor deze geannoteerde objecten. Deze modellen gebruiken retrieval-mechanismen om consistentie over tijd te behouden.
3. Dual-Input Architectuur: Het wereldmodel ontvangt twee soorten input:
   • Visuele input: Gereduceerde pixels (64x64).
   • Object-features: De extracted feature-vectoren van de segmentatiemodel.
4. Latente Representatie: Beide inputs worden gecodeerd naar discrete latent-variabelen via een Categorical VAE.
   • Object-features worden verwerkt via MLP-encoders.
   • Visuele input via CNN-encoders.
5. Spatiaal-Temporale Dynamica: Een transformer-architectuur (gebaseerd op STORM of DreamerV3) modelleert de dynamica. Deze gebruikt:
   • Spatiale attention: Om relaties tussen objecten en tussen objecten en de scène te modelleren.
   • Temporale attention: Om de evolutie van deze tokens over de tijd te voorspellen.
6. Beleidsleer: Het agent leert een beleid (policy) en een waarde-functie (critic) op basis van "imagined trajectories" (gesimuleerde ervaringen) gegenereerd door dit wereldmodel, volgens een actor-critic algoritme.

Belangrijkste Bijdragen

1. OC-STORM Framework: Het is, naar kennis van de auteurs, het eerste werk dat succesvol few-shot, vooraf getrainde objectsegmentatiemodellen integreert in wereldmodellen voor zowel de Atari 100k benchmark als de visueel complexe omgeving van Hollow Knight.
2. Generalisatie zonder Privileged Information: De methode vereist geen toegang tot interne game-state data (zoals hitboxes of coördinaten) en komt uit met minimale annotatie, wat het toepasbaar maakt voor real-world scenario's.
3. Uitgebreide Evaluatie: De auteurs testen verschillende backbones (DreamerV3, STORM), segmentatiemodellen (Cutie, SAM2) en representatiemethoden (vector vs. mask). Ze tonen aan dat vector-gebaseerde representaties superieur zijn aan mask-gebaseerde methoden in dit kader.
4. Ablatie en Analyse: Er wordt diepgaand geanalyseerd over de robuustheid tegen segmentatiefouten, de impact van het aantal annotaties, en de prestaties op continue controle-taken (Meta-World).

Resultaten

De experimenten tonen aan dat OC-STORM aanzienlijk beter presteert dan bestaande baselines:

• Atari 100k Benchmark: OC-STORM (specifiek de variant met Cutie) behaalt state-of-the-art sample-efficiëntie.
  • Het behaalt een gemiddelde Human-Normalized Score (HNS) van 134,8% (tegenover 119,4% voor de beste DreamerV3 variant en 124,6% voor de STORM baseline).
  • De prestaties zijn vooral sterk in games waar beslissingsrelevante informatie lokaal is in objecten.
• Hollow Knight (Boss Fights): In deze visueel complexe omgeving, waar standaard modellen vaak falen, convergeert OC-STORM aanzienlijk sneller en bereikt het hogere scores.
  • Bijvoorbeeld tegen de "Mage Lord" en "Pure Vessel" bossen, waar de standaard STORM agent moeite heeft, slaagt OC-STORM erin om effectieve strategieën te leren.
  • De win-rate tegen de "Hornet Protector" stijgt van 66,7% (STORM) naar 100% (OC-STORM).
• Robuustheid: Het systeem is robuust tegen onvolledige detectie; zelfs als het segmentatiemodel soms objecten mist, degradeert de prestatie slechts geleidelijk.
• Continue Controle: Op de Meta-World benchmark presteert OC-STORM ook beter dan gespecialiseerde methoden zoals MWM, wat aantoont dat de methode overdraagbaar is naar continue controle-taken.

Betekenis en Conclusie

OC-STORM markeert een belangrijke stap in de integratie van moderne computer vision (foundation models voor segmentatie) met versterkende leer.

• Oplossing voor Sample-Inefficiëntie: Door de wereldmodellen te dwingen om te focussen op semantische entiteiten in plaats van pixelruis, wordt de leer-efficiëntie drastisch verbeterd.
• Praktische Toepasbaarheid: De noodzaak van slechts een handvol annotaties maakt deze aanpak haalbaar voor complexe, dynamische omgevingen waar handmatige labeling onmogelijk is.
• Toekomstperspectief: Het werk onderstreept dat het combineren van object-gecentreerde priors met model-based RL een veelbelovende richting is voor het oplossen van complexe visuele taken, zowel in games als in de echte wereld.

De auteurs erkennen wel beperkingen, zoals moeite met het onderscheiden van identieke objecten (duplicaten) en het modelleren van geometrische structuren (zoals muren), maar over het geheel genomen biedt OC-STORM een krachtig en efficiënt framework voor visueel complexe RL-taken.