Parallelizing Model-based Reinforcement Learning Over the Sequence Length
Authors: ZiRui Wang
Yue DENG
Junfeng Long
Yin Zhang
Venue: NeurIPS 2024 poster (2024)
Keywords: Model-based reinforcement learning
world model
parallelization
Primary area: reinforcement_learning
TL;DR: This paper introduces the PaMoRL framework, which parallelizes MBRL over the sequence length, improving training speed and sample efficiency.

Abstract: Recently, Model-based Reinforcement Learning (MBRL) methods have demonstrated stunning sample efficiency in various RL domains.
However, achieving this extraordinary sample efficiency comes with additional training costs in terms of computations, memory, and training time.
To address these challenges, we propose the **Pa**rallelized **Mo**del-based **R**einforcement **L**earning (**PaMoRL**) framework.
PaMoRL introduces two novel techniques: the **P**arallel **W**orld **M**odel (**PWM**) and the **P**arallelized **E**ligibility **T**race **E**stimation (**PETE**) to parallelize both model learning and policy learning stages of current MBRL methods over the sequence length.
Our PaMoRL framework is hardware-efficient and stable, and it can be applied to various tasks with discrete or continuous action spaces using a single set of hyperparameters.
The empirical results demonstrate that the PWM and PETE within PaMoRL significantly increase training speed without sacrificing inference efficiency.
In terms of sample efficiency, PaMoRL maintains an MBRL-level sample efficiency that outperforms other no-look-ahead MBRL methods and model-free RL methods, and it even exceeds the performance of planning-based MBRL methods and methods with larger networks in certain tasks.

