Language-conditioned world model improves policy generalization by reading environmental descriptions (2511.22904v1)

Abstract: To interact effectively with humans in the real world, it is important for agents to understand language that describes the dynamics of the environment--that is, how the environment behaves--rather than just task instructions specifying "what to do". Understanding this dynamics-descriptive language is important for human-agent interaction and agent behavior. Recent work address this problem using a model-based approach: language is incorporated into a world model, which is then used to learn a behavior policy. However, these existing methods either do not demonstrate policy generalization to unseen games or rely on limiting assumptions. For instance, assuming that the latency induced by inference-time planning is tolerable for the target task or expert demonstrations are available. Expanding on this line of research, we focus on improving policy generalization from a language-conditioned world model while dropping these assumptions. We propose a model-based reinforcement learning approach, where a language-conditioned world model is trained through interaction with the environment, and a policy is learned from this model--without planning or expert demonstrations. Our method proposes Language-aware Encoder for Dreamer World Model (LED-WM) built on top of DreamerV3. LED-WM features an observation encoder that uses an attention mechanism to explicitly ground language descriptions to entities in the observation. We show that policies trained with LED-WM generalize more effectively to unseen games described by novel dynamics and language compared to other baselines in several settings in two environments: MESSENGER and MESSENGER-WM.To highlight how the policy can leverage the trained world model before real-world deployment, we demonstrate the policy can be improved through fine-tuning on synthetic test trajectories generated by the world model.

• The paper introduces LED-WM, which integrates dynamics-descriptive language via entity-centric attention to enhance policy generalization.
• It employs a frozen T5 encoder and cross-modal alignment to map language to entity embeddings, ensuring robust performance on compositional grid-world benchmarks.
• Experimental results show strong zero-shot and few-shot generalization across novel linguistic formulations and dynamic configurations in MESSENGER and MESSENGER-WM.

Language-Conditioned World Models for Policy Generalization via Entity-Descriptive Language

Problem Motivation and Overview

The research addresses the integration of dynamics-descriptive language into model-based reinforcement learning (MBRL) agents. While traditional language-conditioned agents focus on linguistic task specifications (“what to do”), this work shifts focus towards environmental descriptions—language that explicates how entities behave, interact, and evolve over time (“how the environment behaves”). The practical motivation is the deployment of autonomous agents capable of adapting to novel or dynamically changing environments, relying solely on language manuals for up-to-date domain knowledge.

The core research question is: To what extent does explicitly grounding descriptive language to environmental entities in a world model facilitate out-of-distribution (OOD) policy generalization? Prior work on language-conditioned world models either bypasses compositional grounding or makes restrictive assumptions, such as accepting substantial inference-time planning latency or requiring extensive expert demonstrations. This research removes those assumptions and demonstrates a model-based RL approach combining entity-language grounding with sample-efficient policy learning.

Environment and Generalization Benchmarks

The paper utilizes two compositional grid-world domains—MESSENGER and MESSENGER-WM—specifically designed for language grounding and generalization assessment.

• MESSENGER comprises multiple generalization settings, from unseen language surfaces to unseen entity-role-movement configurations. The agent navigates a $10 \times 10$ grid with entities (e.g., planes, ferries, researchers) and must interpret a natural language manual to correctly infer roles (messenger, goal, enemy) and select a success trajectory.
• MESSENGER-WM extends evaluation to compositional generalization, requiring the agent to handle entirely novel combinations and attribute assignments of entities, movements, and roles.

The taxonomy of generalization includes:

• Novel language: Synonyms or paraphrasing in manuals.
• Novel entity-role/movement assignments: Entities assuming previously unseen combinations.
• Novel game dynamics: Unseen movement patterns or compositional arrangements.

  Figure 1: An example of dynamics-descriptive language in MESSENGER. The manual describes how each entity moves and interacts, which must be interpreted and grounded to the symbolic observation for successful agent behavior.

LED-WM: Language-Aware Encoder for Dreamer World Model

The main technical contribution is the Language-aware Encoder for Dreamer World Model (LED-WM), built atop DreamerV3. Critically, LED-WM employs an entity-centric attention mechanism to align sentence embeddings from the language manual with entity embeddings in the observation. This mechanism ensures that dynamics-descriptive language is not only incorporated, but explicitly grounded spatially and temporally to the corresponding entities.

Input Processing Breakdown:

• Language manuals are encoded via a frozen T5 encoder, yielding per-entity sentence embeddings.
• Each entity (and the agent) is represented as a learned symbol embedding; time-step information and movement history are incorporated via auxiliary embeddings and temporal arrays.
• For each entity, an MLP maps entity identity and dynamics to a query, which is then aligned via attention to sentence (manual) keys and values.

