Discrete Codebook World Models for Continuous Control

We thank Reviewer iQtr for their feedback. We refer the reviewer to our general response where we outline the improvements to our submission. Here we address specific questions.

W1: This paper lacks comparisons with the more advanced baselines like EfficientZero-V2.

AW1: EfficientZero-V2 proposes improvements on the policy and planning level:

We design two key algorithmic enhancements: a sampled-based tree search for action planning, ensuring policy improvement in continuous action spaces, and a search-based value estimation strategy to more efficiently utilize previously gathered data and mitigate off-policy issues.

In contrast, our paper investigates and proposes changes to the latent embedding and dynamics functions. While it stands to reason that many methods improving upon DreamerV3 and TD-MPC2, such as EfficientZero-V2, can also benefit from our proposed changes, we are unfortunately not able to compare against all possible variations of DreamerV3, which is cited 471 times. Hence, we specifically compare against the baseline implementations of DreamerV3 and TD-MPC2 which use different latent embeddings (continuous vs discrete), as well as perform ablation studies into the different discretization schemes of the latent embedding space. We agree with you that it would be beneficial to develop a kitchen sink approach combining these advancements from different papers in future work, but we think this is out of scope for this paper. However, please see our updated draft for further ablation studies and additional experiments.

Q1: What's your idea about if FSQ-based discretization could be used for action space discretization? (like action sampling) If yes, how will you do that?

AQ1: In this paper we are specifically interested in applying FSQ to the discretization of latent state/observation embeddings. We did not investigate the possibility of action space discretization and leave this to future work.

Q7: In Line 514, why is $|\mathcal{C}|=2^4$ for $L=[5,3]$, when it is claimed $|\mathcal{C}| = \prod L_i$ in Line 224.

AQ7: We invite the reviewer to check Line 514 again (now line 516), which reads as $|\mathcal{C}|\approx 2^4$ and not as $|\mathcal{C}| = 2^4$. Data types are often written as $2^X$ because of the binary nature of computer systems. For example, real numbers are often represented in single precision (32 bits) which are of size $2^{32}$. As such, representing the size of discrete codebooks using $2^X$ makes it easier to compare sizes, which is why it is common in FSQ and VQ-VAE to represent the codebook size as $|\mathcal{C}| \approx 2^X$. We note that the reviewer is correct that the codebook size is given by $|\mathcal{C}|=5 \times 3 =15 \approx 2^4$. In order to prevent this confusion in the future, we have updated Line 516 to read as:

We find that a latent dimension of $d=1024$ with FSQ levels $\mathcal{L}=[5,3]$, which corresponds to a codebook size $|\mathcal{C}| = 5 \times 3 = 15 \approx2^{4}$, performs best in the harder DMC tasks.

Q8: Could you provide additional experiments on more tasks, such as ManiSkill2/Myosuite, which TD-MPC2 used?

AQ8: We have been able to get results in 45 MetaWorld tasks (previously we only had 7) and 5 MyoSuite tasks. We have compared our method with baselines on Myosuite tasks in Section B.10 and Figure 19 and Figure 20. Please also see our general response.

Q9: It would also be beneficial to include experiments on more base algorithms, such as DreamerV3 (see question 2).

AQ9: We have included results for soft actor-critic (SAC) and the original TD-MPC in all DMControl and Meta-World tasks. As such, we are now comparing against TD-MPC, TD-MPC2, SAC, and DreamerV3, in all figures that report baselines.

Q10: The most direct evidence that the proposed discrete latents outperform one-hot and continuous latents would be to replace the latent space in DreamerV3 and TD-MPC2 with your proposed approach. Why wasn't this head-to-head comparison conducted?

AQ10: Please see Appendix B.9 (and our general response) where we have included a section comparing DCWM and DreamerV3.

Q11: Have you compared FSQ and VQ (referred to as dictionary learning in Line 206)?

AQ11: We agree this is an interesting comparison and thank the reviewer for raising it. In response to this question, we have included results for VQ-VAE to Appendix B.7. Please see our response to all reviewers for further details.

Q12: Codebook encoding only outperforms one-hot encoding in one task in Figure 6, yet it is claimed to be more computationally efficient in Line 485. Is there any numerical analysis to support this claim?

