{'fraction_non_alphanumeric': 0.03607335239095347, 'fraction_numerical': 0.009139879042451822, 'mean_word_length': 4.613360986791835, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Discrete Factorial Representations as an Abstraction for Goal Conditioned Reinforcement Learning

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. How to \textit{specify} and \textit{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. Defining 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. We propose to address this by learning factorial representations of goals and processing the resulting representation via a discretization bottleneck, for coarser goal specification, through an approach we call DGRL. 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. 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.


Introduction
Reinforcement Learning is a popular and highly general framework [26,61] focusing on how to select actions for an agent to yield high long-term sum of rewards. An important question is how to control the desired behavior of an RL agent, both during training and evaluation [25]. One way to control this behavior is by specifying a reward signal [53,57]. While this approach is very general, the reward signal can be hard to design and may not be the most informative form of feedback. The credit assignment problem in RL can become difficult when the reward signal is sparse [63,64,37,38,58], such as policy gradients becoming nearly flat in regions where reward is almost never achieved. 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 [21].
One potential way to flexibly specify and ground the desired behavior of RL agents is by training agents that receive a reward when they reach a goal specified explicitly to them [24]. 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 [62,49,45]. 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. It also allows for more flexible and tightly constrained control over its desired behavior [12,5,28]. Finally, the diversity of goals seen during training should help improve both credit assignment and generalization [49,44].
While this framework is promising, it introduces two new challenges: goal grounding [8,2] and goal specification [4]. Goal grounding refers to defining the goal space, and goal specification refers to selecting what goal the agent should try to reach in a given context. 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 specification, which makes the nature of the agent's task fundamentally different. What makes grounding and specifying goals challenging? Consider trying to train a goalconditioned RL agent to pick up various fruits from a table. For example, we may want it to pick up a red apple or a green pear (illustrated in Figure 1). 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). Goal Grounding refers to this challenge of relating high-dimensional observations and the space of relevant goals. Goal Specification refers to picking a suitable goal for the agent to reach and computing an appropriate reward when it is reached. It also implies specifying goals reachable in the agent's current context [40,27]. Goal specification 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,28,23]. Goals specified in language are an excellent fit for these desiderata, as language is a compressed discrete representation which is useful for out-ofdistribution generalization, while being compositional and expressive [20,11,23,18,1,68]. At the same time, connecting language feedback for an agent is non-trivial, requiring special assumptions or a labeling framework [9].
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. To perform this discretization, we use Vector-Quantization [65,48,34] which discretizes a continuous representation using a codebook of discrete and learnable codes. The approach proposed here, called DGRL, serves two complementary purposes. First, it provides a structured representation of the raw visual goals. By representing the visual goals as a composition of discrete codes from a learned dictionary, it simplifies the grounding of unseen goals, i.e., goals not seen during training, to novel compositions of the trained discrete codes. We show empirically that this improves the generalization performance of goal-reaching policies while remaining expressive enough. 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 final goal). In this case, goal-inference is learned end-to-end. The effectiveness of goal-conditioned HRL relies on the specification of semantically meaningful sub-goals. Using factorial discrete sub-goals allows the higher-level policy to specify semantically meaningful objectives to the lower-level policy.

