Flex-$\pi$: A Multi-Stream World-Action Model with Compute Flexibility

Technische Samenvatting van de Preprint: FLEX-π: EEN MULTI-STREAM WORLD-ACTION MODEL MET COMPUTE FLEXIBILITEIT

Probleemstelling

World-action modellen (WAMs) zijn een krachtig alternatief geworden voor vision-language-action modellen (VLAs) door gezamenlijk toekomstige observaties en acties te voorspellen, waardoor ze robuuste ruimtelijke-temporele representaties leren en priors erven van grootschalige videogeneratie. De huidige generieke WAMs voorspellen echter bijna uitsluitend toekomstige RGB-beeldlatenten. Deze latenten, getraind op pure pixelreconstructie, leggen uiterlijke details vast maar slagen er vaak niet in om uit te lijnen met de expliciete 3D-geometrie en object-niveau semantiek die vereist zijn voor precieze robotmanipulatie.

Het direct superviseren van WAMs op geometrie en semantiek brengt doorgaans drie significante kosten met zich mee:

1. Extra Sensoren: Het vereisen van nieuwe hardware-modaliteiten (bijv. dieptesensoren).
2. Nieuwe Pre-training: Het noodzakelijk maken van het trainen van nieuwe visuele priors vanaf nul.
3. Inference Latentie: Het vertragen van implementatie door complexe multi-modale verwerking.

De paper stelt de vraag: Hoe kunnen we de sterke punten van WAMs (gezamenlijke voorspelling en sterke visuele priors) versterken om 3D-geometrie en semantiek te omvatten zonder deze offers te brengen?

Methodologie: FLEX-π

De auteurs introduceren FLEX-π, een WAM met 6 miljard parameters dat gelijktijdig toekomstige RGB-observaties, 3D-pointmaps en object-centrische DINO-semantiek voorspelt. De kerninnovatie is het bereiken van deze "gratis lunch" van multi-modale gronding zonder nieuwe sensoren of pre-training.

1. Gedeelde Latente Ruimte via een Bevroren VAE

FLEX-π maakt gebruik van een enkele, bevroren Variational Autoencoder (VAE) van een vooraf getraind videogeneratiemodel (Wan-2.2).

• RGB Codering: De VAE codeert RGB-beelden zoals gebruikelijk.
• 3D Pointmap Codering: Verrassend genoeg ontdekten de auteurs dat deze zelfde bevroren VAE, die uitsluitend op RGB-pixels is getraind, 3D-pointmaps bijna verliesloos kan coderen en decoderen. Dit maakt het mogelijk om 3D-geometrie in exact dezelfde latente ruimte als RGB te projecteren zonder een nieuwe encoder te trainen.
• Semantische Codering: Object-centrische semantiek wordt afgeleid met behulp van een bevroren DINOv3-encoder.

2. Multi-Stream Mixture-of-Transformers (MoT) Backbone

Het model embedt drie visuele stromen (RGB, Pointmap, DINO) in een gedeelde latente ruimte van tokenstromen.

• Architectuur: Een Mixture-of-Transformers backbone (5 miljard parameters) verwerkt deze stromen. Het gebruikt een gedeelde visuele transformer-backbone maar behoudt onafhankelijke feedforward projectienetwerken voor elke modaliteit.
• Action Expert: Een aparte, smallere Action Expert Transformer (1 miljard parameters) voert cross-attention uit op de visuele stromen om actie-chunks te genereren.
• Causale Gezamenlijke Generatie: Actie-tokens letten op huidige observaties en voeren cross-attention uit op actieve toekomstige visuele tokens. Dit stelt de generatie van acties in staat om te profiteren van de evoluerende representatie van de toekomst door het model.

3. Flexibele Training via Stream Dropout en Cross-Modality Forcing

Om een enkel checkpoint te laten ondersteunen voor elke combinatie van inputs en outputs tijdens inference, gebruikt FLEX-π twee belangrijke trainingsstrategieën:

• Visual Stream Dropout: Tijdens de training worden inputstromen willekeurig gedropt (met een kans van 0,5), wat ervoor zorgt dat het model kan opereren met gedeeltelijke observaties (bijv. alleen RGB).
• Cross-Modality Forcing: Cruciaal is dat zelfs als een stroom wordt gedropt van de input, het model nog steeds wordt gedwongen om die stroom te denoisen en te voorspellen bij de output. Bijvoorbeeld, als de 3D-pointmap ontbreekt in de input, moet het model de toekomstige pointmap synthetiseren op basis van de RGB- en DINO-inputs. Dit dwingt het model om een verenigde representatie te internaliseren waarbij verschijning, geometrie en semantiek wederzijds voorspellend zijn.

