Investigating Online RL in World Models

Uljad Berdica, Kelvin Li, Michael Beukman, Alexander David Goldie, Mattie Fellows, Perla Maiolino, Jakob Nicolaus Foerster

Keywords: World models, Domain Randomization, Offline RL

TL;DR: We train a set of different world models on a dataset of mixed expertise and use them as levels to train an RL agent

Abstract: Significant advances in online reinforcement learning (RL) remain limited by the need for extensive environment interaction or accurate simulators. World models trained on large-scale uncurated *offline data* could provide a training paradigm for generalist AI agents which alleviates the need for task specific simulation environments. Unfortunately, current offline RL methods rely on truncated rollouts that can lead to value overestimation and limit out-of-sample exploration. Additioanlly, common offline RL datasets have been shows to have a bias towards healthy behavior which does not help with the development of generalizable methods. We propose an algorithm and a data curation method that addresses both of these concerns by demonstrating that effective *full-length rollout* training is possible *without hand-crafted penalties* by treating each member of the world model ensemble as a level in the Unsupervised Environment Design (UED) framework. Our method achieves competitive performance even with less transitions than the same online algorithms are traditionally trained on. We find that training a recurrent policy on an ensemble of world models is sufficient to ensure transfer to the original environment and match online PPO performance on standard offline-RL benchmarks while maintaining robust performance on our dataset, where conventional offline RL methods underperform.

Primary Area: reinforcement learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 7854