PaperlayerView exact revision on arXiv

Source-linked AI summary

Language-Guided World Models: A Model-Based Approach to AI Control

Alex Zhang, Khanh Nguyen, Jens Tuyls, Albert Lin, Karthik Narasimhan

arXiv:2402.01695v3cs.CLcs.AIcs.LG

TL;DR

The paper addresses the challenge of grounding language in dynamics well enough to simulate compositionally novel environments. It evaluates language-conditioned trajectory generation with an EMMA-augmented Transformer. That model generalizes in the hardest setting, outperforming baselines and competing with oracle semantic parsing and grounding.

  • Problem

    Recovering dynamics for compositionally novel environments is challenging because parameters are independent and entity-local, complicating exact recovery from observations without inductive biases.

    Source: §3.2 · p. 4

  • Method

    The authors encode language manuals and generate tokenized trajectories with a Transformer augmented by EMMA attention to model environments.

    Source: §1 · p. 2, §4 · p. 5

  • Results

    The EMMA-augmented model generalizes in the hardest setting, substantially outperforms baselines across metrics, and is competitive with oracle semantic parsing and grounding.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    A cautious agent can present plans for human review and revise them from language feedback, illustrating a potential application to safety and transparency.

    Source: §1 · p. 3

    The model still lacks performance, and the grid-world environments severely underrepresent the real world.

    Source: §7 · p. 9

Abstract

from arXiv · show

This paper introduces the concept of Language-Guided World Models (LWMs) -- probabilistic models that can simulate environments by reading texts. Agents equipped with these models provide humans with more extensive and efficient control, allowing them to simultaneously alter agent behaviors in multiple tasks via natural verbal communication. In this work, we take initial steps in developing robust LWMs that can generalize to compositionally novel language descriptions. We design a challenging world modeling benchmark based on the game of MESSENGER (Hanjie et al., 2021), featuring evaluation settings that require varying degrees of compositional generalization. Our experiments reveal the lack of generalizability of the state-of-the-art Transformer model, as it offers marginal improvements in simulation quality over a no-text baseline. We devise a more robust model by fusing the Transformer with the EMMA attention mechanism (Hanjie et al., 2021). Our model substantially outperforms the Transformer and approaches the performance of a model with an oracle semantic parsing and grounding capability. To demonstrate the practicality of this model in improving AI safety and transparency, we simulate a scenario in which the model enables an agent to present plans to a human before execution, and to revise plans based on their language feedback.

1 Introduction

The introduction frames language-steerable world models as a way to make model-based agents more controllable, then addresses the challenge of grounding language and generalizing compositionally.

  • A shared world model lets human changes affect multiple task policies simultaneously, complementing the model-free approach of adapting policies individually.Source: §1 · p. 1
  • Grounding language in environmental dynamics is challenging because descriptions involve rich concepts and may introduce new terms without clear definitions.Source: §1 · p. 2
  • The MESSENGER benchmark tests whether models can learn entity attributes from videos and simulate compositionally novel environments across three increasingly difficult evaluation settings.Source: §1 · p. 2
  • The EMMA-augmented model generalizes in the hardest evaluation setting, substantially outperforms baselines across metrics, and is competitive with oracle semantic parsing and grounding.Source: §1 · pp. 2–3
  • In the most challenging setting, an LWM-based agent achieves average reward three to four times higher than an observational-world-model agent without additional environment interactions.Source: §1 · p. 3
  • The authors advocate modular agents with language-parameterized components and note that LWMs can use human-written texts to reduce later interaction-based fine-tuning.Source: §1 · pp. 1–3

2 Background: world models

This background defines world models and contrasts model-based agents with observational, model-free approaches. It explains how shared models can ease adaptation while observation-based control makes human feedback costly.

  • An MDP specifies states, actions, and transitions; model-based agents use probabilistic world models to approximate dynamics and plan, unlike model-free agents that learn policies directly from data.Source: §2 · p. 3
  • Because task policies share a world model, model-based agents can adapt to changed dynamics using task-agnostic data, while model-free data-collection costs scale with the number of tasks.Source: §2 · p. 3
  • Observational world models condition on interaction histories and adapt through observations, but observations convey abstract intentions poorly and require potentially costly, risky real-world actions.Source: §2 · p. 3

3 Language-guided world models (LWMs)

The section introduces LWMs as world models that use language manuals to model environment dynamics, then formalizes learning from trajectory–manual pairs. It frames compositional generalization in entity-based environments and instantiates it in MESSENGER-WM.

3 Language-guided world models (LWMs)

  • LWMs interpret language descriptions to simulate environment dynamics, allowing people to adapt agent behavior through natural communication and use existing human-written texts.The paper formulates the goal as learning dynamics from manuals paired with trajectories generated in environments parameterized by v.Source: §3 · p. 4, §3.1 · p. 4

3.2 Modeling entity-based environments

  • Compositional novelty means every attribute is individually seen during training, but the full attribute vector is unseen, requiring language-based recovery of independent, entity-local dynamics.The formulation assumes each manual description portrays all attributes of one entity, so the number of descriptions equals the number of entities.Source: §3.2 · p. 4

