Original: Discrete Factorial Representations as an Abstraction for Goal-Conditioned Reinforcement Learning Anonymous Author(s)
Revised: Discrete Factorial Representations as an Abstraction for Goal Conditioned RL

Original: AffiliationAddress

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 ofobjectives.
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 bothreliably 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: Defining goals in the space of noisy andhigh-dimensional sensory inputs poses a challenge for training goal-conditionedagents, 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 bylearning factorial representations of goals and processing the resulting representation via a discretization bottleneck, for coarser specification of goals, through anapproach 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 evaluatingthis method on tasks ranging from maze environments to complex robotic navigation and manipulation tasks.
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-boundingthe expected return on out-of-distribution goals, while still allowing for specifyinggoals 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, 57] focusing on how toselect 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 howto control the desired behavior of an RL agent during both  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 wayto control this behavior is by specifying a reward signal [49, 53].
Revised: One way to control this behavior is by specifying a reward signal [52, 56].

Original: While this approach is verygeneral, the reward signal can be hard to design and may not be the most informative form offeedback.
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 difficult when the reward signal issparse [59, 60, 33, 34, 54], such as policy gradients becoming nearly flat in regions where reward isalmost 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 ahigh reward rather than learning a diverse set of skills for coping with novel challenges [19].
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 flexibly specify and ground the desired behavior of RL agents is by trainingagents that receive a reward when they reach a goal specified explicitly to them [22].
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, andis given a reward only when it reaches the goalthat it was instructed to reach [58, 45, 41].
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: Thisprovides a richer signal for the agent than simply collecting more samples oriented around a singlegoal, 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: Italso allows for more flexible and tightly constrained control over the desired behavior of a learnedagent [11, 4, 27].
Revised: It also allows for more ﬂexible and tightly constrained control over its desired behavior [12, 5, 27].

Original: Moreover, the diversity of goals seen during training should help improve bothcredit assignment and generalization [45, 40].
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 [7, 1] andgoal specification [3].
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 defining the goal space and goal specification refersto 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 rewardedin goal-conditioned RL when it reaches the rewardwhich it was specified to reach, whereas in goalfree RL a reward is provided regardless of any such specification, which makes the nature of theagent’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 goalschallenging?
Revised: What makes grounding and specifying goals challenging?

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

Original: For example, we maywant it to pick up a red apple or a greenpear (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 numberof possible goals of interests may be fairlysmall, such as the set of all valid combinations of fruits and their colors, while thenumber of possible observations of goals isextremely 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 ofrelating high-dimensional observations andthe space of relevant goals.
Revised: Goal Grounding refers to this challenge of relating high-dimensional observations and the space of relevant goals.

Original: Goal Specification refers to picking a suitable goal for the agent toreach 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 goalsreachable in the agent’s current context [36, 26].
Revised: It also implies specifying goals reachable in the agent’s current context [39, 26].

Original: Goal specification could be done either manually by a developer or by another RL agent, such as a high-level agent which generates goals thata lower-level agent  tries to reach [11, 4, 27, 21].
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 specified in language are an excellent fitfor these desiderata, as language is a compressed discrete representation which is useful for out-ofdistribution generalization, while being compositional and expressive [18, 10, 21, 16].
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 sametime, connecting language feedback for an agent is non-trivial (requiring special assumptions or alabeling framework) [8].
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 by itself, or trainedjointly with the downstream RL objective) while forcing the goal representations to be discrete andfactorial .
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 [61, 44, 30] which discretizes acontinuous 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 proposedhere (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 codesfrom a learned dictionary, it makes it easier to ground unseen goals (i.e., goals not seen during training) using (novel) compositions of the  discrete codes learned during the training process.
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 showempirically that this improves  generalization performance of goal-reaching policies to unseen goalswhile 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 toan agent (i.e.. a lower-levelpolicy) to complete the task (i.e., reach the  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 specification of semantically meaningfulsub-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 semanticallymeaningful 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 lie in the state space S , i.e., G = S (or in observation space O ).
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 goalconditioned policy as π ( a | s, g ) (either stochastic or deterministic), and its expected total return as J ( π ) = 
Revised: 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)

Original: E (cid:104) (cid:80) Tt ={+=
+}
Revised: Tt {+=
+}

Original: 0 R ( s t , g, a ) (cid:105) where the goal g is either sampled from a distribution ρ g or providedby another higher level policy π hθ h ( g | s ) .
Revised: 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 ) .

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 ) =

