We thank the reviewer for acknowledging our clear presentation and logical method explanation. We address the conceptual concerns below to clarify our claims and contributions.

for waekness 1) Qualitatively Different Tasks: Our results actually demonstrate stronger task diversity than initially described. In Procgen environments, longer episode horizons represent qualitatively harder tasks, not just extended versions. Figure 4 shows our method discovers 2-4 step horizon increases during training, which translates to more complex tasks as demonstrated in the tn vs. tN visual examples. The progression from simple direct paths to complex multi-obstacle courses (CoinRun), basic corridors to intricate multi-room maze navigation (Maze) demonstrates more obstacles, enemies, and spatial complexities the agent must overcome before reaching the goal. Our PLR mechanism automatically discovers this difficulty progression without manual design, showing that agents learn to solve increasingly complex instances requiring fundamentally different strategies and capabilities. We have also added two more additional Procgen diverse domains for camera ready version—Heist (puzzle-solving) and Miner (physics simulation)—achieving 19.8% and 38.5% improvements respectively.

for weakness 2) Open-Ended Learning: We appreciate the reviewer highlighting that we have not yet achieved open-ended learning in the absolute sense. Our claim should state that our work opens a path towards "open-ended learning in learned world models." which we believe is valid. Many would consider that UED algorithms fall within the general scope of “open-ended learning” and indeed PLR was used in the Adaptive Agents paper using XLand [1] which is undeniably an “open-ended learning” paper. As far as we are aware, we are the first to bring these ideas to learned world models. Of course, the world models we use are narrow in scope (modelling a distribution of procgen environments) but we believe the ideas will transfer to larger-scale foundation world models, which are increasingly modelling vast and diverse, open-ended and even unlimited task spaces. We believe our work could be a key stepping stone to using such models for open-ended learning with world model.

for weakness 3) Method Complexity: While our approach integrates multiple components, each serves a clear purpose: diffusion world models generate diverse trajectories, PLR provides principled curriculum selection, and variable horizons prevent overfitting. The integration is natural, and each component is well-established, making the overall approach interpretable. More importantly, this complexity enables capabilities impossible with simpler methods—automated curriculum discovery in offline settings.

for weakness 4) Computational Cost: We acknowledge this limitation transparently. However, it's important to contextualize these costs: Academic constraints: Our 180 GPU days represent the total experimental budget across multiple environments, seeds, and methods on academic hardware. One-time world model training: The 10-hour world model training amortizes across multiple agent training runs. Scaling benefits: More complex environments would benefit even more from automated curriculum vs. naive sampling, making the computational investment increasingly worthwhile.

for question 1) Dataset Quality Sensitivity: This attribute is indeed a fundamental limitation. Our mixed dataset design (expert+medium+random) specifically addresses this by ensuring broad state coverage. This may seem contrived in the more toy problems we address here, using an academic lab computer. However, we believe this is a proxy for foundation world models trained in industry, where they leverage vast Internet datasets that contain a huge diversity of behaviors. This is what world models are good at, and so we wanted to create a proxy for that setting. We believe this is a reasonable one.

for question 2) OMNI-EPIC Comparison: Thank you for this important reference. We should indeed discuss this related work. TLDR: This is an alternative general approach to get at the same idea, open-ended learning in a “Darwin complete” search space. If you see any talks from Jeff Clune he mentions that the two kinds of Darwin complete search spaces currently known are code or neural networks (i.e. world models). This work is the first to explore the latter for open-ended learning, and OMNI-Epic is using the former. Other differences are that OMNI methods use a foundation model to measure “interestingness” but this is orthogonal. We could use an OMNI like algorithm instead of PLR.

[1] Human-Timescale Adaptation in an Open-Ended Task Space, Bauer et al, ICML 2023

We thank the reviewer for acknowledging our method's novelty and clear presentation. We have conducted additional experiments to address some of the concerns raised. We address the technical concerns below to clarify our design rationale and broader applicability.

for weakness 1) 2D UNet Architecture: The choice of world model architecture is orthogonal to our contribution, which focuses on curriculum learning methodology. We chose 2D UNet following DIAMOND [1] for principled reasons: (1) Computational efficiency - 5 denoising steps enable real-time rollout generation crucial for curriculum learning, and (2) Proven effectiveness - DIAMOND demonstrated superior performance over discrete latent methods on visual domains. While 3D DiT architectures show promise for video generation, they require significantly more computation. Our IMAC framework is architecture-agnostic and could work with 3D DiT or other advances as they become more efficient. We will add exploration of 3D DiT architectures as a potential improvement for future work in our conclusions section.

for weakness 2) We selected PLR over alternatives like ACCEL and ADD for principled reasons: (1) Natural compatibility: PLR operates on state-trajectory pairs, perfectly matching our world model rollout structure; (2) ACCEL requires a choice of mutation function. This is somewhat arbitrary and may not make sense in our setting. It could be the horizon H of the rollouts, but it seems challenging to mutate the latent space. We think that PLR is the simplest method, and it makes sense as the first example of UED in a world model. We will add a discussion about the reviewer’s proposed alternative methods (ACCEL, MaxMC) in the conclusion section.

for weakness 3) Our main contribution is not developing a world model that beats others, but rather demonstrating that autocurricula can be highly effective when combined with modern world models. We believe that our framework is world-model agnostic—we chose our baseline world model [1] as a proven, efficient foundation, but our curriculum learning approach could be applied to other world models, including diffusion-based variants, when reward and next-token predictors are available. Regarding the Diffuser baseline: Diffuser is a planning method that requires environment access for online replanning, making it inappropriate for our offline-only setting. Our contribution is showing how UED principles can work with offline world models, not proposing superior diffusion architectures. We will add clearer framing of our methodological contribution in the introduction section of the camera-ready version to better distinguish our curriculum learning focus from world model architectural advances.

for weakness 4) Experimental Scope: To address concerns about broader applicability, we expanded our evaluation to include Heist (puzzle-solving) and Miner (physics simulation), achieving 38.5% and 19.8% improvements, respectively. These additional results will be included in the camera-ready version. This brings us to seven diverse Procgen environments, demonstrating effectiveness across different cognitive domains. Procgen scores represent meaningful improvements because: (1) these environments are specifically designed to be challenging for generalization, especially for offline settings [2]; (2) our baselines achieve low scores precisely because offline RL in procedural environments with mixed datasets is difficult; and (3) our relative improvements (17-48%) are substantial and consistent across environments. Regarding random exploration data: This highlights an important practical consideration. However: (a) many real-world datasets inherently contain exploratory data (e.g., human demonstration datasets often include suboptimal exploration), (b) we believe that our “random” quality data component serves a similar function to data collected from the Internet, representing how foundation world models are trained, and (c) the key insight is that diverse offline data enables world model training—the specific source matters less than coverage.

for question 1) Return Values: Thank you for this clarification request. Procgen environments have varying return scales across different games, making absolute values less meaningful than relative performance. Our focus is on the generalization gap—the key metric is performance on held-out test levels relative to training performance, which directly measures our method's ability to generalize beyond the training distribution. Regarding our return calculation: we use the mean return from an ensemble of reward prediction heads as stated in the paper, where returns are normalized during training. The absolute values reflect this normalization rather than raw environment scores. What matters is the consistent relative improvement our method achieves across all five environments compared to baseline methods facing identical evaluation conditions.

for question 2) Extension to Reward-Free Settings: This is an excellent question that highlights a promising future direction. Currently, our approach relies on explicit reward signals for agent training. We have not explored reward-free settings in this work. However, this represents a particularly exciting direction for foundation world models trained on large-scale, unlabeled data—exactly the type of massive offline datasets that lack explicit reward annotations. We believe this is a natural and important next step that would significantly expand the applicability of our approach to real-world scenarios in which reward engineering is challenging.

for question 3) Continuous Action Spaces: While we focus on discrete actions, our approach is fundamentally compatible with continuous control. The world model architecture can handle continuous observations, and PLR operates on states/trajectories regardless of action space dimensionality.

[1] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in Atari, 2024.

[2] Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in offline reinforcement learning. 2024.

We thank the reviewer for continued engagement with our submission and for the insightful follow-up questions. We appreciate the opportunity to clarify these important points. We address each of your follow-up points below.

1. On Related Work: We agree that a more comprehensive discussion of related work would better frame our contribution and highlight the modularity of our framework. In the camera-ready version, we will expand our related work section to include a detailed discussion of the alternative UED and architectural components you mentioned. Specifically, we will:

• Discuss alternative UED algorithms like ACCEL [4] and ADD [5], contextualizing our choice of PLR while acknowledging that these other methods could potentially be integrated into our framework, and include a discussion on alternative reward estimation strategies like MaxMC [6].

• Reference recent video generation architectures like 3D DiTs [2, 3], positioning our use of a 2D UNet as a deliberate choice for computational efficiency while noting that our curriculum learning approach is architecture-agnostic. We believe that this discussion will make it clearer that while we chose a specific set of components for our implementation, the core contribution—integrating UED with a generative world model—is a general framework open to other state-of-the-art components.

