Original: E FFICIENT E XPLORATION VIA S TATE M ARGINAL M ATCHING
Revised: E FFICIENT E XPLORATION VIA S TATE M ARGINAL M ATCHING

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: Reinforcement learning agents need to explore their unknown environments to solve the tasks given to them.
Revised: Reinforcement learning agents need to explore their unknown environments to solve the tasks given to them.

Original: The Bayes optimal solution to exploration is intractable for complex environments, and while several exploration methods have been proposed as approximations, it remains unclear what underlying objective is being optimized by existing exploration methods, or how they can be altered to incorporate prior knowledge about the task.
Revised: The Bayes optimal solution to exploration is intractable for complex environments, and while several exploration methods have been proposed as approximations, it remains unclear what underlying objective is being optimized by existing exploration methods, or how they can be altered to incorporate prior knowledge about the task.

Original: Moreover, it is unclear how to acquire a single exploration strategy that will be useful for solving multiple downstream tasks.
Revised: Moreover, it is unclear how to acquire a single exploration strategy that will be useful for solving multiple downstream tasks.

Original: We address these shortcomings by learning a single exploration policy that can quickly solve a suite of downstream tasks in a multi-task setting, amortizing the cost of learning to explore .
Revised: We address these shortcomings by learning a single exploration policy that can quickly solve a suite of downstream tasks in a multi-task setting, amortizing the cost of learning to explore .

Original: We recast exploration as a problem of State Marginal Matching (SMM), where we aim to learn a policy for which the state marginal distribution matches a given target state distribution, which can incorporate prior knowledge about the task.
Revised: We recast exploration as a problem of State Marginal Matching (SMM), where we aim to learn a policy for which the state marginal distribution matches a given target state distribution, which can incorporate prior knowledge about the task.

Original: We optimize the objective by reducing it to a two-player, zero-sum game between a state density model and a parametric policy.
Revised: We optimize the objective by reducing it to a two-player, zero-sum game between a state density model and a parametric policy.

Original: Our theoretical analysis of this approach suggests that prior exploration methods do not learn a policy that does distribution matching, but acquire a replay buffer that performs distribution matching, an observation that potentially explains prior methods’ success in single-task settings.
Revised: Our theoretical analysis of this approach suggests that prior exploration methods do not learn a policy that does distribution matching, but acquire a replay buffer that performs distribution matching, an observation that potentially explains prior methods’ success in single-task settings.

Original: On both simulated and real-world tasks, we demonstrate that our algorithm explores faster and adapts more quickly than prior methods.
Revised: On both simulated and real-world tasks, we demonstrate that our algorithm explores faster and adapts more quickly than prior methods.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Reinforcement learning (RL) algorithms must be equipped with exploration mechanisms to effectively solve tasks with limited reward signals.
Revised: Reinforcement learning (RL) algorithms must be equipped with exploration mechanisms to effectively solve tasks with limited reward signals.

Original: These tasks arise in many real-world applications where providing human supervision is expensive.
Revised: These tasks arise in many real-world applications where providing human supervision is expensive.

Original: The inability of current RL algorithms to adequately explore limits their applicability to long-horizon control tasks.
Revised: The inability of current RL algorithms to adequately explore limits their applicability to long-horizon control tasks.

Original: A wealth of prior work has studied exploration for RL.
Revised: A wealth of prior work has studied exploration for RL.

Original: While, in theory, the Bayes-optimal exploration strategy is optimal, it is intractable to compute exactly, motivating work on tractable heuristics for exploration.
Revised: While, in theory, the Bayes-optimal exploration strategy is optimal, it is intractable to compute exactly, motivating work on tractable heuristics for exploration.

Original: Exploration methods based on random actions have limited ability to cover a wide range of states.
Revised: Exploration methods based on random actions have limited ability to cover a wide range of states.

Original: More sophisticated techniques, such as intrinsic motivation, accelerate learning in the single-task setting.
Revised: More sophisticated techniques, such as intrinsic motivation, accelerate learning in the single-task setting.

Original: However, these methods have two limitations.
Revised: However, these methods have two limitations.

Original: First, they do not explicitly deﬁne an objective to quantify “good exploration,” but rather argue that exploration arises implicitly through some iterative procedure.
Revised: First, they do not explicitly deﬁne an objective to quantify “good exploration,” but rather argue that exploration arises implicitly through some iterative procedure.

Original: Lacking a well-deﬁned optimization objective, it remains challenging to understand what these methods are doing and why they work.
Revised: Lacking a well-deﬁned optimization objective, it remains challenging to understand what these methods are doing and why they work.

Original: Similarly, the lack of a metric to quantify exploration, even if only for evaluation, makes it challenging to compare exploration methods and assess progress in this area.
Revised: Similarly, the lack of a metric to quantify exploration, even if only for evaluation, makes it challenging to compare exploration methods and assess progress in this area.

Original: The second limitation is that these methods target the single-task setting.
Revised: The second limitation is that these methods target the single-task setting.

Original: Because these methods aim to converge to the optimal policy for a particular task, it is challenging to repurpose these methods to solve multiple tasks.
Revised: Because these methods aim to converge to the optimal policy for a particular task, it is challenging to repurpose these methods to solve multiple tasks.

Original: We address these shortcomings by considering a multi-task setting, where many different reward functions can be provided for the same set of states and dynamics.
Revised: We address these shortcomings by considering a multi-task setting, where many different reward functions can be provided for the same set of states and dynamics.

Original: Rather than exploring from scratch for each task, we aim to learn a single, task-agnostic exploration policy that can be adapted to many possible downstream reward functions, amortizing the cost of learning to explore.
Revised: Rather than exploring from scratch for each task, we aim to learn a single, task-agnostic exploration policy that can be adapted to many possible downstream reward functions, amortizing the cost of learning to explore.

Original: This exploration policy can be viewed as a prior on the policy for solving downstream tasks.
Revised: This exploration policy can be viewed as a prior on the policy for solving downstream tasks.

Original: Learning will consist of two phases: during training, we acquire this task-agnostic exploration policy; during testing, we use this exploration policy to quickly explore and maximize the task reward.
Revised: Learning will consist of two phases: during training, we acquire this task-agnostic exploration policy; during testing, we use this exploration policy to quickly explore and maximize the task reward.

Original: Learning a single exploration policy is considerably more difﬁcult than doing exploration throughout the course of learning a single task.
Revised: Learning a single exploration policy is considerably more difﬁcult than doing exploration throughout the course of learning a single task.

Original: The latter is done by intrinsic motivation (Pathak et al., 2017;
Revised: The latter is done by intrinsic motivation (Pathak et al., 2017;

Original: Tang et al., 2017; Oudeyer et al., 2007) and count-based exploration methods (Bellemare et al., 2016), which can effectively explore to ﬁnd states with high reward, at which point the agent can decrease exploration and increase exploitation of those high-reward states.
Revised: Tang et al., 2017; Oudeyer et al., 2007) and count-based exploration methods (Bellemare et al., 2016), which can effectively explore to ﬁnd states with high reward, at which point the agent can decrease exploration and increase exploitation of those high-reward states.

Original: While these methods perform efﬁcient exploration for learning a single task, the policy at any particular iteration is not a good exploration policy.
Revised: While these methods perform efﬁcient exploration for learning a single task, the policy at any particular iteration is not a good exploration policy.

Original: For example, the ﬁnal policy at convergence would only visit the high-reward states discovered for the current task.
Revised: For example, the ﬁnal policy at convergence would only visit the high-reward states discovered for the current task.

Original: What objective should be optimized to obtain a good exploration policy?
Revised: What objective should be optimized to obtain a good exploration policy?

Original: We recast exploration as a problem of State Marginal Matching: given a desired state distribution, we learn a mixture of policies for which the state marginal distribution matches this desired distribution.
Revised: We recast exploration as a problem of State Marginal Matching: given a desired state distribution, we learn a mixture of policies for which the state marginal distribution matches this desired distribution.

Original: Without any prior information, this objective reduces to maximizing the marginal state entropy H
Revised: Without any prior information, this objective reduces to maximizing the marginal state entropy H

Original: [ s ] , which encourages the policy to visit as many states as possible.
Revised: [ s ] , which encourages the policy to visit as many states as possible.

Original: The distribution matching objective also provides a convenient mechanism to incorporate prior knowledge about the task, whether in the form of safety constraints that the agent should obey; preferences for some states over other states; reward shaping; or the relative importance of each state dimension for a particular task.
Revised: The distribution matching objective also provides a convenient mechanism to incorporate prior knowledge about the task, whether in the form of safety constraints that the agent should obey; preferences for some states over other states; reward shaping; or the relative importance of each state dimension for a particular task.

Original: We also propose an algorithm to optimize the State Marginal Matching (SMM) objective.
Revised: We also propose an algorithm to optimize the State Marginal Matching (SMM) objective.

Original: First, we reduce the problem of SMM to a two-player, zero-sum game between a policy player and a density player.
Revised: First, we reduce the problem of SMM to a two-player, zero-sum game between a policy player and a density player.

Original: We ﬁnd a Nash Equilibrium for this game using ﬁctitious play (Brown, 1951), a classic procedure from game theory.
Revised: We ﬁnd a Nash Equilibrium for this game using ﬁctitious play (Brown, 1951), a classic procedure from game theory.

Original: Our resulting algorithm iteratively ﬁts a state density model and then updates the policy to visit states with low density under this model.
Revised: Our resulting algorithm iteratively ﬁts a state density model and then updates the policy to visit states with low density under this model.

Original: Our analysis of this approach sheds light on prior work on exploration.
Revised: Our analysis of this approach sheds light on prior work on exploration.

Original: In particular, while the policy learned by existing exploration algorithms does not perform distribution matching, the replay buffer does, an observation that potentially explains the success of prior methods.
Revised: In particular, while the policy learned by existing exploration algorithms does not perform distribution matching, the replay buffer does, an observation that potentially explains the success of prior methods.

Original: On both simulated and real-world tasks, we demonstrate that our algorithm explores more effectively and adapts more quickly to new tasks than state-of-the-art baselines.
Revised: On both simulated and real-world tasks, we demonstrate that our algorithm explores more effectively and adapts more quickly to new tasks than state-of-the-art baselines.

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

Original: Most prior work on exploration has looked at exploration bonuses and intrinsic motivation.
Revised: Most prior work on exploration has looked at exploration bonuses and intrinsic motivation.

Original: One class of exploration methods uses prediction error of some auxiliary task as an exploration bonus, which provides high (intrinsic) reward in states where the predictive model performs poorly (Pathak et al., 2017; Oudeyer et al., 2007; Schmidhuber, 1991;
Revised: One class of exploration methods uses prediction error of some auxiliary task as an exploration bonus, which provides high (intrinsic) reward in states where the predictive model performs poorly (Pathak et al., 2017; Oudeyer et al., 2007; Schmidhuber, 1991;

Original: Houthooft et al., 2016; Burda et al., 2018).
Revised: Houthooft et al., 2016; Burda et al., 2018).

Original: Another set of approaches (Tang et al., 2017; Bellemare et al., 2016; Schmidhuber, 2010) directly encourage the agent to visit novel states.
Revised: Another set of approaches (Tang et al., 2017; Bellemare et al., 2016; Schmidhuber, 2010) directly encourage the agent to visit novel states.

Original: While all methods effectively explore during the course of solving a single task (Taïga et al., 2019), the policy obtained at convergence is often not a good exploration policy (see Section 4).
Revised: While all methods effectively explore during the course of solving a single task (Taïga et al., 2019), the policy obtained at convergence is often not a good exploration policy (see Section 4).

Original: In contrast, our method converges to a highly-exploratory policy by maximizing state entropy in the training objective (Eq. 2).
Revised: In contrast, our method converges to a highly-exploratory policy by maximizing state entropy in the training objective (Eq. 2).

Original: Many exploration algorithms can be classiﬁed by whether they explore in the space of actions, policy parameters, goals, or states.
Revised: Many exploration algorithms can be classiﬁed by whether they explore in the space of actions, policy parameters, goals, or states.