Original: E (cid:104) (cid:80) Tt ={+=
+}
Revised: E (cid:2) (cid:3) Tt {+=
+}

Original: 0 R ( s t , g, a ) | s 0 = s ; π (cid:105) .
Revised: R ( s t , g, a ) | s 0 = s ; π (cid:4) .

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

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

Original: We consider goal-conditioned settings in which the goalsare specified 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, goalsare 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 timescale 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: Thelower level policy executes primitive actions π lθ l ( a | s t , g t ) to reach the goals specified by the highlevel 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: Thehigher level policy is trained to maximize the external reward i.e., the reward function specified bythe MDP.
Revised: The higher level policy is trained to maximize the external reward i.e. , the reward function speciﬁed by the MDP.

Original: Both the higher and lower level policies can be trained with any standard RL algorithms,such as Deep Q-Learning (DQN) [33] or policy optimization based algorithms [47, 48].
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 theenvironment 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 thesampled 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 thedistribution ρ 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: In this work, we consider both the HRL and goal-conditioned setups, and evaluate the significanceof learninga factorial representation of discrete latent goals in a series of complex goal-conditionedtasks.
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 [61, 44, 30] discretizes the bottleneck representationof 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 anencoder.
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 codebookvector 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 thentasked 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 [28, 55, 50, 31] have demonstratedthe benefits of using a pre-training stage where the representations of raw observations are learnedusing 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 fine-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-supervisedrepresentations have been shown to improve sample efficiency.
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 ofthree parts: (a) learning representations of raw visual observations through self-supervised representation objectives, (b) processing the resulting representations via a learned dictionary of discretecodes, and (c) using the resulting discrete representations for downstream goal-conditioned and HRLtasks.
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: With the learnt discrete goal representations, we describe in Section 5 how they can acceleratelearning 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: The goal representation can be learned atthe same time as the downstream-RL objective or pre-trained with self-supervised learning and thenused as a fixed 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: 3.1 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 andthe 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 observationand maps it to a continuous embedding: z e = ϕ ( o t ) .
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-supervisedtechniques for learning representations.
Revised: Here, we explore two different self-supervised techniques for learning representations.

Original: For simpler environments, we use a simple autoencoder withthe 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[31] which optimizes for a contrastive objective as a proxy to maximize the mutual 
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:  representations via 
Revised: Processing continuous representations via a discrete codebook.

Original: We learn discrete representations by using the vector-quantization method from the VQ-VAE paper [61] and follow the multi-factor setup used in Discrete-Value Neural Communication [30].
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 with z e = CONCATENATE ( c 1 , c 2 , . . .
Revised: , c G , where z e = CONCATENATE ( c 1 , c 2 , . . .

Original: , c G ) , where each segment c i ∈ R m/G (such that m is divisble 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 separately 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 concatenatedto obtain the final discretized vector 
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: z q :
Revised: [Equation] [Equation] DISCRETIZE

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

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

Original: [Equation]
Revised: [Equation]

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

Original: Here, sg refers to a stop-gradient operation that blocks gradients from flowing into e o i , and β is a hyperparameter which controls howstrongly 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 [30], we used a moving average toupdate 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 withan 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 fixedhyperparameter 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:80) G i
Revised: The term (cid:3)

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

Original: 2 2is the commitment loss, which only applies to the target segment c i and trains the encoder that outputs c i to make c i stay close to the chosen discrete latent vector e o i .
Revised: || 22 is often called the commitment loss.

Original: We trained the VQ-quantizationprocess 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 evectors to discretize in a model, the mean of thecodebook and commitment loss across all z e vectorswas 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 ) , whereq ( · )is the whole discretization process using the codebook, L is the codebook size,and Gis the numberof factors per vector.
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 adiscretization bottleneck on the continuous representations resulting from the self-supervised pretraining.
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 try with differentnumber of discrete factors 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: Usingdiscretization with more factors slightly increases computation but reduces the number of modelparameters due to the codebook embeddings being reused across the different factors.3.3 Using representations for downstream RL
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. 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: Defining goals in the space of noisy, high-dimensional sensory inputs posesa challenge for generalization to novel goals because the encoder that maps the goal observations tothe 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 embedthe continuous latent representation into a discrete representation such that the representation of thenovel goal is mapped to the fixed set of latent discrete codes, and hence facilitate generalizationto new combinations of these codes while making it easy for downstream learning to figure outthe meaning of each discrete code.
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. 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 goalembedding we used 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 for training a goal-conditioned policy a t ∼ π lθ l ( a | s t , g t ) or agoal-conditioned 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 episode of training, a goal is sampledfrom the  distributionof goals ρ g  and the agent gets rewarded for reaching the sampled goal.
Revised: At each training episode, a goal is sampled from the goal distribution ρ g , and the agent gets rewarded for reaching it.

