Self-Imitation Learning via Trajectory-Conditioned Policy for Hard-Exploration Tasks

1 INTRODUCTION . Hard-exploration tasks , particularly characterized by sparse environment rewards , are traditionally challenging in reinforcement learning ( RL ) , because the agent must carefully balance the exploration and exploitation when taking a long sequence of actions to receive infrequent non-zero rewards . Demonstration data has been shown to be helpful for tackling hard-exploration problems ( Subramanian et al. , 2016 ) ; many existing methods ( Hester et al. , 2018 ; Pohlen et al. , 2018 ; Aytar et al. , 2018 ; Salimans & Chen , 2018 ) provide the guidance for exploration based on imitation learning of expert demonstrations and achieve strong performances on hard-exploration tasks . However , the reliance on human demonstrations largely limits the general applicability of such approaches . The agent ’ s own past good trajectories with high total rewards are easily accessible ( though imperfect ) alternatives for the human-expert demonstrations . Recent works ( Oh et al. , 2018 ; Gangwani et al. , 2018 ) verify that imitation learning from the agent ’ s previous good trajectories could indirectly drive exploration in certain environments . However , imitation of good experiences within limited directions might hurt exploration in some cases . Specifically , in environments with misleading rewards which may trap the agent in local optima , simply imitating ‘ good ’ trajectories that would accumulate misleading positive rewards may guide the agent to a myopic behavior and hinder it from reaching a higher return in the longer term . Therefore , imitating diverse trajectories would be more desirable to help encourage exploration in diverse directions and avoid being distracted by the misleading rewards . For example , as illustrated in Figure 1 , the agent starts in the bottom left corner where it can easily collect apples near its initial location by random exploration and achieve a small positive reward . If the agent imitates the trajectories around the orange path , it would receive the nearby positive rewards quickly but it is unlikely to collect the gold within a given time limit . Therefore , in order to find the optimal path ( purple ) , it is better to exploit the past experiences in diverse directions ( gray paths ) , instead of focusing only on the trajectories with the myopic behavior . This paper investigates how imitation of diverse past trajectories leads a further exploration and helps avoid getting stuck at a sub-optimal behavior . Our main contributions are summarized as follows : ( 1 ) We propose a novel architecture for a trajectory-conditioned policy that can imitate diverse demonstrations . ( 2 ) We show the importance of imitating diverse past experiences to indirectly drive exploration to different regions of the environment , by comparing to existing approaches on various sparse-reward reinforcement learning tasks with discrete and continuous action space . ( 3 ) We achieve a performance comparable with the state-of-the-art on hard-exploration Atari game of Montezuma ’ s Revenge and Pitfall without using expert demonstrations or resetting to arbitrary states . 2 RELATED WORK . Imitation Learning The goal of imitation learning is to train a policy to mimic a given demonstration . Many previous works achieve good results on hard-exploration Atari games by imitating human demonstrations . DQfD ( Hester et al. , 2018 ) combines the temporal difference updates in Q-learning with the supervised classification of the demonstrators ’ actions . Ape-X DQfD ( Pohlen et al. , 2018 ) extends DQfD with the transformed Bellman operator and temporal consistency loss to improve the efficiency of exploration . Aytar et al . ( 2018 ) learn embeddings from a variety of demonstrations videos and proposes the one-shot imitation learning reward , which inspires the design of reward in our method . All these successful attempts rely on the availability of human demonstrations . In contrast , our method treats the agent ’ s past trajectories as demonstrations . Imitation learning is more difficult when the environment becomes more stochastic because the demonstrations could not account for all possible situations . Our method allows for some flexibility to follow the demonstrations in a soft-order and thus could perform well in the environment with a moderate degree of stochasticity . As discussed in Appendix C.1 , we can easily extend our method to handle a more challenging scenario ( e.g. , where the location of objects could be random ) . Yet , imitation learning in extremely stochastic environments is still an open problem ( Ghosh et al. , 2017 ; Paine et al. , 2019 ) . Self-Imitation Learning a good policy by imitating past experiences has been discussed where the agent is trained to imitate only the high-reward trajectories with the SIL ( Oh et al. , 2018 ) or GAIL objective ( Gangwani et al. , 2018 ) . In contrast , we store the past trajectories ending with diverse states in the buffer , because trajectories with low reward in the short term might lead to high reward in the long term , and thus following a diverse set of trajectories could be beneficial for discovering optimal solutions . Furthermore , our method focuses on explicit trajectory-level imitation while existing methods use sampled state-action pairs from the buffer to update the policy . Gangwani et al . ( 2018 ) proposed to learn multiple diverse policies in a SIL framework using the Stein Variational Policy Gradient . Empirically , their exploration can be limited by the number of policies learned simultaneously and the exploration performance of every single policy , as shown in Appendix F. Exploration Many exploration methods ( Schmidhuber , 1991 ; Auer , 2002 ; Chentanez et al. , 2005 ; Strehl & Littman , 2008 ) in RL tend to award a bonus to encourage an agent to visit novel states . Recently this idea was scaled up to large state spaces ( Tang et al. , 2017 ; Bellemare et al. , 2016 ; Ostrovski et al. , 2017 ; Burda et al. , 2018 ) . We propose that instead of directly taking a quantification of novelty as an intrinsic reward , one can encourage exploration by rewarding the agent when it successfully imitates demonstrations that would lead to novel states and gain the advantages in exploitation , as discussed in Appendix I. Go-Explore ( Ecoffet et al. , 2019 ) also shows the benefit of exploration by returning to promising states . Our method can be viewed in general as an extension of Go-Explore , though we do not need to explicitly divide learning into two phases of exploration and robustification . Go-Explore relies on the assumption that the environment is resettable . Resetting to an arbitrary state is often infeasible in real environments and gives an unfair advantage . When using a perfect goal-conditioned policy instead of a direct ‘ reset ’ function , this variant of Go-Explore may not explore as efficiently as our method , as discussed in Appendix H. Previous works attempted reaching a goal state by learning a set of sub-policies ( Liu et al. , 2019 ) or a goal-conditioned policy in pixel observation space ( Dong et al. , 2019 ) . Gregor et al . ( 2016 ) ; Eysenbach et al . ( 2018 ) ; Pong et al . ( 2019 ) seek a diversity of exploration by maximizing the entropy of mixture skill policies or generated goal states . However , these methods do not show experimental results performing well on sparse-reward environments with a rich observation space like Atari games . Goal-Conditioned Policy Andrychowicz et al . ( 2017 ) ; Nair et al . ( 2017 ) ; Schaul et al . ( 2015a ) ; Pathak et al . ( 2018 ) studied learning a goal-conditioned policy . Our trajectory-conditioned policy could be viewed as a goal-conditioned policy . Similarly to hindsight experience replay ( Andrychowicz et al. , 2017 ) , our approach samples goal states from past experiences . Compared to conditioning on a single final goal state , the state trajectory includes rich intermediate information leading the Algorithm 1 Diverse Self-Imitation Learning with Trajectory-Conditioned Policy Initialize parameter θ for the trajectory-conditioned policy πθ ( at|e≤t , ot , g ) Initialize the trajectory buffer D ← ∅ # Store diverse past trajectories Initialize set of transitions in the current episode E ← ∅ # Store current episode trajectory Initialize set of on-policy samples F ← ∅ # Store data for on-policy PPO update Initialize demonstration trajectory g ← ∅ for each iteration i from 1 to I do for each step t do Observe st = { ot , et } and choose an action at ∼ πθ ( at|e≤t , ot , g ) Execute action at in the environment to get rt , ot+1 , et+1 Store transition E ← E ∪ { ( ot , et , at , rt ) } # Positive reward if agent follows demonstration g # No reward after agent completes g and then takes random exploration Determine rDTSILt by comparing e≤t+1 with g ( Eq . 1 ) Store on-policy sample F ← F ∪ { ( ot , et , at , g , r DTSIL t ) } end for if st+1 is terminal then D ← UpdateBuffer ( D , E ) ( Alg . 2 ) Clear current episode trajectory E ← ∅ g ← SampleDemo ( D , i , I ) ( Alg . 3 ) end if θ ← θ − η∇θLRL # Perform PPO update using on-policy samples ( Eq . 2 ) Clear on-policy samples F ← ∅ θ ← θ − η∇θLSL # Perform supervised learning updates using samples from D for J times ( Eq . 3 ) end for agent to follow a demonstration and reach the goal state even far away from the current state . Our method shares the same motivation as Duan et al . ( 2017 ) which uses an attention model over the demonstration but mainly focuses on the block stacking task . However , our architecture is simpler since it does not use an attention model over the current observation and our method is evaluated on various environments . 3 METHOD . The main idea of our method is to maintain a buffer of diverse trajectories collected during training and to train a trajectory-conditioned policy by leveraging reinforcement learning and supervised learning to roughly follow demonstrations sampled from the trajectory buffer . The demonstration trajectories cover diverse possible directions in the environment . Therefore , the agent is encouraged to explore beyond various visited states in the environment and gradually push its exploration frontier further . In the meantime , we can train the policy to imitate the best trajectories collected to exploit the past good experiences . We put more weights on exploration in the early stage of training , and then increases the probability of imitating the best trajectories ( i.e. , exploitation ) as training goes on . We name our method as Diverse Trajectory-conditioned Self-Imitation Learning ( DTSIL ) . 3.1 BACKGROUND AND NOTATION . In the standard reinforcement learning setting , at each time step t , an agent observes a state st , selects an action at ∈ A , and receives a reward rt when transitioning to a next state st+1 ∈ S , where S is a set of all states and A is a set of all actions . The goal is to find a policy πθ ( a|s ) parameterized by θ that maximizes the expected return Eπθ [ ∑T t=0 γ trt ] , where γ ∈ ( 0 , 1 ] is a discount factor . In our work , we assume a state st includes the agent ’ s observation ot ( e.g. , raw pixel image ) and a high-level abstract state embedding et ( e.g. , the agent ’ s location in the abstract space ) . The embedding et may be learnable from o≤t ( e.g. , ADM ( Choi et al. , 2018 ) could localize the agent in Atari games ) , but in this work , we consider a setting where high-level embedding is provided as a part of st 1 . A trajectory-conditioned policy πθ ( at|e≤t , ot , g ) ( which we refer to as πθ ( ·|g ) in shorthand notation ) takes a sequence of state embeddings g = { eg1 , e g 2 , · · · , e g |g| } as input for a demonstration , where |g| is 1In many important application domains ( e.g . the robotics domain ) , such handcrafted representation is available . Also , learning a good state representation itself is an important open question and extremely challenging especially for hard-exploration and sparse-reward environments , which is not the main focus of this work . Therefore , we assume the availability of the high-level representations as many previous works ( Florensa et al. , 2017 ; Liu et al. , 2019 ; Ecoffet et al. , 2019 ; Plappert et al. , 2018 ) the length of the trajectory g. A sequence of the agent ’ s past state embeddings e≤t = { e1 , e2 , · · · , et } is provided to determine which part of the demonstration has been followed by the agent . Together with the current observation ot , it helps to determine the correct action at to accurately imitate the demonstration . Our goal here is to find a set of optimal state embedding sequence ( s ) g∗ and the policy π∗θ ( ·|g ) to maximize the return : g∗ , θ∗ , arg maxg , θ Eπθ ( ·|g ) [ ∑T t=0 γ trt ] . For robustness we may want to find multiple near-optimal embedding sequences with similar returns and a trajectoryconditioned policy for executing them . In our implementation , we train the trajectory-conditioned policy to imitate the best trajectories . Alternatively , an unconditional stochastic policy could also be trained to imitate the best trajectories , which may further improve generalization and robustness ( see Appendix C.1 for more discussion and experiments ) .

Summary: The paper addresses the challenge of hard exploration tasks. The approach taken is to apply self-imitation to a diverse selection of trajectories from past experience -- practice re-doing the strangest things you've ever done. This is claimed to drive more efficient exploration in sparse-reward problems, leading to SOTA results for Montezuma's Revenge without certain common aides.