# Accelerating Model-Based Reinforcement Learning with State-Space World Models

arXiv: https://arxiv.org/abs/2502.20168 · HTML: https://mlsanity.com/paper/2502.20168/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 60 of its 69 references that are in the map.
- First posted on arXiv 27 Feb 2025.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 60)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Habitat: A Platform for Embodied AI Research](https://mlsanity.com/paper/1904.01201/) (2019)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/1806.10293/) (2018)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Learning Quadrupedal Locomotion over Challenging Terrain](https://mlsanity.com/paper/2010.11251/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Learning agile and dynamic motor skills for legged robots](https://mlsanity.com/paper/1901.08652/) (2019)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Learning by Cheating](https://mlsanity.com/paper/1912.12294/) (2019)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Sim-to-Real: Learning Agile Locomotion For Quadruped Robots](https://mlsanity.com/paper/1804.10332/) (2018)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Asymmetric Actor Critic for Image-Based Robot Learning](https://mlsanity.com/paper/1710.06542/) (2017)
- [Learning Continuous Control Policies by Stochastic Value Gradients](https://mlsanity.com/paper/1510.09142/) (2015)
- [DeXtreme: Transfer of Agile In-hand Manipulation from Simulation to Reality](https://mlsanity.com/paper/2210.13702/) (2022)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [Reinforcement and Imitation Learning for Diverse Visuomotor Skills](https://mlsanity.com/paper/1802.09564/) (2018)

## Built on by (most cited first, 10 of 10)

- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [A Survey: Learning Embodied Intelligence from Physical Simulators and World Models](https://mlsanity.com/paper/2507.00917/) (2025)
- [The Reality Gap in Robotics: Challenges, Solutions, and Best Practices](https://mlsanity.com/paper/2510.20808/) (2025)
- [VideoWorld 2: Learning Transferable Knowledge from Real-world Videos](https://mlsanity.com/paper/2602.10102/) (2026)
- [SkyJEPA: Learning Long-Horizon World Models for Zero-Shot Sim-to-Real Control of Quadrotors](https://mlsanity.com/paper/2606.23444/) (2026)
- [SkyDreamer: Interpretable End-to-End Vision-Based Drone Racing with Model-Based Reinforcement Learning](https://mlsanity.com/paper/2510.14783/) (2025)
- [Learning to Walk With Less: A Dyna-Style Approach to Quadrupedal Locomotion](https://mlsanity.com/paper/2509.06296/) (2025)
- [HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning](https://mlsanity.com/paper/2605.05951/) (2026)
- [Debiased Model-based Representations for Sample-efficient Continuous Control](https://mlsanity.com/paper/2605.11711/) (2026)
- [Belief-State RWKV for Reinforcement Learning under Partial Observability](https://mlsanity.com/paper/2604.09671/) (2026)

## Abstract (arXiv)

Reinforcement learning (RL) is a powerful approach for robot learning. However, model-free RL (MFRL) requires a large number of environment interactions to learn successful control policies. This is due to the noisy RL training updates and the complexity of robotic systems, which typically involve highly non-linear dynamics and noisy sensor signals. In contrast, model-based RL (MBRL) not only trains a policy but simultaneously learns a world model that captures the environment's dynamics and rewards. The world model can either be used for planning, for data collection, or to provide first-order policy gradients for training. Leveraging a world model significantly improves sample efficiency compared to model-free RL. However, training a world model alongside the policy increases the computational complexity, leading to longer training times that are often intractable for complex real-world scenarios. In this work, we propose a new method for accelerating model-based RL using state-space world models. Our approach leverages state-space models (SSMs) to parallelize the training of the dynamics model, which is typically the main computational bottleneck. Additionally, we propose an architecture that provides privileged information to the world model during training, which is particularly relevant for partially observable environments. We evaluate our method in several real-world agile quadrotor flight tasks, involving complex dynamics, for both fully and partially observable environments. We demonstrate a significant speedup, reducing the world model training time by up to 10 times, and the overall MBRL training time by up to 4 times. This benefit comes without compromising performance, as our method achieves similar sample efficiency and task rewards to state-of-the-art MBRL methods.