Original: Discrete Factorial Representations as an Abstraction for Goal Conditioned RL
Revised: Discrete Factorial Representations as an Abstraction for Goal Conditioned RL

Original: Abstract
Revised: Abstract

Original: Goal-conditioned reinforcement learning (RL) is a promising direction for training agents that are capable of solving multiple tasks and reach a diverse set of objectives.
Revised: Goal-conditioned reinforcement learning (RL) is a promising direction for training agents that are capable of solving multiple tasks and reach a diverse set of objectives.

Original: How to specify and ground these goals in such a way that we can both reliably reach goals during training as well as generalize to new goals during evaluation remains an open area of research.
Revised: How to specify and ground these goals in such a way that we can both reliably reach goals during training as well as generalize to new goals during evaluation remains an open area of research.

Original: Deﬁning goals in the space of noisy and high-dimensional sensory inputs poses a challenge for training goal-conditioned agents, or even for generalization to novel goals.
Revised: Deﬁning goals in the space of noisy and high-dimensional sensory inputs poses a challenge for training goal-conditioned agents, or even for generalization to novel goals.

Original: We propose to address this by learning factorial representations of goals and processing the resulting representation via a discretization bottleneck, for coarser goal speciﬁcation, through an approach we call DGRL.
Revised: We propose to address this by learning factorial representations of goals and processing the resulting representation via a discretization bottleneck, for coarser goal speciﬁcation, through an approach we call DGRL.

Original: We show that applying a discretizing bottleneck can improve performance in goal-conditioned RL setups, by experimentally evaluating this method on tasks ranging from maze environments to complex robotic navigation and manipulation.
Revised: We show that applying a discretizing bottleneck can improve performance in goal-conditioned RL setups, by experimentally evaluating this method on tasks ranging from maze environments to complex robotic navigation and manipulation.

Original: Additionally, we prove a theorem lower-bounding the expected return on out-of-distribution goals, while still allowing for specifying goals with expressive combinatorial structure.
Revised: Additionally, we prove a theorem lower-bounding the expected return on out-of-distribution goals, while still allowing for specifying goals with expressive combinatorial structure.

Original:  Introduction
Revised:  Introduction

Original: Reinforcement Learning is a popular and highly general framework [25, 60] focusing on how to select actions for an agent to yield high long-term sum of rewards.
Revised: Reinforcement Learning is a popular and highly general framework [25, 60] focusing on how to select actions for an agent to yield high long-term sum of rewards.

Original: An important question is how to control the desired behavior of an RL agent, both during training and evaluation [24].
Revised: An important question is how to control the desired behavior of an RL agent, both during training and evaluation [24].

Original: One way to control this behavior is by specifying a reward signal [52, 56].
Revised: One way to control this behavior is by specifying a reward signal [52, 56].

Original: While this approach is very general the reward signal can be hard to design and may not be the most informative form of feedback.
Revised: While this approach is very general, the reward signal can be hard to design and may not be the most informative form of feedback.

Original: The credit assignment problem in RL can become difﬁcult when the reward signal is sparse [62, 63, 36, 37, 57], such as policy gradients becoming nearly ﬂat in regions where reward is almost never achieved.
Revised: The credit assignment problem in RL can become difﬁcult when the reward signal is sparse [62, 63, 36, 37, 57], such as policy gradients becoming nearly ﬂat in regions where reward is almost never achieved.

Original: Generalization can also suffer if the agent only learns one way to achieve a high reward rather than learning a diverse set of skills for coping with novel challenges [20].
Revised: Generalization can also suffer if the agent only learns one way to achieve a high reward rather than learning a diverse set of skills for coping with novel challenges [20].

Original: One potential way to ﬂexibly specify and ground the desired behavior of RL agents is by training agents that receive a reward when they reach a goal speciﬁed explicitly to them [23].
Revised: One potential way to ﬂexibly specify and ground the desired behavior of RL agents is by training agents that receive a reward when they reach a goal speciﬁed explicitly to them [23].

Original: In this approach, called Goal-Conditioned RL , a single agent is trained to reach a diverse set of goals, and is given a reward only when it reaches the goal it was instructed to reach [61, 48, 44].
Revised: In this approach, called Goal-Conditioned RL , a single agent is trained to reach a diverse set of goals, and is given a reward only when it reaches the goal it was instructed to reach [61, 48, 44].

Original: This provides a richer signal for the agent than simply collecting more samples oriented around a single goal, as reaching multiple goals requires the agent to learn a more diverse and robust set of skills.
Revised: This provides a richer signal for the agent than simply collecting more samples oriented around a single goal, as reaching multiple goals requires the agent to learn a more diverse and robust set of skills.

Original: It also allows for more ﬂexible and tightly constrained control over its desired behavior [12, 5, 27].
Revised: It also allows for more ﬂexible and tightly constrained control over its desired behavior [12, 5, 27].

Original: Finally, the diversity of goals seen during training should help improve both credit assignment and generalization [48, 43].
Revised: Finally, the diversity of goals seen during training should help improve both credit assignment and generalization [48, 43].

Original: While this framework is promising, it introduces two new challenges: goal grounding [8, 2] and goal speciﬁcation [4].
Revised: While this framework is promising, it introduces two new challenges: goal grounding [8, 2] and goal speciﬁcation [4].

Original: Goal grounding refers to deﬁning the goal space, and goal speciﬁcation refers to selecting what goal the agent should try to reach in a given context.
Revised: Goal grounding refers to deﬁning the goal space, and goal speciﬁcation refers to selecting what goal the agent should try to reach in a given context.

Original: The agent is only rewarded in goal-conditioned RL when reaching the reward it was instructed to reach, whereas in goal-free RL a reward is provided regardless of any such speciﬁcation, which makes the nature of the agent’s task fundamentally different.
Revised: The agent is only rewarded in goal-conditioned RL when reaching the reward it was instructed to reach, whereas in goal-free RL a reward is provided regardless of any such speciﬁcation, which makes the nature of the agent’s task fundamentally different.

Original: What makes grounding and specifying goals challenging?
Revised: What makes grounding and specifying goals challenging?

Original: Consider trying to train a goalconditioned RL agent to pick up various fruits from a table.
Revised: Consider trying to train a goalconditioned RL agent to pick up various fruits from a table.

Original: For example, we may want it to pick up a red apple or a green pear (illustrated in Figure 1).
Revised: For example, we may want it to pick up a red apple or a green pear (illustrated in Figure 1).

Original: The number of possible goals of interests may be fairly small, such as the set of all valid combinations of fruits and their colors, while the number of possible observations of goals is extremely large when working in a rich observation space ( e.g images from a camera).
Revised: The number of possible goals of interests may be fairly small, such as the set of all valid combinations of fruits and their colors, while the number of possible observations of goals is extremely large when working in a rich observation space ( e.g images from a camera).

Original: Goal Grounding refers to this challenge of relating high-dimensional observations and the space of relevant goals.
Revised: Goal Grounding refers to this challenge of relating high-dimensional observations and the space of relevant goals.

Original: Goal Speciﬁcation refers to picking a suitable goal for the agent to reach and computing an appropriate reward when it is reached.
Revised: Goal Speciﬁcation refers to picking a suitable goal for the agent to reach and computing an appropriate reward when it is reached.

Original: It also implies specifying goals reachable in the agent’s current context [39, 26].
Revised: It also implies specifying goals reachable in the agent’s current context [39, 26].

Original: Goal speciﬁcation can be done either manually by a developer or by another RL agent, such as a high-level agent which generates goals a lower-level agent then tries to reach [12, 5, 27, 22].
Revised: Goal speciﬁcation can be done either manually by a developer or by another RL agent, such as a high-level agent which generates goals a lower-level agent then tries to reach [12, 5, 27, 22].

