[Paper Review] Learning to Model the World with Language

Jessy Lin, Yuqing Du+5 more|arXiv (Cornell University)|Jul 31, 2023

Multimodal Machine Learning ApplicationsComputer Science3 citations

TL;DR

Dynalang is a multimodal agent that learns to ground language in visual and temporal world dynamics by predicting future text, images, and rewards, enabling robust language understanding and task performance without relying solely on task-specific rewards. It achieves state-of-the-art results in vision-language navigation, game playing, and home environment tasks by decoupling world modeling from policy learning and enabling pretraining on text and video data.

ABSTRACT

To interact with humans and act in the world, agents need to understand the range of language that people use and relate it to the visual world. While current agents can learn to execute simple language instructions, we aim to build agents that leverage diverse language -- language like "this button turns on the TV" or "I put the bowls away" -- that conveys general knowledge, describes the state of the world, provides interactive feedback, and more. Our key idea is that agents should interpret such diverse language as a signal that helps them predict the future: what they will observe, how the world will behave, and which situations will be rewarded. This perspective unifies language understanding with future prediction as a powerful self-supervised learning objective. We instantiate this in Dynalang, an agent that learns a multimodal world model to predict future text and image representations, and learns to act from imagined model rollouts. While current methods that learn language-conditioned policies degrade in performance with more diverse types of language, we show that Dynalang learns to leverage environment descriptions, game rules, and instructions to excel on tasks ranging from game-playing to navigating photorealistic home scans. Finally, we show that our method enables additional capabilities due to learning a generative model: Dynalang can be pretrained on text-only data, enabling learning from offline datasets, and generate language grounded in an environment.

Motivation & Objective

• To develop agents that understand diverse language beyond simple task instructions, such as world knowledge, situational descriptions, and feedback.
• To address the limitation of reinforcement learning in mapping language directly to actions, which provides a weak signal for complex language use.
• To unify language understanding with future prediction as a self-supervised learning objective for multimodal world modeling.
• To enable pretraining on text-only or video-only datasets without actions or rewards, improving sample efficiency and generalization.
• To allow agents to generate language and plan in latent space by leveraging imagined rollouts from the world model.

Proposed method

• Dynalang uses a multimodal world model that compresses visual and textual observations into a shared latent space for future prediction.
• The world model is trained via self-supervised learning to predict future latent representations of text and images, using online interaction data.
• Action policy is trained separately via reinforcement learning using the latent world model as input, decoupling perception from decision-making.
• The model supports pretraining on text-only or video-only data, enabling transfer learning and improved performance from in-domain language data.
• Language generation is enabled by rolling out the world model in latent space and reconstructing tokens, allowing the agent to speak about the environment.
• Imagined rollouts are used to plan and reason about possible futures, with uncertainty modeled through stochastic sampling of latent states.

Experimental results

Research questions

• RQ1 Can future prediction of text, images, and rewards serve as a unified self-supervised objective for grounding language in visual and environmental dynamics?
• RQ2 How does pretraining on text-only or video-only data improve downstream task performance in zero-shot or few-shot settings?
• RQ3 To what extent can an agent learn to use diverse language—such as environment descriptions, rules, and feedback—for improved task execution?
• RQ4 Can the world model support language generation and planning in latent space without explicit language modeling objectives?
• RQ5 How does the model’s ability to predict future observations improve generalization across tasks with complex language inputs?

Key findings

• Dynalang outperforms task-specific architectures on the Messenger benchmark by reading game manuals and solving the most challenging stage, demonstrating robust language grounding.
• Pretraining on a small amount of in-domain text, such as TinyStories, closes most of the performance gap compared to using T5 embeddings, even surpassing them in some settings.
• The model learns to predict future observations and rewards from language inputs, such as correctly anticipating that picking up the bottle will not be rewarded if it is not the target.
• Imagined rollouts show the model captures uncertainty in world states, with different samples representing plausible futures consistent with language descriptions.
• One-token-per-timestep input outperforms sentence-level encoding, suggesting that fine-grained token-level processing preserves more information than lossy sentence embeddings.
• The agent can generate coherent language by reconstructing tokens from latent rollouts, indicating potential for future use as a multimodal language model.