Training Transition Policies via Distribution Matching for Complex Tasks

1 INTRODUCTION . While Reinforcement Learning ( RL ) has made significant improvements for a wide variety of continuous tasks such as locomotion ( Lillicrap et al. , 2016 ; Heess et al. , 2017 ) and robotic manipulation ( Ghosh et al. , 2018 ) , direct end-to-end training on tasks that require complex behaviors often fails . When faced with a complex task , humans may solve simpler subtasks first and combine them . Analogously , hierachical reinforcement learning ( HRL ) uses multiple levels of policies , where the actions of a higher-level meta-controller represent which lower-level policy to execute ( Sutton et al. , 1999 ; Dayan & Hinton , 1992 ) . A multitude of methods has been proposed to train HRL , such as assigning subgoals for the lower-level policies ( Kulkarni et al. , 2016 ) , employing pre-trained lower-level policies ( Frans et al. , 2018 ) , and off-policy RL ( Nachum et al. , 2018 ) . All HRL methods face a common challenge : switching smoothly from one lower-level policy to another when the meta-controller directs that a switch should occur . For example , running hurdles may require switching between running and jumping repeatedly while maintaining forward momentum . Previous approaches to this challenge ( Frans et al. , 2018 ; Andreas et al. , 2017 ; Lee et al. , 2019 ) involve either retraining existing lower-level policies and/or introducing a new policy to execute the switch . We focus on the latter and call the new policy the transition policy , following Lee et al . ( 2019 ) . A transition policy needs to switch between the states produced by one lower-level policy to an appropriate starting state for another . A natural way to train such a policy is to give a reward depending on whether the transition is successfully executed or not . However , this method induces a sparse reward that is difficult to learn from due to the temporal credit assignment problem ( Sutton , 1984 ) . Instead , we use inverse reinforcement learning ( IRL ) techniques to match the distribution of states produced by the transition policy to the distribution of states expected by the next lower-level policy ( Figure 1 ) . While the goal of IRL is infer an expert ’ s reward function from demonstrations of optimal behavior ( Ng et al. , 2000 ; Ziebart et al. , 2008 ) , some recent approaches to IRL ( e.g. , Finn et al . ( 2016 ) ; Ho & Ermon ( 2016 ) ; Fu et al . ( 2018 ) ) learn a policy that imitates the expert ’ s demonstrations in the process . In particular , Ho & Ermon ( 2016 ) propose a generative adversarial network ( GAN ) structure ( Goodfellow et al. , 2014 ) that matches the distribution of the policy ( generator ) with the distribution of the given data ( expert demonstrations ) . We use this GAN structure for distribution matching , and , consequently , we avoid the problem of designing an explicit reward function for transition policies . Although the transition policy is trained to produce the distribution of states and actions that is expected by the next lower-level policy , not all states in the distribution ’ s support have equal probability of success . To increase the rate of successful execution , we introduce a deep Q-network ( DQN ) ( Mnih et al. , 2013 ) to govern the switch from the transition policy to the next lower-level policy . The DQN has two actions , switch or stay , and a simple reward function where a positive reward is obtained if the next lower-level policy is successfully executed ; otherwise , a negative reward is given . The issue of sparse rewards is present here , but it is less of an issue because of the simple decision space . The main contribution of this paper is a new method for training transition policies for transitioning between lower-level policies . The two parts of our approach , the transition policy that matches the distribution of states and the DQN that determines when to transfer control , are both essential . Without the DQN , the state selected for transition may be familiar to the next policy , but lead to a poor success rate . Without the transition policy , there may be no appropriate state for the DQN to transition from . The proposed method reduces issues with sparse and delayed rewards that are encountered by prior approaches . We demonstrate our method on bipedal locomotion and arm manipulation tasks created by Lee et al . ( 2019 ) , in which the agent requires to have diverse skills and appropriately utilize them for each of the tasks . 2 PRELIMINARIES . The problem is formulated as a infinite-horizon Markov decision process ( MDP ) ( Puterman , 2014 ; Sutton & Barto , 2018 ) defined by the tupleM = ( S , A , P , R , γ , µ ) . The agent takes action at ∈ A for the current state st ∈ S at time t , then the reward function R : S × A → R returns the reward for ( st , at ) , and the next state st+1 is determined by the transition probability P ( st+1|st , at ) . γ is the discount factor and µ denotes the initial states distribution . The agent follows a policy πθ ( at|st ) that produces a distribution of actions for each state . The objective of reinforcement learning ( RL ) is to find the optimal θ that maximizes the expected discounted reward : θ∗ = argmax θ E s0∼µ at∼πθ ( ·|st ) st+1∼P ( ·|st , at ) [ ∞∑ t=0 γtR ( st , at ) ] . ( 1 ) 2.1 INVERSE REINFORCEMENT LEARNING . Whereas standard RL trains a policy to maximize a reward function , inverse reinforcement learning ( IRL ) ( Ng et al. , 2000 ; Ziebart et al. , 2008 ) seeks to learn the reward function from expert demonstrations of an optimal policy . Several approaches to the IRL problem ( Ho & Ermon , 2016 ; Fu et al. , 2018 ) adopt a GAN structure ( Goodfellow et al. , 2014 ) —the policy ( generator ) tries to select actions that the discriminator can not distinguish from the expert ’ s . As a by-product , the generator produces a policy with the same distribution as the expert ’ s at convergence . We use the state-action discriminator ( Fu et al. , 2018 ) rather than the trajectory-centric discriminator ( Finn et al. , 2016 ; Ho & Ermon , 2016 ) : Dψ , φ ( st , at , st+1 ) = exp { fψ , φ ( st , at , st+1 ) } exp { fψ , φ ( st , at , st+1 ) } + πθ ( at|st ) , ( 2 ) where fψ , φ ( st , at , st+1 ) = gψ ( st , at ) +γhφ ( st+1 ) −hφ ( st ) at time t , gψ is the reward approximator , and φ are the parameters of a shaping term hφ , which alleviates unwanted shaping for gψ . This discriminator Dψ , φ is trained to distinguish data from the expert and the generator ’ s policy via binary logistic regression , and the policy πθ is updated with Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) according to the rewards rψ , φ provided by the discriminator : rψ , φ ( st , at , st+1 ) = logDψ , φ ( st , at , st+1 ) − log ( 1−Dψ , φ ( st , at , st+1 ) ) ( 3 ) 2.2 DEEP Q-LEARNING . Deep Q-learning ( Riedmiller , 2005 ; Mnih et al. , 2013 ) is a variant of Q-learning ( Watkins & Dayan , 1992 ) that produces a neural network called a deep Q-network ( DQN ) that aims to learn a deterministic optimal action from among a discrete set of actions for each state . Deep Q-learning suffers from overestimating yt ( Thrun & Schwartz , 1993 ) because of the max term , which impedes the learning process . Double Q-learning ( van Hasselt et al. , 2016 ) addresses this problem by decoupling the networks for choosing the action and evaluating value . The target networkQθ′ makes learning more stable by providing a target yt that is fixed for long periods of time : y′t = rt + γ Qθ′ ( st+1 , argmaxat+1Qθ ( st+1 , at+1 ) ) ( 4 ) The DQN is updated using a loss function such as the mean squared error ( y′t − Qθ ( st , at ) ) 2 . We use train a DQN with double Q-learning to get a signal that tells us when to switch from a transition policy to the next lower-level policy . 3 APPROACH . In this section , we present a methodology for transitioning between lower-level policies in hierarchical reinforcement learning ( HRL ) . We assume that we receive pre-trained lower-level subtask Algorithm 1 Training Transition Policy πa , b 1 : Input : pre-trained policies { πa , πb } , transition interval 2 : Initialize transition policy ( generator ) πa , b and discriminator Dψ , φ Equation ( 2 ) 3 : Collect the states Sπa of πa at the start of the transition interval 4 : Collect the trajectories ( states and actions ) Tπb of πb during the transition interval 5 : Initialize transition policy πa , b and discriminator Dψ , φ 6 : for i = 1 to n do 7 : for j = 1 to m do 8 : Sample s ∼ Sπa and set the state of πa , b to s 9 : Collect the trajectory Tπa , b generated by the current πa , b 10 : end for 11 : Train Dψ , φ with binary logistic regression to distinguish Tπa , b and Tπb 12 : Calculate reward rψ , φ set for all ( st , at , st+1 ) ∈ Tπa , b with Equation ( 3 ) 13 : Optimize πa , b with respect to Tπa , b and rψ , φ 14 : end for policies as input along with transition intervals—windows where a transition between two specific lower-level policies should occur . We seek to maximize performance on a complex task by combining the pre-trained policies without re-training them . To do this , we train a transition policy for each pair of pre-trained policies that we must transition between . The transition policy is activated at the start of each transition interval . Then , a DQN governs the switch from the transition policy to the next prescribed pre-trained policy . The DQN improves the rate of transition success relative to simply switching at the end of the transition interval.1 In Section 3.1 , we describe our training procedure for transition policies . In Section 3.2 , we describe how we collect data and design rewards for DQNs . Our framework is summarized by Figure 2 . 3.1 TRAINING THE TRANSITION POLICY . Let { π1 , π2 , ... , πn } denote n pre-trained policies and let πa , b denote a transition policy that aims to connect policy πa to πb . The meta-controller dictates a transition interval for each transition between pre-trained policies—a period of time when the transition should take place . If we were to simply switch from πa to πb at the start of this interval , πb can easily fail because it has never been trained on the last state produced by πa . To prevent this problem , πa , b needs to be able to start from any of the last states produced by πa and lead to a state from which πb can successfully perform the next task . Which states are favorable for starting πb ? A possible answer is the states produced by πb itself . We train πa , b to have the same distribution as πb with IRL ( Ng et al. , 2000 ; Ziebart et al. , 2008 ; Ho & Ermon , 2016 ; Fu et al. , 2018 ) . We specifically choose adversarial inverse reinforcement learning ( AIRL ) ( Fu et al. , 2018 ) , which has the GAN ( Goodfellow et al. , 2014 ) structure . The training procedure for πa , b is described in Algorithm 1. πa , b is trained with two classes of data : one from πa and one from πb . From πa , we collect the states Sπa of πa at the start of the transition interval . From πb , we collect the trajectories Tπb of πb during the the transition interval . The initial state of the generator πa , b is set to s , sampled from Sπa . Then , πa , b tries to generate trajectories that fool the discriminator D ( Equation 2 ) into classifying it as an example of Tπb . The discriminator D is trained with binary logistic regression to distinguish between the data from πb vs. πa , b . WhenD can no longer distinguish between them , πa , b can be said to have a similar distribution to πb . For example , we prepared a pre-trained Walking forward policy π1 and Jumping policy π2 for one of our experimental environments , Hurdle , manually shaping rewards for these two subtask policies . We use a transition interval of 3–4 meters from the hurdle . Thus , Sπ1 is the last states of π1—the sampled set of states when the agent ’ s location is 4 meters from the hurdle . Tπ2 is the trajectories 1We assume that the environment is static enough that it is not critical that the meta-controller has the opportunity to replan after the transition policy terminates . In principle , such replanning could be allowed by introducing an additional layer of transition policies , i.e. , by training training meta-transition policies that transition between transition policies and actions . Algorithm 2 Training Deep Q Network qa , b 1 : Input : pre-trained policies { πa , πb } , transition policy πa , b , transition interval 2 : Initialize qa , b and replay buffer B 3 : for i = 1 to n do 4 : Set environment to state of πa at the start of the transition interval 5 : while the current state is in the transition interval 6 : at ∼ πa , b ( st ) 7 : Execute environment for ( st , at ) and get st+1 and success-fail-alive signal τ 8 : Run qa , b and get an action aq 9 : if aq = switch 10 : set s = st+1 and break 11 : else if aq = stay 12 : Store ( st , aq , rf , st+1 ) in B if τ = fail or st+1 is outside the transition interval 13 : Store ( st , aq , 0 , st+1 ) in B if τ = alive 14 : end if 15 : end while 16 : while πb has not achieved subgoal and has not failed 17 : at ∼ πb ( st ) 18 : Execute environment for ( st , at ) and get st+1 and signal τ 19 : end while 20 : Store ( s , aq , rs , st+1 ) in B if τ = success 21 : Store ( s , aq , rf , st+1 ) in B if τ = fail 22 : Update qa , b with a minibatch from B according to Equation ( 4 ) 23 : end for of π2 from 3–4 meters from the hurdle . The transition policy π1,2 is initialized with one of the last states of π1 and learns to match the distribution π2 just before the jump is initiated .

Summary: The paper addresses the problem of learning to subsequently execute tasks, and transitioning from a first task to a second task. The work contributes a method to train transition policies, and a method to decide when to start executing the second policy, stopping the transition policy. The approach is evaluated on a few simulated robot locomotion/manipulation tasks.