[2603.22430] Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning

Computer Science > Machine Learning arXiv:2603.22430 (cs) [Submitted on 23 Mar 2026] Title:Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning Authors:Rohan Deb, Stephen J. Wright, Arindam Banerjee View a PDF of the paper titled Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning, by Rohan Deb and 2 other authors View PDF HTML (experimental) Abstract:Offline Reinforcement Learning (RL) aims to learn optimal policies from fixed offline datasets, without further interactions with the environment. Such methods train an offline policy (or value function), and apply it at inference time without further refinement. We introduce an inference time adaptation framework inspired by model predictive control (MPC) that utilizes a pretrained policy along with a learned world model of state transitions and rewards. While existing world model and diffusion-planning methods use learned dynamics to generate imagined trajectories during training, or to sample candidate plans at inference time, they do not use inference-time information to optimize the policy parameters on the fly. In contrast, our design is a Differentiable World Model (DWM) pipeline that enables endto-end gradient computation through imagined rollouts for policy optimization at inference time based on MPC. We evaluate our algorithm on D4RL continuous-control benchmarks (MuJoCo locomotion tasks and AntMaze), and show that exploiting inference-...