Original: D IRECTED -I NFO GAIL: L EARNING H IERARCHICAL P OLICIES FROM U NSEGMENTED D EMONSTRATIONS USING D IRECTED I NFORMATION
Revised: D IRECTED -I NFO GAIL: L EARNING H IERARCHICAL P OLICIES FROM U NSEGMENTED D EMONSTRATIONS USING D IRECTED I NFORMATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: The use of imitation learning to learn a single policy for a complex task that has multiple modes or hierarchical structure can be challenging.
Revised: The use of imitation learning to learn a single policy for a complex task that has multiple modes or hierarchical structure can be challenging.

Original: In fact, previous work has shown that when the modes are known, learning separate policies for each mode or sub-task can greatly improve the performance of imitation learning.
Revised: In fact, previous work has shown that when the modes are known, learning separate policies for each mode or sub-task can greatly improve the performance of imitation learning.

Original: In this work, we discover the interaction between sub-tasks from their resulting stateaction trajectory sequences using a directed graphical model.
Revised: In this work, we discover the interaction between sub-tasks from their resulting stateaction trajectory sequences using a directed graphical model.

Original: We propose a new algorithm based on the generative adversarial imitation learning framework which automatically learns sub-task policies from unsegmented demonstrations.
Revised: We propose a new algorithm based on the generative adversarial imitation learning framework which automatically learns sub-task policies from unsegmented demonstrations.

Original: Our approach maximizes the directed information ﬂow in the graphical model between sub-task latent variables and their generated trajectories.
Revised: Our approach maximizes the directed information ﬂow in the graphical model between sub-task latent variables and their generated trajectories.

Original: We also show how our approach connects with the existing Options framework, which is commonly used to learn hierarchical policies
Revised: We also show how our approach connects with the existing Options framework, which is commonly used to learn hierarchical policies

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Complex human activities can often be broken down into various simpler sub-activities or sub-tasks that can serve as the basic building blocks for completing a variety of complicated tasks.
Revised: Complex human activities can often be broken down into various simpler sub-activities or sub-tasks that can serve as the basic building blocks for completing a variety of complicated tasks.

Original: For instance, when driving a car, a driver may perform several simpler sub-tasks such as driving straight in a lane, changing lanes, executing a turn and braking, in different orders and for varying times depending on the source, destination, trafﬁc conditions etc.
Revised: For instance, when driving a car, a driver may perform several simpler sub-tasks such as driving straight in a lane, changing lanes, executing a turn and braking, in different orders and for varying times depending on the source, destination, trafﬁc conditions etc.

Original: Using imitation learning to learn a single monolithic policy to represent a structured activity can be challenging as it does not make explicit the sub-structure between the parts within the activity.
Revised: Using imitation learning to learn a single monolithic policy to represent a structured activity can be challenging as it does not make explicit the sub-structure between the parts within the activity.

Original: In this work, we develop an imitation learning framework that can learn a policy for each of these sub-tasks given unsegmented activity demonstrations and also learn a macro-policy which dictates switching from one sub-task policy to another.
Revised: In this work, we develop an imitation learning framework that can learn a policy for each of these sub-tasks given unsegmented activity demonstrations and also learn a macro-policy which dictates switching from one sub-task policy to another.

Original: Learning sub-task speciﬁc policies has the beneﬁt of shared learning.
Revised: Learning sub-task speciﬁc policies has the beneﬁt of shared learning.

Original: Each such sub-task policy also needs to specialize over a restricted state space, thus making the learning problem easier.
Revised: Each such sub-task policy also needs to specialize over a restricted state space, thus making the learning problem easier.

Original: Previous works in imitation learning (Li et al., 2017; Hausman et al., 2017) focus on learning each sub-task speciﬁc policy using segmented expert demonstrations by modeling the variability in each sub-task policy using a latent variable.
Revised: Previous works in imitation learning (Li et al., 2017; Hausman et al., 2017) focus on learning each sub-task speciﬁc policy using segmented expert demonstrations by modeling the variability in each sub-task policy using a latent variable.

Original: This latent variable is inferred by enforcing high mutual information between the latent variable and expert demonstrations.
Revised: This latent variable is inferred by enforcing high mutual information between the latent variable and expert demonstrations.

Original: This information theoretic perspective is equivalent to the graphical model shown in Figure 1 (Left), where the node c represents the latent variable.
Revised: This information theoretic perspective is equivalent to the graphical model shown in Figure 1 (Left), where the node c represents the latent variable.

Original: However, since learning sub-task policies requires isolated demonstrations for each sub-task, this setup is difﬁcult to scale to many real world scenarios where providing such segmented trajectories is cumbersome.
Revised: However, since learning sub-task policies requires isolated demonstrations for each sub-task, this setup is difﬁcult to scale to many real world scenarios where providing such segmented trajectories is cumbersome.

Original: Further, this setup does not learn a macro-policy to combine the learned sub-task policies in meaningful ways to achieve different tasks.
Revised: Further, this setup does not learn a macro-policy to combine the learned sub-task policies in meaningful ways to achieve different tasks.

Original: In our work, we aim to learn each sub-task policy directly from unsegmented activity demonstrations.
Revised: In our work, we aim to learn each sub-task policy directly from unsegmented activity demonstrations.

Original: For example, given a task consisting of three sub-tasks — A, B and C, we wish to learn a policy to complete sub-task A, learn when to transition from A to B, ﬁnish sub-task B and so on.
Revised: For example, given a task consisting of three sub-tasks — A, B and C, we wish to learn a policy to complete sub-task A, learn when to transition from A to B, ﬁnish sub-task B and so on.

Original: To achieve this we use a causal graphical model, which can be represented as a Dynamic Bayesian Network as shown in Figure 1 (Right).
Revised: To achieve this we use a causal graphical model, which can be represented as a Dynamic Bayesian Network as shown in Figure 1 (Right).

Original: The nodes c t denote latent variables which indicate the currently active sub-task and the nodes τ t denote the state-action pair at time t .
Revised: The nodes c t denote latent variables which indicate the currently active sub-task and the nodes τ t denote the state-action pair at time t .

Original: We consider as given, a set of expert demonstrations, each of which is represented by τ = { τ 1 , · · · , τ T } and has a corresponding sequence of latent factors
Revised: We consider as given, a set of expert demonstrations, each of which is represented by τ = { τ 1 , · · · , τ T } and has a corresponding sequence of latent factors

Original: c = { c 1 , · · · , c T − 1 } .
Revised: c = { c 1 , · · · , c T − 1 } .

Original: The sub-activity at time t dictates what state-action pair was ccτc tc tτ t-1 ττ t τcτττ tτ t c generated at time t .
Revised: The sub-activity at time t dictates what state-action pair was ccτc tc tτ t-1 ττ t τcτττ tτ t c generated at time t .

Original: As we will discuss in Section 3, extending the use of mutual information to learn sub-task policies from unsegmented demonstrations is problematic , as it requires learning the macro-policy as a conditional probability distribution which depends on the unobserved future .
Revised: As we will discuss in Section 3, extending the use of mutual information to learn sub-task policies from unsegmented demonstrations is problematic , as it requires learning the macro-policy as a conditional probability distribution which depends on the unobserved future .

Original: This unobserved future is unknown during earlier points of interaction (Figure 1).
Revised: This unobserved future is unknown during earlier points of interaction (Figure 1).

Original: To alleviate this, in our work we aim to force the policy to generate trajectories that maximize the directed information or causal information (Massey, 1990) ﬂow from trajectories to latent factors of variation within the trajectories instead of mutual information.
Revised: To alleviate this, in our work we aim to force the policy to generate trajectories that maximize the directed information or causal information (Massey, 1990) ﬂow from trajectories to latent factors of variation within the trajectories instead of mutual information.

