WAM-OPD: On-Policy Distillation for World Action Models

Technische Samenvatting: WAM-OPD

Probleemstelling
World Action Models (WAMs) integreren visuele toekomstvoorspelling met robotactiegeneratie, wat een gestructureerde aanpak biedt voor robotbesturing. Echter, versnelde "student"-modellen afgeleid van deze systemen via distillatie (bijv. Flash-WAM) lijden vaak aan een verlies van taakcapaciteiten. Deze degradatie treedt op omdat versnelde studenten tijdens implementatie te maken kunnen krijgen met toestanden die slecht gerepresenteerd zijn in de offline data die gebruikt is voor hun initiële training. Hoewel online reinforcement learning (RL) dit kan aanpakken, leunt het op ijle binaire beloningen (sparse binary rewards), waardoor krediettoewijzing over lange trajecten kostbaar en moeilijk is. Bestaande on-policy distillatie (OPD) methoden, waarbij een student de toestandsverdeling bepaalt terwijl een docent (teacher) dichte supervisie biedt, zijn toegepast op actie-gebaseerde policies of autoregressieve tokenmodellen, maar zijn niet aangepast aan de specifieke architectuur van video-eerst WAMs.

Een kritieke technische uitdaging bij het toepassen van OPD op WAMs is een conditioneringsmismatch. In een video-eerst WAM is de actiegeneratie geconditioneerd op een voorspeld videoplann ($z_t$). Tijdens implementatie genereert de student zijn eigen videoplann ($z_S$) en conditioneert de actiebranch op dit plan. Als een bevroren docent de geschiedenissen van de student labelt met zijn eigen superieure videoplann ($z_T$), creëert het trainen van de actiebranch van de student op $z_T$ een mismatch met de implementatieconditie ($z_S$). Bovendien, omdat de video- en actiepaden Transformer-backboneblokken delen, kan het updaten van de ene modaliteit onbedoeld de representaties wijs maken die voor de andere modaliteit worden gebruikt.

Methodologie
Het artikel introduceert WAM-OPD, een deployment-consistente post-training recept die ontworpen is om deze problemen op te lossen zonder dat daarvoor ijle (sparse-reward) RL nodig is. De kernmethodologie omvat:

1. Deployment-Consistente Datageneratie: De versnelde student handelt in de omgeving en genereert een geschiedenishistorie ($h \sim d_{\pi_S}$) die zijn werkelijke closed-loop gedrag weerspiegelt.
2. Bevroren Docent-Labeling: Een bevroren, tragere "Teacher"-model (gebaseerd op LingBot-VA) labelt deze door de student gegenereerde geschiedenissen met coherente video ($z_T$) en actie ($a_T$) targets.
3. Deployment-Consistente Trainingsgrafiek: Cruciaal is dat tijdens de trainingsoptimalisatie de student zijn exacte deployment-computatiegrafiek volgt:
   • Het genereert eerst zijn eigen videoplann: $z_S = f^{vid}_{\theta}(h_S)$.
   • Vervolgens voorspelt het acties geconditioneerd op zijn eigen plan (waarbij een stop-gradient wordt toegepast op de plan-tensor): $a_S = f^{act}_{\theta}(h_S, \text{sg}(z_S))$.
   • Dit zorgt ervoor dat de actiebranch leert opereren onder het videoplann dat het daadwerkelijk zal gebruiken tijdens inferentie, waardoor de conditional interface mismatch wordt vermeden.
4. Gezamenlijke Doelstelling (Joint Objective): De trainingsdoelstelling combineert drie componenten die worden bijgewerkt via lichtgewicht rank-8 JointLoRA adapters over de gedeelde 30-block Transformer backbone:
   • Video Loss: Een pseudo-Huber penalty die het door de student gegenereerde videoplann ($z_S$) uitlijnt met het target van de docent ($z_T$).
   • Action Loss: Een pseudo-Huber penalty die de actie van de student ($a_S$) uitlijnt met de target van de docent ($a_T$).
   • Flow-Matching Regularisator: Een hulpverlies bij de high-noise boundary om flow-matching leer-signalen voor de actiestroom te behouden.
   • De totale loss is: $L(\theta) = \lambda_z \rho_\delta(z_S - z_T) + \lambda_a \rho_\delta(a_S - a_T) + \lambda_{FM} \|v^{act}_\theta(x^a_{\sigma=1}, 1, c) - (\epsilon_a - a_T)\|^2_2$.

Belangrijkste Bijdragen

• Identificatie van Conditioneringsmismatch: De auteurs wijzen een specifieke foutmodus aan in video-eerst WAM OPD, waarbij het trainen van acties op een videoplann van een docent een mismatch creëert met de implementatieconditie van de student.
• Joint Video-Action Loss en LoRA Update: Zij definiëren een gezamenlijke doelstelling en een efficiënte parameter-update strategie (JointLoRA) voor gedeelde, flow-gebaseerde WAMs die de deployment-grafiek respecteert en tegelijkertijd video-uitlijning gebruikt om de kloof tussen docent- en studentplannen te verkleinen.
• Empirische Validatie: Het artikel biedt exact-gepaarde evaluaties op twee RoboTwin 2.0 taken, waarmee wordt aangetoond dat hetzelfde recept de succespercentages in beide settings verbetert, zij het met variërende magnitudes.

Resultaten
De auteurs evalueerden WAM-OPD op twee taken in de RoboTwin 2.0 simulatieomgeving: HANDOVER MIC (het overhandigen van een microfoon tussen robotarmen) en PUT OBJECT CABINET (een object in een lade plaatsen). Gebruikmakend van een uitgebrachte één-video/één-actie-stap Flash-WAM als baseline:

• HANDOVER MIC: Het succespercentage verbeterde van 0.0% (baseline) naar 58.3%.
• PUT OBJECT CABINET: Het succespercentage verbeterde van 16.7% (baseline) naar 33.3%.

Deze resultaten werden behaald met acht student-trajecten, 160 door de docent gelabelde contexten en 120 optimizer stappen, geëvalueerd op zes vastgestelde scene-seeds onder twee vaste noise banks.

Betekenis en Claims
Het artikel positioneert WAM-OPD als een veelbelovende post-training interface voor video-eerst WAMs, in plaats van een voltooide, universeel generaliseerbare methode. De auteurs stellen expliciet dat de resultaten een "initieel bewijs van capaciteit" vormen en geen bewijs voor brede of uniforme generalisatie.

De betekenis ligt in het aantonen dat dichte docent-supervisie op student-geïnduceerde geschiedenissen effectief de taakcapaciteiten in versnelde WAMs kan herstellen zonder terug te vallen op ijle (sparse-reward) RL. De auteurs benadrukken dat de variërende verbeterings-magnitudes tussen taken de huidige taakafhankelijkheid onderstrepen. Zij erkennen beperkingen, waaronder de kleine steekproefomvang, het gebruik van een vast traject-pakket (dat alleen strikt on-policy is voor de initiële collector), en het gebrek aan met rekenkracht gematchte vergelijkingen met andere methoden zoals SFT of RL. Het werk dient als een verticale doorsnede die de levensvatbaarheid van deployment-consistente distillatie voor deze specifieke modelfamilie vaststelt.