Original: Common exploration strategies including (cid:15) -greedy and Ornstein–Uhlenbeck noise (Lillicrap et al., 2015), as well as standard MaxEnt RL algorithms (Ziebart, 2010; Haarnoja et al., 2018), explore in the action space.
Revised: Common exploration strategies including (cid:15) -greedy and Ornstein–Uhlenbeck noise (Lillicrap et al., 2015), as well as standard MaxEnt RL algorithms (Ziebart, 2010; Haarnoja et al., 2018), explore in the action space.

Original: Recent work (Fortunato et al., 2017;
Revised: Recent work (Fortunato et al., 2017;

Original: Plappert et al., 2017) shows that adding noise to the parameters of the policy can result in good exploration.
Revised: Plappert et al., 2017) shows that adding noise to the parameters of the policy can result in good exploration.

Original: Most closely related to our work are methods that perform exploration in the space of states or goals (Colas et al., 2018; Held et al., 2017; Nair et al., 2018;
Revised: Most closely related to our work are methods that perform exploration in the space of states or goals (Colas et al., 2018; Held et al., 2017; Nair et al., 2018;

Original: Pong et al., 2019; Hazan et al., 2018).
Revised: Pong et al., 2019; Hazan et al., 2018).

Original: In fact, Hazan et al.
Revised: In fact, Hazan et al.

Original: (2018) consider the same State Marginal Matching objective that we examine and propose a similar algorithm.
Revised: (2018) consider the same State Marginal Matching objective that we examine and propose a similar algorithm.

Original: In relation to Hazan et al.
Revised: In relation to Hazan et al.

Original: (2018), our main contributions are (1) empirically showing that exploration based on state-entropy is competitive with existing state-of-the-art exploration methods, and (2) explaining how existing exploration methods based on prediction error are implicitly maximizing this state-entropy objective.
Revised: (2018), our main contributions are (1) empirically showing that exploration based on state-entropy is competitive with existing state-of-the-art exploration methods, and (2) explaining how existing exploration methods based on prediction error are implicitly maximizing this state-entropy objective.

Original: In Appendix C.1, we also discuss how goal-conditioned RL (Kaelbling, 1993; Schaul et al., 2015) can be viewed as a special case of State Marginal Matching when the goal-sampling distribution is learned jointly with the policy.
Revised: In Appendix C.1, we also discuss how goal-conditioned RL (Kaelbling, 1993; Schaul et al., 2015) can be viewed as a special case of State Marginal Matching when the goal-sampling distribution is learned jointly with the policy.

Original: The problems of exploration and meta-reinforcement learning are tightly coupled.
Revised: The problems of exploration and meta-reinforcement learning are tightly coupled.

Original: Meta-reinforcement learning algorithms (Duan et al., 2016; Finn et al., 2017; Rakelly et al., 2019; Mishra et al., 2017) must perform effective exploration if they hope to solve a downstream task.
Revised: Meta-reinforcement learning algorithms (Duan et al., 2016; Finn et al., 2017; Rakelly et al., 2019; Mishra et al., 2017) must perform effective exploration if they hope to solve a downstream task.

Original: Some prior work has explicitly looked at the problem of learning to explore (Gupta et al., 2018; Xu et al., 2018).
Revised: Some prior work has explicitly looked at the problem of learning to explore (Gupta et al., 2018; Xu et al., 2018).

Original: Our problem statement is similar to meta-learning, in that we also aim to learn a policy as a prior for solving downstream tasks.
Revised: Our problem statement is similar to meta-learning, in that we also aim to learn a policy as a prior for solving downstream tasks.

Original: However, whereas meta-RL requires a distribution of task reward functions, our method will require only a single target state marginal distribution.
Revised: However, whereas meta-RL requires a distribution of task reward functions, our method will require only a single target state marginal distribution.

Original: Due to the simpler problem assumptions and training procedure, our method may be easier to apply in real-world domains.
Revised: Due to the simpler problem assumptions and training procedure, our method may be easier to apply in real-world domains.

Original: Related to our approach are standard maximum action entropy algorithms (Haarnoja et al., 2018;
Revised: Related to our approach are standard maximum action entropy algorithms (Haarnoja et al., 2018;

Original: Kappen et al., 2012;
Revised: Kappen et al., 2012;

Original: Rawlik et al., 2013; Ziebart et al., 2008; Theodorou & Todorov, 2012).
Revised: Rawlik et al., 2013; Ziebart et al., 2008; Theodorou & Todorov, 2012).

Original: While these algorithms are referred to as MaxEnt RL , they are maximizing entropy over actions, not states.
Revised: While these algorithms are referred to as MaxEnt RL , they are maximizing entropy over actions, not states.

Original: These algorithms can be viewed as performing inference on a graphical model where the likelihood of a trajectory is given by its exponentiated reward (Toussaint & Storkey, 2006; Levine, 2018;
Revised: These algorithms can be viewed as performing inference on a graphical model where the likelihood of a trajectory is given by its exponentiated reward (Toussaint & Storkey, 2006; Levine, 2018;

Original: Abdolmaleki et al., 2018).
Revised: Abdolmaleki et al., 2018).

Original: While distributions over trajectories induce distributions over states, computing the exact relationship requires integrating over all possible trajectories, an intractable problem for most MDPs.
Revised: While distributions over trajectories induce distributions over states, computing the exact relationship requires integrating over all possible trajectories, an intractable problem for most MDPs.

Original: A related but distinct class of relative entropy methods use a similar entropy-based objective to limit the size of policy updates (Peters et al., 2010; Schulman et al., 2015).
Revised: A related but distinct class of relative entropy methods use a similar entropy-based objective to limit the size of policy updates (Peters et al., 2010; Schulman et al., 2015).

Original: Finally, the idea of distribution matching has been employed successfully in imitation learning (Ziebart et al., 2008; Ho & Ermon, 2016;
Revised: Finally, the idea of distribution matching has been employed successfully in imitation learning (Ziebart et al., 2008; Ho & Ermon, 2016;

Original: Finn et al., 2016; Fu et al., 2017).
Revised: Finn et al., 2016; Fu et al., 2017).

Original: Similar to some inverse RL algorithms (Ho & Ermon, 2016; Fu et al., 2018), our method iterates between learning a policy and learning a reward function, though our reward function is obtained via a density model instead of a discriminator.
Revised: Similar to some inverse RL algorithms (Ho & Ermon, 2016; Fu et al., 2018), our method iterates between learning a policy and learning a reward function, though our reward function is obtained via a density model instead of a discriminator.

Original: While inverse RL algorithms assume access to expert trajectories, we instead assume access to the density of the target state marginal distribution.
Revised: While inverse RL algorithms assume access to expert trajectories, we instead assume access to the density of the target state marginal distribution.

Original: In many realistic settings, such as robotic control with many degrees of freedom, providing fully-speciﬁed trajectories may be much more challenging than deﬁning a target state marginal distribution.
Revised: In many realistic settings, such as robotic control with many degrees of freedom, providing fully-speciﬁed trajectories may be much more challenging than deﬁning a target state marginal distribution.

Original: The latter only requires some aggregate statistics about expert behavior, and does not even need to be realizable by any policy.
Revised: The latter only requires some aggregate statistics about expert behavior, and does not even need to be realizable by any policy.

Original: In summary, our work uniﬁes prior exploration methods as performing approximate distribution matching, and explains how state distribution matching can be performed properly.
Revised: In summary, our work uniﬁes prior exploration methods as performing approximate distribution matching, and explains how state distribution matching can be performed properly.

Original: This perspective provides a clearer picture of exploration, and this observation is useful particularly because many of the underlying ingredients, such as adversarial games and density estimation, have seen recent progress and therefore might be adopted to improve exploration methods.
Revised: This perspective provides a clearer picture of exploration, and this observation is useful particularly because many of the underlying ingredients, such as adversarial games and density estimation, have seen recent progress and therefore might be adopted to improve exploration methods.

Original:  S TATE M ARGINAL M ATCHING
Revised:  S TATE M ARGINAL M ATCHING

Original: In this section, we propose the State Marginal Matching problem as a principled objective for learning to explore , and offer an algorithm for optimizing it.
Revised: In this section, we propose the State Marginal Matching problem as a principled objective for learning to explore , and offer an algorithm for optimizing it.

Original: We consider a parametric policy π θ ∈ Π (cid:44) { π θ | θ ∈ Θ } that chooses actions a ∈ A in a Markov Decision Process (MDP) M with ﬁxed episode lengths T , dynamics distribution p ( s t +1 | s t , a t ) , and initial state distribution p 0 ( s ) .
Revised: We consider a parametric policy π θ ∈ Π (cid:44) { π θ | θ ∈ Θ } that chooses actions a ∈ A in a Markov Decision Process (MDP) M with ﬁxed episode lengths T , dynamics distribution p ( s t +1 | s t , a t ) , and initial state distribution p 0 ( s ) .

Original: The MDP M together with the policy π θ form an implicit generative model over states.
Revised: The MDP M together with the policy π θ form an implicit generative model over states.

Original: We deﬁne the state marginal distribution
Revised: We deﬁne the state marginal distribution

Original: ρ π ( s ) as the probability that the policy visits state s :
Revised: ρ π ( s ) as the probability that the policy visits state s :

Original: [Equation]
Revised: [Equation]

Original: We emphasize that ρ π ( s ) is not a distribution over trajectories, and is not the stationary distribution of the policy after inﬁnitely many steps, but rather the distribution over states visited in a ﬁnite-length episode.
Revised: We emphasize that ρ π ( s ) is not a distribution over trajectories, and is not the stationary distribution of the policy after inﬁnitely many steps, but rather the distribution over states visited in a ﬁnite-length episode.

Original: 2 We also note that any trajectory distribution matching problem can be reduced to a state marginal matching problem by augmenting the current state to include all previous states.
Revised: 2 We also note that any trajectory distribution matching problem can be reduced to a state marginal matching problem by augmenting the current state to include all previous states.

Original: We assume that we are given a target distribution p ∗ ( s ) over states s ∈ S that encodes our belief about the tasks we may be given at test-time.
Revised: We assume that we are given a target distribution p ∗ ( s ) over states s ∈ S that encodes our belief about the tasks we may be given at test-time.

Original: For example, a roboticist might assign small values of p ∗ ( s ) to states that are dangerous, regardless of the desired task.
Revised: For example, a roboticist might assign small values of p ∗ ( s ) to states that are dangerous, regardless of the desired task.

Original: Alternatively, we might also learn p ∗ ( s ) from data about human preferences (Christiano et al., 2017).
Revised: Alternatively, we might also learn p ∗ ( s ) from data about human preferences (Christiano et al., 2017).

Original: For goal-reaching tasks, we can analytically derive the optimal target distribution (Appendix C).
Revised: For goal-reaching tasks, we can analytically derive the optimal target distribution (Appendix C).

Original: Given p ∗ ( s ) , our goal is to ﬁnd a parametric policy that is “closest” to this target distribution, where we measure discrepancy using the
Revised: Given p ∗ ( s ) , our goal is to ﬁnd a parametric policy that is “closest” to this target distribution, where we measure discrepancy using the

Original: Kullback-Leibler (KL) divergence:
Revised: Kullback-Leibler (KL) divergence:

Original: [Equation]
Revised: [Equation]

Original: This is the same objective as in Hazan et al.
Revised: This is the same objective as in Hazan et al.

Original: Note that we use the reverse-KL (Bishop, 2006), which is mode-seeking (i.e., exploratory).
Revised: Note that we use the reverse-KL (Bishop, 2006), which is mode-seeking (i.e., exploratory).

Original: We show in Appendix C that the policies obtained via
Revised: We show in Appendix C that the policies obtained via

Original: State Marginal Matching provide an optimal exploration strategy for a particular distribution over reward functions.
Revised: State Marginal Matching provide an optimal exploration strategy for a particular distribution over reward functions.

Original: To gain intuition for the State Marginal Matching objective, we decomposed it in two ways.
Revised: To gain intuition for the State Marginal Matching objective, we decomposed it in two ways.

