Mastering Memory Tasks with World Models
Authors: Mohammad Reza Samsami
Artem Zholus
Janarthanan Rajendran
Sarath Chandar
Venue: ICLR 2024 oral (2024)
Keywords: model-based reinforcement learning
state space models
memory in reinforcement learning
Primary area: reinforcement learning
TL;DR: We propose R2I, a model-based agent with enhanced memory capabilities which shines in challenging memory reinforcement learning tasks.

Abstract: Current model-based reinforcement learning (MBRL) agents struggle with long-term dependencies. This limits their ability to effectively solve tasks involving extended time gaps between actions and outcomes, or tasks demanding the recalling of distant observations to inform current actions. To improve temporal coherence, we integrate a new family of state space models (SSMs) in world models of MBRL agents to present a new method, Recall to Imagine (R2I). This integration aims to enhance both long-term memory and long-horizon credit assignment. Through a diverse set of illustrative tasks, we systematically demonstrate that R2I not only establishes a new state-of-the-art for challenging memory and credit assignment RL tasks, such as BSuite and POPGym, but also showcases superhuman performance in the complex memory domain of Memory Maze. At the same time, it upholds comparable performance in classic RL tasks, such as Atari and DMC, suggesting the generality of our method. We also show that R2I is faster than the state-of-the-art MBRL method, DreamerV3, resulting in faster wall-time convergence.

Published as a conference paper at ICLR 2024
MASTERING MEMORY TASKS WITH WORLD MODELS
Mohammad Reza Samsami Artem Zholus 1.3 Janarthanan Rajendran
Sarath Chandar
Mila Quebec Al Institute Université de Montréal SPolytechnique Montréal *CIFAR Al Chair
ABSTRACT
Current mrodel-based reinforcement learning (MBRL) agents struggle with long- term dependencies. This limits their ability to effectively solve tasks involving extended time gaps between actions and outcomes, or tasks demanding the recalling of distant observations to inform current actions. To improve temporal coherence, integrate new family of state space models (SSMs) in world models of MBRL agents to present a new method, Recall to Imagine (R2I). This integration aims to enhance both long-term memory and long- long-horizon credit assignment. Through diverse set of illustrative tasks, we systematically demonstrate that R2I not only establishes a new stateoofthhe-art for challenging memory and credit assignment RL tasks, such as BSuite and POPGym, but also showcases superhuman performance in the complex memory domain of Memory Maze. At the same time, it upholds comparable performance in classic RL tasks, such as Atari and DMC, suggesting the generality of our method. We also show that R21 is faster than the state-of-the- art MBRI. method, DreamerV3, resulting in faster wall-time convergence.
INTRODUCTION
In reinforcement learning (RL), world models (Kalweit & Boedecker, 2017; Ha & Schmidhuber, 2018; Hafner al., 2019b), which capture the dynamics of the environment, have emerged as a powerful paradigm for integrating agents with the ability to perceive (Hafner et al., 2019a; 2020; 2023), simulate (Schrittwieser et al., 2020; Ye et al., 2021; Micheli al., 2023), and plan (Schrittwieser et al., 2020) within the learned dynamics. In current model-based reinforcement learning (MBRL), the agent learns the world model from past experiences, enabling to "imagine" the consequences of its actions (such the future environment rewards and observations) and make informed decisions.
MBRL necessitates learning a world model that accurately simulates the environment's is evolution and future rewards, integrating the agent's 's actions over long horizons. This task is compounded by the credit assignment (CA) problem, where an action's 's impact on future rewards must be evaluated. The agent also may need memorize and recall past experiences to infer optimal actions. The challenge of long-term memory and CA frequently arises as a result of inadequate learning of long-range dependencies (Ni et al., 2023), due constraints in world models' backbone network architecture
More specifically, Recurrent Neural Networks (RNNs; Cho et al. (2014)) are employed in most MBRL methods (Ha & Schmidhuber, 2018: 8; Hafner et al., 2019b;a; 2020; 2023) as the world models' backbone architecture because of their ability to handle sequential data. However, their efficacy is hindered by the vanishing gradients (Bengio et al., 1994; Pascanu et al., 2013). Alternately, due the remarkable achievements of Transformers (Vaswani et al., 2017) in language modeling tasks (Brown et al., 2020; Thoppilan et al., 2022), they have been recently adopted to build world models (Chen et al., 2022; Micheli et al., 2023; Robine et al., 2023). Nonetheless, the computational complexity of Transformers quadratic its input sequence length. Even the optimized Transformers (Dai et al., 2019; Zaheer et al., 2021; Choromanski et al., 2022; Bulatov et al., 2022; Ding et al., 2023) become unstable during training on long sequences (Zhang ct al., 2022). This prohibits Transformers based world models from scaling to long input sequence lengths that might be required in certain RL tasks.
Recent studies have revealed that state space models (SSMs) can effectively capture dependencies in tremendously long sequences for supervised learning (SL) and self- self-ssuer learning (SSL) tasks
'Equal contribution. {mohammade d-raaa samsami, artem. zholus Soolussam 3mi a quebec See our website here: recal 122 magi he, hw..iihhhttte github.

---

Published as a conference paper at ICLR 2024
(Gu et al., 2021a; Nguyen et al., 2022; Mehta et al., 2022; Smith et al., 2023; Wang et al., 2023). More specifically, the S4 model (Gu et al., 2021a) redefined the long-range sequence modeling research landscape by mastering highly difficult benchmarks (Tay et al., 2020). The S4 model is derived from a time-invariant linear dynamical system where state matrices are learned (Gu et al., 2021b). In SL and SSL tasks, it exhibits remarkable capability to capture dependencies extending up 16K in length, surpassing the limitations of all prior methods. Given these achievements and MBRL methods' limitations in solving memory and CA tasks, the adoption of S4 or modified version of it is logical decision. In this paper, we introduce novel method termed Recall to Imagine (R21), which is the first MBRL approach utilizing a variant of S4 (which was previously employed in model-free RL (David et al., 2023; Lu et al., 2024)). This method empowers agents with long-term memory. R2I emerges as a general and computationally efficient approach, demonstrating state-of-the-art (SOTA) performance in a range of memory domains. Through rigorous experiments, we demonstrate that R2I not only surpasses the best-perfone rest-performing basclines but also exceeds human performance in tasks requiring long-term memory credit assignment, all while maintaining commendable performance across various other benchmarks. Our contributions can be summarized as follows:
We introduce R2I, a mremory-enannced MBRL agent built upon DreamerV3 (Hafner et al., 2023) that uses a modification of S4 to handle temporal dependencies. R2I inherits the generality of Dreamer operating with fixed world model hyperparameters on every domain, while also offering an improvement in computational speed of up times.
We demonstrate SOTA performance of the R2I agent in diverse set of memory domains: POPGym (Morad et al., 2023), Behavior Suite (BSuite; Osband et al. (2020)), and Memory Maze (Pasukonis al., 2022). Notably, in the Memory Maze, which is challenging 31) domain with extremely long-term memory needed to be solved, R21 outperforms human.
We investigate R2I's performance in established RL benchmarks, namely Atari (Bellemare et al., 2013) and DMC (Tassa et al., 2018). We show that R2I's improved memory does not compromise performance across different types of control tasks, highlighting its generality.
We conduct ablation experiments to show the impact of the design decisions made for R2I.
2 BACKGROUND
2.1 STATE SPACE MODELS
A recent work (Gu et al., 2021a) has introduced a novel Structured State Space Sequence model (S4). This model has shown superior performance in SL and SSL tasks, compared common deep sequence models, including RNNs, convolutional neural networks (CNNs; lec (1998)), and Transformers. outperforms them in terms of both computational efficiency (Gu et al., 2021b) and the ability to model extremely long-range dependencies (Gu et al., 2020). S4 is specific instance of state space models (SSMs), which can be efficiently trained by using specialized parameterization.
SSMs are derived from linear dynamical system with control variable u(t) E and observation variable y() E R. utilizing state variables <t) E CN for a state size N. The system represented by the state matrix A E CNNN and other matrices B E Nx1 C E CIXN, and D E R RIx1
x' (t) Ar(t) + Bu((), y(t) Cz(() + Du(t).
(1)
Note that these SSMs function on continuous sequences. They can be discretized by step size A to allow discrete recurrent representation:
A + Bun: Yn. Cr,, + Dun,
where A, B, C, and D are discrete-time parameters obtained from the continuous-time parameters and A using methods like zero-order hold and bilinear technique (Smith et al., 2023). These representations
neural network layer, and each SSM is used to process single dimension of the input sequence and map it to a single output dimension. This means that there are separate linear transformations for each input dimension, which are followed by nonlinearity. This allows working with discrete sequence tasks, such as language modeling (Merity et al., 2016), speech classification (Warden, 2018), 8), and pixel-level ID image classification (Krizhevsky et al., 2009).
incorporated
S4 model characterizes A as a matrix with a diagonal plus low-rank (DPLR) structure (Gu ct al., 2021a). One benefit of this "structured" representation is that it helps preserve the sequence history;
2

---

Published as a conference paper at ICLR 2024
S4 employs HiPPO framework (Gu al., 2020) to initialize the matrix A with special DPLR matrices. This initialization grants the SSMs the ability to decompose u(t) into a set of infinitely long basis functions, enabling the SSMs capture long-range dependencies. Further, to make S4 more practical on modern hardware, Gu et al. (2021a) have reparameterized the mapping u1:7, 30 - y1:r, 2r as global convolution, referred to as the convolution mode, thereby avoiding sequential training (as in RNNs). This modification has made S4 faster to train, and as elaborated in Gu et al. (2021b), S4 models can be thought of as a fusion of CNNs, RNNs, and classical SSMs. Smith al. (2023) uses parallel scan (Blelloch, 1990) to compute u1:r: Zo + V1:T,, FIr efficient as convolution mode.
S4 has demonstrated impressive empirical results on various established SL and SSL benchmarks involving long dependencies, and it outperforms Transformers (Vaswani et al., 2017; Dao et al., 2022) in terms of inference speed and memory consumption due to its recurrent inference mode. Moreover, some recent works have focused on understanding S4 models, as well as refining them and augmenting their capabilities (Gupta et al., 2022a; Gu et al., 2022; Mehta et al., 2022; Gupta et al., 2022b; Smith et al., 2023; Ma et al., 2023). We have provided additional details in Appendix B to explain this family of S4 models. For the sake of simplicity in this study, we will be referring to all the S4 model variations as "SSMs". It is worth highlighting that few recent methods optimize the performance of SSMs by integrating them with Transformers (Fu et al., 2023; Zuo et al., 2022; Fathi et al., 2023). This enhances the SSMs by adding powerful local attention-based inductive bias.
2.2 FROM IMAGINATION To ACTION
We frame sequential secision-making problem as a partially observable Markov decision process (POMDP) with observations Or, scalar rewards rr, agent's 's actions dN episode continuation flag c, and discount factor 7 E (0, 1), all following dynamics Os, Tt, Ct ~ p(o,, Fr, Ct Oces, act). The goal of RL is train a policy π that maximizes the expected value of the discounted return Ex [E r'r:].
In MBRL, the agent learns model of the environment's dynamics (i.e., the world model), through an iterative process of collecting data using policy, training the world model on the accumulated data, and optimizing the policy through the world model (Sutton, 1990; Ha & Schmidhuber, 2018). The Dreamer agent (Hafner et al., 2019a) and its subsequent versions (Hafner et al., 2020; 2023) have been impactful MBRL systems that learn the environment dynamics in a compact latent space and learn the policy entirely within that latent space. Dreamer agents consist of three primary components: the world model, which predicts the future outcomes of potential actions, the critic, which estimates the value of each state, and the actor, which learns to take optimal actions.
In Dreamer, an RNN-based architecture called Recurrent State-Space Model (RSSM), proposed by Hafner et al. (2019b), serves as the core of the world model, and it can be described as follows. For every time step 6, it represents the latent state through the concatenation of deterministic state he and stochastic state 2,- Here, h2 is updated using Gated Recurrent Unit (GRU; Cho et al. (2014)), and then is utilized to compute 2g, which incorporates information about the current observation o, and is subsequently referred to the posterior state. Additionally, the prior state ê, which predicts 2t without access to og is computed using h4. By leveraging the latent state (zz, ht), we can reconstruct various quantities such as Ot, Tt, and c,. The RSSM comprises three components: sequence model (he = Ta(hh-1, 2t- (t-1))))) )), a representation model (ze qo go(2, f4, or)), and a dynamics model (&g ~ Po
h4)), where at-1 is the action time step 1. and f denotes the combined parameter vector of all components. In addition to the RSSM, the world model has separate prediction heads for or, rr. G. Within the imagination phase, it harnesses the RSSM simulate trajectories. This is performed through an iterative computation of states ε,, h, and actions ân (ā, | & h4) without the need for observations (except in the initial step). The sequences of ε hy:T, â1.T are used to train the actor and the critic. See Appendix D for more details.
3 METHODOLOGY
We introduce R2I (Recall to Imagine), which integrates SSMs in DreamerV3's world model, giving rise to what term the Structured State-Space Model (S3M). The design of the S3M aims to achieve two primary objectives: capturing long-range relations in trajectories and ensuring fast computational performance in MBRL. S3M achieves the desired speed through parallel computation during training and recurrent mode in inference time, which enables quick generation of imagined trajectories. In Figure 1, a visual representation of R2I is provided, and we will now proceed to describe its design.
3