Original: Using directed information requires us to learn a causally conditioned probability distribution (Kramer, 1998) which depends only on the observed past while allowing the unobserved future to be sequentially revealed.
Revised: Using directed information requires us to learn a causally conditioned probability distribution (Kramer, 1998) which depends only on the observed past while allowing the unobserved future to be sequentially revealed.

Original: Further, since there exists feedback in our causal graphical model i.e. , information ﬂows from the latent variables to trajectories and vice versa, directed information also provides a better upper bound on this information ﬂow between the latent variables and expert trajectories than does the conventional mutual information (Massey, 1990; Kramer, 1998).
Revised: Further, since there exists feedback in our causal graphical model i.e. , information ﬂows from the latent variables to trajectories and vice versa, directed information also provides a better upper bound on this information ﬂow between the latent variables and expert trajectories than does the conventional mutual information (Massey, 1990; Kramer, 1998).

Original: We also draw connections with existing work on learning sub-task policies using imitation learning with the options framework (Sutton et al., 1998; Daniel et al., 2016).
Revised: We also draw connections with existing work on learning sub-task policies using imitation learning with the options framework (Sutton et al., 1998; Daniel et al., 2016).

Original: We show that our work, while derived using the information theoretic perspective of maximizing directed information, bears a close resemblance to applying the options framework in a generative adversarial imitation setting.
Revised: We show that our work, while derived using the information theoretic perspective of maximizing directed information, bears a close resemblance to applying the options framework in a generative adversarial imitation setting.

Original: Thus, our approach combines the beneﬁts of learning hierarchical policies using the options framework with the robustness of generative adversarial imitation learning, helping overcome problems such as compounding errors that plague behaviour cloning.
Revised: Thus, our approach combines the beneﬁts of learning hierarchical policies using the options framework with the robustness of generative adversarial imitation learning, helping overcome problems such as compounding errors that plague behaviour cloning.

Original: In summary, the main contributions of our work include:
Revised: In summary, the main contributions of our work include:

Original: • We extend existing generative adversarial imitation learning frameworks to allow for learning of sub-task speciﬁc policies by maximizing directed information in a causal graph of subactivity latent variables and observed trajectory variables.
Revised: • We extend existing generative adversarial imitation learning frameworks to allow for learning of sub-task speciﬁc policies by maximizing directed information in a causal graph of subactivity latent variables and observed trajectory variables.

Original: • We draw connections between previous works on imitation learning with sub-task policies using options and show that our proposed approach can also be seen as option learning in a generative adversarial setting.
Revised: • We draw connections between previous works on imitation learning with sub-task policies using options and show that our proposed approach can also be seen as option learning in a generative adversarial setting.

Original: • We show through experiments on both discrete and continuous state-action spaces, the ability of our approach to segment expert demonstrations into meaningful sub-tasks and combine sub-task speciﬁc policies to perform the desired task.
Revised: • We show through experiments on both discrete and continuous state-action spaces, the ability of our approach to segment expert demonstrations into meaningful sub-tasks and combine sub-task speciﬁc policies to perform the desired task.

Original:  R ELATED W ORK 2. I MITATION L EARNING
Revised:  R ELATED W ORK 2. I MITATION L EARNING

Original: Imitation Learning (Pomerleau, 1989) aims at learning policies that can mimic expert behaviours from demonstrations.
Revised: Imitation Learning (Pomerleau, 1989) aims at learning policies that can mimic expert behaviours from demonstrations.

Original: Modeling the problem as a Markov Decision Process (MDP), the goal in imitation learning is to learn a policy π ( a | s ) , which deﬁnes the conditional distribution over actions a ∈ A given the state s ∈ S , from state-action trajectories τ
Revised: Modeling the problem as a Markov Decision Process (MDP), the goal in imitation learning is to learn a policy π ( a | s ) , which deﬁnes the conditional distribution over actions a ∈ A given the state s ∈ S , from state-action trajectories τ

Original: = ( s 0 , a 0 , · · · , s T ) of expert behaviour.
Revised: = ( s 0 , a 0 , · · · , s T ) of expert behaviour.

Original: Recently,
Revised: Recently,

Original: Ho & Ermon (2016) introduced an imitation learning framework called Generative Adversarial
Revised: Ho & Ermon (2016) introduced an imitation learning framework called Generative Adversarial

Original: Imitation Learning (GAIL) that is able to learn policies for complex high-dimensional physics-based control tasks.
Revised: Imitation Learning (GAIL) that is able to learn policies for complex high-dimensional physics-based control tasks.

Original: They reduce the imitation learning problem into an adversarial learning framework, for which they utilize Generative Adversarial Networks (GAN) (Goodfellow et al., 2014).
Revised: They reduce the imitation learning problem into an adversarial learning framework, for which they utilize Generative Adversarial Networks (GAN) (Goodfellow et al., 2014).

Original: The generator network of the GAN represents the agent’s policy π while the discriminator network serves as a local reward function and learns to differentiate between state-action pairs from the expert policy π E and from the agent’s policy π .
Revised: The generator network of the GAN represents the agent’s policy π while the discriminator network serves as a local reward function and learns to differentiate between state-action pairs from the expert policy π E and from the agent’s policy π .

Original: Mathematically, it is equivalent to optimizing the following,
Revised: Mathematically, it is equivalent to optimizing the following,

Original: [Equation]
Revised: [Equation]

Original: InfoGAIL (Li et al., 2017) and Hausman et al.
Revised: InfoGAIL (Li et al., 2017) and Hausman et al.

Original: (2017) solve the problem of learning from policies generated by a mixture of experts.
Revised: (2017) solve the problem of learning from policies generated by a mixture of experts.

Original: They introduce a latent variable c into the policy function π ( a | s, c ) to separate different type of behaviours present in the demonstration.
Revised: They introduce a latent variable c into the policy function π ( a | s, c ) to separate different type of behaviours present in the demonstration.

Original: To incentivize the network to use the latent variable, they utilize an information-theoretic regularization enforcing that there should be high mutual information between c and the state-action pairs in the generated trajectory, a concept that was ﬁrst introduced in InfoGAN (Chen et al., 2016).
Revised: To incentivize the network to use the latent variable, they utilize an information-theoretic regularization enforcing that there should be high mutual information between c and the state-action pairs in the generated trajectory, a concept that was ﬁrst introduced in InfoGAN (Chen et al., 2016).

Original: They introduce a variational lower bound
Revised: They introduce a variational lower bound

Original: L 1 ( π, Q ) of the mutual information I ( c ; τ ) to the loss function in GAIL.
Revised: L 1 ( π, Q ) of the mutual information I ( c ; τ ) to the loss function in GAIL.

Original: [Equation]
Revised: [Equation]

Original: The modiﬁed objective can then be given as,
Revised: The modiﬁed objective can then be given as,

Original: [Equation]
Revised: [Equation]

Original: InfoGAIL models variations between different trajectories as the latent codes correspond to trajectories coming from different demonstrators.
Revised: InfoGAIL models variations between different trajectories as the latent codes correspond to trajectories coming from different demonstrators.

Original: In contrast, we aim to model intra-trajectory variations and latent codes in our work correspond to sub-tasks (variations) within a demonstration.
Revised: In contrast, we aim to model intra-trajectory variations and latent codes in our work correspond to sub-tasks (variations) within a demonstration.

Original: In Section 3, we discuss why using a mutual information based loss is infeasible in our problem setting and describe our proposed approach.
Revised: In Section 3, we discuss why using a mutual information based loss is infeasible in our problem setting and describe our proposed approach.

