Venue: NAACL 2024

# Code Models are Zero-shot Precondition Reasoners

Lajanugen Logeswaran\({}^{1}\) Sungryull Sohn\({}^{1}\) Yiwei Lyu\({}^{2}\) Anthony Zhe Liu\({}^{2}\) Dong-Ki Kim\({}^{1}\) Dongsub Shim\({}^{1}\) Moontae Lee\({}^{1}\) Honglak Lee\({}^{1,2}\)

\({}^{1}\)LG AI Research \({}^{2}\)University of Michigan, Ann Arbor

###### Abstract

One of the fundamental skills required for an agent acting in an environment to complete tasks is the ability to understand what actions are plausible at any given point. This work explores a novel use of code representations to reason about action preconditions for sequential decision making tasks. Code representations offer the flexibility to model procedural activities and associated constraints as well as the ability to execute and verify constraint satisfaction. Leveraging code representations, we extract action preconditions from demonstration trajectories in a zero-shot manner using pre-trained code models. Given these extracted preconditions, we propose a precondition-aware action sampling strategy that ensures actions predicted by a policy are consistent with preconditions. We demonstrate that the proposed approach enhances the performance of few-shot policy learning approaches across task-oriented dialog and embodied textworld benchmarks.

+
Footnote †: _Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)_, pages 5681–5697 June 16-21, 2024 ©2024 Association for Computational Linguistics

## 1 Introduction

A key capability for learning an optimal agent policy in sequential decision making settings is understanding the plausibility of different actions. For instance, a dialog agent recommending restaurants needs to have basic information like location and type of food in order to look up its database for potential options. Understanding the necessary conditions for performing an action (e.g., location and type of food are necessary for the database lookup action) is referred to as _precondition inference_ or _affordance learning_ in the literature Ahn et al. (2022); Sohn et al. (2020). Knowledge about preconditions has broad implications for policy learning and safety applications.

Few-shot prompted large language models (LLMs) have demonstrated strong capabilities for task planning Logeswaran et al. (2022); Huang et al. (2022); Micheli and Fleuret (2021); Ahn et al. (2022). However, they lack a systematic mechanism to reason about action preconditions. Some approaches attempt to teach LLMs preconditions by providing examples in the form of code assertions Singh et al. (2022); Liang et al. (2022) or natural language rationales Huang et al. (2022b); Yao et al. (2022). Preconditions or rationales for predicted actions are generated dynamically on-the-fly the during inference. However, it is non-trivial to verify whether these preconditions or rationales

Figure 1: We propose a framework to extract action preconditions from demonstration trajectories leveraging code representation and code models. Extracted preconditions are used to construct an agent policy that predicts actions consistent with preconditions.



are adequate as they are dynamically generated. It is further difficult to guarantee with certainty that predicted actions will be consistent with the preconditions or rationales.

Taking inspiration from classical AI planning literature Aeronautiques et al. (1998); Fikes and Nilsson (1993) and recent work that use code representations for reasoning problems, we leverage code representation to reason about preconditions. Programs are a natural formalism to model event sequences and offer the flexibility to express dependency constraints between events, such as in the form of assertions Liang et al. (2022); Singh et al. (2022). Verifying that a program meets the specifications dictated by the assertions amounts to simply executing the program and verifying that the program ran successfully. The ability to execute and verify constraint satisfaction is a key benefit of code compared to alternative representations such as natural language. Representing preconditions in the form of procedural statements in code further provides transparency, controllability and better generalization to unseen scenarios compared to alternative representations of affordance such as neural network functions. Finally, this also enables us to exploit strong priors captured by code understanding models for policy learning problems.

Armed with a code representation, we formulate precondition inference as a code completion problem and propose a zero-shot approach to infer preconditions of actions given demonstration trajectories leveraging pre-trained code models. Given these inferred preconditions, we propose a precondition-aware action prediction strategy that ensures actions predicted by a policy are consistent with the preconditions. We present extensive analysis and ablations that show the impact of different components of our approach.

In summary, we make the following contributions in this work.

* We propose a new approach that leverages program representations and pre-trained code models to extract action preconditions from expert demonstrations alone in a zero-shot manner.
* Combining inferred preconditions with a precondition-aware action prediction strategy, we demonstrate that the proposed framework leads to better agent policies compared to baselines on task-oriented dialog and embodied textworld benchmarks.

## 2 Problem Setting

We consider a sequential decision making setting where the agent receives a sequence of observations \(o_{i}\in\mathcal{O}\) and performs an action \(a_{i}\in\mathcal{A}\), where we assume discrete observation and action spaces \(\mathcal{O},\mathcal{A}\). The agent's trajectory can be represented as a sequence of observations and actions \(\tau=(o_{1},a_{1},o_{2},a_{2},\dots,o_{n},a_{n})\). We consider a few-shot learning setting where we are given demonstrations \(\mathcal{D}=\{\tau^{1},\ldots,\tau^{n}\}\).

Let \(\tau_{<t}=(o_{1:t},a_{1:t-1})\) represent the history of observations and actions. The goal of learning is two-fold. First, for each action \(a\) we intend to estimate a preconditioner function \(g_{a}(\tau_{<t})\in\{0,1\}\) that informs whether the action is plausible in a given context \(\tau_{<t}\) (i.e., \(g_{a}(\tau_{<t})=0\) represents the action is implausible and \(g_{a}(\tau_{<t})=1\) represents the action is plausible). Given these inferred preconditions \(\mathcal{G}=\{g_{a}|a\in\mathcal{A}\}\), our second goal is to estimate a policy \(\pi(a_{t}|\tau_{<t},\mathcal{D},\mathcal{G})\) that predicts actions \(a_{t}\) consistent with the preconditions \(\mathcal{G}\).

## 3 Approach

We first discuss the code representation in Section 3.1. We then describe our approach to precondition inference and action prediction problems in Sections 3.2 and 3.3, respectively. See Figure 1 for an overview of our approach.

### Representing Agent Trajectories as Programs