Original: In Equation 2, we see that State Marginal Matching is equivalent to maximizing the reward function r ( s ) (cid:44) log p ∗ ( s ) while simultaneously maximizing the entropy of states.
Revised: In Equation 2, we see that State Marginal Matching is equivalent to maximizing the reward function r ( s ) (cid:44) log p ∗ ( s ) while simultaneously maximizing the entropy of states.

Original: Note that, unlike traditional MaxEnt RL algorithms (Ziebart et al., 2008; Haarnoja et al., 2018), we regularize the entropy of the state distribution, not the conditional distribution of actions given states, which results in exploration in the space of states rather than in actions.
Revised: Note that, unlike traditional MaxEnt RL algorithms (Ziebart et al., 2008; Haarnoja et al., 2018), we regularize the entropy of the state distribution, not the conditional distribution of actions given states, which results in exploration in the space of states rather than in actions.

Original: Moreover, Equation 1 suggests that
Revised: Moreover, Equation 1 suggests that

Original: State Marginal Matching maximizes a pseudo-reward r ( s ) (cid:44) log p ∗ ( s ) − log ρ π ( s ) , which assigns positive utility to states that the agent visits too infrequently and negative utility to states visited too frequently (see Figure 1).
Revised: State Marginal Matching maximizes a pseudo-reward r ( s ) (cid:44) log p ∗ ( s ) − log ρ π ( s ) , which assigns positive utility to states that the agent visits too infrequently and negative utility to states visited too frequently (see Figure 1).

Original: We emphasize that maximizing this pseudo-reward is not a RL problem because the pseudo-reward depends on the policy.
Revised: We emphasize that maximizing this pseudo-reward is not a RL problem because the pseudo-reward depends on the policy.

Original: 3. O PTIMIZING THE S TATE M ARGINAL M ATCHING O BJECTIVE
Revised: 3. O PTIMIZING THE S TATE M ARGINAL M ATCHING O BJECTIVE

Original: Optimizing Equation 1 to obtain a single exploration policy is more challenging than standard RL because the reward function itself depends on the policy.
Revised: Optimizing Equation 1 to obtain a single exploration policy is more challenging than standard RL because the reward function itself depends on the policy.

Original: To break this cyclic dependency, we introduce a parametric state density model q ψ ( s ) ∈ Q (cid:44) { q ψ | ψ ∈ Ψ } to approximate the policy’s state marginal distribution, ρ π ( s ) .
Revised: To break this cyclic dependency, we introduce a parametric state density model q ψ ( s ) ∈ Q (cid:44) { q ψ | ψ ∈ Ψ } to approximate the policy’s state marginal distribution, ρ π ( s ) .

Original: We assume that the class of density models Q is sufﬁciently expressive to represent every policy:
Revised: We assume that the class of density models Q is sufﬁciently expressive to represent every policy:

Original: Assumption 1.
Revised: Assumption 1.

Original: For every policy π ∈ Π , there exists q ∈ Q such that D KL ( ρ π ( s ) (cid:107) q ( s ))
Revised: For every policy π ∈ Π , there exists q ∈ Q such that D KL ( ρ π ( s ) (cid:107) q ( s ))

Original: = 0 .
Revised: = 0 .

Original: Under this assumption, optimizing the policy w.r.t.
Revised: Under this assumption, optimizing the policy w.r.t.

Original: this approximate distribution q ( s ) will yield the same solution as Equation 1 (see Appendix A for the proof):
Revised: this approximate distribution q ( s ) will yield the same solution as Equation 1 (see Appendix A for the proof):

Original: Proposition 3.1.
Revised: Proposition 3.1.

Original: Let policies Π and density models Q satisfying Assumption 1 be given.
Revised: Let policies Π and density models Q satisfying Assumption 1 be given.

Original: For any target distribution p ∗ , the following optimization problems are equivalent:
Revised: For any target distribution p ∗ , the following optimization problems are equivalent:

Original: [Equation]
Revised: [Equation]

Original: Solving the new max-min optimization problem is equivalent to ﬁnding the Nash equilibrium of a two-player, zero-sum game: a policy player chooses the policy π while the density player chooses the density model q .
Revised: Solving the new max-min optimization problem is equivalent to ﬁnding the Nash equilibrium of a two-player, zero-sum game: a policy player chooses the policy π while the density player chooses the density model q .

Original: To avoid confusion, we use actions to refer to controls a ∈ A output by the policy π in the traditional RL problem and strategies to refer to the decisions π ∈ Π of the policy player and decisions q ∈ Q of the density player.
Revised: To avoid confusion, we use actions to refer to controls a ∈ A output by the policy π in the traditional RL problem and strategies to refer to the decisions π ∈ Π of the policy player and decisions q ∈ Q of the density player.

Original: The Nash existence theorem (Nash, 1951) proves that such a stationary point always exists for such a two-player, zero-sum game.
Revised: The Nash existence theorem (Nash, 1951) proves that such a stationary point always exists for such a two-player, zero-sum game.

Original: One common approach to saddle point games is to alternate between updating player A w.r.t.
Revised: One common approach to saddle point games is to alternate between updating player A w.r.t.

Original: player
Revised: player

Original: B, and updating player B w.r.t.
Revised: B, and updating player B w.r.t.

Original: player A.
Revised: player A.

Original: However, games such as Rock-Paper-Scissors illustrate that such a greedy approach is not guaranteed to converge to a stationary point.
Revised: However, games such as Rock-Paper-Scissors illustrate that such a greedy approach is not guaranteed to converge to a stationary point.

Original: A slight variant,
Revised: A slight variant,

Original: Input: Target distribution p ∗ ( s )
Revised: Input: Target distribution p ∗ ( s )

Original: Initialize policy π ( a | s ) , density model q ( s ) , and replay buffer B . while not converged do q ( m ) ← arg max q E s ∼B ( m − 1) [log q ( s )]
Revised: Initialize policy π ( a | s ) , density model q ( s ) , and replay buffer B . while not converged do q ( m ) ← arg max q E s ∼B ( m − 1) [log q ( s )]

Original: π ( m ) ← arg max π E s ∼
Revised: π ( m ) ← arg max π E s ∼

Original: ρ π ( s ) [ r ( s )] where r ( s ) (cid:44) log p ∗ ( s ) − log q ( m ) ( s )
Revised: ρ π ( s ) [ r ( s )] where r ( s ) (cid:44) log p ∗ ( s ) − log q ( m ) ( s )

Original: B ( m ) ← B ( m − 1) ∪ { ( s t , a t , s t +1 ) }
Revised: B ( m ) ← B ( m − 1) ∪ { ( s t , a t , s t +1 ) }

Original: Tt =1 with new transitions { ( s t , a t , s t +1 ) }
Revised: Tt =1 with new transitions { ( s t , a t , s t +1 ) }

Original: Tt =1 sampled from π ( m ) return historical policies { π (1) , · · · , π ( m ) } ﬁctitious play (Brown, 1951) does converge to a Nash equilibrium in ﬁnite time (Robinson, 1951;
Revised: Tt =1 sampled from π ( m ) return historical policies { π (1) , · · · , π ( m ) } ﬁctitious play (Brown, 1951) does converge to a Nash equilibrium in ﬁnite time (Robinson, 1951;

Original: Daskalakis & Pan, 2014).
Revised: Daskalakis & Pan, 2014).

Original: At each iteration, each player chooses their best strategy in response to the historical average of the opponent’s strategies.
Revised: At each iteration, each player chooses their best strategy in response to the historical average of the opponent’s strategies.

Original: In our setting, ﬁctitious play alternates between ﬁtting the density model to the historical average of policies (Equation 4), and updating the policy with RL to minimize the log-density of the state, using a historical average of the density models (Equation 5):
Revised: In our setting, ﬁctitious play alternates between ﬁtting the density model to the historical average of policies (Equation 4), and updating the policy with RL to minimize the log-density of the state, using a historical average of the density models (Equation 5):

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: s
Revised: s

Original: [Equation]
Revised: [Equation]

Original: ] where
Revised: ] where

Original: [Equation]
Revised: [Equation]

Original: Crucially, the exploration policy is not the last policy, π m +1 , but rather the historical average policy:
Revised: Crucially, the exploration policy is not the last policy, π m +1 , but rather the historical average policy:

Original: Deﬁnition 3.1.
Revised: Deﬁnition 3.1.

Original: A historical average policy ¯ π ( a | s ) , parametrized by a collection of policies π 1 , · · · , π m , is a policy that randomly samples one of the policy iterates
Revised: A historical average policy ¯ π ( a | s ) , parametrized by a collection of policies π 1 , · · · , π m , is a policy that randomly samples one of the policy iterates

Original: π i ∼ Unif [ π 1 , · · · , π m ] at the start of each episode and takes actions according to that policy for each step in the episode.
Revised: π i ∼ Unif [ π 1 , · · · , π m ] at the start of each episode and takes actions according to that policy for each step in the episode.

Original: A new policy is sampled for the next episode.
Revised: A new policy is sampled for the next episode.

Original: We summarize the resulting algorithm in Algorithm 1.
Revised: We summarize the resulting algorithm in Algorithm 1.

Original: In practice, we can efﬁciently implement
Revised: In practice, we can efﬁciently implement

Original: Equation 4 and avoid storing the policy parameters from every iteration by instead storing sampled states from each iteration.
Revised: Equation 4 and avoid storing the policy parameters from every iteration by instead storing sampled states from each iteration.

Original: 3 We cannot perform the same trick for Equation 5, and instead resort to approximating the historical average of density models with the most recent iterate.
Revised: 3 We cannot perform the same trick for Equation 5, and instead resort to approximating the historical average of density models with the most recent iterate.

Original: Algorithmlooks similar to prior exploration methods based on prediction-error, suggesting that we might use
Revised: Algorithmlooks similar to prior exploration methods based on prediction-error, suggesting that we might use

Original: SMM to understand how these prior methods work.
Revised: SMM to understand how these prior methods work.

Original:  W HY D OES P REDICTION -E RROR E XPLORATION W ORK ?
Revised:  W HY D OES P REDICTION -E RROR E XPLORATION W ORK ?

Original: Exploration methods based on prediction error (Burda et al., 2018; Stadie et al., 2015; Pathak et al., 2017; Schmidhuber, 1991; Chentanez et al., 2005) do not converge to an exploratory policy, even in the absence of extrinsic reward.
Revised: Exploration methods based on prediction error (Burda et al., 2018; Stadie et al., 2015; Pathak et al., 2017; Schmidhuber, 1991; Chentanez et al., 2005) do not converge to an exploratory policy, even in the absence of extrinsic reward.

Original: For example, consider the asymptotic behavior of ICM (Pathak et al., 2017) in a deterministic MDP, such as the Atari games where it was evaluated.
Revised: For example, consider the asymptotic behavior of ICM (Pathak et al., 2017) in a deterministic MDP, such as the Atari games where it was evaluated.

Original: At convergence, the predictive model will have zero error in all states, so the exploration bonus is zero – the ICM objective has no effect on the policy at convergence.
Revised: At convergence, the predictive model will have zero error in all states, so the exploration bonus is zero – the ICM objective has no effect on the policy at convergence.

Original: Similarly, consider the exploration bonus in
Revised: Similarly, consider the exploration bonus in

Original: Pseudocounts (Bellemare et al., 2016): 1 / ˆ n ( s ) , where ˆ n ( s ) is the (estimated) number of times that state s has been visited.
Revised: Pseudocounts (Bellemare et al., 2016): 1 / ˆ n ( s ) , where ˆ n ( s ) is the (estimated) number of times that state s has been visited.

Original: In the inﬁnite limit, each state has been visited inﬁnitely many times, so the Pseudocount exploration bonus also goes to zero — Pseudocounts has no effect at convergence.
Revised: In the inﬁnite limit, each state has been visited inﬁnitely many times, so the Pseudocount exploration bonus also goes to zero — Pseudocounts has no effect at convergence.

