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

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AffiliationAddress email

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

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 bothreliably 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: 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.

Original: We propose to address this bylearning compositional representations of goals and processing the resulting representation via a discretization bottleneck, for coarser specification of goals, throughan approach we call DGRL.
Revised: 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.

Original: We show that applying a discretizing bottleneck canimprove performance in goal-conditioned RL setups, by experimentally evaluating this 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 evaluatingthis method on tasks ranging from maze environments to complex robotic navigation and manipulation tasks.

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-boundingthe expected return on out-of-distribution goals, while still allowing for specifyinggoals with expressive combinatorial structure.

Original:  Introduction
Revised:  Introduction

Original: Reinforcement Learning is a popular and highly general framework [22, 49] 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, 57] focusing on how toselect 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 [21].
Revised: An important question is howto control the desired behavior of an RL agent during both training and evaluation [24].

Original: One wayto control this behavior is by specifying a reward signal [44, 46].
Revised: One wayto control this behavior is by specifying a reward signal [49, 53].

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 verygeneral, the reward signal can be hard to design and may not be the most informative form offeedback.

Original: The credit assignment problem in RL can become difficult when the reward signal issparse [51, 52, 28, 29, 47], such as policy gradients becoming nearly flat in regions where reward isalmost never achieved.
Revised: 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.

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 [16].
Revised: 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].

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 [19].
Revised: 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].

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 goal that it was instructed to reach [50, 40, 36].
Revised: 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 goal that it was instructed to reach [58, 45, 41].

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

Original: Italso allows for more flexible and tightly constrained control over the desired behavior of a learnedagent [10, 3, 24].
Revised: Italso allows for more flexible and tightly constrained control over the desired behavior of a learnedagent [11, 4, 27].

Original: Moreover, the diversity of goals seen during training should help improve bothcredit assignment and generalization [40, 35].
Revised: Moreover, the diversity of goals seen during training should help improve bothcredit assignment and generalization [45, 40].

Original: While this framework is promising, it introduces two new challenges: goal grounding [6, 1] andgoal specification [2].
Revised: While this framework is promising, it introduces two new challenges: goal grounding [7, 1] andgoal specification [3].

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 defining the goal space and goal specification refersto 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 reward which 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 rewardedin goal-conditioned RL when it reaches the reward which 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.

Original: What makes grounding and specifying goalschallenging?
Revised: What makes grounding and specifying goalschallenging?

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 variousfruits 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 maywant it to pick up a red apple or a greenpear (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 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).

Original: Goal Grounding refers to this challenge ofrelating high-dimensional observations andthe space of relevant goals.
Revised: Goal Grounding refers to this challenge ofrelating high-dimensional observations andthe 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 Specification refers to picking a suitable goal for the agent toreach and computing an appropriate reward when it is reached.

Original: It also implies specifying goalsreachable in the agent’s current context [31, 23].
Revised: It also implies specifying goalsreachable in the agent’s current context [36, 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 [10, 3, 24, 18].
Revised: 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].

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 [15, 9, 18, 14].
Revised: 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].

Original: At the sametime, connecting language feedback for an agent is non-trivial (requiring special assumptions or alabeling framework) [7].
Revised: At the sametime, connecting language feedback for an agent is non-trivial (requiring special assumptions or alabeling framework) [8].

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

Original: To perform this discretization, we use Vector-Quantization [53, 39, 26] which discretizes a continuous representation using a codebook of discrete and learnable codes.
Revised: To perform this discretization, we use Vector-Quantization [61, 44, 30] which discretizes acontinuous representation using a codebook of discrete and learnable codes.

Original: The approachproposed here (called DGRL) serves two complementary purposes.
Revised: The approach proposedhere (called DGRL) serves two complementary purposes.

Original: First, it provides a structuredrepresentation 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 discretecodes from a learned dictionary, it makes it easier to ground unseen goals (i.e., goals not seen duringtraining) using (novel) compositions of the discrete codes learned during the training process.
Revised: 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.

