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

TL;DR

This work tackles the challenge of learning from dynamics-descriptive language to improve policy generalization in interactive environments. It introduces LED-WM, a language-aware encoder built on DreamerV3 that grounds environmental descriptions to entities via cross-modal attention, enabling a language-conditioned world model to train policies without planning or expert demonstrations. Empirical results show that LED-WM improves generalization to unseen games in both MESSENGER and MESSENGER-WM, outperforming several baselines in many settings, and its performance can be further enhanced through finetuning on synthetic trajectories generated by the world model. The study demonstrates the value of grounding language to entities within a world-model-based RL framework and highlights directions for improving compositional and dynamic generalization in language-conditioned agents.

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.

Figures (6)

• Figure 1: An example of dynamics-descriptive language in a game play. The observation includes a 10 $\times$ 10 grid-world with three entities represented by their associated symbols: (ferry - ), (plane - ), (researcher - ) and one agent (depicted by ). The observation also has a manual on the right, which describes the dynamics of the game. The agent can navigate the grid using five actions: left, right, up, down, and stay. The agent can only interact with entities when it is in the same grid cell as the entity. The agent's task is to identify roles of all entities from the manual, go to the messenger, then go to the goal, while avoiding the enemy. Shaded icons indicate one possible scenario of entity movement over time. By observing entity movement patterns and grounding language to entities based on their behaviors, the agent can infer the roles assigned to each entity: (ferry-enemy), (plane-messenger), and (researcher-goal). The agent can then execute an appropriate plan to complete the task. The dashed line in the grid shows such a possible plan.
• Figure 2: Overview of our proposed world model LED-WM. The world model input consists of: a language manual $L$, a grid-world observation representing entity and agent symbols, and the current time step $t$. Entity, agent symbols, and time step are encoded using learned embeddings, while $L$ is encoded via a frozen T5 encoder. To represent each entity, we employ a multi-layer perceptron (MLP) that processes the entity embedding and its temporal information, capturing its movement pattern relative to the agent, to produce a query vector. We apply an attention network between the query vectors and the sentence embeddings to align each entity with its corresponding sentence. The resulting vectors are then put into their respective entity positions. This produces a language-grounded grid $G_l$, which is then processed by a CNN. The extracted feature vector is flattened and concatenated with the time embedding to form final observation representation $x_t$.
• Figure 3: An example game of MESSENGER S1. In this game, the entity does not have message at the beginning of the game. Therefore, it goes to the messenger to retrieve the message and ends the game. All entities except the agent are stationary, thus the manual only describes roles associated with entity names.
• Figure 4: An example of a game play within a 10 × 10 grid-world from MESSENGER S2. The observation on the left includes three entities represented by their associated symbols: (ferry - ), (plane - ), (researcher - ) and one agent (depicted by ). The game involves three roles: messenger, goal, and enemy. The agent's task is to identify roles of all entities, locate the messenger, deliver it to the goal, and avoid the enemy. To achieve this objective, the agent must use the manual to infer entity roles based on their described dynamics and observed behavior. In the observation in the example, shaded icons indicate one possible scenario of entity locations over time. By observing entity movement patterns and grounding language to entities based on their according behaviors, the agent can infer the roles are assigned: (ferry-enemy), (plane-messenger), and (researcher-goal). After inferring all entity roles, the agent can execute an appropriate plan to complete the task. The dashed line in the observation shows such a possible plan.
• Figure 5: An example game of MESSENGER S3. To win the game, the agent must infer the roles of entities given the manual. Specifically, the same entity names (e.g. airplane, plane with different roles (e.g. enemy, messenger) must be disambiguated by their movement dynamics (e.g. chasing, fleeing). Note that we have a italicized sentence describing an extraneous entity that is not available in the game observation. We also have synonyms for entity names and roles, e.g., airplane, plane; adversary, enemy. The shaded entities show possible entity locations over time and the dashed line shows a possible path for the agent to win the game.