Original: 2. O PTIONS
Revised: 2. O PTIONS

Original: Consider an MDP with states s ∈ S and actions a ∈ A .
Revised: Consider an MDP with states s ∈ S and actions a ∈ A .

Original: Under the options framework (Sutton et al., 1998), an option, indexed by o ∈ O consists of a sub-policy π ( a | s, o ) , a termination policy π ( b | s, ¯ o ) and an option activation policy π ( o | s ) .
Revised: Under the options framework (Sutton et al., 1998), an option, indexed by o ∈ O consists of a sub-policy π ( a | s, o ) , a termination policy π ( b | s, ¯ o ) and an option activation policy π ( o | s ) .

Original: After an option is initiated, actions are generated by the sub-policy until the option is terminated and a new option is selected.
Revised: After an option is initiated, actions are generated by the sub-policy until the option is terminated and a new option is selected.

Original: Options framework has been studied widely in RL literature.
Revised: Options framework has been studied widely in RL literature.

Original: A challening problem related to the options framework is to automatically infer options without supervision.
Revised: A challening problem related to the options framework is to automatically infer options without supervision.

Original: Option discovery approaches often aim to ﬁnd bottleneck states, i.e. , states that the agent has to pass through to reach the goal.
Revised: Option discovery approaches often aim to ﬁnd bottleneck states, i.e. , states that the agent has to pass through to reach the goal.

Original: Many different approaches such as multiple-instance learning (McGovern & Barto, 2001), graph based algorithms (Menache et al., 2002; S¸ims¸ek et al., 2005) have been used to ﬁnd such bottleneck states.
Revised: Many different approaches such as multiple-instance learning (McGovern & Barto, 2001), graph based algorithms (Menache et al., 2002; S¸ims¸ek et al., 2005) have been used to ﬁnd such bottleneck states.

Original: Once the bottleneck states are discovered, the above approaches ﬁnd options policies to reach each such state.
Revised: Once the bottleneck states are discovered, the above approaches ﬁnd options policies to reach each such state.

Original: In contrast, we propose a uniﬁed framework using a information-theoretic approach to automatically discover relevant option policies without the need to discover bottleneck states.
Revised: In contrast, we propose a uniﬁed framework using a information-theoretic approach to automatically discover relevant option policies without the need to discover bottleneck states.

Original: Daniel et al.
Revised: Daniel et al.

Original: (2016) formulate the options framework as a probabilistic graphical model where options are treated as latent variables which are then learned from expert data.
Revised: (2016) formulate the options framework as a probabilistic graphical model where options are treated as latent variables which are then learned from expert data.

Original: The option policies ( π ( a | s, o ) ) are analogous to sub-task policies in our work.
Revised: The option policies ( π ( a | s, o ) ) are analogous to sub-task policies in our work.

Original: These option policies are then learned by maximizing a lower bound using the Expectation-Maximization algorithm (Moon, 1996).
Revised: These option policies are then learned by maximizing a lower bound using the Expectation-Maximization algorithm (Moon, 1996).

Original: We show how this lower bound is closely related to the objective derived in our work.
Revised: We show how this lower bound is closely related to the objective derived in our work.

Original: We further show how this connection allows our method to be seen as a generative adversarial variant of their approach.
Revised: We further show how this connection allows our method to be seen as a generative adversarial variant of their approach.

Original: Prior work in robot learning has also looked at learning motion primitives from unsegmented demonstrations.
Revised: Prior work in robot learning has also looked at learning motion primitives from unsegmented demonstrations.

Original: These primitives usually correspond to a particular skill and are analogous to options.
Revised: These primitives usually correspond to a particular skill and are analogous to options.

Original: Niekum & Barto (2011) used the Beta-Process Autoregressive Hidden Markov Model (BP-AR-HMM) to segment expert demonstrations and post-process these segments to learn motion primitives which provide the ability to use reinforcement learning for policy improvement.
Revised: Niekum & Barto (2011) used the Beta-Process Autoregressive Hidden Markov Model (BP-AR-HMM) to segment expert demonstrations and post-process these segments to learn motion primitives which provide the ability to use reinforcement learning for policy improvement.

Original: Alternately, Krishnan et al.
Revised: Alternately, Krishnan et al.

Original: (2018) use Dirichlet Process Gaussian Mixture Model (DP-GMM) to segment the expert demonstrations by ﬁnding transition states between linear dynamical segments.
Revised: (2018) use Dirichlet Process Gaussian Mixture Model (DP-GMM) to segment the expert demonstrations by ﬁnding transition states between linear dynamical segments.

Original: Similarly, Ranchod et al.
Revised: Similarly, Ranchod et al.

Original: (2015) use the BP-AR-HMM framework to initially segment the expert demonstrations and then use an inverse reinforcement learning step to infer the reward function for each segment.
Revised: (2015) use the BP-AR-HMM framework to initially segment the expert demonstrations and then use an inverse reinforcement learning step to infer the reward function for each segment.

Original: The use of appropriate priors allows these methods to discover options without a priori knowledge of the total number of skills.
Revised: The use of appropriate priors allows these methods to discover options without a priori knowledge of the total number of skills.

Original: However, unlike the above methods in our proposed approach we also learn an appropriate policy over the extracted options.
Revised: However, unlike the above methods in our proposed approach we also learn an appropriate policy over the extracted options.

Original: We show how this allows us to compose the individual option policies to induce novel behaviours which were not present in the expert demonstrations.
Revised: We show how this allows us to compose the individual option policies to induce novel behaviours which were not present in the expert demonstrations.

Original:  P ROPOSED A PPROACH
Revised:  P ROPOSED A PPROACH

Original: As mentioned in the previous section, while prior approaches can learn to disambiguate the multiple modalities in the demonstration of a sub-task and learn to imitate them, they cannot learn to imitate demonstrations of unsegmented long tasks that are formed by a combination of many small sub-tasks.
Revised: As mentioned in the previous section, while prior approaches can learn to disambiguate the multiple modalities in the demonstration of a sub-task and learn to imitate them, they cannot learn to imitate demonstrations of unsegmented long tasks that are formed by a combination of many small sub-tasks.

Original: To learn such sub-task policies from unsegmented deomonstrations we use the graphical model in Figure 1 (Right), i.e. , consider a set of expert demonstrations, each of which is represented by τ = { τ 1 , · · · , τ T } where τ t is the state-action pair observed at time t .
Revised: To learn such sub-task policies from unsegmented deomonstrations we use the graphical model in Figure 1 (Right), i.e. , consider a set of expert demonstrations, each of which is represented by τ = { τ 1 , · · · , τ T } where τ t is the state-action pair observed at time t .

Original: Each such demonstration has a corresponding sequence of latent variables c
Revised: Each such demonstration has a corresponding sequence of latent variables c

Original: = { c 1 , · · · , c T − 1 } which denote the sub-activity in the demonstration at any given time step.
Revised: = { c 1 , · · · , c T − 1 } which denote the sub-activity in the demonstration at any given time step.

Original: As noted before, previous approaches (Li et al., 2017; Hausman et al., 2017) model the expert sub-task demonstrations using only a single latent variable.
Revised: As noted before, previous approaches (Li et al., 2017; Hausman et al., 2017) model the expert sub-task demonstrations using only a single latent variable.

Original: To enforce the model to use this latent variable, these approaches propose to maximize the mutual information between the demonstrated sequence of state-action pairs and the latent embedding of the nature of the sub-activity.
Revised: To enforce the model to use this latent variable, these approaches propose to maximize the mutual information between the demonstrated sequence of state-action pairs and the latent embedding of the nature of the sub-activity.

