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 considers the problem of solving a complex task requiring different skills by combining "subtask" policies pretrained for each individual skill. One approach here consists in finding ways to transition smoothly between the policy $b$ for a subtask and the policy $b$ for the next subtask. The paper proposes to do that based on training a "transition policy" that starts from a state produced by $a$ and attempts to match the _distribution_ of state-action pairs associated with $b$. Additionally, a DQN controls the exact timing at which the transition policy passes the control to $b$, which improves the success rate of the transition. Experiments are conducted on some simulated arm manipulation and bipedal locomotion tasks with similar results to current baselines on the first group and superior results on the second group.