Original: Similar reasoning can be applied to other methods based on prediction error (Burda et al., 2018;
Revised: Similar reasoning can be applied to other methods based on prediction error (Burda et al., 2018;

Original: Stadie et al., 2015).
Revised: Stadie et al., 2015).

Original: More broadly, we can extend this analysis to stochastic MDPs, where we consider an abstract exploration algorithm that alternates between computing some intrinsic reward and performing RL (to convergence) on that intrinsic reward.
Revised: More broadly, we can extend this analysis to stochastic MDPs, where we consider an abstract exploration algorithm that alternates between computing some intrinsic reward and performing RL (to convergence) on that intrinsic reward.

Original: Existing prediction-error exploration
Revised: Existing prediction-error exploration

Original: [Figure]
Revised: [Figure]

Original: methods are all special cases.
Revised: methods are all special cases.

Original: At each iteration, the RL step solves a fully-observed MDP, which always admits a deterministic policy as a solution (Puterman, 2014).
Revised: At each iteration, the RL step solves a fully-observed MDP, which always admits a deterministic policy as a solution (Puterman, 2014).

Original: Thus, any exploration algorithm in this class cannot converge to a single, exploratory policy.
Revised: Thus, any exploration algorithm in this class cannot converge to a single, exploratory policy.

Original: Despite these observations, prior methods do excel at solving hard exploration tasks.
Revised: Despite these observations, prior methods do excel at solving hard exploration tasks.

Original: We draw an analogy to ﬁctitious play to explain their success.
Revised: We draw an analogy to ﬁctitious play to explain their success.

Original: While these methods never acquire an exploratory policy, over the course of training they will eventually visit all states.
Revised: While these methods never acquire an exploratory policy, over the course of training they will eventually visit all states.

Original: In other words, the historical average over policies will visit a wide range of states.
Revised: In other words, the historical average over policies will visit a wide range of states.

Original: Since the replay buffer exactly corresponds to this historical average over states, these methods will obtain a replay buffer with a diverse range of experience, possibly explaining why they succeed at solving hard exploration tasks.
Revised: Since the replay buffer exactly corresponds to this historical average over states, these methods will obtain a replay buffer with a diverse range of experience, possibly explaining why they succeed at solving hard exploration tasks.

Original: Moreover, this analysis suggests a surprisingly simple method for obtaining an exploration from these prior methods: use a mixture of the policy iterates throughout training.
Revised: Moreover, this analysis suggests a surprisingly simple method for obtaining an exploration from these prior methods: use a mixture of the policy iterates throughout training.

Original: The following section will not only compare
Revised: The following section will not only compare

Original: SMM against prior exploration methods, but also show that this historical averaging trick can be used to improve existing exploration methods.
Revised: SMM against prior exploration methods, but also show that this historical averaging trick can be used to improve existing exploration methods.

Original:  S IMULATED E XPERIMENTS
Revised:  S IMULATED E XPERIMENTS

Original: We used simulated control tasks to determine if SMM learns an exploratory policy, to compare SMM to prior exploration methods, and to study the effect of historical averaging.
Revised: We used simulated control tasks to determine if SMM learns an exploratory policy, to compare SMM to prior exploration methods, and to study the effect of historical averaging.

Original: More details can be found in Appendix D, and code will be released upon publication.
Revised: More details can be found in Appendix D, and code will be released upon publication.

Original: Baselines and Implementation Details : We compare to a state-of-the-art off-policy MaxEnt RL algorithm, Soft Actor-Critic (SAC) (Haarnoja et al., 2018); an inverse RL algorithm, Generative
Revised: Baselines and Implementation Details : We compare to a state-of-the-art off-policy MaxEnt RL algorithm, Soft Actor-Critic (SAC) (Haarnoja et al., 2018); an inverse RL algorithm, Generative

Original: Adversarial Imitation Learning (GAIL) (Ho & Ermon, 2016); and three exploration methods:
Revised: Adversarial Imitation Learning (GAIL) (Ho & Ermon, 2016); and three exploration methods:

Original: • Count-based Exploration (C), which discretizes states and uses − log
Revised: • Count-based Exploration (C), which discretizes states and uses − log

Original: ˆ π ( s )
Revised: ˆ π ( s )

Original: as an exploration bonus.
Revised: as an exploration bonus.

Original: • Pseudo-counts (PC) (Bellemare et al., 2016), which uses the recoding probability as a bonus.
Revised: • Pseudo-counts (PC) (Bellemare et al., 2016), which uses the recoding probability as a bonus.

Original: • Intrinsic Curiosity Module (ICM) (Pathak et al., 2017), which uses prediction error as a bonus.
Revised: • Intrinsic Curiosity Module (ICM) (Pathak et al., 2017), which uses prediction error as a bonus.

Original: We used SAC as the base RL algorithm for all exploration methods (SMM, C, PC, ICM).
Revised: We used SAC as the base RL algorithm for all exploration methods (SMM, C, PC, ICM).

Original: To implement SMM, we deﬁne the target distribution in terms of the extrinsic environment reward: p ∗ ( s ) ∝ exp( r env ( s )) .
Revised: To implement SMM, we deﬁne the target distribution in terms of the extrinsic environment reward: p ∗ ( s ) ∝ exp( r env ( s )) .

Original: We use a variational autoencoder (VAE) to model the density q ( s ) for both SMM and Pseudocounts (PC).
Revised: We use a variational autoencoder (VAE) to model the density q ( s ) for both SMM and Pseudocounts (PC).

Original: For the GAIL baseline, we generated synthetic expert data by sampling expert states from the target distribution p ∗ ( s ) (see Appendix D.2 for details).
Revised: For the GAIL baseline, we generated synthetic expert data by sampling expert states from the target distribution p ∗ ( s ) (see Appendix D.2 for details).

Original: Results for all experiments are averaged over 4-5 random seeds.
Revised: Results for all experiments are averaged over 4-5 random seeds.

Original: We start with a sanity check: Is exploration in state space (as done by SMM) better than exploration in action space (as done by MaxEnt RL, e.g., SAC)?
Revised: We start with a sanity check: Is exploration in state space (as done by SMM) better than exploration in action space (as done by MaxEnt RL, e.g., SAC)?

Original: To study this question, we implemented a 2D
Revised: To study this question, we implemented a 2D

Original: Navigation environment, shown in Figure 2a.
Revised: Navigation environment, shown in Figure 2a.

Original: To evaluate each method, we counted the number of hallways that the agent fully explored (i.e., reached the end) during training.
Revised: To evaluate each method, we counted the number of hallways that the agent fully explored (i.e., reached the end) during training.

Original: Figure 2b shows the state visitations for the three hallway environment, illustrating that SAC only explores one hallway, whereas SMM explores all three.
Revised: Figure 2b shows the state visitations for the three hallway environment, illustrating that SAC only explores one hallway, whereas SMM explores all three.

Original: Figure 2c also shows that SMM consistently explores 60% of hallways, whereas SAC rarely visits more than 20% of hallways.
Revised: Figure 2c also shows that SMM consistently explores 60% of hallways, whereas SAC rarely visits more than 20% of hallways.

Original: The remaining simulated experiments used the Manipulation environment (Plappert et al., 2018), shown in Figure 3a.
Revised: The remaining simulated experiments used the Manipulation environment (Plappert et al., 2018), shown in Figure 3a.

Original: Our ﬁrst experiment evaluates whether the exploration policy acquired by
Revised: Our ﬁrst experiment evaluates whether the exploration policy acquired by

Original: SMM allows us to solve downstream tasks more quickly.
Revised: SMM allows us to solve downstream tasks more quickly.

Original: We deﬁned the target distribution to be uniform over the entire state space (joint + block conﬁguration), with the constraints that we put
Revised: We deﬁned the target distribution to be uniform over the entire state space (joint + block conﬁguration), with the constraints that we put

Original: [Figure]
Revised: [Figure]

Original: [Figure] low probability mass on states where the block has fallen off the table; that actions should be small; and that the arm should be close to the object.
Revised: (d) low probability mass on states where the block has fallen off the table; that actions should be small; and that the arm should be close to the object.

Original: As shown in Figure 3b, SMM adapts substantially more quickly than other exploration methods, achieving a success rate 20% higher than the next best method, and reaching the same level of performance of the next baseline (ICM) in 4x fewer episodes.
Revised: As shown in Figure 3b, SMM adapts substantially more quickly than other exploration methods, achieving a success rate 20% higher than the next best method, and reaching the same level of performance of the next baseline (ICM) in 4x fewer episodes.

Original: SMM without historical averaging attains similar performance as the next best baseline (ICM), suggesting that historical averaging is the key ingredient, while the particular choice of prediction error or VAE is less important.
Revised: SMM without historical averaging attains similar performance as the next best baseline (ICM), suggesting that historical averaging is the key ingredient, while the particular choice of prediction error or VAE is less important.

Original: We provide further ablation studies of SMM in Appendix B.2.
Revised: We provide further ablation studies of SMM in Appendix B.2.

Original: While historical averaging is necessary to guarantee convergence (§ 3.1), most prior exploration methods do not employ historical averaging, raising the question of whether it is necessary in practice.
Revised: While historical averaging is necessary to guarantee convergence (§ 3.1), most prior exploration methods do not employ historical averaging, raising the question of whether it is necessary in practice.

Original: To answer this question, we compare SMM to three exploration methods.
Revised: To answer this question, we compare SMM to three exploration methods.

Original: In Figure 3c, we compare the policy obtained at convergence with the historical average of policy iterates over training for each method.
Revised: In Figure 3c, we compare the policy obtained at convergence with the historical average of policy iterates over training for each method.

Original: We measure how well each method explores by computing the marginal state entropy, which we compute by discretizing the state space.
Revised: We measure how well each method explores by computing the marginal state entropy, which we compute by discretizing the state space.

Original: The results show that SMM maximizes state entropy at least as effectively as prior methods, if not better.
Revised: The results show that SMM maximizes state entropy at least as effectively as prior methods, if not better.

Original: While this comparison is somewhat unfair, as we measure exploration using the objective that SMM maximizes, none of the methods we compare against propose metrics for exploration that we could use instead.
Revised: While this comparison is somewhat unfair, as we measure exploration using the objective that SMM maximizes, none of the methods we compare against propose metrics for exploration that we could use instead.

Original: Furthermore, we see that historical averaging is not only beneﬁcial to SMM, but also improves the exploration of prior methods.
Revised: Furthermore, we see that historical averaging is not only beneﬁcial to SMM, but also improves the exploration of prior methods.

Original: In our ﬁnal simulated experiment, we check whether prior knowledge injected via the target distribution is reﬂected in the policy obtained from State Marginal Matching.
Revised: In our ﬁnal simulated experiment, we check whether prior knowledge injected via the target distribution is reﬂected in the policy obtained from State Marginal Matching.

Original: Using the same Manipulation environment as above, we modiﬁed the target distribution to assign larger probability to states where the block was on the left half of the table than on the right half.
Revised: Using the same Manipulation environment as above, we modiﬁed the target distribution to assign larger probability to states where the block was on the left half of the table than on the right half.

Original: In Figure 3d, we measure whether
Revised: In Figure 3d, we measure whether

Original: SMM is able to achieve the target distribution by measuring the discrepancy between the block’s horizontal coordinate and the target distribution.
Revised: SMM is able to achieve the target distribution by measuring the discrepancy between the block’s horizontal coordinate and the target distribution.

Original: Compared to the SAC baseline, SMM and the
Revised: Compared to the SAC baseline, SMM and the

Original: Count baseline are half the distance to the target distribution.
Revised: Count baseline are half the distance to the target distribution.

Original: No method achieves zero discrepancy, suggesting that future methods could be better at matching state marginals.
Revised: No method achieves zero discrepancy, suggesting that future methods could be better at matching state marginals.

Original:  R EAL -W ORLD E XPERIMENTS
Revised:  R EAL -W ORLD E XPERIMENTS

Original: While almost all research on exploration focus on simulated domains, attributes of the real world such as partial observability, nonstationarity, and stochasticity may make the exploration more challenging.
Revised: While almost all research on exploration focus on simulated domains, attributes of the real world such as partial observability, nonstationarity, and stochasticity may make the exploration more challenging.

