Paper: Harnessing Discrete Representations for Continual Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The paper only considers gridworld tasks, and it is hard to say whether the findings will generalize to more visually rich and complex planning tasks. A paper which only tells the reader something about representation learning in gridworlds presents a much less useful and interesting contribution than one whose insights are expected to generalize to a broader class of environments that are of interest to practitioners. This is the main issue which is lowering my score.

Author response:
**[A1]** We posted an in-depth, public comment addressing this point.

**[W1]** Unfortunately, we were unable to prepare additional experiments in time. While experiments in different environments would certainly be nice to have, there is existing prior work that experiments with discrete representations (including comparisons to continuous representations) in Atari ([Hafner et al.,  2020](https://arxiv.org/abs/2010.02193)), and our results align with the findings in that work. Standalone results in only two environments would generally be a low bar for us too, but the alignment of our results with previous work is what gives us more confidence in the results.

Reviewer comment (Experiments):
The RL task is extremely simple for the fixed representation: while the paper does point out that the discrete representations converge 2-4x faster than the representations obtained from the continuous reward model, both representations seem to converge almost immediately.

Author response:
**[A2]** As you point out, the RL task is not too difficult if a good representation is already learned. Our main goal in the episodic RL section is to build up motivation for doing continual learning where these differences can be significant. Looking at the Door Key environment, the VQ-VAE-based agent takes ~80k steps to converge, and the AE-based agent takes ~250k steps. For an agent that you only ever train once this may not matter, but efficiency is extremely important for an agent that needs to adapt to a constantly changing environment—i.e. continual RL. Our long term vision is to have continually learning agents that are much more efficient than both of these (imagine an agent that could learn in just several episodes!). But for now, we think that saving 170k steps is a good step in the right direction.

Reviewer comment (Evaluation):
The error bars on many of the figures are quite large, and there were some strange artifacts, in particular in Figure 3 where going from 512 to 1024 dimensions dramatically increased the error of one of the models, which are not explained. Given the relatively narrow scope and simplicity of the domains under consideration, this increases my concern about how the findings in this paper will generalize to other environments. I understand that evaluating on many tasks requires significant computational resources, but I worry that the results in this paper might be telling us more about the MiniGrid environment than about the general properties of the representations studied.

Author response:
**[A3]** The larger error bars for the end-to-end method in Figure 3 are produced by a bimodal distribution caused by a few diverging runs. We believe that some of the end-to-end runs diverge because the model has to backpropagate multiple loss terms through both the transition model and encoder over multiple time steps. In contrast, the other methods train the autoencoder and transition models separately and do not backpropagate over multiple time steps, which makes the optimization process much simpler. Your comment has made it apparent that our presentation of the data here was not clear enough. To fix this, we are switching the figure to plotting the median and adding a further explanation to the caption. We will upload a new manuscript within the next 24 hours that reflects these changes.

Our work aims to be a thorough study that improves the understanding of this area, and we think that reporting all results, including "not clean” ones, is an important part of helping others understand the realistic state of these problems. In regards to error bars, we are very careful to only make claims that are supported by statistically significant results throughout the paper, and the largest error bars primarily seem to be from standard RL (end-to-end) rather than the methods we investigate.

**[W3]** The divergence on high-dimensional embeddings is likely a product of an inherent instability in the approach, as we separately fine-tune each of the models with different latent dimensions. Backpropagating through multiple timesteps simply makes optimization much harder, and having multiple targets further complicates that process. While having diverging runs is not ideal, the conclusions of the experiment would not change whether or not the runs existed, as the problem only applies to the standard RL baseline.

Reviewer comment (Experiments):
The gap between FTA and VQ-VAE is much larger in one environment than the other, suggesting that there may be domain-specific properties that influence the relative performance of different representations. It would be useful to have a broader range of evaluation environments, even if these are limited to the minigrid setting to see whether this is the case in general.

Author response:
**[A4]** Assuming you are talking about the model learning setting (if not, let us know and we can follow up), the difference in scale is the result of our evaluation metric. The objective in the crossing environment causes the agent to spread out, whereas the objective in the door key environment concentrates the agent into just a few states. Because of this, the KL divergence is much higher in the door key environment when the agent is off by a few steps. This can be seen in Figures 11 and 12, where there is more overlap in the crossing environment.

Reviewer comment (Experiments):
In the description of the continual RL setting, it wasn't clear to me whether the representations were trained on a variety of shuffled layouts or only a single one. It would be very surprising to me if training a transition model on a single layout led to features that generalized so effectively (with apparently no degradation in performance on OOD inputs), but if this was the case it would merit further discussion.

Author response:
**[A1]** The representations are trained on multiple layouts. To clarify how the continual learning training works, both an autoencoder and a policy are continually trained in tandem as new data comes in, with the exception of an initial delay before the policy starts to learn. We will have an updated manuscript uploaded within the next 24 hours that clarifies this.

**[Q1]** We are happy to provide the full data. Figure 16 in the appendix depicts the data starting from timestep 0. We do not show the early part of training in the main body of the paper because no RL training is happening, and hence, the performance is not changing. Only the autoencoder is trained at the start of the process, and RL updates only begin after a delay. An overview of the training algorithm is described in Algorithm 2 in the Appendix.

Reviewer comment (Experiments):
One important factor influencing the usefulness of model-based representation learning is whether the data collected for the pretraining phase covers the state space. Since the transitions are mostly collected from random policies, this seems like it could be a significant issue in a domain like atari or mujoco, where a random policy samples a very limited state distribution. How dependent are the results in this paper on the pretraining distribution used for the learned representations?

Author response:
**[A2]** In the continual learning section, the representations are trained continually as the policy changes even after the initial pretraining period, so the representation changes as the distribution shifts. We found this to work well. We also tried experiments without the pretraining phase. This also worked, but we found the RL performance to be less stable, suggesting that there is a plasticity issue with the policy not being able to update to the changing representations. We think this could be solved by a number of recent interesting works on the stability-plasticity dilemma (e.g. [Dohare et. al 2023](https://arxiv.org/abs/2306.13812), [Nikishin et. al, 2023](https://openreview.net/forum?id=O9cJADBZT1)), but we have yet to test it.

**[Q2]** This is a very interesting direction to explore, as it really grasps at the heart of continual learning. In the future, we would love to see if our method can help adapt more quickly in those more significant changes. Even in the case of Minigrid, the changes can be quite significant when you consider that the model has only ever seen between 1 and 15 variations of the environment in its lifetime. This is reflected in spikes in reconstruction error seen in subplots (c) and (d) of Figure 6. That being said, there are environments that will undergo even more drastic changes as you mentioned. Perhaps testing on a domain like Atari, going from one Atari game to another, would be an interesting direction for future work.

Reviewer comment (Experiments):
Minor: the paper briefly mentions that it will focus on sample-based models, but the reasoning behind this choice isn't fully explained. Is there a reason why expectation models would not be suitable for representation learning? Or is this motivated by the desire to replicate the setting used in prior empirical work?

Author response:
**[A3]** We use sample-based models primarily to build on top of prior work (e.g. DreamerV3), but also because expectation models tend to be more limited, as they lose information about stochastic environments.

Reviewer comment (Experiments):
The experiments do look at scale, but the architectures / domains considered are very basic, so it is unclear how general these results are and how applicable they are to real-world scenarios.

Author response:
**[A1]** We posted a public comment in regards to the use of only Minigrid environments with an in-depth discussion on what that means for generalization to more complex problems.

Reviewer comment (Experiments):
Also the breadth of experiments is not very large in terms of domains when we consider that providing a comprehensive empirical study is a stated goal of the paper.

Author response:
**[A2]** The other comment mentioned above is also related, but in general, we erred on the side of thoroughness for this work. Keeping to the same domain allowed us to look at many properties of these representations while limiting the number of moving pieces. That also allowed us to perform many runs and have high confidence in the validity of our results.

Reviewer comment (Experiments):
The authors could potentially make the paper better by providing a bit more in terms of direct empirical proof of their interpretations i.e. the role of sparsity.

Author response:
**[A3]** We are combining the response to your final point under weakness with the response to the first question, as they are closely related.

**[A1]** After reading your comment on the limited amount of empirical proof related to sparsity, I ctrl+f searched over the paper and realized how little it was mentioned – thank you for catching this oversight! We are making a couple changes on this point:
1. We are adding some extra discussion to section 3.3.3, as the results of the experiment from that section directly support our interpretations on sparsity. This interpretation of the results was previously missing.
2. As you point out in your comment, the question of how the sparsity affects continual learning performance is very interesting, so we are currently running some additional experiments to find out. The original paper on FTA ([Pan et al., 2021](https://arxiv.org/abs/1911.08068)) actually runs an experiment just like this (Figure 7a), where they vary the level of sparsity. They find that it helps, but there is an optimal amount of sparsity, and too much or too little sparsity doesn’t perform as well. The preliminary results of our currently running experiments are looking to draw very similar conclusions. We are posting an initial response now to allow time for responses, but we will follow up again after we have added the new section on sparsity.

Our interpretations are also influenced by many other works that have found similar connections to sparsity (e.g., [Ratitch & Precup, 2004](https://link.springer.com/chapter/10.1007/978-3-540-30115-8_33), [Liu et al., 2018](https://arxiv.org/abs/1811.06626), [Mocanu et al., 2018](https://www.nature.com/articles/s41467-018-04316-3)).

We are posting an update to let you know that we ran some final experiments on sparsity, and the final draft of the paper we uploaded has a new section on sparsity in the Appendix (A.6). To summarize the results, we found a clear, positive correlation with sparsity and continual learning performance up to a certain point. However, going beyond that point and having too much sparsity again becomes harmful.

Reviewer comment (Evaluation):
The authors mention a few times that there is a connection between the success of discrete representations in continual RL and representation sparsity. Did you try actually measuring this sparsity? I guess the size of each one-hot dictates this sparsity level, so it could be interesting to understand the connection between this and continual learning performance.

Author response:
**[A1]** After reading your comment on the limited amount of empirical proof related to sparsity, I ctrl+f searched over the paper and realized how little it was mentioned – thank you for catching this oversight! We are making a couple changes on this point:
1. We are adding some extra discussion to section 3.3.3, as the results of the experiment from that section directly support our interpretations on sparsity. This interpretation of the results was previously missing.
2. As you point out in your comment, the question of how the sparsity affects continual learning performance is very interesting, so we are currently running some additional experiments to find out. The original paper on FTA ([Pan et al., 2021](https://arxiv.org/abs/1911.08068)) actually runs an experiment just like this (Figure 7a), where they vary the level of sparsity. They find that it helps, but there is an optimal amount of sparsity, and too much or too little sparsity doesn’t perform as well. The preliminary results of our currently running experiments are looking to draw very similar conclusions. We are posting an initial response now to allow time for responses, but we will follow up again after we have added the new section on sparsity.

Our interpretations are also influenced by many other works that have found similar connections to sparsity (e.g., [Ratitch & Precup, 2004](https://link.springer.com/chapter/10.1007/978-3-540-30115-8_33), [Liu et al., 2018](https://arxiv.org/abs/1811.06626), [Mocanu et al., 2018](https://www.nature.com/articles/s41467-018-04316-3)).

We are posting an update to let you know that we ran some final experiments on sparsity, and the final draft of the paper we uploaded has a new section on sparsity in the Appendix (A.6). To summarize the results, we found a clear, positive correlation with sparsity and continual learning performance up to a certain point. However, going beyond that point and having too much sparsity again becomes harmful.

Reviewer comment (Evaluation):
In section 3.3.3 I am a bit unclear about what the conclusion is for these experiments. What are some of the insights you believe they present? Is this actionable in a general sense?

Author response:
**[A2]** The original takeaway from section 3.3.3 was that the implicit design of representing discrete latents with a one-hot encoding scheme is an incredibly important (and often overlooked) choice. Our experiment shows that it is not necessarily the “discreteness” of the representation as a whole that matters, but rather the sparse and element-wise binary nature of the representation that matters. We are adjusting the section to make this clearer, and also to discuss how this is evidence for the importance of sparsity.

Reviewer comment (Experiments):
My only slightly major concern on this work is that all experiments are done on the minigrid environments. The pixel observations for this environment are simple which might have contributed to the easier process of discretization. I think to improve the quality of the paper, it will be important to see if the same conclusions can be made when using more visually complex environments, such as Atari 100k, which is also performed in the dreamer paper.

Author response:
**[A1]** We posted an in-depth, public comment in regards to the use of only Minigrid environments.

Reviewer comment (Reproducibility):
Will be great if the authors can add a pseudocode of the algorithm in the appendix to improve the clarity.

Author response:
**[A2-4]** We are working to update the manuscript to address all of the minor concerns and will have a new revision posted within the next 24 hours. Thank you for the suggestions.

Reviewer comment (Presentation):
Figure 14: Although the trend of the result is clear, this should be plotted with mean and standard deviation to improve readability.

Author response:
**[A2-4]** We are working to update the manuscript to address all of the minor concerns and will have a new revision posted within the next 24 hours. Thank you for the suggestions.

Reviewer comment (Presentation):
Some spacing gaps in the appendix can be reduced, for example using vspace in latex.

Author response:
**[A2-4]** We are working to update the manuscript to address all of the minor concerns and will have a new revision posted within the next 24 hours. Thank you for the suggestions.

Reviewer comment (Experiments):
In section 3.3.2, the authors mentioned that “In the plot, an interesting pattern emerges: the performance of all methods become indistinguishable beyond a certain size of the world model”. Perhaps this might hint that the environment is too simple?

Author response:
**[A1]** Yes, this is a great catch and precisely what we aim to show with section 3.3.2. It doesn’t matter what which of our representation methods is used when the environment is simple because all of the methods will perform near-optimally. As the environment becomes more complex, however, discrete representations shine because they are able to learn more with less modeling capacity. It is this result that motivates the use of these representations in continual learning, as continual reinforcement learning is all about how to maximize reward when the agent cannot perfectly model the world (because the world is continually changing from the agent’s perspective).

Reviewer comment (Writing):
In section 3.3.3, the authors mentioned that “It is alternatively possible that the discrete world model performs better simply because the VQ-VAE learns different information that is more conducive to world modeling.” What other different information do the authors think VQ-VAE is learning?

Author response:
**[A2]** We really mean to point out the following: latent spaces are defined both by the information they represent (information) and by the way in which they represent that information (representation). Our goal is to measure how representation alone affects performance, so we ideally want to control for the informational content of the latent space. In section 3.3.3, we do not really expect that the autoencoders are learning different information, but we need to check this possibility to be thorough. We have reworded part of the section to make this clearer.