Original: This is achieved by adding a lower bound to the mutual information between the latent variables and expert demonstrations.
Revised: This is achieved by adding a lower bound to the mutual information between the latent variables and expert demonstrations.

Original: This variational lower bound of the mutual information is then combined with the the adversarial loss for imitation learning proposed in Ho & Ermon (2016).
Revised: This variational lower bound of the mutual information is then combined with the the adversarial loss for imitation learning proposed in Ho & Ermon (2016).

Original: Extending this to our setting, where we have a sequence of latent variables c , yields the following lower bound on the mutual information,
Revised: Extending this to our setting, where we have a sequence of latent variables c , yields the following lower bound on the mutual information,

Original: [Equation]
Revised: [Equation]

Original: Observe that the dependence of q on the entire trajectory τ precludes the use of such a distribution at test time, where only the trajectory up to the current time is known.
Revised: Observe that the dependence of q on the entire trajectory τ precludes the use of such a distribution at test time, where only the trajectory up to the current time is known.

Original: To overcome this limitation, in this work we propose to force the policy to generate trajectories that maximize the directed or causal information ﬂow from trajectories to the sequence of latent sub-activity variables instead.
Revised: To overcome this limitation, in this work we propose to force the policy to generate trajectories that maximize the directed or causal information ﬂow from trajectories to the sequence of latent sub-activity variables instead.

Original: As we show below, by using directed information instead of mutual information, we can replace the dependence on τ with a dependence on the trajectory generated up to current time t .
Revised: As we show below, by using directed information instead of mutual information, we can replace the dependence on τ with a dependence on the trajectory generated up to current time t .

Original: The directed information ﬂow from a sequence X to Y is given by,
Revised: The directed information ﬂow from a sequence X to Y is given by,

Original: [Equation]
Revised: [Equation]

Original: where H ( Y (cid:107) X ) is the causally-conditioned entropy.
Revised: where H ( Y (cid:107) X ) is the causally-conditioned entropy.

Original: Replacing X and Y with sequences τ and c ,
Revised: Replacing X and Y with sequences τ and c ,

Original: [Equation]
Revised: [Equation]

Original: Here τ 1: t = ( s 1 , · · · , a t − 1 , s t ) .
Revised: Here τ 1: t = ( s 1 , · · · , a t − 1 , s t ) .

Original: A variational lower bound, L 1 ( π, q ) of the directed information,
Revised: A variational lower bound, L 1 ( π, q ) of the directed information,

Original: I ( τ → c ) which uses an approximate posterior q ( c t | c 1: t − 1 , τ 1: t ) instead of the true posterior p ( c t | c 1: t − 1 , τ 1: t ) can then be derived to get (See Appendix A.1 for the complete derivation),
Revised: I ( τ → c ) which uses an approximate posterior q ( c t | c 1: t − 1 , τ 1: t ) instead of the true posterior p ( c t | c 1: t − 1 , τ 1: t ) can then be derived to get (See Appendix A.1 for the complete derivation),

Original: L 1 ( π, q ) =
Revised: L 1 ( π, q ) =

Original: (cid:88)
Revised: (cid:88)

Original: [Equation]
Revised: [Equation]

Original: s 

Revised: s

Original: cs 1 cs 2 cs T-1 c TMLP
Revised: cs 1 cs 2 cs T-1 c TMLP

Original: MLP
Revised: MLP

Original: MLP
Revised: MLP

Original: MLP cccc
Revised: MLP cccc

Original: Ts 

Revised: Ts

Original: cs 1 cs 2 cs T-1 c TMLP
Revised: cs 1 cs 2 cs T-1 c TMLP

Original: MLP MLP
Revised: MLP MLP

Original: MLP a
Revised: MLP a

Original: [Figure]
Revised: [Figure]

Original: Thus, by maximizing directed information instead of mutual information, we can learn a posterior distribution over the next latent factor c given the latent factors discovered up to now and the trajectory followed up to now, thereby removing the dependence on the future trajectory.
Revised: Thus, by maximizing directed information instead of mutual information, we can learn a posterior distribution over the next latent factor c given the latent factors discovered up to now and the trajectory followed up to now, thereby removing the dependence on the future trajectory.

Original: In practice, we do not consider the H ( c ) term.
Revised: In practice, we do not consider the H ( c ) term.

Original: This gives us the following objective,
Revised: This gives us the following objective,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: We call this approach Directed-Info GAIL.
Revised: We call this approach Directed-Info GAIL.

Original: Notice that, to compute the loss in equation 3, we need to sample from the prior distribution p ( c 1: t ) .
Revised: Notice that, to compute the loss in equation 3, we need to sample from the prior distribution p ( c 1: t ) .

Original: In order to estimate this distribution, we ﬁrst pre-train a variational auto-encoder (VAE) (Kingma & Welling, 2013) on the expert trajectories, the details of which are described in the next sub-section.
Revised: In order to estimate this distribution, we ﬁrst pre-train a variational auto-encoder (VAE) (Kingma & Welling, 2013) on the expert trajectories, the details of which are described in the next sub-section.

Original: [Figure]
Revised: [Figure]

Original: 2 (left) shows the design of the VAE pictorially.
Revised: 2 (left) shows the design of the VAE pictorially.

Original: The VAE consists of two multi-layer perceptrons that serve as the encoder and the decoder.
Revised: The VAE consists of two multi-layer perceptrons that serve as the encoder and the decoder.

Original: The encoder uses the current state s t and the previous latent variable c t − 1 to produce the current latent variable c t .
Revised: The encoder uses the current state s t and the previous latent variable c t − 1 to produce the current latent variable c t .

Original: We used the Gumbel-softmax trick (Jang et al., 2016) to obtain samples of latent variables from a categorical distribution.
Revised: We used the Gumbel-softmax trick (Jang et al., 2016) to obtain samples of latent variables from a categorical distribution.

Original: The decoder then takes s t and c t as input and outputs the action a t .
Revised: The decoder then takes s t and c t as input and outputs the action a t .

Original: We use the following objective, which maximizes the lower bound of the probability of the trajectories p ( τ ) , to train our VAE,
Revised: We use the following objective, which maximizes the lower bound of the probability of the trajectories p ( τ ) , to train our VAE,

Original: [Equation]
Revised: [Equation]

Original: Figure 2 (right) gives an overview of the complete method.
Revised: Figure 2 (right) gives an overview of the complete method.

Original: The VAE pre-training step allows us to get approximate samples from the distribution p ( c 1: t ) to optimize equation 4.
Revised: The VAE pre-training step allows us to get approximate samples from the distribution p ( c 1: t ) to optimize equation 4.

Original: This is done by using q to obtain samples of latent variable sequence c by using its output on the expert demonstrations.
Revised: This is done by using q to obtain samples of latent variable sequence c by using its output on the expert demonstrations.

Original: In practice, we ﬁx the weights of the network q to those obtained from the VAE pre-training step when optimizing the Directed-Info GAIL loss in equation 4.
Revised: In practice, we ﬁx the weights of the network q to those obtained from the VAE pre-training step when optimizing the Directed-Info GAIL loss in equation 4.

Original: C ONNECTION WITH OPTIONS FRAMEWORK
Revised: C ONNECTION WITH OPTIONS FRAMEWORK

Original: In Daniel et al.
Revised: In Daniel et al.

Original: the authors provide a probabilistic perspective of the options framework.
Revised: the authors provide a probabilistic perspective of the options framework.

Original: Although, Daniel et al.
Revised: Although, Daniel et al.