Original: The aim of this section is to see if SMM explores effectively on a real-world robotic control task.
Revised: The aim of this section is to see if SMM explores effectively on a real-world robotic control task.

Original: We used the D’Claw (Ahn et al., 2019) robotic manipulator, which is a 3-ﬁngered hand positioned vertically above a handle that it can turn.
Revised: We used the D’Claw (Ahn et al., 2019) robotic manipulator, which is a 3-ﬁngered hand positioned vertically above a handle that it can turn.

Original: For all experiments on the D’Claw , we used a target distribution that places uniform mass over all object angles [ − 180 ◦ , 180 ◦ ] .
Revised: For all experiments on the D’Claw , we used a target distribution that places uniform mass over all object angles [ − 180 ◦ , 180 ◦ ] .

Original: [Figure]
Revised: [Figure]

Original: In a ﬁrst experiment, we trained SMM and other baselines in simulation, and then evaluated the acquired exploration policy on the real robot using two metrics: the total number of rotations (in either direction), and the maximum radians turned (in both directions).
Revised: In a ﬁrst experiment, we trained SMM and other baselines in simulation, and then evaluated the acquired exploration policy on the real robot using two metrics: the total number of rotations (in either direction), and the maximum radians turned (in both directions).

Original: For each method, we computed the average metric across 100 evaluation episodes.
Revised: For each method, we computed the average metric across 100 evaluation episodes.

Original: We repeated this process for 5 independent training runs.
Revised: We repeated this process for 5 independent training runs.

Original: Figure 4b shows that SMM turns the knob more than the baselines, and it turns the knob to a wider range of angles.
Revised: Figure 4b shows that SMM turns the knob more than the baselines, and it turns the knob to a wider range of angles.

Original: To test for statistical signiﬁcance, we used a 1-sided Student’s t-test to test the hypothesis that SMM turned the knob more and to a wider range of angles than SAC.
Revised: To test for statistical signiﬁcance, we used a 1-sided Student’s t-test to test the hypothesis that SMM turned the knob more and to a wider range of angles than SAC.

Original: The p-values were all less than 0.05: p = 0 .
Revised: The p-values were all less than 0.05: p = 0 .

Original: 046 for number of rotations, p = 0 .
Revised: 046 for number of rotations, p = 0 .

Original: 019 for maximum clockwise angle, and p = 0 .
Revised: 019 for maximum clockwise angle, and p = 0 .

Original: 001 for maximum counter-clockwise angle.
Revised: 001 for maximum counter-clockwise angle.

Original: In our second experiment, we investigated whether it was possible to learn an exploration policy directly in the real world, without the need for a simulator.
Revised: In our second experiment, we investigated whether it was possible to learn an exploration policy directly in the real world, without the need for a simulator.

Original: Learning to explore in the real world is quite important, as building faithful simulators of complex systems is challenging.
Revised: Learning to explore in the real world is quite important, as building faithful simulators of complex systems is challenging.

Original: The physical constraints of the real robot make data efﬁciency paramount, suggesting that learning to explore will require an effective exploration strategy.
Revised: The physical constraints of the real robot make data efﬁciency paramount, suggesting that learning to explore will require an effective exploration strategy.

Original: In Figure 4c, we plot the range of angles that the policy explores throughout training.
Revised: In Figure 4c, we plot the range of angles that the policy explores throughout training.

Original: Not only does SMM explore a wider range of angles than SAC, but its ability to explore increases throughout training, suggesting that the SMM objective is correlated with real-world metrics of exploration.
Revised: Not only does SMM explore a wider range of angles than SAC, but its ability to explore increases throughout training, suggesting that the SMM objective is correlated with real-world metrics of exploration.

Original: In summary, the results in this section suggests that exploration techniques may actually be useful in the real world, which may encourage future work to study exploration methods on real-world tasks.
Revised: In summary, the results in this section suggests that exploration techniques may actually be useful in the real world, which may encourage future work to study exploration methods on real-world tasks.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this paper, we introduced a formal objective for exploration.
Revised: In this paper, we introduced a formal objective for exploration.

Original: While it is often unclear what existing exploration algorithms will converge to, our State Marginal Matching objective has a clear solution: at convergence, the policy should visit states in proportion to their density under a target distribution.
Revised: While it is often unclear what existing exploration algorithms will converge to, our State Marginal Matching objective has a clear solution: at convergence, the policy should visit states in proportion to their density under a target distribution.

Original: Not only does this objective encourage exploration, it also provides human users with a ﬂexible mechanism to bias exploration towards states they prefer and away from dangerous states.
Revised: Not only does this objective encourage exploration, it also provides human users with a ﬂexible mechanism to bias exploration towards states they prefer and away from dangerous states.

Original: Upon convergence, the resulting policy can thereafter be used as a prior in a multi-task setting, amortizing exploration and enabling faster adaptation to new, potentially sparse, reward functions.
Revised: Upon convergence, the resulting policy can thereafter be used as a prior in a multi-task setting, amortizing exploration and enabling faster adaptation to new, potentially sparse, reward functions.

Original: The algorithm we proposed looks quite similar to previous exploration methods based on prediction error, suggesting that those methods are also performing some form of distribution matching.
Revised: The algorithm we proposed looks quite similar to previous exploration methods based on prediction error, suggesting that those methods are also performing some form of distribution matching.

Original: However, by deriving our method from ﬁrst principles, we note that these prior methods omit a crucial historical averaging step, without which the algorithm is not guaranteed to converge.
Revised: However, by deriving our method from ﬁrst principles, we note that these prior methods omit a crucial historical averaging step, without which the algorithm is not guaranteed to converge.

Original: Experiments on both simulated and real-world tasks demonstrated how SMM learns to explore, enabling an agent to efﬁciently explore in new tasks provided at test time.
Revised: Experiments on both simulated and real-world tasks demonstrated how SMM learns to explore, enabling an agent to efﬁciently explore in new tasks provided at test time.

Original: In future work, we aim to study connections between inverse RL, MaxEnt RL and state marginal matching, all of which perform some form of distribution matching.
Revised: In future work, we aim to study connections between inverse RL, MaxEnt RL and state marginal matching, all of which perform some form of distribution matching.

Original: Empirically, we aim to scale to more complex tasks by parallelizing the training of all mixture components simultaneously.
Revised: Empirically, we aim to scale to more complex tasks by parallelizing the training of all mixture components simultaneously.

Original: Broadly, we expect the state distribution matching problem formulation to enable the development of more effective and principled RL methods that reason about distributions rather than individual states.
Revised: Broadly, we expect the state distribution matching problem formulation to enable the development of more effective and principled RL methods that reason about distributions rather than individual states.

Original: A P ROOFS
Revised: A P ROOFS

Original: Proof of Proposition 3.1.
Revised: Proof of Proposition 3.1.

Original: Note that
Revised: Note that

Original: [Equation]
Revised: [Equation]

Original: By Assumption 1, D KL ( ρ π ( s ) (cid:107) q ( s ))
Revised: By Assumption 1, D KL ( ρ π ( s ) (cid:107) q ( s ))

Original: = 

Revised: =

Original: for some q ∈ Q , so we obtain the desired result: max π
Revised: for some q ∈ Q , so we obtain the desired result: max π

Original: [Equation]
Revised: [Equation]

Original: B B ETTER M ARGINAL M ATCHING WITH M IXTURES OF M IXTURES
Revised: B B ETTER M ARGINAL M ATCHING WITH M IXTURES OF M IXTURES

Original: This section introduces an extension of SMM, SM4, that incorporates mixure modelling.
Revised: This section introduces an extension of SMM, SM4, that incorporates mixure modelling.

Original: Given the challenging problem of exploration in large state spaces, it is natural to wonder whether we can accelerate exploration by automatically decomposing the potentially-multimodal target distribution into a mixture of “easier-to-learn” distributions and learn a corresponding set of policies to do distribution matching for each component.
Revised: Given the challenging problem of exploration in large state spaces, it is natural to wonder whether we can accelerate exploration by automatically decomposing the potentially-multimodal target distribution into a mixture of “easier-to-learn” distributions and learn a corresponding set of policies to do distribution matching for each component.

Original: Note that the mixture model we introduce here is orthogonal to the historical averaging step discussed before.
Revised: Note that the mixture model we introduce here is orthogonal to the historical averaging step discussed before.

Original: Using ρ π z ( s ) to denote the state distribution of the policy conditioned on the latent variable z ∈ Z , the state marginal distribution of the mixture of policies is
Revised: Using ρ π z ( s ) to denote the state distribution of the policy conditioned on the latent variable z ∈ Z , the state marginal distribution of the mixture of policies is

Original: [Equation]
Revised: [Equation]

Original: (7) where p ( z ) is a latent prior.
Revised: (7) where p ( z ) is a latent prior.

Original: As before, we will minimize the KL divergence between this mixture distribution and the target distribution.
Revised: As before, we will minimize the KL divergence between this mixture distribution and the target distribution.

Original: Using Bayes’ rule to re-write ρ π ( s ) in terms of conditional probabilities, we obtain the following optimization problem:
Revised: Using Bayes’ rule to re-write ρ π ( s ) in terms of conditional probabilities, we obtain the following optimization problem:

Original: [Equation]
Revised: [Equation]

Original: Intuitively, this says that the agent should go to states (a) with high density under the target state distribution, (b) where this agent has not been before, and (c) where this agent is clearly distinguishable from the other agents.
Revised: Intuitively, this says that the agent should go to states (a) with high density under the target state distribution, (b) where this agent has not been before, and (c) where this agent is clearly distinguishable from the other agents.

Original: The last term (d) says to explore in the space of mixture components z .
Revised: The last term (d) says to explore in the space of mixture components z .

Original: This decomposition bears a resemblance to the mutual-information objectives in recent work (Achiam et al., 2018; Eysenbach et al., 2018; Co-Reyes et al., 2018).
Revised: This decomposition bears a resemblance to the mutual-information objectives in recent work (Achiam et al., 2018; Eysenbach et al., 2018; Co-Reyes et al., 2018).

Original: Thus, one interpretation of our work is as explaining that mutual information objectives almost perform distribution matching.
Revised: Thus, one interpretation of our work is as explaining that mutual information objectives almost perform distribution matching.

Original: The caveat is that prior work omits the state entropy term − log ρ π z ( s ) which provides high reward for visiting novel states, possibly explaining why these previous works have failed to scale to complex tasks.
Revised: The caveat is that prior work omits the state entropy term − log ρ π z ( s ) which provides high reward for visiting novel states, possibly explaining why these previous works have failed to scale to complex tasks.

Original: B. A LGORITHMIC S UMMARY
Revised: B. A LGORITHMIC S UMMARY

Original: We summarize the resulting procedure in Algorithm 2, which we refer to as SM4 (State Marginal
Revised: We summarize the resulting procedure in Algorithm 2, which we refer to as SM4 (State Marginal

Original: Matching with Mixtures of Mixtures).
Revised: Matching with Mixtures of Mixtures).

Original: The algorithm (1) ﬁts a density model q ( m ) z ( s ) to approximate the state marginal distribution for each policy π z ; (2) learns a discriminator d ( m ) ( z | s ) to predict which policy π z will visit state s ; and (3) uses RL to update each policy π z to maximize the expected return of its corresponding reward function r z ( s ) (cid:44) log p ∗ ( s ) − log
Revised: The algorithm (1) ﬁts a density model q ( m ) z ( s ) to approximate the state marginal distribution for each policy π z ; (2) learns a discriminator d ( m ) ( z | s ) to predict which policy π z will visit state s ; and (3) uses RL to update each policy π z to maximize the expected return of its corresponding reward function r z ( s ) (cid:44) log p ∗ ( s ) − log

Original: ρ π z ( s ) + log p ( z | s ) − log p
Revised: ρ π z ( s ) + log p ( z | s ) − log p

Original: [Equation]
Revised: [Equation]

Original: ) derived in Equation 8.
Revised: ) derived in Equation 8.