We represent the agent's interaction with the environment as a program. We represent every action and observation as one or more function calls that modify the state of a predefined set of variables \(v^{1}\) that capture a summary of the agent's experience (e.g., observations \(o_{1:t}\) and actions \(a_{1:t-1}\)). Such a representation exists since we assume discrete observation and action spaces (for instance, we can define a separate function for each string in \(\mathcal{O},\mathcal{A}\)). Assume we have defined a set of functions \(\mathcal{F}_{\mathcal{O}},\mathcal{F}_{\mathcal{A}}\) corresponding to observations \(\mathcal{O}\) and actions \(\mathcal{A}\) respectively. Given this representation, a trajectory \(\tau\) can be viewed as a program which consists of a sequence of function calls.2 We present an example program in Figure 1 and further examples in the Appendix.

 



### Precondition Inference

Recall that the precondition function \(g_{a}(\tau_{<t})\in\{0,1\}\) informs whether an action \(a\) is plausible in a given context \(\tau_{<t}\). Equivalently, for the corresponding function in the program representation \(f_{a}\in\mathcal{F}_{\mathcal{A}}\), we seek to identify assertion statements in terms of variables \(v\) (which represent a summary of \(\tau_{<t}\)) and the function arguments of \(f_{a}\). (For example in Figure 1, the preconditioner for the action FindRestaurants is identified as assert user.informed['city'] and user.informed['cuisine']).

We predict preconditions for each action independently of other actions and the process below is repeated for each action \(a\in\mathcal{A}\).3 Our approach to precondition inference consists of the following steps: candidate generation, validation and ranking. We detail these steps next (See Appendix D for an illustration and Appendix F for example outputs from each component of the pipeline).

Footnote 3: Jointly reasoning about preconditions for all functions in \(\mathcal{F}_{\mathcal{A}}\) can be interesting and is left to future work.

Candidate Generation.Given demonstration trajectories \(\mathcal{D}\) we first generate candidate preconditions by prompting a pre-trained code generation model. The prompt consists of (i) a demonstration \(\tau\in\mathcal{D}\) (ii) definitions of functions \(\mathcal{F}_{\mathcal{O}}\) and (iii) definition of function \(f_{a}\) with the assert keyword in its body. The assert prefix forces the model to generate assertion statements (as opposed to arbitrary code). We vary the demonstration program and sample multiple precondition candidates for each demonstration to come up with a pool of plausible candidates \(\mathcal{H}^{\text{initial}}_{a}\). The pre-trained model is expected to understand the appropriate contexts in which a function can be used and use this understanding to come up with assertion statements.

Although pretrained models have strong priors about appropriate assert statements, the above process has some limitations. First, generated statements may not be meaningful due to syntax errors or other deficiencies in the generated expressions. Second, the candidates are obtained based on static analysis of the program alone and the model has to implicitly reason about execution and program state in order to predict accurate preconditions.

Candidate Validation.A key advantages of code compared to alternative representations such as natural language is the ability to execute. We augment the above candidate generation approach with a verification approach where candidates are vetted for validity and consistency with the data. Given a candidate assertion, we verify whether each of the demonstration programs can run successfully (e.g,.. replacing the function body with a candidate assertion and executing a demonstration program). Candidates which led to execution failures are discarded. The remaining candidates \(\mathcal{H}^{\text{valid}}_{a}\) are thus valid and consistent with the data.

Many of the assertions generated will not be useful in practice. For example, trivial assertions such as 'assert True' do not convey any useful information about instances where the function is applicable. Although all candidate assertions at this point are valid, they may be _sub-optimal_ for the purpose of constructing a policy. We seek assertions that help discriminate situations where the function is applicable. We propose a ranking mechanism to identify the most discriminative assertions.

Candidate Ranking.We seek to identify a small set of precondition/assertion candidates from \(\mathcal{H}^{\text{valid}}\) useful for constructing an agent policy. We begin with the observation that if candidates \(h_{1},h_{2}\) are such that \(h_{2}\) is satisfied whenever \(h_{1}\) is satisfied (i.e., \(h_{1}\Rightarrow h_{2}\)), \(h_{1}\) is more desirable as it is more discriminative of the contexts where action \(a\) is applicable (and hence leads to a better policy). As we discuss in the experiments, this assumption leads to high precision solutions and compromises recall. We leave alternate ranking criteria to future work.

However, verifying the above property (\(h_{1}\Rightarrow h_{2}\)) for given precondition candidates is generally

Figure 2: **Precondition Inference Overview**: We first generate multiple precondition candidates by prompting a code model with demonstration programs. We then validate these candidates based on the demonstrations via execution, and rank them to identify the most promising candidates.



intractable. We thus consider an approximation where we examine if the property is satisfied in scenarios that appear in the demonstration trajectories. Consider the precondition function \(g_{a}(\,\cdot\,;h)\) which assumes \(h\) to be the precondition of action \(a\). Define \(C^{h}_{a}=\{(i,j)|g_{a}(r^{i}_{<j};h)=1,\tau^{i}\in\mathcal{D}\}\) to be the set of instances \((i,j)\) where precondition \(h\) is satisfied at time-step \(j\) of demonstration trajectory \(\tau^{i}\). We use \(C^{h_{1}}_{a}\subseteq C^{h_{2}}_{a}\) as a proxy to determine whether \(h_{1}\Rightarrow h_{2}\).

We define the optimal set of assertion statements as in Equation (1). When multiple equivalent candidates exist, we choose one random representative to retain in set \(\mathcal{H}^{\text{opt}}_{a}\). Note that the conjunction of assertions in \(\mathcal{H}^{\text{opt}}_{a}\) is equivalent to that of \(\mathcal{H}^{\text{valid}}_{a}\). However, identifying a small set of assertions is beneficial for interpretability and for providing as a prompt to models with limited context lengths.

