SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Background on Density Ratio Estimation Let  $p$  and  $q$  be two densities, such that  $q$  is absolutely continuous with respect to  $p$ , denoted as  $q < < p$  i.e.  $q(\pmb{x}) > 0$  wherever  $p(\pmb{x}) > 0$ . Then, their ratio is defined as  $r(\pmb{x}) = p(\pmb{x}) / q(\pmb{x})$  over the support of  $p$ . We can estimate this density ratio  $r(\pmb{x})$  by training a binary classifier to distinguish between samples from  $p$  and  $q$  [46, 14, 17]. More recent work [45] has shown one can introduce auxiliary densities  $\{m_i\}_{i=1}^M$  and train a multi-class classifier to distinguish samples between  $M$  classes to learn a better-calibrated and more accurate density ratio estimator. Once trained, the log density ratio can be estimated by  $\log r(\pmb{x}) = \hat{h_p}(\pmb{x}) - \hat{h_q}(\pmb{x})$ , where  $\hat{h_i}(\pmb{x})$  is the unnormalized log probability of the input sample under the  $i^{th}$  density, parameterized by the model.
Learning a Classifier to Visually Ground Task Planning 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\},g)$  trained to distinguish samples amongst the conditional distributions  $p(x_{i,1}|w_i,g),\dots,p(x_{i,1}|w_M,g)$  and the marginal distribution  $p(x_{i,1}|g)$ . Upon convergence, the classifier learns to assign high scores to  $(x_{i,1},w_i,g)$  if  $w_i$  is the subgoal corresponding to the observation  $x_{i,1}$  and task  $g$  and low scores otherwise.
Architecture We parameterize  $f_{\phi}$  as a 4-layer multi-layer perceptron (MLP) on top of an ImageNet-pretrained vision encoder (ResNet-18 [15]) and a frozen pretrained language encoder (Flan-T5-Base [8]). The vision encoder encodes the observation  $x_{i,1}$ , and the text encoder encodes the subgoals  $w_{j}$  and the goal  $g$ . The encoded observation, the encoded subgoals, and the encoded goal are concatenated, and passed through a MLP with 3 hidden layers of sizes 512, 256, and 128. The output dimension for MLP (i.e., number of classes for multi-classification)  $M$  is 6 for paint-block environment, 5 for object-arrange environment and 4 for kitchen-tasks environment.
Choice of Large Language Model We use GPT3.5-turbo [37] as our large language model.
B.2 Visual Planning
Background Diffusion Probabilistic Models [44, 19] learn the data distribution  $h(\pmb{x})$  from a dataset  $\mathcal{D} \coloneqq \{\pmb{x}^i\}$ . The data-generating procedure involves a predefined forward noisig process  $q(\pmb{x}_{k+1} | \pmb{x}_k)$  and a trainable reverse process  $p_{\phi}(\pmb{x}_{k-1} | \pmb{x}_k)$ , both parameterized as conditional Gaussian distributions. Here,  $\pmb{x}_0 \coloneqq \pmb{x}$  is a sample,  $\pmb{x}_1, \pmb{x}_2, \dots, \pmb{x}_{K-1}$  are the latents, and  $\pmb{x}_K \sim \mathcal{N}(\pmb{0}, \pmb{I})$  for a sufficiently large  $K$ . Starting with Gaussian noise, samples are then iteratively generated through a series of "denoising" steps. Although a tractable variational lower-bound on  $\log p_{\phi}$  can be optimized to train diffusion models, Ho et al. [19] propose a simplified surrogate loss:
$$
\mathcal {L} _ {\text {d e n o i s e}} (\theta) := \mathbb {E} _ {k \sim [ 1, K ], \boldsymbol {x} _ {0} \sim h, \epsilon \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I})} [ \| \epsilon - \epsilon_ {\phi} (\boldsymbol {x} _ {k}, k) \| ^ {2} ]
$$
The predicted noise  $\epsilon_{\theta}(\pmb{x}_k, k)$ , parameterized with a deep neural network, estimates the noise  $\epsilon \sim \mathcal{N}(0, I)$  added to the dataset sample  $\pmb{x}_0$  to produce noisy  $\pmb{x}_k$ .
Guiding Diffusion Models with Text Diffusion models are most notable for synthesizing high-quality images [41, 35] and videos [20, 53] from text descriptions. Modeling the conditional data distribution  $q(\pmb{x}|\pmb{y})$  makes it possible to generate samples satisfying the text description  $\pmb{y}$ . To enable conditional data generation with diffusion, Ho and Salimans [18] modified the original training setup to learn both a conditional  $\epsilon_{\phi}(\pmb{x}_k,\pmb{y},k)$  and an unconditional  $\epsilon_{\phi}(\pmb{x}_k,k)$  model for the noise. The unconditional noise is represented, in practice, as the conditional noise  $\epsilon_{\phi}(\pmb{x}_k,\emptyset ,k)$ , where a dummy value  $\emptyset$  takes the place of  $\pmb{y}$ . The perturbed noise  $\epsilon_{\phi}(\pmb{x}_k,\emptyset ,k) + \omega (\epsilon_{\phi}(\pmb{x}_k,\pmb{y},k) - \epsilon_{\phi}(\pmb{x}_k,\emptyset ,k))$  (i.e. classifier-free guidance) is used to later generate samples.
Video Diffusion in Latent Space As diffusion models generally perform denoising in the input space [19], the optimization and inference become computationally demanding when dealing with high-dimensional data, such as videos. Inspired by recent works [40, 53], we first use an autoencoder  $v_{\mathrm{enc}}$  to learn a latent space for our video data. It projects an observation trajectory  $\tau_x$  (i.e., video) into a 2D tri-plane representation [53]  $\tau_z = [\tau_z^T, \tau_z^H, \tau_z^W]$  where  $\tau_z^T, \tau_z^H, \tau_z^W$  capture variations in the video across time, height, and width respectively. We then diffuse over this learned latent space [53].
Latent Space Video Diffusion for Visual Planning Our video diffusion model  $p_{\phi}(\tau_x^i |w_i,x_{i,1})$  generates video  $\tau_x^i$  given a language subgoal  $w_{i}$  and the current observation  $x_{i,1}$ . It is parameterized through its noise model  $\epsilon_{\phi}((\tau_z^i)_k,w_i,x_{i,1},k)\coloneqq \epsilon_{\phi}((\tau_z^i)_k,l_{\mathrm{enc}}(w_i),v_{\mathrm{enc}}(x_{i,1}),k)$  where  $\tau_z^i\coloneqq v_{\mathrm{enc}}(\tau_x^i)$  is the latent representation of video  $\tau_x^i$  over which we diffuse. We condition the noise model  $\epsilon_{\phi}$  on subgoal  $w_{i}$  using a pretrained language encoder  $l_{\mathrm{enc}}$  and on current observation  $x_{i,1}$  using video encoder  $v_{\mathrm{enc}}$ . To use  $v_{\mathrm{enc}}$  with a single observation  $x_{i,1}$ , we first tile the observation along the temporal dimension to create a video.
Architecture We now detail the architectures of different components:
- Video Autoencoder We borrow our architecture for  $v_{\mathrm{enc}}$  from PVDM [53] which uses transformers to project video  $\tau_x \in \mathbb{R}^{T \times H \times W}$  to latent codes  $\tau_z = [\tau_z^T, \tau_z^H, \tau_z^W]$  where  $\tau_z^T \in \mathbb{R}^{C \times H' \times W'}$ ,  $\tau_z^H \in \mathbb{R}^{C \times T \times W'}$ ,  $\tau_z^W \in \mathbb{R}^{C \times H' \times T}$ . Here,  $T = 50$  represents the time horizon of a video,  $H = 48$  represents video height,  $W = 64$  represents video width,  $C = 4$  represents latent codebook dimension,  $H' = 12$  represents latent height, and  $W' = 8$  represents latent width.
- Language Encoder We use Flan-T5-Base [8] as the pretrained frozen language encoder  $l_{\mathrm{enc}}$
- Noise Model We borrow PVDM-L architecture [53] which uses 2D UNet architecture, similar to the one in Latent Diffusion Model (LDM) [40], to represent  $p(\tau_z|\tau_z')$ . In our case,  $\tau_z = v_{\mathrm{enc}}(\tau_x^i)$  and  $\tau_z' = v_{\mathrm{enc}}(x_{i,1})$ . To further condition noise model  $\epsilon_{\phi}$  on  $l_{\mathrm{enc}}(w_i)$ , we augment the 2D UNet Model with cross-attention mechanism borrowed by LDM [40].
For implementing these architectures, we used the codebase https://github.com/sihyun-yu/PVDM which contains the code for PVDM and LDM.
Classifier for Consistency between Visual Planning and Action Planning To ensure consistency between visual planning and action planning, we want to sample observation trajectories that maximizes both conditional observation trajectory likelihood from diffusion and the likelihood of sampled actions given the observation trajectory (see equation 4). To approximate likelihood calculation of action trajectory, we learn a binary classifier  $g_{\psi}$  that models if the observation trajectory leads to a high likelihood action trajectories. Since diffusion happens in latent space and we use gradients from  $g_{\psi}$  to bias the denoising of the video diffusion,  $g_{\psi}(\tau_z^i)$  takes the observation trajectory in latent space. The binary classifier  $g_{\psi}$  is trained to distinguish between observation trajectories in latent space sampled from video dataset  $\tau_z^i = v_{\mathrm{enc}}(\tau_x^i), \tau_x^i \sim \mathcal{D}_{\mathrm{video}}$  (i.e. label of 1) and observation trajectories in latent space sampled from video dataset whose frames where randomly shuffled  $(\tau_z^i)' = v_{\mathrm{enc}}(\sigma(\tau_x^i)), \tau_x^i \sim \mathcal{D}_{\mathrm{video}}$  (i.e. label of 0). Here,  $\sigma$  denotes the random shuffling of frames. To randomly shuffle frames in an observation trajectory (of length 50), we first randomly select 5 frames in the observation trajectory. For each of the selected frame, we randomly permute it with its neighboring frame (i.e. either with the frame before it or with the frame after it). Once  $g_{\psi}$  is trained, we use it to bias the denoising of the video diffusion
$$
\begin{array}{l} \hat {\epsilon} := \epsilon_ {\phi} ((\tau_ {z}) _ {k}, v _ {\mathrm {e n c}} (x _ {t}), k) + \omega (\epsilon_ {\phi} ((\tau_ {z}) _ {k}, v _ {\mathrm {e n c}} (x _ {t}), l _ {\mathrm {e n c}} (w), k) - \epsilon_ {\phi} ((\tau_ {z}) _ {k}, v _ {\mathrm {e n c}} (x _ {t}), k)) \\ - \omega^ {\prime} \nabla_ {(\tau_ {z}) _ {k}} \log g _ {\psi} (1 | (\tau_ {z}) _ {k}) \\ \end{array}
$$
Here,  $\hat{\epsilon}$  is the noise used in denoising of the video diffusion and  $\omega ,\omega^{\prime}$  are guidance hyperparameters.
Classifier Architecture The classifier  $g_{\psi}(\tau_z = [\tau_z^T, \tau_z^H, \tau_z^W])$  has a ResNet-9 encoder that converts  $\tau_z^T, \tau_z^H$ , and  $\tau_z^W$  to latent vectors, then concatenate those latent vectors and passes the concatenated vector through an MLP with 2 hidden layers of sizes 256 and 128 and an output layer of size 1.
B.3 Action Planning
To do action planning, we learn an inverse dynamics model to  $p_{\psi}(a_{i,t}|x_{i,t}, x_{i,t+1})$  predicts 7-dimensional robot states  $s_{i,t} = p_{\psi}(x_{i,t})$  and  $s_{i,t+1} = p_{\psi}(x_{i,t+1})$ . The first 6 dimensions of the robot state represent joint angles and the last dimension of the robot state represents the gripper state (i.e., whether it's open or closed). The first 6 action dimension is represented as joint angle difference  $a_{i,t}[: 6] = s_{i,t+1}[: 6] - s_{i,t}[: 6]$  while the last action dimension is gripper state of next timestep  $a_{i,t}[-1] = s_{i,t+1}[-1]$ .
Architecture We use ViT-B [9] (VC-1 [32] initialization) along with a linear layer to parameterize  $p_{\psi}$ . ViT-B projects the observation  $x_{i,t} \in \mathbb{R}^{48 \times 64 \times 3}$  to 768 dimensional latent vector from which the linear layer predicts the 7 dimensional state  $s_{i,t}$ .
C Training and Evaluation
C.1 Task Planning
Inability to recognize uncommon objects
Yes, screwdriver is in brown box
Is computer hard drive in brown box?
No, the computer hard drive is not in the brown box
Is screwdriver in brown box?
Hallucination leading to wrong spatial reasoning
Is scissor in brown box?
Yes, scissor is in brown box
Vision Language Model
Figure 8: Failure in VLM. to recognize uncommon objects like computer hard drives and occasional hallucination of object presence, leading to incorrect visual reasoning.
Training Objective and Dataset for Learned Classifier We use a softmax cross-entropy loss to train the multi-class classifier  $f_{\phi}(x_{i,1}, \{w_j\}_{j=1}^{M}, g)$  to classify an observation  $x_{i,1}$  into one of the  $M$  given subgoal. We train it using the classification dataset  $\mathcal{D}_{\text{classify}} := \{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$ . The classification dataset for paint-block, object-arrange, and kitchen-tasks consists of 58k, 82k and 50k datapoints respectively.
Vision-Language Model (VLM) as a Classifier We use a frozen pretrained Vision-Language Model (VLM) (MiniGPT4 [58]) as a classifier. We first sample a list of all possible subgoals
$W = \{w_{i}\}_{i = 1}^{M}$  from the LLM given the language goal  $g$ . We then use the VLM to eliminate subgoals from  $W$  that have been completed. For each subgoal, we question the VLM whether that subgoal has been completed. For example, consider the subgoal "Place white block in yellow bowl". To see if the subgoal has been completed, we ask the VLM "Is there a block in yellow bowl?" Consider the subgoal "Place green block in brown box" as another example. To see if the subgoal has been completed, we ask the VLM "Is there a green block in brown box?" Furthermore, if the VLM says "yes" and the subgoal has been completed, we also remove other subgoals from  $W$  that should have been completed, such as "Place white block in green bowl". Once we have eliminated the completed subgoals, we use the domain knowledge to determine which subgoal to execute out of all the remaining subgoals. As an example, if the goal is to "Stack green block on top of blue block in brown box" and we have a green block in green bowl and a blue block in blue bowl, we should execute the subgoal "Place blue block in brown box" before the subgoal "Place green block on blue block". While this process of asking questions from VLM to determine the remaining subgoals and then sequencing the remaining subgoals doesn't require any training data, it heavily relies on the task's domain knowledge.
Failure Modes of VLM We observe two common failure modes of the VLM approach in object-arrange environment and visualize them in Figure 8. First, because the model is not trained on any in-domain data, it often fails to recognize uncommon objects, such as computer hard drives, in the observations. Second, it occasionally hallucinates the presence of objects at certain locations and thus leads to incorrect visual reasoning.
VLM as a Subgoal Predictor We also tried to prompt the VLM with 5 examples of goal  $g$  and subgoal candidates  $\{w_i\}_{i=1}^M$  and then directly use it to generate the next subgoal  $w_i$  given the observation  $x_{i,1}$  and the goal  $g$ . However, it completely failed. We hypothesize that the VLM fails to directly generate the next subgoal due to its inability to perform in-context learning.
Evaluation We evaluate the trained classifier  $f_{\phi}$  and the frozen VLM for subgoal prediction accuracy on 5k unseen datapoints, consisting of observation, goal, candidate subgoals and correct subgoal, generated from test tasks  $T_{\mathrm{test}}$ . We average over 4 seeds and show the results in Figure 7.
C.2 Visual Planning
Ego4D dataset processing We pre-train on canonical clips of the Ego4D dataset which are text-annotated short clips made from longer videos. We further divide each canonical clip into 10sec segments from which we derive 50 frames. We resize each frame to  $48 \times 64$ . We create a pretraining Ego4D dataset of (approximately) 344k short clips, each consisting of 50 frames and a text annotation. We use the loader from R3M [34] codebase (https://github.com/facebookresearch/r3m) to load our pretraining Ego4D dataset.
Training Objective and Dataset We use pixel-level L1 reconstruction and negative perceptual similarity for training the autoencoder  $v_{\mathrm{enc}}$ . We borrow this objective from PVDM [53] paper except we don't use adversarial loss. We keep the language encoder frozen. We use denoising loss in video latent space  $\mathbb{E}_{k\sim [1,K],\tau_z,w,x\sim \mathcal{D},\epsilon \sim \mathcal{N}(O,I)}[\| \epsilon -\epsilon_{\phi}((\tau_z)_k,l_{\mathrm{enc}}(w),v_{\mathrm{enc}}(x),k)\| ^2 ]$  to train the noise model  $\epsilon_{\phi}$ . We replace  $w$  with a null token so that  $\epsilon_{\phi}$  learns both a text-conditional model and an unconditional model. We pretrain the autoencoder  $v_{\mathrm{enc}}$  and the noise model  $\epsilon_{\phi}$  on the processed Ego4D dataset. We then finetune it on our dataset  $D_{\mathrm{video}}\coloneqq \{\tau_x^i,w_i\}$  consisting of approximately 100k observation trajectories of length  $T = 50$  and associated text subgoals.
Classifier Training Objective and Dataset We use a binary cross-entropy loss to train the binary classifier  $g_{\psi}(\tau_z)$  that predicts if the observation trajectory in latent space  $\tau_z = v_{\mathrm{enc}}(\tau_x)$  leads to high-likelihood action trajectory. It is trained using trajectories from video dataset  $\tau_x \sim \mathcal{D}_{\mathrm{video}}$ .
C.3 Action Planning
Training Objective and Dataset We train inverse dynamics  $p_{\psi}$  on a dataset  $\mathcal{D}_{\mathrm{inv}}$ . Since actions are differences between robotic joint states, we train  $p_{\psi}$  to directly predict robotic state  $s_{i,t} = p_{\psi}(x_{i,t})$  by minimizing the mean squared error between the predicted robotic state and ground truth robotic state. Hence,  $\mathcal{D}_{\mathrm{inv}} := \{\tau_x^i, \tau_s^i\}$  consists of 1k paired observation and robotic state trajectories, each having a length of  $T = 50$ , in paint-block and object-arrange domains. In kitchen-tasks domain, it consists of 3.5k paired observation and robotic state trajectories, each having a length of  $T = 50$ .
Evaluation We evaluate the trained  $p_{\psi}$  (i.e., VC-1 initialized model and other related models in Figure 6) on 100 unseen paired observation and robotic state trajectories generated from test tasks  $\mathcal{T}_{\mathrm{test}}$ . We use mean squared error to evaluate our inverse dynamics models. We use 4 seeds to calculate the standard error, represented by the shaded area in Figure 6.
C.4 Hyperparameters
Task Planning We train  $f_{\phi}$  for 50 epochs using AdamW optimizer [31], a batch size of 256, a learning rate of  $1e - 3$  and a weight decay of  $1e - 6$ . We used one V100 Nvidia GPU for training the multi-class classifier.
Visual Planning We borrow our hyperparameters for training video diffusion from the PVDM paper [53]. We use AdamW optimizer [31], a batch size of 24 and a learning rate of  $1e - 4$  for training the autoencoder. We use AdamW optimizer, a batch size of 64, and a learning rate of  $1e - 4$  for training the noise model. During the pretraining phase with the Ego4D dataset, we train the autoencoder for 5 epochs and then the noise model for 5 epochs. During the finetuning phase with  $\mathcal{D}_{\mathrm{video}}$ , we train the autoencoder for 10 epochs and then the noise model for 40 epochs. We used two A6000 Nvidia GPUs for training these diffusion models. We train  $g_{\psi}$  for 10 epochs using AdamW optimizer, a batch size of 256 and a learning rate of  $1e - 4$ . We used one V100 Nvidia GPU for training the binary classifier. During classifier-free guidance, we use  $\omega = 4$  and  $\omega' = 1$ .
Action Planning We train VC-1 initialized inverse dynamics model for 20 epochs with AdamW optimizer [31], a batch size of 256 and a learning rate of  $3e - 5$ . We trained other randomly initialized ViT-B inverse dynamics models and randomly initialized ResNet-18 inverse dynamics models for 20 epochs with AdamW optimizer, a batch size of 256, and a learning rate of  $1e - 4$ . We used one V100 Nvidia GPU for training these inverse dynamics models.
D Implementation Details for Gato and SayCan
In training the visual and action planning for HiP, we use 100k robot videos for visual planner and the inverse dynamics, when trained from scratch, utilizes 10k state-action trajectory pairs. In order to ensure fair comparison, we use 110k datapoints for training Gato [39] and SayCan [2].
Gato We borrow the Gato [39] architecture from Vima Codebase and use it for training a language conditioned policy with imitation learning. We use 110k (language, observation trajectory, action trajectory) datapoints in each of the three domains for training Gato. Furthermore, we provide oracle subgoals to Gato.
SayCan We borrow the SayCan [2] algorithm from the SayCan codebase and adapt it to our settings. Following the recommendations of SayCan codebase, we use CLIPort policies as primitives. CLIPort policies take in top-down RGBD view and outputs pick and place pixel coordinates. Then, an underlying motion planner picks the object from the specified pick-coordinate and places the object at the specified place-coordinate. We train CLIPort policies on 110k (language, observation, action) datapoints in paint-block and object-arrange domain. The SayCan paper uses value function as an affordance function to select the correct subgoal given current observation and high level goal. However, CLIPort policies don't have a value function. The SayCan codebase uses a hardcoded scoring function which doesn't apply to object-arrange domain. To overcome these issues, we use the LLM grounding strategy from ? ]. It uses unnormalized logits over the pixel space given by CLIPort policies as affordance and uses it to ground LLM to current observation and thus predict the subgoal. We then compare SayCan with HiP and other baselines on paint-block and object-arrange domain in Table 1. While SayCan outperforms other baselines, HiP still outperforms it both on seen and unseen tasks of paint-block and object-arrange domain. We couldn't run SayCan on kitchen-tasks domain as there's no clear-cut primitive in that domain. This points to a limitation of SayCan which requires tasks to be expressed in terms of primitives with each primitive paired with an affordance function.
E Additional Ablation Studies
E.1 Consistency between task planning and visual planning
To make task planning consistent with visual planning, we need to select subgoal  $w_{i}^{*}$  which maximizes the joint likelihood (see equation 2) of LLM  $p_{\mathrm{LLM}}(w_i|g)$  and video diffusion  $p_{\phi}(\tau_x^i |w_i,x_{i,1})$ . While
generating videos for different subgoal candidates  $w_{i}$  and calculating the likelihood of the generated video is computationally expensive, we would still like to evaluate its performance in subgoal prediction given it is theoretically grounded. To this end, we first sample  $M$  subgoals  $W = \{w_{j}\}_{j=1}^{M}$  from the LLM. Then, we calculate  $w_{i}^{*} = \arg \max_{w \in W} \log p_{\phi}(\tau_{x}^{i}|w,x_{i,1})$  and use  $w_{i}^{*}$  as our predicted subgoal. Since  $\log p_{\phi}(\tau_{x}^{i}|w,x_{i,1})$  is intractable, we estimate its variational lower-bound as an approximation. We use this approach for subgoal prediction in paint-block environment and compare its performance to that of the learned classifier. It achieves a subgoal prediction accuracy of  $54.3 \pm 7.2\%$  whereas the learned classifier achieves a subgoal prediction accuracy of  $98.2 \pm 1.5\%$  in paint-block environment. Both approaches outperform the approach of randomly selecting a subgoal from  $W$  (i.e., no task plan refinement), which yields a subgoal prediction accuracy of  $16.67\%$  given  $M = 6$ . The poor performance of the described approach could result from the fact that the diffusion model only coarsely approximates the true distribution  $p(\tau_{x}^{i}|w_{i},x_{i,1})$ , which results in loose variational lower-bound and thus uncalibrated likelihoods from the diffusion model. A larger diffusion model could better approximate  $p(\tau_{x}^{i}|w_{i},x_{i,1})$ , resulting in tighter variational lower-bound and better-calibrated likelihoods.
E.2 Consistency between visual planning and action planning