[AAAI 2022] Learning Parameterized Task Structure for Generalization to Unseen Entities

We introduce Parameterized Subtask Graph Inference (PSGI), which infers the structure of hierarchical and compositional tasks in a first order logic manner. We test PSGI experimentally on a cooking domain, mining domain (based on Minecraft), and a simulated AI2Thor domain and show PSGI is more efficient and generalizable than prior work.

Introduction

One of the ultimate goals of AI is to replace the human efforts in many real-world tasks to make human’s life better and more convenient. However, it is extremely challenging to teach AI to solve real-world tasks for various reasons. One of the main challenges comes from its compositional task structure.

(Figure 1) A task of making breakfast has a compositional task structure.

Compositional task structure naturally exists in real-world natural language, video, and decision-making settings (e.g., virtual assistants, robotics, etc.). Real-world tasks are composed of multiple smaller tasks, subtasks, that must be completed in certain orders to complete the overall task. For example, consider a household robot tasked to make breakfast in Figure 1. A meal can be served with different dishes and drinks (e.g., scrambled egg and toasted bread), where each could be considered as a subtask. These can then be further decomposed into smaller subtasks until some base subtask (e.g., pickup egg) is reached. Each subtask can provide the agent with a reward; if only a few subtasks provide a reward, this is considered a sparse reward problem. When the subtask dependencies are complex and reward is sparse, learning an optimal policy can require a large number of interactions with the environment. This is the problem scope we focus on in this work: learning to quickly infer and adapt to varying compositional tasks and complex subtask dependencies.

Motivation

(Figure 2) Previous work cannot generalize to unseen tasks with an unseen entity.

Prior work has tackled learning compositional task structures by modeling dependencies between subtasks in a graph structure [Sohn et al., 2018, 2020, Xu et al., 2018, Huang et al., 2019]. In these settings, during training, the agent tries to efficiently adapt to a task by inferring the latent graph structure, then uses the inferred graph to maximize reward during a test. However, this approach does not scale to the tasks with unseen entities (e.g., egg, cabbage, etc). Prior work tries to infer the structure of subtasks individually – they do not consider the relations between entities in related tasks. E.g. they do not infer that [Place egg on pot] and [Place egg on pan] have similar compositional structures (Figure 2). Instead, we propose to use the parameterized options and subtasks to enable zero-shot task generalization with unseen entities (Figure 3).

(Figure 3) Ours can generalize to unseen tasks by inferring the latent task structure in a first-order logic manner.

Method

(Figure 4) Illustration of parameterized subtask graph inference (PSGI) in a toy cooking environment.

In various tasks, the agent must cook various foods to receive reward.

The key idea behind parameterized subtask graph inference (PSGI) is to model a task using parameterized options and subtasks. This allows PSGI to infer the latent task structure in a first-order logic manner. For example, in the cooking task in Figure 4 we represent all [Pickup object] options using a parameterized option, [Pickup x]. Representing options and subtasks in parameterized form serves two roles:

1. The resulting graph is more compact. There is less redundancy when representing compositional tasks that share a common structure. Hence a parameterized subtask graph requires fewer samples to infer (e.g. relations for [Pickup apple], [Pickup pan], etc. are inferred at once with a parameterized option [Pickup x]).

2. The resulting graph can generalize to unseen subtasks, where unseen subtasks may share a similar structure but are not encountered during adaptation (e.g. [Pickup cabbage] in Figure 4).

In this post, we are going to only cover the inference of the precondition and attribute functions, and refer the readers to our paper(Link) for other details.

Parameterized Precondition Inference via Predicate Logic Induction

Let ${\mathrm{\tau }}_{\mathrm{H}}=\left\{({\mathrm{x}}_1,{\mathrm{e}}_1),({\mathrm{x}}_2,{\mathrm{e}}_2),\dots ,({\mathrm{x}}_{\mathrm{H}},{\mathrm{e}}_{\mathrm{H}})\right\}$ be the adaptation trajectory of the adaptation policy ${\mathrm{\pi }}_{adapt}$after H time steps. Our goal is to infer the maximum likelihood parameterized precondition ${\mathrm{G}}_{\mathrm{prec}}$ given this trajectory ${\mathrm{\tau }}_{\mathrm{H}}$.

