PWM: Policy Learning with Multi-Task World Models

TL;DR

PWM introduces a model-based reinforcement learning algorithm that leverages regularized world models to enable fast, efficient policy learning across multiple tasks with high-dimensional action spaces, outperforming existing methods.

Contribution

It presents a novel approach combining offline pre-training and first-order optimization on regularized world models for multi-task continuous control.

Findings

PWM solves tasks with up to 152 action dimensions.

PWM outperforms methods using ground-truth dynamics.

PWM scales to 80 tasks with 27% higher rewards.

Abstract

Reinforcement Learning (RL) has made significant strides in complex tasks but struggles in multi-task settings with different embodiments. World model methods offer scalability by learning a simulation of the environment but often rely on inefficient gradient-free optimization methods for policy extraction. In contrast, gradient-based methods exhibit lower variance but fail to handle discontinuities. Our work reveals that well-regularized world models can generate smoother optimization landscapes than the actual dynamics, facilitating more effective first-order optimization. We introduce Policy learning with multi-task World Models (PWM), a novel model-based RL algorithm for continuous control. Initially, the world model is pre-trained on offline data, and then policies are extracted from it using first-order optimization in less than 10 minutes per task. PWM effectively solves tasks with up to 152 action dimensions and outperforms methods that use ground-truth dynamics. Additionally, PWM scales to an 80-task setting, achieving up to 27% higher rewards than existing baselines without relying on costly online planning. Visualizations and code are available at https://www.imgeorgiev.com/pwm/.