ICML 2026 · Seoul, South Korea

Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning

Jiayu Chen^1,2,*  ·  Le Xu^3,*  ·  Aravind Venugopal⁴  ·  Jeff Schneider⁴

¹The University of Hong Kong   ²INFIFORCE Intelligent Technology Co., Ltd.   ³Tsinghua University   ⁴Carnegie Mellon University   ^*Equal contribution

📄 arXiv 📝 OpenReview 💻 Code 🖼️ Poster (PDF)

Abstract

Offline reinforcement learning (RL) offers a powerful paradigm for data-driven control. Compared to model-free approaches, offline model-based RL (MBRL) explicitly learns a world model from a static dataset and uses it as a surrogate simulator, improving data efficiency and enabling potential generalization beyond the dataset support. However, most existing offline MBRL methods follow a two-stage training procedure: first learning a world model by maximizing the likelihood of the observed transitions, then optimizing a policy to maximize its expected return under the learned model. This objective mismatch results in a world model that is not necessarily optimized for effective policy learning. Moreover, we observe that policies learned via offline MBRL often lack robustness during deployment, and small adversarial noise in the environment can lead to significant performance degradation. To address these, we propose a framework that dynamically adapts the world model alongside the policy under a unified learning objective aimed at improving robustness. At the core of our method is a maximin optimization problem, which we solve by innovatively utilizing Stackelberg learning dynamics. We provide theoretical analysis to support our design and introduce computationally efficient implementations. We benchmark our algorithm on twelve noisy D4RL MuJoCo tasks and three stochastic Tokamak Control tasks, demonstrating its state-of-the-art performance.

Motivation

Existing offline MBRL methods learn a world model to maximize data likelihood, then freeze it (or lightly adapt it) while optimizing the policy. This causes an objective mismatch — the model is trained to explain the offline data, not to support good policy learning. In practice, this leaves policies brittle: even a modest amount of measurement noise at deployment time can sharply degrade the performance of strong offline RL baselines.

Method

We formulate offline MBRL as a constrained maximin problem: the policy maximizes its return under the worst-case world model within an uncertainty set consistent with the offline data, while the world model is adversarially updated to minimize that return within a trust region.

We solve this as a Stackelberg game — the policy is the leader, the world model is the follower that best-responds to the policy. This is the opposite direction from online MBRL, where the model is typically adapted to serve the leader's interest rather than oppose it. Using implicit differentiation through the follower's best response, we derive primal-dual Stackelberg update rules for \((\theta, \phi, \lambda)\), made practical at scale by:

• Approximating expensive second-order terms with the Fisher Information Matrix, using \(\mathbb{E}[\nabla^2 \log P_\phi] = -\,\mathbb{E}[\nabla \log P_\phi \, \nabla \log P_\phi^{\top}]\) to avoid an explicit Hessian.
• The Woodbury matrix identity to invert the resulting low-rank matrix in time linear in the number of parameters rather than cubic.
• A gradient-mask mechanism so off-policy replay data doesn't go stale as the model is updated jointly with the policy — a subtlety RAMBO's alternating-update scheme overlooks.

The full algorithm, ROMBRL, is given in Appendix J of the paper.

Theory

Theorem 3.1 bounds the performance gap between the policy learned by ROMBRL and the best possible comparator policy, assuming the true environment lies in the uncertainty set \(\Phi\) with probability at least \(1 - \delta/2\):

where \(N\) is the offline dataset size, \(|\Phi|\) the covering number of the uncertainty set, \(C\) a concentrability coefficient, and \(\epsilon\) the uncertainty-set radius — so the suboptimality gap shrinks as the dataset grows, giving a formal robustness guarantee rather than a purely heuristic one. Theorems 3.2 and 3.3 instantiate this bound for tabular MDPs and for continuous MDPs with Gaussian world models respectively, characterizing how \(\epsilon\) should scale with dataset size and dimensionality.

Results

D4RL MuJoCo, under 5% measurement noise at deployment. ROMBRL ranks first on 7/12 tasks and second on 4, with the best average score by a wide margin (Cohen's d ≥ 2 over every baseline).

Robustness vs. clean-environment trade-off (OfflineRL-Kit protocol). ROMBRL matches the best clean-environment performance while losing almost nothing to noise.

Tokamak Control (negative tracking error; higher is better). ROMBRL ranks first on 2/3 targets and second on the third, with the lowest variance across seeds.

* marks the second-best result for each row.

RWRL-style deployment perturbations (dropped/stuck observations, body mass, friction, joint damping). ROMBRL achieves the best score on every perturbation type and the smallest overall performance drop.

Poster

Citation

@article{chen2025rombrl,
  title   = {Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning},
  author  = {Chen, Jiayu and Xu, Le and Venugopal, Aravind and Schneider, Jeff},
  journal = {arXiv preprint arXiv:2505.13709},
  year    = {2025}
}