---

Published as a conference paper at ICLR 2024
Learning the world model
Acting in imagination
fu
fr2
Deterministic Nodes
Stochastic Nodes
Prediction Dends
Information Fusion
Figure/ 1: Graphical representation of R2I. (Left) The world model encodes past experiences, transforming observations and actions into compact latent states. Reconstructing the trajectories serves as learning signal for shaping these latent states. (Right) The policy learns from trajectories based on latent states imagined by the world model. The representation corresponds to the full state policy, and we have omitted the critic for the sake of simplifying the illustration.
3.1 WORLD MODEL DETAILS
Non-recurrent representation model. Our objective when updating the world model is to caleulate S3M deterministic states hEr in parallel by simultaneously feeding all actions a and stochastic state 2t, where 7 represents the length of the entire sequence. We aim carry out this computation as I.,, ZT.T [a(((a 21.T), =D) where is a hidden state and fe is sequence model with SSM network. To achieve this, prior access to all actions a:: and stochastic states ž1:T required. However, we encounter challenge due to the sequential nature of the relationship between the representation model ( he, 01) and sequence model fe 2-1 1): at time step t, the representation model's 's most recent output, denoted 2,-1, is fed into the sequence model, and the resulting output h4 is then used within the representation model to generate 2, Hence, similar to Chen et al. (2022); Micheli et al. (2023); Robine al. (2023); Deng et al. (2023), by eliminating the dependency on fü in the representation model, we transform it to a non-recurrent representation model (= oc). This modification allows us to compute the posterior samples independently for each time step, enabling simultaneous computation for all time steps. By utilizing parallelizable function fe. we can then obtain h1r in parallel. Appendix M includes a systematic analysis to investigate how this modification impacts the performance of the DreamerV3 across a diverse set of tasks. The results indicate that transforming qo ct oe; h4) to (z oz) does not hurt the performance.
Architecture details. Inspired by Dreamer, R2I's world model consists of a representation model, a dynamics model, and a sequence model (together forming S3M). In addition to that, there are three prediction heads: an observation predictor (á Zz, he), reward predictor (f | 2,, he). and an episode continuation predictor (é Ze, ht). At each time step, S3M processes a pair of (ae, 21) to output the deterministic state h4. Inside, it operates over the hidden state Iı, it can be defined as h4 xz = fe ((a-1,, 1t,1),,,,1) Specifically, fo is composed of multiple layers of SSMs, each one calculating outputs according to Equation 2. The outputs are then passed to GeLU (Hendrycks & Gimpel, 2023), which is followed by fully-connected GI.U transformation (Dauphin al., 2017), and finally by LayerNorm (Ba et al., 2016). This follows the architecture outlined by Smith et al. (2023), The deterministic state h2 is the output from the final SSM layer. The set of all SSM layer hidden states is denoted 24. See Appendix B.I for SSMs design details. In image-based environments, we leverage CNN encoder for qo (z ox) and a CNN decoder for (o 24, hz). In contrast, in tabular environments, both ( ot) and (o 24, h4) are MLPs. We include the details on network widths, depths, and other hyperparameters in Appendix E.
Training details. R2I optimizes the following objective:
((0)
E 0
(0, h4 Tt, C1, c,,2) 21) L hy, 0t) (8,
(3)
4

---

Published as a conference paper at ICLR 2024
gred (0, h2, 01. re, Ct; 2x) Speed po 2x, he) + po 2, he) + In 2t, he))
(4)
Lam (0, Ne, 0t) = Bayn max (1, KIl[sg o()) p(z h4) D)
(5)
L (0, (0, 0t) = Beep max [L| go I ot) [|sg(p(= [4))
(6)
h1,T U1.7 - (o((a1,,,2I.)):..
(7)
Here, sg represents the stop gradient operation. This loss, resembling the objective utilized in (Hafner et al., 2023), is derived from Evidence Lower Bound (ELBO), but our objective differs from ELBO in three ways. First, we clip KL-divergence when falls below the threshold of I(Hafner et al., 2020; 2023). Secondly, we use KI KL-balancing (Hafner et al., 2020; 2023) to prioritize the training of the S3M. Third, we use scaling coefficients Speed See. Buyn to adjust the influence of each term in the loss function (Higgins et al., 2017; Hafner et al., 2023). Some works on SSMs recommend optimizing state matrices using a smaller learning rate; however, our experiments indicate that the most effective approach is to the same learning rate used in the rest of the world model.
SSMs Computational Modeling. To enable the parallelizability of world model learning, as outlined in Section 1. have the option to select between two distinct approaches: convolution (Gu et al., 2021a) and parallel scan (Smith et al., 2023). After thorough deliberation, we opted for parallel scan due several compelling reasons. Firstly, as discuss later in Section 3.2, it is essential to pass hidden states 2, to the policy in memory environments, critical finding we empirically analyze in Appendix N. Another consequence of not yielding 2N via convolution mode is that it would necessitate several burn-in steps to obtain correct hidden states, akin to Kapturowski et al. (2019), resulting in quadratic imagination complexity. Furthermore, parallel scan enables scaling of sequence length in batch across distributed devices, capabil- ity not supported by the convolution mode. Table 6 summarizes compu- tational complexities associated with different types of recurrences, includ- ing RNNs, SSMs, and Attention used in studies like Chen et al. (2022).
Inference step
Imagination step
State
Method
Training
Parallel
Reset
Attn
O(L2)
(
O(L + m1)
/
RNN
O(L)
O(1)
x
SSM (Conv)
O(L)
0(1)
O(L)
/
x
SSM (Par.Scan)
O(L)
0(1)
0(1)
/
/
Table 1:
The asymptotic runtimes of different architectures.
L is the sequence length and H is the imagination horizon. The outer loop of the imagination process cannot be paral- lelized. Attention and SSM+Conv accept the full context of (T. + H) burn-in and imagined steps which results in O(L - H)²) step complexity for Attention and ((L) for SSM+Conv. SSMs combine compact recurrence with paral- lel computation reaching the best asymptotical complexity.
Finally, parallel scan can facilitate the resetting of hidden states. When sam- pling sequence from the buffer, it may comprise of multiple episodes; thus, the hidden states coming from terminal states to the initial states new episodes must be reset. This boosts the early training performance, when the episodes may be short. Inspired by Lu et al. (2024), we modify the SSM inference operator to support resetting hidden states. Achieving this is not feasible with convolution mode. Details of our SSMs operator used by the parallel scan is provided in Appendix C.
3.2 ACTOR-CRITIC DETAILS
In the design of Dreamer's world model, is assumed that h2 contains information summarizing past observations, actions, and rewards. Then, h4 is leveraged in conjunction with the stochastic state 2, to reconstruct or predict observations, rewards, episode continuation, actions, and values. Unlike D3eamerV3, which utilizes GRU cell wherein h2 is passed both the reconstruction heads and the next recurrent step, R21 exclusively passes h2 to prediction heads, while SSM's hidden state &e used in the next recurrent update of S3M. This implies that the information stored in he and xt could potentially vary. Empirically, we discovered that this difference can lead to the breakdown of policy learning when using (ā, 22, h4), but it remains intact when we use x (ā, | 2e, 2x) in memory-intensive environments. Surprisingly, we found that incorporating all features into the policy (ā; 2,, (,) is not a remedy. The reason lies in the non-stationarity of these features; their empirical distribution changes over time the world model trains, ultimately leading instability in the policy training process. A similar phenomenon was also observed in Robine et al. (2023). We study the dependency of policy features the performance in Appendix N, where we cover a diverse set of environments: from non non-memory vector- rector-based ones to image- based memory environments. In different environments, we condition the policy and value function on the information from S3M
5

---

Published as a conference paper at ICLR 2024
in the following ways: use the output state policy that takes (2),h) input, the hidden state policy that takes (2e, xx) as input, and the full state policy that takes (2,, he, zy) as imput. To train actor-critic, opt for the procedure proposed in DreamerV3 (Hafner et al., 2023). For detailed description of the actor-critic training process, refer to Appendix D.
4 EXPERIMENTS
We conduct comprehensive empirical study to assess the generality and memory capacity of R2I across a wide range of domains, includ- ing credit assignment, memory-intensive tasks, and non-memory tasks, all while maintaining fixed hyperparameters of the world model. We cover five RL domains: BSuite (Osband et al., 2020), POPGym (Morad et al., 2023), Atari 100K (Lukasz Kaiser et al., 2020), DMC Tassa et al. (2018), and Memory Maze (Pasukonis al.,
Vector Environments
Train Step Time, a Train Step Time.
10a
10-)
10-2
32
128
512
2048
Image Environments
2022). The section is organized as follows. In Sections 4.1 and 4.2, we evaluate R2I's performance in two distinct memory-intenswve settings: simple tabular environments and complex 3D environments. We show that not only does R2I achieve the SOTA performance, but it also surpasses human-level performance in the complex Memory Maze domain. In Section 4.3, we demonstrate that we do not trade the generality for improved memory capabilities. Figure 2 shows R2I's impressive computational efficiency, with a speed increase of up to 9 times compared to its predecessor, DreamerV3, Note that the image environments are representative of Memory Maze, and the vector environments represent POPGym.
10a
10-1
33.2
10-2
128
512 2048
Sequence Length in Batch
DreamerV3
(Ours)
Figure 2: Computational time taken by DreamerV3 and R2I (lower is preferred)
We reuse most of the world model hyperparameters from D3eamerV3. In all environments, we use a First-in First-out (FIFO) replay buffer size of 10M steps train R2I. We found this helps stabilize the world model and prevent overfitting on a small buffer. Also, we vary features that the policy conditioned on (i.e., output state policy π (ā, 2e, he), hidden state policy (āă 2t, xz). or full state policy x xā, 2,, hr, l). Our primary takeaway is to leverage the output state policy in non- memory environments and the full state policy or hidden state policy within memory environments, as explained in Section 3.2. We also found that even in memory environments, the full state policy cannot be preferred over the hidden state policy because of the instability of features - since the world model is trained alongside the policy, the former might change the feature distribution which introduces non cottstationarity for the policy.
4.1 QUANTIFYING MEMORY OF R21
In this section, we study the performance of R2I in challenging memory environments of BSuite and POPGym domains, which are tabu- lar environments, Despite their simplicity, these environments pose a challenge for MBRL al- gorithms since the world model needs learn causal connections over time. While SSMs have shown their ability to handle extremely long- range dependencies in SL and SSL (Gu et al., 2021a), this capability does not necessarily trans- late to MBRL, even though the world model optimizes the same supervised objective. This discrepancy arises from the lifelong nature of world model training. That is, needs to boot- strap its performance from a very small dataset with hugely imbalanced reward "labels" (as opposed datasets on which SSMs shine (Tay et al., 2021)). Additionally,
Discounting chain
Memory Longth
0.8
0.6
0.5
0.4
0.2
0.0
100
Dewernt
(Dum)
Dreamer'i
Figure 3: Success rates of DreamerV3 (which holds the previous SOTA) and R2I BSuite en- vironments. A separate model is trained for ev- cry point on the x-axis. A median value (over 10 seeds) is plotted filling between 25-th and 75-th percentiles. Training curves are in Appendix F.
to big and the continuously growing replay buffer imposes the need quickly learn the newly arrived data which requires an ability for quick adaptation of the world model throughout its optimization. The section's 's goal is to give insight into how extensive R2I's memory capabilities.

---

Published as a conference paper at ICLR 2024
RepeatPrev
Autoencode
Concentration
1.0
Final reward
0.5
0.5
0.5
0.0
0.0
-0.5
-0.5
Easy
Medium
Hard
Easy
Medium
Hard
Easy
Medium
Hard
R21 (Ours)
DreamerV3
PRO+MLP
PRO+GRU
PPO+LSTM
PPO+SAD
Figure 4: R2I results in memory-intensive environments of POPGym. Our method es- tablishes the new SOTA in the hardest memory environments; Autoencode: -Easy, -Medium: um; RepeatPrevious: -Medium, -Hard; Concentrat ration: -Medium. Note that Concentration task that can be partially solved without memory. For PPO+S4D, refer to Appendix S.
Behavior Suite experiments. To study the ability of the R2I model to handle longer episodes, we conduct quantitative experiments within a subset of the BSuite environments. These environ- ments are specifically designed to evaluate an agent's memory capacity and its ability to effectively perform credit assignment. In particular, we carry out experiments within Memory Length and D Discount ing Chain environments. The former focuses on memory, and the latter serves as credit assignment task. In Memory Length environment, the goal is to output an action which is dictated by the initial observation (the episode length i.e., the memory steps number is environment parameter). Essentially, the agent must carry the information from the initial observation throughout the entire episode. In the Discount scounting Chain, the first action (which is categorical) causes reward that is only provided after a certain number of steps, specified by the parameter reward delay.
As depicted in Figure 3, the previous SOTA Dreamer learns the dependencies between actions and rewards in both Discount iomuutinnn ng Chain and Memory Length with reward delays of up to 30 environment steps. Note that every either converged the random seed). We plot the success rate as the fraction of runs that achieved success. R2I excels in both tasks, significantly outperforming in the preservation of its learning ability across wider range of varying environment complexities. In these experiments, we leverage the output state policy (i.c., operating on latent variable 2 and S3M output h2). More details are provided in Appendix E.
maximum reward or failed (based on
POPGym experiments. We perform a study to assess R2I in a more challenging benchmark, namely, POPGym (Morad et al., 2023). This suite offers range of RL. environments designed to assess various challenges related to POMDPs, such as navigation, noise robustness, and memory. Based on Ni al. (2023), we select the three most memory- -intensive environments: Sepeattrevionn Autoencode, and foncentrations These environments require an optimal policy to memorize the highest number of events (i.e., actions or observations) at each time step. Each environment in POPGym has three difficulty levels: Easy, Medi um, and Hard. In the memory environments of this study, the complexity is increased by the number of actions or observations that the agent should keep track of simultaneossly. All environments in this study have categorical observation and action spaces. A detailed explanation of the environments is provided in Appendix G.
As POPGym was not included in the DreamerV3 benchmark, we performed hyperparameter tuning of both DreamerV3 and R21, solely on adjusting the network sizes of both. This is because Dreamer V3 generalist agent that works with fixed set of hyperparameters and in this environment, with sizes primarily influencing its data efficiency. We observed a similar characteristic in R2I. The results of hyperparameter tuning are available Appendix L. For R21, we use the hidden state policy: T (ā: Zz, (b)) as we found it much more performant, especially in memory-intensive tasks (see Appendix N for policy inputs ablations). We train R2I in POPGym environments using a unified and fixed set of hyperparameters. In addition to R2I and DreamerV3, we include model-free baselines from Morad et al. (2023). These include PPO (Schulman et al., 2017) model-free policy with different observation backbones, such as GRU, LSTM, MLP, and MLP with timestep number added feature (PosMLP). PPO with GRU is the best- perf-rrrfo model-free baseline of POPGym while PPO+LSTM is the second best. PPO+MLP and PPO+PosMLP are included for sanity check the better their performance is, the less is the memory needed in the environment.
AA policy without any memory exists that outperforms random policy but underperforms the optimal one.
7

---

Published as a conference paper at ICLR 2024
As illustrated in Figure 4, R2I demonstrates the new SOTA performance, outperforming every baseline in Aut oencode, Easy and Medium um tasks. Note that R21 outperforms all 13 model-free baselines of the POPGym benchmark by a huge margin (we did not include them due to space constraints). R2I also shows consistently strong performance in RepeatPreve ous tasks, setting a new SOTA in both Medí um and Hard (compared to all 13 model-free baselines and DreamerV3). In Concentrat ration, the model-free memory baselines fail to outperform a simple MLP policy, suggesting that they all converge to a non-memory-basee towards better memory policy. Its performance is roughly equal to DreamerV3 in an Easy and slightly better in the Medi um task. As Appendix G suggests, all PeepaatPrevion ous tasks require up to 64 memorization steps, while Aut oencoee Easy and Medi um require up to 104. In Concent rat on Easy and Med um this length up to 208 steps, however, since PPO+MLP shows somewhat good performance, likely less than 208 memorization steps are required. This observation is consistent with the results of the BSuite experiments, which demonstrate that our model is capable of memorizing up approximately 100 steps in time. To summarize, these results indicate that R2I significantly pushes the memory limits.
suboptimal policy. R2I advances this
4.2 EVALUATING LONG-TERM MEMORY IN COMPLEX 3D TASKS
Memory Maze (Pasukonis et al., 2022) presents randomized 3D mazes where the egocentric agent is repeatedly tasked to navigate to one of multiple objects. For optimal speed and efficiency, the agent must retain information about the locations of objects, the maze's ' wall layout, and its own position. Each episode can extend for up to 4K environment steps. An ideal agent equipped with long-term memory only needs to explore each maze once, a task achievable in a shorter time than the episode's 's duration; subsequently, it can efficiently find the shortest path to reach each requested target. This task poses fundamental challenge for existing memory- -ugmeete RL algorithms, which fall significantly behind human performance in these tasks.
In this benchmark, we found that DreamerV3 works equally well as DreamerV2 reported in Pasukonis et al. (2022). Therefore, we use the size configuration of Dreamer outlined in Pasukonis et al. (2022). Note that this baseline also leverages truncated backpropagation through time (TBTT), a technique demonstrated to enhance the preservation of information over time (Pasukonis et al., 2022). We use the "medium memory" size configuration of R21 in this work (see Table 2 in Appendix). We use the stochastic state, and S3M output, and hidden
full state policy (x (1(a) zt, xt) i.e., conditioning states at step t) in this environment. We trained and tested R2I and other methods on 4 existing maze sizes: 9x9, 11x11, 13x13, and 15x15. The difference between them in the number of object rooms and the episode lengths. More difficult maze sizes have more environment steps in the episode making more challenging to execute a successful series of object searches. R2I and other baselines are evaluated after 400M environment steps or two weeks of training. We also compare R2I with IMPALA (Espeholt al., 2018), which is the leading model-free approach (Pasukonis et al., 2022).
As shown in Figure 5, R21 consistently outperforms baseline methods in all of these environments. In 9x9 mazes, it demonstrates performance similar to the Dreamer, while significantly outperformig IMPALA. In 1 11x11, [3x13, and 15x15 mazes, it has remarkably better performance than both baselines. Moreover, it has surpassed human-level abilities in solving 9x9, 11x11, and 13x13 mazes. These results establish R2I SOTA in this complex 3D domain.
MemoryMazee aaae-zz--9 9x9
MemoryMaze 11x11
MemoryMaze. 13x13
Memory HaamoryMazz-155ae
80
Retum
40-
60
40
2.0
10
20
20
0
0
1
1eB
1c8
leB
1eB
Environment
Environment step
Environment stop
Environment step
R21 (Ours)
Dreamer (TBTT)
IMPALA
--... Human
Figure 5: Scores in Memory Maze after 400M environment steps. R2I outperforms baselines across difficulty levels, becoming the domain's 's new SOTA. Due to its enhanced computational efficiency, R2I was trained during fewer number of days compared to Dreamer, as illustrated in Figure 26.
8

---

Published as a conference paper at ICLR 2024
4.3 ASSESSING THE GENERALITY OF R2I IN NON-MEMORY DOMAINS
We conduct a sanity check by assessing R21's
Atari Gamer Mean (26 tasks)
Atari Gamer Median
performance on widely used RL bench-
Return Human normalized
(26 tasks)
marks: Atari (Bellemare et al., 2013) and DMC (Tassa et al., 2018), as parts of the DreamerV3
0.3
score
0.2
benchmark (Hafner et al., 2023). Even though these tasks are nearly fully observable and
0.5
0.1
not necessitate extensive memory to solve (it
200k
400k
200k
400k
often enough to model the dynamics of only
DMC-proprio (18 tasks)
DMC-visual (20 tasks)
the last few steps), evaluating R2I on them is essential as we aim to ensure our agent's ' perfor- mance across wide range of tasks that require different types of control: continuous control (in DMC) and discrete (in Atari).
500
500
0
0.0
0.5
1.0
2
Env step
1e6
Env step
1e6
In all the experiments conducted within Atari 100K (Eukasz Kaiser et al., 2020) and DMC, we fix hyperparameters of the world model. In Atari and the proprio benchmark in DMC, utilize output state policies, as we found them more performant (for ablations with different policy types. see Appendix N). In the visual benchmark in DMC, we use hidden state policy. Note that for continuous control, the policy is trained via differentiating through the learned dynamics. R2I maintains performance similar to DreamerV3 in these domains, as demonstrated in Figure 6, implying that in the majority of standard RL tasks (see Appendix Q), R21 does not sacrifice generality for improved memory capabilities.
R2I (Ours)
Dreamer
Figure 6: Average performance in Atari and DMC. For full training curves, see Appendix K (Atari), Appendix (DMC-P), and Appendix J (DMC-V)
CONCLUSION
In this paper, we introduced R2I, a general and fast model-based approach to reinforcement learning that demonstrates superior memory capabilities. R2I integrates two strong algorithms: DreamerV3, a general-purre purpose MBRL algorithm, and SSMs, a family of novel parallelizable sequence models adept at handling extremely long-range dependencies. This integration helps rapid long torg-term memory and long-horizon credit assignment, allowing R2T to excel across a diverse set of domains, all while maintaining fixed hyperparameters across all domains. Through systematic examination, demonstrated that R21 sets a new state-of-the-an in domains demanding long-term temporal reasoning: outperforms all known baselines by large margin on the most challenging memory and credit assignment tasks across different types of memory (long-term and short-term) and observational complexities (tabular and complex 3D). Remarkably, it transcends human performance in complex 3D tasks. Furthermore, we have demonstrated that R2I achieves computation speeds up to 9 times faster than D3eamerV3.
have
Our study presents the first model-based RI. approach that uses SSMs. While R21 offers benefits for improving memory in RL, it also has limitations, which we leave for future research. For instance, can be explored how R2I can be augmented with attention mechanisms, given that Transformers and SSMs exhibit complementary strengths (Mehta et al., 2022). As mentioned in Section 2.1, hybrid architectures have been introduced in language modeling tasks. Moreover, the sequence length within the training batches for world model learning is not currently extremely long, as is the horizon (i.e., the number of steps) of imagination in actor-critic learning. Future work could focus these aspects to further enhance memory capabilities.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We thank Albert Gu for his thorough and insightful feedback on the SSM part of the project. We also acknowledge The Annotated S4 Blog (Rush & Karamcheti, 2022) and S5 codebase (Smith et al., 2023) which inspired our JAX implementation. We also thank Danijar Hafner, Steven Morad, Ali Srahimi-Kalahrour Michel Ma. Tianwei Ni, Darshan Patil, and Roger Creus for their helpful feedback our method and the early draft of the paper. We thank Jurgis Pasukonis for sharing the data for memory maze baseline plots. This research was enabled by computing resources provided by Mila (mila.quebec), the Digital Research Alliance of Canada (alliancecan.ca) and NVIDIA (nvidia.com) We thank Mila's s IDT team, and especially Olexa Bilamiuk for helping with numerous technical questions during this work and especially for the help in the implementation of the new 1/O efficient RL replay buffer. Janarthanan Rajendran acknowledges the support of the IVADO postdoctoral fellowship. Sarath Chandar is supported by the Canada CIFAR AI Chairs program, the Canada Research Chair in Lifelong Machine Learning, and the NSERC Discovery Grant.
REFERENCES
Gradient based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998.
Mohammed Abbad. Perturbation and stability theory for Markov control problems. University of Maryland, Baltimore County, 1991.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare G Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in Neural Information Processing Systems, 2021.
Jose A Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brand- stetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. Advances in Neural Information Processing Systems, 32, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, jun 2013. doi: 10. 1/jair/.ar...91.. URL https: http///oooo.....111.......// 3912.
Yoshua Bengio, Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157-166, 1994.
Guy E. Blelloch. Prefix sums and their applications. Technical Report CMU-CS-90-190. School of Computer Science, Carnegie Mellon University, November 1990.
William L Brogan. Modern control theory. Pearson education india, 1991.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Vlerbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter. Christopher Hesse, Mark Chen. Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
Aydar Bulatov, Yuri Kuratov, and Mikhail Burtsev. Recurrent memory transformer, 2022.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models, 2022.
Kyunghyun Cho, Bart Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau. Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using mn encoder-decoder for statistical machine translation, 2014.
10

---

Published as a conference paper at ICLR 2024
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2022.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context, 2019.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022.
Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learming, volume 70 of Proceedings of Machine Learning Research, pp. 933-941, PMLR, 06-11 Aug 2017. URL https: http::::/pprceeeings..... mlr. www.ppress/vvv. vv20 dwwphinnphnnini html.
Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision s4: Efficient sequence-based RL via state spaces layers. In The Eleventh International Conference on Learning Representations, 2023, URL https: httph:::/pppeeeieeeeeeeee www.reveeeeeeeeeeeomeeeeeeeoeeem net forum ? id=kgjikvV77woo
Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: Rnns, transformers, and s4, 2023.
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers 1,000,000,000 tokens, 2023.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg. and Koray Kavukcuoglu. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures, 2018.
Kevin Esslinger, Robert Platt, and Christopher Amato. Deep transformer q-networks for partially observable reinforcement learning, arXiv preprint arXiv:2206.010% 01078, 2022.
Mahan Fathi, Jonathan Pilault, Pierre-Luc Bacon, Christopher Pal, Orhan Firat, and Ross Goroshin. Block-state transformer, 2023.
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
Albert Gu. Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state-space layers. Advances in Neural Information Processing Systems, 34, 2021b.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models, 2022.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, 2022a.
Ankit Gupta, Harsh Mchta, and Jonathan Berant. Simplifying and understanding state space models with diagonal linear rnns, 2022b.
David Ha and Jürgen Schmidhuber. World models. 2018. doi: 10.5281/ZEN000.122772 URL
https: / zenodo. wtg/reecod//11773...
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 19Xiv:1912.01603 2019a.

---

Published as a conference paper at ICLR 2024
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.
2019b.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. urXiv preprini arXiv:2010.02193 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.0444 2023.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations. 2017. URL wttph:tt//// http////ppeeeeeeeeeee/eee/eeeeee net wttet//oorummmddyyyfdaaa
Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):5223, 2019.
Joshua Joseph, Alborz Geramifard, John W Roberts, Jonathan How, and Nicholas Roy. Rein- forcement learning with misspecified model classes. In 2013 IEEE International Conference on Robotics and Automation, pp. 939-946. IEEE, 2013.
Gabriel Kalweit and Joschka Boedecker. Uncertainty- imagination for continuous deep rein- forcement learning. In Sergey Levine, Vincent Vanhoucke, and Ken Goldberg (eds.). Proceedings of the 1st Annual Conference Robot Learning, volume 78 of Proceedings of Machine Learn- ing Research, pp. 195-206. PMLR, 13-15 Nov 2017. URL https: /ttpppppeeeng....///p///// mlr.. press/vv] 1et////wwwii111 html.
Steven Kapturowski, Georg Ostrovski, Will Dabney, John Quan, and Remi Munos. Recurrent experience replay in distributed reinforcement learning. In International Conference Learning Representations, 2019. URL https: //tp///pppeeeeeeee/ee/eeeee net/ forum pid=rllyyjjaggi.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523 2020.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention, 2023.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. 13947, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models, 2023.
Thomas M. Moerland, Joost Broekens, Aske Plaat, and Catholijn M. Jonker. Model-based reinforce- ment learning: survey, 2022.
Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym: Benchmarking partially observable reinforcement learning. In The Eleventh International Confer- ence on Learning Representations, 2023. URL //tpp::///peeeeview net /tse//forree forum?id= chDrutUTsOK.
12

---

Published as a conference paper at ICLR 2024
Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah. Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos as multidimensionall signals with state spaces. Advances in neural information processing systems, 35:2846-2800 2861, 2022.
Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in ri? decoupling memory from credit assignment, 2023.
Exgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-oriented model-based reinforcement learning with implicit differentiation, 2021.
Ian Oshand, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvári, Satinder Singh, Benjamin Van Roy, Richard Sutton, David Silver, and Hado van Hasselt. Behaviour suite for reinforcement learning. In International Conference on Learning Representations, 2020. URL https: /ttp:////pppeeeeeeeeeeee wwt.wiiwww.eee.eee..........ee nwt/forumt forum? id=rygf-kSYwH.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference machine learning, pp. 7487-7498. PMLR, 2020.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks, 2013.
Jurgis Pasukonis, Timothy Lillicrap, and Danijar Hafner. Evaluating long-term memory in 3d mazes. arXiv preprint arXiv:2210.133833 2022.
David Raposo, Sam Ritter, Adam Santoro, Greg Wayne, Theophane Weber, Matt Botvinick, Hado van Hasselt, and Francis Song. Synthetic returns for long-term credit assignment. arXiv preprint arXiv:2102.12425, 2021.
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer- world models are happy with 100k interactions. In The Eleventh International Conference on Learning Representations, 2023. URL httph::tp///c https::/ wwt.:iiwwwweeeeeeeeeeeeeeeeeeeee
net htte:///oormmiddddaaaaaeeaa
Alexander Rush and Sidd Karamcheti. The annotated s4. In ICLR Blog Track, 2022. URL https: //icll--bloog-rrack..oooo fmet..gthhbbo22222///555aaaaaac
wttpp/ioo/22///22 1ooii/////0 22///022222233 001222//3////// htt////anotttttddeeeeeee s4/. https:/ficlr-boog-
Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588 (7839):604-66) 609, dee 2020, doi: 10. 00000/44418000000055333 URL https: //cppp:///o...o org/100 (03882Fs41B
020-0300000000000000000000000000 03051-40
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling, 2023.
Richard S Sutton. Integrated architectures for learning, planning, and reacting based approximating dynamic programming. In Machine learning proceedings 1990, pp. 216-224, Elsevier, 1990.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Riedmiller. Deepmind control suite, 2018.
Yi Tay, Mostafa Dehghani. Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers, 2020.
13

---

Published as a conference paper at ICLR 2024
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference Learning Representations, 2021. https:
// //openrevi ew net/ forurrrrmiio i Sateeidd-iyyeeeeeeee
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng. Alicia Jin, Taylor Bos, Leslie Baker. Yu Du, YaGuang Li, Hongrae Lee, Huaixiu Steven Zheng, Amin Ghafouri, Marcelo Menegali, Yanping Huang. Maxim Krikun, Dmitry Lepikhin, James Qin, Dehao Chen, Yuanzhong Xu, Zhifeng Chen, Adam Roberts, Maarten Bosma, Vincent Zhao, Yanqi Zhou, Chung-Ching Chang, Igor Krivokon, Will Rusch, Marc Pickett, Pranesh Srinivasan, Laichee Man, Kathleen Meier-Hellstern, Meredith Ringel Morris, Tulsee Doshi, Renelito Delos Santos, Toju Duke, Johnny Soraker, Ben Zevenbergen, Vinodkumar Prabhakaran, Mark Diaz, Ben Hutchinson, Kristen Olson, Alejandra Molina, Erin Hoffman-John, Josh Lee, Lora Aroyo, Ravi Rajakumar, Alena Butryna, Matthew Lamm, Viktoriya Kuzmina, Joe Fenton, Aaron Cohen, Rachel Bernstein, Ray Kurzweil, Blaise Aguera-Arcan Claire Cui, Marian Croak, Ed Chi, and Quoc Le. Lamda: Language models for dialog applications, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M. Rush. Pretraining without attention, 2023.
Qiang Wang, Bei Li, Tong Xiao, Jingho Zhu, Changliang Li, Derek F. Wong, and Lidia S. Chao. Learning deep transformer models for machine translation. In Anna Korhonen, David Traum, and Lluís Märquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1810-1822, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v 1/P199 1176. URL https: ltlplps://aclaannne /aa/a/aaaaannooooooooaoaaaooooll org/p19-1177.
Pete Warden. Speech commands: dataset for limited-vocabulary speech recognition, 2018.
Michael Widrich, Markus Hofmarcher, Vihang Prakash Patil, Angela Bitto-Nemling, and Sepp Hochreiter. Modern hopfield networks for return decomposition for delayed rewards. In Deep RL Workshop NeurlPS 2021, 2021.
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach. Learn., 8(3-4):229-256, may 1992. ISSN 0885-6125. doi: (0.1007/BF00992666 URL https:
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data, 2021.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences, 2021.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li. Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao, and Jianfeng Gao. Efficient long sequence modeling via state space augmented transformer. arXiv preprint arXiv:2212..88136, 2022.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2020. URL https:/ openreview n forum? D
14

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
A Related Work
16
B Variations of SSMs and Our Design Choices
16
B.1 Design Decisions for SSM
17
Managing Several Episodes a Sampled Sequence
D Training Actor-Critic
18
E Algorithm Hyperparameters
19
F BSuite Training Curves
20
POPGym Environments Details
21
II POPGym Training Curves
22
I DMC-proprio Training Curves
23
J DMC-vision Training Curves
24
Atari 100K Training Curves
25
L Hyperparameter Tuning in POPGym
26
M Impact of non- non-Recurrent Representation Model
27
N Policy Input Ablations
28
Criticality of Full Episode Inclusion Training Batch
30
P Model Likelihood Relation with Agent Performance
31
P.I Comparing R2I with S4WM
31
Q Unveiling the Performance in Standard RL Benchmarks
34
R Algorithm Pseudocode
35
S Model-free SSM POPGym
T Detailed Results in Memory Maze
37
15

---

Published as a conference paper at ICLR 2024
A RELATED WORK
Addressing sequential decision-making problems requires the incorporation of temporal reasoning. This necessitates a crucial element of intelligence memory (Parisotto et al., 2020; Esslinger et al., 2022; Morad et al., 2023). Essentially, memory enables humans and machines retain valuable information from past events, empowering us make informed decisions at present. But it does not stop there; another vital aspect is the ability to gauge the effects and consequences of our past actions and choices on the feedhack/outcome we receive now be it success or failure. This process is termed credit assignment (Hung et al., 2019; Arjona-Medina et al., 2019; Widrich et al., 2021; Raposo et al., 2021). In their recent work, Ni et al. (2023) explore the relationship between these two concepts, examining their interplay. The process of learning which parts of history to remember is deeply connected to the prediction of future rewards. Simultaneously. mastering the skill of distributing current rewards to past actions inherently involves a memory component.
To improve temporal reasoning, we leverage model-based RL (Moerland et al., 2022). The essence of MBRL lies in its utilization of a world model (Kalweit & Boedecker, 2017; Ha & Schmidhuber, 2018; Hafner et al., 2019b). The world model employed to plan a sequence of actions that maximize the task reward (Hafner et al., 2019a; 2020; 2023; Schrittwieser et al., 2020; Ye et al., 2021; Micheli et al., 2023). Our focus on MBRL stems from its potential in sample efficiency, reusability, transferability, and safer planning, along with the advantage of more controllable aspects due to its explicit supervised learning component (i.e., world modeling). Within the realm of MBRL methodologies, we decided to build upon the SOTA DreamerV3 framework (Hafner et al., 2023), which has proven capacities for generalization, sample efficiency, and scalability.
MBRL researchers commonly utilize RNNs (Ha & Schmidhuber, 2018; Hafner et al., 2019a; 2020; 2023) or Transformers (Chen et al., 2022; Micheli et al., 2023; Robine et al., 2023) as the backbone of their world models to integrate temporal reasoning into the agent. However, both of these approaches face limitations when it comes to modeling long-range dependencies, as stated earlier in Section 1. Recent research indicates that SSMs (Gu et al., 2021a;b; Gupta et al., 2022a; Gu et al., 2022; Mehta et al., 2022; Gupta et al., 2022b; Smith et al., 2023; Ma et al., 2023) can replace Transformers, capturing dependencies in very long sequences more efficiently (sub-quadratic complexity) and in parallel. Thus, they can be a good candidate for a backhone architecture of world models. Concurrently with this work, S4WM (Deng al., 2023) has also explored the utilization SSMs for enhancing world models within the context of SSL, subtask in MBRL. While S4WM is an improvement in the task of world modeling, it does not extend to proposing methods for MBRL. Furthermore, enhanced world modeling alone does not inherently ensure increased performance in RL tasks (Nikishin et al., 2021), a point that further elaborated upon in Appendix P. This suggests that despite the progress in world modeling accuracy and efficiency, additional strategies must be integrated to translate these improvements into tangible RL gains, as detailed in Appendix P.I.
B VARIATIONS OF SSMS AND OUR DESIGN CHOICES
Structured State Space Sequence model (S4: Gu et al. (2021a)) is built upon classical state space models, which are commonly used in control theory (Brogan, 1991). As mentioned in the Section 2.1, $4 is built upon three pivotal and independent pillars:
Parametrization of the structure. A pivotal aspect is how we parameterize the matrices. S4 conjugates matrices into a diagonal plus low-rank (DPLR).
Dimensionality of the SSM. S4 utilizes separate linear transformations for each input dimension, known as single-input single-output (SISO) SSMs.
Computational modeling. The third axis revolves around how we model the computation process to parallelize. Gu et al. (2021a) model the computation as global convolution.
Recent works have introduced certain modifications to S4, specifically targeting these pillars. Some research (Gupta et al., 2022a; Smith et al., 2023) demonstrate that employing diagonal approximation for the HiPPO (Gu et al., 2020) yields comparable performance and thus represents the state matrices diagonal, instead of DPLR.
16

---

Published as a conference paper at ICLR 2024
Regarding the second axis, Smith al. (2023) tensorize the 1-D operations into multi-inpol multi- output (MIMO) SSM, facilitating a more direct and simplified implementation of SSMs. Furthermore, Smith et al. (2023) replace the global convolution operation with parallel scan operation, thereby eliminating the complex computation involved the convolution kernel.
The synergy between these axes of variations may significantly influence the agent's s behavior and performance. In the following section, we will demonstrate our process for choosing each option and examine how each axis influences the performance. We further will provide more in-depth information regarding our choices for SSM hyperparameters.
B.1 DESIGN DECISIONS FOR SSM
B.1.1 SSM PARAMETRIZATION
While parameterizing the state matrix as a DPLR matrix could potentially offer a more expressive representation compared to diagonal SSM, our decision was diagonalize This decision was made to simplify control while maintaining high performance. Notably, the works of Gupta et al. (2022a) and Gu et al. (2022) demonstrate the feasibility of achieving S4's level of performance by employing significantly simpler, fully diagonal parameterization for state spaces.
B.1.2 SSM DIMENSIONALITY
We decided to employ MIMO SSMs, a choice influenced by a convergence of factors. Firstly, by utilizing MIMO, the latent size be substantially reduced, optimizing computational resources and accelerating processing. As result, the implementation of efficient parallelization facilitated.
Further, MIMO SSMs eliminate the need for mixing layers and offer the flexibility accommodate diverse dynamics and couplings of input features within each layer. This stands in contrast to independent SISO SSMs, which process each channel of input features separately and then combine them using a mixing layer. Accordingly, the MIMO approach enables distinct processing tailored to individual imput channels, leading enhanced model adaptability and improved feature representation (Smith et al., 2023).
To validate the benefits of this decision, an ablation study was conducted. This comparative analysis showcases the differences between MIMO and SISO SSMs, empirically demonstrating the advantages of the MIMO.
C MANAGING SEVERAL EPISODES IN A SAMPLED SEQUENCE
As mentioned Section handling episode boundaries in the sampled sequences necessitates the world model's capability to reset the hidden state. Given our utilization of parallel scan (Blelloch, 1990) for SSM computational modeling, it follows that adaptations the binary operator Smith et al. (2023) are essential. As a reminder, SSMs unroll a linear time-invariant dynamical system in the following form:
Ar + Bun,
Un Ca,, + Dun
Smith et al. (2023) employ parallel scans for the efficient computation of SSM states, denoted as Parallel scans leverage the property that "associative" operations can be computed in arbitrary order. When employing an associative binary operator denoted as on sequence of L clements
[e1, . ez). parallel scan yields Je1, (ei e (....... (Ei e ee)] )]] It is worth mentioning that, parallel scan can be computed with complexity of (log L). Smith et al. (2023), the elements e; and the operator
defined as follows:
en = (en,0; en = (A, Bun), e, Ej = (ej.0 x 6j,,, Cj,0 6,,1 + Cy,1)
Here, X is the matrix-matric product, and is the matrix- product. To make resettable states, we need to incorporate the "done" flag while ensuring the preservation of the associative property. To do so, every element e, is defined as follows:
en (Cn.0, Cn,1, Cn,2) (A, Bunn dn)
17

---

Published as a conference paper at ICLR 2024
where d,, represents the done flag for the nith element. Now, present the modified binary operator defined as:
Gj ((1 d:) e3,0 6, + 6,2 6j.0=
(1 di)e e,, + e,11
clip(er2 + 6,2, 1))
Where the clip function clips the value between zero and It is important to highlight this reparameterization of the operator defined in Lu et al. (2024), eliminates the need for if/else conditions. Thus, discussed in Lu al. (2024), the modified operator remains associative and retains the desirable properties we aimed to achieve.
Unlike Lu al. (2024),
longer assume hidden state initialized to 20 0; hence, different initialization becomes necessary, specifically eo (o, 0).
D TRAINING ACTOR-CRITIC
The policy and value functions are trained using latent imagination, regardless of their input features. This process aligns with the methodology outlined in Section 2.2. begins with the generation of initial stochastic states 21.7 ~ (0( (1:r) and the computation of the sequence of deterministic hidden states h1.7, #1:T fø((a1:T, 21:T).20) through parallel scan operation. We reuse the previously computed values of 21.r, hir,, and Z1.7 during training, which the parallel scan facilitates, thus eliminating the requirement for several burn-in steps before imagination.
Denoting Zz}; the state of after j steps of imagination, given t context steps, we initialize with Coje = Fz, hojs lt, and P0iC = Zt- We then compute the action āoje =(ā Zojt hojt, e) and after hije, Tije = (( 201r). Fo &lje ~ "po(2 | h t). This process continues for H steps, during which compute (t: 21,Hfts Risltt rewards and continue flags computed via f1.M p(f &i.mee himm)) Gr.rje p(č &I. h1.r)),, respectively. The ultimate goal is to
train the policy to maximize the estimated return that follows:
R2 fyge + Cype ( (1oo))) Foje) + AB2+1)), R = =(3jte Ayg
(8)
To train the actor and critic, we employ Reinforce (Williams, 1992) for scenarios with discrete action spaces, coupled with the method of backpropagation through latent dynamics introduced by Hafner et al. (2020). Our approach adheres to the DreamerV3 protocol training the actor and critic networks (Hafner et al., 2023), which includes utilization of the fixed entropy bonus, the use of twohot regression critic, and the employment of percentile based normalization for the returns.
18

---

Published as a conference paper at ICLR 2024
E ALGORITHM HYPERPARAMETERS
Small
Medium
Small
Memory
Memory
Used in
BSuite and POPgym
Atari and DMC
Memory Maze
h1 size
512
512
2048
3t. size, per layer
512
192
512
SSM layers
3
5
SSM units
1024
512
1024
Table 2: S3M size configurations employed in this study.
Name
Value
FIFO replay buffer size
10"
Batch length, I.
1024
Batch size
4
Nonlinearity
LayerNorm + SiLU
SSM discretization method
bilinear
SSM nonlinearity
GeLU GLU LayerNorm
SSM matrices parameterization
Diagonal MIMO 8
SSM dimensionality parameterization
SSM matrix blocks number (HiPPOs number)
SSM discretization range
(10-3 (10-3.10-1) Multi-categorical
Latent variable
Categorical latent variable numer
32
Categorical classes number
32
Unimix probability
0.01
Learning rate
4
Reconstruction loss weight, Spred Dynamics loss weight, Spree Representation loss weight, Srep
1
0.5
0.1
World Model gradient clipping
1000
Adam epsilon
10-8
Actor Critic Hyperparameters
Imagination horizon
15
Discount 7
0.997
Return X
0.95
Entropy weight
3 10-4
Critic EMA decay
0.98
Critic EMA regularizer
Return normalization scale
Per(R, 95) Per(12,5)
Return normalization decay
0.99 10-5
Adam epsilon
Actor-Critic gradient clipping
100
Table 3: Hyperparameter details for R2I. For an in-depth description, please refer Section 3 and consult Hafner al. (2023).
19

---

Published as a conference paper at ICLR 2024
F
BSUITE TRAINING CURVES
memory steps
memory steps
7 memory steps
0.5
0.5
0.5
0.0
0.5
1c6
memory steps
11 memory steps
13 memory steps
1.0
0.5
0.5
0.5
15 memory steps
1c6
18 memory steps
1c6
21 memory steps
1c6
1.0
1.0
1.0
0.0
0.0
0.0
0.0
1e6
1e6
1e6
26 memory steps
memory steps
41 memory steps
1.0
1.0
1.0
0.0
0.0
0.0
0.0
0.5
1.0
0.0
0.5
0.0
0.5
memory steps
1c6
61 memory steps
1c6
71 memory steps
1e6
1.0
1.0
1.0
0.0
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1e6
1e6
1e6
memory steps
memory steps
101 memory steps
0.5
0.0
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1c6
1e6
DrcamerVa
R21 (Durs)
Figure 7: R2I and DreamerV3 training curves in BSuite Memory Length environment, averaged over 10 independent runs per environment (Median reward plotted).
Reward Delay: step
Reward Delay: steps
Reward Delay: 10 steps
1.10
1.10
1.10
1.05
1.05
1.05
1.00
1.00
1.00
1
1e6
1e6
1e6
Reward Delay: 30 steps
Reward Delay: 100 steps
1.10
1.10
1.05
1.05
1.00
1.00
166
(Durs)
DreamerV3
Figure 8: R2I and D3eamerV3 training curves in BSuite Aascounting Chain environment,
averaged over 10 independent runs per environment (Median reward plotted).
20

---

Published as a conference paper at ICLR 2024
G POPGYM ENVIRONMENTS DETAILS
In this research, we examine three environments within the POPGym domain (Morad et al., 2023), each configured with Easy, Medi um, and Hard levels of difficulty to systematically increase the memory demands for the agents for memory-intensive tasks. In those tasks, the fundamental challenge escalates with the difficulty level, necessitating the memorization and subsequent recall of larger number of observations or actions for the more advanced tiers. Table 4 summarizes the different aspects of these tasks.
RepeatPrevious
Autoencode
Concent rantentraatt
-E
-M
-H
-B
-M
-H
-E
-M
-H
Parallel events to recall
4
32
64
52
104
156
104
208
104
The size of the observation space Memory-less suboptimal policy
4
4
3
3
14
X
Table 4: Memory complexities of environments in POPGym. "Parallel events" refers to the maximum number of actions or observations that the agent needs to recall. The third row indicates whether the task can be partially solved without memory. -E is -aasy; -M is -Medium; -H is -Hard.
Repeat.Preiion In this environment, both observation and action spaces are categori- cal with 4 discrete values. The task requires the agent to replicate an observation d++* = or that occurred & steps prior, at any given time step / + k. The agent at each time step needs to simultaneously remember lt categorical values; accordingly, the value of k depends on the difficulty level, being 4 for Easy, 32 for Medium, and 64 for Hard, regardless of the fixed observation space size (which is 4).
Autoencode. The environment operates in two distinct phases. the "watch" phase, the environment generates discrete observation from set of 4 possible values at each time step: agent actions do not affect the observations during this phase. This phase spans the first half of the episode, lasting T/2 steps, where T is the total length of the episode. The episode lengths are set to 104, 208, and 312 steps for the Easy, Medi um, and Hard difficulty levels, respectively. In the subsequent phase, the agent must recall and reproduce the observations by outputting corresponding actions, which are drawn from the 4-value space, similar the observation space. Essentially, the first step in the second phase should be equal to the first observation in the "watch" phase, the subsequent action in the second phase ought to be the same the second observation from phase one, and so on. The agent receives both a phase indicator and the categorical ID to be repeated in its observation. Note that at each time step, the agent must remember T/2 categorical values; thus, successful completion of this task requires explicit memory mechanism, and any policy performing better than random chance is utilizing some form memory.
Concent rat ion. At each step, the agent receives an observation consisting of multiple categories, each with N distinct values. The number of categories is 52 for Easy and Hard, and 104 for Medi um. The values of N are for Easy and Medi um, and 14 for Hard. The task simulates card game with a deck of cards spread face down. The agent can flip two cards at each step, and if they match, they remain face up; otherwise, they are turned face down again. The agent's 's observation includes the state of the full deck, and the episode length corresponds to the minimal average number of steps required to solve the task optimally as determined by Morad et al. (2023). To outperform random policy, the agent must remember card positions and values to find matches more efficiently. Even without memory, an agent can avoid flipping cards already turned face up, which conserves steps without yielding new information or reward, thereby outperforming random policies at basic level. This is marginal yet conspicuous improvement that shows a memory-less policy can gain over a random policy. Note that due to the episode length constraint, this task cannot be solved without memory, as proven by (Morad et al., 2023).
21

---

Published as a conference paper at ICLR 2024
H POPGYM TRAINING CURVES
RepeatPreviousEasy
RepeatPreviousMediumm
RepeatPreviousHard
1.0
1.0
1.0
0.0
0.0
0.0
0.5
-0.5
-1.0
-1.0
-1.0
2
0
3
2
5
1e7
1e7
AutoencodeEasy
MutoencodeMedium
AutoencodeHard
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
-05
-0.5
-1.0
-1.0
-1.0
1e7
1e7
1e7
ConcentrationEasy
ConcentrationMedium
ConcentrationHard
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
0.5
0.5
0.5
-1.0
-1.0
-1.0
1
2
3
5
0
3
4
1
2
4
1c7
1c7
1c7
DreamerV3
R21
Figure 9: R21 and DreamerV3 training curves in the POPGym benchmark, averaged over 3 indepen- dent runs per environment.
22

---

Published as a conference paper at ICLR 2024
1 DMC-PROPRIO TRAINING CURVES
Finger Spin
Pendulum Swingup
Cartpole Swingup
1000
500
500
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Walker Stand
1e6
Acrobot Swingup
Finger Turn Hard
1000
1000
200
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Cartpole Balance
1e6
Walker Run
1e6
Cartpole Swingup Sparse
1e6
1000
500
500
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1e6
1e6
Hopper Stand
Walker Walk
Cup Catch
1000
1000
1000
0
0.0
0.5
1.0
0.0
0.5
0.0
0.5
1.0
1e6
1e6
1e6
Reacher Hard
Hopper Hop
Finger Turn Easy
1000
1000
200
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1e6
1e6
1e6
Cheetah Run
Cartpole Balance Sparse
Reacher Easy
1000
1000
500
0
0
0.5
1e6
1e6
1e6
Average
500
0
0.0
0.5
1.0
1e6
DreamerV3
R21
Figure 10: R2I and DreamerV3 training curves across proprioceptive environments the DMC benchmark, averaged over 3 independent runs per environment.
23

---

Published as a conference paper at ICLR 2024
J DMC-VISION TRAINING CURVES
Finger Turn Hard
Walker Run
Walker Stand
1000
1000
500
500
0
0
2
1e6
1e6
Cheetah Run
Quadruped Run
Cup Catch
1000
500
500
0
0
2
Finger Spin
1e6
1e5 Cartpole Swingup Sparse
Quadruped Walk
166
500
500
0
0
0
1
2
Cartpole Balance Sparse
Cartpole Balance
1e6
Hopper Hop
1e6
1000
1000
250
500
0
2
le6
1e6
1e6
Acrobot Swingup
Cartpole Swingup
Walker Walk
500
1000
500
0
0
0
2
Reacher Hard le6
le6
16
Hopper Stand
Finger Turn Easy
1000
1000
1000
Reacher Easy
1e6
Pendulum Swingup
Average
1e6
1000
1000
500
0
2
1e6
1e6
R21
DreamerV3
Figure 11: R2I and DreamerV3 training curves across visual environments in the DMC benchmark, averaged over 3 independent runs environment.
24

---

Published as a conference paper at ICLR 2024
K ATARI 100K TRAINING CURVES
Atari 100K benchmark (Lukasz Kaiser et al., 2020) a standard RL benchmark comprising 26 Atari games featuring diverse gameplay mechanics. is designed assess broad spectrum of agent skills, and agents are limited to executing 400 thousand discrete actions within each environment, which is approximately equivalent to 2 hours of human gameplay. To put this in perspective, when there are no constraints on sample efficiency, the typical practice is to train agents for 200M steps.
Allen
Amidar
Assault
Asterix
1000
200
1000
2000
500
100
500
1000
200k
400K
200k
400k
200k
400k
200k Breakout
400k
Bank Heist
Battle Zone
Boxing
1000
20000
100
50
100
500
10000
0
0
200k
400k
0
200k
400k
200k
400k
200k
400k
Chopper Command
Crazy Climber
Demon Attack
Freeway
0.05
1000
500
50000
250
0.00
-0.05
0
200k Frostbite
400k
200k
400k
200k
400k
200k
400k
Gopher
Hero
James Bond
1000
10000
10000
400
500
0
0
200k
400k
200k
400k
200k
400k
200k
400kc
Kangaroo
Krull
Kung Fu Master
Ms Pacman
10000
3000
20000
2000
2500
5000
1080
0
200k
400k
200k
400k
200k
400k
200k
400k
Pong
Private Eye
Qbert
Road Runner
4000
20000
0
2000
2000
-20
0
200k
400k
200k
400k
200k
400K
200k
400k
Seaquest
Up Down
Gamer Mean
Gamer Median
0.3
5000
1.0
500
0.5
0.1
0.0
200k
200k
400k
200k
400k
200k
400k
DreamerV3
Figure 12: R21 and DreamerV3 training curves in the Atari after 400K environment steps, averaged over 3 independent runs per environment
25

---

Published as a conference paper at ICLR 2024
L HYPERPARAMETER TUNING IN POPGYM
Encoder/Decoder depth (layers)
RSSM number units
RSSM hidden size of GRU
Replay Capacity
-0.05
-0.10
Final Return
-0.15
-0.20
-0.25
-0.30
1
256
512
1024
256
512
1024
1M
10M
Figure 13: Results of hyperparameter tuning for DreamerV3 in POPGym memory environments.
To establish robust baseline for comparison, we conducted an extensive hyperparameter search on the bascline DreamerV3 model. The results of this search are delineated in Figure 13, Each hyperpa- rameter configuration of DreamerV3 V3 was trained on 9 memory tasks from POPGym considered in this study. To tune the model, we considered the following hyperparameters:
1. Number of units in RSSM, with the sizes of auxiliary layers ranging from 256, 512, 1024;
2. Hidden size of GRU in RSSM with similar values;
3. Replay buffer size with options, 10 10";
4. Sizes of encoder and decoder.
Each configuration used two random seeds over 10 million environment steps. The rationale underly- ing our choices anchored in DreamerV3's origins as a general agent with all hyperparameters fixed except network sizes and training intensity, which control sample efficiency. To identify optimal configuration, we tuned the model sizes around the configuration proposed for Behavior Suite (since POPGym and Behavior Suite share similar properties). We also checked increased replay buffer size of 10M steps (in addition to the standard 1M of DreamerV3) since it improved the performance for R21.
Figure 13 shows box plot for each hyperparameter with the marginal performance of each hyper- parameter (i.e., for hyperparameter X, we fix but vary other hyperparameters, forming an empirical distribution with median, max, min, etc.). We plot box plot of the empirical distribution for each X. This visualization shows how much each hyperparameter can contribute the overall performance (as it shows the median with all other hyperparameters varied but chosen one fixed) and also what is the best hyperparameter value for the task (since the box plot shows the maximal value). Our results confirm the favorable scaling of DreamerV3 (Hafner et al., 2023) in this new environment bigger networks learn bigger rewards. Therefore, we opt for 1024 RSSM units, an RSSM hidden size of 1024, and 5 layers of encoder/decoder. Unlike for R21, for DreamerV3, bigger replay performs less "stable" since it has a bigger variation (e.g., increasing the network size might have an unexpected influence the final agent performance). Yet, since with 10 million steps the model has the best score, we opt for this option for a fair comparison. The final Dreamer model reported in Figure 4 was trained on the best-found configuration.
"eepeetPrevi ous, Autoencode, Concent rati on for each difficulty: Masy, Medi um, Ha
26

---

Published as a conference paper at ICLR 2024
M IMPACT OF NON-RECURRENT REPRESENTATION MODEL
In this section, we study the ramifications of severing the link between the representation model and the sequence model. As previously discussed in Section 3, disrupting this connection facilitates the parallel computation of SSM sequence features (otherwise, if the representation model is fed by the sequence model's 's previous output, the whole computation must be sequential). Our investigation revolves around the impact of maintaining the recurrent pathway from the sequence model to the representation model.
First, we test how removing this recurrent connection affects the performance in the Memory Maze. We do so in two versions of this benchmark. The first one the standard Memory Maze reported in this work. Subsequent to this, we explore an augmented variant of the benchmark that supplements the task with the coordinates target objects, demanding the agent to process the visual input to reconstruct both the input image and the target positions (Pasukonis et al., 2022).
MemoryMaze-9x2
MemoryMaze-1xx
MemoryMaze-1344
MemoryMaze-15x15
10
Retur
20
20
2
2
5
0
1
2 3
5
2
1e7
1e7
1e7
1e7
Environment step
Environment step
Environment step
Environment step
DreamerV3: qitt ft)
DreamerV3: giz:
02)
Figure 14: Representation model ablation in Memory Maze.
MemoryMaze 9x9
MemoryMaze-111
RemoryMaze-13x11
MemoryMaze-15xx
40
40
Return
20
20
20
0
0
1
2
4
2
Environment step 1e7
1e7
1e7
1e7
Environment step
Environment step
Environment step
DreamerV3 probe head: gize | oz, ftd)
DreamerV3 + probe gizt [o])
Figure 15: Representation model ablation in Memory Maze with probing head.
The empirical evidence, depicted in Figures 14 and 15, suggests that the non-recurrent representa- tion model retains or elevates performance levels. Note that the introduction of a probing network equips the agent with additional memory supervision, as it tasked predict target positions. Hence, in the absence of such memory supervision, the non-recurrent representation model demonstrates enhanced performance, while its presence, it sustains performance without loss.
In other words, the non- recurrent representation model is an inductive bias that favors memory. As shown in Figure 16. in standard RL benchmarks, this recurrent connection has zero influence on performance, confirming that this inductive bias favors memory without any negative side effects.
Atari Gamer Mean (26 taskss)
Atari Gamer Median (26 tasks)
DMC-propria (18 tasks]
0.4
750
500
250
200k
0.0
200k
1.0
giz, [o,)
Figure 16: Representation models in standard RL environments.
27

