Technische Samenvatting: WAM4D

Probleemstelling

World Action Models (WAMs) zijn opgekomen als een veelbelovende benadering voor belichaamde intelligentie door toekomstige visuele observaties en uitvoerbare robotacties gezamenlijk te modelleren. De huidige WAMs opereren echter voornamelijk in 2D-video of latente ruimtes. Hoewel deze modellen visueel plausibele toekomstige rollouts kunnen genereren, falen ze vaak in het vastleggen van de 3D-ruimtelijke beperkingen, de verborgen contactgeometrie en de objectdynamiek die vereist zijn voor precieze robotmanipulatie.

Recente pogingen om 4D-geometrie (bijv. dieptekaarten, puntenwolken) in WAMs te integreren, behandelen dichte geometrische representatie doorgaans als een expliciet inferentiedoel. Deze aanpak introduceert aanzienlijke computationele overhead door kostbare geometrische decodering tijdens actie-inferentie, wat de causale generatie van acties vertraagt. Bovendien kan het expliciet voorspellen van dichte 4D-geometrie het doel van het model verschuiven naar geometrische reconstructie in plaats van het versterken van de causale koppeling tussen videovoorspelling en actiegeneratie. De kernuitdaging is hoe sterke geometrische priors gebruikt kunnen worden om ruimtelijk redeneren en actievoorspelling te verbeteren zonder de latentie en complexiteit van dichte geometrische decodering tijdens implementatie te veroorzaken.

Methodologie: WAM4D

De auteurs stellen WAM4D voor, een snelle 4D world action model dat geometrische foundation priors integreert in een causale video-actie transformer met behulp van een lichtgewicht mechanisme dat alleen tijdens de training wordt gebruikt.

1. Spatial Register Tokens

In plaats van dichte geometrie uit te voeren tijdens de inferentie, introduceert WAM4D spatial register tokens als een compacte brug tussen 2D-video latents en geometrische priors.

• Mechanisme: Tijdens de training fungeren leerbare register tokens als queries die toekomstige geometrische informatie extraheren uit causale historische video-features. Deze registers zijn afgestemd op de ruimtelijke lay-out van de multi-view RGB-mozaïek.
• Distillatie: De registers worden bijgewerkt via depth extraction blocks (cross-attention met historische video tokens) en geprojecteerd naar een pre-trained geometric head (geïnitialiseerd vanuit Depth Anything V3). Het model voorspelt toekomstige dieptesequenties vanuit deze registers.
• Training vs. Inferentie: De diepteverliesfunctie (SmoothL1) backpropageert naar de gedeelde video-actie backbone, waardoor ruimtelijke priors worden gedestilleerd in de historische features die worden gebruikt voor actievoorspelling. Cruciaal is dat tijdens de implementatie de register branch, de depth blocks en de geometric head worden verwijderd, waardoor een lichtgewicht 2D observatie-naar-actie interface overblijft die identiek is aan standaard causale WAMs.

2. Causal Mixture Attention (MoT)

Om niet-causale shortcuts te voorkomen waarbij het model zou kunnen "valsspelen" door naar toekomstige geometrie te kijken om acties te voorspellen, ontwerpen de auteurs een Causal Mixture Attention mechanisme voor de Mixture-of-Transformers (MoT) backbone.

• Zichtbaarheidsregels: De attention mask definieert strikt de zichtbaarheid tussen video-, actie- en geometrie-tokens.
  • Toekomstige actie-tokens kunnen aandacht schenken aan historische video, historische acties en hun eigen ruizige toekomstige staten.
  • Toekomstige actie-tokens zijn gemaskeerd voor toekomstige video-tokens en spatial registers.
  • Spatial registers richten hun aandacht alleen op zichzelf en geldige historische video tokens.
• Resultaat: Dit zorgt ervoor dat geometrische supervisie de causale video-features tijdens de training vormgeeft zonder dat hulp-geometrie tokens tijdens de inferentie aan de actie-voorspeller worden blootgesteld.

3. Architectuur en Training