2. On Model-Based RL Baselines and Diffuser : We thank to reviewer for pushing for more clarity on this point and for correcting our imprecise language regarding Diffuser. We apologize for the confusion. You are correct that Diffuser [7] is trained offline. Our previous explanation was poorly phrased. A more precise reason for not using Diffuser as a direct baseline is the fundamental difference in the problem formulation they address:

• Diffuser and its variants are primarily goal-conditioned planning methods. They excel at generating trajectories to reach a specified goal state and are typically evaluated in goal-reaching tasks. Our method trains a generalist policy designed to maximize returns across a wide distribution of unseen levels, without a specified goal state at test time.

• Adapting a goal-conditioned planner like Diffuser to our general reinforcement learning setting (maximizing returns in Procgen) is non-trivial and would require significant modifications to its framework. That said, we take your broader point that demonstrating competitiveness with other modern model-based methods is important. While our primary contribution is the novel world model based curriculum framework for offline RL, we will add a more thorough discussion of the recent model-based RL landscape in our related work section to better situate our results and clarify the relationship between our approach and methods like Diffuser.

3. On Return Estimation and Evaluation Protocol : We sincerely apologize for the major confusion caused by our wording in the previous rebuttal. We want to clarify this critical point in the most unambiguous way possible:

• All final performance scores reported in our paper (e.g., in Table 1) are calculated using the ground-truth rewards from the actual test environments. The agent is evaluated by running it in the Procgen simulator and recording the true, unadulterated score from the environment.

• The learned ensemble of reward heads is used exclusively during the training phase for a single purpose: to provide a reward estimate for the imagined trajectories generated by our world model. The reward model is never used for reporting final evaluation results. We are very grateful you pointed this out. We will revise the experimental setup section of our paper to explicitly and clearly state that final evaluation is performed using ground-truth environment rewards to prevent any future ambiguity.

We are grateful for the constructive feedback provided in the follup up points. We believe that addressing these points will significantly strengthen our paper. We hope these clarifications have resolved your remaining concerns.

We thank the reviewer for their detailed feedback. We address each concern below and hope to clarify the contributions and scope of our work.

for weakness 1) Task Coverage: We respectfully disagree that our evaluation is insufficient. Procgen environments are specifically designed to test generalization and represent one of the most challenging benchmarks for offline RL methods [1]. Our original five environments test fundamentally different capabilities: CoinRun (platforming), Ninja (combat + platforming), Jumper (precise timing), Maze (planning), and CaveFlyer (continuous control-like navigation). This diversity tests different aspects of agent capabilities. That being said, we are pleased to share that we now have two additional environments, Heist (puzzle-solving + sequential dependencies) and Miner (physics simulation + environmental dynamics), where we get a 38.5% and 19.8% gain w.r.t. the baseline, respectively. This takes us up to seven procgen environments with more diversified task definitions, which we feel is now more comprehensive; we hope you agree!. We will add these additional results to our camera-ready version.

for weakness 2) Technical Differences from DIAMOND: While we build on DIAMOND's diffusion architecture, our key innovations include: (1) Variable-length episode generation vs. DIAMOND's fixed horizons, (2) PLR-based state prioritization for curriculum learning, (3) Ensemble uncertainty estimation for reward/termination prediction, (4) Offline-only training paradigm vs. DIAMOND's online/mixed settings, and (5) Sparse reward optimization for Procgen vs. DIAMOND's dense-reward Atari focus. Our IMAC framework is world-model agnostic—the curriculum learning methodology could enhance any generative world model.

for weakness 3) The baselines mentioned are orthogonal to our method and contribution. We are not seeking “sota”, instead, our research question is "Can curriculum learning improve offline world model training for generalization?" We implemented our method by modifying the state-of-the-art world model approach [2], but the core contribution is the curriculum learning methodology with the offline dataset. World model to world model comparisons would test "which architecture performs best," but our contribution is a methodological approach for offline curriculum learning that can be applied to any world model with reward/termination and next-state predictors. We chose to modify DIAMOND [2] as our implementation vehicle because it represents the current state-of-the-art results for visual RL settings, but the curriculum learning for the offline dataset-trained world model principle is the key innovation. Our consistent 17-48% improvements across all offline RL baselines validate the curriculum learning methodology—demonstrating that this approach can systematically improve offline generalization regardless of the underlying world model architecture.