4. Inference Flexibiliteit

Bij implementatie kunnen gebruikers een afweging maken tussen latentie en prestaties door te selecteren welke stromen zij genereren:

• Action-Only Mode: Het model voorspelt alleen acties (snelste inference).
• Full Joint Generation: Het model voorspelt acties plus toekomstige RGB, pointmaps en DINO-features (hoogste prestaties).
• Partial Generation: Elke subset van stromen kan worden gegenereerd.

Belangrijkste Bijdragen

1. Gratis Geometrische en Semantische Gronding: Het aantonen dat een vooraf getrainde RGB-video VAE 3D-pointmaps verliesloos kan coderen, waardoor WAMs kunnen worden gesuperviseerd op 3D-geometrie en DINO-semantiek zonder nieuwe sensoren of pre-training.
2. Compute Flexibiliteit: Een enkel getraind checkpoint ondersteunt elke subset van input/output-stromen, waardoor praktijkgebruikers bij de implementatie hun optimale punt op de snelheid-prestatie-frontier kunnen kiezen.
3. Cross-Modality Forcing: Een trainingsregime dat het model dwingt om ontbrekende modaliteiten te voorspellen vanuit de geobserveerde modaliteiten, wat de interne representatie en beleidsperformance aanzienlijk verbetert.

Experimentele Resultaten

De auteurs evalueren FLEX-π op simulatiebenchmarks (RoboTwin, LIBERO, LIBERO-Plus) en real-world bimanuele manipulatietaken op een YAM-robot.

• Prestaties in de echte wereld: Op vijf behendige, precieze real-world taken (inclusief "Self-Repair Gripper" en "Soft-Bag Zipping") presteert FLEX-π 2–7× beter dan sterke baselines (π0.5, ManiFlow, Fast-WAM) in succesratio's.
  • Opvallend genoeg presteert de action-only variant van FLEX-π (die geen visuele toekomsten genereert tijdens inference) al beter dan alle baselines en draait sneller dan π0.5.
  • Volledige gezamenlijke generatie verhoogt de prestaties verder en bereikt tot wel een absolute verbetering van 42,7% in taakvoltooiing op moeilijke taken vergeleken met de sterkste baseline.
• Data-efficiëntie: FLEX-π is aanzienlijk meer data-efficiënt. In RoboTwin presteert het in low-data regimes (50 demonstraties) 1,9–4,5× beter dan baselines. Zelfs wanneer getraind op slechts 50% van de real-world data, presteert FLEX-π beter dan baselines die de volledige dataset gebruiken.
• Generalisatie (OOD): FLEX-π vertoont robuustheid tegen out-of-distribution verschuivingen (ongeziene objecten, rommel, distractoren). Terwijl baselines zoals ManiFlow en π0.5 aanzienlijke prestatieverliezen (25–26%) lijden onder distributieverschuiving, daalt FLEX-π met minder dan 5%.
• Ablaties:
  • Het verwijderen van DINO- of Pointmap-inputs tijdens de training vermindert de succesratio's aanzienlijk (met respectievelijk 6,8% en 20%).
  • Het verwijderen van Cross-Modality Forcing schaadt de succesratio's met 21%, wat bevestigt dat het dwingen van het model om ontbrekende modaliteiten te synthetiseren cruciaal is voor het leren van een sterke gedeelde representatie.

Betekenis en Claims

De paper claimt dat FLEX-π aantoont dat een WAM geometrische en semantische gronding in essentie gratis kan verkrijgen. Door gebruik te maken van de latente ruimte van een vooraf getraind videomodel en cross-modality forcing toe te passen, bereikt het model:

• Betere Beleidsperformance: Aanzienlijk hogere succesratio's in complexe, contactrijke en langdurige taken.
• Geen Hardware-overhead: Geen behoefte aan extra sensoren naast RGB.
• Geen Pre-training Overhead: Geen noodzaak om nieuwe visuele priors te trainen.
• Implementatie Flexibiliteit: De mogelijkheid om te draaien in een snelle, action-only modus of een hoog-presterende, full-joint modus met hetzelfde checkpoint.

De auteurs concluderen dat hoewel FLEX-π zeer effectief is, het nog steeds meer fine-tuning epochs vereist om te convergeren dan simpelere modellen, en dat de full-joint modus trager is dan parameter-vergelijkbare VLAs. Ze suggereren dat toekomstige verbeteringen kunnen voortkomen uit het combineren van FLEX-π met sterkere VLM-backbones en toegang tot nog grotere robotdatasets.