Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning

• URL: http://arxiv.org/abs/2603.22430v1
• Date: Mon, 23 Mar 2026 18:05:29 GMT
• Title: Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning
• Abstract summary: We introduce an inference time adaptation framework inspired by model predictive control (MPC)<n>We show that exploiting inference-time information to optimize the policy parameters yields consistent gains over strong offline RL baselines.
• Score: 17.293286354647716
• License: http://creativecommons.org/licenses/by-sa/4.0/
• 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-time information to optimize the policy parameters yields consistent gains over strong offline RL baselines.

Related papers

• Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization [46.32504081845328]
  Model-based reinforcement learning can be effectively supported at scale through the use of world models.<n>We identify a bottleneck: a structural misalignment between search and value learning in existing world model approaches.<n>We propose Model-Based Diffusion Policy Optimization (MBDPO) in world models, a framework that unifies search and policy optimization through diffusion policy representations.
  arXiv  Detail & Related papers  (2026-05-25T19:06:51Z)
• IPD: Boosting Sequential Policy with Imaginary Planning Distillation in Offline Reinforcement Learning [13.655904209137006]
  We propose textbfImaginary Planning Distillation (IPD), a novel framework that seamlessly incorporates offline planning into data generation, supervised training, and online inference.<n>Our framework first learns a world model equipped with uncertainty measures and a quasi-optimal value function from the offline data.<n>By replacing the conventional, manually-tuned return-to-go with the learned quasi-optimal value function, IPD improves both decision-making stability and performance during inference.
  arXiv  Detail & Related papers  (2026-03-04T17:05:39Z)
• Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning [6.189693079685375]
  offline model-based RL (MBRL) explicitly learns a world model from a static dataset.<n>We propose a framework that dynamically adapts the world model alongside the policy.<n>We benchmark our algorithm on twelve noisy D4RL MuJoCo tasks and three Tokamak Control tasks, demonstrating its state-of-the-art performance.
  arXiv  Detail & Related papers  (2025-05-19T20:14:33Z)
• Model-Based Offline Reinforcement Learning with Adversarial Data Augmentation [36.9134885948595]
  We introduce Model-based Offline Reinforcement learning with AdversariaL data augmentation.<n>In MORAL, we replace the fixed horizon rollout by employing adversaria data augmentation to execute alternating sampling with ensemble models.<n>Experiments on D4RL benchmark demonstrate that MORAL outperforms other model-based offline RL methods in terms of policy learning and sample efficiency.
  arXiv  Detail & Related papers  (2025-03-26T07:24:34Z)
• Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models [79.2162092822111]
  We systematically evaluate reinforcement learning (RL) and control-based methods on a suite of navigation tasks.<n>We employ a latent dynamics model using the Joint Embedding Predictive Architecture (JEPA) and employ it for planning.<n>Our results show that model-free RL benefits most from large amounts of high-quality data, whereas model-based planning generalizes better to unseen layouts.
  arXiv  Detail & Related papers  (2025-02-20T18:39:41Z)
• Learning from Random Demonstrations: Offline Reinforcement Learning with Importance-Sampled Diffusion Models [19.05224410249602]
  We propose a novel approach for offline reinforcement learning with closed-loop policy evaluation and world-model adaptation. We analyzed the performance of the proposed method and provided an upper bound on the return gap between our method and the real environment under an optimal policy.
  arXiv  Detail & Related papers  (2024-05-30T09:34:31Z)
• MOTO: Offline Pre-training to Online Fine-tuning for Model-based Robot Learning [52.101643259906915]
  We study the problem of offline pre-training and online fine-tuning for reinforcement learning from high-dimensional observations. Existing model-based offline RL methods are not suitable for offline-to-online fine-tuning in high-dimensional domains. We propose an on-policy model-based method that can efficiently reuse prior data through model-based value expansion and policy regularization.
  arXiv  Detail & Related papers  (2024-01-06T21:04:31Z)
• Statistically Efficient Variance Reduction with Double Policy Estimation for Off-Policy Evaluation in Sequence-Modeled Reinforcement Learning [53.97273491846883]
  We propose DPE: an RL algorithm that blends offline sequence modeling and offline reinforcement learning with Double Policy Estimation. We validate our method in multiple tasks of OpenAI Gym with D4RL benchmarks.
  arXiv  Detail & Related papers  (2023-08-28T20:46:07Z)
• 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)
• Regularizing a Model-based Policy Stationary Distribution to Stabilize Offline Reinforcement Learning [62.19209005400561]
  offline reinforcement learning (RL) extends the paradigm of classical RL algorithms to purely learning from static datasets. A key challenge of offline RL is the instability of policy training, caused by the mismatch between the distribution of the offline data and the undiscounted stationary state-action distribution of the learned policy. We regularize the undiscounted stationary distribution of the current policy towards the offline data during the policy optimization process.
  arXiv  Detail & Related papers  (2022-06-14T20:56:16Z)
• MOPO: Model-based Offline Policy Optimization [183.6449600580806]
  offline reinforcement learning (RL) refers to the problem of learning policies entirely from a large batch of previously collected data. We show that an existing model-based RL algorithm already produces significant gains in the offline setting. We propose to modify the existing model-based RL methods by applying them with rewards artificially penalized by the uncertainty of the dynamics.
  arXiv  Detail & Related papers  (2020-05-27T08:46:41Z)

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.