---

Published as a conference paper at ICLR 2024
N POLICY INPUT ABLATIONS
In this section, we aim to decipher which policy input yields su- perior performance. As reminder, S3M comprises three main components: the deterministic output state ft, leveraged by pre- diction heads to summarize information from previous steps; the stochastic state 4t, representing each single observation; and the hidden state 2y, which is passed between time steps. The policy variants are the output state policy x (ā, Rt, he): the hidden state policy π (āē 2t; Ty), and the full state policy T (ā; 2,, h2, 2t).
Discounting
0.8
0.6
0.4
0.2
0.0
20
40
60
80
100
Rewaidd delay, stops
nta, R..K)
na: 22)
R2I: a.. a..ha.
Figure 17: Actor-critic input abla- tions in D scount ng Chain.
Our main insight that in different environments hr, 2t, and St. behave differently and they may represent different information. Additionally, the evolution of their distributions during world model training hurts the actor-critic training stability. Hence, found each domain has its working configuration. The only shared trait is that memory-intensive environments generally prefer policies that incorporate It. either through x (āe 2t, (āā ze; N4, v2) finding corroborated by our analyses Figures 17 and 18. In the specific context of the Memory Maze, the potency of input feature modification is smaller. As Figure 19 shows, the output state policy performs worse than the other two.
However, there exists an extended ver- sion of the Memory Maze benchmark. This task adds additional information which is the target objects' coordi- nates. Although the agent is still given only an image as an observation, is now tasked to reconstruct both the input image and target positions (Pa- sukonis et al., 2022). As Figure 20 shares, illustrates, all three policy vari- ants achieve comparable reward out- comes, suggesting that the differentials between the input states spreefiiill ( (z,, v:)) and (z,, hz)) configurations do not yield distinct advantages in this augmented task scenario.
(9
memoryy-intensive tacks)
non-memory-intenee tasks)
1.0
0.5
0.5
-0.5
-0.5
-1.0
0.5
1.0
R2). maa 22.
R21 maa 22.h
R2t 14.
Figure 18: Actor-critic input ablations in POPGym.
Concluding our analysis, Figure 21 reveals a nuanced but noteworthy observation: the variations in policy inputs exert marginal, albeit detectable, effect on the final performance. This implies considerable overlap in the informational content encapsulated by (Z,, h2) and (z),
MemoryMaze-9M
MemoryMaze 11x11
MemoryMaze-134m 13x13
MemoryMaze-15x19
Return
20
20
10
10
0
2
1e7
1e7
1e7
Environment step 1e7
Environment step
Environment step
Environment step
R21; n(ar]s;.x)
R21; n(ar]s, x.)
R21; n(ai(s..h)
Figure 19: Memory Maze policy ablations without probing heads. The performances of all three variants are qualitatively different indicating that the information in the S3M output ht and S3M hidden 2, is different. In addition, the output state policy π oā; | 2,, he) exhibits a decline in performance as the training progresses.
Non- Nonn-memory-ineensse environments include: CountRecall, Battleship. MineSweeper. Repeat Fi rst, labyrinthisscap,, iabyrinthsxplore, All three difficulties (Rasy, i Ha rd) for each.
28

