Gradient-based Planning with World Models

TL;DR

This paper explores a gradient-based planning method using differentiable world models, demonstrating improved sample efficiency and performance over traditional gradient-free MPC methods, and introduces a hybrid approach combining policies and gradient-based planning.

Contribution

It presents a novel gradient-based planning approach leveraging differentiable world models and a hybrid model that outperforms pure policy methods in complex tasks.

Findings

Gradient-based planning matches or exceeds performance of gradient-free methods.

Hybrid models combining policies and gradient-based planning outperform pure policy approaches.

Sample efficiency is improved with the proposed gradient-based method.

Abstract

The enduring challenge in the field of artificial intelligence has been the control of systems to achieve desired behaviours. While for systems governed by straightforward dynamics equations, methods like Linear Quadratic Regulation (LQR) have historically proven highly effective, most real-world tasks, which require a general problem-solver, demand world models with dynamics that cannot be easily described by simple equations. Consequently, these models must be learned from data using neural networks. Most model predictive control (MPC) algorithms designed for visual world models have traditionally explored gradient-free population-based optimisation methods, such as Cross Entropy and Model Predictive Path Integral (MPPI) for planning. However, we present an exploration of a gradient-based alternative that fully leverages the differentiability of the world model. In our study, we conduct a comparative analysis between our method and other MPC-based alternatives, as well as policy-based algorithms. In a sample-efficient setting, our method achieves on par or superior performance compared to the alternative approaches in most tasks. Additionally, we introduce a hybrid model that combines policy networks and gradient-based MPC, which outperforms pure policy based methods thereby holding promise for Gradient-based planning with world models in complex real-world tasks.