Title: Parallelizing Model-based Reinforcement Learning Over the Sequence Length
Venue: NeurIPS 2024
Accepted: yes
Keywords: Model-based reinforcement learning

world model

parallelization
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

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 MBRLlevel 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.

1 Introduction

Model-based Reinforcement Learning (MBRL) is widely believed to have the great potential to substantially enhance sample efficiency by training a policy through a learned world model [1, 2, 3]. Previous studies [4, 5, 3, 6] achieve the same asymptotic performance as their model-free counterparts while requiring orders of magnitude less interactions. In particular, some recent works have even achieved human-level efficiency in complex RL domains like Atari [7, 8, 9] and robot control [10, 11].

MBRL methods can be generally divided into two stages: model learning and policy learning. During the model learning stage, a parameterized world model is required to predict the environmental dynamics by constructing specific self-supervised learning tasks. The policy learning stage benefits from synthetic interactions between the policy and the world model, hence on-policy actor-critic methods or planning methods such as Model Predictive Path Integral (MPPI) [12, 13] or Monte-Carlo Tree Search (MCTS) [7, 9] 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 returns [2, 14]. However, these powerful techniques often come with additional computations,

Junfeng Long

Atari100KBenchmark

%

140

575

Score

PaMoRL(Ours)

620-

PaMoRL(Ours)

NormalizedS

120

IRIS

REM

STORM

SR-SPR

Score

600

Score

550

525

PaMoRL(Ours)

100

TWM

580

Human

60

80

SPR

AverageS

560

540

SAC.

Average

500

475

450

CURL

Average

40

SimPLe

520-

425

Model-based

Model-free

DrQ-v2

500

400

0

5

10

15

20

25

30

0

50

100

150

200

0

10

20

30

40

50

60

Training FPS onV100 GPU

Training FPSonV100 GPU