---

Published as a conference paper at ICLR 2024
MemoryMaze-9x3
MemoryMaze-1Ix11
MemoryMaze-13xL5
MemoryMaze-15x4
40
Return
10
20
20
20
5
2
3
4
0
1
2
3
4
1
2
3
1e7
1e7
Environment step
Environment step
Environment step
Environment step
R21 probe: nla:/5p
R21 probe: mla:|s.
R21 probe; reled/5p At)
Figure 20: With probing head in Memory Maze, the agent is tasked predict object locations from he (S3M output), effectively making h2 a Markovian state (i.e. equivalent to hidden state Xiy). Thus, we can see the performance of all three policy variants effectively the same.
Atari Gamer Mean (26 tasks)
Atari Gamer Median
DMC-proprio
(26 tasks)
(18 tasks)
0.4
750
1.0
0.3
500
0.2
0.5
250
0.1
0.0
0
200k
400k
200k
400k
0.0
0.5
1.0
1e6
R2I: n(a, 2z xx)
R21: n(a; Z ho x;)
R21: n(a; h.)
Figure 21: Actor-critic input ablations Atari 100K and DMC domains.
29

---

Published as a conference paper at ICLR 2024
O CRITICALITY FULL EPISODE INCLUSION TRAINING BATCH
In the domain of memory-intensive tasks, sequence length within training batches plays a pivotal role in the final per- formance. Our findings suggest that the success of our model strongly contingent on accommodating the entire episode within single sequence in the batch. This requirement ensures that the model can fully leverage temporal dependencies with- the truncation of episodes. Thanks to parallel scan (Blelloch, 1990) (please refer to sections 2.1 and for elaboration), the batch sequence length scales as efficiently as scaling the batch size. Therefore, we opt train the model on complete episodes the training batch across experimental environments.
Memory Length
1.0
Success Rate
0.5
0.0
75
100
Memory steps
R21: sequence length-256
R21: sequence length=64
In our experiments, depicted in Figure 22, we trained our model with two distinct sequence lengths in batch: 64 and 256 steps. As the episode length increases along the x-axis, the agent's ability to retain and utilize the initial observation to produce the correct action the end is tested. The results indicate clear performance dichotomy: when episodes surpass 64 steps, models trained with batch sequences of this length fail to perform (achieving zero performance), while models trained with 256-step sequences maintain robust performance, even as episode lengths extend. This suggests that encompassing the entire RI. episode within the training batch is crucial for model efficacy.
Figure 22: Ablation study of in- put sequence lengths for models Memory Length.
As shown in Figure 23, the performance scales positively with the sequence length, reinforcing the finding above. Here, test the R2I model across four sequence lengths in batch: 64, 128, 256, and 1024 steps, observing a positive correlation between sequence length and performance. Throughout this range, performance improvements are evident with each incremental increase sequence length, underscoring the importance of longer sequences for successful memory task execution.
MemoryMaze-90
MemoryMaze-11xx
MemoryMaze-134x
MemoryMaze
Retur
20
20
20
10
0
0
1
2
3
1e7
1e7
1e7
1e7
Environment step
Environment step
Environment step
Environment step
R21: sequence length- 1024
R21: sequence length- 128
R21: sequence len
Figure 23: Ablation study of input sequence lengths for models on Memory Maze domain.
30