for weakness 4) World model visualizations: The result is valuable feedback. While we don't have space for extensive visualizations, Figure 4 shows example frames from early (t₀) vs. late (tₙ) training stages, demonstrating how PLR discovers increasingly complex scenarios. These visualizations support our claims about emergent curriculum difficulty. To address the reviewer's visualization request, we will add more visualizations to our supplementary material for the camera-ready version.

for weakness 5&6) References: We will correct the STORM reference and acknowledge the relevant work "Pre-Trained Video Generative Models as World Simulators." Thank you for pointing out these important citations that improve the scholarly completeness of our work.

for question 1) A2C Choice: We chose A2C for two main reasons: (1) Standard actor-critic algorithm suitable for our POMDP setting with sparse rewards, and (2) Computational efficiency given our academic GPU constraints and the substantial cost of world model rollout-based training with PLR. Our approach is algorithm-agnostic and could work with other policy optimization methods (PPO, SAC, etc.)—A2C represents a conservative choice that demonstrates our method's effectiveness even with simpler policy optimization.

[1] Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in offline reinforcement learning. 2024.

[2] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in Atari, 2024.

Dear Reviewer Dygh,

We have carefully addressed all the points you raised in our rebuttal, and we believe the additions and explanations provided will help in evaluating the paper further. There is not much time left for us to address your additional comments; hence, we gladly ask you to evaluate our rebuttal at your earliest time, as your insights are important in ensuring a complete assessment of the work. We value the time you spent in this process and would be pleased to offer any more explanations should they be necessary.

Thank you for your time and consideration.

Authors

We thank the reviewer for the positive evaluation and for recognizing the novelty of combining world models with UED. Below, we answer the reviewer's insightful points:

for weakness 1-2) We specifically chose discrete Procgen environments, as they represent a challenging testbed for generalization, where state-of-the-art offline RL methods consistently struggle [1]. To further demonstrate broad applicability across diverse cognitive domains, we expanded our evaluation to include Heist (puzzle-solving) and Miner (physics simulation), where our method continues to show consistent improvements of 38.5% and 19.8%, respectively. These additional results will be included in the camera-ready version. We believe that the procedural generation aspect and diversity of the challenges make these environments particularly suitable for demonstrating curriculum learning benefits. While extending to continuous control is important future work, we believe establishing the core principle in this challenging discrete domain is a valuable first step. We appreciate the reviewer's constructive feedback and look forward to extending this work to more complex domains in future research.

for question 1) The distribution shift between training data and target tasks highlights an excellent point regarding the world model's potential as a bottleneck. We specifically designed our mixed offline dataset (expert+medium+random trajectories) to mitigate this concern by ensuring broad state space coverage. The random exploration component is particularly crucial here, as it provides diverse states that help the world model generalize beyond expert and medium demonstrations. We trust that our PLR mechanism provides an additional robustness layer—by prioritizing states with high learning potential (positive TD errors), we naturally focus on scenarios where the world model predictions are most informative for the agent, rather than areas where the world model might be uncertain or inaccurate. This method creates a feedback loop that steers training toward regions where the world model is most reliable.

for question 2) Scalability to complex environments. It is reasonable to assume that world models may struggle to model complex environments. That being said, progress on this front has been rapid, for example DreamerV3 modeling Minecraft [2], GAIA-2 modeling photorealistic driving scenes [3] and foundation world models like Genie2 modeling a plethora of embodied game-like worlds [4]. We believe this is now a critical time to begin exploring the use of these more capable world models for training agents, and our work is the first to demonstrate that autocurricula could play a key role

for question 3) Transferability of discovered curricula: This is a fascinating question that touches on whether curricula are policy-specific or more generally useful. While we don't have explicit experiments on this, our results suggest the curricula have some generality—we observe consistent patterns across different random seeds and diverse game tasks where PLR progressively discovers longer, more complex scenarios as training progresses. Investigating explicit curriculum transferability represents an important direction for future work that could further validate the generality of discovered curricula.

[1] Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in 437 offline reinforcement learning. 2024.

[2] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2024.

[3] Lloyd Russell, Anthony Hu, Lorenzo Bertoni, George Fedoseev, Jamie Shotton, Elahe Arani, Gianluca Corrado, GAIA-2: A Controllable Multi-View Generative World Model for Autonomous Driving, 2025

[4] Parker-Holder, P Ball, J Bruce, V Dasagi, K Holsheimer, C Kaplanis, A Moufarek, G Scully, J Shar, J Shi, et al. Genie 2: A large-scale foundation world mode, 2024