Investigating Online RL in World Models

Authors: Uljad Berdica, Kelvin Li, Michael Beukman, Alexander David Goldie, Mattie Fellows, Perla Maiolino, Jakob Nicolaus Foerster
Primary area: reinforcement learning
Keywords: World models, Domain Randomization, Offline RL

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.

Decision: reject
Reviewer scores: 5, 5, 5, 3, 3

Reviews and discussion:
<thread>
creation date:
1729952290094

thread_type:
official_review

<thread_object>
creator: reviewer_ta3z

<content>
summary:
The paper addresses the challenge of training online RL models inside world models. For this they rely on an ensemble of world models to train agents in an online fashion, without any offline penalties. The proposed method is evaluated on robotics tasks including Cartpole, Halfcheetah and Hopper.

strengths:
The paper investigates an interesting question, that is, whether RL agents can be trained online inside world models, without the need for constraints required in offline RL.

weaknesses:
The presentation of this paper can be improved. For example, Figure 1 is confusing and lacks a detailed figure caption to explain what’s going on. Dataset distributions of offline RL datasets are criticised in the Discussion section, which feels out of place (also as this section comes after Related Work). Instead of bringing these results after the main experiments, they could be seen as a motivation before the main experiments (which would require experiments that support the motivation).

It is unclear where the proposed method starts and ends, or what the authors consider as their own method. Instead, 6 different variations (PLR, PLR_PVL, DR, DR_STEP, DR_PROB) of ensemble world models are tested against a single world model (WM). 

In the current form, the paper lacks convincing empirical evidence. The main experiments are conducted on toy tasks like Cartpole, Pendulum, Acrobot, Hopper and Halfcheetah. Across most environments (Figure 4, 6, 7), the ensemble methods (PLR, DR, PLR_PVL, DR_STEP) tend to learn faster than a single model (WM) but reach the same level of performance. Furthermore, there is a lack of offline RL baselines (only CQL is provided). Also, there’s a lack of ablations on components of their method. For example, what is the effect of the size of the ensemble? 

The proposed method relies on an ensemble of world models. Training an ensemble of world models is feasible in the toy tasks considered currently, but raises questions on scalability of the proposed methods to more complex environments, such as visual domains. Consequently, reporting information on the additional cost of training and during evaluation of those models would benefit the paper.

questions:
- What is the size of the world model ensemble used in your experiments? What is the effect of the number of world models? 
- What is the additional cost of training and evaluating using an ensemble? 
- How does your method transfer to other benchmarks (e.g., Atari, Procgen, Meta-World, etc.) with regard to efficiency at training and inference time?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730196478799

thread_type:
official_review

<thread_object>
creator: reviewer_9slj

<content>
summary:
The paper investigates online reinforcement learning directly within world models without conservative constraints typically used in offline RL and proposes to combine ensemble world models and prioritized level replay to tackle this problem.

strengths:
- The motivation to perform online RL directly in world models is timely and relevant.

weaknesses:
1. The method is confusing. What is the $\delta$ in Equation (5)? Why the Equation (5) approximates the regret? The right side of Figure 1 is difficult to understand. I suggest adding pseudocode to clarify the algorithm.

2. The paper structure and writing need significant improvement. For example, the introduction lacks a clear statement of contributions. The relationship between contextual MDP subsection in Preliminaries and the reset of paper is unclear. The Preliminaries section is too long.

3. The dataset construction method appears similar to D4RL. Both record the encountered transitions from randomness to expertise during training. Why the authors claim that D4RL is adversarial in data coverage (line 55) and CQL and TD3+BC does not inform this (line 263)? The data coverage is a common concern when building offline RL datasets, and previous benchmarks typically offer choices with various data coverage (D4RL and Atari DQN-Replay[1]).
 