Original: Weshow empirically that this improves generalization performance of goal-reaching policies to unseengoals while remaining expressive enough.
Revised: We showempirically that this improves generalization performance of goal-reaching policies to unseen goalswhile 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 an agent (i.e.. alower-level policy) 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 an agent (i.e.. a lower-levelpolicy) to complete the task (i.e., reach the goal).

Original: In this case goal-inference is learnedend-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 semanticallymeaningful sub-goals.
Revised: The effectiveness of goal-conditioned HRL relies on the specification of semantically meaningfulsub-goals.

Original: Using compositional discrete sub-goals allows the higher-level policy tospecify semantically meaningful objectives to the lower-level policy.
Revised: Using factorial discrete sub-goals allows the higher-level policy to specify semanticallymeaningful 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
Revised: We consider a goal-conditioned Markov Decision Process, where the goals

Original: G lie in the state space S , i.e., G = S (or in observation space O ).
Revised: G lie in the state space S , i.e., G = S (or in observation space O ).

Original: We denote a goalconditioned policy as π ( a | s, g ) (either stochastic or deterministic), and its expected total return as
Revised: We denote a goalconditioned policy as π ( a | s, g ) (either stochastic or deterministic), and its expected total return as

Original: J ( π ) =
Revised: J ( π ) =

Original: E (cid:104) (cid:80) Tt =
Revised: E (cid:104) (cid:80) 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: 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 ) .

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;
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;

Original: V π ( s, g ) =
Revised: V π ( s, g ) =

Original: E (cid:104) (cid:80) Tt =
Revised: E (cid:104) (cid:80) Tt =

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

Original: As in standard RL, the objective in goal-conditioned
Revised: As in standard RL, the objective in goal-conditioned

Original: RL is to maximize the expected discounted returns induced by the goal-conditioned policy.
Revised: 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 goalsare specified 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, goalsare 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 timescale 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: 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.

Original: Thehigher level policy is trained to maximize the external reward i.e., the reward function specified bythe MDP.
Revised: Thehigher level policy is trained to maximize the external reward i.e., the reward function specified bythe MDP.

Original: Both the higher and lower level policies can be trained with any standard RL algorithms,such as Deep Q-Learning (DQN) [28] or policy optimization based algorithms [42, 43].
Revised: 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].

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 theenvironment 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 thesampled 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 thedistribution ρ 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 significance of learning a compositional representation of discrete latent goals in a series of complexgoal-conditioned tasks.
Revised: In this work, we consider both the HRL and goal-conditioned setups, and evaluate the significanceof learning a factorial representation of discrete latent goals in a series of complex goal-conditionedtasks.

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

Original: VQ-VAE [53, 39, 26] discretizes the bottleneck representationof an auto-encoder by adding a codebook of discrete learnable codes.
Revised: VQ-VAE [61, 44, 30] discretizes the bottleneck representationof an auto-encoder by adding a codebook of discrete learnable codes.

Original: The input is passed through anencoder.
Revised: The input is passed through anencoder.

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 codebookvector 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 thentasked with reconstructing the input from this quantized vector.

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

Original: Several papers [25, 48, 45, 27] 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 [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.

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 fine-tuned on) downstream tasks.

Original: These self-supervisedrepresentations have been shown to improve sample efficiency.
Revised: These self-supervisedrepresentations have been shown to improve sample efficiency.

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

Original: With the learnt discrete goal representations, we describe in Section 5 how they can acceleratelearning in complex navigation and manipulation tasks.
Revised: With the learnt discrete goal representations, we describe in Section 5 how they can acceleratelearning 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: 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.

Original: 3.1 Self-Supervised Goal Representation Learning
Revised: 3.1 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 andthe goal observations.

Original: We denote by ϕ the encoder network that takes as input the raw observationand maps it to a continuous embedding: z e = ϕ
Revised: We denote by ϕ the encoder network that takes as input the raw observationand maps it to a continuous embedding: z e = ϕ

Original: ( o t ) .
Revised: ( o t ) .

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