Original: Thisreward can either bean extrinsic reward which is defined as part of the environment or an intrinsicreward which is defined as 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 define the intrinsic reward as thefraction of discrete factors which match in the respective representations of the goal observationand 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 to reach goals within thedistribution ρ g , or for its generalization capability toreach 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 representationof goals g by conditioning on the states every K time-steps, it can also output a sub-goal s g byconditioning 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 ofgoal-conditioned HRL relies on the specification of semantically meaningful sub-goals.
Revised: The effectiveness of goal-conditioned HRL relies on the speciﬁcation of semantically meaningful sub-goals.

Original: Learnedcodebooks (section 3.2) consisting of a set of discrete codes can be used by a higher level policy tospecify 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 thegoal specified 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 5, we empirically show the benefits of the proposed approach for training goal-reachingpolicies or goal-conditioned value functions, as well asin a goal-conditioned hierarchical RL setup.
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

Original:  Theoretical Analysis
Revised:  Theoretical Analysis

Original: In this section,the goal 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 ) 
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 ) = (cid:3) k p

Original: = (cid:80) k p ( g | g ∈ G k ) p ( g ∈ G k ) with theneighborhood set {G k } k , it improves the overall performance in p ( g ) by increasing the concentration in p ( g | g ∈ G k ) .
Revised: ( 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 ) .

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

Original: To state our result, we define φ θ ( 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 functionby 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 denote the i -th element of Q (by ordering elements of Q with an arbitrary ordering).

Original: We also define[ 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 lowerbound 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:80) k ∈I Q |I k | (cid:16) 1 |I k | (cid:80) i ∈I k 
Revised: where ω ( θ ) = 1 n (cid:3) k ∈I Q |I k | (cid:7)|I k | (cid:3) i ∈I k ϕ θ ( g i ) − E

Original: φ θ ( g i ) − E g ∼ ρ g
Revised: g ∼ ρ g

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

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 strictlypositive since 1 |I k | (cid:80) 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 | (cid:3) 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 degreeof ω ( θ ) , 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 extracost ω ( θ ) goes to zero when the number of goal observations n approaches infinity.
Revised: This extra cost ω ( θ ) goes to zero when the number of goal observations n approaches inﬁnity.  Related Work

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.

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

Original: Experiments
Revised: Several prior works focus on goal-conditioned RL to improve sample efﬁciency in deep RL tasks [38, 40]. These build on ideas that were proposed years back to solve long horizon tasks by hierarchical RL specifying goals [23, 13, 14, 68]. 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: 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: 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: We tackle this problem by proposing

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

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

Revised: RL.

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].

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.

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].

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

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

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].

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].

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

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.

Revised: Most recently, [42, 41] utilized unsupervised representation learning to learn a goal representation space, which can further be used 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.

Revised: We emphasize that

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

Revised:  Experiments

Original: The main goal of our experiments is to show that goal discretization can lead to sample efficientlearning 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 thisby training on environments with a set of goals (such as 8 positions within a gridworld) and thenevaluating the agent’s ability to reach a position within the gridworld which it was not tasked withreaching during training.
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 higherlevel 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 taskof 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 setupalso shows the advantages of DGRL for goal specification .
Revised: This setup also shows the advantages of DGRL for goal speciﬁcation .

