SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jiannan Xiang*, Tianhua Tao*, Yi Gu*, Tianmin Shu, Zirui Wang*, Zichao Yang, Zhiting Hu*  
*UC San Diego, *UIUC, MIT, *JHU, CMU
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).<sup>1</sup>
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  $\mathrm{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].
\*Equal contribution. 1The code is available at https://github.com/szxiangjn/world-model-for-language-model. 2Based on GPT-3.5-turbo.
Figure 1: Examples of tasks requiring embodied knowledge (upper), and an overview of our approach (bottom). In the task examples, blue text indicates the useful information for answering the question.
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 \rightarrow 51.23$  Rouge-L on plan generation task,  $30.41\% \rightarrow 67.01\%$  accuracy on counting task), without suffering performance degradation on the pretraining dataset ( $3.443 \rightarrow 3.537$  perplexity on Pile test subset [12]). Moreover, the small GPT-J-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.
World Model. The term "world model" generally refers 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.
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 household 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
Figure 2: The illustration of goal-oriented planning (left) and random exploration (right) in our training paradigm. In MCTS, the path in orange represents the final plan generated by the planner.
Random Exploration
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  $\mathbf{x}$  is the input (e.g., the initial condition in plan generation) and  $\mathbf{y} = \{y_1,\dots,y_M\}$  is the label (e.g., the stepwise action sequence), we finetune LMs by assigning different weights to different tasks:
$$
\mathcal {L} _ {V} = \sum_ {v \in V} \alpha_ {v} \sum_ {m = 1} ^ {M} \log P \left(y _ {m} \mid \mathbf {y} _ {<   m}, \mathbf {x}\right), \tag {1}
$$
where  $\mathcal{L}$  is the loss function;  $V$  is the task set; and  $\alpha_{v}$  is the weight for task  $v$ . Following Flan-T5 [8],  $\mathbf{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:
$$
F _ {i, i} = \frac {1}{N} \sum_ {j = 1} ^ {N} \left(\frac {\partial \mathcal {L} _ {U} ^ {(j)}}{\partial \theta_ {U , i} ^ {*}}\right) ^ {2}, \tag {2}
$$
$$
\mathcal {L} (\theta) = \mathcal {L} _ {V} (\theta) + \lambda \sum_ {i} F _ {i, i} \left(\theta_ {i} - \theta_ {U, i} ^ {*}\right) ^ {2}, \tag {3}
$$
where  $\mathcal{L}$  is the loss function,  $F$  is the fisher matrix,  $\lambda$  is the hyperparameter,  $i$  and  $j$  are the indices for parameters and data samples, respectively, and  $\theta$  and  $\theta_U^*$  are currently trained parameters and frozen task  $U$  parameters, respectively. Notice that the first term  $\mathcal{L}_V(\theta)$  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^{*} \in \mathbb{R}^{r \times d}$  are the trained weight matrix and frozen weight matrix, respectively; and  $B \in \mathbb{R}^{r \times k}$ ,  $A \in \mathbb{R}^{k \times d}$  are two low-rank matrices with  $k \ll \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  $\theta_{i}$  in Equation 3 is the element of  $W$ , and  $\theta_{U,i}^{*}$  is that of  $W^{*}$ . Therefore,  $\theta_{i} - \theta_{U,i}^{*}$  is the element of  $H$ . We can thus transform Equation 3 into the final formula of EWC-LoRA method:
$$
\mathcal {L} (\theta) = \mathcal {L} _ {V} (\theta) + \lambda \sum_ {i} F _ {i, i} h _ {i} ^ {2}, \tag {4}
$$
where  $h_i = \theta_i - \theta_{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].
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: