Locality Sensitive Sparse Encoding for Learning World Models Online

4. NN training details for the Full-replay experiment.

The Full-replay experiments did not involve resetting the NN’s weights. We have tried resetting the weights but empirically found this required much more compute budget since it trained from scratch every time, which is impractical for a lifelong RL agent. Besides, it raises another design choice about when to reset, which may be not trivial since there are no explicit task boundaries.

Instead, we “finetune” a single model continuously, with the train/holdout split (with a ratio of 8:2) to determine when to stop the training and avoid overfitting. We agree that training without resetting weights is likely to overfit early data sets, but empirically this seems okay. We closely follow the codes of MBPO (https://github.com/jannerm/mbpo), which also store all previous experiences and do periodic training without weights resetting, for our implementation. We have updated the paper and provided more details about NN training in Section 5.1 on page 8.

5. Why isn’t the Full-replay MBRL better than the model-free baseline?

The expectation that Full-replay MBRL should outperform the Model-free baseline is based on an assumption: an accurate world model is learned such that it can synthesize transitions on which planning updates accelerate value/policy learning. However, achieving so is non-trivial. As you have mentioned, successful training requires stopping training properly. Though we have employed a train/holdout split to determine when to stop, our parameters may not be optimally set. Apart from this, there is more brittleness in NN training in this case. For example, as you mentioned, not resetting the weights may make NNs overfit early data; but when to reset is hard to know because there are no explicit task boundaries; even if we know when, resetting and training from scratch is even more time-consuming and impractical for lifelong agents. In our experiments, we choose not to reset the weights for efficiency consideration, and this might be a source of training difficulty that leads to inferior Full-replay results, especially for Acrobot environment.

6. Other minor corrections.

Thank you very much for the detailed feedback! We have updated the paper to address all the mentioned issues.

Thank you again for providing insightful comments which helped us to improve our paper, and we hope our responses and updates in the text were able to address any remaining concerns. Please do let us know if you have any further questions as well as what would be expected for score improvement.

Thank you very much for appreciating the originality of our paper, as well as your detailed feedback and suggestions. We have updated the paper with text changes highlighted in yellow, and elaborate on each of the questions below.

1. There is no experiment highlighting the main issue that the paper is trying to address: NN-based world models suffer from catastrophic interference due to non-stationary data.

Experiments in Section 5 are designed to compare our method with different variants of NN-based world models while keeping the value/policy learning the same. Except for the “Model-free” setting, all other baselines train NN-based world models. From Figure 4, we can observe that NN-based world models trained online (referred to as “Online” in the legend) show inferior performance. This is an indication that they suffer from catastrophic interference due to non-stationary data. As a result, the online learned NN world model could only generate synthetic transitions of poor quality, deteriorating the agent performance. Applying continual learning techniques (such as “SI” and “Coreset”) alleviates the forgetting issue to some extent. And learning world models with our method outperforms other baselines for most environments.

2. Why aren't SoTA MBRL algorithms (e.g., Dreamer/TD-MPC) included for comparison?

First, we hope that we can get aligned to the point that NN-based MBRL suffers from catastrophic forgetting, including the SoTA algorithms. The arguments are following,

1. Catastrophic forgetting is a universal phenomenon for neural networks.
2. Existing end-to-end NN-based MBRL methods, such as Dreamer, TD-MPC, and MuZero, etc. rely heavily on techniques to make data more stationary, for example, maintaining a large replay buffer, updating the target network once in a while, or simultaneously running multiple environments. When these components are removed, they would fail to work.
3. However, when they fail we don't know how much is due to the forgetting in policy/value/model because all components are entangled and trained end-to-end in these methods.

If we're aligned on the point that NN-based MBRL suffers from catastrophic forgetting due to non-stationary data, then what we're studying is that we focus on making the world model fully online, which relies on none of the above techniques. With this scope, we only compare our world models and NN-based world models using the Dyna architecture, because it decouples policy/value learning from model learning, making an apple-to-apple comparison on model learning possible.

3. Can your method work for more difficult tasks like Humanoid or handle image-based observations?

Currently, we have not validated it for more challenging tasks like you mentioned. In Humanoid, the dimension of the state is very high, posing a curse of dimensionality to our method. For image-based observations, finding a compact representation with a non-trainable encoder is non-trivial, and our attempts so far do not yield positive results. We have added a Limitation section on page 9 to address them in the paper.

Thank you very much for highlighting our comprehensive experimental results, as well as your valuable feedback on the paper presentation. We have updated the paper with text changes highlighted in yellow, and elaborate on each of the questions below.

1. Significant discussion on catastrophic forgetting in the introduction but little discussion in the main text and left in the appendix.

Thanks for this suggestion to improve our paper presentation. We have updated the paper such that catastrophic forgetting is emphasized more in the main text, especially at the end of Section 3 on page 6. Due to the space constraints, however, we still use Appendix A for more detailed discussions.

2. Other issues regarding the notations.

We have updated accordingly in the paper.

Thank you again for your feedback to help us refine the paper quality, and we hope our responses and updates in the text were able to address any remaining concerns.

Thank you very much for recognizing that the problem we attempt to tackle (learning models online) is important in RL, as well as your constructive feedback and suggestions. We have updated the paper with text changes highlighted in yellow, and elaborate on each of the questions below. Note that the points below do not match those in the original questions due to reorganization, but they cover all questions.

1. Why is the particular sparse representation chosen?

We adopt linear models to have guaranteed online learning performance without forgetting. However, linear models cannot fit complex functions. Adding a nonlinear random feature turns linear models into universal approximators [1], but this usually requires a very wide layer to have good capacity. We make our high-dimensional representation sparse to achieve efficient incremental update. In addition, we also compare our sparse representation with other encoding methods and we show favorable results (Table 1). The idea of our sparse update is inspired by the model mixer in [2].

We have updated the paper on page 6 with a “concluding remark” to make this point clearer.

2. Why is FTA not included in the experiments?

Fuzzy Tiling Activation (FTA) [3] proposes an effective way to reduce catastrophic forgetting of neural networks in online continual learning problems, by learning sparse activations inside the NNs. In Section 4.2, we only borrow their nonstationary online learning setting (the Piecewise Random Walk) and show our method is more capable than NNs of learning temporally correlated data. However, we did not include experiments to compare FTA and our method in either SL or RL settings, for the reasons below.

In SL, from Figure 3.c of [3], we can see FTA slightly suffers from nonstationarity if the data is highly correlated. The reason could be that FTA only mitigates the catastrophic forgetting issue in NNs, but it does not fully solve the problem. Therefore, we explore a more assured pathway in this work, by which a Follow-the-Leader solution can be obtained online efficiently, aiming to eliminate catastrophic forgetting. As Figure 3 in our paper suggests, Losse-FTL can attain very low error for all levels of correlation, showing its non-forgetting property.

In RL, we have attempted to integrate the official FTA codes into NN-based world models, but we found it hard to make it work well (perhaps due to hyper-parameter tuning, such as the tiling dimension $k$ and the sparsity factor $\eta$, or where the FTA layer should be applied). Nevertheless, we compared our method with other representative continual learning techniques, including Synaptic Intelligence and Coreset, to show its effectiveness.

3. Do you update both your model and representation every environment time step?

Our high-dimensional sparse representation is not learnable with a prefixed random projection matrix. Thus we only update the model (the linear weights) but not the representation.

For all RL experiments, we keep the update interval as $25$ time steps for discrete control tasks and $250$ for continuous control tasks to accelerate the experiments. However, we note that our updates are much more efficient compared to NN-based ones with replay (referring to Figure 5) because the latter needs to train until convergence over accumulated experiences.

4. Does the benefit come from policy learning or model learning?

Thanks for this insightful question. We agree that sparse representation helps continual policy learning [3,4,5], as it could potentially reduce the catastrophic interference in NNs. However, in this work, we utilize the high-dimensional sparse representation only during model learning for Losse-FTL, while keeping policy/value learning the same for apple-to-apple comparison. Therefore, we believe the benefit indeed comes from model learning; better model learning ensures a better quality of synthetic transitions to be used for policy/value learning.

5. Why does the proposed algorithm sometimes even outperform the Full-replay baseline?

We agree that theoretically, Full-replay should be as good as ours in eliminating forgetting. We hypothesize that the reason for our method outperforming Full-replay in some tasks is due to the criterion we adopt for “convergence” – at which point we stop the NN training, while ours is solved incrementally in closed form. Apart from the convergence criterion, there is more brittleness in NN training in this case. For example, not resetting the weights may make NNs overfit early data; but when to reset is hard to know because there are no explicit task boundaries; even if we know when, resetting and training from scratch is even more time-consuming and impractical for lifelong agents.

We have updated the paper on page 8 about “Full-replay” to clarify the experimental settings.

6. Other issues.

Thanks for pointing out the issue of Algorithm 1 and the need for better notations. We have updated them directly in the text. To be more specific on Dyna architecture, we remove the original Algorithm 1 and use a diagram of Dyna in Figure 2 to give a more intuitive illustration in Section 2.2.

It also took us some time to consider the term “world model” in this context. We agree that it can refer to tasks that are more challenging, as envisioned in [6]. We hope to retain it for several reasons. (1) Simply using “model” is common in RL contexts, but we feel it might be confusing since “model” is frequently referred to with different contexts, e.g. linear model. In the title, especially, without some hints about RL, people may not recognize “model” refers to the world model. (2) Other terms like “dynamics model” only describe the dynamics part but not the reward part. We also thought of the “environment model”, however, it is not very commonly seen in literature. In the paper, we also tried to put the term “world model” in the context of MBRL, to avoid misunderstanding. Please let us know if there are more suitable terms.

Thank you again for your detailed feedback and suggestions, and we hope our responses and new experiments were able to address any remaining concerns. Please do let us know if you have any further questions as well as what would be expected for score improvement.

References

[1] Huang, G. B., Chen, L., & Siew, C. K. (2006). Universal approximation using incremental constructive feedforward networks with random hidden nodes. IEEE Trans. Neural Networks.
[2] Knoll, B., & de Freitas, N. (2012). A machine learning perspective on predictive coding with PAQ8. In 2012 Data Compression Conference.
[3] Pan, Y., Banman, K., & White, M. (2021). Fuzzy tiling activations: A simple approach to learning sparse representations online. In ICLR.
[4] Liu, V., Kumaraswamy, R., Le, L., & White, M. (2019). The utility of sparse representations for control in reinforcement learning. In AAAI.
[5] Lan, Q., & Mahmood, A. R. (2023). Elephant Neural Networks: Born to Be a Continual Learner. In ICML Workshop on High-dimensional Learning Dynamics.
[6] LeCun, Y. (2022). A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62.

Dear Reviewer 2Y3F,

Thank you again for your valuable comments and suggestions, which are very helpful to us. We have posted responses to the proposed concerns.

We understand that this is quite a busy period, so we sincerely appreciate it if you could take some time to return further feedback on whether our responses resolve your concerns. If there are any other comments, we will try our best to address them.

Best,

The Authors