Original: (2016) consider separate termination and option latent variables ( b t and o t ), for the purpose of comparison, we collapse them into a single latent variable c t , similar to our framework with a distribution p ( c t | s t ,
Revised: (2016) consider separate termination and option latent variables ( b t and o t ), for the purpose of comparison, we collapse them into a single latent variable c t , similar to our framework with a distribution p ( c t | s t ,

Original: c t − 1 ) .
Revised: c t − 1 ) .

Original: The lower-bound derived in Daniel et al.
Revised: The lower-bound derived in Daniel et al.

Original: (2016) which is maximized using Expectation-Maximization (EM) algorithm can then be written as (suppressing dependence on parameters),
Revised: (2016) which is maximized using Expectation-Maximization (EM) algorithm can then be written as (suppressing dependence on parameters),

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: Note that the ﬁrst term in equation 6 i.e. , the expectation over the distribution log p ( c t | s t , c t − 1 ) is the same as equation 3 of our proposed approach with a one-step Markov assumption and a conditional expectation with given expert trajectories instead of an expectation with generated trajectories.
Revised: Note that the ﬁrst term in equation 6 i.e. , the expectation over the distribution log p ( c t | s t , c t − 1 ) is the same as equation 3 of our proposed approach with a one-step Markov assumption and a conditional expectation with given expert trajectories instead of an expectation with generated trajectories.

Original: The second term in equation 6 i.e. , the expectation over log π ( a t | s t , c t ) is replaced by the GAIL loss in equation 4.
Revised: The second term in equation 6 i.e. , the expectation over log π ( a t | s t , c t ) is replaced by the GAIL loss in equation 4.

Original: Our proposed Directed-Info GAIL can be therefore be considered as the generative adversarial variant of imitation learning using the options framework.
Revised: Our proposed Directed-Info GAIL can be therefore be considered as the generative adversarial variant of imitation learning using the options framework.

Original: The VAE behaviour cloning pretraining step in equation 5 is exactly equivalent to equation 6, where we use approximate variational inference using VAEs instead of EM.
Revised: The VAE behaviour cloning pretraining step in equation 5 is exactly equivalent to equation 6, where we use approximate variational inference using VAEs instead of EM.

Original: Thus, our approach combines the beneﬁts of both behavior cloning and generative adversarial imitation learning.
Revised: Thus, our approach combines the beneﬁts of both behavior cloning and generative adversarial imitation learning.

Original: Using GAIL enables learning of robust policies that do not suffer from the problem of compounding errors.
Revised: Using GAIL enables learning of robust policies that do not suffer from the problem of compounding errors.

Original: At the same time, conditioning GAIL on latent codes learned from the behavior cloning step prevents the issue of mode collapse in GANs.
Revised: At the same time, conditioning GAIL on latent codes learned from the behavior cloning step prevents the issue of mode collapse in GANs.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We present results on both discrete and continuous state-action environments.
Revised: We present results on both discrete and continuous state-action environments.

Original: In both of these settings we show that (1) our method is able to segment out sub-tasks from given expert trajectories, (2) learn sub-task conditioned policies, and (3) learn to combine these sub-task policies in order to achieve the task objective.
Revised: In both of these settings we show that (1) our method is able to segment out sub-tasks from given expert trajectories, (2) learn sub-task conditioned policies, and (3) learn to combine these sub-task policies in order to achieve the task objective.

Original: 4. D ISCRETE E NVIRONMENT
Revised: 4. D ISCRETE E NVIRONMENT

Original: For the discrete setting, we choose a grid world environment which consists of a 15 × 11 grid with four rooms connected via corridors as shown in Figure 3.
Revised: For the discrete setting, we choose a grid world environment which consists of a 15 × 11 grid with four rooms connected via corridors as shown in Figure 3.

Original: The agent spawns at a random location in the grid and its goal is to reach an apple, which spawns in one of the four rooms randomly, using the shortest possible path.
Revised: The agent spawns at a random location in the grid and its goal is to reach an apple, which spawns in one of the four rooms randomly, using the shortest possible path.

Original: Through this experiment we aim to see whether our proposed approach is able to infer sub-tasks which correspond to meaningful navigation strategies and combine them to plan paths to different goal states.
Revised: Through this experiment we aim to see whether our proposed approach is able to infer sub-tasks which correspond to meaningful navigation strategies and combine them to plan paths to different goal states.

Original: Figure 3 shows sub-task policies learned by our approach in this task.
Revised: Figure 3 shows sub-task policies learned by our approach in this task.

Original: The two plots on the left correspond to two of the four different values of the latent variable.
Revised: The two plots on the left correspond to two of the four different values of the latent variable.

Original: The arrow at every state in the grid shows the agent action (direction) with the highest probability in that state for that latent variable.
Revised: The arrow at every state in the grid shows the agent action (direction) with the highest probability in that state for that latent variable.

Original: In the discussion that follows, we label the rooms from 1 to 4 starting from the room at the top left and moving in the clockwise direction.
Revised: In the discussion that follows, we label the rooms from 1 to 4 starting from the room at the top left and moving in the clockwise direction.

Original: We observe that the sub-tasks extracted by our approach represent semantically meaningful navigation plans.
Revised: We observe that the sub-tasks extracted by our approach represent semantically meaningful navigation plans.

Original: Also, each latent variable is utilized for a different sub-task.
Revised: Also, each latent variable is utilized for a different sub-task.

Original: For instance, the agent uses the latent code in Figure 3(a), to perform the sub-task of moving from room 1 to room 3 and from room 2 to room 4 and the code in Figure 3(b) to move in the opposite direction.
Revised: For instance, the agent uses the latent code in Figure 3(a), to perform the sub-task of moving from room 1 to room 3 and from room 2 to room 4 and the code in Figure 3(b) to move in the opposite direction.

Original: Further, our approach learns to successfully combine these navigation strategies to achieve the given objectives.
Revised: Further, our approach learns to successfully combine these navigation strategies to achieve the given objectives.

Original: For example, Figure 3(c, d) show examples of how the
Revised: For example, Figure 3(c, d) show examples of how the

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: macro-policy switches between various latent codes to achieve the desired goals of reaching the apples in rooms 1 and 2 respectively.
Revised: macro-policy switches between various latent codes to achieve the desired goals of reaching the apples in rooms 1 and 2 respectively.

Original: 4. C ONTINUOUS E NVIRONMENTS
Revised: 4. C ONTINUOUS E NVIRONMENTS

Original: To validate our proposed approach in continuous control tasks we experiment with 5 continuous state-action environments.
Revised: To validate our proposed approach in continuous control tasks we experiment with 5 continuous state-action environments.

Original: The ﬁrst environment involves learning to draw circles on a 2D plane and is called Circle-World .
Revised: The ﬁrst environment involves learning to draw circles on a 2D plane and is called Circle-World .

Original: In this experiment, the agent must learn to draw a circle in both clockwise and counter-clockwise direction.
Revised: In this experiment, the agent must learn to draw a circle in both clockwise and counter-clockwise direction.

Original: The agent always starts at (0,0), completes a circle in clockwise direction and then retraces its path in the counter-clockwise direction.
Revised: The agent always starts at (0,0), completes a circle in clockwise direction and then retraces its path in the counter-clockwise direction.

Original: The trajectories differ in the radii of the circles.
Revised: The trajectories differ in the radii of the circles.

Original: The state s ∈ R 2 is the (x,y) co-ordinate and the actions a ∈ R 2 is a unit vector representing the direction of motion.
Revised: The state s ∈ R 2 is the (x,y) co-ordinate and the actions a ∈ R 2 is a unit vector representing the direction of motion.

