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

1. Problem Statement

In embodied agent pipelines (specifically within text-based environments like ScienceWorld), a critical bottleneck exists at the decision-making step: determining which candidate actions are physically feasible given the current state.

• The Challenge: Agents often face large sets of valid actions where many are semantically plausible but physically incorrect (e.g., "cool the water" vs. "heat the water" in a thermodynamics task).
• Limitations of Current Approaches: Standard Supervised Fine-Tuning (SFT) trains models to classify actions independently as valid or invalid. This approach fails to explicitly teach the model to discriminate between a correct action and "hard negatives"—actions that are syntactically similar but physically opposite or incompatible. Consequently, SFT models often struggle to rank the correct action above subtle, physically incorrect distractors.

2. Methodology: Contrastive World Model (CWM)

The authors propose CWM, a framework that reframes action scoring as a contrastive ranking problem rather than a binary classification task.

• Core Architecture:
  • Base Model: Fine-tunes Qwen-2.5-7B-Instruct using LoRA (Low-Rank Adaptation) with rank $r=8$.
  • Scoring Function: $f_\theta(s, a)$ outputs the log-likelihood of an action $a$ given a state $s$.
• Training Objective (InfoNCE Loss):
  • Instead of binary cross-entropy, CWM uses an InfoNCE contrastive loss.
  • For a given state $s$, the model is trained to assign a higher score to the positive action ($a^+$) than to a pool of $N$ negative actions ($a^-_i$) simultaneously.
  • Loss Function: $L_{CWM} = -\log \frac{\exp(f(s, a^+)/\tau)}{\exp(f(s, a^+)/\tau) + \sum \exp(f(s, a^-_i)/\tau)}$.
  • Regularization: A margin regularization term is added to ensure a minimum score gap between the positive and the average of negatives.
• Hard Negative Mining Strategy:
  The training data includes a structured pipeline for mining negatives from interaction logs, categorized by difficulty:
  1. Type 1 (Silent Failures): Actions that do nothing (e.g., pushing an immovable object).
  2. Type 2 (Explicit Rejections): Actions the environment explicitly refuses.
  3. Type 3 (Cross-task Transplants): Actions valid in other contexts but wrong here.
  4. Type 4 (Minimal-Edit Negatives - Evaluation Only): Distractors created by changing a single physics-critical word (e.g., "heat" $\to$ "cool"). These are the hardest cases.
  • Configuration: The model is trained with $N=16$ hard negatives per positive sample, significantly increasing the contrastive signal compared to baselines.

3. Key Contributions

1. CWM Framework: A novel method for fine-tuning LLMs as action scorers using InfoNCE loss with a structured hard-negative mining pipeline.
2. Hard Negative Taxonomy: A classification of negative examples for ScienceWorld, specifically defining "minimal-edit" negatives that test physical reasoning boundaries.
3. Empirical Validation: Two distinct studies demonstrating that contrastive training induces representations that capture physical affordances more faithfully than SFT, particularly in out-of-distribution (OOD) scenarios.

4. Experimental Results

The authors evaluated CWM against an SFT baseline, a Vanilla LLM (zero-shot), and a Random baseline on the ScienceWorld benchmark.

Study 1: Intrinsic Affordance Evaluation

• Dataset: 605 hard-negative test pairs, including a "Minimal-Edit" category (74 instances).
• Key Findings:
  • Precision@1 (P@1): CWM achieved 93.24% on minimal-edit negatives, outperforming SFT (86.49%) by +6.76 percentage points.
  • AUC-ROC: CWM achieved 0.929 vs. SFT's 0.906 on minimal-edit tasks.
  • Score Calibration: While SFT produced larger raw score gaps, CWM achieved better ranking accuracy. This indicates SFT scores are miscalibrated across the negative pool, whereas CWM's normalization produces a more reliable relative ordering.
  • Vanilla LLM: Performed worse than random chance on minimal-edit tasks (48.6%), highlighting that untuned LLMs are easily misled by surface-level similarity in physics tasks.

Study 2: Live Filter Characterization (Agent Loop)

• Metric: Gold Action Retention Rate (GARR@K) and Safety Margin (score difference between gold action and the best non-gold action).
• In-Domain Performance: CWM retained the gold action in the top-10 candidates 60.2% of the time.
• Out-of-Distribution (OOD) Stress Test:
  • On unseen task families (thermometer, chemistry-mix), CWM maintained a significantly better Safety Margin (-2.39) compared to SFT (-3.96).
  • Interpretation: A less negative margin indicates the gold action is ranked closer to the top. SFT's scores collapsed under OOD stress, ranking the correct action much lower than incorrect alternatives, whereas CWM's contrastive training preserved the relative ordering.

5. Significance and Conclusion

• Theoretical Insight: The paper challenges the prevailing view (from general IR literature) that SFT outperforms contrastive learning for LLMs. It demonstrates that in physics-grounded embodied settings, where negatives have principled physical structure and are semantically close to positives, contrastive objectives (InfoNCE) are superior.
• Practical Impact:
  • CWM serves as a robust pre-filter for embodied agents. By pruning the action space to the top-K candidates (e.g., K=10), an agent can reduce its search space by ~80-90% while retaining the correct action in the majority of steps.
  • It significantly reduces the risk of "self-inflicted dead ends" caused by executing physically impossible but syntactically plausible actions.
• Future Work: The authors plan to integrate CWM into full agent loops (e.g., combining with ReAct or DRRN) to measure end-to-end task completion rates and explore causal chain negatives.

In summary, CWM proves that teaching an LLM to rank actions against a pool of hard negatives is more effective than teaching it to classify them independently, leading to more reliable physical reasoning in embodied AI.