The aligned language-conditioned entity representations are spatially placed back into their grid positions, forming a language-aware grid. This grid is further processed by a CNN, whose output is concatenated with the time-step embedding to provide a global observation feature vector for the world model.

Figure 2: LED-WM architecture. Entities, language sentences, and dynamics are encoded and aligned via cross-modal attention, producing a language-grounded spatial grid passed through a CNN for feature extraction.

This modification replaces the standard DreamerV3 encoder, ensuring the world model’s latent states are directly dependent on entity-level grounded descriptions.

Experimental Results

Policy Generalization

The primary evaluation considers the zero-shot and few-shot generalization of policies trained from the world model to unseen test games (with new dynamics or new linguistic formulations). Main baselines include EMMA, CRL (model-free with attention or constraints), Dynalang (language-conditioned model-based, no explicit grounding), and EMMA-LWM (world-model-based with behavior cloning from demonstrations).

Key numerical findings:

• MESSENGER (Stage S1): LED-WM achieves a perfect win rate ($100\%$), outperforming all baselines.
• S2/S2-dev: LED-WM achieves strong generalization ($51.6 \pm 2.7\%$ S2; $96.6 \pm 1.0\%$ S2-dev), competitive with or exceeding baselines; performance drop in S2 is attributed to data bias, where CRL’s explicit debiasing yields an edge.
• S3: LED-WM achieves $34.97 \pm 1.73\%$ win rate, exceeding CRL and EMMA (no curriculum) baselines, demonstrating flexible entity-role-movement grounding.
• MESSENGER-WM: LED-WM surpasses EMMA-LWM (which requires expert demonstrations) in all generalization splits, with average sum of scores: NewCombo: 1.31 ± 0.05, NewAttr: 1.15 ± 0.08, NewAll: 1.16 ± 0.02 (versus best EMMA-LWM: $1.18 \pm 0.10$, $0.96 \pm 0.17$, $0.62 \pm 0.21$).

  Figure 3: Example MESSENGER S1 where successful policy depends exclusively on language-based grounding of entity roles; static entities require symbolic rather than behavioral mapping.

  

  Figure 4: Example MESSENGER S2 where the agent must resolve entity roles via both names and observed movement, evidencing the necessity for joint symbolic and behavioral grounding.

  

  Figure 5: Example MESSENGER S3 requiring disambiguation of roles among entities with identical names but differing movement dynamics, highlighting compositional complexity.

World Model Generalization and Policy Fine-tuning

The world model’s generalization is probed by synthetic rollout-based policy fine-tuning on OOD tasks. If the model can correctly simulate unseen game dynamics, policy fine-tuning on rollouts should yield score improvements.

• S2-dev: Policy fine-tuning via LED-WM’s rollouts yields measurable, statistically significant improvements in average score, supporting the claim that the world model learns a compositional, language-grounded dynamics prior.
• S3: Finetuning narrows the gap to optimal policy, confirming latent space generalization.

Ablations and Comparative Analysis

The paper confirms that world models without explicit entity-language grounding (Dynalang) fail to generalize, while LED-WM captures necessary compositional inductive biases through its cross-modal matching. Nevertheless, in problem settings with extreme data imbalance (e.g., a single dynamics combination in training), further mitigation (as in CRL) is warranted.

Figure 6: Overview of different generalization regimes in MESSENGER and MESSENGER-WM, illustrating the spectrum from surface-level language novelty to full compositional recombination.

Theoretical and Practical Implications

This research establishes that explicit entity-level grounding of language in MBRL architectures substantially enhances generalization to unseen entity configurations and novel dynamics linguistic descriptions. Practically, this enables training of agents on compact compositional curricula, while retaining robust performance in significantly altered task regimes—critical for real-world deployment with non-stationary language and domain dynamics.

On a theoretical level, the design advocates for cross-modal attention mechanisms not only at training time but also as architectural priors, aligning structured language and state observations hierarchically.

Future work directions include:

• Investigating robustness to richer visual inputs and continuous spaces,
• Extending language grounding mechanisms to large multimodal foundation models,
• Addressing spurious correlation and data bias with principled regularization,
• Scaling entity-linguistic grounding within hierarchical world modeling frameworks.

Conclusion

LED-WM, through explicit entity-language attention within a model-based RL paradigm, achieves strong and often superior generalization performance in compositional language-conditioned environments, without requiring expert demonstrations or inference-time planning. The approach is empirically validated against established baselines, both in terms of policy and world model generalization, and highlights the central role of explicit language grounding in advancing generalizable decision-making agents.