{'fraction_non_alphanumeric': 0.025444417701062476, 'fraction_numerical': 0.007186328783137915, 'mean_word_length': 4.833132771760931, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Plan Arithmetic: Compositional Plan Vectors for Multi-Task Control

Autonomous agents situated in real-world environments must be able to master large repertoires of skills. While a single short skill can be learned quickly, it would be impractical to learn every task independently. Instead, the agent should share knowledge across behaviors such that each task can be learned efficiently, and such that the resulting model can generalize to new tasks, especially ones that are compositions or subsets of tasks seen previously. A policy conditioned on a goal or demonstration has the potential to share knowledge between tasks if it sees enough diversity of inputs. However, these methods may not generalize to a more complex task at test time. We introduce compositional plan vectors (CPVs) to enable a policy to perform compositions of tasks without additional supervision. CPVs represent trajectories as the sum of the subtasks within them. We show that CPVs can be learned within a one-shot imitation learning framework without any additional supervision or information about task hierarchy, and enable a demonstration-conditioned policy to generalize to tasks that sequence twice as many skills as the tasks seen during training. Analogously to embeddings such as word2vec in NLP, CPVs can also support simple arithmetic operations -- for example, we can add the CPVs for two different tasks to command an agent to compose both tasks, without any additional training.


Introduction
A major challenge in current machine learning is to not only interpolate within the distribution of inputs seen during training, but also to generalize to a wider distribution. While we cannot expect arbitrary generalization, models should be able to compose concepts seen during training into new combinations. With deep learning, agents learn high level representations of the data they perceive. If the data is drawn from a compositional environment, then agents that model the data accurately and efficiently would represent that compositionality without needing specific priors or regularization. In fact, prior work has shown that compositional representations can emerge automatically from simple objectives, most notably a highly structured distribution such as language. These techniques do not explicitly train for compositionality, but employ simple structural constraints that lead to compositional representations. For example, Mikolov et al. found that a language model trained to predict nearby words represented words in a vector space that supported arithmetic analogies: "king" -"man" + "woman" = "queen" [29]. In this work, we aim to learn a compositional feature space to represent robotic skills, such that the addition of multiple skills results in a plan to accomplish all of these skills.
Many tasks can be expressed as compositions of skills, where the same set of skills is shared across many tasks. For example, assembling a chair may require the subtask of picking up a hammer, which is also found in the table assembly task. We posit that a task representation that leverages this compositional structure can generalize more easily to more complex tasks. We propose learning an Figure 1: Compositional plan vectors embed tasks into a space where adding two vectors represents the composition of the tasks, and subtracting a sub-task leaves an embedding of the remaining sub-tasks needed for the task.
embedding space such that tasks could be composed simply by adding their respective embeddings. This idea is illustrated in Figure 2.
In order to learn these representations without additional supervision, we cannot depend on known segmentation of the trajectories into subtasks, or labels about which subtasks are shared between different tasks. Instead, we incorporate compositionality directly into the architecture of the policy. Rather than conditioning the policy on the static embedding of the reference demonstration, we condition the policy on the difference between the embedding of the whole reference trajectory and the partially completed trajectory that the policy is outputting an action for.
The main contributions of our work are the compositional plan vector (CPV) representation and a policy architecture that enables learning of CPVs without any sub-task level supervision. CPVs enable policies to generalize to significantly longer tasks, and they can be added together to represent a composition of tasks. We evaluate CPVs in the oneshot imitation learning paradigm [11,12,19] on a discrete-action environment inspired by Minecraft, where tools must be picked up to remove or build objects, as well as on a 3D simulated pick-and-place environment. For videos and code, please visit our project website https://sites.google.com/berkeley.edu/compositionalplanvectors/home.

Related Work
For many types of high dimensional inputs, Euclidean distances are often meaningless in the raw input space. Words represented as one-hot vectors are equally distant from all other words, and images of the same scene may have entirely different pixel values if the viewpoint is shifted slightly. This has motivated learning representations of language and images that respect desirable properties. Chopra et al. [3] showed that a simple contrastive loss can be used to learn face embeddings. A similar method was also used on image patches to learn general image features [35]. Word2vec found that word representations trained to be predictive of their neighbor words support some level of addition and subtraction [29,24,21]. More recently, Nagarajan used a contrastive approach in learning decomposable embeddings of images by representing objects as vectors and attributes as transformations of those vectors [30]. These methods motivate our goal of learning an embedding space over tasks that supports transformations such as addition and subtraction. Notably, these methods don't rely on explicit regularization for arithmetic operations, but rather use a simple objective combined with the right model structure to allow a compositional representation to emerge. Our method also uses a simple end-to-end policy learning objective, combined with a structural constraint that leads to compositionality.
Hierarchical RL algorithms learn representations of sub-tasks explicitly, by using primitives or goal-conditioning [4,32,26,13,27,7,2,39], or by combining multiple Q-functions [15,36]. Our approach does not learn explicit primitives or skills, but instead aims to summarize the task via a compositional task embedding. A number of prior works have also sought to learn policies that are conditioned on a goal or task [22,8,23,38,5,14,18,28,6,33], but without explicitly considering compositionality. Recent imitation learning methods have learned to predict latent intentions of demonstrations [16,25]. In the one-shot imitation learning paradigm, the policy is conditioned on reference demonstrations at both test and train time. This problem has been explored with metalearning [12] and metric learning [19] for short reaching and pushing tasks. Duan et al. used attention over the reference trajectory to perform block stacking tasks [11]. Our work differs in that we aim to generalize to new compositions of tasks that are out of the distribution of tasks seen during training. Hausman et al. obtain generalization to new compositions of skills by training a generative model over skills [17]. However, unlike our method, these approach does not easily allow for sequencing skills into longer horizon tasks or composing tasks via arithmetic operations on the latent representation.
Prior methods have learned composable task representations by using ground truth knowledge about the task hierarchy. Neural task programming and the neural subtask graph solver generalize to new tasks by decomposing a demonstration into a hierarchical program for the task, but require groundtruth hierarchical decomposition during training [40,37]. Using supervision about the relations between tasks, prior approaches have uses analogy-based objectives to learn task representations that decompose across objects and actions [31] or have set up a modular architectures over subtasks [1] or environments [9]. Unlike our approach, these methods require labels about relationships. We implicitly learn to decompose tasks without supervising the task hierarchy.

Compositional Plan Vectors
In this paper, we introduce compositional plan vectors (CPVs). The goal of CPVs is to obtain policies that generalize to new compositions of skills without requiring skills to be labeled and without knowing the list of skills that may be required. Consider a task named "red-out-yellow-in" which involves taking a red cube out of a box and placing a yellow cube into the box. A plan vector encodes the task as the sum of its parts: a plan vector for taking the red cube out of the box plus the vector for putting the yellow cube into the box should equal the plan vector for the full task. Equivalently, the plan vector for the full task minus the vector for taking the red cube out of the box should equal the vector that encodes "put yellow cube in box." If the list of all possible skills was known ahead of time, separate policies could be learned for each skill, and then the policies could be used in sequence. However, this knowledge is often unavailable in general and limits compositionality to a fixed set of skills. Instead, our goal is to formulate an architecture and regularization that are compositional by design and do not need additional supervision. With our method, CPVs acquire compositional structure because of the structural constraint they place on the policy. To derive the simplest possible structural constraint, we observe that the minimum information that the policy needs about the task in order to complete it is knowledge of the steps that have not yet been done. That is, in the cube example above, after taking out the red cube, only the "yellow-in" portion of the task is needed by the policy. One property of this representation is that task ordering cannot be represented by the CPV because addition is commutative. If ordering is necessary to choose the right action, the policy will have to learn to decode which component of the compositional plan vector must be done first.
As an example, let v be a plan vector for the "red-out-yellow-in" task. To execute the task, a policy π(o 0 , v) outputs an action for the first observation o 0 . After some number t of timesteps, the policy has successfully removed the red cube from the box. This partial trajectory O 0:t can be embedded into a plan vector u, which should encode the "red-out" task. We would like ( v − u) to encode the remaining portion of the task, in this case placing the yellow block into the box. In other words, π(o t , v − u) should take the action that leads to accomplishing the plan described by v given that u has already been accomplished. In order for both v and v − u to encode the yellow-in task, u must not encode as strongly as v. If v is equal to the sum of the vectors for "red-out" and "yellow-in," then v may not encode the ordering of the tasks. However, the policy π(o 0 , v) should have learned that the box must be empty in order to perform the yellow-in task, and that therefore it should perform the red-out task first.
We posit that this structure can be learned without supervision at the subtask-level. Instead, we impose a simple architectural and arithmetic constraints on the policy: the policy must choose its action based on the arithmetic difference between the plan vector embedding of the whole task and the plan vector embedding of the trajectory completed so far. Additionally, the plan vectors of two halves of the same trajectory should add up to the plan vector of the whole trajectory, which we can write down as a regularization objective for the embedding function. By training the policy and the embedding function together to optimize their objectives, we obtain an embedding of tasks that supports compositionality and generalizes to more complex tasks. In principle, CPVs can be used with any end-to-end policy learning objective, including behavioral cloning, reinforcement learning, or inverse reinforcement learning. In this work, we will validate CPVs in a one-shot imitation learning setting.
One-shot imitation learning setup. In one-shot imitation learning, the agent must perform a task conditioned on one reference example of the task. For example, given a demonstration of how to fold a paper crane, the agent would need to fold a paper crane. During training, the agent is provided with pairs of demonstrations, and learns a policy by predicting the actions in one trajectory by using the second as a reference. In the origami example, the agent may have trained on demonstrations of folding paper into a variety of different creatures. Plan vectors. We define a function g φ (O k:l ), parameterized by φ, which takes in a trajectory and outputs a plan vector. The plan vector of a reference trajectory g φ (O ref 0:T ) should encode an abstract representation of the milestones required to accomplish the goal. Similarly, the plan vector of a partially accomplished trajectory g φ (O 0:t ) should encode the steps already taken. We can therefore consider the subtraction of these vectors to encode the steps necessary to complete the task defined by the reference trajectory. Thus, the policy can be structured as a function parameterized by θ that takes in the trajectory and is learned end-to-end.
In this work we use a fully observable state space and only consider tasks that cause a change in the state. For example, we do not consider tasks such as lifting a block and placing it exactly where it was, because this does not result in a useful change to the state. Thus, instead of embedding a whole trajectory O 0:t , we limit g to only look at the first and last state of the trajectory we wish to embed.
Training. With π defined as above, we learn the parameters of the policy with imitation learning. Dataset D containing N demonstrations paired with reference trajectories is collected. Each trajectory may be a different arbitrary length, and the tasks performed by each pair of trajectories are unlabeled. The demonstrations include actions, but the reference trajectories do not. In our settings, the reference trajectories only need to include their first and last states. Formally, where T i is the length of the i th reference trajectory and H i is the length of the ith demonstration. Given the policy architecture defined in Equation 1, the behavioral cloning loss for a discrete action policy is We also introduce a regularization loss function to improve compositionality by enforcing that the sum of the embeddings of two parts of a trajectory is close to the embedding of the full trajectory. We denote this a homomorphism loss L Hom because it constrains the embedding function g to preserve the structure between concatenation of trajectories and addition of real-valued vectors. We implement the loss using the triplet margin loss from [34] with a margin equal to 1: l tri (a, p, n) = max{||a − p|| 2 − ||a − n|| 2 + 1.0, 0} Finally, we follow James et al. in regularizing embeddings of paired trajectories to be close in embedding space, which has been shown to improve performance on new examples [19]. This "pair" loss L Pair pushes the embedding of a demonstration to be similar to the embedding of its reference trajectory and different from other embeddings, which enforces that embeddings are a function of the behavior within a trajectory rather than the appearance of a state.
for any j = i. We empirically evaluate how these losses affect the composability of learned embeddings. While L Pair leverages the supervision from the reference trajectories, L Hom is entirely self-supervised.
Measuring compositionality. To evaluate whether the representation learned by g is compositional, we condition the policy on the sum of plan vectors from multiple tasks and measure the policy's success rate. Given two reference trajectories O ref i 0:T i and O ref j 0:T j , we condition the policy on The policy is successful if it accomplishes both tasks. We also evaluate whether the representation generalizes to more complex tasks. If the relevant objects are in the box, the agent must first remove the lid to interact with the object and also return the lid to the box in order to complete a task.

Sequential Multitask Environments
We introduce two new learning environments, shown in Figures 3 and 4, that test an agent's ability to perform tasks that require different sequences and different numbers of sub-skills. We designed these environments such that the actions change the environment and make new sub-goals possible: in the 3D environment, opening a box and removing its contents makes it possible to put something else into the box. In the crafting environment, chopping down a tree makes it is possible to build a house. Along with the environments, we will release code to generate demonstrations of the compositional tasks.

Crafting Environment
The first evaluation domain is a discrete-action world where objects can be picked up and modified using tools. The environment contains 7 types of objects: tree, rock, logs, wheat, bread, hammer, axe. Logs, hammers, and axes can be picked up, and trees and rocks block the agent. The environment allows for 6 actions: up, down, left, right, pickup, and drop. The transitions are deterministic, and only one object can be held at a time. Pickup has no effect unless the agent is at the same position as a pickup-able object. Drop has no effect unless an object is currently held. When an object is held, it moves with the agent. Unlike the Malmo environment which runs a full game engine [20], this environment can be easily modified to add new object types and interaction rules. We define 5 skills within the environment, ChopTree, BreakRock, BuildHouse, MakeBread, and EatBread, as illustrated in Figure 3.

3D Pick and Place Environment
The second domain is a 3D simulated environment where a robot arm can pick up and drop objects. Four cubes of different colors, as well as a box with a lid, are randomly placed within the workspace. The robot's action space is a continuous 4-dimensional vector: an (x, y) position at which to close the gripper and an (x, y) position at which to open the gripper. The z coordinate of the grasp is chosen automatically. The observation space is a concatenation of the (x, y, z) positions of each of the 4 objects, the box, and the box lid. We define 3 families of skills within the environment: PlaceInCorner, Stack, and PlaceInBox, each of which can be applied on different objects or pairs of objects. Considering tasks that use 1 to 2 skills, there are 420 different tasks. An example of each skill is shown in Figure 4.

Experiments
Our experiments aim to understand how well CPVs can learn tasks of varying complexity, how well they can generalize to tasks that are more complex than those seen during training (thus demonstrating compositionality), and how well they can handle additive composition of tasks, where the policy is expected to perform both of the tasks in sequence. We hypothesize that, by conditioning a policy on the subtraction of the current progress from the goal task embedding, we will learn a task representation that encodes tasks as the sum of their component subtasks. We additionally evaluate how regularizing objectives improve generalization and compositionality. Implementation. We implement g φ and π θ as neural networks. For the crafting environment, where the observations are RGB images, we use the convolutional architecture in Figure 5. The encoder g outputs a 512 dimensional CPV. The policy, shaded in red, takes the subtraction of CPVs concatenated with features from the current observation and outputs a discrete classification over actions.
For the 3D environment, the observation is a state vector containing the positions of each object in the scene, including the box and box lid. The function g again concatenates the inputs, but here the network is fully connected, and the current observation is directly concatenated to the subtraction of the CPVs. To improve the performance of all models and comparisons, we use an object-centric policy inspired by Devin et al. [10], where the policy outputs a softmaxed weighting over the objects in the state. The position of the most attended object is output as the first coordinates of the action (where to grasp). The object attention as well as the features from the observation and CPVs are passed to another fully connected layer to output the position for placing.
Data generation. For the crafting environment, we train all models on a dataset containing 40k pairs of demonstrations, each pair performs the same task. The demonstration pairs are not labeled with what task they are performing. The tasks are randomly generated by sampling 2-4 skills with replacement from the five skills listed previously. A planning algorithm is used to generate demonstration trajectories. For the 3D environment, we collect 180k trajectories of tasks with 1 and 2 skills. All models are trained on this dataset to predict actions from the environment observations shown in Figure 6a. For both environments, we added 10% of noise to the planner's actions but discarded any trajectories that were unsuccessful. The data is divided into training and validation sets 90/10. To evaluate the models, reference trajectories were either regenerated or pulled from the validation set. Compositions of trajectories were never used in training or validation.

Comparisons.
We compare our CPV model to several one-shot imitation learning models. All models are based on Equation 2, where the policy is function of four images: The naïve baseline simply concatenates the four inputs as input to a neural network policy. The TECNets baseline is an implementation of task embedding control networks from [19], where the embeddings are normalized to a unit ball and a margin loss is applied over the cosine distance to push together embeddings of the same task. The policy in TECNets is conditioned on the static reference embedding rather than the subtraction of two embeddings. For both TECNets and our model, g is applied to the concatenation of the two input observations. We perform several ablations of our model, which includes the CPV architecture (including the embedding subtraction as input the policy), the homomorphism regularization, and the pair regularization. We compare the plain version of our model, where the objective is purely imitation learning, to versions that use the regularizations. CPV-Plain uses no regularization, CPV-Pair uses only L Pair , Table 1: Evaluation of generalization and compositionality in the craftin environment. Policies were trained on tasks using between 1 and 4 skills. We evaluate the policies conditioned on reference trajectories that use 4, 8, and 16 skills. We also evaluate the policies on the composition of skills: "2, 2" means that the embeddings of two demonstrations that each use 2 skills were added together, and the policy was conditioned on this sum. For the naïve model, we instead average the observations of the references, which performed somewhat better. All models are variants on the architecture in Figure 5. The max horizon is three times the average number of steps used by the expert for that length of task: 160, 280, and 550, respectively. Numbers are all success percentages. CPV-Hom uses only L Hom , and CPV-Full uses both. To ablate the effect of the architecture vs the regularizations, we run the same set of comparisons for a model denoted TE (task embeddings) which has the same architecture as TECNets without normalizing embeddings. These experiments find whether the regularization losses produce compositionality on their own, or whether they work in conjunction with the CPV architecture. Results. We evaluate the methods on both domains. To be considered successful in the crafting environment, the agent must perform the same sub-skills with the same types of objects as those seen in the reference trajectory. The results on the crafting environment are shown in Table 1, where we report the mean and standard deviation across 3 independent training seeds. We see that both the naïve model and the TECNet model struggle to represent these complex tasks, even the 4 skill tasks that are in the training distribution. We also find that both the CPV architecture and the regularization losses are necessary for both generalizing the longer tasks and composing multiple tasks. The pair loss seems to help mostly with generalization, while the homomorphism losses helps more with compositionality. CPVs are able to generalize to 8 and 16 skills, despite being trained on only 4 skill combinations, and achieve 76% success at composing two tasks just by adding their embedding vectors. Recall that CPVs are not explicitly trained to compose multiple reference trajectories in this way -the compositionality is an extrapolation from the training. The TE ablation, which does not use the subtraction of embeddings as input to the policy, shows worse compositionality than our method even with the homomorphism loss. This supports our hypothesis that structural constraints over the embedding representation contribute significantly to the learning.
These trends continue in the pick and place environment in Table 2, were we report the mean and standard deviation across 3 independent training seeds. In this environment, a trajectory is successful if the objects that were moved in the reference trajectory are in the correct positions: placed in each corner, placed inside the box, or stacked on top of a specific cube. As expected, TECNet performs well on 1 skill tasks which only require moving a single object. TECNet and the naïve model fail to compose tasks, but the CPV model performs as well at composing two 1-skill tasks as it does when imitating 2-skill tasks directly. As before, the TE ablation fails to compose as well as CPV, indicating that the architecture and losses together are needed to learn composable embeddings.

Discussion
Many tasks can be understood as a composition of multiple subtasks. To take advantage of this latent structure without subtask labels, we introduce the compositional plan vector (CPV) architecture along with a homomorphism-preserving loss function, and show that this learns a compositional representation of tasks. Our method learns a task representation and multi-task policy jointly. Our main idea is to condition the policy on the arithmetic difference between the embedding of the goal task and the embedding of the trajectory seen so far. This constraint ensures that the representation space is structured such that subtracting the embedding of a partial trajectory from the embedding of the full trajectory encodes the portion of the task that remains to be completed. Put another way, CPVs encode tasks as a set of subtasks that the agent has left to perform to complete the full task. CPVs enable policies to generalize to tasks twice as long as those seen during training, and two plan vectors can be added together to form a new plan for performing both tasks.
We evaluated CPVs in a one-shot imitation learning setting. Extending our approach to a reinforcement learning setting is a natural next step, as well as further improvements to the architecture to improve efficiency. A particularly promising future direction would be to enable CPVs to learn from unstructured, self-supervised data, reducing the dependence on hand-specified objectives and reward functions.
connected layer maps this directly to continuous valued (x1, y1). This means that the policy can place at any position in the workspace. The naïve model, TECNet model, and CPV models all use this object-centric policy, then only differ in how the input to the policy.

B Hyperparameters
We compared all models across embedding dimension sizes of [64,128,256, and 512]. In the crafting environment, the 512 size was best for all methods. In the grasping environment, the 64 size was best for all methods. For TECNets, we tested λctr = 1 and 0.1, and found that 0.1 was best. All models are trained on either k-80 GPUs or Titan X GPUs.

C Additional Experiments
We ran a pared down experiment on a ViZDoom environment to show the method working from first person images, as shown in C. In the experiment, the skills are reaching 4 different waypoints in the environment. The actions are "turn left," "turn right," and "go forward." The observation space consists of a first person image observation as well as the (x, y) locations of the waypoints. We evaluate on trajectories that must visit 1 or 2 waypoints (skills), and also evaluate on the compositions of these trajectories. The policies were only trained on trajectories that visit up to 3 waypoints. These evaluations are shown in 3.