Parallelizing Model-based Reinforcement Learning Over the Sequence Length
ZiRui Wang Zhejiang University, China
Yue Deng Zhejiang University, China devindeng@zju. edu. en
Junfeng Long Shanghai Al Laboratory. China junfengac@@mail.com
zi ziseoiwongozju.o edu. en
Yin Zhang'
Zhejiang University, China zhangyin98ezju edu. cn
Abstract
Recently, Model-based Reinforcement Learning (MBRL) methods have demon- strated stumning sample efficiency in various RL domains. However, achieving this extraordinary sample efficiency comes with additional training costs in terms computations, memory, and training time. To address these challenges, the Parallelized Model-based Reinforcement Learning (PaMoRLL framework. PaMoRL introduces two novel techniques: the Parallel World Model (PWM) and the Parallelized Eligibility Trace Estimation (PETE) to parallelize both model learning and policy learning stages of current MBRL methods over the sequence length. Our PaMoRL framework is hardware-efficient and stable, and it can be applied to various tasks with discrete or continuous action spaces using single set of hyperparameters. The empirical results demonstrate that the PWM and PETE within PaMoRL significantly increase training speed without sacrificing inference efficiency. In terms of sample efficiency, PaMoRL maintains an MBRL- level sample efficiency that outperforms other no-look- ahead MBRL methods and model-free RL. methods, and it even exceeds the performance of planning-basee MBRL methods and methods with larger networks certain tasks.
propose
Introduction
Model-based Reinforcement Learning (MBRL) is widely believed to have the great potential to substantially enhance sample efficiency by training policy through learned world model 3 Previous studies [ achieve the same asymptotic performance as their model-free counterparts while requiring orders of magnitude less interactions. In particular, some recent works have even and robot control [TO]
achieved human-level efficiency in complex RL domains like Atari 7 S 9
MBRL methods can be generally divided into two stages: model learning and policy learning. During the model Icarning stage, a parameterized world model required to predict the environmental dynamics by constructing specifie self-supervised learning tasks. The policy learning stage benefits from synthetic interactions between the policy and the world model, hence on- on- on-policy actor-critic methods planning methods such Model Predictive Path Integral (MPPI) [2 [3] or Monte-Carlo Tree Search (MCTS) [7] [
can be used for policy improvement. To obtain better performance, techniques like sequential modeling and ensembling are frequently used in the model learning stage, while the policy learning stage mostly involves the computation of eligibility traces or multi-step However, these powerful techniques often come with additional computations,
returns 2 14
Cooresponding author: Yin Zhang.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Atari 100K Benchmark
DeepMind Control Suite (Proprio)
DeepMind Control Suite (Visual)
Average Human Normakzedd Score (%)
140
PaMaRi (Duns) RRSRR
620
PaMoRL (Oure)
575
120
STORM reaaerVV
500
PakloL (Ours)
JRIS
Average Score
525
100
JWM
500
Anorage
80-
SPR
800
475
P PreamerV3
60-
540
SAC
450
40
SmPla
PreamerV3
423
Model based Modeel-/se free
15
25
100
150
200
20
30
50
Troinng FPS V100
Traning RS V100 GPU
Training FPS V100 GPU
Figure 1: Comparisons on Atari 100k benchmark [16] and DeepMind Control Suite 24 Among these methods, D3eamerV3 [7]. and our PaMoRI. are directly evaluated on an NVIDIA V100 GPU, and IRIS (0). TWM (1]. and REM [25) evaluated an A100 GPU, while other methods are evaluated on a P100 GPU. The extrapolation method employed aligns with the setup used in DreamerV3, where it assumes the P100 is twice as slow and the A100 twice as fast.
memories, and training time. This leads users to carefully consider which specific MBRL method to use or even whether to use an MBRL method based on the computational resources available.
In recent years, numerous endeavors have been made to develop an efficient world model architecture. Recurrent Neural Networks (RNNs) are frequently employed the foundational architecture for world models [15 16
3 6 17]. However, the recurrent nature of RNNs hinders parallelization, leading to slow training speeds. In contrast, transformers have emerged as potential successor, garnering acclaim for their remarkable performance in language modeling tasks and parallelized training paradigm [88. Several attempts have been made to incorporate transformers into world models [19 20
However, the quadratic complexity of transformers w.r.t. sequence length limits their efficiency during training and inference. To achieve an RNN-level inference efficiency, extra tricks such as half-precision training or KV-Cache are required (23 Furthermore, none of the aforementioned works have introduced improvements in the hardware efficiency of policy learning.
21 22
In this paper, we aim to mitigate the curse of computational inefficiency of current MBRL methods and achieve the best of both worlds in terms of hardware efficiency and sample efficiency. The key idea is to fully parallelize the computations of sequential data, which has been a main workhorse of the rapid progress deep learning over the past decade [26]. We achieve this by introducing the parallel scan. Specifically, We delve into two classic and widely implemented parallel scanners 27 28)
which can be applied for parallel training by excluding non-linear dependencies 22 30 Motivated
by recent works in efficient sequential modeling 33 33,, we observe that model architectures like linear attentions and linear RNNs not only enable parallel training but also recurrent inference. We
also observe that the computations of cligibility trace estimation over the sequence length by using parallel scan.
can be naturally parallelized
To this end, we introduce the Parallelized Model-based Reinforcement Learning (PaMoRL) frame- work, which consists of two novel techniques as shown in Figure 2 that can parallelize the current MBRL paradigm over sequence length: (1) the Parallel World Model (PWM) and (2) the Parallelized Eligibility Trace Estimation (PETE). The resulting framework, PaMoRL, hardware-effictt stable. It is compatible with various on-policy RL methods and can be applied to both discrete and continuous control problems using a single set of hyperparameters.
and
We evaluated our PaMoRL framework in the Atari 100K benchmark [6] and the DeepMind Con- trol suite (24). Tasks in these domains include discrete and continuous action spaces, images, and proprioception observations. We choose to follow the DreamerV3 [17] paradigm, which relies on "imagination" for policy learning. The summarized experimental results are shown in Figure [[] The empirical results demonstrate that PaMoRL, despite being a framework that incorporates au- toencoding, still benefits greatly from the implementation of dual parallelization techniques (i.e., PWM and PETE). These techniques substantially enhance training speed, allowing PaMoRL to rival the performance of model-free RL methods without decoders (34). In terms of sample efficiency, PaMoRI. outperforms other look-ahead MBRL methods and model-free RL methods. It is worth mentioning that PaMoRI. even outperforms the planning-basn MBRL methods or methods with much larger networks in certain tasks (8).
2

---

Our contributions can be summarized as follows:
We introduce PaMoRI., novel MBRL framework equipped with PWM and PETE that parallelizes both model and policy learning stages over the sequence length simultaneously.
We evaluate our PaMoRL on the Atari 100k benchmark and DMControl suite with recent methods and obtain excellent results in terms of both sample and hardware efficiency. In addition, we conduct ablation studies on the validity of different modules, scanners, and other components.
To the best of our knowledge, we are the first to point out that the computational process of eligibility traces can be parallelized over the sequence length. This technique can not only accelerate the value estimation process of various MBRL methods but any return- based reinforcement learning methods such as TD-1 [2], Retrace K5
and GAE (36] can benefit from it.
2 Background
Model-based Reinforcement Learning. We follow the paradigm of Partially Observable Markov Decision Process (POMDP) with observations o,, scalar rewards rt, actions ax continuation flag p(or, rt, colocet (-cr)- The objective of the Reinforcement Learning (RL) is to train a policy и that maximizes the return , In Model- Reinforcement Learning (MBRL), the RI agent learns a model of the environmental dynamics through an iterative process that involves collecting data using policy, training a model of the environment based on the accumulated data, and optimizing the policy using the learned model 2 14.
a E (0, 1), discount factor (0, 1), and environmental dynamics Ot, Tt,
As universal parallel algorithm building block, the computations of parallel scan involve repeated application of a binary operator b over sequential data arrays. Previous work describes
Parallel Scan.
good example of computation that seems inherently sequential, but for which there is an efficient parallel algorithm. The scan of (D with initial value ap is defined in Equation T
SCAN c [01, ag..... on], an) = [(1] @ ao), (a2 i a1 + ao), ..., (an 2 an-1... + a1 2 ao)]
(1)
First-order linear recurrences h, (A,, @ h4-1) F, be parallelized over the sequence length with the utilization of parallel scans if the following three conditions are met:
i is associative: (a ① b) @
+ (b ) c).
is remi-associative: there exists binary associative operator such that a (boc) (aOb) Sc.
distributes over @: a @ (b + c) = (a 2 b) @ (a 8 c).
We observe vector addition a + b := a + b. matrix-vector multiplication A & b = A b, and matrix-matrix multiplication 3 B == A B fulfill the aforementioned conditions. This allows the parallel computation of z; == (A, 21-1) + by across time steps t, considering input vectors b, and square matrices A,. Considering the operators required in computing linear attentions [31 and eligibility trace estimations 2 [4] involve only diagonal matrices, the linear recurrence can be foromullated as I; = X @ xt-1 + b, where X₃ the eigenvalues of the diagonal matrices and O an element-win wise multiplication.
3 Methodology
We introduce Parallelized Model-based Reinforcement Learning (PaMoRL.) framework, which facilitates dual parallelization across both model and policy learning stages. By parallelized training and recurrent inference, PaMoRL significantly improves training speed while avoiding additional computation overhead during inference. Figure illlstraaees the overview of our PaMoRL framework, and we will now proceed to elaborate on its details.
3.1 Parallelized World Model Learning.
World Model Architecture Overview. As with other MBRI. methods, our world model trained to predict environmental dynamics. Since observations can be high-dimensional (c.g., images),
3

---

World Model Learning
Imagination (Recurrent)
Eligibility Trace Estimation
(Parallel)
(Parallel)
4
Parallel Scan
Brej
Ent
Parallel Scan
0y
Deterministic Nodes
Stochastic Nodes
Concatenation
Figure 2: Overview of our PaMoRL framework. The symbols used in the figure are explained in Sections 3.J and Section 3.2 The computations of the sequential model's outputs and the TD-A returns allow using parallel scans. In contrast, the imaginations cannot be parallelized over the sequence length because non-linear actor network is required for action sampling.
predict future representations rather than future observations. This reduces accumulating errors and enables massively parallel training with large batch size. The compact representations obtained by an autoencoder and can be utilized to predict future observations, reward, and continuation flags.
To exclude non-linear dependencies for parallel training and obtain better performance, make several modifications to the vanilla Recurrent State-Space Model's (RSSM) 3 17) configurations: (1) differentiating the hidden states from the sequential model's outputs hr, (2) excluding hu from the inputs of the encoder and decoder, (3) eliminating the stochastic states z, from the predictors' inputs, and (4) applying Batch Normalization for the encoder and dynamic predictor's outputs before the distributions are computed. Similar to RSSM, our model consists of six components:
Encoder: Zt ~ z0~go (0(2)))))
Decoder: or (ōel
Sequence model: he, th = (a 21 a+-1)
Dynamics predictor: ân pa
(2)
Reward predictor: fr ~ pe (fellt))
Continue predictor: & ~ pe(Cele))
The encoder and decoder convolutional neural networks (CNN) for image inputs and multi-layer perceptrons (MLPs) for proprioception inputs. The sequence model has multiple stacked residual blocks, each of which consists modified linear attention (31] module and a Gated Linear Unit (GLU) [38] module. The dynamics, reward, and continue predictors are all MLPs. Consistent with previous work 6 [7].
the qo (= (o:) stochastic distribution comprising 32 categories, each with 32 classes, and we take trriigh-tthoough gradients through the sampling step [39].
Sequence Model Architectures As mentioned above, each residual block of our sequence model consists a modified linear attention module and a GL.U module. The vanilla linear attention module, as introduced in previous work [3I], employs 1 + ELU as an element-wise kernel function applied to queries and keys ky taking as input. This configuration allows for its reformulation into an RNN-style recurrent form. However, this version of linear attention is prone unstable convergence during training due to the unbounded gradients [40] Thus, we remove the dependent normalizer, which designed approximate the Softmax operator and use an RMSNorm 41 for stabilize training. Furthermore, we incorporate the token mixing module from RWKV [32,, which accepts inputs W, and previous inputs 1, along with the gating mechanism in Gated Recurrent Unit (GRU) 42 to provide an input-dependent decay rate z: for hidden state The subsequent GLU module selects SiLU as the gating function, taking linear attention output y/ input. By integrating all the modifications, we can derive the entire block of the sequence model as shown in Equation

---

4t k + (u, AL)(
Ut Sigmoid( ue W+) ue W..
9: = Sigmoid (f @ + (1 - 1) w-))
I₁ 9. 2c-1 K2 U/2
y: RM +
h; SiLU W, C yr Wy yt.
The architecture of our modified linear attention satisfies the conditions in Section 2 and can be effectively computed using parallel scans. We can refer Table summarize the computational complexities of various model architectures such as vanilla attention, RNN, SSM, and our modified linear attention in the training, inference, and imagination stages.
Loss Functions. The total loss function of model learning is shown as in Equation 4 where 3prd- Grep, and Bäyn are coefficients to adjust the influence each term in the loss function 43
((0)
=ggo
pred Crred (0, he, oc, Pe, Ces 2() + Step he, 01) + Boyn he, 0z)l
Corod (0) In po (rall))) (mpo(c(/4))) + | - 01112
(0) max(1, (L sg(po he))))
(O) (L || po (2, h+)))
The operation sg(-) represents the stop gradient operation. The KL divergences derived from the Evidence Lower Bound (ELBO). We clip the KI. divergence when falls below the threshold of 1 17] and use the KL-balancing trick to prioritize the training losses (17]
3.2 Policy Learning
The policy learning stage incorporates the actor and critic networks, both of which are MLPs, taking concatenation of 2₁ and h, input state 8t-
Actor: ae aa(a st), Critic: (s)))..
Our policy learning method is in line with DreamerV3 [17 and can be used for both discrete and continuous action spaces. The critic uses TD-A (2 as the its target, as shown in Equation 6 where Å represents the reward predicted by the world model, and é represents the predicted continuation flag.
R = F. + (x0)) [(1 - A)m + AR
(6)
(172)) R4 + [+: + (1 ()(7)))
R
The actor utilizes the Reinforce estimator to compute the actor loss with a fixed entropy regularization term. The complete loss is described by Equation
T
R
(8p( )logT)(())ss) - (H((T((alsss)
L(Ø)
sg( Σ
max 8)
E
((V) D (ow(s)) - sg( Rg(
The hyper-parameter 1 represents the coefficient of the entropy regularization term. The normalization ratio S utilized the actor loss is defined Equation & which is computed as the range between the 95th and 5th percentiles of the TD-À returns R² across the batch.
5

---

Table 1: The step complexities [28] of different architectures, where L is the sequence length and H the imagination horizon. Attention considers the full context with burn-in and imagined steps of OI + H). leading to complexity of ([ + H. It is worth noting that the SSMs in recent works 146 47 do incorporate any gating mechanism or selectivities. Thus, despite SSMs and
linear attentions both achieving the minimum complexity, linear attentions remain more expressive.
Architecture
Training
Inference step
Imagination step
Parallel
Resettable
Sclective
Atten RNN
O
O(12)
O(L + m))
(())
O(1)
(1)
X
SSM (FFT)
(Lgg)
O(1)
0(1)
SSM (Scan)
(IL)
0(1)
0(1)
X
Lin-Atten (Scan)
O(L)
O(1)
O(1)
/
V
S = percentile 95) percentile (R) 5)
(8)
By rearranging Equation we can see that the calculations of both TD-À and Retrace returns also also meet the conditions mentioned in Section Therefore, they can be efficiently computed using parallel scan. This observation also applies to other eligibility trace estimation methods such as GAB (36] and Retrace (35), they still satisfy the aforementioned conditions.
3.3 Parallel Scan Algorithms
In both the model learning stage and the policy learning stage, we use two different parallel scanners: the Kogge-stone scanner [55 and the Odd-even scanner (28]
The Kogge-stone scanner is commonly used in hardware design for adders. It has a computational complexity of (LLogg, L) for sequence length L and step complexity of (log, after full parallelization. This indicates that it has higher computational redundancy, lower running time, and sufficient computational resources, making it suitable for parallel computation in a small batch.
The Odd-even scanner 28 based on the concept binary balanced trees. It has computational complexity of ((2L) for a sequence length L and step complexity of O(2log, after being fully parallelized. Despite theoretically taking more steps than the Kogge-stone scanner, it offers lower computational complexity and more uniform load sharing, making it better suited for large-scale parallel computation. Further details and illustrations are in Appendix B
Experiments
In this section, we aim to evaluate both the sample and training efficiency of PaMoRL framework on the Atari 100K benchmark [16] and the DMControl suite (24]. The tasks include various scenarios with image and proprioception observations and discrete and continuous action spaces.
4.1 Experimental Setup
Atari 100K. Atari 100K consists of 26 video games with discrete action dimensions of up to 18. The 100K samples are equated 400K actual game frames, corresponding to approximately hours of real-time gameplay, with action repeats of 4. The human normalized score is defined as (scorengent - scorerandom) /aaaldmmmm/SSSSuemmammma (Scorehuman scorerandom) ), where scorezandom comes from random policy,
and scorehuman is obtained from human players 48
Deed Control Suite. DeepMind Control Suite consists of various control tasks with continu- action spaces. Referring to the categorizations in Sample MuZero [49] and EfficientZero V2 (9) tasks are divided into easy and hard categories. We followed the experimental setup EfficientZero V2 [ and established two benchmarks, named Proprio Control and Visual Control.
Among them, Proprio Control uses proprioception observations with 50K training samples for easy tasks and 100K for hard tasks, and Visual Control uses image observations with 100K training

---

Table 2: Experimental results the 26 games of Atari 100k after hours of real-time experience and human-normalinn normalized aggregate metrics. Bold and underlined numbers indicate the highest and the second second-highest highest scores, respectively. PaMoRL outperforms other methods regarding the number of superhuman games, mean, and median.
Game
Random
Human
SR-SPR
SimPf
IRIS
STORM
DreamerV3
PaMoRI. (Ours)
Alien
227.8
7127.7
801.5
1015.5
616.9
420
674.6
984
959
1270.6
Amidar
5.8
1719.5
176.3
203.1
88
143
121.8
205
139
264,4
Assault
222.4
742
571
1069.5
527.2
1524.4
682.6
801
706
881.8
853.6
2957.3
BankHeist
14.2
753.1
380.9
472.3
34.2
53.1
466.7
641
649
255 9
BattleZone
2360
37187.5
16651
19398.4
5184.4
13074
5068
13540
12250
23120
Boxing
0.1
12.1
35.8
46.7
9.1
70.1
77.5
80
78
87.9
Breakout
1.7
30.5
17.1
28.8
16.4
83.7
20
T6
31
15.8
ChopperCommand
7387.8
974.8
2201
1246.4
1565
1888
420
2110.7
CrazyClimber
10780.5
35829 4
42923.6
43122.3
62583.6
59324.2
71820.4
667766
97190
84102
DemonAutcck
208.2
Freeway
0
29.6
24.4
24,9
20.3
31,1
24.3
Q
0
33.8
Frosthite
65.2
4334.7
18215
8
254.7
259.1
1475.6
1316
909
3711.4
Gopher
257.6
2412.5
715.2
711.2
771
2236.1
1674.8
8240
3730
5085,2
Hero
1027
30826.4
7019.2
7679.6
2656.6
7037.4
7254
1044
11161
12076.2
Jamebond
302.8
365,4
392.8
125.3
462.7
362.4
445
405
Kangaroo
52
3035
3276.4
3254.9
323.1
838.2
1240
4208
4098
2554.7
Krull
5824.8
7782
KungFuMaster
258.5
22736.3
13192.7
17095.6
17257.2
21759.8
24554.6
26182
21420
24624.7
MsPacman
307.3
6951.6
1313.2
1522.6
(480
1588.4
2673
1327
2201.7
Pong
-20.7
14.6
-5.9
-3
12.8
14.6
18.8
11
18
15.5
Privateliye
24.9
69571.3
124
95.8
58.3
100
86.6
7781
882
4968.6
Obert
163.9
13455
669.1
3850.6
1288.8
745.7
3330.8
4522
3405
4730.3
Roadrunner
11.5
7845
14220.5
13623,5
$640.6
9614.6
9109
17564
15565
24726.7
Seaquest
42054.7
774.4
UpNDwon
533.4
11693.2
28138.5
95501.1
3350.3
3546.2
13981.7
7985
7667
11935.8
>iluman
Median
0%
100%
41.53%
56.07%
14%
29%
51%
42.63%
49%
71.75%
Mean
100%
70.34%
118.8 18888%
44%
105%
96%
122.30%
112%
126.64%
samples for easy tasks and 200K for hard tasks. Each benchmark includes 16 tasks. Action repeats are set to 2, and the maximum episode length is 1000 for both benchmarks, in line with previous studies 17 13 9. We choose various baselines for each domain, which include SAC [50) DrQ- and DreamerV3 17
v2 5
4.2 Experimental Results
In this section, we do not compare our results with look-ahead search methods [52 methods using larger networks (8), main goal in terms of sample efficiency is to improve performance while maximizing the hardware efficiency of existing MBRL methods.
7
12 13 or
Atari 100K. The summarized results are shown in Figure The full results for individual games in the Atari 100k benchmark are elaborated in Table where scores normalized against those of human players. Our PaMoRL framework attains a mean score of 126.64% and a median score of 71.75 %, surpassing the other methods in terms of both mean and median human normalized score. For detailed training curves, please refer to Appendix C Additionally, you can find more results and further discussions, including methods with look-ahead search or larger networks, in Appendix []]
DeepMind Control Suite. Table 3 shows that our method achieves mean score of 661.2 across 16 tasks. As shown in Table 3, our method achieves mean score of 661.2 using proprioception observations and 538.7 using image observations across 16 tasks, surpassing the previous state-of- the-art, DreamerV3, The improvement in sample efficiency is attributed to two key modules: the token mixing module in the PWM, where the extra previous input provides more information to the data-dependent decay rate, and the implementation of RMSNorm, which improves the stability of the learning of the linear attention module, especially in the case of limited data. Our PaMoRL framework consistently demonstrates MBRL-level sample efficiency in tasks with proprioception observations, image observations, and discrete and continuous action spaces. Detailed training curves can be found in Figure 9 and Figure 10 in Appendix D
7

---

Table 3: Experimental results the DeepMind Control suite. Bold and underlined numbers indicate the highest and the second-highest scores, respectively. PaMoRL outperforms other baselines in terms of the number of mean and median scores.
Proprio Control
Vision Control
Task
SAC
DreamerV3
PaMoRL (Ours)
CURL
DrQ-v2
DreamerV3
PaMoRL (Ours)
Cartpole Balance
997.6
839.6
994.7
963.3
965.5
956.4
610.3
Caripolc Balance Sparse
993.1
559
997.4
999.4
1000
813
996.5
Caripole Swingup
861.6
527.7
773.6
765.4
756
374.8
281.9
Cup Catch
949.9
729.6
957.9
932.3
468
947.7
966.3
Finger Spin
900
765.8
835.8
850.2
459.4
633.2
765.3
Pendulum Swingup
158.9
830.4
707.1
144.1
233.3
619.3
26.6
Reacher Easy
744
693.4
761.6
467.9
722.1
441.4
950.2
Reacher Hard
646.5
768
645.9
112.7
202.9
120.4
103.7
Caripole Swingup Spanse
256.6
172.7
542.3
8.8
81.2
392.4
263.6
Cheetah Run
680.9
400.8
313.2
405.
418.4
587.3
935.6
Finger Easy
630.8
560.5
617.1
371.5
286.8
366.6
886.2
Finger Tur Hard
414
474.2
389.7
236.3
268.4
258.5
500.1
Hopper Hop
0.1
9.7
387.5
84.5
26.3
76.3
426.9
Hopper Stand
3.8
296,1
151.5
627.7
290.2
652.5
189.7
Quadruped Run
139.7
289
246.7
170,9
339.4
168
344.8
Quadruped Walk
237.5
256.2
457.9
131.8
311.6
122.6
371.6
Mean
538.4
$10.8
611.2
454,5
426.8
470.7
538.7
Median
638.7
543.4
631.5
388.3
325.5
416.9
463.5
Amidar
UpNDown
Obert
00000
200
00000
Spore
150
0000
Score
2000
4000
1000
0
100
200
300
400
100
200
300
400
100
200
300
400
Alien
Boxing
MsPacman
1250
Scord
750
1500
1000
250
000
100
200
300
400
100
200
300
400
100
200
300
400
PWM
PWM w/o RMSNorm
PWM wio Token Mixing
SSM w/ Scan
DroamorV3
Figure 3: Ablation studies of the effectiveness of each module of PWM, where SSM is equivalent removing the data-dependent decay rate of PWM. We also include vanilla DreamerV3 as a baseline.
4.3
Ablation Study
In this section, we will be conducting ablation studies to evaluate the effectiveness of PWM and PETE in terms of stabilizing training and improving hardware efficiency. For more details, including PyTorch-style pseudo-code, please refer to Appendix G
World Model Design. The results presented in Figure demonstrate the impact of adding or removing the token mixing, RMSNorm, and data-dependent decay rate in various games in the Atari 100K benchmark. To showcase the benefits of token mixing in sequence prediction, we focused on tasks such as Alien, Boxing, and MsPacman. Additionally, we measured the improvement of RMSNorm on training stability by considering tasks like Amidar, UpNDown, and Obert.

---

P(PaMoRL Y)
Mean
Median (1)
JQM (1)
Optimality Gap (1)
REM
PaMoRL
Algorithm Y
IRIS
IRIS
TWM
TWM DreamerV3
DreamerVS
SimPLe
SimPLe II
SPR
SPR
0.3 0.6 0.9 1.2
0.2 0.4 0.6
0.2 0.4 0.6
0.45 0.60 0.75
0.00
0.25
0.50
0.75
1.00
Human Normalized Score
Figure 4: (Left) Atari 100K aggregated metrics with 95% stratified bootstrap confidence intervals of the mean, median, and interquartile mean (IQM) human-normalized scores and optimality gap. (Right) Probabilities of improvement, i.e. how likely it is for our PaMoRL outperform baselines.
Sequence Model Runtime (ms) Eligibility Trace Estimation Runtime (ms)
GPU Memory Utilization (MB)
1400
170
1003
100
1003
400
200
Sequentaal
Kagge-stor
Odd-
Sequential
Kagge-slane
Sequentiel Rolloul
Kogge-stone Scamner
Ralloul
Scanmer
Scanmer
Rollout
Scanner
Scanner
Scanmer
Method
Method
GPU Memory Usage (%)
PaMoRL
Well clockclock lime (Minute)
PaMoRL
Dreamerv3
Dreamerv3
16
10
323
84/18
32/32
16/64
8/128
4/256
2/512
1/1324
64/16
32/32
16364
8/128
42256
2/512
1/1024
Batch Size Sequence Length
Figure 5: (Upper) Comparison of parallel scanners with sequential rollout in terms of runtime for sequence modeling and eligibility trace estimation, as well total GPU memory utilization. (Lower) Wall-clock time vs. GPU memory usage comparison for our PaMoRL method, SSM, and DreamerV3 across various batch size and sequence length combinations.
The findings in Figure 3 indicate that, while the token mixing module has minimal impact on the final performance for tasks where the reward can be accurately predicted from single frame (e.g., Boxing), leads to a performance drop on tasks that require several contextual frames to predict the reward accurately (e.g., Alien and Ms. Pacman). Regarding RMSNorm, removing negatively affects the final performance and increases the instability of the training process.
There are two possible reasons for this difference. First, the gradient is bounded after the original normalizer is removed [40]. Adding RMSNorm further enhances training stability, which is especially important in the setting of limited data and end-to-end training. Second, RMSNorm only rescales the input and maintains the original center of the samples, which allows the module's 's output to maximize the information's 's retention.
Parallel Scanner Selection. Figure 5 shows PWM and PETE's runtime and GPU memory uti- lization single 3090 GPU using different scanners, respectively. Sequence model computation

---

achieves 7.2x and 16.6x speedups compared to sequential rollout using the Kogge-stone and Odd- even scanners, respectively, with sequence length of 64. In this case, the Kogge-stone scanner with the theoretically lowest runtime takes more than the Odd-even scanner in practice. This is because the computation of the sequence model involves the parallelism of both batch and hidden dimensions, which belongs to massively parallel computation, and the Kogge-stone scanner cannot realize full parallelism and thus encounters a bottleneck in computational resources. In contrast, the Odd-even scanner is due to less computational redundancy, which allows the computational process of sequence modeling to achieve full parallelism and thus spends less running time. The PETE uses the Kogge-stone scanner and Odd-even scanner to achieve 3x and 2x speedups, respectively, with sequence length of 16. Since the eligibility trace has dimension of only 1, the Kogge-stone scanner can take full advantage of it. It thus achieves less runtime compared to the Odd-even scanner.
Regarding GPU memory utilization, using the Kogge-stone scanner imposes an additional 6x overhead compared to the sequential rollout, while the Odd-even scanner imposes an additional 2x overhead compared to the sequential rollout. However, the additional GPU memory overhead of parallel computation is not significant compared to the GPU memory overhead of encoder and decoder computation, especially in tasks with image observation.
Therefore, we recommend using the Odd-even scanner for PWM and the Kogge-stone scanner for PETE to achieve maximal speed with acceptable additional GPU memory utilization.
Batch Normalization Trick. World models are commonly learned using variational autoencoders
create concise representations of observations. However, they have some drawbacks, such as the tendency to disregard small moving objects. In Figure 11 in Appendix K. the reconstruction results are compared with and without using Batch Normalization for the Pong and Breakout games in the Atari 100K benchmark. It is observed that Batch Normalization improves the ability to distinguish similar video frames and capture information about small objects by re-centering the samples.
Additionally, Figure 12 demonstrates that PWM benefits from the batch normalization trick, whereas D3eamerV3 does not. This is likely due PWM's decoder solely having stochastic states as inputs, making it challenging for training samples to be distinguished from each other in the early stages of training. leading to "posterior collapse" (53) On the other hand, DreamerV3's decoder mitigates this problem by incorporating additional deterministic states as conditional inputs.
5 Conclusion & Limitations
In this paper, we introduce the PaMoRL. framework, an MBRL method capable of being computed using the parallel scan in both the model learning and policy learning stages. The key breakthrough of PaMoRL is the integration two novel techniques: the Parallelized World Model and Parallelizable Eligibility Trace Estimation. With these techniques, we simultaneously accelerate the training process while maintaining MBRL-level sample efficiency. PaMoRL demonstrates excellent hardware efficiency and training stability in various games tasks in the Atari 100K benchmark and DeepMind Control suite without incurring additional overhead during inference. An important contribution of work is the introduction of modified linear attention module in the MBRL method. Furthermore, we show that eligibility trace estimation computation can be parallelized for the first time.
It's important to acknowledge the limitations of our work. For instance, planning- glasedng-based based MBRL methods cannot parallelize computation over the sequence length, which hinders the incorporation of the most sample-efficient methods within our PaMoRL framework to maximize hardware efficiency. would be interesting to explore using hybrid architectures to enhance PaMoRL by leveraging the strengths of Transformers, RNNs, and SSMs. Additionally, the world model and baselines used for comparison in PaMoRL are trained end-to-end with joint optimization of the image encoder and sequence model. While this end-to-ent training paradigm enables the world model to predict the latent representations, also impacts the scalability of the world model.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by the NSFC project (No. 62072399), Zhejiang Provincial Natural Science Foundation of China under Grant No. LZ23F020009, Chinese Knowledge Center for Engineering Sciences and Technology, MoE Engineering Research Center of Digital Library, China Research Centre on Data and Knowledge for Engineering Sciences and Technology, and the Fundamental Research Funds for the Central Universities (No. 226-2024000170).
References
[1] Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4)::60-163, 1991.
[2] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. Robotica, 17(2):229-22) 235, 1999.
[3] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint urXiv: 1912.0160,0 2019.
[4] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
[5] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
[6] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02 193, 2020.
[7] Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34:25476-25488, 2021.
[8] Max Schwarzer, Johan Samir Obando Ceron, Aaron Courville, Marc G Bellemare, Rishabh Agarwal, and Pablo Samuel Castro. Bigger, better, faster: Human-level atari with human-level efficiency. In International Conference on Machine Learning, pages 30365-30380. PMLR, 2023.
[9] Shengjie Wang, Shaohuai Liu, Weirui Ye, Jiacheng You, and Yang Gao. Efficientzero v2: Mastering discrete and continuous control with limited data. arXiv preprint arXiv:2403.00564 2024.
[10] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Day- dreamer: World models for physical robot learning. In Conference on Robot Learning, pages 2226-2240. PMLR, 2023.
[11] Russell Mendonca, Shikhar Bahl, and Deepak Pathak. Structured world models from human videos. arXiv preprint arXiv:2308. 10901, 2023.
[12] Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint 220.::2203304955. 2022.
[13] Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpc2: Scalable, robust world models for continuous control. arXiv preprint 2.0.:2310..6822 2023.
[14] Richard Sutton and Andrew Barto. Reinforcement learning: An introduction. MIT press, 2018.
[15] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122. 2018.
[16] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model- based reinforcement learning for atari. arXiv preprint
2019.
11

---

[17] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint M.XV::3301.04006 2023.
[18] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[19] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09481, 2022.
[20] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv preprint arXiv:2209.0000 20000000,, 2022.
[21] Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. arXiv preprint arXiv:2303.07 109, 2023.
[22] Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. Storm: Efficient stochastic transformer based world models for reinforcement learning. arXiv preprint arXiv:2310.09615. 2023.
[23] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury. Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. Proceedings Machine Learning and Systems, 5, 2023.
[24] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 1801. 00660000,,,, 2018
[25] Lior Cohen, Kaixin Wang, Bingyi Kang, and Shie Mannor. Improving token-based world models with parallel observation prediction. arXiv preprint arXiv:2402.05643 2024.
[26] Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. arXiv preprint vxXiv:239.112252 2023.
[27] Richard E Ladner and Michael J Fischer. Parallel prefix computation. Journal of the ACM (JACM), 27(4):831-833 1980.
[28] Mark Harris, Shubhabrata Sengupta, and John D Owens. Parallel prefix sum (scan) with cuda. GPU gems, 3(39):851-876, 2007. 2007.
[29] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. arXiv preprini arXiv: / 700.044077 2017.
[30] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208. 000:22088.0433,, 2022.
[31] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International conference machine learning. pages 5156-5166 PMLR, 2020.
[32] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint 231.:2300511344. 2023.
[33] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
[34] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations, arXiv preprint arXiv:2007.05929, 2020.
[35] Rémi Munos, Tom Stepleton, Anna Harutyunyan, and Marc Bellemare. Safe and efficient off-policy reinforcement learning. Advances in neural information processing systems, 29, 2016.
12

---

[36] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv: 150.002380 2015.
[37] Guy E Blelloch. Prefix sums and their applications. 1990.
[38] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
[39] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1308.3432, 2013.
[40] Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. arXiv preprint arXiv:2210.10340, 2022.
[41] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[42] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv: /4/2.3555, 2014.
[43] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerehner. beta-vae: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2016.
[44] Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
[45] Peter M Kogge and Harold S Stone. A parallel algorithm for the efficient solution of general class of recurrence equations. IEEE transactions on computers, 100(8);786-793, 1973.
[46] Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: Rnns, transformers, and s4. arXiv preprint arXiv:2307.02000 2023.
[47] Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Master- ing memory tasks with world models. Second Agent Learning peen--ndeddess Workshop. 2023.
[48] Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Mare Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference on machine learning, pages 1995-2003. PMLR, 2016.
[49] Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and planning in complex action spaces. In International Conference on Machine Learning, pages 4476-4486. PMLR, 2021.
[50] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018
[51] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint 1rXiv:2107.09664 2021.
[52] Julian Schritwieserr, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-6N) 2020.
[53] Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. arXiv preprint arXiv: 15: 1511. 1100449,, 2015.
13

---

[54] Carlos E Garcia, David M Prett, and Manfred Morari. Model predictive control: Theory and practice survey. Automatica, 25(3):335-348, 1989.
[55] Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algo- rithmic framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint 180.0038 2018.
[56] Jian Shen, Han Zhao, Weinan Zhang, and Yong Yu. Model-based policy optimization with unsupervised model adaptation. Advances in Neural Information Processing Systems, 33:2823- 2834, 2020.
[57] Marc Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach policy search. In Proceedings of the 28th International Conference on machine learning (ICML-I1), pages 065-472,22 2011.
(58) Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforeement learning with stochastic ensemble value expansion. Advances in neural information processing systems, 31, 2018.
[59] Vladimir Feinberg, Alvin Wan, lon Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Mosel-based value estimation for efficient model-free reinforcement learning. arXiv preprint arXiv: 1803 00/01, 2018.
[60] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint (XXi::2111.00990 2021.
[61] 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. arXiv preprint arXiv:201 1.04006, 2020.
[62] Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.
[63]
Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arXiv:2209. ntte:2229..9999 /295 2022.
[64] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pages 9355-9366.
PMLR, 2021.
[65] Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
[66] Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint wrXiv:2103.02143 2021.
[67] Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. arXiv preprint arXiv:2210.04224, 2022.
[68] Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling, Advances Neural Information Processing Systems, 36, 2024.
[69] Blazej Osinski, Chelsea Finn, Dumitru Erhan, George Tucker, Henryk Michalewski, Konrad Czechowski, Lukasz Mieczyslaw Kaiser, Mohammad Babaeizadeh, Piotr Kozakowski, Piotr Milos, et al. Model-based reinforcement learning for atari. ICI.R, 1:2, 2020.
14

---

A Related Work
Model-based reinforcement learning (MBRL) methods aim construct a world model of the real environment and utilize the model to enhance the performance of policy. The crucial aspect of MBRL lies in its utilization of a world model. Previous methods in MBRL have employed world models in various ways, including searching for optimal action sequences
54
31, generating synthetic
data [ 55 3 56], or improving the value estimation [57 58
Our PaMoRL framework
builds upon DreamerV3 [17], which embraces the Dyna paradigm
2
14
The world model
utilized to interact with the policy to generate synthetic data, aiding model-free RI. methods in maximizing cumulative task rewards during the policy learning stage.
To maximize sample efficiency, MBRL researchers MBRL researchers have attempted employ 3
Recurrent Neural Networks (RNNs) [S 16,
as the architecture of world models. However, the update of the hidden state of the RNNs involves full matrix multiplication and the presence of nonlinearities within the recurrence hinders parallel computation, resulting in slower training speed. While Transformers provide a viable high- perfforrfomm option, there is quadratic relationship between their computational complexity and sequence length, which introduces additional computational and memory overhead.
6 17 or Transformers 19
21
20 22
Recent research has introduced a novel linear RNN architecture with simplified interaction between hidden states called the Structured State Space sequence model (S4) [60] that surpass both Trans- formers and RNNs in Loong-Range Arena benchmarks (LRA) [61]. The S4 model and its variants are designed to effectively handle tasks involving long-range reasoning and draw inspiration from classi- cal continuous-time linear state space models (SSMs) [62), which are well-established components of control theory. The relationship between the time and the frequency domain implies that SSMs have a convolutional view when the decay rate is data-independent, and therefore, training can be accomplished using the fast Fourier transforms (FFT).
We remark on the importance of incorporating a data-dependent decay rate, which is ignored by current works in SSMs until liquid S4 (63 and Mamba [33. Our PWM builds upon linear attention with a data-dependent decay rate, which does not have the convolutional view and thus cannot use FFT for training but allows the use of parallel scans. The field of linear attentions and linear RNNs exhibits close relationship [1. i.c. linear attentions can be reformulated as linear RNNs during auto-regressive decoding, revealing similarities to the update rules observed in fast weight additive outer products [64| 65 These updated rules can be seen as a special case of element-wise linear recurrence. However, this formulation in linear attention cannot forget irrelevant information, resulting in the attention dilution issue. To address this limitation, gating mechanisms [66] 67 68] can be used facilitate the forgetting of irrelevant information similar to those in traditional RNNs.
The work that is most similar to our PaMoRL is Mamba 33 Both our PaMoRI. and Mamba have data-dependent decay rates and employ parallel scans to speed up the training process. However, there are significant differences between PaMoRL and Mamba in terms of model architectures and hardware preferences. In terms of model architecture, Mamba needs to maintain self-consistency with previous work in the SSM family, and therefore it must adhere to the paradigm of classical state space models, representing continuous differential equations. It needs to be parameterized and discretized using special tricks to achieve the gating mechanism implicitly. On the other hand, we recognize this limitation of Mamba and use a more "simple yet effective" gating mechanism. Regarding hardware preference, Mamba employs special 10-aware parallel scanning algorithm for efficient training, which focuses on reducing the number of reads and writes between SRAM and HBM in the GPU through kernel fusion, and is suitable for improving the training efficiency of the hardware features when the model architecture is determined. In contrast, to satisfy the need for flexibility in MBRL, the parallel scanner we use is inspired by high-performance computing hardware design and focuses more on generality. Our parallel scanning method is compatible with arbitrary model architectures, as long as it satisfies the parallelization conditions mentioned in Section architecture- speeciifc parallel scanning method used by Mamba.
2 as compared to the model
15

---

B Illustrations to Parallel Scan Algorithms
B.1 Kogge-stone Scanner
common example such first-order recurrence problem is time-varying linear system, the system's state at timestep / is 2i, computed from the system's internal dynamical variables de and bt. as shown in Equation 9 Depending on the problem, the variables and by can be real or complex numbers, constants, etc.
J1 b1:
T2 = a2 + b2
D3 03 by.
(9)
31 a:ti-1 + bir
F1, = a1 + br.
Before solving the problem, we can define the function A(m, n) and B(m, ni), as shown in Equa- tion 12
A(m, n) = II 0z.
j=n
(10)
B(m, n) = - II az)b,, where S m.
j=1-
Now we can integrate Equation 9 with Equation 10 to get Equation 11
B(1, = I1 bi
(22 1) x2 d2 + b2 + - B(2,))
B(4, = x4 =
+ -4 ((4,),
(11)
B(2i, 1) F2s
II j=++1
-)
B(2
It can be observed in Equation 11 that B(2i, i + 1) is associated with the computation of A(2i, i 1) but independent from B(i, 1), which means that we can split the computation of B(2i, 1) into two parallel parts. For reasons of notational simplicity, functions A(m, n)) and B(m, n), as shown in Equation
define the tuple Q(m, n) that wraps the
Qmn, n) (A(m, n), B(m, n)), where m.
(12)
Figure 6 shows the operation of the Kogge-stone scanner [log, L iterations the solution to the problem 2i...., Pr be computed.
45 when the sequence length I, = After
B.2 Odd-even Scanner
To avoid the extra computational complexity log, J. generated by the Kogge-stone scanner 45 the
Odd-even scanner [28] uses algorithmic pattern that arises often in parallel computing: balanced
16

---

trees. The idea is to build a balanced binary tree on the input data and start scanning from the root. binary tree with L leaves has log, layers with 2d nodes per layer d [0, L). If we perform one operation on each node, then we will perform (T.) operations in one traversal of the tree. 'The tree we construct not an actual data structure, but rather a concept that we use determine what each thread has to do at each step of the traversal.
The algorithm consists of two phases: up-sweep and down-sweep. During the up-sweep phase, traverse from the leaves to the root of the tree. During the down-sweep phase, we backtrack from the root node up the tree, using the results computed the up-sweep phase. Figure shows the operation of the Kogge-stone scanner [28] when the sequence length L = 8.
Note that since this is exclusive scan (i.e., the sum is not included in the result), last element of the array between phases. This zero is propagated back to the head of the array in the down-sweep phase. This scanning algorithm performs (22T) operations, is very efficient.
out the
(83,8)
Q(7,7)
(66,6)
((5,5)
((4,4)
((3,3)
Q(2.2)
@(1.1)
Step
(88.7)
((7.6)
((6,5)
((5.4)
((4.3)
(33,2)
((2.1)
((1.1)
Step 2
(8.55)
((7,4)
(66.3)
((522)
(4.11)
((3.1)
(
((.1))
Step 3
(88.1)
((7.1)
(66.1)
((5.1)
((4.1)
((31))
(2.1)
((1.))
Figure 6: Illustrations of the operation of the Kogge-stone scanner when the sequence length L 8.
(3.8)
((7.7)
((6,6)
((5,5)
((4,4)
((3,3)
Q(2.2)
((111)
Step
((8,7)
((7.7)
((6,))
((55))
((4.3)
((3,3)
((211)
((11))
Step 2
(8355)
Q(7.7)
((6,))
(55
((4.1)
(33,))
Q(211)
((.1))
Step
Q(8,1)
((7.7)
(66,5)
((5,))
((4.1)
((3,3)
Q(2,1)
Q(1,1)
Zero
Step 4
Q(7.7)
(66,5)
((5
((4.1)
((3.3)
Q(2.1)
((11))
Step 5
((4.1)
((7.7)
((6,5)
((555)
(3.3)
((21))
((11)
Step 6
(6.1))
((7.7)
((4.1)
(555)
((2.1)
((3,3)
((11))
Step 7
((7.1)
@(6.1)
((5.1)
((4.1)
((3.1)
((211)
((11)
Figure 7: Illustrations of the operation of the Odd-even scanner when the sequence length L. = 8.
17

---

C Training Curves of the Atari 100K Benchmark
The Atari 100K benchmark 69 is standard RL benchmark comprising 26 Atari games featuring diverse gameplay mechanics. It is designed to assess broad spectrum of agent skills, and agents are limited 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 to train agents 200M steps.
300
Ba WkHolil
BattoZonne
Booling
Broakout
11.
ChooperCommend
CranyClimber
DomonAttock
Frocwaa
Fuxilile
Jamnesaard
200
Kangaroo
KangFuMMazier
McPacman
650
Pong
PriveteEye
Obert
ReadFunner
100
Seracusst
UpNDown
300
Figure 8: Training curves on the Atari 100k benchmark. The solid line represents the average result over 5 seeds, and the filled area indicates the range between the maximum and minimum results across these seeds.
18

---

D Training Curves of the DeepIMind Control Suite
DeepMind Control suite 24 standard RL benchmark comprising various tasks with continuous action spaces. It supports both image observation and low-dimensional proprioception observation. When there are no constraints on sample efficiency, the typical practice to train agents for millions of steps.
coripole belonce
corppoe: balance spanse
cartpolc Swingup
catch
300
101
101
Snger
pendufum
reacher
reacher
cartpole
sparse
easy
400
hopper, stand
hepper, hop
q quadruped sun
QJ quaddruped
2011
Figure 9: Training curves the DeepMind Control suite with image observations. The solid line represents the average result over seeds, and the filled area indicates the range between the maximum and minimum results these seeds.
19

---

carppoio_ Dal anco
corppolo oartpee_oddeeee..oo sparso
cargaio, SWRTOMC
0up_00101
Enger
penzuaamh SWIrLUp
30-
hoppor_soo stand
hoppor_
quduuaa
quadrupod walk
Figure 10: Training curves the DeepMind Control suite with proprioception observations. The solid line represents the average result over 5 seeds, and the filled area indicates the range between the maximum and minimum results across these 5 seeds.
20

---

E Details of Model Architecture
Table 4: Architecture details of the image encoder. The size of the modules is omitted and can be derived from the shape of the tensors. SiLU refers the sigmoid-weighted linear units used for activation, while Linear represents fully connected layer. Flatten and Reshape operations are employed alter the tensor's indexing method while preserving the data and their original order. Conv denotes a CNN layer characterized by kernel = 4, stride = 2, and padding 1. BN denotes the batch normalization layer.
Module
Output tensor shape
Imput image (ot)
x 64 x 64
Conv +BBN ++SiLU
x 32 x
64 x 16
16
Conv3+BN3+S0 Con44+BN++SILU
128 x 8 x 8
256 x
Flatten
4096
Linear + BN5 + Reshape
32 x 32
Table 5: Architecture details of the image decoder. DeConv denotes a transpose CNN layer characterized by kernel = 4, stride = 2, and padding 1.
Module
Output tensor shape
Random samples (z²)
32 x 32
Flatten
1024
Linear+BNI+SiLU
4096
Reshape
256 X 4 4
DeConv ReConvI+BNIn -SNUU+
128 x 8 x
DeConv2+BN2+SiLU
64 x 16
16
DeConv3+BN3+SiLUu
32 x
32
x 32
DeConv4
3 x
64
X
64
Table 6: Action mixer. Concatenate denotes combining the last dimension of two tensors and merging them into tensor. The variable A represents the action dimension. /) denotes the feature dimension of the sequence model. LN is an abbreviation for layer normalization.
Module
Output tensor shape
Random samples (2²) & Action (a;) Reshape and concatenate
32 x 32, A 1024 A
Linear+LN114iLLJ
D
Lin
D
Table 7: Modules which are pure MLPs. l-layer MLP corresponds to fully connected layer. 255 the size of the bucket of symlog two-hot loss I7] K refers to the dimension of proprioception observations.
Name
MLP layers
lutput//lidden/Ouuppt shape
Encoder (proprio)
3
K/512/D
Decoder (proprio)
3
D/512/K
Dynamic predictor
I
DID/1024
Reward predictor
3
DID/225
Continuation predictor
3
DID/1
Actor network
3
DIDIA
Critic network
3
DIDI225
21

---

Hyperparameters
Table 8: Full hyperparameters. Note that the environment will provide "done" signal when losing life but will continue running until the actual reset oceurs. This life information configuration aligns with the setup used in IRIS [20]. Regarding data sampling, each time, we sample B1 trajectories of length 7 for world model training and sample B2 trajectories of length C' for starting the imaginations.
Hyperparameter
Symbol
Value
Sequence model layers
K
2
Hidden size of query, key, and value Hidden size of sequence model output
64
D
512
World model training batch size
B
16
World model training batch length
T
64
Imagination batch size
B2
1024
Imagination horizon
16
Update world model every environment step
1
Update policy environment env step
1
Scan algorithm for world model training Scan algorithm for policy training
Odd-even Kogge-stone
Gamma
0.997action repeat
X
Lambda
0.95 3 x 10
Entropy coefficiency
7
Optimizer
Adam
World model learning rate
1 x 10-4
World model gradient norm clipping
100.0
Actor-critic learning rate Actor-critic gradient norm
3 x 10 10-5
100.0
Gray image input
False False
Frame stacking
Frame skipping
(Atari) or (DMControl)
Use of life information
True (Atari)
22

---

G Pytorch-style Pseudo-code of Parallel Scan
G.1 Odd-even scanner
def
aed_even.parallel. scan www.nnnnnnnnniiinaaa (inputs, operator):
Odd/Even Parallel Scanner.
Inputs:
inputs: tuple of sequence elements. operator: binary operator function.
Outputs:
outputa: tuple of sequence elements.
**.
10
Length inputs [0] shape [0]
11
if Length < 2:
13
return inputs
15
reduced imputs = operator(
(input [n]]
for input in inputs), (input [1::2] for input in inputs)
18
19
odd_inpus induts = odd_ odd_evee..praaleel.com (reduced, operator)
20
if Length % =- 0: even iven_inputs operator(
22
23
(input -1] for input odd_ add_inputs),
24
(input [2: for input in inputs)
else:
27
even ven_inputs- operator
28
(input for input in odd_inpute inputs), (input [2: for input in inputs)
29
31
32
even inputi [
33
torch. cat ((input [0:1], even_input), dim=0)
for (input, ivputin in zip (inputs, even_inputs)
35
36
37
outputs
38
interleave (odd_ input, iventinput) for (even_ input, odd_input) in zip (ip(een_innuus,, odd_ idd_inputs)
41
return outpute
42
def interleave (odd, even):
padded_ torch cat ((odd, torch. zeros_ (odd [1::]) dim=0)
46
outputs torch. stack ((even, padded, oodded_odddl even.. shape [0]]), dim=1)
47
fntputs.flatten outputs outputs. flatten (0,, 1) [: (odd. shape [0] even. shape [0])]
return outputs
23

---

G.2 Kogge-stone scanner
def
(inputs, operator):
e*a
Kogge-Stone Parallel Scanner.
Inputs:
inputs: tuple of sequence elements.
operator: binary operator function.
Outputa:
outputs: tuple of sequence elements.
Length inputs [0] shape [o]
11
Times math.ceill (math. log2 (Length))
12
for i in range (Times):
interval int ** i)
outputs operator(
16
[:-ttr:ainterva] for input in inputs), (input[internnn :] for input in inputs)
17
inputs
20
torch. cat ((input [iitervaal,, output), dim-0) for (input, output) in zip (inputs, outputs)
21
22
]
return inputs
24

---

H Pytorch-style Pseudo-code of Parallelized Eligibility Trace Estimation
def
travaarrrdd
(revard, value, next_value, p_cont, lam)
Parallel Eligibility Trace Estimations.
ones torch. ones ones_ikee (reward) p.cont, lam P_cont ones, lam ones lam - torch. cat ((1a [1:], ones[:1l] dim-0)
delta reward P_cont value value deipped_delta delta. (lin.ffiii(ddm=(0)))) flipped_lam = (p_cont lam) flip (d)
10
11
residual odd_eve__parrleelssa
[flipped_lam, [iipped_delta]] binary_return. fn)
15
returns value residual [1] flip (dims=
16
return returns
19
def per (reward, value, next_value, p_cont, lam)
20
21
Parallel LD-baa Estimations.
23
ones torch ones (reward) P_cont, lam p-cont ones, lam ones
24
25
delta reward + P-cont next value * (1 lam) last delta[-1:] delta torch. cat ([deltal:-1), last), dim=0)
p_cont[-1:] lam [-m]] * next_value [-1:]
28
29
30
flipped_dell delta = delta. delta.flip flip(dims==0,)) flipped_lam (p_cont lam) (d)m
32
33
returns edd_eve
34
[flipped__am, flipped_deltal,
delta], binary.return.fa
returns returns [1]. flip (d)
return returns
37
38
39 def Dinary_reture (or_nnnncuuuuruur cur_i, cur_j):
coef_i, in_i cur_i coef_j, in_j cur_]
41
42
return coef_i coef_j, coef_j in in_j
25

---

I Additional Comparisons on the Atari 100K Benchmark
Atari 100K benchmark [69] is standard RL benchmark comprising 26 Atari games featuring diverse gameplay mechanics. It is designed to assess a broad spectrum of agent skills, and agents are limited 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 to train agents 200M steps.
In this section, we compare the performance of our PaMoRL framework with planning-based methods such EfficientZero 77] and EfficientZero V2 [ and methods with much larger networks, i.e., BBF [8] on the Atari I 100K benchmark. The full results are shown in Table The PaMoRL framework good as the other methods in terms the number of superhuman games, median score, and average score. However, it leads the pack of 13/26 games in terms of an individual game perspective.
Table 9:
Experimental results on the 26 games of Atari 100k after hours of real- real-time experience and Homan-normalized aggregate metrics. Bold and underlined numbers indicate the highest and the second-highese scores, respectively.
Game
Random
Human
EfficientZero
BBF
EfficientZero V2
PaMoRL (Ours)
Alien
227.8
7127.7
808.5
1173.2
1557.7
1270.6
Amidar
5.8
1719.5
148.6
244.6
184.9
264.4
Assault
222.4
1263.1
2098.5
1757.5
833.8
Asterix
210
8503.3
25557.8
3946.1
61810
2957.3
BankHeist
14.2
753.1
351
732.9
1316.7
225.9
BattleZone
2360
37187.5
13871.2
24459.8
14433.3
23120
Boxing
0.1
12.1
52.7
85.8
75
87.9
Breakout
1.7
30.5
414.1
370.6
400.1
15.8
ChopperComunndd
811
7387.8
1117.3
7549.3
II96.6
2110.7
CrazyClimber
10780.5
35829.4
83940.2
58431.8
112363.3
84102
DemonAttack
152.1
1971
13003.9
13341.4
22773.5
208.2
Freeway
29.6
21.8
25.5
33.8
Frostbite
65.2
4334.7
296.3
2384.8
1136.3
3711.4
Gopher
257.6
2412.5
3260.3
1331.2
3868.7
5085.2
Hero
1027
30826.4
9315.9
7818.6
9705
12076.2
Jamebond
29
302.8
517
1129.6
468.3
405
Kangaroo
3035
724.1
6614.7
1886.7
2554.7
Krull
1598
2665.5
5663.3
8223.4
9080
7273.2
KungFuMaster
258.5
22736.3
30944.8
18991.7
28883.3
24624.7
MsPacman
307.3
6951.6
1281.2
2008.3
2251
2201.7
Pong
-20.7
14.6
20.1
16.7
20.8
15.5
PrivateEye
24.9
69571.3
40.5
99.8
4968.6
Qbert
163.9
13455
14448.5
4447.1
16058.3
4730.3
Roadrunner
11.5
7845
17751.3
33426.8
27516.7
24726.7
Seaquest
68.4
42054.7
1100.2
1232.5
1974
595.2
UpNDwon
533.4
11693.2
17264.2
12101.7
15224.3
11953.8
Games >Human
0
26
14
12
21
9
Median
0%
100%
111.53%
91.71%
123.47%
71.75%
Mean
0%
100%
194.46%
224.74%
267.97%
126.64%
26

---

Runtime of Experiments
Table 10: Average runtime of experiments
Task
Atari 100K
Proprio (easy)
Proprio (hard)
Vision (casy)
Vision (hard)
Runtime
3.5 hours
0.94 hours
1.88 hours
2.74 hours
7.1 hours
K
Effectiveness of Batch Normalization Trick
(8)
Figure 11: Visualizations on Batch Normalization trick in Pong and Breakout.
Breakout
Pong
17.5
Store
10.0
-10
25
-20
400%
400
PWM
PWW wio BN
DreamerV3 w/ BN
DreamerV3
Figure 12: Quantitative results on the effectiveness of the Batch Normalization trick.
27

---

L Video Predictions
Context Imput
Imagination
Enviranment Alien
Real
Model
Context Imput
Imagination
Environment Breakout
Real
Model
1:4
F12
F-20
F228
F:36
F-44
P52
F80
Context Input
Imagination
Environment: Pong
Real
Medal
I+0
1-12 Imagination
1-28
I-01
I+
Impul
Real
Model
F12
Context Imput
Imagmation
Environmant Happer Hop
Real
Macill
mh
112
+00
m06
552
MG0
Context Input
Imagination
Environment Quadruped
Real
Model
Figure 13: Multi-step predictions several environments in Atari games and DeepMind Control suite. The world model utilizes observations and actions as contextual input, enabling the imagina- tion of future events spanning 56 frames in an auto-regressive manner.
28

---

M Initializations in Freeway
The reward function in Freeway sparse since the agent only rewarded when it completely crosses the road. In addition, bumping into cars will drag it down, preventing it from smoothly ascending the highway. This poses an exploration problem for newly initialized agents because a random policy will almost surely never obtain a non-zero reward with 100k frames budget. The solution to this problem actually straightforward and requires stretches time when the "UP" action is oversampled. In this paper, we opted for the simplest strategy of having initialized buffer with fulfilled "UP" actions. Hence, we dont't 't need to lowered the sampling temperature to avoid random walks that would not be conducive to learning in the early stages of training. Consequently, once it received its first few rewards through exploration, our PaMoRL could internalize the sparse reward function in its world model.
Figure 14: A game of Freeway. Cars will bump the player down, making it very unlikely the road and be rewarded for random policies.
29

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly claim the three contributions in the introduction. We also describe the method and show the experimental results in the main paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We state the limitations in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
30

---

Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide sufficient information about the hyper-parameters as well as the details in the Appendix. We also pack our code in the supplemental materials.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
31

---

Answer: [Yes]
Justification: We pack the code in the supplemental materials and the code base of each algorithm is clearly cited in the main paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https /nipss.cc/pulli//uuui guwde//ddee/ee/ wwt.:Suumisssispoooooooooooooooe
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We show sufficient experimental settings and hyper-parameters in the main paper and the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: In each graph, we show the average results among 3 differennt seeds and the std variance of the learning curves.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
32

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We show the hardware for training and state the time used for each training process the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the htt.eerrppppp.pppppe/eeeeaaeeeee
NeurIPS Code of Ethics
Answer: [Yes]
Justification: The experimental environments and the baseline algorithms code-bases are fully officially released, open-sourced, and open-accessed.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there is a special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: We propose parallel model-based RL framework for fast training, so there is no societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
33

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We propose parallel model-based RL framework for fast training, so there is no misuse concerns.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The experimental environments well as the algorithms code-bases properly cited in the experiment settings section.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
34

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
batasets.papersstt
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We provide the code with documents in supplemental materials and will release them.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
35

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
36


Reviews:
[{"review_id": "huhSyuDkxL", "reviewer": "Reviewer_MVJ1", "summary": "Model-based RL algorithms are popular due to their strong data efficiency compared to model-free alternatives. However, most MBRL algorithms learn a recurrent world model that scales linearly in time complexity wrt the input sequence length. This paper proposes several changes to the DreamerV3 algorithm that allows for parallelization of the world model training across the sequence (temporal) dimension, in particular use of the Odd-Even parallel scan algorithm + architectural changes that permit use of parallel scans. Experiments are conducted on Atari 100k + DMControl (proprio + pixels), and results indicate that the proposed method both achieves better data efficiency and hardware efficiency compared to prior work.\n\n**Post-rebuttal update:** I believe that the authors have addressed my main concerns. I have increased my score from 4 -> 5 and soundness from 2 -> 3 to reflect that.", "questions": "I would like the authors to address my comments above. I have a couple of additional follow-up questions:\n- I am not entirely sure how to read Figure 1. There appears to be 3 metrics shown in each subfigure, but only two axes. For example, the right-most figure shows GPU memory utilization, mean score, and training FPS. Can the authors please clarify how to read this figure? Additionally, the figure text is very small and difficult to read; I suggest increasing the font size.\n- It appears that the two ablations (token mixing, RMSNorm) are conducted on different sets of tasks. What is the reasoning behind that? This makes the results seem somewhat cherry-picked to me.\n- Based on the description of the use of batch normalization (L257) + the qualitative results in the appendices, it seems like BN is quite useful. Why do the authors refrain from running any quantitative ablations on the use of BN? Does DreamerV3 benefit equally from BN, given that this is largely orthogonal to the actual contributions of this work?", "limitations": "I believe that the authors address limitations adequately in the last paragraph of the paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The problem studied in this paper is both interesting, timely, and highly relevant to the NeurIPS community. This paper extends (at least in spirit) a series of recent papers from LLM literature (e.g. RWKV, Mamba) to MBRL algorithms that similarly can benefit from parallelization of training while maintaining low memory footprint, since many MBRL methods leverage RNNs. It is refreshing to see this kind of work in the MBRL space.\n- The proposed architectural and algorithmic changes are fairly intuitive and seemingly simple to implement, while there is a big potential upside in terms of training parallelization. While the changes are simple, they are not trivial to come up with. I absolutely consider this a key strength of this paper.\n- Experimental results are promising. It appears that the method is effective in terms of parallelizing model training while maintaining a low memory footprint (mostly due to the choice of parallel scan algorithm). Data efficiency / task performance is comparable or better than previous work on MBRL without planning.", "weaknesses": "While I find the technical contributions of this paper quite intriguing, there's a few things that I find somewhat problematic:\n- The key selling point of the proposed method is its parallelism. However, there's pretty limited evidence that it actually achieves this goal. Sure, runtime (ms) is generally quite a bit lower, but I would have expected to see a comparison in training wall-time vs. sequence length used to train the model. Intuitively, the computational gains would increase with the sequence length vs. vanilla DreamerV3, but I did not see any experiments of this sort.\n- Along the same lines, I find it a bit odd that most experiments emphasize better data efficiency when that is not really a central part of the paper. The authors even state themselves that \"our main goal is to maximize the hardware efficiency of existing MBRL methods, rather than pursuing extreme sample efficiency.\" (L188), so why do the experimental results primarily measure data (sample) efficiency rather than hardware efficiency? It seems like the proposed changes improve data efficiency so I understand why the authors would highlight that, but it seems mostly orthogonal to the actual problem that the authors set out to address.\n- Limited ablations. The authors make multiple changes to the algorithm, but only two (token mixing, RMSNorm) are ablated. I'm left with a fairly poor intuition for how important the various algorithmic changes really are for both data efficiency and hardware efficiency, and it is not even clear to me how the ablations compare to e.g. vanilla DreamerV3 here (Figure 3).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for reviewing our paper and for your many detailed comments. The following are responses to the weaknesses and questions you listed:\n\n### W1: Limited evidence of the key selling point of parallelism in our method.\n**R1**: Thank you very much for your suggestions on the experiments, we have already added the wall-clock times and GPU memory usage (based on an RTX 3090 GPU) for our PaMoRL, SSM (with the data-dependent decay rate removed), and DreamerV3[1] trained at different sequence length **in Figure 4 (Right)** in the PDF attached to \"Common Response\". The experimental results demonstrate that the computational gain of our PaMoRL method in terms of wall-clock time increases with sequence length, while at the same time, there is less than 14% GPU memory overhead (~3.36GB) at the maximum sequence length of 1024 compared to the minimum sequence length of 16.\n\n### W2: Why do the experimental results mainly measure data efficiency rather than hardware efficiency?\n**R2**: We believe that the key to determining whether an MBRL method can maximize hardware efficiency should be to determine whether it has operators that do not allow parallelization over the sequence lengths. For example, if the world model architecture and eligibility trace estimation in a particular MBRL method built upon DreamerV3 satisfies the three conditions mentioned in the \"Parallel Scan\" paragraph in the \"Background\" section of our paper (L95-L97), then in principle, this MBRL algorithm is capable of maximizing the hardware efficiency through the parallel scan algorithm. Furthermore, since the observation and action spaces of the individual tasks in each benchmark do not differ much in dimension and our method does not require an individual hyperparameter for each task, which means that our PaMoRL method is computationally highly consistent across tasks, we do not think that setting up too many experiments to demonstrate the hardware efficiency of our method is necessary.\n\nOn the other hand, the MBRL method has the key selling point of its inherent high data efficiency, and most of the previous works [1-3] that concluded that their methods are highly data-efficient stem from the superior performance in various tasks that measure data efficiency. Therefore we set up various experiments measuring data efficiency to empirically demonstrate that our PaMoRL method can also achieve MBRL-level data efficiency.\n\n### W3 & Q2: Limited ablation of the importance of various modules on data efficiency and hardware efficiency. Reasons for existing ablation experiments on different task sets?\n\n**R3**: According to experimental results from previous work [4, 5], both Token Mixing and RMSNorm modules play a key role in their model training. Therefore, we choose \"Alien\", \"Boxing\" and \"MsPacman\", which are tasks focusing on sequence prediction, to demonstrate the benefits of Token Mixing in sequence prediction. We also choose \"Amidar\", \"UpNDown\" and \"Qbert\", which are scattered tasks, to measure the improvement of RMSNorm on training stability.\n\nIn addition, we supplemented the training curves of the above 6 tasks by removing the PWM of Token Mixing and RMSNorm modules and added SSM (PWM without data-dependent decay rate) and vanilla DreamerV3. You can check the results in Figure 3 in the attached PDF of \"Common Response\". The experimental results show that Token Mixing, RMSNorm, and data-dependent decay rate are all beneficial to data efficiency. The results of the experiments on hardware efficiency can be found in **“R1”**.\n\n### Q1: Clarification in Figure 1.\n\n**A1**: We apologize for the incorrect title of Figure 1 due to a plotting error. We have corrected Figure 1 in the PDF attached to \"Common Response\" and increased the font size.\n\n### Q2: Why are existing ablation experiments being performed on different task sets?\n\n**A2**: Please find the response to this question in **\"R3\"**.\n\n### Q3: Quantitative ablation on the role of batch normalization techniques?\n\n**A3**: Thank you very much for your suggestion, we have added quantitative ablations to **Figure 4 (Left)** in the PDF attached to \"Common Response\". The results show that PWM benefits from the batch normalization trick but not DreamerV3. We believe that this is because PWM's decoder has only stochastic states as inputs, which makes it difficult for training samples to distinguish from each other in the early stages of training, resulting in the notorious \"posterior collapse\" [6]. The DreamerV3 decoder mitigates this problem by having additional deterministic states as conditional inputs.\n\n[1] Danijar Hafner, et al. \"Mastering diverse domains through world models.\" arXiv preprint arXiv:2301.04104 (2023).\n\n[2] Vincent Micheli, et al. \"Transformers are sample efficient world models.\" ICLR (2023).\n\n[3] Max Schwarzer, et al. \"Data-Efficient Reinforcement Learning with Self-Predictive Representations.\" ICLR (2021).\n\n[4] Bo Peng, et al. \"RWKV: Reinventing RNNs for the Transformer Era.\" ACL (2023).\n\n[5] Zhen Qin, et al. \"The Devil in Linear Transformer.\" EMNLP (2022). \n\n[6] Samuel R. Bowman, et al. \"Generating Sentences from a Continuous Space.\" CONLL (2016).", "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": "u9b0XLBNqw", "reviewer": "Reviewer_wbEe", "summary": "The paper proposes a novel framework to parallelize model-based RL, including two improvements parallelizing the world model and parallelizing eligibility traces. They demonstrate the dramatic speed-up of training speed without sacrificing inference efficiency. The proposed method achieves state-of-the-art score performance and reduces the runtime of two important components significantly.", "questions": "NA", "limitations": "NA", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "+ A significant speed-up of model generation runtime and eligibility trace estimation runtime ;\n+ The first paper to point out that the computational process of eligibility traces can be parallelized over the sequence length, which would be super helpful to the RL community to improve the efficiency of algorithms;\n+ Novel modifications to the RSSM module in the Dreamer algorithm to eliminate non-linear dependencies, each residual block of the sequence model consists of a modified linear attention module and a GLU module;\n+ A clear illustration of the proposed parallel algorithm with PyTorch (pseudo-)code;\n+ Significant reduction in runtimes of sequence model and eligibility trace estimation, Sequence model computation achieves 7.2× and 16.6× speedups compared to sequential rollout using the Kogge-stone and Odd-even scan algorithms, respectively, with a sequence length of 64.", "weaknesses": "This is an excellent paper and I didn't find any significant flaws.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your high evaluation of our paper and your detailed feedback on its strengths. We are very pleased to see that you have recognized the innovation and contribution of our proposed PaMoRL method, and we will further improve and enhance our paper. Please feel free to let us know if you have any further suggestions. We look forward to your further feedback.", "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": "U5It4NLsCx", "reviewer": "Reviewer_5drW", "summary": "This paper proposes a new framework named Parallelized Model-based Reinforcement Learning (PaMoRL) to improve the training speed of MBRL methods. PaMoRL employs a parallel scan technique to parallelize world model learning and eligibility trace estimation. With experiments on the Atari and DMC domain, the paper claims that PaMoRL can improve the sample efficiency of MBRL methods while significantly reducing training time.", "questions": "How do you extrapolate the training FPS of methods other than DreamerV3 and PaMoRL from other GPUs? Could you provide some additional details?", "limitations": "Fine", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper tackles a significant problem in MBRL.\n- The method is evaluated over two benchmarks, including a sufficient number of tasks.", "weaknesses": "- The paper lacks contribution. Parallelized World Model and Parallelizable Eligibility Trace Estimation merely employ old methods in this MBRL scenario, which are more like programming tricks than novel techniques.\n- Analysis and theoretical derivation of the parallel scan methods are missing.\n- Experiment results are reported only over three seeds, which can hardly be considered confident.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for reviewing our paper and for your many detailed comments. The following are responses to the weaknesses and issues you listed:\n\n### W1: Parallelized World Model and Parallelizable Eligibility Trace Estimation use only old methods and are more like programming tricks than novel techniques.\n\n**R1**: Our PaMoRL method draws inspiration from recent well-known work in the field of LLM (e.g., Mamba [1], RWKV [2], and Linear Transformer [3]) and follows the three conditions in the \"Parallel Scan\" paragraph in our paper (L95-L97) in the \"Background\" section of the paper to design the novel Parallelized World Model architecture.\nFor Parallelizable Eligibility Trace Estimation, we also find that this widely used return estimation method in the RL domain naturally satisfies the parallelizability condition, further accelerating MBRL training in the policy learning phase. From the perspective of the MBRL domain, we have good reasons to believe that our PaMoRL method uses novel techniques rather than just programming tricks.\n\nMoreover, as you point out, we have empirically demonstrated through sufficient experiments that our PaMoRL method significantly speeds up training while exceeding the baseline in most tasks, achieving the best of both worlds in terms of hardware efficiency and data efficiency, an achievement that we believe represents a significant contribution to the MBRL field.\n\n### W2: Lack of analysis and theoretical derivation of the parallel scanning method.\n\n**R2**: The analysis and theoretical derivation of the two parallel scanning methods (Kaggle-Stone and Odd-Even) used in our paper can be found in Chapter 1.4.1 of [4] and Chapter  39.2.1 of [5], respectively. We thank you for your suggestion and will add more discussion and analysis of the theoretical derivations of these two parallel scanning methods in the background section of the revised version of the paper.\n\n### W3: The experimental results are reported for only three seeds, which can hardly be considered confident.\n\n**R3**: We appreciate your suggestions for the experiment, and for this reason, we aligned the experimental setup with the previous well-known work [6-12] by adding two additional random seeds and took the suggestion from *Reviewer earm* to compare different methods using the metrics recommended in [13] to report results that give reliable comparisons.\n\n### Q1: Additional details of extrapolating training FPS from other GPUs for methods other than DreamerV3 and PaMoRL? \n\n**A1**: Among the other methods in Figure 1, IRIS [7] and REM[9], TWM [10] were evaluated on A100 GPUs, while SimPLe [11], STORM [12], and other model-free RL methods were evaluated on P100 GPUs.\nThe extrapolation methods we employ are consistent with the setup used in DreamerV3 [6], where it is assumed that the P100 is twice as fast as the P100 and the A100 is twice as fast.\n\n[1] Albert Gu, et al. \"Mamba: Linear-Time Sequence Modeling with Selective State Spaces.\" arXiv preprint arXiv:2312.00752 (2023).\n\n[2] Bo Peng, et al. \"RWKV: Reinventing RNNs for the Transformer Era.\" ACL (2023).\n\n[3] Zhen Qin, et al. \"The Devil in Linear Transformer.\" EMNLP (2022). \n\n[4] Blelloch, et al. \"Prefix sums and their applications.\" School of Computer Science, Carnegie Mellon University Pittsburgh (1990).\n\n[5] Harris, et al. \"Parallel prefix sum (scan) with CUDA.\" GPU gems (2007).\n\n[6] Danijar Hafner, et al. \"Mastering diverse domains through world models.\" arXiv preprint arXiv:2301.04104 (2023).\n\n[7] Vincent Micheli, et al. \"Transformers are sample efficient world models.\" ICLR (2023).\n\n[8] Max Schwarzer, et al. \"Data-Efficient Reinforcement Learning with Self-Predictive Representations.\" ICLR (2021).\n\n[9] Cohen, Lior, et al. \"Improving Token-Based World Models with Parallel Observation Prediction.\" ICML (2024).\n\n[10] Jan Robine, et al. \"Transformer-based World Models Are Happy With 100k Interactions.\" ICLR (2023).\n\n[11] Lukasz Kaiser, et al. \"Model-Based Reinforcement Learning for Atari.\" ICLR (2020).\n\n[12] Weipu Zhang, et al. \"STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning.\" NeurIPS (2023).\n\n[13] https://github.com/google-research/rliable", "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": "RKaMZZeskl", "reviewer": "Reviewer_earm", "summary": "This paper introduces a new model-based reinforcement learning method, called PaMoRL, which parallelize the world model training (PWM) and the eligibility trace estimation (PETE). Specifically, the parallelization is achieved by leveraging efficient parallel scan operations. On the commonly used Atari 100K and DM Control benchmarks, PaMoRL achieves competitive performance while enjoying significant improvements in training efficiency.", "questions": "Please see weaknesses above.", "limitations": "Yes, the limitations have been discussed in Section 5.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The writing is generally clear and easy to follow.\n* Accelerating eligibility trace computation with parallel scan is a novel design, which can also benefits model-free RL methods.\n* The experiments are comprehensive and solid.", "weaknesses": "* The core idea of this paper is about parallel scans. However, similar techniques have been used in Mamba [37] to accelerate computation. I noticed the authors cited Mamba but they didn't give much discussions on the connections.\n* What's the benefit of the proposed modified linear attention compared to SSM with scan? Appendix B mentioned linear attention is more expressive. Is there any quantitative evidence to show the benefit?\n* Could the authors report the results using metrics recommended in [i]? This will give a more reliable comparison between different methods.\n* A recent work [ii] also studies how to make world model training fast. How does PaMoRL compare to it?\n* How does the GPU overhead brought by scanners scale with different configurations (say, sequence length or other parameters)? Would the overhead become unacceptable when we change the setting or it is always less than say, X%, of the total GPU memory used?\n\n[i] https://github.com/google-research/rliable  \n[ii] Cohen, Lior, et al. \"Improving Token-Based World Models with Parallel Observation Prediction.\" arXiv preprint arXiv:2402.05643 (2024).\n\nMinor issues:\n* Labels and ticks in Figure 1 can be made larger to improve readability (same for Figure 3 (b,c) and the result figures in Appendix).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for reviewing our paper and for your many detailed comments. The following are responses to the weaknesses and questions you listed:\n\n### W1: More discussion of the connection to Mamba[1].\n\n**R1**: Both our PaMoRL method and Mamba use parallel scanning algorithms for acceleration, and both have data-dependent decay rates (i.e., the selective mechanism described in Mamba). However, Mamba uses a special IO-aware parallel scanning algorithm for efficient training, which focuses on reducing the number of reads and writes between SRAM and HBM in the GPU through kernel fusion, i.e., carefully dividing the parameters that need to be saved and those that need to be recomputed, and the method is suitable for targeting the hardware characteristics to improve the training efficiency in the presence of a defined model architecture The method is suitable for improving the training efficiency for hardware characteristics when the model architecture is determined. Our PaMoRL method uses a generalized parallel scanning algorithm that does not require any manual partitioning of parameters and is therefore applicable to a wide range of model architectures. We thank you for your suggestions and will add more discussion of the connection to the Mamba in the \"Background\" section of the revised version of our paper.\n\n### W2: What are the benefits of improved linear attention compared to SSM with scanning? Is there any quantitative evidence that this linear attention is more expressive?\n\n**R2**: Compared to SSM with scanning, improved linear attention has a data-dependent decay rate (i.e. gating mechanism), which allows the world model to filter out irrelevant information and remember relevant information indefinitely. We added SSM training curves to the ablation experiments in Figure 3 in the PDF attached to \"Common Response\", and the results show that Improved Linear Attention outperforms SSM on multiple tasks. Notably, DreamerV3 [2], which also has a gating mechanism, outperforms SSM in most environments, again illustrating the importance of data-dependent decay rates.\n\n### W3: Can the authors report results using the metrics recommended in [3]?\n\n**R3**: We appreciate your suggestion, and in Figure 2 in the attached PDF from \"Common Response\" we use the metrics recommended in [3] to report the results.\n\n### W4: Comparison between a recent work [4] and PaMoRL?\n\n**R4**: Recent work [4] focuses on how to accelerate the Token-based MBRL algorithm, where the world model is accelerated using the RetNet [5] architecture for accelerated training and also accelerates the process of predicting the next observation for the world model using block parallelization. In contrast, PaMoRL does not need to predict multiple fine-grained tokens, its world model architecture increases the data-dependent decay rate compared to RetNet, and it can use a recurrent pattern with minimal computational overhead when predicting the next observation. We have added comparisons to [4] in both Figure 1 and Figure 2 in the PDF attached to \"Common Response\", and our experiments conclude that our PaMoRL method matches and exceeds both the mean, median, and interquartile mean (IQM) human-normalized scores and the optimality gap both match and exceed the method proposed in [4], and also have higher training speed since PaMoRL does not need to predict multiple tokens.\n\n### W5: How does the GPU overhead from parallel scanning scale with different configurations? Does the overhead become unacceptable when changing settings?\n\n**R5**: We have added experiments on GPU memory usage vs. wall clock time as sequence length varies in **Figure 4 (Right)** in the PDF attached to \"Common Response\" (based on an RTX 3090 GPU), and the results show that parallel scanning incurs additional GPU memory usage as sequence length increases, however, at a maximum sequence length of 1024 compared to a minimum of 16, there is only less than 14% (~3.36GB) additional GPU memory overhead.\n\n### W6: The labels and markers in Figure 1 could be made larger to improve readability (as are the result plots in Figure 3 (b, c) and the Appendix).\n\n**R6**: We appreciate your suggestions. We will increase the font size of Figures 1 and Figures 3 in the revised version of the paper. You can find the corrected versions of Figures 1 and Figure 3 in the PDF attached to \"Common Response\".\n\n[1] Albert Gu, et al. \"Mamba: Linear-Time Sequence Modeling with Selective State Spaces.\" arXiv preprint arXiv:2312.00752 (2023).\n\n[2] Danijar Hafner, et al. \"Mastering diverse domains through world models.\" arXiv preprint arXiv:2301.04104 (2023).\n\n[3] https://github.com/google-research/rliable\n\n[4] Cohen, Lior, et al. \"Improving Token-Based World Models with Parallel Observation Prediction.\" ICML (2024).\n\n[5] Yutao Sun, et al. \"Retentive Network: A Successor to Transformer for Large Language Models.\" arXiv preprint arXiv:2307.08621 (2023).", "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:
## Common Response

We thank all reviewers for their valuable feedback, reviewers (*MVJ1*, *wbEe*, *earm*) for recognizing the efficiency and novelty of PaMoRL, and reviewers (*MVJ1*, *5drW*, *earm*) for the promising results and comprehensiveness of the paper's experiments. We summarize the main updates and some frequently asked questions below:

### Why do the experimental results primarily measure data efficiency rather than hardware efficiency?

We believe that the core of determining whether an MBRL method can maximize hardware efficiency should be whether it contains operators that do not allow parallelization over the sequence length. The operators that make up the world model and eligibility trace estimation of our method fully satisfy the three conditions mentioned in the "Parallel Scan" paragraph in the "Background" section of our paper (L95-L97), and thus our method can maximize hardware efficiency through the parallel scan algorithm. In addition, since the observation and action spaces of individual tasks in each benchmark do not differ much in dimension and our method does not need to adjust different hyperparameters for different tasks in the same benchmark, this means that our PaMoRL method is computationally highly consistent across tasks, and therefore we do not think that setting up too many experiments to prove the hardware efficiency of our method is necessary.

On the other hand, the MBRL method has the key selling point of its inherent high data efficiency and most of the previous works [1-3] that concluded that their methods are highly data-efficient stem from the superior performance in multiple tasks that measure data efficiency. We, therefore, consider it necessary to set up various experiments measuring data efficiency to empirically demonstrate that our PaMoRL method can achieve MBRL-level data efficiency.

### Additions to the ablation studies.

We accepted the suggestions of *Reviewer earm* and *Reviewer 5drW* to add two additional random seeds, to achieve alignment with experimental setups in previous well-known work [1-3], to compare different methods using the metrics recommended in [4] to report results, and to add comparisons with recent work [5]. Our PaMoRL matches or exceeds the baseline on mean, median, interquartile mean (IQM) human-normalized scores and optimality gap.
We also accepted *Reviewer MVJ1*'s and *Reviewer earm*'s suggestions to increase the wall-clock time and GPU memory usage (based on an RTX 3090 GPU) for our PWM, SSM (equivalent to PWM without the data-dependent decay rate), and DreamerV3[1] trained under different sequence length settings. The experimental results demonstrate that the computational gain of our PaMoRL method on wall-clock time increases with sequence length, while at the same time, there is less than 14% GPU memory overhead (~3.36 GB) at the maximum sequence length of 1024 compared to the minimum sequence length of 16, which represents the superior hardware efficiency of our PaMoRL method.

In addition, we also selected "*Alien*", "*Boxing*" and "*MsPacman*", which are tasks focusing on sequence prediction, and "*Amidar*", "*UpNDown*", and "*Qbert*", which focus on observation dispersion, for RMSNorm, Token Mixing, and Gating modules, and compare them with DreamerV3. The experimental results show that these modules all play a key role in improving data efficiency.

### Clarification of Figure 1.

We apologize for the incorrect title of Figure 1 due to a plotting error. Figure 1 in our attached PDF puts a corrected figure with an increased font size. We have accepted *Reviewer 5drW*'s suggestion to detail the GPUs used for each baseline training, and have elaborated on the details of the method for extrapolating the training speeds of the different GPU models to the NVIDIA V100 GPUs, which is consistent with the setups used in DreamerV3, where it is assumed that the P100 is twice as fast as the P100, and the A100 is twice as fast as it is.

[1] Danijar Hafner, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).

[2] Vincent Micheli, et al. "Transformers are sample efficient world models." ICLR (2023).

[3] Max Schwarzer, et al. "Data-Efficient Reinforcement Learning with Self-Predictive Representations." ICLR (2021).

[4] https://github.com/google-research/rliable

[5] Cohen, Lior, et al. "Improving Token-Based World Models with Parallel Observation Prediction." ICML (2024).

Meta-review: This paper initially did not convince all Reviewers, who expressed concerns about its novelty and empirical evaluation. The discussion period has been useful to clarify these concerns and resulted to some Reviewers' increasing their scores, resulting to unanimously positive reviews. Remarkably, the speed-up provided by the presented method w.r.t. other baselines has been praised by the Reviewers, making this a significant contribution.

I encourage the Authors to incorporate Reviewers' feedback in the final version.
Decision: Accept (poster)