• Backbone: Gebouwd op de causale video-actie backbone van LingBot-VA, gebruikmakend van een video VAE (Wan2.2) en flow-matching doelen voor zowel de video- als de actiestroom.
• Doelstelling: De totale loss combineert conditional flow matching losses voor video ($\mathcal{L}_{video}$) en actie ($\mathcal{L}_{action}$) met een diepte-supervisie loss ($\mathcal{L}_{depth}$) toegepast op de register branch.
• Implementatie: Het model gebruikt een drie-view mozaïek (kop + twee polscamera's) en voorspelt absolute end-effector poses.

Belangrijkste Bijdragen

1. WAM4D Model: Een snelle 4D world action model die geometrische foundation priors overdraagt naar causale video-actie representaties via een compacte spatial-register interface, waardoor dichte geometrische decodering tijdens de inferentie wordt vermeden.
2. Causal Mixture Attention: Een nieuw aandachtmechanisme voor MoT backbones dat modality-specifieke zichtbaarheid definieert, waardoor geometrische supervisie tijdens de training mogelijk is terwijl een lichtgewicht, causaal beleidspad behouden blijft voor implementatie.
3. Empirische Validatie: Uitgebreide evaluatie op de RoboTwin 2.0 simulatie suite en real-world AstriBot S1 taken, wat een verbeterde ruimtelijke consistentie en succespercentages aantoont.

Experimentele Resultaten

Het paper evalueert WAM4D tegen sterke baselines inclusief VLA policies ($\pi_0$, $\pi_{0.5}$) en recente video-actie modellen (Motus, LingBot-VA, Fast-WAM).

• RoboTwin 2.0 (Simulatie): WAM4D bereikte een succespercentage van 93,8% in schone settings en 89,9% in gerandomiseerde settings, waarmee het state-of-the-art baselines evenaart of overtreft (bijv. LingBot-VA: 92,9% clean, 91,6% rand). Het behield een concurrerende inferentie-latentie (~525ms) en lager VRAM-gebruik (9.71 GiB) vergeleken met modellen zoals Motus (1516ms) en LingBot-VA (843ms).
• Real-World (AstriBot S1): Op vier manipulatietaken (bord tillen, fles plaatsen, dop verwijderen, LEGO sorteren) behaalde WAM4D het hoogste gemiddelde sub-actie succespercentage (90%), waarmee het $\pi_{0.5}$ (74%) en Fast-WAM (80%) overtrof.
• Ablatie-studies:
  • Registerplaatsing: Middle-layer registers (lagen 12, 14, 16, 18) boden de beste balans tussen geometrische distillatie en visuele feature preservatie.
  • Geometric Head: Het gebruik van een trainable pretrained geometric head (geïnitialiseerd vanuit Depth Anything V3) leverde significant betere resultaten op dan willekeurige initialisatie of een vaste pretrained head, wat de noodzaak bevestigt om sterke geometrische priors aan te passen.
  • Causale Beperkingen: Het verwijderen van de causale masking (bidirectionele registers) verbeterde de succespercentages licht, maar verslechterde de geometrie-metrieken en verhoogde de complexiteit, wat de unidirectionele vormgeving rechtvaardigt.

Betekenis en Claims

Het paper claimt dat WAM4D succesvol de trade-off tussen ruimtelijke consistentie en inferentie-efficiëntie in world action modeling aanpakt. Door geometrie te behandelen als een training-time readout in plaats van een inference-time output, bereikt WAM4D het volgende:

• Het verbetert het vermogen van het model om te redeneren over de omvang van objecten, verborgen oppervlakken en contactgeometrie.
• Het behoudt de lichtgewicht implementatiekenmerken van 2D video-actie modellen, waardoor het geschikt is voor real-time robotbesturing.
• Het demonstreert dat geometrische foundation priors effectief kunnen worden gedestilleerd in causale video-features om manipulatieprestaties te verbeteren zonder dat hiervoor expliciete 3D-reconstructie tijdens de uitvoering nodig is.

De auteurs erkennen beperkingen en merken op dat het model geen expliciet langetermijngeheugen voor objecten heeft, wat kan leiden tot inconsistenties in de identiteit tijdens lange autoregressieve rollouts wanneer objecten worden afgeschermd. Ze benadrukken echter dat dit de beleidssucces niet in gevaar brengt in closed-loop controle, aangezien het model continu verse observaties van de omgeving ontvangt.