Training FPS onV100 GPU

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 as 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 a potential successor, garnering acclaim for their remarkable performance in language modeling tasks and parallelized training paradigm [18]. Several attempts have been made to incorporate transformers into world models [19, 20, 21, 22]. 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.

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 in 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 [29, 30]. Motivated by recent works in efficient sequential modeling [31, 32, 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 eligibility trace estimation [2, 14] can be naturally parallelized over the sequence length by using parallel scan.

To this end, we introduce the Pa rallelized Mo del-based R einforcement L earning ( PaMoRL ) framework, which consists of two novel techniques as shown in Figure 2 that can parallelize the current MBRLparadigm over sequence length: (1) the P arallel W orld M odel ( PWM ) and (2) the P arallelized E ligibility T race E stimation ( PETE ). The resulting framework, PaMoRL, is hardware-efficient and 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.

We evaluated our PaMoRL framework in the Atari 100K benchmark [16] and the DeepMind Control 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 1. The empirical results demonstrate that PaMoRL, despite being a framework that incorporates autoencoding, still benefits greatly from the implementation of dual parallelization techniques (i.e., PWMand 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, PaMoRL outperforms other no-look-ahead MBRL methods and model-free RL methods. It is worth mentioning that PaMoRL even outperforms the planning-based MBRL methods or methods with much larger networks in certain tasks [8].

Our contributions can be summarized as follows:

We introduce PaMoRL, a 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λ [2], Retrace [35] 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 t , scalar rewards r t , actions a t , continuation flag c t ∈ { 0 , 1 } , discount factor γ ∈ (0 , 1) , and environmental dynamics o t , r t , c t ∼ p ( o t , r t , c t | o <t , a <t ) . The objective of the Reinforcement Learning (RL) is to train a policy π that maximizes the return ∑ ∞ t =1 γ t -1 r t . In Model-based Reinforcement Learning (MBRL), the RL agent learns a model of the environmental dynamics through an iterative process that involves collecting data using a policy, training a model of the environment based on the accumulated data, and optimizing the policy using the learned model [1, 2, 14].

Parallel Scan. As a universal parallel algorithm building block, the computations of parallel scan involve repeated application of a binary operator ⊕ over sequential data arrays. Previous work[37] describes scan as a good example of a computation that seems inherently sequential, but for which there is an efficient parallel algorithm. The scan of ⊕ with initial value a 0 is defined in Equation 1.

First-order linear recurrences h t := ( A t ⊗ h t -1 ) ⊕ x t can be parallelized over the sequence length with the utilization of parallel scans if the following three conditions are met:

⊕ is associative: ( a ⊕ b ) ⊕ c = a ⊕ ( b ⊕ c ) .

⊗ is semi-associative: there exists a binary associative operator ⊙ such that a ⊗ ( b ⊗ c ) = ( a ⊙ b ) ⊗ c .

⊗ distributes over ⊕ : a ⊗ ( b ⊕ c ) = ( a ⊗ b ) ⊕ ( a ⊗ c ) .

We observe vector addition a ⊕ b := a + b , matrix-vector multiplication A ⊗ b := A · b , and matrix-matrix multiplication A ⊙ B := A · B fulfill the aforementioned conditions. This allows the parallel computation of x t := ( A t · x t -1 ) + b t across time steps t , considering input vectors b t and square matrices A t . Considering the operators required in computing linear attentions [31] and eligibility trace estimations [2, 14] involve only diagonal matrices, the linear recurrence can be re-formulated as x t := λ t ⊙ x t -1 + b t , where λ t is the eigenvalues of the diagonal matrices and ⊙ is an element-wise multiplication.

3 Methodology

We introduce our Pa rallelized Mo del-based R einforcement L earning ( 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 2 illustrates 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 MBRL methods, our world model is trained to predict environmental dynamics. Since observations can be high-dimensional (e.g., images), we

Deterministic Nodes

Stochastic Nodes

Concatenation

C

World Model Learning

(Parallel)

𝑧

!"#

𝑎

!"#

𝑥

!"#

ℎ

!"#

̂ 𝑟

!"#

̂ 𝑐

!"#

𝑜

!"#

" 𝑜

!"#

C

Enc

Dec

̂ 𝑧

!

𝑥

$

𝑧

!

𝑎

!

𝑥

!

ℎ

!

̂ 𝑟

!

̂ 𝑐

!

𝑜

!

" 𝑜

!

C

Enc

Dec

̂ 𝑧

!%#

…

𝑧

&

𝑎

&

𝑥

&

ℎ

&

̂ 𝑟

&

̂ 𝑐

&

𝑜

&

" 𝑜

&

C

Enc

Dec

̂ 𝑧

&%#

Parallel

Scan

Imagination

(Recurrent)

𝑧

!

𝑎

!

𝑥

!

ℎ

!

̂ 𝑟

!

̂ 𝑐

!

𝑜

!

C

Enc

̂ 𝑧

!%#

𝑎

!%#

𝑥

!%#

ℎ

!%#

̂ 𝑟

!%#

̂ 𝑐

!%#

C

̂ 𝑧

'

𝑎

'

𝑥

'

ℎ

'

̂ 𝑟

'

̂ 𝑐

'

C

…

…

Eligibility Trace Estimation

(Parallel)

ℎ

!

̂ 𝑧

!%#

C

ℎ

!%#

𝑅

(

!

𝑅

(

!%#

̂ 𝑐

!%#

̂ 𝑟

!%#

) 𝑣

!%)

+

𝛿

!%#

𝜆

!%#

…

…

̂ 𝑐

!

̂ 𝑟

!

) 𝑣

!%#

+

𝛿

!

𝜆

!

Parallel

Scan

̂ 𝑧

!%)

C

…

…

𝑒

!

𝑒

"

…

𝑒

!

𝑒

"

…

𝑒

!

𝑒

"

…

𝑒

!

𝑒

"

…

𝑒

!

𝑒

"

…

𝑒

!

𝑒

"

…

𝑒

!

𝑒

"

…

predict future representations rather than future observations. This reduces accumulating errors and enables massively parallel training with a large batch size. The compact representations are 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, we make several modifications to the vanilla Recurrent State-Space Model's (RSSM) [3, 6, 17] configurations: (1) differentiating the hidden states x t from the sequential model's outputs h t , (2) excluding h t from the inputs of the encoder and decoder, (3) eliminating the stochastic states z t 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:

z t ∼ q θ ( z t | o t )

Decoder:

ˆ o t ∼ p θ (ˆ o t | z t )

Sequence model:

h t , x t = f θ ( x t - 1 , z t - 1 , a t - 1 )

Dynamics predictor: ˆ z t ∼ p θ (ˆ z t | h t ) (2)

Reward predictor:

ˆ r t ∼ p θ (ˆ r t | h t )

Continue predictor:

ˆ c t ∼ p θ (ˆ c t | h t )

The encoder and decoder use 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 of a 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, 17], we set the q θ ( z t | o t ) as a stochastic distribution comprising 32 categories, each with 32 classes, and we take straight-through gradients through the sampling step [39].