Original: For simpler environments, we use a simple autoencoder withthe reconstruction objective.
Revised: For simpler environments, we use a simple autoencoder withthe reconstruction objective.

Original: For more complex environments, we use the Deep InfoMax approach[27] which optimizes for a contrastive objective as a 
Revised: For more complex environments, we use the Deep InfoMax approach[31] which optimizes for a contrastive objective as a proxy to maximize the mutual

Original: 3.2 Processing continuous representations viaa discrete codebook
Revised: representations via

Original: We quantize the continuous representation into discrete latent variables, following the codebook method from[26], which generalizes the vector quantization used in VQ-VAE [53] to a list of discretized codes instead of a single discretized code.
Revised: 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].

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 with z e = CONCATENATE ( c 1 , c 2 , . . .

Original: , c G ) , where each segment c i ∈ R m/G (such that mG ∈ N + ).
Revised: , c G ) , where each segment c i ∈ R m/G (such that m is divisble by G ).

Original: Each continuous segment c i is mapped separately to a discretized latentspace 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 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):

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 concatenatedto obtain the final discretized vector

Original: z q :
Revised: z q :

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: o i )

Original: || The training procedure closely follows both [26] and [53].
Revised: The training procedure closely follows both [30] and [61].

Original: Here, sg refers to a stop-gradient operation that blocks gradients from flowing into its argument, and β is a hyperparameter which controlsthe reluctance to change the code.
Revised: 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: Unlike [30], we used a moving average toupdate the code embeddings rather than learning them directly as parameters.

Original: We update e o i with an exponential moving average to encourageit to become close to the selected output segment c i .
Revised: We update e o i withan 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 beequal to ηe o i + (1 − η ) c i , where the value of η is a fixed hyperparameter controlling how quicklythe 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 fixedhyperparameter controlling how quickly the moving average updates.

Original: The term (cid:80) Gi
Revised: The term (cid:80) G i

Original: ||
Revised: ||

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

Original: ||
Revised: ||

Original: 22 is the commitment loss, which onlyapplies to the target segment c i and trains the encoder
Revised: 2 2is the commitment loss, which only applies to the target segment c i and trains the encoder

Original: that outputs c i to make c
Revised: that outputs c i to make c

Original: i stay close to thechosen discrete latent vector e
Revised: i stay close to the chosen discrete latent vector e

Original: o i .
Revised: o i .

Original: We trained the codebook together with other parts of the modelby gradient descent.
Revised: We trained the VQ-quantizationprocess 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 thecodebook and commitment loss across all z e vectors was used.
Revised: When there were multiple z evectors to discretize in a model, the mean of the codebook and commitment loss across all z e vectorswas 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 q ( · )is the whole discretization process using the codebook, L is the codebook size, and G is the number of factors per vector.
Revised: ( z e , L, G ) , where q ( · )is the whole discretization process using the codebook, L is the codebook size, and G is the numberof factors per vector.

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 adiscretization bottleneck on the continuous representations resulting from the self-supervised pretraining.

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 try with differentnumber of discrete factors 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: 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

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

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

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 goalembedding we used 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 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 ) .

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

Original: Thisreward can either be an extrinsic reward which is defined as part of the environment or an intrinsicreward which is defined as part of the algorithm.
Revised: Thisreward can either be an extrinsic reward which is defined as part of the environment or an intrinsicreward which is defined as 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 define the intrinsic reward as thefraction of discrete factors which match in the respective representations of the goal observationand the state observation.

Original: At test time, the agent can either be evaluated to reach goals within thedistribution ρ g , or for its generalization capability to reach goals not seen during training.
Revised: At test time, the agent can either be evaluated to reach goals within thedistribution ρ g , or for its generalization capability to reach 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 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 ) .

Original: The effectiveness ofgoal-conditioned HRL relies on the specification of semantically meaningful sub-goals.
Revised: The effectiveness ofgoal-conditioned HRL relies on the specification 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: 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.

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 thegoal specified 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 as in a goal-conditioned hierarchical RL setup.
Revised: In section 5, we empirically show the benefits of the proposed approach for training goal-reachingpolicies 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, 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, 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 )

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: = (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 ) .

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 goalvalues g ∈ G k for each neighborhood G k .

