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

• URL: http://arxiv.org/abs/2502.20168v1
• Date: Thu, 27 Feb 2025 15:05:25 GMT
• Title: Accelerating Model-Based Reinforcement Learning with State-Space World Models
• Authors: Maria Krinner, Elie Aljalbout, Angel Romero, Davide Scaramuzza,
• Abstract summary: Reinforcement learning (RL) is a powerful approach for robot learning.<n>However, model-free RL (MFRL) requires a large number of environment interactions to learn successful control policies.<n>We propose a new method for accelerating model-based RL using state-space world models.
• Score: 18.71404724458449
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Echo Chamber: RL Post-training Amplifies Behaviors Learned in Pretraining [74.83412846804977]
  Reinforcement learning (RL)-based fine-tuning has become a crucial step in post-training language models. We present a systematic end-to-end study of RL fine-tuning for mathematical reasoning by training models entirely from scratch.
  arXiv  Detail & Related papers  (2025-04-10T17:15:53Z)
• Neural Motion Simulator: Pushing the Limit of World Models in Reinforcement Learning [11.762260966376125]
  A motion dynamic model is essential for efficient skill acquisition and effective planning. We introduce the neural motion simulator (MoSim), a world model that predicts the future physical state of an embodied system. MoSim achieves state-of-the-art performance in physical state prediction.
  arXiv  Detail & Related papers  (2025-04-09T17:59:32Z)
• Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient [9.519619751861333]
  We propose a state space model (SSM) based world model based on Mamba.<n>It achieves $O(n)$ memory and computational complexity while effectively capturing long-term dependencies.<n>This model is accessible and can be trained on an off-the-shelf laptop.
  arXiv  Detail & Related papers  (2024-10-11T15:10:40Z)
• Finetuning Offline World Models in the Real World [13.46766121896684]
  Reinforcement Learning (RL) is notoriously data-inefficient, which makes training on a real robot difficult. offline RL has been proposed as a framework for training RL policies on pre-existing datasets without any online interaction. In this work, we consider the problem of pretraining a world model with offline data collected on a real robot, and then finetuning the model on online data collected by planning with the learned model.
  arXiv  Detail & Related papers  (2023-10-24T17:46:12Z)
• HarmonyDream: Task Harmonization Inside World Models [93.07314830304193]
  Model-based reinforcement learning (MBRL) holds the promise of sample-efficient learning. We propose a simple yet effective approach, HarmonyDream, which automatically adjusts loss coefficients to maintain task harmonization.
  arXiv  Detail & Related papers  (2023-09-30T11:38:13Z)
• Simplified Temporal Consistency Reinforcement Learning [19.814047499837084]
  We show that a simple representation learning approach relying on a latent dynamics model trained by latent temporal consistency is sufficient for high-performance RL. Our approach outperforms model-free methods by a large margin and matches model-based methods' sample efficiency while training 2.4 times faster.
  arXiv  Detail & Related papers  (2023-06-15T19:37:43Z)
• Model-Based Reinforcement Learning with Multi-Task Offline Pretraining [59.82457030180094]
  We present a model-based RL method that learns to transfer potentially useful dynamics and action demonstrations from offline data to a novel task. The main idea is to use the world models not only as simulators for behavior learning but also as tools to measure the task relevance. We demonstrate the advantages of our approach compared with the state-of-the-art methods in Meta-World and DeepMind Control Suite.
  arXiv  Detail & Related papers  (2023-06-06T02:24:41Z)
• SAM-RL: Sensing-Aware Model-Based Reinforcement Learning via Differentiable Physics-Based Simulation and Rendering [49.78647219715034]
  We propose a sensing-aware model-based reinforcement learning system called SAM-RL. With the sensing-aware learning pipeline, SAM-RL allows a robot to select an informative viewpoint to monitor the task process. We apply our framework to real world experiments for accomplishing three manipulation tasks: robotic assembly, tool manipulation, and deformable object manipulation.
  arXiv  Detail & Related papers  (2022-10-27T05:30:43Z)
• On the Feasibility of Cross-Task Transfer with Model-Based Reinforcement Learning [45.73223325256312]
  We investigate whether internal models learned by modern model-based RL algorithms can be leveraged to solve new, distinctly different tasks faster. We propose Model-Based Cross-Task Transfer (XTRA), a framework for sample-efficient online RL with scalable pretraining and finetuning of learned world models.
  arXiv  Detail & Related papers  (2022-10-19T17:57:06Z)
• A Unified Framework for Alternating Offline Model Training and Policy Learning [62.19209005400561]
  In offline model-based reinforcement learning, we learn a dynamic model from historically collected data, and utilize the learned model and fixed datasets for policy learning. We develop an iterative offline MBRL framework, where we maximize a lower bound of the true expected return. With the proposed unified model-policy learning framework, we achieve competitive performance on a wide range of continuous-control offline reinforcement learning datasets.
  arXiv  Detail & Related papers  (2022-10-12T04:58:51Z)
• Information Theoretic Model Predictive Q-Learning [64.74041985237105]
  We present a novel theoretical connection between information theoretic MPC and entropy regularized RL. We develop a Q-learning algorithm that can leverage biased models.
  arXiv  Detail & Related papers  (2019-12-31T00:29:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.