---

Published as a conference paper at ICLR 2024
P MODEL LIKELIHOOD RELATION WITH AGENT PERFORMANCE
In MBRL, the agent typically goes through a process of learning a world model and employs it for planning or optimizing the policy. Therefore, one might naturally assume higher likelihood of the world model leads to improved policy performance. Nevertheless, this assumption does hold true in all cases. When the world model has constrained representational capacity or its class is misspecified, its higher likelihood may not necessarily translate into better agent performance (Joseph et vall,20133 Lambert et al., 2020; Nikishin et al., 2021). 1). In other words, the inaccuracies in the world model will only minimally affect the policy if it happens to be perfect. Conversely, in the case of an imperfect model, these inaccuracies can lead to subtle yet significant impacts the agent's ss overall performance (Abbad, 1991). For example, Nikishin et al. (2021) propose method wherein the model likelihood is less than random baseline while the agent achieves higher returns. This implies that the learned world model may not always have to lead to predictions that closely align with the actual states, which essential for optimizing policies.
P.1 COMPARING R2I WITH S4WM
MemoryMaze-9k 9x9
MemoryMaze-11x1
MemoryMaze-13x13
MemoryMaze- 15x15
Episode return
10
20
20
5
0
0
3
2
3
1e7
1e7
1e7
1e7
10
10
10
10
Image loss
5
0
3
1
4
2
5
3
Env step
1e7
Env step
1e7
Env step
1e7
Env step
1e7
DreamerV3: giz) hi)
DreamerV3: giz;
Figure 24: Online RL plots of the DreamerV3 model with recurrent and non-recurrent representation models (for details see Appendix M).
In parallel with our study, Deng et al. (2023) introduce $4-based (Gu et al., 2021a) world model, called S4WM. The incorporation of S4 in the Dreamer world model has proven beneficial, as evidenced by the achievement of higher accuracy with lower Mean Squared Error (MSE). These results were obtained using offline datasets across various tabular environments and the Memory Maze (Pasukonis et al., 2022). However, their work primarily focuses world modeling: for instance, there no performance report in terms of obtained rewards. As previously discussed, a higher model likelihood does not necessarily guarantee better performance (i.e., higher rewards).
Figure 24 illustrates the results of DreamerV3 (which also serves as the predecessor to S4WM) with two distinct configurations for the world model on Memory Maze. One agent (depicted in green) incorporates non-recurrent representation model, while the other (depicted in black) uses recurrent one. Deng et al. (2023) utilize image loss shown, despite the agent with a non-recurrent representation model having a higher image loss, manages to achieve higher rewards. This result related the experimental results of S4WM in the following manner. In S4WM. the authors conduct experiments with the datasets collected by scripted policies within environments like Memory Maze such as Four-rooms, right-rooms, and Ten-rooms. Also, they evaluate image generation metrics on heldout dataset (compared to what the world model was trained on) to assess the in-context learning proficiency the world model across unseen levels. Similarly, the outcomes depicted in Figure 24 provide online RL metrics within memory maze environments. that is progressively updated. Therefore, successive batches may contain data from episodes within newly generated mazes that the model has not encountered before, reinforcing our assertion that these results indeed comparable.
of the key indicators for model likelihood. As
noteworthy that here, the agent is trained with a replay buffer
31

