Code Models are Zero-shot Precondition Reasoners (2024)

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.

## 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 @xcite . 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 @xcite @xcite ., partial program specification demonstration trajectories precondition inference def demonstration1(): user.inform('city') system.request('cuisine') user.inform('cuisine') system.findrestaurants() ... def demonstration2(): system.request('city') system.request('cuisine') user.inform('cuisine') ... class user: def __init__(self): self.@xmath0, slot): self.informed[slot] = true class system: def __init__(self, user): self.@xmath1: assert def findrestaurants(self): assert self.user.informed['city'] assert self.user.informed['cuisine'] def test(): system.request('city') user.inform('city') system.request('cuisine') user.inform('cuisine') system. findrestaurants() 1 inform -inform the value for a slot. request -request the value of a slot. offer -offer a certain value for a slot. system: request(city) user: inform(city) system: request(cuisine) domain description example trajectory action preconditions next @xmath2: user.inform('city') system.request('cuisine') user.inform('cuisine') system.findrestaurants() ... def demonstration2(): system.request('city') system.request('cuisine') user.inform('cuisine') ... class user: def __init__(self): self.@xmath3, slot): self.informed[slot] = true class system: def __init__(self, user): self.@xmath4: assert def findrestaurants(self): assert self.user.informed['city'] assert self.user.informed['cuisine def test(): system.request('city') user.inform('city') system.request('cuisine') user.inform('cuisine') system. findrestaurants() 1 … def demonstration(): system.request('city') user.inform('city') system.request('cuisine') user.inform('cuisine') system.findrestaurants() ... class user: def __init__(self): self.@xmath5.@xmath6... def inform(self, slot): self.informed[slot] = true class system: def __init__(self, user): self.@xmath7: assert def findrestaurants(self): assert self.user.informed['city'] assert self.user.informed['cuisine'] def test(): user.inform('cuisine') system.request('city') user.inform('city') system. findrestaurants() 1 def demonstration(): system.request('city') user.inform('city') system.request('cuisine') user.inform('cuisine') system.findrestaurants() ... class user: def __init__(self): self.@xmath8.@xmath9... def inform(self, slot): self.informed[slot] = true class system: def __init__(self, user): self.@xmath10: assert def findrestaurants(self): assert self.user.informed['city'] assert self.user.informed['cuisine'] def test(): user.inform('cuisine') system.request('city') user.inform('city') system. findrestaurants() 1 … def demonstration(): system.request('city') user.inform('city') system.request('cuisine') user.inform('cuisine') system.findrestaurants() system.set_query_status(true) system.offer('restaurant') system.offer('city') system.inform_count() user.request_alts() system.offer('restaurant') system.offer('city') user.request('has_live_music') system.inform('has_live_music') user.inform_intent('reserverestaurant') user.select() system.request('time') user.inform('time') system.confirm('restaurant') system.confirm('city') system.confirm('time') system.confirm('party_size') system.confirm('date') user.inform('date') user.negate() system.confirm('time') system.confirm('date') user.inform('time') user.negate() system.confirm('time') system.confirm('date') user.affirm() system.reserverestaurant() system.set_query_status(true) system.notify_success() user.request('serves_alcohol') system.inform('serves_alcohol') user.thank_you() user.goodbye() system.goodbye() 1 def demonstration(): system.request('city') user.inform('city') system.request('cuisine') user.inform('cuisine') system.findrestaurants() ... class user: def __init__(self): self.@xmath11.@xmath12... def inform(self, slot): self.informed[slot] = true class system: def __init__(self, user): self.@xmath13: assert def findrestaurants(self): assert self.user.informed['city'] assert self.user.informed['cuisine def test(): user.inform('cuisine') system.request('city') user.inform('city') system. findrestaurants() 1 def demonstration(): system.request('city') user.inform('city') system.request('cuisine') user.inform('cuisine') system.findrestaurants() ... class user: def __init__(self): self.@xmath14.@xmath15... def inform(self, slot): self.informed[slot] = true class system: def __init__(self, user): self.@xmath16: assert def findrestaurants(self): assert self.user.informed['city'] assert self.user.informed['cuisine'] def test(): user.inform('cuisine') system.request('city') user.inform('city') system. findrestaurants() 1 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. taking inspiration from classical ai planning literature @xcite 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 @xcite . 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. ## problem setting we consider a sequential decision making setting where the agent receives a sequence of observations o @xmath17, where we assume discrete observation and action spaces o, a. the agent's trajectory can be represented as a sequence of observations and actions τ = (o 1 , a 1 , o 2 , a 2 , . . . , o n , a n ). we consider a few-shot learning setting where we are given demonstrations @xmath18, . . . , τ n }. let τ <@xmath19: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 precondition function g a (τ <t ) ∈ {0, 1} that informs whether the action is plausible in a given context τ <t (i.e., g a (τ <t ) = 0 represents the action is implausible and g a (τ <t ) = 1 represents the action is plausible). given these inferred preconditions @xmath20, our second goal is to estimate a policy π(a t |τ <t , d, g) that predicts actions a t consistent with the preconditions g. ## 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. ## 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 @xcite and starcoder 16b @xcite 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. ## related work programs as policies. there exist prior work that advocate viewing robot policies as code/programs @xcite . 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 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 @xcite @xcite . these work find that code models have strong capability to reason about event sequences with minimal supervision compared to language models. ## 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 codecompletion 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.