Title: Language Models Meet World Models: Embodied Experiences Enhance Language Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Language Model

World Model

Embodied Experience
Abstract: 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 elastic weight consolidation (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).

Language Models Meet World Models: Embodied Experiences Enhance Language Models

Abstract

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 elastic weight consolidation (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

1 Introduction

Language Models (LMs) have demonstrated impressive performance on a wide range of natural language processing tasks [34, 48, 4, 7, 54]. In particular, recent studies show that LMs can assist decision-making for embodied tasks [1, 18, 25, 45, 19], demonstrating a certain level of understanding of the physical world. However, such understanding is not robust enough for many reasoning and planning tasks in physical environments. As shown in Figure 1, even the latest large LMs like ChatGPT 2 can still make mistakes in seemingly simple inquiries, such as counting objects in a location. We hypothesize that this is because current LMs trained merely with large-scale text corpora are devoid of embodied experiences such as navigating in an environment, interacting with objects, and sensing as well as tracking the world state. Consequently, they lack robust and comprehensive embodied knowledge necessary for reasoning and planning associated with physical environments. A related line of research finetunes LMs in order to improve specific embodied tasks, resulting in task-specialized models [6, 58, 21, 57].

It was on the

sofa.

（e.g. GPT-J-6B

）

Track & Count

Track & Permanence

1. Walk to living room..

2. Sit on sofa.

3.

Stand up.

4.

Walk to living room.

…

1.

Walk to

living room.

2.

Sit on sofa.

3.

Watch TV.

There are

two

items.

There are

three

items.

John put a book on the desk.

[Irrelevant Actions] Mary took the

book.

[Irrelevant Actions] She

placed it

on the sofa

…

Where was the book?

Emily found a desk and placed the

cell

phone

on top of it.

[Irrelevant Actions]

, …

putting the

lime

down next to the cell

phone.

[Irrelevant Actions]

She finally

put an

apple

on the desk.

How many items are there on the desk?

It was on the

desk

.

Embodied Experiences

How to

watch TV

?

Finetune

Goal-Oriented

Planning

Random

Exploration

1

Tracking (of book)

On bed, in bedroom

In hand, in bedroom

In hand, in living room.

Action

Walk to bedroom.

Grab book.

Walk to living room.

2

3

3

1

2

1

2

3

ChatGPT

ChatGPT

ChatGPT

LM

Activity: Watch TV

Walk to living room.

Sit on sofa.

Watch TV.

1

2

3

Question:

LM

LM

Answer:

E2WM

training paradigm

World Model

In this paper, we aim to inject diverse fundamental embodied knowledge and skills into pretrained LMs, while retaining the models' generality. We introduce a novel training paradigm for LMs - finetuning with Embodied Experiences from World Models (E2WM). Here, world models are embodied simulators that emulate physical interactions in real-world environments (e.g., VirtualHome [36]). They provide LMs with the opportunity to comprehend object interactions within the environment and to execute actions, thus enabling a level of active engagement previously unattainable. These world models serve as a simplified and cost-effective replica of our real world that can significantly augment the conventional pretraining paradigm. We anticipate that finetuning LMs on embodied experiences gathered from world models can enhance their embodied knowledge-and, with the preserved model generality-consequently strengthen their abilities to solve a broad range of embodied tasks.

In this work, we consider a diverse range of fundamental knowledge and skills for embodied tasks, including tracking objects, planning to complete given goals, recognizing other agents' behaviors, etc . To this end, we introduce two ways to collect embodied experiences from world models that give rise to the desired knowledge and skills: goal-oriented planning and random exploration (Figure 1). Specifically, goal-oriented planning aims to gather experiences associated with planning and goal-oriented agent behaviors, while random exploration focuses on accumulating experiences that involve object and world state tracking. In goal-oriented planning, models are given the goal ( e.g. , IN(dust, trash can) ) for a specific activity ( e.g. , Clean Floor ), and is supposed to generate a plan to complete it. To find the plan, we devise Monte Carlo Tree Search (MCTS) [5, 44] to explore the world model. Then the process will be stored as an embodied experience. In random exploration, one or more agents are deployed in the world model to execute random actions, while the locations and the movements of all the objects are tracked simultaneously.

After collecting the embodied experiences, we use them to construct a set of fine-tuning tasks (e.g., plan generation, activity recognition, and tracking). Crucially, to finetune LMs on the collected embodied experiences while retaining their original general knowledge and capabilities, we propose to incorporate the classical Elastic Weight Consolidation (EWC) [22] into our training paradigm. By regularizing the finetuning loss, EWC aims to preserve the important LM parameters from pretraining. We show that EWC is substantially more effective than the popular KL regularization [35, 28, 31]. We further introduce efficient low-rank updates by harmonizing the recent Low-Rank Adaptation (LoRA) [16] with the EWC regularizer. This results in the new EWC-LoRA update rule that greatly reduces training costs and makes our E2WM paradigm accessible to cheap hardware (GPUs).

We instantiate a world model using a virtual household simulator, VirtualHome [36, 37], and apply our method to GPT-Neo-1.3B [3], GPT-J-6B [49], OPT-13B [59], and LLaMA-13B [48] models. To test the generalizability of the finetuned LMs, we evaluate them on a variety of unseen tasks which

demand similar embodied knowledge required to solve the training tasks. Additionally, we assess the models' performance on the original pretraining data to determine the extent to which their core language modeling abilities are retained. Experiments show that our method significantly improves the baselines on both seen and unseen tasks ( e.g. , 34 . 31 → 51 . 23 Rouge-L on plan generation task, 30.41 % → 67 . 01% accuracy on counting task), without suffering performance degradation on the pretraining dataset ( 3 . 443 → 3 . 537 perplexity on Pile test subset [12]). Moreover, the small GPTJ-6B, OPT-13B, and LLaMA-13B models finetuned with our E2WM paradigm even outperforms ChatGPT on many of the tasks. The experimental results demonstrate the effectiveness of E2WM as a promising fine-tuning mechanism to enhance pretrained LMs with generalizable embodied knowledge and skills.

2 Related Work

World Model. The term 'world model' generally refer to a computational representation of the physical world, capable of simulating changes in the world's state in response to various actions. For instance, humans possess an internal world model that aids in predicting the outcomes of specific actions during the planning process. Recent research induces world models from large LMs for robust human-like reasoning [15]. In this work, we employ a simulator equipped with a physics engine to serve as our world model, effectively emulating real-world conditions. In the field of embodied AI, various world models are built to replicate the real world and serve as virtual test environments for assessing robotic agents before real-world deployment. For example, VirtualHome [36, 37] is a simulated 3D household environment implemented by Unity3D game engine. AI2-THOR [23] consists of near photo-realistic 3D indoor scenes and has richer object attributes and interaction types. Other indoor household World Models include VRKitchen [13], CHALET [56], MINOS [41], House3D [53], etc. Besides, MineCraft is a more challenging and open-ended world model, which has a large number of objectives and a large-scale task hierarchy [14, 27, 20]. In this paper, we use VirtualHome as our world model.

Language Model Grounding. A significant number of recent works focused on grounding language models to world models [1, 24, 38, 47, 55]. Some of them freeze LMs and leverage certain prompting strategies or specifically-designed modules. For example, Zero-Shot Planner [18] prompts LMs to generate activity plans and translate them into admissible actions. Mind's eye [29] prompts LMs to do simulations with physical engines to answer physical reasoning questions. SayCan [1] uses a learned affordance function to assist LMs in selecting valid actions. DEPS [50] prompts LMs to describe, explain and generate action plans, incorporated with a learned selector module to choose the most efficient path. There are also other works finetuning LMs towards better downstream task performance. For example, Li et al. [25] finetune LMs with supervised learning for interactive decision making, and Carta et al. [6] ground LMs with online reinforcement learning. Different from these works aiming to optimize LMs for specific tasks in the target environments, our work instead focuses on improving the language model itself by acquiring knowledge from world models.

Language Model Regularization. To facilitate the acquisition of new knowledge and skills without losing LMs' language modeling abilities , regularization is often introduced during finetuning. One popular method is adding KL penalty [35, 28, 46, 52, 33, 60], which leverages KL divergence between the output probability of the currently trained model and the original model to regularize the LM in an RL manner, i.e. , by computing policy gradients. For example, InstructGPT uses KL penalty to mitigate over-optimization of the reward model [35], and Liu et al. [28] add KL regularization for training a commonsense knowledge generator. In this work, we instead use elastic weight consolidation (EWC) for regularization. Our empirical results demonstrate that EWC is more effective than applying KL penalty for retaining language modeling abilities and generality of LMs.

3 Approach

In this work, we propose a new training paradigm, namely finetuning with Embodied Experiences from World Models (E2WM), to inject embodied knowledge into LMs without sacrificing its generality and language modeling abilities. The world model we use is VirtualHome [36, 37], a multi-agent simulator for househould activities. In VirtualHome, an executable action step can be simplified as the format of [action] <arg> , e.g. , [Grab] <apple> . The world state of VirtualHome consists of objects and their relations ( e.g. , apple on table). Details about VirtualHome can be found in Appendix

Activity: Watch TV

Question:

How to watch TV? TV and

sofa is in living room…

Answer:

Walk to living room. Sit

on sofa. Turn on TV.

Question:

Given a plan: Walk to living

room. Sit on sofa. Turn on TV.

What is the task?

Answer: Watch TV.

1

2

3

7

5

4

8

9

6

3

Take pillow

5

Walk to living room

Put pillow on table

9

1

Grab pillow

Walk to bathroom

7

2

Give pillow to

4

Grab apple

6

Put apple on table

Walk to bedroom

8

Question:

Tom grabbed pillow. Tom

gave pillow to … How many

objects are on the table?

Answer:

Two. They are pillow and apple.

Question:

Tom grabbed pillow. Tom walked

to kitchen … What is the order of

rooms where pillow appears?

Answer:

Bedroom, kitchen, living room

Goal-Oriented Planning

Random Exploration

Plan Generation

Activity Recognition

Counting

Object Path Tracking

Walk to bathroom

Grab toothpaste

Walk to bedroom

Walk to living room

Sit on Sofa

Walk to table

Grab apple

Turn on TV

Walk to living room

MCTS

-0.1

-0.1

+1.9

-0.1

+1.9

+1.9

+1.9

-0.1

-0.1

A.1. We first describe how to gather embodied experiences in the world model in Section 3.1. Then in Section 3.2 we demonstrate how to finetune LMs by utilizing collected experiences, as well as our proposed method EWC-LoRA for efficient knowledge generalization.

3.1 Collecting Embodied Experiences from World Model

LMs pretrained on large scale human-written text corpus often have difficulties in solving basic reasoning and planning in physical environments. This is because that LMs lack necessary embodied knowledge and experiential understanding of the physical world. To address the problem, we propose to leverage world models to collect diverse embodied experiences for enhancing LMs. Specifically, to inject different types of embodied knowledge into LMs, we introduce two ways to gather experiences: goal-oriented planning and random exploration. Figure 2 illustrates the two methods.

Goal-oriented Planning. One important embodied skill is to plan and complete a specific goal, e.g. , placing tableware properly to set up the table. To endow LMs with this ability, we propose goal-oriented planning. The approach aims to generate experiences that are goal-oriented, thus are useful to facilitate the acquisition of skills and task planning abilities for executing a range of activities in the world model. To do that, we collect various activities and their corresponding goals. Formally, the goal for an activity in the world model is defined as a set of predicates describing the target world state. For instance, an activity can be set up table , and its goal can be ON(fork, table);ON(plate, table) , which means that fork and plate should be put on the table to fulfill the activity. More details about predicates and goal definitions can be found in Appendix A.2. As shown in Figure 2, in goal-oriented planning, we devise a Monte Carlo Tree Search (MCTS) planner to search through the action space and find a plan, i.e. , a sequence of actions, to achieve the goal. The key to successful MCTS is the reward design. At each time step, if at least one goal predicate is satisfied, the MCTS planner will get a reward of +2, and the achieved goal predicates will be removed from the goal. This ensures that the planner does not repeatedly execute the same action to receive rewards, but rather focuses on achieving the remaining unfulfilled goals. Besides, it will get a -0.1 penalty after each time step to discourage the planner from doing actions irrelevant to fulfilling the goals. Finally, we store the planning process as an embodied experience.

Random Exploration. In real-world scenarios, humans not only acquire new knowledge by finishing tasks, but also learn by just randomly exploring the surroundings, e.g. , randomly observing/tracking objects and knowing their properties. To mimic this learning process, we propose another approach, namely random exploration. By simply exploring in the world model, embodied experiences emerge that involve advanced cognitive abilities including object permanence and tracking, as agents observe and track the consistent existence of objects even when they are out of sight. Then the experiences are gathered for finetuning LMs later. Specifically, the approach deploys one or multiple agents in the world model wandering aimlessly and randomly executing actions. As illustrated in Figure 2, multiple agents are in the same environment, interacting with each other or executing different actions

on the same objects, which simulates complex situations. During the exploration, the moving paths and the final locations of all the objects in the world model are recorded. Then the whole process is captured as an embodied experience.

3.2 Finetuning LMs with Embodied Experiences

There are multiple ways to utilize collected embodied experiences for LMs finetuning, such as supervised learning and reinforcement learning. In this work, we use them with supervised learning for simplicity and efficiency. Specifically, goal-oriented planning experiences are compiled into data examples in two formats: plan generation and activity recognition . As shown in Figure 2, in plan generation, the model is required to generate a stepwise action sequence to fulfill an activity, given the state of some relevant objects as the initial condition. In activity recognition, the model needs to recognize the activity name given its plan. Experiences obtained from random exploration are also transformed into two self-supervision tasks: counting and object path tracking . Examples of the two tasks can be seen in Figure 2. Specifically, for counting, the LM is tasked with identifying the number and name of the objects at a specific location after the agents performed relevant and irrelevant actions and arranged objects randomly. In object path tracking, the model is tasked to output the moving path of an object that is picked up by different agents and moved to different rooms at different times. All the tasks are trained with cross-entropy loss. Suppose that x is the input ( e.g. , the initial condition in plan generation) and y = { y 1 , ..., y M } is the label ( e.g. , the stepwise action sequence), we finetune LMs by assigning different weights to different tasks:

where L is the loss function; V is the task set; and α v is the weight for task v . Following Flan-T5 [8], x is a prompt formatted to contain a task instruction and sampled in-context demonstrations. We provide all prompts in Appendix A.3.

Efficient Finetuning with Preserved Generality. However, there are two key problems for simply finetuning LMs. The first one is that LMs will easily overfit to the downstream tasks, leading to performance degradation on other tasks. This deviates from our goal that the model should generalize acquired knowledge across various tasks. Another problem is that finetuning the entire LM is resource-intensive and time-consuming, especially when the LM is extremely large. To overcome the problem and facilitate continual and efficient knowledge acquisition with world models, we propose to finetune only a small number of weights using low-rank adaptors (LoRA) [16] with elastic weight consolidation (EWC) [22], which we refer to as EWC-LoRA.

EWC is a regularization-based method typically used in the area of continual learning [9]. It calculates a fisher matrix [11] to estimate the importance of each parameter for a task and then uses it to regularize the training on a new task. The regularization term helps to constrain the parameter updates for the new task to avoid forgetting the previous knowledge. Let U be the pretraining task set, and V be the finetuning task set. Following [2], we have:

where L is the loss function, F is the fisher matrix, λ is the hyperparameter, i and j are the indices for parameters and data samples, respectively, and θ and θ ∗ U are currently trained parameters and frozen task U parameters, respectively. Notice that the first term L V ( θ ) in Equation 3 is calculated in Equation 1, and the second term is the EWC regularizer. In Equation 2, the fisher matrix is calculated by averaging the sum of squares of the gradients from the task U , which indicates the significance of each parameter to the task U . Then the matrix is used in Equation 3 to weigh the shift of model parameters when training on V . By using EWC, the LM learns to adapt to new tasks without catastrophic forgetting on the pretraining task, which forces it to understand and digest new knowledge from the finetuning tasks instead of overfitting to them.

However, EWC is both time- and memory-inefficient. First, it requires finetuning the entire set of large LM's parameters. Moreover, the approach involves creating a frozen original model and a fisher matrix that is the same size as the LM, leading to a memory overhead of three times the original size.

This makes it particularly challenging to apply to larger models. To alleviate the problem, we propose to combine EWC with low-rank adaptors (LoRA), a parameter-efficient tuning method. LoRA freezes the pretrained model weights and injects two trainable low-rank matrices into each layer of the model. Suppose that W, W ∗ ∈ R r × d are the trained weight matrix and frozen weight matrix, respectively; and B ∈ R r × k , A ∈ R k × d are two low-rank matrices with k ≪ min( r, d ) . Then the formula for LoRA can be written as W = W ∗ + BA . Suppose that H is flattened BA . Notably, we found that θ i in Equation 3 is the element of W , and θ ∗ U,i is that of W ∗ . Therefore, θ i -θ ∗ U,i is the element of H . We can thus transform Equation 3 into the final formula of EWC-LoRA method:

where h i = θ i -θ ∗ U,i is the i -th element of H . One of the benefits of this rewriting is that we no longer need to store the trained LM weight matrixes as what vanilla EWC does, which saves plenty of memory space. Besides, we only need to update B and A during the finetuning, which also lowers memory requirements and leads to much faster training speed. Surprisingly, as shown later, we empirically found that adding LoRA into EWC can further mitigate the issue of catastrophic forgetting and overfitting. This aligns with the previous conclusion that limiting the dimension of the optimization problem can alleviate catastrophic forgetting [32].

4 Experiments

Training Details. For goal-oriented planning, we collected activities and their corresponding target goals with data from RobotHow [36], a housework activity knowledge base created in VirtualHome. We applied our method to GPT-Neo-1.3B [3], GPT-J-6B [49], OPT-13B [59], and LLaMA-13B [48]. To save computing resources, we use Int8 technique [10]. Both of the models were trained with the AdamW optimizer [30]. All the hyperparameters are chosen according to the performance on a held out set. We used one NVIDIA GeForce RTX 3090 for training. More details can be found in Appendix A.4.

4.1 Downstream Evaluation Tasks

We developed various downstream evaluation tasks for each type of embodied knowledge, including both the training tasks as well as novel tasks unseen during training used for generalization evaluation. Additionally, we evaluate our models on bAbI [51], a dataset for testing multiple types of knowledge and abilities including embodied knowledge, logic reasoning, linguistic knowledge, etc. We select the bAbI tasks related to embodied knowledge for our evaluation. We evaluate all the unseen tasks including bAbI under few-shot settings, specifically 2-10 shots, by providing a few in-context exemplars in the prompts. We discuss more details of the tasks below.

Plan Generation. To evaluate planning ability, we construct downstream tasks using human-written plans from RobotHow. Specifically, we have three tasks:

Plan Generation Evaluation. In this task, the model needs to generate a plan for a housework activity. It is similar to the training task but uses human-written plans as the ground truth instead of the collected experiences. We include activities unseen during training to test the generalizability of the model. Inspired by the previous study showing that LMs can easily be distracted by irrelevant context [43], we also created samples having states of unrelated objects in the context ( e.g. , TV is on for activity Make Coffee ) to confuse LMs. In summary, this results in four settings: Vanilla Seen, Vanilla Unseen, Confusing Seen, and Confusing Unseen. We have 175/54/135/43 examples for four settings, respectively. We use Rouge-L [26] as the metric.

Housework QA. This is a multi-choice QA task, which is unseen during training. It asks which choice is the relevant object to finish a household activity, e.g., which object is relevant to making coffee? It has 261 examples in total, and we use accuracy as the metric. When evaluating, we provide 10 in-context exemplars in the prompts, so this task is evaluated as a 10-shot learning task.

Negation Housework QA. This is similar to Housework QA but inquires about the irrelevant object, e.g., which object is irrelevant to making coffee? It is more challenging than the vanilla QA because LMs that simply memorize the word co-occurrence in the training data may succeed in the vanilla QA but will fail in the negation QA. This task has 162 examples and uses accuracy as the metric. We provide 10 in-context exemplars in the prompts.

Activity Recognition. We developed two multi-choice QA tasks with the human-written plans and the state changes from RobotHow to test the knowledge gained from activity recognition:

Activity Recognition QA. In this task, a human-written plan is given and the model needs to choose the correct activity name. An example of the question is 'Given a plan: Walk to living room. Sit on sofa. Turn on TV. What is the name of this activity?' . And the answer should be Watch TV . The task has 549 examples. We use accuracy as the metric.

Activity Inference QA. In this task, we use the final state of the world model as input for the model to infer the activity name. For example, the input can be 'Tom is sitting on the sofa and facing the TV. The TV is on. What is a possible activity he is doing?' , and the answer is 'Watch TV' . We have 262 examples for this task and use accuracy to measure the performance. We provide 10 in-context exemplars in the prompts.

Counting. We gathered random exploration experiences to construct Counting QA for evaluating skills learned from the counting task. The model is required to answer the number of objects in a specific location. For example, a query can be 'Tom put an apple on the table. Tom turned on the TV. Tom put a cup on the table. How many objects are there on the table?' . We can see that there will be irrelevant actions like turn on TV to confuse the model and make the question more challenging. We collected 194 samples for the task and used accuracy as the evaluation metric. We provide 5 in-context exemplars in the input prompts.

Object Path Tracking. We developed two downstream tasks for the object path tracking training task, namely Object Path Tracking Evaluation and Object Location QA.

Object Path Tacking Evaluation. This evaluation task is the same as the training task, where the model is required to generate the full moving path of an object. An example is 'Tom walked to the kitchen. Tom grabbed the apple. Mary walked to the bedroom. Tom walked to the living room. What is the order of the rooms where the apple appears?' . This question typically includes multiple agents and many irrelevant actions, which makes it difficult to track the object. This task contains 200 examples. Following Huang et al. [18], we evaluate the performance by calculating the length of the longest common subsequence (LCS) between the ground truth and the generated path, normalized by the maximum length of the two.

Object Location QA. In this task, the model is asked about the location of an object before/after it moves to another location, e.g., where is the apple before/after the kitchen? This task has 200 examples with accuracy as the metric. We provide 2 in-context exemplars in the prompts.

A previous study on prompting multiple QA questions [39] introduces two prompting methods, multiple choice prompt and cloze prompt, and two normalization methods, length, and unconditioned normalization. For all the multi-choice QA tasks, we choose the combination of prompting and normalization methods which yields the best performance on a held out set.

To further verify the effectiveness of our method, we evaluate our finetuned GPT-Neo and GPT-J on the bAbI dataset. Specifically. we select 8 test sets from bAbI that align with the abilities covered in our collected embodied experiences. We include the description of each test set in Appendix A.5. For all the bAbI tasks, we do 2-shot learning by providing 2 in-context exemplars in the input prompts.

Besides downstream tasks, we also want to ensure that our approach does not hurt language modeling performance of the models. Therefore, following previous work [42], we evaluate the perplexity on a subset of Pile [12] test set, which is the pretraining dataset for GPT-Neo and GPT-J. We sampled 5000 examples from Pile test set for evaluation.

4.2 Results

Constructed Evaluation Tasks. Results for all the downstream evaluation tasks are shown in Figure 3 and Figure 4. For all the models, we compare the results obtained after finetuning with world model against those of the original base models. For GPT-J, we also include a finetuned model without EWC-LoRA as a baseline. Detailed numbers of the results can be found in Appendix A.6. We also conduct human evaluations for GPT-J on the plan generation task, which can be found in Appendix A.7. In general, the models trained with the world model significantly outperform the baselines on various downstream tasks. Our method is not only effective for small 1.3B model, but can also scale to larger 6B and 13B models. Specifically, our finetuned GPT-J and LLaMA-13B

0

20

40

60

80

100

Vanilla

UnSeen

Confusing

Seen

Confusing

Unseen

Object Path

Tracking

Housework

QA

Negation

Housework

QA

Activity

Recognition

Activity

Inference

Counting

Object

Location QA

Plan Generation

(Rouge-L)

(LCS)

Performance (%)

GPT-Neo (Ours)

GPT-J (Ours)

GPT-Neo (Base)

GPT-J (Base)

GPT-J (FT)

ChatGPT

0

20

40

60

80

100

Vanilla

UnSeen

Confusing

Seen

Confusing

Unseen

Object Path

Tracking

Housework

QA

Negation

Housework

QA

Activity

Recognition

Activity

Inference

Counting

Object

Location QA

Plan Generation

(Rouge-L)

(LCS)

Performance (%)

OPT-13B (Ours)

LLaMA-13B (Ours)

OPT-13B (Base)

LLaMA-13B (Base)

ChatGPT

0

20

40

60

80

100

Single

Supporting Fact

Two Supporting

Fact

Three

Supporting Fact

Counting

Lists/Sets

Simple

Negation

Time Reasoning

Positional

Reasoning

Accuracy (%)

GPT-Neo (Ours)

GPT-J (Ours)

GPT-Neo (Base)

GPT-J (Base)

ChatGPT

Table 2: Results of different regularization methods. The abbreviations in the Task column stand for the corresponding evaluation tasks for four training tasks. We use asterisk ∗ to mark the perplexity of base models.

with world model even achieve better performance than ChatGPT as a much larger LM on most of the 11 tasks. Besides, we can see the world model improves LMs on both seen and unseen tasks. This demonstrates that our model absorbs the knowledge for goal-oriented planning and random exploration instead of memorizing the seen experiences. Specifically, the better plan generation performance under the "Confusing" setting indicates that the world model improves the ability of LMs to avoid being interfered with by irrelevant contexts. On both Housework QA and Negation Housework QA, our models surpass the baselines, showing that our models also acquire knowledge about the necessary objects for completing a housework activity. Results on other downstream tasks also prove the effectiveness of our method. For example, on both Activity Recognition Evaluation and Activity Inference, our approach improves over the baselines significantly. Moreover, improvements can be observed in the downstream tasks regarding random exploration. On Counting and Object Location QA tasks, our LLaMA-13B trained with the world model even surpasses ChatGPT.

bAbI Tasks. To further verify the effectiveness of our method, we evaluate our finetuned models on the bAbI dataset. The results are shown in Figure 5. We can see that our approach significantly outperforms the base LMs. Notably, after finetuned with VirtualHome experiences, GPT-J surpasses the much stronger ChatGPT on the most challenging tasks. Specifically, it outperforms ChatGPT on the Three Supporting Fact task, where the model is required to use three supporting facts from the context to answer a question like ' where was the apple before the kitchen? ', and Lists/Sets task, which asks the model to give the answers in the form of a list, e.g. , the answer for ' What is Daniel holding? ' is ' apple, milk '. These results prove that our approach enables LMs to acquire the knowledge and skills inherent in embodied experiences, rather than simply overfitting to the training environment.

Language Modeling. In addition to verifying improved performance on the downstream tasks, we also report results on Pile test subset to ensure the preservation of the generality and language modeling abilities of LMs. From the experimental results shown in Table 1, we can see that our approach only causes a negligible increase in perplexity over the base models. This demonstrates the effectiveness of EWC-LoRA to preserve the generality and linguistic competence of LMs. To verify the generality on other NLP tasks, we also include the results on SuperGLUE [40] in Appendix A.8.

4.3 Comparison of Different Regularization Methods

We compare our proposed EWC-LoRA with EWC and LoRA. Besides, we also include the baseline using KL penalty as regularization.The experimental results are shown in Table 2. We also include the results of four evaluation tasks. Notice that we do not include the results of GPT-J with pure EWC and KL, since they are overly memory-intensive or time-consuming. EWC requires an original model and a fisher matrix other than the trained model, which triples the memory usage, making

it hard to be applied to large models like GPT-J-6B. Besides, KL penalty term is computed by L KL = E ( x , y ) ∼ P θ ∗ [ -log ( P θ ∗ ( y | x ) /P θ ( y | x ))] , thus it requires sampling from the model output probability, which is time-consuming. On the contrary, EWC-LoRA is both memory- and timeefficient. In Table 2, we can see that EWC-LoRA achieves the lowest perplexity compared to other methods, while still significantly outperforming the base LMs. Compared with pure EWC, applying pure LoRA greatly decreases perplexity, which is consistent with the previous conclusion that limiting the dimension of the optimization problem can mitigate catastrophic forgetting [32]. EWC-LoRA further decreases perplexity, making it extremely close to the original perplexity, while achieving comparable performance with LoRA on downstream tasks. This demonstrates the effectiveness of EWC-LoRA. Besides, We can find that combing LoRA with KL will greatly increase perplexity while not achieving better downstream performance. Overall, our proposed EWC-LoRA achieves the best trade-off between the perplexity and the downstream performance, which outperforms baselines significantly while almost not increasing the perplexity on the pretraining dataset.

4.4 Ablation Studies

To study the contribution of each training task, we conducted an ablation study by removing one training task every time. We use GPT-Neo-1.3B as the base model. We include the results on tasks seen during training in Table 3. Results on all the tasks can be found in Appendix A. We can see that the removal of a training task with similar ability leads to a notable decrease in the model's performance on downstream tasks. For example, the performance of plan generation drops significantly when plan generation is removed from the training tasks. Similarly, the removal of activity recognition or object path tracking from the training tasks leads to a decline in performance in their respective downstream tasks. We conclude that our gathered embodied experience has a tremendous contribution to teaching the corresponding reasoning ability by finetuning. Interestingly, Counting QA performance shows an increase when counting is omitted from the training tasks, possibly because the ability of counting can be inferred from other training tasks.

5 Conclusion & Future Work

We proposed a new training framework that uses world models to enhance language models. It first collects embodied experiences from world models through both goal-oriented planning and random exploration. The experiences are then compiled into appropriate formats for LMs finetuning. We further introduce EWC-LoRA, which not only facilitates parameter-efficient tuning but also alleviates catastrophic forgetting and enables knowledge generalization. We show the strong performance of our method on a large number of downstream evaluation tasks.

This work demonstrates the advantage of panoramic learning with all forms of experience [17]. On the other hand, the present work is limited to a single household environment as the world model. In the future, we intend to study how to integrate embodied experiences from different world models and generalize knowledge learned from each world model to different domains.

Acknowledgements. This project is partially supported by DARPA ECOLE HR00112390063.

References

Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. Do as i can, not as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691 , 2022.

Kristjan Arumae, Qing Sun, and Parminder Bhatia. An empirical investigation towards efficient multi-domain language model pre-training. arXiv preprint arXiv:2010.00784 , 2020.

Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, March 2021.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Cameron B Browne, Edward Powley, Daniel Whitehouse, Simon M Lucas, Peter I Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. A survey of monte carlo tree search methods. IEEE Transactions on Computational Intelligence and AI in games , 4(1):1-43, 2012.

Thomas Carta, Clément Romac, Thomas Wolf, Sylvain Lamprier, Olivier Sigaud, and PierreYves Oudeyer. Grounding large language models in interactive environments with online reinforcement learning. arXiv preprint arXiv:2302.02662 , 2023.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetuned language models. arXiv preprint arXiv:2210.11416 , 2022.

Matthias De Lange, Rahaf Aljundi, Marc Masana, Sarah Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. A continual learning survey: Defying forgetting in classification tasks. IEEE transactions on pattern analysis and machine intelligence , 44(7):33663385, 2021.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm.int8(): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339 , 2022.

B Roy Frieden. Science from Fisher information , volume 974. Citeseer, 2004.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 , 2020.

Xiaofeng Gao, Ran Gong, Tianmin Shu, Xu Xie, Shu Wang, and Song-Chun Zhu. Vrkitchen: an interactive 3d virtual environment for task-oriented learning. arXiv , abs/1903.05757, 2019.

William H Guss, Brandon Houghton, Nicholay Topin, Phillip Wang, Cayden Codel, Manuela Veloso, and Ruslan Salakhutdinov. Minerl: a large-scale dataset of minecraft demonstrations. In Proceedings of the 28th International Joint Conference on Artificial Intelligence , pages 2442-2448, 2019.

Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu. Reasoning with Language Model is Planning with World Model. NeurIPS , 2023.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Zhiting Hu and Eric P. Xing. Toward a 'Standard Model' of Machine Learning. Harvard Data Science Review , 4(4), oct 27 2022. https://hdsr.mitpress.mit.edu/pub/zkib7xth.

Wenlong Huang, Pieter Abbeel, Deepak Pathak, and Igor Mordatch. Language models as zero-shot planners: Extracting actionable knowledge for embodied agents. In International Conference on Machine Learning , pages 9118-9147. PMLR, 2022.

Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tompson, Igor Mordatch, Yevgen Chebotar, Pierre Sermanet, Tomas Jackson, Noah Brown, Linda Luu, Sergey Levine, Karol Hausman, and brian ichter. Inner monologue: Embodied reasoning through planning with language models. In Karen Liu, Dana Kulic, and Jeff Ichnowski, editors, Proceedings of The 6th Conference on Robot Learning , volume 205 of Proceedings of Machine Learning Research , pages 1769-1782. PMLR, 14-18 Dec 2023.

Anssi Kanervisto, Stephanie Milani, Karolis Ramanauskas, Nicholay Topin, Zichuan Lin, Junyou Li, Jianing Shi, Deheng Ye, Qiang Fu, Wei Yang, et al. Minerl diamond 2021 competition: Overview, results, and lessons learned. NeurIPS 2021 Competitions and Demonstrations Track , pages 13-28, 2022.

Yash Kant, Arun Ramachandran, Sriram Yenamandra, Igor Gilitschenski, Dhruv Batra, Andrew Szot, and Harsh Agrawal. Housekeep: Tidying virtual households using commonsense reasoning. In Computer Vision - ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXXIX , page 355-373, Berlin, Heidelberg, 2022. Springer-Verlag.

James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences , 114(13):3521-3526, 2017.

Eric Kolve, Roozbeh Mottaghi, Winson Han, Eli VanderBilt, Luca Weihs, Alvaro Herrasti, Matt Deitke, Kiana Ehsani, Daniel Gordon, Yuke Zhu, et al. Ai2-thor: An interactive 3d environment for visual ai. arXiv preprint arXiv:1712.05474 , 2017.

Liunian Harold Li, Mark Yatskar, Da Yin, Cho-Jui Hsieh, and Kai-Wei Chang. Visualbert: A simple and performant baseline for vision and language. arXiv preprint arXiv:1908.03557 , 2019.

Shuang Li, Xavier Puig, Chris Paxton, Yilun Du, Clinton Wang, Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, et al. Pre-trained language models for interactive decision-making. Advances in Neural Information Processing Systems , 35:31199-31212, 2022.

Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out , pages 74-81, 2004.

Zichuan Lin, Junyou Li, Jianing Shi, Deheng Ye, Qiang Fu, and Wei Yang. Juewu-mc: Playing minecraft with sample-efficient hierarchical reinforcement learning. arXiv preprint arXiv:2112.04907 , 2021.

Jiacheng Liu, Skyler Hallinan, Ximing Lu, Pengfei He, Sean Welleck, Hannaneh Hajishirzi, and Yejin Choi. Rainier: Reinforced knowledge introspector for commonsense question answering. arXiv preprint arXiv:2210.03078 , 2022.

Ruibo Liu, Jason Wei, Shixiang Shane Gu, Te-Yen Wu, Soroush Vosoughi, Claire Cui, Denny Zhou, and Andrew M Dai. Mind's eye: Grounded language model reasoning through simulation. arXiv preprint arXiv:2210.05359 , 2022.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Ximing Lu, Sean Welleck, Jack Hessel, Liwei Jiang, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, and Yejin Choi. Quark: Controllable text generation with reinforced unlearning. Advances in neural information processing systems , 35:27591-27609, 2022.

Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation , volume 24, pages 109-165. Elsevier, 1989.

Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christopher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, et al. Webgpt: Browser-assisted question-answering with human feedback. arXiv preprint arXiv:2112.09332 , 2021.

OpenAI. Gpt-4 technical report. ArXiv , abs/2303.08774, 2023.

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems , 35:27730-27744, 2022.

Xavier Puig, Kevin Ra, Marko Boben, Jiaman Li, Tingwu Wang, Sanja Fidler, and Antonio Torralba. Virtualhome: Simulating household activities via programs. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 8494-8502, 2018.

Xavier Puig, Tianmin Shu, Shuang Li, Zilin Wang, Yuan-Hong Liao, Joshua B. Tenenbaum, Sanja Fidler, and Antonio Torralba. Watch-and-help: A challenge for social perception and human-{ai} collaboration. In International Conference on Learning Representations , 2021.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

Joshua Robinson, Christopher Michael Rytting, and David Wingate. Leveraging large language models for multiple choice question answering. arXiv preprint arXiv:2210.12353 , 2022.

Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich. Superglue: Learning feature matching with graph neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4938-4947, 2020.

Manolis Savva, Angel X Chang, Alexey Dosovitskiy, Thomas Funkhouser, and Vladlen Koltun. Minos: Multimodal indoor simulator for navigation in complex environments. arXiv preprint arXiv:1712.03931 , 2017.

Timo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu, Maria Lomeli, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. arXiv preprint arXiv:2302.04761 , 2023.

Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed Chi, Nathanael Schärli, and Denny Zhou. Large language models can be easily distracted by irrelevant context. arXiv preprint arXiv:2302.00093 , 2023.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016.

Ishika Singh, Valts Blukis, Arsalan Mousavian, Ankit Goyal, Danfei Xu, Jonathan Tremblay, Dieter Fox, Jesse Thomason, and Animesh Garg. Progprompt: Generating situated robot task plans using large language models. In Workshop on Language and Robotics at CoRL 2022 , 2022.

Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems , 33:3008-3021, 2020.

Alessandro Suglia, Qiaozi Gao, Jesse Thomason, Govind Thattai, and Gaurav Sukhatme. Embodied bert: A transformer model for embodied, language-guided visual task completion. arXiv preprint arXiv:2108.04927 , 2021.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax , May 2021.

Zihao Wang, Shaofei Cai, Anji Liu, Xiaojian Ma, and Yitao Liang. Describe, explain, plan and select: Interactive planning with large language models enables open-world multi-task agents. arXiv preprint arXiv:2302.01560 , 2023.

Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M Rush, Bart Van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks. In 4th International Conference on Learning Representations, ICLR 2016 , 2016.

Jeff Wu, Long Ouyang, Daniel M Ziegler, Nisan Stiennon, Ryan Lowe, Jan Leike, and Paul Christiano. Recursively summarizing books with human feedback. arXiv preprint arXiv:2109.10862 , 2021.

Yi Wu, Yuxin Wu, Georgia Gkioxari, and Yuandong Tian. Building generalizable agents with a realistic and rich 3d environment. arXiv preprint arXiv:1801.02209 , 2018.

Jiannan Xiang, Zhengzhong Liu, Yucheng Zhou, Eric Xing, and Zhiting Hu. ASDOT: Any-shot data-to-text generation with pretrained language models. In Findings of the Association for Computational Linguistics: EMNLP 2022 , pages 1886-1899, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.

Jiannan Xiang, Xin Wang, and William Yang Wang. Learning to stop: A simple yet effective approach to urban vision-language navigation. In Findings of the Association for Computational Linguistics: EMNLP 2020 , pages 699-707, 2020.

Claudia Yan, Dipendra Misra, Andrew Bennnett, Aaron Walsman, Yonatan Bisk, and Yoav Artzi. Chalet: Cornell house agent learning environment. arXiv preprint arXiv:1801.07357 , 2018.

Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. ReAct: Synergizing reasoning and acting in language models. In International Conference on Learning Representations (ICLR) , 2023.

Rowan Zellers, Ari Holtzman, Matthew E Peters, Roozbeh Mottaghi, Aniruddha Kembhavi, Ali Farhadi, and Yejin Choi. PIGLeT: Language grounding through neuro-symbolic interaction in a 3d world. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) , pages 2040-2050, 2021.

Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593 , 2019.