---

Published as a conference paper at ICLR 2024
Table 5: Architectural comparison between R2I and S4WM
Feature
R21
S4WM
Computational modeling
Parallel sean
Conv mode
Layer normalization
Postnorm
Prenorm
SSM in posterior and prior
Shared
Not shared?
Post-SSM transformation
GLU transformation
Linear transformation
SSM dimensionality
MIMO
SISO
SSM parameterization
Diagonal
DPLR
SSM discretization
Bilinear
? (Unspecified)
Handling resets technique
Reset to zero or learnable vectors
Not implemented
Exposure of hidden states
Provided by parallel scan
Not implemented
Policy training mode
DreamerV3's objective
None
Same as DreamerV3 but without free info in KI.
World model objective
Same DreamerV3
This highlights the idea that, in order to enhance long-term temporal reasoning in the context of MBRL and tackle memory-intensive tasks, simply maximizing world model likelihood is insufficient. That is why S4WM incorporates completely different set of hyperparameters and design decisions, as detailed in Table 5. Here is comprehensive explanation of each design decision that sets S4WM apart from R2I:
1. Computational modeling. R2I uses parallel scan (Blelloch, 1990; Smith et al., 2023) while S4WM uses global convolution (Gu et al., 2021 la) for training. One difference between these two techniques is that parallel scan uses SSM to compute u1.7, Zo - 91.T, Tır while convolution mode uses SSM to compute ,1:7, 3o > VI.T. Fr. Yet, discussed in Section 3.2 and empirically shown in Appendix N, providing the policy with the SSMs hidden states F1.T crucial for the agent's performance. This is not feasible with the convolution mode, as it does not yield these states.
2. Layer normalization. R2I applies post-normalization (i.e., after the SSM layers and right after the merge of the residual connections), whereas S4WM employs pre-normalization (i.e., before the SSM layer but subsequent to the residual branching). treenoomalization tends to be more conducive for training models with significant depth, whereas post- normalization can offer a modest enhancement in generalization when fewer number of layers are used (Wang et al., 2019). Given that our model must continuously generalize on data newly added to the replay buffer (as it performs online RL.), post-normalization emerges as the more intuitive selection for Online RL. This choice is confirmed by our preliminary results.
3. Shared SSM network. R2I employs a shared SSM network for prior prediction and posterior stochastic and deterministic models' inference, whereas S4WM investigates both shared and non-shared options, ultimately reporting the non-shared option as superior for image prediction. In this study, we have not explored the non-shared option; however, anticipate that the non-shared version might prove less efficient for policy learning due to potential feature divergence between the posterior and prior SSM networks. This is because the policy is trained using features from the prior model but deployed in the environment with features from the posterior model.
4. Post-SSM transformation. R2I uses a GLU transformation (Dauphin et al., 2017), while S4WM uses a linear fully connected transformation. In our preliminary experiments, we explored both options and found that the GLU transformation yields better empirical results, particularly in terms of the expected return.
'eere we refer to the WWWM-FulPPsterrian model. Even though it is not the main model, the S4WM- FullPosterior showed superior image prediction performance.
32

