Weakly-Supervised Learning of Disentangled and Interpretable Skills for Hierarchical Reinforcement Learning
Authors: Wonil Song
Sangryul Jeon
Hyesong Choi
Kwanghoon Sohn
Dongbo Min
Venue: ICLR 2022 Submitted (2022)
Keywords: Reinforcement learning
Variational autoencoder
Disentangled representation learning

Abstract: Hierarchical reinforcement learning (RL) usually requires task-agnostic and interpretable skills that can be applicable to various downstream tasks. While many recent works have been proposed to learn such skills for a policy in unsupervised manner, the learned skills are still uninterpretable. To alleviate this, we propose a novel WEakly-supervised learning approach for learning Disentangled and Interpretable Skills (WEDIS) from the continuous latent representations of trajectories. We accomplish this by extending a trajectory variational autoencoder (VAE) to impose an inductive bias with weak labels, which explicitly enforces the trajectory representations to be disentangled into factors of interest that we intend the model to learn. Given the latent representations as skills, a skill-based policy network is trained to generate similar trajectories to the learned decoder of the trajectory VAE. Additionally, we propose to train a policy network with single-step transitions and perform the trajectory-level behaviors at test time with the knowledge on the skills, which simplifies the exploration problem in the training. With a sample-efficient planning strategy based on the skills, we demonstrate that our method is effective in solving the hierarchical RL problems in experiments on several challenging navigation tasks with a long horizon and sparse rewards.

Under review as a conference paper at ICLR 2022
WEAKLY-SUPERVISHH LEARNING OF DISENTANGLED AND INTERPRETABLE SKILLS FOR HIERARCHICAL REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Hierarchical reinforcement learning (RL) usually requires skills that can be appli- cable to various downstream tasks. While many recent works have been proposed to learn such skills in supervised or unsupervised manners, the learned skills are often entangled, which makes it difficult to interpret those skills. To alleviate this, we propose a novel WEakly-supervisee entangled and Interpretable Skills (WEDIS) from the continuous latent represen- tations of trajectories. We accomplish this by extending a trajectory variational autoencoder (VAE) to impose an inductive bias with weak labels, which explicitly enforces the trajectory representations to be disentangled into factors of interest that we intend the model to learn. Given the latent representations as skills, a skill-based policy network is trained to generate similar trajectories to the learned decoder of the trajectory VAE. Additionally, we propose train a policy network with single-step transitions and perform the trajectory-level behaviors at test time with the knowledge on the skills, which simplifies the training procedure for the policy. With sample-efficient planning strategy based on the skills, we show that our method is effective in solving the hierarchical RL problems in experiments on several challenging navigation tasks with long horizon and sparse rewards.
learning approach for Icarning Dis-
INTRODUCTION
Deep reinforcement learning (RI) has achieved great success for various applications, ranging from playing games (Mnih et al. [013 Silver et al., 2016) to complex locomotion and robots control (Lillierap et al., 2015 Schulman et al. 2015 2017 Haarnoja et al. 2017). However, several chal- lenges such sparse rewards or inadaptability to unlearned tasks still hinder its practical usages in real-world problems. To alleviate these challenges, hierarchical RL (Sutton et al. 1999 Dietterich] 2000) has been studied where agent pre-learns reusable skills from prior experiences and hier- archically solve higher-level problems by combining the skills. Two issues need to be resolved for the successful deployment of the hierarchical RL; how learn useful skills and how to effectively make use of the skills for various downstream tasks.
possible approach for skills that can be applicable various downstream tasks is learn without task-specific rewards (Eysenbach et al. 2018). Another way to achieve the useful skills is to make them predictable. To learn those skills, (Co-Reyes al. 2018 Sharma et al. 2019) proposed to combine model-free and model-based RL approaches, where s baiel-based predictive model, dy- namics model over the latent space, is trained together with a saiel-based policy network.By By using the predictive model for model-based planning during testing time, these works showed to efficiently solve various downstream tasks without the need to learn additional higher-level policies. However, since they did not consider how the skill is embedded into the latent space, the factors consisting of the skill often are entangled when the skill a continuous latent variable. Compared to the en- tangled one, the skill consisting of disentangled factors has several advantages in its applicability in that the factors can be separately interpreted and handled.
In this paper, we introduce novel "'Eakly-supervvi and Interpretable Skills (WEDIS) from the continuous latent representations of trajectories that are composed of several generative factors, e.g., speed, direction, and curvature. To this end, we propose
learning approach for learning Disentangled

---

Under review as a conference paper at ICLR 2022
weakly-supervised trajectory variational autoencoder TWET-VAE) model that is an extension of the trajectory VAE (Co-Reyes et al., 2018 consisting of a recurrent neural network (RNN). We leverage the weak labels (Margonis et al. 2020p to enforce an inductive bias on the model, which explicitly enforces the trajectory representations to be disentangled into factors of interest that we intend the model to learn. To train the WET-VAE, we first synthetically generate a trajectory dataset by the combination of several factors of interest, because the trajectories obtained by an online exploration are likely to contain meaningless samples such as random walks. With the trajectory dataset, the WET- VAE model is trained apart from policy network. It is worthy of noting that while this is similar to imitation learning, our data acquisition is much simpler than collecting expert demonstration.
Sequentially, we train a skill-based policy network with the WET-VAE fixed. Given the latent rep- resentations as skills, the skill-based policy network is trained to generate similar trajectories with the decoder of the WET-VAE by minimizing the KI. divergence between two trajectory distribu- tions. However, training policy to generate trajectory given a skill is difficult since it is unlikely to explore the corresponding trajectory in the training procedure. Instead, we propose to train the policy network with the single-step transitions and perform the trajectory-level behaviors in the test time, which can be achieved with the knowledge of the learned skills. This simplifies the training procedure of the policy, and also allows for sample-efficient large-scale planning strategy with the scaled trajectories. In experiments in Mujoco Ant environment, we show that our disentangled and interpretable skills are effective in solving challenging sparse reward and long-horizon problems in 2D navigation in mazes.
2 RELATED WORKS
Numerous approaches (Sutton et al. 1999: Bacon et al. 2017 Florensa et al. 2017 Hausman al. 2018 Haarmoja et al. | Eysenbach et al. 2018; Shankar et al. 2019; Shankar & Gupta 2020 Co-Reyes et al. 2018 Sharma et al. 2019) have explored on learning reusable skills in RL to solve challenging long-horize horizon or sparse reward problems. (Sutton et al. 1999) pioneered a way to control higher-level abstraction by introducing an option-framework, which learns low-level primitives in top-down manner. (Bacon et al. (017)) proposed an option-critic architecture that learns sub-policies of options. Also, several works (Florensa et al. 2017 Hausman et al., 2018 Haarnoja et al.) 2018) introduced to learn skills with multiple tasks in a bottom-up manner. However, designing reward functions still requires expert knowledge and such task-specific rewards may limit generalization ability of the agent to the downstream tasks. To overcome this issue, recent works (Eysenbach et al.) 2 Co-Reyes al. 2018: Sharma et
2018 Achiam al., 2018
2019 Campos et al. (0200) proposed an unsupervised framework that does not require a hand-specified reward function.
Model-based RL methods (Levine et al. 2016 Nagabandi et al., 2018 Chua et al., 2018 Ha & Schmidhuber 2018) aim to learn dynamics model of the environment. While these works are capable of solving unlearned tasks without the needs of an additional learning via planning through the dynamics model, they are often the risk of falling into over-fitting due to huge capacity of the required data to explore the environment. Instead of learning the underlying dynamics, some methods (Co-Reyes et al. 2018 Sharma et al. 2019) attempted combine the model-free and model-based RL for learning a skill-based predictive model and a skill-based policy. Despite the improved results, they still suffer from the lack of the interpretability of the skills.
Learning disentangled latent representations of factors of variation within dataset is beneficial to variety of downstream tasks such as few-shot classification and data generation, thanks to the interpretability of the disentangled factors. (Higgins et al., 2016) proposed 3-VAE, an unsupervised method learn the disentangled representations by modifying the weight of the KL-divergence term of the VAE (Kingma & Welling 2013; Rezende et al. 2014) greater than one. Afterwards, while several variants <Kim & Mnih |018| Chen et al.] 2018) improved the 3-VAE by introducing a total correlation (TC) term, (Locatello llo et al. 2019a) pointed out the inherent limitation of the purely unsupervised approaches and emphasized the need of an inductive bias. Recent works (Locatello) et al. 2019b Shu et al., 2019: Locatello et al., 2020; Margonis et al.) 2020) proposed various forms of weak supervision to encourage the inductive bias to learn the disentangled representations. While there are various categories
the weak labels, we used them in terms of ones that 1) are roughly divided into fewer classes and 2) can be obtained with programming by using the knowledge on the factors without the need for manual laheling.
2

