Original: S ELF -I MITATION L EARNING VIA T RAJECTORYC ONDITIONED P OLICY FOR H ARD -E XPLORATION T ASKS
Revised: S ELF -I MITATION L EARNING VIA T RAJECTORYC ONDITIONED P OLICY FOR H ARD -E XPLORATION T ASKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Imitation learning from human-expert demonstrations has been shown to be greatly helpful for challenging reinforcement learning problems with sparse environment rewards.
Revised: Imitation learning from human-expert demonstrations has been shown to be greatly helpful for challenging reinforcement learning problems with sparse environment rewards.

Original: However, it is very difﬁcult to achieve similar success without relying on expert demonstrations.
Revised: However, it is very difﬁcult to achieve similar success without relying on expert demonstrations.

Original: Recent works on self-imitation learning showed that imitating the agent’s own past good experience could indirectly drive exploration in some environments, but these methods often lead to sub-optimal and myopic behavior.
Revised: Recent works on self-imitation learning showed that imitating the agent’s own past good experience could indirectly drive exploration in some environments, but these methods often lead to sub-optimal and myopic behavior.

Original: To address this issue, we argue that exploration in diverse directions by imitating diverse trajectories, instead of focusing on limited good trajectories, is more desirable for the hard-exploration tasks.
Revised: To address this issue, we argue that exploration in diverse directions by imitating diverse trajectories, instead of focusing on limited good trajectories, is more desirable for the hard-exploration tasks.

Original: We propose a new method of learning a trajectory-conditioned policy to imitate diverse trajectories from the agent’s own past experiences and show that such self-imitation helps avoid myopic behavior and increases the chance of ﬁnding a globally optimal solution for hard-exploration tasks, especially when there are misleading rewards.
Revised: We propose a new method of learning a trajectory-conditioned policy to imitate diverse trajectories from the agent’s own past experiences and show that such self-imitation helps avoid myopic behavior and increases the chance of ﬁnding a globally optimal solution for hard-exploration tasks, especially when there are misleading rewards.

Original: Our method signiﬁcantly outperforms existing self-imitation learning and count-based exploration methods on various hard-exploration tasks with local optima.
Revised: Our method signiﬁcantly outperforms existing self-imitation learning and count-based exploration methods on various hard-exploration tasks with local optima.

Original: In particular, we report a state-of-the-art score of more than 20,000 points on Montezumas Revenge without using expert demonstrations or resetting to arbitrary states.
Revised: In particular, we report a state-of-the-art score of more than 20,000 points on Montezumas Revenge without using expert demonstrations or resetting to arbitrary states.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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.
Revised: 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.

Original: Demonstration data has been shown to be helpful for tackling hard-exploration problems (Subramanian et al., 2016);
Revised: Demonstration data has been shown to be helpful for tackling hard-exploration problems (Subramanian et al., 2016);

Original: many existing methods (Hester et al., 2018; Pohlen et al., 2018; Aytar et al., 2018;
Revised: many existing methods (Hester et al., 2018; Pohlen et al., 2018; Aytar et al., 2018;

Original: Salimans & Chen, 2018) provide the guidance for exploration based on imitation learning of expert demonstrations and achieve strong performances on hard-exploration tasks.
Revised: Salimans & Chen, 2018) provide the guidance for exploration based on imitation learning of expert demonstrations and achieve strong performances on hard-exploration tasks.

Original: However, the reliance on human demonstrations largely limits the general applicability of such approaches
Revised: However, the reliance on human demonstrations largely limits the general applicability of such approaches.

Original: The agent’s own past good trajectories with high total rewards are easily accessible (though imperfect) alternatives for the human-expert demonstrations.
Revised: The agent’s own past good trajectories with high total rewards are easily accessible (though imperfect) alternatives for the human-expert demonstrations.

Original: Recent works (Oh et al., 2018;
Revised: Recent works (Oh et al., 2018;

Original: Gangwani et al., 2018) verify that imitation learning from the agent’s previous good trajectories could indirectly drive exploration in certain environments.
Revised: Gangwani et al., 2018) verify that imitation learning from the agent’s previous good trajectories could indirectly drive exploration in certain environments.

Original: However, imitation of good experiences within limited directions might hurt exploration in some cases.
Revised: However, imitation of good experiences within limited directions might hurt exploration in some cases.

Original: Speciﬁcally, 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.
Revised: Speciﬁcally, 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.

Original: Therefore, imitating diverse trajectories would be more desirable to  encourage exploration in diverse directions and avoid being distracted by the misleading rewards.
Revised: Therefore, imitating diverse trajectories would be more desirable to help encourage exploration in diverse directions and avoid being distracted by the misleading rewards.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Therefore, in order to ﬁnd 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.
Revised: Therefore, in order to ﬁnd 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.

Original: This paper investigates how imitation of diverse past trajectories leads a further exploration and helps avoid getting stuck at a sub-optimal behavior.
Revised: This paper investigates how imitation of diverse past trajectories leads a further exploration and helps avoid getting stuck at a sub-optimal behavior.

Original: Our main contributions are summarized as follows: (1) We propose a novel architecture for a trajectory-conditioned policy that can imitate diverse demonstrations.
Revised: Our main contributions are summarized as follows: (1) We propose a novel architecture for a trajectory-conditioned policy that can imitate diverse demonstrations.

Original: (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.
Revised: (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.

Original: (3) We achieve a performance comparable with the state-of-the-art on hard-exploration Atari game of Montezuma’s
Revised: (3) We achieve a performance comparable with the state-of-the-art on hard-exploration Atari game of Montezuma’s

Original: Revenge and Pitfall without using expert demonstrations or resetting to arbitrary states.
Revised: Revenge and Pitfall without using expert demonstrations or resetting to arbitrary states.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Imitation Learning The goal of imitation learning is to train a policy to mimic a given demonstration.
Revised: Imitation Learning The goal of imitation learning is to train a policy to mimic a given demonstration.

Original: Many previous works achieve good results on hard-exploration Atari games by imitating human demonstrations.
Revised: Many previous works achieve good results on hard-exploration Atari games by imitating human demonstrations.

Original: DQfD (Hester et al., 2018) combines the temporal difference updates in Q-learning with the supervised classiﬁcation of the demonstrators’ actions.
Revised: DQfD (Hester et al., 2018) combines the temporal difference updates in Q-learning with the supervised classiﬁcation of the demonstrators’ actions.

Original: Ape-X DQfD (Pohlen et al., 2018) extends DQfD with the transformed Bellman operator and temporal consistency loss to improve the efﬁciency of exploration.
Revised: Ape-X DQfD (Pohlen et al., 2018) extends DQfD with the transformed Bellman operator and temporal consistency loss to improve the efﬁciency of exploration.

Original: Aytar et al.
Revised: Aytar et al.

Original: (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.
Revised: (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.

Original: All these successful attempts rely on the availability of human demonstrations.
Revised: All these successful attempts rely on the availability of human demonstrations.

Original: In contrast, our method treats the agent’s past trajectories as demonstrations.
Revised: In contrast, our method treats the agent’s past trajectories as demonstrations.

Original: Imitation learning is more difﬁcult when the environment becomes more stochastic because the demonstrations could not account for all possible situations.
Revised: Imitation learning is more difﬁcult when the environment becomes more stochastic because the demonstrations could not account for all possible situations.

Original: Our method allows for some ﬂexibility to follow the demonstrations in a soft-order and thus could perform well in the environment with a moderate degree of stochasticity.
Revised: Our method allows for some ﬂexibility to follow the demonstrations in a soft-order and thus could perform well in the environment with a moderate degree of stochasticity.

Original: 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).
Revised: 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).

Original: Yet, imitation learning in extremely stochastic environments is still an open problem (Ghosh et al., 2017; Paine et al., 2019).
Revised: Yet, imitation learning in extremely stochastic environments is still an open problem (Ghosh et al., 2017; Paine et al., 2019).

Original: 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).
Revised: 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).

Original: 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 beneﬁcial for discovering optimal solutions.
Revised: 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 beneﬁcial for discovering optimal solutions.

Original: Furthermore, our method focuses on explicit trajectory-level imitation while existing methods use sampled state-action pairs from the buffer to update the policy.
Revised: Furthermore, our method focuses on explicit trajectory-level imitation while existing methods use sampled state-action pairs from the buffer to update the policy.

Original: Gangwani et al.
Revised: Gangwani et al.

Original: (2018) proposed to learn multiple diverse policies in a SIL framework using the Stein Variational
Revised: (2018) proposed to learn multiple diverse policies in a SIL framework using the Stein Variational

Original: Policy Gradient.
Revised: Policy Gradient.

Original: 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.
Revised: 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.

Original: Exploration Many exploration methods (Schmidhuber, 1991; Auer, 2002; Chentanez et al., 2005;
Revised: Exploration Many exploration methods (Schmidhuber, 1991; Auer, 2002; Chentanez et al., 2005;

Original: Strehl & Littman, 2008) in RL tend to award a bonus to encourage an agent to visit novel states.
Revised: Strehl & Littman, 2008) in RL tend to award a bonus to encourage an agent to visit novel states.

Original: Recently this idea was scaled up to large state spaces (Tang et al., 2017; Bellemare et al., 2016;
Revised: Recently this idea was scaled up to large state spaces (Tang et al., 2017; Bellemare et al., 2016;

Original: Ostrovski et al., 2017; Burda et al., 2018).
Revised: Ostrovski et al., 2017; Burda et al., 2018).

Original: We propose that instead of directly taking a quantiﬁcation 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 beneﬁt of exploration by returning to promising states.
Revised: We propose that instead of directly taking a quantiﬁcation 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 beneﬁt of exploration by returning to promising states.

Original: Our method can be viewed in general as an extension of
Revised: Our method can be viewed in general as an extension of

Original: Go-Explore, though we do not need to explicitly divide learning into two phases of exploration and robustiﬁcation.
Revised: Go-Explore, though we do not need to explicitly divide learning into two phases of exploration and robustiﬁcation.

Original: Go-Explore relies on the assumption that the environment is resettable.
Revised: Go-Explore relies on the assumption that the environment is resettable.

Original: Resetting to an arbitrary statecan result in over two orders of magnitude reduction in sample complexity and thus give an unfair advantage. More importantly, such resetting is often infeasible in real environments.
Revised: Resetting to an arbitrary state is often infeasible in real environments and gives an unfair advantage.

Original: When using a perfect goal-conditioned policy instead of a direct ‘reset’ function, this variant of Go-Explore may not explore as efﬁciently 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).
Revised: When using a perfect goal-conditioned policy instead of a direct ‘reset’ function, this variant of Go-Explore may not explore as efﬁciently 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).

Original: Gregor et al.
Revised: Gregor et al.

Original: Eysenbach et al.
Revised: Eysenbach et al.

Original:  Pong et al.
Revised: ; Pong et al.

Original: (2019) seek a diversity of exploration by maximizing the entropy of mixture skill policies or generated goal states.
Revised: (2019) seek a diversity of exploration by maximizing the entropy of mixture skill policies or generated goal states.

