SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zirui Zhao Wee Sun Lee David Hsu
National University of Singapore {ziruiz, leews, dyhsu}@comp.nus.edu.sg
Large-scale task planning is a major challenge. Recent work exploits large language models (LLMs) directly as a policy and shows surprisingly interesting results. This paper shows that LLMs provide a commonsense model of the world in addition to a policy that acts on it. The world model and the policy can be combined in a search algorithm, such as Monte Carlo Tree Search (MCTS), to scale up task planning. In our new LLM-MCTS algorithm, the LLM-induced world model provides a commonsense prior belief for MCTS to achieve effective reasoning; the LLM-induced policy acts as a heuristic to guide the search, vastly improving search efficiency. Experiments show that LLM-MCTS outperforms both MCTS alone and policies induced by LLMs (GPT2 and GPT3.5) by a wide margin for complex, novel tasks. Further experiments and analyses on multiple tasks—multiplication, travel planning, object rearrangement—suggest minimum description length (MDL) as a general guiding principle: if the description length of the world model is substantially smaller than that of the policy, using LLM as a world model for model-based planning is likely better than using LLM solely as a policy.<sup>1</sup>
Consider, for example, an autonomous robot butler in a household environment. The human user sits in the living room and asks the robot to "Put fruits into the fridge." The robot looks for fruits, such as apples, peaches, etc., which may be on a table in the dining room, on the kitchen counter, but unlikely in a wardrobe in the bedroom. To complete the task, the robot has to consider fruits' likely locations as well as their spatial relations, traverse these locations efficiently, and finally put them into the fridge. A household environment typically contains hundreds of movable items and locations, resulting in a huge search space that makes the task very challenging for the robot.
Recently, multiple attempts exploiting pre-trained large language models (LLMs) show surprisingly interesting results for such tasks [22, 18, 2, 19, 5, 44]. Their underlying idea is simple: treat the LLM as a policy and query it directly for the next actions, given the history of past actions and observations. We call this strategy  $L$ -Policy, which exploits LLMs' vast commonsense knowledge to circumvent the challenge of searching a very large space. In our example task, L-Policy may simply instruct the robot to move to the hallway and then to the kitchen. Even though LLMs are trained on internet-scale data, this strategy shows limits in generalization [10, 4], especially when encountering uncommon, complex tasks. Alternatively, we may use LLMs' knowledge to build a world model and apply a planning algorithm to the model. The world model may contain, e.g., a belief over the target object's location, which biases the search and drastically improves search efficiency. We call this strategy  $L$ -Model. L-Model's performance depends on two critical preconditions: the accuracy of the world model and the efficiency of the planning algorithm. The former is a question of sample complexity for learning, and the latter is that of computational complexity.
<sup>1</sup>The code and supplementary materials are available at https://llm-mcts.github.io.
This paper presents LLM-MCTS (shown in Fig 1), which combines the ideas of L-Model and L-Policy for large-scale task planning. Like L-Model, LLM-MCTS uses an LLM to build a commonsense world model; it then uses the model to perform Monte Carlo Tree Search (MCTS) [8] online for the next actions. During the tree search, LLM-MCTS chooses the promising action branches heuristically by querying the LLM. This is similar in spirit to L-Policy. While L-Policy commits to the actions chosen by the LLM for execution, LLM-MCTS uses these choices only as a search heuristic.
We evaluate LLM-MCTS in VirtualHome [24], a standard household activity simulation platform widely used in earlier work [18, 22, 25, 33]. The evaluation set consists of 800 randomly generated large-scale, partially observable object rearrangement tasks. In each task, a robot aims to fetch a common household item with an unknown location and place it in a designated container. Our main experimental findings are summarized below:
F1. L-Model performs poorly. There are two possible reasons. One is model inaccuracy: the robot has an incorrect belief of the target object's location. The other is huge search space size, beyond the reach of even the state-of-the-art MCTS algorithm. Further experiments indicate that search space size is the main cause.
F2. L-Policy performs reasonably with both GPT2 and GPT3.5, but the performance degrades quickly for novel, complex tasks. This generally corroborates with earlier results [22, 18, 2, 19].
F3. LLM-MCTS outperforms L-Model. This clearly shows the benefit of using the LLM as a heuristic policy to guide the search.
F4. LLM-MCTS outperforms L-Policy, especially for novel, complex tasks. LLM-MCTS basically combines L-Model and L-Policy. Since the L-Model performs very poorly on its own, why does the combination outperform L-Policy? One explanation is that search space size is the main cause of L-Model's poor performance. The LLM-induced world model is sufficiently accurate; tree search with this world model, when limited to the neighbourhood of the LLM-induced policy, provides improved performance over L-Policy.
The explanation for (F4) begs a new question: with an efficient planning algorithm for large search space, would L-Model outperform L-Policy? To answer this question, we study two related, simpler tasks: multiplication of two large numbers and multi-hop travel planning. We discuss multiplication here and travel planning in Section 4.1. A decimal number is described as a sequence of  $n$  digits,  $(d_{n - 1}, d_{n - 2}, \dots, d_0)$ . There are two methods of implementing multiplication with an LLM. The first one corresponds to L-Policy. We represent the multiplication function as a table. Each row or column corresponds to a number. The table entry is the multiplication of two numbers, obtained by querying an LLM. Experimentally, GPT4 performs single-digit multiplication perfectly with  $100\%$  accuracy, 2-digit multiplication with  $99\%$  accuracy, 4-digit multiplication with merely  $4\%$  accuracy, and fails almost completely on 5-digit multiplication [10]. The second approach uses LLM-derived small single-digit multiplication tables, which GPT4 performs with  $100\%$  accuracy. To multiply multi-digit numbers, it applies the long multiplication algorithm with the single-digit table. This method corresponds to L-Model. While long multiplication differs from planning in algorithmic details, it plays the same role in L-Model and is highly efficient. Clearly, this second method achieves  $100\%$  accuracy for arbitrarily large numbers, provided that the single-digit multiplication table is accurate. So, the L-Model outperforms L-Policy for the multiplication task, contrary to the finding for object rearrangement tasks.
How do we choose between L-Model and L-Policy then? One idea is the minimum description length (MDL) principle. Theoretical analysis suggests that a hypothesis with a shorter description length has a smaller generalization error and is preferred [28]. For multiplication, the method corresponding to L-Policy uses a large table of  $O(10^{2n})$  entries. It takes  $O(n10^{2n})$  bits to represent it. The method corresponding to the L-Model uses a single-digit multiplication table of constant size. The long multiplication algorithm can be encoded in any reasonable programming language with constant size. So, the total representation size is constant. According to MDL, the L-Model has a smaller generalization error than the L-Policy for multiplication, with sufficiently large  $n$ . This is fully consistent with experimental results [10]. The analysis of travel planning provides further evidence (Section 4.1).
In summary, LLM-MCTS combines the ideas of L-Model and L-Policy, outperforming either alone for complex task planning, particularly, object rearrangement (Sections 2 and 3). To choose between L-Model and L-Policy, MDL provides a useful guiding principle (Section 4). In essence, simplicity is preferred, a well-known general principle in machine learning.
Figure 1: Overview of LLM-MCTS. For each simulation in the MCTS, we sample from the commonsense belief to obtain an initial state of the world and use the LLM as heuristics to guide the trajectory to promising parts of the search tree.
2 LLM-MCTS: Monte Carlo planning with commonsense knowledge
We aim to solve task-planning problems in large-scale domains with partial observation. One example is object rearrangement tasks [3] in household environments. It is a meaningful and challenging problem with a large-scale and long-term planning horizon. It has many practical implications in everyday life [3, 34, 39, 20, 17], such as setting the table, tidying up the room, loading the dishwasher, etc.. To solve the problem, we present LLM-MCTS (shown in Fig 1), which combines the L-Model and L-Policy for large-scale planning. It uses LLM to build a commonsense world model to perform MCTS for reasoned planning and uses L-Policy to guide the MCTS and reduce the large search space.
2.1 Task planning
We focus on task-planning problems with partial observation and large-scale domains. The problem can be formulated as a Partially Observable Markov Decision Process (POMDP):  $(S, A, \Omega, T, O, R, \gamma)$ . The state space  $S$  defines the state of the robot and its environment. The action space  $A$  defines the action that the robot can do.  $\Omega$  is the observation space.  $T$  defines the transition function of states, which we assume to be given.  $O$  is the observation function that provides partial information about a state.  $R(s, a)$  is the reward function determined by the action  $a$  taken at the state  $s$ . The discount factor is specified by  $\gamma$ . The history trajectory  $h_t$  at time step  $t$  consists of a sequence of executed actions and received observations up to time  $t - 1$ ,  $h_t = (o_0, a_0, o_1, a_1, \ldots, o_{t-1}, a_{t-1})$ . The objective is to find an optimal policy  $\pi^*(h_t)$  that maximizes the expected cumulative rewards  $\pi^*(h_t) = \arg \max_{a \in A} \mathbb{E}\left[\sum_{t=0}^{\infty} \gamma^i R(s_{t+i}, a_{t+i}) | a_t = a\right]$ .
In this work, we focus on the object rearrangement task, though our approach is a general method for large-scale task planning. Object rearrangement is a representative embodied AI task [3, 34, 39, 20, 17] with various daily applications, such as setting the table, tidying up the room, loading the dishwasher, and more. It is a challenging task [3] as the robot must navigate, locate target objects and positions, and execute multi-step planning. As with hundreds of items and containers in the domain, identifying target objects can be challenging. Even with known state information, it requires long-horizon planning to achieve the goal. In addition, the vast number of domain objects leads to a large action space, given the actions are to interact with objects. The vast action space produces an exponentially large search tree, making the planning extremely challenging.
Following the approach of [22, 25, 33], we model the task as a POMDP as outlined above. The state  $S$  comprises variables denoting the positions of the robot, movable items, and containers. Actions  $A$  encompass five predefined actions from VirtualHome, parameterized by object/container/rooms: (1) pick(object), where the robot collects an observed, proximate object; (2) place(object,placement), allowing the robot to set a picked object nearby or inside an open container; (3) open/container) and (4) close/container), for interacting with an observed, nearby containers; and (5) move(room/object/container), where the robot relocates within a room or near an observed object/container. Given our assumption of the robot's familiarity with house structures and the manageable size of the house, the robot can move directly to designated rooms. The deterministic transition  $T$  is pre-defined by actions. Partial observation  $O$  enables the robot to discern object/container positions within its room or an opened container at its location. The objective is to reorganize household items based on verbal instructions, represented by a reward of  $R$  for achieving the desired item arrangement.
2.2 LLM as a commonsense world model
A commonsense prior belief of states can improve the effectiveness of object and location searches by prioritizing the search to appropriate locations. Our approach utilizes LLM's commonsense knowledge to generate the initial belief of states, which is updated with each action and observation in the real world. MCTS samples from the belief in simulation to estimate the value of the action.
Initial belief of state. We use object-centric state representation and categorize the objects in the house as moveable objects (e.g., apples), containers (e.g., fridge), and surfaces (e.g., kitchen table). The states of a moveable object might be inside the containers or on the surfaces. The containers and surfaces should be inside a room. Similar to [22, 25], we maintain the belief in object-centric graphs, where nodes are objects and edges describe abstract-level relationships (e.g., apples are inside fridge, fridge is inside kitchen) between objects and rooms. Details are in the Appendix C.2.
Assume a dataset  $\mathcal{D}$  is accessible, containing expert actions and observations in similar household environments to solve daily tasks. LLMs can use the observations in the data to know what are the objects in the house and predict their positions, forming the commonsense belief of the state. To achieve this, we find all the objects, containers, and surfaces that appeared in the dataset  $\mathcal{D}$  to form a list of objects  $\mathcal{D}_{\mathrm{obj}}$  using a unique name for all of them. To approximate  $b(s_0)$ , we ask the LLMs to sample the positions of objects  $M$  times. For each sample, we ask the LLM to predict the position of objects using  $\mathcal{D}_{\mathrm{obj}}$  and a fixed prompt. For instance, we ask LLM to complete "The containers in the apartment are: fridge, ..., The surfaces in the apartment are: kitchen counter, ..., Question: what are the possible positions of strawberry? Answer: inside fridge, inside pantry... Question: what are the possible positions of apple? Answer:" We use three prompt examples to provide example formats of the response. The exact prompts we used are provided in the appendix. As the responses from LLM are free-form natural language, we have to precisely map those expressions to  $\mathcal{D}_{\mathrm{obj}}$  for consistent state representation. Thus, we encode the names of objects in the LLM's response into embeddings using sentence-BERT  $f(\cdot)$  [26] and examine their cosine similarity to the unique name of objects in  $\mathcal{D}_{\mathrm{obj}}$ : CosineSim  $(e_i,e) = \frac{f(e_i)f(e)}{\|f(e_i)\|\|f(e)\|}$ , where  $e$  is the name of objects, containers, or surfaces in the LLM's response, and  $e_i \in \mathcal{D}_{\mathrm{obj}}$  are the unique names in the object list. We select the most similar expressions in  $\mathcal{D}_{\mathrm{obj}}$  to form the sampled state. For example, when querying the position of an apple, the LLM's response is "on the kitchen table," we use the above technique to translate "the kitchen table" to "kitchenable," a unique name in  $\mathcal{D}_{\mathrm{obj}}$ .
Goal. Similar to [40], we use LLMs to translate the natural language goal into a formal goal for MCTS. We use a fixed set of prompt examples for LLM to interpret natural language goals, such as "put one apple into the fridge" is translated as a tuple "(apple, inside, fridge)." For compositional instructions, it will translate it into multiple tuples, such as "put one apple on the kitchen table and one plate inside the dishwasher" is translated as "(apple, on, kitchentable), (plate, inside, dishwasher)." We precisely map the LLM-generated goal into the admissible expressions in  $\mathcal{D}_{\mathrm{obj}}$  for search using the same representation as the state. In MCTS, the goal is used to identify the reward. As the representations are the same, we can directly check whether the object's state is the same as the goal by string matching. If the goal is reached, it will receive a large positive reward, or 0 otherwise.
2.3 LLM as a heuristic policy
We use LLMs to play the role of  $\pi(a|h)$  in PUCT to guide the action selection in the simulation procedure. In this procedure, the LLM takes as input the examples in the dataset, the goal description, the current observation, and the history of actions, and then outputs the suggested action plan (e.g., "Next actions: move to the kitchen, open the fridge, ..."). Similar to [22], the observations and goal description are translated into English sentences. As the answer of LLM is from the conditional distribution of the following words given the context, it can also be viewed as a commonsense policy of actions to take conditioned on the context of tasks, observations, and completed actions. However, direct implementation and access to the probability value of the GPT-3.5 is not available. Thus, we propose an empirical policy distribution  $\hat{\pi}$  that uses sampling to approximate the policy distribution.
We sample the LLM for  $M$  times to approximate the policy probability distribution. For each sample, we query the LLM with prompt and trajectory history  $h$  and receive an answer of the following actions to take  $\alpha_{i}\sim \mathrm{LLM}(h,\mathrm{prompt})$ , where  $\alpha_{i}$  is the first action of the answer. The prompt examples are retrieved from the dataset according to the similarity to the current language instruction  $\ell$ . We use [26] to translate the instructions in the dataset  $\ell_i\in \mathcal{D}$  into embedding and examine their
Algorithm 1 LLM-MCTS
1: procedure SEARCH(h,b,T,N) 20: if  $\gamma^d < \epsilon$  or done = True then  
2:  $n\gets 0$  21: return 0  
3: while  $n < N$  do 22: end if  
4:  $s\sim b(s)$  23: if  $h$  is not in  $\mathcal{T}$  then  
5: SIMULATE(s,h,False,0, $\mathcal{T}$ ) 24:  $\mathcal{T}\gets \mathcal{T}\cup h,N(h)\gets 0$   
6:  $n\gets n + 1$  25:  $\forall a\in A,N(h,a)\gets 0,Q(h,a)\gets 0$   
7: end while 26: return ROLLOUT(s,h,done,d)  
8: return  $\mathrm{argmax}_{a\in A}Q(h,a)$  27: end if  
9: end procedure 28:  $\hat{\pi} (a|h)\leftarrow \text{QUERYLLMPOLICY}(h)$   
10: procedure ROLLOUT(s,h,done,d) 29:  $a^{*}\gets \operatorname {argmax}_{a\in A}Q(h,a) + c\hat{\pi} (a|h)\frac{\sqrt{N(h)}}{N(h,a) + 1}$   
11: if  $\gamma^d < \epsilon$  or done = True then  
12: return 0 30:  $(s',o,r,\text{done})\sim \mathcal{G}(s,a^{*})$   
13: end if 31:  $h^\prime \gets \text{PUSHBACK}(h,[a^*,o]),d^\prime \gets d + 1$   
14:  $a\sim \pi_{\text{rollout}}(h,\cdot)$  32:  $R\gets r + \gamma \cdot \text{SIMULATE}(s',h',\text{done},d',\mathcal{T})$   
15:  $(s',o,r,\text{done})\sim \mathcal{G}(s,a)$  33:  $N(h,a^{*}) += 1,N(h) += 1$   
16:  $h^\prime \gets \text{PUSHBACK}(h,[a^*,o]),d^\prime \gets d + 1$   
17: return  $r + \gamma \cdot \text{ROLLOUT}(s,h',\text{done},d')$  34:  $Q(h,a^{*})\gets Q(h,a^{*}) + \frac{R - Q(h,a^{*})}{N(h,a^{*})}$   
18: end procedure 35: return R  
19: procedure SIMULATE(s,h,done,d, $\mathcal{T}$ ) 36: end procedure
cosine similarity to the current instruction:  $\operatorname{CosineSim}(\ell_i, \ell)$ . In experiments, we use a subset of  $\mathcal{D}$  to show its performance when restricted to a small training set. We select the top  $K$  similar instructions and use the corresponding expert trajectories as a  $K$ -shot prompt. However, the answer  $\alpha_i$  is a free-formed natural language sentence that cannot be mapped to admissible actions for the agent directly. To ensure that the action can be executed, we follow the method in prior works [18] to represent the actions and admissible actions by embeddings from [26] and evaluate their cosine similarity  $\operatorname{CosineSim}(\alpha_i, a)$ . The empirical policy distribution is formulated as follows:  $\hat{\pi}(a|h) = \lambda \frac{1}{|A|} + (1 - \lambda)\operatorname{Softmax}\{\sum_{i=1}^{M} \operatorname{CosineSim}(\alpha_i, a) - \eta\}$ , where  $\eta$  is the average value of  $\sum_{i} \operatorname{CosineSim}(\alpha_i, a)$  and  $|A|$  is the size of the admissible action space.  $\lambda$  is a hyper-parameter that adds randomness to the belief, as the sampled actions from LLM could be very deterministic. Therefore, the empirical policy distribution is a mixture of approximated policy from LLM and uniform distribution. The example prompts are provided in Appendix F.
2.4 Monte Carlo tree search
We integrate the commonsense world belief and policy from LLM in MCTS, presented in Alg 1. For each simulation, MCTS samples a state from the belief  $b(s)$  at the root (line 4). It independently samples one position for each object to construct a state  $s$ . This sampled state  $s$  is then employed in the simulation, generating a new tree trajectory. An action  $a^*$  is chosen during the simulation based on the  $Q$  value, visit counts, and LLM policy (lines 28 and 29). The observation and transition function, denoted as  $\mathcal{G}$  (lines 15 and 30), predict the next state  $s'$  given the selected action  $a^*$  and the sampled state  $s$ , thus progressing to the subsequent step in the simulation (lines 30 and 31). When encountering leaf nodes in the tree, MCTS expands the tree and performs a random rollout for the corresponding node (lines 23 to 26). A uniform policy is employed to sample actions in the rollout, and the discounted reward is then returned (lines 14 to 17). Upon completing the task or reaching the maximum depth, the accumulated rewards are backpropagated, updating each node's estimated  $Q$  value (lines 32 to 35). Following  $N$  simulations, the output action is determined based on the estimated  $Q$  value (lines 3 to 8). Upon completion of the search process, the agent will execute an action and receive a new observation. For simplicity, we assume that the observation and transition functions are deterministic and known. In cases where an object is detected, its corresponding position within the belief will be updated with the observed position. Conversely, if the object remains undetected at certain positions, the belief regarding its presence in those positions will be rendered null, denoted by a zero value.
VirtualHome. We proceed with our experiments in the VirtualHome [24], a large household simulated environment with a large domain, partial observations, and large action space. It contains hundreds of interactive items and containers with various types of rooms. It is a well-suited platform for evaluating embodied decision-making for solving daily tasks in household environments.