Preliminaries
Goal-conditioned RL. We consider a goal-conditioned Markov Decision Process, where the goals g ∈ G live in the state space S, i.e., G = S. We denote a goal-conditioned policy as π(a|s, g) (either stochastic or deterministic), and its expected total return as J(π) = E T t=0 R(s t , g, a) where the goal g is either sampled from a distribution ρ g or provided by another higher level policy π h θ h (g|s). The value function V π is additionally conditioned on goals, and is trained to predict the expected sum of future rewards conditioned on states and goals; V π (s, g) = E T t=0 R(s t , g, a) | s 0 = s; π . As in standard RL, the objective in goal-conditioned RL is to maximize the expected discounted returns induced by the goal-conditioned policy. Figure 2: Summary of Proposed DGRL Model for improving goal grounding and goal specification by making goal representations discrete and factorial. We learn a latent representation for both observations and goals using a self-supervised learning method (sec. 3). We convert the learnt latent representation into discrete latents based on a VQ-VAE quantization bottleneck with multiple factor outputs (sec. 3). We use the resulting discrete representations for downstream RL tasks: (i) to train a goal-conditioned policy or value function, and (ii) in the context of goal-conditioned hierarchical reinforcement learning (sec. 3.1).
Hierarchical Reinforcement Learning. We consider goal-conditioned settings in which the goals are specified in the observation space. In the hierarchical reinforcement learning (HRL) setup, goals are provided by a higher level policy π h θ h (g|s t ). 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. The lower level policy executes primitive actions π l θ l (a|s t , g t ) to reach the goals specified by the high-level policy and is trained to maximize the intrinsic reward provided by the high-level policy. The higher level policy is trained to maximize the external reward i.e., the reward function specified by the MDP. Both the higher and lower level policies can be trained with any standard RL algorithms, such as Deep Q-Learning (DQN) [37] or policy optimization based algorithms [51,52,31]. Alternately, one can also consider another setup for goal-conditioned RL, where the goals are provided by the environment g 1 , . . . g L , and are part of the state or observation space. 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. 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. We consider both the HRL and goal-conditioned setups, and evaluate the significance of learning factorial representation of discrete latent goals in a series of complex goal-conditioned tasks.
Vector Quantized Representations. VQ-VAE [65,48,34] discretizes the bottleneck representation of an auto-encoder by adding a codebook of discrete learnable codes. The input is passed through an encoder. 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. The decoder is then tasked with reconstructing the input from this quantized vector.
Self-supervised learning of representations. Several papers [32,59,54,35] have demonstrated the benefits of using a pre-training stage where the representations of raw states are learned using self-supervised objectives in a task-agnostic fashion. After the pre-training stage, the representations can be used for (and potentially also fine-tuned on) downstream tasks. These self-supervised representations have been shown to improve sample efficiency.

Discrete Goal-Conditioned Reinforcement Learning (DGRL)
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. We later describe, in Section 6, how discrete goal representations can accelerate learning in complex navigation and manipulation tasks. 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 fixed representation for RL.
Self-Supervised Goal Representation Learning. One can use any off-the shelf self-supervised method for learning representations of the raw state and the goal observations. We denote by φ the encoder network that takes as input the raw state and maps it to a continuous embedding: z e = φ(s). Here, we explore two different self-supervised techniques for learning representations. For simpler environments, we use a simple autoencoder with its standard reconstruction objective. For more complex environments, we use the Deep InfoMax approach [35] which optimizes for a contrastive objective as a proxy to maximizing the mutual information between representations of nearby states in the same trajectory.
Processing continuous representations via a discrete codebook. We learn discrete representations by using the vector-quantization method from the VQ-VAE paper [65], and follow the multifactor setup used in Discrete-Value Neural Communication [34]. The discretization process for each vector z e ∈ H ⊂ R m is described as follows. First, vector z e is divided into G segments c 1 , c 2 , . . . , c G , where z e = CONCATENATE(c 1 , c 2 , . . . , c G ), and each segment c i ∈ R m/G (implying that m is divisible by G). 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): These discrete codes, which we call the factors of the continuous representation z e , are concatenated to obtain the final discretized vector z q : The loss for vector quantization is: L discretization = β G G i ||c i − sg(e oi )|| 2 2 . The training procedure closely follows both [34] and [65]. Here, sg refers to a stop-gradient operation that blocks gradients from flowing into e oi , and β is a hyperparameter which controls how strongly we move the codes toward the encoded values. Unlike [34], we used a moving average to update the code embeddings rather than learning them directly as parameters. We update e oi with an exponential moving average to encourage it to become close to the selected output segment c i . This update sets the new value of e oi to be equal to ηe oi + (1 − η)c i , where the value of η is a fixed hyperparameter controlling how quickly the moving average updates. 2 The term G i ||c i − sg(e oi )|| 2 2 is often called the commitment loss. We trained the VQ-quantization process together with other parts of the model by gradient descent. 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.
Summary. The multiple steps described above can be summarized by z q = q(z e , L, G), where L is the codebook size, G the number of factors per vector, and q(·) the whole discretization process. 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. The number of factors G is a hyper-parameter. In our experiments, we explored different values: G = 1, 2, 4, 8, 16, and found that G = 16 worked the best. 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.

