Replay Buffer with Local Forgetting for Adaptive Deep Model-Based Reinforcement Learning
Authors: Ali Rahimi-Kalahroudi
Janarthanan Rajendran
Ida Momennejad
Harm van Seijen
Sarath Chandar
Venue: Submitted to ICLR 2023 (2023)

Abstract: One of the key behavioral characteristics used in neuroscience to determine whether the subject of study---be it a rodent or a human---exhibits model-based learning is effective adaptation to local changes in the environment. In reinforcement learning, however, recent work has shown that modern deep model-based reinforcement-learning (MBRL) methods adapt poorly to such changes. An explanation for this mismatch is that MBRL methods are typically designed with sample-efficiency on a single task in mind and the requirements for effective adaptation are substantially higher, both in terms of the learned world model and the planning routine. One particularly challenging requirement is that the learned world model has to be sufficiently accurate throughout relevant parts of the state-space. This is challenging for deep-learning-based world models due to catastrophic forgetting. And while a replay buffer can mitigate the effects of catastrophic forgetting, the traditional first-in-first-out replay buffer precludes effective adaptation due to maintaining stale data. In this work, we show that a conceptually simple variation of this traditional replay buffer is able to overcome this limitation. By removing only samples from the buffer from the local neighbourhood of the newly observed samples, deep world models can be built that maintain their accuracy across the state-space, while also being able to effectively adapt to changes in the reward function. We demonstrate this by applying our replay-buffer variation to the classical Dyna method, as well as to recent methods such as PlaNet and  DreamerV2, showing for the first time that deep model-based methods are able to achieve effective adaptation.

Under review as a conference paper at ICLR 2023
REPLAY BUFFER WITH LOCAL FORGETTING FOR ADAPTIVE DEEP MODEL-BASED REINTORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
One of the key behavioral characteristics used in neuroscience to determine whether the subject of study-be a rodent or human-cxhibits model-base based learning is effective adaptation to local changes in the environment. In reinforce- ment learning, however, recent work has shown that modern deep model-based reinforeementt-eamiing (MBRL) methods adapt poorly to such changes. An ex- planation for this mismatch that MBRL methods are typically designed with sample-efficiency on a single task in mind and the requirements for effective adap- tation are substantially higher, both in terms of the learned world model and the planning routine. One particularly challenging requirement is that the learned world model has to be sufficiently accurate throughout relevant parts of the state- space. This is challenging for deep-learning-basee trophic forgetting. And while replay buffer can mitigate the effects of catas- trophic forgetting, the traditional first-in-frsttouut replay buffer precludes effective adaptation due to maintaining stale data. In this work, we show that a conceptually simple variation of this traditional replay buffer is able to overcome this limitation. By removing only samples from the buffer from the local neighbourhood of the newly observed samples, deep world models can be built that maintain their accu- racy across the state-space, while also being able to effectively adapt to changes in the reward function. We demonstrate this by applying our replay-buffe varia- tion to a deep version of the classical Dyna method, such as PlaNet and DreamerV2, demonstrating that deep model- based methods can adapt effectively well to local changes in the environment.
world models due to catas-
well to recent methods
INTRODUCTION
Recent work has shown that modern deep MBRL methods adapt poorly to local changes in the environment (Van Seijen et al. 2020 Wan et al.) 2022), despite this being a key characteristic of model-based learning in humans and animals (Daw et al. 2011]. The analysis by Wan et al. (2022) revealed that there are broadly two causes for this lack of adaptivity: an insufficient world model or insufficient planning. And the former one is an especially challenging one to overcome when deep- learning-based world models are considered. The core of this challenge lies in the fact that adaptivity requires world model that is accurate across the relevant state-space, as small change in reward or transition function can change the trajectory of the optimal policy entirely. By contrast, to achieve high single-task sample-effit sapplee-fffiieney---er common metric used in MBRL research it is sufficient that the world model is accurate along the current behavior policy.
For deep world models, accuracy across the state-space hard to achieve and maintain, even with sufficient exploration. The reason is that collected samples are strongly correlated, and, at the fi- nal stages of learning, mostly come from states along the trajectory of the optimal policy. Due to catastrophic forgetting, the quality of the predictions further away from this trajectory quickly de- grades. A common strategy to counter this is to use replay buffer. By randomly sampling from large replay buffer and using these samples to update the world model. the effects of catastrophic forgetting are greatly reduced. However, the traditional disadvantage that it hinders effective adaptation, as out-of-date samples interfere with the new data.
infiitt-ouuuu (FIFO) replay buffer has the

---