Original: Goals speciﬁed in language are an excellent ﬁt for these desiderata, as language is a compressed discrete representation which is useful for out-ofdistribution generalization, while being compositional and expressive [19, 11, 22, 17, 1, 67].
Revised: Goals speciﬁed in language are an excellent ﬁt for these desiderata, as language is a compressed discrete representation which is useful for out-ofdistribution generalization, while being compositional and expressive [19, 11, 22, 17, 1, 67].

Original: At the same time, connecting language feedback for an agent is non-trivial, requiring special assumptions or a labeling framework [9].
Revised: At the same time, connecting language feedback for an agent is non-trivial, requiring special assumptions or a labeling framework [9].

Original: We propose to learn the goal representations with self-supervised learning (either trained on their own, or jointly with the downstream RL objective) while forcing them to be discrete and factorial .
Revised: We propose to learn the goal representations with self-supervised learning (either trained on their own, or jointly with the downstream RL objective) while forcing them to be discrete and factorial .

Original: To perform this discretization, we use Vector-Quantization [64, 47, 33] which discretizes a continuous representation using a codebook of discrete and learnable codes.
Revised: To perform this discretization, we use Vector-Quantization [64, 47, 33] which discretizes a continuous representation using a codebook of discrete and learnable codes.

Original: The approach proposed here, called DGRL, serves two complementary purposes.
Revised: The approach proposed here, called DGRL, serves two complementary purposes.

Original: First, it provides a structured representation of the raw visual goals.
Revised: First, it provides a structured representation of the raw visual goals.

Original: By representing the visual goals as a composition of discrete codes from a learned dictionary, it simpliﬁes the grounding of unseen goals, i.e. , goals not seen during training, to novel compositions of the trained discrete codes.
Revised: By representing the visual goals as a composition of discrete codes from a learned dictionary, it simpliﬁes the grounding of unseen goals, i.e. , goals not seen during training, to novel compositions of the trained discrete codes.

Original: We show empirically that this improves the generalization performance of goal-reaching policies while remaining expressive enough.
Revised: We show empirically that this improves the generalization performance of goal-reaching policies while remaining expressive enough.

Original: Second, the learned discrete codes can be used by another agent (like a higher-level policy in hierarchical RL) to specify sub-goals to a lower-level policy, and eventually complete the task ( i.e. , reach the ﬁnal goal).
Revised: Second, the learned discrete codes can be used by another agent (like a higher-level policy in hierarchical RL) to specify sub-goals to a lower-level policy, and eventually complete the task ( i.e. , reach the ﬁnal goal).

Original: In this case, goal-inference is learned end-to-end.
Revised: In this case, goal-inference is learned end-to-end.

Original: The effectiveness of goal-conditioned HRL relies on the speciﬁcation of semantically meaningful sub-goals.
Revised: The effectiveness of goal-conditioned HRL relies on the speciﬁcation of semantically meaningful sub-goals.

Original: Using factorial discrete sub-goals allows the higher-level policy to specify semantically meaningful objectives to the lower-level policy.
Revised: Using factorial discrete sub-goals allows the higher-level policy to specify semantically meaningful objectives to the lower-level policy.

Original:  Preliminaries
Revised:  Preliminaries

Original: Goal-conditioned RL.
Revised: Goal-conditioned RL.

Original: We consider a goal-conditioned Markov Decision Process, where the goals g ∈ G live in the state space S , i.e. , G = S .
Revised: We consider a goal-conditioned Markov Decision Process, where the goals g ∈ G live in the state space S , i.e. , G = S .

Original: We denote a goal-conditioned policy as π ( a | s, g ) (either stochastic or deterministic), and its expected total return as J ( π ) =E (cid:2) (cid:3)
Revised: We denote a goal-conditioned policy as π ( a | s, g ) (either stochastic or deterministic), and its expected total return as J ( π ) =

Original: [-Tt =
-]
Revised: [-Tt =
-]E

Original: R ( s t , g, a ) (cid:4) where the goal g is either sampled from a distribution ρ g or provided by another higher level policy π hθ h ( g | s ) .
Revised: [Equation] # where the goal g is either sampled from a distribution ρ g or provided by another higher level policy π h θ h ( g | s ) .

Original: The value function V π is additionally conditioned on goals, and is trained to predict the expected sum of future rewards conditioned on states and goals; V π ( s, g ) = 
Revised: The value function V π is additionally conditioned on goals, and is trained to predict the expected sum of future rewards conditioned on states and goals; V π ( s, g ) = E ! "

Original: E (cid:2) (cid:3)
Revised: [Equation]

Original: [-Tt =
 R ( s t , g, a ) |-] s 0{+
+} = s ; π (cid:4) .
Revised: [-Tt =
 R ( s t , g, a ) |-] s {+
+} = s ; π # .

Original: As in standard RL, the objective in goal-conditioned RL is to maximize the expected discountedreturns induced by the goal-conditioned policy.
Revised: As in standard RL, the objective in goal-conditioned RL is to maximize the expected discounted

Original: Hierarchical Reinforcement Learning.
Revised: Hierarchical Reinforcement Learning.

Original: We consider goal-conditioned settings in which the goals are speciﬁed in the observation space.
Revised: We consider goal-conditioned settings in which the goals are speciﬁed in the observation space.

Original: In the hierarchical reinforcement learning (HRL) setup, goals are provided by a higher level policy π hθ h ( g | s t ) .
Revised: In the hierarchical reinforcement learning (HRL) setup, goals are provided by a higher level policy π h θ h ( g | s t ) .

Original: The higher level policy operates at a coarser time scale and chooses a goal g t ∼ π hθ h ( g | s t ) to reach for the lower level policy every K steps.
Revised: The higher level policy operates at a coarser time scale and chooses a goal g t ∼ π h θ h ( g | s t ) to reach for the lower level policy every K steps.

Original: The lower level policy executes primitive actions π lθ l ( a | s t , g t ) to reach the goals speciﬁed by the high-level policy and is trained to maximize the intrinsic reward provided by the high-level policy.
Revised: The lower level policy executes primitive actions π l θ l ( a | s t , g t ) to reach the goals speciﬁed by the high-level policy and is trained to maximize the intrinsic reward provided by the high-level policy.

Original: The higher level policy is trained to maximize the external reward i.e. , the reward function speciﬁed by the
Revised: The higher level policy is trained to maximize the external reward i.e. , the reward function speciﬁed by the

Original: MDP.
Revised: MDP.

Original: Both the higher and lower level policies can be trained with any standard RL algorithms, such as Deep Q-Learning (DQN) [36] or policy optimization based algorithms [50, 51, 30].
Revised: Both the higher and lower level policies can be trained with any standard RL algorithms, such as Deep Q-Learning (DQN) [36] or policy optimization based algorithms [50, 51, 30].

Original: Alternately, one can also consider another setup for goal-conditioned RL, where the goals are provided by the environment g 1 , . . .
Revised: Alternately, one can also consider another setup for goal-conditioned RL, where the goals are provided by the environment g 1 , . . .

Original: g L , and are part of the state or observation space.
Revised: g L , and are part of the state or observation space.

Original: At each episode of training, one of the goals is sampled from the distribution of goals ρ g and the policy is trained to reach the sampled goal.
Revised: At each episode of training, one of the goals is sampled from the distribution of goals ρ g and the policy is trained to reach the sampled goal.

Original: At test time, the agent can be evaluated either on its ability to reach goals within the distribution ρ g , or for its out-of-distribution generalization capability to reach new kinds of goals.
Revised: At test time, the agent can be evaluated either on its ability to reach goals within the distribution ρ g , or for its out-of-distribution generalization capability to reach new kinds of goals.