---

Under review as a conference paper at ICLR 2022
3 PRELIMINARIES
Consider a Markov decision process (MDP) (S.A, Po. 7), where is set of states, A is a set of action, P: P:8 x A S + IR is transition probability distribution, r: 8 x A + RR is reward function, Po : S + R is an initial state distribution and 7 E (0, 1) is a discount factor. We denote stochastic policy as S A Rt RL has goal of maximizing the expected discounted sum of rewards for an episode horizon Hg:
HB ((T)) E _ ('r a:)
Variational autoencoder (VAE) optimizes variational the lower bound of the marginal likelihood of dataset. Given observed datapoint x, the variational lowerbound is defined as
log (x)) 2 L(0, Φ : 2) (x1=)) (x] DKL( (=<[)))p))))
where p(2) is prior distribution of a latent variable the decoder (s 2) is generative model given a latent & parameterized by 6, and the encoder qs (=(22z) (z) is an approximate posterior distribution parameterized by Φ. In Equation 2 the first term is the reconstruction term of the autoencoder, and the second term is the KI. divergence regularization. In our work, we will focus on the aspect of the generative model of the decoder.
WEAKLY SUPERVISED LEARNING OF DISENTANGLED AND INTERPRETABLE SKILL (WEDIS)
Our framework consists of three stages; 1) generating trajectory training data with factors of interest 2) training the WET-VAE model, whose decoder is used for the predictive model and 3) training policy network to generate the similar trajectories with the predictive model conditioned on skills. The generation process of the trajectory dataset explained in Appendix A.1.I due to the lack of space. As a notation, we will use superscript for factors and subscript for time steps. The WEDIS algorithm is summarized in Figure 2
4.1 LEARNING DISENTANGLED AND INTERPRETABLE REPRESENTATIONS OF TRAJECTORY
To learn the temporally extended behaviors, (Co- Reyes et al., 2018) proposed trajectory VAE model consisting of the RNN architecture. The tra- jectory VAE learns latent representations of trajec- tories, which will be used as skills for policy, However, this model, which learns the represen- tations in the unsupervised manner, does not con- sider which factors of variation of a trajectory are embedded in the latent space. Thus, the factors that are often entangled make the interpretation of the representations difficult, exposing limitations in further applicability of the learned skills. To ad- dress this, propose weakly-supervised trajec- tory VAE (WET-VAE) model that leverages an in- ductive bias in the form of weak supervision (Mar-) gonis et al. 2020) to explicitly enforce the model to learn the disentangled representations consisting of desired factors, yielding interpretable skills.
EOTM
E
ESTRE
ESTTK
D2
Encodar
Decoder
Computation graph for the WET- VAE model. For weak supervision, we add set of weak labels " MM the model. These labels encourage the model to learn disentangled trajectory representation that consists of corresponding factors. We assume the fixed initial state so at the origin.
Figure 1:
Consider a latent-variable generative model p(T(s))
to generate trajectory given a latent variable We assume the fixed initial state so at the ori- gin as when given other initial states we can obtain the next states with a linear translation based on the initial states such that p(s)so. 2) p(s sol2). Considering M factors of interest to gen- erate trajectories, the weak supervision can be provided by simply adding a set of M weak la- bels y {v) ,..,... M to the generative model, where each label " is one-hot encoded vector
3

---