\[\mathcal{H}^{\text{opt}}_{a}=\{h\in\mathcal{H}^{\text{valid}}_{a}\mid\bar{ \#}h^{\prime}\in\mathcal{H}^{\text{valid}}_{a}\text{ s.t. }C^{h^{\prime}}_{a}\subset C^{h}_{a}\} \tag{1}\]

### Precondition-aware Action Prediction

We pose action prediction as a code completion problem where given a partial agent trajectory, a code model is tasked with suggesting possible next actions (as functions from \(\mathcal{F}_{\mathcal{A}}\), with appropriate arguments). Given any policy which predicts the next action given past actions and observations, we consider a simple approach to augment the policy with precondition knowledge. Next action candidates are sampled from the policy until an action consistent with the preconditions is found or a maximum number of attempts is exceeded. The first attempt uses greedy sampling and subsequent attempts use random sampling. If random sampling does not yield an action consistent with preconditions, we default to the greedy action. We sample an action by generating tokens until a newline token is encountered. See Appendix E for pseudocode.

## 4 Experiments

We attempt to answer the following key questions in our evaluation: (i) Is it possible to extract information about action preconditions from demonstrations of agent behavior? (ii) Is such inferred precondition information useful for building better agent policies?

We use the python programming language as the code representation in our experiments due to its simplicity and popularity, as well as the recent release of code models that primarily focus on python. We use the CodeGen 2B (Nijkamp et al., 2022) and StarCode 16B (Li et al., 2023) open-source pretrained code models in our experiments. Specifically, we use the versions of these models which were further pre-trained on python code after initial pre-training on many programming languages.

### Benchmarks

Task Oriented Dialog Benchmark.The decision making component of a task-oriented dialog system is a dialog manager which takes as input a sequence of utterances represented as dialog acts and predicts the next action. In this setting, a user and a system (agent) take turns to speak and user utterances constitute the agent's observations. We use the SGD dataset (Rastogi et al., 2020) in our experiments. There are 11 user acts and 11 system acts defined in the dataset, each of which takes either no argument, a single slot argument or a single intent argument. We define analogous functions \(\mathcal{F}_{\mathcal{O}}\) and \(\mathcal{F}_{\mathcal{A}}\) corresponding to each of these acts. We take an object-oriented approach and group these functions respectively under a user class and a system class. We define variables \(v\) corresponding to each user action that records whether the action was performed (e.g., informed_slot[], requested_slot[] in Figure 1). We experiment with 10 domains (schemas) from the SGD dataset. 10 instances are used as demonstrations and 50 instances for testing. For evaluation purposes we manually define the ground-truth preconditions for each system action based on prior knowledge about the dialog acts. Note that these are only used for evaluation and no supervision is provided to the model about ground-truth preconditions.

Embodied Textworld.We experiment with the ALFororld embodied textworld benchmark Shridhar et al. (2020) which involves an agent interacting with an environment to perform object interaction tasks. The observations and actions are natural language statements and the agent is expected to perform a task specified using a natural language instruction (e.g., _move the keys to the table_). There are 9 types of actions which include interaction and navigation actions. Each of these action types take an object argument and/or a receptacle argument and we define analogous functions \(\mathcal{F}_{\mathcal{A}}\). In addition, we define auxiliary functions \(\mathcal{F}_{\mathcal{O}}\) for adding/removing an object from the agent's inventory and updating the set of objects visible to the agent. We define three variables \(v\) that summarize the agent's 

experience: the set of visible objects, the agent's inventory and states of objects in the environment (e.g., open/closed). The dataset has 6 task types. We use 2 instances from each task type as demonstrations and the standard test set of the benchmark for testing. The benchmark provides ground-truth preconditions for the actions in a PDDL (Planning Domain Definition Language) representation, which we convert to python assertions for evaluation. We provide the complete program specification for these benchmarks in Appendix A.

Precondition Inference Metrics.Recall that a precondition \(g_{a}(\tau_{<t})\in\{0,1\}\) informs whether an action \(a\in\mathcal{A}\) is plausible given context \(\tau_{<t}=(o_{1:t},a_{1:t-1})\). Define \(C_{a}=\{(i,j)|g_{a}(\tau^{i}_{<j})=1,\tau^{i}\in\mathcal{D}^{\text{test}}\}\) to be the set of instances \((i,j)\) where the precondition of action \(a\) is satisfied at time-step \(j\) of test trajectory \(\tau^{i}\in\mathcal{D}^{\text{test}}\). We define precision and recall metrics for a particular action \(a\) as in Equation (2), where \(C^{\text{pred}}_{a},C^{\text{gt}}_{a}\) respectively correspond to the predicted and ground-truth preconditions. The \(F_{1}\) score is defined as the harmonic mean of precision and recall. These metrics are macro-averaged across all actions \(a\in\mathcal{A}\) to obtain the final metrics.

\[\text{Prec}=\frac{|C^{\text{pred}}_{a}\cap C^{\text{gt}}_{a}|}{|C^{\text{ pred}}_{a}|},\text{Rec}=\frac{|C^{\text{pred}}_{a}\cap C^{\text{gt}}_{a}|}{|C^{ \text{gt}}_{a}|} \tag{2}\]

Baselines.We first consider a binary classifier which we call the _Neural Precondition_ baseline. If we have labeled data of the form \((\tau_{<t},a,l)\), where \(l\in\{0,1\}\) indicates whether action \(a\) is plausible given context \(\tau_{<t}\), we can train a binary classifier \(h_{a}\) that predicts if the precondition for \(a\) is satisfied. However, in our setting, we only have positive instances where we know that the precondition of an action was satisfied, i.e., \(\mathcal{D}^{\text{pos}}=\{(\tau_{<t},a_{t},1)|\tau\in\mathcal{D},t\leq n\}\), where \(\mathcal{D}\) is a dataset of expert trajectories. In order to construct negative instances, we assume that the precondition for every future action \(a_{t^{\prime}};t^{\prime}>t\) that appeared in the trajectory \(\tau\) apart from the ground-truth action \(a_{t}\) was not satisfied. Based on this assumption we construct negative instances \(\mathcal{D}^{\text{neg}}=\{(\tau_{<t},a_{t^{\prime}},0)|t^{\prime}>t,a_{t^{ \prime}}\neq a_{t},\tau\in\mathcal{D}\}\) and train a binary classifier \(h\) on \(\mathcal{D}^{\text{pos}}\cup\mathcal{D}^{\text{neg}}\). In practice, we uniformly sample instances from \(\mathcal{D}^{\text{pos}}\) and \(\mathcal{D}^{\text{neg}}\) to train the classifier. We also consider a _Prompting_ baseline where we prompt a code model with a partial program specification and predict assertion statements with greedy decoding.

Results.In Table 1 we compare the precondition inference performance of different methods. The precision for our models are high since the ranking criteria we adopted encourages high precision solutions. For example, in the _Restaurants_ domain of the SGD dataset, the ground truth preconditioner for the action INFORM(slot) is requested_slot[slot].4 However, the predicted precondition is (requested_slot[slot] and query_success==True). The models pick up the fact that query_success==True in all instances the INFORM(slot) action appeared in the demonstration trajectories. In this case precision is \(1.0\) since the ground truth precondition will be satisfied whenever the predicted precondition is satisfied. However, the recall is lower (\(0.57\)) since the predicted and ground-truth preconditions may not agree when query_success!=True.

Footnote 4: ‘self’ and ’user’ omitted for brevity

The Neural Precondition baseline generally underperforms other methods. The assumption that actions in the future are invalid for the current time-step does not always hold. Comparing the labels in the generated data (e.g., \(\mathcal{D}_{\text{pos}},\mathcal{D}_{\text{neg}}\)) with the corresponding ground-truth labels (computed using ground-truth preconditions), we observe that the (precision, recall) of the generated labels are respectively (1, 0.74) and (1, 0.60) for the SGD and Alfworld datasets.5 When ground-truth precondition labels are used for training the classifier, performance improves (F1 scores of 0.81 and 0.67 in SGD and Alfworld, respectively), but generalization to unseen scenarios is still limited.

Footnote 5: Precision is 1.0 since labels in \(\mathcal{D}_{\text{pos}}\) are guaranteed to be correct. Recall is lower since this is not the case for \(\mathcal{D}_{\text{neg}}\).

Table 2 presents qualitative prediction results for

\begin{table}
\begin{tabular}{l c c c c c c} \hline \hline Model & \multicolumn{3}{c}{SGD} & \multicolumn{3}{c}{Alfworld} \\ \cline{2-9}  & Prec & Rec & \(F_{1}\) & Prec & Rec & \(F_{1}\) \\ \hline Neural Prec. & 0.77 & 0.65 & 0.70 & 0.64 & 0.42 & 0.51 \\ Prompt.(2B) & 0.73 & 0.75 & 0.74 & 0.59 & **0.89** & 0.71 \\ Prompt.(16B) & 0.73 & 0.77 & 0.75 & 0.69 & 0.68 & 0.68 \\ \hline Ours(2B) & 0.91 & 0.69 & 0.78 & **1.0** & 0.85 & **0.92** \\ Ours(16B) & **0.92** & **0.78** & **0.84** & **1.0** & 0.81 & 0.90 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Precondition Inference performance on the SGD and Alfvén benchmarks. Metrics are precision, recall and \(F_{1}\) score (All metrics higher † the better). The prompting baseline and our approach are evaluated with CodeGen 2B and StarCoder 16B models.

 

all action types in the _Buses_ domain of the SGD benchmarks. Note that our models receive no supervision about preconditions and predict these candidates by only observing expert demonstrations. We present qualitative prediction results for the Alfworld benchmark in Appendix C. Next, we analyze how these inferred action precondition can help build better agent policies.

### Precondition-aware Agent Policy

Metrics.In the SGD benchmark, we evaluate policy performance using actions in the demonstrations as reference. Since the agent needs to predict multiple actions in a turn, we compute \(F_{1}\) score treating demonstration actions as ground-truth. For the embodied textworld task, a simulation environment is available, and we define **success rate (SR)** which measures how often the agent successfully completed the given task. We also define **precondition compatibility (Cmp.)**, which measures how often the predicted action is compatible with (i.e., does not violate) the ground-truth preconditions.

Baselines.We consider the following baseline policy approaches.

* **BC**: We consider a behavioral cloning baseline where a code model is fine-tuned (with LoRA) on the given expert demonstrations with a language modeling training objective.
* **Act**Yao et al. (2022): A few-shot prompting baseline where training demonstrations are provided as a prompt to a pre-trained model and the model predicts the next action.
* **ReAct**Yao et al. (2022): A variation of the _Act_ prompting strategy where the agent predicts a chain-of-thought natural language rationale before predicting each action. We design 'think' prompts similar to the original work and insert them as code comments in the demonstrations. At every step the agent generates an optional code comment ('think' step) and predicts an action.

The StarCoder 16B model is used as the backbone model for all the policy methods unless specified otherwise. We use preconditions predicted by the StarCoder 16B model in all experiments.

Results.We present the main results in Table 3. The table shows the performance of each baseline policy (BC, Act, ReAct) and the performance when each of these methods are augmented with preconditions inferred by our approach. We observe that the precondition-aware action sampling strategy combined with predicted preconditions lead to consistent improvements over each baseline pol

\begin{table}
\begin{tabular}{l l c r} \hline \hline Ground-truth precondition & Predicted precondition & Prec & Rec & \(F_{1}\) \\ def INFORM(self, slot): & def INFOORM(self, slot): & 1 & 1 & 1 \\ assert self.user.requested\_slot & assert slot in self.user.requested\_slot & & 1 & 1 & 1 \\ der REQUEST(self, slot): & & 1 & 1 & 1 \\ assert std.user.informed\_slot[slot] == False & & 1 & 1 & 1 \\ der GOODWE(self): & & 1 & 1 & 1 \\ assert self.user.no\_more & & 1 & 1 & 1 \\ der OFPR(self, slot): & & 1 & 1 & 1 \\ assert self.over\_success == True & & & 1 & 1 & 1 \\ der INFORM_COUNT(self): & & 1 & 1 & 1 \\ assert self.query\_success == True & & & & \\ assert self.user.selected & & & & \\ der CONFIG(self, slot): & & 0.49 & 1 & 0.66 \\ assert self.user.infected\_slot[’fromLocation’] & & & 0.96 & 0.7 & 0.81 \\ assert self.user.informed\_slot[’freeLypers’] & & & & \\ assert self.user.informed\_slot[’reurles’] & & & & \\ der NOTIFYSUCCESS(self): & & 1 & 0.44 & 0.62 \\ assert self.user.selected & & & & \\ assert not self.user.no\_more & & & & \\ \hline def RQQMR(self): & & 1 & 1 & 1 \\ assert self.user.selected & & & & \\ def FinObs(self): & & & & \\ assert not self.user.no\_more & & & & \\ assert set.user.informed\_slot[’fromLocation’] & & & & \\ assert self.user.informed\_slot[’leaving\_date’] & & & & \\ assert self.user.informed\_slot[’leaving\_date’] & & & & \\ \hline assert self.user.informed\_slot[’fromlocation’] & & & 1 & 1 & 1 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Ground-truth and predicted preconditions for actions in the _Buses_ domain of the SGD benchmark, along with Precision, Recall scores for predictions.

 

icy. In particular, we observe that the proposed precondition-based reasoning approach helps models generate actions that are more accurate and consistent with ground-truth preconditions.

On the Alfworld benchmark, our approach improves the success rate of the React agent from 44% to 48% and its precondition compatibility from 81% to 96%. This shows the synergistic potential of our approach and recent advances in prompting for policy learning problems such as generating natural language rationales before predicting actions.

These results shows that explicitly reasoning about preconditions helps build better policies. Although few-shot prompting enables language/code models to perform tasks with limited supervision, they are generally limited by the number of demonstrations that can fit in the context window. Condensing the information in multiple trajectories in a small set of rules (e.g., preconditions) can help overcome this limitation.

### Ablations and Analysis

We perform a series of ablations to understand the impact of different components in our pipeline.

How to incorporate precondition information?In Table 4 we analyze the impact of (i) including precondition information in the prompt and (ii) the precondition-aware action sampling strategy. Overall, while the inclusion of precondition information in the model prompt generally helps, our proposed precondition-aware action prediction strategy yields more consistent and significant improvements (e.g., \(0.64\) to \(0.68\) with predicted preconditions and \(0.73\) with ground-truth preconditions in the 1-shot setting). In addition, it helps predict actions that are more consistent with their preconditions (e.g., precondition compatibility improves to \(>0.9\) for both 2B and 1GB models).

Amount of Supervision.Figure 3 shows model performance for varying amounts of supervision. Due to it's maximum context size of 2048 tokens, CodeGen 2B can only accommodate upto 4 demonstrations. Knowledge about preconditions is particularly helpful when the number of demonstrations is small. Even as we increase the number of demonstrations, models continue to benefit from explicitly reasoning about preconditions. Furthermore, performance with ground-truth preconditions shows that improvements in quality of preconditions lead to improvements in policy performance.

Model scale.We observe that the small model benefits more from precondition knowledge compared to the big model regardless of the amount of supervision. Enhancing the reasoning capabilities of small models is important as big models demand higher costs and computation. Leveraging precondition information and execution-based verification is a promising strategy to enhance small models.

## 5 Related Work

Programs as Policies.There exist prior work that advocate viewing robot policies as code/programs Liang et al. (2022); Singh et al. (2022). Similar to natural language based prompting, LLMs are prompted with examples of programs corresponding to example tasks and are required to generate programs for a query task. The programs can be rich and composed of functions supported by the target robot API or third party library functions. Code comments help break down high-level task

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline Model & Prec. & Prec. & 1-shot & 10-shot \\ size & prompt & sample & \(F_{1}\uparrow\) Cmp.\(\uparrow\) & \(F_{1}\uparrow\) Cmp.\(\uparrow\) \\ \hline
2B & ✗ & ✗ & 0.51 & 0.73 & - & - \\
2B & ✓ & ✗ & 0.52 & 0.75 & - & - \\
2B & ✓ & ✓ & 0.58 & **0.92** & - & - \\ \hline
16B & ✗ & ✗ & 0.64 & 0.78 & 0.89 & 0.97 \\
16B & ✓ & ✗ & 0.65 & 0.77 & 0.90 & 0.97 \\
16B & ✓ & ✓ & **0.68** & 0.90 & **0.91** & **0.99** \\ \hline \hline \end{tabular}
\end{table}
Table 4: Ablation to study the effect of (a) providing preconditions as part of the prompt (column 2) and (b) precondition-aware action prediction strategy (column 3) in the SGD benchmark. The 2B model can only accommodate upto 4 demonstrations and hence is not evaluated in the 10-shot setting.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Approach & \multicolumn{2}{c}{SGD} & \multicolumn{2}{c}{Alfworld} \\ \cline{2-5}  & \(F_{1}\) & Comp. & SR & Comp. \\ \hline BC & 0.88 & 0.96 & 0.11 & 0.53 \\ BC+Ours & **0.89** & **0.98** & 0.14 & 0.74 \\ \hline Act & 0.84 & 0.92 & 0.06 & 0.56 \\ Act+Ours & 0.85 & 0.97 & 0.23 & **0.96** \\ \hline ReAct & - & - & 0.44 & 0.81 \\ ReAct+Ours & - & - & **0.48** & **0.96** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Policy performance on SGD and Alfworld benchmarks. The performance metrics are \(F_{1}\) score, task success rate (SR) and precondition compatibility (Cmp.) (All metrics higher \(\uparrow\) the better).

 

into subtasks and assertions are used to take environment feedback into account and provide an error recovery mechanism. These prior work assume that precondition information is specified as part of the prompt. In contrast, our work attempts to discover such information from action trajectories.

Programs have also been used as a representation for structured prediction tasks in NLP (Wang et al., 2022; Madaan et al., 2022; Zhang et al., 2023). These work find that code models have strong capability to reason about event sequences with minimal supervision compared to language models.

Reasoning with Execution.Prior work has attempted to augment capabilities of language models with programs and execution. Liu et al. (2022) combine LLMs with a physics simulator to answer physics questions. Gao et al. (2022) interleave natural language chain-of-thought statements with program statements to perform calculations for arithmetic reasoning problems. Analogous to these work, we find that the ability to verify via execution improves the performance of LMs.

Reasoning via Prompting.Chain-of-thought prompting (Wei et al., 2022) has emerged as a powerful technique for getting language models to perform step-by-step reasoning. These ideas have also been applied to planning problems (Huang et al., 2022b; Yao et al., 2022) where agents are taught how to come up with rationales for predicted actions with few-shot demonstrations. Compared to chain-of-thought rationales, which are dynamically generated on the fly and are problem-specific, we seek to identify a set of rules that capture action preconditions. This provides more controllability over the action generation process and the rules can also be vetted/edited to achieve desired behavior.

Subtask Graph Framework.Subtask Graphs are a modeling framework to learn subtask preconditions (John et al., 2018, 2020; Jang et al., 2023; Logeswaran et al., 2023). Preconditions are modeled as boolean expressions involving a pre-defined set of boolean subtask variables which represent whether a subtask has been completed or not. In contrast to boolean expressions as the class of functions to model preconditions, the program representation we consider has the flexibility to represent a broader set of scenarios. We draw inspiration from these works to formulate and evaluate the precondition inference component in our approach.

Planning with Symbolic Problem Specifications.Symbolic problem specifications such as PDDL have been considered in classical AI planning Aeronautiques et al. (1998). Recent approaches have augmented LLMs with PDDL problem specifications and symbolic solvers Liu et al. (2023); Silver et al. (2023). A PDDL specification exposes the preconditions and effects of actions, described in terms of a set of predicates that describe environment state. It typically assumes deterministic, fully observable environments with known preconditions and effects and known goal state. On the other hand, our approach is more flexible and is applicable in scenarios where partial knowledge about the task and environment are available.

## 6 Conclusion

This work presented a novel approach to reason about action preconditions using programs for learning agent policies in a sequential decision making setting. We proposed to use programs as a representation of the agent's observations and actions and showed that precondition inference and action prediction can be formulated as code-completion problems. By leveraging the strong priors of pre-trained code models, we proposed a novel approach to infer action preconditions from demonstration trajectories without any additional supervision. With the predicted preconditions, our precondition-aware action prediction strategy enables the agent to predict actions that are consistent with the preconditions and lead to better task completion compared to baselines. Our study opens an exciting new direction to reason about action preconditions by leveraging code models.

Figure 3: Ablation showing policy performance (F1 score) in the SGD benchmark when varying (a) number of demonstrations (1 to 10), (b) model scale (2B, 1GB) and (c) precondition information available to policy model (None/Predicted/Ground-truth).

 



### Limitations

Designing program specifications requires domain knowledge, which is a limitation of our work. For instance, we define variables such as the agent's inventory and set of visible objects in the embodied textworld benchmark and variables that capture which slots were requested/informed in the dialog benchmark. It would be interesting to generalize our approach to broader scenarios by identifying the key variables of interest with less manual intervention. For example, this may be possible with few-shot prompting where an LLM is shown example program specifications for a few domains and is expected to come up with program specifications for new domains. It would be interesting to further explore the use of preconditions to more actively direct the agent towards task completion.

## References

* the planning domain definition language. _Technical Report, Tech. Rep._
* Ahn et al. (2022) Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. 2022. Do as i can, not as i say: Grounding language in robotic affordances. _arXiv preprint arXiv:2204.01691_.
* Fikes and Nilsson (1993) Richard E Fikes and Nils J Nilsson. 1993. Strips, a retrospective. _Artificial intelligence_, 59(1-2):227-232.
* Gao et al. (2022) Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. 2022. Pal: Program-aided language models. _arXiv preprint arXiv:2211.10435_.
* Huang et al. (2022a) Wenlong Huang, Pieter Abbeel, Deepak Pathak, and Igor Mordatch. 2022a. Language models as zero-shot planners: Extracting actionable knowledge for embodied agents. In _International Conference on Machine Learning_, pages 9118-9147. PMLR.
* Huang et al. (2022b) Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tompson, Igor Mordatch, Yevgen Chebotar, et al. 2022b. Inner monologue: Embodied reasoning through planning with language models. _arXiv preprint arXiv:2207.05608_.
* Jang et al. (2023) Yunseok Jang, Sungryull Sohn, Lajanugen Logeswaran, Tiange Luo, Moontae Lee, and Honglak Lee. 2023. Multimodal subtask graph generation from instructional videos. _arXiv preprint arXiv:2302.08672_.
* Li et al. (2021) Raymond Li, Loublna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chin, Qian Liu, Evgenii Zheltonozhskii, Terry Yue Zhuo, Thomas Wang, Olivier Dehaene, Mistigin Davadori, Joel Lamy-Poirier, Joao Monteiro, Oleh Shliazkho, Nicolas Gontier, Nicholas Meade, Armel Zebanez, Ming-Ho Yee, Logesh Kumar Umapathi, Jian Zhu, Benjamin Lipkin, Muhathasm Olbolukov, Zhiruo Wang, Rudra Murthy, Jason Stillerman, Siva Sankalg Patel, Dmitty Abukhanov, Marco Zocca, Mannam Dey, Zhiham Zhang, Nour Fahmy, Urvashi Bhattacharyya, Wenhao Yu, Suwayam Singh, Sasha Luccioni, Paulo Villegas, Maxim Kunakov, Fedor Zhdanov, Manuel Romero, Tony Lee, Nadav Timor, Jennifer Ding, Claire Schlesinger, Hailey Schoelkopf, Jan Ebert, Tri Dao, Mayank Mishra, Alex Gu, Jennifer Robinson, Carolyn Jane Anderson, Brendan Dolan-Gavitt, Danish Contractor, Siva Reddy, Daniel Fried, Dzmitry Bahdanau, Yacine Jernite, Carlos Munoz Ferrandis, Sean Hughes, Thomas Wolf, Arjun Guha, Leandro von Werra, and Harm de Vries. 2023. Starcooler: may the source be with you!
* Liang et al. (2022) Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karool Hausman, Brian Ichter, Pete Florence, and Andy Zeng. 2022. Code as policies: Language model programs for embodied control. _arXiv preprint arXiv:2209.07753_.
* Liu et al. (2023) Bo Liu, Yuqian Jiang, Xiaohan Zhang, Qiang Liu, Shiqi Zhang, Joydeep Biswas, and Peter Stone. 2023. Llm+ p: Empowering large language models with optimal planning proficiency. _arXiv preprint arXiv:2304.11477_.
* Liu et al. (2022) Ruibo Liu, Jason Wei, Shixiang Shane Gu, Te-Yen Wu, Sorousn Vosoughi, Claire Cui, Denny Zhou, and Andrew M Dai. 2022. Mind's eye: Grounded language model reasoning through simulation. _arXiv preprint arXiv:2210.05359_.
* Logeswaran et al. (2022) Lajanugen Logeswaran, Yao Fu, Moontae Lee, and Honglak Lee. 2022. Few-shot subgoal planning with language models. _arXiv preprint arXiv:2205.14288_.
* Logeswaran et al. (2023) Lajanugen Logeswaran, Sungryull Sohn, Yunseok Jang, Moontae Lee, and Honglak Lee. 2023. Unsupervised task graph generation from instructional video transcripts. _arXiv preprint arXiv:2302.09173_.
* Madan et al. (2022) Aman Madan, Shuyan Zhou, Uri Alon, Yiming Yang, and Graham Neubig. 2022. Language models of code are few-shot commonsense learners. _arXiv preprint arXiv:2210.07128_.
* Micheli and Fleuret (2021) Vincent Micheli and Francois Fleuret. 2021. Language models are few-shot butlers. _arXiv preprint arXiv:2104.07972_.
* Nijkamp et al. (2022) Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, and Caiming Xiong. 2022. Codegen: An open large language model for code with multi-turn program synthesis. _arXiv preprint arXiv:2203.13474_

Abhinav Rastogi, Xiaoxue Zang, Srinivas Sunkara, Raghav Gupta, and Pranav Khaitan. 2020. Towards scalable multi-domain conversational agents: The schema-guided dialogue dataset. In _Proceedings of the AAAI Conference on Artificial Intelligence_, volume 34, pages 8689-8696.
* Shridhar et al. (2020) Mohit Shridhar, Xingdi Yuan, Marc-Alexandre Cote, Yonatan Bisk, Adam Trischler, and Matthew Hausknecht. 2020. Alfordl: Aligning text and embodied environments for interactive learning. _arXiv preprint arXiv:2010.03768_.
* Silver et al. (2023) Tom Silver, Soham Dan, Kavitina Srinivas, Joshua B Tenenbaum, Leslie Pack Kaelbling, and Michael Katz. 2023. Generalized planning in pddl domains with pretrained large language models. _arXiv preprint arXiv:2305.11014_.
* Singh et al. (2022) Ishika Singh, Valts Bluksis, Arsalan Mousavian, Ankit Goyal, Danfei Xu, Jonathan Tremblay, Dieter Fox, Jesse Thomason, and Animesh Garg. 2022. Progprompt: Generating situated robot task plans using large language models. _arXiv preprint arXiv:2209.11302_.
* Sohn et al. (2018) Sungryull Sohn, Junhyuk Oh, and Honglak Lee. 2018. Hierarchical reinforcement learning for zero-shot generalization with subtask dependencies. _Advances in neural information processing systems_, 31.
* Sohn et al. (2020) Sungryull Sohn, Hyunjae Woo, Jongwook Choi, and Honglak Lee. 2020. Meta reinforcement learning with autonomous inference of subtask dependencies. _arXiv preprint arXiv:2001.00248_.
* Wang et al. (2022) Xinguyao Wang, Sha Li, and Heng Ji. 2022. Code4struct: Code generation for few-shot structured prediction from natural language. _arXiv preprint arXiv:2210.12810_.
* Wei et al. (2022) Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. 2022. Chain of thought prompting elicits reasoning in large language models. _arXiv preprint arXiv:2201.11903_.
* Yao et al. (2022) Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. 2022. React: Synergizing reasoning and acting in language models. _arXiv preprint arXiv:2210.03629_.
* Zhang et al. (2023) Li Zhang, Hainiu Xu, Yue Yang, Shuyan Zhou, Weiqiu You, Manni Arora, and Chris Callison-Burch. 2023. Causal reasoning of entities and events in procedural texts. _arXiv preprint arXiv:2301.10896_.

 

Program Specification

We present the program specification for the _Restaurants_ domain in the SGD benchmark in Figure 4 and the Alfworld benchmark in Figure 5.

``` fromcollectionsimportdefaultcl classUER: def__init_(self); self.informed_intent=defaultdict(lambda:False) self.informed_slot=defaultdict(lambda:False) self.requested_slot=defaultdict(lambda:False) self.no_more=false self.selected~false self.affirmed=False self.affirm_intent=False self.negative_intent= False self.request_alternatives=False def INFOGMBLINTENT(self, intent): self.informed_intent[intern]=True defNEQATLE_INTMENT(self): self.negate_internet=True defAFIRLHINTENT(self): self.arfim_intent=True defREQUSTALTS(Self): self.request_alternatives=True defINGGH(self, slot): self.informed_slot[slot]=True defREQUST(self, slot): self.requested_slot[slot]=True defGOODEV(self): self.no_more=True defTHANK_VOU(self): self.no_more=True defSLECT(self): self.selected=True defAFIFIM(self): self.affirmed=True defNGGARE(self): self.arfirmed=False classSVSTEM: def__init_(self, user): self.user=user self.query_success=Home defINFORM(self, slot): assertself.user.requested_slot[slot] defREGUST(self, slot): assertnotself.user.informed_slot[slot] defCOOKBYE(self): assertself.user.non_more asertelf.user.no_more deffindestuantaus(self): assertself.user.informed_slot['citiy'] assertself.user.informed_slot['citise'] assertself.user.informed_intendestuant*] defReserveSteurant(self): assertself.user.selected orself.user.affirmed defOFFER(self, slot): assertself.user.word_successorself.user.affirmed defNFORM_COUNT(self): assertself.query_success defoffer_INTER(self, intent): assertself.user.selected defCONFIRION(self, slot): assertself.user.informed_slot['time'] assertself.user.informed_slot['city'] assertself.user.selectedor\ sell.user.setledor\ self.user.informed_slot['restaurant_name'] defNOITLY_SUCCESS(self): assertself.query_success defNOTIFY_FAILUHEN(self): assertnotself.query_success defREQ_WRITE(self): assertself.user.selectedorself.user_no_more/ notself.query_success defset_query_status(self, status): self.query_success=status defNOR(self): assertself.query_success=Status defNOR(self): assert.user.selectedorself.user.confconferences defNOR(self): assert.user.add 

from collections import defaultdict class Environment:  def__init_(self):  self.object_states = defaultdict(Lambda: False)  def set_property(self, obj, property, value):  self.object_states[(obj, property)] = value  def get_property(self, obj, property):  returnself.object_states[(obj, property)]

class Agent:  def__init_(self, env):  self.env == env  self.inventory = None  self.visible_objects = set()  def add_inventory(self, obj):  self.inventory = obj  define remove_inventory(self, obj):  self.inventory =None  def update_visible_objects(self, 'args):  self.visible_objects.update(list(args))  def is_visible(self, obj):  returnobj in self.visible_objects  def goto(self, recep):  assert self.is_visible(recep)  def open(self, recep):  assert self.is_visible(recep)  assert self.env.get_property(recep, 'open') == False  def close(self, recep):  assert self.is_visible(recep)  assert self.env.get_property(recep, 'open') == True  def take(self, obj, recep):  assert self.is_visible(obj)  assert self.is_visible(recep)  assert self.injury == None, 'Can only hold one object at a time'  def put(self, obj, recep):  assert self.inventory == obj, 'Need to be holding the object'  assert self.is_visible(recep)  def clean(self, obj, recep):  assert self.is_visible(recep)  assert self.inerometry == obj, 'Need to be holding the object'  assert 'sink' in recep

 def heat(self, obj, recep):  assert self.is_visible(recep)  assert self.inventory == obj, 'Need to be holding the object'  assert "microwave" in recep

 def cool(self, obj, recep):  assert self.is_visible(recep)  assert self.inventory == obj, 'Need to be holding the object'  assert 'ridge' in recep

 def topsgle(self, obj):  assert self.is_visible(obj)

Figure 5: Program specification of the Alfworld benchmark. Note that the assertion statements are assumed to be unknown and only used for evaluation.

 

Example Trajectories

Figure 6 shows an example trajectory from the _Restaurants_ domain in the SGD benchmark and the corresponding program representation. Figures 7 and 8 show an example trajectory from the _pick and place_ task of the Alfworld benchmark in its original text representation and the corresponding program representation.

 

* [16] M. C. Dickson, J. R. Kern, and A. S. Hocker, "The _Practical Design of the Standard Model_", 2007.

[MISSING_PAGE_POST]

 

Qualitative Prediction Results for Alfworld

Figure 9 shows precondition prediction results for the Alfworld benchmark.

Ground-Truth Preconditions def goto(self, recep): assert self.is_visible(recep) def open(self, recep): assert self.is_visible(recep) assert self.env.get_property(recep, 'open') == False def close(self, recep): assert self.is_visible(recep) assert self.env.get_property(recep, 'open') == True def take(self, obj, creep): assert self.is_visible(obj) assert self.is_visible(recep) assert self.inventory = None def put(self, obj, recep): assert self.inventory == obj assert self.is_visible(recep) def clean(self, obj, recep): assert self.is_visible(recep) assert self.inventory == obj assert 'sink' in recep def heat(self, obj, creep): assert self.is_visible(recep) assert self.inventory == obj assert 'microwave' in recep def cool(self, obj, creep): assert self.is_visible(recep) assert self.memory == obj assert 'ridge' in recep def toggle(self, obj): assert self.is_visible(obj) assert self.inventory is not Rene

Figure 9: Ground-truth and predicted preconditions for actions in the Alfworld benchmark.

 



## Appendix D Precondition Inference

Figure 10 illustrates the pipeline for action precondition generation.

## Appendix E Precondition-aware Action Prediction

Algorithm 1 presents the pseudocode for our precondition-aware action sampling strategy.

```
Inputs:query,max_attempts verified\(\leftarrow\)False attempts\(\gets\)0 prediction\(\leftarrow\)GreedySample(query) while not verified and attempts<max_attemptsdo ifattemps>0then prediction\(\leftarrow\)RandomSample(query) program\(\leftarrow\)Query+prediction verified\(\leftarrow\)VerifyProgram(program) attempts\(\leftarrow\)attempts+1 ifnotverifiedthen prediction\(\leftarrow\)GreedySample(query) Outputs:prediction,verified
```

**Algorithm 1** Sample Next Action with Verification

Figure 10: **Precondition Inference Overview**: We first generate multiple precondition candidates by prompting a code model with demonstration programs. We then validate these candidates based on the demonstrations via execution, and rank them to identify the most promising candidates.



Precondition Inference Example

Below we present an example showing the intermediate outputs of the precondition inference pipeline for action INFORM in the _Restaurants_ domain of the SGD benchmark. The INFORM action takes a slot argument and the function syntax is INFORM(self, slot) (See the full program specification in Appendix A).

* Generate candidate preconditions \({\cal H}^{\text{initial}}_{a}\) by prompting a code model with demonstrations not self.user.requested_slot[slot], "Requested slot' sclf.query_success:=None (self.query_success == True) (self.query_success) self.query_success is not None (self.user.informed_slot[slot]) self.query_success is None (slot in self.user.informed_slot) slot to all self.user.revised_slot.keys() self.user.affirmed = True (lastinct(slot, ____.)) isinstance(slot, str) slot1:~date' set.query_success in (True, False) slot1:~serves_alcohol' self.user.requested_slot[slot]
* Identify valid candidates \({\cal H}^{\text{valid}}_{a}\) based on execution against the demonstration programs self.query_success!=None (self.query_success == True) (self.query_success) self.query_success is not None isinstance(slot, str) slot1:~date' self.query_success in (True, False) self.user.requested_slot[slot]
* Identify candidates that are functionally equivalent if Cluster 0 self.query_success!=None self.query_success is not None self.query_success in (True, False)
#Cluster | (self.query_success == True) (self.query_success)
#Cluster 2 isinstance(slot,str) slot1= 'date'
#Cluster 3 self.user.requested_slot[slot]
* Identify the precondition clusters that satisfy Equation (1)
#Cluster 1 (subsumes both cluster 0 and 2) (self.query_success == True) (self.query_success)
#Cluster 3 self.user.requested_slot[slot]
* Choose single representative (randomly) from each cluster to construct \({\cal H}^{\text{opt}}_{a}\) (self.query_success == True) self.user.requested_slot[slot]