Skip to content

Language-conditioned world model improves policy generalization by reading environmental descriptions

🕒 Published (v1): 2025-11-28 06:13 UTC · Source: Arxiv · Venue: NEURIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LED-WM extends the DreamerV3 model-based RL framework with a cross-modal attention encoder that explicitly grounds language descriptions of environment dynamics to symbolic entities in a grid observation. Policies trained on the resulting language-conditioned world model generalize to unseen games with novel entity roles, movement behaviors, and paraphrased language—without requiring expert demonstrations or inference-time planning. On MESSENGER-WM, LED-WM outperforms EMMA-LWM across all compositional generalization settings despite the latter using expert supervision.

Problem

Agents that condition on language describing environment dynamics (e.g., "the floor is slippery") rather than task instructions must generalize to games with unseen entity-role and movement combinations. Existing model-based approaches either require expert demonstrations (EMMA-LWM) or expensive MCTS inference-time planning (Reader), while model-free baselines (Dynalang) fail to generalize to unseen dynamics because they lack explicit language-to-entity grounding.

Method

LED-WM replaces DreamerV3's observation encoder with a Language-aware Encoder for Dreamer (LED):

  1. Language encoding: A frozen T5 encoder produces sentence embeddings for each entity description in the language manual.
  2. Entity encoding: Each entity symbol in the 10×10 grid gets a learned embedding; its temporal position history relative to the agent is captured via a dot-product alignment score (Eq. 2) tracking movement direction vs. agent-relative direction.
  3. Cross-modal attention: Per-entity queries are formed by an MLP over [entity_embedding, temporal_history]; keys/values come from sentence embeddings. Scaled dot-product attention aligns each entity to its description sentence.
  4. Language-grounded grid: Attention output vectors are placed back at each entity's spatial position, producing a language-aware grid G_l, then processed by a CNN.
  5. The final representation feeds into DreamerV3's RSSM; the decoder is dropped and multi-step reward/continue prediction is adopted. Policy is learned from world-model rollouts—no planning, no demonstrations.
  6. Fine-tuning via synthetic rollouts: Given a trained LED-WM and a test game, 60 synthetic trajectories are generated and used to fine-tune the policy if estimated value falls below a threshold.

Key Contributions

  • Demonstrates that Dynalang (MBRL without explicit entity grounding) fails to generalize to unseen MESSENGER games (win rates: S1 0.03, S2 0.04, S3 0.03 vs. LED-WM's 100/51.6/35.0).
  • Proposes LED-WM, an attention-based language-conditioned world model that achieves strong policy generalization without expert demonstrations or inference-time planning.
  • Shows that LED-WM's world model generalizes to unseen dynamics: fine-tuning on synthetic test trajectories improves policy performance (S2-dev improvement statistically significant at 95% confidence; S3 mean improves from −0.11 to −0.01).
  • Outperforms EMMA-LWM on all three MESSENGER-WM compositional settings (NewCombo: 1.31 vs. 0.98; NewAttr: 1.15 vs. 0.29; NewAll: 1.16 vs. 0.13).

Results

  • MESSENGER S1: LED-WM 100% win rate vs. EMMA (w/ curriculum) 88%, CRL 85%.
  • MESSENGER S2: LED-WM 51.6% vs. CRL 95% (CRL wins by addressing training-set bias); Dynalang 0.04%.
  • MESSENGER S2-dev: LED-WM 96.6% vs. EMMA (w/ curriculum) 45%.
  • MESSENGER S3: LED-WM 35.0% vs. EMMA (w/ curriculum) 22%, CRL 10%.
  • MESSENGER-WM NewCombo/NewAttr/NewAll: LED-WM 1.31/1.15/1.16 vs. best EMMA-LWM variant 0.98/0.29/0.62 (avg. sum of scores).
  • Fine-tuning on synthetic rollouts: S2-dev improves from 1.4478 → 1.4513; S3 from −0.11 → −0.01.

Limitations

  • Underperforms CRL on MESSENGER S2, where the training set has only a single entity-movement combination and spurious identity-role correlations; LED-WM lacks an explicit bias-mitigation mechanism.
  • Absolute policy improvement from synthetic fine-tuning is modest, especially on S3.
  • Evaluation confined to small, discrete 10×10 grid-world environments (MESSENGER/MESSENGER-WM); generalization to continuous or visually rich settings is not demonstrated.
  • Relies on the MESSENGER benchmark's structured entity-role language templates; robustness to fully free-form natural language dynamics descriptions is untested.
  • Tested only with MESSENGER S2 dynamics (chasing/fleeing/stationary); more diverse training dynamics may be needed for harder compositional splits.

Relevance to Agentic AI / LLM Agents

This work directly addresses how agents can exploit dynamics-descriptive natural language—not just goal instructions—to adapt behavior to unseen environments, a core requirement for generalizable real-world agents. The LED-WM architecture demonstrates that injecting language grounding into the world model (rather than only the policy) enables compositional generalization without costly inference-time planning or expert demonstrations, lowering deployment barriers. The fine-tuning-via-synthetic-rollouts capability shows a viable path for pre-deployment world-model adaptation, relevant to sim-to-real transfer in embodied and robotic agents. For researchers tracking LLM-grounded agents, this paper sits at the intersection of MBRL, language-conditioned world models, and compositional generalization—extending the DreamerV3 lineage toward language-aware agents.