A Appendix

A.1 VirtualHome

The complete format of an executable action step in VirtualHome is <char{char_id}> [Action] <Object> (Object_id) . Specifically, char_id specifies which agent to execute the action when multiple agents are in the world model at the same time. Action should be a supported atomic action in VirtualHome. Object is the object with which the agent interacts. Each object in the environment is assigned an Object_id to distinguish it from others of the same object class. We designed a template for each action to transform them into natural text for LMs finetuing. The full list of executable actions can be found in Table 4. Note that in the list, we omit <char{char_id}> and (Object_id) for simplicity.

A.2 Acitivity Goal And Predicate

The goal of an household activity in VirtualHome consists of several predicates. Each predicate represents a condition of one object or a relation between two objects. For example, OPEN(coffe

maker) means the coffee maker is open, and ON(apple, table) means an apple is on the table. The goal is only achieved when all the predicates are achieved. We collected activities and goals from RobotHow.

A.3 Data Format and Prompts

Following Chung et al. [8], we use instructions with in-context exemplars as prompts. Specifically, the instruction, the question context, and the answer will be provided in each exemplar, and the full prompt will contain multiple such exemplars for in-context learning. The format of the data and the exemplar for each task is provided below.

A.3.1 Plan Generation

Data Example

In-context Exemplar

