Learning Task Decomposition with Ordered Memory Policy Network
Authors: Yuchen Lu
Yikang Shen
Siyuan Zhou
Aaron Courville
Joshua B. Tenenbaum
Chuang Gan
Venue: ICLR 2021 Poster (2021)
Keywords: Task Segmentation
Hierarchical Imitation Learning
Network Inductive Bias

Abstract: Many complex real-world tasks are composed of several levels of subtasks. Humans leverage these hierarchical structures to accelerate the learning process and achieve better generalization. In this work, we study the inductive bias and propose Ordered Memory Policy Network (OMPN) to discover subtask hierarchy by learning from demonstration. The discovered subtask hierarchy could be used to perform task decomposition, recovering the subtask boundaries in an unstructured demonstration. Experiments on Craft and Dial demonstrate that our model can achieve higher task decomposition performance under both unsupervised and weakly supervised settings, comparing with strong baselines. OMPN can also be directly applied to partially observable environments and still achieve higher task decomposition performance. Our visualization further confirms that the subtask hierarchy can emerge in our model 1.

Published as a conference paper at ICLR 2021
LEARNING TASK DECOMPOSITION WITH ORDERED MEMORY POLICY NETWORK
Yuchen Lu & Yikang Shen University of Montreal, Mila Montreal, Canada
Siyuan Zhou Peking University Beijing, China
Aaron Courville
University of Montreal, Mila, CIFAR Montreal, Canada
Joshua B. Tenenbaum MIT BCS, CBMM, CSAIL Cambridge, United States
Chuang Gan
MIT-IBM Watson Al Lab Cambridge, United States
ABSTRACT
Many complex real-world tasks are composed of several levels of sub-tasks. Hu- mans leverage these hierarchical structures to accelerate the learning process and achieve better generalization. In this work, we study the inductive bias and pro- pose Ordered Memory Policy Network (OMPN) to discover subtask hierarchy by learning from demonstration. The discovered subtask hierarchy could be used to perform task decomposition, recovering the subtask boundaries in an unstruc- tured demonstration. Experiments on Craft and Dial demonstrate that our model can achieve higher task decomposition performance under both unsupervised and weakly supervised settings, comparing with strong baselines. OMPN can also be directly applied to partially observable environments and still achieve higher task decomposition performance. Our visualization further confirms that the subtask hierarchy can emerge in our model !
INTRODUCTION
Learning from Demonstration (LfD) is a popular paradigm for policy learning and has served as a warm-up stage in many successful reinforcement learning applications (Vinyals et al., 2019; Silver et al., 2016). However, beyond simply imitating the experts' behaviors, an intelligent agent's 's crucial capability is to decompose an expert's behavior into set of useful skills and discover sub-tasks. The discovered structure from expert demonstrations could be leveraged to re-use previously learned skills in the face of new environments (Sutton al., 1999; Gupta et al., 2019; Andreas al., 2017). Since manually labeling sub-task boundaries for each demonstration video is extremely expensive and difficult scale up, it is essential to learn task decomposition unsupervisedly, where the only supervision signal comes from the demonstration itself.
This question of discovering a meaningful segmentation of the demonstration trajectory is the key focus of Hierarchical Imitation Learning (Kipf et al., 2019; Shiarlis et al., 2018; Fox et al., 2017; Achiam et al., 2018) These works can be summarized finding the optimal behavior hierarchy so that the behavior can be better predicted (Solway et al., 2014). They usually model the sub-task structure as latent variables, and the subtask identifications are extracted from learnt posterior. In this paper, we propose novel perspective to solve this challenge: could we design a smarter neural network architecture, so that the sub-task structure can emerge during imitation learning? To be specific, we want design recurrent policy network such that examining the memory trace at each time step could reveal the underlying subtask structure.
Drawing inspiration from the Hierarchical Abstract Machine (Parr & Russell, 1998), we propose that each subtask can be considered as a finite state machine. A hicrarchy of sub-tasks can be represented different slots inside the memory bank. At each time step, subtask can be internally updated with the new information, call the next-level subtask, or return the control to the previous level subtask. If our designed architecture maintains hierarchy of sub-tasks operating the described manner,
Project page: https:://rreeee memo ..g.thu...bb...... 1o/

---

Published as a conference paper at ICLR 2021
Maketindge
Maketindge
Makelingge
Masedraagg
Makeetindge
Mikweiiddge
GatMastaiial
Gethlatera oi
Gelliiterial
CalMistriial
Build
Su Libiddge
OedBass
Gediiood
GoFaactoy
GerWooo
GerWood
GetCiness
GetCiease
GoFlactory
(a)
(h)
Figure (a) A simple grid world with the task "make bridge", which can be decomposed into multi-level subtask structure. (b) The representation of subtask structure within the agent memory with horizontal update and vertical expansion at each time step. The black arrow indicates a copy operation. The expansion position is the memory slot where the vertical expansion starts and is marked blue.
then subtask identification can be as easy as monitoring when the low-level subtask returns control to the higher-level subtask, or when the high-level subtask expands to the new lower-level subtask.
We give an illustrative grid- world example in Figure 1. In this example, there different in- gredients like grass for the agent to pickup. There is also a factory where the agent can use the ingredients. Suppose the agent wants to complete the task of building bridge. This task can be decomposed into a tree-like, multi-level structure, where the root task is divided into GetMMaterial and BuildBridge. vide sketch on how this subtask structure should be represented inside the agent's memory during each time step. The memory would be divided into different levels, corresponding structure. When t = 1, the model just starts with the root task, Make Bridge, and vertically expands into GetMaterian which further vertically expands into GetW ood. The vertical expansion corre- sponds to planning or calling the next level subtasks. The action is produced from the lowest-level memory. The intermediate Gaerial is copied for < 3, but horizontally updated when GetW ood is finished. The horizontal update can be thought of as internal update for each aterial vertically expands into a different child GetGrass. The root task is always copied until aterial is finished at f 4. As result, M Bridge goes through one horizontal update at = and then expands into BuildBridge and GoF actory. We can identify the subtask boundaries from this representation by looking at the change of expansion position, which is defined to be the memory slot where vertical expansion happens. E.g., from / / 2 to /. = 3, the expansion position goes from the lowest level to the middle level, suggesting the com- pletion of the low-level subtask. From = 4 to t = 5. the expansion position goes from the lowest level to the highest level, suggesting the completion of both low-level and mid-level subtasks.
Ret.Materid can be further divided into GetGrass and GetW ood. We pro-
the subtask
3,
subtask, and the updated Getriaarrall
Driven by this intuition, we propose the Ordered Memory Policy Network (OMPN) to support the subtask hierarchy described in Figure 1. We propose use a bottom-up recurrence and a top- down recurrence to implement horizontal update and vertical expansion respectively. Our proposed memory- update rule further maintains a hierarchy among memories such that the higher-level mem- ory can store longer-term information. At each time step, the model would softly decide the ex- pansion position from which to perform vertical expansion based on a differentiable stick-breaking process, that model can be trained end-to-end. We demonstrate the efffcctiveness of our ap- proach with multi-task behavior cloning. We perform experiments on both grid-world as well as more challenging robotic tasks. We show that OMPN is able perform task decomposition in both an unsupervised and weakly supervised manner, comparing favorably with strong baselines. Mean- while. OMPN still maintains the similar, if not better, performance on behavior cloning in terms of sample complexity and returns. Our ablation study shows the contribution of each component in our architecture. Our visualization further confirms that the subtask hierarchy emerges in our model's expanding positions.
2 ORDERED MEMORY POLICY NETWORK
We describe our policy architecture given the intuition described above. Our model is a recurrent policy network p(a"|s", M²) where M E Ruxm is block of memory while each memory has
2