4. The checkpoint frequency seems uniform based on Figure 2, despite claims of heuristic selection. At the 5th ckpt, the agent has basically converged. The distribution perhaps change little after the subsequent sampling. Should the sampling frequency be increased between the first few ckpts?

5. The world model architecture (simple MLP with current state and action as input) seems overly simplistic compared to state-of-the-art approaches like RSSM, Transformer, or diffusion models. Since the authors assumes that learning is done within a generalist world model (line 58 and 87), the experiment results of MLP with no historical trajectory are not convincing. I even suspect that using a single SOTA architecture world model can already solve this problem.

6. Based on Figures 4-6, the DR methods perform well, and the proposed PLR methods do not show clear advantages.

7. Judging from the rendering results of the `ref` and `cite` commands, this paper does not use the ICLR 2025 template!

[1] Agarwal R, Schuurmans D, Norouzi M. An optimistic perspective on offline reinforcement learning. International conference on machine learning. PMLR, 2020: 104-114.

questions:
1. What is the size of the collected dataset in terms of transitions?
2. Could you provide pseudocode for the algorithm to aid understanding?
3. Why were existing implementations of CQL[1] and SAC_N[2] not used?
4. There are numerous typos that hinder readability. Such as incorrect citation format (line 51 and 53), "This setting Furthermore, they showed ..." (line 131), "Therefore, the policy’s minimizes ..." (line 136), "Prioritized Level Replay as described in with an ..." (line 305) and so on.
5. The abstract states that "training inside world models is usually studied in the context of offline RL." However, many online RL algorithms (e.g., MBPO, Dreamer, TD-MPC, IRIS, TWM, STORM) train agents entirely with imagined trajectories (i.e., within world model). On the contrary, most of the offline model-based RL seems to be based on the Dyna framework, that is, using both offline datasets and world model imagination trajectories to train the agent.

[1] https://github.com/young-geng/JaxCQL

[2] https://github.com/Howuhh/sac-n-jax
</content>
</thread_object>

</thread>
<thread>
creation date:
1730667648506

thread_type:
official_review

<thread_object>
creator: reviewer_rqd6

<content>
summary:
This paper combined Unsupervised Environment Discovery (UED) and world model ensembles to provide a method for offline model-based RL that is sample-efficient. It treats different world models as levels within the PLR curriculum method in UED. It evaluates the method on vector-based cartpole, hopper, half cheetah, and pendulum tasks.

strengths:
- the paper performs an interesting combination of UED and world models, considering each world model as a level in UED
- the method doesn't rely on online tuning in the environment (and uses held out world models to tune hyperparameters)
- the results demonstrate sample-efficiency gains compared to CQL and a vanilla world model baseline

weaknesses:
- The differences with other model-based offline RL methods like MOPO [1], MBPO [2], Planning with Diffusion [3] is not clear, especially since the world model in the experiments in this paper, unlike David Ha's work, is just a simple MLP. 
- It would be helpful to compare with other offline MBRL methods like the ones above as well as world model based approaches like IRIS [4] and Dreamer [5], particularly on more challenging environments with image observations than the vector observation based  locomotion environments. Detailing the differences with these world model based methods would also be useful
- Other papers to discuss in the related work include sample-efficient BC approaches that can work with very few demonstrations like ROT [6] and MCNN [7]. In summary, this work could use more comparisons (or atleast discussions comparing) to other works on offline MBRL, world models, and sample-efficient behavior cloning as well as evaluations in more challenging environments with image observations.

Minor comments:
- the return plotted for different methods is not normalized --- this makes it hard to determine its performance between random and expert and hard to compare with other papers
- confusingly, "inside world model" is referred to as "simulation" and "in simulation" is referred to as "real world"
- Appendix A.3 is empty

[1] T Yu, et al, MOPO: Model-based Offline Policy Optimization, NeurIPS 20

[2] M Janner, et al, When to Trust Your Model: Model-Based Policy Optimization, NeurIPS 19