3.3 The MESSENGER-WM benchmark

  • MESSENGER-WM uses a visually simple grid-world game with entity attributes and varied language descriptions to focus on world-modeling challenges beyond realistic image generation.The benchmark accompanies game trajectories with manuals and evaluates the three increasingly difficult compositional-generalization settings.Source: §3.3 · pp. 4–5
  • MESSENGER-WM tests compositional generalization with NewCombo, NewAttr, and NewAll settings, while its entity dynamics challenge models to learn attribute independence and route attributes to the correct entities.The settings progress from novel identity combinations, to new attributes on familiar identity triplets, to both forms of novelty.Source: §3.3 · p. 5

4 Modeling approach

The approach represents MESSENGER states as entity-disentangled grids and casts world modeling as manual-to-trajectory sequence generation. An EMMA-style attention stage routes language-derived attributes to entity tokens before trajectory decoding.

State representation

  • MESSENGER states use an idealized entity-disentangled grid, but models must still recognize manual attributes and associate them with the correct entity.Learning entity-disentangled representations directly from pixel-based environments remains future work.Source: §4 · p. 5

World modeling as sequence generation

  • The model encodes a language manual and decodes a trajectory as a sequence of tokens, using an encoder-decoder Transformer.Trajectory tuples are converted into token blocks and concatenated across time into one sequence.Source: §4 · pp. 5–6

Model inputs

  • The manual is encoded from individual descriptions, while trajectory blocks represent actions, entity identities and locations, rewards, and termination signals.The manual descriptions are first converted to hidden vectors by pretrained BERT and a Transformer encoder.Source: §4 · pp. 5–6

Entity mapper with multi-modal attention

  • An EMMA variant identifies each entity’s description, extracts its attribute words, and uses dot-product attention to add those features to the entity’s identity token.Identity words supply keys, while words describing other attributes supply values.Source: §4 · p. 6

Model training

  • Training minimizes cross-entropy for next-token prediction but excludes action-token positions and the first block’s tokens because inference supplies them.The targets are the tokenized ground-truth trajectories.Source: §4 · p. 6

5 Experiments

The experiments compare EMMA-LWM with observational, Transformer, ChatGPT-grounded, and oracle models on MESSENGER, then test language-guided plan adaptation. EMMA-LWM leads non-oracle baselines in world-model evaluation and improves returns in the plan-adaptation setting.

5.1 Baselines

  • The baselines include a no-text observational model, a BERT-encoding Transformer, a ChatGPT-grounded GPTHard model, and OracleParse with ground-truth information extraction.The observational baseline zeros out the manual representation; GPTHard uses ChatGPT to ground descriptions to entities.Source: §5.1 · p. 6

5.2 Results

  • EMMA-LWM consistently outperforms all baselines on NewAttr and NewAll ground-truth trajectories, nearing OracleParse; it also leads non-oracle models across imaginary-trajectory metrics.The ground-truth evaluation uses cross-entropy loss; imaginary-trajectory evaluation measures reward and termination prediction and player-to-entity distance similarity.Source: §5.2 · pp. 7–8, Table 1 · p. 7, Table 2 · p. 8
  • Standard Transformer performance varies sharply across initializations, while GPTHard's underperformance is attributed to imperfect identity parsing and EMMA-LWM's joint extraction of identities and attributes.Some Standard runs match EMMA-LWM, while others perform as poorly as Observational; ChatGPT identity accuracy is reported at around 90%.Source: §5.2 · p. 7

5.3 Application: agents that discuss plans with humans

  • The plan-discussion simulation uses randomly initialized policies barred from environment interaction, with unseen-environment world models and imitation-learning feedback.A human can provide action corrections or overwrite plans, and language feedback is supplied once before policy adaptation.Source: §5.3 · p. 8
  • With EMMA-LWM, policies improve average return in every evaluation setting, though considerable gaps remain against OracleParse.The comparison represents language feedback that can modify the world model versus observational-model learning with imitation feedback alone.Source: §5.3 · pp. 8–9, Table 3 · p. 9

6 Related work

Related work spans world models, language-based adaptation, and compositional generalization. This paper distinguishes its approach by using language to control an environment model and studying intricate interactions among entities.

World models

  • Prior world models advanced from feed-forward and recurrent networks to Transformers and support model-based reinforcement learning, but language control of world models remains underexplored.Earlier language-conditioned models studied emergent language or incorporated language features without using language to control the world model.Source: §6 · p. 9

Language-based adaptation

  • Unlike instruction-following and language-feedback work that directly trains policies, this work uses text to enhance an environment model that captures its transition function.Prior work also uses environment-dynamics descriptions to improve policy learning, while game-manual agents focus on playing games.Source: §6 · p. 9

Compositional generalization for language-guided world models

  • MESSENGER studies more intricate compositional concepts involving interactions among multiple entities than prior LWM work focused on straightforward color and object mappings.Lin et al. (2024) did not demonstrate generalizability in MESSENGER.Source: §6 · p. 9

7 Conclusion

