Source-linked AI summary
Language Models Meet World Models: Embodied Experiences Enhance Language Models
Jiannan Xiang, Tianhua Tao, Yi Gu, Tianmin Shu, Zirui Wang, Zichao Yang, Zhiting Hu
TL;DR
Language models often lack robust physical-world reasoning because text-only training omits embodied experiences. The paper finetunes them on experiences gathered in a world model using goal-oriented planning and random exploration, with EWC-LoRA preserving generality and efficiency. Across downstream evaluations, the approach improves base models and sometimes surpasses ChatGPT, while remaining limited by its single household environment.
Problem
Language models remain unreliable on physical-world reasoning and planning because their text-only training lacks embodied experiences.
Method
E2WM collects goal-oriented and exploratory experiences in VirtualHome, then finetunes language models using EWC-LoRA.
Results
The method significantly improves base models on seen and unseen embodied tasks, with GPT-J-6B, OPT-13B, and LLaMA-13B outperforming ChatGPT on many tasks.
Takeaways & Limitations
World-model experiences provide a training mechanism for acquiring generalizable embodied knowledge and skills in language models.
Takeaways & Limitations
The work uses a single household environment as its world model and leaves integration across world models and domains for future study.
Abstract
from arXiv · showhide
While large language models (LMs) have shown remarkable capabilities across numerous tasks, they often struggle with simple reasoning and planning in physical environments, such as understanding object permanence or planning household activities. The limitation arises from the fact that LMs are trained only on written text and miss essential embodied knowledge and skills. In this paper, we propose a new paradigm of enhancing LMs by finetuning them with world models, to gain diverse embodied knowledge while retaining their general language capabilities. Our approach deploys an embodied agent in a world model, particularly a simulator of the physical world (VirtualHome), and acquires a diverse set of embodied experiences through both goal-oriented planning and random exploration. These experiences are then used to finetune LMs to teach diverse abilities of reasoning and acting in the physical world, e.g., planning and completing goals, object permanence and tracking, etc. Moreover, it is desirable to preserve the generality of LMs during finetuning, which facilitates generalizing the embodied knowledge across tasks rather than being tied to specific simulations. We thus further introduce the classical (EWC) for selective weight updates, combined with low-rank adapters (LoRA) for training efficiency. Extensive experiments show our approach substantially improves base LMs on 18 downstream tasks by 64.28% on average. In particular, the small LMs (1.3B, 6B, and 13B) enhanced by our approach match or even outperform much larger LMs (e.g., ChatGPT).
1 Introduction
Current language models can assist with embodied tasks but remain unreliable on simple physical-world reasoning and planning because text-only training omits embodied experiences. E2WM finetunes models on experiences from world models while preserving generality, improving seen and unseen embodied tasks.
- Motivation: Text-trained language models remain unreliable on simple physical-world reasoning and planning, including counting objects in a location.The paper attributes this limitation to missing experiences such as navigating, interacting with objects, and sensing.
- Approach: E2WM finetunes pretrained language models with embodied experiences collected from world models such as VirtualHome.World models provide simulated physical interactions in a simplified, cost-effective environment.
- Approach: Goal-oriented planning and random exploration collect complementary experiences for planning, object tracking, and other embodied skills.Goal-oriented planning targets goal-directed behavior, while random exploration accumulates object-interaction experiences.
- Approach: Collected experiences are converted into finetuning tasks, while EWC regularizes updates to preserve important pretrained language-model parameters.The paper reports EWC as more effective than KL regularization for retaining general capabilities.
- Results: Finetuned GPT-J-6B, OPT-13B, and LLaMA-13B models can outperform ChatGPT on many evaluated tasks.The experiments compare models trained with the VirtualHome world model against original baselines and assess seen and unseen tasks.
- Results: 34.31 →51.23 Rouge-L on plan generation and 30.41% →67.01% accuracy on counting show improvements over baselines without pretraining-data degradation.Pile test-subset perplexity changes from 3.443 to 3.537.
2 Related Work
Related work covers world models that simulate physical environments, methods for grounding language models in those environments, and regularization techniques for preserving language-model capabilities during finetuning.
- World Model: World models computationally represent physical-world state changes in response to actions and can serve as virtual test environments for embodied agents.Examples include VirtualHome, AI2-THOR, VRKitchen, CHALET, MINOS, House3D, and MineCraft.
- Language Model Grounding: Prior grounding methods connect language models to world models through prompting, learned modules, or action-affordance mechanisms.Examples include Zero-Shot Planner, Mind’s eye, and SayCan.
- Language Model Regularization: Finetuning regularization commonly uses a KL penalty between the current and original model outputs to reduce loss of language-model abilities.The passage situates this approach in reinforcement-learning-style policy-gradient regularization.
3 Approach
E2WM uses VirtualHome to collect embodied experiences through goal-oriented planning and random exploration, then finetunes language models on tasks derived from those experiences. EWC-LoRA limits parameter updates to improve knowledge retention and finetuning efficiency.
- Collecting Embodied Experiences: E2WM uses VirtualHome as a world model for collecting embodied experiences involving physical interactions and executable actions.The simulator represents household activities and world states, allowing agents to interact with objects and execute actions.
- Collecting Embodied Experiences: Goal-oriented planning uses MCTS to search action sequences that achieve activity goals defined by target-state predicates.For example, setting up a table can require fork and plate predicates to be satisfied on the table.
- Collecting Embodied Experiences: Random exploration records object interactions, movements, and locations to provide experiences for object permanence and tracking.The process captures observations of objects as agents explore and move through the environment.
- Finetuning LMs with Embodied Experiences: Collected experiences are converted into supervised tasks for plan generation, activity recognition, counting, and object path tracking.Plan generation and activity recognition use goal-oriented experiences, while counting and path tracking use random-exploration experiences.
- Efficient Finetuning with Preserved Generality: EWC-LoRA combines EWC with low-rank adapters to update fewer parameters and reduce the computational and memory burden of full-model finetuning.LoRA freezes pretrained weights and trains two low-rank matrices, while the approach addresses EWC’s three-times-model-size memory overhead.
- Efficient Finetuning with Preserved Generality: EWC regularizes finetuning by weighting parameter changes according to their importance to pretraining, helping reduce catastrophic forgetting.The Fisher matrix estimates parameter importance, and the regularizer constrains shifts during finetuning.
4 Experiments
The experiments evaluate world-model-finetuned language models on constructed embodied tasks, unseen-task generalization, bAbI, language modeling, regularization, and ablations. Across these evaluations, the approach improves downstream performance while largely preserving language-modeling ability, with task-specific training experiences contributing to corresponding skills.
- Evaluation Tasks: The evaluation covers plan generation, household-object QA, activity recognition and inference, counting, object-path tracking, bAbI embodied-knowledge tasks, and Pile perplexity.Unseen activities and irrelevant contextual information are included to test generalization and distraction resistance.
- Constructed Evaluation Tasks: Finetuned GPT-J and LLaMA-13B outperform ChatGPT on most of the 11 constructed tasks, including LLaMA-13B gains on counting and object-location QA.The approach improves both seen and unseen tasks, including under confusing settings with irrelevant context.
- bAbI Tasks: GPT-J surpasses ChatGPT on the most challenging bAbI tasks, including Three Supporting Fact and Lists/Sets, while the method underperforms base models only on Two Supporting Fact.These results are presented as evidence that the models acquire embodied knowledge rather than simply overfitting to the training environment.
- Comparison of Different Regularization Methods: EWC-LoRA achieves the best trade-off between downstream performance and perplexity among the compared regularization methods.EWC is memory-intensive for large models, KL regularization is time-consuming, and combining KL with LoRA increases perplexity without improving downstream performance.
- Ablation Studies: Removing a training task generally decreases performance on the corresponding downstream ability, whereas omitting counting can increase Counting QA performance.The ablation attributes corresponding skill gains to the gathered embodied experiences while noting that counting may be inferred from other tasks.
5 Conclusion & Future Work
The framework enhances language models with world-model experiences collected through goal-oriented planning and random exploration, then compiles them for fine-tuning. EWC-LoRA supports parameter-efficient tuning, mitigates catastrophic forgetting, and enables knowledge generalization, while the current study remains limited to one household environment.
- 5 Conclusion & Future Work: The framework collects embodied experiences from world models through goal-oriented planning and random exploration.These experiences are compiled into formats suitable for language-model fine-tuning.
- 5 Conclusion & Future Work: EWC-LoRA facilitates parameter-efficient tuning while alleviating catastrophic forgetting and enabling knowledge generalization.
- 5 Conclusion & Future Work: The method demonstrates strong performance across a large number of downstream evaluation tasks.
- 5 Conclusion & Future Work: The present work is limited to a single household environment as its world model.Future work aims to integrate experiences from different world models and generalize learned knowledge across domains.
A.1 VirtualHome
VirtualHome represents executable actions with an agent identifier, an atomic action, an object, and that object's class-specific identifier. Natural-language templates transform these actions for language-model fine-tuning.
- A.1 VirtualHome: Each executable VirtualHome action specifies an agent, an atomic action, an object, and an Object_id.The agent identifier selects which agent acts when multiple agents are present.
- A.1 VirtualHome: char_id identifies which agent executes an action when multiple agents occupy the world model.
- A.1 VirtualHome: Action must be a supported atomic action in VirtualHome.
- A.1 VirtualHome: Object denotes the object with which the agent interacts.
- A.1 VirtualHome: Each environment object receives an Object_id to distinguish objects sharing the same class.
- A.1 VirtualHome: Action templates transform executable actions into natural text for language-model fine-tuning.The complete executable-action list appears in Table 4, where agent and object identifiers are omitted for simplicity.
A.2 Acitivity Goal And Predicate
A VirtualHome household activity goal is a conjunction of predicates describing object conditions or relations. Activities and goals were collected from RobotHow.
- A.2 Acitivity Goal And Predicate: A household activity goal consists of several predicates.
- A.2 Acitivity Goal And Predicate: Each predicate represents an object's condition or a relation between two objects.
- A.2 Acitivity Goal And Predicate: OPEN(coffee maker) denotes that the coffee maker is open.
- A.2 Acitivity Goal And Predicate: ON(apple, table) denotes that an apple is on the table.
- A.2 Acitivity Goal And Predicate: The goal is achieved only when all predicates are satisfied.
- A.2 Acitivity Goal And Predicate: Activities and goals were collected from RobotHow.
A.3 Data Format and Prompts
The prompts use in-context exemplars containing instructions, question contexts, and answers. Full prompts include multiple exemplars, with task-specific data formats and examples provided below.
- A.3 Data Format and Prompts: The prompting scheme follows Chung et al. by using instructions with in-context exemplars.
- A.3 Data Format and Prompts: Each exemplar contains an instruction, question context, and answer.
- A.3 Data Format and Prompts: The full prompt contains multiple exemplars for in-context learning.
- A.3 Data Format and Prompts: The data format and exemplar are provided separately for each task.
- A.3 Data Format and Prompts: Instructions, contexts, and answers are supplied within every exemplar.
- A.3 Data Format and Prompts: Task-specific exemplars operationalize the prompts for in-context learning.
A.4 Hyperparameters
The experiments use task-specific training settings for learning rate, batch size, task weights, epochs, EWC regularization, and LoRA configuration.
- GPT-Neo-1.3B and GPT-J-6B use a learning rate of 8 × 10−5 and a batch size of 20.
- The task weights are 1.0 for plan generation, 0.7 for activity recognition, 1.0 for counting, and 1.0 for object path tracking.
- GPT-Neo-1.3B is trained for 3 epochs with EWC coefficient λ = 0.5, while GPT-J-6B is trained for 5 epochs with λ = 2.
- Training takes 40 minutes for GPT-Neo and 220 minutes for GPT-J, using LoRA rank 8 and coefficient 32.
A.5 bAbI Dataset
The bAbI evaluation includes eight tasks testing supporting-fact reasoning, counting, set/list handling, negation, temporal reasoning, and positional reasoning, with examples illustrating their formats.
- The dataset contains 8 embodied-knowledge tasks, including supporting-fact, counting, Lists/Sets, negation, time, and positional reasoning.
- The supporting-fact tasks require answering questions using one, two, or three facts from the context.
- Example tasks ask where Mary is, where the football is, and where an apple was before the kitchen, using state changes in the context.
- The Lists/Sets example asks which objects Daniel is holding after he picks up and drops items.
- The Time Reasoning example asks where Julie went after the park and where she was before it.
- Table 5 provides examples for the bAbI tasks.
A.6 Results of Main Experiments and Ablation Studies
The paper reports experimental results on constructed downstream tasks, bAbI, SuperGLUE, and ablation studies, alongside human evaluation of planning and a broader-impact warning about harmful inputs.
- Main experiments and ablations: The experiments report results on constructed downstream tasks, bAbI, and ablation studies, with results organized in Tables 6–8.
- Human evaluations: GPT-J plan-generation accuracy rises from 24.0% for the base model to 62.4% after finetuning in human evaluation.
- SuperGLUE results: The finetuned model matches and sometimes outperforms the baseline on appropriate SuperGLUE tasks, indicating retained general language capability.
- Broader impact: The trained language model remains susceptible to unintended outputs when faced with harmful or adversarial inputs.