Original: We consider both the HRL and goal-conditioned setups, and evaluate the signiﬁcance of learning factorial representation of discrete latent goals in a series of complex goal-conditioned tasks.
Revised: We consider both the HRL and goal-conditioned setups, and evaluate the signiﬁcance of learning factorial representation of discrete latent goals in a series of complex goal-conditioned tasks.

Original: Vector Quantized Representations.
Revised: Vector Quantized Representations.

Original: VQ-VAE [64, 47, 33] discretizes the bottleneck representation of an auto-encoder by adding a codebook of discrete learnable codes.
Revised: VQ-VAE [64, 47, 33] discretizes the bottleneck representation of an auto-encoder by adding a codebook of discrete learnable codes.

Original: The input is passed through an encoder.
Revised: The input is passed through an encoder.

Original: The output of the encoder is compared to all the vectors in the codebook, and the codebook vector closest to the continuous encoded representation is fed to the decoder.
Revised: The output of the encoder is compared to all the vectors in the codebook, and the codebook vector closest to the continuous encoded representation is fed to the decoder.

Original: The decoder is then tasked with reconstructing the input from this quantized vector.
Revised: The decoder is then tasked with reconstructing the input from this quantized vector.

Original: Self-supervised learning of representations.
Revised: Self-supervised learning of representations.

Original: Several papers [31, 58, 53, 34] have demonstrated the beneﬁts of using a pre-training stage where the representations of raw states are learned using self-supervised objectives in a task-agnostic fashion.
Revised: Several papers [31, 58, 53, 34] have demonstrated the beneﬁts of using a pre-training stage where the representations of raw states are learned using self-supervised objectives in a task-agnostic fashion.

Original: After the pre-training stage, the representations can be used for (and potentially also ﬁne-tuned on) downstream tasks.
Revised: After the pre-training stage, the representations can be used for (and potentially also ﬁne-tuned on) downstream tasks.

Original: These self-supervised representations have been shown to improve sample efﬁciency.
Revised: These self-supervised representations have been shown to improve sample efﬁciency.

Original:  Discrete Goal-Conditioned Reinforcement Learning (DGRL)
Revised:  Discrete Goal-Conditioned Reinforcement Learning (DGRL)

Original: In this section, we provide technical details on the proposed framework, DGRL, which consists of three parts: (a) learning representations of raw visual observations through self-supervised representation objectives, (b) processing the resulting representations via a learned dictionary of discrete codes, and (c) using the resulting discrete representations for downstream goal-conditioned and HRL tasks.
Revised: In this section, we provide technical details on the proposed framework, DGRL, which consists of three parts: (a) learning representations of raw visual observations through self-supervised representation objectives, (b) processing the resulting representations via a learned dictionary of discrete codes, and (c) using the resulting discrete representations for downstream goal-conditioned and HRL tasks.

Original: We later describe, in Section 6, how discrete goal representations can accelerate learning in complex navigation and manipulation tasks.
Revised: We later describe, in Section 6, how discrete goal representations can accelerate learning in complex navigation and manipulation tasks.

Original: We emphasize that these representations can be learned at the same time as the downstream-RL objective or pre-trained with self-supervised learning, and then used as a ﬁxed representation for RL.
Revised: We emphasize that these representations can be learned at the same time as the downstream-RL objective or pre-trained with self-supervised learning, and then used as a ﬁxed representation for RL.

Original: Self-Supervised Goal Representation Learning
Revised: Self-Supervised Goal Representation Learning

Original: One can use any off-the shelf self-supervised method for learning representations of the raw state and the goal observations.
Revised: One can use any off-the shelf self-supervised method for learning representations of the raw state and the goal observations.

Original: We denote by φ the encoder network that takes as input the raw state and maps it to a continuous embedding: z e = φ ( s ) .
Revised: We denote by φ the encoder network that takes as input the raw state and maps it to a continuous embedding: z e = φ ( s ) .

Original: Here, we explore two different self-supervised techniques for learning representations.
Revised: Here, we explore two different self-supervised techniques for learning representations.

Original: For simpler environments, we use a simple autoencoder with its standard reconstruction objective.
Revised: For simpler environments, we use a simple autoencoder with its standard reconstruction objective.

Original: For more complex environments, we use the Deep InfoMax approach [34] which optimizes for a contrastive objective as a proxy to maximizing the mutual information between representations of nearby states in the same trajectory.
Revised: For more complex environments, we use the Deep InfoMax approach [34] which optimizes for a contrastive objective as a proxy to maximizing the mutual information between representations of nearby states in the same trajectory.

Original: Processing continuous representations via a discrete codebook.
Revised: Processing continuous representations via a discrete codebook.

Original: We learn discrete representations by using the vector-quantization method from the VQ-VAE paper [64], and follow the multifactor setup used in Discrete-Value Neural Communication [33].
Revised: We learn discrete representations by using the vector-quantization method from the VQ-VAE paper [64], and follow the multifactor setup used in Discrete-Value Neural Communication [33].

Original: The discretization process for each vector z e ∈ H ⊂ R m is described as follows.
Revised: The discretization process for each vector z e ∈ H ⊂ R m is described as follows.

Original: First, vector z e is divided into G segments c 1 , c 2 , . . .
Revised: First, vector z e is divided into G segments c 1 , c 2 , . . .

Original: , c G , where z e = CONCATENATE ( c 1 , c 2 , . . .
Revised: , c G , where z e = CONCATENATE ( c 1 , c 2 , . . .

Original: , c G ) , and each segment c i ∈ R m/G (implying that m is divisible by G ).
Revised: , c G ) , and each segment c i ∈ R m/G (implying that m is divisible by G ).

Original: Each continuous segment c i is mapped independently to a discretized latent vector e ∈ R L × ( m/G ) where L is the size of the discrete latent space ( i.e. , an L -way categorical variable):
Revised: Each continuous segment c i is mapped independently to a discretized latent vector e ∈ R L × ( m/G ) where L is the size of the discrete latent space ( i.e. , an L -way categorical variable):

Original: [Equation]
Revised: [Equation]

Original: These discrete codes, which we call the factors of the continuous representation z e , are concatenated to obtain the ﬁnal discretized vector z q :[Equation] [Equation] DISCRETIZE
Revised: These discrete codes, which we call the factors of the continuous representation z e , are concatenated to obtain the ﬁnal discretized vector z q :

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: The loss for vector quantization is: L discretization
Revised: The loss for vector quantization is: L discretization

Original: [Equation]
Revised: [Equation]

Original: The training procedure closely follows both [33] and [64].
Revised: The training procedure closely follows both [33] and [64].

Original: Here, sg refers to a stop-gradient operation that blocks gradients from ﬂowing into e o i , and β is a hyperparameter which controls how strongly we move the codes toward the encoded values.
Revised: Here, sg refers to a stop-gradient operation that blocks gradients from ﬂowing into e o i , and β is a hyperparameter which controls how strongly we move the codes toward the encoded values.

Original: Unlike [33], we used a moving average to update the code embeddings rather than learning them directly as parameters.
Revised: Unlike [33], we used a moving average to update the code embeddings rather than learning them directly as parameters.

Original: We update e o i with an exponential moving average to encourage it to become close to the selected output segment c i .
Revised: We update e o i with an exponential moving average to encourage it to become close to the selected output segment c i .

Original: This update sets the new value of e o i to be equal to ηe o i +(1 − η ) c i , where the value of η is a ﬁxed hyperparameter controlling how quickly the moving average updates.
Revised: This update sets the new value of e o i to be equal to η e o i +(1 − η ) c i , where the value of η is a ﬁxed hyperparameter controlling how quickly the moving average updates.