Sequence Model Architectures. As mentioned above, each residual block of our sequence model consists of a modified linear attention module and a GLU module. The vanilla linear attention module, as introduced in previous work [31], employs 1 + ELU as an element-wise kernel function applied to queries q t and keys k t taking u t as input. This configuration allows for its reformulation into an RNN-style recurrent form. However, this version of linear attention is prone to unstable convergence during training due to the unbounded gradients [40]. Thus, we remove the time-dependent normalizer, which is designed to 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 u t and previous inputs u t -1 , along with the gating mechanism in Gated Recurrent Unit (GRU) [42] to provide an input-dependent decay rate g t for hidden state x t . The subsequent GLU module selects SiLU as the gating function, taking linear attention output y t as input. By integrating all the modifications, we can derive the entire block of the sequence model as shown in Equation 3.

The architecture of our modified linear attention satisfies the conditions in Section 2 and can be effectively computed using parallel scans. We can refer to Table 1 to 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 β pred , β rep, and β dyn are coefficients to adjust the influence of each term in the loss function [43, 17].

The operation sg ( · ) represents the stop gradient operation. The KL divergences are derived from the Evidence Lower Bound (ELBO). We clip the KL divergence when it falls below the threshold of 1 [6, 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 z t and h t as input state s t .

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λ [2] as the its target, as shown in Equation 6, where ˆ r t represents the reward predicted by the world model, and ˆ c t represents the predicted continuation flag.

The actor utilizes the Reinforce estimator [44] to compute the actor loss with a fixed entropy regularization term. The complete loss is described by Equation 7.

The hyper-parameter η represents the coefficient of the entropy regularization term. The normalization ratio S utilized in the actor loss is defined in Equation 8, which is computed as the range between the 95 th and 5 th percentiles of the TDλ returns R λ t across the batch.

By rearranging Equation 6, we can see that the calculations of both TDλ and Retrace returns also also meet the conditions mentioned in Section 2. Therefore, they can be efficiently computed using parallel scan. This observation also applies to other eligibility trace estimation methods such as GAE [36] and Retrace [35], as 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 [45] and the Odd-even scanner [28].

The Kogge-stone scanner [45] is commonly used in hardware design for adders. It has a computational complexity of O ( L log 2 L ) for sequence length L and a step complexity of O ( log 2 L ) 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] is based on the concept of binary balanced trees. It has a computational complexity of O (2 L ) for a sequence length L and a step complexity of O (2 log 2 L ) 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.

4 Experiments

In this section, we aim to evaluate both the sample and training efficiency of our 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 to 400K actual game frames, corresponding to approximately 2 hours of real-time gameplay, with action repeats of 4. The human normalized score is defined as ( scoreagent -scorerandom ) / ( scorehuman -scorerandom ) , where scorerandom comes from a random policy, and scorehuman is obtained from human players [48].