Original: The only difference from Algorithm 1 is that we learn a discriminator d ( z | s ) , in addition to updating the density models q z ( s ) and the policies π z ( a | s ) .
Revised: The only difference from Algorithm 1 is that we learn a discriminator d ( z | s ) , in addition to updating the density models q z ( s ) and the policies π z ( a | s ) .

Original: Jensen’s inequality tells us that maximizing the log-density of the learned discriminator will maximize a lower bound on the true density (see Agakov (2004)):
Revised: Jensen’s inequality tells us that maximizing the log-density of the learned discriminator will maximize a lower bound on the true density (see Agakov (2004)):

Original: [Equation]
Revised: [Equation]

Original: Input: Target distribution p ∗ ( s )
Revised: Input: Target distribution p ∗ ( s )

Original: Initialize policy π z ( a | s ) , density model q z ( s ) , discriminator d ( z | s ) , and replay buffer B . while not converged do for z = 1 , · · · , n do (cid:46) (1) Update density model for each policy π z .
Revised: Initialize policy π z ( a | s ) , density model q z ( s ) , discriminator d ( z | s ) , and replay buffer B . while not converged do for z = 1 , · · · , n do (cid:46) (1) Update density model for each policy π z .

Original: q ( m ) z ← arg max q E { s | ( z (cid:48) ,s )
Revised: q ( m ) z ← arg max q E { s | ( z (cid:48) ,s )

Original: ∼B ( m − 1) ,z (cid:48)
Revised: ∼B ( m − 1) ,z (cid:48)

Original: =
Revised: =

Original: z } [log q ( s )]
Revised: z } [log q ( s )]

Original: d ( m ) ← arg max d E ( z,s ) ∼B ( m − 1) [log d ( z | s )]
Revised: d ( m ) ← arg max d E ( z,s ) ∼B ( m − 1) [log d ( z | s )]

Original: (cid:46) (2) Update discriminator.
Revised: (cid:46) (2) Update discriminator.

Original: for z = 1 , · · · , n do r ( m )
Revised: for z = 1 , · · · , n do r ( m )

Original: z ( s ) (cid:44) log p ∗ ( s ) − log q ( m )
Revised: z ( s ) (cid:44) log p ∗ ( s ) − log q ( m )

Original: z ( s ) + log d ( m ) ( z | s ) − log p ( z ) π ( m )
Revised: z ( s ) + log d ( m ) ( z | s ) − log p ( z ) π ( m )

Original: z ← arg max π E ρ π ( s ) (cid:104) r
Revised: z ← arg max π E ρ π ( s ) (cid:104) r

Original: ( m ) z ( s ) (cid:105) (cid:46) (3) Update each policy π z .
Revised: ( m ) z ( s ) (cid:105) (cid:46) (3) Update each policy π z .

Original: Sample latent skill z ( m ) ∼ p
Revised: Sample latent skill z ( m ) ∼ p

Original: ( z )
Revised: ( z )

Original: Sample transitions { ( s t , a t , s t +1 ) }
Revised: Sample transitions { ( s t , a t , s t +1 ) }

Original: Tt =1 with π ( m ) z ( a | s )
Revised: Tt =1 with π ( m ) z ( a | s )

Original: B ( m ) ← B ( m − 1) ∪ { ( z ( m ) , s t , a t , s t +1 ) }
Revised: B ( m ) ← B ( m − 1) ∪ { ( z ( m ) , s t , a t , s t +1 ) }

Original: Tt =
Revised: Tt =

Original: return {{ π (1)1 , · · · , π (1) n } , · · · , { π ( m ), · · · , π ( m ) n }} z
Revised: return {{ π (1)1 , · · · , π (1) n } , · · · , { π ( m ), · · · , π ( m ) n }} z

Original: ) to predict will
Revised: ) to predict will

Original: In our implementation, the density model q z ( s ) is a VAE that inputs the concatenated vector { s, z } of the state s and the latent skill z used to obtain this sample s ; and the discriminator is a feedforward
Revised: In our implementation, the density model q z ( s ) is a VAE that inputs the concatenated vector { s, z } of the state s and the latent skill z used to obtain this sample s ; and the discriminator is a feedforward

Original: MLP.
Revised: MLP.

Original: The algorithm returns the historical average of mixtures of policies (a total of n · m policies).
Revised: The algorithm returns the historical average of mixtures of policies (a total of n · m policies).

Original: Our implementation uses a uniform categorical distribution for the prior p ( z ) , and does not implement the update for p ( z ) .
Revised: Our implementation uses a uniform categorical distribution for the prior p ( z ) , and does not implement the update for p ( z ) .

Original: Note that the updates for each z can be conducted in parallel.
Revised: Note that the updates for each z can be conducted in parallel.

Original: A distributed implementation would emulate broadcast-collect algorithms (Lynch, 1996), with each worker updating the policy independently, and periodically aggregating results to update the discriminator d ( z | s ) .
Revised: A distributed implementation would emulate broadcast-collect algorithms (Lynch, 1996), with each worker updating the policy independently, and periodically aggregating results to update the discriminator d ( z | s ) .

Original: Such a distributed implementation has the appealing property that each compute node would explore a different part of the state space.
Revised: Such a distributed implementation has the appealing property that each compute node would explore a different part of the state space.

Original: While there has been some work on multi-agent coordinated exploration (Parisotto et al., 2019) and concurrent exploration (Dimakopoulou & Van Roy, 2018), it remains a fairly unexplored area (pun intended) and we believe that SMM with Mixtures of Mixtures offers a simple approach to this problem.
Revised: While there has been some work on multi-agent coordinated exploration (Parisotto et al., 2019) and concurrent exploration (Dimakopoulou & Van Roy, 2018), it remains a fairly unexplored area (pun intended) and we believe that SMM with Mixtures of Mixtures offers a simple approach to this problem.

Original: B. A BLATION S TUDY
Revised: B. A BLATION S TUDY

Original: To understand the relative contribution of each component in the SM4 objective (Equation 8), we compare SM4 to baselines that lack conditional state entropy
Revised: To understand the relative contribution of each component in the SM4 objective (Equation 8), we compare SM4 to baselines that lack conditional state entropy

Original: H π z [ s ] = − log
Revised: H π z [ s ] = − log

Original: ρ π z ( s ) , latent conditional action entropy log p ( z | s ) , or both (i.e, SAC).
Revised: ρ π z ( s ) , latent conditional action entropy log p ( z | s ) , or both (i.e, SAC).

Original: In Figure 5a, we plot the training time performance on the Navigation task with 3 halls of length 50.
Revised: In Figure 5a, we plot the training time performance on the Navigation task with 3 halls of length 50.

Original: We see that SM4 relies heavily on both key differences from SAC.
Revised: We see that SM4 relies heavily on both key differences from SAC.

Original: In Figure 5b, we study the effect of mixture modelling on test-time exploration in the Manipulation environment.
Revised: In Figure 5b, we study the effect of mixture modelling on test-time exploration in the Manipulation environment.

Original: After running SMM/SM4 with a uniform distribution, we count the number of episodes required to ﬁnd an (unknown) goal state.
Revised: After running SMM/SM4 with a uniform distribution, we count the number of episodes required to ﬁnd an (unknown) goal state.

Original: We run each method for the same number of environment transitions; a mixture of three policies does not get to take three times more transitions.
Revised: We run each method for the same number of environment transitions; a mixture of three policies does not get to take three times more transitions.

Original: We ﬁnd that increasing the number of mixture components increases the agents success.
Revised: We ﬁnd that increasing the number of mixture components increases the agents success.

Original: However, the effect was smaller when using historical averaging.
Revised: However, the effect was smaller when using historical averaging.

Original: Taken together, this result suggests that efﬁcient exploration requires either historical averaging or mixture modelling, but might not need both.
Revised: Taken together, this result suggests that efﬁcient exploration requires either historical averaging or mixture modelling, but might not need both.

Original: [Figure]
Revised: [Figure]

Original: C C HOOSING p ∗ ( s ) FOR G OAL -R EACHING T ASKS
Revised: C C HOOSING p ∗ ( s ) FOR G OAL -R EACHING T ASKS

Original: In general, the choice of the target distribution p ∗ ( s ) will depend on the distribution of test-time tasks.
Revised: In general, the choice of the target distribution p ∗ ( s ) will depend on the distribution of test-time tasks.

Original: In this section, we consider the special case where the test-time tasks correspond to goal-reaching derive the optimal target distribution p ∗ ( s ) .
Revised: In this section, we consider the special case where the test-time tasks correspond to goal-reaching derive the optimal target distribution p ∗ ( s ) .

Original: We consider the setting where goals g ∼ p g ( g ) are sampled from some known distribution.
Revised: We consider the setting where goals g ∼ p g ( g ) are sampled from some known distribution.

Original: Our goal is to minimize the number of episodes required to reach that goal state.
Revised: Our goal is to minimize the number of episodes required to reach that goal state.

Original: We deﬁne reaching the goal state as visiting a state that lies within an (cid:15) ball of the goal, where both (cid:15) > 0 and the distance metric are known.
Revised: We deﬁne reaching the goal state as visiting a state that lies within an (cid:15) ball of the goal, where both (cid:15) > 0 and the distance metric are known.

Original: We start with a simple lemma that shows that the probability that we reach the goal at any state in a trajectory is at least the probability that we reach the goal at a randomly chosen state in that same trajectory.
Revised: We start with a simple lemma that shows that the probability that we reach the goal at any state in a trajectory is at least the probability that we reach the goal at a randomly chosen state in that same trajectory.

Original: Deﬁning the binary random variable z t (cid:44) 

Revised: Deﬁning the binary random variable z t (cid:44)

Original: ( (cid:107) s t − g (cid:107) ≤ (cid:15) ) as the event that the state at time t reaches the goal state, we can formally state the claim as follows:
Revised: ( (cid:107) s t − g (cid:107) ≤ (cid:15) ) as the event that the state at time t reaches the goal state, we can formally state the claim as follows:

Original: Lemma C.1.
Revised: Lemma C.1.

Original: [Equation]
Revised: [Equation]

Original: where t ∼ Unif [1 , · · · , H ]
Revised: where t ∼ Unif [1 , · · · , H ]

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We start by noting the following implication:
Revised: We start by noting the following implication:

Original: [Equation]
Revised: [Equation]

Original: Thus, the probability of the event on the RHS must be at least as large as the probability of the event on the LHS:
Revised: Thus, the probability of the event on the RHS must be at least as large as the probability of the event on the LHS:

Original: [Equation]
Revised: [Equation]

Original: Next, we look at the expected number of episodes to reach the goal state.
Revised: Next, we look at the expected number of episodes to reach the goal state.

Original: Since each episode is independent, the expected hitting time is simply
Revised: Since each episode is independent, the expected hitting time is simply

Original: [Equation]
Revised: [Equation]

Original: Note that we have upper-bounded the hitting time using Lemma C.1.
Revised: Note that we have upper-bounded the hitting time using Lemma C.1.

Original: Since the goal g is a random variable, we take an expectation over g :
Revised: Since the goal g is a random variable, we take an expectation over g :

Original: [Equation]
Revised: [Equation]

Original: We can rewrite the RHS using p ∗ ( s ) to denote the target state marginal distribution:
Revised: We can rewrite the RHS using p ∗ ( s ) to denote the target state marginal distribution:

Original: [Equation]
Revised: [Equation]

Original: We will minimize F , an upper bound on the expected hitting time.
Revised: We will minimize F , an upper bound on the expected hitting time.

Original: Lemma C.2.
Revised: Lemma C.2.

Original: The state marginal distribution p ∗ ( s ) ∝ (cid:112)
Revised: The state marginal distribution p ∗ ( s ) ∝ (cid:112)

Original: ˜
Revised: ˜

Original: p
Revised: p

Original: ( s ) minimizes F ( p ∗ ) ,
Revised: ( s ) minimizes F ( p ∗ ) ,

Original: where ˜
Revised: where ˜

Original: [Equation]
Revised: [Equation]

Original: d
Revised: d

Original: [Equation]
Revised: [Equation]