Original:  The term (cid:3)
Revised: 2 The term "

Original: Gi || c i − sg( e o i )
Revised: Gi || c i − sg( e o i )

Original: || 22 is often called the commitment loss.
Revised: || 22 is often called the commitment loss.

Original: We trained the VQ-quantization process together with other parts of the model by gradient descent.
Revised: We trained the VQ-quantization process together with other parts of the model by gradient descent.

Original: When there were multiple z e vectors to discretize in a model, the mean of the commitment loss across all z e vectors was used.
Revised: When there were multiple z e vectors to discretize in a model, the mean of the commitment loss across all z e vectors was used.

Original: Summary.
Revised: Summary.

Original: The multiple steps described above can be summarized by z q = q
Revised: The multiple steps described above can be summarized by z q = q

Original: ( z e , L, G ) , where L is the codebook size, G the number of factors per vector, and q ( · ) the whole discretization process.
Revised: ( z e , L, G ) , where L is the codebook size, G the number of factors per vector, and q ( · ) the whole discretization process.

Original: We train the representations for both the state and goal observations with this discretization bottleneck applied to the continuous representations resulting from the self-supervised training.
Revised: We train the representations for both the state and goal observations with this discretization bottleneck applied to the continuous representations resulting from the self-supervised training.

Original: The number of factors G is a hyper-parameter.
Revised: The number of factors G is a hyper-parameter.

Original: In our experiments, we explored different values: G = 1 , 2 , 4 , 8 , 16 , and found that G = 16 worked the best.
Revised: In our experiments, we explored different values: G = 1 , 2 , 4 , 8 , 16 , and found that G = 16 worked the best.

Original: Discretizing with more factors slightly increases computation but reduces the number of model parameters due to the codebook embeddings being reused across the different factors.
Revised: Discretizing with more factors slightly increases computation but reduces the number of model parameters due to the codebook embeddings being reused across the different factors.

Original: 3.1 Using representations for downstream RL
Revised: 3.1 Using representations for downstream RL

Original: We use the discrete representations for downstream RL tasks: (i) to train a goal-conditioned policy, and (ii) in the context of hierarchical reinforcement learning.
Revised: We use the discrete representations for downstream RL tasks: (i) to train a goal-conditioned policy, and (ii) in the context of hierarchical reinforcement learning.

Original: Goal-conditioned RL.
Revised: Goal-conditioned RL.

Original: Deﬁning goals in the space of noisy, high-dimensional sensory inputs poses a challenge for generalization to novel goals because the encoder that maps the goal observations to the low dimensional latent representation may fail to generalize.
Revised: Deﬁning goals in the space of noisy, high-dimensional sensory inputs poses a challenge for generalization to novel goals because the encoder that maps the goal observations to the low dimensional latent representation may fail to generalize.

Original: One way to address this is to embed the continuous latent representation into a discrete representation such that the representation of the novel goal is mapped to the ﬁxed set of latent discrete codes.
Revised: One way to address this is to embed the continuous latent representation into a discrete representation such that the representation of the novel goal is mapped to the ﬁxed set of latent discrete codes.

Original: This facilitate generalization to new combinations of these codes while making it easy for downstream learning to ﬁgure out the meaning of each discrete code.
Revised: This facilitate generalization to new combinations of these codes while making it easy for downstream learning to ﬁgure out the meaning of each discrete code.

Original: In this setup, instead of feeding the continuous state and goal embeddings to the agent, we use their discretized versions, thus grounding goal representations in the input space.
Revised: In this setup, instead of feeding the continuous state and goal embeddings to the agent, we use their discretized versions, thus grounding goal representations in the input space.

Original: We use the resulting representations to train a goal-conditioned policy a t ∼ π lθ l ( a | s t , g t ) or a goalconditioned action value function Q ( s t , a t , g t ) .
Revised: We use the resulting representations to train a goal-conditioned policy a t ∼ π l θ l ( a | s t , g t ) or a goalconditioned action value function Q ( s t , a t , g t ) .

Original: At each training episode, a goal is sampled from the goal distribution ρ g , and the agent gets rewarded for reaching it.
Revised: At each training episode, a goal is sampled from the goal distribution ρ g , and the agent gets rewarded for reaching it.

Original: This reward can either be extrinsic , i.e. , part of the environment, or intrinsic , i.e. , part of the algorithm.
Revised: This reward can either be extrinsic , i.e. , part of the environment, or intrinsic , i.e. , part of the algorithm.

Original: In DGRL, we deﬁne the intrinsic reward as the fraction of discrete factors which match in the respective representations of the goal observation and of the state observation.
Revised: In DGRL, we deﬁne the intrinsic reward as the fraction of discrete factors which match in the respective representations of the goal observation and of the state observation.

Original: At test time, the agent can either be evaluated on reaching goals within the distribution ρ g , or for its generalization capability to goals not seen during training.
Revised: At test time, the agent can either be evaluated on reaching goals within the distribution ρ g , or for its generalization capability to goals not seen during training.

Original: Hierarchical RL.
Revised: Hierarchical RL.

Original: The higher level policy g t ∼ π hθ h ( g | s t ) outputs a continuous representation of goals g by conditioning on the states every K time-steps, it can also output a sub-goal s g by conditioning on both states s and environment goals g , i.e. , π hθ h ( s g | s t , g t ) .
Revised: The higher level policy g t ∼ π h θ h ( g | s t ) outputs a continuous representation of goals g by conditioning on the states every K time-steps, it can also output a sub-goal s g by conditioning on both states s and environment goals g , i.e. , π h θ h ( s g | s t , g t ) .

Original: The effectiveness of goal-conditioned HRL relies on the speciﬁcation of semantically meaningful sub-goals.
Revised: The effectiveness of goal-conditioned HRL relies on the speciﬁcation of semantically meaningful sub-goals.

Original: Learned codebooks (Section 3) consisting of a set of discrete codes can be used by a higher level policy to specify which goal to reach to a lower level policy.
Revised: Learned codebooks (Section 3) consisting of a set of discrete codes can be used by a higher level policy to specify which goal to reach to a lower level policy.

Original: The use of learned codebooks ensures that the goal speciﬁed by the higher level policy is grounded in the space of raw-observations.
Revised: The use of learned codebooks ensures that the goal speciﬁed by the higher level policy is grounded in the space of raw-observations.

Original: In Section 6, we empirically show the beneﬁts of the proposed approach for training goal-reaching policies or goal-conditioned value functions, as well as 
Revised: In Section 6, we empirically show the beneﬁts of the proposed approach for training goal-reaching policies or goal-conditioned value functions, as well as in a goal-conditioned hierarchical RL setup.

Original:  Theoretical Analysis
Revised:  Theoretical Analysis

Original: In this section, discretization is shown to improve generalization to novel goals by enhancing the concentration of the goal distribution within each neighborhood of discretized goal values; i.e. , by decomposing the goal probability p ( g ) into p ( g ) = (cid:3) k p
Revised: In this section, discretization is shown to improve generalization to novel goals by enhancing the concentration of the goal distribution within each neighborhood of discretized goal values; i.e. , by decomposing the goal probability p ( g ) into p ( g ) = "

Original:  ( g | g ∈ G k ) p ( g ∈ G k ) with the neighborhood set {G k } k , it improves the overall performance in p ( g ) by increasing the concentration in p ( g | g ∈ G k ) .
Revised: k p ( g | g ∈ G k ) p ( g ∈ G k ) with the neigh- borhood set {G k } k , it improves the overall performance in p ( g ) by increasing the concentration in p ( g | g ∈ G k ) .

