SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anurag Ajay\*†§, Seungwook Han \*†‡§ Yilun Du \*§, Shuang Li §, Abhi Gupta §, Tommi Jaakkola §, Josh Tenenbaum §, Leslie Kaelbling §, Akash Srivastava ‡, Pulkit Agrawal †§  
Improbable AI Lab† MIT-IBM Watson AI Lab‡ Massachusetts Institute Technology§ https://hierarchical-planning-foundation-model.github.io/
To make effective decisions in novel environments with long-horizon goals, it is crucial to engage in hierarchical reasoning across spatial and temporal scales. This entails planning abstract subgoal sequences, visually reasoning about the underlying plans, and executing actions in accordance with the devised plan through visual-motor control. We propose Compositional Foundation Models for Hierarchical Planning (HiP), a foundation model which leverages multiple expert foundation model trained on language, vision and action data individually jointly together to solve long-horizon tasks. We use a large language model to construct symbolic plans that are grounded in the environment through a large video diffusion model. Generated video plans are then grounded to visual-motor control, through an inverse dynamics model that infers actions from generated videos. To enable effective reasoning within this hierarchy, we enforce consistency between the models via iterative refinement. We illustrate the efficacy and adaptability of our approach in three different long-horizon table-top manipulation tasks.
Consider the task of making a cup of tea in an unfamiliar house. To successfully execute this task, an effective approach is to reason hierarchically at multiple levels: an abstract level, e.g. the high level steps needed to heat up the tea, a concrete geometric level e.g., how we should physically navigate to and in the kitchen, and at a control level e.g. how we should actuate our joints to lift a cup. It is further important that reasoning at each level is self-consistent with each other – an abstract plan to look in cabinets for tea kettles must also be physically plausible at the geometric level and executable given the actuations we are capable of. In this paper, we explore how we can create agents capable of solving novel long-horizon tasks which require hierarchical reasoning.
Large "foundation models" have become a dominant paradigm in solving tasks in natural language processing [36, 47, 7], computer vision [26], and mathematical reasoning [27]. In line with this paradigm, a question of broad interest is to develop a "foundation model" that can solve novel and long-horizon decision-making tasks. Some prior works [39, 6] collected paired visual, language and action data and trained a monolithic neural network for solving long-horizon tasks. However, collecting paired visual, language and action data is expensive and hard to scale up. Another line of prior works [10, 28] finetune large language models (LLM) on both visual and language inputs using task-specific robot demonstrations. This is problematic because, unlike the abundance of text on the Internet, paired vision and language robotics demonstrations are not readily available and are expensive to collect. Furthermore, finetuning high-performing language models, such as GPT3.5/4 [37, 36] and PaLM [7], is currently impossible, as the model weights are not open-sourced.
* denotes equal contribution. Authors are also affiliated with Computer Science and Artificial Laboratory (CSAIL). Correspondence to aajay@mit.edu, swhan@mit.edu and yilundu@mit.edu.
Figure 1: Compositional Foundation Models for Hierarchical Planning. HiP uses a task model, represented using a LLM, to create an abstract plan, a visual model, represented using a video model, to generate an image trajectory plan, and an ego-centric action model to infer actions from the image trajectory.
The key characteristic of the foundation model is that solving a new task or adapting to a new environment is possible with much less data compared to training from scratch for that task or domain. Instead of building a foundation model for long-term planning by collecting paired language- vision-action data, in this work we seek a scalable alternative - can we reduce the need for a costly and tedious process of collecting paired data across three modalities and yet be relatively efficient at solving new planning tasks? We propose Compositional Foundation Models for Hierarchical Planning (HiP), a foundation model that is a composition of different expert models trained on language, vision, and action data individually. Because these models are trained individually, the data requirements for constructing the foundation models are substantially reduced (Figure 1). Given an abstract language instruction describing the desired task, HiP uses a large language model to find a sequence of sub-tasks (i.e., planning). HiP then uses a large video diffusion model to capture geometric and physical information about the world and generates a more detailed plan in form of an observation-only trajectory. Finally, HiP uses a large pre-trained inverse model that maps a sequence of ego-centric images into actions. The compositional design choice for decision-making allows separate models to reason at different levels of the hierarchy, and jointly make expert decisions without the need for ever collecting expensive paired decision-making data across modalities.
Given three models trained independently, they can produce inconsistent outputs that can lead to overall planning failure. For instance, a naive approach for composing models is to take the maximum-likelihood output at each stage. However, a step of plan which is high likelihood under one model, i.e. looking for a tea kettle in a cabinet may have zero likelihood under a seperate model, i.e. if there is no cabinet in the house. It is instead important to sample a plan that jointly maximizes likelihood across every expert model. To create consistent plans across our disparate models, we propose an iterative refinement mechanism to ensure consistency using feedback from the downstream models [28]. At each step of the language model's generative process, intermediate feedback from a likelihood estimator conditioned on an image of the current state is incorporated into the output distribution. Similarly, at each step of the video model generation, intermediate feedback from the action model refines video generation. This iterative refinement procedure promotes consensus among the different models and thereby enables hierarchically consistent plans that are both responsive to the goal and executable given the current state and agent. Our proposed iterative refinement approach is computationally efficient to train, as it does not require any large model finetuning. Furthermore, we do not require access to the model's weights and our approach works with any models that offer only input and output API access.
In summary, we propose a compositional foundation model for hierarchical planning that leverages a composition of foundation models, learned separately on different modalities of Internet and egocentric robotics data, to construct long-horizon plans. We demonstrate promising results on three long-horizon tabletop manipulation environments.
2 Compositional Foundation Models for Hierarchical Planning
We propose HiP, a foundation model that decomposes the problem of generating action trajectories for long-horizon tasks specified by a language goal  $g$  into three levels of hierarchy: (1) Task planning – inferring a language subgoal  $w_{i}$  conditioned on observation  $x_{i,1}$  and language goal  $g$ ; (2) Visual
Figure 2: Planning with HiP. Given a language goal  $g$  and current observation  $x_{t}$ , LLM generates next subgoal  $w$  with feedback from a visual plausibility model. Then, Diffusion uses observation  $x_{t}$  and subgoal  $w$  to generate observation trajectory  $\tau_{x}$  with feedback from an action feasibility model. Finally, action planning uses inverse dynamics to generate action  $a_{t}$  from current  $x_{t}$  and generated observation  $x_{t+1}$  (action planning).
planning - generating a physically plausible plan as a sequence of image trajectories  $\tau_{x}^{i} = \{x_{i,1:T}\}$ , one for each given language subgoal  $w_{i}$  and observation at first timestep  $x_{i,1}$ ; (3) Action planning - inferring a sequence of action trajectories  $\tau_{a}^{i} = \{a_{i,1:T-1}\}$  from the image trajectories  $\tau_{x}^{i}$  executing the plan. Figure 2 illustrates the model architecture and a pseudocode is provided in Algorithm 1.
Let  $p_{\Theta}$  model this hierarchical decision-making process. Given our three levels of hierarchy,  $p_{\Theta}$  can be factorized into the following: task distribution  $p_{\theta}$ , visual distribution  $p_{\phi}$ , and action distribution  $p_{\psi}$ . The distribution over plans, conditioned on the goal and an image of the initial state, can be written under the Markov assumption as:
$$
p _ {\Theta} \left(W, \left\{\tau_ {x} ^ {i} \right\}, \left\{\tau_ {a} ^ {i} \right\} \mid g, x _ {1, 1}\right) = \underbrace {\left(\prod_ {i = 1} ^ {N} p _ {\theta} \left(w _ {i} \mid g\right)\right)} _ {\text {t a s k p l a n n i n g}} \underbrace {\left(\prod_ {i = 1} ^ {N} p _ {\phi} \left(\tau_ {x} ^ {i} \mid w _ {i} , x _ {i , 1}\right)\right)} _ {\text {v i s u a l p l a n n i n g}} \underbrace {\left(\prod_ {i = 1} ^ {N} \prod_ {t = 1} ^ {T - 1} p _ {\psi} \left(a _ {i , t} \mid x _ {i , t} , x _ {i , t + 1}\right)\right)} _ {\text {a c t i o n p l a n n i n g}} (1)
$$
We seek to find action trajectories  $\tau_{a}^{i}$ , image trajectories  $\tau_{x}^{i}$  and subgoals  $W = \{w_{i}\}$  which maximize the above likelihood. Please see Appendix A for a derivation of this factorization. In the following sub-sections, we describe the form of each of these components, how they are trained, and how they are used to infer a final plan for completing the long-horizon task.
2.1 Task Planning via Large Language Models
Given a task specified in language  $g$  and the current observation  $x_{i,1}$ , we use a pretrained LLM as the task planner, which decomposes the goal into a sequence of subgoals. The LLM aims to infer the next subgoal  $w_{i}$  given a goal  $g$  and models the distribution  $p_{\mathrm{LLM}}(w_i|g)$ . As the language model is trained on a vast amount of data on the Internet, it captures powerful semantic priors on what steps should be taken to accomplish a particular task. To adapt the LLM to obtain a subgoal sequence relevant to our task, we prompt it with some examples of domain specific data consisting of high-level goals paired with desirable subgoal sequences.
However, directly sampling subgoals using a LLM can lead to samples that are inconsistent with the overall joint distribution in Eqn (1), as the subgoal  $w_{i}$  not only affects the marginal likelihood of task planning but also the downstream likelihoods of the visual planning model. Consider the example in Figure 2 where the agent is tasked with packing computer mouse, black and blue sneakers, pepsi box and toy train in brown box. Let's say the computer mouse is already in the brown box. While the subgoal of placing computer mouse in brown box has high-likelihood under task model  $p_{\theta}(w_i|g)$ , the resulting observation trajectory generated by visual model  $p_{\phi}(\tau_x^i |w_i,x_{i,1})$  will have a low-likelihood under  $p_{\phi}$  given the subgoal is already completed. Next, we describe how we use iterative refinement to capture this dependency between language decoding and visual planning to sample from Eqn (1).
Consistency with Visual Planning To ensure that we sample subgoal  $w_{i}$  that maximizes the joint distribution in Eqn (1), we should sample a subgoal that maximizes the following joint likelihood
$$
w _ {i} ^ {*} = \underset {w _ {i}} {\arg \max } p _ {\text {L L M}} \left(w _ {i} \mid g\right) p _ {\phi} \left(\tau_ {x} ^ {i} \mid w _ {i}, x _ {i, 1}\right), \tag {2}
$$
i.e. a likelihood that maximizes both conditional subgoal generation likelihood from a LLM and the likelihood of sampled videos  $\tau_x^i$  given the language instruction and current image  $x_{i,1}$ . One way to determine the optimal subgoal  $w_{i}^{*}$  is to generate multiple  $w_{i}$  from LLM and score them using the likelihood of videos sampled from our video model  $p_{\phi}(\tau_x^i |w_i,x_{i,1})$ . However, the video generation process is computationally expensive, so we take a different approach.
Algorithm 1 Decision Making with HiP
1: Models: Large language model  $p_{\text{LLM}}$ , Subgoal classifier  $f_{\phi}$ , Noise model of diffusion  $\epsilon_{\phi}$ , Observation trajectory classifier  $g_{\psi}$ , Inverse dynamics  $p_{\psi}$   
2: Hyperparameters: Guidance scales  $\omega, \omega'$ , Denoising diffusion steps  $K$   
3: Input: Current observation  $x_{t}$ , Language goal  $g$   
4: # Task Planning  
5: for  $i = 1 \ldots M$  do  
6: Generate subgoal  $w_{i} \sim p_{\text{LLM}}(w_{i}|g)$   
7: end for  
8: Collect candidate subgoals  $W \leftarrow \{w_{i}\}_{i=1}^{M}$   
9: # Iterative Refinement from Visual Planning  
10:  $w \leftarrow \arg \max_{w} f_{\phi}(x_{t}, W, g)$   
11: # Visual Planning  
12: Initialize  $(\tau_{x})_{K} \sim \mathcal{N}(0, I)$   
13: for  $k = K \ldots 1$  do  
14: # Iterative Refinement from Action Planning  
15:  $\hat{\epsilon} \leftarrow \epsilon_{\phi}((\tau_{x})_{k}, x_{t}, k) + \omega(\epsilon_{\phi}((\tau_{x})_{k}, x_{t}, w, k) - \epsilon_{\phi}((\tau_{x})_{k}, x_{t}, k)) - \omega' \nabla_{(\tau_{x})_{k}} \log g_{\psi}(1|(\tau_{x})_{k})$   
16:  $(\tau_{x})_{k-1} \leftarrow \text{Denoise}((\tau_{x})_{k}, \hat{\epsilon})$   
17: end for  
18: # Action Planning  
19: Extract  $(x_{t}, x_{t+1})$  from  $(\tau_{x})_{0}$   
20: Execute  $a_{t} \leftarrow p_{\psi}(a_{t}|x_{t}, x_{t+1})$
The likelihood of video generation  $p_{\phi}(\tau_x^i |w_i,x_{i,1})$  primarily corresponds to the feasibility of a language subgoal  $w_{i}$  with respect to the initial image  $x_{i,1}$ . Thus an approximation of Eqn (2) is to directly optimize the conditional density
$$
w _ {i} ^ {*} = \underset {w _ {i}} {\arg \max } p \left(w _ {i} \mid g, x _ {i, 1}\right). \tag {3}
$$
We can rewrite Eqn (3) as
$$
w _ {i} ^ {*} = \underset {w _ {i}} {\arg \max } \log p _ {\mathrm {L L M}} (w _ {i} | g) + \log \left(\frac {p (x _ {i , 1} | w _ {i} , g)}{p (x _ {i , 1} | g)}\right)
$$
We estimate the density ratio  $\frac{p(x_{i,1}|w_i,g)}{p(x_{i,1}|g)}$  with a multi-class classifier  $f_{\phi}(x_{i,1},\{w_j\}_{i = 1}^M,g)$  that chooses the appropriate subgoal  $w_{i}^{*}$  from candidate subgoals  $\{w_{j}\}_{j = 1}^{M}$  generated by the LLM. The classifier implicitly estimates the relative log likelihood estimate of  $p(x_{i,1}|w_i,g)$  and use these logits to estimate the log density ratio with respect to each of the  $M$  subgoals and find  $w_{i}^{*}$  that maximizes the estimate [45]. We use a dataset  $\mathcal{D}_{\mathrm{classify}}\coloneqq \{x_{i,1},g,\{w_j\}_{j = 1}^M,i\}$  consisting of observation  $x_{i,1}$ , goal  $g$ , candidate subgoals  $\{w_j\}_{j = 1}^M$  and the correct subgoal label  $i$  to train  $f_{\phi}$ . For further architectural details, please refer to Appendix B.1.
2.2 Visual Planning with Video Generation
Upon obtaining a language subgoal  $w_{i}$  from task planning, our visual planner generates a plausible observation trajectory  $\tau_{x}^{i}$  conditioned on current observation  $x_{i,1}$  and subgoal  $w_{i}$ . We use a video diffusion model for visual planning given its success in generating text-conditioned videos [20, 53]. To provide our video diffusion model with a rich prior for physically plausible motions, we pretrain it  $p_{\phi}(\tau_x^i |w_i,x_{i,1})$  on a large-scale text-to-video dataset Ego4D [13]. We then finetune it on the task-specific video dataset  $\mathcal{D}_{\mathrm{video}}\coloneqq \{\tau_x^i,w_i\}$  consisting of observation trajectories  $\tau_{x}^{i}$  satisfying subgoal  $w_{i}$ . For further architectural details, please refer to Appendix B.2.
However, analogous to the consistent sampling problem in task planning, directly sampling observation trajectories with video diffusion can lead to samples that are inconsistent with the overall joint distribution in Eqn (1). The observation trajectory  $\tau_x^i$  not only affects the marginal likelihood of visual planning, but also the downstream likelihood of the action planning model.
Consistency with Action Planning To ensure observation trajectories  $\tau_{x}^{i}$  that correctly maximize the joint distribution in Eqn (1), we optimize an observation trajectory that maximizes the following joint likelihood
$$
\left(\tau_ {x} ^ {i}\right) ^ {*} = \underset {\tau_ {x} ^ {i}} {\arg \max } p _ {\phi} \left(\tau_ {x} ^ {i} \mid w _ {i}, x _ {i, 1}\right) \prod_ {t = 1} ^ {T - 1} p _ {\psi} \left(a _ {i, t} \mid x _ {i, t}, x _ {i, t + 1}\right), \tag {4}
$$
i.e. an image sequence that maximizes both conditional observation trajectory likelihood from video diffusion and the likelihood of sampled actions  $\tau_{a}^{i}$  given the observation trajectory  $\tau_{x}^{i}$ .
To sample such an observation trajectory, we could iteratively bias the denoising of video diffusion using the log-likelihood of the sampled actions  $\prod_{t=1}^{T-1} \log p_{\psi}(a_{i,t} | x_{i,t}, x_{i,t+1})$ . While this solution is principled, it is slow as it requires sampling of entire action trajectories and calculating the corresponding likelihoods during every step of the denoising process. Thus, we approximate the sampling and the likelihood calculation of action trajectory  $\prod_{t=1}^{T-1} p_{\psi}(a_{i,t} | x_{i,t}, x_{i,t+1})$  with a binary classifier  $g_{\psi}(\tau_x^i)$  that models if the observation trajectory  $\tau_x^i$  leads to a high-likelihood action trajectory.
We learn a binary classifier  $g_{\psi}$  to assign high likelihood to feasible trajectories sampled from our video dataset  $\tau_x^i \sim \mathcal{D}_{\mathrm{video}}$  and low likelihood to infeasible trajectories generated by randomly shuffling the order of consecutive frames in feasible trajectories. Once trained, we can use the likelihood  $\log g_{\psi}(1|\tau_x^i)$  to bias the denoising of the video diffusion and maximize the likelihood of the ensuing action trajectory. For further details on binary classifier, please refer to Appendix C.2.
2.3 Action Planning with Inverse Dynamics
After generating an observation trajectory  $\tau_{x}^{i}$  from visual planning, our action planner generates an action trajectory  $\tau_{a}^{i}$  from the observation trajectory. We leverage egocentric internet images for providing our action planner with useful visual priors. Our action planner is parameterized as an inverse dynamics model [1, 38] that infers the action  $a_{i,t}$  given the observation pair  $(x_{i,t}, x_{i,t+1})$ :
$$
a _ {i, t} \sim p _ {\psi} \left(a _ {i, t} \mid x _ {i, t}, x _ {i, t + 1}\right)
$$
Training To imbue the inverse dynamics  $p_{\psi}$  with useful visual priors, we initialize it with VC-1 [32] weights, pretrained on ego-centric images and ImageNet. We then finetune it on dataset  $\mathcal{D}_{\mathrm{inv}} \coloneqq \{\tau_x^i, \tau_a^i\}$  consisting of paired observation and action trajectories by optimizing:
$$
\max  _ {\psi} \mathbb {E} _ {\tau \in \mathcal {D} _ {\text {i n v}}} \left[ \log p _ {\psi} \left(a _ {i, t} \mid x _ {i, t}, x _ {i, t + 1}\right) \right]
$$
For further architectural details, please refer to Appendix B.3.
3 Experimental Evaluations
We evaluate the ability of HiP to solve long-horizon planning tasks that are drawn from distributions with substantial variation, including the number and types of objects and their arrangements. We then study the effects of iterative refinement and of pretraining on overall performance of HiP. We also compare against an alternative strategy of visually grounding the LLM without any task-specific data. In addition, we study how granularity of subgoals affects HiP's performance, ablate over choice of visual planning model and analyze sensitivity of iterative refinement to hyperparameters (Appendix E).
3.1 Evaluation Environments
We evaluate HiP on three environments, paint-block, object-arrange, and kitchen-tasks which are inspired by combinatorial planning tasks in Mao et al. [33], Shridhar et al. [43] and Xing et al. [50] respectively.
- paint-block: A robot has to manipulate blocks in the environment to satisfy language goal instructions, such as stack pink block on yellow block and place green block right of them. However, objects of correct colors may not be present in the environment, in which case, the robot needs to first pick up white blocks and put them in the appropriate color bowls to paint them. After that, it should perform appropriate pick-and-place operations to stack a pink block on the yellow block and place the green block right of them. A new task  $\mathcal{T}$  is generated by randomly selecting 3 final colors (out of 10 possible colors) for the blocks and then sampling a relation (out of 3 possible relations) for each pair of blocks. The precise locations of individual blocks, bowls, and boxes are fully randomized across different tasks. Tasks have  $4\sim 6$  subgoals.
- object-arrange: A robot has to place appropriate objects in the brown box to satisfy language goal instructions such as place shoe, tablet, alarm clock, and scissor in brown box. However, the environment may have distractor objects. Furthermore, some objects can be dirty, indicated by a lack of texture and yellow color. For these objects, the robot must first place them in a blue cleaning box and only afterwards place those objects in the brown box. A new task  $\mathcal{T}$  is generated
"Stack green block on top of yellow block and place orange block to the right of the stack"
place white block in yellow bowl
place yellow block in brown box
place white block in green bowl
place green block on yellow block
place white block in orange bowl
place orange block on right of yellow block
"Pack hammer, frypan, toy school bus, and butterfinger chocolate into the brown box"
place dirty object in the clean box
place hammer in the brown box
place frypan in the brown box
place toy school bus in the brown box
place butterfinger chocolate in brown box
"Open microwave, move kettle out of the way, light the kitchen area, and open upper right drawer"
open the microwave
move the kettle to the back stove
slide the upper right drawer
Figure 3: Example Executions. Example long-horizon generated plans on tasks in paint-block, object-arrange, and kitchen-tasks domains.
by randomly selecting 7 objects (out of 55 possible objects), out of which 3 are distractors, and then randomly making one non-distractor object dirty. The precise locations of individual objects and boxes are fully randomized across different tasks. Tasks usually have  $3 \sim 5$  subgoals.
- kitchen-tasks: A robot has to complete kitchen subtasks to satisfy language goal instructions such as open microwave, move kettle out of the way, light the kitchen area, and open upper right drawer. However, the environment may have objects irrelevant to the subtasks that the robot must ignore. Furthermore, some kitchen subtasks specified in the language goal might already be completed, and the robot should ignore those tasks. There are 7 possible kitchen subtasks: opening the microwave, moving the kettle, switching on lights, turning on the bottom knob, turning on the top knob, opening the left drawer, and opening the right drawer. A new task  $\mathcal{T}$  is generated by randomly selecting 4 out of 7 possible kitchen subtasks, randomly selecting an instance of microwave out of 3 possible instances, randomly selecting an instance of kettle out of 4 possible instances, randomly and independently selecting texture of counter, floor and drawer out of 3 possible textures and randomizing initial pose of kettle and microwave. With  $50\%$  probability, one of 4 selected kitchen subtask is completed before the start of the task. Hence, tasks usually have  $3 \sim 4$  subtasks (i.e. subgoals).