Learning to Reason: Distilling Hierarchy via Self-Supervision and Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) aims to compute the optimal control policy while an agent interacts with the environment . Recent advances in deep learning enable RL frameworks to utilize deep neural networks to efficiently represent and learn a policy having a flexible and expressive structure . As a result , we ’ ve been witnessing RL agents that already achieved or even exceeded human-level performances in particular tasks ( Mnih et al. , 2015 ; Silver et al. , 2017 ) . The core of intelligence , however , is not just to learn a policy for a particular problem instance , but to solve various multiple tasks or immediately adapt to a new task . Given that a huge computational burden makes it unrealistic to learn an individual policy for each task , an agent should be able to reason about its action . If predictions about consequences of actions are available , e.g. , by using an internal model ( Ha & Schmidhuber , 2018 ; Kaiser et al. , 2019 ) , an intelligent agent can plan a sequence of its actions . Involving planning procedures in a control policy could provide adaptiveness to an agent , but it is often not trivial to learn such a prediction & planning framework : First , it is difficult to obtain the exact internal dynamic model directly represented in high-dimensional state ( observation ) space . Model errors inevitably become larger in the high-dimensional space and are accumulated along the prediction/planning horizon . This prohibits planning methods from producing a valid prediction and so a sensible plan . Second , and perhaps more importantly , planning methods can not help but relying on some dynamic programming or search procedures , which quickly become intractable for problems with high degrees of freedom ( DOFs ) because the size of search space grows exponentially with DOFs , i.e. , the curse of dimensionality ( LaValle , 2006 ) . Crucial evidence found in the cognitive science field is that there exists a certain type of hierarchical structure in the humans ’ motor control scheme addressing the aforementioned fundamental difficulty ( Todorov & Ghahramani , 2003 ; Todorov , 2004 ) . Such a hierarchical structure is known to utilize two levels of parallel control loops , operating in different time scales ; in a coarser scale , the high-level loop generates task-relevant commands for the agent to perform a given task , and then in a finer time scale , the low-level loop maps those commands into control signals while actively reacting to disturbances that the high-level loop could not consider ( e.g. , the spinal cord ) ( Todorov & Ghahramani , 2003 ) . Because the low-level loop does not passively generate control signals from high-level commands , the high-level loop is able to focus only on the task-relevant aspects of the environment dynamics that can be represented in a low-dimensional form . Consequently , this hierarchical structure allows us for efficiently predicting and planning the future states to compute the commands . Motivated by this evidence , we propose a framework , termed `` DISH '' , that DIStills a Hierarchical structure for reasoning and control . As depicted in Fig . 1 , the proposed framework has two levels of hierarchy . The high-level loop represents an agent ’ s current state as a lowdimensional latent state and generates/reasons task-relevant high-level commands by predicting and planning the future in the latent space . The low-level loop receives the high-level commands as well as the current states and maps them into the high-dimensional control signal . Two different types of learning are required to build such a framework : ( i ) a lowdimensional latent representation for an internal model should be obtained from agent ’ s own experiences via self-supervised learning ; ( ii ) a control policy should be learned while interacting with the environment via reinforcement learning . We combined these two learning problems by transforming a multitask RL problem into generative model learning using the control-inference duality ( Levine , 2018 ; Todorov , 2008 ; Rawlik et al. , 2012 ) . In this perspective , an agent equipped with a low-level control policy is viewed as a generative model that outputs trajectories according to high-level commands . Reasoning the high-level commands is then considered as a posterior inference problem ; we introduce a low-dimensional internal model to make this inference tractable . We demonstrate that the proposed framework can learn the compact representation ( 3-dimensional latent states for a humanoid robot having 90-dimensional states ) and the control policy while solving a small number of imitation tasks , and the learned planning and control scheme is immediately applicable to new tasks , e.g. , navigation through a cluttered environment . 2 RELATED WORK . Hierarchical RL : To apply task-specific policies learned from individual RL problems to various tasks , hierarchical structures are often considered where each learned policy serves as a low-level controlller , i.e. , as a `` skill '' , and a high-level controller selects which skills to perform in the context the agent lies at ( Peng et al. , 2018 ; 2019 ; Merel et al. , 2019a ; Lee et al. , 2019 ) . Peng et al . ( 2018 ; 2019 ) trained robust control policies for imitating a broad range of example motion clips and integrated multiple skills into a composite policy capable of executing various tasks . Merel et al . ( 2019a ) similarly trained many imitation policies and utilized them as individual skills that a high-level controller chooses based on the visual inputs . Lee et al . ( 2019 ) included transition policies which help the agent smoothly switch between the skills . Another line of approaches is using continuous-valued latent variables to represent skills ( Co-Reyes et al. , 2018 ; Gupta et al. , 2018 ; Eysenbach et al. , 2019 ; Florensa et al. , 2017 ; Hausman et al. , 2018 ) . Co-Reyes et al . ( 2018 ) proposed an autoencoder-like framework where an encoder compresses trajectories into latent variables , a state decoder reconstructs trajectories , and a policy decoder provides a control policy to follow the reconstructed trajectory . Gupta et al . ( 2018 ) ; Eysenbach et al . ( 2019 ) ; Florensa et al . ( 2017 ) also introduced latent variables to efficiently represent various policies . Instead of using one static latent variable , Merel et al . ( 2019b ) proposed a framework that encodes expert ’ s demonstrations as latent trajectories and infers a latent trajectory from an unseen skill for one-shot imitation . Haarnoja et al . ( 2018a ) proposed a hierarchical structure for RL problems where marginalization of low-level actions provides a new system for high-level action . In their framework , policies at all levels can be learned with different reward functions such that a high-level policy becomes easier to be optimized from the marginalization . Note that the above hierarchical RL approaches train the high-level policy by solving another RL problem ; because the individual skill or the latent variables compress dynamics of the agent , variations of them provide efficient exploration for the high-level RL . Our framework also considers low-dimensional and continuous latent trajectories to represent various policies . Rather than learning a high-level policy , however , our framework learns an internal model with which the high-level module performs reasoning ; the agent can efficiently reason its high-level commands by searching the low-dimensional latent space with the learned internal model . The learned planning/control structure is then directly applicable to new sets of tasks the agent hasn ’ t met during training . Only a few recent works ( Hafner et al. , 2019 ; Sharma et al. , 2019 ) incorporated reasoning processes into high-level modules , but neither of them exploits low-dimensional latent space for planning ( Sharma et al. , 2019 ) nor low-dimensional commands ( Hafner et al. , 2019 ) . Our ablation study in Section 4.1 shows the effectiveness of utilizing both latent states and commands and , to our best knowledge , DISH is the first framework doing so . Model-based RL & Learning to Plan : Model-based RL algorithms attempt to learn the agent ’ s dynamics and utilize the planning and control methods to perform tasks ( Williams et al. , 2017 ; Deisenroth et al. , 2015 ; Chua et al. , 2018 ) . Williams et al . ( 2017 ) ; Chua et al . ( 2018 ) utilized deep neural networks to model the dynamics and adopted the model predictive control method on the learned dynamics ; Deisenroth et al . ( 2015 ) used the Gaussian processes as system dynamics , which leads to the efficient and stable policy search . Though these methods have shown impressive results , they are not directly applicable to systems having high DOFs because high-dimensional modeling is hard to be exact and even advanced planning and control methods are not very scalable to such systems . One exceptional work was proposed by Ha & Schmidhuber ( 2018 ) , where the variational autoencoder and the recurrent neural network are combined to model the dynamics of the observation . They showed that a simple linear policy w.r.t the low-dimensional latent state can control the low DOFs agent , but ( i ) high-DOFs systems require a more complicated policy structure to output highdimensional actions and ( ii ) reasoning ( or planning ) by predicting the future is essential to solve a set of complex tasks . On the other hand , Ha et al . ( 2018a ; b ) trained the low-dimensional latent dynamics from expert ’ s demonstrations and generated motion plans using the learned dynamics ; the high-dimensional motion plans were able to be computed efficiently , but the control policy for executing those plans was not considered . Some recent works have attempted to build the policy network in such way that resembles the advanced planning and optimal control methods : Tamar et al . ( 2016 ) encoded the value iteration procedures into the network ; Okada et al . ( 2017 ) ; Amos et al . ( 2018 ) wired the network so as to resemble the path-integral control and the iterative LQR methods , respectively . The whole policy networks are trained end-to-end and , interestingly , system dynamics and a cost function emerge during the learning procedure . However , these methods were basically designed just to mimic the expert ’ s behaviors , i.e. , addressing inverse RL problems , and also tried to find the control policy directly in the ( possibly high-dimensional ) state space . 3 DISH : DISTILLING HIERARCHY FOR PLANNING AND CONTROL . 3.1 MULTITASK RL AS LATENT VARIABLE MODEL LEARNING . Suppose that a dynamical system with states s ∈ S is controlled by actions a ∈ A , where the states evolve with the stochastic dynamics p ( sk+1|sk , ak ) from the initial states p ( s1 ) . Let r̃k ( sk , ak ) denote a reward function that the agent wants to maximize with the control policy πθ ( ak|sk ) . Reinforcement learning problems are then formulated as the following optimization problem : θ∗ = arg max θ Eqθ ( s1 : K , a1 : K ) [ K∑ k=1 r̃k ( sk , ak ) ] , ( 1 ) where the controlled trajectory distribution qθ is given by : qθ ( s1 : K , a1 : K ) ≡ p ( s1 ) K∏ k=1 p ( sk+1|sk , ak ) πθ ( ak|sk ) . ( 2 ) By introducing an artificial binary random variable ot , called the optimality variable , whose emission probability is given by exponential of a state-dependent reward , i.e . p ( Ok = 1|sk ) = exp ( rk ( sk ) ) , and by defining an appropriate action prior p ( a ) and corresponding the uncontrolled trajectory distribution , p ( s1 : K , a1 : K ) ≡ p ( s1 ) ∏K k=1 p ( sk+1|sk , ak ) p ( ak ) , we can view the above RL problem as a probabilistic inference problem for a graphical model in Fig 2 ( a ) . The objective of such an inference problem is to find the optimal variational parameter , θ , such that the controlled trajectory distribution qθ ( s1 : K , a1 : K ) fits the posterior distribution p ( s1 : K , a1 : K |O1 : K = 1 ) best . More detailed derivations of this duality can be found in Appendix A.2 or in the tutorial paper ( Levine , 2018 ) . Rather than solving one particular task , i.e. , one reward function , agents are often required to perform various tasks . Let T be a set of tasks , and πθ∗t ( ak|sk ) be the optimal policy for t th task , i.e. , θ∗t = arg max θt Eqθt ( s1 : K , a1 : K ) [ K∑ k=1 r̃ ( t ) k ( sk , ak ) ] , ∀t ∈ T . ( 3 ) For high DOF systems , where policies πθt represent a mapping from a high-dimensional state space to a high-dimensional action space , individually optimizing each policy is computationally too expensive . Instead of doing so , we can assume that tasks the agent needs to perform require similar solution properties and consequently the optimal policies have some sort of common structures . We can then introduce a low-dimensional latent variable h ( t ) that , by compressing a particular aspect of πθt over all the policies , each policy can be conditioned on as πθ ( ak|sk , h ( t ) ) . Such a hierarchical structure is depicted as Fig . 2 ( b ) , where h can be interpreted as high-level commands . We can then define the uncontrolled and the task t ’ s controlled trajectory distributions as p ( s1 : K , a1 : K , h1 : K ) ≡ p ( s1 ) K∏ k=1 p ( sk+1|sk , ak ) p ( ak ) p ( hk ) , ( 4 ) q ( t ) θ ( s1 : K , a1 : K , h1 : K ) ≡ p ( s1 ) K∏ k=1 p ( sk+1|sk , ak ) πθ ( ak|sk , hk ) q ( t ) ( hk|sk ) , ( 5 ) receptively . In other words , the control policy πθ is shared across all the tasks , actively mapping high-level commands h , into actual actions , a . Only high-level commands vary with the given task specifications . In the perspective of control as inference , a corresponding inference problem now has two parts : one for the policy parameter θ and another for the task-specific commands h. Note that , if a high-level policy π̄θ ( h|s ) is used to compute high-level commands , the learning problem then becomes the standard Hierarchical RL ( HRL ) . We instead introduce a reasoning module to generate high-level commands which infers the optimal h for a given task t and a current state s by predicting futures . As often used in many HRL methods , the high-level module of the proposed framework operates in a coarser time scale than the low-level policy does . Similar to the latent model learning in Appendix A.3 and the control-inference duality in Appendix A.2 , we can derive the following lower-bound of optimality likelihood L ( t ) for a task t : log pθ ( O ( t ) 1 : K = 1 ) = log ∫ p ( O ( t ) 1 : K = 1|s1 : K ) p ( τ ) q ( t ) θ ( τ ) q ( t ) θ ( τ ) dτ ≥ E q ( t ) θ ( τ ) [ K∑ k=1 r ( t ) k ( sk ) − log πθ ( ak|sk , hk ) p ( ak ) − log q ( t ) ( hk|sk ) p ( hk ) ] ≡ L ( t ) ( θ , q ) , ( 6 ) where τ ≡ ( s1 : K , a1 : K , h1 : K ) . This suggests a novel learning scheme of the hierarchical policy in Equation 5 : ( i ) For a given task t and a fixed low-level policy πθ , high-level commands hk are computed via variational inference . This inference procedure q ( h|s ) should take predictions about future rewards into account to generate h , which can be interpreted as planning . To do so , we build an internal model via self-supervised learning and perform planning with the internal model . ( ii ) With the planning module equipped , a low-level policy πθ ( a|s , h ) generates control actions a as in RL problems , which can be trained using standard deep RL algorithms ( Schulman et al. , 2017 ; Haarnoja et al. , 2018b ) .

Summary: This paper proposes a latent variable model to perform imitation learning. The authors propose the model in the control-as-inference framework and introduce two additional latent variables: one that represents a latent state (z) and another that represents a latent action (h). For the generative model, the authors use a sequence latent variable model. For inferring the latent action, the authors use a particle filter. For inferring the states, the authors use an "Adaptive path-integral autoencoder," though it was unclear where the controls "u" come from. (I assume u is the same as the actions, at which point inferring the states amounts to rollout the policy in the sequence latent variable model). The authors compare to not having the latent states and/or not having the latent actions, and demonstrate that they get better imitation learning scores.