Original: Notice that in Circle-World, the expert trajectories include two different actions (for clockwise and anti-clockwise direction) for every state ( x , y ) in the trajectory, thus making the problem multi-modal in nature.
Revised: Notice that in Circle-World, the expert trajectories include two different actions (for clockwise and anti-clockwise direction) for every state ( x , y ) in the trajectory, thus making the problem multi-modal in nature.

Original: This requires the agent to appropriately disambiguate between the two different phases of the trajectory.
Revised: This requires the agent to appropriately disambiguate between the two different phases of the trajectory.

Original: Further, to show the scalability of our approach to higher dimensional continuous control tasks we also show experiments on Pendulum, Inverted Pendulum, Hopper and Walker environments, provided in OpenAI Gym (Brockman et al., 2016).
Revised: Further, to show the scalability of our approach to higher dimensional continuous control tasks we also show experiments on Pendulum, Inverted Pendulum, Hopper and Walker environments, provided in OpenAI Gym (Brockman et al., 2016).

Original: Each task is progressively more challenging, with a larger state and action space.
Revised: Each task is progressively more challenging, with a larger state and action space.

Original: Our aim with these experiments is to see whether our approach can identify certain action primitives which helps the agent to complete the given task successfully.
Revised: Our aim with these experiments is to see whether our approach can identify certain action primitives which helps the agent to complete the given task successfully.

Original: To verify the effectiveness of our proposed approach we do a comparative analysis of our results with both GAIL (Ho & Ermon, 2016) and the supervised behavior cloning approaching using a VAE.
Revised: To verify the effectiveness of our proposed approach we do a comparative analysis of our results with both GAIL (Ho & Ermon, 2016) and the supervised behavior cloning approaching using a VAE.

Original: To generate expert trajectories we train an agent using Proximal Policy Optimization (Schulman et al., 2017).
Revised: To generate expert trajectories we train an agent using Proximal Policy Optimization (Schulman et al., 2017).

Original: We used 25 expert trajectories for the Pendulum and Inverted Pendulum tasks and 50 expert trajectories for experiments with the Hopper and Walker environments.
Revised: We used 25 expert trajectories for the Pendulum and Inverted Pendulum tasks and 50 expert trajectories for experiments with the Hopper and Walker environments.

Original: Figures 4(a, b, c) show results on the Circle-World environment.
Revised: Figures 4(a, b, c) show results on the Circle-World environment.

Original: As can be seen in Figure 4(a, b), when using two sub-task latent variables, our method learns to segment the demonstrations into
Revised: As can be seen in Figure 4(a, b), when using two sub-task latent variables, our method learns to segment the demonstrations into

Original: [Table]
Revised: [Table]

Original: two intuitive sub-tasks of drawing circles in clockwise and counterclockwise directions.
Revised: two intuitive sub-tasks of drawing circles in clockwise and counterclockwise directions.

Original: Hence, our method is able to identify the underlying modes and thus ﬁnd meaningful sub-task segmentations from unsegmented data.
Revised: Hence, our method is able to identify the underlying modes and thus ﬁnd meaningful sub-task segmentations from unsegmented data.

Original: We also illustrate how the learned sub-task policies can be composed to perform new types of behavior that were unobserved in the expert data.
Revised: We also illustrate how the learned sub-task policies can be composed to perform new types of behavior that were unobserved in the expert data.

Original: In Figure 4(c) we show how the sub-task policies can be combined to draw the circles in inverted order of direction by swapping the learned macro-policy with a different desired policy.
Revised: In Figure 4(c) we show how the sub-task policies can be combined to draw the circles in inverted order of direction by swapping the learned macro-policy with a different desired policy.

Original: Thus, the sub-task policies can be utilized as a library of primitive actions which is a signiﬁcant beneﬁt over methods learning monolithic policies.
Revised: Thus, the sub-task policies can be utilized as a library of primitive actions which is a signiﬁcant beneﬁt over methods learning monolithic policies.

Original: We now discuss the results on the classical Pendulum environment.
Revised: We now discuss the results on the classical Pendulum environment.

Original: Figure 4(d) shows the sub-task latent variables assigned by our approach to the various states.
Revised: Figure 4(d) shows the sub-task latent variables assigned by our approach to the various states.

Original: As can be seen in the ﬁgure, the network is able to associate different latent variables to different sub-tasks.
Revised: As can be seen in the ﬁgure, the network is able to associate different latent variables to different sub-tasks.

Original: For instance, states that have a high velocity are assigned a particular latent variable (shown in blue).
Revised: For instance, states that have a high velocity are assigned a particular latent variable (shown in blue).

Original: Similarly, states that lie close to position 0 and have low velocity (i.e. the desired target position) get assigned another latent variable (shown in green).
Revised: Similarly, states that lie close to position 0 and have low velocity (i.e. the desired target position) get assigned another latent variable (shown in green).

Original: The remaining states get classiﬁed as a separate sub-task.
Revised: The remaining states get classiﬁed as a separate sub-task.

Original: Figure 5 shows the results on the higher dimensional continuous control, Hopper and Walker, environments.
Revised: Figure 5 shows the results on the higher dimensional continuous control, Hopper and Walker, environments.

Original: Figure 5(a) shows a plots for sub-task latent variable assignment obtained on these environments.
Revised: Figure 5(a) shows a plots for sub-task latent variable assignment obtained on these environments.

Original: Our proposed method identiﬁes basic action primitives which are then chained together to effectively perform the two locomotion tasks.
Revised: Our proposed method identiﬁes basic action primitives which are then chained together to effectively perform the two locomotion tasks.

Original: Figure 5(b) shows that our approach learns to assign separate latent variable values for different action primitives such as, jumping, mid-air and landing phases of these tasks, with the latent variable changing approximately periodically as the agent performs the periodic hopping/walking motion.
Revised: Figure 5(b) shows that our approach learns to assign separate latent variable values for different action primitives such as, jumping, mid-air and landing phases of these tasks, with the latent variable changing approximately periodically as the agent performs the periodic hopping/walking motion.

Original: Finally, in Table 1 we also show the quantitative evaluation on the above continuous control environments.
Revised: Finally, in Table 1 we also show the quantitative evaluation on the above continuous control environments.

Original: We report the mean and standard deviations of the returns over 300 episodes.
Revised: We report the mean and standard deviations of the returns over 300 episodes.

Original: As can be seen, our approach improves the performance over the VAE pre-training step, overcoming the issue of compounding errors.
Revised: As can be seen, our approach improves the performance over the VAE pre-training step, overcoming the issue of compounding errors.

Original: The performance of our approach is comparable to the state-of-the-art
Revised: The performance of our approach is comparable to the state-of-the-art

Original: GAIL (Ho & Ermon, 2016).
Revised: GAIL (Ho & Ermon, 2016).

Original: Our method moreover, has the added advantage of segmenting the demonstrations into sub-tasks and also providing composable sub-task policies.
Revised: Our method moreover, has the added advantage of segmenting the demonstrations into sub-tasks and also providing composable sub-task policies.

Original: We further analyze our proposed approach in more detail in the Appendix.
Revised: We further analyze our proposed approach in more detail in the Appendix.

Original: In Appendix A.4 we visualize the sub-tasks in a low-dimensional sub-space.
Revised: In Appendix A.4 we visualize the sub-tasks in a low-dimensional sub-space.

Original: Also, in Appendix A.5 we show results when using a larger dimensional sub-task latent variable.
Revised: Also, in Appendix A.5 we show results when using a larger dimensional sub-task latent variable.