Original: Intuitively, this is because the discretization removes varieties of possible goal values g ∈ G k for each neighborhood G k .
Revised: Intuitively, this is because the discretization removes varieties of possible goal val- ues g ∈ G k for each neighborhood G k .

Original: To state our result, we deﬁne ϕ θ ( g ) =
Revised: To state our result, we deﬁne ϕ θ ( g ) =

Original: E s 

Revised: E s

Original: [ V π ( s 0 , g )] , where θ ∈ R m is the vector containing model parameters learned through n goals observed during training phase, g 1 , . . .
Revised: [ V π ( s 0 , g )] , where θ ∈ R m is the vector containing model parameters learned through n goals observed during training phase, g 1 , . . .

Original: , g n .
Revised: , g n .

Original: We denote the discretization of g by q ( g ) , and the identity function by id as id( g ) =
Revised: We denote the discretization of g by q ( g ) , and the identity function by id as id( g ) =

Original: g .
Revised: g .

Original: Let Q = { q ( g ) : g ∈ G} and ˆ d be a distance function.
Revised: Let Q = { q ( g ) : g ∈ G} and ˆ d be a distance function.

Original: We use Q i to denote the i -th element of Q (by ordering elements of Q with an arbitrary ordering).
Revised: We use Q i to de- note the i -th element of Q (by ordering elements of Q with an arbitrary ordering).

Original: We also deﬁne [ n ] = { 1 , . . . , n } , G k = { g ∈ G : k = arg min i ∈ [ |Q| ] ˆ d ( q ( g ) , Q i ) } , I k
Revised: We also deﬁne [ n ] = { 1 , . . . , n } , G k = { g ∈ G : k = arg min i ∈ [ |Q| ] ˆ d ( q ( g ) , Q i ) } , I k

Original: = { i ∈ [ n ] : g i ∈
Revised: = { i ∈ [ n ] : g i ∈

Original: G k } , and
Revised: G k } , and

Original: I Q = { k ∈ [ |Q| ] : |I k | ≥ 1 } .
Revised: I Q = { k ∈ [ |Q| ] : |I k | ≥ 1 } .

Original: We denote by c a constant in ( n, θ, Θ , δ, S ) .
Revised: We denote by c a constant in ( n, θ , Θ , δ , S ) .

Original: The following theorem (proof in Appendix D) shows that the goal discretization improves the lower bound of the expected sum of rewards for unseen goals
Revised: The following theorem (proof in Appendix D) shows that the goal discretization improves the lower bound of the expected sum of rewards for unseen goals

Original: E g ∼ ρ g
Revised: E g ∼ ρ g

Original: ( ϕ θ ◦ ς )( g ))] by the margin of ω ( θ ) :
Revised: ( ϕ θ ◦ ς )( g ))] by the margin of ω ( θ ) :

Original: Theorem 1.
Revised: Theorem 1.

Original: For any δ > 0 , with probability at least 1 − δ , the following holds for any θ ∈ R m and ς ∈ { id , q } :
Revised: For any δ > 0 , with probability at least 1 − δ , the following holds for any θ ∈ R m and ς ∈ { id , q } :

Original: [Equation]
Revised: [Equation]

Original: where ω ( θ ) = 1 n (cid:3) k ∈I Q |I k | (cid:7)|I k | (cid:3) i ∈I k ϕ θ ( g i ) − E
Revised: where ω ( θ ) = 1 n " k ∈ I Q |I k | & |I k | " i ∈ I k ϕ θ ( g i ) − E

Original: g ∼ ρ g
Revised: g ∼ ρ g

Original: [ ϕ θ ( g ) | g ∈ G k ] (cid:8) .
Revised: [ ϕ θ ( g ) | g ∈ G k ] ’ .

Original: Moreover, for any compact Θ ⊂ R m , if ϕ θ ( g ) is continuous at each θ ∈ Θ for almost all g and is dominated by a function χ as | ϕ θ ( g ) | ≤ χ ( g ) for all θ ∈ Θ with E g [ χ ( g )] < ∞ , then the following holds:
Revised: Moreover, for any compact Θ ⊂ R m , if ϕ θ ( g ) is continuous at each θ ∈ Θ for almost all g and is dominated by a function χ as | ϕ θ ( g ) | ≤ χ ( g ) for all θ ∈ Θ with E g [ χ ( g )] < ∞ , then the following holds:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Detailed proof provided in the Appendix D
Revised: Detailed proof provided in the Appendix D

Original: Without the goal discretization, we incur an extra cost of ω ( θ ) , which is expected to be strictly positive since 1 |I k | (cid:3) i ∈I k ϕ θ ( g i ) is maximized during training while E g ∼ ρ g [ ϕ θ ( g ) | g ∈ G k ] is not.
Revised: Without the goal discretization, we incur an extra cost of ω ( θ ) , which is expected to be strictly positive since 1 |I k | " i ∈ I k ϕ θ ( g i ) is maximized during training while E g ∼ ρ g [ ϕ θ ( g ) | g ∈ G k ] is not.

Original: Thus, the goal discretization can improve the expected sum of rewards for unseen goals by the degree of ω ( θ ) , which measures the concentration of the goal distribution in each neighborhood.
Revised: Thus, the goal discretization can improve the expected sum of rewards for unseen goals by the degree of ω ( θ ) , which measures the concentration of the goal distribution in each neighborhood.

Original: This extra cost ω ( θ ) goes to zero when the number of goal observations n approaches inﬁnity.
Revised: This extra cost ω ( θ ) goes to zero when the number of goal observations n approaches inﬁnity.

Original:  Related Work
Revised:  Related Work

Original: Learning with multiple hierarchies has long been proposed in the RL literature, where goal conditioned HRL implements high level planning and low level control using sub-goals.
Revised: Learning with multiple hierarchies has long been proposed in the RL literature, where goal conditioned HRL implements high level planning and low level control using sub-goals.

Original: Often in goal conditioned HRL, the higher level policy speciﬁes goals which may not have good speciﬁcation and grounding .
Revised: Often in goal conditioned HRL, the higher level policy speciﬁes goals which may not have good speciﬁcation and grounding .

Original: Several prior works focus on goal-conditioned RL to improve sample efﬁciency in deep
Revised: Several prior works focus on goal-conditioned RL to improve sample efﬁciency in deep

Original: RL tasks [38, 40].
Revised: RL tasks [38, 40].

Original: These build on ideas that were proposed years back to solve long horizon tasks by hierarchical RL specifying goals [23, 13, 14, 68].
Revised: These build on ideas that were proposed years back to solve long horizon tasks by hierarchical RL specifying goals [23, 13, 14, 68].

Original: The goal is to learn to solve sub-goals provided to the policy, by learning to predict a sequence of actions that can reach each of the sub-goals [66, 49, 41].
Revised: The goal is to learn to solve sub-goals provided to the policy, by learning to predict a sequence of actions that can reach each of the sub-goals [66, 49, 41].

Original: Additionally, in existing HRL literature, distance measures are often used based on goal-conditioned value functions, allowing to measure distances between states and the sequence of sub-goals to reach [16, 71], for planning [42], or exploration [29].
Revised: Additionally, in existing HRL literature, distance measures are often used based on goal-conditioned value functions, allowing to measure distances between states and the sequence of sub-goals to reach [16, 71], for planning [42], or exploration [29].

Original: Since the set of goals speciﬁed in the state space can be arbitrary, an additional constraint is often also learnt to tie the distribution of selected goals to those the lower level policy can reach [72].
Revised: Since the set of goals speciﬁed in the state space can be arbitrary, an additional constraint is often also learnt to tie the distribution of selected goals to those the lower level policy can reach [72].

Original: We tackle this problem by proposing
Revised: We tackle this problem by proposing