A.3.2 Housework QA

Data Example

In-context Exemplar

A.3.3 Negation Housework QA

Data Example

In-context Exemplar

A.3.4 Activity Recognition

Data Example

In-context Exemplar

A.3.5 Activity Inference

Data Example

In-context Exemplar

A.3.6 Counting

Data Example

In-context Exemplar

Given a sequence of actions in a house, and a question about what items are located in a specific place. Answer the number of items and list the items.

A.3.7 Counting QA

Data Example

In-context Exemplar

A.3.8 Object Path Tracking

Data Example

In-context Exemplar

A.3.9 Object Location QA

Data Example

In-context Exemplar

A.4 Hyperparameters

For both GPT-Neo-1.3B and GPT-J-6B, we use a learning rate of 8 × 10 -5 and a batch size of 20. The weights for plan generation, activity recognition, counting, and object path tracking are 1.0, 0.7, 1.0, and 1.0, respectively. We trained GPT-Neo-1.3B for 3 epochs with the EWC coefficient λ = 0 . 5 in Equation 4. For GPT-J-6B, we trained it for 5 epochs with λ = 2 . With our approach, it takes 40 minutes to train a GPT-Neo and 220 minutes to train a GPT-J. We used a rank of 8 and coefficient of 32 for LoRA's hyperparameters.