Original: s is a smoothed version of the target density.
Revised: s is a smoothed version of the target density.

Original: Before presenting the proof, we provide a bit of intuition.
Revised: Before presenting the proof, we provide a bit of intuition.

Original: In the case where (cid:15) → 0 , the optimal target distribution is p ∗ ( s ) ∝ (cid:112) p g ( s ) .
Revised: In the case where (cid:15) → 0 , the optimal target distribution is p ∗ ( s ) ∝ (cid:112) p g ( s ) .

Original: For non-zero (cid:15) , the policy in Lemma C.2 is equivalent to convolving p g ( s ) with a box ﬁlter before taking the square root.
Revised: For non-zero (cid:15) , the policy in Lemma C.2 is equivalent to convolving p g ( s ) with a box ﬁlter before taking the square root.

Original: In both cases, we see that the optimal policy does distribution matching to some function of the goal distribution.
Revised: In both cases, we see that the optimal policy does distribution matching to some function of the goal distribution.

Original: Note that ˜ p ( · ) may not sum to one and therefore is not a proper probability distribution.
Revised: Note that ˜ p ( · ) may not sum to one and therefore is not a proper probability distribution.

Original: Proof.
Revised: Proof.

Original: We start by forming the Lagrangian:
Revised: We start by forming the Lagrangian:

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst derivative is
Revised: The ﬁrst derivative is

Original: [Equation]
Revised: [Equation]

Original: Note that the second derivative is positive, indicating that this Lagrangian is convex, so all stationary points must be global minima:
Revised: Note that the second derivative is positive, indicating that this Lagrangian is convex, so all stationary points must be global minima:

Original: [Equation]
Revised: [Equation]

Original: Setting the ﬁrst derivative equal to zero and rearranging terms, we obtain
Revised: Setting the ﬁrst derivative equal to zero and rearranging terms, we obtain

Original: [Equation]
Revised: [Equation]

Original: Renaming ˜ s ↔ s , we obtain the desired result.
Revised: Renaming ˜ s ↔ s , we obtain the desired result.

Original: C. C ONNECTIONS TO G OAL -C ONDITIONED RL
Revised: C. C ONNECTIONS TO G OAL -C ONDITIONED RL

Original: Goal-Conditioned RL (Kaelbling, 1993; Nair et al., 2018; Held et al., 2017) can be viewed as a special case of State Marginal Matching when the goal-sampling distribution is learned jointly with the policy.
Revised: Goal-Conditioned RL (Kaelbling, 1993; Nair et al., 2018; Held et al., 2017) can be viewed as a special case of State Marginal Matching when the goal-sampling distribution is learned jointly with the policy.

Original: In particular, consider the State Marginal Matching with a mixture policy (Algorithm 2), where the mixture component z maps bijectively to goal states g .
Revised: In particular, consider the State Marginal Matching with a mixture policy (Algorithm 2), where the mixture component z maps bijectively to goal states g .

Original: In this case, we learn goal-conditioned policies of the form π ( a | s, g ) .
Revised: In this case, we learn goal-conditioned policies of the form π ( a | s, g ) .

Original: We start by swapping g for z in the SMM objective with Mixtures of
Revised: We start by swapping g for z in the SMM objective with Mixtures of

Original: Policies (Equation 8):
Revised: Policies (Equation 8):

Original: [Equation]
Revised: [Equation]

Original: The second term p ( g | s ) is an estimate of which goal the agent is trying to reach, similar to objectives in intent inference (Ziebart et al., 2009; Xie et al., 2013).
Revised: The second term p ( g | s ) is an estimate of which goal the agent is trying to reach, similar to objectives in intent inference (Ziebart et al., 2009; Xie et al., 2013).

Original: The third term π ( s | g ) is the distribution over states visited by the policy when attempting to reach goal g .
Revised: The third term π ( s | g ) is the distribution over states visited by the policy when attempting to reach goal g .

Original: For an optimal goal-conditioned policy in an inﬁnite-horizon MDP, both of these terms are Dirac functions:
Revised: For an optimal goal-conditioned policy in an inﬁnite-horizon MDP, both of these terms are Dirac functions:

Original: [Equation]
Revised: [Equation]

Original: In this setting, the State Marginal Matching objective simply says to sample goals g ∼ π ( g ) with probability equal to the density of that goal under the target distribution.
Revised: In this setting, the State Marginal Matching objective simply says to sample goals g ∼ π ( g ) with probability equal to the density of that goal under the target distribution.

Original: [Equation]
Revised: [Equation]

Original: Whether goal-conditioned RL is the preferable way to do distribution matching depends on (1) the difﬁculty of sampling goals and (2) the supervision that will be provided at test time.
Revised: Whether goal-conditioned RL is the preferable way to do distribution matching depends on (1) the difﬁculty of sampling goals and (2) the supervision that will be provided at test time.

Original: It is natural to use goal-conditioned RL in settings where it is easy to sample goals, such as when the space of goals is small and ﬁnite or otherwise low-dimensional.
Revised: It is natural to use goal-conditioned RL in settings where it is easy to sample goals, such as when the space of goals is small and ﬁnite or otherwise low-dimensional.

Original: If a large collection of goals is available apriori, we could use importance sampling to generate goals to train the goal-conditioned policy (Pong et al., 2019).
Revised: If a large collection of goals is available apriori, we could use importance sampling to generate goals to train the goal-conditioned policy (Pong et al., 2019).

Original: However, many real-world settings have high-dimensional goals, which can be challenging to sample.
Revised: However, many real-world settings have high-dimensional goals, which can be challenging to sample.

Original: While goal-conditioned RL is likely the right approach when we will be given a test-time task, a latent-conditioned policy may explore better in settings where the goal-state is not provided at test-time.
Revised: While goal-conditioned RL is likely the right approach when we will be given a test-time task, a latent-conditioned policy may explore better in settings where the goal-state is not provided at test-time.

Original: D A DDITIONAL E XPERIMENTS & E XPERIMENTAL D ETAILS D. E NVIRONMENT D ETAILS
Revised: D A DDITIONAL E XPERIMENTS & E XPERIMENTAL D ETAILS D. E NVIRONMENT D ETAILS

Original: We summarize the environment parameters for Navigation (Figures 2, 5a), Manipulation (Figures 3, 5b, 7, 8, 9, 10), and D’Claw (Figure 4) in Table 1.
Revised: We summarize the environment parameters for Navigation (Figures 2, 5a), Manipulation (Figures 3, 5b, 7, 8, 9, 10), and D’Claw (Figure 4) in Table 1.

Original: Navigation : Episodes have a maximum time horizon of 100 steps.
Revised: Navigation : Episodes have a maximum time horizon of 100 steps.

Original: The environment reward is
Revised: The environment reward is

Original: [Equation]
Revised: [Equation]

Original: for any
Revised: for any

Original: [Equation]
Revised: [Equation]

