Paper: Locality Sensitive Sparse Encoding for Learning World Models Online (ICLR-2024)

Reviewer comment (Experiments):
There is no comparison to state-of-the-art MBRL strategies that use world models, e.g. Dreamer/TD-MPC. As such, there is no experiment highlighting the main issue that the paper is trying to address: that NN-based world models suffer from catastrophic interference due to non-stationary data.

Author response:
**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.

Reviewer comment (Experiments):
It is unclear whether the proposed method can scale to difficult environments like humanoid or handle image-based observations. Even if the method cannot do this, it should be addressed in the text.

Author response:
**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.

Reviewer comment (Experiments):
The paper should clarify what benefits the proposed method has over other world model methods like Dreamer. For example, it would be valuable to show a state-of-the-art world model method failing because it cannot handle non-stationary data.

Author response:
**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.

Reviewer comment (Reproducibility):
There is insufficient detail on the training of the NNs in the experiments. In particular, did the full-replay experiment involve resetting the neural network's weights? If so, what initialization was used? When was the NN training stopped? Was the data split into train/validation sets and used to stop training when the validation loss stopped improving?

Author response:
**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.

Reviewer comment (Reproducibility):
The implementation details are important to ensure the NNs don't overfit on early data sets, but the paper does not provide enough information to clarify this.

Author response:
**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.

Reviewer comment (Evaluation):
It is unclear why the full-replay strategy, which is model-based, does not appear to have better sample efficiency than the model-free experiment. This raises the question of whether high-dimensional features sacrifice sample efficiency in favour of formulating a linear model that can handle non-stationary data.

Author response:
**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.

Reviewer comment (Writing):
In the first paragraph of Section 2.2, the reward function is defined as \( R(\mathbf{s}, \mathbf{a}, \mathbf{s}') \) but in the optimal policy equation it is written as \( R(\mathbf{s}, \mathbf{a}) \), which is inconsistent.

Author response:
**6. Other minor corrections.**

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

Reviewer comment (Writing):
In the third paragraph of Section 2.2, "We firstly formulate" should be corrected to "We first formulate".

Author response:
**6. Other minor corrections.**

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

Reviewer comment (Writing):
In Section 2.1, the sentence "When the input is a convex set \( \mathcal{S} \), the prediction a vector \( \mathbf{w}_{t} \in \mathcal{S} \)" does not read properly and should be clarified.

Author response:
**6. Other minor corrections.**

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

Reviewer comment (Writing):
The symbol \( \rho \) is used to denote both the initial state distribution and the dimension of the grids in Section 3.2, which may cause confusion.

Author response:
**6. Other minor corrections.**

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

Reviewer comment (Presentation):
The value of \( \delta \) in Fig. 1 is not given.

Author response:
**6. Other minor corrections.**

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

Reviewer comment (Writing):
In the first sentence of the abstract, "model-based RL has better sample efficiency" is stated without clarification of "better than what?". Presumably, it is better than its model-free counterparts, but this should be explicit.

Author response:
**6. Other minor corrections.**

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

Reviewer comment (Presentation):
It is unusual to end the paper with a section titled "Summary". It is recommended to change this to "Conclusion".

Author response:
**6. Other minor corrections.**

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

Reviewer comment (Experiments):
Is there any reason for why the particular sparse representation learning method is chosen? Furthermore, in the experiments part, FTA should also be compared as a baseline. It is unclear why you compare it in a supervised learning setting but omit in a RL setting. The performance on a SL setting does not invalidate/validate another. As an empirical paper, I think a rigorous comparison is necessary.

Author response:
**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.

Reviewer comment (Reproducibility):
Could you clarify do you update both your model and representation every environment time step?

Author response:
**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.

Reviewer comment (Experiments):
There is a critical weakness in the paper: the paper claims to develop a sparse representation-based approach for model learning, but it is not justified the reported benefits come from the use of the sparse representation for policy learning or for model learning. Note that the former has been extensively studied. In general, a full replay method should be the best in mitigating catastrophic forgetting, but the empirical results reported that the proposed algorithm can sometimes even outperform full replay. That raises a natural question that the benefit mainly comes from the policy learning part by using sparse representation, rather than the proposed model learning part.

Author response:
**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.

Reviewer comment (Presentation):
In Algorithm 1, it is better to be specific and use the title Dyna architecture, rather than MBRL, as there are numerous MBRL algorithms and not everyone is as Algorithm 1 described.

Author response:
**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.

Reviewer comment (Writing):
The term “world model” might intrigue the readers to see much more challenging tasks than the paper presented, this can be seen by other papers using such terms. It is better to rephrase it to be more precise.

Author response:
**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.