Original: However, these methods do not show experimental results performing well on sparse-reward environments with a rich observation space like Atari games.
Revised: However, these methods do not show experimental results performing well on sparse-reward environments with a rich observation space like Atari games.

Original: Goal-Conditioned Policy Andrychowicz et al.
Revised: Goal-Conditioned Policy Andrychowicz et al.

Original: ; Nair et al.
Revised: ; Nair et al.

Original: ; Schaul et al.
Revised: ; Schaul et al.

Original: (2015a);
Revised: (2015a);

Original: Pathak et al.
Revised: Pathak et al.

Original: (2018) studied learning a goal-conditioned policy.
Revised: (2018) studied learning a goal-conditioned policy.

Revised: Our trajectory-conditioned policy could be viewed as a goal-conditioned policy.

Original: Similarly to hindsight experience replay (Andrychowicz et al., 2017), our approach samples goal states from past experiences.
Revised: Similarly to hindsight experience replay (Andrychowicz et al., 2017), our approach samples goal states from past experiences.

Original: Compared to conditioning on a single  goal state, the state trajectory can lead theagent to follow a demonstration to 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.
Revised: Compared to conditioning on a single ﬁnal goal state, the state trajectory includes rich intermediate information leading the

Original: Initialize parameter θ for the trajectory-conditioned policy π θ ( a t | e ≤ t , o t , g )
Revised: Initialize parameter θ for the trajectory-conditioned policy π θ ( a t | e ≤ t , o t , g )

Original: Initialize the trajectory buffer D ← ∅ # Store diverse past trajectories
Revised: Initialize the trajectory buffer D ← ∅ # Store diverse past trajectories

Original: Initialize set of transitions in the current episode E ← ∅ # Store current episode trajectory
Revised: Initialize set of transitions in the current episode E ← ∅ # Store current episode trajectory

Original: Initialize set of on-policy samples F ← ∅ # Store data for on-policy PPO update
Revised: Initialize set of on-policy samples F ← ∅ # Store data for on-policy PPO update

Original: Initialize demonstration trajectory g ← ∅ for each iteration i from 1 to I do for each step t do
Revised: Initialize demonstration trajectory g ← ∅ for each iteration i from 1 to I do for each step t do

Original: Observe s t = { o t , e t } and choose an action a t ∼ π θ ( a t | e ≤ t , o t , g ) Execute action a t in the environment to get r t , o t +1 , e t +1 Store transition E ← E ∪ { ( o t , e t , a t , r t ) } # Positive reward if agent follows demonstration
Revised: Observe s t = { o t , e t } and choose an action a t ∼ π θ ( a t | e ≤ t , o t , g ) Execute action a t in the environment to get r t , o t +1 , e t +1 Store transition E ← E ∪ { ( o t , e t , a t , r t ) } # Positive reward if agent follows demonstration

Original: g # No reward after agent completes g and then takes random exploration
Revised: g # No reward after agent completes g and then takes random exploration

Original: Determine r DTSIL t by comparing e ≤ t +1 with g (Eq. 1) Store on-policy sample F ← F ∪ (cid:8) ( o t , e t , a t , g, r DTSIL t ) (cid:9) end for if s t +1 is terminal then
Revised: Determine r DTSIL t by comparing e ≤ t +1 with g (Eq. 1) Store on-policy sample F ← F ∪ (cid:8) ( o t , e t , a t , g, r DTSIL t ) (cid:9) end for if s t +1 is terminal then

Original: D ← UpdateBuffer
Revised: D ← UpdateBuffer

Original: ( D , E ) (Alg. 2)
Revised: ( D , E ) (Alg. 2)

Original: Clear current episode trajectory E ← ∅ g ← SampleDemo( D , i , I ) (Alg. 3) end if θ ← θ − η ∇ θ L RL # Perform PPO update using on-policy samples (Eq. 2)
Revised: Clear current episode trajectory E ← ∅ g ← SampleDemo( D , i , I ) (Alg. 3) end if θ ← θ − η ∇ θ L RL # Perform PPO update using on-policy samples (Eq. 2)

Original: Clear on-policy samples F ← ∅ θ ← θ − η ∇ θ L SL # Perform supervised learning updates using samples from D for J times (Eq. 3) end for 
Revised: Clear on-policy samples F ← ∅ θ ← θ − η ∇ θ L SL # 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.

Original:  M ETHOD
Revised:  M ETHOD

Original: 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.
Revised: 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.

Original: The demonstration trajectories cover diverse possible directions in the environment.
Revised: The demonstration trajectories cover diverse possible directions in the environment.

Original: Therefore, the agent is encouraged to explore beyond various visited states in the environment and gradually push its exploration frontier further.
Revised: Therefore, the agent is encouraged to explore beyond various visited states in the environment and gradually push its exploration frontier further.

Original: In the meantime, we can train the policy to imitate the best trajectories collected to exploit the past good experiences.
Revised: In the meantime, we can train the policy to imitate the best trajectories collected to exploit the past good experiences.

Original: 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.
Revised: 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.

Original: We name our method as Diverse Trajectory-conditioned Self-Imitation Learning (DTSIL).
Revised: We name our method as Diverse Trajectory-conditioned Self-Imitation Learning (DTSIL).

Original: 3. B ACKGROUND AND N OTATION
Revised: 3. B ACKGROUND AND N OTATION

Original: In the standard reinforcement learning setting, at each time step t , an agent observes a state s t , selects an action a t ∈ A , and receives a reward r t when transitioning to a next state s t +1 ∈ S , where S is a set of all states and A is a set of all actions.
Revised: In the standard reinforcement learning setting, at each time step t , an agent observes a state s t , selects an action a t ∈ A , and receives a reward r t when transitioning to a next state s t +1 ∈ S , where S is a set of all states and A is a set of all actions.

Original: The goal is to ﬁnd a policy π θ ( a | s ) parameterized by θ that maximizes the expected return E π θ [ (cid:80) T t =

Revised: The goal is to ﬁnd a policy π θ ( a | s ) parameterized by θ that maximizes the expected return E π θ [ (cid:80) Tt =

Original: γ t r t ] , where γ ∈ (0 , 1] is a discount factor.
Revised: γ t r t ] , where γ ∈ (0 , 1] is a discount factor.

Original: In our work, we assume a state s t includes the agent’s observation o t (e.g., raw pixel image) and a high-level abstract state embedding e t (e.g., the agent’s location in the abstract space).
Revised: In our work, we assume a state s t includes the agent’s observation o t (e.g., raw pixel image) and a high-level abstract state embedding e t (e.g., the agent’s location in the abstract space).

Original: The embedding e t 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 s t 1 .
Revised: The embedding e t 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 s t 1 .

Original: A trajectory-conditioned policy π θ ( a t | e ≤ t , o t , g ) (which we refer to as π θ ( ·| g ) in shorthand notation) takes a sequence of state embeddings g = { e g 1 , e g 2 , · · · , e g | g | } as input for a demonstration, where | g | is the length of the trajectory g .
Revised: A trajectory-conditioned policy π θ ( a t | e ≤ t , o t , g ) (which we refer to as π θ ( ·| g ) in shorthand notation) takes a sequence of state embeddings g = { e g 1 , e g 2 , · · · , e g | g | } as input for a demonstration, where | g | is the length of the trajectory g .

Original: A sequence of the agent’s past state embeddings e ≤ t = { e 1 , e 2 , · · · , e t } is provided to determine which part of the demonstration has been followed by the agent.
Revised: A sequence of the agent’s past state embeddings e ≤ t = { e 1 , e 2 , · · · , e t } is provided to determine which part of the demonstration has been followed by the agent.

Original: Together with the current observation o t , it helps to determine the correct action a t to accurately imitate the demonstration.
Revised: Together with the current observation o t , it helps to determine the correct action a t to accurately imitate the demonstration.

Original: Our goal here is to ﬁnd a set of optimal state embedding sequence(s) g ∗ and the policy π ∗ θ ( ·| g ) to maximize the return: g ∗ , θ ∗ (cid:44) arg max g,θ E π θ ( ·| g )
Revised: Our goal here is to ﬁnd a set of optimal state embedding sequence(s) g ∗ and the policy π ∗ θ ( ·| g ) to maximize the return: g ∗ , θ ∗ (cid:44) arg max g,θ E π θ ( ·| g )

Original: [ (cid:80) Tt =

Revised: [ (cid:80) Tt =

Original: γ t r t ] .
Revised: γ t r t ] .

Original: For robustness we may want to ﬁnd multiple near-optimal embedding sequences with similar returns and a trajectoryconditioned policy for executing them.
Revised: For robustness we may want to ﬁnd multiple near-optimal embedding sequences with similar returns and a trajectoryconditioned policy for executing them.

Original: In our implementation, we train the trajectory-conditioned policy to imitate the best trajectories.
Revised: In our implementation, we train the trajectory-conditioned policy to imitate the best trajectories.

Original: Alternatively, an unconditional stochastic policy could also be trained to imitate the best trajectories, which may further improve generalization and robustness (see
Revised: Alternatively, an unconditional stochastic policy could also be trained to imitate the best trajectories, which may further improve generalization and robustness (see

Original: Appendix C.1 for more discussion and experiments).
Revised: Appendix C.1 for more discussion and experiments).

Original: 3. O RGANIZING T RAJECTORY B UFFER
Revised: 3. O RGANIZING T RAJECTORY B UFFER

Original: We maintain a trajectory buffer D
Revised: We maintain a trajectory buffer D

Original: = { ( e (1) , τ (1) , n (1) ) , ( e (2) , τ (2) , n (2) ) , · · · } of diverse past trajectories.
Revised: = { ( e (1) , τ (1) , n (1) ) , ( e (2) , τ (2) , n (2) ) , · · · } of diverse past trajectories.

Original: For each embedding-trajectory-count tuple ( e ( i ) , τ ( i ) , n ( i ) ) , τ ( i ) is the best trajectory ending with a state with the high-level representation e ( i ) , and n ( i ) is the number of times the cluster represented by this state embedding e ( i ) has been visited during training.
Revised: For each embedding-trajectory-count tuple ( e ( i ) , τ ( i ) , n ( i ) ) , τ ( i ) is the best trajectory ending with a state with the high-level representation e ( i ) , and n ( i ) is the number of times the cluster represented by this state embedding e ( i ) has been visited during training.

Original: To maintain a compact buffer, similar state embeddings within the tolerance threshold th can be clustered together, and the existing entry is replaced if an improved trajectory τ ( i ) ending with a near-identical state is found.
Revised: To maintain a compact buffer, similar state embeddings within the tolerance threshold th can be clustered together, and the existing entry is replaced if an improved trajectory τ ( i ) ending with a near-identical state is found.

Original: When given a new episode E = { ( o 0 , e 0 , a 0 , r 0 ) , · · · , ( o T , e T , a T , r T ) } , all of the state embeddings e t (1 ≤ t ≤ T ) in this episode E are considered as follows (similarly to Ecoffet et al.
Revised: When given a new episode E = { ( o 0 , e 0 , a 0 , r 0 ) , · · · , ( o T , e T , a T , r T ) } , all of the state embeddings e t (1 ≤ t ≤ T ) in this episode E are considered as follows (similarly to Ecoffet et al.

Original: (2019)), because the buffer should maintain all of the possible paths available for future exploration to avoid missing any possibility to ﬁnd an optimal solution.
Revised: (2019)), because the buffer should maintain all of the possible paths available for future exploration to avoid missing any possibility to ﬁnd an optimal solution.

Original: If the Euclidean distance between e t and any state embedding e ( i ) in the buffer is larger than th (i.e e t does not belong to any existing cluster in the buffer), ( e t , τ ≤ t , 1) is directly pushed into the buffer, where τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , · · · , ( o t , e t , a t , r t ) } is the agent’s partial episode ending with e t .
Revised: If the Euclidean distance between e t and any state embedding e ( i ) in the buffer is larger than th (i.e e t does not belong to any existing cluster in the buffer), ( e t , τ ≤ t , 1) is directly pushed into the buffer, where τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , · · · , ( o t , e t , a t , r t ) } is the agent’s partial episode ending with e t .

Original: If there exists e ( k ) similar to e t (i.e., e ( k ) and e t belong to the same cluster within threshold th ) and the partial episode τ ≤ t is better (i.e., higher return or shorter trajectory) than the stored trajectory τ ( k ) , τ ( k ) is replaced by the current trajectory τ ≤ t , and e ( k ) is replaced by e t to represent this cluster of state embeddings.
Revised: If there exists e ( k ) similar to e t (i.e., e ( k ) and e t belong to the same cluster within threshold th ) and the partial episode τ ≤ t is better (i.e., higher return or shorter trajectory) than the stored trajectory τ ( k ) , τ ( k ) is replaced by the current trajectory τ ≤ t , and e ( k ) is replaced by e t to represent this cluster of state embeddings.

Original: The full algorithm in pseudo-code is described in Appendix A.1.
Revised: The full algorithm in pseudo-code is described in Appendix A.1.

Original: 3. S AMPLING D EMONSTRATIONS
Revised: 3. S AMPLING D EMONSTRATIONS

Original: When learning a trajectory-conditioned policy π , demonstration trajectories are sampled from the buffer D .
Revised: When learning a trajectory-conditioned policy π , demonstration trajectories are sampled from the buffer D .

Original: We record the count n ( i ) of how many times the cluster represented by this state embedding e ( i ) is visited.
Revised: We record the count n ( i ) of how many times the cluster represented by this state embedding e ( i ) is visited.

Original: In the exploration mode, we set the probability of sampling each trajectory as 1 / √ n ( i ) .
Revised: In the exploration mode, we set the probability of sampling each trajectory as 1 / √ n ( i ) .

Original: This is inspired by the count-based exploration bonus (Strehl & Littman, 2008; Bellemare et al., 2016) and the idea of rank-based prioritization (Schaul et al., 2015b; Ecoffet et al., 2019):
Revised: This is inspired by the count-based exploration bonus (Strehl & Littman, 2008; Bellemare et al., 2016) and the idea of rank-based prioritization (Schaul et al., 2015b; Ecoffet et al., 2019):

Original: we prioritize a trajectory that ends with a less frequently visited state because this leads the agent to reach rarely visited regions in the state space and is more promising for discovering novel states.
Revised: we prioritize a trajectory that ends with a less frequently visited state because this leads the agent to reach rarely visited regions in the state space and is more promising for discovering novel states.

Original: On the other hand, in the imitation mode, we sample the best trajectories stored in the buffer for imitation learning.
Revised: On the other hand, in the imitation mode, we sample the best trajectories stored in the buffer for imitation learning.

Original: These trajectories are used to train the policy to converge to a high-reward behavior (Aytar et al., 2018; Ecoffet et al., 2019).
Revised: These trajectories are used to train the policy to converge to a high-reward behavior (Aytar et al., 2018; Ecoffet et al., 2019).

Original: To balance between exploration and exploitation, we decrease the probability of taking the exploration mode and exploit the best experiences more as training goes on.
Revised: To balance between exploration and exploitation, we decrease the probability of taking the exploration mode and exploit the best experiences more as training goes on.

Original: The algorithm is described in Appendix A.2.
Revised: The algorithm is described in Appendix A.2.

Original: 3. L EARNING T RAJECTORY -C ONDITIONED P OLICY
Revised: 3. L EARNING T RAJECTORY -C ONDITIONED P OLICY

Original: Imitation Reward Given a demonstration trajectory g = { e g 0 , e g 1 , · · · , e g | g | } , we provide reward signals for imitating g .
Revised: Imitation Reward Given a demonstration trajectory g = { e g 0 , e g 1 , · · · , e g | g | } , we provide reward signals for imitating g .

Original: At the beginning of an episode, the index u of the last visited state embedding in the demonstration is initialized as u = − 1 .
Revised: At the beginning of an episode, the index u of the last visited state embedding in the demonstration is initialized as u = − 1 .

Original: At each step t , if the agent’s new state s t +1 has an embedding e t +1 and it is the similar enough to any of the next ∆ t state embeddings starting from the last visited state embedding e gu in the demonstration (i.e., (cid:107) e t +1 − e gu (cid:48) (cid:107) <
Revised: At each step t , if the agent’s new state s t +1 has an embedding e t +1 and it is the similar enough to any of the next ∆ t state embeddings starting from the last visited state embedding e gu in the demonstration (i.e., (cid:107) e t +1 − e gu (cid:48) (cid:107) <

Original: th where u < u (cid:48) ≤ u +∆ t ), then it receives a positive imitation reward r im , and the index of the last visited state embedding in the demonstration is updated as u ← u (cid:48) .
Revised: th where u < u (cid:48) ≤ u +∆ t ), then it receives a positive imitation reward r im , and the index of the last visited state embedding in the demonstration is updated as u ← u (cid:48) .

Original: This encourages the agent to visit the state embeddings in the demonstration in a soft-order so that the agent could explore around the demonstration and the demonstration plays a role to guide the agent to the region of interest in the state embedding space.
Revised: This encourages the agent to visit the state embeddings in the demonstration in a soft-order so that the agent could explore around the demonstration and the demonstration plays a role to guide the agent to the region of interest in the state embedding space.

Original: To summarize, the agent receives a reward r DTSIL t deﬁned as
Revised: To summarize, the agent receives a reward r DTSIL t deﬁned as

Original: [Equation]
Revised: [Equation]

Revised: im if ∃ u (cid:48) , u < u (cid:48)

Revised: ≤ u + ∆ t , such that (cid:107) e gu (cid:48) − e t +1 (cid:107) <

Revised: th otherwise, (1)

Revised: [Figure]

Original:  where f ( · ) is a monotonically increasing function (e.g., clipping (Mnih et al., 2015)).
Revised: [Figure] [Figure] [Table] +8n50erJxc4+9ZK/YEUvZB3bKPrMzNmGCXbOf7Bf7HX2MimgV6ZvUaLCpec5uRXT1F85F3XY=</latexit> where f ( · ) is a monotonically increasing function (e.g., clipping (Mnih et al., 2015)).

Original: Figureillustrates the updates of u during an episode when the agent visits a state whose embedding is close to state embeddings in the demonstration g .
Revised: Figureillustrates the updates of u during an episode when the agent visits a state whose embedding is close to state embeddings in the demonstration g .

Original: Policy Architecture For imitation learning with diverse demonstrations, we design a trajectoryconditioned policy π θ ( a t | e ≤ t , o t , g ) that should imitate any given trajectory g .Inspired by neural [Figure]
Revised: Policy Architecture For imitation learning with diverse demonstrations, we design a trajectoryconditioned policy π θ ( a t | e ≤ t , o t , g ) that should imitate any given trajectory g .

Original: [Figure] [Figure] [Table] +8n50erJxc4+9ZK/YEUvZB3bKPrMzNmGCXbOf7Bf7HX2MimgV6ZvUaLCpec5uRXT1F85F3XY=</latexit> machine translation methods (Sutskever et al., 2014; Bahdanau et al., 2014), one can view the demonstration as the source sequence and view the incomplete trajectory of the agent’s state representations as the target sequence.
Revised: Inspired by neural machine translation methods (Sutskever et al., 2014; Bahdanau et al., 2014), one can view the demonstration as the source sequence and view the incomplete trajectory of the agent’s state representations as the target sequence.

Original: We apply a recurrent neural network (RNN) and an attention mechanism to the sequence data to predict actions that would make the agent follow the demonstration.
Revised: We apply a recurrent neural network (RNN) and an attention mechanism to the sequence data to predict actions that would make the agent follow the demonstration.

Original: As illustrated in Figure 3, RNN computes the hidden features
Revised: As illustrated in Figure 3, RNN computes the hidden features

Original: h gi for each state embedding e gi ( 0 ≤ i ≤ | g | ) in the demonstration and derives the hidden features h t for the agent’s state representation e t .
Revised: h gi for each state embedding e gi ( 0 ≤ i ≤ | g | ) in the demonstration and derives the hidden features h t for the agent’s state representation e t .

Original: Then the attention weight α t is computed by comparing the current agent’s hidden features h t with the demonstration’s hidden features h gi ( 0 ≤ i ≤ | g | ).
Revised: Then the attention weight α t is computed by comparing the current agent’s hidden features h t with the demonstration’s hidden features h gi ( 0 ≤ i ≤ | g | ).

Original: The attention readout c t is computed as an attention-weighted summation of the demonstration’s hidden features to capture the relevant information in the demonstration trajectory and to predict the action a t .
Revised: The attention readout c t is computed as an attention-weighted summation of the demonstration’s hidden features to capture the relevant information in the demonstration trajectory and to predict the action a t .

Original: The more details of policy architecture are described in Appendix B.
Revised: The more details of policy architecture are described in Appendix B.

Original: Reinforcement Learning Objective With the reward deﬁned as r DTSIL t (Equation 1), the trajectoryconditioned policy π θ can be trained with a policy gradient algorithm (Sutton et al., 2000):
Revised: Reinforcement Learning Objective With the reward deﬁned as r DTSIL t (Equation 1), the trajectoryconditioned policy π θ can be trained with a policy gradient algorithm (Sutton et al., 2000):

Original: [Equation]
Revised: [Equation]

Original:  where (cid:98) 
Revised: A t ] , where (cid:98) A

Original: [Equation]
Revised: [Equation]

Original:  where the expectation E π θ indicates the empirical average over a ﬁnite batch of on-policy samples and n denotes the number of rollout steps taken in each iteration.
Revised: = [Equation] where the expectation E π θ indicates the empirical average over a ﬁnite batch of on-policy samples and n denotes the number of rollout steps taken in each iteration.

Original: We use Proximal Policy Optimization (PPO)
Revised: We use Proximal Policy Optimization (PPO)

Original: (Schulman et al., 2017) as an actor-critic policy gradient algorithm for our experiments.
Revised: (Schulman et al., 2017) as an actor-critic policy gradient algorithm for our experiments.

Original: Supervised Learning Objective To improve trajectory-conditioned imitation learning and to better leverage the past trajectories, we propose a supervised learning objective.
Revised: Supervised Learning Objective To improve trajectory-conditioned imitation learning and to better leverage the past trajectories, we propose a supervised learning objective.

Original: We sample a trajectory τ = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) · · · } ∈ D , formulate the demonstration g = { e 0 , e 1 , · · · , e | g | } and assume the agent’s incomplete trajectory is the partial trajectory g ≤ t = e ≤ t = { e 0 , e 1 , · · · , e t } for any 1 ≤ t ≤ | g | .
Revised: We sample a trajectory τ = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) · · · } ∈ D , formulate the demonstration g = { e 0 , e 1 , · · · , e | g | } and assume the agent’s incomplete trajectory is the partial trajectory g ≤ t = e ≤ t = { e 0 , e 1 , · · · , e t } for any 1 ≤ t ≤ | g | .

Original: Then a t is the ‘correct’ action at step t for the agent to imitate the demonstration.
Revised: Then a t is the ‘correct’ action at step t for the agent to imitate the demonstration.

Original: Our supervised learning objective is to maximize the log probability of taking such actions:
Revised: Our supervised learning objective is to maximize the log probability of taking such actions:

Original: [Equation]
Revised: [Equation]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In the experiments, we aim to answer the following questions: (1) How well does the trajectoryconditioned policy imitate the diverse demonstration trajectories?
Revised: In the experiments, we aim to answer the following questions: (1) How well does the trajectoryconditioned policy imitate the diverse demonstration trajectories?

Original: (2) Does imitation of the past diverse experience enable the agent to further explore more diverse directions and guide the exploration to ﬁnd the trajectory with a near-optimal total reward?
Revised: (2) Does imitation of the past diverse experience enable the agent to further explore more diverse directions and guide the exploration to ﬁnd the trajectory with a near-optimal total reward?

Original: (3) Can our proposed method aid in avoiding myopic behaviors and converge to near-optimal solutions?
Revised: (3) Can our proposed method aid in avoiding myopic behaviors and converge to near-optimal solutions?

Original: 4. I MPLEMENTATION D ETAILS

Original: Our algorithm begins with an empty buffer D .

Original: We initialize the demonstration as a list of zero vectors, and the agent performs random exploration to collect trajectories to ﬁll the buffer D .

Original: In practice, the sampled demonstration trajectory g = { e g 0 , e g 1 , · · · , e g | g | } could be lengthy.

Original: We present a part of the demonstration as the input into the policy, similarly to translating a paragraph sentence by sentence.

Original: Speciﬁcally, we ﬁrst input { e g 0 , e g 1 , · · · , e gm } ( m ≤ | g | ) into the policy.

Original: When the index of the agent’s last visited state embedding in the demonstration u belongs to { m − ∆ t, · · · , m } , we think that the agent has accomplished this part of the demonstration, and switch to the next part

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: { e gu , e gu +1 , · · · , e gu + m } .

Original: We repeat this process until the last part of the demonstration.

Original: If the last part { e gu , e gu +1 , · · · , e g | g | } is less than m + 1 steps long, we pad the sequence with zero vectors.

Original: A reward function f ( r t ) =

Original: r t is used on the Apple-Gold, Deep Sea and MuJoco domain, and f ( r t )

Original: = 2 · clip( r t , 0 , 1) on other environments.

Original: r im = 0 .

Original: 1 is the reward to encourage imitation.

Original: More details about hyperparameters and the environments can be found in the Appendix D.

Original: We compare our method with the following baselines: (1) PPO:
Revised: We compare our method with the following baselines: (1) PPO:

Original: Proximal Policy Optimization (Schulman et al., 2017); (2) PPO+EXP: PPO with reward f ( r t ) + λ/ (cid:112) N ( e t ) , where λ/ (cid:112) N ( e t ) is the count-based exploration bonus, N ( e ) is the number of times the cluster which the state representation e belongs to was visited during training and λ is the hyper-parameter controlling the weight of exploration term; (3) PPO+SIL: PPO with Self-Imitation Learning
Revised: Proximal Policy Optimization (Schulman et al., 2017); (2) PPO+EXP: PPO with reward f ( r t ) + λ/ (cid:112) N ( e t ) , where λ/ (cid:112) N ( e t ) is the count-based exploration bonus, N ( e ) is the number of times the cluster which the state representation e belongs to was visited during training and λ is the hyper-parameter controlling the weight of exploration term; (3) PPO+SIL: PPO with Self-Imitation Learning

Original: (Oh et al., 2018).
Revised: (Oh et al., 2018).

Revised: More details about the implementation can be found in the Appendix.

Original: 4. A PPLE -G OLD D OMAIN
Revised: 4. A PPLE -G OLD D OMAIN

Original: The Apple-Gold domain (shown in Figure 1) is a simple grid-world environment with misleading rewards that can lead the agent to local optima.
Revised: The Apple-Gold domain (shown in Figure 1) is a simple grid-world environment with misleading rewards that can lead the agent to local optima.

Original: An observation consists of the agent’s location ( x t , y t ) and binary variables showing whether the agent has gotten the apple or the gold.
Revised: An observation consists of the agent’s location ( x t , y t ) and binary variables showing whether the agent has gotten the apple or the gold.

Original: A state is represented as the agent’s location and the cumulative positive reward: e t = ( x t , y t , (cid:80) ti =1 max( r i , 0)) , indicating the location of the agent and the collected objects.
Revised: A state is represented as the agent’s location and the cumulative positive reward: e t = ( x t , y t , (cid:80) t i =1 max( r i , 0)) , indicating the location of the agent and the collected objects.

Original: As shown in Figure 4a, PPO, PPO+SIL, and PPO+EXP agents are stuck with the sub-optimal policy of collecting the two apples.
Revised: As shown in Figure 4a, PPO, PPO+SIL, and PPO+EXP agents are stuck with the sub-optimal policy of collecting the two apples.

Original: In Figure 4b, PPO+EXP agent could occasionally explore further and gather the gold with total reward 8.5.
Revised: In Figure 4b, PPO+EXP agent could occasionally explore further and gather the gold with total reward 8.5.

Original: However, the agent does not replicate the good trajectory due to the negative reward along the optimal path and network forgetting about the good experiences.
Revised: However, the agent does not replicate the good trajectory due to the negative reward along the optimal path and network forgetting about the good experiences.

Original: DTSIL marches forward on the right side of the maze and achieves the highest total reward 8.5 within the time limit.
Revised: DTSIL marches forward on the right side of the maze and achieves the highest total reward 8.5 within the time limit.

Original: Figure 4c shows the number of different state embeddings found during training.
Revised: Figure 4c shows the number of different state embeddings found during training.

Original: In Figure 4d, we show the average success ratio of the imitation during training.
Revised: In Figure 4d, we show the average success ratio of the imitation during training.

Original: It is deﬁned as follows: for a given demonstration g = { e g 0 , e g 1 , · · · , e g | g | } , let u be the index of the last visited state embedding in g when the agent’s current episode terminates, then the success ratio of imitating g is u | g | (i.e., the portion of trajectory imitated).
Revised: It is deﬁned as follows: for a given demonstration g = { e g 0 , e g 1 , · · · , e g | g | } , let u be the index of the last visited state embedding in g when the agent’s current episode terminates, then the success ratio of imitating g is u | g | (i.e., the portion of trajectory imitated).

Original: Ideally, we want the success ratio to be 1.0, which indicates that the trajectory-conditioned policy could successfully follow any given demonstration from the buffer.
Revised: Ideally, we want the success ratio to be 1.0, which indicates that the trajectory-conditioned policy could successfully follow any given demonstration from the buffer.

Original: At 5M steps, the trajectories with the optimal total reward 8.5 are found, and our trajectory-conditioned policy eventually imitates them well with a success ratio around 1.0.
Revised: At 5M steps, the trajectories with the optimal total reward 8.5 are found, and our trajectory-conditioned policy eventually imitates them well with a success ratio around 1.0.

Original: [Figure]

Original: Figure 5 visualizes a learning process.
Revised: Figure 5 visualizes a learning process.

Original: PPO+SIL fails on this task because the agent quickly exploits good experiences of collecting the apples and the buffer is ﬁlled with the trajectories in the  nearby region.
Revised: PPO+SIL fails on this task because the agent quickly exploits good experiences of collecting the apples and the buffer is ﬁlled with the trajectories in the [Figure] [Table] nearby region.

Original: On the contrary, DTSIL maintains a buffer of diverse trajectories which are used as demonstrations to guide the agent to explore different regions and discover an optimal behavior.
Revised: On the contrary, DTSIL maintains a buffer of diverse trajectories which are used as demonstrations to guide the agent to explore different regions and discover an optimal behavior.

Original: 4. A TARI M ONTEZUMA ’ S R EVENGE AND P ITFALL
Revised: 4. A TARI M ONTEZUMA ’ S R EVENGE AND P ITFALL

Original: We evaluate our method on the hard-exploration game Montezuma’s Revenge and Pitfall in the Arcade
Revised: We evaluate our method on the hard-exploration game Montezuma’s Revenge and Pitfall in the Arcade

Original: Learning Environment (ALE) (Bellemare et al., 2013; Machado et al., 2017).
Revised: Learning Environment (ALE) (Bellemare et al., 2013; Machado et al., 2017).

Original: The environment setting is the same as Mnih et al.
Revised: The environment setting is the same as Mnih et al.

Original: There is a random initial delay resulting in stochasticity in the environment.
Revised: There is a random initial delay resulting in stochasticity in the environment.

Original: The observation is a frame of raw pixel images, and the state representation e t = (room t , x t , y t , (cid:80) ti =1 max( r i , 0)) consists of the agent’s ground truth location (obtained from
Revised: The observation is a frame of raw pixel images, and the state representation e t = (room t , x t , y t , (cid:80) ti =1 max( r i , 0)) consists of the agent’s ground truth location (obtained from

Original: RAM) and the accumulated positive environment reward, which implicitly indicates the objects the agent has collected 2 .
Revised: RAM) and the accumulated positive environment reward, which implicitly indicates the objects the agent has collected 2 .

Original: It is worth noting that even with the ground-truth location of the agent, on these two infamously difﬁcult games, it is highly non-trivial to explore efﬁciently and avoid local optima without relying on expert demonstrations or being able to reset to arbitrary states.
Revised: It is worth noting that even with the ground-truth location of the agent, on these two infamously difﬁcult games, it is highly non-trivial to explore efﬁciently and avoid local optima without relying on expert demonstrations or being able to reset to arbitrary states.

Original: In addition to the agent’s location information, many complicated elements such as moving entities, traps, and the agent’s inventory are included in the state.
Revised: In addition to the agent’s location information, many complicated elements such as moving entities, traps, and the agent’s inventory are included in the state.

Original: Therefore, these Atari games with agent’s location information are still much more challenging than the grid world environments.
Revised: Therefore, these Atari games with agent’s location information are still much more challenging than the grid world environments.

Original: Empirically, as summarized in Table 1, the previous SOTA baselines using the agent’s ground truth location information even fails to achieve high scores.
Revised: Empirically, as summarized in Table 1, the previous SOTA baselines using the agent’s ground truth location information even fails to achieve high scores.

Original: Using the state representation e t , we introduce a variant ‘DTSIL+EXP’ that adds a count-based exploration bonus r + t
Revised: Using the state representation e t , we introduce a variant ‘DTSIL+EXP’ that adds a count-based exploration bonus r + t

Original: = 1 / (cid:112) N ( e t ) to Eq.1 for faster exploration 3 .
Revised: = 1 / (cid:112) N ( e t ) to Eq.1 for faster exploration 3 .

Original: As shown in Figure 6a and 6e, in the early stage, the average episode reward of DTSIL+EXP is worse than PPO+EXP because our policy is trained to imitate diverse demonstrations rather than directly maximize the environment reward.
Revised: As shown in Figure 6a and 6e, in the early stage, the average episode reward of DTSIL+EXP is worse than PPO+EXP because our policy is trained to imitate diverse demonstrations rather than directly maximize the environment reward.

Original: Contrary to PPO+EXP, DTSIL+EXP agent is not eager to myopically follow the high-reward path since the path with a relatively low score in the short term might lead to higher rewards in the long term.
Revised: Contrary to PPO+EXP, DTSIL+EXP agent is not eager to myopically follow the high-reward path since the path with a relatively low score in the short term might lead to higher rewards in the long term.

Original: On Montezuma’s Revenge, for example, with two keys in hand, PPO+EXP agent often opens a nearby door and loses the chance of opening the last two doors of the ﬁrst level 4 .
Revised: On Montezuma’s Revenge, for example, with two keys in hand, PPO+EXP agent often opens a nearby door and loses the chance of opening the last two doors of the ﬁrst level 4 .

Original: As training continues, DTSIL+EXP successfully discovers trajectories to pass the ﬁrst level with a total reward of more than 20,000, as shown in Figure 6b.
Revised: As training continues, DTSIL+EXP successfully discovers trajectories to pass the ﬁrst level with a total reward of more than 20,000, as shown in Figure 6b.

Original: While gradually increasing the probability of imitating 
Revised: While gradually increasing the probability of imitating the best trajectories in the buffer by sampling them as demonstrations, the average episode reward

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: the best trajectories in the buffer by sampling them as demonstrations, the average episode reward could increase to surpass 20,000 in Figure 6a.
Revised: could increase to surpass 20,000 in Figure 6a.

Original: On Pitfall, the positive reward is much sparser and most of the actions yield small negative rewards that would discourage getting a high total reward in the long term.
Revised: On Pitfall, the positive reward is much sparser and most of the actions yield small negative rewards that would discourage getting a high total reward in the long term.

Original: However, our method still stores the trajectories with negative rewards, encourages the agent to visit these novel regions and then discovers good paths with positive rewards as illustrated in Figure 6f.
Revised: However, our method still stores the trajectories with negative rewards, encourages the agent to visit these novel regions and then discovers good paths with positive rewards as illustrated in Figure 6f.

Original: Therefore we are able to eventually reach average episode reward over 0 in Figure 6e, without expert demonstrations.
Revised: Therefore we are able to eventually reach average episode reward over 0 in Figure 6e, without expert demonstrations.

Original: Table 1 compares our proposed method with previous works without using any expert demonstration or resetting to an arbitrary state, where our approach signiﬁcantly outperforms the other approaches which make use of the same information from RAM about the agent’s location.
Revised: Table 1 compares our proposed method with previous works without using any expert demonstration or resetting to an arbitrary state, where our approach signiﬁcantly outperforms the other approaches which make use of the same information from RAM about the agent’s location.

Original: In the Appendix, we present more experimental results on other interesting environments with discrete action space such as Deep Sea (Osband et al., 2019).
Revised: In the Appendix C.3, we present more experimental results on other interesting environments with discrete action space such as Deep Sea (Osband et al., 2019).

Original: 4. M UJOCO
Revised: 4. M UJOCO

Original: We evaluate DTSIL on continuous control tasks.
Revised: We evaluate DTSIL on continuous control tasks.

Original: We adapt the maze environment introduced in (Duan et al., 2016) to construct a set of challenging tasks, which require the point mass agent to collect the key, open the door with the same color and ﬁnally reach the treasure to get a high score.
Revised: We adapt the maze environment introduced in (Duan et al., 2016) to construct a set of challenging tasks, which require the point mass agent to collect the key, open the door with the same color and ﬁnally reach the treasure to get a high score.

Original: One key cannot be re-used once it was used before to open a door with the same color, which makes the agent to be easily trapped.
Revised: One key cannot be re-used once it was used before to open a door with the same color, which makes the agent to be easily trapped.

Original: A visualization of these environments is shown in Figure 7.
Revised: A visualization of these environments is shown in Figure 7.

Revised: The agent’s initial location is randomly sampled from a Gaussian distribution as in standard MuJoco tasks (Brockman et al., 2016).

Original: The observation is the agent’s location and range sensor reading about nearby objects.
Revised: The observation is the agent’s location and range sensor reading about nearby objects.

Original: The state representation is e t =
Revised: The state representation is e t =

Original: ( x t , y t , (cid:80) ti =1 r i ) .
Revised: ( x t , y t , (cid:80) ti =1 r i ) .

Original: As shown in the ﬁrst maze of Figure 7, the agent can easily get the blue key near its initial location and open the blue door in the upper part.
Revised: As shown in the ﬁrst maze of Figure 7, the agent can easily get the blue key near its initial location and open the blue door in the upper part.

Original: However, the optimal path is to bring the key to open the blue door in the bottom and obtain the treasure, reaching an episode reward of 9.
Revised: However, the optimal path is to bring the key to open the blue door in the bottom and obtain the treasure, reaching an episode reward of 9.

Original: In the second maze, the agent should bring the blue key and pick up the green key while avoiding opening the blue door in the upper part.
Revised: In the second maze, the agent should bring the blue key and pick up the green key while avoiding opening the blue door in the upper part.

Original: Then, the green and blue key can open the two doors at the bottom of the maze, which results in the total reward of 12.
Revised: Then, the green and blue key can open the two doors at the bottom of the maze, which results in the total reward of 12.

Original: The learning curves in Figure 7 show that PPO, PPO+EXP, and PPO+SIL may get stuck at a sub-optimal behavior, whereas our policy eventually converges to the behavior achieving the high episode reward.
Revised: The learning curves in Figure 7 show that PPO, PPO+EXP, and PPO+SIL may get stuck at a sub-optimal behavior, whereas our policy eventually converges to the behavior achieving the high episode reward.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: This paper proposes to learn diverse policies by imitating diverse trajectory-level demonstrations through count-based exploration over these trajectories.
Revised: This paper proposes to learn diverse policies by imitating diverse trajectory-level demonstrations through count-based exploration over these trajectories.

Original: Imitation of diverse past trajectories can guide the agent to rarely visited states and encourages further exploration of novel states.
Revised: Imitation of diverse past trajectories can guide the agent to rarely visited states and encourages further exploration of novel states.

Original: We show that in a variety of environments with local optima, our method signiﬁcantly improves self-imitation learning (SIL).
Revised: We show that in a variety of environments with local optima, our method signiﬁcantly improves self-imitation learning (SIL).

Original: It avoids prematurely converging to a myopic solution and learns a near-optimal behavior to achieve a high total reward.
Revised: It avoids prematurely converging to a myopic solution and learns a near-optimal behavior to achieve a high total reward.

Original: Appendix A D ETAILED D ESCRIPTION OF A LGORITHMS A. A LGORITHM OF T RAJECTORY B UFFER U PDATE
Revised: Appendix A D ETAILED D ESCRIPTION OF A LGORITHMS A. A LGORITHM OF T RAJECTORY B UFFER U PDATE

Original: In Algorithm 2, we summarize how to process the collected episode and store the diverse trajectories in the trajectory buffer.
Revised: In Algorithm 2, we summarize how to process the collected episode and store the diverse trajectories in the trajectory buffer.

Original: Algorithm 2 Update Trajectory Buffer
Revised: Algorithm 2 Update Trajectory Buffer

Original: Input: the trajectory buffer D
Revised: Input: the trajectory buffer D

Original: = { ( e (1) , τ (1) , n (1) ) , ( e (2) , τ (2) , n (2) ) , · · · }
Revised: = { ( e (1) , τ (1) , n (1) ) , ( e (2) , τ (2) , n (2) ) , · · · }

Original: Input: the current episode E = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) , · · · , ( o T , e T , a T , r T ) }
Revised: Input: the current episode E = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) , · · · , ( o T , e T , a T , r T ) }

Original: Input: the threshold th for high level state embedding # Consider all the states in E for each step t do # Consider state s t and partial episode τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , · · · , ( o t , e t , a t , r t ) } if there exists ( e ( k ) , τ ( k ) , n ( k ) )
Revised: Input: the threshold th for high level state embedding # Consider all the states in E for each step t do # Consider state s t and partial episode τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , · · · , ( o t , e t , a t , r t ) } if there exists ( e ( k ) , τ ( k ) , n ( k ) )

Original: ∈ D where (cid:107) e
Revised: ∈ D where (cid:107) e

Original: ( k ) − e t (cid:107) <
Revised: ( k ) − e t (cid:107) <

Original: th then # Compare partial episode τ ≤ t with stored trajectory τ ( k ) if τ ≤ t has higher total reward or reaches the same total reward with less steps then τ ( k ) ← τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) , · · · , ( o t , e t , a t , r t ) } e ( k ) ← e t end if n ( k ) ← n ( k ) +else
Revised: th then # Compare partial episode τ ≤ t with stored trajectory τ ( k ) if τ ≤ t has higher total reward or reaches the same total reward with less steps then τ ( k ) ← τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) , · · · , ( o t , e t , a t , r t ) } e ( k ) ← e t end if n ( k ) ← n ( k ) +else

Original: D ← D ∪ ( e t , τ ≤ t , 1) where τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) , · · · , ( o t , e t , a t , r t ) } end if end for return D
Revised: D ← D ∪ ( e t , τ ≤ t , 1) where τ ≤ t = { ( o 0 , e 0 , a 0 , r 0 ) , ( o 1 , e 1 , a 1 , r 1 ) , · · · , ( o t , e t , a t , r t ) } end if end for return D

Original: A. A LGORITHM OF S AMPLING D EMONSTRATIONS
Revised: A. A LGORITHM OF S AMPLING D EMONSTRATIONS

Original: In Algorithm 3, we summarize how to sample the demonstrations from the trajectory buffer for exploration or imitation.
Revised: In Algorithm 3, we summarize how to sample the demonstrations from the trajectory buffer for exploration or imitation.

Original: Considering the current iteration i and the total number of iterations I , the probability of sampling demonstration for imitation to learn good behavior is iI and the probability of sampling demonstration from exploration is 1 − iI .
Revised: Considering the current iteration i and the total number of iterations I , the probability of sampling demonstration for imitation to learn good behavior is iI and the probability of sampling demonstration from exploration is 1 − iI .

Original: Algorithm 3 Sample Demonstration Trajectories
Revised: Algorithm 3 Sample Demonstration Trajectories

Original: Input: the trajectory buffer D
Revised: Input: the trajectory buffer D

Original: =
Revised: =