where ${\mathrm{f}}_{\mathrm{Gprec}}$is the precondition function, ${\mathrm{E}}^{\left(\mathrm{i}\right)}$is the entity set of the i-th option in the given task, ${\mathrm{x}}_{\mathrm{t}}$is the subtask completion vector indicating if each subtask is completed at time t, and ${\mathrm{e}}_{\mathrm{t}}$is the subtask eligibility vector indicating if each subtask is eligible (i.e., precondition has been satisfied) at time t. Similar to prior research^[2], this problem can be translated into an inductive logic programming (ILP) problem. We infer the eligibility ${\mathrm{e}}_{\mathrm{t}}$ (Boolean output) of some option O with some entities $\mathrm{E}=\left\{{\mathrm{X}}_1,{\mathrm{X}}_2,\dots \right\}$, from Boolean input formed by all possible completion values $\left\{{{\mathrm{x}}_{\mathrm{t}}}^{\mathrm{i}}\right\}$, and all attribute values ${{\left\{{\mathrm{f}}_{\mathrm{i}}\left(\mathrm{X}\right)\right\}}^{\mathrm{i}}}_{\mathrm{X}\in \mathrm{E}}$. For more detail, we refer to our full paper.

Zero-shot Learning Entity Attributes

To enable parameterized representation, we also learn the attributes of the components in the compositional tasks. These attributes are used to indicate differences in the structures of parameterized options and subtasks. For example, in the cooking task in Figure 4, not every object can be picked up with Pickup, so the inferred attribute $f_{pickupable}\left(x\right)$is a precondition to [Pickup x]. Similarly, in a more complex cooking task, some object x may need to be sliced, before it can be boiled (e.g. cabbage), but some do not (e.g. egg). As there is no supervision involved, we formulate inferring attribute as a zero-shot learning problem. Note the inferred attributes that are preconditions for options should not only construct an accurate predicate subtask graph for options seen in the adaptation phase, but also unseen options. During the adaptation phase, the agent will encounter a set of seen entities E. We construct candidate attributes from E using our smoothness assumption, where similar entities result in similar preconditions. We generate candidate attributes based on similarity using the given entity (word) embedding $\varphi$, such as bag-of-words^[5]. Let $\mathrm{C}=\left\{{\mathrm{C}}_1,{\mathrm{C}}_2,...\right\}$be an exhaustive set of clusters generated from E using entity embedding $\varphi$. Then, we define a candidate attribute function from each cluster.

${\hat{f}}_i\left(X\right)≔\mathrm{𝕀}[X\in C_i]$

To infer the attribute of an unseen entity $X\in E$, we use a 1-Nearest Neighbor classifier that uses the attributes of the nearest seen entity.

${\hat{f}}_i\left(X\right)≔\mathrm{𝕀}[X^{*}\in C_i]$

Where $X^{*}={\mathrm{argmin}}_{X^{\text{'}}\in E}\mathrm{dist}\left(\varphi \right(X),\varphi (X^{\text{'}}\left)\right)$.

Experiment

We tested our agent on two 2D grid-world domains (Cooking and Mining) and one realistic household domain based on AI2Thor^[6]. The environment implements the tasks of cooking different dishes in different house layouts with different objects.

(Figure 5) (left) AI2Thor and (right) grid-world domains used to test our PSGI agent.

In an experiment, we aim to answer the following two research questions:

1. Does PSGI efficiently learn the latent structure of hierarchical and compositional tasks compared to prior work?
2. Can PSGI generalize the latent structure of unseen subtasks and options using the inferred attributes?

(Figure 6) The inferred parameterized subtask graph by PSGI after 2000 timesteps in the Cooking environment.

Options are represented in rectangular nodes. Subtask completions and attributes are in oval nodes.

(Figure 7) The adaptation curves in the Cooking, AI2Thor, and Mining domains.

In this challenging transfer learning tasks, our PSGI was able to infer task structure in parameterized form and generalize to unseen subtask (Figure 6). Moreover, PSGI demonstrated higher sample efficiency thanks to sharing the knowledge of precondition and effect across different subtasks (Figure 7). Meanwhile, other baselines do not learn any meaningful policy except MSGI+ in Cooking and AI2Thor. This demonstrates that the parameterized subtask graph enables PSGI to share the experience of similar subtasks and options (e.g., pickup X on Y for all possible pairs of X and Y) such that the sample efficiency is increased by roughly the factor of number of entities compared to using subtask graph in MSGI+.

Future work

In this work, we only considered the compositionality that can be learned via direct interaction with environment in RL setting. In fact, compositionality can be learned from various sources of information such as visuals or texts. In LG AI Research, we aim to investigate how we can better discover and exploit such latent compositional task structure for sample efficient learning and strong generalization.

▶Learning Parameterized Task Structure for Generalization to Unseen Entities (Link)

- Code (Link)