Technische Samenvatting: Causale Object-Centrische Modellen voor Planning met Monte Carlo Tree Search

Probleemstelling

Model-gebaseerd reinforcement learning (MBRL) heeft aanzienlijk succes behaald door de dynamiek van omgevingen te leren om geïmageerde ervaringen te genereren, waardoor de noodzaak voor interactie met de echte wereld wordt verminderd. Het leren van nauwkeurige wereldmodellen blijft echter uitdagend in hoogdimensionele, niet-stationaire omgevingen die bestaan uit meerdere interagerende objecten.

Bestaande visuele RL-benaderingen vertrouwen doorgaans op Convolutional Neural Network (CNN) encoders die een enkele holistische representatie van een invoerbeeld produceren. Deze monolithische representaties falen vaak in het vastleggen van object-niveau structuren en interacties, die cruciaal zijn voor besluitvorming. In complexe scènes kan de aanwezigheid van dynamische achtergronden, irrelevante entiteiten of kleine maar taakrelevante objecten de prestaties van de agent aanzienlijk verslechteren. Terwijl mensen de wereld waarnemen als bestaande uit discrete entiteiten, wat efficiënt redeneren mogelijk maakt, worstelen standaard RL-agents om deze entiteiten te isoleren zonder expliciete supervisie of instance segmentatie-masks.

Methodologie: COMET

De auteurs introduceren COMET (Causal Object-centric Model for Efficient Tree search), een model-gebaseerd reinforcement learning algoritme dat Monte Carlo Tree Search (MCTS) uitvoert binnen een slot-gestructureerde latente ruimte. COMET integreert drie primaire componenten:

1. Bevroren Object-Centrische Encoder

COMET maakt gebruik van een bevroren, unsupervised object-centrische encoder (bijv. SLATE, DINOSAUR, of Slot Contrast) om visuele observaties in te mappen naar een set object-centrische latente representaties, of "slots" ($\bar{s}_t = \{s^1_t, \dots, s^n_t\}$).

• Slot Initialisatie: Om de permutatie-invariantie en de stochastische aard van slot attention in statische beeldmodellen aan te pakken, initialiseren de auteurs de slot-representaties op tijdstip $t+1$ met behulp van de slots van tijdstip $t$. Dit zorgt voor temporele consistentie van object-toewijzingen binnen een episode.
• Bevriezen: De encoder is getraind via een random policy en blijft bevroren tijdens de RL-fase om stabiele object-niveau features te bieden.

2. Transformer-gebaseerd Wereldmodel met Action-Slot Fusie

De kern van COMET is een transformer-gebaseerd wereldmodel dat opereert op deze slots. In tegenstelling tot standaard MuZero-achtige benaderingen die een monolithische state vector verwerken, verwerkt COMET een sequentie van object-slots.

• Action-Object Binding: Een nieuw mechanisme fuseert acties met object-slots. In plaats van alle object-latents te comprimeren tot een enkele vector om de volgende staat te voorspellen (wat een bottleneck creëert), wordt de action embedding onafhankelijk geconcateneerd met elke slot. Deze slot-geconditioneerde action embeddings worden door een gedeelde MLP projector en vervolgens door een transformer-backbone gestuurd om de volgende state slots ($\hat{s}_{t+1}$) en rewards te voorspellen.
• Architectuur: Het model volgt een UniZero-stijl trainingspipeline met gebruik van een nanoGPT-gebaseerde transformer-backbone. Het voorspelt toekomstige slots en rewards op basis van de huidige set slots en de uitgevoerde actie.

3. Object-Causale Attention voor Policy en Value

De policy en value heads maken gebruik van een aangepaste transformer-architectuur uitgerust met object-causale attention.