DeepMind Control Suite. DeepMind Control Suite consists of various control tasks with continuous 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 of EfficientZero V2 [9] 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

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], DrQv2 [51], and DreamerV3 [17].

4.2 Experimental Results

In this section, we do not compare our results with look-ahead search methods [52, 7, 12, 13] or methods using larger networks [8], as our main goal in terms of sample efficiency is to improve performance while maximizing the hardware efficiency of existing MBRL methods.

Atari 100K. The summarized results are shown in Figure 4. The full results for individual games in the Atari 100k benchmark are elaborated in Table 2, where scores are 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 I.

DeepMind Control Suite. Table 3 shows that our method achieves a mean score of 661.2 across 16 tasks. As shown in Table 3, our method achieves a mean score of 661.2 using proprioception observations and 538.7 using image observations across 16 tasks, surpassing the previous state-ofthe-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.

Amidar

UpNDown

6000

250

一

10000

5000

200

8000

4000

Score

150

Score

6000

3000

100

4000

2000

50

2000

1000

0

0

0

100

200

300

400K

0

100

200

300

400K

0

100

200

300

400K

Alien

Boxing

MsPacman

1500

100

2500

1250

80

2000

Score

1000

Score

60

Score

750

40

1500

500

20

1000

250

0

009

100

200

300

400K

0

100

200

300

400K

0

100

200

300

400K

PWM

PWMw/oRMSNorm

PWM w/oToken Mixing

SSMw/Scan

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 3 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 Qbert .

Eligibility Trace Estimation Runtime (ms)

1.2

140

1400-

120

1.0

1200

100

80

1000

80

90

008

60

009

t0

40

400

20

0.2

200

0

0.0

Rollout

Rollout

Rollout

F009

30

Method

Method

PaMoRL

Wall-clock Time (Minute)

PaMoRL

25

SSMw/Scan

500

SSMw/Scan

20

400

15

300

200

0

64/16

32/32

16/64

8/128

4/256

2/512

1/1024

64/16

32/32

16/64

8/128

4/256

2/512

1/1024

BatchSize/SequenceLength

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 a single frame (e.g., Boxing ), it 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 it 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 output to maximize the information's retention.

Parallel Scanner Selection. Figure 5 shows PWM and PETE's runtime and GPU memory utilization on a single 3090 GPU using different scanners, respectively. Sequence model computation

achieves 7 . 2 × and 16 . 6 × speedups compared to sequential rollout using the Kogge-stone and Oddeven scanners, respectively, with a 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 3 × and 2 × speedups, respectively, with a sequence length of 16. Since the eligibility trace has a 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 6 × overhead compared to the sequential rollout, while the Odd-even scanner imposes an additional 2 × 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 to 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 DreamerV3 does not. This is likely due to 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 of 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 or tasks in the Atari 100K benchmark and DeepMind Control suite without incurring additional overhead during inference. An important contribution of our work is the introduction of a 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-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. It 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-end training paradigm enables the world model to predict the latent representations, it also impacts the scalability of the world model.

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-2024-00170).

References

Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin , 2(4):160-163, 1991.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. Robotica , 17(2):229-235, 1999.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 , 2019.

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.

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.

Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193 , 2020.

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.

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.

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.

Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Daydreamer: World models for physical robot learning. In Conference on Robot Learning , pages 2226-2240. PMLR, 2023.

Russell Mendonca, Shikhar Bahl, and Deepak Pathak. Structured world models from human videos. arXiv preprint arXiv:2308.10901 , 2023.

Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.04955 , 2022.

Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpc2: Scalable, robust world models for continuous control. arXiv preprint arXiv:2310.16828 , 2023.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122 , 2018.

Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Modelbased reinforcement learning for atari. arXiv preprint arXiv:1903.00374 , 2019.

Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09481 , 2022.

Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample efficient world models. arXiv preprint arXiv:2209.00588 , 2022.

Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. arXiv preprint arXiv:2303.07109 , 2023.

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.

Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems , 5, 2023.

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.00690 , 2018.