---

Published as a conference paper at ICLR 2024
5. SSM dimensionality. R2I employs MIMO (Smith et al., 2023), while S4WM utilizes SISO (Gu et al., 2021a). Although the superiority of one over the other is not distinctly evident, MIMO provides more fine-grained control over the number of parameters in the SSM.
6. SSM parameterization. R2I adopts diagonal parameterization (Gupta et al., 2022a) of the SSM, while S4WM employs a Diagonal-Phis Llaa--Pll-LL--aaaaa (DPLR) parameterization (Gu et al., 2021a). In our experiments, we found that both approaches perform comparably well in terms of agent performance. However, DPLR results in a significantly slower computation of imagination steps- approxiaatty to 3 times slower- therefore, we opted for the diagonal parameterization (See Appendix B.1.1).
7. SSM discretization. R2I utilizes bilinear discretization for its simplicity. S4WM does not specify the discretization method used; however, the standard choice for discretization within the S4 model typically relies on Woodbury's 's Identity (Gu et al., 2021a), which necessitates matrix inversion at every training step relatively costly operation. This is the reason decided early on in this work not to proceed with $4.
8. World model training objective. R2I uses the same objective introduced DreamerV3. In contrast, S4WM utilizes the ELBO objective in the formulation of DreamerV3, with the sole difference being that free information (Hafner et al., 2023) is not incorporated in the introduced in DreamerV3 as remedy to the policy
KL term. Notably, free information overfitting (it is evidenced by the ablations in Hafner et al. (2023)).
9. Episode reset handling technique. R2I leverages the parallel sean operator introduced by Lu et al. (2024) and modifies it to allow a learnable zero vector when RL episode is reset, a technique that was introduced in DreamerV3, While S4WM is also based on the DreamerV3 framework, it overlooks this modification and lacks the mechanism to integrate episode boundaries into the training batch. This adaptation is critical under the non-stationarity of the policy being trained; it primarily aimed at enhancing the efficiency of policy training; as the policy distribution changes, resulting in adjusted episode lengths, S4WM will be required modify its sequence length accordingly-eithee increasing decreasing it as necessary.
33

---

Published as a conference paper at ICLR 2024
Q UNVEILING THE PERFORMANCE IN STANDARD RL BENCHMARKS
To facilitate a more comprehensive comparison between R21 and DreamerV3, we build performance profiles using the RLiable package (Agarwal et al., 2021). In particular, we use these profiles to illustrate the relationship between the distribution of final performances in R2I and DreamerV3, as depicted in Figure 25. In Atari 100K (Lukasz Kaiser et al., 2020), the performance profiles of DreamerV3 and R2I show remarkable similarity, indicating that their performance levels are almost identical. Looking at the proprioceptive and visual benchmarks in DMC (Tassa et al., 2018), we observe that R2I exhibits drop in the fraction of runs with returns between 500 and 900. Outside of this interval, the difference between R2T and DreamerV3 diminishes significantly. Considering these observations, across the majority of tested environments, aligning with our assertion that R2I does not compromise performance despite its enhanced memory capabilities. Note that this is non-trivial property since these are significantly diverse environments, featuring continuous control (DMC), discrete control (Atari), stochastic dynamics (e.g., because of the sticky actions in Atari), deterministic dynamics (in DMC), sparse rewards (e.g.. there are several environments in DMC built specifically with explicit reward sparsification), and dense rewards. Our findings confirm that the performance in the majority of these cases is preserved.
conclude that R2I and DreamerV3 display comparable levels of performance
DreamerV3
DreamerVs
1.00
Fraction of runs with score
1.00
0.75
0.75
0.50
Fraction of nuns
0.50
0.25
0.25
0.00
0.00
250
500
750
1000
250
Return (+)
Return (x)
(a) Aggregated performance profiles for R21 and 33eamerV3 across 20 visual environments in the DMC benchmark, averaged over independent runs per environment.
(b) Aggregated performance profiles for R2I and Dreamer V3 across 18 proprioceptive environments in the DMC benchmark, averaged over indepen- dent runs per environment.
DreamerV3
A 1.00
Fraction of ruuns score
0.75
0.50
0.25
0.00
0.00
0.25
1.30 1.75 2.00 Human Normalized Score (T)
0.30
0.75
1.00
1:25
(c) Aggregated performance profiles for R2I and D3eamerV3 across 26 environments in the Atari 100K benchmark, averaged over independent runs per environment.
Figure 25: Comprehensive performance comparison profiles of DreamerV3 and R21 across three standard RL domains: Atari 100K, DMC-vision, and DMC-proprio. These visualizations highlight the comparative performance of the R2I and DreamerV3 V3 across diverse spectrum of tasks.
34

---