Original: e (1) , τ (1) , n (1) ) , ( e (2) , τ (2) , n (2) ) , · · · }
Revised: e (1) , τ (1) , n (1) ) , ( e (2) , τ (2) , n (2) ) , · · · }

Original: Input: current iteration i , total number of iterations I .
Revised: Input: current iteration i , total number of iterations I .

Original: # With probability iI , run the imitation mode; with probability 1 − iI , run the exploration mode if random number ∼ U [0 , 1] is smaller than iI then # sample the top-K trajectories reaching near-optimal score in the buffer g ← { e 0 , e 1 , · · · , e | g | } for all ( o t , e t , a t , r t ) ∈ τ best else
Revised: # With probability iI , run the imitation mode; with probability 1 − iI , run the exploration mode if random number ∼ U [0 , 1] is smaller than iI then # sample the top-K trajectories reaching near-optimal score in the buffer g ← { e 0 , e 1 , · · · , e | g | } for all ( o t , e t , a t , r t ) ∈ τ best else

Original: Calculate probability distribution p ← [ 1 √ n (1) , 1 √ n (2) , · · · ] p ← p (cid:80) j p j
Revised: Calculate probability distribution p ← [ 1 √ n (1) , 1 √ n (2) , · · · ] p ← p (cid:80) j p j

Original: Sample ( e, τ, n ) ∼ Categorical ( D , p )
Revised: Sample ( e, τ, n ) ∼ Categorical ( D , p )

Original: g ← { e 0 , e 1 , · · · , e | g | } for all ( o t , e t , a t , r t ) ∈ τ end if return g
Revised: g ← { e 0 , e 1 , · · · , e | g | } for all ( o t , e t , a t , r t ) ∈ τ end if return g

Original: B D ETAILS OF N ETWORK A RCHITECTURE 
Revised: B D ETAILS OF N ETWORK A RCHITECTURE AND T RAINING P ROCESS

Original: In the trajectory-conditioned policy (Figure 3), we ﬁrst embed the input state e t (or e gi ) with a fully-connected layer with 64 units.
Revised: In the trajectory-conditioned policy (Figure 3), we ﬁrst embed the input state e t (or e gi ) with a fully-connected layer with 64 units.

Original: Next,  RNN with gated recurrent units (GRU) computes the feature h t (or h gi ) with 128 units.
Revised: Next, a RNN with gated recurrent units (GRU) computes the feature h t (or h gi ) with 128 units.

Original: The attention weight α t is calculated based on the Bahdanau attention mechanism (Bahdanau et al., 2014).
Revised: The attention weight α t is calculated based on the Bahdanau attention mechanism (Bahdanau et al., 2014).

Original: The concatenation of the attention readout c t , the hidden feature of agent’s current state h t , and convolutional features from the observation are used to predict π ( a t | e ≤ t , o t , g ) with a linear layer.
Revised: The concatenation of the attention readout c t , the hidden feature of agent’s current state h t , and convolutional features from the observation are used to predict π ( a t | e ≤ t , o t , g ) with a linear layer.

Original: For experiments on the Apple-Key domain, Toy Montezuma’s Revenge, and Mujoco, the features from o t are not required for the policy.
Revised: For experiments on the Apple-Key domain, Toy Montezuma’s Revenge, and Mujoco, the features from o t are not required for the policy.

Original: However, on the Atari games such as Montezuma’s Revenge, it is necessary to take the raw observation o t as input into policy because the location information in e ≤ t solely could not let the agent to take temporal context into account (e.g. avoiding moving skulls and passing laser gates).
Revised: However, on the Atari games such as Montezuma’s Revenge, it is necessary to take the raw observation o t as input into policy because the location information in e ≤ t solely could not let the agent to take temporal context into account (e.g. avoiding moving skulls and passing laser gates).

Original: With the raw observation o t with shape 84 × 84 × 4 as input, three convolutional layers are used to encode o t and then the convolutional feature is ﬂattened.
Revised: With the raw observation o t with shape 84 × 84 × 4 as input, three convolutional layers are used to encode o t and then the convolutional feature is ﬂattened.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Revised: During training, our algorithm begins with an empty buffer D .

Revised: We initialize the demonstration as a list of zero vectors.

Original: C A DDITIONAL E XPERIMENTS C. G ENERALIZATION AND R OBUSTNESS IN S TOCHASTIC E NVIRONMENTS
Revised: With such an input demonstration, the agent performs random exploration to collect trajectories to ﬁll the buffer D . In practice, the sampled demonstration trajectory g = { e g 0 , e g 1 , · · · , e g | g | } could be lengthy. We present a part of the demonstration as the input into the policy, similarly to translating a paragraph sentence by sentence. Speciﬁcally, we ﬁrst input { e g 0 , e g 1 , · · · , e gm } ( m ≤ | g | ) into the policy.

Revised: When the index of the agent’s last visited state embedding in the demonstration u belongs to { m − ∆ t, · · · , m } , we think that the agent has accomplished this part of the demonstration, and switch to the next part { e gu , e gu +1 , · · · , e gu + m } .

Revised: We repeat this process until the last part of the demonstration.

Revised: If the last part { e gu , e gu +1 , · · · , e g | g | } is less than m + 1 steps long, we pad the sequence with zero vectors.

Revised: A reward function f ( r t ) =

Revised: r t is used on the Apple-Gold, Deep Sea and MuJoco domain, and f ( r t )

Revised: = 2 · clip( r t , 0 , 1) on other environments.

Revised: r im = 0 .

Revised: 1 is the reward to encourage imitation.

Revised: More details about hyperparameters and the environments can be found in the Appendix D.

Revised: C A DDITIONAL E XPERIMENTS C. G ENERALIZATION AND R OBUSTNESS IN S TOCHASTIC E NVIRONMENTS

Original: We evaluate our method on environments with different levels of stochasticity.
Revised: We evaluate our method on environments with different levels of stochasticity.

Original: For Apple-Gold domain, in the environments with random initial location of the agent (Figure 9) [Figure] or with sticky action (Figure 10), our DTSIL still outperforms the baselines 
Revised: For Apple-Gold domain, in the environments with random initial location of the agent (Figure 9), or with sticky action (Figure 10), our DTSIL still outperforms the baselines and achieves near-optimal total episode reward.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: In a more challenging scenario when the location of the objects can be random, previous works using expert demonstrations (e
Revised: In a more challenging scenario when the location of the objects can be random, previous works using expert demonstrations (e

Original: [Figure]
Revised: [Figure]

Original: g. Aytar et al.
Revised: g. Aytar et al.

Original: (2018)) would also struggle.
Revised: (2018)) would also struggle.

Original: Our method can be easily extended to handle this difﬁculty by ditilling the behavior of good trajectories we collected to train an unconditional policy robust to the stochasticity.
Revised: Our method can be easily extended to handle this difﬁculty by ditilling the behavior of good trajectories we collected to train an unconditional policy robust to the stochasticity.

Original: For example, on the Apple-Gold domain with pixel observation (Figure 11), the location of the gold could be random in the upper middle part of the maze.
Revised: For example, on the Apple-Gold domain with pixel observation (Figure 11), the location of the gold could be random in the upper middle part of the maze.

Original: We ﬁrst explore for a sufﬁciently large number of timesteps (e
Revised: We ﬁrst explore for a sufﬁciently large number of timesteps (e

Original: [Figure]
Revised: [Figure]

Original: g. 10M timesteps) with the trajectory-conditioned policy to collect good trajectories and then train an unconditional policy by distilling the behavior, as shown in Figure 11, to always collect the gold.
Revised: g. 10M timesteps) with the trajectory-conditioned policy to collect good trajectories and then train an unconditional policy by distilling the behavior, as shown in Figure 11, to always collect the gold.

Original: We could see DTSIL with unconditional policy training is able to generalize in the stochastic environment.[Figure] g e R
Revised: We could see DTSIL with unconditional policy training is able to generalize in the stochastic environment.

Original: [Figure]
Revised: [Figure]

Original: C. E XPERIMENTS ON T OY M ONTEZUMA R EVENGE
Revised: C. E XPERIMENTS ON T OY M ONTEZUMA R EVENGE

Original: We evaluate our method on a more challenging domain, Toy Montezuma’s Revenge (Roderick et al., 2018), which requires a more sophisticated strategy to explore the environment.
Revised: We evaluate our method on a more challenging domain, Toy Montezuma’s Revenge (Roderick et al., 2018), which requires a more sophisticated strategy to explore the environment.

Original: As shown in Figure 12, there are 24 rooms similar to the layout of the ﬁrst level of Atari Montezuma’s Revenge, with a discrete grid for each room.
Revised: As shown in Figure 12, there are 24 rooms similar to the layout of the ﬁrst level of Atari Montezuma’s Revenge, with a discrete grid for each room.

Original: The agent should navigate the labyrinth to locate the keys, unlock the doors and reach the goal (the treasure room).
Revised: The agent should navigate the labyrinth to locate the keys, unlock the doors and reach the goal (the treasure room).

Original: The observation is represented by the agent’s location and cumulative episode reward.
Revised: The observation is represented by the agent’s location and cumulative episode reward.

Original: The state representation e t = (room t , x t , y t , (cid:80) ti =1 r i ) is the same as the observation.
Revised: The state representation e t = (room t , x t , y t , (cid:80) ti =1 r i ) is the same as the observation.

Original: The learning curve of the averaged episode reward in Figure 13 shows that PPO, PPO+SIL, and
Revised: The learning curve of the averaged episode reward in Figure 13 shows that PPO, PPO+SIL, and

Original: PPO+EXP could not learn a policy to reach the goal.
Revised: PPO+EXP could not learn a policy to reach the goal.

Original: The PPO+EXP agent occasionally ﬁnds a trajectory with the total reward of 11,200 reaching the treasure room, but fails to exploit this experience.
Revised: The PPO+EXP agent occasionally ﬁnds a trajectory with the total reward of 11,200 reaching the treasure room, but fails to exploit this experience.

Original: On the other hand, our method learns a good behavior of not only reaching the goal room, but also collecting all of the keys to achieve an optimal total reward of 11,600.
Revised: On the other hand, our method learns a good behavior of not only reaching the goal room, but also collecting all of the keys to achieve an optimal total reward of 11,600.

Original: [Figure]
Revised: [Figure]

Original: A v e r a g
Revised: A v e r a g

Original: e
Revised: e

Original: I m i t a t i o n S u cc e ss R a t i o
Revised: I m i t a t i o n S u cc e ss R a t i o

Original: C. E XPERIMENTS ON D EEP S EA
Revised: C. E XPERIMENTS ON D EEP S EA

Original: As introduced in Osband et al.
Revised: As introduced in Osband et al.

Original: (2019), the deep sea problem is implemented as an N × N grid with a one-hot encoding for state.
Revised: (2019), the deep sea problem is implemented as an N × N grid with a one-hot encoding for state.

Original: The agent begins each episode in the top left corner of the grid and descends one row per timestep.
Revised: The agent begins each episode in the top left corner of the grid and descends one row per timestep.

Original: Each episode terminates after N steps, when the agent reaches the bottom row.
Revised: Each episode terminates after N steps, when the agent reaches the bottom row.