The paper introduces world models adapted through natural language and outlines their advantages over observational models. It also identifies limited model performance and the gap between grid-world experiments and the real world as boundaries on the work.

7 Conclusion

  • The paper introduces Language-Guided World Models, which can be adapted through natural language, and outlines advantages over traditional observational world models.The authors frame the work as a step toward enhancing controllability and addressing compositional generalization.Source: §7 · p. 9
  • The model remains limited in performance, and the grid-world environments severely underrepresent the real world.The authors present the work as an initial contribution intended to inspire further work on compositional generalization.Source: §7 · p. 9

A GPTHard model

GPTHard uses ChatGPT to identify entities mentioned in game manuals and routes extracted values to identity embeddings. It uses parsed identities rather than learning soft attention for identity grounding.

A GPTHard model

  • GPTHard generates description values and routes them directly to identity embeddings instead of learning soft attention.The approach asks ChatGPT to extract identity, role, and movement information but uses only the parsed identity in GPTHard.Source: A GPTHard model · p. 13
  • GPTHard prompts ChatGPT to identify entities in manual descriptions, while the OracleParse comparison uses ground-truth parses.The reported ChatGPT identity-parsing accuracy is the fraction of games in each split where all three identities are correctly identified.Source: A GPTHard model · p. 13
  • The ChatGPT prompt requests a character identity, movement type, and role type for each description.The examples use role labels such as dangerous enemy, secret message, and essential objective, and movement labels such as chasing, fleeing, and stationary.Source: A GPTHard model · p. 13

B Dataset

The dataset uses five rule-based behavior policies and generates trajectories for training and evaluation. Table 4 reports dataset statistics, including ChatGPT’s accuracy in identifying all three game identities by split.

B Dataset

  • Five rule-based policies generate trajectories: survive, win, suicide, obtain the message, and act randomly.The win policy obtains the message and then runs to the goal; it does not plan to avoid the enemy.Source: B Dataset · p. 13
  • Training uses 66 trajectories per game, with policies sampled uniformly; evaluation uses five trajectories per game, one per policy.Source: B Dataset · p. 13
  • Table 4 reports MESSENGER data statistics, including the fraction of games in each split where ChatGPT identifies all three identities correctly.Source: Table 4 · p. 14
  • Table 5 lists the training hyperparameters.Source: Table 5 · p. 14

C Training details

The implementation compares standard cross-attention with EMMA and uses a non-default initialization scheme because the default PyTorch initialization did not support compositional generalization. Experiments were run on single NVIDIA RTX2080 devices.

C Training details

  • The Transformer implementation is largely based on the IRIS codebase.Source: C Training details · p. 14
  • The Standard baseline uses cross-attention, while the proposed model uses EMMA.Source: C Training details · p. 14
  • The default PyTorch initialization did not suffice for compositional generalization, so the authors adopted an initialization scheme from the IRIS codebase.Source: C Training details · p. 14
  • Experiments primarily used single NVIDIA RTX2080 GPUs; Table 1 models trained for 24 GPU hours across five seeds.For Tables 3 and 6, five world models were trained on 90 games with three downstream policy-training strategies, at 12 GPU hours per game.Source: C Training details · p. 14

D Imitation learning experiments

The imitation-learning experiments use an EMMA-based policy trained from recent observations and compare expert policies across online imitation learning and filtered behavior cloning. Evaluation covers randomly selected test environments and all stage-two starting configurations.

D Imitation learning experiments

  • The learning policy processes the three most recent observations with a convolution-then-MLP encoder and is trained for 2,000 batches.It uses the world models’ optimizer hyperparameters.Source: D Imitation learning experiments · p. 14
  • Online imitation learning uses the rule-based win policy as its expert; filtered behavior cloning uses converged and not-fully-converged EMMA policies as near-optimal and suboptimal experts.The two filtered-BC checkpoints were trained for 10,000 and 2,000 iterations, respectively.Source: D Imitation learning experiments · p. 14
  • Evaluation randomly selects 10 environments per split and tests each policy for 48 real-environment episodes covering all 24 stage-two initial configurations.Source: D Imitation learning experiments · p. 14

E Extended results

In real-environment plan-discussion experiments, EMMA-LWM outperforms all baselines. It also maintains an advantage across ground-truth trajectory prefix lengths, including when compared with an observational model given a much longer history.

E Extended results

  • EMMA-LWM outperforms baselines at every ground-truth prefix length, while losses generally decrease as prefixes grow.Figure 4 plots cross-entropy losses with 95% t-value confidence intervals.Source: E Extended results · pp. 14–15, Figure 4 · p. 15
  • Across splits, one-step EMMA-LWM outperforms Observational conditioned on one third of a ground-truth trajectory.The authors interpret this comparison as evidence that EMMA-LWM leverages textual information.Source: E Extended results · p. 15
  • EMMA-LWM outperforms all baselines in average returns across all plan-discussion settings.Table 6 reports means over 30 randomly selected test environments per split, with 95% t-value confidence intervals.Source: E Extended results · p. 15, Table 6 · p. 16
Loading 2402.01695v3…