Published as a conference paper at ICLR 2024
R ALGORITHM PSEUDOCODE
Algorithm I: Recall to Imagine (R21), full state policy training
Initialize empty FIFO Replay Buffer D;
// is set to the total env. steps in batch to form at least
one batch to start training
Prefill with environment steps following the random policy; while not converged do for train step = 1..C do
World model training (parallel mode)
Draw n training trajectories
= ~ D, ==
// All variables are expected to incl ude batch dimension // which is omi mitted for simplicity
// See Section 3.1
Encode observations batch 21.7 ~ (=(=): 0rT);;
Compute S3M hidden states and outputs with parallel scan h1.T, Fi. & fo((112,,21 21:1),
// Note that I1.T include hidden states from all SSM layers Reconstruct rewards, observations, continue flags
fi.r, Ô1.r, ci.r I p(o, 2,, hp)(r 2,, h.)p(c; 44),;;;
Compute objective using Eq. (3), (4), (5), (6) and optimize WM parameters 8; // Actor-critis training (recurrent mode)
// See Section 2.2 and Appendix D for details
Initialize imagination 2oj1:7, Moo1: A0|1:T +- I1.T, h,,,, for imagination step = 0.. - do
a:ji.r ~ (ā Eq,.T, A.p.rr
// Note that this is a one-step inference of SSM in // recurrent mode
(141]::;;2+411::. + ([[(appp.,,)pp.))..
3441|117 ~ pa(2 h++1)1:);;
P141 ~
pppir-p(pr p(r & (asqprr);
441j1 ~ p(c 144p::h.411pp.
end
Estimate returns Ri 1:7 using Eq. and &Immee Aumpes
Update actor-critic using estimated returns according to the rule from Hafner et al. (2023);
end
Collect N steps in the environment using the World Model and the policy and store it to D; // the ratio C/N is training intensity which should be // chosen according to the desired data-eff fifecyy
end
Algorithm presents the pseudo-code for training the R2I in an online RL fashion. Note that the imagination phase is initialized from every possible step in the training batch. For example, hIgHI a tensor shape [bat ch.size, T, H, dim], which represents the imagination output after 1 to H imagination steps starting from every possible time point (in total, there are T of them). This concept was initially proposed by Dreamer (Hafner et al., 2019a) and subsequently adopted by DreamerV2 (Hafner et al., 2020) and DreamerV3 (Hafner al., 2023). We employ the same idea since the hidden states of the SSM are outputted by the parallel sean algorithm, allowing the use of all hidden states compute imagination steps.
35

---

Published as a conference paper at ICLR 2024
S MODEL-FREE SSM IN POPGYM
Figure 4 presents the performance of both model-free and model-based approaches the most memory-intensive environments in POPGym (Morad al., 2023). terms of model-based method- ologies, Dreamer is included, while model-free baselines comprise Proximal Policy Optimization (PPO; Schulman et al. (2017)) executed with variety of network architectures, such as MLP, GRU, LSTM, and S4D (Gu et al., 2022). S4D stands as an SSM approach characterized by a SISO frame- work (Gu al., 2021a), diagonal parameterization for SSM, and an efficient convolution mode for parallelization. The primary distinction between S4 (Gu et al., 2021a and S4D lies in their parameterization methods; S4D's diagonal approach, as opposed to S4's DPLR parameterization, which makes S4D faster and potentially more stable.
According to the findings reported the POPGym paper (Morad et al., 2023), the PPO+S4D configuration emerged as the least performant among the thirteen model-free baselines. It frequently encountered numerical difficulties such as not-a-numbers (NaNs), which can be attributed to gradient explosions. Please note that due to S4D's s reliance convolution mode, it is incapable of managing multiple episodes within sequence.
36

---

Published as a conference paper at ICLR 2024
T DETAILED RESULTS IN MEMORY MAZE
The system achieves a total through-
Human Normalized
Oracle Normalized
Wall time
A 100
Env
Score
Steps
put of approximately 350 frames per second (FPS), leveraging two
time
9x9
33.55
127%
96%
17M
0.6
1.2
NVIDIA A100 GPUs with 40GB of memory, along with 40 environment workers. The training intensity is 51 replayed steps per 1 sampled step or 82 environment steps per one gradient update of the world model and actor- critic (since the batch size is 4096). To optimize utilization, training is dis- tributed across the two A100 GPUs via batch-wise data parallelism. Additionally, we interleave training with environment sampling where the rollout worker lives on of the training GPUs.
1Ix11
51.96
117%
89%
66M
2.2
4.4
13x13
58.14
104.7%
78%
206M
6.9
13.8
15x15
40.86
60%
46%
Avg
46.13
102%
77%
Table 6: Algorithm statistics. Steps, Wall time, and A100
time are environment steps, wall time days, and A100 days until the superhuman performance.
MemoryMaze-90
MemoryMaze 11x11
MemoryMaze-13mm
MemoryMaze-15x13
60
40
60
Retum
40
20
10
20
20
10
10
0
100
15
10
15
days
days
days
days
R21 (Ours)
Dreamer (TBTT)
IMPALA
Human
Figure 26: Wall time durations for training agents over several days in Memory Maze. Throughout the training period, Dreamer consistently requires equal or greater wall time compared to R21, as its superior speed allows it to accumulate more environment steps within the same timeframe, as depicted in Fig 5.
MemoryMaze-9Kx
MemoryMaze 11x11
MemoryMaze-11
MemoryMaze-15x13
Retum
40
60
40
20
1.0
20
20
10
20
1.0
20
10
days
days
days
A100 days
R21 (Ours)
Dreamer (TBTT)
IMPALA
Human
Figure 27: Training duration on A100 GPUs over several days.
37


Reviews:
[{"review_id": "W47fWeUTFL", "reviewer": "Reviewer_Emy3", "summary": "This paper introduces a new Model-based RL algorithm Recall to Imagine (R2I) that upgrades DreamerV3 by non-trivially incorporating S4 networks in substitutions of the RSSMs that have been commonly used since PlaNet. The paper provides abundant evidence to support the main claim that R2I outperforms DreamerV3 and is significantly more computationally efficient", "questions": "Correctly addressing the two points above is what can change the most my opinion. Additionally, there are a couple of minor things I noticed:\n* Appendix N Figure 14, the second plot is missing the expected error\n* Section 3.1 line 4 where x_t is \"a\" hidden state and f_θ is a sequence model with \"a\"  SSM\nnetwork", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "While employing S4 as an alternative to RSSMs was a natural step to come, I think the authors do a very good work here. The paper is very well presented and easy to follow. Authors do a very detailed theoretical analysis of the approach, building up intuition and making it very easy to follow. It is remarkable that they even include in-depth discussions on why they discarded alternative features when designing their algorithm, which should be a more common practice. The empirical analysis is also extensive and provides solid ground with the numerous ablations and the varied array of benchmarks where the R2I vs Dreamer comparison is drawn. R2I proves to be significantly faster to train while being a better performer, specially in challenging long-term dependencies\n\nI believe that this paper will be very relevant for the ICLR community.", "weaknesses": "My biggest concern is the lack of an explicit literature review/ related work section, which I would suggest to include in the appendix. Specifically, I believe that it is of special relevance to include a more in-depth comparison between R2I and S4WM -currently briefly mentioned in the conclusions- since both combine DreamerV3 with SSMs.\n\nAlso, I noticed there are no mentions about making the code available. Thus, it would be specially benefitting for reproducibility if authors include a summary of the alg in pseudocode at the appendix.\n\n-- After Rebuttal --\n\nAuthors addressed very well all my concerns, the paper now presents clearly its differences with respect prior work and with a well documented code reproducibility should be easy to reproduce. I believe this will be a very relevant work for the RL community", "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": "xn2F3yqpf5", "reviewer": "Reviewer_pPqM", "summary": "This work proposes Recall to Imagine (R2I), a Model-Based Reinforcement Learning (MBRL) agent that integrates the Dreamer framework with State Space Models (SSMs), in order to alleviate the well-known challenges of long-term dependencies regarding memory and credit assignment. This integration works by replacing the GRU-based representation model with the SSM, which enables parallel predictions and improves the capabilities for capturing long-term dependencies. The paper presents a rich empirical analysis in a variety of memory environments, achieving state-of-the-art results in them while incurring a small drop in performance in the standard benchmarks. Furthermore, the paper provides an extensive ablation analysis of diverse design choices and hyperparameters in this integration.", "questions": "- In Appendix G (BSuite environment), is there any hypothesis on why sometimes harder environments (longer memory steps) present better performance than easier ones? For instance, R2I’s performance on 31 memory steps is better than 15 memory steps. Similarly, the performance in 81 memory steps seems better (or more stable) than 41 memory steps.\n\n\n\n\n\n\n\n\n\n\n===================== **POST-REBUTTAL** ==================================\n\nDear authors,\n\nThanks for putting so much effort on addressing my concerns. I believe they led to substantial improvements in an already good work, so I am raising my scores towards acceptance.\n\nSpecifically:\n\n- I appreciate the efforts on formulating hypotheses on why the different input variations work differently across environments. I believe the raised hypotheses are valid and potential venues for future work. As I mentioned before, this seems to be a difficult problem that requires careful investigation. It would be interesting to bring this rebuttal discussion into the paper, as I believe this could also be an important question for other readers. This is also valid to the point related to Appendix G.\n\n- I also would like to thank you for adding the new PPO + SSMs baseline. Turns out to be very different from what I expected, and given the discussion on Appendix T, I think it should require careful tuning to work with PPO (which is, indeed, a very sensible algorithm). But this is an argument in favor of the proposed method and perhaps another open question to be addressed (i.e., how to make SSM to work with PPO)\n\n- Thanks for addressing the tone of the claim in Section 4.3 and providing more evidence regarding it. I agree that given the diversity of environments, it is hard to ensure that all of them will attain the same performance, and the new wording also better reflects the presented results.\n\n- Lastly, the Appendix Q is great. I think this was one common concern from many reviewers and it was well addressed with many details. For sure the strongest reason to increase the scores.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper addresses the problem of long-term dependencies for World Models, which is an issue for RNNs and Transformers on handling sequences in representational models. Therefore, it is very relevant to the community. \n\n    - Furthermore, employing SSMs to replace the aforementioned backbones for learning temporal dependencies is sound and well-motivated.\n\n- The proposed architecture establishes new state-of-the-art performance for several tasks in the considered environments (BSuite, POPGym, and Memory Maze), with a noticeable improvement in computational efficiency, as shown in Figure 2.\n\n- The work brings extensive and insightful ablation studies in many design decisions in the R2I architecture. These are presented as a systematic evaluation in Appendices M to P and helps understanding how the proposed method works.", "weaknesses": "- Despite the conducted ablation in Appendix O, the question on why the different input variations work differently across environments still remains open. And the raised hypothesis on “feature instability” sounds vague and not properly backed up by a solid argument. It would be great to provide a better understanding of this challenge to give more clarity on how the method works, but I understand this is a difficult open problem that demands a careful investigation.\n\n- The work adopts SSMs to address the challenge of handling long-term dependencies in RL (memory and credit assignment). Nevertheless, it does not motivate the employment of Dreamer (or, more generally, Model-Based RL). I think it is important to describe and motivate why Dreamer was used instead of Model-Free RL, as it is not clear why MBRL would be better than MFRL for these memory tasks (unless there is another motivation besides asymptotic performance, such as sample efficiency).\n\n    - In the same line, Figure 4 brings some “memory-augmented” Model-Free baselines, but it lacks “PPO + SSMs”. This baseline would definitely clarify my concern. If there is no constraint in the sample budget, it is possible (perhaps expected) that the MFRL agent would perform better.\n\n- In Section 4.3, the claim of “not sacrificing generality” is questionable. There is a small drop in performance. For instance, in Appendix J (DMC-proprio), DreamerV3 is (at least) slightly better in 6 tasks.  In Appendix K (DMC-Vision), 8 environments. In Appendix L (Atari), 10 environments. I suggest rephrasing the claim to account what is observed in the Appendices.\n\n- The work from Deng et.al [1], proposing S4WM, looks very similar to the proposed one. Indeed, both works propose replacing the RSSM with SSMs with the same motivation: improving memory capabilities. The work on S4WM was publicly released approximately 2.5 months before this submission, which can be seen as concurrent work. Nevertheless, I believe the work is almost overlooked by the proposed paper, which has a small citation in Section 5. Given the similarity, it would be crucial to provide a more detailed comparison contrasting both works, in terms of methodology and evaluation, perhaps in the Introduction or in a separate Appendix.\n\n**Minor Concerns**\n\n- In Appendix H, task Autoencode: the episode length for the Hard task is 156. Is that right? I believe it is supposed to be 256.\n\n**References**\n\n[1] Deng et. al. Facing off World Model Backbones: RNNs, Transformers, and S4. NeurIPS, 2023.\n\n**Summary of the Review**\n\nThe work brings an effective improvement on World Models due to a well-motivated employment of State Space Models. This validates this architecture and extends its effectiveness in dealing with sequences in RL. The work does a great job in empirical analysis, anticipating many questions and answering them with extensive experiments and ablations. On the other hand, I believe the paper could be improved if the aforementioned concerns were addressed. Nevertheless, these concerns are not critical enough to prevent acceptance.", "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": "ajND9EiCWR", "reviewer": "Reviewer_fFhp", "summary": "This paper introduces state space models (SSM), in particular S4, into world models in the framework of model-based RL to improve its long-term memory and long-horizon credit assignment, as well as computational efficiency. Specifically, RSSM in Dreamer is replaced with SSM (S4), resulting in the proposed R2I agents. Design decisions to do so are carefully chosen, and empirical studies demonstrate improved performance in memory-demanding domains, including POPGym, bsuite, and Memory Maze.", "questions": "The authors should properly resolve my concerns mentioned in the weakness part.\n\nThere are also some minor questions:\n\n1. Dreamer is compared in Memory Maze tasks. Does this Dreamer baseline include the TBTT technique proposed by Pasukonis et al., which improves the memory of RSSM?\n2. Why not include Transformer-based world models as baselines? Transformers are also widely believed to well model long-horizon dependencies.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. Improved MBRL performance with S4-based world models has been validated in memory-demanding domains.\n2. Careful designs of S4-based world models, including non-recurrent representation model and SSM computational modeling.\n3. Extensive experiments in a number of domains.\n4. Well written with sufficient experimental details.", "weaknesses": "1. Limited contribution. It is notable that there already exists a S4-based world model, namely S4WM [1]. Despite minor design choices, the major difference of this paper is that it conducts MBRL experiments while S4WM only conducts world model learning (e.g. imagination and reward prediction). However, in my humble opinion, it is not surprising that improvements in long-term memory can lead to improved MBRL performance in memory-demanding domains.\n2. Three kinds of actor and critic inputs are introduced, namely, output state, hidden state and full state, which results in a critical design choice to be tuned for each domain. Although the authors provide some takeaways to select between them, it is not always true. For instance, output state policy is utilized in memory-demanding environments, Bsuite, while hidden state policy is used in non-memory environments, DMC.\n3. The authors claim that R2I does not sacrifice generality for improved memory capabilities. However, there is a clear trend in Figure 6, that R2I performs worse than Dreamer in standard RL tasks. \n4. Some inaccurate statements. For example, the authors say R2I's 'objective differs from ELBO in three ways', but to my knowledge, these three points are all borrowed from DreamerV3 but without explicitly being mentioned in the text.\n\n[1] Deng et al. Facing off world model backbones: Rnns, transformers, and s4.", "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 (oral)