DYNAMIC COMPUTATION GRAPHS
Neural networks that compute over graph structures are a natural fit for problems in a variety of domains, including natural language (parse trees) and cheminformatics (molecular graphs). However, since the computation graph has a different shape and size for every input, such networks do not directly support batched training or inference. They are also difficult to implement in popular deep learning libraries, which are based on static data-flow graphs. We introduce a technique called dynamic batching, which not only batches together operations between different input graphs of dissimilar shape, but also between different nodes within a single input graph. The technique allows us to create static graphs, using popular libraries, that emulate dynamic computation graphs of arbitrary shape and size. We further present a high-level library1 of compositional blocks that simplifies the creation of dynamic graph models. Using the library, we demonstrate concise and batch-wise parallel implementations for a variety of models from the literature.
1 INTRODUCTION
Training deep neural networks directly on minimally pre-processed corpora has led to many recent performance breakthroughs, mainly on problems in domains such as vision (Krizhevsky et al., 2012) and natural language (Bahdanau et al., 2015) where the inputs can be cast as dense n-dimensional arrays (henceforth tensors), or sequences of tensors. These successes exploit the effectiveness of training via gradient descent on mini-batches of tens to hundreds of inputs, implemented using the parallel SIMD capabilities of modern GPUs (Oh & Jung, 2004) and multi-core CPUs (Vanhoucke et al., 2011). This, in turn has led to a proliferation of libraries making it easier to train and deploy such models, by expressing them in terms of differentiable data-flow graphs over tensors (Abadi et al., 2016; Theano Development Team, 2016; Collobert et al., 2011).
However, there is also a long history of neural networks that compute over structures such as parse trees (Pollack, 1990), logical terms (Goller & Kuchler, 1996), and molecular graphs (Bianucci et al., 2000). In these models, each distinct input has a different computation graph structure; we say that they use dynamic computation graphs (DCGs). Such models continue to be developed and have recently yielded superior results on problems such as sentiment classification and semantic relatedness (Tai et al., 2015; Li et al., 2015), question-answering (Andreas et al., 2016), and screening of chemical compounds (Kearnes et al., 2016). Despite these successes, most practitioners avoid DCGs for implementation reasons. For example, Bowman et al. (2016) assert that “because TreeRNNs use a different model structure for each sentence ... efficient batching is impossible in standard implementations”. Moreover, even if efficient batching were possible in principle, current libraries such as TensorFlow (Abadi et al., 2016) assume that the data-flow graph is static (i.e. is the same for each input) and impose a significant cost to graph construction, which makes it infeasible to build a new graph for each input.
Section 2 introduces dynamic batching, which enables efficient batching for training and inference with DCGs. Dynamic batching runs DCGs efficiently with existing libraries that only support static data-flow graphs; e.g. the same static graph can run a TreeRNN over any parse tree. We present empirical results for our implementation in TensorFlow. Section 3 presents a combinator library for concisely implementing models with DCGs using dynamic batching. Section 4 concludes.
1The library is called TensorFlow Fold and lives at http://github.com/tensorflow/fold.
2 DYNAMIC BATCHING
In deep learning libraries like TensorFlow, computations are manually batched. The computation is expressed as a static graph of mathematical operations, such as y = σ(x · w + c), which are polymorphic in batch size; an input x of dimensions (b, n) will yield an output of dimensions (b,m), where b is the batch size. With DCGs, the graph of operations is not static, but is assumed to be different for every input, so multiple inputs no longer naturally batch together in the same way. The dynamic batching algorithm overcomes this difficulty. Given a set of computation graphs as input, each of which has a different size and topology, it will rewrite the graphs by batching together all instances of the same operation that occur at the same depth in the graph. The rewriting process inserts additional concat and gather operations to move data between the batched operations; the indices to gather encode the topology of the original input graphs.
We distinguish between individual operations appearing as nodes in the underlying data-flow graph, such as addition or matrix-multiply, and small sub-graphs that conceptually act as functions over tensors, such as a feed-forward layer or LSTM cell. We refer to the former as “ops”, and to the latter as “operations.” Operations, (i.e. sub-graphs), form the building-blocks from which neural networks with DCGs are composed; dynamic batching schedules operations, not ops. Our algorithm requires that all operations which might be used be specified in advance, and it enumerates them for scheduling purposes. For example, a binary TreeRNN for NLP parse trees has two operations: embedding table lookups for words at the leaves of the tree, and RNN cells for the non-terminals.
The inputs and outputs of operations have tensor types. Each input or output may have a different type, but all types must be fixed and fully specified in advance. A tensor type consists of a shape, x1, . . . xn, together with a scalar data type (e.g. float32). The inputs to an operation shall be tensors of dimension (b, x1, . . . xn), where b is the batch size and x1 . . . xn is the shape of corresponding input tensor type. The outputs must all be tensors of dimension (b, y1, . . . ym), where y1, . . . ym is the shape of the corresponding output tensor type. Operations must be polymorphic with respect to the batch size, because the batch size will change each time the operation is invoked, depending on the topologies of the input graphs. However, their tensor types are fixed, so that it is possible to assign a known tensor type to each edge in the input computation graph.
The dynamic batching algorithm takes a directed acyclic computation graph as input. A batch of multiple input graphs can be treated as a single disconnected graph. Source nodes are constant tensors, and non-source nodes are operations. Edges connect one of the outputs of a node to one of the inputs of another node. Scheduling is performed using a greedy algorithm:
• Assign a depth to each node in the graph. Nodes with no dependencies (constants) are assigned depth zero. Nodes with only dependencies of depth zero are assigned depth one, nodes whose dependencies have a maximum depth of one get assigned depth two, etc.
• Insert pass-through (identity) operations so that an operation at depth d + 1 only refers to results at depth d.
• Batch together all nodes invoking the same operation at the same depth into a single node.
• Concatenate all outputs which have the same depth and tensor type. The order of concatenation corresponds to the order in which the dynamic batching operations were enumerated.
• Assign a label (d, t, i) to each edge in the original graph, where d is the depth, t is the tensor type, and i is the integer index for that edge into the (concatenated) outputs for d, t. The schedule for the graph consists of the indices i for all edges, which are grouped together by depth and operation.
In our TensorFlow implementation, each dynamic operation is instantiated once in the static data-flow graph. The inputs to each operation are tf.gather ops, and the outputs are fed into tf.concat ops, as described above. These TensorFlow ops are then placed within a tf.while_loop. Each iteration of the loop will evaluate all of the operations at a particular depth. The loop maintains state variables for each tensor type t, and feeds the output of concat for tensor type t and iteration d into the input of the gathers at tensor type t and iteration d+ 1. The indices for gather at iteration d are drawn from the edge labels i for depth d in the schedule. The initial values for the state variables at iteration/depth 0 are the constants in the input graph.
Dynamic batching allows us to construct a static TensorFlow graph that contains a single instance of each operation, yet can emulate input graphs of arbitrary size and topology where operations may appear an arbitrary number of times. The TensorFlow concat, gather, and while_loop ops are all differentiable, so gradients calculations and back-propagation do not require any additional code.
For example, a binary TreeRNN as described above yields a TensorFlow data-flow graph with a tf.while_loop whose body is shown on the left of Figure 1. Here each gather has an additional input (the indices for the given op at the given depth) which picks out which elements the operations are to be called with. The long downward arrows are the pass-throughs. The algorithm consumes a tree such as the one shown on the right of Figure 1 and turns it into inputs for the gather operations at each depth (here depth is the loop counter for the tf.while_loop.)
2.1 EXPERIMENTAL RESULTS
We have implemented dynamic batching as part of a new library, TensorFlow Fold, and designed a synthetic speed benchmark to compare it with manual batching in native TensorFlow. The benchmark uses the same underlying kernels and execution engine in both cases. Native TensorFlow cannot batch together trees of different shapes so, for testing purposes, we use a batch of random binary trees, all of which have the same shape. These test results thus represent a best-case scenario, in which all operations can be batched together perfectly. For the manual batching tests, we construct a static data-flow graph of operations corresponding to the shape of the tree. For the dynamic batching tests, we traverse each tree to construct a schedule, as described above.
The leaves of the tree are lookups into an embedding table, while the non-terminals implement a variant of the Tree-LSTM (Tai et al., 2015) equations. The tree size is 128, with a state size of 1024 for the LSTM. The CPU tests were run on a Dell z620 workstation with dual 8-core Intel Xeon processors (32 hardware threads), and the GPU tests were done using a consumer Nvidia GeForce GTX-1080 card. We compare manual batching, dynamic batching where all trees have the same shape, and dynamic batching where each tree has a different shape (the column marked “full dynamic”). There is no measurable penalty for dealing with trees of different shapes.
The test results shown in Table 1 emphasize the importance of batching, especially on GPUs. TensorFlow will launch a GPU kernel for every node in the tree, so there is a fixed overhead, proportional to the size of the tree, that dominates execution for small batch sizes. TensorFlow does not begin to saturate the GPU until relatively large batch sizes – 1024 or higher. The difference in speed between fully-batched and unbatched is over 160x.
Dynamic batching has less kernel invocation overhead because the data-flow graph is smaller. Dynamic batching instantiates each operation only once, and invokes it once for each depth, so the number of kernel invocations is log(n), rather than n, where n is tree size. Dynamic batching thus achieves substantial speedups even at batch size 1, because it batches operations at the same depth within a single tree.
However, the extra concat and gather ops that dynamic batching inserts do have a cost. The “cost ratio” column above shows the ratio between dynamic and manual batching, in the case where all trees in the batch have the same shape. The cost is only 20% for inference on GPUs with batch-size 1, but rises to 60% for training with backpropagation. The cost is mainly visible at large batch sizes, because it is balanced by the benefit of within-tree batching at smaller sizes.
Even with the cost, dynamic batching yields a 120x speedup over using a batch size of 1 on GPU, and 28x on CPU. The “speedup ratio” column above shows the ratio between the per-tree time for dynamic batching on random shapes (“full dynamic”), versus manual batching with a batch size of 1. Note that using a batch size of 1 is not actually feasible for TensorFlow, because TensorFlow has a large graph construction overhead, which is not included in these measurements, but it may apply to other libraries that lack such overhead.
3 A COMBINATOR LIBRARY FOR NEURAL NETWORKS
In addition to dynamic batching, the TensorFlow Fold library provides a set of combinators that simplify the task of constructing neural networks for DCGs. Our goal here is to show how dynamic batching enables implementing deep learning models (which are growing ever more complex) at a higher level of abstraction than manual batching. This in turn facilitates a more rapid feedback loop for trying out novel model variants, and thus obtaining superior results.
The design of the library was inspired by functional programming techniques such as parser combinators (Hutton & Meijer, 1996) and arrows (Hughes, 2000). In a combinator library computations are structured compositionally, by plugging together simpler computations in various ways. The basic unit of computation in TensorFlow Fold is a block, essentially a function from input to output. In a typical DCG model, the input is a graph or tree of some kind, and the output is a vector, which can be attached to a loss for training.
For example, consider a model where the inputs are sequences of words, of varying lengths, and the output is a sentence vector. Our library provide several different ways of handling sequences. Given a simpler block f that operates on elements of the sequence, or g on pairs of elements, we define the following combinators:
• Map(f): yields [f(x1), f(x2), . . . f(xn)]. Applies f to each element of the sequence, e.g. embedding each of the words of a sentence into RN .
• Fold(g, z): yields g(. . . g(g(z, x1), x2), . . . xn). Applies g sequentially in a leftward chain, e.g. running an RNN over a sequence. By default z = 0.
• Reduce(g): yields g(Reduce([x1, . . . xbn/2c]),Reduce([xbn/2c+1, . . . xn])). Applies g in a balanced tree,2 e.g. max or sum-pooling over the elements.
Note that it is not necessary to pad or truncate sequences to the same length; dynamic batching handles sequences of differing lengths.
3.1 TYPE SYSTEM
Blocks are statically typed; each block has an input type and an output type. Types are inferred where possible, but must be explicitly specified in some cases. A type is one of the following:
• Input denotes objects in the host language (Python), such as trees and dictionaries.
• Tensordtype,shape denotes tensors of a particular dtype and shape. 3
• Tuple(t1, . . . tn), denotes a tuple of values of types t1, . . . tn.
• Sequence(t), denotes a sequence of elements of type t, of any length.
• Void is the unit type.
For example Sequence(Sequence(Tuple(Tensorfloat32,[],Tensorint8,[3,4]))) denotes jagged arrays whose elements are pairs (float32, int83×4).
3.2 BLOCKS AND COMBINATORS
Blocks are composed hierarchically; a block expression is always a tree. The non-terminals in the tree are combinators such as Map and Fold, which take simpler blocks as arguments. The leaves of the tree are atomic blocks, which include the following:
• Scalar: Input → Tensor Convert a Python scalar to a tensor.
• Tensor: Input → Tensor Convert a NumPy array to a tensor.
• Function(h): [Tensor or Tuple(Tensor , . . .)]→ [Tensor or Tuple(Tensor , . . .)] Defines an operation h (see Section 2) over tensors. Operations with multiple inputs and outputs use tuples of tensors.
• InputTransform(h): Input → Input Applies a user-defined Python function h to pre-process the input.
In addition to the the sequence combinators described above, important combinators in the library include the following:
• b1 >> b2: Function composition; the output of b1 is fed to the input of b2.
• Record({l1 : b1 , . . . ln : bn}: Input → Tuple(t1, . . . tn) Takes a Python dictionary or tuple as input, and applies each block bi to the field labeled li, to yield an object of type ti. Returns a tuple of the results for all fields.
• OneOf(b1, . . . bn): Input → t Conditionally dispatches on its input to one of the blocks b1, . . . bn.
• Optional(b): Input → t Applies b if the input is not None, otherwise returns zeros. A special case of OneOf.
• AllOf(b1, . . . bn): t0 → Tuple(t1, . . . tn) Passes its input of type t0 to each of the blocks b1, . . . bn, returning a tuple of results.
2Reduce uses a balanced tree rather than a chain in order to minimize computation depth and provide more opportunities for batching.
3Note that the leading batch size for tensors is not part of the shape of the corresponding Tensor type.
3.3 PIPELINES
Assume we have a set of (text, label) pairs as input and wish to predict the label from the text. The text consists of words, and we want to use an array of pretrained word embeddings (word_matrix) and corresponding dictionary mapping words to indices (word_idx). We call word_idx.get(word) to obtain the index of word in word_matrix, or None if word is unknown.
We start by creating a block which embeds each word into a continuous space:
word2vec = (InputTransform(word_idx.get) >> Optional(Scalar('int32')) >> Function(Embedding(initializer=word_matrix)))
This block uses an InputTransform to get the index of a word, which is passed to an Optional block that converts the scalar index to a tensor (or 0 if None). This in turn gets passed to an Embedding operation, which performs a lookup into an embedding table.
With word2vec in hand, we can define text2vec, which embeds sentences:
split = InputTransform(str.split) rnn_cell = Concat() >> Function(FC(d, activation=tf.nn.relu)) text2vec = split >> Map(word2vec) >> Fold(rnn_cell, Zeros(d))
We use an InputTransform to split the string into words. Then we map the words to vectors with word2vec, and combine the word vectors with a simple RNN, which uses a single fully connected layer FC with d hidden units. The Zeros block defines the initial state for the RNN.
Assume there are n labels; we use a linear layer with n outputs to get unscaled logits:
text2logits = text2vec >> Function(FC(n, activation=None))
For training, we create a Record block to convert the label to a tensor as well, and calculate loss:
record = Record([('text', text2logits), ('label', Scalar('int32'))]) loss = record >> Function(tf.nn.sparse_softmax_cross_entropy)
Finally, we create a Compiler, which validates a block, performs type-checking, and sets up dynamic batching in TensorFlow. Outputs of a compiled block are available as TensorFlow tensors, so training now proceeds as it would for any other TensorFlow model:
compiler = Compiler.create(loss) cross_entropy = Compiler.output_tensors[0] train_op = tf.train.AdamOptimizer().minimize(cross_entropy)
3.4 COMPLEX COMPOSITIONS
Recently, Raffel & Ellis (2016) have introduced an attention model for feed-forward neural networks. The model generalizes average-pooling and is defined as:
et = a(ht), αt = exp(et)∑T
k=1 exp(ek) , c = T∑ t=1 αtht (1)
where a is a learnable function.
In this model, the block architecture is not a simple pipeline (i.e. a composition using >>) but instead forms a directed acyclic graph, as illustrated in Figure 2. A Composition block allows blocks to be composed into DAGs. The model code and details may be found in Appendix A.
3.5 RECURSIVE DEFINITIONS
N -ary Tree-LSTMs (Tai et al., 2015, sec. 3.2) generalize LSTMs from 1 toN previous states. In Tai et al. (2015, sec. 5.1) they are applied to classify sentences from the Stanford Sentiment Treebank. This corpus consists of binarized constituency parse trees of one-sentence movie reviews, where every node has a sentiment label. At the leaves of the tree, words are mapped to word-embedding vectors which serve as the input to a binary tree-LSTM with 0 for the previous states. At the internal nodes, the LSTM takes 0 as input, and previous states from its two children. More formally,
hword = TreeLSTM(Embedding(word), 0, 0) (2) hleft,right = TreeLSTM(0, hleft, hright) (3)
where TreeLSTM(x, hleft, hright) is a learnable function corresponding to Tai et al. (2015) eqs. 9-14 with N = 2. Since a tree is a recursive data type, a model that processes trees must be recursively defined, as illustrated by the cycle in Figure 2. A ForwardDeclaration allows the creation of recursive models:
expr = ForwardDeclaration() word = AllOf(Record([('word', word2vec)]), Zeros((state_size, state_size)) pair = AllOf(Zeros(embedding_size), Record([('left', expr()), ('right', expr())])) expr_def = (OneOf(key_fn=len, case_blocks=[(1, word), (2, pair)]) >> TreeLSTM(state_size)) expr.resolve_to(expr_def)
A forward declaration like expr is not itself a block, but may be called (using the expr() syntax) to create references – i.e. blocks which refer to the declaration. The subsequent call to resolve_to then updates all the references to refer to expr_def.
The word2vec block is as defined in Section 3.3.
3.5.1 EXPERIMENTAL RESULTS
Here we briefly report on some experiments with our implementation of N -ary Tree-LSTMs for sentiment analysis. While we set a new state-of-the-art, that is not really the point here. Our models are not particularly original, and could certainly be implemented without using TensorFlow Fold. What Fold does is to enable simpler and more concise definitions (see Table 3), along with faster execution, thus making it easier to rapidly explore novel model variants.
We used constituency Tree-LSTMs with tuned Glove vectors for word embedding, which achieved the best results of all sentiment models presented in Tai et al. (2015). In addition to this specific model, we have explored several novel variants.4 In particular, Tai et al. (2015) employed non-
4Unsuccessful variants included standard LSTMs (i.e. having only a single forget gate) accepting pooled histories from their children, and models based on character rather than word-level embeddings.
recurrent dropout and L2 weight regularization. We eliminated weight regularization in favor of the recurrent dropout scheme introduced by Semeniuta et al. (2016) and increased the LSTM state size from 150 to 300, leaving all other hyperparameters unchanged.
Results are shown in Table 2, including the best previously reported results. Fine-grained accuracy is measured for all trees and calculated based on the five possible labels. Binary accuracy is measured only for trees with non-neutral sentiment, and is based on negative vs. positive classification. The numbers in parentheses are standard deviations. Tai et al. (2015) report five independent runs, our results are based on thirty independent runs.5 Noting the small size of this dataset (8544/1101/2210 trees for train/dev/test), we further evaluated an ensemble consisting of these thirty independently trained models; this variant sets a new state-of-the-art on both subtasks.
3.6 GRAPH CONVOLUTIONS
As a final example, we have used the Fold library to implement the graph convolution model introduced by Kearnes et al. (2016) for molecules, which are represented as undirected graphs of atoms. The code is more complex than our previous examples because it involves nested Composition blocks, and is given in Appendix B.
4 DISCUSSION
Neural architectures with dynamic computation graphs suffer from inefficient batching and poor tooling. Dynamic batching solves the former problem in full generality, we believe for the first time. The SPINN architecture (Bowman et al., 2016) is an alternative stack-based approach that also enables efficient batching with DCGs, but it is limited to binary trees, and requires padding/truncation to handle trees of different sizes. The Fold library addresses the tooling problem by providing a high-level combinator library which is intended to make it easy for practitioners to rapidly develop and iterate on architectures with DCGs.
The experimental results presented in section 2.1 quantify the impact of dynamic batching. The impact of the combinator library is harder to demonstrate quantitatively. One way to approach this (with a large grain of salt) is by comparing lines of code, which we do in Table 3, vs. the original author’s sources. See Appendix C for details on the comparison protocol. Of course, a very short implementation is suboptimal if it comes at the cost of flexibility. The results in Section 3.5.1 show that models from the literature can be reimplemented in Fold, then extended to achieve superior performance. We suspect that other models with DCGs will have quite a bit of “head room” as well, due to simply having less work done tuning them compared with more mainstream architectures.
5Munkhdalai & Yu (2016a;b) do not report standard deviations or number of runs.
A FEED-FORWARD ATTENTION
The feed-forward attention model from Section 3.4 may be implemented in Fold as follows:
attention = Composition() with attention.scope():
h = attention.input exp_e = Map(a >> Function(tf.exp)).reads(h) z = (Sum() >> Broadcast()).reads(exp_e) alpha = ZipWith(Function(tf.div)).reads(exp_e, z) c = (ZipWith(Function(tf.mul)) >> Sum()).reads(alpha, h) attention.output.reads(c)
Within a composition scope, blocks may be wired together with reads, provided no directed cycles are formed. The input and output properties are used to define the overall inputs and outputs of the composition block. This example introduces several additional block types:
• Sum is a specialization of Reduce that performs elementwise addition. • ZipWith is a variant of Map that accepts n sequences as input and applies an n-ary function f elementwise (stopping when the end of the shortest input sequence is reached).
• Broadcast creates a Sequence(t) from a single t, repeating the same element endlessly.
B GRAPH CONVOLUTIONS
This section implements the graph convolution model introduced by Kearnes et al. (2016), for molecules represented as undirected graphs of atoms. There are real-valued feature vectors for each atom and for each distinct pair of atoms. For a molecule having N atoms, we index its atom feature vectors as ai ∈ Rn for 1 ≤ i ≤ N . We index its pair feature vectors as pi,j ∈ Rm for 1 ≤ i, j ≤ N , where pi,j = pj,i and pi,i = 0. The core of the graph convolution model is the weave module, which combines atom-level and pair-level features using six learnable functions (typically fully connected ReLU layers). The weave module can be stacked arbitrarily to create deep graph convolution models. Denoting inputs and outputs by x and y superscripts respectively, the weave module is:
ayi = fA(fA→A(a x i ), N∑ j=1 fP→A(p x i,j)) (4)
pyi,j = fP (fA→P (a x i , a x j ) + fA→P (a x j , a x i ), fP→P (p x i,j)) (5)
where fA, fP , fA→A, fA→P , fP→A and fP→P are learnable functions.
It is noteworthy that the ax → py calculation involves a nested scan over the atoms; for each ai we must calculate fA→P (axi , a x j ) + fA→P (a x j , a x i ) for all 1 ≤ j ≤ N :
a_i_to_p = Composition() with a_i_to_p.scope():
a_x_i = Broadcast().reads(a_i_to_p.input[0]) a_x = a_i_to_p.input[1] f_i_j = ZipWith(Concat() >> f_a_p).reads(a_x_i, a_x) f_j_i = ZipWith(Concat() >> f_a_p).reads(a_x, a_x_i) p = ZipWith(Sum()).reads(f_i_j, f_j_i) a_i_to_p.output.reads(p)
The input to the a_i_to_p composition block is (axi , ax). It has the type
Tuple(Tensorfloat32,[n],Sequence(Tensorfloat32,[n])).
We broadcast axi over a x twice in succession to compute fA→P (axi , a x j ) and fA→P (a x j , a x i ) for all 1 ≤ j ≤ N , yielding f_i_j and f_j_i, which are length-n sequences of vectors. We join and sum
each of these vectors elementwise to obtain the ultimate output of the block, which is also a length-n sequence of vectors. The overall weave module may now be implemented as follows:
weave = Composition() with weave.scope():
a_x = weave.input[0] p_x = weave.input[1] a_to_a = Map(f_a_a).reads(a_x) p_to_a = Map(Map(f_p_a) >> Sum()).reads(p_x) a_y = ZipWith(Concat() >> f_a).reads(a_to_a, p_to_a) a_to_p = ZipWith(a_i_to_p).reads(a_x, Broadcast().reads(a_x)) p_to_p = Map(Map(f_p_p)).reads(p_x) p_y = ZipWith(ZipWith(Concat() >> f_p)).reads(a_to_p, p_to_p) weave.output.reads(a_y, p_y)
The input to weave is (ax, px). It has the type
Tuple(Sequence(Tensorfloat32,[n]),Sequence(Sequence(Tensorfloat32,[m]))).
The calculation may be understood as follows:
• a_to_a maps over ax with fA→A, going from Sequence(Tensor) to Sequence(Tensor).
• p_to_a maps over px with fA→P and sums along the inner dimension, reducing from Sequence(Sequence(Tensor)) to Sequence(Tensor).
• a_y zips a_to_a and p_to_a with fA, going from Tuple(Sequence(Tensor),Sequence(Tensor)) to Sequence(Tensor).
• a_to_p broadcasts ax over itself with a_i_to_p, expanding from Sequence(Tensor) to Sequence(Sequence(Tensor)).
• p_to_p maps over px with fP→P , going from Sequence(Sequence(Tensor)) to Sequence(Sequence(Tensor)).
• p_y zips a_to_p and p_to_p with fP , going from Tuple(Sequence(Sequence(Tensor)),Sequence(Sequence(Tensor))) to Sequence(Sequence(Tensor)).
C CALCULATING LINES OF CODE
Our protocol for calculating lines6 of code is as follows:
• Define the functional unit of comparison as an input-output mapping.
• Prepare a single file that implements this functionality and nothing else.
• Remove import statements, abstract base classes, logging, file i/o, and validation logic.
• Count lines of code, ignoring blank lines and comments.7.
FEED-FORWARD ATTENTION
The functional unit of comparison is creating the model for the variable-length experiment described in Raffel & Ellis (2016, sec. 2.3). This includes the loss and accuracy calculations, but does not include the training loop or the creation of training data. The original implementation8 is in Python and uses Theano and Lasagne. The TensorFlow Fold implementation is more concise, partly due to differences between TensorFlow and Lasagne. Fold itself reduces implementation complexity by eliminating the need for manual batching, e.g. x.sum(axis=1) where batching is explicit over axis 0, vs. x >> Sum(), which is implicitly batched.
6All of the implementations we examine are formatted with 80-column lines excepting the Tree-LSTM implementation, which has a few lines that are slightly longer; we still count these as single lines.
7The calculations were performed with cloc (https://github.com/AlDanial/cloc). 8Commit e8fce3e from https://github.com/craffel/ff-attention.
TREE-LSTM
The functional unit of comparison is creating a (binary) constituency Tree-LSTM and running an epoch of training for the fine-grained sentiment classification task as described in Tai et al. (2015, sec. 5.1). This does not include loading the word embeddings or dataset, which are provided as inputs. The original implementation9 is in Lua and uses Torch. Lua terminates blocks with the end keyword; we do not count these lines. Here, the use of Python and TensorFlow leads to substantially more concise code than with Lua and Torch. Unlike the previous example manual batching plays no role here, because the original implementation computes gradients and losses one tree at a time. Fold reduces complexity here by using a OneOf block to distinguish between leaves and internal nodes, rather than a recursive function that explicitly traverses the tree.
GRAPH CONVOLUTION
The functional unit of comparison is creating a single weave module as described in Kearnes et al. (2016, sec. 3.3). The original implementation10 is in Python and uses TensorFlow. Here, both implementations use the same language and deep learning library. Fold helps by eliminating the need for manual batching, as in the first example. This is particularly apparent in the atoms-to-pairs calculation, which requires making n “copies” of an n × d matrix x to get an n × n × d tensor. In native TensorFlow the first dimension is batch, and the copying is explicit, as reshape(tile(x, [1, n, 1]), [batch_size, n, n, d]). In Fold, x >> Broadcast() suffices, because the number of copies needed is determined lazily by subsequent computations.
9Commit b02ad49 from https://github.com/stanfordnlp/treelstm. 10Provided by Kearnes et al. (2016).

The paper presents a novel strategy to deal with dynamic computation graphs. They arise, when the computation is dynamically influenced by the input data, such as in LSTMs. The authors propose an `unrolling' strategy over the operations done at every step, which allows a new kind of batching of inputs.
The presented idea is novel and the results clearly indicate the potential of the approach. For the sake of clarity of the presentation I would drop parts of Section 3 ("A combinator library for neural networks") which presents technical details that are in general interesting, but do not help the understanding of the core idea of the paper. The presented experimental results on the "Stanford Sentiment Treebank" are in my opinion not supporting the claim of the paper, which is towards speed, than a little bit confusing. It is important to point out that even though the presented ensemble "[...] variant sets a new state-of-the-art on both subtasks" [p. 8], this is not due to the framework, not even due to the model (comp. lines 4 and 2 of Tab. 2), but probably, and this can only be speculated about, due to the ensemble averaging. I would appreciate a clearer argumentation in this respect.
Update on Jan. 17th:
after the authors update for their newest revision, I increase my rating to 8 due to the again improved, now very clear argumentation.


We describe a framework for multitask deep reinforcement learning guided by policy sketches. Sketches annotate each task with a sequence of named subtasks, providing high-level structural relationships among tasks, but not providing the detailed guidance required by previous work on learning policy abstractions for RL (e.g. intermediate rewards, subtask completion signals, or intrinsic motivations). Our approach associates every subtask with its own modular subpolicy, and jointly optimizes over full task-specific policies by tying parameters across shared subpolicies. This optimization is accomplished via a simple decoupled actor–critic training objective that facilitates learning common behaviors from dissimilar reward functions. We evaluate the effectiveness of our approach on a maze navigation game and a 2-D Minecraft-inspired crafting game. Both games feature extremely sparse rewards that can be obtained only after completing a number of high-level subgoals (e.g. escaping from a sequence of locked rooms or collecting and combining various ingredients in the proper order). Experiments illustrate two main advantages of our approach. First, we outperform standard baselines that learn task-specific or shared monolithic policies. Second, our method naturally induces a library of primitive behaviors that can be recombined to rapidly acquire policies for new tasks.
1 INTRODUCTION
This paper describes a framework for learning composable deep subpolicies in a multitask setting, guided only by abstract policy sketches. We are interested in problems like the ones shown in Figure 1, with collections of tasks that involve sparse rewards and long-term planning, but which share structure in the form of common subgoals or reusable high-level actions. Our work aims to develop models that can learn efficiently from these sparse rewards and rapidly adapt to new tasks, by exploiting this shared structure and translating success on one task into progress on others. Our approach ultimately induces a library of high-level actions directly from symbolic annotations like the ones marked K1 and K2 in the figure.
This approach builds on a significant body of research in reinforcement learning that focuses on hierarchical representations of behavior. In these approaches, a high-level controller learns a policy over high-level actions—known variously as options (Sutton et al., 1999), skills
(Konidaris & Barto, 2007), or primitives (Hauser et al., 2008)—which are themselves implemented as policies over low-level actions in the environment. While one line of research (e.g. Daniel et al. (2012)) investigates learning hierarchical policies without any supervision, such hierarchies are empirically difficult to learn directly from unconstrained interaction (Hengst, 2002). The bulk of existing work instead relies on additional information (in the form of intermediate rewards, subtask completion signals, or intrinsic motivations) that guide the learner toward useful high-level actions. While effective, these approaches depend on state representations simple or structured enough that suitable reward signals can be effectively engineered by hand.
Here we focus on multitask learning of hierarchical policies from a weaker form of supervision: at training time, each task (τ1 and τ2 in Figure 1) is annotated with a sketch (K1 andK2) consisting of a sequence of high-level action symbols (b1, b2 and b3)—with no information about how these actions should be implemented. Our approach associates each such high-level action with its own lowlevel subpolicy, and jointly optimizes over concatenated task-specific policies by tying parameters across shared subpolicies. Our thesis is that even the minimal information about high-level policy structure contained in a sketch provides enough of a learning signal to induce general, reusable subpolicies. Crucially, sketches are totally ungrounded in the representation of the world—they require no intervention in a simulator or environment model.
The present work may be viewed as an extension of recent approaches for learning compositional deep architectures from structured program descriptors (Andreas et al., 2016; Reed & de Freitas, 2015). Here we focus on learning in interactive environments with reinforcement training signals. This extension presents a variety of technical challenges. Concretely, our contributions are:
• A general paradigm for multitask, hierarchical, deep reinforcement learning guided by abstract sketches of task-specific policies.
• A concrete agent architecture for learning in this paradigm, featuring a modular model structure and multitask actor–critic training objective.
We evaluate our approach on two families of tasks: a maze navigation game (Figure 3a), in which the agent must navigate through a sequence of locked doors to reach a target room; and a 2-D Minecraft-inspired crafting game (Figure 3b), in which the agent must acquire particular resources by finding raw ingredients, combining them together in the proper order, and in some cases building intermediate tools that enable the agent to alter the environment itself. In both games, the agent receives a reward only after the final goal is accomplished. For the most challenging tasks, involving sequences of four or five high-level actions, a task-specific agent initially following a random policy essentially never discovers the reward signal.
We evaluate a modular agent architecture trained with guidance from policy sketches under several different data conditions: (1) when learning the full collection of tasks jointly via reinforcement, (2) in a zero-shot setting where a policy sketch is available for the held-out task, and (3) in a adaptation setting, where sketches are hidden and the agent must learn a policy over high-level actions. In all cases, our approach substantially outperforms standard policy optimization baselines.
2 RELATED WORK
The agent representation we describe in this paper belongs to the broader family of hierarchical reinforcement learners described in the literature. As detailed in Section 3, our subpolicies may be viewed as a relaxation of the options framework first described by Sutton et al. (1999). A large body of work describes techniques for learning options and related abstract actions, in both single- and multitask settings. For learning the implementation of options, most techniques rely on intermediate supervisory signals, e.g. to encourage exploration (Kearns & Singh, 2002) or completion of predefined subtasks (Kulkarni et al., 2016). An alternative family of approaches employs either posthoc analysis of already-learned policies to extract reusable sub-components (Stolle & Precup, 2002; Konidaris et al., 2011). Techniques for learning options with less guidance than the present work include Bacon & Precup (2015) and Vezhnevets et al. (2016), and other general hierarchical policy learners include Daniel et al. (2012), Bakker & Schmidhuber (2004) and Menache et al. (2002).
Once a library of high-level actions exists, agents are faced with the problem of learning high-level (typically semi-Markov) policies that invoke appropriate high-level actions in sequence (Precup,
2000). The learning problem we describe in this paper is in some sense the direct dual to the problem of learning these high-level policies. There, the agent begins with an inventory of complex primitives and must learn to model their behavior and select among them; here we begin knowing the names of appropriate high-level actions but nothing about how they are implemented, and must infer implementations (but not, initially, high-level plans) from context. We expect that our approach could be coupled with a generic learner of options policies to provide a general mechanism for hierarchical RL; we leave this for future work.
Our approach is also inspired by a number of recent efforts toward compositional reasoning and interaction with structured deep models. Such models have been previously used for tasks involving question answering (Iyyer et al., 2014; Andreas et al., 2016) and relational reasoning (Socher et al., 2012), and more recently for multi-task, multi-robot transfer problems (Devin et al., 2016). In this work—as in existing approaches employing dynamically assembled modular networks—taskspecific training signals are propagated through a collection of composed discrete structures with tied weights. Here the composed structures specify time-varying policies rather than feedforward computations, and their parameters must be learned via interaction rather than direct supervision. Another closely related family of models includes neural programmers (Neelakantan et al., 2015) and programmer–interpreters (Reed & de Freitas, 2015), which generate discrete computational structures but require supervision in the form of output actions or full execution traces.
A closely related line of work is the Hierarchical Abstract Machines (HAM) framework introduced by Parr & Russell (1998). Like our approach, HAMs begin with a representation of a high-level policy as an automaton (or a more general computer program; Andre & Russell, 2001) and use reinforcement learning to fill in low-level details. Variations on this architecture have considered a number of control constructs beyond the scope of the current paper (e.g. concurrency and recursion; Marthi et al., 2004). However, because these approaches attempt to learn a single representation of the Q function for all subtasks and contexts, they require extremely strong formal assumptions about the form of the reward function and state representation (Andre & Russell, 2002) that the present work avoids by decoupling the policy representation from the value function.
Our approach also bears some resemblance to the instruction following literature in natural language processing. Existing work on instruction following falls into two broad categories: approaches that require a highly structured (typically logical) action and world representations (Chen & Mooney, 2011; Artzi & Zettlemoyer, 2013; Andreas & Klein, 2015; Tellex et al., 2011), and approaches that require detailed supervision of action sequences or dense reward signals essentially equivalent to full action traces (Branavan et al., 2009; Vogel & Jurafsky, 2010; Mei et al., 2016). By contrast, the framework we describe here involves no formal or logical language for describing plans, and no supervised action sequences. Additionally, the modular model described in this paper natrually supports adaptation to tasks where no sketches are available, while all existing instruction following models learn a joint policy over instructions and actions, and are unable to function in the absence of instructions.
3 LEARNING MODULAR POLICIES
We consider a multitask reinforcement learning problem arising from a family of infinite-horizon discounted Markov decision processes in a shared environment. This environment is specified by a tuple (S,A, P, γ), with S a set of states, A a set of low-level actions, P : S × A × S → R a transition probability distribution, and γ a discount factor. Each task τ ∈ T is then specified by a pair (Rτ , ρτ ), with Rτ : S → R a task-specific reward function and ρτ : S → R an initial distribution over states. For a fixed sequence {(si, ai)} of states and actions obtained from a rollout of a given policy, we will denote the empirical return starting in state si as qi := ∑∞ j=i γ
jR(sj). In addition to the components of a standard multitask RL problem, we assume that tasks are annotated with sketches Kτ , each consisting of a sequence (bτ1, bτ2, . . .) of high-level symbolic labels drawn from a fixed vocabulary B. Our model associates each of these symbols with a randomly initialized modular subpolicy. By sharing each subpolicy across all tasks annotated with the corresponding symbol, our approach naturally learns the shared abstraction for the corresponding subtask, without requiring any information about the grounding of that task to be explicitly specified by annotation.
3.1 MODEL
We exploit the structural information provided by sketches by constructing for each symbol b a corresponding subpolicy πb. At each timestep, a subpolicy may select either a low-level action a ∈ A or a special STOP action. We denote the augmented state space A+ := A ∪ {STOP}. While this framework is agnostic to the implementation of subpolicies, we are especially interested in the case where subpolicies are specified by deep networks. As shown in Figure 2, the experiments in this paper represent each πb as a neural network whose input is a representation of the current state, and whose output is a distribution over A+. While all action spaces in our experiments are discrete, it is straightforward to instead allow this last layer to parameterize a mixed distribution over an underlying continuous action space and the STOP action. These subpolicies may be viewed as options of the kind described by Sutton et al. (1999), with the key distinction that they have no initiation semantics, but are instead invokable everywhere, and have no explicit representation as a function from an initial state to a distribution over final states (instead implicitly using the STOP action to terminate).
Given a sketch, a task-specific policy Πτ is formed by concatenating its associated subpolicies in sequence. In particular, the high-level policy maintains a subpolicy index i (initially 0), and executes actions from πbi until the STOP symbol is emitted, at which point control is passed to πbi+1 . We may thus think of Πτ as inducing a Markov chain over the state space S × B, with transitions given by:
(s, bi) → (s′, bi) with probability ∑ a∈A πbi(a|s) · P (s′|s, a)
→ (s, bi+1) with probability πbi(STOP|s) Note that Πτ is semi-Markov with respect to projection of the augmented state space S ×B onto the underlying state space S. We denote the complete family of task-specific policies Π := ⋃ τ{Πτ}, and let each πb be an arbitrary function of the current environment state parameterized by some weight vector θb. The learning problem is to optimize over all θb to maximize the sum of expected discounted rewards J(Π) := ∑ τ J(Πτ ) := ∑ τ Esi∼Πτ [∑ i γ iRτ (si) ] across all tasks τ ∈ T .
3.2 POLICY OPTIMIZATION
Here that optimization is accomplished via a simple decoupled actor–critic method. In a standard policy gradient approach, with a single policy π with parameters θ, we compute gradient steps of the form (Williams, 1992):
∇θJ(π) = ∑ i ( ∇θ log π(ai|si) )( qi − c(s) ) , (1)
where the baseline or “critic” c can be chosen independently of the future without introducing bias into the gradient. Recalling our previous definition of qi as the empirical return starting from si, this form of the gradient corresponds to a generalized advantage estimator (Schulman et al., 2015) with λ = 1. Here c achieves close to the optimal variance (Greensmith et al., 2004) when it is set exactly equal to the state-value function Vπ(si) = Eπqi for the target policy π starting in state si.
The situation becomes slightly more complicated when generalizing to modular policies built by sequencing subpolicies. In this case, we will have one subpolicy per symbol but one critic per task. This is because subpolicies πb might participate in a number of composed policies Πτ , each associated with its own reward function Rτ . Thus individual subpolicies are not uniquely identified with value functions, and the aforementioned subpolicy-specific state-value estimator is no longer well-defined. We extend the actor–critic method to incorporate the decoupling of policies from value functions by allowing the critic to vary per-sample (that is, per-task-and-timestep) depending on the reward function with which the sample is associated. Noting that∇θbJ(Π) = ∑ t:b∈Kτ ∇θbJ(Πτ ), i.e. the expected reward across all tasks in which πb participates, we have:
∇θJ(Π) = ∑ τ ∇θJ(Πτ ) = ∑ τ ∑ i ( ∇θb log πb(aτi|sτi) )( qi − cτ (sτi) ) , (2)
where each state-action pair (sτi, aτi) was selected by the subpolicy πb in the context of the task τ .
Now minimization of the gradient variance requires that each cτ actually depend on the task identity. (This follows immediately by applying the corresponding argument in Greensmith et al. (2004) individually to each term in the sum over τ in Equation 2.) Because the value function is itself unknown, an approximation must be estimated from data. Here we allow these cτ to be implemented with an arbitrary function approximator parameterized by a vector ητ . This is trained to minimize a squared error criterion, with gradients given by
∇ητ [ 1
2 ∑ i (qi − cτ (si))2 ] = ∑ i ( ∇ητ cτ (si) )( qi − cτ (si) ) . (3)
Alternative forms of the advantage estimator (e.g. the TD residual Rτ (si) + Vτ (si+1) − γVτ (si) or any other member of the GAE family) can be easily substituted by simply maintaining one such estimator per task. Experiments (Section 4.3) show that conditioning on both the state and the task identity results in noticeable performance improvements, suggesting that the variance reduction provided by this objective is important for efficient joint learning of modular policies.
Algorithm 1 DO-STEP(Π, curriculum) 1: D ← ∅ 2: while |D| < D do 3: τ ∼ curriculum(·) . sample task τ from curriculum (Section 3.3) 4: d = {(si, ai, bi = Kτ,i, qi, τ), . . .} ∼ Πτ . do rollout 5: D ← D ∪ d
6: for b ∈ B, τ ∈ T do 7: d = {(si, ai, b′, qi, τ ′) ∈ D : b′ = b, τ ′ = τ} 8: θb ← θb − αD ∑ d ( ∇ log πb(ai|si) )( qi − cτ (si) ) . update policy
9: ητ ← ητ − βD ∑ d ( ∇cτ (si) )( qi − cτ (si) ) . update critic
The complete procedure for computing a single gradient step is given in Algorithm 1. (The outer training loop over these steps, which is driven by a curriculum learning procedure, is described in the following section and specified in Algorithm 2.) This is an on-policy algorithm. In each step, the agent samples tasks from a task distribution provided by a curriculum (described in the following subsection). The current family of policies Π is used to perform rollouts in each sampled task, accumulating the resulting tuples of (states, low-level actions, high-level symbols, rewards, and task identities) into a dataset D. Once D reaches a maximum size D, it is used to compute gradients w.r.t. both policy and critic parameters, and the parameter vectors are updated accordingly. The step sizes α and β in Algorithm 1 can be chosen adaptively using any first-order method.
3.3 CURRICULUM LEARNING
For complex tasks, like the one depicted in Figure 3b, it is difficult for the agent to discover any states with positive reward until many subpolicy behaviors have already been learned. It is thus a better use of the learner’s time to focus on “easy” tasks, where many rollouts will result in high reward from which appropriate subpolicy behavior can be inferred. But there is a fundamental tradeoff involved here: if the learner spends too much time on easy tasks before being made aware of the existence of harder ones, it may overfit and learn subpolicies that no longer generalize or exhibit the desired structural properties.
To avoid both of these problems, we use a curriculum learning scheme (Bengio et al., 2009) that allows the model to smoothly scale up from easy tasks to more difficult ones while avoiding overfitting. Initially the model is presented with tasks associated with short sketches. Once average reward on all these tasks reaches a certain threshold, the length limit is incremented. We assume that rewards across tasks are normalized with maximum achievable reward 0 < qi < 1. Let Êrτ denote the empirical estimate of the expected reward for the current policy on task t. Then at each timestep, tasks are sampled in proportion to 1 − Êrτ , which by assumption must be positive. Experiments show that both components of this curriculum learning scheme improve the rate at which the model converges to a good policy (Section 4.3).
The complete curriculum-based training procedure is specified in Algorithm 2. Initially, the maximum sketch length `max is set to one, and the curriculum initialized to sample length-1 tasks uniformly. (Neither of the environments we consider in this paper feature any length-1 tasks; in this case, observe that Algorithm 2 will simply advance to length-2 tasks without any parameter updates.) For each setting of `max, the algorithm uses the current collection of task policies Π to compute and apply the gradient step described in Algorithm 1. The rollouts obtained from the call to DO-STEP can also be used to compute reward estimates Êrτ ; these estimates determine a new task distribution for the curriculum. The inner loop is repeated until the reward threshold rmin is exceeded, at which point `max is incremented and the process repeated over a (now-expanded) collection of tasks.
4 EXPERIMENTS
As described in the introduction, we evaluate the performance of our approach in two environments: a maze navigation game and a crafting game. Both games involve nontrivial low-level control: agents must learn to avoid obstacles and interact with various kinds of objects. But the environments also feature hierarchical structure: rewards are accessible only after the agent has completed two to five high-level actions in the appropriate sequence.
In all our experiments, we implement each subpolicy as a multilayer perceptron with ReLU nonlinearities and a hidden layer with 128 hidden units, and each critic as a linear function of the current state. Each subpolicy network receives as input a set of features describing the current state of the environment, and outputs a distribution over actions. The agent acts at every timestep by sampling from this distribution. The gradient steps given in lines 8 and 9 of Algorithm 1 are implemented using RMSPROP (Tieleman, 2012) with a step size of 0.001 and gradient clipping to a unit norm. We take the batch size parameter D in Algorithm 1 to be 2000, and set γ = 0.9 in both environments. For curriculum learning, the improvement threshold rgood is set to 0.8.
4.1 ENVIRONMENTS
The maze environment (Figure 3a) corresponds closely to the the “light world” described by Konidaris & Barto (2007). The agent is placed in a discrete world consisting of a series of rooms, some of which are connected by doors. Some doors require that the agent first pick up a key to open them. For our experiments, each task corresponds to a goal room (always at the same position relative to the agent’s starting position) that the agent must reach by navigating through a sequence of intermediate rooms. The agent has one sensor on each side of its body, which reports the distance to keys, closed doors, and open doors in the corresponding direction. Sketches specify a particular sequence of directions for the agent to traverse between rooms to reach the goal. Mazes are sampled with random sizes and random decisions about whether to connect rooms with open doors, locked doors, or no doors. The sketch always corresponds to a viable traversal from the start to the goal position, but other (possibly shorter) traversals may also exist.
The crafting environment (Figure 3b) is inspired by the popular game Minecraft, but is implemented in a 2-D grid world. The agent may interact with some objects in the world by facing them
Algorithm 2 TRAIN-POLICIES() 1: Π = INIT() . initialize subpolicies randomly 2: `max ← 1 3: loop 4: rmin ←∞ 5: curriculum(·) = Unif(T ′) . initialize `max-step curriculum uniformly 6: T ′ = {τ ∈ T : |Kτ | ≤ `max} 7: while rmin < rgood do 8: DO-STEP(Π, curriculum) . update parameters (Algorithm 1) 9: Z = ∑ t∈T ′ [1− Êrτ ] 10: curriculum(t) = 1[τ ∈ T ′](1− Êrτ )/Z ∀τ ∈ T 11: rmin ← minτ Êrτ 12: `max ← `max + 1
and executing a special INTERACT action. Interacting with raw materials initially scattered around the environment causes them to be added to an inventory. Interacting with different crafting stations causes objects in the agent’s inventory to be combined or transformed into other objects. Each task in this game corresponds to some crafted object the agent must produce; the most complicated goals require the agent to also craft intermediate ingredients, and in some cases build tools (like a pickaxe and a bridge) to reach ingredients located in initially inaccessible regions of the environment.
A complete listing of tasks and sketches is given in Appendix A.
4.2 MULTITASK LEARNING
The primary experimental question in this paper is whether the extra structure provided by policy sketches alone is enough to enable fast learning of coupled policies across tasks. To evaluate this, we compare our modular approach to two policy gradient baselines—one that learns an independent policy for each task and one that learns a joint policy across all tasks—as well as a critic-only Q reader baseline. For the independent model, task-specific policies are represented by networks with the same structure as the modular subpolicies. The joint model conditions both on these environment features, as well as a feature vector encoding the complete sketch. The Q reader forms the same joint state and action space described in Section 3.1, and learns a single feedforward network to map from both environment states and representations of action symbols onto Q values. This baseline can be viewed either as a chain-structured hierarchical abstract machine with a learned state abstractor (Andre & Russell, 2002), or as a standard instruction following baseline from the natural language processing literature (Vogel & Jurafsky, 2010).
Learning curves for baselines and the modular model are shown in Figure 4. It can be seen that in both the maze domain and the crafting domain, our approach substantially outperforms the baselines: it induces policies with substantially higher average reward and converges more quickly than the policy gradient baselines. It can further be seen in Figure 4c that after policies have been learned on simple tasks, the model is able to rapidly adapt to more complex ones, even when the longer tasks involve high-level actions not required for any of the short tasks (Appendix A).
Having demonstrated the overall effectiveness of our approach, our remaining experiments explore (1) the importance of various components of the training procedure, and (2) the learned models’ ability to generalize or adapt to held-out tasks. For compactness, we restrict our consideration on the crafting domain, which features a larger and more diverse range of tasks and high-level actions.
4.3 ABLATIONS
In addition to the overall modular parameter-tying structure induced by our sketches, the key components of our training procedure are the decoupled critic and the curriculum. Our next experiments investigate the extent to which these are necessary for good performance.
To evaluate the the critic, we consider three ablations: (1) removing the dependence of the model on the environment state, in which case the baseline is a single scalar per task; (2) removing the dependence of the model on the task, in which case the baseline is a conventional generalized advantage estimator; and (3) removing both, in which case the baseline is a single scalar, as in a vanilla policy gradient approach. Results are shown in Figure 5a. Introducing both state and task dependence into the baseline leads to faster convergence of the model: the approach with a constant baseline achieves less than half the overall performance of the full critic after 3 million episodes. Introducing task and state dependence independently improve this performance; combining them gives the best result.
We also investigate two aspects of our curriculum learning scheme: starting with short examples and moving to long ones, and sampling tasks in inverse proportion to their accumulated reward. Experiments are shown in Figure 5b. We again see that both components are essential for good performance. Sampling uniformly across all tasks of the target length results in slow convergence.
4.4 ZERO-SHOT AND ADAPTATION LEARNING
In our final experiments, we consider the model’s ability to generalize to new tasks unseen at training time. We consider two evaluation conditions: a zero-shot setting, in which the model is provided a sketch for the new task and must immediately achieve good performance, and a adaptation setting, in which no sketch is provided and the model must learn the form of a suitable sketch by interacting with the new task.
Model MT 0-S Ad.
Independent .44 – <.1 Joint .49 <.1 – Modular .89 .77 .76
Table 1: Model performance under various evaluation conditions. MT is the multitask training condition described in Section 4.2, while 0-S and Ad. are respectively the zero-shot and adaptation experiments described in Section 4.4.
We hold out two length-four tasks from the full inventory used in Section 4.2, and train on the remaining tasks. For zeroshot experiments, we simply form the concatenated policy described by the sketches of the held-out tasks, and repeatedly execute this policy (without learning) in order to obtain an estimate of its effectiveness. For adaptation experiments, we consider ordinary reinforcement learning over B rather than A, implementing the high-level learner with the same agent architecture as described in Section 3.1. Note that the Independent baseline cannot be applied to the zero-shot evaluation, while the joint baseline cannot be applied to the adaptation baseline (because it depends on pre-specified sketch fea-
tures). Results are shown in Table 1. The held-out tasks are sufficiently challenging that the baselines are unable to obtain more than negligible reward, while the modular model does comparatively well.
5 CONCLUSIONS
We have described an approach for multitask learning of neural network policies guided by symbolic policy sketches. By associating each symbol appearing in a sketch with a modular neural subpolicy, we have shown that it is possible to build agents that share behavior across tasks in order to achieve success in tasks with sparse and delayed rewards. This process induces an inventory of reusable and interpretable subpolicies which can be employed for zero-shot generalization when further sketches are available, and hierarchical reinforcement learning when they are not. Our work suggests that these sketches, which are easy to produce and require no grounding in the environment, provide an effective scaffold for learning hierarchical policies from minimal supervision. We have released our code at http://github.com/jacobandreas/psketch.
ACKNOWLEDGMENTS
JA is supported by a Facebook Graduate Fellowship and a Huawei / Berkeley AI fellowship.
A TASKS AND SKETCHES
The complete list of tasks, sketches, and symbols is given below. Tasks marked with an asterisk∗ are held out for the generalization experiments described in Section 4.4, but included in the multitask training experiments in Sections 4.2 and 4.3.
Goal Sketch
Maze environment goal1 left left goal2 left down goal3 right down goal4 up left goal5 up right goal6 up right up goal7 down right up goal8 left left down goal9 right down down goal10 left up right
Crafting environment make plank get wood use toolshed make stick get wood use workbench make cloth get grass use factory make rope get grass use toolshed make bridge get iron get wood use factory make bed∗ get wood use toolshed get grass use workbench make axe∗ get wood use workbench get iron use toolshed make shears get wood use workbench get iron use workbench get gold get iron get wood use factory use bridge get gem get wood use workbench get iron use toolshed use axe

The paper presents an approach to learning shared neural representations of temporal abstractions in hierarchical RL, based on actor-critic methods. The approach is illustrated in two tasks: gridworld with objects and a simplified Minecraft problem).  The idea of providing symbolic descriptions of tasks and learning corresponding "implementations" is potentially interesting and the empirical results are promising.  However, there are two main drawbacks of the current incarnation of this work.  First, the ideas presented in the paper have all been explored in other work (symbolic specifications, actor-critic, shared representations).  While related work is discussed, it is not really clear what is new here, and what is the main contribution of this work besides providing a new implementation of existing ideas in the context of deep learning. The main contribution if the work needs to be clearly spelled out.  Secondly, the approach presented relies crucially on curriculum learning (this is quite clear from the experiments).  While the authors argue that specifying tasks in simplified language is easy, designing a curriculum may in fact be pretty complicated, depending on the task at hand.  The examples provided are fairly small, and there is no hint of how curriculum can be designed for larger problems. Because the approach is sensitive to the curriculum, this limits the potential utility of the work. It is also unclear if there is a way to provide supervision automatically, instead of doing it based on prior domain knowledge.
More minor comments:
- The experiments are not described in enough detail in the paper. It's great to provide github code, but one needs to explain in the paper why certain choices were made in the task setup (were these optimized? What's this the first thing that worked?) Even with the code, the experiments as described are not reproducible
- The description of the approach is pretty tangled with the specific algorithmic choices. Can the authors step back and think more generally of how this approach can be formalized?  I think this would help relate it to the prior work more clearly as well.