[3] M Janner, et al, Planning with diffusion for flexible behavior synthesis, ICML 22

[4] V Micheli, et al, Transformers are Sample-Efficient World Models, ICLR 23

[5] D Hafner, et al, Mastering Diverse Domains through World Models 

[6] S Haldar, et al, Watch and match: Supercharging imitation with regularized optimal transport, CoRL 22

[7] K Sridhar, et al, Memory-Consistent Neural Networks for Imitation Learning, ICLR 24

questions:
Please see weaknesses above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730809712144

thread_type:
official_review

<thread_object>
creator: reviewer_dkos

<content>
summary:
The paper explores the potential of using uncurated offline data to train world models that can serve as a training ground for reinforcement learning. The primary goal is to enable the transfer of learned policies from these world models to the real world, thereby reducing the reliance on task-specific simulation environments. The authors demonstrate that by ensembling multiple independently trained world models, they can achieve robust transfer to the real world, even when the offline datasets are much smaller than those typically used in offline RL.

strengths:
The paper presents a novel approach for training RL agents using world models derived from large-scale, uncurated offline data. By employing an ensemble of world models trained on the same dataset and leveraging them to create learning curricula through the Unsupervised Environment Design method, this work introduces a fresh perspective to RL.

weaknesses:
- In my view, the main issue with this paper is that it somewhat exaggerates the contributions of the proposed method. The title, "Investigating Online RL in World Models," is slightly misleading, as the study addresses an offline RL problem. Additionally, the term "world models" could be confusing. While the paper discusses many existing visual world model approaches (such as Ha and Schmidhuber's NIPS 2018 paper and recent interactive video generation studies), it does not actually work with visual data, instead focusing on fully observable MDPs in low-dimensional state spaces. I suggest the authors consider replacing the term "world models" to more accurately reflect the context.

- The organization of the paper is also disjointed, with an imbalanced structure. For example, Chapter 2 uses considerable space for background information, while the methods section in Chapter 3 is relatively brief. This structure makes it challenging for readers to fully grasp the paper's core contributions.

- In methodology, the authors treat world models trained on offline data of varying quality as different levels in an unsupervised environment design approach. I recommend that the authors discuss the motivation and rationale for this choice, explaining why this training method would lead to a robust and transferable policy.

- While the paper outlines an ambitious story, it lacks sufficient experimental support: 
(1) The authors claim to use the D4RL dataset but do not provide comprehensive experimental results across different tasks. The focus on the relatively simple Hopper task is insufficient to support their claim.
(2) The paper lacks comparisons with recent offline RL methods, as well as more detailed model analysis, such as examining the impact of the number of world models --- The authors mention training 100 world models in line 77, which seems excessive for a simple Hopper task and introduces considerable training overhead, which raises questions about the method's practical use in real-world applications.

questions:
- Could the authors provide additional experiments on visual offline RL tasks to demonstrate the world model's generalizability in open-world or visual environments?
- I strongly suggest that the authors further discuss the impact of the number of world models on the experimental results and clarify the necessity of using 100 world models.
- If possible, please refine the methods section to more clearly highlight the core contributions of the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730836737311

thread_type:
official_review

<thread_object>
creator: reviewer_pply

<content>
summary:
The paper presents a novel approach to online reinforcement learning. It is different from the typical approach to online interactions directly with the underlying environment (sim or real). Instead, this work studies RL methods within learned world models, attempting to mitigate common pitfalls of offline RL (reward exploitation, skewed datasets, etc.) and avoid the costly and sometimes infeasible samplings directly from the environments.

Under the settings where authors collected a more uniformly distributed dataset (in terms of state/action coverage compared to D4RL), this work trained PPO from an ensemble of independently trained world models, using Domain Randomization techniques (DR) and Unsupervised environment design (UED) methods. On the curated dataset, the experiment results suggest significant improvements over offline RL methods (CQL & SACn).

Overall, this work is novel and the results presented in the paper offers new insights to training RL agents purely from world models. If there are more experimental evidence from different tasks or environments to support the the claim that "full roll-out training inside world models is possible", and clarify the questions I have below, I would be willing to raise the score.

strengths:
1. Originality:
- While the approach of using an ensemble of world models to reduce over-fitting and exploitation has been previously studied for model-based RL methods, this work differentiates itself by training entirely within the world model and on full-length roll-outs without any penalty terms inside learnt models.

2. Quality:
- The experiments are well-designed with clear assumptions and comparisons against multiple baselines. The use of different data scales and detailed ablation studies on the components of their method provides a thorough validation of their claims.

3. Clarity:
- The paper is well-organized with informative figures and tables. Especially figure 9 and 10, they explained how the curated dataset differs from the d4rl dataset. The writing is generally easy to follow. 

4. Significance:
- The proposed assumptions, settings, and methods are valuable to the RL research community as it shows preliminary positive results on smaller scale world models, which could potentially serve as basis for training RL agents within larger and more capable world models on more complex tasks.

weaknesses:
1. While the paper presents results from multiple tasks (pendulum, half-cheetah, cartpole, hopper), there is a lack of extensive testing across a wider variety of environments and tasks. This raises some questions about the robustness and generalizability of the proposed method beyond the tested scenarios. Perhaps some tasks such as ant-maze or robot arm manipulation ones.

2. The ensembles of world models seems essential to the proposed method. A sweep over the numbers of world models in the ensemble v.s. tasks' performance could reveal more information on how many world models are needed to achieve certain level of task performance.

3. Minor typos:
- Line 131: This setting Furthermore
- Line 509: the type (of) increasingly available large-scale datasets

questions:
1. "Each of the baselines is tuned by doing a grid search of the ranges documented in their respective papers:
- Is the grid search conducted over their original dataset and transferred to the curated dataset? Or is the grid search done on the curated dataset?

2. What would happen if the proposed method is trained and evaluated on the d4rl datasets? In the current paper, we see baseline methods do not perform as well as the proposed method under the paper's settings. The authors argued d4rl dataset is biased towards the baseline methods and provided visualizations. However, it would still be interesting to see how it impacts the proposed method.

3. "We open source all our code and data to facilitate further work in this exciting direction."
- I couldn't find a link to the code repo, or find any supplementary materials.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732026535610

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Balanced the paper structure with a focus on clarity, more implementation details, results and ablations

comment:
Dear reviewers,

We are very grateful for the detailed feedback. Thank you for pointing out how our work:

 >"suggest significant improvements over offline RL methods " 

> "novel approach"

> "reducing the reliance on task-specific simulation environments"

>"sample-efficient" and "interesting combination of UED and world models"



 In addition to the individual responses, we now address shared remarks . The updated paper PDF has been submitted with the most notable changes highlighted in $\color{blue}{blue}$ for your convenience. 


## Clarifying the contribution

We agree that the original introduction did not clearly state our contributions. We have updated to reflect the following: 

1. We focus on having online Reinforcement Learning breakthroughs **like zero-shot transfer with UED** and **successful algorithms like PPO**  transfer to  domains beyond the highly structured ones common in this area of research

2. We inspect the current progress of offline dataset for tasks in RL and identify notable pathologies in model-based methods [1] stemming from **health-biased datasets** and **truncated rollouts** in addition to the well-knownd drawback in model-free offline methods like out-of-sample exploration and the reliance on full sequences

3. **We address** the truncation pathologies by introducing a Policy Optimization with World Model Ensemble Rollouts algorithm, the pseudocode for which has been added under in section 3 as well as a more detailed caption for Figure 1 illustrating the various training options our algorithm covers.

4. **We address** the bias in the dataset by curating our own without any prior assumptions and using checkpoints from a policy trained from initialization up to convergence. The importance of not using checkpoints after convergence and *not* biasing towards more expert states is very important to our claims and has been addressed in our writing in section 4.1 

5. We address your concern and **test our method the D4RL benchmark**. Since we are using an online RL algorithm like PPO, we compare our results to the traditional online training baselines [2]. The tables below show that we outperform (Hopper) or match (Halfcheeta) online PPO using the same number of environment interactions **(1M or more)** equal to the dataset's transition count **(1M max)**. 

## D4RL results 
### Mean Estimates with std error

#### Hopper

| Algorithm | Transition Count | Estimate ± Range |
|-----------|--------|-----------------|
| DR_PROB | 750k| 2532.23 ± 440.96 |
| DR_PROB | 1M | 1918.84 ± 181.51 |
| DR | 750k | 2340.36 ± 478.00 |
| DR | 1M | 1593.71 ± 121.53 |
| WM | 750k | 309.40 ± 131.78 |
| WM | 1M | 138.12 ± 75.73 |
| DR_STEP | 750k| 1583.82 ± 473.90 |
| DR_STEP | 1M | 2032.70 ± 511.87 |
| PLR_PVL | 750k | 2263.78 ± 193.64 |
| PLR_PVL | 1M | 1704.73 ± 302.25 |

#### HalfCheetah

| Algorithm | Ratio | Estimate ± Range |
|-----------|--------|-----------------|
| DR_PROB | 750k | 2049.58 ± 507.44 |
| DR_PROB |1M| 1654.08 ± 599.01 |
| DR | 750k | 1675.61 ± 677.12 |
| DR | 1M | 3843.04 ± 545.58 |
| WM | 750k | -581.86 ± 91.81 |
| WM | 1M | -287.47 ± 148.11 |
| DR_STEP | 750k | 808.88 ± 222.86 |
| DR_STEP | 1M | 980.68 ± 97.41 |
| PLR_PVL | 750k| 2527.19 ± 911.77 |
| PLR_PVL | 1M | 4937.60 ± 335.30 |

## Additional Results 

We would like to **bring additional results** to your attention like 

a) the effective use of UED as world model/level selector to prevent reward hacking during training in Section 5.1 