Lior Cohen, Kaixin Wang, Bingyi Kang, and Shie Mannor. Improving token-based world models with parallel observation prediction. arXiv preprint arXiv:2402.05643 , 2024.

Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. arXiv preprint arXiv:2309.12252 , 2023.

Richard E Ladner and Michael J Fischer. Parallel prefix computation. Journal of the ACM (JACM) , 27(4):831-838, 1980.

Mark Harris, Shubhabrata Sengupta, and John D Owens. Parallel prefix sum (scan) with cuda. GPU gems , 3(39):851-876, 2007.

Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. arXiv preprint arXiv:1709.04057 , 2017.

Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933 , 2022.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning , pages 5156-5165. PMLR, 2020.

Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305.13048 , 2023.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 , 2023.

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.

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.

John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. Highdimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438 , 2015.

Guy E Blelloch. Prefix sums and their applications. 1990.

Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 , 2020.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

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.

Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems , 32, 2019.

Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555 , 2014.

Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International conference on learning representations , 2016.

Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems , 12, 1999.

Peter M Kogge and Harold S Stone. A parallel algorithm for the efficient solution of a general class of recurrence equations. IEEE transactions on computers , 100(8):786-793, 1973.

Fei Deng, Junyeong Park, and Sungjin Ahn. Facing off world model backbones: Rnns, transformers, and s4. arXiv preprint arXiv:2307.02064 , 2023.

Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Mastering memory tasks with world models. In Second Agent Learning in Open-Endedness Workshop , 2023.

Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference on machine learning , pages 1995-2003. PMLR, 2016.

Thomas Hubert, Julian Schrittwieser, Ioannis 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.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645 , 2021.

Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588(7839):604-609, 2020.

Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. arXiv preprint arXiv:1511.06349 , 2015.

Carlos E Garcia, David M Prett, and Manfred Morari. Model predictive control: Theory and practice-a survey. Automatica , 25(3):335-348, 1989.

Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint arXiv:1807.03858 , 2018.

Jian Shen, Han Zhao, Weinan Zhang, and Yong Yu. Model-based policy optimization with unsupervised model adaptation. Advances in Neural Information Processing Systems , 33:28232834, 2020.

Marc Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11) , pages 465-472, 2011.

Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sampleefficient reinforcement learning with stochastic ensemble value expansion. Advances in neural information processing systems , 31, 2018.

Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint arXiv:1803.00101 , 2018.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 , 2021.

Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:2011.04006 , 2020.

Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.

Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arXiv:2209.12951 , 2022.

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.

Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation , 4(1):131-139, 1992.

Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143 , 2021.

Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. arXiv preprint arXiv:2210.04243 , 2022.

Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. Advances in Neural Information Processing Systems , 36, 2024.

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. ICLR , 1:2, 2020.

A Related Work

Model-based reinforcement learning (MBRL) methods aim to 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, 4, 3], generating synthetic data [1, 55, 3, 5, 56], or improving the value estimation [57, 4, 58, 59]. Our PaMoRL framework builds upon DreamerV3 [17], which embraces the Dyna paradigm [1, 2, 14]. The world model is utilized to interact with the policy to generate synthetic data, aiding model-free RL methods in maximizing cumulative task rewards during the policy learning stage.

To maximize sample efficiency, MBRL researchers MBRL researchers have attempted to employ Recurrent Neural Networks (RNNs) [15, 16, 3, 6, 17] or Transformers [19, 21, 20, 22] 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-performance option, there is a quadratic relationship between their computational complexity and sequence length, which introduces additional computational and memory overhead.

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 Transformers and RNNs in Long-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 classical 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 a close relationship [31], i.e. 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 to 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 PaMoRL 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 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, 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 2, as compared to the model architecture-specific parallel scanning method used by Mamba.