A.5 bAbI Dataset

We include 8 tasks from bAbI that test embodied knowledge. They are: One Supporting Fact, Two Supporting Fact, Three Supporting Fact, Counting, Lists/Sets, Simple Negation, Time Reasoning, Positional Reasoning. Examples for each task are shown in Table 5.

Task 1: Single Supporting Fact

Mary went to the bathroom.

John moved to the hallway. Mary travelled to the office.

Where is Mary? A:office

Task 3: Three Supporting Facts

John picked up the apple.

John went to the office.

John went to the kitchen.

John dropped the apple.

Where was the apple before the kitchen? A:office

Task 5: Lists/Sets

Daniel picks up the football.

Daniel drops the newspaper.

Daniel picks up the milk.

What is Daniel holding? milk, football

Task 7: Time Reasoning

In the afternoon Julie went to the park. Yesterday Julie was at school. Julie went to the cinema this evening. Where did Julie go after the park? A:cinema Where was Julie before the park? A:school

Task 2: Two Supporting Facts

John is in the playground.

John picked up the football.

Bob went to the kitchen.

Where is the football? A:playground

Task 4: Counting

Daniel picked up the football.

Daniel dropped the football.

Daniel got the milk.

Daniel took the apple. A: office

How many objects is Daniel holding? A: two

