Imagined Autocurricula
Authors: Ahmet H. Güzel
Matthew Thomas Jackson
Jarek Luca Liesen
Tim Rocktäschel
Jakob Nicolaus Foerster
Ilija Bogunovic
Jack Parker-Holder
Venue: NeurIPS 2025 poster ()
Keywords: World Models
Reinforcement Learning
UED
Generalization
Primary area: reinforcement_learning
TL;DR: Unsupervised Environment Design to create automatic curricula over world model-generated environments, showing strong generalization on unseen procedural tasks while training exclusively within offline learned world models.

Abstract: Training agents to act in embodied environments typically requires vast training data or access to accurate simulation, neither of which exists for many cases in the real world. Instead, world models are emerging as an alternative–leveraging offline, passively collected data, they make it possible to generate diverse worlds for training agents in simulation. In this work, we harness world models to generate “imagined” environments to train robust agents capable of generalizing to novel task variations. One of the challenges in doing this is ensuring the agent trains on useful generated data. We thus propose a novel approach IMAC (Imagined Autocurricula) leveraging Unsupervised Environment Design (UED), induces an automatic curriculum over generated worlds. In a series of challenging, procedurally generated environments, we show it is possible to achieve strong transfer performance on held-out environments having trained only inside a world model learned from a narrower dataset. We believe this opens the path to utilizing larger-scale, foundation world models for generally capable agents.

Imagined Autocurricula
Ahmet H. Güzel*
Matthew T. Jackson University of Oxford
University College London Centre
Jarek L. Liesen
Tim Rocktäschel
Jakob N. Foerster University of Oxford
University of Oxford
University College London AI Centre
Ilija Bogunovic
Jack Parker-Holder University College London AI Centre
University College London Al Centre
Abstract
Training agents to act in embodied environments typically requires vast training data or access accurate simulation, neither of which exists for many cases the real world. Instead, world models are emerging as an allernatie--everrainng offline, passively collected data, they make possible to generate diverse worlds for training agents in simulation. In this work, we harness world models to gen- erate "imagined" environments to train robust agents capable of generalizing novel task variations. One of the challenges in doing this is ensuring the agent trains on useful generated data. We thus propose a novel approach IMAC (Imag- ined Autocurricula) leveraging Unsupervised Environment Design (UED), which induces an automatic curriculum over generated worlds. In a series of challeng- ing, procedurally generated environments, we show is possible to achieve strong transfer performance on held-out environments having trained only inside a world model learned from narrower dataset. We believe this opens the path to utilizing larger-scale, foundation world models for generally capable agents.
Introduction
Despite sigmificant advancements in AI, we remain far from generally capable agents [1]. In the 2010s, progress towards this goal was driven by breakthroughs in deep reinforcement learning (RL,, (2]) where agents such as AlphaGoo [3] showed was possible to discover new knowledge beyond human capabilities, like the famous move37. However, scaling RL in simulation to produce more general intelligence remains bottlenecked by the lack sufficiently diverse and real-world simula- tors [4. 5]. By contrast, more recent years have been dominated by agents trained from vast quanti- ties of Internet data. Here, agents have a broader knowledge of the real world but lack of ability to discover new, superhuman behaviors since they are largely trained to mimic the training data.
World models are emerging as a promising approach leverage the best of both. They "stand on the shoulders of giants" by training on vast Internet data [6. 7, 8, 9, 10]. but then enable agents to train with RL inside their generated (or imagined) environments. Thus, world models make it possible for agents to imagine potential outcomes without direct environment interaction, allowing them to train across diverse scenarios with substantially fewer real experiences [11, 12, 13, 14, 15, 16, 17]. In theory, large foundation world models could then be used to generate sufficiently rich environments for agents to learn superhuman behaviors for a swathe of embodied settings. Progress world models research has typically fallen into two categories.
"Corresponding author (ahmet.guzel.co 23 (3@uul.ac.kk)
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Rolout sata
Denolser traming
buffer
X
[5, r's
T
RewardEnd training
PLR Update and Sampling
Offine Dataset Buffeer
5
Diverse Data Collection
Learned World Model
Figure I: Architecture overview: (1) Offline Data Collection: Diverse state-action-reward- state
tuples are collected via policy Th to form the dataset. (2) World Model Training: A state transition denoiser and reward/termination predictor are trained on the dataset and frozen for consistency. (3) Agent Training: The agent learns through imagine-based autocurriculum using Prioritized Level Replay, which balances between sampling from existing experiences and generating new rollouts to update the buffer.
The first seeks to train world models from single (online offline) environments, achieving strong and sample-efficiet performance [18, 19, 17]. The latter sees more general, so-called foundation world models trained from large-scale internet data [6. 7, 9, 10], but these models have not been shown to be useful for training generalist agents.
In this paper, we attempt to bridge this gap. We focus on the offline RL setting, where agents must learn from pre-collected datasets without additional environment interaction [20]. We focus on procedurally generated environments, meaning our offline data contains sequences (action labelled videos) from variety of different levels. We then seek to transfer the learned agents to new, unseen levels from the same environment, This presents a significant challenge for existing methods, which often fail to generalize to new tasks when training purely from offline data [21,22] We believe world models can be the solution to this problem-we first train diffusion world model from the offline samples and then use it to generate new, imagined worlds. Crucially, naively generating data from world models can lead to ineffective training data, making agent learning inefficient. This ineffi- ciency stems from both the quality of the worlds themselves and the challenge of setting appropriate hyperparameters like imagined episode length across diverse tasks. To address this challenge, leverage the Unsupervised Environment Design paradigm (UED) [23])- 23heewweree a teacher proposes levels that maximize student regret. Rather than training fixed or random sets of imagined worlds, we create an Imagined Autocurricula, or IMAC. Importantly, our framework world-mode agnos- diffusion world model applied to any world model architecture that provides reward and next-token predictions. Our con- tribution lies can effectively guide agent training within learned world models from offline mixed dataset.
tic-while
foundation,
autocurriculum approach be
advancing world model architectures, but in demonstrating how UED principles
Our approach IMAC uses Prioritized Level Replay (PLR, [24]) as UED algorithm, which we show provides a natural complement to the learned world model-th the world model generates diverse po- tential training trajectories "imagined environments. while PLR strategically selects subsequent training tasks from these imagined rollouts. Figure illustrates the overall architecture of our ap- the
proach. As show, this prioritization process naturally induces automatic curriculum generated, imagined worlds, meaning the agent is exposed to increasingly challenging training tasks. As testbed for this paradigm, we leverage seven Procgen environments with high-dimensional vi- sual inputs, where both model-free [25] and model-based approaches [26] struggle to transfer knowl- edge unseen scenarios [21, 22]. The resulting agents have significantly stronger generalization performance than state-of-the-ar baselines, demonstrating strong transfer performance on held-out environments having trained only inside world model learned from a narrower dataset. We believe this opens the path to utilizing larger-scar scale, foundation world models for generally capable agents.
Our primary contribution is the demonstration that world models offer the path to training diverse offline datasets and subsequently generalizing to new, unseen tasks. Importantly, this is made possi- Imagined Autocurrieulum; rather than fixing hyperparameters randomizing
ble by training settings, each which could be problematic for learning. We believe this is the first example of open-ended learning in learned world models, and could enable significant progress towards gener- alist agents with ever more powerful world models in the future.

---