B Illustrations to Parallel Scan Algorithms

B.1 Kogge-stone Scanner

A common example of such a first-order recurrence problem is a time-varying linear system, the system's state at timestep t is x t , computed from the system's internal dynamical variables a t and b t , as shown in Equation 9. Depending on the problem, the variables a t and b t can be real or complex numbers, constants, etc.

Before solving the problem, we can define the function A ( m,n ) and B ( m,n ) , as shown in Equation 12.

Now we can integrate Equation 9 with Equation 10 to get Equation 11.

It can be observed in Equation 11 that B (2 i, i +1) is associated with the computation of A (2 i, i +1) but independent from B ( i, 1) , which means that we can split the computation of B (2 i, 1) into two parallel parts. For reasons of notational simplicity, we define the tuple Q ( m,n ) that wraps the functions A ( m,n ) and B ( m,n ) , as shown in Equation 12.

Figure 6 shows the operation of the Kogge-stone scanner [45] when the sequence length L = 8 . After ⌈ log 2 L ⌉ iterations the solution to the problem x 1 , . . . , x T can be computed.

B.2 Odd-even Scanner

To avoid the extra computational complexity of log 2 L generated by the Kogge-stone scanner [45], the Odd-even scanner [28] uses an algorithmic pattern that arises often in parallel computing: balanced

trees. The idea is to build a balanced binary tree on the input data and start scanning from the root. A binary tree with L leaves has log 2 L layers with 2 d nodes per layer d ∈ [0 , L ) . If we perform one operation on each node, then we will perform O ( L ) operations in one traversal of the tree. The tree we construct is not an actual data structure, but rather a concept that we use to 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, we 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 in the up-sweep phase. Figure 7 shows the operation of the Kogge-stone scanner [28] when the sequence length L = 8 .

Note that since this is an exclusive scan (i.e., the sum is not included in the result), we zero out the 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 O (2 L ) operations, so it is very efficient.

Q(8,8)

Q(7,7)

Q(6,6)

Q(5,5)

Q(4,4)

Q(3,3)

Q(2,2)

Q(1,1)

Step 1

Q(8,7)

Q(7,6)

Q(6,5)

Q(5,4)

Q(4,3)

Q(3,2)

Q(2,1)

Q(1,1)

Step 2

Q(8,5)

Q(7,4)

Q(6,3)

Q(5,2)

Q(4,1)

Q(3,1)

Q(2,1)

Q(1,1)

Step 3

Q(6,1)

Q(5,1)

Q(4,1)

Q(3,1)

Q(2,1)

Q(1,1)

Q(8,8)

Q(7,7)

Q(6,6)

Q(5,5)

Q(4,4)

Q(3,3)

Q(2,2)

Q(1,1)

Step 1

Q(8,7)

Q(7,7)

Q(6,5)

Q(5,5)

Q(4,3)

Q(3,3)

Q(2,1)

Q(1,1)

Step 2

Q(8,5)

Q(7,7)

Q(6,5)

Q(5,5)

Q(4,1)

Q(3,3)

Q(2,1)

Q(1,1)

Step 3

Q(7,7)

Q(6,5)

Q(5,5)

Q(4,1)

Q(3,3)

Q(2,1)

Q(1,1)

Zero

Step 4

0

Q(7,7)

Q(6,5)

Q(5,5)

Q(4,1)

Q(3,3)

Q(2,1)

Q(1,1)

Step 5

Q(4,1)

Q(7,7)

Q(6,5)

Q(5,5)

0

Q(3,3)

Q(2,1)

Q(1,1)

Step 6

Q(6,1)

Q(7,7)

Q(4,1)

Q(5,5)

Q(2,1)

Q(3,3)

0

Q(1,1)

Step 7

Q(6,1)

Q(5,1)

Q(4,1)

Q(3,1)

Q(2,1)

Q(1,1)

0

C Training Curves of the Atari 100K Benchmark