Original: To state our result, we define φ θ ( g ) =
Revised: To state our result, we define φ θ ( g ) =

Original: E s 

Revised: E s

Original: [ V π ( s 0 , g )] ,where θ ∈ R m is the vector containing model parameters learned through  goals observed duringtraining 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 ofgoal g by q ( g ) , and the identity functionby id as id( x ) =x .
Revised: We denote the discretization of g by q ( g ) , and the identity functionby id as id( g ) =

Original:  Let Q = { q ( g ) : g ∈ G} and ˆ d be a distance function.
Revised: g . 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 define[ 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 E) 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 lowerbound 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:80) k ∈I Q |I k | (cid:16) 1 |I k | (cid:80) i ∈I k

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

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

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

Original: ≤ χ ( g ) for all θ ∈ Θ with E g [ χ ( g )] < ∞ , then the following holds:
Revised: ≤ χ ( 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 E
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
Revised: Without the goal discretization, we incur an extra cost of ω ( θ ) , which is expected to be strictlypositive since 1 |I k | (cid:80) i ∈I

Original: k φ θ ( g i ) is maximized during training while E g ∼ ρ g [ φ θ ( g ) | g ∈ G k ] is not.
Revised: 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 degreeof ω ( θ ) , 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 extracost ω ( θ ) goes to zero when the number of goal observations n approaches infinity.

Original:  Experiments
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 efficientlearning 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 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.

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 higherlevel 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 taskof 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 setupalso shows the advantages of DGRL for goal specification .

Original: A secondary goal of our experiments isto show that using many discrete factors is often critical for optimal performance, which proves thevalue of compositionality in grounding goals .
Revised: 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 .

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 the goal-conditioned RL community.
Revised: Experimentally, we analyse DGRL on several challenging goal-conditioned testbeds that have previously been used in the goal-conditioned 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, wedemonstrate improvements on five distinct goal-conditioned RL tasks.

Original: We considere maze navigation tasks where images are used as observations and we show improved generalization to novelgoals.
Revised: We consider maze navigationtasks 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 procedurallygenerated hard exploration Minigrid environments [8] and find that it outperforms state-of-the-artexploration baselines.
Revised: 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.

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 andmanipulation tasks, where goals come from a high-level controller.

Original: Finally, we showthat discrete representations also significantly improve sample efficient learning on a challengingvision-based robotic manipulation environment.
Revised: Finally, we show that discreterepresentations also significantly improve sample efficient learning on a challenging vision-basedrobotic manipulation environment.

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

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-levelobservation space showing the position of the agent and the goal in a gridworld

Original: We consider twomazes spiral and single-loop topology.

Original: Experiment setup given in Appendix C.2. 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: 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.

Original: Becausethe gridworld is small the random rollout policy achieves good coverage of the state space, so we found this was sufficient for learning a good goal representation.
Revised: 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.

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 specific goal israndomly sampled from a distribution of goals, and the DQN agent is trained to reach the specifiedgoal 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 trainingdistribution, 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 for exploration of the goal-DQNagent.

Original: Since we learn a discrete compositional 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 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.

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 thecosine 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: Figure 3 showsthat DGRL significantly outperform a continuous baseline goal DQN agent, when trained on eitherfour goals or eight goals.

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

Original: [Figure]
Revised: [Figure]

Revised: u

Original: [Figure]
Revised: [Figure]

Original: In the previous experiment

Original: [Figure]
Revised: [Figure]

Original:  we evaluated the generalization ability of DGRL by showing that learning discrete compositional 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 specifies goals using the learned codebook and the goal-conditioned agent is tasked with reaching the goals specified by the goal generating agent.
Revised: 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.

Original: We test various settings, where the goal generating agents is parameterized as anadversarial teacher [4], 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 adversarialteacher [5], 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 [4]and [38] and evaluate DGRL on procedurally generated MiniGrid environments [8].
Revised: We follow the experimental setup of [5]and [43] and evaluate DGRL on procedurally generated MiniGrid environments [9].

Original: In [4], a goalgenerating teacher proposes goals to train a goal-conditioned “student” policy. 
Revised: In [5], a goalgenerating teacher proposes goals to train a goal-conditioned “student” policy. We integrate DGRL Model KCmedium

Original: We integrate DGRLon top of AMIGO [4] and compare DGRL on a hard exploration task with state-of-the-art exploration baselines.
Revised: [Table] [Figure] on top of AMIGO [5] 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.3.

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.[Table] [Figure]
Revised: 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.

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 [58], where discrete goals in the state space areprovided by a higher level policy.
Revised: 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.

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 ϕ ( · ) ofthe goals, then discretize the representation with a learned codebook.

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

Original: HRAC [58] agent (details in Appendix C.1).
Revised: HRAC [68] agent (details in Appendix C.2).

Original: Figure 6 shows an illustration of the KeyChest environment and a performance comparison of DGRL with different group factors G .
Revised: Figure 6 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 compositionality.
Revised: 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.

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

Original: We employed DGRL on three different continuous controltasks:
Revised: We employed DGRL on three different continuous controltasks:

Original: AntMazeSparse, AntFall and AntPush task.
Revised: AntMazeSparse, AntFall and AntPush task.

Original: We emphasize that these tasks are the morechallenging counterparts of AntGather and AntMaze tasks, typically used in the hierarchical RLcommunity [30, 32].
Revised: We emphasize that these tasks are the morechallenging counterparts of AntGather and AntMaze tasks, typically used in the hierarchical RLcommunity [35, 37].

Original: Figure 7 provides illustration of these tasks.
Revised: Figure 7 provides illustration of these tasks.

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

Original: Details of the experimental setup areprovided in Appendix.
Revised: Details of the experimental setup areprovided in Appendix 5.

Original: Figure 7 shows that specifying the goals using the learned codebook helps
Revised: Figure 7 shows that specifying the goals using the learned codebook helps

Original: DGRL achieve a higher success rate compared to the HRAC baseline.
Revised: DGRL achieve a higher success rate compared to the HRAC baseline.

Original: [Figure]

Original: 260 Timesteps (x1000)

Original: S u ccess

Original: R a t e ( % )

Original: Success Rate on U-Shaped Ant Environment DGRL, G=

Original: DGRL, G=8 DGRL, G=16 RIS 0.

Original: 0.8 Episodes 1e 0.

Original: S u cc e ss

Original: R a t e ( % )

Original: AntMazeMedTrainEnv-v1 DGRL, G=4 DGRL, G=

Original: DGRL, G=16 RIS

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 [5, 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 [6, RIS]: a U-shapedmaze, and an S-shaped maze (the S-shaped maze is shown in Figure 8).

Original: The ant navigating in themaze is trained to reach any goal in the environment.
Revised: The ant navigating in themaze 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 inan extended temporal setting with a difficult configuration, we compare the success rate of DGRLintegrated on top of RIS with several baselines.

Original: We emphasize the difficulty of these tasks, whereexisting baselines like soft actor critic [16, SAC] and temporal difference models [37, TDM] failcompletely.
Revised: We emphasize the difficulty of these tasks, whereexisting baselines like soft actor critic [19, SAC] and temporal difference models [42, TDM] failcompletely.

Original: Results in Figure 8 show that DGRL improves the sample efficiency over the RIS baseline.
Revised: Results in Figure 8 show that DGRL improves the sample efficiency over the RIS baseline.

Original: Additional experimental setup and environment configurations are provided in Appendix C.4.
Revised: Additional experimental setup and environment configurations are provided in Appendix C.5. [Figure] [Figure] shaped

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

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

Original: This manipulation task is adapted from [34], wherethe baseline RIS is already shown to be superior to previous goal conditioning methods.
Revised: This manipulation task is adapted from [39], wherethe 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 taskof the agent is to control a 2-DoF robotic arm from image input and move a puck positioned onthe 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 theagent’s success at placing the puck in desired positions in a temporally extended configuration.

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

Original: Results in Figure 9 show that DGRL improves the sample efficiency over the
Revised: Results in Figure 9 show that DGRL improves the sample efficiency over the

Original: RIS baseline.
Revised: RIS baseline.

Original: Details of the experimental setup is provided in Appendix C.5.
Revised: Details of the experimental setup is provided in Appendix C.6.

Original: [Figure]
Revised: [Figure]

Revised: DGRL with

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

Revised: G=16, while

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

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 compositional .
Revised: Our work provides direct evidence that performance of goal-conditioned RL can be improved whenthe representations of the goals are both discrete and factorial .

Original: We show that an instantiationof this idea using multi-factorcompositional discretization significantly improves performance on adiverse set of benchmarks.
Revised: We show that an instantiation ofthis idea using multi-factor discretization significantly improves performance on a diverse set ofbenchmarks.

Original: Moreover, DGRL provides especially large improvements in generalization when an agent is tasked with seeking out goals not seen during training.

Original: An interesting questionthat arises from our work is how to theoretically ground and specify goals, which might be helpfulfor efficient structured exploration in tasks where goal seeking is crucial.
Revised: 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
Revised: Checklist

Original: The checklist follows the references.
Revised: The checklist follows the references.

Original: Please read the checklist guidelines carefully for informationon how to answer these questions.
Revised: 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] .
Revised: 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.
Revised: 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:
Revised: For example:

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

Original: [Yes] See Section ?? .•
Revised: [Yes] See Section ?? .•

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

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

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

Original: [N/A]
Revised: [N/A]

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

Original: Note that the
Revised: Note that the

Original: Checklist section does not count towards the page limit.
Revised: Checklist section does not count towards the page limit.

Original: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
Revised: In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?

Original: [Yes] .
Revised: [Yes] .

Original: Main claims are justified with a set of illustrative examples and experiments; both standard and toy examples included for experimentsto justify the key claims. (b) Did you describe the limitations of your work?
Revised: 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?

Original: [Yes]
Revised: [Yes]

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A]
Revised: [N/A]

Original: Besides
Revised: Besides

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

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

Original: We do not conduct
Revised: We do not conduct

Original: any experiments that could be in itself harmful for society but understand that
Revised: any experiments that could be in itself harmful for society but understand that

Original: the
Revised: the

Original: applicability of this method is very broad.
Revised: applicability of this method is very broad.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

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

Original: [Yes]
Revised: [Yes]

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

Original: [Yes]
Revised: [Yes]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes] .
Revised: [Yes] .

Original: Code will be
Revised: Code will be

Original: provided along with the supplementary materials for reproducibility.
Revised: provided along with the supplementary materials for reproducibility.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: (b) Did you specify all the training details (e g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes] .
Revised: [Yes] .

Original: More details to be included in supplemetary section(c)
Revised: More details to be included in supplemetary section

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

Original: ments multiple times)?
Revised: ments multiple times)?

Original: [Yes] .
Revised: [Yes] .

Original: All experiments are run across 3 to 5 random seeds.
Revised: 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.g.,
Revised: (d) Did you include the total amount of compute and the type of resources used (e

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

Original: [Yes] .
Revised: [Yes] .

Original: This will be mentioned in
Revised: This will be mentioned in

Original: supplementary and acknowledgements
Revised: supplementary and acknowledgements

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [N/A]
Revised: [N/A]

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

Original: [N/A]
Revised: [N [Table] A]

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

Revised: [Table]

Original: you’re using/curating?
Revised: you’re using/curating?

Original: [N/A]
Revised: [N/A]

Revised: [Table]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifi-
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifi-

Revised: [Table]

Original: able information or offensive content?
Revised: able information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Revised: [Table]

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

Original: if
Revised: [Table]

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Revised: [Table]

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

Revised: [Table]

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?
Revised: spent on participant compensation?

Original: [N/A]
Revised: [N/A]