Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: Summary: 
Paper 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. 

Strengths:
Paper 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. 

Discussion:
My reservations are with the evaluations. 
- 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.
- 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.
review_title: Clear, easy to follow, missing comparisons. 

final_score: rating: 6: Marginally above acceptance threshold
review_content: ### Summary
This 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.

### Strengths
- The idea of using the ordered memory module to learn the subtask structure in multiple levels is novel and intuitive.
- The subtask decomposition results show superior task alignment of the proposed method over baseline methods in both unsupervised and weakly supervised learning settings.
- The paper is well-written and easy to follow.

### Weaknesses
- 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?
- 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)?
- 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.
- 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?

### Questions and additional feedback
- It would be easier to understand Figure 3 if each map is rendered with icons, not with words.
- Adding sub-captions for four graphs in Figure 4 will make it easy to read.
- In Figure 5, what is the difference between (a,b) and (c,d)? It is unclear from the text and figure.

### Overall assessment
The 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.
review_title: Interesting memory architecture for task decomposition but requires more thorough experiments

final_score: rating: 4: Ok but not good enough - rejection
review_content: #######################################################################

Summary:

In 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.

#######################################################################

Reasons for score:

While 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.

#######################################################################

Pros:

1. I found the introduction to be clearly written (if lightly referenced) and appreciated the simple example

2. The architectural construction itself would seem to be reasonably simple, and therefore potentially more broadly applicable in the future.

#######################################################################

Cons:

1. 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.

2. The experimental section could be strengthened

(1) nit: Fig 3 would be much easier to read if you used images rather than letters in your plots

(2) nit: In general all axes should be labeled, experiments separated by title, methods by colour etc

(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)

#######################################################################

Questions during rebuttal period:

Q1: 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?

Q2: could you say a bit more about the degenerate behaviour observed when the explicit pi_end was excluded?

Q3: the alignment accuracy for different thresholds in Fig9 are not monotonic - do you have any intuition for why this is?

#######################################################################

Some typos:

(1) nit: bolding key / leading values in table has probably become standard now

(2) Fig 6 axes values
review_title: Review: Learning task decomposition with order-memory policy network

final_score: rating: 6: Marginally above acceptance threshold
review_content: *Summarize what the paper claims to contribute. Be positive and generous.*

The 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. 

*List strong and weak points of the paper. Be as comprehensive as possible.*

**Strong points**
- 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.
- 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.

**Weak points**

- 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.
- 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?
- For experiments in Craft, it is not mentioned how the experience is collected. Please include a clear description or note mentioning this.

*Clearly state your recommendation (accept or reject) with one or two key reasons for this choice. Provide supporting arguments for your recommendation.*

I 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).

*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. *

Please 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. 

*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.*

- 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.
- Page 3: “The score fti can be interpreted **as** the probability that subtask i is completed at time t.”
- 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.”
- Page 5: “However ONLSTM does not ~~provides~~ provide mechanism to achieve the top-down and bottom-up recurrence.”

review_title: Interesting model, but the paper requires some minor changes

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)