Original: In each state there is a random but ﬁxed mapping between actions
Revised: In each state there is a random but ﬁxed mapping between actions

Original: A = { 0 , 1 } and the transitions ‘left’ and ‘right’.
Revised: A = { 0 , 1 } and the transitions ‘left’ and ‘right’.

Original: At each timestep there is a small cost r = − 0 .
Revised: At each timestep there is a small cost r = − 0 .

Original: 01 /N of moving right, and r = 0 for moving left.
Revised: 01 /N of moving right, and r = 0 for moving left.

Original: However, should the agent transition right at every timestep of the episode it will be rewarded with an additional reward of +1.
Revised: However, should the agent transition right at every timestep of the episode it will be rewarded with an additional reward of +1.

Original: This presents a particularly challenging exploration problem for two reasons.
Revised: This presents a particularly challenging exploration problem for two reasons.

Original: First, following the ‘gradient’ of small intermediate rewards leads the agent away from the optimal policy.
Revised: First, following the ‘gradient’ of small intermediate rewards leads the agent away from the optimal policy.

Original: Second, a policy that explores with actions uniformly at random has probability 2 − N of reaching the rewarding state in any episode.
Revised: Second, a policy that explores with actions uniformly at random has probability 2 − N of reaching the rewarding state in any episode.

Original: We compare DTSIL and baselines on deep sea environments with 10 × 10 grid and and 30 × 30 grid.
Revised: We compare DTSIL and baselines on deep sea environments with 10 × 10 grid and and 30 × 30 grid.

Original: The state embedding we use here is exactly the observation.
Revised: The state embedding we use here is exactly the observation.

Original: The result is shown in Figure 14.
Revised: The result is shown in Figure 14.

Original: On the ﬁrst environment, it is easy for all of the methods to converge to the optimal behavior.
Revised: On the ﬁrst environment, it is easy for all of the methods to converge to the optimal behavior.

Original: The second one is much more challenging to ﬁnd the optimal trajectory maximizing total reward.
Revised: The second one is much more challenging to ﬁnd the optimal trajectory maximizing total reward.

Original: Therefore, PPO and PPO+SIL fails at such environment due to the hard exploration.
Revised: Therefore, PPO and PPO+SIL fails at such environment due to the hard exploration.

Original: PPO+EXP could not always explore to ﬁnd the good behavior and exploit it efﬁciently within 12M timesteps.
Revised: PPO+EXP could not always explore to ﬁnd the good behavior and exploit it efﬁciently within 12M timesteps.

Original: DTSIL successfully discovers the right way and imitate to converge to the optimal behavior.
Revised: DTSIL successfully discovers the right way and imitate to converge to the optimal behavior.

Original: [Figure]
Revised: [Figure]

Original: D H YPERPARAMETERS
Revised: D H YPERPARAMETERS

Original: The hyper-parameters for our proposed method used in each experiment are listed in Table 2.
Revised: The hyper-parameters for our proposed method used in each experiment are listed in Table 2.

Original: On
Revised: On

Original: Mujoco environment, RL loss alone worked well so we did not include SL loss for behavior cloning.
Revised: Mujoco environment, RL loss alone worked well so we did not include SL loss for behavior cloning.

Original: On the other environments when action prediction in behavior cloning is poor [Table] we set a large J for quickly learning to imitate demonstrations.
Revised: On the other environments when action prediction in behavior cloning is poor, we set a large J for quickly learning to imitate demonstrations.

Original: When action prediction is accurate enough [Table] we de-emphasize behavior cloning to enhance exploration around the demonstration.
Revised: When action prediction is accurate enough, we de-emphasize behavior cloning to enhance exploration around the demonstration.

Original: ∆ t inﬂuences how ﬂexibly the demonstration should be followed.
Revised: ∆ t inﬂuences how ﬂexibly the demonstration should be followed.

Original: In our experiment, we have ∆ t < m due to the limit of the length m of the input demonstration part.
Revised: In our experiment, we have ∆ t < m due to the limit of the length m of the input demonstration part.

Original: When the demonstration is longer and harder to follow, we would want larger ∆ t to generously provide imitation reward.
Revised: When the demonstration is longer and harder to follow, we would want larger ∆ t to generously provide imitation reward (More detailed ablation study of the hyper-parameter ∆ t is in Appendix J).

Original: On Atari games, there are much more different trajectories stored in the buffer, so we sample top-100 trajectories as demonstration for imitation of best experiences.
Revised: On Atari games, there are much more different trajectories stored in the buffer, so we sample top-100 trajectories as demonstration for imitation of best experiences.

Original: On the other environments, the total number of trajectories is much smaller, so we only take top-10 or top-1.
Revised: On the other environments, the total number of trajectories is much smaller, so we only take top-10 or top-

Original: [Table]
Revised: [Table]

Original: E E NVIRONMENT S ETTING
Revised: E E NVIRONMENT S ETTING

Original: For each experiment we conducted, we list the detailed environment setting in Table 3.
Revised: For each experiment we conducted [Table] we list the detailed environment setting in Table 3.

Original: There is stochasticity in the environments of Apple-Gold domain, Toy MontezumaRevenge, Atari
Revised: There is stochasticity in the environments of Apple-Gold domain, Toy MontezumaRevenge, Atari

Original: [Table]
Revised: [Table]

Original: and Mujoco.
Revised: and Mujoco.

Original: On Atari games, we use setting of the random initial delay introduced in (Mnih et al., 2015).
Revised: On Atari games, we use setting of the random initial delay introduced in (Mnih et al., 2015).

Original: On
Revised: On

Original: Mujoco domain, the agent’s initial location is randomly sampled from a Gaussian distribution, as in standard MuJoco tasks in OpenAI Gym (Brockman et al., 2016).
Revised: Mujoco domain, the agent’s initial location is randomly sampled from a Gaussian distribution, as in standard MuJoco tasks in OpenAI Gym (Brockman et al., 2016).

Original: [Table]
Revised: [Table]

Original: F C OMPARISON WITH L EARNING D IVERSE P OLICIES BY SVPG
Revised: F C OMPARISON WITH L EARNING D IVERSE P OLICIES BY SVPG

Original: While the code for the Stein variational policy gradient (SVPG) in Gangwani et al.
Revised: While the code for the Stein variational policy gradient (SVPG) in Gangwani et al.

Original: (2018) has not yet been released, we replicate the method in Gangwani et al.
Revised: (2018) has not yet been released, we replicate the method in Gangwani et al.

Original: (2018) to learn diverse policies.
Revised: (2018) to learn diverse policies.

Original: Their experiments focus on continuous control tasks with relatively simple observation spaces with limited local optimal branches in the state space.
Revised: Their experiments focus on continuous control tasks with relatively simple observation spaces with limited local optimal branches in the state space.

Original: We learn 8 diverse policies in parallel following their method on our Apple-Gold domain with discrete action space.
Revised: We learn 8 diverse policies in parallel following their method on our Apple-Gold domain with discrete action space.

Original: Figure 15 shows a visualization of the learning progress: the 8 policies learn to cover different regions of the environment.
Revised: Figure 15 shows a visualization of the learning progress: the 8 policies learn to cover different regions of the environment.

Original: The method explores better than PPO+SIL, but the exploration of each individual agent is not strong enough to ﬁnd the optimal path to achieve the highest episode reward.
Revised: The method explores better than PPO+SIL, but the exploration of each individual agent is not strong enough to ﬁnd the optimal path to achieve the highest episode reward.

Original: [Figure]
Revised: [Figure]

Original: G M AP OF A TARI M ONTEZUMA ’ S R EVENGE AT THE F IRST L EVEL
Revised: G M AP OF A TARI M ONTEZUMA ’ S R EVENGE AT THE F IRST L EVEL

Original: On Montezuma’s Revenge, there are multiple levels and each level consists of 24 rooms.
Revised: On Montezuma’s Revenge, there are multiple levels and each level consists of 24 rooms.

Original: A map of
Revised: A map of

Original: Atari Montezuma’s Revenge at the ﬁrst level is shown in Figure 16.
Revised: Atari Montezuma’s Revenge at the ﬁrst level is shown in Figure 16.

Original: It is challenging to bring two keys to open the two doors in room 17 behind the treasure in room 15, where the agent can pass to the next level.
Revised: It is challenging to bring two keys to open the two doors in room 17 behind the treasure in room 15, where the agent can pass to the next level.

Original: H S TUDY OF E XPLORATION E FFICIENCY
Revised: H S TUDY OF E XPLORATION E FFICIENCY

Original: To evaluate the efﬁciency of exploration, we compare our method with the “exploration phase” in the
Revised: To evaluate the efﬁciency of exploration, we compare our method with the “exploration phase” in the

Original: Go-Explore algorithm
Revised: Go-Explore algorithm

Original: (Ecoffet et al., 2019).
Revised: (Ecoffet et al., 2019).

Original: The idea behind Go-Explore is to reset the agent to any interesting state sampled from the buffer of state embeddings, and then explore further using random actions.
Revised: The idea behind Go-Explore is to reset the agent to any interesting state sampled from the buffer of state embeddings, and then explore further using random actions.

Original: To study the exploration efﬁciency of our method, we modify the Go-Explore code such that we could not reset to any arbitrary states in the environment.
Revised: To study the exploration efﬁciency of our method, we modify the Go-Explore code such that we could not reset to any arbitrary states in the environment.

Original: Similarly to (Ecoffet et al., 2019), we use the state representation (level t , room t , x t , y t , k t ) where k t is the number of keys the agent holds and ( x t , y t ) is in a 9 × 9 grid division of the frame, and the sampling weight 1 √ n ( i ) to sample goal states from the buffer (It worth to note that the state representation and goal-state sampling function recommended in Go-Explore paper is more complicated than this setting).
Revised: Similarly to (Ecoffet et al., 2019), we use the state representation (level t , room t , x t , y t , k t ) where k t is the number of keys the agent holds and ( x t , y t ) is in a 9 × 9 grid division of the frame, and the sampling weight 1 √ n ( i ) to sample goal states from the buffer (It worth to note that the state representation and goal-state sampling function recommended in Go-Explore paper is more complicated than this setting).

Original: In the Go-Explore method without using the direct ‘reset’ function and with a perfect goal-conditioned policy to visit any state sampled from the buffer, the agent could precisely advance to the goal state by following the stored trajectory.
Revised: In the Go-Explore method without using the direct ‘reset’ function and with a perfect goal-conditioned policy to visit any state sampled from the buffer, the agent could precisely advance to the goal state by following the stored trajectory.

Original: The total steps taken in the environment are counted by summing the number of steps taken to follow the stored trajectories and the number of steps taken to explore.
Revised: The total steps taken in the environment are counted by summing the number of steps taken to follow the stored trajectories and the number of steps taken to explore.

Original: [Figure]
Revised: [Figure]

Original: t ) .
Revised: t ) .

Original: In Figure
Revised: In Figure

Original: [Figure]
Revised: [Figure]

Original: we show the average number of rooms found and the number of different state representations found during training.
Revised: we show the average number of rooms found and the number of different state representations found during training.