Using representations for downstream RL
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.
Goal-conditioned RL. Defining 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 2 Note that this could also be thought of as a gradient step on eo i taken in the direction ci − eo i . 4 the low dimensional latent representation may fail to generalize. 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 fixed set of latent discrete codes. This facilitate generalization to new combinations of these codes while making it easy for downstream learning to figure out the meaning of each discrete code. 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.
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 ). At each training episode, a goal is sampled from the goal distribution ρ g , and the agent gets rewarded for reaching it. This reward can either be extrinsic, i.e., part of the environment, or intrinsic, i.e., part of the algorithm. In DGRL, we define the intrinsic reward as the fraction of discrete factors which match in the respective representations of the goal observation and of the state observation. 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. Hierarchical RL. 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 ). The effectiveness of goal-conditioned HRL relies on the specification of semantically meaningful sub-goals. 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. The use of learned codebooks ensures that the goal specified by the higher level policy is grounded in the space of raw-observations.
In Section 6, we empirically show the benefits of the proposed approach for training goal-reaching policies or goal-conditioned value functions, as well as in a goal-conditioned hierarchical RL setup.

Theoretical Analysis
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) = k p(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 ). Intuitively, this is because the discretization removes varieties of possible goal values g ∈ G k for each neighborhood G k . To state our result, we define where θ ∈ R m is the vector containing model parameters learned through n goals observed during training phase, g 1 , . . . , g n . We denote the discretization of g by q(g), and the identity function by id as id(g) = g. Let Q = {q(g) : g ∈ G} andd be a distance function. We use Q i to denote the i-th element of Q (by ordering elements of Q with an arbitrary ordering). We also define We denote by c a constant in (n, θ, Θ, δ, S).
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 E g∼ρg [(ϕ θ • ς)(g))] by the margin of ω(θ): Theorem 1. For any δ > 0, with probability at least 1 − δ, the following holds for any θ ∈ R m and ς ∈ {id, q}: 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: Proof. Detailed proof provided in the Appendix D Without the goal discretization, we incur an extra cost of ω(θ), which is expected to be strictly positive since 1 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. This extra cost ω(θ) goes to zero when the number of goal observations n approaches infinity.

Related Work
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. Often in goal conditioned HRL, the higher level policy specifies goals which may not have good specification and grounding. Several prior works focus on goal-conditioned RL to improve sample efficiency in deep RL tasks [39,41]. These build on ideas that were proposed years back to solve long horizon tasks by hierarchical RL specifying goals [24,13,14,69]. 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 [67,50,42]. 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 [17,72], for planning [43], or exploration [30]. Since the set of goals specified 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 [73]. We tackle this problem by proposing DGRL, for better grounding and specification of sub-goal representations.
In previous works, mutual information based objectives have been proposed for goal conditioned RL. They perform goal-based representation learning, in order to improve stability of training goalconditioned value functions [39,41], or to provide goal representations allowing to identify decision states for better exploration [19]. However, for most of these settings, the sub-goals are based on an external reward and are lacking in terms of specification, which can lead to inefficient training. [33,73] have proposed approaches that penalize the high level controller for generating sub-goals that are too difficult for the lower level policies, through the use of additional constrained objectives [73]. We highlight that our proposed DGRL can be generically applied to any goal conditioned RL literature for better grounding and specification of the sub-goals. Furthermore, recent work has shown significance of learning representations through self supervised objectives in RL [3,55], often as a pre-training phase [56,70], which can help for both exploration [36] and control [71].
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]. Previous works [33] have often used the entire observation space as goals, which is not scalable for complex tasks. Other works have used a pre-defined space of sub-goals as domain knowledge [39], or self-play for sub-goal representations [60] to reduce the complexity of goal space design. Most recently, [43,42] utilized unsupervised representation learning to learn a goal representation space, which can further be used for planning and control. In this work, we show additionally that using a bottleneck can further lead to factorial representation of goals, while helping with goal specification via learning a latent space of discrete goals usable for planning and control. We emphasize that DGRL can be integrated on any existing goal conditioned approach that utilizes learning a sub-goal representation.

