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

arXiv: https://arxiv.org/abs/2505.13709 · HTML: https://mlsanity.com/paper/2505.13709/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 35 of its 54 references that are in the map.
- First posted on arXiv 19 May 2025.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 35)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs](https://mlsanity.com/paper/2006.10814/) (2020)
- [Pessimistic Model-based Offline Reinforcement Learning under Partial Coverage](https://mlsanity.com/paper/2107.06226/) (2021)
- [RAMBO-RL: Robust Adversarial Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2204.12581/) (2022)
- [Robust Reinforcement Learning using Offline Data](https://mlsanity.com/paper/2208.05129/) (2022)
- [The Value Equivalence Principle for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2011.03506/) (2020)
- [Distributionally Robust Model-Based Offline Reinforcement Learning with Near-Optimal Sample Complexity](https://mlsanity.com/paper/2208.05767/) (2022)
- [RORL: Robust Offline Reinforcement Learning via Conservative Smoothing](https://mlsanity.com/paper/2206.02829/) (2022)
- [Safe Policies for Reinforcement Learning via Primal-Dual Methods](https://mlsanity.com/paper/1911.09101/) (2019)
- [A Game Theoretic Framework for Model Based Reinforcement Learning](https://mlsanity.com/paper/2004.07804/) (2020)
- [Revisiting Design Choices in Offline Model-Based Reinforcement Learning](https://mlsanity.com/paper/2110.04135/) (2021)
- [Mismatched No More: Joint Model-Policy Optimization for Model-Based RL](https://mlsanity.com/paper/2110.02758/) (2021)
- [Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation](https://mlsanity.com/paper/2106.03273/) (2021)
- [Robust Reinforcement Learning as a Stackelberg Game via Adaptively-Regularized Adversarial Training](https://mlsanity.com/paper/2202.09514/) (2022)
- [Deep Generative Models for Offline Policy Learning: Tutorial, Survey, and Perspectives on Future Directions](https://mlsanity.com/paper/2402.13777/) (2024)
- [A Unified Framework for Alternating Offline Model Training and Policy Learning](https://mlsanity.com/paper/2210.05922/) (2022)
- [The Virtues of Laziness in Model-based RL: A Unified Objective and Algorithms](https://mlsanity.com/paper/2303.00694/) (2023)
- [Bayes Adaptive Monte Carlo Tree Search for Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2410.11234/) (2024)
- [TOM: Learning Policy-Aware Models for Model-Based Reinforcement Learning via Transition Occupancy Matching](https://mlsanity.com/paper/2305.12663/) (2023)
- [A Unified View on Solving Objective Mismatch in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2310.06253/) (2023)
- [Full Shot Predictions for the DIII-D Tokamak via Deep Recurrent Networks](https://mlsanity.com/paper/2404.12416/) (2024)
- [Uncertainty-based Offline Variational Bayesian Reinforcement Learning for Robustness under Diverse Data Corruptions](https://mlsanity.com/paper/2411.00465/) (2024)

## Built on by (most cited first, 2 of 2)

- [Offline Reinforcement Learning for Rotation Profile Control in Tokamaks](https://mlsanity.com/paper/2605.05857/) (2026)
- [Offline Reinforcement Learning for Plasma Control in Nuclear Fusion: Codebase and Benchmark](https://mlsanity.com/paper/2606.07550/) (2026)

## Abstract (arXiv)

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.