AQ12: One-hot, label, and codebook encodings have different impacts on computational efficiency in neural networks. First, one-hot encodings increase the input dimension as they create new columns for each category. In our setting, we have $d=512$ latent dimensions and $15$ discrete values so the one-hot encoding is of size $512\times15=7680$. In contrast, label encodings are of size $d=512$. Codebook encodings offer a trade-off as they are off size $d \times b = 512 \times 2 =1024$. Second, one-hot encodings create sparse data as there are not many non-zero values in the encoding. In contrast, both label and codebook encodings create dense data as they have many non-zero values. Due to the increased dimensionality and sparsity, one-hot encoding may lead to longer training times compared to label encoding. The table below summarizes the differences between the encodings. We have also updated Figure 7 (now Figure 6) to compare both sample efficiency and computational efficiency of one-hot, label, and codebook encodings in the hard Dog Run and Humanoid Walk tasks. It shows that one-hot encodings result in slower training.

Q5: I did not understand why "we did not configure the transition model to accept the label and one-hot encodings as this resulted in the agent being unable to learn."

AQ5: This is an important point so thank you for highlighting that this is not clear. We have updated Section 5.3 to provide more clarity.

Our world model consists of NNs for the dynamics $P_{\phi}(\mathbf{c}' \mid \mathbf{c}, \mathbf{a})$, reward $R_{\xi}(\mathbf{c}, \mathbf{a})$, critic $Q_{\psi}(\mathbf{c}, \mathbf{a})$, and prior policy $\pi\_{\eta}(\mathbf{c})$, which all make predictions given the discrete codebook encoding $\mathbf{c}=\mathbf{e}\_{\text{codes}}$. We compare DCWM's codebook encoding $\mathbf{e}\_{\text{codes}}=\mathbf{c}^{i} \in \mathcal{C}$ to (i) label encoding $e\_{\text{labels}} = i \in \{1,\dots|\mathcal{C}|\}$ and (ii) one-hot encoding $\mathbf{e}\_{\text{one-hot}} =\mathbf{v} \in \{0,1\}^{|\mathcal{C}|}$ given $\textstyle \sum_{i=1}^{|\mathcal{C}|} v_{i} =1$. In Fig. 6, we evaluate what happens when we replace the codebook encoding $\mathbf{c}$ with either one-hot $\mathbf{e}\_{\text{one-hot}}$ or label $\mathbf{e}\_{\text{label}}$ encodings. In these experiments, we did not modify the dynamics $P\_{\phi}(\mathbf{c}' \mid \mathbf{c}, \mathbf{a})$, that is, the dynamics continued to make predictions using the codebook encoding $\mathbf{c}$ and did not use the one-hot or label encodings. This is because we found that using one-hot and label encodings for the dynamics led to the agent being unable to learn. This suggests that our codebook encoding is essential when leveraging a discrete latent space for learning continuous control. More specifically, we evaluated the following experiment configurations:

1. Codes (ours): All components used codes: $P\_{\phi}(\mathbf{c}' \mid \mathbf{c}, \mathbf{a})$, reward $R_{\xi}(\mathbf{c}, \mathbf{a})$, critic $Q\_{\psi}(\mathbf{c}, \mathbf{a})$ and prior policy $\pi_{\eta}(\mathbf{c})$.

2. Labels: Dynamics model used codes $P_{\phi}(\mathbf{c}' \mid \mathbf{c}, \mathbf{a})$ whilst reward $R_{\xi}(\mathbf{e}\_{\text{labels}}, \mathbf{a})$, critic $Q_{\psi}(\mathbf{e}\_{\text{labels}}, \mathbf{a})$ and prior policy $\pi_{\eta}(\mathbf{e}\_{\text{labels}})$ used labels $\mathbf{e}\_{\text{labels}}$ obtained from the code's index $i$ in the codebook.

3. One-hot: Dynamics model used codes $P_{\phi}(\mathbf{c}' \mid \mathbf{c}, \mathbf{a})$ whilst reward $R_{\xi}(\mathbf{e}\_{\text{one-hot}}, \mathbf{a})$, critic $Q_{\psi}(\mathbf{e}\_{\text{one-hot}}, \mathbf{a})$ and prior policy $\pi_{\eta}(\mathbf{e}\_{\text{one-hot}})$, used one-hot $\mathbf{e}\_{\text{one-hot}}$ obtained by applying PyTorch's torch.nn.functional.one_hot to the label encoding.

Please let us know if anything remains unclear or if you have any suggestions for improvement.

Q6: I did not understand the sentence "our codebook is configured with b = 2 channels and the label encoding incorrectly assumes an ordinal structure through both." How does the hyperparameter $b$ in your proposed codebook encoding affect label encoding?

AQ6: We would like to refer you to the new "Preliminaries" section (Section 3) where we discuss one-hot, label, and codebook encodings, and how they imply ordinal structure. To avoid this confusion in the future, we have also added the following text to Section 5.3:

The label encoding (blue) struggles to learn in the Humanoid Walk task and is often less sample efficient than the alternative encodings. This is likely because the label encoding is not expressive enough to model the multi-dimensional ordinal structure of our codebook. Let us provide intuition via a simple example. Our codebook has $b=2$ channels, so two different codes may take the form $e\_{\text{codes}}(A)=[0.5, -0.5]$ and $e\_{**codes**}(B)=[0, 0.5]$. As a result, our codebook encoding can model ordinal structure in both of its channels, i.e., $e\_{\text{codes}}(A)_1>e\_{\text{codes}}(B)_1$ whilst $e\_{\text{codes}}(A)_2 < e\_{\text{codes}}(B)_2$. The corresponding label encoding would encode this as $e\_{\text{labels}}(A)=1$ and $e\_{\text{labels}}(B)=2$, which incorrectly implies that $B>A$. In short, the label encoding cannot model the multi-dimensional ordinal structure of the codebook $\mathcal{C}$.

We are glad to see that you found our paper insightful and experiments informative. Please, see below our answers to your raised questions:

Q1: The claim in the abstract that discrete latent spaces typically underperform in continuous control tasks needs specification...

AQ1: We thank the reviewer for this comment. First, we would like to point out that we compare our approach to both DreamerV3 and TD-MPC2 as baselines as they utilize world models with different latent state embeddings. Furthermore, a sole comparison between TD-MPC2 to DreamerV3 was already presented by [Hansen et al., 2024], thus, we found it informative and beneficial for the reader to perform a comparison of our presented approach to both methods, especially due to their use of different encoding schemes. We have clarified our statement about the loss of performance when using discrete latent spaces - we did not mean to downplay DreamerV3's achievements unfairly. As you can see in Figures 5/12 and Line 462, we performed an in-depth evaluation of different latent spaces and found the proposed discrete latent space to perform better than continuous latent spaces, even though it stands to reason that discrete codes may lead to a decrease in performance due to information loss.

Q2: A preliminary section clarifying different discrete encodings is necessary; I find the distinction between label encoding and one-hot encoding unclear for the first-time reading.

AQ2: Thank you for bringing this to our attention. We have now added a Preliminaries section (Section 3) which introduces and compares one-hot, label, and codebook encodings. We think this is a nice addition to the paper but please let us know if anything remains unclear.

Q3: The methods section extensively describes FSQ, which is not originally proposed here. Also, the methods section contains many elements that directly follow TD-MPC2, which may confuse non-expert readers regarding what is novel. Clear differentiation is needed. For example, TD-MPC2 exactly uses an ensemble of five critics with randomly sampled two for bootstrap, but this is not explicitly mentioned by the authors in Line 277.

AQ3: We clearly indicate and reference to the reader that FSQ is prior work, please see Line 208. As the discretization scheme is an important aspect of the paper, we find it appropriate to discuss its background and intuition in more detail. Nevertheless, the reviewer is correct that we intentionally kept aspects of our algorithm similar to TD-MPC2 in an attempt to isolate the impact of our world model's latent space. More specifically, we used an ensemble of five critics, the same NN architectures, and we used MPPI for decision-time planning. To improve clarity, we added a further discussion on the specific contributions and differentiations to other model-based RL methods in Section 4.1.

Q4: Can you elaborate on what you mean by "codebook encoding captures similarity between observations" and how this contrasts with one-hot encodings in Line 248?

AQ4: We thank the reviewer for raising this question as this is an important point that we want to make sure is clear in our paper. We have added further clarification regarding this statement in the new "Preliminaries" section (Section 3) which we include here for convenience:

• One-hot encoding Given categories $A$, $B$, $C$, and $D$, a one-hot encoding would take the form $A=[1,0,0,0]$, $B=[0,1,0,0]$, $C=[0,0,1,0]$ and $D=[0,0,0,1]$, respectively. This encoding creates equidistant vectors in the feature space, as the Euclidean distance between any two one-hot vectors is always $\sqrt{2}$.
• Label encoding Given categories $A$, $B$, $C$, and $D$, label encoding would result in $A=1$, $B=2$, $C=3$, and $D=4$ respectively.
• Codebook encoding Given categories $A$, $B$, $C$, and $D$, a codebook might encode them as $A=[0.5,0.5]$, $B=[0.5,-0.5]$, $C=[-0.5,0.5]$, and $D=[-0.5,-0.5]$ respectively.

Mathematically, if we have an ordinal relationship $A < B < C$, label and codebook encodings can ensure $|e(A) - e(B)| < |e(A) - e(C)|$, where $e(x)$ is the encoding function. One-hot encoding, however, results in $|e(A) - e(B)| = |e(A) - e(C)| = \sqrt{2}$ for all distinct pairs, eliminating any notion of ordering. Whilst this may be beneficial in some scenarios, e.g. when modelling distinct categories like fruits, it means that they cannot capture the inherent ordering in continuous data. In contrast, label and codebook encodings can capture ordinal relationships.

Since no rebuttal is provided, I decided to keep my score currently.

We thank the reviewer for their comments and are glad that they found the paper well-written and the presented ablation experiments sufficient.

W1: The test for multi-tasking learning is missing. It would be great if the authors could consider extend their studies to multi-task learning and study the differences in codebooks between different control tasks.

AW1: We agree that extending DCWM to the multi-task RL setting and learning different codebooks is interesting, however, we intentionally left this for future work as our current paper is focused solely on discrete codebook encodings for world models. As we detail in our paper, the use of discrete code embeddings increases the performance of world models. Evaluating the use of a task-dependent variable code embedding, as you propose, is beyond the current scope of this paper especially given the short time window available for running additional experiments in the discussion period. However, we appreciate your idea for this possible extension.

W2: The stochastic transition dynamics in the latent space could be better described and explained.

AW2: We thank the reviewer for bringing this to our attention. In response, we have updated Section 4.1 to include further detail on the stochastic latent transition dynamics. We include our update here for your convenience:

As we have a discrete latent space, we formulate the transition dynamics to model the distribution over the next latent state $\mathbf{c}'$ given the previous latent state $\mathbf{c}$ and action $\mathbf{a}$. That is, we model stochastic transition dynamics in the latent space. We denote the probability of the next latent state $\mathbf{c}'$ taking the value of the $i^{\text{th}}$ code $\mathbf{c}^{i}$ as $p_{i} = P\_{\phi}(\mathbf{c}'=\mathbf{c}^{i} \mid \mathbf{c}, \mathbf{a})$. This results in the next latent state following a categorical distribution $\mathbf{c}' \sim \mathrm{Categorical} \left( p_{1},\ldots,p\_{|\mathcal{C}|} \right)$. We use a standard classification setup, where we use an MLP to predict the logits $\mathbf{l} = \{l_{1}, \ldots, l\_{|\mathcal{C}|}\}$. Note that logits are the raw outputs from the final layer of the neural network (NN), which represent the unnormalized probabilities of the next latent state $\mathbf{c}'$ taking the value of each discrete code in the codebook $\mathcal{C}$. The logit for the $i^{th}$ code is given by $l_{i} = d_{\phi,i} (\mathbf{c}, \mathbf{a}) \in \mathbb{R}$. We then apply the softmax operation to obtain the probabilities $p_{i}$ of the next latent state taking each discrete code in the codebook $\mathcal{C}$, i.e., $p\_{i} = \text{softmax}\_{i}(\mathbf{l})$. DCWM uses the discrete codes $\mathbf{c}$ as its latent state to make dynamics, reward $R\_{\xi}(\mathbf{c},\mathbf{a})$, value $\mathbf{q}\_{\psi}(\mathbf{c},\mathbf{a})$, and policy $\pi\_{\eta}(\mathbf{c})$ predictions. Our hypothesis is that learning these components with a discrete latent space will be more efficient than with a continuous latent space.

Please let us know if we have adequately improved our explanation and if you have any further suggestions to help us improve it.

Q1: The authors kept most hyperparameters fixed across all tasks, what parameters will affect the performance of some special tasks?

AQ1: Due to limited computing resources we did not fine-tune hyperparameters for each task, but instead aimed to use the same set of hyperparameters across all tasks, similar to prior work on model-based RL with world models. However, in more complex tasks where training takes longer, e.g., the Dog and Humanoid tasks, noise scheduling affects the training. Therefore, we reduced the exploration noise later.

Q2: Is the codebook a direct application of quantization-aware training (QAT)?

AQ2: While there are possible parallels between embedding-discretization and quantization-aware training, the latter is usually concerned with robustifying neural networks to quantization effects at the deployment stage. Hence, we would argue that both approaches differ both in their application and desired outcomes.

Q3: How to determine the weights of the expected code? Is this a hyperparameter of the algorithm?

AQ3: We would like to ask the reviewer to clarify this question if possible. We are unsure which weights you refer to. As a clarification, we do not discretize the weights of the neural network, but we discretize the embedding space - this is a different approach to quantization-aware training.

We are happy to see that you found our paper excellent, insightful and in a good state for publication. Please find below our clarifications to your questions:

Q1: I imagine that taking the expected code might sometimes result in an invalid state, as discussed here: "Whilst the expected value of a discrete variable does not necessarily take a valid discrete value, we find it effective in our setting." Is that simply then fed into the dynamics model (after normalization etc.) and the reward model MLPs?

AQ1: Yes you are correct, we simply feed the expected code as the input to the dynamics, reward, Q-function and prior policy.

Q2: Relatedly, when you normalize the code to [-1, 1], you claim it improves performance. Is there a reference for this or was this empirically observed? Would you be able to speak to why this is required? (One thought I had was it normalizes the different quantization levels, i.e. the two channels would otherwise have values ranging from [-5, 5] and [-3, 3]).

AQ2: The insight that normalization of the code to the range of [-1, 1] improves performance was indeed raised in Mentzer et al. (2023). We will add a mention of this in the new version of the paper.

Please also see our general response (and revised draft) where we have provided new experiments to strengthen our submission.

Keeping my score, as the authors have not provided any additional response.

We thank Reviewer BJj9 for their valuable feedback. We are glad they found the paper well-written and that our experimental evaluation provides evidence that the proposed algorithmic details are critical to its success.

W1: The one question that I have relates to the use of REDQ to reduce bias in the TD learning. Since the baseline does not seem to use this method, it begs the question of how much of the demonstrated performance is due to this component. It would be nice to see how DCWM compares to a version that uses the standard one or two Q functions to make sure the performance gap is in fact due to the choice of latent spaces.

Q1: What does the algorithm performance look like without using REDQ?

AW1: First of all, let us highlight that we did not use REDQ to obtain a performance improvement over TD-MPC2. On the contrary, we intentionally used the same strategy for the critic as was used in the TD-MPC2 paper. They used an ensemble of five critics and sub-sample two critics before calculating either the mean or the minimum of the two sub-sampled Q-values. So once again, let us clarify that we use the same critic strategy as the TD-MPC2 baseline. Nevertheless, to avoid confusion in the future, we have updated Line 293 to explicitly state that this was done in TD-MPC2:

"We also reduce bias in the TD target by following REDQ (Chen et al., 2021) and learning an ensemble of $N_{q}=5$ critics, as was done in TD-MPC2."

The reviewer is right that results using the standard double Q learning would provide further insights into what aspects of our method contribute to its performance. As such, we have run experiments using the standard double Q learning setup and report results in the Reacher Hard, Walker Walk, Dog Run and Humanoid Walk tasks. This experiment provides further insights into what parts of our algorithm contribute to its high performance, so we have added the figure to Appendix B.8. We thank the reviewer for bringing this up and hope this answers the question.

Q2: What are the computational costs or benefits to this method over TD-MPC in terms of runtime?

AQ2: We thank the reviewer for the good question. Our method (DCWM) has a similar runtime to TD-MPC2 when using our default hyperparameters. For example, in the Dog Run task, TD-MPC2 takes roughly 22.5 hours to reach 2M time steps while our method (DCWM) takes 28.5 hours. We will update the paper to include this.

We thank Reviewer BJj9 for their feedback. Please let us know if you have any further questions or if you have any suggestions that can help us improve our paper's score.

We are glad to see that you found the paper well-written and that you appreciated our research into the use of discrete latent spaces for model-based reinforcement learning with world models.

Q1: I think a comparison of the method with other approaches that are not based on the same embeddings ideas can be added.

W1: What it misses I think is a comparison of the method with other approaches that are not based on the same embeddings ideas. For example, I think MAMBA and Hungry hungry hippos (H3) apply to similar scenarios.

AW1: We want to highlight that we indeed consider, in our experiments and ablation studies, different baselines using different embedding spaces from continuous (TD-MPC2) to the use of discrete one-hot encodings used in DreamerV3. For comparisons between model-based reinforcement learning methods utilizing latent embeddings for planning/rollouts and RL methods not using latent embeddings we would like to refer you to the studies performed in e.g. [Hansen et al. 2023] and [Hafner et al., 2023]. Please, also see our ablation of different encoding styles in the latent embedding in Figure 7.

Lastly, we are unsure how a comparison of language-modelling methods such as Hungry hungry hippos (H3), which is used to encode text sequences in LLMs, would benefit this paper which tackles continuous control tasks in reinforcement learning settings. Extending the presented approach to language tasks appears to us to be out-of-scope and a serious extension of the presented approach, warranting a new paper. Therefore, we would like to invite the reviewer to motivate or clarify this comparison in a follow-up if possible to help us understand the benefits of adding H3.

Hi, we noticed you adjusted your score. Can you elaborate on why you decreased the score? We would love the opportunity to improve our paper based on your feedback.

The discussion period will be coming to an end soon and we would greatly appreciate your feedback. We have made significant improvements based on the feedback raised by all reviewers, specifically:

• Improved empirical study evaluating different embeddings (Figure 5) We have increased the number of tasks in our comparison of different latent space embeddings (Figure 5) from 4 tasks to 20 tasks (10 DMControl and 10 Meta-World). This significantly improves the statistical significance of our empirical study where we compare the impact of different latent space embeddings. The results show that DCWM's discrete latent space combined with using Gumbel-softmax sampling when making multi-step dynamics predictions, offers improved sample efficiency over continuous latent spaces.
• Shown that DCWM improves TD-MPC2 In Figure 21, we have shown that replacing TD-MPC2's continuous latent space with our discrete and stochastic latent space improves performance. This is an interesting result that highlights our findings are transferrable outside of our specific setup.
• Improved empirical study We have increased the number of Meta-World tasks that we evaluate DCWM on from 7 to 45, significantly improving the statistical significance of our empirical study. Please see Figures 3 and 14.
• Comparison of our FSQ vs VQ-VAE In Figure 15, we have shown the impact of using vector quantization (VQ) in place of our finite scalar quantization (FSQ). The results indicate that the VQ-VAE approach requires environment-specific adjustments which undermine its general applicability compared to our approach of using FSQ.

These changes have substantially strengthened our paper's contributions. We'd like to draw your attention to our in depth discussion with Reviewer HZL5, who initially gave our paper a lower score but has since increased after seeing these improvements. We believe these enhancements strengthen our paper significantly and we'd greatly appreciate it if you could reassess your score in light of these updates. We would be grateful for any additional feedback you may have, and once again, we'd like to thank you for your time and efforts.

We thank the reviewers for their continued input. We have been working hard to incorporate your constructive feedback and we have now updated and improved our manuscript once again. Below we summarize the improvements:

• How does TD-MPC2 perform with DCWM's latent space? (Appendix B.11) Reviewer HZL5 highlighted (and Reviewer Nctn agreed) that the best way to evaluate our quantized latent space was to directly replace the latent space in TD-MPC2 for a head-to-head comparison. We have now evaluated TD-MPC2 with DCWM's latent space in 10 DMControl and 10 Meta-World tasks. The results show that using DCWM inside TD-MPC2 improves sample efficiency in the 10 DMControl tasks, whilst in the 10 Meta-World tasks, the performance of all methods seems about equal. This suggests that discrete stochastic latent spaces can be advantageous in the context of world models for continuous control. This is an interesting result that we believe motivates further research into discrete and stochastic latent spaces for world models. We are happy with the new results and we thank Reviewer HZL5 for the suggestion. Please see Appendix B.11 and Figure 21 for further details and the full write-up. If reviewers agree, we are happy to move Figure 21 into the main text. However, for the ease of the rebuttal, we thought it would be easiest to add a new appendix solely dedicated to the new results.
• Increased number of tasks for ablations (Figure 5) We have increased the number of tasks in the latent space comparison (Figure 5) from 4 tasks to 20 tasks (10 DMControl and 10 Meta-World). This significantly improves the statistical significance of our empirical study where we compare the impact of different latent space designs. The results show that DCWM's discrete latent space combined with using Gumbel-softmax sampling when making multi-step dynamics predictions, offers improved sample efficiency over continuous latent spaces. Perhaps surprisingly, continuous stochastic latent spaces do not offer the same benefit. Please see Section 5.2 and Figure 5 for further details.
• Combined and moved figures To make space for these new results, we moved the DMControl and Meta-World aggregate performance metrics plots to Appendix B.1 and have added a figure in the main text to summarize the aggregate metrics for both DMControl and Meta-World (Figure 3).

Once again, we would like to thank all the reviewers for their time. Please let us know if there are any remaining questions or concerns. We would be happy to run any additional experiments that you think would improve our paper.