Original: A video of our results on Hopper and Walker environments can be seen at https://sites.google.com/view/directedinfo-gail .
Revised: A video of our results on Hopper and Walker environments can be seen at https://sites.google.com/view/directedinfo-gail .

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Learning separate sub-task policies can help improve the performance of imitation learning when the demonstrated task is complex and has a hierarchical structure.
Revised: Learning separate sub-task policies can help improve the performance of imitation learning when the demonstrated task is complex and has a hierarchical structure.

Original: In this work, we present an algorithm that infers these latent sub-task policies directly from given unstructured and unlabelled expert demonstrations.
Revised: In this work, we present an algorithm that infers these latent sub-task policies directly from given unstructured and unlabelled expert demonstrations.

Original: We model the problem of imitation learning as a directed graph with sub-task latent variables and observed trajectory variables.
Revised: We model the problem of imitation learning as a directed graph with sub-task latent variables and observed trajectory variables.

Original: We use the notion of directed information in a generative adversarial imitation learning framework to learn sub-task and macro policies.
Revised: We use the notion of directed information in a generative adversarial imitation learning framework to learn sub-task and macro policies.

Original: We further show theoretical connections with the options literature as used in hierarchical reinforcement and imitation learning.
Revised: We further show theoretical connections with the options literature as used in hierarchical reinforcement and imitation learning.

Original: We evaluate our method on both discrete and continuous environments.
Revised: We evaluate our method on both discrete and continuous environments.

Original: Our experiments show that our method is able to segment the expert demonstrations into different sub-tasks, learn sub-task speciﬁc policies and also learn a macro-policy that can combines these sub-task.
Revised: Our experiments show that our method is able to segment the expert demonstrations into different sub-tasks, learn sub-task speciﬁc policies and also learn a macro-policy that can combines these sub-task.

Original: A A PPENDIX A. D ERIVATION FOR D IRECTED -I NFO L OSS
Revised: A A PPENDIX A. D ERIVATION FOR D IRECTED -I NFO L OSS

Original: The directed information ﬂow from a sequence X to Y is given by:
Revised: The directed information ﬂow from a sequence X to Y is given by:

Original: [Equation]
Revised: [Equation]

Original: where H ( Y (cid:107) X ) is the causally-conditioned entropy.
Revised: where H ( Y (cid:107) X ) is the causally-conditioned entropy.

Original: Replacing X and Y with the sequences τ and c give,
Revised: Replacing X and Y with the sequences τ and c give,

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Here τ 1: t = ( s 1 , · · · , a t − 1 , s t ) .
Revised: Here τ 1: t = ( s 1 , · · · , a t − 1 , s t ) .

Original: The lower bound in equation 7 requires us to know the true posterior distribution to compute the expectation.
Revised: The lower bound in equation 7 requires us to know the true posterior distribution to compute the expectation.

Original: To avoid sampling from p ( c t | c 1: t − 1 , τ 1: t ) , we use the following,
Revised: To avoid sampling from p ( c t | c 1: t − 1 , τ 1: t ) , we use the following,

Original: [Equation]
Revised: [Equation]

Original: where the last step follows from the causal restriction that future provided variables ( c t ) do not inﬂuence earlier predicted variables ( τ 1: t consists of states up to time t .
Revised: where the last step follows from the causal restriction that future provided variables ( c t ) do not inﬂuence earlier predicted variables ( τ 1: t consists of states up to time t .

Original: c t does not effect state s t )
Revised: c t does not effect state s t )

Original: [Equation]
Revised: [Equation]

Original: Putting the result in equation 8 in equation 7 gives,
Revised: Putting the result in equation 8 in equation 7 gives,

Original: [Equation]
Revised: [Equation]

Original: Directed Info-GAIL
Revised: Directed Info-GAIL

Original: VAE pre-training
Revised: VAE pre-training

Original: Environment
Revised: Environment

Original: Epochs
Revised: Epochs

Original: Batch Size posterior λ
Revised: Batch Size posterior λ

Original: Epochs
Revised: Epochs

Original: Batch Size
Revised: Batch Size

Original: Discrete  0.
Revised: Discrete  0.

Original: Circle-World  0.
Revised: Circle-World  0.

Original: Pendulum (both)  0.
Revised: Pendulum (both)  0.

Original: Hopper-v   0.
Revised: Hopper-v   0.

Original: Walker2d-v   0.
Revised: Walker2d-v   0.

Original: [Figure]
Revised: [Figure]

Original: Thus, by maximizing directed information instead of mutual information, we can learn a posterior distribution over the next latent factor c given the latent factors discovered up to now and the trajectory followed up to now, thereby removing the dependence on the future trajectory.
Revised: Thus, by maximizing directed information instead of mutual information, we can learn a posterior distribution over the next latent factor c given the latent factors discovered up to now and the trajectory followed up to now, thereby removing the dependence on the future trajectory.

Original: In practice, we do not consider the H ( c ) term.
Revised: In practice, we do not consider the H ( c ) term.

Original: This gives us the objective,
Revised: This gives us the objective,

Original: [Equation]
Revised: [Equation]

Original: In practice, we ﬁx q from the VAE pre-training and only minimize over the policy π in equation 4.
Revised: In practice, we ﬁx q from the VAE pre-training and only minimize over the policy π in equation 4.

Original: A. I MPLEMENTATION D ETAILS
Revised: A. I MPLEMENTATION D ETAILS

Original: Table 2 lists the experiment settings for all of the different environments.
Revised: Table 2 lists the experiment settings for all of the different environments.

Original: We use multi-layer perceptrons for our policy (generator), value, reward (discriminator) and posterior function representations.
Revised: We use multi-layer perceptrons for our policy (generator), value, reward (discriminator) and posterior function representations.

Original: Each network consisted of 2 hidden layers with 64 units in each layer and ReLU as our non-linearity function.
Revised: Each network consisted of 2 hidden layers with 64 units in each layer and ReLU as our non-linearity function.

Original: We used Adam (Kingma & Ba, 2014) as our optimizer setting an initial learning rate of 3 e − 4 .
Revised: We used Adam (Kingma & Ba, 2014) as our optimizer setting an initial learning rate of 3 e − 4 .

Original: Further, we used the Proximal Policy Optimization algorithm (Schulman et al., 2017) to train our policy network with (cid:15) = 0 .
Revised: Further, we used the Proximal Policy Optimization algorithm (Schulman et al., 2017) to train our policy network with (cid:15) = 0 .

Original: For the VAE pre-training step we set the VAE learning rate also to 3 e − 4 .
Revised: For the VAE pre-training step we set the VAE learning rate also to 3 e − 4 .

Original: For the Gumbel-Softmax distribution we set an initial temperature τ = 5 . 0 .
Revised: For the Gumbel-Softmax distribution we set an initial temperature τ = 5 . 0 .

Original: The temperature is annealed using using an exponential decay with the following schedule τ = max(0 .
Revised: The temperature is annealed using using an exponential decay with the following schedule τ = max(0 .

Original: 1 , exp − kt ) , where k = 3 e − 3 and t is the current epoch.
Revised: 1 , exp − kt ) , where k = 3 e − 3 and t is the current epoch.

Original: A. C IRCLE -W ORLD SMOOTHING
Revised: A. C IRCLE -W ORLD SMOOTHING

Original: In the Circle-World experiment, we added another loss term L s to VAE pre-training loss L V AE , which penalizes the number of times the latent variable switches from one value to another.
Revised: In the Circle-World experiment, we added another loss term L s to VAE pre-training loss L V AE , which penalizes the number of times the latent variable switches from one value to another.

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: Figure 6 shows the segmentation of expert trajectories with and without the L s term.
Revised: Figure 6 shows the segmentation of expert trajectories with and without the L s term.