Original: Even if we assume that there is a perfect goal-conditioned policy in Go-Explore to guide the agent to follow the stored trajectory exactly and visit the goal state, the learning curves demonstrate that our method is more efﬁcient for exploring diverse state representations and consequently visits several rooms.
Revised: Even if we assume that there is a perfect goal-conditioned policy in Go-Explore to guide the agent to follow the stored trajectory exactly and visit the goal state, the learning curves demonstrate that our method is more efﬁcient for exploring diverse state representations and consequently visits several rooms.

Original: This is because our method uses the countbased exploration bonus to encourages the exploration around and beyond the stored trajectories and the imitation reward allows the agent to follow the demonstrations in a soft-order
Revised: This is because our method uses the countbased exploration bonus to encourages the exploration around and beyond the stored trajectories and the imitation reward allows the agent to follow the demonstrations in a soft-order.

Original: [Figure]
Revised: [Figure]

Original: In addition, we notice that, comparing with the state embedding (room t , x t , y t , (cid:80) ti =1 r i ) , the state embedding (level t , room t , x t , y t , k t ) makes the size of embedding space smaller so that the exploration could be more efﬁcient.
Revised: In addition, we notice that, comparing with the state embedding (room t , x t , y t , (cid:80) ti =1 r i ) , the state embedding (level t , room t , x t , y t , k t ) makes the size of embedding space smaller so that the exploration could be more efﬁcient.

Original: Such state representation conﬂates similar states while not conﬂating states that are meaningfully different.
Revised: Such state representation conﬂates similar states while not conﬂating states that are meaningfully different.

Original: Therefore, our method could reach a higher average score around 29,817.
Revised: Therefore, our method could reach a higher average score around 29,817.

Original: Here, the baseline PPO+Exp is essentially the CoEX method introduced in Choi et al.
Revised: Here, the baseline PPO+Exp is essentially the CoEX method introduced in Choi et al.

Original: (2018) with state embedding extracted from RAM, therefore DTSIL performs better than PPO+CoEX+RAM no matter the state embedding is (room t , x t , y t , (cid:80) t i =1 r i ) or (level t , room t , x t , y t , k t ) .
Revised: (2018) with state embedding extracted from RAM, therefore DTSIL performs better than PPO+CoEX+RAM no matter the state embedding is (room t , x t , y t , (cid:80) ti =1 r i ) or (level t , room t , x t , y t , k t ) .

Original: I S TUDY OF A DVANTAGE IN E XPLOITATION
Revised: I S TUDY OF A DVANTAGE IN E XPLOITATION

Original: We compared our method DTSIL with PPO+EXP in the main text.
Revised: We compared our method DTSIL with PPO+EXP in the main text.

Original: PPO+EXP encourages exploration to novel states by providing auxiliary rewards to the agent, while our method rewards the agent when it successfully follow the demonstrations which leads to novel states.
Revised: PPO+EXP encourages exploration to novel states by providing auxiliary rewards to the agent, while our method rewards the agent when it successfully follow the demonstrations which leads to novel states.

Original: In order to understand more about the difference between these two mechanism, we propose a variant of our method denoted as “DTSIL-combine”.
Revised: In order to understand more about the difference between these two mechanism, we propose a variant of our method denoted as “DTSIL-combine”.

Original: In this variant, we do not separate the exploration mode and imitation mode.
Revised: In this variant, we do not separate the exploration mode and imitation mode.

Original: However, we always sample the top-K best trajectories with highest total reward (cid:80) | τ | t =

Revised: However, we always sample the top-K best trajectories with highest total reward (cid:80) | τ | t =

Original: ( f ( r t ) + λ √ N ( e t ) ) .
Revised: ( f ( r t ) + λ √ N ( e t ) ) .

Original: The PPO+EXP baseline directly optimize the objective (cid:80) | τ | t =

Revised: The PPO+EXP baseline directly optimize the objective (cid:80) | τ | t =

Original: ( f ( r t ) + λ √ N ( e t ) ) via the reinforcement learning algorithm while this variant indirectly optimizes such objective by imitating the best trajectories with highest value of (cid:80) | τ | t =

Revised: ( f ( r t ) + λ √ N ( e t ) ) via the reinforcement learning algorithm while this variant indirectly optimizes such objective by imitating the best trajectories with highest value of (cid:80) | τ | t =

Original: ( f ( r t ) + λ √ N ( e t ) ) .
Revised: ( f ( r t ) + λ √ N ( e t ) ) .

Original: We investigate the different performance of these two methods on the Apple-Gold domain.
Revised: We investigate the different performance of these two methods on the Apple-Gold domain.

Original: [Figure]
Revised: [Figure]

Original: With λ = 10 which is the best hyper-parameter we found after searching λ = 5 , 10 , 20 , 50 for
Revised: With λ = 10 which is the best hyper-parameter we found after searching λ = 5 , 10 , 20 , 50 for

Original: PPO+EXP on Apple-Gold domain, in Figure 19c we could notice both PPO+EXP agent and DTSILcombine agent have 3 out of 5 runs ﬁnding the optimal trajectory with episode rewards 8.

Revised: PPO+EXP on Apple-Gold domain, in Figure 19c we could notice both PPO+EXP agent and DTSILcombine agent have 3 out of 5 runs ﬁnding the optimal trajectory with episode rewards 8.

Original: and the other 2 runs get stuck at the sub-optimal behavior.
Revised: and the other 2 runs get stuck at the sub-optimal behavior.

Original: However, in Figure 19b, DTSIL-combine is better at optimizing the objective (cid:80) | τ | t =

Revised: However, in Figure 19b, DTSIL-combine is better at optimizing the objective (cid:80) | τ | t =

Original: f ( r t )
Revised: f ( r t )

Original: + 10 √ N ( e t ) averaging over the agent’s recent episodes and therefore it achieves higher environment environment reward as training goes on.
Revised: + 10 √ N ( e t ) averaging over the agent’s recent episodes and therefore it achieves higher environment environment reward as training goes on.

Original: As shown in Figure 19a, DTSIL-combine agent reproduce the good trajectories to collect the gold in 3 out of 5 runs while PPO+EXP agent is trapped at the behavior collecting the apples.
Revised: As shown in Figure 19a, DTSIL-combine agent reproduce the good trajectories to collect the gold in 3 out of 5 runs while PPO+EXP agent is trapped at the behavior collecting the apples.

Original: The main reason might be that DTSIL-combine agent never forgets the good experience of collecting gold, and always select such good trajectories as demonstrations to guide the agent, while PPO+EXP might forget the good trajectories occasionally found or fails to exploit them before the exploration bonus vanishes.
Revised: The main reason might be that DTSIL-combine agent never forgets the good experience of collecting gold, and always select such good trajectories as demonstrations to guide the agent, while PPO+EXP might forget the good trajectories occasionally found or fails to exploit them before the exploration bonus vanishes.

Original: The importance of exploitation of the good experience to help the agent reproduce high-reward trajectories is also discussed in Oh et al.
Revised: The importance of exploitation of the good experience to help the agent reproduce high-reward trajectories is also discussed in Oh et al.

Revised: J A BLATION S TUDY OF H YPER - PARAMETER ∆ t

Revised: We study the effect of the hyper-parameter ∆ t in the various environments.

Revised: At each step, we provide m = 10 (considering the computational burden, we selected the value of m as 10 for all experiments) steps of state embeddings from the demonstration trajectory as input into the trajectory-conditioned policy.

Revised: Then we evaluate whether the agent has visited any of the next ∆ t state embeddings from the lastly visited state embedding in the demonstration.

Revised: It worth to note that the last visited state embedding must be included in the part of policy input.

Revised: Thus, we should only compare the agent’s current state embedding with the state embeddings in the demonstration segment provided into the policy for imitation reward.

Revised: Therefore, the proper value of ∆ t should be less than m =

Revised: and we consider ∆ t = 2 , 4 , 8 .

Revised: On the simple environments such as Apple-Gold and Deep Sea, we could see in Figure 20a and 20b the different values of ∆ t do not inﬂuence the performance a lot.

Revised: As shown in Figure 20c and 20d, on more challenging environments where the demonstration trajectory is longer and it’s harder to learn to imitate the demonstration, it’s better to set a larger value of ∆ t and we can provide imitation reward to the agent and encourage the imitation learning more generously.

Revised: In general, allowing the agent some ﬂexibility of imitation by setting ∆ t close to m works well.

Revised: In summary, ∆ t = 8 is a proper value for all of our primary experiment environments.

Revised: [Figure]

Revised: A v e r a g e

Revised: R e w a r d

Revised: [Figure]

Revised: [Figure]

Revised: K S TUDY OF THE S TOCHASTICITY IN THE E NVIRONMENT

Revised: As listed in Table 3, we consider stochastic environments for all the primary experiments, including the Apple-Gold domain, Montezuma’s Revenge, Pitfall, and Mujoco maze.

Revised: We set the initial state of the agent to be random.

Revised: On the Apple-Gold domain, the agent takes 3 random steps in the left bottom part of the maze before the episode starts.

Revised: On Montezuma’s Revenge, the mechanism of random initial no-ops is one of the standard ways to introduce stochasticity in the environment (Machado et al., 2017), as previous work (Mnih et al., 2015).

Revised: To show the difﬁculty in policy learning introduced by the stochastic environments, we tried to memorize and repeat the action sequence in the demonstration trajectory and check whether the agent could successfully visit the state of interest by following the demonstration.

Revised: In Figure 21c and 22c, it is clear that the success ratio in imitation is much lower than DTSIL, because the agent could not successfully follow the demonstration by just repeating the action sequence step by step.

Revised: Especially on Montezuma’s Revenge, with random initial no-ops, the state of the enemy and electricity beam when the agent starts moving is randomized.

Revised: Thus it could not successfully avoid death by just repeating the stored action sequence.

Revised: Obviously, when the environment is completely deterministic, repeating the action sequence could perfectly guide the agent to the state of interest.

Revised: However, with the stochasticity of the random initial state, just memorizing the action sequence is sufﬁcient to make the agent revisit novel regions.

Revised: Thus, the agent could not revisit the novel regions as efﬁciently as DTSIL to discover better trajectories and converge to a better total episode reward, as shown in Figure 21a, 21b,22a, 22b .

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: L R ANDOM E XPLORATION WITH E PSILON -G REEDY P OLICY

Revised: In this section, we investigate an additional baseline method, random exploration with epsilon-greedy policy, which is a traditional exploration method in reinforcement learning problems.

Revised: We consider combining epsilon-greedy policy with PPO or DQN framework and run experiments on Apple-Gold,

Revised: Deep Sea, and Montezuma’s Revenge.

Revised: Our implementation is based on OpenAI baselines (Dhariwal et al., 2017).

Revised: As shown in Figure 23, 24, 25, we searched different values of the hyper-parameters for the scheduling of the epsilon, though the performance is not better than DTSIL.

Revised: Especially, random exploration performs poorly on Montezuma’s Revenge.

Revised: The average score is less than 100, which is consistent with the experimental results from previous works (Mnih et al., 2015; Schulman et al., 2017).

Revised: [Figure]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]