b) Effectively solving the classical control tasks in a few episodes worth of transitions 

c) We **qualitatively** and **quantitatively** show in Section 5.6 that our recurrent policy learns to identify different world models and distinguish them from the real environment. This re-affirms how our method leads to diverse world models and useful abstraction for curating a training curricula. 

d) Appendix A3 shows that our JAX implementation of the supervised world models training is very scaleable due to the use of vectorized mapping and the genreation of multiple levels in parallel.  This is on top of the already proven 100x speedups of PureJaxRL. Our baselines are also based on jax implementations

We believe our paper presents a useful algorithms for the community to use and build upon in addition to observations on the state of the field. We are open to improving this submission to reflect that as efficiently as possible and make sure the potential of this work is fully conveyed.

___

[1] Sims et al. *The Edge-of-Reach Problem in Offline Model-Based Reinforcement Learning*

[2] ICLR blog with baseline table https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/

[3] https://github.com/luchris429/purejaxrl

[4] https://github.com/nissymori/JAX-CORL
</content>
</thread_object>

</thread>
<thread>
creation date:
1733306503232

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
A final grounding

comment:
# Dear reviewers,

We are very thankful for your constructive feedback in revising the paper, especially the subset who engaged with our responses. We have improved the quality of the paper by: 

**1.** clarifying the contributions 

