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

Model-based RL (MBRL) simultaneously learns a policy and a world model thatcaptures the environment’s dynamics and rewards. The world model can eitherbe used for planning, for data collection, or to provide first-order policy gradientsfor training. Leveraging a world model significantly improves sample efficiencycompared to model-free RL. However, training a world model alongside the pol-icy increases the computational complexity, leading to longer training times thatare often intractable for complex real-world scenarios. In this work, we proposea new method for accelerating model-based RL using state-space world models.Our approach leverages state-space models (SSMs) to parallelize the training ofthe dynamics model, which is typically the main computational bottleneck. Ad-ditionally, we propose an architecture that provides privileged information to theworld model during training, which is particularly relevant for partially observableenvironments. We evaluate our method in several real-world agile quadrotor flighttasks, involving complex dynamics, for both fully and partially observable envi-ronments. We demonstrate a significant speedup, reducing the world model train-ing 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 achievessimilar sample efficiency and task rewards to state-of-the-art MBRL methods.