The Atari 100K benchmark [69] is a 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 to executing 400 thousand discrete actions within each environment, which is approximately equivalent to 2 hours of human gameplay. To put this in perspective, when there are no constraints on sample efficiency, the typical practice is to train agents for 200M steps.

Alien

Amidar

Assault

Asterix

1600

1200

4000

1400

300

3500

1200

250

1000

3000

800

150

200

800

2500

2000

600

100

009

1500

200

400

50

200

1000

500

100

200

300

400K

100

200

300

400K

100

200

300

400 K

100

200

300

400K

BankHeist

BattleZone

Boxing

Breakout

25000

250

s

200

150

20000

15000

12.5

10.0

15.0

10000

7.5

5.0

50

5000

2.5

100

200

300

400K

100

200

300

400 K

100

200

300

400 K

100

200

300

400K

ChopperCommand

CrazyClimber

DemonAttack

Freeway

8

2500

80000

250

2000

60000

200

25

20

1500

40000

150

1000

20000

100

500

50

100

200

300

400K

100

200

300

400K

100

200

300

400 K

100

200

300

400 K

Frostbite

Gopher

Hero

Jamesbond

6000

12000

14000

500

3000

5000

10000

400

2000

4000

8000

300

3000

6000

200

1000

2000

4000

1000

2000

100

0

100

200

300

400K

100

200

00C

400K

0

100

200

300

400 K

100

200

300

400K

Kangaroo

Krull

KungFuMaster

MsPacman

4000

8000

25000

2500

7000

3000

6000

20000

2000

2000

5000

15000

1500

4000

10000

1000

1000

3000

5000

500

2000

100

200

300

400K

100

200

00E

400K

100

200

300

400 K

100

200

300

400K

Pong

PrivateEye

RoadRunner

5000

0009

25000

4000

4000

20000

0000

3000

15000

2000

2000

10000

-10 -

1000

1000

5000

20

100

200

300

400K

-1000

100

200

300

400K

100

200

300

400 K

100

200

300

400K

Seaquest

UpNDown

700

14000

600

12000

500

10000

8000

300

6000

200

4000

100

2000

100

200

300

400K

0

0

100

200

300

400K

D Training Curves of the DeepMind Control Suite

DeepMind Control suite [24] is a 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 is to train agents for millions of steps.

cartpole_balance

cartpole_balance_sparse

cartpole_swingup

cup_catch

1000

1000

300

1000

800

800

250

800

600

200

600

009

400

150

400

400

200

100

200

20

40

60

80

100K

0

20

40

60

80

100K

50

0

20

40

60

80

100K

D

40

60

80

100K

finger_spin

pendulum_swingup

reacher_easy

reacher_hard

60

1000

008

009

800

50

40

800

600

400

30

600

400

20

400

200

10

200

200

0

0

0

20

40

60

80

100K

20

40

60

80

100K

20

40

60

80

100K

20

40

60

80

100K

cartpole_swingup_sparse

cheetah_run

finger_turn_easy

finger_turn_hard

1000

1000

300

800

800

800

200

600

600

600

400

400

400

100

200

200

200

0

0

50

100

150

50

100

150

200K

0

50

100

150

200K

100

150

200K

hopper_stand

hopper_hop

quadruped_run

quadruped_walk

600

200

400

500

150

400

300

400

100

300

200

200

200

100

100

0

50

50

100

100

150

200K

100

150

200K

50

150

200K

50

100

150

200K

cartpole_balance

cartpole_balance_sparse

cartpole_swingup

cup_catch

1000

1000

900

800

800

800

700

600

600

600

400

400

500

300

400

200

200

200

00

20

30

40

20

20

30

finger_spin

pendulum_swingup

reacher_easy

reacher_hard

1000

1000

800

800

800

800

600

600

400

400

00

200

200

30

20

20

30

40

cartpole_swingup_sparse

cheetah_run

finger_turn_easy

finger_turn_hard