Under review as a conference paper at ICLR 2022
for each factor. The idea is that latent representation 2 E RM which generate the trajecto- ries based on the M disentangled generative factors, should also be able to reconstruct the factors. Assuming that the trajectory and the factors that are represented as the multiple labels satisfy ditional independence with respect to a given 2, the generative model is extended with the labels p(T, w|2) p(T]=))(v"(<>)- (/y
MM (z). Then, the variational lower bound of the marginal joint distribution P(T, y) can be formulated as follows:
L(O, T. y)
E qu(z v.0) loo (T 8/2)] - ()
M
gole 7.9) log po 81:2-1,
log po po (xt(go(EIT,
(3) (e [11, $2, ST-1 2). Since (a(y [2) can be under- stood as classifier for each factor, the factors should be distinctly embedded in latent representa- tion z for high classification probability. As a result, this enforces a disentangled representation of
where pe(T 2) (e(s) (2) 2)
the factors.
Practically, the scales of the values the log-likelioods the states and labels are different due to the difference in dimensionality. To fill the gap, introduce balancing weight y inspired by (Margonis et al.) 2020). We also use a weight > to emphasize the KL divergence term for better disentanglement in the spirit of the 3AV (Higgins al. 2016). Then, the final objective function becomes:
L(0,0 2: 8,7)
Σ log (1"))12)) - DDi v)l)((=) (4)
E (
log (o: (1 + 7
The WET-VAE model is trained to maximize Equation 4 To handle the sequential data, the RNN architecture with LSTMs as in Figure [ model can learn the disentangled representations of trajectories that consist of factors of variation contributing over different time steps.
With the support of the weak supervision, this
4.2
LEARNING POLICY WITH LEARNED SKILLS
4.2.1
OPTIMIZATION OF POLICY NETWORK TO IMITATE PREDICTIVE MODEL
After training of the WAT,VAE, its decoder po (T(z) is used as skill-based predictive model of tra- jectories. To learn the disentangled and interpretable skills, the policy network is trained to generate similar trajectories with the predictive model based on the skills. It can be accomplished by opti- mizing the policy to minimize the KI. divergence between the distributions of trajectories generated by the two networks.
min EeccalDxaooocttllllo((((11
where và is a parameter of the skiel-based policy network of the trajectories generated by the policy interacting with the environment, which can be computed as (T(2) = f-.. S. I p(s, at) ) (0.[9o:t1.2)) da-- -da Since inferring (T(2) is intractable due to the lack of knowledge on the transition probability distribution s a.) of the underlying dynamics the environment, Equation cannnot optimized by direct backprop- agation with respect to vh. Instead, following (Co-Reyes et al. 2018). we optimize the skill-based policy network with the RL method by rewriting (see Appendix A.2.2 the equation as follows:
Th (a/T, 2) and pe (T|2) is the distribution
max (p.( (lo po)] - (()())))))
where (T(T, 2) = (T())))((z) and hnd is the entropy of the trajectories generated by pol- icy given skills. By using the log-likelihood log po (T(22) calculated by the predictive model for trajectory explored by the policy as a reward function, we can optimize Equation with conven- tional model-free RL algorithms (Schulman al. 2015 2017 Haarnoja et al. 2017 with entropy regularization. In our implementation, we adopted the soft actor-critic (SAC) (Haarnoja et al. 2017) algorithm that includes the entropy regularization a part of optimization.
(T[=))
4

---

Under review as a conference paper at ICLR 2022
Note that in contrast to CCo-Reyes et al. 2018) that simultaneously optimizes both the trajectory VAE and the policy network with trajectory samples collected from an exploration, our method increases stability in the training since the pre-trained decoder provides reliably fixed rewards.
4.2.2 TRAINING POLICY WITH SINGLE-STEP TRANSITIONS
Given a skill, a policy network needs to explore the trajectory similar to the predictive model to receive reward in Equation [ However, it is very unlikely find all trajectories corresponding to each skill by an exploration, which makes the training difficult. To overcome this, we propose to train the policy network with single-step transitions instead of the full trajectories, and perform the trajectory-level behaviors at test time. This can be achieved by exploiting the known knowledge on the factors since the factors of our disentangled skills are interpretable.
Consider a latent E Z = (Zsingle, m C IR IMM++MM M2 that can be separated to 2 single Z single C RR Mi contributing to single-step transition and 2 mmulti Zmuuti S R. contributing over multiple time steps. For example, while speed and direction can be included in Zsingle curva- ture and acceleration can included in gmults, transitions, can generate the same trajectory by a latent zingle in combination with malli That is, with the known relations the factors, we can compute zzingle of the next time step, which has the same effect of the execution of the full latent To this end, we introduce relation functions to relate the relevant factors. For a m th factor n of zsingle the relation function for the factor can be expressed as 24+1 ft (z4") muddti and properly chosen by their relationships. For instance, speed
trajectory consists of a sequence of single-step
can set a relation function for the speed factor together with the acceleration factor such as 4.11 fapeed (2,peed _200)= 2, 2,peed weight 2000 where the weight can be properly chosen by testing the decoder and policy heuristically after training.
In this way, we train the policy network with ziigle over single-step transitions in Equation 6 and make use of the full skill z with the relation functions at test time. As the single-step transition does require handle the temporal information, we use a feedforward network for the policy instead of RNN.
Algorithm 1: WEDIS
Traectory Generator
Require: M generative factors, weak labeling criteria Generate: Trajectory dataset by combination of the factors D [ ( (80,...;87 8T) ) J==
Factors of Interest
Trajectories Dataset
Initialize: parameters e (encoder), 0 (decoder),
leg. Speed, Direction, Curvature) (Generated)
/ (actor), vb (critic), z (value function)
Training Batch
(a) Generation of Training Trajectories
while Training do // WET-VAE training
Sample trajectory batch B ~
Compute labels {y- ,.... yd }
Optimize e and 0 by Equation
Encoder
Decoder
end
Test: Check factors in while Training do
// policy training
Conditioning
Reprasentation
Latent
Imitating
[Skil]
Sample 2 single from prior
(b) Training of WET-VAE
Set 2 = singge 2 meiti) where muuti = (0,....))
Execute actions
(a)s, single)
Collect the samples in replay buffer R
Sample transition batch {&e, ar 81+188 R
Compute reward log pa ($t+1 Optimize v. & and e with the SAC by Equation 6
Policy
end
Trajectories (Actual)
Test: Find relation functions for each factor in 2 single
(c) Training of Policy
Figure 2: Illustration of WEDIS algorithm. Two networks are trained separately. After the training procedure of the two networks, we check the embedded factors by the latent traversal and find the relation functions by heuristically testing the WET-VAE decoder and policy.
5

---

Under review as a conference paper at ICLR 2022
5 PLANNING WITH TRAJECTORY SCALING
Algorithm 2: MPC with trajectory scaling
Require: predictive model po(7 6, z), policy Th (a) single) M relation functions f for each n in géngge,, reward function, episode horizon Hg. primitive horizon Hg, trajectory length T, planning horizon Hp, initial state 30, sample size for i +- to H T') do 80 f- 8)
z²
p(t)a.2)
z
g
(a) Predicted trajectory for
Sample latent sequences {2,..., ZHp In from a distribution
H2
for j + to Hp do // for all N samples [ [s)..... ST) ~ ( 23)
H2
Compute distances between the states, {Ai, AT]
[ (50 + H2 A..... 90 + H₂ (A1 +....... AT)] H1 8T // the last state of "3"
Staling by H2
80
H2
end
Evaluate the rewards of the scaled trajectory sequences
Choose the first latent 21 of the best trajectory (2 (ingle zmmutt) ) 21
(b) Scaled trajectory for planning
for & + 1 to T do
x(als,a, cingte
for & + 1 to Hg do Execute action Th (a stt - 1) (No- 13. zaingle) end
H2 steps
H2 steps
Executing
(2m) mm1 < {2 multe) ) 3 m single x- (2 ) m=
x [ migle mmitti
aalss, magin
end
2t in
f(-1-1 1
pingle muattt
Hg stops
8o 8T-Hz end
[c) Actual trajectory
Figure 3: Illustration trajectory scaling. This trajectory scaling provides large scale planning that is beneficial to solving long-horizon problems.
After the sequential training of the predictive model and policy network, we can solve hicrarchical RL problems without any additional learning at test time via model-based planning. For the model- based planning, we employ the model- predeipree control (MPC) method (Garcia et al., 1989). With dynamics model po ($ St, a) of the environment, the MPC planner generates several numbers of trajectory samples т {Sg, .j...., $H,} for finite planning horizon Hp, and then evaluates each trajectory according a reward function given for task. After choosing the first action of the best trajectory maximizing the planning rewards, the agent executes the chosen action and the MPC planner repeats iteratively this procedure predictive model pe (T 22), we can follow the MPC strategy over the latent space instead of the action space A. In this strategy, we first sample sequences of latents {=1, 22. ZHpp) } and evaluate the latent sequences based on rewards of the trajectories generated by the predictive model given the latents, where each latent generates sub-trajectory of length T.
the next state until the episode horizon Hg. Using the
Even though we can perform the trajectory-level planning based on this predictive model, the scale of the movement lengths from actions of an agent is limited due to its inherent design specification, which generates trivial trajectories with small scales. To perform planning with a meaningful scale of trajectories, we propose trajectory scaling method. In contrast to a RNN policy that processes temporal information at every time step, our feedforward policy network can consistently perform the single-step transition given skill single from 2. By using this, we hold on single for H, steps and change it by the relation functions at every H step, which generates scaled trajectories by H, times. the case of the predicted trajectories p+(T]2) that are simulated for planning, we can simply scale the trajectories by scaling each transition distance. This trajectory scaling method provides large-scale planning strategy that can plan Hp x H, x T steps and execute actions over H. x T steps for each skill, which is beneficial solving long- horizon problems with high sample-efficiene, by reducing the maximum planning horizon to Hg/ (H, x T). 3 summarizes the MPC procedure with trajectory scaling.

---

Under review as a conference paper at ICLR 2022
EXPERIMENTS
In this section, first provide training details and qualitative results of the WET-VAE and the policy. Then, we test our algorithm on goal navigation tasks in mazes in the Mujoco Todorov et al. (012); BBookmman al. (2016) Ant environment (Figure [4], which are challenging hierarchical RL problems. For details of the generation process of the trajectory dataset and labeling criteria, please see Appendix A.I.I The videos of the learned skills and maze navigation are available at http://cooooieec/c///////eeeee
6.1 QUALITATIVE RESULTS
6.1.1 TRAINING OF THE WET- VAE NETWORK
With weak supervision
Without weak supervision
(Direction)
(No
22 (Spred)
(Speed)
(Curvature)
23 [S+0+C)
32
[
Tc)
[d] =
Figure 4: The latent traversals for each skill. (a) ~ (c) show the results of the WET-VAE model, and (d) shows the results the trajectory VAE that does not use any weak supervision, where both
models are trained on trajectories with T = 5. Each row represents the trajectories generated by latent variables changing values corresponding dimensions of the latent variable from -1.5 1.5, while keeping the others to zeros.
Using the trajectory dataset generated with 7 = 5, we first trained the WET-VAE with a three- dimensional continuous latent variable z We the prior distribution of the latent variable an isotropic unit Gaussian distribution p(=) = (o0 1). Also, for the posterior, we set the distribution Gaussian distribution with diagonal covariance matrix.
To demonstrate the benefit of the weak supervision in training the WET-VAE, we plot the latent traversals of trajectories of the decoder Figure 4 For each traversal, the latents are sampled by changing values of one dimension from -1.5 to 1.5, while keeping values of the other dimensions to zeros. Figure 4 shows single-step transitions based on each factor. While 2 and 22 clearly capture the factors on the direction and speed, it looks like that 23 has almost no effect. Since 2 involves the factor of the curvature, this result reasonable. On the other hand, it is shown that z3 controls the curvature of trajectories when the trajectory is generated with T 5 in Figure Ap. These results imply that the three factors are suitably disentangled according to each dimension Furthermore, even when the trajectory length for generation becomes greater than the length of the training data, the contributions of the learned factors consistently remain as shown in Figure with T 10. Note that the weak supervision successfully disentangle factors that contribute
the latent z.

---

Under review as a conference paper at ICLR 2022
different time steps. where the factors of speed and direction affect the each single step, and the curvature factor affects across the multiple time-steps.
On the other hand, Figure 4d shows that some factors of skills learned from the trajectory VAE, which is trained without weak supervision, are entangled. While the speed factor seems to be disen- tangled in the second dimension z², all factors seems to be still entangled in the third dimension 2². Moreover, no factor is embedded in the first dimension 2 as shown in the first row, implying that this element has effect on the generation of the trajectories.
6.1.2 TRAINING OF THE POLICY NETWORK
The policy network is trained in the Mujoco Ant environment with the rewards calculated by the WET-VAE decoder, acting as the predictive model. Since the latent consists of the factors of speed, direction and curvature, it is separated to zasngle
(z=peed 22ir)) and 2 nnnusss Thus in Equation 6 the policy is conditioned on z siingl and the predictive model is conditioned on (ziingle, 0J to provide log-likelshoods for T 1. Also, instead of the full states of the agent, we restrict the states to compute the rewards to two dimensional x y coordinates (Sharma et al. 2019). which enables the policy to be trained with the predictive model trained with 2D trajec- tories. For the relation functions, we set an identity function as f fspeed (z² 12ppedd 2²0")) for speed and f (2 (2- 2 0.2 - 2 cur for direction, where we found the weight by heuristic test after training both networks.
zewr
In Figure 5 we plot latent traversals of trajectories based on each skill, where each row show the trajectories generated by the predictive model and the policy network, respectively. We plot the speed and direction trajectories with Hz 100 and 1, and the curvature trajectory with Hz 40 and T 5. We compare our algorithm against unsupervised skill learning methods (SeCTAR (Co-Reyes et al. 2018), DADS (Sharma et al. 2019)), which are also model-based methods over the latent space. We set Hz = 100 for DADS, and
50 for SeCTAR. According to the algorithmic design, DADS fixed a trajectory length T' 1, and SeCTAR fixed horizon Hy 1. For WEDIS, given each skill, it is shown that the actual trajectories are similar to the corresponding predictive trajectories, which implies that policy acquired disentangled and interpretable skills. On the other hand, for SeCTAR and DADS, the interpretation of each dimension of the skill is unclear.
SeCTAR
DADS
WEDIS (Ours)
Actual Pradicted Path
(Direction)
(Speed)
[Curvature)
Figure 5: The latent traversals of the decoder for each skill. (Top) Predicted trajectories by the predictive model. (Bottom) Actual trajectories of the agent by the policy network. The colors of (red, green, blue, cyan, yellow) correspond to the values of (-1,-0.5, O, 0.5, 1) of latents, respectively.
6.2 QUANTITATIVE RESULTS
To evaluate the performance on the hierarchical RL prob- lems, we set two challenging maze navigation tasks as shown in Figure 6, which cannot be solved with a single skill, but require to appropriately combine several skills. Also, due to the existence of traps that are close to the goals to some extent, these problems require long-horizon reason- ing even for dense rewards cases. addition to SeCTAR and DADS, compare with two ablations, single-step transition (WEDIS-S, T and no trajectory scaling (WEDIS-NTH -NS, H = 1). All of these algo- rithms solve the hierarchical RL. problems with zero-shot via the MPC over the latent space. The
Maze
Maze
Figure 6: Two navigation tasks. The goals are denoted as the red circles.
8

---

Under review as a conference paper at ICLR 2022
task reward is given by negative distance to goal position r(s) = =- g1.- With the task reward and an additional collision penalty, we set the planning reward as rd rd((s) = r()) up 1 (smenlissoo)) for dense reward and rs)s) = w, [((())<<e) - up (==collssoon) for sparse reward, where 1 is an indicator function that outputs one if the condition is true and zero otherwise. W, and up are reward weight and a penalty weight, respectively, and e an acceptable error.
6.2.1 NAVIGATION VIA MODEL-BASE) PLANNING
In the MPC planning, we evaluated the performance on both dense reward and sparse reward cases. As an evaluation metric, we used the averaged distance between the goal and the final position over ten trials of planning. In Figure 7 we plot the evaluation results according to sample size used for the MPC for comparing the sample-efficiency. WEDIS shows the outstanding results for all the cases, even with less planning samples. Especially, the results on the sparse reward demonstrate that WEDIS is effective solving long-horizon problems. WEDIS-S and DADS are superior than WEDIS-NTS, implying larger scale is required for effective planning, despite of the benefit of the trajectory-level prediction. In addition, while both perform temporally-extende over single-step transitions, WEDIS-S outperforms DADS, which implies benefits of the disentan- glement. For example, the controllability wall. For SeCTAR, it has poor performance since the skills do not follow predictions well. Some examples of the MPC results are shown in Figure X where WEDIS plots smooth trajectories over large scales. For the more details of the planning parameters, please Appendix A.1.4
behaviors
speed can be beneficial to avoiding collisions to the
Maze: (Dense)
Mazel (Saare))
Naze2 (Densel
Maze2 iSparsel
DADS
NED
Figure 7: The evaluation results of the MPC the navigation problems. The dashed gray lines denote goal- goal-distances from the origin. WEDIS outperforms the unsupervised skill learning methods.
SeCTAR (T 40, =1)
DADS (T 1,11, 20)
WEDIS:S (T 1, WEDIS-NTS (T 20,11,
WEDIS =,
Figure 8: Examples of the MPC results on sparse rewards (Top) Maze (Bottom) Maze The blue and red lines represent the predictive paths and the actual paths, respectively, and the cross marks indicate the goal positions. The dashed gray lines denote error between the last positions of predictive and actual trajectories each planning step.
7 CONCLUSION
This work has introduced method to learn continuous skills from the disentangled and interpretable representations of trajectories. To do that, we proposed a WET-VAE model by extending the trajec- tory VAE with weak labels. Using the interpretability, single-step transitions and perform the trajectory-leve behaviors at test time, which simplifies the exploration problem and provides effective large-scale planning strategy. In the experiments of challenging navigation tasks, we demonstrated that our method outperforms the unsupervised skill learning methods with higher sample-effiiinncy
proposed to train the policy network with

---

Under review as a conference paper at ICLR 2022
8
REPRODUCIBILITY STATEMENT
For the reproducibility of the proposed method, we will open our code at https://sies..ooggee http:///w..oooooeeo////////////a We also provide the full hyperparameters in Appendix A.]
REFERENCES
Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. arXiv preprint arXi::1807.10299 2018.
Puerre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Victor Campos, Alexander Trott, Caiming Xiong, Richard Socher, Xavier Giró-i Nieto, and Jordi Torres, Explore, discover and learn: Unsupervised discovery of state-covering skills. In Interna- tional Conference on Machine Learning, pp. 1317-1327. PMLR, 2020.
Ricky TQ Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentangle- ment in variational autoencoders. arXiv preprint arXiv:/802.04942. 049442 2018.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep rein- forcement learning in handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018. URL
neur wwt.nnep....ppppee////....e pctppappee/211//1eee.. Paper, pdf
John Co-Reyes, YuXuan Liu, Abhishek Gupta, Benjamin Eysenbach, Pieter Abbeel, and Sergey Levine. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajec- tory embeddings. In International Conference on Machine Learning, pp. 1009-1018, PMLR, 2018.
Thomas G Dietterich. Hierarchical reinforcement learning with the maxq value function decompo- sition. Journal of artificial intelligence research, 13:227-303, 2000.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint wrXiv:1802.06070 2018.
Carlos Florensa, Yan Duan, and Pieter Abbeel. Stochastic neural networks for hierarchical rein- forcement learning. arXiv preprint arXiv:1704.03012 2017.
Carlos E Garcia, David M Prett, and Manfred Morari. Model predictive control: Theory and prac- tice-a a survey, Automatica, 25(3):335-348, 1989.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso-
ciates, Inc., 2018. URL httpps://prroceedgg....ee neure http://ts....ccc///////020/1 cctpappee/c1///iieee.e [2252166622233330222222222222222
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. 2017.
Tuomas Haaroja, Kristian Hartikainen, Pieter Abbeel, and Sergey Levine. Latent space policies for hicrarchical reinforcement learning. In International Conference on Machine Learning, pp- 1851-1860. PMLR, 2018.
10

---

Under review as a conference paper at ICLR 2022
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. In International Conference on Learn ing Representations, 2018.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. 2016.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning. pp. 2649-2658. PMLR, 2018.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint
1eXiv::412.69660 2014.
Diederik Kingma and Max Welling, Auto-encoding variational bayes. arXiv:13126111, 2013.
arXiv preprint
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. The Journal of Machine Learning Research, 17(1):1334-137, 2016.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint 1rXiv:1509.00 02971, 2015.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. international conference on machine learning, pp. 4124. PMLR, 2019a.
Francesco Locatello, Michael Tschannen, Stefan Bauer, Gunnar Rätsch, Bernhard Schölkopf, and Olivier Bachem. Disentangling factors of variation using few labels. arXiv preprint 1rXiv:1905.01258, 2019b.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly- -speevied disentanglement without compromises. In International Confer- ence on Machine Learning, pp. 6348 6359.. PMLR, 2020.
Vasilis Margonis, Athanasios Davvetas, and Iraklis A Klampanos. Wela-vae: Learning alternative disentangled representations using weak labels. arXiv preprint arXiv:2008.0995 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602. 2013.
Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE Interna- tional Conference on Robotics and Automation (ICRA), pp. 7559-7566. IEEE, 2018
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286. PMLR, 2014.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-1897, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.063447 2017.
Tanmay Shankar and Abhinay Gupta. Learning robot skills with temporal variational inference. In International Conference on Machine Learning, pp. 8624-8633, PMLR, 2020.
Tanmay Shankar, Shubham Tulsiani, Lerrel Pinto, and Abhinav Gupta. Discovering motor programs by recomposing demonstrations. In International Conference Learning Representations, 2019.

---

Under review as a conference paper at ICLR 2022
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. arXiv preprint arXiv:1907.0163: 2019.
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disen- tanglement with guarantees. arXiv preprint arXiv:1910.09772. 2019.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche,
Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game go with deep neural networks and tree search. nature, 529(7587):4844455 2016.
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: frame- work for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181- 211, 1999.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012.
A APPENDIX
A.1 IMPLEMENTATION DETAILS
A.1.1 TRAJECTORY DATASET AND WEAK LABELS
To encourage the disentangled representation learning, we created a trajectory dataset that consists of one-hundred thousand numbers of synthetically generated trajectories. Each trajectory is gener- ated from the combination of several generative factors of interest that we intend the model to learn. For 2D navigation, the trajectories consist of sequences of two-dimensional data points of x and coordinates. As the factors, we set three factors of speed, direction, and curvature the trajectories. For the speed, limited the length of each single-step transition up to one, maintaining consis- tently over multiple steps. To allow for only the positive speed and direction, we limited the direction ranging the angles of right-sided semicircle. For the curvature, ranging from -30 to 30 degrees at each transition step. Also, to include straight lines together, we set the angle for curvature which means a sequence of five transition steps that are composed of six points including an initial point. The details of the factors to generate trajectories used the main experiments and additional experiments in the appendices are summarized in Table
rotated the paths with angle
with 0.2 probability. Finally, set the trajectory length T 5,
Factors
Sampling Distribution
Speed Direction
U--1,1) (ID), (O. 1) (2D) U(-n/2, (semi). U(-n, (full)
Curvature
U 7 π /6)
Acceleration
s/5), s sampled speed
Table The generative Factors
While the training proceeds, weak labels are assigned to the batches as one-hot encoded vectors for each factor, which are computed by programming based on the knowledge of the factors. We roughly define the labeling criteria to divide the factors (speed, direction, and curvature) into four (0.25 interval), four (45 degree interval), and three (straight, right, left) segments. For instance, given a trajectory with (speed = 0.1, direction = 40 degree, curvature angle = -10 degree), then the trajectory assigned with label values of one, two, and three for each factor.
A.1.2 WET-VAE
We used one hidden layer with size of (64, 64) for all units in Figure 1 At each transition step, all of the state and one-hot encoded multiple labels are concatenated as an input. To consider state- agnostic transitions, infer distances from the current states and next states instead of directly
12

---

Under review as a conference paper at ICLR 2022
Direction
Curvature
Samples the trajectory dataset
Figure 9: The labeling criteria and the trajectories in the dataset.
predicting next states. Therefore, we feed the distances as state inputs for the encoding LSTM, and the state decoder also outputs mean and variance for the distances. For training losses, we used the negative log-likelihood for the state outputs and the cross-entropy for the label outputs. For activa- tion functions of the output layers, we used linear functions for the mean and log-variance of state decoder, and softmax functions for the label decoders. In Equation 4 set the hyperparame- ters as y = 10, = 9. We trained the model using Adam optimizer (Kingma & Ba 2014) with the fixed learning rate of 2e-4, the batch size of 2048, and 300 epochs.
A.1.3 POLICY NETWORK
For the policy optimization, we used the SAC algorithm (Haarnoja et al, 2017). To learn the position-agnostic behaviors, we used the state input s to the policy π ( 2), excluding the po- sitional information of (x. y) coordinates in Ant and x coordinate in Half-Cheetah
in DADS (Sharma al. 2019p. While we trained the WET-VAE with the maximum speed of one, we com- puted the likeliodd log Pp(T[z) for the reward by reducing the mean by 1/10 and the variance the Ant. For reduced them 1/5 and 1/20, respectively. In the SAC, we used the entropy tem- perature of 0.01. We trained the model using Adam optimizer (Kingma & Ba 2014) with the fixed learning rate of 3e-4, the batch size of 256, and 64 iterations per epoch for totally 30000 epochs. For other hyperparameters, we used the discount factor of 0.99, target smoothing coefficient of 0.01, the latent sampling distribution the unit Gaussian distribution, and two hidden layers of (512, 512) for all of the actor, critic and value function in Ant, and (256, 256) in Half- Chlf-Cheetah.
by 1/20 in Equation 6 considering the allowable length of one-step movement Half-Cheetall
A.1.4 PLANNING
We set primitive horizons H₃ = 20 for DADS and H2 40 for WEDIS and WEDIS-S, which determined by comparing the average speed. For SeCTAR and WEDIS-NTS, set trajectory lengths as T = 40 and T = 20, respectively. According to the primitive horizons and trajectory lengths,
episode horizons HE 300 for DADS and HE 600 for others. For dense reward tasks, we used planning horizons Hp = 4 for all methods except for Hp 3 of WEDIS. In contrast, for sparse reward tasks we used the maximum planning horizons Hg/ (Hg T') for all methods. These planning parameters are summarized in Table 2
For sampling distribution of the skills, DADS that trained from the uniform distribution used uniform distribution from 1 to 1, and others that were trained from the normal Gaussian distribution used a uniform distribution from -1.5 to 1.5 except for the dimension of speed factor of WEDIS and its ablations. For WEDIS and the ablations, we limit the minimum speed as the one sampled from zero since too low speed is useless for the navigation. Note that this strategy is also an advantage of method thanks to the disentanglement and interpretability of the skills. For the predictions, used the mean output from the VET-VAE decoder as the trajectories without sampling. The actions of agent were sampled from the distributions given by the policy.
We used the collision penalty weight wp = 100, the reward weight w, = 5, and the acceptable error 5 in the sparse reward cases. In Figure the goals are positioned the Cartesian coordinate of (24, 12). As the straight paths are blocked on the walls the mazes, the effective distance to goals 36.
13

---

Under review as a conference paper at ICLR 2022
Algorithms
H
H
T
Steps/ Skill
Hp (Dense Sparse)
Sampling Distribution
SeCTAR
600
1
40
40
4/ 15
U(-1,1)
DADS
300
20
20
4/ 15
U 1.5)
WEDIS-S
600
40
40
4/ 15
U(-1.5, 1.5), Uspond
(11,500)
WEDIS-NTS
600
20
20
4 4/30
U(-1.5, 1.5), U Uspeed
-1.5, 0)
WEDIS
600
40
5
200
3/3
U(-1.5, 1.5), Uspeed
(1.550
Table 2: The MPC parameters
A.2 DERIVATION
A.2.1
VARIATIONAL LOWER BOUND
log po( v) log
PO(T, y. z)dz
log
PO(T, vl2)p(2)dz
/ pO M p))))()
log
Co(2|T, v) (z|T, v)dz
/ c0(2|T, v) log (T, yl2)p(z) dz
(0(2|T.8)
00(=|T, v) log (T y1z) + / go (=(T, v) log
p(z) dz
[=]2]T, 7, y) (x log (T g1æ)] - (x = L(0, ò :T:TTJU :.
where the inequality is from Jensen's inequality. The log-likelihood log pe (r, y/2) is decomposed as below.
log (T, y|z) = log (p<(T)2)po
log (s., 822)) .... nM [2)]
log (s
po (sr su:r-1 (1')(1 po() [2)
log I po 8t 11 2)
II po [z)]
T
M
log po (sg
+
log po "0
Therefore, the objective function of WET-VAE becomes
((0,0;T.1) = E (lgg x|=)] DxL( x)l(()))
(=/ra))) log (0 +
log ( [2]] -x
(9)
By adding the additional weights to this equation, finally obtain Equation 4
14

---

Under review as a conference paper at ICLR 2022
A.2.2 POLICY REWARD
Pe (T|2) E ( = =EpCs (x(1(2)log
pT( drl
p(z) S. Pu(T(2) log pe (T(z)
Pe PT(T(z)
drdz
( dr dz
p(=))>(T|2) log
po [T(
p(2) Po log po (r(
p(z) p()) pol
p(z) /. pe (e(() p
H(Pp((T22) Bpstraa)logoooll
(10)
As the above equation is an objective function for a minimization problem, we can obtain Equation 6 by changing the sign.
A.3
ADDITIONAL EXPERIMENTAL RESULTS
A.3.1 WET-VAE
In this section, we first check the effects of the two weights, 3 and 7 in Equation 4 show the results of the three additional experiments to learn other representations: coordinate representation and Cartesian coordinate representation using proper weak labels for each 1D. 3) disentanglement of
Then, we learning polar
representation in 2D. 2) disentanglement of speed and acceleration four factors of speed, direction, curvature and acceleration in 2D. For each of the three additional experiments, we generated the trajectory dataset again by following Table [
Each column in Figure 10 shows the latent traversals of models with all the same hyper-parameters except for B and y. All models are trained with same dataset used Section (e), we did not reuse the model in Section 6.1.I] but reproduced it again for this experiment. These results imply that the proper choices of both 3 and y are required for the training of the WET-VAE.
For model in
Figure II shows the latent traversals of representations of a two-dimensional latent variable R learned with the polar coordinate and Cartesian coordinate. For the polar coordinate representation, used weak labels of radius divided into four segments with 0.25 interval from zero to one, and angle divided into four segments with an interval of 90 degrees from zero to 360 degrees. For the Cartesian coordinate representation, we used weak labels of and y coordinates divided into four segments with 0.5 interval from -1 to 1. For comparison, we also provide the results of learning Cartesian coordinate in an unsupervised manner without weak labels. The three models in Figure II are trained with the same trajectory dataset generated by the combination of the radius (speed) and angle (direction). It is worthy of noting that the uses of different (r and y coordinates) or no labels can result in different representations from the true generative factors (radius and angle). This implies the importance of the inductive bias to learn the desired factors.
To disentangle speed and acceleration factors in 1D, we set the weak labels divided into four seg- ments from -1 to for speed, and three segments of zero, positive and negative values the ac- celeration. Figure 12 shows that the WET-VAE Iearns the disentangled representation of two- dimensional latent variable E ? with speed and direction factors ID. In this case, the factor contributing single-step transition 2 singge is the speed factor sppeed and the factor contributing over multiple time-steps neuls is the acceleration factor 2 noc On the other hand, when there is no weak supervision, the model does not learn the meaningful representations.
In Figure B it is shown that the WET-VAE can learn the disentangled representations with multiple
factors in both z single and z malli where 2 single
(2 apedd dir) and zmulth (2 ace 2 zear). ). In this
15

---

Under review as a conference paper at ICLR 2022
case, the model learns the Nour-dimensional latent representations E IR with weak labels for each factor, where we set the full-cirele range for the direction.
As shown in Figure II and Figure 13 even if the model trained with dataset of the full-circle angles with the weak supervision, the empty regions appear at some angles. The empty regions were also shown differently at each training this reason, we trained the model with the semicircle angles in the main experiments to completely cover the only positive directions even though the model trained with full circle angles can cover the larger range of the direction.
with the same hyperparameters and dataset. In
14
[a] 1, =0
(6)
1c 9, y
(8)
1,
le|
10
Figure 10: The latent traversals from -1.5 to 1.5. (a) (c) show the results from the models trained with all the same hyper-parameters except for and y. The trajectories in all figures are plotted with T 5.
Polar coordinate
Cartesian coordinate
Without supervision
1.00
1.00
0.7
0.55
025
025
035
-025
-0.50
-0.50
1-0.50
-1.00
-1.00
00
05
-1.00
075
0.55
8.50
008
8.00
0.0
-0.50
-0.75
-075
-0.25
-1.
-05
60
05
18
-18
-65
6.
05
1.0
-10
-05
00
0.5
15
T=1
T=1
T
Figure 11: The latent traversals from -1.5 to 1.5 (Left) polar coordinate representation (Middle) Cartesian coordinate representation (Right) representation learned without weak supervision.
16

---

Under review as a conference paper at ICLR 2022
With weak supervision
Without weak supervision
22 (Accaleration x² (Speed)
T=1
7
7
Figure 12: The latent traversals from -1 to WET-VAE can learn the disentangled representations of the acceleration and speed in ID. For visibility, we represent the values of the latents at the axis and use different colors each stme-sep for 5. Since the speed is zero for 2! 0, the first two figures of the second rows are plotted with z² 1
-1-
-1
44
z² (Curvature)
2 (Acceleration]
2 (Direction)
(Speed)
Figure 13:
The latent traversals from -1.5 to 1.5. WET-VAE can learn the disentangled represen- tations of the curvature, acceleration, direction and speed. For visibility, we plot the acceleration figure with the points with different colors for each time-step. instead of the arrowed lines.
A.3.2 POLICY NETWORK ON ANT OF FULL CIRCLE ANGLES AND
Figure 14: Mujoco environments. (Left) Ant (Right) Half-Cheetall
We present the additional experimental results for the Ant and the Calf-Cheetah of Mujoco in Figure Figure 15 shows that the Ant learns the disentangled skills of the speed and direction of the full
14
circle angles in 2D. Figure 16 it is shown that the Half-Cheetah learns the disentangled speed- acceleration skills in 1D.
To train the policy network for the Half-Cheetah, the ID speed-acceleration model Figure 12 is used as the predictive model. After training, the relation function was set as fapord (2 speed 2 ace speed
f
0.1
a.0
A.3.3 MORE PLANNING EXAMPLES
We provide the additional MPC results of the main experiment for the Maze1 and Maze 2 in Figure 17 Additionally, presents results to show that method sample-efficient with the examples in S-shaped maze and obstacles in Figure 18 In each planning procedure, we used a sequence just skills for the maze and just single skill for the obstacle, respectively.
17

---

Under review as a conference paper at ICLR 2022
4-
3
2
2
0
0-
-2
-2-
-3-
-4
-2
0
-4 -2
z'
(Speed)
z² (Direction)
Figure 15: The latent traversals in Ant trained with speed and the direction ranging the full circle angles.
(Acieleration)
32 (ppedd)
26
20
10
T-1
T-6
T-6
Figure 16: The latent traversals from -1 to Half-Cheetah. (Top) predicted trajectories by the predictive model. (Bottom) actual trajectories by the Half-Cheetah. For visibility, we represent the values of the latents at the axis and use different colors at each time-step for T =
Figure 17: Additional MPC results. (Top) Mazel. (Bottom) Maze2.
Figure 18: Additional MPC results. (Top) S-shaped maze. We use T = and H. = 50 for sequence two skills. (Bottom) Obstacles. T' and H, = 50 for single skill.
18


Reviews:
[{"review_id": "nGcwDLt5-Pe", "reviewer": "Reviewer_onQA", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n* Interesting visualizations of learned interpretable skills on the Ant navigation task. Additionally, the paper showed that combining interpretable skills with model-based planning (MPC) for long-horizon tasks can result in better performance on maze navigation.\n\nWeaknesses:\n\n(1) The data requirements for the proposed method seem impractical for most real applications. The method requires a synthetically generated dataset (in their experiment, they used 100K trajectories), which are generated by varying the generative factors. If I understood the dataset generation procedure correctly, generating this dataset (of trajectories with varying speeds/directions/acceleration) requires that you already have expert policies that can run at various speeds/directions/accelerations.\n\n1a) How does the dataset size & quality affect the quality of the learned disentanglement, and performance of the learned policy?\n\n1b) (Clarification) Does the 100K trajectory dataset consist of the full observations from Ant/Half-Cheetah (e.g., including joint velocities)? If so, were expert policies and a simulation environment needed to synthetically generate this dataset?\n\n1c) (Clarification) Did the other methods (SeCTAR, DADS) make use of the same 100K trajectory dataset (for fairer comparison), and if so, how?\n\n(2) The idea of using weakly-supervised disentangled representations in order to learn an “interpretable” policy behavior, and comparing against unsupervised representation learning methods such as VAE, has also already been studied in a prior work [1], which was not discussed/cited in the paper. The method in [1] also learns disentangled latent encodings Z, which is used to condition the policy, leading to “interpretable” skills. \n\n[1] “Weakly-Supervised RL for Controllable Behavior”, NeurIPS 2020 https://arxiv.org/abs/2004.02860\n\n(3) In general, I felt that the writing & presentation could be improved.\n\n3a) I had trouble understanding what “latent traversals of trajectories of the decoder” mean (in Figures 4, 5, 10, 11, 13, 15, 16). For example, in Figure 4, could you clarify what the X- and Y-axis labels are, and what each line represents? I am guessing that each line corresponds to a different latent variable (by varying one dimension, while keeping the rest fixed), and the X/Y axes are positions of the ant?\n\n3b) Could you elaborate on *why* the disentanglement helps with MPC planning & result in higher performance (Fig 7)? Could it possibly be because the learned skills are more distinguishable from one another? Can this be verified empirically somehow?\n\nMinor typos:\nPage 3: “decesion\" -> “decision”\nPage 3: “We will focus on the aspect of the generative model of the decoder”: I wasn’t sure what this sentence means. Did you perhaps mean: “we will focus on the generative modelling aspect of the decoder”?\nPage 3: “which learns the representations in the unsupervised manner” ->  “which learns the representations in an unsupervised manner”\nPage 5: “in the Equation 6” -> “in Equation 6”\nPage 7: “based on the each factor” -> “based on each factor”\nFigure 3: “solving the long-horizon problems” -> “solving long-horizon problems”\nFigure 15: “the direction ranging the full circle angles” could be better worded.\nFigure 16: “half-cheetah” -> “Half-Cheetah”\n", "summary_of_the_paper": "The paper proposes a method that learns disentangled skill representations, and shows qualitative & quantitative results on Mujoco Ant navigation. (1) They first synthetically generate a trajectory dataset by the combination of different factors. (2) They train a trajectory VAE (Co-Reyes ‘18) that enforces the learned trajectory representations to be disentangled using weak labels (Margonis ‘20). (3) Then, keeping the pre-trained trajectory VAE fixed, they learn a skill-based policy to generate similar trajectories to the learned decoder of the trajectory VAE by minimizing the KL divergence between the trajectory distributions.", "summary_of_the_review": "My main concern is regarding the practicality of the proposed method (see Weakness #1 above). I also feel that discussion of prior work and overall paper presentation could be improved.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "jfEllfls-s-", "reviewer": "Reviewer_FYox", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I found the method introduced in the paper to be somewhat complex, and the experimental section to not justifying this complexity. The overall ingredients (VAE on trajectories, train skill policies on a resulting reward signal) are among the current techniques used in skill discovery (e.g., OPAL (https://arxiv.org/abs/2010.13611), EDL (http://arxiv.org/abs/2002.03647)). While the authors show gains over purely unsupervised skill learning methods (DADS, SeCTAR) on the downstream navigation tasks, my main concern is the way that supervision is introduced in WEDIS. I have trouble understanding the motivation for an auto-encoder setup when reference trajectories can be generated beforehand. Why not simply pre-train a policy with a dense reward to explicitly follow these trajectories, conditioned explicitly on the salient factors (direction, speed, curvature)? This should at least be a baseline in the paper. I can see how the training setup in WEDIS could, in principle, allow for more generality, but then again it requires the ability to generate reference trajectories (which in turn would make it possible to train policies to follow them directly). The experiments also lack a comparison to plain Soft Actor-Critic.\n\nFinally, I don't really agree with the premise that current skill learning methods are fundamentally hindered by a lack interpretability. While it's true that the availability of a model enables the usage of standard planning methods like MPC, several works also showed effective hierarchical control with RL, where explicit semantics are not strictly required (e.g., DIAYN, OPAL (https://arxiv.org/abs/2010.13611) or NPMP (https://arxiv.org/abs/1811.11711)).\n\nThe writing would benefit from an overhaul wrt grammar.", "summary_of_the_paper": "The paper at hand proposes a new framework for pre-training skill policies (WEDIS), and use them for control in a hierarchical setup with MPC. The main idea is that skill policies should follow a set of generated trajectories based on some salient factors. For this, the authors train a VAE that will then provide both the control variables (latent encoding) and a predictive model of the trajectory (decoder). The skill policy is trained to match the trajectories of the predictive model, and the predictive model is then used for MPC.", "summary_of_the_review": "In my view, the experiments in the paper fail to justify the complexity of the method that was introduced. With the amount of supervision required, simpler and potentially more effective approaches come to mind but are missing in the comparison.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "iZxk2MphfDO", "reviewer": "Reviewer_CPCV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "My main concerns with the work are its lack of clarity and how many decisions related to the design of the algorithm are not well justified.\n\n- As a non-expert, I found the paper hard to digest. In the process of reviewing this paper, I read some other already published works and found them easier to follow and better motivated. For instance, a clean definition of what is meant by \"skill\" would have been helpful.\n- The weak-labels used to improve disentanglement are computed from the synthetic trajectories. From what I understand, they contain a lot of information about the direction, speed and curvature. If the goal of the WET-VAE is to extract direction, curvature and speed from the trajectories, why not using supervised learning to do that? For instance, one could train a neural network to output direction, speed and curvature using the supervision of the computed labels y. Then, one could use this encoder trained with supervision to train a generative model conditioned on these.\n- The advantage of disentanglement is not obvious here (apart from interpretation). It seems to be used only to separate z into z^single and z^multi. But it is not clear if this is useful or why it should be useful. This connects to the following points related to the MPC approach.\n- In Algorithm 2 when executing the policy, I did not understand why it was important to exclude z^{multi} from the latent representation fed to the policy. An ablation study would be convincing. \n- In Algorithm 2, the role of the hand-crafted functions f is unclear to me. My understanding is the following: when executing a rescaled plan, it is used to update skill z from one rescale step to another. What is the advantage of doing so instead of just using z*_1 across the subtrajectory?\n- In Algorithm 2, at each step, planning in the latent space is performed to pick the best action. What is the distribution from which the potential trajectories z_1, …, z_{H_p} are sampled? Is it uniform?\n- In some sense, both the artificially generated sequences and the scaling factor H_z must be chosen by the user. It seems H_z is thus redundant since the appropriate sequence length could have been chosen  prior to the sequence generation, no? This question also makes me doubt the usefulness of the hand-crafted functions f, which I suspect are useful only because of the rescaling.\n\n", "summary_of_the_paper": "This work proposes to learn representations of artificially generated trajectories using a VAE-based approach that relies on weak labels to improve disentanglement. Then, a skill-based policy (i.e. a policy conditioned on the learned sequence representation) is trained to imitate the learned model via soft actor-critic. Moreover, the paper introduces a hierarchical planning strategy that explores at the “aggregated sequence level”, i.e. the representation space, and allows for “trajectory rescaling”. The authors show empirically both the disentanglement of WET-VAE and the advantage of their hierarchical approach.  ", "summary_of_the_review": "Given the lack of clarity and the multiple algorithm design decisions that are not sufficiently justified, I cannot recommend acceptation. I cannot comment on the novelty of the contribution as I am not well versed into hierarchical reinforcement learning.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4mo2fD3mcyI", "reviewer": "Reviewer_BoDo", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n- The related work section provides a comprehensive overview of relevant prior works.\n\n- The provided planning visualizations clearly show the differences between the learned skill spaces and how they influence downstream learning.\n\n- The appendix provides detailed instructions for generating the used dataset and reproducing the policy training.\n\n\n## Weaknesses\nI have a number of major concerns with respect to the proposed problem formulation and approach:\n\n1. The proposed approach is compared to unsupervised skill learning approaches (like DIAYN, SeCTAR, DADS), but it assumes access to an automatically pre-generated dataset of training trajectories. This directly violates the premise of unsupervised skill-*discovery* approaches: if we were able to generate meaningful trajectories in the first place, we would not need to train an agent to discover them. This becomes more apparent if we move from simple 2D ant navigation problems as used in the paper's evaluation to more complex problems like kitchen manipulation with a robot arm: there it would be very challenging to pre-define object interaction primitives etc, which is why unsupervised skill discovery approaches were proposed in the first place. Thus, I believe the problem formulation of the paper is fundamentally flawed.\n\n2. The proposed approach does not learn disentangled skills with \"weak supervision\" but assumes fully annotated training trajectories with all relevant factors of variation given for every time step. Thus I see the proposed method as a fully supervised approach to learning a disentangled latent space instead with strong assumptions on the available supervision. Defining this set of \"interesting factors of variation\" in the first place requires considerable human effort and can be non-trivial in harder problems.\n\n3. The paper does not motivate clearly *why* it would be beneficial to train a latent space with disentanglement of the given factors of variation. This should be much clearer motivated in the introduction. Currently it only vaguely mentions that a lack of disentanglement and interpretability of skills \"limits their applicability\".\n\n4. The explanation of the model is in part unclear. I am specifically confused about section 4.2.2 (see questions below). This also makes understanding the downstream usage of the model challenging.\n\n5. Section 5 introduces an approach which the paper calls \"trajectory scaling\". From my understanding it simply describes the common RL practice of \"action repeat\" which is eg used by default in many Atari games. Although it is less common in MPC-based planning approaches this connection should be made in the paper.\n\n6. The experimental comparisons to *truly* unsupervised approaches like SeCTAR and DADS are not fair in my understanding: these approaches do not have access to a pre-collected dataset of demonstrations and thus need to learn to explore the environment while learning the skill policy. This is a much harder task, but this distinction is never acknowledged in the experimental section.\n\n7. The writing of the paper in part impairs its clarity. Many formulations stay vague or are not put in appropriate context, eg from the abstract without any further explanation of \"single-step transitions\" or \"trajectory-level behaviors\": \"Additionally, we propose to train a policy network with single- step transitions and perform the trajectory-level behaviors at test time with the knowledge on the skills, which simplifies the exploration problem in the training.\" In other instances the writing does not seem to explain all the details necessary to understand the introduced concepts (see point (4) above).\n\n\n## Questions\n- I am confused why we need to train single-step transition models for reward computation instead of training with the pre-trained trajectory-level models as reward-computers. Is this to get a denser reward function? How does the model in Figure 1 change when we split the latent variable into z_single and z_multi? How does \"heuristically\" finding the relation functions between z_single and z_multi work exactly?\n", "summary_of_the_paper": "The paper proposes an approach for learning interpretable skill embeddings by leveraging an automatically generated dataset of trajectories annotated with factors of variation like direction, speed and curvature which should be used to structure the learned skill space. It shows that these factors can indeed be used to learn a latent space with disentangled factors of variation which can be used via an MPC-based planning method for solving downstream tasks in a OpenAI Gym ant environment.", "summary_of_the_review": "In my opinion the assumptions made in the paper contradict the premise of the unsupervised skill discovery problem. Further, the approach is not clearly described, the experimental evaluation is not fair due to differing assumptions and the overall writing of the paper needs improvement. Thus, I do not recommend acceptance of the submission.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors present a method for learning disentangled skill representation that uses weak supervision. The reviewers mentioned that the paper tackles an important problem, delivers interesting and novel visualizations of the learned skills, and positions the paper well in the context of related work. The reviewers also point out several points of criticism: the complexity of the method, lack of convincing comparisons to baselines that utilize the same amount of data and the quality of writing, among others. I encourage the authors to address these points in the future version of the paper.
Decision: Reject