Original: A secondary goal of our experiments isto show that using many discrete factors is often critical for optimal performance, which proves thevalue 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 challenging goal-conditioned testbeds that have previously been used in thegoal-conditioned 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, wedemonstrate improvements on five distinct goal-conditioned RL tasks.
Revised: DGRL in principle can be applied to any existing downstream goal-conditioned RL tasks. We demonstrate improvements on ﬁve such tasks.

Original: We consider maze navigationtasks 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-generatedhard exploration Minigrid environments [9] and find that it outperforms state-of-the-art explorationbaselines.
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 andmanipulation 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 discreterepresentations also significantly improve sample efficient learning on a challenging vision-basedrobotic 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.

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.

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

Original: Learning to Reach Diverse and Novel Goals.
Revised: [Figure] representation. We ﬁnd that different factors capture information of different semantic nature. More precisely, factor 1 tends to encode the shape of the digit, while factor 2 specialized in its color.

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

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

Original:  We study a gridworld navigation task in which anagent is trained to reach a goal from a small finite set of training goals, and during evaluation istasked with reaching a novel goal unseen during training.
Revised: Learning to Reach Diverse and Novel Goals. 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-levelobservation 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.

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

Original:  For this task, we train a goal-conditioned Deep Q-Learning (DQN) agent, and use a pre-trainedrepresentation ϕ ( · ) where the encoder is trained using data from a random rollout policy.
Revised: Experiment setup is given in Appendix C.2. 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: Becausethe gridworld is small the random rollout policy achieves good coverage of the state space, so wefound this was sufficient 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 specific goal israndomly sampled from a distribution of goals, and the DQN agent is trained to reach the specifiedgoal 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 trainingdistribution, 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 for exploration of the goal-DQNagent.
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 explorationbonus based on the discrete latent codebooks; i.e., we embed the states and goals using the learnedcodes and then compute an intrinsic exploration bonus based on the fraction of learned factors thatmatch.
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 thecosine 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: Figure 3 showsthat DGRL significantly outperform a continuous baseline goal DQN agent, when trained on eitherfour 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 (Figure4) and demonstrate improved generalization to novel goals. [Figure]
Revised: We evaluate generalization to 4 novel goals unseen during training (Figure 5) and demonstrate improved generalization.

Original: u [Figure] [Figure] 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, weconsider various setups in which a goal generating agent specifies goals using the learned codebookand the goal-conditioned agent is tasked with reaching the goals specified by the goal generatingagent.
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 adversarialteacher [5], 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 [5]and [43] and evaluate DGRL on procedurally generated MiniGrid environments [9].
Revised: We follow the experimental setup of [6] and [46] and evaluate DGRL on procedurally generated MiniGrid environments [10].

Original: In [5], a goalgenerating 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 integrateDGRL
Revised: We integrate

Original: Model

Original: KCmedium [Table] [Figure] on top of AMIGO [5] 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.3.
Revised: Appendix C.4.

Original: Note that unlike RIDE and RND, we do not provide an additional exploration bonusto DGRL, and find that DGRL can still solve this hard exploration task more efficiently.
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 stateaction KeyChest maze navigation task, following [68], where discrete goals in the state space areprovided 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 ϕ ( · ) ofthe 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 [68] agent (details in Appendix C.2).
Revised: HRAC [72] agent (details in Appendix C.3).

Original: Figure 6 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 = 8or G = 16 ) improves the sample efficiency of the goal reaching agent, providing evidence for thebenefits 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. [Table]

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

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

Original: We emphasize that these tasks are the morechallenging counterparts of AntGather and AntMaze tasks, typically used in the hierarchical RLcommunity [35, 37].
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 7 provides illustration of these tasks.
Revised: Figure 8 provides an illustration.

Original: We evaluate goal discretizationby 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 areprovided in Appendix 5.
Revised: Details of the experimental setup are provided in Appendix 6.

Original: Figure 7 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 [6, RIS]: a U-shapedmaze, and an S-shaped maze (the S-shaped maze is shown in Figure 8).
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 themaze 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 inan extended temporal setting with a difficult configuration, we compare the success rate of DGRLintegrated 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 difficulty of these tasks, whereexisting baselines like soft actor critic [19, SAC] and temporal difference models [42, TDM] failcompletely.
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 8 show that DGRL improves the sample efficiency 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 configurations are provided in Appendix C.5. [Figure] [Figure] shaped
Revised: Additional experimental setup and environment conﬁgurations are provided in Appendix C.6.

