AIM: Intent-Aware Unified world action Modeling with Spatial Value Maps

Probleemstelling

Bestaande wereld-actie-modellen (Unified World Action Models) voor robotica, die vaak zijn gebaseerd op vooraf getrainde videogeneratietools, kampen met een fundamentele structurele mismatch. Hoewel videomodellen uitstekend zijn in het voorspellen van hoe een scène visueel evolueert (wat de scène eruit zal zien), missen ze vaak het vermogen om te redeneren over waar een robot moet interageren en waarom die interactie nuttig is voor de taak.

• De beperking: Het direct decoderen van acties uit toekomstige visuele representaties (RGB-frames) dwingt het model om manipulatie-intenties impliciet af te leiden uit een dichte, visuele representatie die vaak ruis bevat (zoals achtergronddetails) die irrelevant is voor de volgende motorcommando's.
• Het gevolg: Dit vereist aanzienlijke robot-specifieke training om betrouwbare acties te genereren, wat de generalisatie en efficiëntie beperkt, vooral bij complexe, contact-gevoelige taken.

Methodologie: AIM (Intent-Aware Unified World Action Modeling)

De auteurs introduceren AIM, een model dat deze kloof overbrugt door een expliciete ruimtelijke interface in te voeren tussen toekomstvoorspelling en actiegeneratie.

1. Architectuur en Kernidee
In plaats van acties direct uit toekomstige visuele tokens te decoderen, voorspelt AIM een uitgelijnde ruimtelijke waardenkaart (Spatial Value Map).

• Ruwheid: Deze kaart encodeert de interactiestructuur die nodig is voor de taakuitvoering (bijv. waar een object moet worden gegrepen of geplaatst).
• Factorisatie: Het model factoriseert de voorspellingsverdeling als: eerst voorspellen van toekomstige visuele staten én waardenkaarten, en vervolgens decoderen van controleacties geconditioneerd op deze waardenkaarten, niet op de ruwe RGB-frames.
• Basis: Het model is gebouwd op de vooraf getrainde videogenerator Wan2.2-TI2V-5B en gebruikt een Mixture-of-Transformers architectuur.

2. Intent-Causale Aandacht (Intent-Causal Attention)
Een cruciaal onderdeel van de architectuur is een speciaal ontworpen masker voor de gedeelde zelf-aandacht (self-attention) lagen:

• Visuele stream: Kan kijken naar het verleden, instructies en toekomstige frames.
• Waarden-kaart stream: Kan kijken naar het verleden en toekomstige visuele frames.
• Actie-branch: Mag alleen kijken naar het verleden, de huidige observatie en de voorspelde waardenkaart.
• Effect: Dit zorgt ervoor dat de actie-branch toekomstinformatie uitsluitend via de intentie-gerichte waardenrepresentatie ontvangt, en niet direct via de visuele tokens. Dit schept een duidelijke scheiding tussen scenische evolutie en manipulatie-intentie.

3. Training en Self-Distillation RL
Het trainingsproces bestaat uit twee fasen:

• Fase I (Supervised Learning): Het model wordt getraind om toekomstige RGB-frames, waardenkaarten en acties gezamenlijk te genereren. De loss-functie combineert flow-matching voor video en waardenkaarten met een inverse-dynamica loss voor de acties.
• Fase II (Self-Distillation Reinforcement Learning): Na het supervisorische trainen wordt een RL-fase ingevoerd waarbij alleen de actie-branch wordt geoptimaliseerd, terwijl de video- en waarden-branches bevroren blijven.
  • Beloning: De actie-branch krijgt een dichte beloning (dense reward) gebaseerd op of de geprojecteerde actie-locatie valt binnen de hoge-waarde gebieden van de gevroren waardenkaart, gecombineerd met een spaarzame taak-succesbeloning.
  • Doel: Dit stabiliseert het visuele wereldmodel en optimaliseert de controle zonder catastrofale drift in de toekomstvoorspellingen.

4. Dataset
Om dit te ondersteunen, hebben de auteurs een simulatiedataset van 30.000 manipulatietrajecten opgezet (RoboTwin 2.0 benchmark). Deze dataset bevat gesynchroniseerde multi-view video's, acties en handmatige waardenkaart-annotaties die zijn gegenereerd op basis van contactgebeurtenissen en fysieke toestanden in de simulator.

Belangrijkste Bijdragen

1. Intent-bewust Model: Introductie van een expliciete ruimtelijke waardenkaart-interface die de koppeling tussen wereldvoorspelling en actiegeneratie verbetert door taakrelevante interactiestructuur te coderen.
2. Unificatie Framework: Een geïntegreerd trainingsframework dat gezamenlijke frame-waarde-generatie, intent-causale aandacht en post-training self-distillation RL combineert.
3. Prestaties en Dataset: Het creëren van een grote schaal dataset en het behalen van state-of-the-art resultaten op de RoboTwin 2.0 benchmark.

Resultaten

De experimenten op de RoboTwin 2.0 benchmark tonen aanzienlijke verbeteringen aan:

• Algemene Succesratio (SR): AIM bereikt een gemiddelde succesratio van 94,0% in de "Easy" setting en 92,1% in de "Hard" setting.
• Vergelijking: Dit presteert significant beter dan bestaande baselines zoals Motus, Fast-WAM, GigaWorld en LingBot-VA.
  • Bijvoorbeeld: Een verbetering van +5,3% t.o.v. Motus in de Easy setting.
• Specifieke Verbeteringen: De verbeteringen zijn het meest opvallend bij contact-gevoelige en stap-afhankelijke taken (zoals het plaatsen van objecten, het scannen van objecten of het draaien van schakelaars), waar nauwkeurige ruimtelijke lokalisatie van interactiegebieden cruciaal is.
• Interpreteerbaarheid: De waardenkaarten focussen zich op betekenisvolle interactiegebieden in plaats van algemene visuele aandacht, wat aantoont dat het model de juiste ruimtelijke intenties leert.

Betekenis en Conclusie

Dit werk demonstreert dat het expliciet modelleren van ruimtelijke intenties via waardenkaarten een effectieve brug vormt tussen visuele wereldmodellen en robotcontrole. Door de actiegeneratie te koppelen aan een "waarom en waar"-representatie in plaats van alleen "wat", kan AIM betrouwbare acties genereren zonder zware robot-specifieke hertraining. De methode biedt een schaalbaar en interpreteerbaar kader voor robuuste robotmanipulatie in complexe omgevingen.