Experiments
The main goal of our experiments is to show that goal discretization can lead to sample efficient learning and generalization to novel goals, in goal-conditioned RL. 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. Second, we consider hierarchical goal-conditioned RL, in which a higher-level agent generates goals that a lower-level agent is tasked with reaching. In this case, the task of reaching novel goals occurs organically as the higher-level model selects new goals. This setup also shows the advantages of DGRL for goal specification. 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.
We evaluate our proposed method DGRL by integrating it into existing state-of-the-art goalconditioned and hierarchical RL tasks. Experimentally, we analyse DGRL on several challeng-ing testbeds that have previously been used in the RL community. DGRL in principle can be applied to any existing downstream goal-conditioned RL tasks. We demonstrate improvements on five such tasks. We consider maze navigation where images are used as observations and we show improved generalization to novel goals. We integrate DGRL to an existing goal-conditioned baseline for navigating procedurally-generated hard exploration Minigrid environments [10] and find that it outperforms state-of-the-art exploration baselines. We also show improvements with DGRL on continuous control (Ant) navigation and manipulation tasks, where goals come from a high-level controller. Finally, we show that discrete representations also significantly improve sample efficient learning on a challenging vision-based robotic manipulation environment.
Demonstrating Factorized Representation Learning We first pick a color-mnist supervised learning example to support the idea that DGRL can learn factorized or compositional representations.  representation. We find 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. This empirically suggests the emergence of "factorization" in the learnt representations. Further experimental details are provided in section C.1.
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 finite set of training goals, and during evaluation is tasked with reaching a novel goal unseen during training. This is a navigation task with a pixel-level observation space showing the position of the agent and the goal in a gridworld. We consider two mazes spiral and single-loop topology. 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. Because the 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. At each episode, a specific goal is randomly sampled from a distribution of goals, and the DQN agent is trained to reach the specified goal for that episode. During evaluation, we test the learned agent on goals either from the training distribution, or not seen during training. Furthermore, for this task, we additionally use an intrinsic reward to promote exploration of the goal-DQN agent. 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. 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. Figure 4 shows that DGRL significantly outperforms a continuous baseline goal DQN agent, when trained on either four goals or eight goals. We evaluate generalization to 4 novel goals unseen during training ( Figure 5) and demonstrate improved generalization.   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. Now, we consider various setups in which a goal generating agent specifies goals using the learned codebook and a goal-conditioned agent is tasked with reaching the goals specified by the goal generating agent. 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.
Procedurally Generated MiniGrid Exploration Task. We follow the experimental setup of [6] and [47] and evaluate DGRL on procedurally generated MiniGrid environments [10]. In [6], a goal-generating teacher proposes goals to train a goal-conditioned "student" policy. We integrate DGRL on top of AMIGO [6] and compare DGRL on a hard exploration task with state-of-the-art exploration baselines. Experimental results are summarized in Table 1 and more details provided in Appendix C.4. Note that unlike RIDE and RND, we do not provide an additional exploration bonus to DGRL, and find that DGRL can still solve this hard exploration task more efficiently.
Goal Grounding in KeyChest Maze Navigation Domain. We consider a simple discrete state action KeyChest maze navigation task, following [73], where discrete goals in the state space are provided by a higher level policy. For this task, to integrate DGRL, we learn an embedding φ(·) of the goals, then discretize the representation with a learned codebook. We compare with a baseline HRAC [73] agent (details in Appendix C.3). Figure 7 shows an illustration of the KeyChest environment and a performance comparison of DGRL with different group factors G. Using fewer factors (G = 4) performs worse than the HRAC baseline, whereas using a larger number of factors (G = 8 or G = 16) improves the sample efficiency of the goal reaching agent, providing evidence for the benefits of factorization.    Ant Manipulation Control Domains. We employed DGRL on three different continuous control tasks: AntMazeSparse, AntFall and AntPush. We emphasize that these tasks are the more challenging counterparts of AntGather and AntMaze tasks, typically used in the hierarchical RL community [39,41]. Figure 8 provides an illustration. We evaluate goal discretization by integrating DGRL to the state-of-the-art HRAC baseline. Details of the experimental setup are provided in Appendix 6. Figure 8 shows that specifying the goals using the learned codebook helps DGRL achieve a higher success rate compared to the HRAC baseline.
Ant Navigation Maze Tasks. 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). The ant navigating in the maze is trained to reach any goal in the environment. The agent is evaluated for generalization in an extended temporal setting with a difficult configuration, we compare the success rate of DGRL integrated on top of RIS with several baselines. We emphasize the difficulty of these tasks, where existing baselines like soft actor critic [21,SAC] and temporal difference models [46, TDM] fail completely. Results in Figure 9 show that DGRL improves the sample efficiency over the RIS baseline. Additional experimental setup and environment configurations are provided in Appendix C.6. : Performance comparison with the success rate of reaching goal positions during evaluation in an extended temporal configuration of the U-shaped and maze-shaped Ant navigation tasks. We find 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 [7]. The RIS baseline is based on raw data provided by the authors.
Vision Based Robotic Manipulation. 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. This manipulation task is adapted from [43], where the baseline RIS is already shown to be superior to previous goal conditioning methods. The task of the agent is to control a 2-DoF robotic arm from image input and move a puck positioned on the table. The Sawyer task is designed for training and generalization. At test time, it evaluates the agent's success at placing the puck in desired positions in a temporally extended configuration. This is a challenging vision-based complex motor task, since test time generalization requires temporally extended reasoning. Results in Figure 10 show that DGRL improves the sample efficiency over the RIS baseline. Details of the experimental setup are provided in Appendix C.7. Conclusion 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. We show that an instantiation of this idea using multi-factor discretization significantly improves performance on a diverse set of benchmarks.
Limitations and Future Work An interesting question that arises from our work is how to theoretically ground and specify goals, which might be helpful for efficient structured exploration in tasks where goal seeking is crucial. Additionally, while we demonstrate that the factorial representations learnt by DGRL can be beneficial, it would be interesting to explore whether we can also enforce compositionality in the latent embeddings. 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.
In section 5, we discuss extended related work. In section C we include details of our experiment setup and additional ablation analysis evaluating DGRL. Section D contains detailed proof for theorem 1. Finally, we provide our algorithm, DGRL that requires minimal changes to existing goal conditioned RL setup, in section E, with sample code snippet provided in F

