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

TL;DR

This work tackles robustness in offline model-based reinforcement learning by formulating a constrained maximin objective that jointly optimizes a policy and adversarially adapts the world model. It introduces ROMBRL, a Stackelberg-dynamics-based method that updates the policy as the leader and the world model as the follower, with theoretical bounds on suboptimality and practical techniques (Woodbury inverses, Fisher-based Hessian approximations, gradient masking) to ensure tractable training. Theoretical guarantees tie the worst-case policy performance to an uncertainty budget and data size, while empirical results demonstrate state-of-the-art robustness and performance on noisy D4RL MuJoCo tasks and stochastic Tokamak Control benchmarks. The combination of rigorous analysis, efficient optimization, and strong empirical validation supports ROMBRL as a principled, scalable approach for robust offline control with uncertain dynamics. This framework has practical significance for deploying offline RL in real-world, noise-sensitive environments such as autonomous robotics and fusion energy systems.

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.

Figures (4)

• Figure 1: Average scores of different offline RL algorithms on nine D4RL MuJoCo tasks (corresponding to the tasks in Table \ref{['table:1']} excluding those with the random data type), before and after applying random noise to state transitions. The noise is modeled as zero-mean Gaussian with a standard deviation equal to 5% of the state change, simulating common measurement noise.
• Figure 2: Evaluation results on D4RL MuJoCo. The figure shows the progression of evaluation scores over training epochs for the proposed algorithm and baseline methods. Solid lines indicate the average performance across multiple random seeds. For clarity of presentation, the curves have been smoothed using a sliding window, and confidence intervals are omitted.
• Figure 3: Illustration of the Tokamak Control tasks. The RL controller is trained to apply actuators, such as power and torque, based on the current plasma state, with the goal of driving the plasma toward a target profile. For practical reasons, the real tokamak is replaced with an ensemble of dynamics models trained on operational data from a real device -- DIII-D. These models are used to generate data for offline RL and evaluate the trained policies.
• Figure 4: Evaluation results on Tokamak Control tasks. The figure shows the progression of episodic tracking errors over training epochs for the proposed algorithm and baseline methods. Solid lines indicate the average performance across multiple random seeds. For clarity of presentation, the curves have been smoothed using a sliding window, and confidence intervals are omitted.

Theorems & Definitions (14)

• Definition 1: LSE
• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Lemma 1
• Lemma 2
• proof
• Lemma 3
• Lemma 4