Identifying Selections for Unsupervised Subtask Discovery
Authors: Yiwen Qiu
Yujia Zheng
Kun Zhang
Venue: NeurIPS 2024 poster (2024)
Keywords: Imitation Learning
Causality
Selection
Primary area: reinforcement_learning
TL;DR: This paper addresses the issue of subtask discovery from a causal perspective in Imitation Learning by identifying subgoals as selections

Abstract: When solving long-horizon tasks, it is intriguing to decompose the high-level task into subtasks. Decomposing experiences into reusable subtasks can improve data efficiency, accelerate policy generalization, and in general provide promising solutions to multi-task reinforcement learning and imitation learning problems. However, the concept of subtasks is not sufficiently understood and modeled yet, and existing works often overlook the true structure of the data generation process: subtasks are the results of a *selection* mechanism on actions, rather than possible underlying confounders or intermediates. Specifically, we provide a theory to identify, and experiments to verify the existence of selection variables in such data. These selections serve as subgoals that indicate subtasks and guide policy. In light of this idea, we develop a sequential non-negative matrix factorization (seq- NMF) method to learn these subgoals and extract meaningful behavior patterns as subtasks. Our empirical results on a challenging Kitchen environment demonstrate that the learned subtasks effectively enhance the generalization to new tasks in multi-task imitation learning scenarios. The codes are provided at this [*link*](https://anonymous.4open.science/r/Identifying\_Selections\_for\_Unsupervised\_Subtask\_Discovery/README.md).

Identifying Selections for Unsupervised Subtask Discovery
Yiwen Qiu
Yujia Zheng
Carnegie Mellon University
Carnegie Mellon University
Pittsburgh, PA 15213
Pittsburgh, PA 15213
yiweng@andrew. cdu
yujiazh2@andrew cmu .mu.edu
Kun Zhang
Carnegie Mellon University, MBZUAI Pittsburgh, PA 15213
kunz1@andrew. cmu.edu edu
Abstract
When solving long-horizon tasks, intriguing to decompose the high-level task into subtasks. Decomposing experiences into reusable subtasks can improve data efficiency, accelerate policy generalization, and in general provide promising solutions to multi-task reinforcement learning and imitation learning problems. However, the concept of subtasks not sufficiently understood and modeled yet, and existing works often overlook the true structure of the data generation process: subtasks are the results of a selection mechanism on actions, rather than possible underlying confounders or intermediates. Specifically, we provide a theory to identify, and experiments to verify the existence of selection variables such data. These selections serve subgoals that indicate subtasks and guide policy. In light of this idea, we develop a sequential non-ne non-negative matrix factorization (seq- NMF) method to learn these subgoals and extract meaningful behavior patterns subtasks. Our empirical results on a challenging Kitchen environment demonstrate that the learned subtasks effectively enhance the generalization to new tasks in multi-task imitation learning scenarios. The codes are provided at this link.
Introduction
Being able reuse learned skills from past experiences and conduct hierarchical planning (LeCun [2022], Hafner et al. [2022], Gehring et al. [2021]) is crucial for tackling real-world challenging tasks such as driving: it meaningful to let higher levels of abstractions perform longer-term prediction (of subgoals), while lower levels of policy perform shorter-term actions. The concept of breaking down a task into subtasks can be illustrated through the example commut ing to New York. You at home, and the overall task, commuting to New York, can be decomposed into smaller, manageable subtasks. This includes subtasks like walking out of the house, getting into the car, driving and catching an airplane.
Even more granular, each of these subtasks can be further broken down into smaller actions: k ng out of the house involves standing up, grabbing the luggage and wa king to the door. This method decomposing task into subtasks fits our intuition of how humans perform actions and helps to simplify complex tasks, making them more manageable.
"Corresponding author.
'We predominately use the concept of subtask in this paper, which has similar semantics as that of option or skill in other literature.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

For artificial intelligence (AI) to match the ability of humans in terms of understanding the event structures (Zacks al. [2001]. Baldassano et al. [2017]) and learning to perform complex and long- horizon tasks by reinforcement learning (RI.) (Sutton and Barto [2018]) or imitation learning (II.) (Hussein et al. [2017]), it is natural to ask this question: with an abundance of past collected either human or robotics experiences, how can one extract reusable subtasks, such that we can use them to solve future unseen new complex tasks? Current RL is well known for its sample inefficiency, and learning decomposition of subtasks serves as the basis for perform complex tasks via planning. The benefit is straightforward: extracting reusable disentangled (Denil et al. [2017]) temporal-extended common structure can enhance data efficiency and accelerate learning of new tasks (Thrun and O' Sullivam [1996]. Florensa et al. [2017]. Griffiths et al. [2019], Jing et al. [2021]). This is the main motivation for subtask discovery, the problem we aim to thoroughly investigate in this paper.
In subtask discovery, the criterion to segment these subtasks is vital, and should be consistent with the true generating process. However, most prior works did not explore the ground-truth generative structure to identify that criterion- instead, they simply focus on designing heuris- tic algorithms (McGovern and Barto [2001]. Simşek and Barto [2004], Wang et al. [2014], Paul et al. (2019]) or maximizing the likelihood of the data sequences that is built upon intuitive graphical models (Krishnan et al. (2017), Kipf et al. [2019], Sharma al. [2019b], Yu et al. (2019]). Relying on a structure that conflicts with the true generating process may mislead the segmentation as well as worsen the performance of those downstream IL tasks. Thus, we argue that it is helpful to consider the causal structure underlying the sequence, specifically, the selec- tion variables as subgoals to discover subtasks.
picnicking
go to a fnovic
Figure I: Example of subgoals as selections. One subgoal to "go picnicking", another subgoal is to "go to a movie". In order to "go picnicking", you need to go shopping first and then drive to the park; in order to "g to movie". you need to check the movie information online first and then get the tickets. The actions caused us to accom- plish the subtasks, and we essentially select the actions based on (conditioned on) the subgoals want to achieve. On the contrary, weather is a con- founder of the states and actions: changing our ac- tions would not influence the weather, but actions influence whether we can achieve the subgoals.
Selective inclusion of data points is prevalent
in practical scenarios. We provide a short illustration of selection in Fig. to help distinguish the confounder and the selection case. We are at home, and have two subgoals to select from: "pienicking" and "movie", and go back home again. One way to look at selection is recognizing its preferential nature of including certain data (Heckman [1979]): achieving a subgoal involves performing subtask, which consists of a sequence of states and actions that must follow a unique pattern. Another way to look at selection by considering interventions (Eberhardt and Scheines [2007]) which is a widely used concept in the literature on causal discovery. Intervention involves assigning values to a single variable: in our case, the actions lead us to achieve a subgoal, and intervening on the actions would lead us to achieve others. In other words, the subgoal functions as a selection for the actions. Overlooking the selection structure is adopting an inappropriate inductive bias, and can distort our understanding of the data. Therefore, the incorporation of identifying and modeling the selection structure to uncover the true causal process is crucial for understanding subtasks. A more comprehensive literature review recent advances in understanding selections and subtask discovery can be found in Appx. A.
The main contributions of this paper are as follows: (1) We show that one can identify the selection structure without interventional experiments, and confirm its existence through our experiments. (2) According to the identified structure, and based on the formal definition of subtasks, we propose a novel sequential non- -egaateee matrix factorization (MeQ-NMF) method to learn subgoals selections. (3) We demonstrate that the learned subtasks can be leveraged to improve the performance of imitation learning in solving new tasks.
2 Preliminary
Imitation Learning. In standard Imitation Learning, we collected a set of trajectories = (Tn IN from an expert in a Markov Decision Process (MDP). An MDP is defined by AS,A, R. po. 7), with S as the state space, A as the action space, P : S x A x S ++ [0, 1] as the transition probability, R: S x A + I as the reward function, Po the initial state distribution, and y as the discount factor.
2

---

Suppose we are given N expert trajectories Tn generated in an MDP with unknown reward. Each trajectory Tn = (st, "t.".. | is composed of a series of states &t E S and actions a, E A. The goal of the agent is to learn a policy ((a, se) : 8 x A + [0, 1] that mimics the expert's behavior. Option Framework. For multi-task settings, learning hicrarchical policy that extracts basic skills has been proven useful in solving tasks composed of multiple subtasks. Sutton et al. [1999] proposed the option framework, which hierarchical reinforcement learning framework that decomposes task into set of temporally extended options. An option O is defined as tuple (I, x, 18), where I I S is the initiation set, π is the policy, and 9 : s - [0, 1] the termination function deciding whether the current option terminates. The execution of a policy within an option is semi-Markov decision process (SMDP). By sequentially performing SMDPs, that is, taking an option O₂ until terminates, and then taking the next option O₂ until it terminates, the agent can learn hierarchical policy that executes in the overall MDP.
Causal Graph and Common Assumptions. In Bayesian network, the distribution IP over a set of variables is assumed to be Markov w.r.t. to a directed acyclic graph (DAG) G = {V, E. where V is the set of vertices and & is the set of edges, and the DAG G is faithful to the data. The Markov condition, and faithfulness assumption are defined in Def. and Def. 2, respectively.
Definition 1. (Markov Condition (Spirtes et al. [2001], Pearl [2009])) Given a DAG & and distribution P over the variable set V, every variable X in V is probabilistically independent of its non-descendants given its parents in G.
Definition 2. (Faithfulness Assumption (Spirtes et al. [2001], Pearl (2009]) There are no indepen- dencies between variables that are not entailed by the Markov Condition.
Combining the Markov Condition and Faithfulness Assumption, we can use d-separation criterion (denoted as X IT Y Z) to read all the conditional independencies from a given DAG G: Definition 3. (d-separation (Spirtes al. [2001], Pearl [2009])) Two sets of nodes X and Y in G is said to be d-separated by a set of nodes Z C V if and only if: for every path p that connects one node i in X to one node j in Y, either (1) p contains
such that m is in Z, or (2) p contains a collider m, i.e. i + m < j such that m and all descendants of m are not in Z.
Problem Formulation Given the above context, we formulate the considered imitation learning problem as follows: we have distribution of tasks PT( and corresponding set of expert trajectories D= = (T,, 1 and we aim to let the imitater learn a policy that can be transferred to new tasks that follow a different distribution (T) Each task sampled from (T) should be generated by a MDP and composed of a sequence of option {O,, ... }. where O = (I,, TT 33);)- We use O == - O, to denote all J options, and &, = {Se, a It= a sub-sequence of states and actions (st, a4) from any trajectory Tn Each trajectory can be partitioned into sub-sequences E_p with maximum length L. Unlike traditional MIL (Seyed Ghasemipour et al. [2019], Yu et al. [2019]), we assume a shift in the task distribution, i.e. (T) * PT( but they only share the same option set, and we expect the agent to leverage the past learned subtasks to carry out new tasks efficiently.
ILL
Roadmap for solving the "subtask discovery" problem The question raised above on how to discover useful subtasks for policy generalization on new tasks can be answered two-fold. First, it is essential to adopt an accurate understanding of subtasks that aligns with a true data generation process: it is a matter of comprehending temporal data dynamics (Matsubara et al. [2014], Chen et al. [2018), Dyn), to which we provide an answer in Sec. 3.1. Second, based on that understanding, we need to design a learning algorithm that can extract subtasks from expert demonstrations, as discussed in Sec. 3.2. Finally, the learned subtasks should be used to facilitate policy training so that the policy can quickly adapt to new tasks by alternating between subtasks, which we address in Sec. 4.
3 Subgoal as Selection
The first step is to understand what a subtask is. We propose to understand subtask by building the causal graph and distinguishing different potential structures. We assert that subtask is indicated by a selection variable denoted as gt (subgoal), and will distinguish from confounder and intermediate node (Sec. 3.1). Then, we give formal definition of subtasks as sub-sequences that can be representative of common behavior patterns and avoid uncertainties in policy (Sec. 3.2). With the understanding of selection (Sec. 3.1) and formal definition (Sec. 3.2), we propose novel sequential non-negative matrix factorization (seq-NMF) method that aligns with these two ideas to leam subtasks from multi-task expert demonstrations (Sec. 3.3).
Overview
3

---

3.1 Identifying Selections in Data
To better understand subtasks, we build DAG G (V. E} to represent the generation process of trajectory by setting each St, as vertices (2), and add edges (E) by connecting St - St+1, at + 1 to represent the transition function P. There are three potential patterns follows:
Definition 4. (Confounder, selection and intermediate node) Ct is confounder if St Cc - at. gt is a selection if St + St <- at, m, is an intermediate node St - M. + at.
In other words, the causal Dynamic Bayesian Network (DBN) (Murphy [2002]) of a series of (Se, ag, :.. is one of the following three scenarios:
82
83
m1
m2
81
$2
83
01
0/2
$1
8/
83
01
0/
91
92
01
022
(a) Confounders
(b) Selections
(c) Intermediates
Figure 2: Three kinds of dependency patterns of DAGs that we aim to distinguish. Structure (1) models the confounder case St f Ct + at, structure (2) models the selection case St + ge at, and structure (3) models the mediator case St - m, + ag. all three scenarios, the solid black arrows (-) indicate the transition function that is invariant across different tasks. The dashed arrows (+) indicate dependencies between nodes d, and d++1- We take them to be direct adjacencies the main paper, and for potentially higher-order dependencies,
refer to Appx. B.4.
Selection implies that we can only observe the data points for which the selection criterion is met, such as reaching some subgoal j. c.g. gt (f) = 1. As a consequence, the trajectory distribution p(st, ae) is actually the conditional one: p(s. at gt). Understanding subgoals selections that are always conditioned on enables us to explain the relationship between the states and actions: the dependency between S, and a, is the result of the given selection, i.e. V. Ka gt. We argue that the policy (aa (t) does not indicate a direct causal relation between states and action, but rather an inference from states actions, and the dependency is built by subgoal selections.
We start by proposing a conditional independence test (CI test) based condition in Prop. 1 as a sufficient condition for recognizing selections:
Proposition 1. (Sufficient condition) Assuming that the graphical representation is Markov and faithful to the measured data, st y at dr. then de is selection variable, i.c., de = gt. under the assumption that:
(confounder, selection, and intermediate nodes can not co-exist) At each time step, de can only be one of [Ce, gt, m.). (For a relaxation of this assumption, see Appx. B.4).
2. (consistency in time series) At every time step, de plays the same role as one of {ct, gt, m,).
For necessary and sufficient condition, we have the following proposition:
Proposition 2. (Necessary and sufficient condition) de is selection variable (de := gt) if and only if s Kaa
de and d, _L a++1 d,+1-
Since our goal is to identify the presence of selection, the necessity aspect is not our primary focus, but it is still worthwhile to consider. For a condition that weaker but has real-world implications, we have the following necessary condition:
Propasition 3. (Necessary condition) If d, is selection variable (d, = gt), then $,41 1 1 g: s,, a,. (Such independency does not hold true for confounders case which discussed in Appx. B.3)
This proposition implies that 81+1 only determined by St. a, and solely relies on the transition function in the environment. The hidden variable g, would not provide much additional information for the prediction of $1+1- On the other hand, if we consider the confounder case, it is unrealistic that changing the subtask would result in changing the transition function, the confounder case entails.

---

Such is strong evidence for us to assert that subtasks should be considered as selections. The proof of Prop. Prop. 2 and Prop. are provided in Appx. B.1, B.2. B.3, respectively.
Remark 1. As relaxation of Prop. 1, we do not assert gt Et and m, to be mutually exclusive. As long as we have St. J at. de and de I at++ d then the selection mechanism holds. There can also be confounders and intermediate nodes, but for simplicity, we do not consider them in the main paper and leave the combination of multiple Ct; gt and m, to Appx. B.4.
Remark There a parallel between the selection variable and the reward signal reinforcement learning from probabilistic inference view (Levine (2018)). See discussion in Appx. E.
Based on Prop. 2, we verify that there is indeed selection in the trajectories, as indicated by experi- ments (Sec. 5.1). These selection variables serve as subgoals that can facilitate imitation learning.
3.2 Definition of the Subtask
In Prop. I, we provide the sufficient condition to identify selection (subgoal), that is St 1 at gt. In other words, the current action a, is affected by both the current state S, and the current subgoal ge- Only St alone is not sufficient to determine action, but also the current subgoal gt guides the agent's s action. e.g. When arriving at a crossroad, and you are deciding whether to turn left or right, it is only when a subgoal (the left road or the right road) is selected, then the -ondiilooodd policy Tg |, St. is uniquely determined. Therefore, we can define subtasks as sub-sequences that can: (1) be representative of common behavior patterns (because g& guided the selection of a (2) avoid uncertainties in policy (a, 8t). that is, different distributions of action predictions given state. The formal definition goes as in Def. 5.
Definition 5. Subtask is defined as a set of J options O = =O,}/= s.t. for some partition of
trajectories, & = (st, at. At, at. )_1) as a sub-sequence of (s) from any trajectory Tn and the maximum length of the sub-sequence, and also the maximum lag of each subtask:
min J 1, s.t.
(VC,) (30,) Ep Oj , and if (3s) E Sp. Ep + p()
(I)
that Ty(
S 5:) # ( St si), then s; O3.5 2
We require every sub-sequence to be generated from some option ~ jj. which means that its first state (0) Ij and the last state is termination state B, (-1() = 1, and the actions in En are
generated by Tz. [a Se). Importantly, when there are multiple policies available at one state 8,, i.e. (y St si) # (as St = si), then these policies should correspond different options, Oj and Oy (j # j'). in order to avoid unmodeled uncertainty for the imitator.
By definition, subtasks are our way of recovering minimal number of options from trajectories, and we view sub-sequences in the data as instantiations of options, because they are generated from corresponding SMDPs. Then, each Tn should be some combination of those sub-sequences. On the one hand, we aim avoid excessively granular partitions that result in large number of one-step options, because it is essential to capture long-term patterns. On the other hand, we seek a sufficient number of subtasks avoid policy ambiguity; when multiple options are available for a single state, we should be able to differentiate them by setting distinct subgoals. Consequently, in situations with varying distributions of action predictions, we employ different subtasks capture these differences. This involves selecting subgoals to predict actions that a single policy cannot suffice. For example, determining to turn left a crossroad as one subgoal and turning right as another.
Justification of necessity By this definition of subtasks, reinforce the necessity of learning subtasks from the perspective of avoiding policy uncertainties in multiple trajectories. Since the trajectories are collected by a variety of human or robot experts under different tasks, they are likely to exhibit different optimal policies. If there is an ambiguity about which policy to imitate, i.c., there are multiple optimal action predictions at hand, we can comprehend it as there is latent variable g& affecting the prediction of a, namely a subgoal that helps to select action. Learning one policy distribution p(a, st) from a mixture of different policies ignores the variety of behaviors, and focuses merely on the marginal policy x x(a, 5e) rather than subgoal conditioned policy TT (a, Se. gt), leading to unsatisfactory results. To our knowledge, we are the first ones to give a definition that eliminates the ambiguity in the definition of subtasks. While previous works provide only general and vague definitions by human intuition (MeGovern and Barto [2001], Şimşek and
Barto [2004], Wang et al. [2014], Paul et al. [2019], Krishnan et al. [2017], Kipf et al. [2019], Sharma
5

---

et al. [2019b], Yu et al. [2019]), e.g, semantically meaningful sub- sub-trajectories, we explicitly express subtasks as patterns that capture all policy uncertainties exhibited in the dataset. These uncertainties are then mitigated by policy selections, i.e. the subgoal-conditioned policy.
3.3 Learning Subtasks with Seq-NMF
In Sec. 3.1, we give the conditions to identify subgoals as selections (verified in experiments in Sec. 5.1). Sec. 3.2, we provide formal definition of subtasks. Combining these concepts, we
connect subgoals to subtasks: subgoals are multi-dimensional binary variables G2 E {0,1)", while subtasks are sets of options O that generate sub-sequences ε, {Se a I1=1" In particular,
SL
use G. (Ep) = (Ep ~ O2) to indicate whether the subgoal of a sub-sequence is generated from an SMDP captured by Jj where 1 (.) is the identify function. Similarly, we use H.!) ([se. at])
(st E I to indicate whether 8t is an initial state of j.
122,
X
a5
O
Sp-
x
93
94
G
h
H
(a) Causal model for expert trajectories.
(b) Illustration of seq-NMF
Figure 3: Figure (a) is the causal model for expert trajectories, which is further abstracted as the matrices in Figure (b), which can be learned by a seq-NMF algorithm. both figures, data matrix X is the aggregated (se; a ) and H E {0, J represents the binary subgoal matrix.
This formulation can be intuitively understood as follows: subtasks represent the behavior patterns one might select perform and are thus exhibited in the expert trajectories. In contrast, the subgoal is the selection variable itself, indicating whether or not this behavior pattern has been executed.
Method: Sequential Non-Negative Matrix Factorization Learning such binary coefficient matrix and feature pattern is closely related the non-negative matrix factorization (NMF) (Lee and Seung [1999]), which focuses on decomposing complex dataset into set of simpler, interpretable components while maintaining non-negativity constraints. thorough review of NMF is in Appx. G.
In our setting, instead of using vector to represent pattern, we want the pattern to be tem- porally extended, sharing the merit of those works of extensions of NMF (Smaragdis [2004, 2006], Mackevicius et al. [2019]). We set X; = (st; ag) and concatenate X across time to form data matrix X. We then transform the problem of learning subgoals into matrix fac- torization problem: identifying the repeated patterns within sub-sequences. The entire data matrix (trajectories) can be reconstructed with subtasks representing temporally extended pat- terns, and binary indicator representing which option is selected at each time step. We define JxL as tiree-dimensionl tensor, with subtask patterns h {0, 1| as corresponding indicator binary
O, O, E R DxL, and H h h2 matrix, where D d, + d, is the dimension of Xt. and L is the maximum length of each subtask pattern. We construct the decomposition as:
O O Og
L-1
X O H. where(O H)
Σ
Σ o Odee H
H(--))
and a convolution operator that aggregates the patterns across time lag L. Then the optimization problem is transformed into:

---

(0", H.) arg min O,H
| - X| + R)
L-1 0 dfj
s.t. Xat =
((-
j=1
|| || is the Frobenius norm, and R. is the regularizor. In order for it to fit our framework proposed in Sec. 3.1 and Sec. 3.2, we need three terms in the regularizor: R = Rbin + R1 + Rsim with
corresponding learning rates Abin A1 and Xaim-
The first term Riin corresponds to the binary nature of selection. Because the set {0, 1 is not convex,
remove the constraint H {0, 1} to regularizer Rbin Abin | (1 H) forcing the subgoal to be binary, where @ is the clement-wise product. Because of sparsity of subtasks in its definition, require minimal number of subgoals, which should be an L. penalty term. We use the L1 penalty R2 = A ||| to approximate such sparsity since solving the LO regularized problem is NP-hard. Finally, we should have distinct features of subtasks: distinct common patterns should be distinguished as the same subtask, i.e. there should not be similar or duplicated patterns between
( + X SH
any two different subtasks O, and Oj.. We use R Asmm
avoid such redundancy, where S is a T x T smoothing matrix: Siz = 1 when | ii- K< L and Sij = 1 otherwise. Specifically, the first term (o -xx x)
1,i5j
calculates the overlap of data X with subtask pattern j at each
DE OdjeX ((++e) and * is the transposed convolution operator. Then by multiplying the loadings SH within the time shift of L, we obtain the correlation between different patterns' overlapping with the data. If the correlation is high, it means that the two patterns have similar power in explaining the data at time I. Diagonal entries are omitted. We provide the detailed discussion on subtasks ambiguity in Appx. F.
time step f, where (o*x), je
Optimization. In terms of optimization, we derive multiplication rules which have been proven to be more efficient in solving problems with non-ne negative constraints Lee and Seung [1999]), rather than relying standard gradient descent. The detailed derivation is provided in Appx. G and the overall algorithm of seq-NMF is described in Appx. C.I.
X + H
O
*
X
O + O
H
(4)
X
H
BR 20
O
X
JR 7
4 Transfering to New Tasks
After learning subtasks from demonstrations, it is intuitive utilize them by augmenting the action space with the subgoal selection. We learn new policy that takes in both the state and subgoal as input, same in other literature (Sharma et al. [2019b]. Kipf et al. [2019], Jing al. [2021], Jiang et al., Chen et al. [2023]).
Among all the different ways to perform II., such Behavioral Cloning (BC) (Bain and Sammut [1995]), Inverse Reinforcement Learning (IRL) (Abbeel and Ng [2004], Ng et al. [2000], Ziebart et al. [2008], Finn al. [2016]), and Generative Adversarial Imitation Learning (GAIL) (Ho and Ermon [2016], Fu et al. [2018]), we adopt a GAlL-based approach that matches the occupancy measure between the learned policy and the demonstrations through adversarial learning to seek the optimal policy. The overall objective is:
min max "
E (or log(1 Do(st, Gz. 9:)) + E( (o G.)),
(5)
where Tg is the augmented policy, and Dy S x A J + [0, 1] is parametric discriminator that aims at distinguishing between the samples generated by the learned policy and the demonstrations. The policy is trained via PPO (Schulman et al. [2017]). Algorithms for the overall training procedure and the execution of policy can be found in Appx. C.2. Comparison with other II. algorithms is provided in our experiment in See. 5.3.
7

---

Experiments
The goals of our experiments are closely tied our formulation in Sec. 3. Sec. 5.1, we verify the existence of selection in data. In Sec. 5.2, we evaluate the effectiveness of seq-NMF in recovering subgoals selections. In Sec. 5.3, we demonstrate the learned subgoals and subtasks are transferable to new tasks by re-composition.
5.1 Verifying Subgoals as Selections
We first verify the theory in Sec. 3 that selections can be identified in data by the following CI tests: (1) whether St L at gt holds, (2) whether gt 1' a++1 gt holds, and (3) whether $1+1 1 g: S; a holds. Our empirical results provide an affirmative answer to all these questions, suggesting that selections do exist, and they can serve as subgoals.
Synthetic Color Dataset We follow the didactic experiment in Jiang et al. and construct color se- quences similarly. The dataset consists of repeating patterns of repetitive color, either with 3 or 10 steps of time lag in each pattern, which Fig. is an illus- tration. Details about the construction and the CI test results are elaborated in Appx. D.I.
Color 3
Color 10
Figure Patterns in Color-3 and -10.
Driving Dataset In the driving environment, there are two tasks to finish. As shown Fig. 5, both cars start at the left end, either facing up or down. The first task is to drive to the right end following the yellow path, while the second task is to follow the blue path. Each state is represented by tuple (xc y, 0) E R (coordinates and orientation), and each action is the angular velocity A R. We collected 100 trajectories in total, 50 for each task.
Cl test
Single-step
Multi-step
St
at
gt
0.00
0.00
gt
a
gr+1
le-7
0.008
St+1
gt
St, a,
0.40
0.65
Figure 5: Two tasks in Driving environment.
Table 1: P-values for CI tests Driving.
CI test results We show our results in Tab. I for the Driving dataset. In short, the answer for (1) whether St J a, ge (2) whether gi A St. and (3) S 1 g: St, a are all yes. For Driving Dataset, as is shown in Tab. 1, single-step denotes that we are treating S; at different time steps as different variables, and calculate the mean of p-values across time. Multi-step denotes that we are aggregating at different times steps together as single variable, and sample a subset of data every time, calculate the mean of the p-values across subsets. Note that the variance p-values in Tab. are less than 0.001, so we only record the mean here. By testing CI condition (1) as sufficient, and CI conditions (1) and (2) together as both necessary and sufficient, we verify that subgoals are essentially selection variables. Moreover, by CI condition (3), we also validate that the next state should be determined solely by the previous state and action, and particularly depends on the physical transition, but not influenced by the subgoals.
5.2 Evaluating the Effectiveness of Seq-NMF
Next, we evaluate the seq-NMF in recovering selections and discovering subtasks. Results in the Driving dataset are in Fig. 6 while those for Color are in Appx. D.4. The y-axis represents the dominance of each subtask in explaining the whole sequence. We plot two sequences and each lasts for around 110 steps. Our algorithm finds the "crossing point" and automatically partitions every trajectory into three sub- tasks (before reaching the first crossing point, the middle, and after reaching the second crossing point).
0.8
Stftaak
Sutraak
Sutraak-
100
120
140
1200
150
200
Time stop
Figure 6: Seq-NMF result on Driving.

---

5.3 Transfering to New Tasks
Kitchen Dataset The evaluation of the imitation learning performed on challenging Kitchen environment from D4RL(Fu et al. [2020]) benchmark. Each agent required to perform a sequence of subtasks including manipulating the microwave, kettle, cabinet, etc. Each task is composed of subtasks and the composition is unknown to the agent. We use the demonstrations provided by (Gupta et al. [2019]) for reproducibility, which only contain state and action pairs but reward. We the demonstrations of two tasks for training, and require the agent to accomplish target task that has different composition of subtasks than any one of the demonstrations, but the units of subtasks have been performed. Detailed description is included in Appx. D.2.
Figure 7: Kitchen environment
Baseline methods We compare our method with the following stte-off-hhe-at (SOTA) hierarchical imitation learning methods to prove its efficacy: (1) H-AIRL (Chen et al. [2023]) is a variant of the
meta-hierarchical imitation learning method proposed in Chen et al. [2023] that doesn't incorporate the task context, and is learning an option-based hierarchical policy just as in our setting. (2) Directed-info GShLmSharma et al. [2019b]) and (3) Option-GAlL (Jing al. [2021]) are two other competitive baselines that are proven to be effective in solving multi-task imitation learning.
175
150
Retann
100
40
Ous DI-GAIL Option-GAIL
25
20
-20
400
600
200
800
Tetal Steps (x10*)
Total Stops (x10*)
(a) Transferring kettle + bottom burner
(b) Transferring microwave - bottom burner + top burner + slide cabinet.
+ top burner + hinge cabinet.
Figure Results for solving new tasks the Kitchen environment w.r.t. training steps.
Results on new tasks with 4-subtasks We show the results on new tasks with 4-subtasks in Fig. We use the episodic accumulated return as the metric. Training is repeated 5 times randomly for each algorithm, with the mean shown as solid line and the standard deviation as a shaded area. We observe that our method outperforms all the baselines in both tasks. The agent trained with the selection-based subtasks can quickly adapt to the new task with different subtask compositions, achieving higher return. The results show that our method can effectively transfer the subtask knowledge learned from the demonstrations to new tasks.
Results on new tasks with 5-subtasks (generalized) We conduct additional experiments by considering distribu- tion shift problem that involves longer- horizon tasks, and plot the results in Fig. 9. Specifically, under the Kitchen environment, we keep the same train- ing set of tasks (each task is composed of sequential manipulation subtasks), and tests the method's 's generalizability to new task with different permutation of one more subtask, i.e. 5 sequential manipulation subtasks. Such generaliza-
250
200
150
Ours
50
DI-GAIL
H--IRL
Option GAIL
0
200
400
600
800
Total Steps (x10*)
Figure 9: (Generalized) Transfering to microwave -> bot- tom burner - top burner - slide cabinet hinge cabinet. tion to longer- longer-horizon horizon tasks is not taken into consideration by some of the existing works (Chen

---

et al. [2023]). and our empirical results show that our formulations are able to deal with such more challenging distribution shift problem.
By detecting the existence of selections in data, and recognizing them as subgoals, we can effectively learn the subtasks that are useful for future use-subg are the selection indicators for the subtask patterns, Such procedure is in contrast to other methods that neglect the real structure in data and merely adopt a maximization of the likelihood objective, which might be misleading for recovering meaningful patterns. Also, their joint optimization of up to five networks (Chen et al. [2023]) makes model instability major concern. Our method, on the other hand, learns subgoals directly from demonstrations, disentangling it from the policy training. This makes the learning more stable and interpretable, and the learned subtasks can easily adapt to new scenarios.
Conclusion
In short, we target at subtask decomposing problem. While previous research has not sufficiently analyze the concept of subtasks, which might lead to an inappropriate inference of subgoals, we propose to view subtasks outcomes of selections. We first verify the existence of selection variables in the data based on our theory. From this perspective, we recognize subgoals as selections and develop a sequential non-negative matrix factorization (seq-NMF) method for subtask discovery. We rigorously evaluate the algorithm on various tasks and demonstrate the existence of selection and the effectiveness of the method. Finally, our empirical results in challenging multi-task imitation learning setting further show that the learned subtasks significantly enhance generalization to new tasks, suggesting exciting directions on uncovering the causal process in the data, also showing a new perspective on improving the transferability of policy.
The main limitation in this work lies in that we not yet able to deal with the scenarios where there might be multiple factors at work, such as the case where there are both underlying confounders and selections. Confounders might be related to other types of distribution shift, e.g. change in the system dynamics, robot embodiment, etc. In future work, we will investigate the causal process in other contexts, and aiming at providing more general framework for subtask discovery.
Acknowledgments and Disclosure of Funding
We thank Clark Glymour, Xiangchen Song, Biwei Huang, Zeyu Tang, Guangyi Chen, Jiaqi Sun, Yiding Jiang, for the valuable discussions and their support, and we thank the anonymous reviewers for their suggestions. This material is based upon work supported by NSF Award No. 2229881, 1, AI Institute for Societal Decision Making (AI-SDM), the National Institutes of Health (NIH) under Contract 101HL115900, and grants from Salesforce, Apple Inc., Quris AI, and Florin Court Capital. Funding to attend this conference was provided by the CMU GSA/Provost Conference Funding.
References
Dynamical Deep Generative Latent Modeling of 3D Skeletal Motion I International Journal of Computer Vision. https:/link-ssringer-com www.../m..dm..oc..or........
1007/s 2222266232226662222222222222000- 8.
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In ICML, 2004.
Michael Bain and Claude Sammut. A framework for bchavioural cloning. In Machine Intelligence 15, 1995.
Christopher Baldassano, Janice Chen, Asieh Zadbood. Jonathan Pillow, Uri Hasson, and Kenneth A Norman. Discovering event structure in continuous narrative perception and memory. Neuron, 95(3):709-721, 2017.
Elias Bareinboim, Jin Tian, and Judea Pearl. Recovering from selection bias in causal and statistical inference. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 28, 2014.
Ritwik Bera, Vinicius G. Goecks, Gregory M. Gremillion, John Valasck, and Nicholas R. Waytowich. PODNet: A Neural Network for Discovery of Plannable Options, February 2020.
Jiayu Chen, Dipesh Tamboli, Tian Lan, and Vaneet Aggarwal. Multi-task Hierarchical Adversarial Inverse Reinforcement Learning. In Proceedings of the 40th International Conference on Machine Learning, pages 4895 4920. PMLR, July 2023.
10

---

Leihao Chen. Onno Zoeter, and Joris M Mooij. Modeling latent selection with structural causal models. arXiv preprint arXiv:2401.06925, 2024.
Liangzhe Chen, Sorour E. Amiri, and B. Aditya Prakash. Automatic Segmentation of Data Sequences. AAAI, 32 (1), April 2018. ISSN 2374-3468, 115995399, doi: 10.1609/aaai, v3211.11815.
Juan D. Correa, Jin Tian, and Elias Bareinboim. Identification causal effects in presence of selection bias. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages
2019.
Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion imitation learning. Advances in neural information processing systems, 32, 2019.
Mísha Denil, Sergio Gómez Colmenarejo, Serkan Cabi, David Saxton, and Nando De Freitas. Programmable agents. arXiv preprint arXiv: arXiv: /ee../006.0...00 06333, 2017.
Frederick Eberhardt and Richard Scheines. Interventions and causal inference. Philosophy of science, 74(5): 981-995, 2007.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. ICML, 2016.
Carlos Florensa, Duan, and Pieter Abbeel. Stochastic neural networks for hierarchical reinforcement learning. arXiv preprint arXiv:/704.0000 2017.
Patrick Forré and Joris M Mooij. Causal calculus in the presence of cycles, latent confounders and selection bias. In Uncertainty in Artificial Intelligence, pages 71-80. PMLR, 2020.
Roy Fox, Sanjay Krishnan, Ion Stoica, and Ken Goldberg. Multi-Level Discovery of Deep Options, October 2017.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforcement learning. In ICLR, 2018.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, Sergey Levine, D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Jonas Gehring, Gabriel Synnaeve. Andreas Krause, and Nicolas Usunier. Hierarchical skills for efficient exploration. Advances in Neural Information Processing Systems, 34:1 1553--1166,, 2021.
Thomas L Griffiths, Frederick Callaway, Michael B Chang, Erin Grant, Paul M Krueger, and Falk Lieder. Doing more with less: meta-reasee reasoning and meta-learning humans and machines. Current Opinion in Behavioral Sciences, 29:24-30, 2019.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay Policy Learning: Solving Hori-oo Tasks via Imitation and Reinforcement Learning, October 2019.
Danijar Hafner, Kuang-Huei Lee, Ian Fischer, and Pieter Abbeel. Deep Hierarchical Planning from Pixels, June 2022.
James J Heckman. Sample selection hias as specification emor. Econometrica: Journal of the econometric society, pages 153-161, 1979.
Miguel A Hernán, Sonia Bermández-Diaz, and James M Robins. A structural approach to selection bias. Epidemiology, pages 615-625, 2004.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In NeurlPS, 2016.
Ahmed Hussein, Mohamed Medhat Gaber. Eyad Elyan, and Chrisina Jayne, Imitation learning: survey of learning methods. ACM Computing Surveys (CSUR), 5055):1-35, 2017.
Yiding Jiang, Evan Zheran Liu, Benjamin Eysenbach, J Zico Kolter, and Chelsea Finn. Learning Options Compression. page 26.
Mingxuan Jing, Wenbing Huang. Fuchun Sun, Xiaojian Ma. Tao Kong, Chuang Gan, and Lei Li. Adversarial Option- Aware Hicrarchical Imitation Learning, June 2021.
Taesup Kim, Sungjin Ahn, and Yoshua Bengio. Variational temporal abstraction. Advances in neural information processing systems, 32, 2019.

---

Thomas Kipf, Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Sanchez-Gonzalez, Edward Grefenstette. Pushmeet Kohli, and Peter Battaglia. CompILE: Compositional Imitation Learning and Execution, May 2019.
Sanjay Krishnan, Roy Fox, Ion Stoica, and Ken Goldherg. DDCO: Discovery of Deep Continuous Options for Robot Learning from Demonstrations, October 2017.
Yann LeCun. A Path Towards Autonomous Machine Intelligence Version 0.9.2, 2022-06-27. 2022.
Daniel D. Lee and H. Schastian Scung. Learning the parts of objects by non-negative matrix factorization.
Nature, 401 (6755);88877) Octoher 1999. ISSN 1476-4687. doi: 10. 1033/4454
Sang-Hy Lee. Learning compound tasks without task-specific knowledge via imitation and self- superviuee learning. 2020.
Tabitha E Lcc, Jialiang Alan Zhao, Amrita S Sawhney, Siddharth Girdhar, and Oliver Kroemer. Causal reasoning in simulation for structure and transfer learning of robot manipulation policies. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 4776-4782. IEEE, 2021.
Sergey Levine. Reinforcement Learning and Control Probabilistic Inference: Tutorial and Review, May 2018.
Emily L Mackevicius, Andrew H Bahle, Alex H Williams, Shijic Gu, Natalia Denisenko, Mark S Goldman, and Michale S Fec. Unsupervised discovery of temporal sequences in high- ditheminnniionl datascts, with applications to neuroscience. eLife, 8:e38471, February 2019. ISSN 2050-084X. doi: 10.7554/eLifc.38411.
Yasuko Matsubara, Yasushi Sakurai, and Christos Faloutsos. AutoPlait: Automatic mining of co-evolving time sequences. In Proceedings of the 2014 ACM SIGMOD International Conference on Management Data, SIGMOD '14, pages 904, New York, NY, USA, June 2014, Association for Computing Machinery. ISBN 978-1 178-1-4503-2376-11 doi:
Amy McGovern and Andrew G Barto. Automatic Discovery of Subgoals in Reinforcement Learning using Diverse Density. 2001.
Kevin Patrick Murphy. Dynamic bayesian networks: representation, inference and learning. University of Californa, Berkeley, 2002.
Andrew Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning. In JCML, 2000.
Sujoy Paul, Jeroen Vanbaar, Amit Soy-Chowdhury Learning from Trajectories via Subgoal Discovery. In Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Judea Pearl. Causality. Cambridge university press. 2009.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint /rXiv:/707.06::: 2017.
Seyed Kamyar Seyed Ghasemipour, Shixiang Shane Gu, and Richard Zemel. Smile: Scalable meta inverse reinforcement learning through context-conditional policies. Advances Neural Information Processing Systems, 32, 2019.
Arjun Sharma, Mohit Sharma, Nicholas Rhinehart, and Kris Kitani. DIRECTED-INFO GAII:: LEARN- ING HIERARCHICAL POLICIES FROM UNSEGMENTED DEMONSTRATIONS USING DIRECTED INFORMATION. 2019a.
Arjun Sharma, Mohit Sharma, Nicholas Rhinehart, and Kris M. Kitani. Directed-Info GAIL: Learning Hierar- chical Policies from Unsegmented Demonstrations using Directed Information, March 2019b.
Özgür Şimşek and Andrew G. Barto. Using relative novelty to identify useful temporal abstractions in reinforce- ment learning. In Twenty-First International Conference Machine Learning ICML '04, page 95, Banff, Alberta, Canada, 2004. ACM Press. doi: 01144/11153301111111
Özgür Şimşek, Alicia P. Walfe, and Andrew G. Barto. Identifying useful subgoals in reinforcement learning by local graph partitioning. In Proceedings of the 22nd International Conference Machine Learning ICML '05. pages 816-823, Boun, Germany, 2005, ACM Press. ISBN 978- 1-59593--800111 doi: 10.1145/110231. 1102454.
Paris Smaragdis. Non-negetive matrix factor deconvolution: extraction of multiple sound sources from mono- phonic inputs. In Independent Component Analysis and Blind Signal Separation: Fifth International Conference, ICA 2004, Granada, Spain, September 22 24,2,, 2004. Proceedings 5, pages 494 499. Springer, 2004.
12

---

Paris Smaragdis. Convolutive speech bases and their application to supervised speech separation. IEEE Transactions Audio, Speech, and Language Processing, 15(1):1-12, 2006.
Peter Spirtes, Christopher Meek, and Thomas Richardson. Causal inference in the presence of latent variables and selection bias. In Proceedings of the Eleventh conference Uncertainty artificial intelligence, pages 499-506, 1995.
Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 1 12(1122:181-211, 1999.
Daniel Tanneherg, Kai Ploeger, Elmar Rueckert, and Jan Peters. SKID RAW: Skill Discovery From Raw Trajectories. IEEE Robot. Autom. Lett., (0044666-4703, July 2021. ISSN 2377-3766, 2377-3774, doi: R0.10000/AA.22213066889
Sebastian Thrun and Joseph O'Sullivan. Discovering structure in multiple learning tasks: The te algorithm. In ICML, volume 96, pages 489-497. Citeseer, 1996.
Philip Versteeg, Joris Mooij, and Cheng Zhang. Local constraint based causal discovery under selection bias. In Conference on Causal Learning and Reasoning, pages 840-860. PMLR, 2022.
Hao Wang. Shunguo Fan, Jinhua Song, Yang Gao, and Xingguo Chen, Reinforcement learning transfer based on subgoal discovery and subtask similarity. IEEE/CAA Journal of Automatica Sinica, 1(3):257-266, July 2014. ISSN 2329-9274. doi: 10.1 10.100//4..221147044663
Lantao Yu, Tianhe Yu, Chelsea Finn, and Stefano Ermon. Meta-inverse reinforcement learning with probabilistic context variables. Septemher 2019.
Jeffrey Zacks, Barbara Tversky, and Gowri Iyer. Perceiving, remembering, and communicating structure in events. Journal of experimental psychology: General, 10299, 2001.
Jesse Zhang, Karl Pertsch, Jicfan Yang, and Joseph Lim. Minimum description length skills for accelerated reinforcement learning. Self-Supervision for Reinforcement Learning Workshop-ICIA volume 2021, 2021.
Jiji Zhang. On the completeness of orientation rules for causal discovery the presence of latent confounders and selection bias. Artificial Intelligence, 177(11917) 2008
Kun Zhang, Jiji Zhang, Biwei Huang, Bernhard Schölkopf, and Clark Glymour. On the identifiability and estimation functional causal models in the presence outcome-dependent selection. In UAI, 2016.
Yujia Zheng, Zeyu Tang. Yiwen Qiu, Bernhard Schölkopf, Kun Zhang. Detecting and identifying selection structure sequential data. ICMI., 2024.
Yifeng Zhu, Peter Stone, and Yuke Zhu. Bottom-up skill discovery from unsegmented demonstrations for long-horizon robot manipulation. IEEE Robotics and Automation Letters, 7(2)4126 4133, 2022.
Brian Ziebart, Andrew Maas, Andrew Bagnell, and Anind Dey. Maximum entropy inverse reinforcement learning. In AAAI, 2008
13

---

A Related Works
Skill or option discovery While there multiple terms across several lines of research, e.g. skills options, they all refer the same problem and focus on learning set of subtasks that be used solve complex task via hierarchical planning. The problem also our main focus here. The subgoals are defined heuristic as bottle-neck regions in tasks, for example, by finding states that are visited frequently by the expert ( McGovem and Barto [2001], Şimçek and Barto [2004]), finding minimal cut of the state transition graph Simzek al. [2005),, or by clustering and other similarity measures Wang al. [2014], Paul et al. [2019]). Recently, the skills are learned through modeling skills as latent variables and maximizing likelihood. For example, there are deep option learning methods Krishnan et al. [2017]. Fox et al. [2017]. Bera et al. [2020]) hierarchical clustering ( Zhu et al. [2022]). information theory-induced methods ( Sharma et al. [2019a]. Lee [2020]), segmentation and abstractions modeling Kipf et al. [2019]. Tanneberg et al. [2021]] ID, and methods that leverage other criteria such minimal description length Zhang et al. [2021]. Jiang et al.). However, these methods not provide clear definition of skills or options, and the learned partitions face challenges their interpretability. We can address the problem by identifying sub-goal nub-goals as selections that enhance interpretability.
22
21
22
Z3
bu
b1
21
22
22
31
.22
22
24
(a) Graphical model in DI-GAIL (Sharma et al. [2019a])
(b) Graphical model in ComPILE (Kipf et al. (2019]), where be represents the "boundary"
Figure 10: Two examples of the graphical models used in other literature.
For example, Fig. 10 shows two different graphical models proposed in DI-GAIL and ComPILE that are representative. State and action are aggregated into one single variable aı fsi, az) and the subtask variable is denoted as z, in both cases. Their corresponding methods do not take into consideration of the true generation process, thus might be unable to reflect the real data structure and lead to biased inference.
Multi-task IL (MIL) Current MIL methods aim at training policy that can be conveniently adapted to multiple tasks, often done by incorporating task context variable to condition on, such as (Seyed Ghasemipour al. [2019], Yu al. [2019]). However, the task variable not always available and these methods neglect the hierarchical structure in accomplishing the task. Then Chen et al. [2023] takes into account skill/option framework to build a hierarchical policy along with the context variable. However, this approach inherently suffers from optimization problems and require demanding hyperparameter tuning process, due to the overall complexity of the framework, while provide solution with disentangled optimization procedures that is efficient to deploy.
We also build rich literature the following topics, which our insights provide important implications:
Causal RL and IL This line of research introduces causal ideas to improve IL problems, by conditioning the imitation policy on the causal parents (De Haan et al. [2019]) by learning compact causal structure between the states, actions and rewards (Lee et al. [2021]). However, the selection structure under this context not sufficiently investigated, and we provide this new perspective.
Selection bias In causal literature, previous works mainly focus on understanding selection as a distortion of data, and aim at alleviate the selection effect (Spirtes et al. [1995]. Hernán et al. [2004], Zhang [2008]. Bareinboim al. [2014], Zhang Chen [2024]). However, though controlling the selection bias is important, they fail to view the selection as a source of information that can facilitate learning and inferencing. Zheng et al. [2024] propose methods to discover the selection structure in the sequential data. We explore the underlying selection structure in imitation learning settings and fill the missing gaps in comprehensively understanding the selection process in the data.
[2016]. Correa et al. [2019], Forré and Mooij [2020]. Versteeg al. [2022],
14

---

B Proofs
B.1 Proof of Proposition
Proposition 1. (Sufficient condition) Assuming that the graphical representation is Markov and faithful to the measured data, st aa, d,, then d, is selection variable, i.e., de == gt, under the assumption that:
1. (confounder, selection, and intermediate nodes can not co-exist) At each time step, de can only be one of (et, gt. m.). J. (For relaxation of this assumption, see Appx. B.4).
2. (consistency in a time series) At every time step. de plays the same role as one of {ce, gt. me).
We can use d-separation Pearl [2009] to distinguish the d, == ge case from the other two kinds of dependencies.
Va as de if de g: because there is path + gt at, y de- But if d, de m, because 8t at Ct in the confounder case and St as m, in the intermediate node case, then 8t 1 a: de.
B.2 Proof of Proposition
Proposition 2. (Necessary and sufficient condition) di is a selection variable (d, = gt) if and only if + ar de and d, 1 a d1+1-
In Prop. I, we have already proved that St y as de is sufficient to show that d. = ge. For the necessary part, we need to show that if di is selection, then entails that d, 1 a de+1. Because ge Ka a 1t+1. by d-separation, g: X a gt+1.
B.3 Proof of Proposition
Proposition 3. (Necessary condition) If de is selection variable (d, = gt). then St+11 1 g: St, ac. (Such independency does not hold true for confounders case which is discussed in Appx. B.3)
When conditioning on St and ag, then gt and st+1 are d-separated. But for a confounder case, C₂ and St-1 are not. The difference is essentially because that we have - gt+1 <- Se41 where there is unshielded collider g between gt and S that blocks the path, while ce c++1 S not blocked. Therefore, Se+1
Ht St: at, but St-1 1 K ct St, at.
B.4 Relaxation of Assumptions
We extend our theory in Sec. 3.1 by two types of relaxations of our assumptions. We provide more general graphical model that represents the data generation process in Fig. 11. There are two parts that will be discussed this section: (1) the inclusion co-existence ge. Ct and m. (2) the higher-order structure behind those variables (e.g. the set of higher-order confounders and selections are denoted as Uc and Us, respectively, in Fig. 11).
B.4.1 Co-existance of Ct, and m
first point we assert this section that is possible di does not take only one (Ce, Et. me). the main paper, we only discuss the pure case for simplicity, however, is also possible that d, can take multiple characters of {Ce, gt, m.). We argue that even if the data is generated under multiple hidden variables. for example, both confounders and selections selection still holds, and propose the following Prop. 4.
play, our sufficiency criteria for recognizing the existence of
Proposition 4. (Sufficient condition under multiple types of hidden variables) Assuming that the graphical representation is Markov and faithful the measured data, s, / a, de, then g: must exist in the hidden structure, i.c., E di (we use di to denote the set of hidden variables), under the modified assumptions that
(one or more structures confounder, selection, intermediate node) At each time step, de can be a subset of {ct, St, m. i.c. de C {ct, St, m.).
2. (consistency in time series) Similar in Prop. 1.
Proof In the graphical model of Fig. 11, both confounders ((c,) and selections ([g.) are exhibited. Then conditioning on de does not change the key CI criteria we proposed in Prop. 1: st Ka, de, because long as the path St + g: + at is d-connected (when two variables are not d-separated, they are di d-connected) when conditioning on gt. the dependency between s; and a: is preserved, regardless of whether there are confounders or intermediate nodes. If g: is not present, i.c. d; S {Cr, m. then there is no dependency between and a when conditioning on dt.
15

---

Implications Such a relaxation that includes multiple types of hidden variables is reasonable in real-world scenarios. While the observed data generated by cobbool-condiioone policy, is still possible that there some hidden confounder that influences the joint data distribution. For example, when you are selecting which way to go to school, you may be influenced by the weather. Given the same state and action, say, you decide to walk out of home, changing the weather will change the distribution of the next state: it causes you to act otherwisc, e.g. stay home. On the other hand, changing your actions would not influence the weather. In this case, the weather is confounder that is distinct from the subgoal, but still influences the data generation process. The main proposition Prop. one that provides the sufficient condition to identify selections, and Prop. 4 implies that we could not eliminate the possibility of confounders and intermediates at play together with selections. Learning the hidden confounder or intermediate structures along with selections is another exciting topic that we leave for future work.
B.4.2 Higher-order Structure Behind Data
In the main paper, we assume that the true causal graph of the data takes one of the three scenarios in Fig. 2, with no hidden high-level structures. A second relaxation we make the inclusion of higher order structures. We argue that if there are higher-order structures in the data, and loosen the assumption of a direct adjacency between de and de+1, then Prop. 1 still holds true.
Uc
C1
C2
C3
Illustration Take the case of UG (the set of higher-order goals) for example. The idea that there potentially could be hierarchical subgoals is natural: we take the example of commuting to New York in Sec. 1 as an illustration. The first subtask walking out of the house is further decomposed as standing up, grabbing the luggage ge and walking to the door. Each one of the four behav- iors can also be understood as smaller subtasks that involve multiple more primitive actions: for walking, you need move your left leg, right hand, then right leg, left hand, etc. On a grander scale, commut ing to New York could possibly be a part of a bigger task of going on business trip: commuting to New York, going to a meeting,
81
82
83
a1
Q2
a3
91
92
93
9T
In this example, walk walking to the door would be the first level subgoals (g)). Then it contributes to a higher level of subgoal: commuting to New York, which we represent gi E Ua. Then this second-order subgoal further contributes to third-order subgoal. the final goal going on a business trip, which we represent as gi Ug. All these subgoals has pattern of g: + ge g:- + ..., where the higher levels of subgoals of lower levels of subgoals (achieving the lower-order goals causes us to achieve the higher-order goals). The consecutive subgoals are dependent of each other because of conditioning on higher-order subgoals: gt 1 gt g.. because conditioning d-connected. Noting that this relaxation of high-order structures is practical in the real-world setting, we provide a corresponding proposition in Prop.::
giving presentation.
UG
Figure 11: Graphical model with relaxed as- sumptions. (1) We allow the co-existance of gt and c, (similarly for ge and m ). (2) We also assume there are potential higher-order underlying confounders Uc in the data that create the dependencies between Ct and Ct+1. and underlying selections Us in the data that create the dependencies between ge and Ct+1-
always seleetive
g. will make gi and g
Proposition 5. (Higher-order structure) Assuming that there are higher-order structures behind ct2 St and mes and direct edges between de and d are not necessarily present. Then Prop. is still valid.
The oigher-order structure does not change the key CI criteria proposed in Prop. 1 and Prop. 2. When conditioning de, if de is confounder intermediate, then it d-separates st and a, and we will obtain St P at. Only when de is selection/subgoal, would a collider on the path between St and at. Thus, only in the case of a selection, we will obtain St a Now we have proved that the sufficient condition in (1) still holds.
Proof
This relaxation for the high-order structure implies that the direct adjacency between de and d is not necessary for the sufficiency criteria to hold.
16

---

C
Algorithms
C.1 Algorithm for Seq-NMF
Algorithm SeqNMF for learning subtasks
Input: Data matrix X, number of subtasks J, maximum time delay in subtask L, regularization strength Abins A1, Aaim
Output:
O, H
Normalized X into [0, 1]
Initialize O and randomly
Set =1 while (i < maxiter) and (A cost 2 tolerance) do
Update H using multiplicative update from Eqn. (4)
Renormaliie so maximum value H is 1, and normalize O accordingly. Update O using multiplicative update from Eqn. (4)
i=i+]
end while
Set Atin, À1 Asim to zero, do one final unregularized multiplicative update of O and Return: O, H
C.2 Algorithms for Transfering to New Tasks
Algorithm 2 Imitation Learning in new tasks with subgoals
Input: Subtask patterns O, expert demonstrations (St, ar, "..
Initialize the policy Tg. discriminator Dy- for each training episode do
Generate M trajectories {8., ār, B.} by exploring in the target environment with Ty and with
Algo. 3.
Update Dg by minimizing LIL based on {st, a,, ge and {st, āg gr}
Train T. by PPO (Schulman al. [2017]), based on {šr, a, g: and Dg which defines the reward RTL-
end for
Algorithm Executing subgoal-conditioned policy
Input: subgoal-conditionel policy (a St, ge), subtask patterns O, initial state &p. (O denotes the last state in subtask j that has value and can has a smaller index than L.)
Output: Collected sequence {š, â, g.).
Procedure Executing
policy Te
while task is not done do
Query the potential initial subgoal go arg miny 1|50 0 ||.
Set current subgoal
go-
while subtask has terminated do
Take action a₂ = arg maxa we (ā st,
Observe next state 8t+1
Terminate if ||s o 1 Se
end while
Query the most potential next subgoal g = arg min, ||s 0(1)
end while
Return: Collected sequence (8,, ā 8
17

---

D Experimental Details
D.I Data Generation for Synthetic Color Dataset
For Color-3, there are 3 types of patterns: 3 consecutive red, 3 consecutive yellow, and 3 consecutive blue. Each pattern is either generated independently (Simple) or correlated with the previous pattern (Conditional). For Color-10, there are 2 patterns in the dataset: consecutive red + consecutive yellow 4 consecutive blue) and consecutive blue 3 consecutive yellow 4 consecutive red). Subgoals indicate the type of pattern to take. One sequence has a length of T = 300 for both Color-3 and Color- 10, and we collected 100 sequences for cach.
Because we want to conduct CI tests on the datasct, we need low-dimensional variables. Each state is represented by a variable with mean of 1, 3.3 for red, yellow, and blue respectively, and an additive Gaussian noise. Each subgoal is discrete variable indicating which color is selected for the next time step, taking one of the values of {0, 1, 2]. repeating the same value for steps plus noise. Each action the difference between the current state and the subgoal, i.c, at = St gt. then taking the action means that the next state is updated with $141 = St + 0; + where ~ (0, 0.01). In Color-3 (Simple). each pattern is selected independently, and in Color-3 (Conditional), also repeated for 3 steps. The conditional comes by re-coloring the yellow pattern generated in the Color-3 (Simple) dataset to purple if its precedent color is yellow or blue.
added
more color purple to it, represented by mean value of 4. Every pattern
D.2 Data Configuration in Kitchen
Transferring to new tasks with 4 subtasks We use kettle + bottom burner + top burner + slide cabinet and microwave + bottom burner + top burner + hinge cabinet in the demonstrations for training, and use two different tasks for testing: (a) kettle + bott burner + top burner + hinge cabinet, and (b) mi crowave + bottom burner + top burner slide cabinet. The difference in the target tasks' subtask composition is marked in blue.
Transferring to new tasks with 5-subtasks (generalized) To answer the question of whether the algorithm able to generalize a longer horizon, we it on a target task of subtasks to make it a more challenging scenario. Specifically, we use bottom burner + top burner + s slide cabinet + hinge cabinet nettnet,,, and microwave + bottom burner + top burner + hinge cabinet in the demonstrations for training. For testing, we require the agent to manipulate: mi crowave - bottom burner top burner + slide cabinet -> hinge cabinet in order.
D.3
CI Tests on Color Dataset.
For Color Dataset, we calculate the mean of p-values across sequences, as is shown in Tab. 2. The answer for (1) whether St J. at ge (2) whether gt Y ac+1 ge. and (3) St+1 Jl g: St, a, are all yes.
CI test
Color-3
Color-10
(1) St 1 at gt
0.00
0.00
(2) gt 1 a
g++1
0.00
0.00
(3) 81+1 1 g:
St, ag
0.99
0.99
Table 2: P-values for CI tests in Color-3 and -10.
D.4 Supplementary Results for Seq-NMF
Visualization of the learned matrices. The direct visualization of matrix H and O of NcF on all three datasets are shown in Fig. 13. In cach subfigure, matrix H is shown at the top, with cach "row" of H representing subgoal (0 or 1, indicating whether a subgoal is selected not). and matrix O is shown the left, with each "column" of O representing subtask. Different subtasks are marked with different colors. The matrix in middle is the convolutional product of O and H. which is the trajectory matrix
Fig. (a), we show the results on the Color-10 10 dataset. Two significant subtask patterns are shown on the left (Ired, yellow, blue) and (blue, yellow, red]), each pattern lasts for 10 steps. On the top. H has spikes that indicate these patterns' corresponding appearance in the data matrix. Fig. 13 (b) shows the results on the Driving dataset and (c) shows the results on the Kitchen dataset, subtasks and subgoals are shown similarly.
18

---

(a) seq- scq-NMF result Color-10.
(b) seq-NMF result Driving
Figure 12: Results on the Kitchen dataset on new tasks.
(a) seq-NMF result of Kitchen
Figure 13: Results the Kitchen dataset tasks.
Conversion from to G We convert learned binary indicator matrix and subtask patterns O the subgoal G as follows:
Σl H gt. = E/=1 E1-1 Ozi HzC-o)
(6)
Visualization of the learned subtask patterns. We above equation to turn the learned matrices into "options". and mark every time-bin with the corresponding option. The result of Driving is shown in the main paper in Fig. 6. As for the Color-10 data which consists of two patterns, detect these subtask patterns co-evolving in the whole sequence, marked red blue, respectively. Each subtask pattern for 10 steps, which aligns with the true data generation process. When there are consecutive same subtask patterns, boundaries are not marked here but be found in the learned matrix H.
Subtask partition
0.8
0.6
0.4
Subtask Subtask
Numerical Evaluation. We compare the results our method on the Color-3 (Simple) and Color-3 (Conditional) datasets to VTA (Kim et al. [2019]) and LOVE (Jiang et al.) Tab. We show the precision, recall, and Fl methods recognizing the correct boundaries subtasks, which the indicator terms in H of our method. Our method not designed for high-dimen dimensional data, so we constructed our own dataset, but for the result of LOVE, directly take the results
20
40
60
80
100
Time step
of both
Figure 14: seq-NMF result on Color- 10. The dominance of each subtask in explaining 10 sequences.
19

---

reported in Jiang et al.. Note that, the raw data of color frames
be simply compressed to single dimension by various dimension reduction methods, after which we can use the compressed data to conduct our experiments.
Color-3 (Simple)
Color-3 (Conditional)
VTA
LOVE
Ours
VTA
LOVE
Ours
Precision
0.87 + 0.19
0.99 + 0.01
0.99 + 0.00
0.84 + 0.22
0.99 + 0.01
0.99 + 0.00
Recall
0.79 + 0.13
0.85 = 0.03
1.00 0.00
0.82 0.16
0.83 + 0.06
1.00 0.00
F1
0.82 + 0.13
0.91 + 0.02
0.99
+ 0.00
0.83 - 0.19
0.90 + 0.03
0.99 + 0.00
Table 3: Effect of our method on the Color-3 (Simple) and Color-3 (Conditional) datasets (5 seeds) in terms of the precision, recall and Fl score for recovering the correct boundaries of subtasks.
20

---

D.5 Hyperparameters
For the experiments in seq-NMF, we the following hyperparameters in Tab. 4. We use the same hyper- parameters Asm, λι, and Asa for all three datasets. Such hyperparameters penalty-wiid when their values are approximately at the same level. Our results chosen by comparing the three hyperparameters chosen this way are relatively robust such that one set of Asi. As. and Apia can all our three setting. As for the subtasks, we assume knowing the real number of subtasks for best performance, however we assume the algorithm also be able to detect subtask patterns as long number of pattern duration I. is also the approximated upper bound for the length of subtask, and the table provide one reasonably good estimation.
have upper bound estimation of the number of subtasks. The
Hyperparameter
Value
Asim
0.0001
0.001
Abin
0.01
Number of subtasks (Color
3
Number of subtasks (Color 10)
2
Number of subtasks (Driving)
5
Number of subtasks (Kitchen)
5
(Color 3)
3
L (Color 10)
10
(Driving)
40
L (Kitchen)
100
maxlter
300
start bos_loss iter
30
Table 4: Hyperparameters in seq-NMF
For the experiments in transfering new tasks, we use the following hyperparameters Tab, We use the same hyperparameters as Chen et al. [2023] fair comparison.
Hyperparameter
Value
4096
n__poch
2000
#Policy Config
activation
Relu
hidden dim (policy)
(256, 256)
log sa__clmp_pPo_ policy
(-20., 0.)
policy
3.c-4
dim_c
#PPO Config
hidden dim (critic)
(256, 256)
Ir critic
3.e-4
use_ use_gae
True
gamma
0.99
gae_tau
0.95
clip_eps
0.2
mini batch_ size
1024
#IL Config
hidden_die dim (discriminator)
(256, 256)
Table 5: Hyperparameters hierarchical IL
21

---

E
Relation to the Probablistic Inference View of Reinforcement Learning
$1
$2
$3
81
82
Q1
0/
92
9T
T1
T2
TT
Uc
R
(a)
(b)
Figure 15: In Fig (a), we show the DAG built upon our subgoal framework, and in Fig (b), we add reward nodes as "optimality variables" that are always conditioned to be true for control under the framework by (Levine [2018]).
standard task reinforcement learning can be characterized by reward function r (s., a.). Solving an RL problem is searching for policy that infers the most probable action sequence most probable action distributions from states, such that:
π arg ((a (r (sc. a.)]
(7)
This optimization procedure means that we seek policy that maximizes the expected reward. The question is. how we represent this problem in probabilistic graphical model (PGM), such that estimating (at st) from the PGM is equivalent solving the optimization problem maximizing the expected accumulated reward, as above?
Levine [2018] build the graphical model by laying out the states and actions that form the backbone, and add reward nodes as "optimality variables" that are conditioned to be true for an optimal policy, is shown in Fig. 15. The optimality variable 2, 1 denotes that the action is optimal, and 0 denotes that the action is not optimal. Then, would like to show that when conditioning on the optimality variable 2s, the optimal action sequence the one maximizes expected reward:
(r mI:T) x p (T,))T) = ( II (2, sc ax) (8t=1 Se, at)
= psi) II exp (r (si, a:)) ( (5141 s1, ar)
(8)
(si) II ( st)
exp Σ
(sc, ar)
when we choose the conditional distribution over to he (2, = a.) exp (r (s,, ae)). The first term the right-hand side of Eqn. completely determined hy the system's dynamics, which would be constant the system is deterministie. Then, the probability observing trajectory in proportion to exponential of the accumulated reward:
T (st, (s.
( ZI.T) 1[p(T) exp Er(s,
where indicator 1 [p(r) * 0] indicate that the trajectory is feasible one. Then, maximizing the posterior to get a sequence of optimal actions in this graphical model equivalent to maximizing the expected reward, which gives the distribution of (a. st, Zi 1).
Implication we take the selection variable to be scalar, then the posterior conditional policy we derive from the DAG in Fig. 15(a) gives (a. st, 9. 1) (can be understood as achieving single goal). Similar to the above analysis, such posterior gives the optimal policy standard reinforcement learning framework.
22

---

So, we successfully build the analogy of the selection structure identified in Sec. 3.1 to probabilistic inference view of reinforcement learning. It strengthens the connection between the subgoal framework and the reinforcement learning task maximizing the expected reward, and provides new perspective to understand the subtask discovery problem.
Subtask Ambiguity
There are two types of ambiguities that we care about here in subtask partition, as is shown in Fig. 16.
The Type A ambiguity (Fig. 16(a)) is that different subtask patterns share similar parts, as is marked with red. for both O₂ and Og during first two subtasks, because they are
This would result a high value of both highly correlated to the data matrix during that time period. Then multiplying it by H (with time shift within L) yields high value of this overall regularization term.
The Type B ambiguity (Fig. 16(b)) appears when subtask patterns are more granularly partitioned than desire, The pattern of light blue to dark bluc exhibited as one consistent pattern across data, but could be recognized different finer patterns shown the box. This contradict with our intuition that subtasks should be sparse enough, we can suppress this ambiguity by adding sparsity penalty Ri.
== 122.7T
1,2,
Sp....
0
O
Bp...
h0
P:
2)
(a) Type ambiguity: different subtask patterns share similar parts. Can suppressed by Raim
Type B ambiguity: subtask patterns granularly partitioned than desired. Can sup- pressed by R, = [H|]-
[I(O X)SH [|ss-
Figure 16: Different Types of ambiguities
G Derivation of Seq-NMF
At its core, NMF seeks approximate data matrix V the product of two smaller non-negative matrices, W base patterns or components and coefficient matrix that describes how these patterns combine approximate original data. This can he expressed mathematically as:
h
hg
R W H, where
[wi
W2
WK]
H
(10)
hg
G.1 Multiplicative Update for Standard NMF
In standard NME, where the optimization problem we try solve is:
(W. H arg min L(W, H), W.H
C(W H " パ V
(1I)
W W,H
23

---

the optimization problem is convex w.rl W, H separately, but not altogether. By alternatingly updating W and H. obtain the following gradient descent steps (additive update):
a
L(W H) = XH - XH
@w
(12)
C(W. H) = W X W X.
UH
Thus, gradient descent steps for W and H are:
W W
VWW
XH
XH
H H T W
W x'x
with 7/w and 7T being the learning rates. To avoid resulting in negative values, it is more desirable to use multiplicative updates. By setting:
W
WHH H
7J8
W WH
the equivalent multiplicative updates becomes (Lee and Seung [1999]):
XH
XH
W - W
W x
WH
XH
(13)
W
X
W X
H H
H
W WH
WTX
G.2 Multiplicative update for seq-NMF
Then, we incorporate the time element. Applying the above derivation for standard NMF for each time-bin delay f, we get the following update rules for convNMF (Smaragdis [2004]):
X
H
W. W.ex
X
XH
H Σ. W
W X
H
Σ, WT X
Where the operator + shifts matrix in the direction by timebins, which implies delay by E timebins, <- shifts matrix the <- direction by timebins. When L 1, convNMP reduced to standard NMF
G.3
Derivation of Regularizer Terms in Standard NMF
Adding the regularizer terms in standard NMF, we get the following additive update rules similar to Eqn. (12):
dC
OR
XH
XH
+
oW
OW
ØR
W
X
W
X
ØH
0H
We set:
W
T/W
XHU
3W
H
W X
9R
and turns into multiplicative update:
XH
W 6 W
1/W aW
W X
XH
+
SW
(14)
ac
W
X
H
H
TIH HH
H
W
X
PR OH
24

---

G.4
Derivation of Regularizer Terms in Seq-NMF (Ours)
Applying Eqn. (14) at each time-bin, the following multiplicative update rules for convNMF:
W W.W.W [
X H +aw a... .
(15)
W + X
H
W
2R SH
For each regularizer Rbin, R₁ and Ruim in our formulation in Sec. 3.3, separately derive the corresponding gradient used for multiplicative update (we change the matrix W to O, keeping aligned with the main part of paper).
BRain
Abin
H
To G
H H
Too
H)
0H
ØR1
A1
1)H
(16)
0H
0Ram
Avim (1 1J0 & XS
ÖH
URiin
d0
8R1
0
(17)
80
ÜRai
Amm X SH (1 -1-I)
+/
a
Therefore, the overall updating rules are:
X
H
O.. < O. X
XH KH. +deim X SHT(1
(1-1)
(18)
x
H X
0 Abin HOTPC To
④ To) A IJH + Awi (1 1J0 XS
H Hardware
All experiments were conducted on either NVIDIA L40, GeForce RTX 3080 Ti, or Mac MI chip with 16GB
of RAM.
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and introduction are clearly aligned with the theoretical (in Sec. 3 and 4) and experimental results (in Sec. 5) presented in the paper. The introduction provides a concise summary of what the reader can expect from the paper.
Guidelines:
The answer NA means that abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match thcorctical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper discusses its limitations in the last section, addressing a problem related to a relaxation of the key assumptions (the existence of multiple structures). It highlights the conditions under which the results are valid and the potential implications of these liriitations. We include brief discussion of this limitation in Appendix well.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create separate "Limitations" scction their paper.
The paper should point out any strong assumptions and how robust results are violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well- sell-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-ter system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms how they scale with dataset size.
applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
26

---

Justification: We include theoretical results that require assumptions and proofs in our main section (Sec. 3). We start by common assumptions in causal literature and analyze the possibility in imitation learning setting. We provided the complete proofs of the proposition made the main paper in the Appendix B.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, if they appear the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All information needed to reproduce the main experimental results is fully disclosed in our experiments section (Sec. 5). This includes detailed descriptions of the datasets used, the experimental setup, and the methods employed, including all baselines. The paper also provides specifics on hyperparameters, full algorithms, adopted codespaces in our Appendix. Our code is also provided in the supplemental material.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to hosted model (c.g., the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution a way to access this model for reproducing the results a way to reproduce the model (e.g.. with an open-source dataset or instructions for how to construct the dataset).
new model (e.g., large language model), then there should either be
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome to describe the particular way they provide for reprodueibility. In the case of closed-source models, it may he that access to the model is limited in some way (e.g.. to registered users). but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data code, with sufficient instructions to faithfully reproduce the main experimental results, as described supplemental material?
Answer: [Yes]
27

---

Justification: We provide open access to the code, along with README file to help reproduce the main experimental results in the supplementary material for all three parts of our experiments, reported Sec. 5.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https: httppppp.cppplll..ccppppc
guides CodeS ubbmii ssi ooooooooooiiooooo icy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" acceptable answer. Papers cannot rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmank).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / /// nips. cc/
publ // des/ CCdeeubmi ss on cyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Betting/Detaile
Question: Does the paper specify all training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The paper specifies all training and test details. For the handicraft data, described how to generate data; or we provide the information of benchmark that is used. We also include the choice of hyperparameters, the types of optimizers used, and the criteria for their selection in our Appendix.
Guidelines:
The answer means that the paper does not include experiments.
experimental setting should be presented in the core the paper to level of detail that necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [Yes]
The paper reports means and variances of the experiments. We made sure that for all the experiments conducted throughout the paper, averaged across at least 5 runs make sure that the results are reliable.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error hars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the enror bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
28

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful symmetric error bars that would yield results that are out of range (e.g. negative error rates). error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
show in tables or figures
8. Experiments Compute Resources
Question: each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer:
[Yes]
Justification: The compute resources used for the experiments are clearly described in the last section of the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn make into the paper).
9. Code Of Ethics
Question:
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code Ethics https:://///. //np:rr//eurrrr iss. cc/ publ c/Ethi csGui i del iine?
Answer: [Yes]
Justification:
The paper adheres the ethical guidelines set forth by NeurIPS. We ensured that research is conducted responsibly, with considerations for potential biases, fairness, and the broader impact on society.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., if there is a special considcration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
This paper presents work whose goal is to advance the field of Reinforcement Learning and Imitation Learning. There are many potential socictal consequences of our work, none which we feel must be specifically highlighted here.
Justification:
Guidelines:
The answer means that there is no socictal impact of the work performed.
the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance). fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there a direct path to any negative applications, the authors should point out. For example, the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
legitimate to point out that an improvement
29

---

The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have heen put in place for responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The work doesn't 't pose any risks. We either generated our own data or used justified reliable benchmark.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards allow for controlled use of the model, for example by requiring that adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code. data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We credited the author for the code package and benchmark dataset that have been used in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, possible, include URL.
The name of the license (c.g.. CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g.. website), the copyright and terms of service of that source should provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papersw swi theode. com dataseets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: include the code we used in experiments to supplemental materials. We also added the documentation to help understand the code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training. license, limitations.
30

---

The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, rememher to anonymize your assets (if applicahle). You can either create an anonymized URL include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least minimum wage country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: We do not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere to the NeurlPS Code Ethics and the guidelines for their institution.
initial submissions, not include any information that would break anonymity applica- ble), such as the institution conducting the review.
31


Reviews:
[{"review_id": "UHgbFKl4KP", "reviewer": "Reviewer_WHCR", "summary": "This paper studies the problem of decomposing expert trajectories (in the context of imitation learning) into sub-trajectories corresponding to subtasks. First, the authors introduce a causal framework to understand and explain what subtasks mean in this context. Subtasks are then defined to be variables that reduce uncertainty in the observed expert actions (or selections in the language of causal analysis). Motivated by this definition, a matrix factorization based task decomposition algorithm is presented. Experiments on multiple environments  demonstrate that the algorithm is effective at discovering subtasks from expert trajectories.", "questions": "- In the second line of Equation (1), should the \"forall\" and \"exists\" terms be switched? Should it read as \"for all sub-trajectories, there is exists a subtask...\" ?\n- The use of convolution in the matrix factorization algorithm suggests that it allows multiple subtasks to be \"active\" at a given step. Could you provide some intuition behind this?", "limitations": "I don't think there are any major limitations that need to be discussed in the paper.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- There are many novel aspects to the paper. The authors provide an argument to view subtasks as selection variables and this insight is used to develop the task decomposition algorithm presented in the paper. The task decomposition algorithm also seems novel and the use of matrix factorization here is very apt.\n- Experiments presented in the paper indicate that the proposed approach is effective at discovering useful subtasks. The experiments also show that the subtask decomposition algorithm enables learning policies (in the context of imitation learning) that generalize well to new and unseen tasks. In this context of transfer learning, the proposed approach outperforms many existing state-of-the-art methods.\n- I believe that the paper offers interesting insights relevant to the NeurIPS community and RL researchers working on hierarchical and compositional RL. These insights have the potential to inspire new directions of research.", "weaknesses": "One primary weakness is that some parts of the paper lack clarity. For instance, the mathematical objective in Equation (1) is not very clear  although the overall idea and intuitive definition is clear from the text. Defining the notations early on in the paper would make it much easier to read and understand.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are truly grateful for the time you have taken to review our work and for your insightful comments. Your valuable feedback has significantly enhanced the clarity of our manuscript.\n\n**Q1**: *“One primary weakness is that some parts of the paper lack clarity. For instance, the mathematical objective in Equation (1) is not very clear although the overall idea and intuitive definition is clear from the text. Defining the notations early on in the paper would make it much easier to read and understand.”*\n\n**A1**: We sincerely thank you for your great feedback and suggestions. Following your suggestion, we have carefully gone through the paper to make sure the notations are friendly. Specifically, in the updated version of the manuscript, we have defined the notations early before using them:\n> “**Problem Formulation and Notations**: Given the above context, we formulate the considered imitation learning problem as follows: we have a distribution of tasks $\\mathcal{P}_e(\\mathcal{T})$ and a corresponding set of expert trajectories $\\mathcal{D}=\\\\{\\tau_n\\\\}\\_{n=1}^N$, and we aim to let the imitater learn a policy that can be transferred to new tasks that follow a different distribution $\\mathcal{P}_i(\\mathcal{T})$. Each task sampled from $\\mathcal{P}\\_{\\cdot}(\\mathcal{T})$ should be generated by a MDP and composed of a sequence of option $\\{\\mathcal{O}_j, \\cdots\\}$, where $\\mathcal{O}_j =\\langle \\mathcal{I}_j, \\pi_j, \\beta_j \\rangle _j$. We use $\\mathbb{\\mathcal{O}}=\\bigcup\\_{j=1}^J \\mathcal{O}_j$ to denote all J options, and ${\\xi\\_{p}=\\\\{\\mathbf{s_t},\\mathbf{a_t}, ...\\\\}\\_{t=1}^{\\leq L} }$ as a sub-sequence of states and actions $(\\mathbf{s_t}, \\mathbf{a_t})$ from any trajectory $\\tau_n$. Each trajectory can be partitioned into sub-sequences ${\\xi\\_{p}}$ with maximum length L.”\n\n**Q2**: *“In the second line of Equation (1), should the \"forall\" and \"exists\" terms be switched? Should it read as \"for all sub-trajectories, there exists a subtask...\"?”*\n\n\n**A2**: We highly appreciate your insightful question. We agree with your interpretation of \"for all sub-trajectory, there exists a subtask such that the sub-trajectory is sampled from the option, ...\" and adjust the \"for all\" and \"exist\" notation accordingly. You are completely correct and it is better to present the notations in the way you suggested.\n\n**Q3**: *“The use of convolution in the matrix factorization algorithm suggests that it allows multiple subtasks to be \"active\" at a given step. Could you provide some intuition behind this?”*\n\n**A3**: We sincerely thank you for the great question. We fully agree with you that, in practice, it is possible to have undesired duplicate activations (which could be due to optimization errors like local optimal solutions, etc.), while the intuition is that subtasks should have distinct patterns and be distinctively activated.\n\nAt the same time, it might be worth noting that our formulation has already discouraged such situations. More specifically, given that our main insight is to view subgoals as binary selections, we enforce the learned coefficient matrix also to be $0/1$. Then if multiple subtasks are activated at the same time, the sparsity constraint $R_1$ as well as the redundancy constraint $R_\\text{sim}$ would penalize those competing factors. Therefore, the goal of our regularization is to forbid the cases that you kindly mentioned, of which the effect has been empirically verified in our experiments and we observed no duplicated activations.", "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": "", "extra_scores": {}}, {"review_id": "J0cea7KG2u", "reviewer": "Reviewer_imJV", "summary": "The paper studies the subtask decomposing problem. The paper proposes a formal definition of subtasks as the outcome of selections. The proposed seq-NMF is introduced to learn the subgoals and extract subtasks, conforming with the proposed theory. The experimental results show strong results on transferring to new tasks.", "questions": "See weaknesses.", "limitations": "Yes, the limitations are discussed in the paper.", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written and easy to follow.\n- The paper offers a formal definition of subtasks, to address the challenge of interpretability that has not been discussed in the prior works.\n- The insight of the definition of segmented subtasks should be consistent with the true generating process is novel and interesting.\n- The paper validates the proposed theory through the unseen long-horizon manipulation tasks and the results present a strong performance of the proposed method.", "weaknesses": "- The paper only considers state sequences to extract the subgoals and subtasks. In real-world settings, the states are not available. So it is not clear if the proposed theory is also applicable to the real world with various modalities, e.g., visual perception. It would be nice to explore and discuss the possibility.\n- The generalization of the proposed definition of the subtask is rather specific and the shift of task distribution is not significant. So the generalization of the proposed method might be less convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your thorough review and the constructive insights you provided, which will undoubtedly enhance the quality of our manuscript. In response to your comments, we have included several new discussions in the revised manuscript, and additional experiments, particularly addressing more generalized multi-task learning scenarios.\n\n\n\n\n**Q1**: *“The paper only considers state sequences to extract the subgoals and subtasks. In real-world settings, the states are not available. So it is not clear if the proposed theory is also applicable to the real world with various modalities, e.g., visual perception. It would be nice to explore and discuss the possibility.”*\n\n\n**A1**: The point that you raised is insightful and inspiring, and application to various data modalities will surely be an important aspect. It requires different techniques, and there are certainly a lot of practical issues to consider. In this work, we mainly focus on a theoretical analysis of the understanding of subtasks, and on rigorously validating our interpretation. \n\n\n* *Possible solutions*: When states are not available, it is still possible to extract useful representation from raw input data. For example, the CI test for high-dimensional data can be performed on the low-dimensional embedding level, followed by matrix decomposition and hierarchical imitation learning. In order to obtain the embeddings that are beneficial for downstream tasks, we essentially turn to a representation learning problem that learns a compact way to encode the sufficient feature for subtask identification and transfer, while mitigating the effect of redundant features. Various techniques can be applied to facilitate this representation extraction process [1] [2]. It is intuitively applicable that the proposed method adapts to more data modalities. We believe that our results will provide inspiration to the community for exploring more general scenarios, including multi-modalities data. It is a direction that is worth further investigation.\n\n\n\t\n[1] Huang, Biwei, et al. \"Adarl: What, where, and how to adapt in transfer reinforcement learning.\" arXiv preprint arXiv:2107.02729 (2021).\n\n\n[2] Liu, Yuren, et al. \"Learning world models with identifiable factorization.\" Advances in Neural Information Processing Systems 36 (2023): 31831-31864.\n\n\n**Q2**: *“The generalization of the proposed definition of the subtask is rather specific and the shift of task distribution is not significant. So the generalization of the proposed method might be less convincing.”*\n\n**A2**: We sincerely thank you for pointing out this possible limitation. In order for a wider generalization, we need to consider more practical issues. Your insight inspires us to provide a more general answer to the problem such that we are able to handle more distribution shift scenarios. \n\n\nIn light of your suggestion, we have made the following efforts, including new experiments, in the updated manuscript:\n\n\n* *More experimental results on distribution shift*: We conducted additional experiments by considering a distribution shift problem that involves longer-horizon tasks, and incorporated the results in the uploaded pdf. Specifically, under the Kitchen environment, we keep the same training set of tasks (each task is composed of $4$ sequential manipulation subtasks), and test the method’s generalizability to a new task composed of one more subtask, i.e. $5$ random sequential manipulation subtasks. Such generalization to longer-horizon tasks is not taken into consideration by some of the existing works [3], and our empirical results show that our formulations are able to deal with such a more challenging distribution shift problem. Please refer to the **Global Response** for experiment details.\n\n\n* *Rationale for adopting the current generalization cases*: We totally agree with you that our settings do not encompass all kinds of task distribution shifts, but focus on the shift where new tasks are randomly composed of seen subtasks. \nThis assumption on task distribution is also widely used in other subtask/skill discovery literature. For instance,  BUDS [1] inquires about its ability to solve new task variants that require different subtask combinations; LOVE [2] tests on the grid-world environment where each subtask corresponds to picking up a specific object; MH-AILR [3] also uses the Kitchen environment, in which the tasks require the sequential completion of four specific subtasks. In short, all works consider the generalization in cases where the training set of subtasks provides full support of the target set of subtasks; our generalization experiment in the Kitchen environment is not more limited than other works.\n\n* *Possible solutions*: We provide preliminary solutions to such a problem in more types of task distribution shifts: for example, utilizing RL combined with IL when the target subtask is out-of-distribution of the training set; learning subtasks on a more granular level such that more primitive actions can serve to facilitate the discovery of new subtasks. We believe that our results could serve as a theoretical basis, which may prove helpful for many future explorations by the community. \n\nMeanwhile, we take this work as a starting point, and we aim to develop a principled, easily-extensible way to resolve the subtask-discovery problem. We thank you for your great advice, which definitely inspired us to do so.\n\n\n[1] Yifeng Zhu, et al. Bottom-up skill discovery from unsegmented demonstrations for long-horizon robot manipulation. IEEE Robotics and Automation Letters, 7(2):4126–4133, 2022.\n\n\n[2] Yiding Jiang, Evan Zheran Liu, Benjamin Eysenbach, J Zico Kolter, and Chelsea Finn. Learning Options via Compression.\n\n\n[3] Jiayu Chen, Dipesh Tamboli, Tian Lan, and Vaneet Aggarwal. Multi-task Hierarchical Adversarial Inverse Reinforcement Learning. In Proceedings of the 40th International Conference on Machine Learning, pages 4895–4920. PMLR, July 2023.", "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": "", "extra_scores": {}}, {"review_id": "laCgbcMaeu", "reviewer": "Reviewer_1Ld6", "summary": "This paper addresses offline subtask discovery from a causal perspective by identifying subgoals as selections, targeting at solving long-horizon tasks and acquiring transferrable skills. The algorithm design is well-motivated and shows superior performance in offline subtask discovery.", "questions": "Please see the weakness part.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "(a) The causal-graph-based algorithm design is well motivated and theoretically solid.\n\n(b) The presentation is elegant and easy to follow.\n\n(c) Superiority of the proposed algorithm is shown through both quantitative and qualitative study.", "weaknesses": "(a) Kitchen is a challenging benchmark, but each of its tasks consists of a sequence of subtasks. Consequently, subtask discovery from this well-structured offline data is relatively straightforward. It would be beneficial to demonstrate the subtask discovery capability of the proposed algorithm in tasks that lack clearly defined, semantically obvious subtask structures.\n\n(b) This work also focuses on multi-task learning, but it is confined to scenarios where the tasks in the set are merely different compositions of the same set of subtasks. More generalized multi-task learning can be future directions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are profoundly thankful for your valuable feedback and the time you have dedicated to reading it, as they will surely improve the quality of this manuscript. In light of your suggestion, we incorporated additional discussions, as well as experiments to demonstrate the generalizability of our method.\n\n**Q1**: *“Kitchen is a challenging benchmark, but each of its tasks consists of a sequence of subtasks. Consequently, subtask discovery from this well-structured offline data is relatively straightforward. It would be beneficial to demonstrate the subtask discovery capability of the proposed algorithm in tasks that lack clearly defined, semantically obvious subtask structures.”*\n\n**A1**: We completely agree with you that, as you mentioned, Kitchen is a well-structured dataset thus it is easier to consider. Further, we aim to take it as a starting point to develop a principled way to tackle subtask discovery problems. More practical issues will surely need to be discussed in the future.\n\nIn our setting, we deal with multi-task learning problems when new tasks are arbitrarily composed by learned subtasks. Like others, we do not use the skill structure annotations in either training or testing, so the semantically meaningful \"subtasks\" used in the paper are simply a way to refer to the learned segmentations. It is not necessary to have them pre-defined but it is good to have them well-structured for easier evaluation.\n\nAt the same time, we do not rely on a more restricted subtask structure than other work. we adopt a common setting in the literature regarding the subtask discovery problem. By conducting a comprehensive literature review on all the datasets that subtask-discovery work considers, we found that in order to evaluate the methods that are proposed and also for better understanding, it is common that we use tasks that are naturally composed of semantically meaningful and clear subtasks.\n\nFor example, the latest works: [1] consider a grid-world setting BabyAI [2] that uses well-structured language (e.g.GoTo, Pickup) to indicate subgoals to accomplish, and ALFRED [3] environment that asks the agent to perform household tasks, both involve clearly defined subtasks like navigating in the kitchen, putting into the microwave, etc. Another work [4] considers MetaWorld, training on tasks that are composed of a sequence of subtasks like grasp, and move.\n\nOverall, we do agree with your insight that the semantically clear-defined subtask structure is a straightforward scenario. Meanwhile, this kind of subtask structure is also employed by all related works, and we think that the Kitchen environment is a reasonable choice for this first step of exploring a theoretical framework for subtask-discovery, and your great suggestion inspires us to see how the algorithm would perform in other environments.\n\n[1] Fu, Haotian, et al. \"Language-guided Skill Learning with Temporal Variational Inference.\" In Forty-first International Conference on Machine Learning (ICML 2024).\n\n[2] Chevalier-Boisvert, et al. Babyai: A platform to study the sample efficiency of grounded language learning. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.\n\n[3] Shridhar, M., et al. ALFRED: A benchmark for interpreting grounded instructions for everyday tasks. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, pp. 10737–10746. Computer Vision Foundation / IEEE, 2020a.\n\n[4] Zheng, Ruijie, et al\"PRISE: LLM-Style Sequence Compression for Learning Temporal Action Abstractions in Control.\" In Forty-first International Conference on Machine Learning (ICML 2024).\n\n\n**Q2**: *“This work also focuses on multi-task learning, but it is confined to scenarios where the tasks in the set are merely different compositions of the same set of subtasks. More generalized multi-task learning can be a future direction.”*\n\n**A2**: We sincerely appreciate you bringing this to our attention. Generalization to new scenarios that encompass unseen subtasks, i.e. when no discovered subtasks are reusable, can be challenging and deserves further investigation. \n\nIn light of your great suggestion, we have made the following efforts, including new experiments in our revision:\n\n* *Possible solutions*: We provide some preliminary solutions to such a problem that you have insightfully pointed out in more generalized scenarios: for example, considering the ability to discover new subtasks when the target subtask is out-of-distribution of the training set; learning subtasks more granularly such that more primitive actions can serve to facilitate new tasks learning. \n\n* *Additional experiments*: We conducted additional experiments under in Kitchen, and tested the method’s generalizability to a new task composed of one more subtask, i.e. $5$ sequential manipulation subtasks. Such generalization to longer-horizon tasks is challenging and is not taken into account by our MIL baseline[3]. Please refer to the uploaded pdf for the result and the **Global Response** for experimental details.\n\nWe deal with scenarios where all training and testing tasks share the same set of subtasks, while with different (unseen) compositions, just like other related works [1] [2] [3], and have validated our method through challenging experiments. We are greatly thankful for your suggestion, as it undoubtedly encourages us to pursue further investigation.\n\n[1] Yifeng Zhu, et al. Bottom-up skill discovery from unsegmented demonstrations for long-horizon robot manipulation. IEEE Robotics and Automation Letters, 7(2):4126–4133, 2022.\n\n[2] Yiding Jiang, et al. Learning Options via Compression.\n\n[3] Jiayu Chen, et al. Multi-task Hierarchical Adversarial Inverse Reinforcement Learning. In Proceedings of the 40th International Conference on Machine Learning, pages 4895–4920. PMLR, July 2023.", "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": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all the reviewers for your dedicated time and insightful comments. We are happy to see that the novelty of this work is well-recognized by all reviewers, which lies in the idea of identifying subtasks as selections, as well as the soundness of the corresponding algorithms. We added one additional experiment to test the generalization ability of the method to long-horizon tasks, details of which are provided below, and the result is presented in the figure of the uploaded pdf.

**Additional experiment in Kitchen**:
In Sec. 5.3, we have tested the performance of the proposed algorithm on a new task in the Kitchen environment [1] that is composed of $4$ subtasks with an unforeseen combination, while trained on expert demonstrations [2] with matching horizon ($4$ subtasks). A further question is, is the algorithm able to generalize to more scenarios? 

In order to answer this question, we test on a target task of $5$ subtasks to make it a more challenging generalization scenario. Specifically, we use |bottom burner| $\rightarrow$ |top burner| $\rightarrow$ |slide cabinet| $\rightarrow$  |hinge cabinet| , and |microwave| $\rightarrow$ |bottom| |burner| $\rightarrow$ |top burner| $\rightarrow$ |hinge cabinet| in the demonstrations for training. For testing, we require the agent to manipulate: |microwave| $\rightarrow$ |bottom burner| $\rightarrow$ |top burner| $\rightarrow$ |slide cabinet| $\rightarrow$ |hinge cabinet| in order. We show the result (accumulated return w.r.t. training steps) in the additional figure.
The baseline method H-AIRL[3] has only been evaluated on a 4-subtask horizon, and our algorithm is demonstrated to have better performance compared with all baselines. This suggests that the subtasks that we learned from the demonstrations as selections show better generalization ability even in accomplishing challenging long-horizon tasks, further validating the superiority of our method.

[1] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.

[2] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay Policy Learning: Solving Long-Horizon Tasks via Imitation and Reinforcement Learning, October 2019.

[3] Jiayu Chen, Dipesh Tamboli, Tian Lan, and Vaneet Aggarwal. Multi-task Hierarchical Adversarial Inverse Reinforcement Learning. In Proceedings of the 40th International Conference on Machine Learning, pages 4895–4920. PMLR, July 2023

Meta-review: This paper introduces a new definition of subgoal as a selection variable in causal inference and proposes a sequential non-negative matrix factorization method to discover subgoals. The empirical result on the Kitchen environment shows that the proposed method can learn useful subtasks from expert demonstrations and outperform relevant baselines as a result. 

All of the reviewers found that the proposed definition of subgoal is well-motivated and convincing, and the proposed method based on this definition is also technically sound. Although there were minor concerns around the tasks in the Kitchen environment not challenging enough, all of the reviewers unanimously agreed that the paper is novel and strong enough to be presented at NeurIPS. Thus, I recommend to accept this paper.
Decision: Accept (poster)