Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning

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.

• The paper presents an inference-time MPC that dynamically optimizes policy parameters using differentiable diffusion world models, addressing key offline RL challenges.
• It leverages gradient-based adaptation with recursive Jacobian propagation through imagined rollouts, validated on continuous-control benchmarks.
• Empirical results show significant performance gains over strong baselines in MuJoCo and AntMaze tasks, highlighting practical benefits for offline RL applications.

Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning

Motivation and Background

The paper addresses offline reinforcement learning (RL), where the agent must learn optimal policies from a fixed dataset without further environment interaction. In classical offline RL, an agent learns a policy or value function using static data and then deploys it unchanged at inference time. Distribution shift, scarcity of accurate $Q$-values, and challenges in long-horizon credit assignment persist in this paradigm. To circumvent such limitations, generative world models, particularly diffusion-based approaches, have been proposed, but prior works mainly use these models for training augmentation or sampling plausible trajectories at inference without policy adaptation.

This work introduces inference-time adaptation by leveraging a differentiable world model (DWM) in a model-predictive control (MPC) framework. The central premise is to optimize the policy at inference based on imagined rollouts, thereby exploiting the information in the current state and the world model to refine policy parameters in real-time.

Differentiable World Model Architecture

The DWM pipeline consists of three modules:

1. A diffusion-based state-transition sampler $f_\theta$ trained to represent $p_\theta(s_{t+1} | s_t, a_t)$ via conditional reverse diffusion with Gaussian noise.
2. A reward model $r_\xi$ trained by supervised regression to estimate scalar rewards for state-action pairs.
3. A terminal value function $Q_\phi$, used to provide a surrogate terminal value at the end of rollouts, and a policy $\pi_\psi$ optimized to achieve high $Q$ while regularized toward the dataset behavior.

All modules are differentiable with respect to their inputs, enabling end-to-end gradient-based optimization. For any fixed realization of noise, the diffusion sampler admits deterministic, differentiable mapping from $(s_t, a_t)$ to $s_{t+1}$.

Figure 1: Inference-time MPC using a diffusion world model, showing the learned components and the gradient flow during rollouts from current environment state.

Inference-Time Model Predictive Control

At inference, given the current environment state $s_t$, the MPC framework operates as follows:

• Sample $M$ noise sequences.
• For each, roll out a trajectory of fixed horizon $H$ using the current policy and the diffusion-based dynamics model.
• Accumulate a surrogate return comprising predicted rewards and a terminal value from the critic.
• Compute the Monte Carlo estimate of the expected return across sampled rollouts.
• Perform several steps of gradient ascent, updating $\psi$ (policy parameters) by backpropagating through the full computation graph formed by rollouts.
• Execute the first action of the adapted policy in the real environment and repeat.

This mechanism enables per-step adaptation of the policy conditioned on the encountered state, using gradient flow from the surrogate MPC objective through both the policy and the world model.

Gradient Backpropagation and Chain Rule Recursion

The paper formalizes gradient propagation through rollouts produced by the diffusion sampler, expressing sensitivities via chain rule recursions. Explicit Jacobian recursions are derived for computing how policy parameters $\psi$ perturb the rollout trajectory and, ultimately, the surrogate return. The diffusion sampler gradients are computed via recursive differentiation of the reverse-step map, enabling efficient and tractable computation even with deep generative models.

Empirical Evaluation

The approach is evaluated on D4RL continuous-control benchmarks, including MuJoCo locomotion (HalfCheetah, Hopper, Walker2d) and AntMaze tasks. Offline, world model components are trained on dataset transitions and rewards. The diffusion model's prediction RMSE and reward model accuracy improve as training progresses, yielding robust dynamics and reward predictions.

Figure 2: RMSE of diffusion-based state prediction across training steps for three environments, demonstrating consistent monotonic error reduction.

Figure 3: RMSE of the reward model across training steps, showing decreasing error for all benchmark environments.

Inference-time adaptation by the proposed MPC-DWM framework consistently outperforms strong baselines (TD3+BC, IQL, CQL, SAC-RND, ReBRAC). Notably, average normalized scores of 85.33 (MuJoCo) and 85.07 (AntMaze) surpass prior state-of-the-art results, with improvements over strong pre-trained policies in 16 out of 18 MuJoCo and all AntMaze datasets. On medium-replay tasks, performance gains are most pronounced, indicating robust adaptation via the DWM at inference.

Contradicting prior generative planning and model-based offline RL, this work demonstrates that direct optimization of policy parameters at inference, rather than just selection of action sequences, produces tangible statistical benefits.

Practical and Theoretical Implications

The methodology is broadly applicable to offline RL domains with limited online interaction or safety constraints, such as recommender systems, robotics, and healthcare. The framework augments deployment phase computation in RL agents, leveraging the differentiability of modern generative world models to optimize policies dynamically.

Theoretically, the work bridges generative modeling and differentiable control, suggesting that local optimization at inference via learned dynamics outperforms relying solely on pre-trained policies. The explicit recursion for gradient propagation expands tractable optimization in deep generative models, with potential implications for other sequential decision-making domains.

Limitations and Future Work

Current implementation is computationally intensive due to rollout unrolling and backpropagation, especially as horizon length increases. Future development should consider multi-step generative models (e.g., flow matching, mean flow objectives) to reduce inference cost. Decoupling policy improvement from BPTT via flow Q-learning or leveraging offline-to-online updating schemes may enhance robustness and efficiency.

Conclusion

The paper establishes that inference-time model-predictive adaptation of policy parameters through differentiable diffusion world models significantly improves agent performance in offline RL. The framework achieves superior returns on challenging benchmarks, validating the utility of exploiting local dynamics and reward information at deployment. Extensions to multi-step prediction, scalable optimization, and offline-to-online policy refinement are natural directions for advancing the paradigm (2603.22430).