Task 6: Simple Negation

Sandra travelled to the office.

Fred is no longer in the office.

Is Fred in the office? A:no

Is Sandra in the office? A:yes

Task 8: Positional Reasoning

The triangle is to the right of the blue square. The red square is on top of the blue square. The red sphere is to the right of the blue square.

Is the red sphere to the right of the blue square? A:yes

Is the red square to the left of the triangle? A:yes

A.6 Results of Main Experiments and Ablation Studies

Experimental results on our constructed downstream tasks are shown in Table 6, and the results on bAbI are shown in Table 7. We also show the results of ablation studies in Table 8.

A.7 Human Evaluations

We conduct human evaluations on plan generation for GPT-J model. Following Huang et al. [18] we asked 3 people to annotate whether each task can be completed using a generated plan. We randomly sampled 150 tasks and asked each person to annotate 50 of them. The Results show that the base GPT-J model can only achieve 24.0% accuracy, while the finetuned model can achieve 62.4%. The higher planning accuracy demonstrates the superior task planning ability of our model.

A.8 SuperGLUE Results

We evaluate the base GPT-J model and our finetuned model on appropriate SuperGLUE tasks, e.g. , that can be formulated as a multi-choice QA task without prompt engineering. Our model's performance matches and even outperforms the baseline, showing our model retains the general language capability.