**2.** adding experiments with datasets of similar shape and size, but different distribution to the one we generated ourselves

**3.** adding ablations on key components of our method in section **5.5** and **A6**

**4.** quantifying the computational advantage of our method by training every network, end-to-end on a single GPU in Appendix **A3**.

We have addressed all your concerns to the best of our understanding, which is reflected in both the anonymous repo updates *and* the paper itself, with the edits labeled in $\color{blue}{blue}$. During this rebuttal period, we have noticed a recurring disagreement likely resulting from the **novelty of our approach** and the imperfect phrasing in parts of the initial submission. 

# The use of *world models*

Machine Learning rapidly evolves through new paradigms, exemplified by the 2018 'world models' paper [1]. This seminal paper introduced an agent training pipeline that is used to model video game environments where the observations are pixel-based. The Dreamer line of work [2,3,4] uses the same concept of visual inputs and achieves impressive results in Atari and visual MuJoCo. The IRIS paper uses a *more sample-efficient* Transformer architecture for the world model which set a new SOTA. Their paper states that **"the quality of the world model is the cornerstone of their approach"** and that **"the agent will learn a suboptimal policy if the world model is flawed"**. This line of work is essentially **online** and our method deals with situations where this world model IS NOT possible. 

Recent work commonly ties *'world models'* to pixel-based environments like Atari, but this narrow interpretation overlooks the term's broader origins. The authors would like to point to the fact that the terminology for the seminal *world models* paper [1] is based on the earlier work by Schmidhuber et al. [5] that clearly states a world model is what helps a controller learn efficiently. If an agent is only proprioceptive and its observations are the same dimensionality as the original test environment (MuJoCo, Brax and classical control), then the use of the term *world model* is also correct and does not warrant using pixels.  

