Weakly-Supervised Learning of Disentangled and Interpretable Skills for Hierarchical Reinforcement Learning

1 INTRODUCTION . Deep reinforcement learning ( RL ) has achieved great success for various applications , ranging from playing games ( Mnih et al. , 2013 ; Silver et al. , 2016 ) to complex locomotion and robots control ( Lillicrap et al. , 2015 ; Schulman et al. , 2015 ; 2017 ; Haarnoja et al. , 2017 ) . However , several challenges such as 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 an agent pre-learns reusable skills from prior experiences and hierarchically solve higher-level problems by combining the skills . Two issues need to be resolved for the successful deployment of the hierarchical RL ; how to learn useful skills and how to effectively make use of the skills for various downstream tasks . A possible approach for skills that can be applicable to various downstream tasks is to 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 et al. , 2018 ; Sharma et al. , 2019 ) proposed to combine model-free and model-based RL approaches , where a skill-based predictive model , a dynamics model over the latent space , is trained together with a skill-based policy network.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 is a continuous latent variable . Compared to the entangled 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 a novel WEakly-supervised learning approach for learning Disentangled 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 a weakly-supervised trajectory variational autoencoder ( WET-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. , 2020 ) 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 a 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 representations as skills , the skill-based policy network is trained to generate similar trajectories with the decoder of the WET-VAE by minimizing the KL divergence between two trajectory distributions . However , training a policy to generate a 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 a 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 et al. , 2018 ; Haarnoja et al. , 2018 ; 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-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 a top-down manner . ( Bacon et al. , 2017 ) 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 a generalization ability of the agent to the downstream tasks . To overcome this issue , recent works ( Eysenbach et al. , 2018 ; Achiam et al. , 2018 ; Co-Reyes et al. , 2018 ; Sharma et al. , 2019 ; Campos et al. , 2020 ) 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 a 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 at the risk of falling into over-fitting due to a 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 to 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 a 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 β-VAE , an unsupervised method to 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 , 2018 ; Chen et al. , 2018 ) improved the β-VAE by introducing a total correlation ( TC ) term , ( Locatello 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 on 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 labeling . 3 PRELIMINARIES . Consider a Markov decision process ( MDP ) ( S , A , P , r , ρ0 , γ ) , where S is a set of states , A is a set of action , P : S × A × S → R+ is a transition probability distribution , r : S × A→ R is a reward function , ρ0 : S → R+ is an initial state distribution and γ ∈ ( 0 , 1 ) is a discount factor . We denote a stochastic policy as π : S × A→ R+ . RL has a goal of maximizing the expected discounted sum of rewards for an episode horizon HE : η ( π ) = Eπ [ HE∑ t=0 γtr ( st , at ) ] ( 1 ) Variational autoencoder ( VAE ) optimizes variational the lower bound of the marginal likelihood of dataset . Given an observed datapoint x , the variational lowerbound is defined as : log pθ ( x ) ≥ L ( θ , φ ; x ) = Eqφ ( z |x ) [ log pθ ( x |z ) ] −DKL ( qφ ( z |x ) ‖p ( z ) ) , ( 2 ) where p ( z ) is a prior distribution of a latent variable z , the decoder pθ ( x |z ) is a generative model given a latent z parameterized by θ , and the encoder qφ ( z |x ) 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 KL divergence regularization . In our work , we will focus on the aspect of the generative model of the decoder . 4 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 a policy network to generate the similar trajectories with the predictive model conditioned on skills . The generation process of the trajectory dataset is explained in Appendix A.1.1 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 , ( CoReyes et al. , 2018 ) proposed a trajectory VAE model consisting of the RNN architecture . The trajectory VAE learns latent representations of trajectories , which will be used as skills for a policy . However , this model , which learns the representations in the unsupervised manner , does not consider 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 address this , we propose a weakly-supervised trajectory VAE ( WET-VAE ) model that leverages an inductive bias in the form of weak supervision ( Margonis et al. , 2020 ) to explicitly enforce the model to learn the disentangled representations consisting of desired factors , yielding interpretable skills . Consider a latent-variable generative model p ( τ |z ) to generate a trajectory τ given a latent variable z . We assume the fixed initial state s0 at the origin 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|s0 , z ) = p ( s − s0|z ) . Considering M factors of interest to generate trajectories , the weak supervision can be provided by simply adding a set of M weak labels y = { y1 , ... , yM } to the generative model , where each label ym is one-hot encoded vector for each factor . The idea is that a latent representation z ∈ RM , which can generate the trajectories 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 conditional independence with respect to a given z , the generative model is extended with the labels p ( τ , y |z ) = p ( τ |z ) p ( y1|z ) · · ·p ( yM |z ) . Then , the variational lower bound of the marginal joint distribution p ( τ , y ) can be formulated as follows : L ( θ , φ ; τ , y ) = Eqφ ( z |τ , y ) [ log pθ ( τ , y |z ) ] −DKL ( qφ ( z |τ , y ) ‖p ( z ) ) = Eqφ ( z |τ , y ) [ T∑ t=1 log pθ ( st |s1 : t−1 , z ) + M∑ m=1 log pθ ( y m |z ) ] −DKL ( qφ ( z |τ , y ) ‖p ( z ) ) , ( 3 ) where pθ ( τ |z ) = pθ ( s1|z ) pθ ( s2|s1 , z ) · · · pθ ( sT |s1 , s2 , ... , sT−1 , z ) . Since pθ ( ym |z ) can be understood as a classifier for each factor , the factors should be distinctly embedded in a latent representation z for high classification probability . As a result , this enforces a disentangled representation of the factors . Practically , the scales of the values of the log-likelihoods of the states and labels are different due to the difference in dimensionality . To fill the gap , we introduce a balancing weight γ inspired by ( Margonis et al. , 2020 ) . We also use a weight β > 1 to emphasize the KL divergence term for better disentanglement in the spirit of the β-VAE ( Higgins et al. , 2016 ) . Then , the final objective function becomes : L ( θ , φ ; τ , y , β , γ ) = Eqφ ( z |τ , y ) [ T∑ t=1 log pθ ( st |s1 : t−1 , z ) + γ · M∑ m=1 log pθ ( y m |z ) ] − β ·DKL ( qφ ( z |τ , y ) ‖p ( z ) ) ( 4 ) The WET-VAE model is trained to maximize Equation 4 . To handle the sequential data , we use the RNN architecture with LSTMs as in Figure 1 . With the support of the weak supervision , this model can learn the disentangled representations of trajectories that consist of factors of variation contributing over different time steps .

Summary: 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.