Research map / Reinforcement learning / Visual and representation RL
Parallelizing Model-based Reinforcement Learning Over the Sequence Length
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PaMoRL, which parallelizes both world-model learning (Parallel World Model with linear attention and parallel scans) and policy learning (Parallelized Eligibility Trace Estimation) over sequence length in DreamerV3-style MBRL, reaching strong Atari 100K and DMControl sample efficiency with faster training.
Problem
Model-based RL methods achieve high sample efficiency but incur high computation, memory, and training-time costs; RNN-based world models are sequential and transformer world models scale quadratically with sequence length.
Approach
Replace the recurrent/attention sequence model with a modified linear attention (no time-dependent normalizer, RMSNorm, RWKV-style token mixing, GRU-style data-dependent decay) that admits parallel scan training and recurrent inference, and compute TD-lambda/eligibility-trace returns with a parallel scan so both model and policy learning are parallelized over sequence length.
Key results
- PaMoRL achieves a mean human-normalized score of 126.64% and median 71.75% on Atari 100K, outperforming DreamerV3 (112%/49%) and other no-look-ahead baselines in mean and median.
- On DMControl, PaMoRL reaches a mean of 661.2 with proprioception and 538.7 with image observations, versus DreamerV3's 611.2 and 538.7 respectively (proprio/vision comparison as reported).
- Parallel scanning gives 7.2x (Kogge-stone) and 16.6x (Odd-even) speedups over sequential rollout for sequence model computation at sequence length 64.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add the token mixing module to PWM (ablation) | PaMoRL · Alien and MsPacman | performance | Helped |
| add the token mixing module to PWM (ablation) | PaMoRL · Boxing | final performance | No clear effect |
| add RMSNorm to the sequence model (ablation) | PaMoRL · Atari 100K | final performance and training stability | Helped |
| add the data-dependent decay rate (ablation) | PaMoRL sequence model · Atari 100K | performance | Mixed |
| use batch normalization in PWM (ablation) | PaMoRL · Atari 100K | world-model learning | Helped |
| use the Kogge-stone scanner for parallel sequence modeling (ablation) | PaMoRL PWM · Sequence length 64 | runtime speedup: 7.2× | Helped |
| use the Odd-even scanner for parallel sequence modeling (ablation) | PaMoRL PWM · Sequence length 64 | runtime speedup: 16.6× | Helped |
| use the Kogge-stone scanner for eligibility trace estimation (ablation) | PaMoRL PETE · Sequence length 16 | runtime speedup: 3× | Helped |
| use the Odd-even scanner for eligibility trace estimation (ablation) | PaMoRL PETE · Sequence length 16 | runtime speedup: 2× | Helped |
| use the Kogge-stone scanner instead of sequential rollout (ablation) | PaMoRL · Parallel computation | GPU memory utilization: 6× | Hurt |
| use the Odd-even scanner instead of sequential rollout (ablation) | PaMoRL · Parallel computation | GPU memory utilization: 2× | Hurt |
| use PaMoRL instead of the strongest reported baseline | Atari 100K · Atari 100K | mean human-normalized score: 118.84% -> 126.64% | Helped |
Surprising findings
- Kogge-stone scanner has the lowest theoretical step complexity but runs slower than Odd-even scanner in practice for the world model due to limited parallel resource use.
- Without Batch Normalization the decoder suffers posterior collapse in PWM, whereas DreamerV3 does not need the trick.
Limitations
- Planning-based MBRL methods cannot be parallelized over sequence length, so the most sample-efficient planning methods cannot be combined with PaMoRL.
- End-to-end training of encoder and sequence model limits world model scalability.
- Performance is below EfficientZero V2 and BBF on several metrics (superhuman count, median, mean) on Atari 100K.
Experimental setup
Datasets: Atari 100K DeepMind Control Suite
Benchmarks: Atari 100K DeepMind Control Suite
Metrics: human-normalized mean score human-normalized median score number of superhuman games DMControl mean score DMControl median score wall-clock runtime GPU memory usage
Baselines: DreamerV3 IRIS TWM STORM SPR SR-SPR SimPLe SAC DrQ-v2 CURL EfficientZero EfficientZero V2
Builds on: DreamerV3 RSSM RWKV token mixing Mamba TD-lambda
Closest papers by meaning
- SRL: Scaling Distributed Reinforcement Learning to Over Ten Thousand Cores (ICLR 2024 · similarity 0.89)
- Scalable Offline Model-Based RL with Action Chunks (ICLR 2026 · similarity 0.89)
- Sample-Efficient Automated Deep Reinforcement Learning (ICLR 2021 · similarity 0.89)
- Towards General-Purpose Model-Free Reinforcement Learning (ICLR 2025 · similarity 0.88)
- BroRL: Scaling Reinforcement Learning via Broadened Exploration (ICML 2026 · similarity 0.88)
- ESRL: Efficient Sampling-Based Reinforcement Learning for Sequence Generation (AAAI 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.