Locality Sensitive Sparse Encoding for Learning World Models Online
Authors: Zichen Liu
Chao Du
Wee Sun Lee
Min Lin
Venue: ICLR 2024 poster (2024)
Keywords: model-based rl
online learning
incremental learning
catastrophic forgetting
Primary area: reinforcement learning

Abstract: Acquiring an accurate world model $\textit{online}$ for model-based reinforcement learning (MBRL) is challenging due to data nonstationarity, which typically causes catastrophic forgetting for neural networks (NNs). From the online learning perspective, a Follow-The-Leader (FTL) world model is desirable, which optimally fits all previous experiences at each round. Unfortunately, NN-based models need re-training on all accumulated data at every interaction step to achieve FTL, which is computationally expensive for lifelong agents. In this paper, we revisit models that can achieve FTL with incremental updates. Specifically, our world model is a linear regression model supported by nonlinear random features. The linear part ensures efficient FTL update while the nonlinear random feature empowers the fitting of complex environments. To best trade off model capacity and computation efficiency, we introduce a locality sensitive sparse encoding, which allows us to conduct efficient sparse updates even with very high dimensional nonlinear features. We validate the representation power of our encoding and verify that it allows efficient online learning under data covariate shift. We also show, in the Dyna MBRL setting, that our world models learned online using a $\textit{single pass}$ of trajectory data either surpass or match the performance of deep world models trained with replay and other continual learning methods.

Published as a conference paper at ICLR 2024
LOCALITY SENSITIVE SPARSE ENCODING FOR LEARNING WORLD MODELS ONLINE
Zichen Liu" 13
Chao Du
Wee Sun Lee"
Min Lin'
Sea Al Lab
National University of Singapore
ABSTRACT
Acquiring an accurate world model online for model-based reinforcement learn- ing (MBRL) is challenging due data nonstationarity, which typically causes catastrophic forgetting for neural networks (NNs). From the online learning per- spective, a Teeleeer (FTL) world model is desirable, which optimally fits all previous experiences at each round. Unfortunately, NN-based models need re- re-training on all accumulated data every interaction step achieve FTL. which is computationally expensive for lifelong agents. In this paper, we revisit models that can achieve FTL with incremental updates. Specifically, our world model is a linear regression model supported by nonlinear random features. The linear part ensures efficient FTL update while the nonlinear random feature empowers the fitting of complex environments. To best trade off model capacity and compu- tation efficiency, we introduce a locality sensitive sparse encoding, which allows to conduct efficient sparse updates even with very high dimensional nonlinear features. We validate the representation power of our encoding and verify that it allows efficient online learning under data covariate shift. We also show, in the Dyna MBRL setting, that our world models learned online using a single pass of trajectory data either surpass match the performance of deep world models trained with replay and other continual learning methods.
INTRODUCTION
World models have been demonstrated to enable sample-efficient model-based reinforcement learn- ing (MBRL) (Sutton, 1990; Janner et al., 2019; Kaiser et al., 2020; Schrittwieser et al., 2020), and are deemed as a critical component for next-generation intelligent agents (Sutton al., 2022; LeCun, 2022). Unfortunately, learning accurate world models in an incremental online manner is challeng- ing, because the data generated from agent-enviromment interaction is nonstationarily distributed, due to the continually changing policy, state visitation, and even environment dynamics. For neural network (NN) based world models, the nonstationarity could lead to catastrophic forgetting (Mc- Closkey & Cohen, 1989; French, 1999), making models inaccurate at recently under-visited regions (as illustrated the top row Figure 1). Planning with inaccurate models is detrimental as the model emrors can compound due to rollouts, resulting in misleading agent updates (Talvitie, 2017; Jafferjee et al., 2020; Wang et al., 2021; Liu et al., 2023). To attain a world model of good accu- racy over the observed data coverage, NN-based methods often maintain all collected experiences from the start of the environment interaction and perform periodic re-training, possibly at every step, resulting in growing computation cost for lifelong agents.
In this paper, we aim to develop a world model, in the Dyna MBRL architecture (Sutton, 1990; 1991), that learns incrementally without forgetting prior knowledge about the environment. We notice that re-training NN-based world models every step on previous observations till convergence resembles the concept of Dollow-The-Leadeee (FTL. in online learning (Shalev-Shwartz al., 2012). While re-training the NN on all previous data is prohibitively expensive, especially under the stream- ing setting in RL, specific types of models studied in online learning can achieve incremental FTL with only a constant computation cost. To this end, revisit the classic idea of learning a lin- ear regressor on top of pon-linear random features. The loss function of such models is quadratic with respect to the parameters, satisfying the online learning requirement. It may seem retrograde choice given all the success of deep learning, and the findings that a shallow NN needs to be ex- ponentially large match the capability of deep one (Eldan & Shamir, 2016; Telgarsky, 2016).

---

Published as a conference paper at ICLR 2024
G
Unvisited regions under current policy
Visited regions
policy
St. when d(see / (+) >
St when d(sees. / [e). s.
Ours
TT0
Figure 1: This Gridworld environment requires the agent to navigate from the start position ("S")) to the goal location ("G"") with shortest path. The tabular Q-lcarning agent starts from random policy тo and improves to get better policies T. Ty. leading to narrower state visitation towards the optimal trajectories (the yellow regions). Due such distributional shift, the NN-based model (top) may forget recently under-visited regions, even though has explored there before. The red circles indicate erroneous predictions where the Euclidean distance between the ground truth next state and the prediction is greater than a threshold (6 0.05). In contrast, our method (bottom) learns online, and at cach step incrementally computes the optimal solution over all accumulated data, thus is resilient to forgetting.
Nevertheless, believe it is worth exploring under the RL context for three reasons: (1) In RL., data is streamed online and highly non-stationary. advocating for models capable incremental FTL. (2) Many continuous control problems have moderate number of dimensions that could fall within the capability of shallow models (Rajeswaran et al., 2017). (3) Sparse features can be employed to further enlarge the model capacity without extra computation cost (Knoll de Freitas, 2012).
Inspired by Knoll & de Freitas (2012), we propose an expressive spurse non-linear feature repre- sentation which we call locality sensitive sparse encoding. Our encoder generates high- nigh-dimensional sparse features with random projection and soft binning. Exploiting the sparsity, we further develop an efficient algorithm for online model learning, which only updates a small subset of weights while continually tracking solution to the FTL objective. World models learned online with our method are resilient to forgetting compared to those based on NNs (see the bottom row of Figure for intuition). We empirically validate the representational advantage of our encoding over other non-linear features, and demonstrate that our method outperforms NNs the online supervised learning setting (Orabona, 2019; Hoi al., 2021) well as the model bbaeed reinforcement learning setting (Sutton, 1990) with models learned online.
PRELIMINARIES
In this section, we first recap the protocol of online learning and introduce the strategy. Then we revisit Dyna, a classic MBRL architecture, where we apply our method to learn world models online.
2.1 ONLINE LEARNING
Online learning refers to a learning paradigm where the learner needs to make a sequence of accurate predictions given knowledge about the correct answers for all prior questions (Shalev- Shwartz et al., 2012). Formally, at round i, the online learner is given a question x E X and asked to provide an answer it, which we denote as (x(x). letting h, E H. X + y be model in the hypothesis class built by the learner. After predicting the answer, the learner will be revealed the ground truth y. E y and suffers loss (x+), yr). The learner's goal to adjust its model to achieve the lowest possible regret relative to H defined as (H)
max EH EEE ((hx(x)), - ET, (((((x)). (xx)..))) When the input is a convex set S, the pre- diction is a vector w, E S, and the loss ye) is convex, the problem is cast as online convex optimization. Absorbing the target into the loss, (w(:)) ((w,, the regret with respect
2

---

Published as a conference paper at ICLR 2024
competing hypothesis (a vector u) is then defined as ( Σو= le(we) - - (u). One intuitive strategy is for the learner to predict a vector w, at any online round such that it achieves minimal loss over all past rounds. This strategy is usually referred to as Seoow-The-Ledee (FTL):
Vt, arg min WES
((w).
2.2 MODEL-BASED REINFORCEMENT LEARNING WITH DYNA
Reinforcement learning problems are usually formulated with the standard Markov Decision Pro- cess (MDP) M A, P, R 7, Po), where S and A denote the state and action spaces, P(s'ls, a) the Markovian transition dynamics, R(s, a, s') the reward function, 7 E (0, 1) the discount factor, and Po the initial state distribution. The goal of RL is to learn the optimal policy that maximizes the discounted cumulative reward: 'e" arg max,, Ex, E== ?" RR (s. St,8)+1)
s ~ Po]. Model-based RL solves the optimal policy with learned world model. Existing NN-based MBRL methods (Schrittwieser et al., 2020; Hansen et al., 2022; Hafner et al., 2023) have shown state-of-the- art performance on various domains but rely heavily on techniques to make data more stationary, such as maintaining large replay buffer or periodically updating the target network. With these components removed, they would fail to work. When they fail, however, it is unclear how much is entan-
due to the forgetting in policies, value functions, world models because all components gled and trained end-to-end these methods. Therefore, we focus on Dyna, which learns a world model (or simply model in RL literature) alongside learning the hase agent' in decoupled manner.
Figure illustrates the Dyna architecture that we employ in this paper study online model learning. The base agent will act in the environment and collect environment experiences, which can be used for learning the world model to mimic the environment's behavior. With such model learned, the agent then simulate with it to synthesize model experiences for planning updates. Search control defines how the agent queries the model, and some commonly used ones include predecessor (Moore & Atkeson, 1993), on-policy (Janner et al., 2019) and hill-climbing (Pan et al., 2019). Next, we define how to learn and plan with the model.
Base Agent
acting
planning updates
direct RI. updates
Simulation
Experience
search
model learning
control
World Model
Model learning. We first formulate how to learn the dy- namics model m(s'|s,
(s,ls, a) with environment experiences. Instead of directly predicting the next state, model the dynamics the change of the successive states š' = s At m(As]s, a), where At is the discrete time step. This technique enables more stable predictions and widely adopted (Chua et al., 2018; Janner et al., 2019; Yu et al., 2020). The dynamics can either be modeled as a stochastic or a determin- istic function, with the former predicting the stateddependent variance besides the mean to capture aleatoric uncertainty in the environment (Chua et al.. 2018). We focus on learning a deterministic dynamics model since we are mainly interested in environments without aleatoric uncertainty. As shown in Lutter et al. (2021), deterministic dynamics models can obtain comparable performance with stochastic counterparts. Even when the environment exhibits stochasticity, generating rollouts with fixed variance is also as effective as learned variance (Nagabandi et al., 2020). Therefore, dynamics model learning is regression problem with the objective
Figure 2: The Dyna architecture.
integrator which estimates
(...a..)) [| a) As|
(2)
where Danv is the environment experiences and As = s' s. The reward model learning follows a similar regression objective and we will omit it throughout this paper for brevity.
Model planning. In Dyna-style algorithms, the learned model is meant for the agent to interact with, as substitution of the real environment. When the agent "plans" with the model. it queries the model at (5, ā) E S x A to synthesize model experiences D, = ($, a, f, S'), which serves as a data source for agent learning. There are various ways to determine how query the model. For example,
We use the term base agent to refer other components than the model, e.g., the value function and policy.
3

---

Published as a conference paper at ICLR 2024
the classic Dyna-Q (Sutton, 1990) conducts one-step simulations state-action pairs uniformly
sampled from prior experiences. More recently, MBPO (Janner et al., 2019) generalizes the one-step simulation to mhort-horizon on- on-policy rollout branched from observed states and provides theoretical analysis for the monotonic improvement of model-based policy optimization.
LEARNING WORLD MODELS ONLINE
The agent-environment interaction generates temporally correlated data stream (st, a, Tt, 1t At each time step, the world model observes
[Sr, [s, E RR predicts the next state and suffers quadratic loss (f( [x+) Yell2: where S and A are the dimensionalities of state and action spaces, Yr = ($r41 - S:) R is the state difference target. In most model-based RL algorithms built with neural networks, the model learning is conducted in a supervised offline manner. with periodic re-training over the whole dataset (see the objective in Equation 2). Such methods achieve FTL for NNs but are inefficient due to the growing size of the dataset. We develop in this section an online method for learning world models with efficient incremental update on every single environment step.
3.1 ONLINE FOLLOW-THEE
MODEL LEARNING
Our world model m RDx on an expressive and sparse non-linear feature o(x) E R.. The FTL objective in Equation thus converts to least squares loss per time step:
((x) W employs a linear function approximation with weights W
Vt, W.
arg min 1 W- -Y--1l WERDXS
(3)
where 4, = [Ø(x),..... o (x, T E R denotes non-linear features for observations accumu- lated until time step Y, yy..... y R is the corresponding targets, and || || denotes the Frobenius norm. We can expect our world model with the weights satisfying Equation to have good regret, as the regret of FTL for online linear regression problems with quadratic loss is Regrety (W) = (log T') (Shalev-Shwartz et al., 2012; Gaillard et al., 2019). As a consequence, our online learned world model will be immune to forgetting, rendering suitable for lifelong agents.
In the following sections, we present how to attain a feature representation that is expressive enough for modeling complex transition dynamics (Section 3.2), as well as how its sparsity helps to achieve efficient incremental update (Section 3.3).
3.2 FEATURE REPRESENTATION
We construct high-dimensional sparse features using random projection followed by soft binning. The input vector x, E R is first projected into feature space a x - Px, where P E (S+A)) is a random projection matrix sampled from a multivariate Gaussian distribution so that the transfor- mation approximately preserves similarity in the original space (Johnson & Lindenstrauss, 1984). Afterwards, each element of (xx) is binned in soft manner by locating its neighboring edges as indices and computing the distances between all edges as values. We denote the binning operation as b : Rd + Rd where D > d. Compared to naive binning which produces binary one-hot vec- tors and loses precision, ours has greater discriminative power by generating multi-hot real-valued representations.
We provide the following example to illustrate. Assume we are binning d (xx) = 1.7 into a 1-d grid with edges [0, 2,2,3) 3). The naive binning would generate vector [0, 1, 0] to indicate the value falls into the central bin. In comparison, the soft binning first locates the neighbors and 2 as indices, and computes distances to the neighbors as values, thus forming a vector [0, 0.7, 0.3, 0). Our soft binning can discriminate two inputs falling inside the same bin, such as 01 (xx) = 1.7 and a (x,) = 1.2, while naive binning fails to do so.
Altogether, our feature encoder a bo a resembles Locality Sensitive Hashing (LSH) with random projections (Charikar, 2002) but is more expressive thanks soft binning. We term our feature representation as Locality sensitive sparse encoding, hence Losse We depict the overall Losse process in Figure
4

---

each Published as a conference paper at ICLR 2024
We note that the raw output of b can be nulli-dimeensioal before flattening (e.g. 2- dimensional in Figure 3), and we use p to de- note its dimensionality. A p-dimensional grid has À evenly spaced bins along each axis, thus producing a highly sparse feature vector of length X per grid. Finally, we can stack grids with different random projection direc- tions get diverse features. Losse has guar- anteed sparsity as stated in Remark 3.1.
N
X
((x))
b
Remark 3.1 (Sparsity guarantee of Losse). For any p. X, E Zy0, and any input vec- tor x. Φ(x) outputs vector whose number of nonzero entries ||o(x)]lo satisfies l|o(x)llo S k22, and the proportion of nonzero entries in ((x) is at most (2)
Figure 3: Locality sensitive sparse encoding. a(.) projects input vectors into a random feature space, and b(-) softly bins oxx)) into multiple p-dimensional grids, which are flattened and stacked into a high-dimensions sparse encoding (xx)).
To give a concrete example, if we set p 3 and À = 10, then [[o(x)]lo is at most 8K while the dimension of Φ(x) can be as high as 1000s. The high-dimensional feature brings us more fitting capacity, while the bounded sparsity permits efficient model updates with constant overheads, will show in the following section.
3.3 EFFICIENT SPARSE INCREMENTAL UPDATE
no-regret world model after steps of interaction can be updated online by computing the weights We+t (Equation 3), which can be solved in closed form using the normal equation:
W A, B,, where A, = & p. and B, = Φ Yr are two memory matrices. Hence, obtain simple algorithm (Algorithm
learn the model online at every time step.
Algorithm 2 Sparse online model learning
Algorithm Online model learning
1: t == 0, Ao E RDD = 0, Bo E DxS 0
1: = 0, Ao E DxD Bu E IR DxS
while True do
2: while True do
3: t t ++1
3:
[ 1 +1
4:
8 f nonzero.index
4:
A; + A + (xx))((xx))
5:
At.ss +- A, + th (x. (x))o,(x;)
5:
B, + B. + o(x))y
6: B <- B, + (x()))
6: W A, BB
7: W.+1. f- A (Br,s A W1.)
7: end while
8: end while
Updating the memory matrices A,, B, is relatively cheap, but computing weights (line of Algo- rithm 1) involves matrix inversion, which is computationally expensive. Though W can be cursively updated using Cherman-Morrisee formula (Sherman Morrison, 1950) as done in Least- Squares TD (Bradtke & Barto, 1996), when the feature becomes very nigh-dimensional to gain expressiveness, it can become inefficient or even impractical.
Fortunately, the incremental updates ((x:))(x;) T and o(x:)yT are highly sparse by Losse, so that can efficiently update small subset of weights. As shown Algorithm 2, first record the indices where the softly binned random features fall into (line 4), which are used to locate those "activated" entries the current step and update only small block of the memory matrices (line 5-6). Most importantly, the matrix inversion for updating the weight subset can be conducted much smaller sub-matrix efficiently while ensuring optimality (line 7). The cardinality of the index set is bounded by the number of non-zero features (k2*), giving constant compute cost regardless of the number of data samples. We next derive how our sparse weight update rule solves Equation in block-wise manner, which yields optimal fitting for all data including the current update step.
We first re-index the sparse vector such that it is the concatenation of two vectors Φ(x) (x),,, (x)). where all zero entries go to the former and the latter is densely real-valued. We discard the subscript time step for notational convenience. Intuitively, the observation
5

---

Published as a conference paper at ICLR 2024
will only update subset of weights associated with its non-zero feature entries, which we denote as W, E RRx where K' is the dimension of (x). Then W; E [ is the complement weight matrix. Correspondingly, the accumulative memories
and
T can be
decomposed as
A
Az.
B3 B
A
B =
(4)
Ass
Ass
Then the objective in Equation becomes
arg argmin min
+ W
( Y/F)
W
(5)
W W
Ags
A. A.s
W W
W; W
B
Bargmin min W
Th
-
Asz
B,
where (:-)p p denotes the Frobenius inner product. Treating W constant since is not affected by current data and optimizing W, only gives
arg min (TrTTA A W Tr(W. A. W.) + 2TT(W A. W W
(6)
2T (.)
The new objective in Equation is quadratic, so we differentiate with respect W, and set the derivative to zero to obtain the solution of the sub-matrix
W, = A (B, A Wz).
which gives us the sparse update rule in Algorithm 2.
Concluding remark. So far, we have presented an efficient algorithm that learns world models online without forgetting. In our algorithm, linear models are exploited so that we can achieve no-regret online learning with Follow-hhe-Leader To enlarge the model capacity, we devise high-dimensional nonlinear random feature encoding, Losse, turning linear models into universal approximators (Huang et al., 2006). Moreover, the sparsity guarantee of Losse permits efficient updates of the shallow but wide model. We name method Losse-FTL. We note that feature sparsity has been utilized to mitigate catastrophic forgetting for decades (McCloskey Cohen, 1989; French, 1991; Liu et al., 2019; Lan & Mahmood, 2023). Our work differs from them in that not only does sparsity help reduce feature interference, but the incremental closed-form solution also guarantees the optimal fitting of all observed data, thus eliminating forgetting. More discussion on feature sparsity can be found in Appendix A. Another related line of research to ours is analytic class-incremental learning (Zhuang al., 2022; 2023; 2024), which we discuss in Appendix B due to space constraints. The implementation details of Losse-FTL can be found in Appendix E.
EMPIRICAL RESULTS ON SUPERVISED LEARNING
In this section, we will showcase empirical results to validate the expressiveness of Losse (Sec- tion 4.1) and the online learning capability of Losse-FTL. (Section 4.2), both in supervised learning settings.
4.1 COMPARING FEATURE REPRESENTATIONS
We compare our locality sensitive sparse encoding with other feature encoding techniques, including Random Fourier Features (Rahimi & Recht, 2007), Random ReLU Features (Sun al., 2018), and Random Tile Coding, which combines random projection and Tile Coding (Sutton & Barto, 2018). We note that all these methods including ours, random projections are used to construct high- dimensional features to acquire more representation power. However, they have different properties with regards to the sparsity and feature values, as summarized in Table Losse designed to be sparse and real-valued. Sparse encoding allows for more efficient incremental update (Algorithm 2), while distance- based real values provide better generalization. Although ReLU also produces sparse features, it does not guarantee the sparsity level, which is dependent the sign of its outputs.
We test different encoding methods on an image denoising task, where the inputs are MNIST (Deng, 2012) images with Gaussian noise, and the outputs clean images. The flattened noisy images

---

Published as a conference paper at ICLR 2024
first encoded by different methods produce feature vectors, on which a linear layer is applied to predict the clean images of the same size. We use mini-batch stochastic gradient descent to optimize the weights of the linear layer until convergence. To keep the online update efficiency similar, we keep the same number of non-zero entries for all encoding methods. The mean squared emror the test set for different patch sizes reported in Table 1. We observe that Losse achieves the lowest error across all patch sizes, justifying its representational strength over other non-linear feature encoders. When compared with the NN baseline, which is strong function approximator but not an efficient online learnrr, linear models with Losse work better when the patch size is 36 or smaller, suggesting its sufficient capacity for problems with a moderate number of dimensions, such as locomotion or robotics (Todorov et al., 2012; Zhu et al., 2020). See Appendix C.I for more experimental details.
0.15
Mi-bartch
Sparse Real-value
9
16
25
36
49
Mi-corine
NN
0.36
0.35
0.34
0.34
0.33
Squared
0.10
Loss--TL
Fourier ReLU
0.32
0.39
1,00
1.40
0.05
/
0.35
0.34
0.37
0.39
0.43
Tile Code
0.36
0.39
0.61
0.73
Losse
0.28
0.29
0.31
0.34
0.40
0.0
0.2
0.0
Table 1: Properties of different encoding methods and Figure 4: Mean squared errors on the stream their mean squared errors on the image demoising task on learning task of different correlation levels. Solid different patch sizes. All numbers are scaled by 10
lines and shaded areas correspond to the means and stand errors of 30 runs.
4.2 ONLINE LEARNING WITH COVARIATE SHIFT
Next, we consider a supervised stream learning setting, where we can precisely control the level of covariate shift in the training data, and test the online learning capability of Losse-FTI. against the neural network counterpart. Similar to Pan et al. (2021), we create a synthetic data stream with observations sampled from a non-stationary input distribution. Specifically, the synthetic data is generated according to a Piecewise Random Walk. At each time step, the observation is sampled from a Gaussian X, (S,,,, 5²). where /3 is fixed and S, drifts every steps. The drifting follows Gaussian first order auto-regressive random walk S = (1 e).S, + Zr, AL mod = 0, where e E (0, 1] and Z₂ N(00 o²) with fixed a. For X, = Ft, the target is defined as yt = sin (2mx²).
As proven in Pan et al. (2021), X, can share the same equilibrium distribution but possess different levels of temporal correlation if B, e, and a are properly chosen. It turns out the three scalars can be uniquely determined by single parameter d E 1), which we call correlation level. When = 0, the generated data recovers i.i.d. property. Given the data stream { (t),)))))n we fit model online and measure the mean squared error the end of learning a holdout test set, which contains independent samples across all S,. As a comparison, we use neural networks to fit online well in batch and plot the errors under different correlation levels in Figure 4. The results show that our method outperforms neural networks in two aspects. First, when the data i..dd (d = 0), Losse-FTL achieves much lower error than NN-online. This indicates neural networks with gradient descent have low sample efficiency, even on stationary data. Training NNs using batch samples alleviates the issue and reaches similar performance to ours. Second, our model consistently attains very low error across all correlation levels, showing its capacity of non- forgetting online learning. In contrast, neural networks learned with both batch and online updates incur increasingly high errors when d is large, indicating catastrophic forgetting in the presence of data nonstationarity. More experimental details can be found in Appendix C.2.
5 EMPIRICAL RESULTS ON REINFORCEMENT LEARNING
In this section, we will demonstrate that world models built with Losse-FTL can be accurately learned online, outperforming several NN-based world model baselines and improving the data efficiency of RI. agents. We first introduce the settings and baselines in Section 5.1, and then present our results in Section 5.2.
7

---

Published as a conference paper at ICLR 2024
Griswarid
Mauntain Car
Aurobot
0.8
0.8
Normafoed edeedt returm
0.8
0.6
Modai-free
00.
Oniine
0.4
0.4
Corrent
Full-rephy
10:2
0.0
10.0
0.0
0
10k
20k
30k
40k
50k
10k
20k
30k
40k
50k
0
10k
20k
30k
40k
50k
Inverted Pendulums
Reacher
Hopper
1.0
Nomaltee wileode stors
0.8
0.8
0.6
0.5
0.4
0.2
0.2
0.0
0.0
0
5k
10k
15k
25k
50k
75k
100k
0
25k
50k
75k
100k
Figure Learning curves showing normalized episode return. We compare our method with five baselines on (top) discrete control and (bottom) continuous control benchmarks. Solid curves depict the means of multiple runs with different random seeds, while shaded areas represent standard errors.
5.1 SETTINGS AND BASELINES
We employ the Dyna MBRL architecture and restrict our study to the model learning part, keep- ing the base agent's value (and policy) learning untouched. We compare Losse-FTL world models learned online with their NN-based counterparts, using the same model-free base agent. We intro- duce the baselines below. The algorithm and more experimental details can be found in Appendix D.
Model-free. The model-free base agent used in Dyna. We use two strong off-policy agents: DQN (Mnih et al., 2015) for discrete action space and SAC (Haarnoja et al., 2018) for continuous action space. All other settings below are MBRL coupled with either DQN SAC.
Full-replay. We use neural networks to build the world model and train the NNs over all collected data till convergence. Before the training loop starts, the data split into train and holdout with a ratio of 4:1. The loss on the holdout set is measured after each epoch. We stop training if the loss does improve over consecutive epochs. The NNs trained continuously without resetting the weights (Nagabandi et al., 2020). This is the closest setting to FTL for NNs, but with growing computational cost as the data accumulates.
Online. We train NN-based world models in a stream learning fashion, where only mini-batch of 256 transitions is kept for training. However, we still update the model for 250 gradient descent steps on each mini-batch, unlike our method which only needs a single pass of the interaction trajectory.
SI. This similar to Online, but the NN is trained with Synaptic Intelligence (SI) (Zenke et al., 2017), regularization-bant based continual learning (CI.) method to overcome catastrophic forgetting. However, there is no explicit task boundary defined in the RL process, thus we adapt SI to treat each training sample
task.
Coreset. This refers to rehearsal rehearsal-based CL methods (Lopez- Paz Ranzato, 2017; Chaudhry et al., 2019) that keep important subset of experiences for replay. There are different ways to decide whether to keep or remove a data point. We use Reservoir sampling (Vitter, 1985), which uniformly samples subset of items from stream of unknown length. Similar to Full-replay, data split and early stopping strategies are adopted for each training epoch.
5.2 RESULTS
The learning curves are presented in Figure 5, where we compare Losse-FTL with aforementioned baselines 6 environments. We first revisit the discrete control Gridworld environment, which is used for illustration Figure 1. This environment is a variant of the one introduced by Peng & Williams (1993) test Dyna-style planning. The agent is required navigate from the starting lo-
8

---

Published as a conference paper at ICLR 2024
cation to the goal position without being blocked by the barrier. A small offset is added on each step to make the state space continuous. In such navigation tasks, intuitively. the world model needs to learn from extremely nonstationary data. When the agent evolves from randomly behaving nearly optimal, its state visitation changes dramatically, making all NN-based world models completely fail to learn due catastrophic forgetting, unless experiences are fully replayed (see the first plot in Figure 5). The poorly learned models are unable to provide useful planning, leading to almost zero performance. In contrast, world model based Losse-FTL can be learned accurately and improves the data efficiency over Model-free.
In the other two general discrete control tasks, Mountain Car and Acrobot, the data nonstationarity still exists but may be less severe than that in navigation tasks. The results show that Losse-FT FTL con- sistently outperforms all baselines and brings sample efficiency improvement.
For continuous control tasks from Gym Mujoco (Todorov et al., 2012), we can observe that Online NN fails to capture the dynamics for all environments, resulting in even worse performance than the Model-free baseline. While SI and Coreset can help reduce forgetting some extent, their perfor- mances vary across tasks and are inferior to Full-replay in all cases. Losse-FTL, however, achieves on-par even better performance than Full-replay, which is a strong baseline that continuously trains a deep neural network on all collected data until convergence to pursue FTL. These positive results verify the strength of Losse-FTL for building an online world model.
Model update efficiency. Among contin- ual learning methods with neural networks, rehearsal- ones are closer to the FTL objec- tive and usually demonstrate strong performance (Chaudhry et al., 2019; Boschini et al., 2022). It seems that Coreset can offer good balance be- tween computation cost and non-forgetting per- formance with a properly set buffer size. We hence ablate Coreset with different replay sizes and compare their performance with Losse-FTL. Figure 6 shows both the sample efficiency and wall-clock efficiency on the Gridworld environment. Coreset with a small replay size leads to poor asymptotic performance while being relatively more efficient. Increasing the replay size indeed contributes to better results at the cost more computation. In comparison, Losse-FTL updates the model purely online without any replay buffer and achieves the best performance with the best wall-clock efficiency.
500
010
210
500
510)
Figure 6: (Left) Sample and (Right) wall-clock ef- ficiency comparison between world models learned with Losse-FTL and Coreset of different replay sizes.
LIMITATION AND FUTURE WORK
Currently, we have not validated our method on problems with very high- sigh-dimensional state space such as Humanoid, or tasks with image observations. They are expected to challenging for linear models to directly work on. Extending Losse-FTL. to handle such large-scale problems, perhaps using pre-trained models obtain compact sparse encoding, a potential avenue for future work.
7 CONCLUSION
In this work, we investigate the problem of online world model learning using Dyna. We first demonstrate that NN-based world models suffer from catastrophic forgetting when used for MBRL because the data collected in the RL process is innately nonstationary. As a result, re-training over all previous data is adopted by common MBRL methods, which is of low efficiency. Through the lens of online learning, we uncover the implicit connection between NN re-training and FTL,, which allows us to formulate online model learning with linear regressors to achieve incremental update. We devise locality sensitive sparse encoding (Losse), a high-dimensional non-linear random feature generator that can be coupled with linear layer to produce models with greater capacities. We provide sparsity guarantee for Losse, controlled by two parameters, and X, which can be adjusted develop an efficient sparse update algorithm. We test our method in both supervised learning and MBRL settings, with the presence of data nonstationarity. The positive results verify that Losse-FTL capable of learning accurate world models online, which enhances both the sample and computation efficiency of reinforcement learning.

---

Published as a conference paper at ICLR 2024
REFERENCES
Matteo Boschini, Lorenzo Bonicelli, Pietro Buzzega, Angelo Porrello, and Simone Calderara. Class- incremental continual learning into the extended der-verse. IEEE Transactions on Pattern Analy- sis and Machine Intelligence, 2022.
Steven J Bradtke and Andrew G Barto. Linear least-squares algorithms for temporal difference learning. Machine Learning, 1996.
Moses Charikar. Similarity estimation techniques from rounding algorithms. Proceedings of the Thiry-fourth Anmual ACM Symposium Theory of Computing. 2002.
Arslan Chaudhry, Marcus Rohrbach, Mohamed Elhoseiny, Thalaiyasingam Ajanthan, P Dokania, P TOrr, and M Ranzato. Continual learning with tiny episodic memories. In Workshop on Multi- Task and Lifelong Reinforcement Learning at ICMI,, 2019.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in a handful of trials using probabilistic dynamies models. Advances in Neural Information Processing Systems, 2018.
Li Deng. The mnist database of bandwritten digit images for machine learning research. IEEE Signal Processing, 2012.
Ronen Eldan and Ohad Shamir. The power of depth for feedforward neural networks. In Conference on Learning Theory, 2016.
Robert M French. Using semi-distributed representations to overcome catastrophic forgetting in connectionist networks. In Proceedings of the 13th Annual Cognitive Science Society Conference, 1991.
Robert M French. Catastrophic forgetting in connectionist networks. Trends in Cognitive Sciences, 1999.
Pierre Gaillard, Sébastien Gerchinovitz, Malo Huard, and Gilles Stoltz. Uniform regret bounds over Rdd for the sequential linear regression problem with the square loss. In Algorithmic Learning Theory, 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning. 2018
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint (xxii:2301.04100 2023.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In ICML, 2022.
J Fernando Hernandez-Gardi and Richard S Sutton. Learning sparse representations incrementally in deep reinforcement learning. In Workshop on Continual Learning at NewrIPS, 2019.
Steven CH Hoi, Doyen Sahoo, Jing Lu, and Peilin Zhao. Online learning: A comprehensive survey. Neurocomputing, 2021.
Guang-Bin Huang, Lei Chen, Chee Kheong Siew, et al. Universal approximation using incremental constructive feedforward networks with random hidden nodes. IEEE Trans. Neural Networks, 2006.
Taher Jafferjee, Ehsan Imani, Erin Talvitie, Martha White, and Micheal Bowling. Hallucinating value: A pitfall of dyna-style planning with imperfect environment models. arXiv preprint arXiv:2006.00 04363, 2020.
Michael Janner. Justin Fu. Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in Neural Information Processing Systems, 2019.
10

---

Published as a conference paper at ICLR 2024
William B Johnson and Joram Lindenstrauss. Extensions of lipschitz mappings into hilbert space. Contemporary Mathematics, 1984.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. In International Conference on Learning Representations, 2020.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Byron Knoll and Nando de Freitas. A machine learning perspective on predictive coding with paq8. In Data Compression Conference, 2012.
Qingfeng Lan and A Rupam Mahmood. Elephant neural networks: Born to be continual learner. In Workshop on High-dimensional Learning Dynamics at ICMI., 2023.
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 2022.
Vincent Liu, Raksha Kumaraswamy, Lei Lc, and Martha White. The utility of sparse representa- tions for control in reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2019.
Zichen Liu, Siyi Li, Wee Sun Lee, Shuicheng Yan, and Zhongwen Xu. Efficient offline policy optimization with learned model. In International Conference Learning Representations, 2023.
David Lopez-Paz and Mare' Aurelio Ranzato. Gradient episodic memory for continual learning. Advances in neural information processing systems, 2017.
Michael Lutter, Leonard Hasenelever, Arunkumar Byravan, Gabriel Dulac-Arolld, Piotr Trochim, Nicolas Heess, Josh Merel, and Yuval Tassa. Learning dynamics models for model predictive agents. arXiv preprint arXiv:2109. /4311, 2021.
Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of Learning and Motivation. 1989.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforeement learning. Nature, 2015.
Andrew W Moore and Christopher G Atkeson. Prioritized sweeping: Reinforcement learning with less data and less time. Machine Learning, 1993.
Anusha Nagabandi, Kurt Konolige, Sergey Levine, and Vikash Kumar. Deep dynamics models for learning dexterous manipulation. In Conference on Robot Learning, 2020.
Francesco Orabona. A modern introduction to online learning. arXiv preprint arXiv:1912.13211 2019.
Yangchen Pan, Hengshuai Yao, Amir-man mmss-ud Farahmand, and Martha White. Hill climbing on value estimates for search-control in dyna. 2019.
Yangchen Pan, Kirby Banman, and Martha White. Fuzzy tiling activations: A simple approach to learning sparse representations online. In International Conference on Learning Representations, 2021.
Jing Peng and Ronald J Williams. Efficient learning and planning within the dyna framework. Adaptive Behavior, 1993.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in Neural Information Processing Systems, 2007.
11

---

Published as a conference paper at ICLR 2024
Aravind Rajeswaran, Kendall Lowrey, Emanuel V Todorov, and Sham M Kakade. Towards general- ization and simplicity in continuous control. Advances in Neural Information Processing Systems, 2017.
Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 2020.
Shai Shalev-Shwartz et al. Online learning and online convex optimization. Foundations and Trends in Machine Learning, 2012.
Jack Sherman and Winifred Morrison. Adjustment of an inverse matrix corresponding to a change in one element of given matrix. The Annals of Mathematical Statistics, 1950.
Yitong Sun, Anna Gilbert, and Ambuj Tewari. On the approximation properties of random relu features. arXiv preprint arXiv:1810.0453 2018.
Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximat- ing dynamic programming. In Machine Learning Proceedings. 1990.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 1991.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. 2015 8.
Richard S Sutton, Michael Bowling, and Patrick M Pilarski. The alberta plan for ai research. arXiv preprint arXiv:2208. 1tt::22208.117333 17173, 2022.
Erik Talvitie. Self-correcting models for model-based reinforcement learning. Proceedings of the AAAI Conference Artificial Intelligence, 2017.
Matus Telgarsky. Benefits of depth in neural networks. In Conference on Learning Theory, 2016.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In IEEE International Conference on Intelligent Robots and Systems, 2012.
Jeffrey S Vitter. Random sampling with a reservoir. ACM Transactions on Mathematical Software, 1985.
Jianhao Wang. Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-hased imagination. Advances in neural information processing systems, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 2020.
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In International Conference on Machine Learning, 2017.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martin, Abhishek Joshi, Soroush Nasiri- any, and Yifeng Zhu. robosuite: A modular simulation framework and benchmark for robot learning. In arXiv preprint arXiv:2009. /2293, 2020.
Huiping Zhuang, Zhenyu Weng, Hongxin Wei, Renchunzi Xie, Kar-Ann Toh, and Zhiping Lin. ACIL: Analytic class-ineremental learning with absolute memorization and privacy protection. In Advances in Neural Information Processing Systems, 2022.
Huiping Zhuang, Zhenyu Weng, Run He, Zhiping Lin, and Ziqian Zeng. GKEAL: Gaussian kernel embedded analytic learning for few-shot class incremental task. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, 2023.
Huiping Zhuang, Run He, Kai Tong, Ziqian Zeng, Cen Chen, and Zhiping Lin. DS-AL: A dual- stream analytic learning for exemplar-fe free class inss-inccrmeentl learning. Proceedings of the AAAI Conference Artificial Intelligence, 2024.
12

---

Published as a conference paper at ICLR 2024
A ON FEATURE SPARSITY
Feature sparsity has been pursued reduce forgetting dating back McCloskey Cohen (1989); French (1991). The core idea is that catastrophic forgetting occurs due to the overlap of distributed representations in NNs, hence it can be reduced by reducing the overlap. Sparse features are lo- cal features with less overlap, so they should alleviate the forgetting issue. Recent works have demonstrated positive results by enforcing sparse hidden features NNs either by regularization (Liu et al., 2019; Hernandez-Garcia & Sutton, 2019) or by deterministic activation functions (Pan et al., 1021; Lan & Mahmood, 2023). However, we emphasize that, in our work, norgeetgg is ensured by FTL and the high- dimeensiiiii sparse encoding is mainly for representational capacity and computational efficiency, though it also reduces feature interference as a byproduct.
We first conduct theoretical analysis similar to Lan & Mahmood (2023) to understand how sparsity contributes to non-forget,ing, then present empirical results to show sparsity alone is not sufficient to guarantee forgettgt......
Assume we are regressing a scalar target given an input vector using squared loss: { x. y) = (x)(x) y)², where f is our predictor with parameters w. When new data sample (xe: yt
arrives, we want to adjust the parameters w to minimize E(J, Xt, ye) with gradient descent:
w'
Aw -a Xx, Ht)
(8)
The change of the parameters will lead to the change of the predictor's 's output, and by the first-order Taylor expansion:
Sw [w(()) R V ( (w' - -w)
a w V
-av (e Ve). Ww)w V Swffw(e (x:)))
* and analyze (ww)w(x), V Ww (xx)). For linear models as in our case, (x) = w 6(x), thus ww V (w]w((x)) (o(x), o(x))). Then we could investigate the following question: after adjust the parameters (Equation 8) based on the new sample, what's the output difference (Equation 9) at any x for x that is similar or dissimilar to x,?
We assume non-zon zero loss so V
1. For similar to Xg, expect (o(x), ((x,)) * 0 to yield improvement and generaliza- tion. This will hold since a proper feature encoder Φ(-) should approximately preserve the similarity in the original space.
2. For dissimilar to X,, we hope (d(x), (xx))) as 0 to reduce interference (therefore less forgetting). The claim is that if Φ(-) produces sparse features, it is highly likely that the inner product is approximately zero, hence meeting our expectation.
However, argue that feature sparsity is not sufficient to mitigate forgetting, because the proba- bility that there is no single overlap between representations low even if the feature sparse. In Figure 7, show the mean squared errors for the piecewise random walk stream learning task intro- duced in Section 4.2. We use our Losse representation with different sparsity levels as features and compare sparse FTL update with gradient descent update. We observe that the gradient de- scent update relying on feature sparsity still suffers from catastrophic forgetting until X = 30, which corresponds
l000-dimensional feature with 40 nonzero entries an extremely high sparsity level. Nevertheless, increasing sparsity indeed helps sparse features to reduce interference when us- ing gradient descent (comparing the absolute values of MSES across different A). In contrast, Losse- FTL ensures non-forgetting incrementally by finding the optimal firting for all experienced samples, even with the presence of feature overlap. Notably, when Losse gets more sparse, we can observe the enror for Losse-FTL decreases, verifying our design that the high dimensional sparse feature brings more fitting capacity while maintaining the same computational cost with our sparse update.
B ANALYTIC CONTINUAL LEARNING
Our work is closely related to a paradigm called analytic continual learning, where the closed form solution, if obtainable, is employed in training continual learner. A recent line of research explores
13

---

Published as a conference paper at ICLR 2024
1m10
1-25
0.03
Mee Squaed Eroe
0.02
0.01
3-25
X-30
Loste GD
LTs
Moi Squardt Enor
0.00115
0.0010
0.0005
0.00
0.2
0.4
0.6
0.8
0.0
0.6
0.8
Figure 7: Mean squared errors on piecewise random walk stream learning. The results compare the FTL update to the gradient descent update using Losse features with different sparsity levels. Note that gradient descent (GD) update
mini-batch of 50 samples for each update, resulting in lower emoos than FTL's when the feature dimension gets higher.
this topic focusing on class-incremental learning (Zhuang et al., 2022; 2023; 2024). In particular, Zhuang et al. (2022) employ a pre-trained neural network to extract feature vectors x after process called feature expansion, and then learn the weights of fully connected layer WFON on the frozen features for iaasemnncremetll classification. The weights at round can solved analytically in recursive manner utilizing the Woodbury matrix identity:
(k) FCN
w FON"
RR X (
X
(fc) W (k-1) FCN
R& (fo) Ytrain
(10)
where Y!rain is the label for class k, and Rg
EE/ X (fo)T
x (fe)
+ +nI) is the regularized
feature auto-correlation matrix (similar to the inverse of A, $ db in our formulation), which also has a recursive form:
Ra R Rr 1 x
(feT I x (fe)
Rr X (fo)T
X (fe) R
(11)
N₂ N.
Though Zhuang et al. (2022) have achieved state-of-the-art class-incremental learning performance, there are few limitations. First, a deep network must be pre-trained and frozen to serve as an ex- pressive feature encoder, which is simply infeasible in online settings. Second, Equation 11 involves inverting matrix whose size grows with the number samples, limiting its scalability. Third, learning system requires the knowledge of task boundaries for matrix decomposition, which may not be accessible in many roaldworld scenarios. In contrast, our work develops an expressive fea- ture encoder that does not rely pre-training, proposes constant-time update rule exploiting the feature sparsity, and applies to reinforcement learning where the nonstationarity appears without explicit task boundaries.
C SUPERVISED LEARNING DETAILS
In this section, we provide more details regarding the supervised learning experiments on image denoising (Appendix C.I) and piecewise random walk (Appendix C.2).
C.I IMAGE DENOISING
We add pixel-wise Gaussian noise with = 0, o = 0.3 to normalized images from the MNIST
dataset (Deng, 201 2), and create train and test splits with ratio 9 1. To test the feature representation
14

---

Published as a conference paper at ICLR 2024
different difficulties, crop patches from the center with different sizes, ranging from 2 to 7 x 7. Figure shows some sample pairs. For different feature encoding methods compared
3
Figure 8: Sample pairs for the image denoising task.
in Section 4.1, limit the number of non-zero feature entries up 80 for fair comparison. This means 2-d binning with 20 grids for Losse-FTL. The number of bins affects the performance slightly, due different granularity for generalization. Hence we sweep over (5, 8, We observe that the scale of the standard deviation of the Gaussian random projection matrix affects the final performance significantly for random ReLU and random Frouier features, so we select the best one from a sweep over 10.1, 0.3, 0.5, 1, 5. 10). The linear layer is optimized using Adam (Kingma & Ba, 2015) with a learning rate 0.0001. We repeated experiments for settings for times and report the mean scores.
C.2 PIECEWISE RANDOM WALK
Recall that observations in the process are sampled from Gaussian X, (S,, (2) with {Se JeN being a Gaussian first order ruto-regressive random walk S:41 (1 c)S, , + Ze. where e (0 and Z4 ~ N 22). The equilibrium distribution (Xe)ee is also a Gaussian with E[X,] =
and variance 5² 3 + zz- We can use a single parameter [0, 1), which we call correlation level, control the variance:
c
VI
B 2
σ² d
B
8 = (1 d)
2
where B gives high probability bound. See Pan et al. (2021) for detailed derivation.
We visualize sampling trajectories of X, for = 1 E [0, 1) in Figure 9.
For neural networks trained in Section 4.2, we use 2-layer MLP with 50 hidden units. For NN-Batch, we take 50 samples from each X₁ and use mini-batch gradient descent to update the weights. Adam optimizer is used with the best learning rate swept from (5x 10-6 1 10-5 5x ,.... .,. 1 x 1x10 2)-
For Losse-FTL, we use 2-d binning with 10 bins for each grid and use 10 grids to construct the final feature. We all experiments for 50 independent seeds and report the means and standard errors.
D REINFORCEMENT LEARNING DETAILS
Learning algorithm. We present the detailed Dyna-style algorithm based on Losse-FTL in Algo- rithm 3, where the world models are learned online. For NN-based settings, line 4 of Algorithm replaced by training neural network (optionally with replay buffer or other continual learning techniques).
Settings for discrete control tasks. We use K = p0 2, and À 10 for Losse-FTL, and 3-layer MLPs with 32 hidden units for neural networks. For Coreset we maintain buffer with size 100, and for SI we use the best & and e from a grid search. The DQN parameters are updated using real data with interval of interactions. In MBRL, 16 planning steps conducted after each real data update. Both real data updates and planning updates use mini batch of 32. The learning rate of DON is swept from (1 10-2 10-3 10-3 ,..., 10-5) } for all configurations, All world
15

---

end for Published as a conference paper at ICLR 2024
0121
U.4
0-
200
400
0
200
400
200
400
d 0.62
0.83
1.38
0
200
100
0
200
100
200
400
Figure Visualization of piecewise random walk sampling trajectories with different correlation levels d.
Algorithm Dyna MBRL with Losse-FTL
Require: environment M, agent π, world model f. state search control c, model unroll length k
model experiences Dm
1: for epoch {1,..., F) do
2: for interaction E (1...., ......,K) K] do
3: s,a,s', rollout (M, x)
4: Update / with (s, a, s',r) (Algorithm 2)
D incremental model updates
5: end for
6: for planning E [L....,NN) do 5 8: ä. ModelUnroll(f.s.m
7:
D uniform state sampling chort-horizen on-policy unroll
9:
Dm < ( end for
10:
11:
for learning E {l,.... G} do
Update DQN or SAC batch of model data Dm
13: end for
models are updated every 25 environment steps to accelerate experiments. We run all experiments for 30 independent seeds and report the means and standard errors.
Settings for continuous control tasks. We use and 4-layer MLPs with 400 hidden units for neural networks following Janner et al. (2019). Our implementation follows closely the official codebase of MBPO (Janner al., 2019), and also adopts the same hypae-pasameters for both MBPO and SAC. The replay size of Coreset 5000, and the hyper-parameters for SI are chosen from grid search. All world models are updated every 250 environment steps to accelerate experiments. We run all experiments for 5 independent seeds and report the means and standard errors.
300, = 2, and X = 10 for Losse-FTL.
E IMPLEMENTATION DETAILS
We provide more implementation details of Losse-FTL in this section. Referring to Figure 3, the input is first pre-processed to be bounded between [-3,3], 3), and then projected with a matrix
httpp:://gt http://tttto..o.....oooo.aao jttt:////////nne///////a
16

---

Published as a conference paper at ICLR 2024
containing values sampled from an isotropic Gaussian with fl 0 and Σ I, where e is the fan- in dimension, i.c., 5 + A. This ensures the projected values (xx)) are bounded with high probability and facilitates more even bin utilization. In computing the sparse update at line of Algorithm 2, we use (Ass + +eI) for some small > codes in Listing 1.
0 to ensure its invertibility. We provide Jax-based Python
import math
from typing import Nomediuple, Tuple
import chex
import jax
import jax. numpy as Jnp
import numpy as np
9
class ToaseFarama (NamedTuple):
11
cound: jax Array
projection: Aexay
xtx: Aaxay
sty: Asray
jax.Array
15
18
def ste__d_a indddddidexxiidiceee. offsets, n_fess, bin_dim, bibis):: "*"Compute the fisttened __grids_pe:_lsh ln_bins 1) ea bin_dim
weight matrix.
indices jup. reshape (indices, (-1, bin_dim, r_seat))
offsets jnp. reshape (offsets, (-1, bin_dim, r_featt)
indices inp. atack ([ (["rdi indicca + 1), axis--11 8 (-1, bin_die, n_ reat,, 2] values jap.stack multiplier prpp-power
offsete, offsetsl, axis-
3 [21, din r_feat, 2)
25
+,, jup.arange (bin_dim multiplier(t, Nona, None]
-1, -1)
# shape a (-t, f_feat, (2, bin_die
shape_suffix, [tuple(<p) for p in sp.split inp.eys (bin_dim, dtypermp. int32) bin_dimi]
29
insices sum sum.jnp.iin. [anijjjp.reiile. reshape (indices(:, 1], (+l, n_test, *sufrix)) for i, suffix An enumerate (shape_sut suffix!)
30
values math, prod (jap. reshape (values(:, ib, (-1, feata, vauffix)) for i, suffix in enumerate (shape_ suffixi)
# both indices and vaives has the shape (-1, feaet, *8, =8fn_dmm_ now.
indices +s exppexppdd_die
n_grids_per
arppaeange in_ feath, axis-tuple trange (-bin_dim, 11)
# expand 1 dim in the front and bin_dim in the back. rephreshape (indices, (-1, a_feat 2xm)inddi)) values = reshape (values, (-1, feat 2esbin_dim)) return indices, values
class
41
"**Linesr regressor with 10cality Sensitive Sparse Encoding (Losse).
Ne update the linear weights online sparsoly following Algorithm. 2 in the paper, i.o., compoting the incremental closed-form solation based on nexly og oata pointe.
def
init_
47
self, inout dimut_dims: Suplejint, int], mum nun_features: int, nun_bins: int,
bin_dim: int,
52
eps: float,
self. festtursss fum_ueesures self.mun_bing num_bins sel sell.bin_dim bin_dim
self.inout_dims inout_dims self.eps eps
n_edges - nun_bins +
n_grids_pes__shh (a_edges
self,d n_gridb_per_lsh num features
def init (ae-t, input_dia _im inolf.inout_dii output_dim self atd 1 jap. (inputt___dd projection std www.addm.ttnattteeeooo..... rng,
jax. randon, Paannon
1osaeParams:
17

---

Published as a conference paper at ICLR 2024
(input_dim,, salf-num_Footurse self.bin_sim),
return LosseParems
count prant-jap-oo jap.erray
projeation-prajctee xtx=jnp, seros veros[(selfe)
76
self.d,), dtype- intrrojjeettnnttt....... .dtype)
jny zeros (self.d, ostpus_ prope-projcction. dttpee........e............ w-jnp. zeros (er.s ffsel outpot_dim), prope-projection on.dtyp))
78
def u
self,
83
parama: LosseFarame,
jex.Array,
jax.Array,
LozseParams:
(x, yh, (i,)) assert non-batched
indices, volues salf ind.ff_idicess_rrti sar__andvvlluee (paroms.project projection, x)
parama sela._ (parama, indi ces, values, parama self self._update__ (e__uddaee_(pazaant) indices)
90
params
def predict (self, pazams: LossePerame, x: :ex.y: indices, volues output paroms. (idiices] valuea[.... Nona]
ind___idddee_d_d_e (parems. carrams.rroecttinn
95
return sumput.sum(]
97
def andiices_end_vevvee
self, projectien: jax.Arrey, jex.Array,
100
101
102
3:
Ja9. matmul projection) h jax.an. an. sigmoid(hh jap.pilppi... 0, self.num_bin indices floor .astypeljppppppp... Lnt int offsets indices, values _to_td_index(
104
106
107
indices
109
indices,
offsets,
111
self het_inum__aaurree
self.bin_dine niinsssm.bins,
113
114
'ndices, values
116
def
_update_memory
118
self,
params: TasseFarans,
indices: jax.Array,
121
values: Aex.Array,
jax.Array,
123
LosseParams:
chex..eseer__ecpaaaaaaeee shtppp____pppppeiiiiiii (indices, values, y), prefix_lea-]]
125
stx_indices (indices salf,d) [t, None] + indicos None, :]
xtx_ind'cen
xtx xtx_ www.iniiee....atteee
xty_indices indices. flatten!)
128
xt* xtx_updates values 5, Nona] values(:, None;
xtx_updates stx_updates.
130
sty_updates values :, None] y[:, None, :]
131
sty_updates sey_:pdaaes..xehapee (-1, y.shape [-1]]
132
return
(
133
xtx-parame.xte attccamm.xxx.taaaee (xtxxin [xtx_indicea] .adddaaddddd (xtx_upatee)) xty-perane [xty_indiees] add add(xty_upaaaees)) count-params, count. + +... ahahaee [8],
134
135
137
def _update_ki
139
sel
140
parama:
LosseParane,
141
indices: jax.arrsy,
142
>> LosseParams:
indices. (1attten))))))
144
sub_indices indices lindices self,d) [t, None] + indizes (None, :] sub
sub_xtx xtx jap.reshapp (parama. xtx [sub_inddoos; [indi ces.cee.....e.. ahspee(0]) aub_xty parama xty xty [indicea]
146
147
sub_xtx
148
sub jnp.matmul(
149
reshape (pacems. xtx, (selC,d, self.d) [indices],
pazams.N,
18

---

Published as a conference paper at ICLR 2024
151
152
b sub_xty xtaw + Jnp. poppmmmmmll (sub_xxxx, porans sesaassw((sa ((iddices)
a_norm parana count self.eps " Jnp.eye (ae)
154
b_norm b count
135
solution jnp. winalg.solveieee.........oo.e solve nelnomm, b_norm)
return repllee pr-params (indices] set (solution)
157
losse Tosse 159
inout_dima- (1, 11r
161
Fum_
162
num sum_bing=13,
163
bin_dime2,
164
epa-le-5,
166
167
lcsse. init
jax. (losse.inig) init)
168
lesse. lesse.update jax. (losse. update, denate_argnums- (0,)) donate to avoid copy
169
lcsse.predict jax.
Listing Jax-based Python codes for learning an online linear regressor with Losse-FTL.
19


Reviews:
[{"review_id": "kQVoyHQ5Cg", "reviewer": "Reviewer_2Y3F", "summary": "The paper introduces a model-based reinforcement learning approach that utilizes a sparse representation-to-representation model. The use of sparse representation aims to address the challenge of catastrophic forgetting in a reinforcement learning (RL) setting, where data generation constantly shifts. The architecture employed for model-based reinforcement learning is Dyna. The proposed method involves learning nonlinear sparse features and building a model based on this sparse representation. To enhance computational efficiency, a method for updating model weights using sparse representation is presented. Empirical experiments are included to demonstrate the effectiveness of this approach.", "questions": "see above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper targets an important topic of learning a model in RL \n\n2. I do not see many related works of building sparse representation-based models.", "weaknesses": "I will list below main weaknesses for improvements, centred around the main contribution of the paper. \n\n1. 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 invlidate/validate another. As an empirical paper, I think a rigorous comparison is necessary. \n\n2. Could you clarify do you update both your model and representation every environment time step? \n\n3. 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 extensive 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. \n\n4. other issues. \n\nAlg 1. it is better to be specific, use the title Dyna architecture, rather than MBRL, as there are numerous MBRL algorithms and not everyone is as Alg 1 described. \n\nAlg 2, line 4 & 5: shouldn’t it be outer product? Please specify the dimension of the matrices capital Phi and letter phi. This is nontrivial as it affects the understanding of the algorithm. \n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TzMwvS942k", "reviewer": "Reviewer_LJy1", "summary": "The paper presents a method for the online learning of a world model for model-based reinforcement learning (MBRL). To obtain efficient updates to the world model, the world model is expressed as a linear combination of a set of spare features. This efficiency allows online learning at a constant computational cost.", "questions": "In figure (1): d(s_(t+1), f(st, at)) and \\delta were not defined in the caption or anywhere obvious in the main text.\n\nIn eq (3), does || . ||^2_F denote the Frobenius matrix norm? It is only stated so after eq (5). It helps to have the notation introduced earlier here. Especially since “F” is the dimension of the feature space", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The work is well motivated in the introduction and a sufficient and clear background is provided for non-expert readers in the preliminaries section. The algorithms, definitions, etc. are mathematically rigorously presented. \n\nA comprehensive set of experiments has been conducted demonstrating the efficacy of Losse-FTL", "weaknesses": "Significant discussion around catastrophic forgetting was mentioned in the introduction but little discussion is presented in the main text and left in the appendix.\n\n“Example 3.1” could be a regular paragraph. Formatting this as an Example does not improve readability and is, in fact, the only Example in the entire paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "76XC5M0kbk", "reviewer": "Reviewer_CrA3", "summary": "This paper presents a world model for model-based reinforcement learning (RL) which can be learned online and does not\nrequire full retraining on all previous data.\nThe authors highlight that training world models is subject to issues arising in continual learning.\nThat is, each sequential data set can be interpreted as coming from a new task and thus the world model needs to be retrained\nafter each agent-environment interaction.\nThis is because the data collected via agent-environment interaction is non-stationary.\nThey propose a world model based on a linear regression model which uses high-dimensional nonlinear features.\nImportantly, the linear model can be updated given new data whilst retaining good predictive performance on old data.\nThey compare their feature encoding method to other feature encoding techniques in an image denoising task on MNIST.\nThey then evaluate their method's ability to handle training data covariate shift in an artificial online learning experiment.\nFinally, they evaluate their method's ability to combat non-stationary data in model-based RL.", "questions": "- What are the limitations of the proposed method? Can it handle image observations? Can it scale to difficult environments like Humanoid?\n- Why haven't you compared to any other world model algorithms? E.g. Dreamer, TD-MPC.\n- How was the NN full-replay experiment implemented? 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "This paper addresses an important problem in model-based RL, which is likely a problem that must be solved for developing lifelong agents.\nThe method for updating the linear model online is simple but appears to be effective.\nThis is also the first time I have seen a non-trainable encoder used for world models.\nIt is very common to see world models with NN encoders and NN transition models operating in the encoder's latent space.\nTypically the dynamic model operates on a latent state which is lower dimensional than the high-dimensional observations.\nPerhaps I am not aware of the relevant literature, but this seems like an interesting and original idea.", "weaknesses": "This paper has two main weaknesses.\nFirstly, there is no comparison to state-of-the-art MBRL strategies that use world models, e.g. Dreamer/TD-MPC.\nAs such, there is no experiment highlighting the main issue that the paper is trying to address:\nthat NN-based world models suffer from catastrophic interference due to non-stationary data.\nSecond, all of the RL experiments are in simple RL environments.\nFrom the current results, it is impossible to know how practically useful this world model is.\nThere is no discussion about its limitations nor is there a comparison to other model-based RL algorithms that use world models.\nSure the proposed method works on some simple RL environments but can it scale to difficult environments like humanoid and can it handle image-based observations?\nIt is OK if the method cannot do this but it should be addressed in the text.\nMoreover, it should be made clear what benefit it does have over other world model methods (like Dreamer).\nFor example, I'd like to see a state-of-the-art world model method (like Dreamer) performing poorly/failing because it cannot handle non-stationary data.\n\nI also have questions regarding the training of the NNs in the experiments.\nDid the full-replay experiment involve resetting the neural network's weights? If so, what initialization was used?\nWhen was the NN training stopped? Was the data split into train/validation sets and used to stop training when the\nvalidation loss stopped improving?\nThe paper needs more details to explain exactly how this was implemented.\nIn my experience, these steps are important to ensure the NNs don't overfit on early data sets.\n\nI am also unsure why the full-replay strategy (which is model-based), does not appear to have better sample\nefficiency than the model-free experiment. Am I missing something here?\nPerhaps this is an interesting point for discussion. Do the high-dimensional features sacrifice sample efficiency\nin favour of formulating a linear model which can handle the non-stationary data?\nI'm not sure if this is correct.\nMy main point here is that the paper has not answered all of my questions about the method.\n\nThe experiments tell the first part of a nice story.\nTable 1 compares to other encoding methods and Fig. 3 clearly shows how the method handles covariate shift better than NNs.\nFig. 4 also acts as a nice ablation for comparing the method to other CL strategies within the same set-up.\nHowever, the experiments section lacks a comparison to other MBRL strategies which use world models.\nIn particular, there is no experiment highlighting the main issue the paper is trying to address.\nThat is, there is no MBRL experiment failing due to the non-stationarity of the training data.\n\nMinor comments and corrections:\n- The abstract is very long. I would recommend shortening it.\n- Sections shouldn't lead straight into subsections (Section 4/4.1, 5/5.1, B/B.1, C/C.1). There should be text explaining what the reader can expect to read in the section.\n- In the first paragraph of Section 2.2, the reward function is defined as $R(\\mathbf{s}, \\mathbf{a}, \\mathbf{s}')$ but then in the optimal policy equation you use $R(\\mathbf{s}, \\mathbf{a})$.\n- Fourth line of Section 2.2 the initial state distribution is $\\rho$ but earlier it is $\\rho_0$.\n- Third paragraph of Section 2.2 - \"We firstly formulate\" should be \"We first formulate\".\n- Section 2.1 - \"When the input is a convex set $\\mathcal{S}$, the prediction a vector $\\mathbf{w}_{t} \\in \\mathcal{S}$\". This sentence doesn't read properly.\n- $\\rho$ is used to denote the initial state distribution and to denote the dimension of the grids in Section 3.2.\n- What is the value of $\\delta$ in Fig. 1?\n- What are $\\pi_0$, $\\pi_t$ and $\\pi_{t'}$?\n- The first sentence of the abstract says model-based RL has better sample efficiency. Better than what? It's model-free counterparts?\n- It is unusual to end the paper with a section titled \"Summary\". I recommend changing this to \"Conclusion\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)