Original: We observed that without adding the smoothing penalty, the VAE learns to segment the expert trajectories into semi-circles as shown in Figure 6(a).
Revised: We observed that without adding the smoothing penalty, the VAE learns to segment the expert trajectories into semi-circles as shown in Figure 6(a).

Original: While a valid solution, this does not match with the intuitive segmentation of the task into two sub-tasks of drawing circles in clockwise and counter-clockwise directions.
Revised: While a valid solution, this does not match with the intuitive segmentation of the task into two sub-tasks of drawing circles in clockwise and counter-clockwise directions.

Original: The smoothing term can be thought of as a prior, forcing the network to change the latent variable as few times as possible.
Revised: The smoothing term can be thought of as a prior, forcing the network to change the latent variable as few times as possible.

Original: This helps reach a solution where the network switches between latent variables only when required.
Revised: This helps reach a solution where the network switches between latent variables only when required.

Original: Figure 6(b) shows an example of segmentation obtained on expert trajectories after smoothing.
Revised: Figure 6(b) shows an example of segmentation obtained on expert trajectories after smoothing.

Original: Thus, adding more terms to the VAE pre-training loss can be a good way to introduce priors and bias solutions towards those that match with human notion of sub-tasks.
Revised: Thus, adding more terms to the VAE pre-training loss can be a good way to introduce priors and bias solutions towards those that match with human notion of sub-tasks.

Original: A. PCA VISUALIZATION OF SUB - TASKS
Revised: A. PCA VISUALIZATION OF SUB - TASKS

Original: In Figure 7, we show the plots expert states, reduced in dimensionality using Principal Component
Revised: In Figure 7, we show the plots expert states, reduced in dimensionality using Principal Component

Original: Analysis (PCA), in Hopper and Walker environments.
Revised: Analysis (PCA), in Hopper and Walker environments.

Original: States are color coded by the latent code assigned at these states.
Revised: States are color coded by the latent code assigned at these states.

Original: We reduced the dimension of states in Hopper from 11 to 2 and in Walker from 17 to 3.
Revised: We reduced the dimension of states in Hopper from 11 to 2 and in Walker from 17 to 3.

Original: These low dimensional representations are able to cover ∼ 90% of variance in the states.
Revised: These low dimensional representations are able to cover ∼ 90% of variance in the states.

Original: As can be seen in the ﬁgure, states in different parts of the space get assigned different latent variables.
Revised: As can be seen in the ﬁgure, states in different parts of the space get assigned different latent variables.

Original: This further shows that our proposed approach is able to segment trajectories in such a way so that states that are similar to each other get assigned to the same segment (latent variable).
Revised: This further shows that our proposed approach is able to segment trajectories in such a way so that states that are similar to each other get assigned to the same segment (latent variable).

Original: A. U SING L ARGER C ONTEXT
Revised: A. U SING L ARGER C ONTEXT

Original: For the following discussion we will represent a k -dimensional categorical variable as belonging to ∆ k − 1 simplex.
Revised: For the following discussion we will represent a k -dimensional categorical variable as belonging to ∆ k − 1 simplex.

Revised: To observe how the dimensionality of the sub-task latent variable affects

Original: To observe how the dimensionality of the sub-task latent variable affects our proposed  approach we show results with larger dimensionality for the categorical latent variable c t .
Revised: our proposed [Figure] [Table] approach we show results with larger dimensionality for the categorical latent variable c t .

Original: Since
Revised: Since

Original: DirectedInfo-GAIL infers the sub-tasks in an unsupervised manner, we expect our approach to output meaningful sub-tasks irrespective of the dimensionality of c t .
Revised: DirectedInfo-GAIL infers the sub-tasks in an unsupervised manner, we expect our approach to output meaningful sub-tasks irrespective of the dimensionality of c t .

Original: Figure 8 shows results for using a higher dimensional sub-task latent variable.
Revised: Figure 8 shows results for using a higher dimensional sub-task latent variable.

Original: Precisely, we assume c t to be a 8-dimensional one hot vector, i.e., c t ∈ ∆ 7 .
Revised: Precisely, we assume c t to be a 8-dimensional one hot vector, i.e., c t ∈ ∆ 7 .

Original: As seen in the above ﬁgure, even with a larger context our approach identiﬁes similar basic action primitives as done previously
Revised: As seen in the above ﬁgure, even with a larger context our approach identiﬁes similar basic action primitives as done previously

Original: when c t ∈ ∆ 3 .
Revised: when c t ∈ ∆ 3 .

Original: This shows that despite larger dimensionality our approach is able to reuse appropriate context inferred previously.
Revised: This shows that despite larger dimensionality our approach is able to reuse appropriate context inferred previously.

Original: We also visualize the context values for the low-dimensional state-space embedding obtained by PCA.
Revised: We also visualize the context values for the low-dimensional state-space embedding obtained by PCA.

Original: Although not perfectly identical, these context values are similar to the visualizations observed previously for c t ∈ ∆ 3 .
Revised: Although not perfectly identical, these context values are similar to the visualizations observed previously for c t ∈ ∆ 3 .

Original: Thus our proposed approach is able, to some extent, infer appropriate sub-task representations independent of the dimensionality of the context variable.
Revised: Thus our proposed approach is able, to some extent, infer appropriate sub-task representations independent of the dimensionality of the context variable.

Revised: A. R ESULTS ON O PEN AI R OBOTICS E NVIRONMENT

Revised: We further performed experiments on the FetchPickandPlace-v1 task in OpenAI Gym.

Revised: In each episode of this task, the object and goal locations are selected randomly.

Revised: The robot then must ﬁrst reach and pick the object, and then move it to the goal location.

Revised: We trained agents using both our proposed Directed-Info GAIL and the baseline GAIL approaches.

Revised: We used 500 expert demonstrations.

Revised: While our method was able to learn to segment the expert demonstrations into the Pick and Place sub-tasks correctly, as can be seen in Figure 9 and the videos at https://sites.google.com/view/directedinfo-gail/home#h.p_ 4dsbuC5expkZ , neither our approach, nor GAIL was able to successfully complete the task.

Revised: In our preliminary results, we found that the robot, in both our proposed approach and GAIL, would reach the object but fail to grasp it despite repeated attempts.

Revised: To the best of our knowledge, no other work has successfully trained GAIL on this task either.

Revised: Our preliminary experiments suggested that stronger supervision may be necessary to teach the agent the subtle action of grasping.

Revised: In order to provide this supervision, we additionally trained the policy to minimize the L2 distance between the policy action and the expert action on states in the expert demonstrations.

Revised: At every training step, we compute the discriminator and policy (generator) gradient using the Directed-Info

Revised: GAIL (or in the baseline, GAIL) loss using states and actions generated by the policy.

Revised: Along with this gradient, we also sample a batch of states from the expert demonstrations and compute the policy gradient that minimizes the L2 loss between actions that the policy takes at these states and the actions taken by the expert.

Revised: We weigh these two gradients to train the policy.

Revised: Table 3 shows the returns computed over 100 episodes.

Revised: Adding the L2 measure as an additional loss led to signiﬁcant improvement.

Revised: Our proposed approach Directed-Info GAIL + L2 loss outperforms the baselines.

Revised: Moreover, we believe that this quantitative improvement does not reﬂect the true performance gain obtained using our method.

Revised: The reward function is such that a correct grasp but incorrect movement (e.g. motion in the opposite direction or dropping of the object) is penalized more than a failed grasp.

Revised: Thus, the reward function does not capture the extent to which the task was completed.

Revised: Qualitatively, we observed a much more signiﬁcant difference in performance between the proposed approach and the baseline.

Revised: This can be seen in the sample videos of the success