B Demonstrating Factorized Representations in a Robot Experiment with Visual Background Distractors
We demonstrate the ability of DGRL to learn factorial representations on real world robot data, where the robot arm moves in presence of background video distractors [29]. Further details on the robot arm data collection are provided below. The data contains rich temporal background noise. We first learn a representation with a simple auto-encoder, following by the discretization bottleneck of DGRL, and then reconstruct the image with different discrete factors. Figure 11 demonstrates factorization in the learnt representation tweaking the different factors used in the discretization bottleneck. In particular, we observe some form of "compositionality" emerging as the decoder was never trained on some of the combinations of factors, for instance (person in the background + orange lamp) and (person in the background + arm to the left).
Experiment Details: In this task, the robot arm moves on top of a grid layout, containing 9 different positions. We denote these as the true states. We use two cameras to take images, for the dataset, one from the front side of the robot and the other with a top down view from above. We collect a dataset containing pixel based observations only, where the images consist of the robot arm along with the background distractors. Inspired by the exogenous noise information setup [16], we setup the robot task while there is a TV playing a video in the background, with other flashing lights nearby. The offline dataset consists of 6 hours of robot data, with 14000 samples from the arm, taking high level actions of move left, right, up and down. A sample point image is collected after each action, and the background distractors changes significantly, due to video and lighting in the background. The goal of the experiment is to predict accurately the ground truth state position by learning latent representations with DGRL.
Experiment Results: We evaluate the ability of DGRL to accurately reconstruct the image, by learning the latent state representation while also ignoring the background distractors. This is denoted as the Image Noise, where we compare DGRL with and without VIB, alongside a baseline agent which only learns a representation. For learning latent representations, we use a multi-step inverse dynamics model [29]. In addition, we compare the ability of DGRL to accurately predict the ground truth states, denoted by State Accuracy solely from the observations, as a classification task. This is challenging since the learnt representation needs to predict ground states while ignoring the irrelevant background information. Furthermore with the learnt model, we predict the time-step for each observation as an additional metric to determine effectiveness of DGRL. The time-step is an indicator of the background noise that appeared in each sample; and with Temporal Noise, we evaluate DGRL to predict the time step while ignoring irrelevant information from observations. Experiment results in Figure 11 shows that the use of VIB helps improve the ability of DGRL to remove noise from the representation, while being able to almost perfectly predict the ground truth state of the robot.

C Experiment Details and Additional Results
For all our experiments, we use existing open-source implementations of the baseline algorithms. We mostly integrate DGRL with HRAC [73] and RIS [7], which are state of the art goal conditioned algorithms on the complex Ant navigation and robotic manipulation tasks. We use the same hyperparameters and default configurations as used in the baselines. For the simpler maze tasks such as from MiniGrid and KeyChest, we also use existing setups used by previous algorithms [6,73]. All our experiment results are based on 3 random seeds, and we provide our implementation for reproducibility.  Figure 11: An autoencoder trained with 8 discrete codes to reconstruct images of a real-robotic arm with independently changing distractors (TV and lamp). The first two rows show the reconstruction of the initial images. In the middle two rows, we change one of the discrete factors to match its value in the first image, in the resulting reconstruction, the lamp has turned orange. On the bottom two rows, we change another discrete factor, to match its value from the first image, and observe that it made the robotic arm point left 17

C.1 Experiment Details on Color-MNIST Dataset
We include a brief description of the experiment details, used for the color-mnist example to demonstrate factorization. The pixel-based input is first passed through an encoder (a two-layer neural network) to obtain its latent representation with the dimension of 30, we then quantize the continuous representation into two groups of discrete codes, where the codebook size is 256. In the training procedure, two groups of the discrete codes are then concatenated to obtain the discretized representation, and finally passed through a decoder (another two-layer neural network), where we used reconstruction loss (MSE loss) combining with the loss for vector quantization to train the network. While in the testing procedure, we used zero vector to substitute one group of the discrete codes, and then obtain the reconstructed image by concatenating it with the other group and passing through the decoder.  In this section, we provide additional details for the visual maze tasks. We provide environment visualizations, for the spiralworld and loopworld environments in figure 12. For this task, we use 6 × 6 gridworlds, where the agent receives pixel based full observations, of size 80 × 80 × 3, of the environment. The agent starts deterministically in the environment from one corner of the loop, and the goals are provided in the state space, where agent receives observations of goals. The agent receives a negative reward of −1 at every time step, and a reward of +5 when reaching any of the goals.

C.2 Visual MazeWorlds
We train a goal conditioned Q learning agent, with epsilon greedy exploration, where the agent receives goal observations in addition to state observations. We use a simple 2 layer architecture for the value function, which additionally conditions on the goals. In our experiments, we use 4-goals and 8-goals environment, where the agent is trained to reach each of the goals, sampled at every episode. For pixel based observations, we learn a representation φ, which can be trained via any self-supervised representation learning objective. For our experiments, we use the Deep InfoMax algorithm [35] for training the representation learner φ. Since the environment is quite simple, we found that even pre-training φ with random rollout policies is often good enough for training the value functions, as suggested by our experimental results. For DGRL, we additionally use a discrete bottleneck on top of the embedding, and experimentally show that DGRL can significantly outperform a baseline goal DQN agent.

C.3 KeyChest Domain
For prototype and motivation of DGRL, we also used a simple discrete state and action KeyChest environment, as shown in figure 7. This task is inspired from the HRAC algorithm [73], where the agent starts stochastically in the environment, and the goal is to pick up a key and open a chest. We follow a HRL setup in this task, where the higher level policy provides discrete goals in the state space for the lower level policy to reach. Since the environment has injected stochasticity, such a task requires both low-level control conditioned on the goals, provided by a higher level planner.
The environment has a 3-dimensional state space where the first two represent the position of the agent, while the third dimension represents whether the agent has picked up the key or not. The reward function is sparse, such that the agent only receives a reward +1 for picking up the key, and a reward of +5 if it can open the chest; otherwise zero rewards elsewhere. This hints to a hard exploration task, where specification of goal plays a key role. The lower level policy receives a goal reaching reward only, upon reaching the goal state. The higher level policy receives rewards from the environment directly. For the goal reaching reward, we use a standard Euclidean distance between states and goals, and the goal is achieved by the lower level policy if this distance is below a threshold of 0.5.
For DGRL, we apply a discrete bottleneck by first learning a representation of the discrete goal in state space, using an encoder. We then apply a discrete bottleneck on the learnt embedding, followed by a decoder that maps the learnt discrete goal embedding back to the original state space. For baseline [73], the algorithm uses the raw goal states provided by the higher level policy. Both the higher and lower level policies are learnt with actor-critic algorithms. The high level policy is trained based on the task reward, whereas the lower level policy is rewarded for reaching the goals, or nearby regions of the goals, provided by the higher level policy.