• Causale Scores: Voor elke latente slot $h^k_t$ schat het netwerk een causale score $\alpha^k_t \in [0, 1]$ in, die de waarschijnlijkheid vertegenwoordigt dat het bijbehorende object relevant is voor de huidige beslissing.
• Attention Modulatie: Deze scores worden gebruikt om de geschaalde dot-product attention te moduleren. Een geleerde gewichtsmatrix $W_t$ tilt een vooraf gedefinieerde causale graafstructuur naar het token-niveau, waardoor het attention-mechanisme kan concentreren op taakrelevante entiteiten (causale objecten) terwijl irrelevante entiteiten (niet-causale objecten) worden onderdrukt. Dit stelt de agent in staat om de planning te richten op de specifieke subset van objecten die momenteel met de omgeving interageren.

Belangrijkste Bijdragen

1. COMET Algoritme: De introductie van een MCTS-gebaseerd object-centrisch MBRL algoritme dat bevroren object-niveau representaties combineert met een transformer-gebaseerd wereldmodel voor planning in een gestructureerde latente ruimte.
2. Action-Slot Fusion Mechanisme: Een nieuwe methode om acties te binden aan object-centrische slots binnen een verenigde transformer-backbone. Dit maakt object-centrische wereldmodellering en policy/value voorspelling mogelijk zonder de objectinformatie te laten instorten in een enkele vector.
3. Object-Causale Attention: De integratie van geleerde causale scores om token-interacties te moduleren, waardoor de policy tijdens besluitvorming dynamisch kan focussen op taakrelevante entiteiten.
4. Empirische Evaluatie: Uitgebreide evaluatie over acht visueel en dynamisch diverse taken, inclusief de Object-Centric Visual RL benchmark, ManiSkill, Robosuite en VizDoom.

Experimentele Resultaten

COMET werd geëvalueerd tegen verschillende baselines:

• Object-Centrische Baselines: OCRL (model-free PPO met object encoder), OCDreamer (object-centric RSSM), en SOLD (unsupervised slot-based dynamics).
• Monolithische Baselines: UniZero (standaard MuZero-stijl latente planning).

Prestatiebevindingen:

• Sample Efficiëntie: Over de diverse reeks taken heen bereikt COMET een hogere Mean Normalized Score tijdens de vroege stadia van training vergeleken met zowel sterke monolithische MCTS-gebaseerde methoden als andere object-centrische RL-baselines.
• Taakspecifieken:
  • In visueel eenvoudige taken (bijv. Object Goal, Object Comparison), maakt COMET gebruik van hoogwaardige object-representaties en causale attention om relevante objecten accuraat te identificeren, wat leidt tot snellere convergentie.
  • In continuous control taken (bijv. Block Lifting, Cube Pushing), vertoont COMET consistente prestaties, hoewel het niet universeel beter presteert dan alle baselines op elke specifieke metriek (het overtrof bijvoorbeeld de concurrenten niet op Block Lifting).
• Generalisatie: De methode toont robuustheid over omgevingen met variërende visuele complexiteit, van 2D-vormen op monochrome achtergronden tot complexe 3D robotmanipulatie en first-person shooter scenario's.

Betekenis en Claims

Het artikel stelt dat het expliciet modelleren van het belang van individuele objecten voor besluitvorming de policy-learning verbetert in object-georiënteerde omgevingen. Door een expliciete object-niveau inductieve bias te introduceren in MuZero-stijl latente planning, stelt COMET het model in staat om te redeneren over individuele entiteiten en hun interacties, in plaats van de scène te behandelen als een holistische vlek.

De auteurs beweren dat COMET effectief het probleem aanpakt van het representeren van observaties in complexe, multi-object omgevingen. Door object-representaties te ontkoppelen en causale attention te gebruiken om relevantie te filteren, bereikt het algoritme een hogere sample-efficiëntie en meer gefocuste besluitvorming. Het werk suggereert dat object-centrische inductieve biases een levensvatbaar en effectief pad zijn voor het verbeteren van model-gebaseerde planning in visuele RL, met name wanneer instance segmentatie niet beschikbaar is of wanneer unsupervised representation learning vereist is.