Original: [ n
Revised: [ n

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where s xy is the xy-position of the agent.
Revised: where s xy is the xy-position of the agent.

Original: We used a uniform target distribution over the end of all m halls, so the environment reward at training time is r env ( s )
Revised: We used a uniform target distribution over the end of all m halls, so the environment reward at training time is r env ( s )

Original: = 1 m if the robot is close enough to the end of any of the halls.
Revised: = 1 m if the robot is close enough to the end of any of the halls.

Original: We used a ﬁxed hall length of 10 in Figures 2b and 2c, and length 50 in Figure 5a.
Revised: We used a ﬁxed hall length of 10 in Figures 2b and 2c, and length 50 in Figure 5a.

Original: All experiments used m = 3 halls, except in Figure 2c where we varied the number of halls { 3 , 5 , 7 } .
Revised: All experiments used m = 3 halls, except in Figure 2c where we varied the number of halls { 3 , 5 , 7 } .

Original: Manipulation .
Revised: Manipulation .

Original: We used the simulated Fetch Robotics arm 5 implemented by Plappert et al.
Revised: We used the simulated Fetch Robotics arm 5 implemented by Plappert et al.

Original: (2018) using the MuJoCo simulator Todorov et al.
Revised: (2018) using the MuJoCo simulator Todorov et al.

Original: The state vector s ∈ R 28 includes the xyzcoordinates s obj , s robot ∈ R 3 of the block and the robot gripper respectively, as well as their velocities, orientations, and relative position s obj − s robot .
Revised: The state vector s ∈ R 28 includes the xyzcoordinates s obj , s robot ∈ R 3 of the block and the robot gripper respectively, as well as their velocities, orientations, and relative position s obj − s robot .

Original: At the beginning of each episode, we spawn the object at the center of the table, and the robot gripper above the initial block position.
Revised: At the beginning of each episode, we spawn the object at the center of the table, and the robot gripper above the initial block position.

Original: We terminate each episode after 50 environment steps, or if the block falls off the table.
Revised: We terminate each episode after 50 environment steps, or if the block falls off the table.

Original: We considered two target state marginal distributions.
Revised: We considered two target state marginal distributions.

Original: In Manipulation-Uniform , the target density is given by
Revised: In Manipulation-Uniform , the target density is given by

Original: [Equation]
Revised: [Equation]

Original: where α 1 , α 2 , α 3 > 0 are ﬁxed weights, and the rewards
Revised: where α 1 , α 2 , α 3 > 0 are ﬁxed weights, and the rewards

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: 1 ( s
Revised: 1 ( s

Original: [Equation]
Revised: [Equation]

Original: is on the table surface )
Revised: is on the table surface )

Original: [Equation]
Revised: [Equation]

Original: correspond to (1) a uniform distribution of the block position over the table surface (the agent receives +0 reward while the block is on the table), (2) an indicator reward for moving the robot gripper close to the block, and (3) action penalty, respectively.
Revised: correspond to (1) a uniform distribution of the block position over the table surface (the agent receives +0 reward while the block is on the table), (2) an indicator reward for moving the robot gripper close to the block, and (3) action penalty, respectively.

Original: The environment reward is a weighted sum of the three reward terms: r env ( s ) (cid:44) 20 r goal ( s ) + r robot ( s ) + 0 .
Revised: The environment reward is a weighted sum of the three reward terms: r env ( s ) (cid:44) 20 r goal ( s ) + r robot ( s ) + 0 .

Original: 1 r action ( s ) .
Revised: 1 r action ( s ) .

Original: At test-time, we sample a goal block location g ∈ R 3 uniformly on the table surface, and the goal is not observed by the agent.
Revised: At test-time, we sample a goal block location g ∈ R 3 uniformly on the table surface, and the goal is not observed by the agent.

Original: In Manipulation-Half , the target state density places higher probability mass to states where the block is on the left-side of the table.
Revised: In Manipulation-Half , the target state density places higher probability mass to states where the block is on the left-side of the table.

Original: This is implemented by replacing r goal ( s ) with a reward function that gives a slightly higher reward +0.1 for states where the block is on the left-side of the table.
Revised: This is implemented by replacing r goal ( s ) with a reward function that gives a slightly higher reward +0.1 for states where the block is on the left-side of the table.

Original: D’Claw .
Revised: D’Claw .

Original: The D’Claw robot (Ahn et al., 2019; Zhu et al., 2019) 6 controls three claws to rotate a valve object.
Revised: The D’Claw robot (Ahn et al., 2019; Zhu et al., 2019) 6 controls three claws to rotate a valve object.

Original: The environment consists of a 9-dimensional action space (three joints per claw) and a 12-dimensional observation space that encodes the joint angles and object orientation.
Revised: The environment consists of a 9-dimensional action space (three joints per claw) and a 12-dimensional observation space that encodes the joint angles and object orientation.

Original: We ﬁxed each
Revised: We ﬁxed each

Original: [Table]
Revised: [Table]

Original: episode at 50 timesteps, which is about 5 seconds on the real robot.
Revised: episode at 50 timesteps, which is about 5 seconds on the real robot.

Original: In the hardware experiments, each algorithm was trained on the same four D’Claw robots to ensure consistency.
Revised: In the hardware experiments, each algorithm was trained on the same four D’Claw robots to ensure consistency.

Original: We deﬁned the target state distribution to place uniform probability mass over all object angles in [ − 180 ◦ , 180 ◦ ] .
Revised: We deﬁned the target state distribution to place uniform probability mass over all object angles in [ − 180 ◦ , 180 ◦ ] .

Original: It also incorporates reward shaping terms that place lower probability mass on states with high joint velocity and on states with joint positions that deviate far from the initial position (see (Zhu et al., 2019)).
Revised: It also incorporates reward shaping terms that place lower probability mass on states with high joint velocity and on states with joint positions that deviate far from the initial position (see (Zhu et al., 2019)).

Original: D. GAILI MPLEMENTATION D ETAILS
Revised: D. GAIL

Original: GAIL assumes access to expert demonstrations, which SMM and the other exploration methods do not require.
Revised: GAIL assumes access to expert demonstrations, which SMM and the other exploration methods do not require.

Original: To compare GAIL with the exploration methods on a level footing, we used states sampled from p ∗ ( s ) to train GAIL
Revised: To compare GAIL with the exploration methods on a level footing, we used states sampled from p ∗ ( s ) to train GAIL

Original: [Table]
Revised: [Table]

Original: and restricted the GAIL discriminator input to states only (no actions).
Revised: and restricted the GAIL discriminator input to states only (no actions).

Original: For D’Claw (Fig. 4), we sampled the valve object angle uniformly in [ − 180 ◦ , 180 ◦ ] .
Revised: For D’Claw (Fig. 4), we sampled the valve object angle uniformly in [ − 180 ◦ , 180 ◦ ] .

Original: For
Revised: For

Original: Manipulation-Uniform (Fig. 3c), we sampled object positions s object uniformly on the table surface, and tried two different sampling distributions for the gripper position s robot (see Fig. 6).
Revised: Manipulation-Uniform (Fig. 3c), we sampled object positions s object uniformly on the table surface, and tried two different sampling distributions for the gripper position s robot (see Fig. 6).

Original: For both environments, all other state dimensions were sampled uniformly in [ − 10 , 10] .
Revised: For both environments, all other state dimensions were sampled uniformly in [ − 10 , 10] .

Original: Since the state samples from p ∗ ( s ) may not be reachable from the initial state, the policy may not be able to fool the discriminator.
Revised: Since the state samples from p ∗ ( s ) may not be reachable from the initial state, the policy may not be able to fool the discriminator.

Original: To get around this problem, we also tried training GAIL with the discriminator input restricted to only the state dimensions corresponding to the object position or gripper position ( Manipulation ), or the object angle ( D’Claw ).
Revised: To get around this problem, we also tried training GAIL with the discriminator input restricted to only the state dimensions corresponding to the object position or gripper position ( Manipulation ), or the object angle ( D’Claw ).

Original: We summarize these GAIL ablation experiments in Fig.
Revised: We summarize these GAIL ablation experiments in Fig.

Original: In our experiments, we used the best GAIL ablation model to compare against the exploration baselines in Figures 3c and 4.
Revised: In our experiments, we used the best GAIL ablation model to compare against the exploration baselines in Figures 3c and 4.

Original: We observed that GAIL training is more unstable compared to the exploration baselines.

Original: Thus, for

Original: GAIL, we did not take the ﬁnal iterate (e.g., policy at convergence) but instead used early termination (e.g., take the best iterate according to the state entropy metric).

Original: If we simply took the ﬁnal iterate, then GAIL performed very poorly (with zero H

Original: π [ s object ] for both Manipulation and D’Claw ).

Original: D. VAE D ENSITY M ODEL
Revised: D. VAE D ENSITY M ODEL

Original: In our SMM implementation, we estimated the density of data x as p ( x ) ≈ decoder (ˆ x = x | z = encoder ( x )) .
Revised: In our SMM implementation, we estimated the density of data x as p ( x ) ≈ decoder (ˆ x = x | z = encoder ( x )) .

Original: That is, we encoded x to z , reconstruction ˆ x from z , and then took the likelihood of the true data x under a unit-variance Gaussian distribution centered at the reconstructed ˆ x .
Revised: That is, we encoded x to z , reconstruction ˆ x from z , and then took the likelihood of the true data x under a unit-variance Gaussian distribution centered at the reconstructed ˆ x .

Original: The log-likelihood is therefore given by the mean-squared error between the data x and the reconstruction ˆ x , plus a constant that is independent of x : log q ( x ) =
Revised: The log-likelihood is therefore given by the mean-squared error between the data x and the reconstruction ˆ x , plus a constant that is independent of x : log q ( x ) =

Original: 1 2 (cid:107) x − ˆ x (cid:107) 2 2 + C .
Revised: 1 2 (cid:107) x − ˆ x (cid:107) 2 2 + C .

Original: [Figure]

Original:  C OMPUTATIONAL C OMPLEXITY
Revised: D. C OMPUTATIONAL C OMPLEXITY

Original: We compare the wall-clock time of each exploration method in Table 2.
Revised: We compare the wall-clock time of each exploration method in Table 2.

Original: The computational cost of our method is comparable with prior work.
Revised: The computational cost of our method is comparable with prior work.

Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: D. A LGORITHM H YPERPARAMETERS
Revised: D. A LGORITHM H YPERPARAMETERS

Original: We summarize hyperparameter settings in Table 3.
Revised: We summarize hyperparameter settings in Table 3.

Original: All algorithms were trained for 1e5 steps on
Revised: All algorithms were trained for 1e5 steps on

Original: Navigation , 1e6 steps on Manipulation , 1e6 steps on D’Claw Sim2Real, and 1e5 steps on D’Claw hardware.
Revised: Navigation , 1e6 steps on Manipulation , 1e6 steps on D’Claw Sim2Real, and 1e5 steps on D’Claw hardware.

Original: Loss Hyperparameters .
Revised: Loss Hyperparameters .

Original: SAC reward scale controls the action entropy reward w.r t. the extrinsic reward.
Revised: SAC reward scale controls the action entropy reward w.r.t. the extrinsic reward.

Original: Count coeff controls the intrinsic count-based exploration reward w.r.t.
Revised: Count coeff controls the intrinsic count-based exploration reward w.r.t.

Original: the extrinsic reward and SAC action entropy reward.
Revised: the extrinsic reward and SAC action entropy reward.

Original: Similarly, Pseudocount coeff controls the intrinsic pseudocount exploration reward.
Revised: Similarly, Pseudocount coeff controls the intrinsic pseudocount exploration reward.

Original: SMM coeff for H [ s | z ] and H [ z | s ] control the different loss components (state entropy and latent conditional entropy) of the SMM objective in Equation 8.
Revised: SMM coeff for H [ s | z ] and H [ z | s ] control the different loss components (state entropy and latent conditional entropy) of the SMM objective in Equation 8.

Original: Historical Averaging .
Revised: Historical Averaging .

Original: In the Manipulation experiments, we tried the following sampling strategies for historical averaging: (1) Uniform : Sample policies uniformly across training iterations.
Revised: In the Manipulation experiments, we tried the following sampling strategies for historical averaging: (1) Uniform : Sample policies uniformly across training iterations.

Original: Exponential : Sample policies, with recent policies sampled exponentially more than earlier ones.
Revised: Exponential : Sample policies, with recent policies sampled exponentially more than earlier ones.

Original: Last :
Revised: Last :

Original: Sample the N latest policies uniformly at random.
Revised: Sample the N latest policies uniformly at random.

Original: We found that Uniform worked less well, possibly due to the policies at early iterations not being trained enough.
Revised: We found that Uniform worked less well, possibly due to the policies at early iterations not being trained enough.

Original: We found negligible difference in the state entropy metric between Exponential vs. Last , and between sampling 5 vs.historical policies, and we also note that it is unnecessary to keep checkpoints from every iteration.
Revised: We found negligible difference in the state entropy metric between Exponential vs. Last , and between sampling 5 vs.historical policies, and we also note that it is unnecessary to keep checkpoints from every iteration.

Original: Network Hyperparameters .
Revised: Network Hyperparameters .

Original: For all algorithms, we use a Gaussian policy with two hidden layers with Tanh activation and a ﬁnal fully-connected layer.
Revised: For all algorithms, we use a Gaussian policy with two hidden layers with Tanh activation and a ﬁnal fully-connected layer.

Original: The Value function and Q-function each are a feedforward MLP with two hidden layers with ReLU activation and a ﬁnal fully-connected layer.
Revised: The Value function and Q-function each are a feedforward MLP with two hidden layers with ReLU activation and a ﬁnal fully-connected layer.

Original: Each hidden layer is of size 300 (SMM, SAC, ICM, C, PC) or 256 (GAIL).
Revised: Each hidden layer is of size 300 (SMM, SAC, ICM, C, PC) or 256 (GAIL).

Original: The same network conﬁguration is used for the SMM discriminator, d
Revised: The same network conﬁguration is used for the SMM discriminator, d

Original: ( z | s ) , but with different input and output sizes.
Revised: ( z | s ) , but with different input and output sizes.

Original: The SMM density model, q ( s ) , is modeled by a VAE with encoder and decoder networks each consisting of two hidden layers of size (150, 150) with ReLU activation.
Revised: The SMM density model, q ( s ) , is modeled by a VAE with encoder and decoder networks each consisting of two hidden layers of size (150, 150) with ReLU activation.

Original: The same VAE network conﬁguration is used for Pseudocount.
Revised: The same VAE network conﬁguration is used for Pseudocount.

Original: GAIL Hyperparameters :
Revised: GAIL Hyperparameters :

Original: The replay buffer is ﬁlled with 1e4 random actions before training, for training stability.
Revised: The replay buffer is ﬁlled with 1e4 random actions before training, for training stability.

Original: We perform onediscriminator update per SAC update. Other hyperparameter settings, such as batch size for both discriminator and policy updates, are summarized in Table 3. 
Revised: We perform one settings, such as batch size for both discriminator and policy updates, are summarized in Table 3. We observed that GAIL training is more unstable compared to the exploration baselines. Thus, for GAIL, we did not take the ﬁnal iterate (e

Revised: [Figure]

Original:  D. V ISUALIZING THE M ANIPULATION E NVIRONMENT
Revised: g., policy at convergence) but instead used early termination (e [Figure] g., take the best iterate according to the state entropy metric). D. V ISUALIZING THE M ANIPULATION E NVIRONMENT

Original: We visualize where different methods push the block in the Manipulation environment.
Revised: We visualize where different methods push the block in the Manipulation environment.

Original: More precisely, we visualize the log state marginal log
Revised: More precisely, we visualize the log state marginal log

Original: ρ π z ( s ) over block XY-coordinates s =
Revised: ρ π z ( s ) over block XY-coordinates s =

Original: ( x, y ) in
Revised: ( x, y ) in

Original: Figures 7 and 8.
Revised: Figures 7 and 8.

Original: In Figure 9, we plot goals sampled at test-time, colored by the number of episodes each method required to push the block to that goal location.
Revised: In Figure 9, we plot goals sampled at test-time, colored by the number of episodes each method required to push the block to that goal location.

Original: Blue dots indicate that the agent found the goal quickly.
Revised: Blue dots indicate that the agent found the goal quickly.

Original: We observe that SMM has the most blue dots, indicating that it succeeds in exploring a wide range of states at test-time.
Revised: We observe that SMM has the most blue dots, indicating that it succeeds in exploring a wide range of states at test-time.

Original: [Figure]
Revised: [Figure]

Original: [Figure]

Original: [Figure]
Revised: [Figure]

Original: environment reward is a weighted sum of three terms: r goal ( s ) (+0 if object is on table, -1 otherwise), r robot ( s ) (+ if robot gripper is close to block), and r action (action penalty term), with weights -20, 1, 0.1 respectively (see Appendix D.1).

Original: The exploration methods (ICM, Count, SMM) also optimize an auxilliary exploration loss, which makes the agent more likely to move around the block.
Revised: (ICM, Count, SMM) also optimize an auxilliary exploration loss the agent more likely to move around the block

Original: Compared to SAC, this causes worse due
Revised: SAC

Original: [Table]
Revised: [Table]