Original: DGRL, for better grounding and speciﬁcation of sub-goal representations.
Revised: DGRL, for better grounding and speciﬁcation of sub-goal representations.

Original: In previous works, mutual information based objectives have been proposed for goal conditioned
Revised: In previous works, mutual information based objectives have been proposed for goal conditioned

Original: RL.
Revised: RL.

Original: They perform goal-based representation learning, in order to improve stability of training goalconditioned value functions [38, 40], or to provide goal representations allowing to identify decision states for better exploration [18].
Revised: They perform goal-based representation learning, in order to improve stability of training goalconditioned value functions [38, 40], or to provide goal representations allowing to identify decision states for better exploration [18].

Original: However, for most of these settings, the sub-goals are based on an external reward and are lacking in terms of speciﬁcation, which can lead to inefﬁcient training.
Revised: However, for most of these settings, the sub-goals are based on an external reward and are lacking in terms of speciﬁcation, which can lead to inefﬁcient training.

Original: [32, 72] have proposed approaches that penalize the high level controller for generating sub-goals that are too difﬁcult for the lower level policies, through the use of additional constrained objectives [72].
Revised: [32, 72] have proposed approaches that penalize the high level controller for generating sub-goals that are too difﬁcult for the lower level policies, through the use of additional constrained objectives [72].

Original: We highlight that our proposed DGRL can be generically applied to any goal conditioned
Revised: We highlight that our proposed DGRL can be generically applied to any goal conditioned

Original: RL literature for better grounding and speciﬁcation of the sub-goals.
Revised: RL literature for better grounding and speciﬁcation of the sub-goals.

Original: Furthermore, recent work has shown signiﬁcance of learning representations through self supervised objectives in RL [3, 54], often as a pre-training phase [55, 69], which can help for both exploration [35] and control [70].
Revised: Furthermore, recent work has shown signiﬁcance of learning representations through self supervised objectives in RL [3, 54], often as a pre-training phase [55, 69], which can help for both exploration [35] and control [70].

Original: In the context of goal conditioned RL, it can be a challenging problem as it additionally requires learning reliable representations of goals in parallel, purely from high dimensional observations [15].
Revised: In the context of goal conditioned RL, it can be a challenging problem as it additionally requires learning reliable representations of goals in parallel, purely from high dimensional observations [15].

Original: Previous works [32] have often used the entire observation space as goals, which is not scalable for complex tasks.
Revised: Previous works [32] have often used the entire observation space as goals, which is not scalable for complex tasks.

Original: Other works have used a pre-deﬁned space of sub-goals as domain knowledge [38], or self-play for sub-goal representations [59] to reduce the complexity of goal space design.
Revised: Other works have used a pre-deﬁned space of sub-goals as domain knowledge [38], or self-play for sub-goal representations [59] to reduce the complexity of goal space design.

Original: Most recently, [42, 41] utilized unsupervised representation learning to learn a goal representation space, which can further be used for planning and control.
Revised: Most recently, [42, 41] utilized unsupervised representation learning to learn a goal representation space, which can further be used for planning and control.

Original: In this work, we show additionally that using a bottleneck can further lead to factorial representation of goals, while helping with goal speciﬁcation via learning a latent space of discrete goals usable for planning and control.
Revised: In this work, we show additionally that using a bottleneck can further lead to factorial representation of goals, while helping with goal speciﬁcation via learning a latent space of discrete goals usable for planning and control.

Original: We emphasize that
Revised: We emphasize that

Original: DGRL can be integrated on any existing goal conditioned approach that utilizes learning a sub-goal representation.
Revised: DGRL can be integrated on any existing goal conditioned approach that utilizes learning a sub-goal representation.

Original:  Experiments
Revised:  Experiments

Original: The main goal of our experiments is to show that goal discretization can lead to sample efﬁcient learning and generalization to novel goals, in goal-conditioned RL.
Revised: The main goal of our experiments is to show that goal discretization can lead to sample efﬁcient learning and generalization to novel goals, in goal-conditioned RL.

Original: First, we directly study this by training on environments with a set of goals (such as 8 positions within a gridworld) and then evaluating the agent’s ability to reach a position within the gridworld which it was not trained to reach.
Revised: First, we directly study this by training on environments with a set of goals (such as 8 positions within a gridworld) and then evaluating the agent’s ability to reach a position within the gridworld which it was not trained to reach.

Original: Second, we consider hierarchical goal-conditioned RL, in which a higher-level agent generates goals that a lower-level agent is tasked with reaching.
Revised: Second, we consider hierarchical goal-conditioned RL, in which a higher-level agent generates goals that a lower-level agent is tasked with reaching.

Original: In this case, the task of reaching novel goals occurs organically as the higher-level model selects new goals.
Revised: In this case, the task of reaching novel goals occurs organically as the higher-level model selects new goals.

Original: This setup also shows the advantages of
Revised: This setup also shows the advantages of

Original: DGRL for goal speciﬁcation .
Revised: DGRL for goal speciﬁcation .

Original: A secondary goal of our experiments is to show that using many discrete factors is often critical for optimal performance, which proves the value of factorization in grounding goals .
Revised: A secondary goal of our experiments is to show that using many discrete factors is often critical for optimal performance, which proves the value of factorization in grounding goals .

Original: We evaluate our proposed method DGRL by integrating it into existing state-of-the-art goalconditioned and hierarchical RL tasks.
Revised: We evaluate our proposed method DGRL by integrating it into existing state-of-the-art goalconditioned and hierarchical RL tasks.

Original: Experimentally, we analyse DGRL on several challeng- ing testbeds that have previously been used in the RL community.
Revised: Experimentally, we analyse DGRL on several challeng- ing testbeds that have previously been used in the RL community.

Original: DGRL in principle can be applied to any existing downstream goal-conditioned RL tasks.
Revised: DGRL in principle can be applied to any existing downstream goal-conditioned RL tasks.

Original: We demonstrate improvements on ﬁve such tasks.
Revised: We demonstrate improvements on ﬁve such tasks.

Original: We consider maze navigation where images are used as observations and we show improved generalization to novel goals.
Revised: We consider maze navigation where images are used as observations and we show improved generalization to novel goals.

Original: We integrate DGRL to an existing goal-conditioned baseline for navigating procedurally-generated hard exploration Minigrid environments [10] and ﬁnd that it outperforms state-of-the-art exploration baselines.
Revised: We integrate DGRL to an existing goal-conditioned baseline for navigating procedurally-generated hard exploration Minigrid environments [10] and ﬁnd that it outperforms state-of-the-art exploration baselines.

Original: We also show improvements with DGRL on continuous control (Ant) navigation and manipulation tasks, where goals come from a high-level controller.
Revised: We also show improvements with DGRL on continuous control (Ant) navigation and manipulation tasks, where goals come from a high-level controller.

Original: Finally, we show that discrete representations also signiﬁcantly improve sample efﬁcient learning on a challenging vision-based robotic manipulation environment.
Revised: Finally, we show that discrete representations also signiﬁcantly improve sample efﬁcient learning on a challenging vision-based robotic manipulation environment.

Original: Demonstrating Factorized Representation Learning We ﬁrst pick a color-mnist supervised learning example to support the idea that DGRL can learn factorized or compositional representations.
Revised: Demonstrating Factorized Representation Learning We ﬁrst pick a color-mnist supervised learning example to support the idea that DGRL can learn factorized or compositional representations.

Original: Figure 3 displays reconstructed images from a trained decoder operating on a discretized 2-factor
Revised: Figure 3 displays reconstructed images from a trained decoder operating on a discretized 2-factor

Original: [Figure]
Revised: [Figure]

Original: representation.
Revised: representation.

Original: We ﬁnd that different factors capture information of different semantic nature.
Revised: We ﬁnd that different factors capture information of different semantic nature.