---

Published as a conference paper at ICLR 2021
dimension We M, to refer the ith slot of the memory, M [M1 M2, ... My). The highest-level memory is M, while the lowest-level memory is M1. Each memory can be thought of as the representation of a subtask. We use the superseript to denote the time step
At each time step, our model will first transform the observation s' E to x' E RT. This can be achieved by domain-specific observation encoder. Then have ordered-memory module M', O OM M( to generate the next memory and the output. The output Of sent into feed-forward neural net to generate the action distribution.
2.1 ORDERED MEMORY MODULE
Up
2-1
Butum Up
Mr
M,
.6
4
M/
.6
C
C:
Top Down
M;
R(C))
== 1,0,0
9,4.0
(a)
(b)
F-1
Butom
MakiBridge
MakeBBiddge
M6
Tap Down
M6
k(CS)
A(CL))
GetMatorial
BulldBridge
M.
A(H(O))
A(M())))
[0.0.3]
GetGrass
GoFactory
bed
L-6
(c)
(d)
Figure 2: Dataflow of how M will be updated in M for three memory slots when the expansion position is a (a) low, (b) middle, or (c) high position. Blue arrows and red arrows corresponding to the vertical expansions and horizontal updates. (d) is snapshot t = from the grid-world exam- ple Figure lb. The subtask-update behavior corresponds to the memory-update when the expansion position is at the high position.
The ordered memory module first goes through a bottom-up recurrence. This operation implements the horizontal update and updates each memory with the new observation. We define C to be the updated memory:
C. F(C 2" M(
for i 1, ..... where Co = x' and F is cell function. Different from our mental diagram, we make it an recurrent process since the high-level memory might be able to get information from the updated lower-level memory in addition the observation. In our experiment recurrence will help model perform better than in task decomposition. For each memory, we also generate a score f from 0 with ft = G(z², M.) for i I...,, n. The score ff can be interpreted as the probability that subtask i is completed at time f.
find that such
In order to properly generate the final expansion position, we would like to insert the inductive bias that the higher-level subtask is expanded only if the higher-level subtask is not completed while all the lower-level subtasks are completed, as is shown in Figure 1b. As a result we use stick-breaking process as follows:
(L - f MI 11
1<isn i-1
#
ff
Finally we have the expansion position प. / Σ म' as a properly normalized distribution over memories. We can also define the ending probability as the probability that every subtask is finished.
(1)
annd
3

---

Published as a conference paper at ICLR 2021
Then we use top-down recurrence on the memory to implement the vertical expansion. from M 0, have
Starting
M h(##+11 C++1 (1 ) +=)
where न E2n in the following way:
my. me Ejsi Ty. and h can be any cell function. Then we update the memory
M' =M= (1-
11
where the output read from the lowest-level memory O Mj. For better understanding purpose, we show in Figure 2 how M will be updated into M with n = 3, when the expansion position high, middle and low position respectively. The memory higher than the expansion position will be preserved, while the memory at and lower than the expansion position will be over-written. We also take the snapshot of = from our the early example Figure 1b and show that the subtask-update behavior corresponds to our medory- when the expansion position is at the high position.
Although we show only the discrete case for illustration, the vector n' is actually continuous. As result, the whole process fully differentiable and can be trained end-to-end. More details can be found in the appendix A.
The memory depths n hyper-parane paper-parameter here. n is too small, we might not have enough capacity to cover the underlying structure in the data, If is too large, might impose some extra optimization difficulty. In our experiments we investigate the effect of using different n.
2.2 UNSUPERVISED TASK DECOMPOSITION WITH BEHAVIOR CLONING
We assume the following setting. We firstly have training phase to perform behavior cloning an unstructured demonstration dataset with state-action action pairs. Then during the detection phase, the user would specify the number of subtasks K for the model to produce the task boundaries.
We firstly describe our training phase. We develop a unique regularization technique which can help model learning the underlying hierarchy structure. Suppose have action space A. We first augment this action space with A'' = A U (done), where done is a special action. Then we can modify the action distribution accordingly:
p'(a'ls)) = [p(a'|s')(1 (1-mTd)) a' E (T
a' done
Then for each demonstration trajectory T = {s', a a')T=1, we transformed into 7' (sT+1 -T+11 done). which is essentially telling the model to output done only after the end of the trajectory. This process can be achieved on both discrete and continuous action space without heavy human involvement described in Appendix A. Then we will maximize Σ log l' (a'ls" on T'. We find that including and into the loss is crucial prevent our model degenerating into only using the lowest-level memory, since it provides the signal to raise the expansion position at the end of the trajectory, benefiting the task decomposition performance. We also justify this in our ablation study.
Since the expansion position should be high the low-level subtasks are completed, we can achieve unsupervised task decomposition by monitoring the behavior of n². To be specific, we define Tang Ex ē as the expected expansion position. Given Taug, consider the following methods to recover the subtask boundaries.
Top-K In this method we choose the time steps of K largest Kang to detect the boundary, where K the desired number sub-tasks given by the user during the detection phase. We find that this method is suitable for the discrete action space, where there is a very clear boundary between subtasks.
Thresholding In this method we standardize the Tang into nang from 0 to I, and then compute a Boolean array 1 ((Tayy > thres), where thres from to 1. We retrieve the subtask boundaries from the ending time step of each True segments. We find this method is suitable for continuous control settings, where the subtask boundaries are more ambiguous and smoothed out across time steps. We also design algorithm automatically select the threshold in Appendix B.
4

---

Published as a conference paper at ICLR 2021
3 RELATED WORK
Our work is related to option discovery and bierarchical imitation learning. The existing option discovery works have focused on building probabilistic graphical model on the trajectory, with op- tions as latent variables. DDO (Fox et al., 2017) proposes an iterative EM-like algorithm to discover multiple level of options from the demonstration. DDO was later applied in the continuous action space (Krishnan et al., 2017) and program modelling (Fox et al., 2018). Recent works like com- pILE (Kipf et al., 2019) and VALOR (Achiam al., 2018) also extend this idea by incorporating more powerful inference methods like VAE (Kingma & Welling, 2013). Lee (2020) also explore unsupervise task decompostion via imitation, but their method is not fully end-to-end, requires an auxiliary self-supervision loss, and does not support multi-level structure. Our work focuses on the role of neural network inductive bias in discovering re-usable options or subtasks from demonstra- tion. We do not have an explicit "inference" stage in our training algorithm to infer the option/task ID from the observations. Instead, this inference "stage" is implicitly designed into our model ar- chitecture via the stick-breaking process and expansion position. Based on these considerations, we choose compILE as the representative baseline for this field of work.
Our work is also related to Hierarchical RL ( (Vezhnevets et al., 2017; Nachum et al., 2018; Bacon et al., 2017). These works usually propose an architecture that has a high-level controller to output goal, while the low-level architecture takes the goal and outputs the primitive actions. However, these works mainly deal with the control problem, and do not focus on learning task decomposition from the demonstration. Recent works (Gupta et al., 2019; Lynch et al., 2020) also include hierarchi- cal imitation learning stage as a way to pretraining low-level policies before apply Hierarchical RL for finetuning, however they do not produce the task boundaries and therefore are not comparable to our works. Moreover, their hierarchical IL algorithm exploits the fact that the goal can be described as a point in the state space, so that they are able to re-label the ending state of an unstructured demonstrations as fake goal to train the coul-conditioned policy. Mcanwhile our approach is de- signed to be general. In addition to the option framework, our work is closely related to Hierarchical Abstract Machine (HAM) (EI Hihi & Bengio, 1996). Our concept of subtask is similar to the finite state machine (FSM). The horizontal update corresponds to the internal update of the FSM, while the vertical expansion corresponds to calling the next level of the FSM. Our stick-breaking process is also a continuous realization of the idea that low-level FSM transfers control back to high-level FSM at completion.
Recent work (Andreas et al., 2017) introduces the modular policy networks for reinforcement learn- ing so that it can be used decompose complex task into several simple subtasks. In this setting, the agent is provided a sequence of subtasks, called sketch, at the beginning. Shiarlis et al. (2018) propose TACO to jointly learn sketch alignment with action sequence, as well as imitating the tra- jectory. This work can only be applied in the "weakly supervised" setting, where they have some information like the sub-task sequence. Nevertheless, we also choose TACO (Shiarlis et al., 2018) as one of our baselines.
Incorporating varying time-scale for each neuron to capture hierarchy not a new idea (Chung et al., 2016; El Hihi & Bengio, 1996; Koutnik et al., 2014). However, these works do not focus on recovering the structure after training, which makes these methods less interpretable. Shen et al. (2018) introduce Ordered Neurons and show that they can induce syntactic structure by examin- ing the hidden states after language modelling. However ONLSTM does not provide mechanism to achieve the top-down and bottom-up recurrence. Our model is mainly inspired by the Ordered Mem- ory (Shen et al., 2019). However, unlike previous work our model is decoder expanding from root task to subtasks, while the Ordered Memory is an encoder composing constituents into sentences. Recently Mittal et al. (2020) propose to combine top-down and bottom-up process. However their main motivation is handle uncertainty in the sequential prediction and they do not maintain hierarchy of memories with different update frequencies.
4 EXPERIMENT
We would like evaluate whether OMPN is able to jointly learning task decomposition during behavior cloning. We would like to answer the following questions in our experiments.
Q1: Can OMPN be applied in both continuous and discrete action space?
5

---

Published as a conference paper at ICLR 2021
Q2: Can OMPN be applied in both unsupervised, as well as, weakly supervised setting for task decomposition?
Q3: How much does each component helps the task decomposition?
Q4: How does task decomposition performance change with different hyper-parameters e.g., memory dimension and memory depths n??
4.1 SETUP AND METRICS
For the discrete action space, we use a grid world enviromment called Craft adapted from Andreas et al. (2017)'. At the beginning of each episode, an agent is equipped with task along with the sketch, e.g. makecloth = (getgrass, go actory). The original environment is fully observable. To further test our model, we make it also support partial observation by providing a self-centric window. For the continuous action space, we have robotic setting called Dial (Shiarlis al., 2018) where a JACO 6DoF manipulator interact with a large number pad". For each episode, the sketch is a sequence of numbers be pressed. More details on the demonstration can be found in Appendix E.
We experiment in both unsupervised and weakly supervised settings. For the unsupervised setting, did not provide any task information. For the weakly supervised setting, we provide the subtask sequence, or sketch, in addition to the observation. We use nosketch and sketch to denote these two settings respectively. We choose compILE to be our unsupervised bascline while TACO to be our weakly supervised baseline.
We use the ground -tuth K to get the best performance of both our models and the baselines. This setting is consistent with the previous literature (Kipf et al., 2019). The details about task decompo- sition metric can be found in the appendix C.
4.2 TASK DECOMPOSITION RESULTS
Craft
Dial
Full
Partial
Align Ace
F1(tol=1)
Align Ace
FI(tol=1)
Align Acc
NoSketch
OMPN
93(1.7)
95(1.2)
84(6)
89(4.6)
87(4.0)
compILE
86(1.4)
97(0.8)
54(1.4)
57(4.8)
45(5.2)
OMPN
97(1.2)
98(0.9)
72(6.2)
83(8.1)
82(7.4)
Sketch
TACO
90(3.6)
66(2.2)
98(0.1)
Table 1: Alignment accuracy and F1 scores with tolerance on Craft and Dial. The results averaged over five runs, and the number in the parenthesis is the standard deviation.
In this section, we mainly address the question Q1 and Q2. Our main results for task decomposition in Craft and Dial are in Table 1. In Craft, we use the TopK detection algorithm. Our results show that OMPN is able to outperform baselines in both unsupervised and weakly-supervised settings with higher Fl scores and alignment accuracy. In general, there is a decrease of the performance for models when moving from full observations to partial observations. Nevertheless, compared with the baselines, we find that OMPN suffers less performance drop than the baselines. The F1 score results with different K other than the ground truth is in Table and Table 6.
In Dial, we use the automatic threshold method described in Appendix B. Our results is able to outperform compILE for the unsupervised setting, but not better than TACO when the sketch information is given. In general, our current model does not benefit from the additional sketch information, and we hypothesize that it is because we use a very trivial way of incorporating sketch information by simple concatenating it with the observation. A more effective way would be using attention over the sketch sequence to properly feed the related task information into the model. The Fl score results with different バ other than the ground truth is in Table 8.
'nttps:/ http://hhh...cmococoooaoaaraae// hhttps://gf http:////....tbcoooo/o//o/aoaaaa http://ww...aaaieeeeeeeeeeee/eaa

---

Published as a conference paper at ICLR 2021
*
Factory
Iron
Way Grass
slot2
Toolshed
slot1
Agent
Wood
Workbench
Figure 3: Visualization of the learnt expanding positions Craft. We present and the action sequence. The four subtasks are GetW ood, GoT oolshded, GetGrass and GoW orkbench. In the action sequence. "u" is either picking up/using the object. Each ground truth subtask is highlighted with an arrow of different colors.
4.3 QUALITATIVE ANALYSIS
In this section, we provide some visualization of the learnt expanding positions to qualitatively evaluate whether OMPN can operate with subtask hierarchy.
Targ
fnal
UESONT
lower
3
Init
Init
332
250
time step
Figure 4: Visualization of the learnt expanding positions of Dial domain. The task is to press the number sequence [9, 8, 6,3). We plot the marg- Our threshold selection algorithm produce upper bound and and lower bound, and the final threshold is computed as the average. The task boundary detected as the last time step of a segment above the final threshold. The frames at the detected task boundary show that the robot just finishes each subtask.
We firstly visualize the behavior of expanding positions for Craft in Figure 3. We find that at the end of each subtask, the model learns switch higher expansion position when facing the target object/place, while within each subtask, the model learns to maintain the lower expansion position. This is our desired behavior described Figure 1. What is interesting is that although the ground truth hierarchy structure might be only two-levels, given some redundant depths. In this example, the model combines the first two subtasks into one group and the last two subtasks into another group. More results can be found in Appendix F.
model is able to learn multiple level hierarchy
In Figure 4, we show the qualitative result in Dial. We find that, instead of having a sharp sub- task boundary, the boundary between subtasks is more ambiguous with high expanding positions

---

Published as a conference paper at ICLR 2021
across multiple time steps, and this motivates design of the threshold algorithm. This happens because we use the observation to determine the expanding position. In Dial, the state difference less significant due small time skip and continuous space, the expanding position the boundaries might be high for multiple time steps. Also just like Craft, the number of subtasks naturally emerge from the peak patterns. We also provide addition qualitative result Kitchen released by Gupta et al. (2019) in Figure 5. More results on Kitchen can be found in Appendix I.
aeg
A
B
Init
205
time step
Figure 5: Visualization of the learnt expanding positions of Kitchen. The subtasks are Microwave, Bottom Knob, Hinge Cabinet and Slider. We show the upper bound, lower bound and final threshold produced by our detection algorithm.
4.4
ABLATION STUDY
Align, Acc.
Align, Acc.
Align. Acc.
1.0
0.7
0.7
0.5
0.5
0.4
No Bottomop Recurr
Rottommu
0.4
No Bottomup Recur
0.5
Buttornup
Bottornup Recurr
No Done
No Dane
0.4
No Done
Topdown
CRIPN
0.2
0.3
0.0
0.2
044
066
0.8
1.0
0.0
0.5
1.0
1.5
2.0
o.n
0.4
0.6
0.8
1.0
framez
167
fromes
1w7
fromes
167
(a) NoSketch + Cralt(Full)
(b) NoSketch + Crafi(Partial)
(c) NoSketch Dial
Align. Acc.
Align. Acc.
Align. Acc.
3.0
0.9
0.8
0.9
0.7-
0.8
0.7
0.6
0.7
0.6
0.5
Bottomup
Bottomup
0.5
No Bottomup Recum
No Bothomup Recurr
0.3
Done
0.4
Topdown
Topdown
CMPN
0.2
DMPN
0.3
0.6
0.25
0.35
1.00
1.25
107
le5
frames
167
(d) Sketch + Cralt(Full)
(e) Sketch + Craft(Partial)
(f) Sketch + Dial
Figure 6: Ablation study for task alignment accuracy for all settings.
8

---

Published as a conference paper at ICLR 2021
In this section, aim to answer Q3 and Q4. We perform the ablation study well as some hyper-parameter analysis. The results are in Figure 6. We summarize our findings as below:
For No Bottomup and No Topdown, process from the model. We find that both hurt the alignment accuracy and removing bottom-up process hurts more. This is expected since bottom-up recurrence updates the subtasks with the latest observations before predicting the termination scores, while the top-down process is more related to outputting actions.
remove completely either the bottom-up or the top-bottom
For No Bottomup Recurr, we remove the recurrence in the bottom-up process by making C F(O, x², M(-1) so as to preserve the same number parameters as OMPN. Although this hurts the alignment accuracy least, the existence of the performance drop confirms our intuition that the outputs of the lower-level subtasks are beneficial for the higher-level subtasks to predict termination scores, resulting in better task decomposition.
For No Done, we use the OMPN architecture but remove the Tend from the loss. We find that the model is still able to learn the structure based on the inductive bias some degree, but the alignment accuracy is much worse.
We also perform hyper-parameter analysis and see its effect on the task decomposition results in Appendix H. We find that our task decomposition results is robust to the memory depths and memory size m most cases.
4.5 BEHAVIOR CLONING
We show the behavior cloning results in Figure 7 and the full results are in Figure 10. For Craft, the sketch information necessary to make the task solvable, so we don't see much difference when there is no sketch information. On the contrary, with full observation and sketch information (Figure 7a), this setting might be too easy to show the difference since memory-less MLP can also achieve almost perfect success rate. As result, only when the environment moves to more challenging but still solvable setting with partial observation (Figure 7b), OMPN outperform LSTM on the success rate.
For Dial, we find that behaviour cloning alone is not able to solve the task and all our models never generate the maximum return due to the exposure bias. This is consistent with the previous literature on applying behavior cloning to robotics tasks with continuous states (Laskey et al., 2017). More details be found in Figure 13.
Rate
Success Rate
0.9
0.5
LS
000%
c10-
0.4
0.2
0.2
SC3
0.0
CHPN
0.0
0M/P
3.5
3.5
2.5
Tomes
Frimes
(a) Craft Full
(b) Craft Partial
(c) Dial
Figure 7: The behavior cloning results when sketch information is provided. For Craft, we define success the completion of four subtasks. For Dial, the total maximum return is 4.
CONCLUSION
In this work, we investigate the problem learning the subtask hierarchy from the demonstration trajectory. We propose novel Ordered Memory Policy Network (OMPN) that can represent the subtask structure and leverage to perform unsupervised task decomposition. Our experiments show that OMPN learns to recover the subtask boundary in both unsupervised and weakly supervised settings with behavior cloning. In the future, we plan to develop novel control algorithm based on the inductive bias for faster adaptation compositional combinations of the subtasks.

---

Published as a conference paper at ICLR 2021
Acknowledgement This work is in part supported by the Center for Brain, Minds, and Machines (CBMM, funded by NSF STC award CCF-1231216), and IBM Re- search.
REFERENCES
Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. arXiv preprint arXiv:1807.1029: 2018. 5,5
Jacob Andreas, Dan Klein, and Sergey Levine. Modular multitask reinforcement learning with policy sketches. International Conference on Machine Learning, pp. 166-175, 2017. 1,5,
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Thirty-Firsi AAAI
Conference on Artificial Intelligence, 2017.
Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural net- works. arXiv preprint arXiv: 100966609000000 2016.
Salah EI Hihi and Yoshua Bengio. Hierarchical recurrent neural networks for long-term dependen- cies. In Advances neural information processing systems, pp. 493-499, 1996.
Roy Fox, Sanjay Krishman, lon Stoica, and Ken Goldberg. Multi-level discovery of deep options.
arXiv preprint /0tt::/703.06600
2017. 1,5
Roy Fox, Richard Shin, Sanjay Krishnan, Ken Goldberg, Dawn Song, and Ion Stoica. Parametrized hierarchical procedures for neural programming. /CLR 2018, 2018.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy tasks via imitation and reinforcement learning. arXiv preprint
learning: Solving
arXiv:1910 1956, 2019. 1, 5, 8,
Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114 2013. 5
arXiv preprint
Thomas Kipf, Yujia Li, Hanjun Dai, Vinieius Zambaldi, Alvaro Sanchez-Gonzalez, Edward Grefen- stette, Pushmeet Kohli, and Peter Battaglia. Compile: Compositional imitation learning and execution. In International Conference on Machine Learning. pp. 3418-3428. PMLR, 2019. 1, 5,
Jan Koutnik, Klaus Greff, Faustino Gomez, and Juergen Schmidhuber. A clockwork In Inter- national Conference on Machine Learning, pp. 1863-1871, 2014. 5
Sanjay Krishnan, Roy Fox, Ion Stoica, and Ken Goldberg. Ddco: Discovery of deep continuous options for robot learning from demonstrations, arXiv preprint 17000:17 2017.
Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. Dart: Noise injection for robust imitation learning. arXiv preprint /703.0/770.09955 2017.
Sang-Hyun Lee. Learning compound tasks without task-speeific knowledge via imitation and self- supervised learning. In International Conference on Machine Learning, 2020. 5
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on Robot Learning, pp. 1113- 1132. PMLR. 2020. 5
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning to combine top-down and bottom- up signals in recurrent neural networks with attention over modules. arXiv preprint arXiv:2006.166 16981, 2020.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. In Advances in Neural Information Processing Systems, pp. 3303-3313, 2018. 5
10

---

Published as a conference paper at ICLR 2021
Ronald Parr and Stuart Russell. Reinforcement learning with hierarchies of machines. In Advances in neural information processing systems, pp. 1043-1049, 1998.
Yikang Shen, Shawn Tan, Alessandro Sordoni, and Aaron Courville. Ordered neurons: Integrating tree structures into recurrent neural networks. In International Conference on Learning Repre- sentations, 2018.
Yikang Shen, Shawn Tan, Arian Hosseini, Zhouhan Lin, Alessandro Sordoni, and Aaron C Courville. Ordered memory. Advances Neural Information Processing Systems, pp. 5037- 5048, 2019. 5, 12
Kyriacos Shiarlis, Markus Wulfmeier, Sasha Salter, Shimon Whiteson, and Ingmar Posner. Taco: Learning task decomposition via temporal alignment for control. In International Conference on Machine Learming, pp. 4654 4663, 2018. 5, 6, 13, 15
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam,, Marc Lanctot, et al. Mastering the game go with deep neural networks and tree search. nature, 529(7587):484 489, 2016. 1
Alec Solway, Carlos Diuk, Natalia Córdova, Debbie Yee, Andrew Barto, Yael Niv, and Matthew M Botvinick. Optimal behavioral hierarchy. PLOS Comput Biol, 10(8)e 2014. 1
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: frame- work for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181- 211, 1999.
Alexander Sasha Vezhnevets, Simon Osindero, Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu. Feudal networks for hierarchical reinforcement learning. arXiv preprint DrXiiv:7033... 161, 2017.
Oriol Vinyals, Igor Babuschkin, Wojciech Czarnecki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)3350 354, 2019.

---

Published as a conference paper at ICLR 2021
A OMPN ARCHITECTURE DETAILS
We use the gated recursive cell function from Shen et al. (2019) in the top-down and bottom up recurrence. We use two-layer MLP to compute the score fü for the stick-breaking process. For the initial memory M we send the environment information into the highest slot while keep the rest of the slots to be zeros. unsupervised setting, then the every slot is initialized as zero. At the first time step, we also skip the bottom-up process and hard code x such that the memory expands from the highest level. This used make sure that first time step, we could propagate our memory with the expanded subtasks from root task. In our experiment, our cell functions does not share the parameters. We find that to be better than shared-parametar
We set the number of slots to be 3 in both Craft and Dial, and each memory has dimension 128. We use Adam optimizer to train our model with 8 0.9, 82 = 0.999. The learning rate is 0.001 in Craft and 0.0005 in Dial. We set the length of BPTT to be 64 in both experiments. We clip the gradients with L2 norm 0.2. The observation has dimension 1076 in Craft and 39 in Dial. We use linear layer to encode the observation. After reading the output O concatenate it with the observation x' and send them to a linear layer to produce the action.
In section 2, we describe that augment the action space into A {done} and we append the az ) with one last step, which is T {st+1, done]. This can be easily done the data is generated by letting expert agent interact with the environment in Algorithm 1. If you do not have the luxury of environment interaction, then you can simply let 87 = ST, 47+1 done. We find that augmenting the trajectory in this way does change the performance in our Dial experiment, since the task boundary is smoothed out across time steps for continuous action space, but hurts the performance for Craft, since the final action of craft is usually USE, which can change the state a lot.
trajectory r = (x,,
Algorithm 1: Data Collection with Gym API
env= gym. (mm.maaa((nmm))
done False
obs env.reset()
traj []
repeat
action bot.act(obs)
nextobs, reward, done env.step(action)
traj.append((obs, action, reward))
obs= nextobs
until done is True;
traj (ojs done action)
B THRESHOLDING ALGORITHM
We here provide the peseudo-code for our threshold algorithm. Algorithm 2, we detect boundary by recording the time steps which goes from above to below the given threshold. We return the final K time steps where K is given. We assume the last subtask ends when the episode ends.
We design Algorithm 3 to automatically select the threshold. Our intuition is that, the optimal threshold should pass through K peaks, where K is the provided by the users, As result, we pick the highest peak and the lowest valley as the upper bound and lower bound for the threshold, and we return the final threshold the middle of these two.
C TASK DECOMPOSITION METRIC
C.I F1 SCORES WITH TOLERANCE
For each trajectory, we are given set of ground truth task boundary gt of length L which is the number of subtasks. The algorithm also produce I. task boundary predictions. This can be done in
12

---

Published as a conference paper at ICLR 2021
Algorithm 2: Get boundary from a given threshold
Data: Standardized average expanding positions Tang with length L, the number of subtasks K, and threshold T
Result: list of K split points.
preds =
prec = False
for t in range( L) do
curr = T [t] > T
if pret curr then
continue
else
if prev and not curr then
preds.append(( 1)
prev curr
apeddx. 1)
Return preds -1][: K][:: -1]
Algorithm 3: Automatic threshold selection
Data: Standardized average expanding positions Tang with length
Result: A final threshold T diff TTvn [:-1
upid = for in Iange(1, len(diff)) if diff[i] 0 and diff[i-1] 0]
lowid [i for i in range(1, len(diff)) if diff[i] >> 0 and diff'i 1] < 0]
upval [uvi[upid]]
lowval "Tagg [lowid]
upper upval.maxx)
lower lowval.min()
final = (upper + lower) 2
Return final
OMPN by setting the correct K in lopK boundary detection. For complLE, we the number of segments to be equal to N. Nevertheless, our definition of Fl can be extended to arbitaray number of predictions.
Σie match(preds,, gt,, tol)
precision
#predictions)
Eis match(gts, preds,, tol)
precision
#ground truth
where the match defined as
match(x, y, tol) = ly tol & x S S tol]
where [ is the Iverson bracket. The tolerance
C.2 TASK ALIGNMENT ACCURACY
This metric is taken from Shiarlis et al. (2018). Suppose we have a sketch of subtasks [b1, 62, 63, 64] and we have the ground truth assignment ferue (et (Etrue)=-1- Similar have the predicted alignment Epred. The alignment accuracy is simply
Σistred struel
For OMPN and compILE, we obtain the task boundary first and construct the alignment as a result. For TACO, follow the original paper to obtain the alignment.
13

---

Published as a conference paper at ICLR 2021
D BASELINE
D.1 COMPILE DETAILS
latent
[concrete, gaussian]
prior
[0.3, 0.5,0.7] [0.05, 0.1, 0.2]
kl_b
zl Z
[0.05, 0.1, 0.2]
Table 2: complLE hyperparameter search.
Our implementation of complLE is taken from the author github". However, their released code only work for toy digit sequence example. As result we modify the encoder and decoder respectively for our environments. During our experiment, we perform the following hyper-parameter sweep on the baseline in Table 2. Although the authors use latent be concrete during their paper, find that gaussian perform better in our case. We find that Gaussian with prior = 0.5 performs the best in Craft. For Dial, these configurations perform equally bad.
We show the task alignments of complLE for Craft in Figure 8. It seems that complLE learn the task boundary one-off. However, since the subtask ground truth can be ad hoc, this brings the question how should we decide whether our model is learning structure that makes sense or not? Further investigation in building better benchmark/metrri is required.
act:
tru:4 4 4 4
4 4 4
4 4
2
2
2
2
dec:4
4
2
act:: 6 + +
+
tru:4 4 4 4
4
dec:4 4
act:4
tru:4 4
4
2
dec:4 4 4 4
4
4
2
2
2
act:4
tru:4
dec:4 4 4 4 4 4
2
2 2 2
2 2
Figure 8: Task alignment results of complL. on Craft.
D.2 TACO DETAILS
dropout [0.2, 0.4, 0.6, 0.8]
decay
[0.2, 0.4, 0.6, 0.8]
Table TACO hyperparameter search.
We the implementation from author github and modifly it into pytorch. Although the author also conduct experiment on Craft and Dial, they did not release the demonstration dataset they use. As result, cannot directly prediction of STOP and apply a linear decaying schedule during training. The hyperparameter search in table 3. We find the best hyperparameter to be 0.4, 0.4 for Craft. For Dial, the result is not sensitive to the hyperparameters.
their numbers from the paper. We also apply dropout the
"https://gf thub. httt:::///gthhbb....//////////// cotp///////oo//ccomco// 'https://e hhtb.ccom//yrricccshhiiaa
14

---

Published as a conference paper at ICLR 2021
E DEMONSTRATION GENERATION
We use a rule-based agent to generate the demonstration for both Craft and Dial. For Craft, we train on 500 episodes each on ArkeAre, MkkShears and Made Each of these task further composed of four subtasks. For Dial, we generate 1400 trajectories for imitation learning with each sketch being digits.
For Craft with full observations, we design a shortest path solver to go to the target location of each subtask. For Craft with partial observation, we maintain an internal memory about the currently seen map. If the target object for the current subtask is not seen on the internal memory, we perform left- to-right, down-to-top exploration until the target object appears inside the memory. Once the target object is seen, it defaults to the behavior in full observations. For Dial, we use the hand-designed controller in Shiarlis et al. (2018) to generate the demonstration.
F CRAFT
train our models on MakeBed, MakeAre and MakeShears. The detail of their task decom- position is in Table 4. We show the behavior cloning results for all settings in Figure 10. We display more visuzliation of task decomposition results from OMPN Figure 9.
makebed makeaxe makeshears
get wood, make at toolshed, get grass, make at workbench get wood, make at workbench, get iron, make at toolshed get wood, make at workbench, get iron, make at workbench
Table 4: Details of training tasks decomposition.
6172
9011
305-
9074
6613
100
913
417
631
5.0
6111
812
202
Figure 9: More results on æ Craft. The model is able to robustly switch to a higher expanding position the end of subtasks. The model will also sometimes discover multi-level hierarchy.
We show the results of task decomposition when the given K is different in table and table 6. We find that when you increase the K, the recall increases while the precision decreases.
15

---

Published as a conference paper at ICLR 2021
0.5
Success
0.55
Success
Success Rabe
0.2
2.
0.43
0.45
M6.
0.2
CA4
0.63
ORB7%
50
0.2
0.4
08
2.0
00
0.5
3.0
2.0
0.5
3.5
2.0
(a) Full + NoSketch
(b) Full + Sketch
(c) Partial NoSketch
(d) Partial Sketch
Figure 10: Behavior Cloning for Craft.
Full NoSketch
Full Sketch
(toll=)
Pre(tol=1)
Rec(tol=1)
F1(tol=1)
Pre(tol=1)
Rec(tol=1)
K=2
62(2.4)
93(3.7)
47(1.8)
62(2.4)
93(3.4)
46(1.8)
K=3
81(2)
95(2.4)
71(1.8)
81(2.2)
95(2.3)
71(2.0)
K=4
95(1.2)
95(1.8)
95(1.8)
98(0.9)
98(1.6)
97(2.1)
K=5
92(1.6)
87(2.8)
99(0.5)
93(6.4)
87(1.1)
99(0.5)
K=6
88(2.9)
78(3.9)
100(0)
88(0.6)
79(0.9)
99(0.2)
Table 5: Parsing results for full observations with different
Parital NoSketch
Partial Sketch
Fi(tol=1)
Pre(tol=1)
Rec(tol=1)
Fi(tol=1)
Pre(tol=1)
Rec(tol=1)
K=2
64(1.8)
96(2.7)
48(1.4)
57(3.7)
88(5.6)
43(2.8)
K=3
83(1.9)
96(1.9)
72(1.9)
74(.99)
88(4.6)
64(3.5)
K=4
89(4.))
97(1.6)
82(2.7)
83(8.1)
88(4.4)
80(4.6)
K=5
93(0.8)
90(1.9)
96(2.7)
89(3.2)
85(3.7)
94(3.8)
K=6
90(1.9)
85(1.7)
98(2.3)
87(2.1)
813.4)
97(2.3)
Table 6: Parsing results for partial observations with different K
16

---

Published as a conference paper at ICLR 2021
G DIAL
We show in Table 7 the task alignment result for different thresholds. We can see that optimal fixed threshold is around 0.4 0.5 for Dial, and our threshold selection algorithm could produce competitive results. We demonstrate more expanding positions in Figure 11. We also show the failure cases, where our selected threshold fails recover the skill boundary. Nevertheless, we can see that the peak pattern exists. We show the task alignment curves for different thresholds well as the auto-selected threshold Figure 12.
Align. Acc. at different threshold
0.2
0.3
0.4
0.5
0.6
0.7
Auto
OMPN + nocnv 57(12.8)
76(9.5)
87(5.7)
89(1.4)
81(5.7)
70(9.7)
87(4)
OMPN + sketch 71(11.6)
81(10.6)
85(7.4)
84(6.6)
76(7.7)
60(6.4)
84(5.7)
Table 7: Task alignment accuracy for different threshold in Dial. The result for automatic threshold selection is in the last column.
Tra threahold
basr
0.5
28
2.-
00
20
100
340
dm 5t0p
Figure 11: More task decomposition visualization in Dial. Our algorithm recovers the skill boundary for the first four trajectories fails in the last two. Nevertheless, one can see that our model still display the peak patterns for each subtask, and more advanced thresholding method could be designed to recover the skill boundaries.
OMMPN
OMPN
0.6
0.5
2032
0.4
0.00
0.25
0.50
0.75
1.00
1.25
1.50
0.00
0.25
0.75
3.00
1.25
1.50
22800
steps
(a) Dial NoSketch
(b) Dial Sketch
Figure 12: The learning curves of task alignment accuracy for different thresholds as well as the automatically selected one.
17

---

Published as a conference paper at ICLR 2021
The behavior cloning results are in Figure 13.
Return
Return
MLP OMPN
MLP OMPN
3125
045
5875-
0.3
2329
1.00
+0.0
.046
steps
steas
(a) Dial NoSketch
(b) Dial + Sketch
Figure 13: The behavior cloning results in Dial domain.
We show the task decomposition results when the is mis-specified in Table
NoSketch
Sketch
(tol=1)
rec(tol=1)
pre(tol=1)
fl(tol=1)
rec(tol=1)
pre(tol=1)
K = 2
59(2.9)
45(2.1)
90(4.4)
58(1.4)
44(1.1)
88(2.1)
K = 3
73(5.5))
63(5.1)
85(6.2)
72(3.8)
63(3.4)
84(4.5)
K 4
84(7.1)
81(7.1)
84(6.9)
81(4.8)
81(4.6)
82(5.1)
K = 5
86(5.1)
91(3.6)
83(6.3)
83(5.2)
86(5.4)
82(5.3)
K 6
87(4.2)
93(1.5)
82(6.2)
84(5.2)
88(5.5)
81(5.2)
K = 7
87(4.1)
93(1.4)
82(6.4)
84(5.2)
89(6)
81(5.2)
K 8
86(4.1)
93(1.6)
82(6.6)
84(5.2)
90(6..1)
81(5.1)
Table 8: Fl score, recall and precision with tolerance 1 computed at different k.
18

---

Published as a conference paper at ICLR 2021
H HYPERPARAMETER ANALYSIS
Craft Full
Craft Partial
Dial
NoSketch
Sketch
NoSketch
Sketch
NoSketch
Sketch
π = 3, m =
93(1.7)
97(1.2)
84(6)
72(6.2)
87(4)
82(7.4)
N = 2, m = 64
96(1.4)
96(1.4)
87(3.3)
77(1.2)
88(3.2)
82(5.7)
4, m 64
96(0.6)
96(2.5))
88(3.1)
78(10.7)
86(9.8)
84(10)
3. m 32
92(4.2)
91(10.3)
88(6)
74(5.5)
88(4.9)
83(3.7)
3. m = 128
96(1.5)
97(1.2)
86(3)
75(5.7)
87(2.4)
83(6.2)
Table 9: Task alignments accuracy for different memory dimension m and depths n. The default setting is the first row. The next two rows change the depths, while the last two rows change the memory dimension. The number in the parenthesis is the standard deviation.
We discuss the effect of hyper-parameters on the task decomposition results. We summarize the results in Table We show the detailed alignment accuracy curves for Craft in Figure 14 and for Dial in Figure 15.
1.2
Aign
Align.
Allgn Acc
0.8
5.7'
37
0.6-
3.6
2.
35
255
0xpth
6c5th
Septh
Oepth
2.
0.3
Mepth
depth
0.2
cepth
0epe7
deptt
depth
famnes
forres
181
Fanis
fammes
(a) Full NoSketch
(b) Full Sketch
Partial NoSketch
(d) Partial Sketch
Algn
Align.Acc Acc.
1.0
Align.
Align. Acc.
5.7
0.0
26
85
'idden 123
hrdden 170
10bden 170
Feddec 125
Nddden
Nidden 22
105den
feddon
laddem
hobker
1akden
Nidden
hames
S07
tamr
lnl
ler
107
(e) Full NoSketch
(I) Full Sketch
(g) Partial + NoSketch
(h) Partial
Sketch
Figure 14: Task alignment accuracy with varying hyperparameters in Craft. We change the depths in the first row and memory size in the second row.
Align Ace.
Align Acc.
Algh. 8oc.
Algn, Acc.
0.5
3.9
0.7
5.7
0.5
55
5A
08202
54
123
0.7
4ep114
6
610 1225 D00
5.7
Li0
Jib
1.0
200
922 USD
LiU 1:25 130
(a) Dial + NoSketch
(b) Dial + NoSketch
(c) Dial + Sketch
(d) Dial + Sketch
Figure 15: Task alignment accuracy in Dial. In (a) and (c) we change the depths while in (b) and (d) we change the memory size.
19

---

Published as a conference paper at ICLR 2021
I QUALITATIVE RESULTS ON KITCHEN
We display the qualitative results for Kitchen. We use the demonstration released by Gupta et al. (2019). Since they do not provide ground truth task boundaries, we provide the qualitative results. We use a hidden size of 64 and the memory depths of We set the learning rate to be 0.0001 with Adam and the BPTT length be 150. The input dimension is 60 and action dim is 9. We train our model in the unsupervised (NoSketch) setting. We show the visualization of our expanding position, in similar fashion of the Dial domain, in Figure and Figure 17. We show the visualizations for other trajectories in Figure 18
T Tavg
fral
D
upper
08
100
233
time step
Figure 16: Visualization of the learnt expanding positions of Kitchen. The subtasks are Microwave, Bottom Knob, Hinge Cabinet and Slider.
Tavg
hit
upper
lower
Init
time step
Figure 17: Visualization the learnt expanding positions of Kitchen. The subtasks are Kettle, Bottom Knob, Hinge Cabinet and Slider.
20

---

Published as a conference paper at ICLR 2021
20-
final
uppeer
upper
08
35
bwer
50-
00:
230
105
166
276
time step
time step
final
final
upper
upper
0.
50
0.2
52-
100
150
202
150
203
time step
time step
final
findl upper
upper
0.
lower
50
lowe
62-
52-
100
129
150
130
315
time step
time step
Figure 18: More task decomposition visualization Kitchen. Our algorithm seems to recover the skill boundary for the first four trajectories but fails in the last two. Nevertheless, one can see that our model still display the peak patterns for each subtask, and more advanced post-processing thresholding method might be able to recover the task boundary.
21


Reviews:
[{"review_id": "y4NJasJzlcp", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "*Summarize what the paper claims to contribute. Be positive and generous.*\n\nThe authors introduce a new architecture, the Ordered Memory Policy Network (OMPN) that has an explicit inductive bias for modelling a hierarchy of sub-tasks. They claim that OMPN discovers task decomposition on demonstration from imitation learning, both in an unsupervised and weakly supervised setting, and its performance compares favourably with that of strong baselines on the Craft and Dial tasks. The paper also includes a small ablation study that illustrates the effect on performance of removing either bottom-up or top-down recurrence from the network. There is some analysis of performance, showing that the model can indeed identify the sub-task boundaries in both the Craft and Dial tasks. \n\n*List strong and weak points of the paper. Be as comprehensive as possible.*\n\n**Strong points**\n- I find the architecture that the authors propose interesting and novel. There is a comprehensive comparison with other approaches in the literature and a fair quantitative comparison with relevant baselines on two benchmarks.\n- The authors provide a simple ablation study to show how important are the top-down or bottom-up recurrent connections in their model. They also provide some visualisation to show how the model segments the sequence into sub-tasks.\n\n**Weak points**\n\n- I found section 2.1 and 2.2 a bit dense and slow to read. I appreciate it might be difficult to communicate a novel and complex model effectively to every reader, but though I could understand the various components, I did not feel I had gained much insight into why this model makes sense. Perhaps the authors could try to communicate their motivation for the design choices in a revised version of the paper - but I leave this up to the authors.\n- Though the paper provides a quantitative comparison with baselines on two datasets, it would be great to know what the authors think are the limitations of their proposed model. Will it scale? In what cases/condition will it not work?\n- For experiments in Craft, it is not mentioned how the experience is collected. Please include a clear description or note mentioning this.\n\n*Clearly state your recommendation (accept or reject) with one or two key reasons for this choice. Provide supporting arguments for your recommendation.*\n\nI recommend that the paper is conditionally accepted. The paper presents a novel method that outperforms other baselines on a difficult problem and does provide some analysis as to how the model works and some ablation results. Overall, the paper is well written, but there are a few points that I believe need to be addressed, if the paper is to be accepted. Please see minor comments below for small typos and necessary additions. Another condition is that the authors specify how the trajectories in Craft were collected (see last point of weak points above).\n\n*Ask questions you would like answered by the authors to help you clarify your understanding of the paper and provide the additional evidence you need to be confident in your assessment. *\n\nPlease respond to my points above. In addition, I would be interested to understand a bit more the insight of how the model works and what motivates the choices you made. Also, I believe that Figure 3 requires some more explanation, as I am not confident I understand everything that’s meant to communicate. \n\n*Provide additional feedback with the aim to improve the paper. Make it clear that these points are here to help, and not necessarily part of your decision assessment.*\n\n- Figure 4: Please indicate which task id/level correspond to each subplot, or means to reproduce it. Either to the plots or the figure caption.\n- Page 3: “The score fti can be interpreted **as** the probability that subtask i is completed at time t.”\n- Page 4: “We find this method is suitable for continuous control ~~settubgs~~ settings, where the subtask boundaries are more ambiguous and smoothed out ~~accross~~ across time steps.”\n- Page 5: “However ONLSTM does not ~~provides~~ provide mechanism to achieve the top-down and bottom-up recurrence.”\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "xnT71qyLzFJ", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#######################################################################\n\nSummary:\n\nIn this paper the authors propose a new method for task decomposition. First a new neural architecture is proposed which includes a set of 'memories' who's operation is inspired by the HAMs architecture in which individual memories correspond to subtasks which can be internally updated, call the next-level subtask in the stack, or return control to the operating subtask. After training, subtask boundaries can then be extracted by inspecting the values of these memories.\n\n#######################################################################\n\nReasons for score:\n\nWhile I think the proposed method is interesting concept there seems to me to be insufficient evidence to suggest that it is general. The authors seem to explicitly control the depth of the memory stack and the expected number of subtasks in subsequent analysis but do not present alternate results. I think the current work would be notably stronger given broader experimental support.\n\n#######################################################################\n\nPros:\n\n1. I found the introduction to be clearly written (if lightly referenced) and appreciated the simple example\n\n2. The architectural construction itself would seem to be reasonably simple, and therefore potentially more broadly applicable in the future.\n\n#######################################################################\n\nCons:\n\n1. There were a number of unjustified (or weakly justified) design choices made throughout the paper. While I understand that this is to some extent inevitable when implementing new methods, I think the paper could be notably bolstered by discussing these further. For example I am not sure how the number of memories 'n', or the dimension of those memories 'm' ought to be chosen in general, or what the effect of different choices might be. Similarly, how was the threshold value selected? Simply by visual inspection? It would also be valuable to expound on the rational for making the bottom-up relation recurrent.\n\n2. The experimental section could be strengthened\n\n(1) nit: Fig 3 would be much easier to read if you used images rather than letters in your plots\n\n(2) nit: In general all axes should be labeled, experiments separated by title, methods by colour etc\n\n(3) Experiments are insufficient to justify claims (for example, in the craft world partially observable setting, the observed performance improvement over LSTM seem to be significant - it'd have been great to see that further explored... for the top-k choice, why not show the results for other values of k or suggest methodology for determining this value automatically)\n\n#######################################################################\n\nQuestions during rebuttal period:\n\nQ1: if a practitioner does not know the correct value of 'n' and 'k' ahead of time, how would they use this approach to discover the task decomposition? Or rather, how should these values be chosen?\n\nQ2: could you say a bit more about the degenerate behaviour observed when the explicit pi_end was excluded?\n\nQ3: the alignment accuracy for different thresholds in Fig9 are not monotonic - do you have any intuition for why this is?\n\n#######################################################################\n\nSome typos:\n\n(1) nit: bolding key / leading values in table has probably become standard now\n\n(2) Fig 6 axes values", "extra_scores": {}, "extra_text": {}}, {"review_id": "xiPEuhEThyq", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \nPaper introduces Ordered Memory Policy Network (OMPN) with an objective to learn sub-task decomposition and hierarchy from demonstrations in the context of Imitation Learning under unsupervised and weakly supervised settings. Authors approach the problem of uncovering the task substructure from the architecture design lens where the goal is to design architectures that have the right inductive biases leading to the discovery of task sub-structures. The proposed solution views sub-tasks as finite state machines represented as memory banks that are updated via top-down and bottom-up recurrences. \n\nStrengths:\nPaper is well polished and easy to follow. OMPN is shown to be effective in two domains. OMPN advantage in partially observable environments is also effectively demonstrated. \n\nDiscussion:\nMy reservations are with the evaluations. \n- The choice of baselines also seems a bit narrow. Comparisons to related methods dealing with sub-task decomposition and organization are missing. For example, a comparison to the Relay Policy Learning[1] (which the paper claims to be the most related work) is missing. Another closely related method Learning Latent Plans from Play[2] (https://learning-from-play.github.io/) is also missing.\n- Selected tasks have a very shallow task hierarchy and sparse task structure. The ceiling of the approach and limitations aren't clearly outlined and are less evident. This is also evident from the tasks considered in [1] and [2]. These methods have shown to be effective with play demonstrations in rich scenes such as kitchen and study-table scenes where are underlying task-structure is much more convoluted to uncover. Without similar comparisons, it's hard to evaluate the strength of OMPN. I'd strongly advise on including experiments under similar settings which will make the submission really strong.", "extra_scores": {}, "extra_text": {}}, {"review_id": "HwpQVkPwEi", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThis paper proposes a hierarchical memory structure, where each layer of memory stores information about the corresponding level of sub-task and uses this structure for better behavioral cloning. The memory module mimics the call-return mechanism, in which the higher-level memory (sub-task) changes to the next one when the lower-level memory (sub-task) is completed. With this inductive bias, the memory module and policy can be trained end-to-end from a sequence of state-action pairs, and learn effective task decomposition in multiple levels. By looking at changes in each level of the memory module, which means changes in sub-tasks, we can find the task decomposition. The experimental results on the Craft World (grid world) and Jaco arm dialing environments demonstrate that the proposed method can learn better task decomposition in both unsupervised and weakly supervised settings.\n\n### Strengths\n- The idea of using the ordered memory module to learn the subtask structure in multiple levels is novel and intuitive.\n- The subtask decomposition results show superior task alignment of the proposed method over baseline methods in both unsupervised and weakly supervised learning settings.\n- The paper is well-written and easy to follow.\n\n### Weaknesses\n- In Section 4.1, the paper claims that the proposed method outperforms the LSTM baseline due to its superior ability to store longer-term information. But, OMPN outperforms the LSTM baseline only under the weakly supervised setting. Is there any good explanation of why the proposed method does not show improved performance in the unsupervised setting?\n- The learning curves and final performance in the Dial task are missing in the paper. Does it achieve higher performance compared to baselines (MLP, LSTM policies)?\n- It is not clear whether the hierarchy in memory is important or not. For example, in Figure 3, the horizontal expansion does not happen in slot 3, which can mean 2 level-hierarchy could be enough for solving the task. Ablation study on a different number of levels in the memory module can help to understand the importance of the hierarchical memory.\n- In Figure 6, when the robot arm presses the number 4, the expected expansion position is high for multiple steps, which means the high-level task is likely to be changed over multiple steps. But this result seems not desirable. Is there any explanation why this still makes sense?\n\n### Questions and additional feedback\n- It would be easier to understand Figure 3 if each map is rendered with icons, not with words.\n- Adding sub-captions for four graphs in Figure 4 will make it easy to read.\n- In Figure 5, what is the difference between (a,b) and (c,d)? It is unclear from the text and figure.\n\n### Overall assessment\nThe proposed method is novel and intuitive and shows improved task decomposition on the Craft World and Dial environments. However, the paper needs a more thorough analysis of the results, and an additional ablation study can make the claims stronger. In summary, the reviewer thinks this paper can be accepted with a few more experiments and analysis.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents the Order-Memory Policy Network (OMPN), an architecture for modelling a hierarchy of sub-tasks and discovering task decompositions from demonstration data. Results are presented on a compositional grid-world task (Craft) and on a simulated robotics task (Dial).

The reviewers agree that the proposed method is novel and interesting, that the paper addresses an important problem, and that it is well-written. One main criticism by the reviewers, the lack of experimental evaluation of different hyperparameter choices, such as the depth of the memory stack and the expected number of subtasks, has to a large part already been addressed in the revision by the authors. The total number of hyperparameters that need to be tuned, however, is quite large and the authors are encouraged to revise their claim "Our central message is that OMPN is a general off-the-shelf model for task decompositions" in this light. The paper is borderline, and could clearly benefit from a revised, stronger presentation and more extensive experimental evaluation, but I am confident that the authors can use the time until the camera-ready version is due to address some of the remaining feedback by the reviewers, and hence I think that this paper can be accepted.

The authors are further encouraged to take the following additional reviewer feedback into account, which was brought up during the internal discussion period:
1) The complexity of the proposed method could be better justified by more thoroughly investigating the effectiveness of using a multi-level hierarchy (e.g., by running experiments on more complicated and hierarchical tasks with multiple branches).
2) Further strengthening down-stream performance evaluation, such as in imitation learning (in addition to the already presented behavioral cloning results) and/or reinforcement learning, would further strengthen the paper and demonstrate that the discovered decomposition is indeed useful.

Decision: Accept (Poster)