**Our work** does not rely on pixel-perfect reconstructions and it specifically addresses the challenges that arise when these architectures are not used and compounding error **becomes** an issue. We implement a pixel-based version of our method in our repo and qualitatively demonstrate that the pixel-based world models are not a relevant setting to test our method. We will make sure to add more quantitative plots that crystallize this important delineation, since our work's codebase **already** covers pixel-based agents and models. 

Furthermore, in addition to providing evidence for how our method solves these problems and how our PPO trained actor-critic transfers, we show that even with the same data, our algorithm is sufficient to create distinct world models that serve as *levels* or *contexts*. This demonstrates that our world model selection is relevant and can pave the way for the use of UED algorithms in more general settings.  

# Why this matters
In conclusion, we believe that the lack of focus on Dreamer-like, pixel-based world models **is not grounds for rejection**, neither is the lack of a clearly better-performing UED algorithm. That is in fact **a strength of our algorithm** which can easily switch between methods through boolean flags. Other environments requested like AntMaze would only strengthen the paper and we will finish running experiments on them, even though they are not the most commonly used environment in **accepted offline RL** papers. We are grateful that you have raised an extremely important point that will be added to the discussion section with the relevant plots and evaluations. 

A historic strength of this venue and our field is showing that work should not be judged by misconceptions induced by the historical context and recent trends, but value in opening future research avenues, and in our case, bringing online RL methods to a more general offline setting without penalties. 

[1] Ha, David, and Jürgen Schmidhuber. *World models*, 2018

[2] Hafner, Danijar, et al. *Dream to control: Learning behaviors by latent imagination*, ICLR 2020

[3] Hafner, Danijar, et al. *Mastering atari with discrete world models*, ICLR 2021

[4] Micheli, Vincent, Eloi Alonso, and François Fleuret. *Transformers are sample-efficient world models*, ICLR 2023

[5] Schmidhuber, Jürgen. *On learning to think: Algorithmic information theory for novel combinations of reinforcement learning controllers and recurrent neural world models*, 2015
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper explores the potential of using offline data to train world models that can serve as a training ground for reinforcement learning. While there are some strengths of this approach (e.g., training with only a small amount of offline data), it does appear like the main claim "full rollout training is possible" is an overclaim. MoReL uses full rollout training and so does any other approach based on model-based planning. I think while no reviewer brought up this connection, such a connection is important to ascertain the relevance and significance of this paper. I also think we need a broader set of domains, since the bar for this research should be high (including Atari, and yes Atari is standard in offline RL literature). Unfortunately we are not able to accept this paper right now.