Original: More precisely, factor 1 tends to encode the shape of the digit, while factor 2 specialized in its color.
Revised: More precisely, factor 1 tends to encode the shape of the digit, while factor 2 specialized in its color.

Original: This empirically suggests the emergence of “factorization” in the learnt representations.
Revised: This empirically suggests the emergence of “factorization” in the learnt representations.

Original: Further experimental details are provided in section C.1.
Revised: Further experimental details are provided in section C.1.

Original: Learning to Reach Diverse and Novel Goals.
Revised: Learning to Reach Diverse and Novel Goals.

Original: We study a gridworld navigation task in which an agent is trained to reach a goal from a small ﬁnite set of training goals, and during evaluation is tasked with reaching a novel goal unseen during training.
Revised: We study a gridworld navigation task in which an agent is trained to reach a goal from a small ﬁnite set of training goals, and during evaluation is tasked with reaching a novel goal unseen during training.

Original: This is a navigation task with a pixel-level observation space showing the position of the agent and the goal in a gridworld.
Revised: This is a navigation task with a pixel-level observation space showing the position of the agent and the goal in a gridworld.

Original: We consider two mazes spiral and single-loop topology.
Revised: We consider two mazes spiral and single-loop topology.

Original: Experiment setup is given in Appendix C.2.
Revised: Experiment setup is given in Appendix C.2.

Original: For this task, we train a goal-conditioned Deep Q-Learning (DQN) agent, and use a pre-trained representation φ ( · ) where the encoder is trained using data from a random rollout policy.
Revised: For this task, we train a goal-conditioned Deep Q-Learning (DQN) agent, and use a pre-trained representation φ ( · ) where the encoder is trained using data from a random rollout policy.

Original: Because the gridworld is small the random rollout policy achieves good coverage of the state space, so we found this was sufﬁcient for learning a good goal representation.
Revised: Because the gridworld is small the random rollout policy achieves good coverage of the state space, so we found this was sufﬁcient for learning a good goal representation.

Original: At each episode, a speciﬁc goal is randomly sampled from a distribution of goals, and the DQN agent is trained to reach the speciﬁed goal for that episode.
Revised: At each episode, a speciﬁc goal is randomly sampled from a distribution of goals, and the DQN agent is trained to reach the speciﬁed goal for that episode.

Original: During evaluation, we test the learned agent on goals either from the training distribution, or not seen during training.
Revised: During evaluation, we test the learned agent on goals either from the training distribution, or not seen during training.

Original: Furthermore, for this task, we additionally use an intrinsic reward to promote exploration of the goal-DQN agent.
Revised: Furthermore, for this task, we additionally use an intrinsic reward to promote exploration of the goal-DQN agent.

Original: Since we learn a discrete factorial representation of the goal, we compute an exploration bonus based on the discrete latent codebooks; i.e. , we embed the states and goals using the learned codes and then compute an intrinsic exploration bonus based on the fraction of learned factors that match.
Revised: Since we learn a discrete factorial representation of the goal, we compute an exploration bonus based on the discrete latent codebooks; i.e. , we embed the states and goals using the learned codes and then compute an intrinsic exploration bonus based on the fraction of learned factors that match.

Original: For the baseline goal-DQN agent, we provide an additional reward bonus based on the cosine distance between continuous embeddings of the state observation and goal.
Revised: For the baseline goal-DQN agent, we provide an additional reward bonus based on the cosine distance between continuous embeddings of the state observation and goal.

Original: Figureshows that DGRL signiﬁcantly outperforms a continuous baseline goal DQN agent, when trained on either four goals or eight goals.
Revised: Figureshows that DGRL signiﬁcantly outperforms a continuous baseline goal DQN agent, when trained on either four goals or eight goals.

Original: We evaluate generalization to 4 novel goals unseen during training (Figure 5) and demonstrate improved generalization.
Revised: We evaluate generalization to 4 novel goals unseen during training (Figure 5) and demonstrate improved generalization.

Original: In the previous experiment we evaluated the generalization ability of DGRL by showing that learning discrete factorial representations of goals can improve generalization to novel goals.
Revised: In the previous experiment we evaluated the generalization ability of DGRL by showing that learning discrete factorial representations of goals can improve generalization to novel goals.

Original: Now, we consider various setups in which a goal generating agent speciﬁes goals using the learned codebook and a goal-conditioned agent is tasked with reaching the goals speciﬁed by the goal generating agent.
Revised: Now, we consider various setups in which a goal generating agent speciﬁes goals using the learned codebook and a goal-conditioned agent is tasked with reaching the goals speciﬁed by the goal generating agent.

Original: We test various settings, where the goal generating agents is parameterized as an adversarial teacher [6] or as a higher-level policy in the case of hierarchical RL.
Revised: We test various settings, where the goal generating agents is parameterized as an adversarial teacher [6] or as a higher-level policy in the case of hierarchical RL.

Original: Procedurally Generated MiniGrid Exploration Task.
Revised: Procedurally Generated MiniGrid Exploration Task.

Original: We follow the experimental setup of [6] and [46] and evaluate DGRL on procedurally generated MiniGrid environments [10].
Revised: We follow the experimental setup of [6] and [46] and evaluate DGRL on procedurally generated MiniGrid environments [10].

Original: In [6], a goal-generating teacher proposes goals to train a goal-conditioned “student” policy.
Revised: In [6], a goal-generating teacher proposes goals to train a goal-conditioned “student” policy.

Original: We integrate
Revised: We integrate

Original: DGRL on top of AMIGO [6] and compare DGRL on a hard exploration task with state-of-the-art exploration baselines.
Revised: DGRL on top of AMIGO [6] and compare DGRL on a hard exploration task with state-of-the-art exploration baselines.

Original: Experimental results are summarized in Table 1 and more details provided in
Revised: Experimental results are summarized in Table 1 and more details provided in

Original: Appendix C.4.
Revised: Appendix C.4.

Original: Note that unlike RIDE and RND, we do not provide an additional exploration bonus to DGRL, and ﬁnd that DGRL can still solve this hard exploration task more efﬁciently.
Revised: Note that unlike RIDE and RND, we do not provide an additional exploration bonus to DGRL, and ﬁnd that DGRL can still solve this hard exploration task more efﬁciently.

Original: Goal Grounding in KeyChest Maze Navigation Domain.
Revised: Goal Grounding in KeyChest Maze Navigation Domain.

Original: We consider a simple discrete state action KeyChest maze navigation task, following [72], where discrete goals in the state space are provided by a higher level policy.
Revised: We consider a simple discrete state action KeyChest maze navigation task, following [72], where discrete goals in the state space are provided by a higher level policy.

Original: For this task, to integrate DGRL, we learn an embedding φ ( · ) of the goals, then discretize the representation with a learned codebook.
Revised: For this task, to integrate DGRL, we learn an embedding φ ( · ) of the goals, then discretize the representation with a learned codebook.

Original: We compare with a baseline
Revised: We compare with a baseline

Original: HRAC [72] agent (details in Appendix C.3).
Revised: HRAC [72] agent (details in Appendix C.3).

Original: Figure 7 shows an illustration of the KeyChest environment and a performance comparison of DGRL with different group factors G .
Revised: Figure 7 shows an illustration of the KeyChest environment and a performance comparison of DGRL with different group factors G .

Original: Using fewer factors ( G = 4 ) performs worse than the HRAC baseline, whereas using a larger number of factors ( G =or G = 16 ) improves the sample efﬁciency of the goal reaching agent, providing evidence for the beneﬁts of factorization. 
Revised: Using fewer factors ( G = 4 ) performs worse than the HRAC baseline, whereas using a larger number of factors ( G =or G = 16 ) improves the sample efﬁciency of the goal reaching agent, providing evidence for the beneﬁts of factorization. [Figure] [Table] [Figure]

