CWM: Contrastive World Models for Action Feasibility Learning in Embodied Agent Pipelines

Probleemstelling

In de architectuur van embodied agents (fysieke of gesimuleerde agenten die handelingen uitvoeren) is een betrouwbare actie-feasibility scorer (een module die bepaalt welke acties fysiek uitvoerbaar zijn) een kritieke bottleneck. Voordat een agent kan plannen of redeneren, moet het weten welke kandidaat-acties fysiek mogelijk zijn in de huidige staat.

Bestaande benaderingen gebruiken vaak Supervised Fine-Tuning (SFT) om deze scorers te trainen. De beperking van SFT is dat het elke kandidaat-actie onafhankelijk behandelt (binair: geldig of ongeldig). Dit leert het model niet expliciet om te discrimineren tussen acties die fysiek correct zijn en acties die subtiel verkeerd zijn (bijvoorbeeld acties die semantisch vergelijkbaar zijn maar fysiek onmogelijk). Hierdoor verspilt een agent stappen aan het testen van onmogelijke interacties of loopt het vast in doodlopende paden.

Methodologie: Contrastive World Model (CWM)

De auteurs stellen CWM voor, een methode om een Large Language Model (LLM) te fine-tunen als een actiescorer met behulp van een InfoNCE contrastieve loss-functie, in plaats van standaard SFT.

• Kernidee: Het model wordt getraind om een geldige actie ($a^+$) in de score-ruimte verder weg te duwen van ongeldige acties ($a^-$), met speciale nadruk op "hard negatives". Hard negatives zijn semantisch zeer vergelijkbare acties die fysiek incompatibel zijn (bijv. "verwarm water" vs. "koel water").
• Model Architectuur: Ze fine-tunen het Qwen-2.5-7B model met LoRA (Low-Rank Adaptation). De scorer $f_\theta$ geeft een score op basis van de log-likelihood van een actie gegeven een staat.
• Trainingsdoel (Loss Function):
  • InfoNCE Loss: Het model moet de positieve actie hoger scoren dan een pool van $N$ negatieve acties tegelijkertijd (softmax-genormaliseerd), in plaats van paren onafhankelijk te classificeren.
  • Margin Regularisatie: Een extra term zorgt voor een minimum scorekloof tussen de positieve actie en het gemiddelde van de negatieve acties.
• Hard Negative Mining: De auteurs gebruiken een gestructureerde pipeline om negatieve voorbeelden te selecteren uit interactielogs van de ScienceWorld omgeving. Ze onderscheiden vier typen:
  1. Silent failures: Acties die niets doen (makkelijk).
  2. Explicit rejections: Acties die de omgeving expliciet weigert (gemiddeld moeilijk).
  3. Cross-task transplants: Acties die in andere taken geldig zijn, maar hier niet passen (moeilijk).
  4. Minimal-edit negatives (Type 4): Acties met één woordverandering die het fysieke resultaat omkeren (zeer moeilijk). Voor training worden 16 negatieve voorbeelden per positief voorbeeld gebruikt.

Belangrijkste Bijdragen

1. CWM Architectuur: Een nieuwe aanpak voor het fine-tunen van LLMs als actiescorers via InfoNCE loss met een gestructureerde pipeline voor het vinden van harde negatieve voorbeelden.
2. Taxonomie van Harde Negatieven: Een classificatie van negatieve voorbeelden in ScienceWorld, met name de focus op "minimal-edit" acties die fysieke tegenstellingen simuleren.
3. Empirische Validatie: Twee uitgebreide studies die aantonen dat contrastief leren superieur is aan SFT in fysiek gestructureerde omgevingen.

Resultaten

De evaluatie vond plaats op de ScienceWorld benchmark (tekstgebaseerde simulaties van wetenschappelijke taken).

1. Intrinsieke Affordance Evaluatie (Harde Negatieve Test):

• Op een testset van 605 paren met minimale edits (één woordverandering) behaalde CWM een Precision@1 van 93,24%, vergeleken met 86,49% voor SFT (+6,76 procentpunten).
• De AUC-ROC was hoger voor CWM (0,929) dan voor SFT (0,906).
• Analyse: Hoewel SFT soms een grotere absolute scorekloof produceerde, was de rangschikking van CWM betrouwbaarder. SFT faalde vaak bij het onderscheiden van fysiek tegenovergestelde acties die er oppervlakkig hetzelfde uitzagen.

2. Filterkarakterisering in Live Agent Loop:

• In een live scenario werd gemeten hoe vaak de "gouden" (expert) actie binnen de top-K gekozen acties bleef.
• Out-of-Distribution (OOD) Stress Test: Op taken die niet in de training zaten, presteerde CWM beter in termen van Safety Margin (het verschil in score tussen de beste actie en de gouden actie).
  • CWM Safety Margin: -2,39
  • SFT Safety Margin: -3,96
• Een minder negatieve waarde betekent dat de gouden actie dichter bij de top staat. Dit toont aan dat CWM robuuster is wanneer de actie-ruimte dicht en onbekend is; SFT neigt hier de gouden actie lager te rangschikken ten opzichte van concurrerende, maar onjuiste, opties.

Betekenis en Conclusie

Het paper concludeert dat voor het scoren van acties in fysiek gestructureerde omgevingen de trainingsdoelstelling cruciaal is.

• SFT vs. Contrastief Leren: Waar SFT goed werkt voor duidelijk onderscheidbare acties, faalt het bij "harde" negatieven die fysiek dicht bij de waarheid liggen. De InfoNCE loss dwingt het model om een relatieve rangschikking te leren binnen een pool van kandidaten, wat leidt tot representaties die fysieke haalbaarheid (affordances) nauwkeuriger weergeven.
• Toepassing: CWM fungeert als een effectieve pre-filter voor planningsalgoritmen. Het kan de zoekruimte voor een agent met 80-90% verkleinen terwijl de juiste actie in de meeste gevallen behouden blijft, zelfs in nieuwe, onbekende taken.
• Toekomst: De auteurs plannen de integratie van CWM in volledige agentenloops (bijv. gecombineerd met Reinforcement Learning of ReAct) om de totale taakvoltooiing te verbeteren.

Kortom, CWM biedt een fundamentele verbetering voor embodied agents door het vermogen om fysieke onmogelijkheden te onderscheiden van fysieke mogelijkheden, zelfs wanneer het verschil slechts één woord bedraagt.