C.4 MiniGrid Environments
For our experiments, we follow the setup of [6] and [47] and evaluate DGRL on a simple procedurally generated MiniGrid environment [10]. The minigrid environments are a suite of hard exploration testbeds in RL where the task is designed such that exploration and representation of the visual observations can be disentangled. Following [6], we use the KCharder, KeyCorrS4R3 environment, which requires finding a key that can unlock a door which blocks a room. This door needs to be opened by the agent so as to reach the goal. In our experiments, we use the same learning rates, network architecture and other hyperparameters as in the AMIGO paper [6]. We use the opensourced implementation provided by the authors (for more details, see [6]), and simply integrate DGRL on top of the AMIGO baseline, learning a factorial representation for the goal observations.

C.5 Ant Manipulation tasks
We employ an ant robot with a continuous 8-dimensional action space for all three Ant manipulation tasks. Each episode terminates at 500 time steps in all three tasks.
AntMazeSparse AntMazeSparse is a challenging navigation task with sparse rewards. This environment has a continuous state space including current position, velocity, target location and the current timestep t. The agent is provided with a sparse reward by +1 only if the Euclidean distance between the agent and the target location is smaller than 1, where the target location is set at (2.0, 9.0) in the center corridor.
AntPush This environment has the same state space as AntMazeSparse task. The difference is that this environment has a movable block which the agent can interact with. To successfully reach the target position (0.0, 19.0), the agent must push the large block to the side to clear the path to the target location. The success of the agent is defined as having the Euclidean distance of 5 from the target position.
AntFall This environment extends the navigation to three dimensions. Similar to the AntPush task, the environment still has a movable block, while the agent must move the block into a chasm instead of pushing it aside, so that it may walk over it without falling to reach the target position.
The target position is fixed to (0.0, 27.0, 4.5) in this environment.
For our experiments, we use the setup provided by the HRAC baseline [73]. In HRAC, an additional adjacency constrained is trained, along with the lower and higher level policies, such that the goals provided by higher controller are within a constrained region of the state space, that can be reachable by the lower level policy. The higher and lower level policies are both trained based on actor-critic algorithms, with separate replay buffers for each policy. The replay buffer for the higher level policy stores one every K transitions.
We integrate DGRL on top of the HRAC baseline setup, where we apply the discrete goal bottleneck based on the output embeddings from higher level policies. Compared to HRAC, DGRL with HRAC would condition the lower level policies on the discrete embeddings of goals. Figure 13 provides results for all the Ant manipulation tasks. Here we provide the results obtained for all the 5 different Ant environments, and we experiment for a range of discrete factors from 2 to 32. We find that while the best performing factor G is not consistent for all tasks, in general factors of 4, 8 and 16 typically outperform 2 and 32 factors.

C.6 Ant Navigation Tasks
Following the setup in RIS [7], we employ an ant robot with a continuous 8-dimensional action space for all four Ant navigation tasks. reasoning. Figure 14(c) shows an additional result on the U-shaped Ant maze navigation environment. All these environments have a continuous state space including the current position, orientation, the joint angles and the velocities. The goal is considered reached when the Euclidean distance from the target position of the environment is less than 0.5. The agent gets a −1 reward at each time step until the goal is reached. The U -shaped maze has a size of 7.5 × 18, the S-shaped maze has a size of 12 × 12, while the π-shaped maze and ω-shaped maze share the same size of 16 × 16. In the training stage, target and initial position of the agent are sampled randomly at the beginning of each episode. At evaluation, the initial state and the target position are fixed, as illustrated in Figure 14, to test the performance of the agent on challenging configurations that require temporally extended