Revised: [Table]

Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: Ant Manipulation Control Domains.
Revised: Ant Manipulation Control Domains.

Original: We employed DGRL on three different continuous control tasks: AntMazeSparse, AntFall and AntPush.
Revised: We employed DGRL on three different continuous control tasks: AntMazeSparse, AntFall and AntPush.

Original: We emphasize that these tasks are the more challenging counterparts of AntGather and AntMaze tasks, typically used in the hierarchical RL community [38, 40].
Revised: We emphasize that these tasks are the more challenging counterparts of AntGather and AntMaze tasks, typically used in the hierarchical RL community [38, 40].

Original: Figure 8 provides an illustration.
Revised: Figure 8 provides an illustration.

Original: We evaluate goal discretization by integrating DGRL to the state-of-the-art HRAC baseline.
Revised: We evaluate goal discretization by integrating DGRL to the state-of-the-art HRAC baseline.

Original: Details of the experimental setup are provided in Appendix 6.
Revised: Details of the experimental setup are provided in Appendix 6.

Original: Figure 8 shows that specifying the goals using the learned codebook helps DGRL achieve a higher success rate compared to the HRAC baseline.
Revised: Figure 8 shows that specifying the goals using the learned codebook helps DGRL achieve a higher success rate compared to the HRAC baseline.

Original: Ant Navigation Maze Tasks.
Revised: Ant Navigation Maze Tasks.

Original: We consider Ant navigation tasks that require extended temporal reasoning, following the setup in Reinforcement learning with Imagined Subgoals [7, RIS]: a U-shaped maze, and an S-shaped maze (the S-shaped maze is shown in Figure 9).
Revised: We consider Ant navigation tasks that require extended temporal reasoning, following the setup in Reinforcement learning with Imagined Subgoals [7, RIS]: a U-shaped maze, and an S-shaped maze (the S-shaped maze is shown in Figure 9).

Original: The ant navigating in the maze is trained to reach any goal in the environment.
Revised: The ant navigating in the maze is trained to reach any goal in the environment.

Original: The agent is evaluated for generalization in an extended temporal setting with a difﬁcult conﬁguration, we compare the success rate of DGRL integrated on top of RIS with several baselines.
Revised: The agent is evaluated for generalization in an extended temporal setting with a difﬁcult conﬁguration, we compare the success rate of DGRL integrated on top of RIS with several baselines.

Original: We emphasize the difﬁculty of these tasks, where existing baselines like soft actor critic [20, SAC] and temporal difference models [45, TDM] fail completely.
Revised: We emphasize the difﬁculty of these tasks, where existing baselines like soft actor critic [20, SAC] and temporal difference models [45, TDM] fail completely.

Original: Results in Figure 9 show that DGRL improves the sample efﬁciency over the RIS baseline.
Revised: Results in Figure 9 show that DGRL improves the sample efﬁciency over the RIS baseline.

Original: Additional experimental setup and environment conﬁgurations are provided in Appendix C.6.
Revised: Additional experimental setup and environment conﬁgurations are provided in Appendix C.6.

Original: Vision Based Robotic Manipulation.
Revised: Vision Based Robotic Manipulation.

Original: Finally, we assess DGRL on a hard vision-based robotic manipulation task, and use the same setup as in Section 6 to integrate DGRL with the state-ofthe-art RIS baseline on the Sawyer task in Figure 10.
Revised: Finally we assess DGRL on a hard vision-based robotic manipulation task, and use the same setup as in Section 6 to integrate DGRL with the state-ofthe-art RIS baseline on the Sawyer task in Figure 10.

Original: This manipulation task is adapted from [42], where the baseline RIS is already shown to be superior to previous goal conditioning methods.
Revised: This manipulation task is adapted from [42], where the baseline RIS is already shown to be superior to previous goal conditioning methods.

Original: The task of the agent is to control a 2-DoF robotic arm from image input and move a puck positioned on the table.
Revised: The task of the agent is to control a 2-DoF robotic arm from image input and move a puck positioned on the table.

Original: The Sawyer task is designed for training and generalization.
Revised: The Sawyer task is designed for training and generalization.

Original: At test time, it evaluates the agent’s success at placing the puck in desired positions in a temporally extended conﬁguration.
Revised: At test time, it evaluates the agent’s success at placing the puck in desired positions in a temporally extended conﬁguration.

Original: This is  challenging vision-based complex motor task time  requires  extended reasoning.
Revised: This is a challenging vision-based complex motor task, since test time generalization requires temporally extended reasoning.

Original: Results in Figure 10 show that DGRL improves the sample efﬁciency over the
Revised: Results in Figure 10 show that DGRL improves the sample efﬁciency over the

Original: RIS baseline.
Revised: RIS baseline.

Original: Details of the experimental setup are DGRL with
Revised: Details of the experimental setup are provided in Appendix C.

Original: RIS (with a larger number of factors, G= and
Revised: RIS (with a larger number of factors, G= and

Original: G=16, while
Revised: G=16, while

Original: G=4 fails), improvesover the RIS baseline
Revised: G=4 fails), improves RIS baseline

Original:  Discussion 
Revised:  Discussion Conclusion

Original: Conclusion

Original: Our work provides direct evidence that performance of goal-conditioned RL can be improved when the representations of the goals are both discrete and factorial .
Revised: Our work provides direct evidence that performance of goal-conditioned RL can be improved when the representations of the goals are both discrete and factorial .

Original: We show that an instantiation of this idea using multi-factor discretization signiﬁcantly improves performance on a diverse set of benchmarks.
Revised: We show that an instantiation of this idea using multi-factor discretization signiﬁcantly improves performance on a diverse set of benchmarks.

Original: Limitations and Future Work
Revised: Limitations and Future Work

Original: An interesting question that arises from our work is how to theoretically ground and specify goals, which might be helpful for efﬁcient structured exploration in tasks where goal seeking is crucial.
Revised: An interesting question that arises from our work is how to theoretically ground and specify goals, which might be helpful for efﬁcient structured exploration in tasks where goal seeking is crucial.

Original: Additionally, while we demonstrate that the factorial representations learnt by DGRL can be beneﬁcial, it would be interesting to explore whether we can also enforce compositionality in the latent embeddings.
Revised: Additionally, while we demonstrate that the factorial representations learnt by DGRL can be beneﬁcial, it would be interesting to explore whether we can also enforce compositionality in the latent embeddings.

Original: Finally, exploring in more details the structure of the discrete factors forming the goals, their coverage, and to which extent they semantically capture the underlying factors of the environment is a promising research avenue.
Revised: Finally, exploring in more details the structure of the discrete factors forming the goals, their coverage, and to which extent they semantically capture the underlying factors of the environment is a promising research avenue.

Original: Acknowledgement
Revised: Acknowledgement

Original: The authors would like to thank Nicolas Heess, John Langford, Yonathan Efroni, Manan Tomar,
Revised: The authors would like to thank Nicolas Heess, John Langford, Yonathan Efroni, Manan Tomar,

Original: Dipendra Misra, Akshay Krishnamurthy, Pierre-Yves Oudeyer, Harm Van Seijen and Doina Precup for valuable discussions and insightful comments related to this work.
Revised: Dipendra Misra, Akshay Krishnamurthy, Pierre-Yves Oudeyer, Harm Van Seijen and Doina Precup for valuable discussions and insightful comments related to this work.

Original: Hongyu Zang and Xin Li were partially supported by NSFC under Grant 62276024.
Revised: Hongyu Zang and Xin Li were partially supported by NSFC under Grant 62276024.