SlowGuess/NAACL_2022_OCR huggingface dataset:
Lajanugen Logeswaran*, Yao Fu†, Moontae Lee*†, Honglak Lee*‡ LG AI Research*, University of Illinois at Chicago†, University of Michigan‡
Pre-trained language models have shown successful progress in many text understanding benchmarks. This work explores the capability of these models to predict actionable plans in real-world environments. Given a text instruction, we show that language priors encoded in pre-trained models allow us to infer fine-grained subgoal sequences. In contrast to recent methods which make strong assumptions about subgoal supervision, our experiments show that language models can infer detailed subgoal sequences from few training sequences without any fine-tuning. We further propose a simple strategy to re-rank language model predictions based on interaction and feedback from the environment. Combined with pretrained navigation and visual reasoning components, our approach demonstrates competitive performance on subgoal prediction and task completion in the ALFRED benchmark compared to prior methods that assume more subgoal supervision.
Developing autonomous agents that can complete specific tasks given goal descriptions embodies human-level intelligence. Successful agents in this setting require multiple reasoning capabilities including natural language understanding, visual reasoning, and acting over long temporal horizons. Training black-box models that map instructions and observations to suitable actions has proven to be difficult due to challenges in interpreting and reasoning with multimodal information, especially in the absence of strong supervision. Thus, generalization in this setting demands effective strategies for planning, exploration, and incorporating feedback from the environment.
Generalization in human agents, on the other hand, stems from our ability to naturally brainstorm
Figure 1: Overview of subgoal prediction approach. (a) A pretrained language model prompted with a sequence of training instances, i.e., (instruction, subgoal sequence) pairs, and a query instruction predicts top-k hypotheses using beam search. (b) These predictions are then re-ranked by incorporating information about the environment.
abstract subgoals, better calibrating executable actions and their sequences. Planning at the level of subgoals instead of low-level actions allows us to better adapt to unfamiliar settings. We posit that language supervision can help realize such planning capabilities effectively in artificial agents. First, text is a natural API for interacting with intelligent agents that act in the real world to complete tasks. Knowledge available in the form of text corpora, descriptions and instructions can be exploited to build better agents (Branavan et al., 2012; Zhong et al., 2020). Second, strong language priors are useful to reason about causal sequences of events (Li et al., 2021). Language priors can further inform about object affordances (e.g. an apple is sliceable, whereas a table is not) and other contextual knowledge (e.g. a slicing task is more likely to be performed in a kitchen than a bathroom) (Chen
*Correspondence to lllajan@lgresearch.ai
Human Language Technologies, pages 5493 - 5506
et al., 2020). Recent advances have demonstrated that large language models are able to capture such priors, as evidenced by their strong capabilities in language understanding and beyond (Devlin et al., 2019; Radford et al., 2019; Brown et al., 2020; Bommasani et al., 2021). This leads to the natural question of whether priors learned by language models can help reason about subgoals.
We study the ability of language models to reason about plans composed of a sequence of intermediate goals for completing basic object manipulation tasks in a household environment specified using text instructions. In particular, we use the in-context learning ability (Brown et al., 2020) of large pre-trained language models to reason about subgoals. In contrast to prior methods that finetune language models to predict subgoals/actions (Jansen, 2020; Yao et al., 2020), we show that they can predict subgoal sequences effectively without any fine-tuning. We teach the model how instructions translate into subgoal sequences by constructing a prompt using few examples. Given the prompt and a query instruction the language model predicts likely subgoal sequences (see Figure 1 for an illustration).
While language models are capable of generating strong hypotheses, we observe that these predictions may not be directly usable by agents acting in real environments. First, they suffer from calibration issues: Language models have a tendency to repeat content from the prompt (Zhao et al., 2021). We show that mutual-information inspired metrics help mitigate calibration issues and lead to better ranking of model generated hypotheses.
Second, real-world agents have to update their beliefs and predictions based on interaction and feedback from the environment. Without such feedback we cannot expect the predicted plan to be executable in the environment. We execute plans proposed by the language model in the environment using a pre-trained low-level policy and collect feedback about task success/failure. We use this feedback as a learning signal to train a ranking model that re-ranks language model predictions. In contrast to prior methods that rely on strong subgoal supervision and task level expert trajectories, we show that combining subgoal predictions with a pre-trained subgoal execution policy leads to a strong embodied agent baseline.
We make the following contributions in this work. We show that
- Large language models can predict subgoals from text instructions with very little supervision using in-context learning.
- Incorporating a small amount of feedback from interaction with the environment such as agent state and task success/failure outcome improves language model predictions.
- Combining predicted subgoals with a pre-trained low-level policy for navigation and visual reasoning leads to a simple modular agent policy that performs well on an embodied learning setting.
Language models for planning and interaction The use of language models for planning and action prediction has been explored in prior work. Jansen (2020) fine-tuned a language model to predict subgoal sequences for text instructions from the ALFRED benchmark. Micheli and Fleuret (2021) take a similar approach, but show that imitation learning with few instances combined with reinforcement learning produces models that work well on the ALFWorld benchmark (Shridhar et al., 2021). Yao et al. (2020) demonstrate a similar approach for interactive fiction games (Hausknecht et al., 2020). In contrast to these prior methods, our approach does not assume strong supervision and we demonstrate generalization with limited training examples. Furthermore, in order to exploit the generalization capabilities of large language models, we do not finetune these models and instead use their in-context learning ability. Finally, our approach allows us to build policies that inherit the strong generalization capabilities of these large pre-trained models such as compositional generalization.
Large language models and few-shot learning Brown et al. (2020) showed that pre-trained large language models have few-shot learning capabilities. Given a few examples  $\{(x_i, y_i = f(x_i))\}$  that define a task  $f$  such as classification or translation and a query instance  $x^q$ , prompting a language model with a string such as " $x_1 = y_1; x_2 = y_2; \ldots; x_n = y_n; x^q = "$  leads to meaningful completions by the language model  $y^q \approx f(x^q)$ . This few-shot learning capability of language models has since then been studied and improved upon with approaches like prefix engineering (Schick and Schütze, 2021), prompt tuning (Li and Liang, 2021), model calibration (Zhao et al., 2021) and other methods (Min et al., 2021a). We adopt a similar approach for few-shot subgoal inference.
We assume that subgoal supervision is available for a small number of training tasks and use the language model to infer subgoals for unseen tasks.
Instruction following There is rich literature on agents that follow language instructions (Branavan et al. (2009); Mei et al. (2016); Fried et al. (2018); Suhr et al. (2019) inter alia<sup>1</sup>). Recent developments in simulated environments and benchmarks with human annotated instructions have driven progress in embodied agents that learn from text instructions (Shridhar et al., 2020; Kolve et al., 2017). Successful agents in these settings require multiple reasoning capabilities including language understanding, visual reasoning and learning to act over long time-horizons. Recent embodied learning literature exploit subgoal supervision, pre-trained visual reasoning components and pre-trained transformer models to do well on the task (Singh et al., 2020; Suglia et al., 2021; Zhang and Chai, 2021; Corona et al., 2021; Blukis et al., 2022). Unlike these methods, we do not assume access to strong subgoal supervision or task level expert supervision. We combine language model predictions with pretrained low-level navigation and interaction policies to obtain a competitive agent policy.
Few-shot semantic parsing Subgoal inference from text instructions can be considered a semantic parsing problem where the subgoal sequences serves as a formal representation of text. Shin et al. (2021) show that few-shot semantic parsers can be derived from language models and demonstrate their applicability on text-to-SQL (Finegan-Dollak et al., 2018) and SCAN (Lake and Baroni, 2018) benchmarks. Furrer et al. (2020) and Herzig et al. (2021) further study the compositional generalization ability of such semantic parsers. In our work we make use of ideas introduced in these works such as dynamic prompt creation, constrained decoding and intermediate representations.
We first consider subgoal inference as a semantic parsing problem where a text instruction needs to be translated to a sequence of subgoals and propose an approach to few-shot subgoal inference based on pre-trained language models in Section 3.1. We extend this setting to an agent acting in a simulated environment which can execute these subgoals, observe feedback, and improve upon language model
predictions for more accurate subgoal inference in Section 3.2.
3.1 Few-shot subgoal inference
Subgoals We are interested in a particular subclass of instruction following problems which involve performing a sequence of object interactions in an embodied environment. Each object interaction requires navigating to a particular object and performing an action on it. A task is considered successfully completed if the state of objects in the end satisfy a set of task-specific constraints (for instance, objects that need to be sliced/warmed/cooled/cleaned have the appropriate state change). It is thus natural to define a subgoal as one or more sequence of object interactions. A subgoal  $g$  is specified as  $g = (b, o) \in \mathcal{B} \times \mathcal{O}$  where  $b \in \mathcal{B} = \{\text{Pickup, Clean, Heat, ..}\}$  is one of a predefined set of abstract actions and  $o \in \mathcal{O} = \{\text{Apple, Microwave, DeskLamp, Ottoman, ..}\}$  is an object category.
Subgoal inference problem Given a text instruction  $\tau$ , subgoal inference seeks to predict a sequence of subgoals  $\tau \mapsto g = (g^{(1)}, \dots, g^{(n)})$ . To perform in-context learning with a language model, we consider a representation  $v(g)$  of  $g$  that looks like natural text, where  $v$  is a pre-defined invertible mapping. Such representations have been referred to in the literature as verbalizers (Min et al., 2021a), intermediate representations (Herzig et al., 2021) and canonical representations (Shin et al., 2021), where the purpose is to represent the output in a format the language model understands. In a slight abuse of notation, we will use  $g$  to refer to either a subgoal sequence or its textual representation  $v(g)$  depending on the context.
Generating subgoals We assume that a small amount of training data  $\{(\tau_1, g_1), \dots, (\tau_n, g_n)\}$  is given. The language model is prompted with a comma separated concatenation of the training examples, each in the format "  $\tau_i = g_i$  ", followed by a query  $\tau$ , formatted as "  $\tau =$ ". We assume that the probability of a hypothesis  $h$  (i.e., text representation of a subgoal sequence) can be modeled as in Equation (1), where  $h_i$  is the  $i^{\text{th}}$  token of  $h$  and the token probabilities are derived from the language model.
$$
p (h | \tau) = \prod_ {i} p _ {\mathrm {L M}} \left(h _ {i} \mid h _ {<   i}, \tau , \left\{\tau_ {j}, g _ {j} \right\} _ {j = 1} ^ {n}\right) \tag {1}
$$
We use beam search to identify the top-k hypotheses according to  $p(h|\tau)$ . Generated hypotheses are
<sup>1</sup>See Luketina et al. (2019) for a comprehensive survey
LM predictions
1. slice apple, heat in microwave, put in shelf.
2. slice apple, heat in microwave, put in dining table.
3. pick up apple, put in shelf.
Subgoals completed: slice apple, heat in microwave
Task
Move a heated slice of apple to the white table with shelving below
Figure 2: Re-ranking language model predictions with interaction and feedback from the environment. Given the task, language model predictions, completed subgoals and the agent state we come up with a ranked list of subgoal sequences. The agent then executes the next subgoal from the highest ranked plan. The completed subgoal is added to the partial plan and the process continues until stop subgoal is encountered. During training the agent receives a positive reward if the task is successfully completed, which we use as supervision to train the ranking model.
Re-ranked predictions
1. slice apple, heat in microwave, put in dining table.
2. slice apple, heat in microwave, put in shelf.
(Put,diningtable)
Execute subgoal
Update subgoals
constrained to be valid representations of subgoal sequences by considering only tokens which lead to a valid partial prefix of a subgoal sequence at each step of beam search.
Re-ranking predictions Recent studies have found that language models have popularity and recency biases: the tendency to repeat content mentioned in the prompt, especially content appearing later in the prompt (Zhao et al., 2021). They considered a simple approach to mitigate such biases in model predictions for classification tasks by comparing the likelihood of an output label with and without the query. In contrast to this 'direct model' which models the probability of a label given the input  $p(y|x)$ , Min et al. (2021a) showed that a 'channel model' which models  $p(x|y)$  leads to better, more stable models.
Inspired by these observations, we propose to use  $p(\tau|h)$  to score hypotheses in addition to  $p(h|\tau)$ . Mutual Information based ranking metrics are a natural candidate and they have been explored in the text generation literature (Li et al., 2016; Li and Jurafsky, 2016). We generate multiple hypotheses from the model using  $p(h|\tau)$  and the generated hypotheses are re-scored using the weighted mutual information metric  $(1 - \lambda)\log p(h|\tau) + \lambda\log p(\tau|h)$  where  $\lambda$  is a hyperparameter<sup>2</sup>. To compute  $p(\tau|h)$ , we again use the language model prompted with "  $g_1 = \tau_1, \dots, g_n = \tau_n, h =$ " as the query and compute the conditional probability of  $\tau$ . We expect this paradigm of generating a set of strong hypotheses, followed by accurate re-ranking is more generally applicable to other few-shot language understanding problems.
3.2 Agent policy and incorporating environment feedback
We next consider building an agent that acts in a visual environment to complete tasks given text instructions. While Section 3.1 treated the language model as a knowledge extraction system, in the real world plans need to be updated based on interaction and feedback from the environment. We thus propose a method to improve language model predictions based on environment interaction. Since our goal is to learn the planning component of the agent, we assume a pre-trained low-level policy is provided and optimize over the space of plans. Jointly learning both components is beyond the scope of this work and left as future work.
We assume that a representation of the agent state  $s$  is available. The state representation captures information about the environment (e.g. objects present and their locations) estimated based on the agent's visual observations. As the agent explores the environment and collects new observations the state representation is updated. Assuming that a low-level policy  $\pi_L$  pre-trained to execute a given subgoal is provided, our goal is to train a high-level policy  $\pi_H$  which proposes the subgoals to be executed by the low-level policy. More formally, the high-level policy models  $\pi_H(g^{(t)}|\tau, s_t, g^{(< t)})$  where  $\tau$  is a text instruction,  $s_t$  is the state representation and  $g^{(< t)}$  is the sequence of subgoals completed so far at high-level time-step  $t^3$ .
While the language model can generate compelling subgoal hypotheses, it doesn't take into account information about the environment. For instance, knowledge about the type of room the
2Appendix A details the connection to Mutual Information.
3 Alternatively, this can be framed as a POMDP in a hierarchical reinforcement learning setting.
agent is in (kitchen, bathroom, etc.) and the objects present in it are useful to infer the kind of tasks and subgoals that can be performed. We propose to re-rank hypotheses generated by the language model based on information from the environment to construct  $\pi_H$ . The plans generated by the language model are executed in the environment using  $\pi_L$ . The success/failure outcomes of these plan executions are used to construct a labeled dataset of instructions  $\tau$ , plans  $g$  and agent state  $s$ . A supervised ranking model  $f(g,\tau,s;\theta)$  is trained using this data to re-rank the language model predictions. We represent the ranking model as  $f(g,\tau,s;\theta) = \theta^T\mathrm{concat}(f^{\mathrm{state}}(s),f^{\mathrm{text}}(\tau,g))$  where  $f^{\mathrm{state}}(s)$  is a state embedding,  $f^{\mathrm{text}}(\tau,g)$  is a joint encoding of  $\tau$  and  $g$  produced by a text encoder and  $\theta$  is a parameter vector. Although a text embedding can be derived from the language model, we use a BERT encoder in favor of obtaining a smaller dimensional representation ( $f^{\mathrm{text}} = \mathrm{BERT}_{\mathrm{CLS}}$ ). See Appendix C for more details.
During inference, an instruction  $\tau$  is given, and we use the procedure in Section 3.1 to generate top-k hypotheses. At each step, hypotheses inconsistent with the sequence of subgoals executed so far are pruned and the remaining hypotheses are re-ranked based on the current agent state using  $f$ . The agent attempts the next subgoal proposed by the top hypothesis. The process ends when the stop subgoal is predicted. See Figure 2 for an illustration and Appendix D for more details about the training and inference algorithms.
We use data from the ALFRED benchmark proposed by Shridhar et al. (2020) in our experiments. The ALFRED task requires an agent to execute instructions specified in text to accomplish basic tasks in an embodied environment. A given task is described using a high-level language directive as well as low-level step-by-step instructions (We only use the high-level description). The dataset consists of 7 task types (and 11 fine-grained types), and has more than 20k natural language task descriptions collected from human annotators. In addition, expert demonstrations computed by a planner are also made available. Tasks require acting over many time-steps, with an average of 50 actions, and the longest tasks require  $100+$  steps.
The ground truth subgoal sequences in the
dataset consist of both navigation subgoals and object interaction subgoals. We discard the navigation subgoals and only retain the interaction subgoals for the following reasons. First, the interaction subgoals are sufficient for an agent to successfully complete the task. Second, predicting navigation subgoals from the text instruction alone may not always be possible as they often depend on the scene layout.
Subgoal representation A subgoal  $g^s$  is specified as  $g^s = (b, o) \in \mathcal{B} \times \mathcal{O}$  where  $|\mathcal{B}| = 7$  and  $|\mathcal{O}| = 80$ . We define a textual representation  $v(b)$  of each action type (e.g.  $v(\text{Pickup}) = \text{'pick up'}$ ,  $v(\text{Heat}) = \text{'heat in'}$ ). The object types  $o$  are identified by a text string  $v(o)$  in the dataset and we directly use them as the text representation with minimal pre-processing (e.g.  $v(\text{apple}) = \text{'apple'}$ ,  $v(\text{desklamp}) = \text{'desk lamp'}$ ). The subgoal is represented as  $v(g^s) = v(b) v(o)'$  (e.g.  $v((\text{Pickup}, \text{apple})) = \text{'pick up apple'}$ ). A subgoal sequence  $g = (g^{(1)}, \dots, g^{(n)})$  is represented as  $v(g) = v(g^{(1)}), \dots, v(g^{(n)})'$ . Text representations of all subgoals are given in Appendix B. Note that there are many plausible choices for the representation  $v$  and a different set of choices can lead to different results.
Metrics We use top-k recall to evaluate the ability of language models to generate plans from instructions by comparing against ground truth plans. In addition, we also evaluate the performance of an agent acting in the AI2-Thor (Kolve et al., 2017) simulator to complete tasks using task success rate (the percentage of tasks successfully completed).
4.2 Few-shot subgoal inference
We construct a training set of  $N = 22$  instances by randomly choosing two instances per fine-grained task type. The language model is prompted with a concatenation of these training examples and the query instance. We perform constrained beam search decoding with a beam size of 10 to generate subgoal sequences. At each step of beam search, only tokens which lead to a valid partial prefix of a subgoal sequence are considered. All model generated hypotheses thus correspond to valid subgoal sequences. We evaluate models on the valid-seen split of the dataset which has 800 instances.
Table 2 shows subgoal inference results categorized by task type. We use publicly available pre-trained transformer language models GPT2-XL (Radford et al., 2019) and GPT-J (Wang and Komat-
Top-k recall	Ranking criteria	GPT2-XL	GPT-J
N=11	N=22	N=33	N=11	N=22	N=33
k=10	log p(h|τ)	47.93	58.54	59.51	64.02	71.10	72.07
k=1	log p(h|τ)	16.71	23.29	23.41	30.12	36.10	35.85
k=1	log p(τ|h)	19.88	22.07	25.61	34.15	34.39	35.00
k=1	1/2 log p(h|τ) + 1/2 logp(τ|h)	29.88	32.20	33.17	44.02	47.80	49.63
Table 1: Re-ranking model generated hypotheses using different criteria. The first section shows top-10 recall of generated hypotheses (using  $p(h|\tau)$ ). The second section shows top-1 recall after re-ranking these hypotheses using different criteria. Results are shown for GPT2-XL and GPT-J models when the number of training instances  $N$  is varied.
Task	GPT2-XL	GPT-J
top-1	top-10	top-1	top-10
look at obj in light	34.04	80.85	42.55	80.85
pick and place simple	25.35	59.15	48.59	73.24
pick two obj and place	28.23	68.55	26.61	70.16
pick heat then place	27.10	71.03	47.66	85.98
pick cool then place	30.95	65.87	39.68	80.16
pick clean then place	11.61	50.00	34.82	73.21
pick place movable	6.09	17.39	12.17	35.65
Overall	23.29	58.54	36.10	71.10
Table 2: Top-k recall for subgoal sequences predicted by GPT2-XL and GPT-J models categorized by task type.
suzaki, 2021) via the HuggingFace library (Wolf et al., 2020), which respectively have 1.5B and 6B parameters, in our experiments. The first six of the seven task types have two object arguments each. The pick place movable task type has three object arguments and hence a lower recall than the other task types. The top-10 recall of GPT2-XL and GPT-J are respectively  $59\%$  and  $71\%$ , which shows that large language models have strong ability to reason about plans from few training examples.