C.7 Robotic Manipulation Task
We consider the visual robot manipulation task Sawyer, from the multiworld environments of [43]. Our setup is based entirely on the experiment details and code provided in the open-source codebase of RIS [7]. The goal of the agent is to operate a 2D position control and manipulation task. The observations of the environment are based on 84 × 84 RGB image of the environment. DGRL is trained using a discrete bottleneck based on the learnt representation of the pixel based observations of the environment. We compare the Sawyer environment based on an existing RIS [7] baseline which is already shown to outperform other baselines on this task. For more details on the suite of multiworld environments, including the Sawyer manipulation task, see [43]. In this setting, we can additionally check for generalization. At test time, the cumulative returns of the agent are provided  Figure 14: Visualization of Ant navigation environment, which are standard goal conditioned HRL tasks. We integrate DGRL on top of the existing state of the art RIS baseline [7]. Performance comparison of success rate on AntU-shaped environment (right). We compare DGRL with different factors G and the RIS baseline [7] which has already been shown to outperform other goal conditioned baselines such as LEAP [43] on the Ant navigation tasks.
when evaluated on a slightly different task, such as a hard configuration where we evaluate the policy and bottleneck based on learnt representations in the given task. We consider a goal-conditioned Markov decision process, defined by states s ∈ S, goals g ∈ G, actions a ∈ A, a reward function r(s, a, g), a transition dynamics p(s |s, a), a maximum horizon T , the initial state distribution ρ 0 , and the goal distribution ρ g . The objective in goal-conditioned RL is to obtain a policy π(a|s, g) to maximize the expected sum of rewards E g∼ρg,((st,at)) T t=1 [ T t=1 r(s t , a t , g)] where the sequence of state-action pairs ((s t , a t )) T t=1 is sampled according to s 0 ∼ ρ 0 , a t ∼ π(a t |s t , g), and s t+1 ∼ p(s t+1 |s t , a t ).
To study the phenomenon of the goal observations, we define where θ ∈ R m is the vector containing model parameters learned through goals observed during training phase, g 1 , . . . , g n . Define 1{a = b} = 1 if a = b and 1{a = b} = 0 if a = b. Moreover, ϕ θ • ς represents the composition of functions ϕ θ and ς.
In our proof, we will use the following previous result: Lemma 1 (Bretagnolle-Huber-Carol inequality). [66, Proposition A.6.6] If X 1 , . . . , X K are multinomially distributed with parameters m and p 1 , . . . , p K , then forM > 0, Proof. We decompose the goal probability p(g) into p(g) = k p(g|g ∈ G k )p(g ∈ G k ) with the neighborhood set {G k } k and analyze the concentrations of random variables in terms of p(g|g ∈ G k ) and p(g ∈ G k ). Let ς ∈ {id, q}. We define M = sup g∈G ϕ θ (g). Then, Using this, we decompose the difference as Substituting these into equation (2) yields = k E g∼ρg [(ϕ θ • ς)(g)|g ∈ G k ] |I k | n − Pr(g ∈ G k ) By using Lemma 1 by setting δ = 2 K exp − mM 2 2 and solving for M , we have that for any δ > 0, with probability at least 1 − δ, k |I k | n − Pr(z ∈ G k ) ≤ 2 ln(2 |Q| /δ) n ≤ 2|Q| ln(2/δ) n , where the last inequality follows from the fact that 1/δ ≤ 1/δ |Q| as δ ∈ (0, 1) and |Q| ≥ 1. Here, notice that the term of k |I k | n − Pr(z ∈ G y k ) does not depend on θ. Moreover, note that for any (f, h, M ) such that M > 0 and B ≥ 0 for all X, we have that P(f (X) ≥ M ) ≥ P(f (X) > M ) ≥ P(Bf (X) + h(X) > BM + h(X)), where the probability is with respect to the randomness of X. Thus, by combining this and equation (3), we have that for any δ > 0, with probability at least 1 − δ, the following holds for all θ, (4) Therefore, for any δ > 0, with probability at least 1 − δ, the following holds for any θ ∈ R m and ς ∈ {id, q}: Define u n (S, θ) = max k∈I Q 1 |I k | i∈I k ϕ θ (g i ) − E g∼ρg [ϕ θ (g)|g ∈ G k ]. Then, since k∈I Q |I k | = n, ω(θ) ≤ u n (S, θ) 1 n k∈I Q |I k | = u n (S, θ).
For each k ∈ [|Q|], if p(g ∈ G k ) = 0, then the probability of the event of |I k | ≥ 1 is zero. Thus, by taking union bounds, with probability one, for all k ∈ I Q , |I k | → ∞ as n → ∞. Therefore, by using the uniform law of large numbers (Theorem 2 of [22]) and union bounds over k ∈ I Q (noticing that |I Q | ≤ |Q| is finite), we have that sup θ∈Θ |u n (S, θ)|

E Algorithm in Goal Conditioned RL
We present the entire algorithm of DGRL built on top of RIS in algoirithm 1.

F Code Snippet of DGRL
We show a simple code snippet of algorithm 1 below.