Under review as a conference paper at ICLR 2023
To address the challenge of catastrophic forgetting, while also avoiding interference from out-of-date samples, we propose variation of the traditional FIFO replay buffer. Instead of removing the oldest sample from the replay buffer once the buffer is full, the oldest sample in the local neighbourhood of the new sample is removed. This conceptually simple idea naturally leads to a replay buffer approximately spread out equally across the space space, while local changes accounted for quickly. Consequently, updating the deep world model with samples drawn randomly from this replay buffer results a world model that is approximately accurate across the state-space at each moment in time. We call this replay buffer variation LoFo (Local Forgetting) replay buffer. One practical challenge to our proposed variation is that a locality-function needs be learned that determines whether or not sample from the replay buffers falls within the local neighborhood of a newly observed sample. We train this locality function using contrastive learning (Hadsell et al.) 2006 Dosovitskiy et al. 2014 Wu et al. 2018) during the initial stages of learning, after which it is fixed and used as basis for the LoFo replay buffer.
whose samples
We demonstrate the effectiveness of the LoFo replay buffer by combining it with deep version of the classical Dyna method and measuring its adaptivity on a variation of the MountainCar task well mini-grid domain, using the same Local Change Adaptation (LoCA) setup as used by Wan et al. (2022). We then test the limits of our approach by applying the same idea to both PlaNet (Hafner et al. 2019b) and DreamerV2 (Hafner al. 2020). which world models based on recurrent networks and are intended for continuous-action domains. Experiments with these modified methods on variations of the MuJoCo Reacher domain demonstrate that LoFo replay buffer can substantially improve adaptivity of more advanced deep MBRL methods as well.
2 BACKGROUND: LOCAL CHANGE ADAPTATION (LoCA) SETUP
Inspired by work from neuroscience on detecting model-basee behavior in humans and animals (Daw) ct al. 201T), Van Seijen et al. (2020) proposed the Local Change Adaptation (LoCA) regret as
Training
T2
Evaluation
perimental setup and metric to evaluate model based behavior of RL methods. LoCA regret measures how quickly a method adapts its policy after observ- ing a local environment change and tries to estimate how close a method is to idcal model-based behav- ior. Wan et al. (0022) improved the LoCA setup by making it simpler, less sensitive to hyperparameters and easily applicable to stochastic environments. evaluates if an RL method can reach close to op- timal performance after observing local environ- ment change and makes binary classification of methods that effectively adapt to local changes and those that cannot. In this work use this improved LoCA setup (Figure T for all our evaluations. It sists of two main components, the task configuration and the experiment configuration.
Phase
(Task A)
Phase
(Task B)
Zone Boundary Initial State Distribution
State (T1, T2)
Figure I: LoCA Setup: During training, in Phase 1, the agent is randomly initialized throughout the state-space and has to solve Task A which has higher reward in T1 com- pared to T2. In Phase 2, the agent is initial- ized within the one-way T1-zone and the ward function changes to Task B with TI's reward changing and having lower reward than T2, hence changing the optimal policy. During evaluation the agent is initialized ran- domly throughout the entire state-space in both the phases.
The task configuration considers an environment with two tasks, namely A and B, which differ only in their reward functions. method's adaptivity determined by analyzing how effectively it can adapt from task to task A learning environment suited for the LoCA setup is made up of two terminal states, namely TI and T2. The reward function for both tasks is always zero except when transitioning to ei- ther of the terminal states. For task A, the agent receives a reward of 4 upon transitioning to TI and 2 upon transitioning to T2. For task B, however, transition TI yields reward of 1 (instead of as in task A), and transitioning to T2 results the same reward as in task A, which is 2. In addition, both of these tasks share the same discount factor 0 < 7 1. Finally, the transition dynamics for
2

---

Under review as a conference paper at ICLR 2023
local area around T1 is such that it is impossible for agent to move out, once entered. This is called the Tl-zone.
The change between the two tasks is local and is only in their reward functions. Specifically, tasks A and B share the same environment dynamics, except that a transition to TI results in different rewards between the two tasks. In cognitive neuroscience this is also known as reward transfer, or retrospective reward revaluation (Momennejad et al. 2017). While the difference in tasks A and B only local, their optimal policies are completely different for most states, i.e., the optimal policy for task A tries to reach TI, while for task B, it tries to reach T2 except when the agent is within the T1-zone, where the optimal policy tries to reach TI.
The experiment configuration consists of two different training phases. These two phases differ in the task and initial state distribution. Throughout Phase I, the environment dynamics are set that of task A's, and the initial state is drawn uniformly random from the entire state-space. When the first phase finishes, the task changes to task B, and Phase begins. The initial state for the second phase is drawn uniformly random from only states within the Tl-zone; hence, the agent's presence is limited to a local area around TI. During each phase, the agent is periodically evaluated for some number of episodes, and the initial states for each evaluation episode is drawn uniformly random over the entire state-space. The average overall return is reported for each evaluation and compared to the average return of the corresponding optimal policy in that phase. Lastly, the agent does not receive any signal when the task changes.
For method to reach optimal performance in Phase 2, it has to adapt the new reward function corresponding to task B, i.e., it has to change its policy from pointing towards T1 to pointing towards T2 for most of the state-space. The method has to perform this adaptation the reward change while only observing states within the local region around T1 where the reward change has occurred. Given sufficient amount of time train in Phase followed by Phase 2. method is classified as adaptive if it is able to achieve close to optimal performance in both Phase I and Phase 2. If method achieves near-optimal performance Phase but not Phase 2, it is classified as a non- adaptive method. And lastly, no assessment is made when a method fails to reach close to optimal performance in Phase
3 CHALLENGES IN ADAPTIVE DEEP MBRL
Using the LoCA setup. Wan et al. 2022 observed that current popular deep MBRL methods, such as PlaNet (Hafner et al. 201 9b) and D2eamerV2 (Hafner et al., 2019a 2020) fail to adapt to local environment changes. Their analysis revealed that the key reason for their failure to adapt is their inability to build and maintain correct world models when environment changes occur. A small local change in reward or transition function can change the optimal policy entirely. For an MBRL method to adapt to such environment changes and update its policy, it needs to maintain a world model that is sufficiently accurate throughout the relevant parts of the state-space. It is not sufficient to have a world model that is accurate only along the current behavior/optimal policy, as is the case for current methods such PlaNet and Dreamer, developed for single non-changing task settings.
Current deep MBRI. methods that use deep neural networks as function approximators, such as PlaNet and DreamerV2 store their recent experience a large fists-n-first-ouu (FIFO) replay buffer. In a FIFO replay buffer, when the buffer gets full, as new transition samples are added to the buffer, the oldest samples are removed. Transitions are then sampled from the replay buffer to train the world model, which is then used for planning and updating the agent's policy. Instead of training the model directly from the stream of data that the agent is observing, having a replay buffer and sampling transitions from it to train the deep neural network based model, helps break the strong correlation between the stream of data that the agent is observing. This correlation hurtful for the i..d., assumptions made by the stochastic pradient-descent based optimizers typically used to update the parameters of the neural networks. Replaying past experience also helps mitigate forgetting predictions about states that the agent does not visit frequently.
When a change occurs in the environment, using a large replay buffer results in interference of the old oft-of- data with the new data, resulting in an incorrect world model. For example, in the LoCA setup, when the agent enters Phase 2 with task B, there is a local change in the reward function, specifically the reward corresponding T1 changes. The agent which is restricted to the
3

---

Under review as a conference paper at ICLR 2023
TI-zone in Phase 2 is able to observe this local change. However, the replay buffer from which transitions are sampled to update the model still has lots of transitions from Phase with the stale incorrect reward corresponding to TI from task A. This interference from the old data results in an incorrect world model, which affects planning and thereby renders the method not adaptive to the local change in Phase 2.
In a FIFO replay buffer, the current samples in the replay buffer are determined by the current behavior policy. Therefore, as time progresses, the old data will all be removed from the buffer. This will happen faster if we make the replay buffer smaller. While it might avoid the interference from old stale data after some time, would result in catastrophic forgetting of model predictions corresponding to states not recently visited. For example, in the LoCA setup, when the replay buffer small, after entering Phase 2, the stale transitions around from phase will be removed from the replay buffer over time, stopping the interference problem. This, however, also removes other still relevant and correct transitions from other parts of the state-space, which are now replaced by transitions from just within the Tl-zone. A model trained with such replay buffer forgets the learned model for states outside the TI-zone, which again affects the planning and renders the method not adaptive to the changes in Phase 2.
Therefore, the core challenge in building adaptive deep MBRL methods is to address the eeerfereence-orgettinn dilemma and maintain world models that are sufficiently accurate throughout the relevant state-space when tasks change, irrespective of the current behavior policy.
4 LOCAL FORGETTING (LOFO) REPLAY BUFFER FOR ADAPTIVE DEEP MBRL
We propose LoFo (Local Forgetting) replay buffer, conceptually simple variation of the traditional FIFO replay buffer that is able to address the core challenges in building an adaptive deep MBRL method. Instead of removing the oldest samples from the replay buffer once it is full as done in the traditional FIFO replay buffer, in a LoFo replay buffer, the oldest samples in only the local neighbourhood of the new samples are removed.
When a change occurs in the environment and the agent observes that change, adding the new samples to the replay buffer and removing the oldest samples from only the local neighbourhood of the new data instead of the oldest samples from the entire replay buffer facilitates removal of the potentially incorrect and stale data sooner. This helps mitigate the interference of the old out-of-date samples with the new samples when tasks change.
Since only the samples in a local neighbourhood the new samples are removed, old samples in other parts of the state-space including those that have not been visited recently remain in the replay buffer. This leads to a replay buffer whose samples are approximately spread out equally throughout the entire relevant state-space, irrespective of the current behavior policy. Updating the world model with samples drawn from this LoFo replay buffer results in a world model that is approximately accurate throughout the relevant state-space avoiding the problem of catastrophic forgetting of predictions related to states not recently visited.
LoFo replay buffer can therefore address the interference-forgetingg dilemma and result in learning a world model that is sufficiently accurate throughout the relevant state-space even when there are environment changes. The current deep MBRL methods can then utilise the correct world models to update and adapt their policies, resulting in an adaptive deep MBRL method.
5 LoFo REPLAY BUFFER WITH CONTRASTIVE STATE LOCALITY
An instantiation of LoFo replay buffer requires detinition of a local neighbourhood to ob- served new sample and a way to determine which samples in the replay buffer are within that local neighbourhood. While there could be several ways to do this, in this work, learn state lo- cality function using a contrastive learning technique which is then used to both define the local neighbourhood and also identify samples within that local neighbourhood.
Using contrastive learning (Hadsell et al. 2006 Dosovitskiy et al., 2014 Wu et al. 2018 we learn neural embedding representation of states such that states that temporally closer, i.e., reachable
4

---

Under review as a conference paper at ICLR 2023
with fewer actions also closer in their neural embedding representation. Specifically, we learn an embedding function v = [s),, where fe is a deep neural network that maps state s to an embedding vector v. This embedding function induces a distance metrie in the state-space, such as (8; 83) (*+) (=) for states & and 8j. We train the embedding function such that states that temporally closer, have a smaller distance between the embeddings learnt and states that are temporally farther, have a relatively larger distance between them in the embedding space.
Let s' represent state that is one action away from state s, and represent a of states that not. Let ID = {8, 8', 88 be a dataset of their collection. We train the embedding function to minimize the following loss function:
L(D)=
Σ (0,0",3)ED
[1fo(s) fo(s')1 (8- 1-8-Estai (ss)llfo(s) fo (3)))
where /3 > is hyperparameter. This loss function trains the embedding of states and s' that are temporally next to each other to be closer by minimising their distance towards zero, while pushing the embedding of states 8 that are not temporally next to 8 to be on average farther, with a cumulative squared distance value close to B, positive number.
In our experiments, we collect trajectories using a random behavior policy at the beginning of train- ing and use samples from that trajectory to form the dataset D {s, s', 5) to train our embedding function fe. The learnt embedding function is then fixed and the distance between the state embed- dings is used as a proxy for state locality. In complex environments with exploration challenges, a random behavior policy might be insufficient to cover the relevant state-space needed to learn a good state locality estimate. It is important future work to figure out good ways to learn state locality in such settings.
The samples stored in the replay buffer generally of the form (s, a, s'), representing a transition an agent experienced by taking an action a, from the state 8 and moving to a state s' and obtaining reward r. In LoFo replay buffer, upon taking an action in the environment, the generated experience sample is used to first gather all the samples in the local neighbourhood of the new sample. This is done by estimating the state locality using the distance do 8;) = (s) (s)) between state s in the new sample and states } corresponding to starting states in all of the samples in the replay buffer. The samples in the replay buffer whose starting state 's distance to do(s, 8;) < Dioent are selected. Droeat is scalar hyperparameter that determines the size of the local neighbourhood around any given state.
If the number of samples within the local neighbourhood is equal to or above threshold Niocat then the oldest sample in that local neighbourhood is removed. Nooca is a positive integer hyperparameter that determines the maximum number of samples that are stored in the replay buffer within any local neighbourhood of a given sample. The new sample is now added to the LoFo replay buffer.
6 ADAPTIVE DEEP DYNA-Q
Wan et (2022) successfully built an adaptive linear Dyna-Q MBRL method. However, they showed that hrep-learning-baee version of the same approach failed to achieve adaptivity. They further identified inferior learned models resulting from catastrophic forgetting or interference from stale data as the key reason for the failure. In this section we show that replacing the traditional FIFO replay buffer used in Wan et al. (022) with the LoFo replay buffer makes deep Dyna-Q adaptive.
We evaluate deep Dyna-Q with LoFo replay buffer on the LoCA setup of two domains. First, the LoCA setup of the MountainCar domain (MountainCarl.oCa
used in Wan et al. (2022). Second,
the LoCA setup of a variant of the simple Mini-grid domain (M)))G)idLoCA), Our empirical sults show that deep Dyna-Q with a LoFo replay buffer is successfully able to adapt to the local environment changes in the LoCA setup in both the domains.
MountainCarLoCA setup (Figure 24 consists of variation of the classical MountainCar domain, with
cart having move up a hill. There are two terminal states, Tl at the top of the hill and T2 corresponding to the cart being at the bottom of the hill with velocity close to zero. Figure 3a shows the learning curves of different versions of deep Dyna-Q on MountainCarl.oCA
5

---

Under review as a conference paper at ICLR 2023
(Velocity
T2 (Volxity 0)
(a) MountainCar
(b) MiniGrid
(c) Reacher
(d) RandomReacher
Figure 2: Illustration of environments corresponding to the different domains used in our experi- ments. The dashed white lines (not visible to the agent) show the boundary of T1 Tl-zone.
Phase
Phose 2
Phose 1.
Phuse 2
Dacoutted Retura
Discounted Retarn
FIPO (sizee Se6)
FIFO (mm=-166)
1.5
FIPO (stz--25ed)
Lefb
FIPO (stz)=366)
optimal
1.0
FTFO SeirenSed
LoFa optimal
0.5-
7IFO ( Well
Ha sheil
1.5
3.0
4.5
0.3
0.6
0.9 Total Steps
1.2
15
Total Steps
x10*
x10
(a) MountainCarI.oCA
(b) MiniGridLoCA
Figure 3: Plots showing the learning curves of deep Dyna-Q with LoFo replay buffer (adaptive) and FIFO replay buffers of different sizes (not adaptive) on (a) MountainCarL.oC. and (b) Mini- GridLoCA. Each learning curve is average discounted return over ten runs, and the shaded area represents the standard error. The maximum possible return in each Phase is represented by solid black line. (a) Note that the Handcrafted-LoFF refers variant of the LoFo replay buffer that, instead of learning the state locality function, uses handcrafted locality function (See Appendix A.4).
The best replay buffer size for the baseline deep Dyna-Q Q with FIFO replay buffer is 4.5e6, the one that stores all the samples seen far. We observe that while all the methods reach close optimal performance in Phase 1, only the method with LoFo replay buffer is able to adapt to reward change that happens in Phase and reach close to optimal performance in Phase 2, making an adaptive MBRL method as per the LoCA evaluation.
The two-dimensional state-space of the MountainCar domain allows visualization of properties as- sociated with states throughout the state-space a 2D plot. Figure 4 shows the 2D histogram of the states across the state-space whose transition samples are stored in the LoFo replay buffer and the traditional FIFO replay buffer at the end of Phase 1 and Phase 2. We observe that in FIFO replay buffer, at the end of Phase 2, almost all of the samples in the buffer are just from a small region in the state-space, that corresponds the T1-zone. Samples of states from other parts of the state-space that were present in Phase 1 have mostly been removed from the buffer. This leads to catastrophic forgetting of model prediction for states in parts of the state-space outside the TI-zone. In LoFo replay buffer however, even at the end of Phase 2, samples are maintained from across the entire relevant state-space, enabling maintaining accurate world model for the entire relevant state-space. Also, LoFo replay buffer stores only about 3e4 samples in total at the end of Phase 2, compared to 4.5e6 samples stored the FIFO replay buffer.
The MiniGridL setup has an image based high- high-dimensional input. The environment (Figure 2bp an 8 x RGB grid world with two green colored terminal states, T1 at the top left corner with 2 x 2 Tl-zone and T2 at the bottom-right corner. The agent is red triangle which can choose to go straight, turn left or turn right. We use a deep convolution neural network to encode the image observations.
Figure 3 shows the learning curves of different versions of deep Dyna-Q Q on MiniGridLoCA. We compare deep Dyna-Q using LoFo replay buffer with the variants that FIFO replay buffers with different buffer sizes, ranging from buffers that store all the samples seen (size = 1.8e6) to the buffer that stores only many samples as LoFo replay buffer uses (size 2.56e4). We observe that only

---

Under review as a conference paper at ICLR 2023
LoFo- Phase 1
LoFo Phase 2
FIFO Phase
FIFO Phase 2
Teels
veloodsy
veluett
Indodity
Hest
w+0
Geedt
I--00
Swett
100
pucitoo
postiim
position
posciion
Figure 4: Histogram of the states across the state-space of MountainCarLoCA environment whose transition samples are stored in the LoFo replay buffer and the traditional FIFO replay buffer (size = =)e6) at the end of Phase and Phase 2. While LoFo replay buffer maintains samples from states throughout the relevant state-space in both Phase I and Phase 2, FIFO replay buffer removes almost all samples from states other than the TI-zone by the end of Phase 2.
deep Dyna-Q Q with LoFo replay buffer is able to adapt to changes in Phase 2 and converge to close to optimal performance. The method achieves adaptivity while storing only around 2.56e4 samples in the replay buffer, two orders of magnitude less samples compared to the best performing FIFO replay buffer based method which is of size 1.8e6. Note that among the methods that FIFO replay buffers, two methods with largest buffer sizes, one with a buffer of size 9c5 and the other that stores all the samples with a size of 1.8e6 are initially able to reach close to optimal performance in Phase 2, but their performance degrades over time afterwards. This is because, the total number of distinct states in the MiniGridLoCA environment are less (256 to be precise) and therefore samples from all over the state-space from Phase 1 stay longer in a large replay buffer. There is a sweet spot when there are still samples from throughout the sate- from Phase 1, while the proportion of samples from Phase 2 that capture the reward change in the Ti-zone are much higher than the stale rewards of TI from Phase avoiding any serious interference. This can lead to a performance that is temporarily close to optimal performance but quickly degrades as samples from Phase I are removed in the FIFO buffer and catastrophic forgetting happens. Additional details on all the deep Dyna-Q experiments are provided in Appendix A and
ADAPTIVE PLANET AND DREAMERV2
To see if more complex methods can be made adaptive as well by using LoFo replay buffer, we applied it to the deep MBRL methods PlaNet (Hafner et al. 2019b) and DreamerV2 (Hafner al.) 2020) and evaluated their performance. PlaNet and Dreamer V2 use recurrent model for reward and transition predictions which require sample-sequences to make updates to rather than individual samples. Therefore, we first need to extend the concept of the LOFo replay buffer to sample- sequences. On a high level, what we aim to achieve is that updates made to the world model come from states spread out more-or-less. equally across the state-space. In the case of sample-sequencee sequences, we try to approximate this by ensuring that the start-state of each sequence is drawn approximately random across the state-space.
Our approach to achieve good coverage of sequence start-states is to use two separate buffers: a state-buffer which is curated similarly to the LoFo replay buffer, and a trajectory-buffer storing observed trajectories. For now, assume that the trajectory-buffer stores all observed trajectories. maximum sample-size
In the Appendix D we show how the trajectory-buffer can be bounded of B - N, where B is the size of the state-buffer and N is the size of a sample-sequence used for updates. When a new sample (st, Tt, $(+1) is observed, it appended to the trajectory-buffer while state is added to the state-buffer. Crucially, if the state-buffer contains more states from the local neighborhood of state 8, than some threshold amount, the oldest state from this neighborhood is removed from the state-buffer. Each state in the state-buffer points to copy of itself in the trajectory-bufie When a state 8; is removed from the state-buffer, the corresponding reward r; in the trajectory buffer is replaced by None, indicating that these rewards should no longer be used for training reward- preadd-peddittoos. The world model is updated by first sampling a random state su from the btafer and then drawing sample sequence of size N from the trajectory-bufee starting with state s,. Whenever the reward is N None, the loss of the corresponding reward prediction is set to O.
7

---

Under review as a conference paper at ICLR 2023
We evaluate modified versions of PlaNet and DreamerV2 the LoCA setup applied to two domains. First, a variation on the Reacher domain, called ReacherLoCA, introduced by Wan et al. (2022). And second, more complex extension of ReacherLoCA which refer to the Random- izedReacherl.co LoCA. For both domains, we observe that loFo replay buffer results in adapting to the local reward change in T1 and learning a sufficiently accurate reward model throughout the relevant state-space in Phase 2, reaching close to optimal performance.
The ReacherLoCA setup consists of a variation of the Reacher domain Tassa et al. (2018). It is a continuous-action domain with 64* x 64 RGB images as observations. The environment has two targets corresponding to Ti and T2 fixed the top right and the bottom left quadrants respectively (Figure (c). The agent controls the angular velocity of two connected bars to reach target and (Figure 2dl setup is an extension of the ReacherL.oCA where the location of targets in the Reacher environ- ment can vary each episode along a circle around the center, while still being opposite to each other.
remain at the target till the episode ends (1000 time steps). The RandomizedReachert
The target T2 is colored differently from Tl for the agent to be able to differentiate between them.
Undivozatted Returs
1000
Panse
DremasV22
Phose
Utlleeconnee Rotuns
1000
Plase:
Pluse 2
Paaiet
[ Retum
$000
Droamee22
Paace
Dromae Lods
3000
Palet + Lafu
3000
Dioamee2222 a Lofe sptimal
oithrss
optimal
2000
2000
2000
1000-
1000
1000-
0.5
1.5
2.0
Total Steps
Total Stops
Total Stegs
(a) ReacherLoCA
(b) ReacherL.oCA
(c) RandomizedRRachhrL.CCA
Figure 5: Plots showing the learning curves of D2eamerV2 and PlaNet with a LoFo replay buffer (adaptive) and with FIFO replay buffers (not adaptive) on b) ReacherL and c) Randomize- dReacherlooCA. Each learning curve is an average undiscounted return over ten runs, and the shaded area represents the standard error. The maximum possible return in each phase is represented by a solid black line.
Optimall Phuse
LoFo- Phose
Larre FIFO Phese
Small PIPO Phase 2
60
0.2
Figure 6: Visualization of the estimated rewards from the DreamerV2 agent's reward model the end of Phase 2. Each point on the heatmap represents the agent's position in the Reacher environ- ment.
Figure shows the learning curves of DreamerV2 and PlaNet on ReacherLoCA, and DreamerV2 on RandomizedReachett
We observe that in both the setups all the methods reach close to optimal performance in Phase However, in Phase 2 only DreamerV2 and PlaNet with LoFo replay buffer is able to adapt to the environment change, making them adaptive deep MBRL methods. DreamerV2 and PlaNet with the traditional FIFO replay buffer fail to adapt in Phase 2.
In Figure o we visualize the reward predictions of the different DreamerV2 methods. We observe that DreamerV2 that uses the LoFo replay buffer for learning its reward model has adapted its reward predictions for target T1 (top right) correctly to around +1 in Phase 2. When we use large FIFO replay buffer, we observe that the reward for target Tl at the end of Phase 2 is overestimated to around +2.5 because of the interference of stale samples from Phase 1. On the other hand when we use a small FIFO replay buffer, DreamerV2's reward prediction at the end of Phase 2 for Tl is accurate around +1, but the model has completely forgotten the reward for target T2, and other parts of the state-space outside the TI-zone. Additional details on all the PlaNet and DreamerV2 experiments are provided in Appendix and D
8

---

Under review as a conference paper at ICLR 2023
8 RELATED WORK
Besides the relation to the work of Wan et al. (2022). discussed in Section 6 this work has con- nections with several other lines of work. The method used for learning the state locality function ( in this work is inspired by the work of Hartikainen et al. (2019). They propose a method to learn -specii-peccfiic dustance-function that a measure of the expected number of time steps to reach the goal state from a given state. And they use this to shape the reward function.
Several works have proposed modifications to the traditional FIFO replay buffer to handle various challenges. The work that is most related to ours is that of PPrruswalkam et al. (2022). They focused on the setting of self-supervised learning from a continual stream of unsupervised data. And to address catastrophic forgetting they proposed to use a replay buffer that removes the most correlated samples when the buffer is full. For the purpose of adaptation this is not a good strategy, however, as incorrect, out-of-date data can remain in the buffer for long time.
LOFO replay buffer provides a mechanism for adding and removing samples from replay buffer. This determines the distribution of samples that stored in the replay buffer. Various works in the past have developed different strategies for sampling or weighting samples in the replay buffer, in order to make the training more efficient. These include prioritizing samples that 1) have high temporal difference (TD) error (Schaul et al. 2013), 2) are closest to the current state the agent is in (Sun et al. 2020 and 3) have low TD target uncertainty (Kumar et al.) 2020 Lee et al. [021]. These methods are complementary to our work and can be implemented on a LoFO replay buffer.
9 LIMITATIONS AND FUTURE WORK
The general strategy behind the replay buffer variation we presented boils down to forgetting samples that are spatially close, but temporally far from currently observed samples. We believe this to be a good general strategy to achieve effective adaptation in a non-stationary world. However, the specific implementation of this principle will differ depending on the problem type and the non- stationarity considered. In this paper, we considered only reward non-stationarity and domains where exploration is easy. In this scenario, learning locality function during the initial learning phase and keeping it fixed thereafter is sufficient. However, when the transition dynamics is non- stationary as well, the locality function needs to be maintained and updated across time as well, as the distances between individual states can change over time.
Now that we have shown that adaptive deep MBRL is possible in principle, a logical next step for future work is scaling up these methods to larger and harder domains. Because, while DreamerV2 has shown to be able to achieve good single-task performance on such domains, it is not a given that replay buffer variation is sufficient for making it adaptive for such domains well. For example, the ReacherLoCA we considered has a fairly small decision horizon (i.e., how far an agent needs to plan ahead to construct a good policy-for ReacherLoCA it takes on average about 25 actions to reach a goal state). Achieving adaptivity for longer decision horizons puts higher demands on the planning routine as well (see Wan et al. (2022) for planning-relae related pitfalls that impede adaptivity), so it might be needed to make changes to Dreamer V2's planning routine as well.
10 CONCLUSION
In this work, we considered one of the key features of model-based behavior: the ability to adapt to local environmental changes. In previous work, it has shown that tabular and linear model-based methods are able to achieve this form of adaptivity. However, deep model-based methods struggle due to an interplay between catastrophic forgetting and interference. To address the challenges with deep model-based methods, we proposed the LoFo replay buffer, whose samples are approximately spread equally across the state-space. Furthermore, we conducted various experiments to show that utilizing the LoFo replay buffer with deep MBRL methods can make them adapt effectively to local this type
changes in the reward function. This is the first time-to the best of our knowledge- of adaptivity has been shown for deep MBRL methods. This is an important step towards more practical real-world application of RI., since the stationary assumption does not always apply there.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for gymnasium. lnttps://github.
2018.
Nathaniel D Daw, Samuel Gershman, Ben Seymour, Peter Dayan, and Raymond Dolan. Model- based influences on humans' choices and striatal prediction errors. Neuron, 69(6):104-1115,, 2011.
Alexey Dosovitskiy, Jost Tobias Springenberg. Martin Riedmiller, and Thomas Brox. Discrimi- native unsupervised feature learning with convolutional neural networks. Advances in neural information processing systems, 27, 2014.
Raia Hadsell, Sumit Chopra, and Yann LeCun. Dimensionality reduction by learning an invariant mapping. In 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recogni- tion (CVPR'O0) volume 2, pp. 1735-1742. IEEE, 2006.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2566 PMLR, 2019b.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02193 2020.
Kristian Hartikainen, Xinyang Geng, Tuomas Haarnoja, and Sergey Levine. Dynamical distance learning for semi-supervised and unsupervised skill discovery. arXiv preprint urXiv: urXiv:1907.0822 1907.08255, 2019.
Aviral Kumar, Abhishek Gupta, and Sergey Levine. Discor: Corrective feedback in reinforcement learning via distribution correction. In Advances in Neural In- formation Processing Systems, volume 33, pp. 18560-18572. Curran Associates, Inc., 2020. URL https://procedings nww.igs.............. ccpps///cccppeer2////eeee pdf
Kimin Lee, Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Sunrise:
A simple unified frame- work for ensemble learning in deep reinforcement learning. In Proceedings of the 38th Inter- national Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 6131-6141, PMLR, 18-24 Jul 2021. URL intpps://proceedii press/vi33 /9/ ee http://eelgg.ht... html
mlr.
I Momennejad, Russek, H Cheong, Botvinick, N Daw, and Gershman. The successor representation in human reinforcement learning. Nature Human Behaviour, 199):680~ 692, September 2017.
Senthil Purushwalkam, Pedro Morgado, and Abhinav Gupta. The challenges of continuous self- supervised learning, arXiv preprint arXiv:2203.12710 2022.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay, arXiv preprint arXiv:1511.0599 2015.
Peiquan Sun, Wengang Zhou, and Houqiang Li. Attentive experience replay. Proceedings of the AAAI Conference on Artificial Intelligence, 34004)55900-5977, Apr. 2020. doi: 10. 1609/aaai. v34:044..34444 60449, URL 7049
www.aaai..rriiineaeeeeeeeeeoeeee Caapp/AAA//artiin
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li. Diego de Las Casas, David Bud- den, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq. et al. Deepmind control suite. arXiv preprint arXiv: 1801 000690 2018.
10

---

Under review as a conference paper at ICLR 2023
Hann Van Seijen, Hadi Nekoei, Evan Racah, and Sarath Chandar. The loca regret: A consistent metric to evaluate model-based behavior in reinforcement learning. In Ad- vances in Neural Information Processing Systems, volume 33, pp. 6562-6572. Curran Asso- ciates, Inc., 2020. URL http///ttts//procedingg...ss neurip ps. cc/ppaper/202///11e. 4Bdb71587d8 45
pp
Yi Wan, Ali Kalahiod Janarthanan Rajendran, Ida Momennejad, Sarath Chandar, and Harm H Van Seijen. Towards evaluating adaptivity of model-based reinforcement learning meth- ods. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 22536-22561. PMLR, 17-23 Jul 2022. URL
https://proceednns mlr www.apees//veee 62 wwn2dd. html
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non- parametric instance discrimination. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 3733-3742, 2018.

---

Under review as a conference paper at ICLR 2023
A
ADDITIONAL DETAILS FOR THE EXPERIMENTS ON MOUNTAINCARLOCA SETUP
A.1
EXPERIMENT SETUP
Uniform distribution over the entire state- state-space Uniform distribution over
Phase training
Phase 1 evaluation
Initial distributions
a small region Uniform distribution states within TI To-e
Phase 2 training
Uniform distribution over a small region
Phase 2 evaluation
Training
Phase steps
1.5 10
steps
Phase steps
3 x 10
Maximum number of steps
500
before an episode terminates
Other
Training steps between two evaluations
10a
details
Number runs Number of evaluation episodes
10
10
Table 1: Experiment setup for testing the Deep Dyna-Q the MountainCarl.oCA domain.
The CoLAAainCarL.oCA was first introduced by Van Seijen et al. (0220, which a variant of the well-known MountainCar environment. TI is located at the top of the mountain (position > 0.5, and velocity > 0), and T2 is located at the valley ((position + 0.52) + 100 x velocity² < 0.07?). The T1-zone contains all the states within 0.4 S position I 0.5 and 0 I velocity I 0.07. The discount factor for this environment is À = 0.99. And lastly, for each evaluation, the agent is initialized roughly at the middle
and T2 (-0.2 S position S -0.1 and -0.01 S velocity S 0.01). Table I shows the experiment setup we used to evaluate the deep Dyna-Q agent's 's adaptivity under the LoCA setup.
A.2 HYPERPARAMETERS
Table shows the final values of the hyperparameters that were used in the LOFO replay buffer for generating the results presented in Figure 3a We searched over 3 E {1) Diocal E (0.01, 0.005), and Niocal {1, 5, 10).
We used the same version of the deep Dyna-Q algorithm used in Wan et al. 2022 However, instead of having separate neural networks for different actions in each part the model (dynamics, reward, and termination models). we use just one network and concatenate the action with the output of the middle layer (the one that has 63 output units) and then feed it to the next layer. Table 3 summarizes the important hyperparameters for the deep Dyna-Q method on the MountainCarLoCA. Both the deep Dyna-Q agent that used the LOFO replay buffer and the baseline agents in Figure Ba used these hyperparameters, and they only differ in the choice of the replay buffer. Wan et al. (2022) found that using the same replay buffer size for learning the model and planning worked the best. Therefore, followed the same.
A.3 ADDITIONAL RESULTS FOR THE DEEP DYNA-Q
We compare the reward models of the Dyna-Q agents that use the LoFo replay buffer and the traditional FIFO replay buffer by visualizing their prediction of rewards over all states in the Moun- tainCarl.oCC .oCA domain at the end of each phase in Figure 7 We observe that the reward model of the agent that uses the LoFo replay buffer can estimate the rewards correctly in both phases. However, this is not true for the agent that uses the traditional FIFO replay buffer since it fails to predict the correct rewards associated with T2 due to catastrophic forgetting.
12

---

Under review as a conference paper at ICLR 2023
Embedding network
MLP:[64 x 64 x 64 x 16..
architecture
Activation Function: tanh
Optimizer
Adam, learning rate: 10-4
B
10
Number of negative samples
128
Mini-batch size
32
Total number of random steps for creating dataset D Number of training epochs
100000
D Diocal Niocal
0.005
Table Hyperparameters used for the LoFo replay buffer on the MountainCarl.oCA domain.
MLP with tanh,
Dynamics model
[64 64 x 63 x 64 x 64).
MLP with tanh,
Neural networks
Reward model
664 64 x 63 x 64 x 64), MLP with tanh,
Termination model
64 x 64 x 63 x 64 x 64).
MLP with kanh,
vation-valuc estimator
64 x 64 x 64 x 64].
Adam,
Value optimizer
learning rate: x 10-6
Optimizer
Adam,
Model optimizer
learning rate: 10-5 Epsilon greedy
Exploration parameter
< 0.5
Number of random steps
50000
Other details
before training
Target network update frequency
500
Number of model learning steps
5
Number of planning steps
Mini-batch size of model learning Mini-batch size of planning
32
32
Table 3: Hyperparameters used for the deep Dyna-Q on the MountainCarl.oCA domain.
LoFo Phose
LoFo Phase 2
FIFO Phase
FIFO Phase
velkeity
valcotty
1.0
10
Figure 7: Visualization of the estimated rewards from the deep Dyna-Q agent's reward model at the end of each phase. In each heatmap, the axis represents the agent's 's position, and the axis represents its velocity.
13

---

Additionally, we tested another variant of the LOFO replay buffer that uses a handerafted state locality function instead of using the learnt contrastive state locality function. Since the Mountain- CarLoCA domain's 's state-space is 2-dimensional, where the first dimension indicates the position Under review as a conference paper at ICLR 2023
A.4 HANDCRAFTED STATE LOCALITY FUNCTION
and the second one the velocity of the car (s : (at, v)), we can scale the range of the velocities to match the scale of the positions and use the euclidean distance as the locality function:
Chanderafted (8i, 85) / (1: - xz)² + 150 x (u; v2)²
The learning curves of the deep Dyna-Q agent that uses the LoFo replay buffer with the handcrafted state locality function the MountainCarl.oCa domain is presented in Figure Ba We observe that the LoFo replay buffer with the learnt state locality function able match the performance of the buffer with the handcrafted state locality function.
B ADDITIONAL DETAILS FOR THE EXPERIMENTS ON THE MINIGRIDLOCA SETUP
B.1 EXPERIMENT SETUP
Uniform distribution over
Phase training
the entire state-space
Uniform distribution over
Phase 1 evaluation
the entire space
Initial distributions
Uniform distribution over
Phase training
states within T1-zone
(2 x 2 subgrid)
Uniform distribution over the entire state-space
Phase 2 evaluation
Training
Phase steps
x 10
steps
Phase steps
1.5 106
Maximum number of steps
100
before an episode terminates
Other details
Training steps between evaluations Number of runs Number of evaluation episodes
10
10
10
Table 4: Experiment setup for testing the Deep Dyna-Q method on the MiniGridLoCAA domain.
The implementation of the MiniGridLoCA domain is done using the MiniGrid python package (Chevalier-Boisse Boisvert al. 2018) Table shows the experiment setup we used to evaluate the deep Dyna-Q Q agent's 's adaptivity under the LoCA setup.
B.2
HYPERPARAMETERS
Table 5 shows the final hyperparameters setting used in the LoFo replay buffer for generating the results presented Figure 3b Furthermore, we only searched over Niocal (50, 75, 100, 150). While we started with the same B and Drocat as the MountainCarl.oCA we found that by decreasing the Diocet to 0.001, the learned state locality function can make a clear distinction between all possible states (256).
The algorithmic design of the deep Dyna-Q agent we used for the MiniGridLoCA domain is the same that used for the MountainCarLoCA. We only changed the neural network architecture for the model and the action-value estimator. Table summarizes the architecture of the neural networks. Note that first encode given state a low-dimensional vector for the various parts of the model (dynamics, reward, and termination model). Then, we concatenate the given action to the resulting vector and feed it to the MIP layers. This is mainly due the fact that did not want
14

---

Under review as a conference paper at ICLR 2023
CNN: (Channels:[32 x 64 x 64
Embedding network
Kernel Sizes:[8 3 x 3]
architecture
Strides: [4 x 2)),
Followed by MLP: [112 x 16),
Activation Function: relu
Optimizer
Adam, learning rate: 10
B
10
Number of negative samples Mini-batch size
128
32
Total number of random steps for creating dataset II) Number training epochs
25000
Diocal
0.001
Niocal
Table 5: Hyperparameters used for the LoFo replay buffer on the MiniGridLoCA domain.
CNN:
(Channels: (22 x 64 x 64]
Kernel Sizes: 3 x 3]
Strides: x x 2)).
Dynamics model
Followed by Transposed CNN:
(Channels: (64 x 32 x 3]
Kernel Sizes: x 6 x 5
Strides: | x 3)).
Activation Function: relu
CNN:
(Channels: (32 x 64 x 64
Kernel Sizes: 88 3 3]
Reward model
Strides: x 2 x 2),
Neural networks
Followed by MLP: (5122,
Activation Function: relu
CNN:
(Channels: x 64 x 64
Kernel Sizes: 3 x 3)
Termination model
Strides: 2 2)),
Followed by MLP:
Activation Function: relu
CNN:
(Channels:(32 x 64 x 64]
Kernel Sizes: (8 x 3 x 3]
Action-value estimator
Strides: 2),
Followed by [512[512],
Activation Function: relu
Table 6: Neural networks' architecture for the deep Dyna-Q on the MiniGridL.oCA domain.
keep separate networks for each action as it was done in Wan et al. (2022). Finally, Table 7 shows the final hyperparameters used to generate the learning curves Figure 3b when the agent used the LoFo replay buffer and the traditional FIFO replay buffer.
B.3 ADDITIONAL RESULTS FOR THE DEEP DYNA-Q
Figure shows the learning curves of our search over different values of Niocal for the deep Dyna-Q agent using the LoFo replay buffer. Since the locality function makes a clear distinction between the states of the MiniGridL CACA, using specific Nioctl means that the LoFo replay buffer stores
15

---

Under review as a conference paper at ICLR 2023
Adam, learning rate: 6.25 x 10
Value optimizer
Optimizer
Adam,
Model optimizer
learning rate: 10 10-4
Epsilon greedy 0.5
Exploration parameter
Number of random steps
2000
Other details
before training
Target network update frequency
5000
Number of model learning steps
Number of planning steps
Mini-batch size of model learning Mini-batch size of planning
128
128
Table 7: Final hyperparameters for the deep Dyna-Q on the MiniGridLoCA domain.
Phase
Phase 2
Discounted Return
Naaal 100
M
3
Niaral 150
optimal
Nura 50
2-
0
0.3
0.6
0.9
1.2
1.5
1.8
Total Steps
x106
Figure 8: Learning curves for the Dyna-Q agent that uses the LoFo replay buffer with different Noooal values (Drocat = 0.001), each averaged over 10 random seeds.
Phase
Phase 2
Discounted Return
Dusmi 0.001
D 0.05
3
Donai 0.005
Dicat 0.1
Donai 0.01
optimal
2
0.3
0.6
0.9
1.2
1.5
1.8
Total Steps
100
Figure 9: Learning curves for the Dyna-Q agent that uses the LoFo replay buffer with different Diocal values ( NNocat 100), each averaged over 10 random seeds.
exactly Niocat samples each state. Hence, an interesting observation from Figure 8 is that higher Niocal results in higher performance in Phase 2.
Figure I shows the learning curves of our search over different values of Diocal for the deep Dyna-Q agent using the LoFo replay buffer when Niocat = 100. We observe that small values Dioeal result in successful adaptation in Phase 2.
Additionally, we visualized the estimated rewards from the deep Dyna-Q Q agent's reward model when using the LoFo replay buffer and the traditional FIFO replay buffer the end of each phase in Figure 1O To create each heatmap, we first generated the estimated rewards by placing the agent in given state (a specific cell and direction) and taking the going straight action. This gave four
16

---

Under review as a conference paper at ICLR 2023
LoFo
Phase
LoFo Phase 2
FIFO Phase
FIFO
Phase 2
Figure 10: Visualization of the estimated rewards from the deep Dyna-Q agent's reward model at the end of each phase. Each point on the heatmap represents the agent's position 8 x grid of the MiniGridLoCA environment.
different 2D heatmaps, each representing a particular direction. Then, we summed the estimated rewards over the four possible directions and created the final heatmap.
From Figure 10 we observe the deep Dyna-Q agent that uses the LoFo replay buffer can predict the correct summation of the rewards during each phase. In contrast, the agent that uses the traditional FIFO replay buffer can only make accurate predictions in Phase The inability of this agent to correctly predict the summation of the rewards is due forgetting. Note that we used the agent with the traditional FIFO replay buffer of size 3e5 in Figure 10
C ADDITIONAL DETAILS FOR THE EXPERIMENTS ON THE REACHERLOCA SETUP
C.I EXPERIMENT SETUP
Uniform distribution
Phase training
the entire state-space
Uniform distribution
Phase evaluation
the entire states outside TI-zone
Initial distributions
Unitorm distribution
Phase 2 training
states within T1-zone
Uniform distribution over
Phase 2 evaluation
the entire states outside T1-zone
Training steps
Phase steps
3 x 10
Phase 2 steps
x 105
Number of steps
1000
before an episode terminates
Other details
Training steps between Number of runs
evaluations
15000
10
Number of evaluation episodes
Table 8: Experiment setup for testing the PlaNet method the ReacherL.oCA domain.
17

---

Under review as a conference paper at ICLR 2023
Uniform distribution over the entire state-space Uniform distribution over the entire states outside Tl Tl-zone Uniform distribution over states within T1-zone Uniform distribution over the entire states outside T1-zone
Phase training
Phase evaluation
Initial distributions
Phase training
Phase 2 evaluation
Training steps
Phase I steps
10
Phase 2 steps
1.5 x 106
Number of steps
1000
before episode terminates
Other
Training steps between two evaluations
10000
details
Number of runs
10
Number evaluation episodes
Table Experiment setup for testing the DreamerV2 method on the ReacherLoCA domain.
MablesSaand] show the experiment setup we used to evaluate the PlaNet and the DreamerV2 agents' adaptivity under the LoCA setup respectively. It is worth mentioning that Wan et al. (2022) stayed as close as possible to the original Reacher domain (Tassa et al. 2018) when creating the ReacherLoCA to facilitate reusing the best hyperparameters previously used much possible. A transition to the targets in the ReacherLoCA domain does not terminate the episode and the agent keeps receiving rewards until 1000 timesteps. Note that having targets instead of terminal states does affect the requirements the LoCA setup.
C.2 HYPERPARAMETERS
CNN:
(Channels:{33 (32 x x 128 256] Kernel Sizes:/4 x x 4 x 4]
Embedding network architecture
Strides: [2 x 2
x 2))
Followed by MLP::512 x 64, 32],
Activation Function: relu
Optimizer
Adam, learning rate: 10-4
B Number of negative samples Mini-bate size Total number of random steps for creating dataset D Number of training epochs
50
128
32
105
5
D Diocal Niocat
0.05
10
Table 10: Hyperparameters used for the LoFo replay buffer the ReacherL.oCA domains.
Table 10 shows the values of the hyperparameters used in the LoFo replay buffer for generating the results presented in Figures 5a and (5)) Furthermore, we searched over different values of Niocal E 52 10, 20, 40, 80}. For both PlaNet and DreamerV2 agents the best setting in experiments is Niocal = 10.
For the rest of the hyperparameters, use the hyperparameter values used by Wan al. (2022) for the DreamerV2 agent for both when it uses the LoFo replay buffer and the traditional replay buffer. Furthermore, instead of having entropy regularizer for the exploration, add noise to actions (the default value used in the original Dreamer method (Hafner et al. (019a)) since the ReacherL domain is relatively simple environment.
18

---

Under review as a conference paper at ICLR 2023
For the PlaNet agent, perform a hyperparameter search only for the learning rate (Ir E {3
10 10-4, 4, 1004,10-+1 +)) The best learning rate for the agent that uses the LoFo replay buffer is 10
Other hyperparameters are the same as the best setting in Wan et al. (2022).
C.3 ADDITIONAL RESULTS FOR PLANET
Optimal Phase
LoFo- Phase
Large FIFO
Phase 2
0.1
0.0
0.0
1L0
Figure 11: Visualization of the estimated rewards from the PlaNet agent's reward model at the end of Phase 2. Each heatmap's x and y axes represent the agent's 's position in the ReacherLoCA domain.
In Figure 11
we visualize the reward predictions of the PlaNet agents that use the LoFo replay buffer and the traditional FIFO replay buffer. We also added the true reward visualization (as opti- mal) for reference. The results are similar to those of DreamerV2's s (Figure
D ADDITIONAL DETAILS FOR THE EXPERIMENTS ON THE RANDOMIZEDRECCHERLOC.. SETUP
D.1 EXPERIMENT SETUP AND HYPERPARAMETERS
Uniform distribution over the entire state-space Uniform distribution over the entire states outside T1-zone Uniform distribution over states within T1-zone
Phase training
Phase evaluation
Initial
distributions
Phase training
Uniform distribution over the entire states outside Ti-zone
Phase 2 evaluation
Training
Phase steps
1.5 10
steps
Phase 2 steps
3.5 x 106
Number of steps
1000
before episode terminates
Other details
Training steps between two evaluations
10000
Number of runs
10
Number of evaluation episodes
8
Table 11: Experiment setup for testing the DreamerV2 method on the RandomizedRcacheLL.CCA domain.
In the RandomizeddRaccheLLLCAA. the location of the red target (TI) is randomly taken from the circle centered the center of the state-space (the dotted black circle in Figure 2d). And then, the green target (T2) is placed at the opposite end of the circle. Table TableeTihhe shows the experiment setup used to evaluate the DreamverV2 agent's adaptivity under the LoCA setup.
The best hyperparameter setting for the LoFo replay buffer is exactly what is mentioned in Table 10 except that Nineal = (we searched among Ni E {1, 2, 10)). Otherwise, we use the same hyperparameters for the DreamerV2 agent as used in the ReacherLooC. setup.
19

---

Under review as a conference paper at ICLR 2023
Phase
Phase 2
1000
Undiscounted Retum
Niseal
Muxt optinall
Mixal
3000
Nosal 2
2000
1000
0
1.5
3.0 Total Steps
4.0
5.0
x10
Figure 12: Learning curves for the Dreamer agent that uses the LoFo replay buffer with different Nocol values (Droeat = 0.05), each averaged over 10 random seeds.
Phase
Phase 2
4000
Undiscounted Retarn
Dussi 0.05
Maat 0.1
Drsst 0.075
optimal
3000
2000
1000
1.5
3.0 Total Steps
4.0
5.0
10
Figure 13:
Learning curves for the Dreamer V2 agent that uses the LoFo replay buffer with different Dioeat values (Ntocat = 2), each averaged over random seeds.
D.2 ADDITIONAL RESULTS FOR THE DREAMER
Figure 2 and 13 shows the learning curves of our search over different values of Nocat and Dioeat for the DreamerV2 agent using the LOFO replay buffer. Similar the MiniGridLoCA domain, observe that small values of Diocal result in successful adaptation in Phase 2 for the Randomize- dReacherLoCA.
20

---

Under review as a conference paper at ICLR 2023
E LOFO REPLAY BUFFER WITH RECURRENT MODELS
In this section, we show that trajectory-buffer can bounded to maximum sample-size of B x N in practice, where B is the size of the state-buffer and N is the size of sample-sequence used for updates.
As mentioned in Section 7 upon removing 8; from the state-buffer, r; is trajectory-buffe and the agent never uses sample-sequence starting from Now. if we look closely the braject only sample-sequences starting with states sy where & (i N, i] contains the (8;, 0j, 8g, Ak None, 8i+1). Conceptually, we can remove (st, Qz Fú, 8(+1) from the N, are removed from the state-buffer, because in this case there we
None in the
trajectory-buffer
remains no sample-sequence containing (s;, Ti 8;-1) for training the agent. By doing argue that no N consecutive N one rewards can found in the trajectory buffer. Because in that case, the last None reward belongs to no valid sample-sequence, and therefore, it should have been removed from the trajectory- buffer..
Given that for each state in the state-buffer, we know their corresponding reward in the trajectory- buffer is not N None, the worst-case scenario, there can be at most N None rewards after such samples in the trajectory-buffer Hence, by counting them well, each sequence in the trajectory- buffer starting from a state the state-buffer can be at most of the length N. Therefore, the total number of samples stored in the trajectory-buffer would be most B N.
21


Reviews:
[{"review_id": "dirwHJelCls", "reviewer": "Reviewer_54Vu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposed a replay buffer strategy to help model-based RL models effectively adapt to local changes in the environment. When local changes happen, the traditional first-in-first-out (FIFO) data buffer will interfere with the model training due to the out-of-data samples still being stored in the buffer. To fix this, the authors proposed to update the buffer by only replacing the data in the local neighborhood of the newly observed samples. Contrastive learning method is used to learn the representation of states and to force consecutive states to have smaller Euclidean distance, thus neighborhood can be determined by applying a threshold on the distance measure.", "summary_of_the_review": "This paper proposed a replay buffer strategy to help model-based RL models effectively adapt to local changes in the environment. However, the extreme cases designed in the experiments can not provide a general use for the proposed method.", "strength_and_weaknesses": "### Strength\nThe proposed method is easy to understand and efficient. Experiments showed expected results.\n\n### Weakness\n- Major\n    I doubt whether the Local Change Adaptation (LoCA) is a suitable testbed for the proposed purpose. In phase 2, the agent is trapped around the T1 area, thus bringing a sharp distribution shift in the collected data during this phase. Data collected during this phase have a very poor state coverage, which hurts the world model learning severely. This designed small local change in the environment results in a huge shift in data distribution, which is a significant effect for MBRL that learns a policy during imagination. A more realistic case would be the case where only the reward function changes, while the agent movement not being limited in any way. \n- Minor\n    - Could the authors explain or provide experiment results analyzing the sensitivity of the hyper-param, $D_{local}$ and $N_{local}$ used in the experiments?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "GL3Kyd2Lzp", "reviewer": "Reviewer_oofa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper studies model-based reinforcement learning in the context of changing environments. More specifically, the algorithms are evaluated in the Local Change Adaptation (LoCA) setting. Well known methods like Dreamer struggle in this setting. The paper proposes a simple method that modifies the first-in-first-out rule of the replay buffer and instead removes old data if it is very similar to a newly added state. The similarity is measured by a function learned through contrastive learning. The proposed method shows improved results over the baseline.", "summary_of_the_review": "The proposed method is interesting and the quality of the paper is very good. However, the limited evaluation makes the paper less significant. Still, it is a good first step and hence the paper is for me just over the acceptance threshold.", "strength_and_weaknesses": "The paper is very well and clearly written. The proposed idea is rather simple - which is not a negative aspect - and makes conceptually sense to solve the studied problem. It is shown that the proposed method increases performance in the environments used for evaluation. However, as also mentioned by the authors, the environments are rather simple and it is not clear if the results would transfer to more interesting environments. Currently, the similarity function is trained on random trajectories collected at the beginning of the training. This will not work for more complicated environments, where most parts of the state space have not been observed in the random trajectories. This point limits the significance of the work. Further, while domain adaptation is an important topic it is only studied in the context of one specific benchmark that accounts for scenarios of domain adaptation but not for others.\n\n\nAre the random trajectories used to train the similarity function counted as environment steps in the shown plots? \n\nMinor things:\np.5 the 5th to last row has an unnecessary space before a comma \nI find it irritating that there is textof the main manuscript between Fig.2 and Fig. 3. I suggest the Figures should be just underneath each other. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4MUdKfukeK", "reviewer": "Reviewer_G6gQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a simple method for adaptation of model-based RL agent to new environments. Starting from the replay buffer filled with transitions from the old environment, the method works by gradually adding in the new transition from the new environment while removing data that are close to the states in the new transitions. In this way, the old data near the newly explored states in the new environment would not negatively hurt the learning of the world model while the old experience far away could still provide a good prior for the world model in the unknown states to prevent catastrophic forgetting. The closeness between two states required in this method is obtained by contrastive learning on a set of transitions collected by a random behavior policy in the new environment. The paper shows that the proposed method can improve upon the standard first-in-first-out (FIFO) replay buffer of different size on top of various existing model-based RL approaches in multiple custom tasks.", "summary_of_the_review": "While the proposed adaptation approach is conceptually simple and the idea of local forgetting is appealing, it is unclear to me if the proposed approach is practically useful since the paper is missing the baseline of simply retraining a new RL agent from scratch in the new environment (Phase 2). In addition, the paper is missing important discussions and comparisons of the proposed method in relation to existing continual methods (which could serve as potentially strong baselines). Therefore, I would not recommend for acceptance of the paper at its current state. \n\nA set of more thorough experiments with harder tasks where Phase 2 is difficult to solve without the experience of Phase 1 could potentially make the paper much stronger.", "strength_and_weaknesses": "*Strength*\n- The proposed method is conceptually simple and can be combined with existing model-based RL methods \n- The paper is well written with great clarity on the proposed method that seems reproducible. \n\n*Weaknesses*\n- I have doubts over the difficulties of the tasks tested in the paper. In particular, it is unclear to me that whether agents need to use the knowledge from the Phase 1 (environment with old reward function) to perform well in Phase 2 (environment with new reward function) at all. I cannot find the performance of training the agent from scratch directly in Phase 2 (e.g., in Figure 3 and Figure 5). \n- The related work section is very terse and there are not many discussions how this work is similar or different from prior approaches in the continual learning literature (e.g., [1, 2, 3]). Some of them could be valid baselines. For example, instead of training the world model on the full replay buffer, we can train on a small subset of the most informative data from the Phase 1 environment (e.g., coreset [2]), or even discard the data from the Phase 1 environment completely and fine-tune the world model with parameter regularization on the data from the new environment only (e.g., EWC [1]).\n- In the last paragraph, the authors wrote that \"this is the first time–to the best of our knowledge that adaptivity for the deep MBRL methods has been shown.\" -- I think this statement is incorrect (see [4])\n\n[1] Kirkpatrick, James, et al. \"Overcoming catastrophic forgetting in neural networks.\" Proceedings of the national academy of sciences 114.13 (2017): 3521-3526.\n[2] Nguyen, Cuong V., et al. \"Variational continual learning.\" arXiv preprint arXiv:1710.10628 (2017).\n[3] Rolnick, David, et al. \"Experience replay for continual learning.\" Advances in Neural Information Processing Systems 32 (2019).\n[4] Nagabandi, Anusha, Chelsea Finn, and Sergey Levine. \"Deep online learning via meta-learning: Continual adaptation for model-based RL.\" arXiv preprint arXiv:1812.07671 (2018).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject