ROMBRL: Policy-Driven World Model Adaptation for Robust Offline Model-based RL

Official implementation for "Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning" (ICML 2026), by Jiayu Chen, Le Xu, Aravind Venugopal, and Jeff Schneider.

• 📄 Paper: arXiv:2505.13709 · OpenReview
• 🌐 Project page: agentic-intelligence-lab.github.io/ROMBRL
• 📝 Blog post: agentic-intelligence-lab.org/blog
• 🖼️ Poster: docs/assets/poster.pdf

TL;DR

Offline model-based RL (MBRL) usually learns a world model and a policy in two separate stages — fit the model to maximize data likelihood, then optimize the policy against the fixed model. This objective mismatch leaves policies brittle to deployment-time noise. ROMBRL jointly adapts the world model with the policy under a single constrained maximin objective, solved via Stackelberg learning dynamics (policy as leader, world model as adversarial follower), with a formal suboptimality bound. It achieves state-of-the-art robustness on D4RL MuJoCo and stochastic Tokamak Control benchmarks, at almost no cost to clean-environment performance.

Repository Structure

ROMBRL/
├── D4RL/     # D4RL MuJoCo experiments — Tables 1, 2, and 4
└── Fusion/   # Tokamak Control experiments — Table 3

Each folder is a self-contained fork of OfflineRL-Kit extended with our method (rombrl2/rombrl policies) and baselines.

Data Availability

• D4RL (Tables 1, 2, 4): uses the public D4RL MuJoCo datasets — fully reproducible.
• Tokamak Control (Table 3): uses operational data from the DIII-D tokamak, which is proprietary and not redistributed in this repository. We are unable to release it until we obtain the necessary approvals; see Fusion/README.md. The Fusion/ code (dynamics model, environment, RL pipeline) is provided for reference and can be run once you have access to equivalent data.

Getting Started

Setup is folder-specific. The D4RL and Tokamak Control experiments use different environment and simulator stacks, and their dependency sets are mutually incompatible — create a separate Python 3.9 environment for each. Each folder's README is the authoritative source for its installation steps, data placement, and run commands:

• D4RL/README.md — MuJoCo 2.1.0 binary, ctree extension build, checkpoint download, and how a run is configured.
• Fusion/README.md — dynamics-toolbox (install order matters), ctree extension build, and the data paths that must be edited before running.

Reproducing the Paper's Results

Paper Table	Experiment	Instructions
Tables 1, 2, 4	D4RL MuJoCo under measurement noise and RWRL-style deployment perturbations	D4RL/README.md
Table 3	Tokamak Control tracking tasks	Fusion/README.md

One thing worth knowing before you start: the D4RL/ scripts select their task and seed through a variable at the bottom of each run_XXX.py, not through --task/--seed. See D4RL/README.md for the details.

Citation

The official PMLR proceedings for ICML 2026 (volume 306) have not been posted yet; please cite the arXiv version for now:

@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}
}

Acknowledgements

This work was funded in part by the Department of Energy Fusion Energy Sciences under grant DE-SC0024544. The codebase builds on OfflineRL-Kit.

License

Released under the MIT License.