400

800

700

350

600

300

600

500

250

600

200

600

300

0SL

400

200

100

200

100

50

60

60

40

09

hopper_stand

hopper_hop

quadruped_run

quadruped_walk

400

700

350

200

000

300

250

150

500

400

150

100

200

300

100

50

100

200

100

E Details of Model Architecture

F Hyperparameters

G Pytorch-style Pseudo-code of Parallel Scan

G.1 Odd-even scanner

G.2 Kogge-stone scanner

H Pytorch-style Pseudo-code of Parallelized Eligibility Trace Estimation

I Additional Comparisons on the Atari 100K Benchmark

Atari 100K benchmark [69] is a 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 to executing 400 thousand discrete actions within each environment, which is approximately equivalent to 2 hours of human gameplay. To put this in perspective, when there are no constraints on sample efficiency, the typical practice is to train agents for 200M steps.

In this section, we compare the performance of our PaMoRL framework with planning-based methods such as EfficientZero [7] and EfficientZero V2 [9] and methods with much larger networks, i.e., BBF [8] on the Atari 100K benchmark. The full results are shown in Table I. The PaMoRL framework is not as good as the other methods in terms of 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.

J Runtime of Experiments

K Effectiveness of Batch Normalization Trick

(a)

(b)

O

O

Breakout

Pong

30

17.5

15.0

20

12.5

10

Score

Score

10.0

0

7.5

5.0

-10

2.5

1

-20

0

100

200

300

400K

0

100

200

300

400K

PWM

PWM w/o BN

DreamerV3w/BN

L Video Predictions

Imagination

Environment:Alien

Real

Model

t=0

t=4

t=12

t=28

t=44

t=52

t=60

Imagination

Environment:Breakout

Real

Model

t=0

t=4

t=12

t=28

t=44

t=52

t=60

Imagination

Environment:Pong

Real

Model

t0

t=4

t=12

t=28

=36

t=44

t=52

t=60

Imagination

Environment:CupCatch

Real

Model

t=4

t=12

t=28

=36

=44

t=52

t=60

Imagination

Environment:HopperHop

Real

Model

t=0

t=4

t=12

t=28

t=44

t=52

t=60

Imagination

Environment:QuadrupedRun

Real

Model

t0

t=4

t=12

t=28

t=44

t=52

t=60

M Initializations in Freeway

The reward function in Freeway is sparse since the agent is 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 a 100k frames budget. The solution to this problem is actually straightforward and requires stretches of time when the "UP" action is oversampled. In this paper, we opted for the simplest strategy of having an initialized buffer with fulfilled "UP" actions. Hence, we dont'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.

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 in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

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 not include experiments.

If the paper includes experiments, a No answer to 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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so '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 to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is 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, hyperparameters, 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 a 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 different 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 are accompanied by error bars, confidence 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.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We show the hardware for training and state the time used for each training process in 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 as 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: We propose a 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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We propose a 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 a 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 to 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 we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.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 as well as the algorithms code-bases are 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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

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 not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new 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 new 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. Crowdsourcing 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 as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 not involve crowdsourcing nor research with human subjects.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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).\n\nQUESTIONS:\nPlease see weaknesses above.\n\nLIMITATIONS:\nYes, the limitations have been discussed in Section 5.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

comments: 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.

STRENGTHS:
- The paper tackles a significant problem in MBRL.
- 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.
- Analysis and theoretical derivation of the parallel scan methods are missing.
- Experiment results are reported only over three seeds, which can hardly be considered confident.

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-10

comments: 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.

STRENGTHS:
+ A significant speed-up of model generation runtime and eligibility trace estimation runtime ;
+ 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;
+ 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;
+ A clear illustration of the proposed parallel algorithm with PyTorch (pseudo-)code;
+ 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.

QUESTIONS:
NA

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nModel-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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nWhile 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).\n\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nI believe that the authors address limitations adequately in the last paragraph of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}