A.9 Broader Impact

Like other generation systems, the language model trained by our approach is susceptible to producing unintended output when confronted with harmful input, such as unethical text or input intended for adversarial attacks. Therefore, we strongly advise against utilizing our approach outside of controlled research environments until these risks have been mitigated. It is important to note that a thoughtless deployment of our method could potentially enable malicious exploitation of the underlying language models. Thus, precautions, such as implementing a filtering mechanism, must be taken.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new collection of data and tasks from an embodied environment to enhance the embodied reasoning ability of pre-trained language models. The paper also designs a fine-tuning strategy that combines the advantages of EWC and LoRA for stable fine-tuning. The experiments show that the proposed method can efficiently fine-tune GPT-J and outperform much larger language models.\n\nSTRENGTHS:\nThis paper designs two ways of data collection in the embodied environment that address goal-specific and general cases in real-world applications.\n\nThis paper designs a comprehensive reasoning task set based on the collected data to evaluate the reasoning performance from different aspects (plan generation, object tracking, etc.)\n\nThis paper adopts a suitable parameter updating strategy (EWC_LoRA) for language model fine-tuning, which demonstrates time and memory efficiency.\n\nWEAKNESSES:\nThe paper lacks experiments on why EWC is needed. The motivation for including EWC is to avoid overfitting downstream tasks or catastrophic forgetting on the pretraining task, but LoRA has a similar purpose. From the experiments in Table 2, EWC_LoRA has a lower performance than LoRA except for a slight improvement in perplexity.\n\nQUESTIONS:\n1. Related to weakness: It will be good to see whether there are any other benefits gained from EWC.\n\n2. How do you ensure that the answers of the negation Housework QA are truly irrelevant?\n\n3. The activity recognition performance of all methods is quite similar. Does this indicate that the activity types are easily distinguishable without requiring a full description of the experience, for example, just keywords at the end are sufficient for the task? Have you tried to include confusing activities in the evaluation set?\n\n4. How do you introduce the confusion term (counting task, confusing unseen of plan generation)? Do you have any guideline or do you just randomly inject phrases?\n\nLIMITATIONS:\nThe author mentioned the limitations of only collecting embodied experience from one world model.\n\nsocietal impact: did not find negative impacts, probably if the data collection is extended to the real world in the future, it should avoid collecting personal or sensitive experiences.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a training paradigm termed \"finetuning with Embodied Experiences from World Models (E2WM)\" to enhance the abilities of language models (LMs) in reasoning and planning tasks associated with physical environments. The authors argue that LMs trained solely on large-scale text corpora lack the embodied knowledge necessary for robust performance in such tasks. To address this, they propose leveraging world models, specifically the VirtualHome simulator, to collect diverse embodied experiences and use them to construct fine-tuning tasks, such as plan generation, activity recognition, counting, and object path tracking. The proposed fine-tuning method incorporates Elastic Weight Consolidation (EWC) regularization with low-rank adaptation (LoRA) to preserve the models' generality and avoid catastrophic forgetting and make the training process more efficient.\n\nSTRENGTHS:\n- The paper introduces an approach to enhance off-the-shelf LLMs with embodied knowledge. The experiments involve both goal-oriented planning and random exploration. in collecting embodied experiences. While EWC and LoRA are existing methods (hence the paper may be perceived as with limited technical novelty), this work shows that their incorporation into the fine-tuning process helps retain the models' general knowledge and capabilities while adapting them to new tasks. \n- The related work seems sufficient, albeit a section on catastrophic forgetting, e.g., [1] would be a plus.\n- The paper provides an evaluation of the finetuned LMs on both seen and unseen tasks, and several ablation studies. The generalizability of the models is also tested, as well as their performance on the original pretraining data to ensure core language modeling abilities are retained. \n\n[1] Korbak, Tomasz, Hady Elsahar, German Kruszewski, and Marc Dymetman. \"Controlling conditional language models without catastrophic forgetting.\" In International Conference on Machine Learning, pp. 11499-11528. PMLR, 2022.\n\nWEAKNESSES:\n- Lack of comparison: Only two GPT-based LLMs are used. Would be valuable to consider a more diverse set of recent LLMs. The paper also does not compare the proposed E2WM paradigm with existing methods for enhancing LMs with embodied knowledge, such as [1].\n- Scalability to larger models: Although the EWC-LoRA approach is designed to improve efficiency, it is not clear how well it scales to larger LM architectures. Further investigation into the scalability of the proposed method would be beneficial.\n- Limitations on benchmarks: Evaluation is performed on bAbI but there are several embodied benchmarks and simulators that could be considered. Doing so would allow evaluation in realistic embodied settings, and compare other embodied models proposed (such as the Episodic Transformer [2] or [1]).\n- Experimental Results: The number of tasks and models evaluated is relatively small and the presentation of the results could be clearer, e.g., adding numbers to bars in Figs. 3 and4. It seems that ChatGPT outperforms the fine-tuned models, so overall the model capacity appears to be an important factor, and the proposed fine-tuning shows marginal improvements in some of the tasks. \n\n[1] Lin, B.Y., Huang, C., Liu, Q., Gu, W., Sommerer, S. and Ren, X., 2023, June. On grounded planning for embodied tasks with language models. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 37, No. 11, pp. 13192-13200).\n[2] Pashevich, Alexander, Cordelia Schmid, and Chen Sun. \"Episodic transformer for vision-and-language navigation.\" In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 15942-15952. 2021.\n\nQUESTIONS:\nHave the authors considered an evaluation on actual embodied benchmarks? \nWhat is the motivation for choosing Babi as the benchmark dataset?\n\nLIMITATIONS:\nThere is a very brief discussion on limitations as part of the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to incorporate world models into large language models to enable understanding of object permanence and planning capabilities that are missing in text-only models. Specifically, the authors utilize embodied environments (VirtualHome) to collect training examples including goal-oriented planning (plan generation, activity recognition) and random exploration (counting and object plan tracking). To avoid overfitting, this paper finetunes GPT-Neo-1 3B and GPT-J-6B models using LoRA and EWC. Evaluated on seen and unseen tasks, results show that the proposed method after finetuning outperforms the corresponding text-only baselines, and match or outperform ChatGPT performance under few-shot setting.\n\nSTRENGTHS:\n1. This paper is motivated by an interesting limitation in text-only language models, which is the lack of embodied knowledge in pre-training. The proposed fine-tuning dataset, along with EWC-LoRA method, improve the model performance on challenging embodied reasoning and understanding tasks, without dramatically increasing the original language model perplexity.\n2. The paper includes some interesting ablation studies to discuss how much improvement is observed from each fine-tuning task.\n\nWEAKNESSES:\n1. The contribution and findings of the paper is limited. Although the limitation of the LLMs is well motivated in the paper, fine-tuning on similar distribution to improve model performance and thus outperform baselines is expected. Furthermore, fine-tuning with LoRA to improve efficiency, as well as EWC to reduce overfitting have been well studied. Given that the (relatively weak) baseline is the original text-only model (and a few-shot prompted ChatGPT model), it is not convincing that the proposed method (either finetuning data or method) is generalizable to broader embodied tasks (e.g., compared to an embodied or multi-modal model).\n2. Some important details are missing and the evaluation is not very convincing. The paper is phrased to incorporate world model into LLMs, but it is confusing (without explicit explanation) that both training and evaluation are not with the embodied environment. Rather, only partial of the data is collected using VirtualHome. More importantly, it is not clear how the data are constructed (for example, how to sample \"irrelevant actions\" for counting QA and statistics of the training and eval sets) and what the held out set is for training the model. For evaluation, I understand the automatic eval is the most convenient metic, but have been widely adapted, metrics such as Rough-L for tasks like plan generation may not reflect model performance on complex text generation tasks. See more details in the questions below.\n\nQUESTIONS:\n1. Have you done other evaluation on the base language model (i.e., NLP benchmarks such as superglue) apart from perplexity? \n2. Why do you think EWC only outperform LoRA (in table 2)? As a regularization method, EWC constraint parameter updates and thus should understand LoRA by intuition.\n3. Do you have a naive fine-tuning baseline to compare to?\n4. Do you have detailed analysis on when and why the model improves performance on both seen and unseen tasks?\n\nLIMITATIONS:\nThe paper briefly mentioned limitations.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new method to improve the embodied planning capability of large language models (LLM) via adding both goal-oriented planning and random exploration data from a world model/simulator as well as the EWC-LORA regularizer that prevents catastrophic forgetting of pre-training tasks. The new regularizer EWC-LORA proposed by the authors is not only time and memory-efficient compared prior KL constraint-based regularization and also leads to negligible perplexity increase and better downstream task performance. The authors show superior performance of the method over base LLMs in constructed embodied planning and activity description tasks based on the VirtualHome simulator and also the bAbI a dataset for testing multiple types of knowledge and abilities including embodied knowledge, logic reasoning, linguistic knowledge, etc. Note that the method can only outperform much bigger LLMs such as ChatGPT in many of the scenarios.\n\nSTRENGTHS:\n1. The authors presents a simple yet effective way to improve LLM's embodied planning capability without sacrificing the abilities during pre-training under a small/negligible compute/memory cost. Such idea is neat and of great practical importance. I think the method will be quite significant for the future directions of making LLMs better at embodied tasks.\n\n2. The authors have done extensive experiments in both customized and existing benchmarks/datasets to show that the proposed method can outperform base models without fine-tuning with embodied data. It is impressive that the authors show that the small models with embodied planning data fine-tuning can outperform large LLMs such as ChatGPT in many scenarios.\n\n3. The authors also performed detailed ablation studies to show the importance of the proposed EWC-LORA regularizer via comparing to KL divergence and EWC only. The ablation studies on including various data mixture also make the paper more complete.\n\nWEAKNESSES:\n1. I think the authors should compare to stronger baselines that also consider improving language model's embodied decision-making and reasoning capabilities such as [1,2]. Improvement over base models is good but not that surprising and convincing.\n\n2. The authors should consider other embodied decision-making benchmarks such as ALFWorld, which is used in [1. 2]. This would provide a clearer picture of the comparison between the method and prior approaches. It would also add more tasks in the empirical evaluation, which can further validate the method.\n\n[1] Shinn, Noah, Beck Labash, and Ashwin Gopinath. \"Reflexion: an autonomous agent with dynamic memory and self-reflection.\" arXiv preprint arXiv:2303.11366 (2023).\n\n[2] Yao, Shunyu, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. \"React: Synergizing reasoning and acting in language models.\" arXiv preprint arXiv:2210.03629 (2022).\n\nQUESTIONS:\nPlease address the comments listed in the section above.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to enhance language models by finetuning them on “embodied experiences”, which are textual data generated by household activity simulator, which is referred to as “world models” in the paper. Through evaluations on several tasks (e.g., planning, object tracking), it is shown that finetuning on these “embodied experiences” leads to better performance compared to larger models (e.g., ChatGPT) that are not trained on these data. To minimize catastrophic forgetting and improve training efficiency, the paper also proposes to combine “elastic weight consolidation” and “low-rank adapters” to finetune the language models.\n\nSTRENGTHS:\n- The premise of the paper is interesting, novel, and promising — because language models are not trained on embodied data, they might be less robust to those scenarios concerning interactions with the environments (i.e., the “embodied settings”).\n- The idea of using household activity simulator as a “world model” is interesting and likely significant in the context of enhancing language models.\n- Moreover, the authors conduct thorough experiments that support the central claim.\n- The writing and presentation of the paper are also clear.\n\nWEAKNESSES:\nDespite the strengths of the paper, below are some concerns for the problem settings and evaluations:\n\n- While it is shown on GPT-Neo and GPT-J (which are relatively small language models in today’s standard) that the proposed approach improves their capabilities on results such as bAbI, as also indicated in the paper evaluations, larger models (i.e. ChatGPT) which are not trained on these embodied data also attain similar performance. Because the constructed tasks only require scene context in the text form, it is unclear whether similar delta will be seen on larger models. Or put this in other words: will this improvement diminish with larger scale training even without embodied experiences? The reviewer would like to note that due to practical reasons, it is understandable that experiments like these may not be done for larger LMs, but it is worth further discussions or experiments in the paper to support the claim.\n- Another weakness is that the training tasks need to be hand-curated for the “embodied experience”, which includes planning, activity recognition, counting, and object path tracking. This is unlike how autoregressive LMs are trained, which only requires one main self-supervised task of next token prediction. This brings to the question of how scalable the proposed approach is, as the broader “embodied experiences” include tasks of much higher diversity, in addition to including other modalities such as vision.\n\nQUESTIONS:\nAn important baseline is ChatGPT, which is not trained on “embodied experiences”, but it doesn’t seem like the prompt for querying ChatGPT is provided. In contrast, the paper notes that they used few-shot prompting for the smaller LMs. If they are given the same prompt, would the performance differ?\n\nLIMITATIONS:\nThe limitations are not adequately discussed. See the comments in the “weakness” section above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-11"}