2 Preliminaries
2.1 (Offline) Reinforcement Learning
To model environments, we consider Partially Observable Markov Decision Processes (POMDPs) [2], which are defined by tuple A, R, S2, 0, S, A, and Ω2 respectively represent the set of states, actions, and observations. (s [0t; 81) the conditional transition distribution over next state 81+1 E S given previous action E A and state 8 E S. Each transition produces an observation 0s+1 E sampled from 0(044118441 az), and reward T4+1 R sampled from R(r++1|$4+1) dt, 8t).
An agent interacting with the POMDP observes on and maintains a belief state be based on the history of observations and actions. The agent then selects an action a, according to its policy. Re- inforcement learning aims to find the optimal policy - a conditional distribution a( [he) that max- imizes the expected discounted return E |Et= y'r,), where is the reward received at time step t,
y E [0. 1) is the discount factor, and h, represents the history of observations and actions up time In offline reinforcement learning, the agent cannot interact directly with the POMDP, but is given a dataset of previously recorded interactions (0t, at, 1 0t+1). There are no restrictions on how D is collected, so it may include data from mix of policies with varying levels of expertise and state space coverage. Generalizing to states that differ from those D is the main challenge offline RL.
World Models
A world model [27] is parameterized model of an environment that can be used to train rein- forcement learning agents and autoregressively generate synthetic data. World models are able to generate imaginary trajectories by sampling from the following joint distribution:
T-1 II =(aa)s)))((sss))) (()
Creating world model for POMDP reMurrequires modeling the transition and observation distributions and O, and reward distribution R. Given a dataset of interactions, we optimize the model parameters to minimize expected loss
min E, [To( oi) + Ln(fe,rt)). f,0.k
where o and ft the world model predictions of o and r, given at, 0t-1, and Lo: I appro- priate loss functions. In practice, instead of predicting the next observation based on the prediction for the next state, can predict it directly from the previous observation and action buffer, learning state transition dynamics implicitly via a hidden state (e.g., that of a recurrent neural network). Ad- ditionally, it is common to predict the residual in the observation Λθ, ō.+1 oc, which biases the world model towards temporal consistency in the autoregressive prediction.
2.3 Diffusion Models
Inspired by non-equilibrium thermodynamics, diffusion models [28] generate data by reversing gradual noising process. The forward process corrupts a data sample Xy ~ ((xp) into pure Gaussian noise Xy ~ N IO, through Markov chain of Gaussian transitions:
q(x, x,-1) = N(x; V1 1111-xx00 x,,
(3)
where {3,17- is a fixed variance schedule. The generative model learns the reverse pro- cess (x;-1 xı) by approximating the score function, i.c., the gradient of the log posterior Vx, log 9(x; Xo). In practice, neural network (x(x,,)) t) is trained to predict the noise added during the forward process, and the training objective is typically:
Caipll Exo.,.c [||c
co t)1²]
(4)
where IO, and x, Vo, xo vI aie e with II'' (1 :3).
3

---

This score matching framework, introduced by [29], enables estimation of score models directly from empirical data without requiring explicit knowledge of the true underlying score function. While standard diffusion models typically function unconditional generators of a data distribution Pana(1) (x), adapting them for sequential racision-making requires conditioning on past information to predict future states.
To employ diffusion models world models, we extend Eæ model the environment's dynamics p(Z++1 [=ge; [ age),, where Ft++1 is the next observation, while &gr and agr denote the history of past observations and actions. This adaptation is particularly relevant for POMDPs, where the true environmental state must be inferred from observation and action sequences. The modified score model Eg is trained by extending the standard diffusion objective to:
((0) co( T, .0.00
(5)
where aft1 is the noised version of the next observation noise level T, and is the noise added during the forward process.
To generate subsequent observation 2t+1, we iteratively solve the reverse-time Stochastic Dif- ferential Equation (SDE)-o) or an equivalent Ordinary Differential Equation (ODE)-guided by our score model. While various numerical solvers can be employed, there exists a practical trade-off: sample quality typically improves with increased Number of Function Evalua- tions (NFE), directly impacting the computational cost of the diffusion-based world model during inference.
2.4 Prioritized Level Replay
When training agents to solve procedurally generated environments, previous methods have com- monly relied on continuously and uniformly sampling new levels during training. This can be sample-inefficient since many sampled levels do not contribute significantly to the agent's learn- ing progress. As an alternative, Prioritized Level Replay (PLR) [24] has been proposed, where levels are replayed based on their potential to induce greater learning progress in the agent. In orig- inal PLR implementation, it maintains a buffer of previously encountered levels, including a score that represents an estimate of learning potential and a staleness factor that represents how long ago that level was sampled. There are many possible choices for computing the score - a simple yet effective one is given by the TD-error d, rt + (((++1) V(s)). During training, we randomly select between sampling randomly generated level with a selected episode length range randomly and sampling one from the buffer. When sampling from the buffer, level 4; is selected for replay with probability
4 (1-p) (l(( Assen + p- Pc(( |Neen, C,c),
(6)
which mixture distribution of Ps and Po with parameter p. They prioritize levels with high scores and high staleness, respectively. Here Ageen is the replay buffer, S is the level scores, C' is the last time each level was sampled, and is how many levels were sampled since the beginning of training. In our world model setting, we save both the initial state and the desired imagination horizon h in the buffer. When the algorithm samples an entry from the buffer, it uses the same initial state to begin imagination and runs the world model forward for exactly h steps. By interpreting world models as procedural generators of agent experience, we can directly apply PLR prioritize the most valuable imagined experiences. This yields a novel method for model- offline RI. that efficienlly leverages imagined trajectories.
3 Imagined Autocurricula
Our approach consists of three key components: (1) diffusion world model trained from diverse offline data, (2) using this model to generate imagined random episode length rollouts for agent training, and (3) implementing an auto curriculum over these imagined environments.
3.1 Diffusion World Model Training
First, we collect diverse visual offline dataset of different behaviors (expert, medium, and random) via behavior policy as D = ((0,, 0n Ti, o Then we train the diffusion model as conditional

---

generative model of environment dynamics, p( a<t), handling the POMDP setting where the true Markovian state is unknown and must be approximated from history. Our training objec- tive follows Equation 5. Following DIAMOND [11], we utilize the Elucidating Diffusion Models (EDM) formulation [30] rather than the traditional DDPM [28] approach. This critical design choice enables our world model to remain stable over long time horizons with significantly fewer denois- ing steps. During training, we sample trajectory segments from our dataset and apply noise using a r-level perturbation kernel to create training pairs. To incorporate temporal context, we maintain buffer containing L = 4 previous observations and actions as conditioning information. Our diffu- sion world model implements standard 2D U-Net architecture [31] the core denoising network. The temporal information flows through two pathways: past observations are concatenated channel- wise with the noisy observation being denoised, while action information is integrated via adaptive group normalization layers [32] within residual blocks [33]. We iteratively solve the reverse SDE using Fuler's s method with exactly 5 denoising steps. We also follow their approach of using full- image observations (64x64 3-channel for Procgen) rather than discrete latent space representations, as they demonstrated superior performance over other state-of-the-arr world models that use discrete latent space. Given the similar image dimensions between Atari and our Procgen environment, we leveraged their extensively validated hyperparameter settings for our diffusion model. Their compre- hensive ablation study showed these parameters achieve both notable image generation quality and computational efficiency, making them an ideal foundation for our work. This approach allows our diffusion model to capture fine visual details critical for reinforcement learning for Procgen agents, such as small rewards and enemies, while maintaining temporal consistency across long sequences.
3.2 Agent Training in Imagined Environments
With our diffusion world model trained, we implement auxiliary predictors for reward and termi- nation signals using ensemble of E' prediction heads to capture uncertainty during imagined rollouts. While the base architecture follows Alonso et al. [1I], employing a separate model Rgg with standard CNN [34] and LSTM [35] layers to handle partial observability, we extend their ap- proach through our ensemble of prediction heads, which provides crucial uncertainty estimates dur- ing long-hov hong-oorizon planning in POMDP setting. This model is trained on the same offline dataset used for the world model, optimizing for accurate reward and episode termination predictions, with the agent using the mean of the ensemble predictions during training. Building upon these diffusion and auxiliary models, the RL agent follows an actor-critic architecture parameterized by a shared CNN-LSTM backbone with separate policy and value function heads. The agent is trained using Advantage Actor-Crite (A2C) [36] with X-returns for advantage estimation as in [37]. The pol- icy head Fæ is optimized to maximize expected returns while the value network Va minimizes the temporal difference error between predicted values and the computed A-returns.
During training, we generate imagined trajectories by first sampling an initial state 8o from our di- verse offline dataset, then autoregressively producing sequences of states, rewards, and termination signals by rolling out our policy through the diffusion denoiser model and ensemble auxiliary pre- dictors. Unlike the fixed horizon approach in [11], we randomly sample the imagined episode length between minimum and maximum horizon bounds to incorporate greater diversity in training expe- riences and prevent the agent from exploiting fixed-horizon dynamics. Finally, we compute policy gradients and value updates based on these imagined experiences, enabling effective policy learning without additional environment interaction.
3.3 Autocurriculum
Our IMAC implementation follows the PLR framework [24], maintaining a replay buffer of previ- ously encountered initial states so with associated priority scores. For each training iteration, with probability 0.5, we train the agent using the current policy and world model on initial states sampled from this prioritized buffer. With the remaining 0.5 probability, we select random initial state and a random imagined episode length, T. to explore the environment space and update the buffer. The priority score for each initial state su computed using the agent's 's temporal difference errors, dg. which serve as a proxy for learning potential. The specific formula for the score is:
score(su)
1
Σ t= k=1
(a))* max 8x)
T
5

---

Here, T is the length of the imagined episode (which is randomly chosen during the exploration phase), 0g the temporal difference error at timestep k, Y is the discount factor, and is the trace decay parameter. The inner sum, - max(0, 6x), represents a truncated, discounted, and X-weighted sum of future positive TD errors from timestep t to T. We specifically focus posi- tive TD errors through the max(0, 6x) operation because they indicate instances where the agent's s value estimate
too low, suggesting unexpectedly good outcomes that are particularly valuable for learning. The outer sum and division by 7 average this sum over all timesteps / in the episode. Based on this prioritization mechanism, we believe that our approach creates three key advantages for generalization: First, by focusing on the scenarios with the highest learning potential, the agent spends more computation on informative experiences rather than those where little can be learned. Second, the random horizon setting creates diversity in training experiences, preventing the agent from exploiting fixed-length dynamics. Finally, the prioritization mechanism naturally discovers an emergent curriculum that gradually increases in difficulty as the agent improves. The emer- gent curriculum effect is particularly important for procedurally generated environments, where the difficulty distribution can be complex and unknown a priori. Our IMAC approach discovers this difficult landscape dynamically during training. focusing on the boundary of the agent's current capabilities- PLR authors call "threshold levels" that provide the highest learning signal. As the agent improves, this boundary shifts to more challenging scenarios, creating a curriculum that scales with the agent's s abilities without requiring manual design or environment-specife knowledge. Importantly, unlike previous autocurriculum learning approaches that require direct control of the environment generation process, our method works entirely within the imagined andom-lengtt roll- outs from our world model, which we call our imagined autocurricula, making applicable to any environment where a world model can be trained, including those where the procedural generation process is black box. Details of model architecture, algorithm, training hyperparameters used, and dataset details for this work are given in Appendices A and B, C, and D.
4 Experiments
4.1 Procgen Benchmark
For comprehensive evaluation of our approach, we use a challenging subset of the Procgen Bench- mark [38], collection of procedurally generated environments designed specifically to test general- ization in reinforcement learning. Unlike the Atari benchmark used in previous world model studies [11, 16, 15, 17, 39]. Procgen environments are procedurally generated and thus inherently test an agent's ability to generalize to unseen levels with the same underlying mechanics.
We created a mixed offline dataset capturing diverse interaction patterns to represent the type of diverse data we may see in Internet videos. Similar to "mixed" datasets in existing benchmarks [40, 22] it consists of three complementary components: Expert trajectories (Dexpen) from fully trained PPO agents, Medium-quality trajectories (Dediiim) from partially trained agents reaching approximately 50% of expert performance, and Random exploration trajectories (Dramdom) to ensure broad state space coverage. This mixture forms dataset D Drandom U Dexpen U Dmediuu totaling I million transitions per game across 200 training levels. We focus on seven environments from Procgen: CoinRun, Ninja, Jumper, Maze, and CaveFlyer, selected based on their poor performance in previous offline RL benchmarks [21] and their sparse reward structure. For each environment, we use 200 levels for training our world model while evaluating generalization on unseen test lev- els (201-co). After data collection, we trained the world model and reward/ternination predictors (requiring approximately 10 hours on an NVIDIA RTX 4090), then froze these components for gen- erating imagined rollouts during agent training. Each agent seed's training required approximately 4 days on a single RTX 4090, with experiments conducted across 3 random seeds for 5 million total steps each, resulting in 180 GPU days of computation.
We compare our approach with several established offline RL methods including Behavior Cloning (BC), Conservative Q-Leaming (CQL) [41], Implicit Q-Learning (IQL) [42]. Batch-Constrrinee Q- learning (BCQ) [43], Behavior- Constrained Transformers (BCT), and Decision Transformers (DT) [44]. We also conducted an ablation study comparing our approach against world model baselines, evaluating the impact of both fixed (15 steps) and variable (between and 22 steps) imagined episode lengths.

---

4.2
Procgen Benchmark Results
Table I: Generalization results: All values represent the mean return over three random seeds +
standard deviation when transferring agents held out levels on the Procgen benchmark.
Method
CoinRun
Ninja
Jumper
Maze
CaveFlyer
Heist
Miner
BC
5.31 + 0.31
4.22 + 0.21
3.92 + 0.42
4.03 + 0.29
3.12 + 0.17
2.1 + 0.22
4.68 + 0.11
CQL
5.12 + 0.22
3.29 + 0.32
2.24
1.82
1.89 0.27
0.3
0.32 + 0.09
0.09
IQL
4.96 + 0.20
3.11
2.35
1.98
0.27 +
0.5 0.08
BCQ
4.33 0.14
3.24
2.19
1
2.61
0.14
0.42 0.13
BCT
5.11
0.28
3.72
3.98
1: 0.19
2.76
0.23
0.10
DT
5.03 0.18
4.03
3.84
4.32 0.11
3.04
0.32
1.82
1.11
iMac
6.10
5.12
5.78
5.41
0.19
3.87 0.22
2.91 0.15
5.61 = 0.13
Table 2: World model ablation: Performance evaluation on the Procgen benchmark all values
represent return over three random seeds + standard deviation
Method
CoinRun
Ninja
Jumper
Maze
CaveFlyer
Heist
Miner
WM Fixed 4.76 + 0.19 3.78 = 0.21 3.71 + 0.19 3.92 = 0.32 2.94 + 0.13 2.05 + 0.14 3.18 + 0.10
WM Random 5.36 + 0.15 4.33 = 0.18 4.11 0.11
4.62 = 0.15 3,02 + 0.19 2.19 - 0.18 4.55 + 0.15
iMac
6.20 1 0.14 5.15 1 0.12 5.78 1 0.21 5.41 1 0.19 3.87 1 0.22 2.91 1 0.15 5.61 1 0.13
The results in Table demonstrate that IMAC method consistently outperforms state-of-the-art of- fline RL algorithms across all evaluated Procgen environments. Our approach achieves improve- ments of 17% on CoinRun, 19% on Ninja, 48% on Jumper. 35% on Maze, 24% on CaveFlyer, 38% on Heist, and 19% on Miner compared to the best-performing model-free baseline for each environment.
Table further demonstrates IMAC's effectiveness through ablation studies comparing different horizon-setting strategies. While both fixed and random episode length world model approaches show strong performance, our IMAC method incorporating prioritized initial state selection consis- tently delivers improved resuiis-achieving up 56% improvement over the fixed horizon baseline on Jumper and 38% on Maze. These results provide compelling evidence that procedural general- ization world models benefits significantly from our autocurriculum approach, which automati- cally discovers and focuses training on the most informative scenarios. The substantial performance advantage across diverse environments demonstrates that IMAC effectively addresses the core chal- lenge of generalization from limited offline datasets in procedurally generated environments. We provide all the hyperparameters used in model-free baseline algorithms in Appendix F.
To provide a comprehensive statistical analysis our approach's 's effectiveness, Figure 2 presents all evaluated environments, offering quan-
the aggregated mean performance improvements titative evidence of the consistent generalization benefits achieved by our method compared to the baselines presented in both tables.
IMac
BCT
WM Random
COL
IQL
3.5
4.5
Test Performance
Test Performance
Figure 2: Mean test performance across seven Procgen environments. Left: Comparison of model- free baselines against our proposed IMAC method. Right: World model ablation study comparing WM variants (Fixed and Random) against the full iMac approach
7

---

R ramdom
RPLR
Rrandom
LPLR
Lyandag
0.12
56
0.-
IN
0.5
Imasined Eptoode Becum
02
0.1
-0.08
-2.4
8M
114
241 300
4H
SM BM
14 2M 411 5M
Figure 3: Evolution of performance and episode length during training with our Imagined Autocur- ricula approach. Left: Mean episodic return of the random baseline (Raand))) showing steady improvement over training steps. Right: The difference between PLR and Random approaches in terms of both return (RpLR Rondom) and imagined episode length (LPLR LRandom).
Our results provide compelling evidence for the effectiveness of the emergent curriculum created by our Imagined Autocurricula approach. As shown Figure 3, initially PLR selects similar episode lengths as random sampling while achieving higher returns for CoinRun. However, after approxi- mately 4M steps, significant shift occurs-PL. -PLR begins prioritizing substantially longer episodes while temporarily accepting lower comparative returns. This demonstrates the emergent curricu- lum effect: our prioritization mechanism automatically discovers and focuses on more challenging. temporally-extendee scenarios aps the agent's s capabilities improve.
Figure 4: Left: Difference in imagined episode length between PLR and Random baselines across seven Procgen environments. Right: Example frames from early f, and late fx training stages, with red boxes representing agents and orange boxes representing rewards.
The visual examples in Figure 4 further illustrate this curriculum effect across all environments. Comparing initial states to with later states tx. we observe clear progression in complexity from simple, direct paths to complex obstacle arrangements for CoinRun, and from straightforward cor- ridors intricate maze structures requiring extended planning in Maze. The image pairs show representative states from early and late training stages, illustrating the progression from simpler to more complex scenarios. This progressive increase in horizon differences correlates directly with our method's performance on test environments, supporting our hypothesis that learning temporally extended behaviors is crucial for generalization in procedurally generated environments. What's particularly notable is that this curriculum emerges naturally from our prioritization mechanism without any explicit difficulty programming or environment-speciffe knowledge. Our method dis- covers this effective learning progression while training exclusively on imagined rollouts from fixed world model trained on offline data, validating our core claim that prioritizing states based on their learning potential substantially improves generalization capabilities.
5 Related Work
Generalization in Reinforcement Learning has been extensively studied, though primarily in on-
line settings where agents can actively collect new data during training [45, 46, 47,38, 48, 49,50, 51, 52, 53, 54, 55, 56, 57, 58, 59]. Benchmarks for evaluating generalization capabilities have evolved from environment suites like Procgen [38] and NetHack [41] for online settings, to more recent offline alternatives such as V-D4RI. and OFFLINE PROCGEN [22, 21J. These offline benchmarks

---

present unique challenges, particularly in visual domains where balancing dataset scale and acces- sibility becomes Brucial-lrre--caall datasets like Atari, StarCraft, and MineRL contain millions of samples but require substantial computational resources [60, 61, 62], while more accessible alterna- tives offer 100,000 to million samples per environment. Despite this progress, generalization in offline RL remains relatively unexplored compared to its online counterpart, especially when dealing with procedurally generated environments where model-free methods have shown significant limi- tations. Our work addresses this gap by leveraging world models to improve generalization across diverse procedurally generated environments without requiring additional environment interaction beyond the initial offline dataset.
World Models have evolved significantly since the first introduction of the concept of reinforcement learning within an imagined neural network environment [27]. Early applications to Atari games through SimPL.e [39] established the Atari 100k benchmark, highlighting the sample efficiency benefits of model-based approaches. The field advanced substantially with DreamerV2 [17], which pioneered learning directly from the latent space of recurrent state space model (RSSM), while DreamerV3 [12] achieved human-level performance across diverse domains using consistent hyperparameters. Recent architectural innovations include transformer-based approaches such as TWM [16] and STORM [32], which adapt DreamerV2's and DreamerV3's RSSM frameworks respectively, employing different tokenization strategies. IRIS [15] represents an alternative approach, constructing language of image tokens through discrete autoencoder and composing these tokens temporally using an autoregressive transformer. Most recently. DIAMOND [11] introduced diffusion models as an alternative to discrete latent representations for world modeling, demonstrating that improved visual fidelity can lead to better agent performance, particularly for tasks where fine visual details are critical. Recent advances in video generation, including 3D DIT architectures like Cog VideoX [63], Cosmos [9], and realistic world generators like Genie 3 [8] offer promising alternatives for world modeling. While we employ 2D UNet for computational efficiency, our curriculum learning framework is architecture- agnostic and could leverage these more advanced architectures as they become computationally more feasible. There has also been work using world models to generalize to unseen tasks [64, 65], however these were focused on domain randomization setting (e.g. our "random" " baseline) rather than an automated curriculum. By contrast, our work builds upon these foundations and leverages PLR to generate an emergent curriculum, which subsequently improves generalization unseen tasks.
Curriculum Learning in RL applies the principle of gradually increasing task difficulty to re- inforcement learning, improving both training efficiency and generalization capabilities [66, 67]. Determining learning priorities is central to effective curriculum design, with various measures in- cluding reward signals [68], learning progress [69], TD-errors [70], and state novelty [71]. Recent advances in automated curriculum learning [72] have reduced reliance on manual design, with meth- ods like Prioritized Level Replay (PLR) [24] demonstrating that targeting environments with the highest learning potential creates an emergent curriculum that progressively challenges the agent the frontier of its capabilities. While we employ PLR in our implementation, our framework is compatible with other UED algorithms including ACCEL [73]. which uses evolution strategies for curriculum design; ADD [74], which leverages diffusion models for adversarial environment genera- tion. These methods could potentially be integrated into our framework, highlighting the modularity of our approach. Our work applies this curriculum-bart based approach to world models, creating an imagined autocurriculum that operates entirely within generated rollouts from a fixed world model trained on offline data, enabling effective generalization without requiring explicit environment con- trol or online interaction.
6 Limitations
While our approach demonstrates promising results, we acknowledge certain constraints inherent to offline world modeling. The quality of generated scenarios naturally depends on training data diver- sity, though our mixed-expertise dataset mitigates this concern in practical settings. Implementation of IMAC requires substantial computational resources, and its effectiveness varies across environ- ment types. Future work could address these limitations through efficiency optimizations, domain- adaptive world modeling techniques, automated parameter tuning, and improved uncertainty quan- tification. Additional promising directions include innovative data augmentation strategies and ex- tensions to continuous control domains with higher-dimensional state spaces-al all building upon our demonstrated foundation. We provide computational cost analysis for our method in Appendix G.

---

7 Conclusion
We introduced IMAC, a novel approach combining iiffusion-basee world models with automatic curriculum learning to train general agents. Our method leverages offline data to generate diverse imagined environments and employs Prioritized Level Replay to create emergent curriculum that adapts to the agent's 's capabilities. Experimental results across seven Procgen environments demon- strate that IMAC consistently outperforms tfte-of-the-ar offline RL algorithms and world model baselines, achieving improvements of up to 48% compared to model-free approaches and 56% over exed-horizon world models. The emergent curriculum proves particularly effective for environ- ments requiring extended planning horizons and complex navigation. Our work demonstrates that strong transfer performance held-out environments is possible using only imagined trajectories from a world model trained on limited dataset. This finding opens a promising path toward uti- lizing larger-scale foundation world models for developing generally capable agents that can handle novel task variations.
Future work should address the current limitations of our approach by exploring techniques to reduce computational requirements, investigating novel curriculum learning methodologics with world model architectures with stronger uncertainty quantification, and developing methods to au- tomatically adapt prioritization parameters based on environment characteristics. Additionally, ex- tending IMAC continuous control domains and environments with higher complexity extremely sparse rewards represents an important direction for expanding the applicability of our method.
8
Acknowledgements
This work supported by the Engineering and Physical Sciences Research Council (EPSRC) EP/S021566/1.
References
[1] Meredith Ringel Morris, Jascha Sohl-dickstein, Noah Fiedel, Tris Warkentin, Allan Dafoe, Aleksandra Faust, Clement Farabet, and Shane Legg. Levels of AGI: Operationalizing Progress on the Path to AGI. In The International Conference on Machine Learning, 2024.
[2] Richard Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018.
[3] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanetot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukeuoglu, Thore Graepel, and Demis Hassabis. Mas- tering the game of Go with deep neural networks and tree search, January 2016.
[4] Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, Nat McAlcese, Nathalie Bradley-Schmice
Nathaniel Wong, Nicolas Porcel, Roberta Raileanu, Steph Hughes-Fitt, Valentin Dalibard, and Wojciech Marian Czarnecki. Open-Ended Learning Leads to Generally Capable Agents, July 2021.
[5] Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmicg, Michael Chang, Natalie Clay, Adrian Col- lister, Vibhavari Dasagi, Lucy Gonzalez, Karol Gregor, Edward Hughes, Sheleem Kashem, Maria Loks- Loks-Thompson. Hannah Openshaw. Jack Parker-Holder. Shreya Pathak, Nicolas Perez- Nieves, Nemanja Rakicevic, Tim Rocktäschel, Yannick Schroecker, Jakub Sygnowski, Karl Tuyls, Sarah York, Alexander Zacherl, and Lei Zhang. Timesclee Adaptation in Open-Ended Task Space, January 2023.
[6] J Parker- Ball, Bruce, V Dasagi, K Holsheimer, C Kaplanis, A Moufarek, Scully, J Shar, Shi, et al. Genie 2: A large-scale foundation world model, 2024.
10

---

[7] Jake Bruce, Michael Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi. Edward Hughes, Matthew Lai. Aditi Mavalankar, Richie Steigerwald, Chris Apps, Yusuf Aytar, Sarah Bechtle, Feryal Behbahani, Stephanie Chan, Nicolas Heess, Lucy Gonzalez, Simon Osindero, Sherjil Ozair, Scott Reed, Jingwei Zhang. Konrad Zolna, Jeff Clune, Nando de Freitas, Satinder Singh, and Tim Rocktäschel. Genie: Generative Interactive Environments. In The International Conference on Machine Learning, 2024.
[8] P.ilipJ. J. Ball, Jakob Bauer, Frank Belletti, Bethanie Brownfield, Ariel Ephrat, Shlomi Fruchter, Agrim Gupta, Kristian Holsheimer, Aleksander Holynski, Jiri Hron, Christos Kaplanis, Mar- jorie Limont, Matt McGill, Yanko Oliveira, Jack Parker-Holder, Frank Perbet, Guy Scully, Jeremy Shar, Stephen Spencer, Omer Tov, Ruben Villegas, Emma Wang. Jessica Yung, Cip Bactu, Jordi Berbel, David Bridson, Jake Bruce, Gavin Buttimore, Sarah Chakera, Bilva Chan- dra, Paul Collins. Alex Cullum, Bogdan Damoc, Vibha Dasagi, Maxime Gazeau, Charles Gbadamosi, Woohyun Han, Ed Hirst, Ashyana Kachra, Lucie Kerley, Kristian Kjems, Eva Knoepfel, Vika Koriakin, Jessica Lo. Cong Lu, Zeb Mehring, Alex Moufarek, Henna Nand- wani, Valeria Oliveira, Fabio Pardo, Jane Park, Andrew Pierson, Ben Poole. Helen Ran, Tim Salimans, Manuel Sanchez. Igor Saprykin, Amy Shen, Sailesh Sidhwani, Duncan Smith, Joe Stanton, Hamish Tomlinson, Dimple Vijaykumar, Luyu Wang, Piers Wingfield, Nat Wong, Keyang Xu, Christopher Yew, Nick Young. Vadim Zubov, Douglas Eck, Dumitru Erhan, Ko- ray Kavukcuoglu, Demis Hassabis, Zoubin Gharamani, Raia Hadsell, Aäron van den Oord, Inbar Mosseri, Adrian Bolton, Satinder Singh, and Tim Rocktäschel. Genie 3: A new frontier for world models. Google Deepmind, 2025.
[9] NVIDIA, Niket Agarwal, Arslan Ali, Maciej Bala, Yogesh Balaji, Erik Barker, Tiffany Cai, Prithvijit Chattopadhyay, Yongxin Chen, Yin Cui, Yifan Ding, Daniel Dworakowski, Jiao- jiao Fan, Michele Fenzi, Francesco Ferroni, Sanja Fidler, Dieter Fox, Songwei Ge, Yunhao Ge, Jinwei Gu, Siddharth Gururani, Ethan He, Jiahui Huang, Jacob Huffman, Pooya Jan- naty, Jingyi Jin, Seung Wook Kim, Gergely Klár, Grace Lam, Shiyi Lan, Laura Laxie Anqi Li, Zhaoshuo Li, Chen-Hsuan Lin, Tsung-Yi Lin, Huan Ling, Ming-Ye Yu Liu, Xian Liu, Alice Luo, Qianli Ma, Hanzi Mao, Kaichun Mo, Arsalan Mousavian, Seungjun Nah, Sri- harsha Niverty, David Page. Despoina Paschalidou, Zeeshan Patel, Lindsey Pavao, Morteza Ramezanali, Fitsum Reda, Xiaowei Ren, Vasanth Rao Naik Sabavat, Ed Schmerling, Stella Shi, Bartosz Stefaniak, Shitao Tang, Lyne Tchapmi, Przemek Tredak, Wei-Cheng Tseng, Jibin Varghese, Hao Wang, Haoxiang Wang. Heng Wang, Fing-Chun Wang, Fangyin Wei, Xinyue Wei, Jay Zhangjie Wu, Jiashu Xu, Wei Yang. Lin Yen-Chen, Xiaohui Zeng, Yu Zeng, Jing Zhang. Qinsheng Zhang, Yuxuan Zhang, Qingqing Zhao, and Artur Zolkowski. Cosmos world foundation model platform for physical ai, 2025.
[10] Sherry Yang. Yilun Du, Seyed Kamyar Seyed Ghasemipour, Jonathan Tompson, Leslie Pack Kaelbling, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. In The Twelfth International Conference Learning Representations, 2024.
[11] 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.
[12] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse do- mains through world models, 2024.
[13] Quentin Garrido, Mahmoud Assran, Nicolas Ballas, Adrien Bardes, Laurent Najman, and Yann LeCun. Learning and leveraging world models in visual representation learning, 2024.
[14] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Ken Goldberg, and Pieter Abbeel. Day- dreamer: World models for physical robot learning, 2022.
[15] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers sample-efficient world models, 2023.
[16] Jan Robine, Tobias Uelwer, and Stefan Harmeling. Smaller world models for reinforcement learning, 2023.
[17] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models, 2020.

---

[18] Danijar Hafner, Timothy Lillicrap. Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels, 2018.
[19] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Represen- tations, 2020.
[20] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems, 2020.
[21] Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in offline reinforcement learning, 2024.
[22] Cong Lu, Philip Ball, Tim G. J. Rudner, Jack Parker-Holder, Michael A. Osborne, and Yee Whye Teh. Challenges and opportunities in offline reinforcement learning from visual observations, 2023.
[23] Michael Dennis, Natasha Jaques, Eugene Vinitsky, Alexandre Bayen, Stuart Russell, Andrew Critch, and Sergey Levine. Emergent complexity and zero-shot transfer via unsupervised en- vironment design, 2020.
[24] Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. Interna- tional Conference on Machine Learning, pages 4940-4950. PMLR, 2021.
[25] Rafael Figueiredo Prudencio, Marcos R. O. A. Maximo, and Esther Luna Colombini. A survey on offline reinforcement learning: Taxonomy, review, and open problems, August 2024.
[26] Haoyang He. A survey on offline model-based reinforcement learning, 2023.
[27] David Ha and Jürgen Schmidhuber. World models, 2018
[28] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020.
[29] Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching, 2005.
[30] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- generative models, 2022.
[31] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation, 2015.
[32] Ruimao Zhang, Zhanglin Peng, Lingyun Wu, Zhen Li, and Ping Luo. Exemplar normalization for learning deep representation, 2020.
[33] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
[34] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Backpropagation applied to handwritten zip code recognition, 1989.
[35] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory, November 1997.
[36] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy Lill- icrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning, 2016.
[37] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation, 2015,
[38] Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning, 2020.
[39] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Kon- rad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mo- hiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model-based reinforcement learning for atari, 2024.
12

---

[40] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4/rI}: Datasets for deep data-driven reinforcement learning, 2021.
[41] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning, 2020.
[42] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning, 2021.
[43] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration, 2019.
[44] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, 2021.
[45] Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations, 2018.
[46] Marlos C. Machado, Marc G. Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents, 2017.
[47] Charles Packer, Katelyn Gao, Jemej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning, 2019.
[48] Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. survey of zero-shot generalisation in deep reinforcement learning, 2023.
[49] Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa, Julian Togelius, and Sebastian Risi. Illuminating generalization in deep reinforcement learning through proce- dural level generation, 2018.
[50] Alex Nichol, Vicki Pfau, Christopher Hesse, Oleg Klimov, and John Schulman. Gotta learn fast: A new benchmark for generalization in 2018
[51] Heinrich Küttler, Nantas Nardelli, Alexander Miller, Roberta Raileanu, Marco Selvatici, Edward Grefenstette, and Tim Rocktäschel. The nethack learning environment, 2020.
[52] Emmanuel Bengio, Joelle Pineau, and Doina Precup. Interference and generalization in tem- poral difference learning, 2020.
[53] Martin Bertran, Natalia Martinez, Mariano Phielipp, and Guillermo Sapiro. Instance based generalization in reinforcement learning, 2020.
[54] Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan P Adams, and Sergey Levine. Why generalization in rl difficult: Epistemic pomdps and implicit partial observability, In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Ad- vances Neural Information Processing Systems, volume 34, pages 25502-25515. Curran Associates, Inc., 2021.
[55] Jiafei Lyu, Le Wan, Xiu Li, and Zongqing Lu. Understanding what affects generalization gap visual reinforcement learning: Theory and empirical evidence, 2024.
[56] Andy Ehrenberg, Robert Kirk, Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. A study of off-policy learning in environments with procedural content generation. In ICLR Workshop on Agent Learning in Open-Endedness, 2022.
[57] Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal. Learning dy- namics and generalization in reinforcement learning, 2022.
[58] Mhairi Dunion, Trevor Mclnroe, Kevin Sebastian Luck, Josiah Hanna, and Stefano V. Al- brecht. Conditional mutual information for disentangled representations in reinforcement learning. In Conference Neural Information Processing Systems, 2023.
13

---

[59] Abdulaziz Almuzairee, Nicklas Hansen, and Henrik I. Christensen. A recipe for unbounded data augmentation in visual reinforcement learning, 2024.
[60] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning, 2020.
[61] Oriol Vinyals, Timo Ewalds, Sergey Bartunov, Petko Georgiev, Alexander Sasha Vezhnevets, Michelle Yeo, Alireza Makhzani, Heinrich Küttler, John Agapiou, Julian Schrittwieser, John Quan, Stephen Gaffney, Stig Petersen, Karen Simonyan, Tom Schaul, Hado van Hasselt, David Silver, Timothy Lillicrap, Kevin Calderone, Paul Keet, Anthony Brunasso, David Lawrence, Anders Ekermo, Jacob Repp, and Rodney Tsing. Starcraft ii: new challenge for reinforce- ment learning, 2017.
[62] Linxi Fan, Guanzhi Wang, Yunfan Jiang, Ajay Mandlekar, Yuncong Yang, Haoyi Zhu, Andrew Tang, De-An Huang, Yuke Zhu, and Anima Anandkumar. Minedojo: Building open-ended embodied agents with internet-scale knowledge, 2022.
[63] Zhuoyi Yang. Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang. Jiazheng Xu, Yuanming Yang. Wenyi Hong, Xiaohan Zhang, Guanyu Feng, Da Yin, Yuxuan Zhang, Weihan Wang, Yean Cheng, Bin Xu, Xiaotao Gu, Yuxiao Dong, and Jie Tang. Cogvideox: Text-to-video diffusion models with an expert transformer, 2025.
[64] Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts. Augmented world mod- els facilitate zero-shot dynamics generalization from single offline environment. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 619-629, PMLR, 18-24 Jul 2021.
[65] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-awan dynamics model for generalization in model-based reinforcement learning. In International Conference on Machine Learning, 2020.
[66] Yoshua Bengio. Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In International Conference on Machine Learning, 2009.
[67] Tambet Matüssnn, Avital Oliver, Taco Cohen, and John Schulman. Teacher-student curriculum learning, 2017.
[68] Marein Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay, 2018.
[69] Rémy Portelas, Cédric Colas, Katja Hofmann, and Pierre- Yves Oudeyer. Teacher algorithms for curriculum learning of deep in continuously parameterized environments. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura, editors, Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pages 835-853. PMLR, 30 Oct-01 Nov 2020.
[70] Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay, 2016.
[71] Nikolay Savinov, Anton Raichuk, Raphaël Marinier, Damien Vincent, Marc Pollefeys, Timo- thy Lillicrap, and Sylvain Gelly. Episodic curiosity through reachability, 2019.
[72] Alex Graves, Mare G. Bellemare, Jacob Menick, Remi Munos, and Koray Kavukcuoglu. Au- tomated curriculum learning for neural networks, 2017.
[73] Jack Parker-Holder, Minqi Jiang, Michael Dennis, Mikayel Samvelyan, Jakob Foerster, Ed- ward Grefenstette, and Tim Rocktäschel. Evolving curricula with regret-based environment design, 2023.
[74] Hojun Chung, Junseo Lee, Minsoo Kim, Dohyeong Kim, and Songhwai Oh. Adversarial environment design via regret-guided diffusion models, 2024.
14

---

[75] Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning, 05 2019.
[76] Yuxin Wu and Kaiming He. Group normalization, 2018.
[77] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. function approximation in reinforcement learning, 2017.
linear units for neural network
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The contributions outlined in the introduction are fully demonstrated through our experimental results.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: Limitations of our work are discussed in Section 6.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: The paper contains theoretical results.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [es]]
Justification: Our introduced domain fully detailed Section 3, while all hyperparame- ters used are provided in the Supplemantary Material.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The pipeline uses already open-sourced codebases, and the remainder of the code will be released GitHub shortly after this submission.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: This information partly stated in the experiments section and full details included in the supplementary material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments? Answer: [Yes]
Justification: All results are run over multiple seeds with error bars reported.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
15

---

A Model Architectures
Diffusion World Model The diffusion model Da employs a standard 2D U-Net architecture [31] for next-frame prediction. The model conditions history window of 4 previous frames and their corresponding actions, along with the diffusion timestep т. We implement observation conditioning through frame stacking along the channel dimension, while action and diffusion time conditioning are incorporated via adaptive group normalization layers [32] within the U-Net's residual blocks.
Reward and Termination Predictor The reward/termination model Rg utilizes a shared back- bone architecture with separate prediction heads for reward and termination signals. Crucially, we employ an ensemble of 10 prediction heads to enable uncertainty quantification during imagina- tion rollouts. The model processes sequences of frames and actions through convolutional residual blocks [33] followed by an LSTM cell [36, 35]. During inference, the ensemble predictions are aggregated to compute both mean estimates and uncertainty measures, providing robustness against prediction errors long-horizon imagined trajectories. Prior to imagination, we perform burn-in [75] using the conditioning frames and actions to properly initialize the LSTM's hidden and cell states.
Actor-Criie Network The policy Ta and value network Vo share common feature extraction backbone with separate output layers. We refer to this combined architecture as the actor-critic network (x, V) though technically V represents state-value function rather than an action-value critic. The network processes individual frames through convolutional encoder consisting of four residual blocks with 2x2 max-pooling (stride 2) between blocks. Each residual block contains a main path with group normalization [76], SiLU activation [77]. and 3x3 convolution (stride 1, padding 1). The convolutional features then processed by LSTM cell to maintain temporal context. Similar to the reward model, we initialize the LSTM states through burn-in on conditioning frames before beginning the imagination procedure.
Please refer Table below hyperparameter values.
Table 3: Architecture Parameters for iMAC.
Hyperparameter
Value
Diffusion Model (D₂)
Observation conditioning mechanism
Frame stacking
Action conditioning mechanism
Adaptive Group Normalization Adaptive Group Normalization
Diffusion time conditioning mechanism
Residual blocks layers
[2, 2, 2, 2] [64, 64, 64, 64] 256
Residual blocks channels
Residual blocks conditioning dimension
deward/Termination Model (R&))
Action conditioning mechanism
Adaptive Group Normalization
Residual blocks layers
[2, 2, 2, 2]
Residual blocks channels
[32, 32, 32, 32]
Residual blocks conditioning dimension
128 512
LSTM dimension
Actor-Crite Model (no and Vo)
Residual blocks layers
[1, 1, 1] [32, 32, 64, 64]
Residual blocks channels
LSTM dimension
512
16

---

B
World Model Training Parameters
Table 4: Hyperparameters for IMAC.
Hyperparameter
Value
Training loop
Number of epochs
1000
Training steps per epoch
100
Batch size
32
Environment steps per epoch
100
Epsilon (greedy) for collection
0.01
RL hyperparameters
Imagination horizon (H)
15, Random:[[52)
Discount factor (r)
0.985
Entropy weight (7)
0,001
X-returns coefficient (A)
0.95
Sequence construction during training For Dg. number of conditioning observations and actions (L)
For Ry, burn-in length (Bx), set to practice
For Ryg, training sequence length (BR + H)
Fixed:19, Random[9-26)
For T and Vgs, burn-in length (Bx,v), set to L in practice
Optimization
Optimizer
AdamW W
Learning rate
4e-5
Epsilon
1e-8
Weight decay (Do)
1e-2
Weight decay (Rp)
le-2
Weight decay (To and V&)
5e-5
Diffusion Sampling
Method
Fuler
Number of steps
Environment
Image observation dimensions
64x64x3
Action space
Discrete (up to 18 actions)
Frameskip
Max noop
30
Termination on life loss
True
Reward clipping
{0. 1}
PLR
Stalaness Factor
0.1 2500
Buffer Size
17

---

C Baseline World Model and Actor-Critic Training Algorithm
Algorithm 1: iMac (Sequential Offline Training)
Procedure training_lowe loops_ooppuuaa D) I| D is a fixed offline dataset;
Phase Train Diffusion Model: for epochs diffusion do do for steps_ _isfsssfiiionn model do do
end
end
// Freeze Diffusion Model Dæ; Phase Train Reward/Termination Model; for epochs anochs_reward do do for steps steps_reward reward end model do do update reward end model(D); end end
// Freeze Reward/Termination Model Rg: Phase Train Actor-Critic with Frozen Models; for epochs actor _orrticcccccc do do for steps_ actorr_ critic do do update actor coitticittt) // Uses frozen Da and Rs: end end
Procedure update _iiitttdddfiiinn____ diffusiio__moell()) model(
Sample sequence ( 1-1,,,. 2e: Or, 22+1) D;
Sample log(a) ~ ~N((Feean F2ra) // log-normal sigma distribution from EDM;
Define // default identity schedule from EDM;
Sample 2/41 N(xx 02I // Add independent Gaussian noise;
Compute =0 (2( T, 2/ L+1 -+
1......... 20 az);
Compute reconstruction loss L(0) || # 12
Update Ds:
Procedure update mondate_rewatte reward_end (nd_model())
Sample indices := [t....,t+ L H 1} burn-in imagination horizon; Sample sequence (1%) 0i, 0g, rs, de)ier ~ D;
Initialize h = c /// LSTM hidden and cell states;
for i T do
Compute i;, d, h = (x ais h c);
end
Compute L(v) Eiez CE(r,, sign(r:)) + CE(d,, d.) CE: cross-entropy loss;
Update Re:
Procedure update_actor latatattrrrrticcc D))
Sample initial buffer (x) 1+1- at 14--+1...... 2) D; Burn-in buffer with Rg. TT. and Vo to initialize LSTM states;
for == to
do
Sample a; (a(a)l=)) Sample reward r; and termination d, with frozen Rep
Sample next observation 2 by simulating reverse diffusion process with frozen Dg:
end
Compute (z( for i tt.... + H: Compute RI. losses Lv (Ø) and L. Update To and V4:
18

---

Algorithm 2: iMac: Actor-Critic Training with PLR (Prioritized Level Replay)
Procedure train train_actor aacttrrctttt wwwttc__i.iiie B)
|/ D is the offiine dataset, B is the prioritized buffer; for epochs acttr_crri _rrtic do do
I/ Update buffer with probability 0.5; random() < 0.5 then then update PLR buffer(D, B); end
for steps actor corticc do do
// Train actor-critic using only buffer data; update_acse accto_crrtticcec faciiiiffrmm_bae (oom_buffer(B)) end end
Procedure update_PPas PL. (T.R__fffe((1 D, B)
// Sample horizon length randomly:
Sample initial sequence (x (=t= 1+ .............................. , D;
Burn-in buffer with Rg: F. and Va initialize LSTM states; // Imagine trajectory using frozen models; Initialize trajectory -= ():
for i = to do do
Sample a; ~ Ta(a)l>?):
Sample reward and termination d, with frozen RØ:
Sample next observation #l41 by simulating reverse diffusion process with frozen Dg: Vo(x));
Compute TD-error 8; = r. (
T U {< 0e Tu d,, 2. 22+ 6;)]:
if d, = 1 then then
break;
end end
// Calculate PLR score as positive mean TD-error; Compute PLR score
mean (max(6;, 0) for 6 in T);
// Update buffer if score is higher than any existing entry;
if 8, min(s: b E B) then then Replace lowest-scoring entry in with (7,87); end
Procedure update pondate_actor_____ aataacctrrree actor_critic_ffiieee__ _wwi_ccrtffoooo from buffer( B)
// Sample trajectory from buffer based on priorities;
Sample trajectory 7 ~~ B according to PLR scores; // Extract states, actions, rewards, etc.;
Extract (x.,, a,, rs, d,, x2 from т: I/ Compute values for all states in trajectory; Compute (z for all states in r; |/ Compute advantage estimates and returns;
Compute advantages A, and returns G2 for all transitions: // Update actor and critic:
Compute RL losses Ly (Ø) (Vo(2)) -;))::
Compute (o) = - Σ, log (a(a))2)) Ai
Update a and Vj:
19

---

D Dataset Collection Details
We constructed our offline dataset by collecting trajectories from 200 procedurally generated levels for each of the five Procgen environments (CoinRun, Ninja, Jumper, Maze, and Cavellyerr, totaling million environment steps per game. We explicitly constrained data collection to levels 0-199 for each environment, ensuring complete isolation from the test levels (200+) used for evaluation, thereby preventing any potential data leakage and maintaining a strict train-test split that properly truly unseen procedurally generated
assesses generalization the benchmark protocol from the latest work [21], we created a mixed-quality dataset with equal proportions (1/3 each) of expert, medium, and random trajectories. Expert trajectories were col- lected from fully-trained PPO agents [36], medium trajectories from agents achieving 50% of expert performance, and random trajectories from uniformly random action selection. Expert trajectories were collected from PPO agents trained until convergence (when learning plateaued) rather than for a fixed number of steps, ensuring optimal performance for each game's 's unique characteristics. This mixed composition introduces additional challenges compared to standard offline RL benchmark [21] by reducing the proportion of successful episodes. The dataset is stored standard (s. a, r, s²) format where observations are 64x64x3 RGB images, actions a vary by dnvironment-specitn action spaces, and rewards are caanss-appparingg only at episode termination in all five selected en- vironments. This extreme sparsity in reward signals presents a particularly challenging scenario for world model learning, as the model must learn long-horizon dynamics without intermediate reward guidance. The hyperparameters of behavior policy are provided at Appendix F.
and extending
E
Analysis
E.1
Emergent Curriculum Effect
Our results demonstrate a clear emergent curriculum effect through the Prioritized Level Replay mechanism. As shown in Figure 2, the PLR approach initially selects similar episode lengths to random sampling while achieving higher returns. However, after approximately 4M training steps, significant shift occurs where PLR begins prioritizing substantially longer episodes (up 2.4x longer in some environments) while temporarily accepting lower comparative returns. This behavior indicates that the autocurriculum naturally discovers and focuses on more challenging, temporally- extended scenarios as the agent's capabilities improve, without requiring any manual curriculum design or environment-specific knowledge.
F,2
Component Contributions
The ablation study in Table 2 reveals the critical importance of combining both PLR and random horizon sampling. While fxed-horizon world models show reasonable performance (e.g., 3.71 on Jumper), incorporating random horizons improves results by 11% (4., and adding PLRR --ased prioritization yields an additional 41% improvement (5.78). This suggests that the success of AA stems from the synergistic effect of (1) diverse training experiences through random horizons pre- venting overfitting to fixed-length dynamics, and (2) intelligent prioritization that focuses computa- tional resources on the most informative imagined trajectories.
E.3
Generalization Mechanisms
The superior generalization of IMAC can be attributed its ability generate and prioritize diverse imagined experiences beyond the offline dataset distribution. By training exclusively on imagined rollouts from a world model learned on 200 levels, our agents achieve strong performance on unseen test levels (201+), with improvements ranging from 17% to 48% over model-free baselines. This suggests that the combination of world models and autocurricula enables agents to explore a richer space of potential scenarios than what exists in the original offline data, effectively augmenting the training distribution in a principled manner.
20

---

Baseline Algorithm Hyperparameters
We conducted comprehensive hyperparameter optimization for all baseline algorithms using grid search following the previous benchmark work [21]. All model-free offline RL methods (BCQ, CQL, IQL) and behavior cloning (BC) shared identical encoder architectures, utilizing ResNet for visual feature extraction from the 64x64x3 Procgen observations. Experiments
executed
NVIDIA RTX 4090 GPUs with 24GB memory. For behavior cloning (BC), we explored batch sizes {64, 128, 256, 512] and learning rates (5x10 10 x 10 6 10 $3. The Q-learning based methods (BCQ. CQL. IQL) employed dual-network architectures with target networks, where investigated both Polyak averaging and direct weight copying strategies. For direct copying, we tested update frequencies {1. 100, 1000), while for Polyak averaging, we examined (0.005, 0.5, 0.99). Algorithm-specine hyperparameters included: BCQ's s action selection threshold (0.3, 0.5, 0.7); CQL's regularization coefficient a E {0.5, 1.0, 4.0, 8.0); and IQL's temperature E {3.0, 7.0, 10.0} with expectile weights E (0.7, 0.8, 0.9). For sequence modeling approaches (DT, BCT), we additionally optimized context lengths E (5, 10, 30, 50). Decision Transformer utilized return-to-go multipliers 11,5), where the maximum return was determined from training data and scaled accordingly. Both DT and BCT maintained a dropout rate of 0.1 following [44]. Table presents the final selected hyperparameters achieving optimal performance on our IM-step mixed dataset.
Table 5: Final hyperparameter configurations for baseline algorithms on IM mixed dataset
Algorithm
Hyperparameter
Value
Learning Rate
0.0005
BC
Batch Size
256
Learning Rate
0.0005
Batch Size
512
BCT
Context Length
30
Eval Return Multiplier
0
Learning Rate
0.0005
Batch Size
512
DT
Context Length
10
Eval Return Multiplier
5
Learning Rate
0.0005
Batch Size Target Model Weight Update
256
BCQ
Direct copy
Target Update Frequency Threshold
1000
0.5
Learning Rate
0.0005
Batch Size Target Model Weight Update
256
CQL
Direct copy
Target Update Frequency Alpha
1000
4.0
Learning Rate
0.0005
Target Model Weight Update Batch Size
Direct copy
512
IQL
Target Update Frequency
100
Temperature
3.0
Expectile
0.8
21

---

G
Training Time Profile
Table 6 presents a comprehensive breakdown of computational time requirements for training IMAC game per seed.
Table 6: Detailed computational time analysis for IMAC training pipeline. Measurements were obtained using NVIDIA RTX 4090 GPU with hyperparameters detailed in previous sections. These timing results serve as representative benchmarks, with actual runtimes varying based on hardware configuration, environment setup, and specific training conditions.
Component
Time (hours)
World Model Training (Offfine)
Diffusion model
6.8
Reward/Terminttio model
5.4
Agent Training
Total (1000 epochs)
75.0
Per epoch (avg)
0.075
Total Training Time
87.2
22


Reviews:
[{"review_id": "RjFJJgD8fK", "reviewer": "Reviewer_jStg", "summary": "The paper introduces IMAC (Imagined Autocurricula), a novel approach combining world models with automatic curriculum learning to train rl policies. The key innovation is using Prioritized Level Replay (PLR) to create an emergent curriculum over \"imagined\" environments generated by a world model trained solely on offline data. The paper trains a diffusion world model from raw offline data, then uses it to generate diverse rl training scenarios. PLR automatically selects which imagined rollouts to prioritize based on learning potential, creating a natural curriculum that adapts to the agent's improving capabilities. This results in SOTA results vs. comparable models.", "questions": "- How well does this approach work under distribution shift between the input data used to learn the world model and the tasks that need to be completed? In a way the world model is now the weakest link. How well does it generalize?\n- Does the complexity of this approach scale with the complexity of the environment? It would be great to see some experiments where the environment is an ever more difficult function to learn, and seeing how the model becomes a bottleneck, and if this is feasible to scale to highly complex real world robotics environments.\n- What happens if we use the curriculum developed by one policy in one training run and use it to train another model from scratch? Is the online nature of the curriculum's development important or Is the curriculum good even when consumed offline.", "limitations": "None", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- The paper presents a successful combination of learned world models with unsupervised environment design (UED), opening a new research direction for training generally agents with online data, without simulators..\n- IMAC demonstrates substantial improvements (up to 48% on Jumper, 35% on Maze) over both model-free offline RL methods and fixed-horizon world model baselines, with consistent gains across all tested environments.\n- The approach automatically discovers effective curricula without manual design, naturally progressing from simple to complex scenarios.\n\nWeaknesses:\n- The approach only evaluates on 5 discrete Procgen environments with relatively simple visual observations; generalization to continuous control, high-dimensional state spaces, or real-world robotics remains unexplored.\n- Additional ablations on the bottlenecks posed by the world model, and scalability would be nice.", "was_revised": true, "final_justification": "I maintain my rating of accept.", "rebuttal": "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:\n\n**_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.\n\n**_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.\n\n**_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**\n\n**_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.\n\n[1] Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in 437 offline reinforcement learning. 2024.\n\n[2] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2024.\n\n[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\n\n[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"}, {"review_id": "dKab5PDNOg", "reviewer": "Reviewer_Dygh", "summary": "The paper proposes Imagined Autocurricula, a model-based offline reinforcement learning method that employs Prioritized Level Replay to exploit valuable states for agent learning. By focusing on initial states with the highest learning potential and utilizing a dynamic imagining horizon, this method demonstrates superior generalization to held-out levels in five Procgen environments.", "questions": "Q1) Any reason for choosing A2C as the agent learning algorithm?", "limitations": "Please see the Weaknesses part.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "# Strengths\n\nS1) The proposed method is well-motivated.\n\nS2) The writing is clear and easy to follow.\n\nS3) Compared to other baselines, the proposed method achieves a notable performance gain on five challenging Procgen environments.\n\n# Weaknesses\n\nW1) **Task coverage.** In Line 35-36, the authors claim that they attempt to bridge the gap between previous MBRL methods that focus on single environments and those utilizing Internet-pretrained world models for training generalist agents. However, the proposed method is trained and evaluated on only five tasks with similar goal-reaching objectives, which is not sufficient to convincingly demonstrate its generalization capabilities. Additionally, the proposed method is also only able to generalize to different levels within the same Procgen environment, a setting that has been explored by previous methods.\n\nW2) **Differences to previous works.** It appears that the proposed world model primarily follows the previous DIAMOND world model. It would be helpful if the differences could be elaborated.\n\nW3) **Lacking competitive baselines.** As a model-based method, the proposed method is mainly compared with some vanilla baselines in Table 1. While more advanced baselines (like [1-4]) and world models (like DreamerV3, TWM, STORM, IRIS, and DIAMOND mentioned by the authors in the related work).\n\nW4) **Missing visualization.** It would be great to show some predictions from the world model.\n\nW5) **Wrong reference.** In Line 312, the reference for STORM points to an unrelated paper. The correct one should be [5].\n\nW6) **A related reference.** [6] also proposes prioritized imagination for MBRL.\n\n---\n\n[1] Rethinking Decision Transformer via Hierarchical Reinforcement Learning. ICML 2024\n\n[2] Elastic Decision Transformer. NeurIPS 2023\n\n[3] IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies. arXiv 2023\n\n[4] AlignIQL: Policy Alignment in Implicit Q-Learning through Constrained Optimization. arXiv 2024\n\n[5] STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning. NeurIPS 2023\n\n[6] Pre-Trained Video Generative Models as World Simulators. arXiv 2025", "was_revised": true, "final_justification": "Some of my concerns have been addressed by the rebuttal. However, I still find that the evaluation setting and selected baselines are not sufficiently significant and convincing. Therefore, I modified my rating to borderline accept.", "rebuttal": "We thank the reviewer for their detailed feedback. We address each concern below and hope to clarify the contributions and scope of our work.\n\n**_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.\n\n**_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.\n\n**_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.**\n\n**_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.**\n\n**_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.\n\n**_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.\n\n[1] Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in offline reinforcement learning. 2024.\n\n[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."}, {"review_id": "UAPf5VPBN3", "reviewer": "Reviewer_Rrfo", "summary": "This paper presents an offline model-based RL method that integrates unsupervised environment design with a generative world model, to train agents capable of generalizing to unseen tasks.", "questions": "- In table 1, what is the maximum achievable return for each task? It is unclear whether the reported scores represent meaningful improvements, as the absolute values appear relatively low.\n- Can the proposed approach be extended to settings without explicit reward signals, which are common in real-world applications?\n- Is the method applicable to continuous action spaces, or is it limited to discrete environments?", "limitations": "yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n- The paper introduces a well-motivated and novel approach that leverages generative world models to train agents with strong generalization capabilities. Applying unsupervised environment design when generating imaginary rollouts seems like a powerful idea and has the potential to be used in large scale foundation world models.\n- The paper is clearly written. It is easy to understand its main idea and algorithm details.\n- The proposed method shows significant improvements in generalization performance.\n\n### Weaknesses\nWhile the proposed method combines existing components, its contribution would be clearer with a more in-depth justification of the core design choices. In particular, explaining why these components were selected over other viable alternatives would help clarify the novelty and applicability of the method.\n- Justification for the use of a 2D UNet architecture\n  - The authors adopt 2D UNet architecture following DIAMOND [1]. However, recent advances in video models often use 3D DiT architecture [2, 3]. Although an exhaustive architectural comparison may be infeasible, referencing and briefly discussing recent alternatives would help contextualize this design choice.\n- Justification for using prioritized level replay (PLR)\n  -  There exist recent UED algorithms that may be applicable in the proposed framework. For instance, ACCEL [4] could be implemented by varying the noise schedule, and ADD [5] could be integrated naturally by guiding the diffusion world model. \nAdditionally, although the authors employ PVL for prioritization score, alternative regret estimation strategies such as MaxMC [6] are available and sometimes yield stronger performance. A discussion of these alternatives would enhance the technical depth of the paper.\n- Limited baselines\n  - Among baselines, only transformer-based methods can be seen as model-based approaches. Given that the proposed method leverages a diffusion world model, incorporating baselines such as Diffuser [7] and its variants would allow for a more comprehensive evaluation.\n- Experiments\n  - While the authors conduct extensive experiments in Procgen benchmark, I'm concerned about whether the proposed method works well in other domains. Specifically, while the authors utilize random exploration data to ensure broad state coverage, there are many cases where random exploration data is unavailable (e.g. real world robotics tasks). \n\n### References\n\n[1] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, François Fleuret, \"Diffusion for World Modeling: Visual Details Matter in Atari\", NeurIPS, 2024.  \n[2] Zhuoyi Yang, Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng, Da Yin, Yuxuan Zhang, Weihan Wang, Yean Cheng, Bin Xu, Xiaotao Gu, Yuxiao Dong, Jie Tang, \"CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer\", ICLR, 2025.  \n[3] NVIDIA, \"Cosmos World Foundation Model Platform for Physical AI\", arXiv, 2025.  \n[4] Jack Parker-Holder, Minqi Jiang, Michael Dennis, Mikayel Samvelyan, Jakob Foerster, Edward Grefenstette, Tim Rocktäschel, \"Evolving Curricula with Regret-Based Environment Design\", ICML, 2022.  \n[5] Hojun Chung, Junseo Lee, Minsoo Kim, Dohyeong Kim, Songhwai Oh, \"Adversarial Environment Design via Regret-Guided Diffusion Models\", NeurIPS, 2024.  \n[6] Minqi Jiang, Michael Dennis, Jack Parker-Holder, Jakob Foerster, Edward Grefenstette, Tim Rocktäschel, \"Replay-Guided Adversarial Environment Design\", NeurIPS, 2021.  \n[7] Michael Janner, Yilun Du, Joshua B. Tenenbaum, Sergey Levine, \"Planning with Diffusion for Flexible Behavior Synthesis\", ICML, 2022.", "was_revised": true, "final_justification": "- Addressed issue: related work and general applicability\n\nThe initial submission lacked discussion of recent work on diffusion model architectures and UED algorithms. The authors have clarified that they will include these missing references and discussions in the revised related work section.\nFurthermore, the authors plan to address the potential applicability of the proposed method to more general settings, such as continuous action spaces and real-world robotic tasks without explicit rewards, in the conclusion section.\n\n- Unresolved issue: lack of model-based RL baselines\n\nThe concern regarding the lack of model-based RL baselines remains unresolved. However, the proposed method demonstrates strong empirical performance compared to model-free baselines, and the approach is clearly distinct from existing methods. I therefore find the contribution to be valid and meaningful.\n\nOverall, the integration of world models and UED is novel and provides a promising direction for addressing the limitations of both components. Based on these strengths and the issues that were satisfactorily addressed, I have increased my score.", "rebuttal": "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.\n\n**_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.\n\n**_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.** \n\n**_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.\n\n**_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.**\n\n**_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.\n\n**_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. \n\n**_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.\n\n[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.\n\n[2] Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in offline reinforcement learning. 2024."}, {"review_id": "9ma9jhT5Yc", "reviewer": "Reviewer_NSyY", "summary": "This paper introduces Imagined Autocurricula (iMAC), a novel method for training reinforcement learning agents. The method leverages offline data to generate diverse imagined environments using a trained diffusion model and employs Prioritized Level Replay to create an emergent curriculum that adapts to the agent’s capabilities.", "questions": "The performance of IMAC is fundamentally tied to the quality of the generative world model, which in turn depends on the offline dataset used for its training. The paper uses a \"mixed\" dataset of expert, medium, and random trajectories. How does IMAC's performance degrade if the world model is trained on a less diverse or lower-quality dataset?\n\nThe authors should discuss the OMNI-EPIC paper, as it appears to make similar claims. A discussion comparing the two approaches would help to clarify the novel contributions of this paper.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\nThe paper is clear and well-written. The introduction sets up the problem context and motivation effectively. The proposed method is explained logically, with a helpful architectural overview in Figure 1.\n\n## Weaknesses\n\nThe quality of the diffusion model depends on PPO being able to solve the tasks during data collection. It is unclear if the final agent is able to solve qualitatively different tasks after being trained with iMac.\n\nThe paper claims to demonstrate \"open-ended learning in learned world models\". This is a powerful claim, as open-ended learning often implies the generation of tasks with ever-increasing and potentially unbounded complexity. Given that the world model is fixed after training on a finite dataset, the space of generated environments is necessarily constrained by what the model has learned.\n\nA potential weakness is the complexity of the method with many moving pieces.\n\nAnother potential weakness is the very high computational cost—180 GPU days for the main experiments —which may be a barrier to reproducibility and adoption. However, the authors are transparent about this in the limitations section.", "was_revised": true, "final_justification": "The authors have adequately addressed my concerns.", "rebuttal": "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.\n\n**_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.**\n\n**_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.\n\n**_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.\n\n**_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.\n\n**_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.\n\n**_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.\n\n[1] Human-Timescale Adaptation in an Open-Ended Task Space, Bauer et al, ICML 2023"}]

Author rebuttal:

Meta-review: The paper introduces a model-based offline RL framework that trains a diffusion world model on mixed offline data, then generates imagined rollouts and applies Prioritized Level Replay to form an emergent curriculum that adapts as policy competence grows. Experiments on several Procgen environments indicate consistent gains over relevant baselines. The rebuttal further clarifies differences from prior world-model work, moderates claims around open-endedness, and commits to broader related-work coverage.

The paper provides a well-motivated integration of world models with UED and demonstrates strong empirical results on the Procgen domains used. The two outstanding weaknesses are the lack of additional model-based baselines and the limited evaluation setting. Balancing these, the contribution feels solid and the evidence within scope is persuasive.

The overall decision is to accept, with the authors encouraged to take into account reviewer feedback and update the paper as promised in order to produce the strongest version of the work for the camera-ready.
Decision: Accept (poster)