Original: -shaped that integrating DGRL with RIS can lead to more sample efficient convergence on these tasks while baselines such as SAC and TDM (not shown) fail completely on both AntU and AntMaze as reported by [6].

Original: The RIS baseline is based on raw data provided by the authors.

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

Original: Finally, we assess DGRL on a hard vision-based roboticmanipulation task, and use the same setup as in section 5 to integrate DGRL with the state-of-theart RIS baseline on the Sawyer task in Figure 9.
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 [39], wherethe 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 taskof the agent is to control a 2-DoF robotic arm from image input and move a puck positioned onthe 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 theagent’s success at placing the puck in desired positions in a temporally extended configuration.
Revised: At test time, it evaluates the agent’s success at placing the puck in desired positions in a temporally extended conﬁguration.

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

Original: Results in Figure 9 show that DGRL improves the sample efficiency 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 is provided in Appendix C.6. [Figure]
Revised: Details of the experimental setup are

Original: DGRL with
Revised: DGRL with

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), improves over the RIS baseline
Revised: G=4 fails), improves over the RIS baseline

Revised:  Discussion

Original: Conclusion
Revised: Conclusion

Original: Our work provides direct evidence that performance of goal-conditioned RL can be improved whenthe 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 ofthis idea using multi-factor discretization significantly improves performance on a diverse set ofbenchmarks.
Revised: We show that an instantiation of this idea using multi-factor discretization signiﬁcantly improves performance on a diverse set of benchmarks.

Original: An interesting question that arises from our work is how to theoretically ground andspecify goals, which might be helpful for efficient structured exploration in tasks where goal seekingis crucial.

Original: Checklist

Original: The checklist follows the references.

Original: Please read the checklist guidelines carefully for informationon how to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:

Original: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? .•

Original: Did you include the license to the code and datasets?

Original: [No] The code and the data areproprietary.•

Original: Did you include the license to the code and datasets?

Original: [N/A]

Original: Please do not modify the questions and only use the provided macros for your answers.

Original: Note that the

Original: Checklist section does not count towards the page limit. In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below. For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
Revised: Limitations and Future Work

Original: [Yes] .

Original: Main claims are justified with a set of illustrativeexamples and experiments; both standard and toy examples included for experimentsto justify the key claims.(b) Did you describe the limitations of your work? [Yes] (c) Did you discuss any potential negative societal impacts of your work? [N/A]
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: Besides

Original: the general impact of machine learning methods we do not foresee and identify any

Original: particular and immediate negative societal impact of our work.

Original: We do not conduct

Original: any experiments that could be  initself harmful for society but understand that the applicability of this method is very broad. (d) Have you read the ethics review guidelines and ensured that your paper conforms to
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: them?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes]

Original: If you ran experiments... (a) Did you include the code, data, andinstructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] .
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: Code will be

Original: provided along with the supplementary materials for reproducibility.

Original: (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters,

Original: how they

Original: were chosen)?

Original: [Yes] .

Original: More details to be included in supplemetary section

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)?

Original: [Yes] .

Original: All experiments are run across 3 to 5 random seeds.

Original: (d) Did you include the total amount of compute and the type of resources used (e

Original: g.,type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] . This will be mentioned in supplementary and acknowledgements If you are using existing assets (e
Revised: Acknowledgement

Original: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators? [N/A] (b) Did you mention the license of the assets? [N/A]
Revised: The authors would like to thank Nicolas Heess, John Langford, Yonathan Efroni, Manan Tomar,

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N

Original: [Table]

Original: A]

Original: (d) Did you discuss whether and how consent was obtained from people whose data

Original: [Table]

Original: you’re using/curating?

Original: [N/A]

Original: [Table]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifi- [Table] able information or offensive content? [N/A]
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: If you used crowdsourcing or conducted research with human subjects...

Original: [Table]

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: [Table]

Original: applicable?

Original: [N/A]

Original: [Table]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: [Table] Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: Hongyu Zang and Xin Li were partially supported by NSFC under Grant 62276024.

Original: spent on participant compensation?

Original: [N/A]